[CI] Fix broken CI (#6599)

Revert
4fb3d5e1b2
it breaks E2E Test

- vLLM version: v0.15.0
- vLLM main:
d7e17aaacd
This commit is contained in:
wangxiyuan
2026-02-06 17:23:58 +08:00
committed by GitHub
parent 19b5d44ea8
commit 06c0aed124
17 changed files with 1147 additions and 947 deletions

View File

@@ -1,14 +1,11 @@
from typing import Any
from typing import Any, Dict, Optional
from vllm.model_executor.models.utils import extract_layer_index
from vllm_ascend.distributed.parallel_state import get_shard_weight_group
from vllm_ascend.ops.layer_shard_linear import (
is_hidden_layer,
post_process_after_loading_for_shard_weight_series,
reach_layer_for_shard_weight_series,
register_layer_to_shard_weight_series,
)
is_hidden_layer, post_process_after_loading_for_shard_weight_series,
reach_layer_for_shard_weight_series, register_layer_to_shard_weight_series)
from vllm_ascend.utils import flashcomm2_enable, o_shard_enable
@@ -29,7 +26,7 @@ class Flashcomm2OShardManager:
"""
def __init__(self):
self._shard_layers: dict[int, Any] = {}
self._shard_layers: Dict[int, Any] = {}
def flashcomm2_oshard_enable(self):
return flashcomm2_enable() and o_shard_enable()
@@ -55,10 +52,12 @@ class Flashcomm2OShardManager:
self._shard_layers[layer_idx] = layer
register_layer_to_shard_weight_series(
series_name="o_proj", group=get_shard_weight_group(), layer=layer, prefetch_step=prefetch_step
)
series_name="o_proj",
group=get_shard_weight_group(),
layer=layer,
prefetch_step=prefetch_step)
def get_layer(self, layer_idx: int) -> Any | None:
def get_layer(self, layer_idx: int) -> Optional[Any]:
"""Safely retrieves a registered layer by its index.
Args: