diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index 32cd5f9b..87970090 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/afd_plugin/compat/npu/feature_validation.py @@ -4,6 +4,7 @@ from __future__ import annotations +import logging from typing import TYPE_CHECKING from afd_plugin.config import ( @@ -18,6 +19,8 @@ from afd_plugin.connectors.base import ConnectorExtraInfo +logger = logging.getLogger(__name__) + def fail_if_unsupported_npu_afd_features( vllm_config: VllmConfig, @@ -111,7 +114,13 @@ def _fail_if_unsupported_npu_async_moe_ubatching_features( num_ubatches: int, split: str, ) -> None: - from afd_plugin.connectors.npu.async_cam import ASYNC_MOE_REQUEST_SPLIT + from afd_plugin.connectors.npu.async_cam import ( + ASYNC_MOE_REQUEST_SPLIT, + ASYNC_MOE_TOKEN_SPLIT, + ) + from afd_plugin.v1.worker.npu.ubatch_utils import ( + enable_token_balanced_async_moe_split, + ) parallel_config = vllm_config.parallel_config if not afd_config.compute_gate_on_attention: @@ -123,15 +132,35 @@ def _fail_if_unsupported_npu_async_moe_ubatching_features( "async_moe_ubatching currently supports exactly two stages; " f"got async_moe_num_ubatches={num_ubatches}", ) - if split != ASYNC_MOE_REQUEST_SPLIT: + if split not in (ASYNC_MOE_REQUEST_SPLIT, ASYNC_MOE_TOKEN_SPLIT): raise RuntimeError( - "async_moe_ubatching currently supports only request-boundary split; " + "async_moe_split must be 'request' or 'token'; " f"got async_moe_split={split!r}", ) + # The Attention process owns stage planning and SP layout conversion. + # FFN processes consume the resulting stage payloads, so their independent + # TP/CP topology must not be used to validate the Attention split policy. + if afd_config.is_ffn_server: + return if int(parallel_config.decode_context_parallel_size) > 1: raise RuntimeError( "async_moe_ubatching does not support decode context parallel metadata yet", ) + token_split_capable = enable_token_balanced_async_moe_split(vllm_config) + if split == ASYNC_MOE_TOKEN_SPLIT: + if not token_split_capable: + raise RuntimeError( + "async_moe_split='token' requires a non-PCP Attention DP+TP/SP " + "topology (Attention tensor_parallel_size > 1, no " + "prefill/decode context parallel)", + ) + elif split == ASYNC_MOE_REQUEST_SPLIT and token_split_capable: + logger.warning( + "async_moe_ubatching runs on a non-PCP DP+TP/SP topology with " + "async_moe_split='request'; request lengths can be skewed, " + "consider async_moe_split='token' for token-balanced " + "microbatches", + ) __all__ = ["fail_if_unsupported_npu_afd_features"] diff --git a/afd_plugin/connectors/npu/async_cam.py b/afd_plugin/connectors/npu/async_cam.py index bc5195ab..bc334aff 100644 --- a/afd_plugin/connectors/npu/async_cam.py +++ b/afd_plugin/connectors/npu/async_cam.py @@ -18,9 +18,11 @@ The supported deployment requires ``async=true``, eager execution, Ascend CAM operator packages, and matching topology/configuration on every rank. vLLM native DBO, ACL graph execution, and decode are not supported. -Optional AFD-managed MoE ubatching is a separate two-stage request-boundary -pipeline. See ``docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md`` for configuration, -rank derivation, launch guidance, and the full limitations. +Optional AFD-managed MoE ubatching is a separate two-stage pipeline. It uses +request boundaries without sequence parallelism and TP-aligned token stages +when sequence parallelism is active. See +``docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md`` for configuration, rank +derivation, launch guidance, and the full limitations. """ from __future__ import annotations @@ -65,6 +67,7 @@ CAM_COMM_ID = 0 ATTN_RANKS_PER_DP_CONFIG_KEY = "attn_ranks_per_dp" ASYNC_MOE_REQUEST_SPLIT = "request" +ASYNC_MOE_TOKEN_SPLIT = "token" _AFD_ASYNC_EXTRA_CONFIG_FIELDS: Final[frozenset[str]] = frozenset( { @@ -85,10 +88,14 @@ class AFDAsyncExtraInfo(ConnectorExtraInfo): Attributes: dynamic_quant: Dynamic quantization mode accepted by CAM operators. - attn_ranks_per_dp: Number of Attention ranks in each data-parallel group. - async_moe_ubatching: Whether request-boundary async MoE ubatching is used. + attn_ranks_per_dp: Number of Attention ranks in each data-parallel + replica. This is the CAM Attention grouping width and is + independent of the FFN process's local TP size. + async_moe_ubatching: Whether two-stage async MoE ubatching is used. async_moe_num_ubatches: Number of stages used by async MoE ubatching. - async_moe_split: Boundary at which async MoE work is split. + async_moe_split: Boundary at which async MoE work is split: "request" + for request boundaries, "token" for TP-aligned token-balanced + split points (non-PCP DP+TP/SP topologies only). """ dynamic_quant: int = 0 @@ -233,8 +240,9 @@ def __init__( Communication resources are created collectively by ``init_afd_connector``. ``afd_role_rank`` must already contain the - process's role-local DP/PCP-derived offset; ``attn_ranks_per_dp`` is - used as the CAM Attention TP width. + process's role-local distributed-rank-derived offset; + ``attn_ranks_per_dp`` is used as the CAM Attention grouping width and + does not describe the FFN process's local TP topology. """ super().__init__(rank, local_rank, vllm_config, afd_config) self._initialized = False diff --git a/afd_plugin/model_executor/models/forward_context.py b/afd_plugin/model_executor/models/forward_context.py index b8f33c78..8e2718fd 100644 --- a/afd_plugin/model_executor/models/forward_context.py +++ b/afd_plugin/model_executor/models/forward_context.py @@ -18,7 +18,44 @@ ASYNC_MOE_UBATCH_METADATA_KEY: Final[str] = "afd_async_moe_ubatch_metadata" -class AsyncMoeUbatchMetadata(TypedDict): +class _AsyncMoeUbatchMetadataOptional(TypedDict, total=False): + """Optional SP-local fields for ``AsyncMoeUbatchMetadata``. + + A separate ``total=False`` base class is used instead of + ``typing.NotRequired`` because the plugin supports Python 3.10, where + ``NotRequired`` is only available via the external ``typing_extensions`` + package. Keeping the base class avoids adding a new runtime dependency. + """ + + # SP-local stage layout: when SP shards the full batch across TP ranks, + # transpose it into equal per-stage rank shards before attention and + # restore the original layout after the async MoE pipeline. + use_sp_stage_resharding: bool + sp_local_stage_slices: UBatchSlices + stage_actual_token_counts: list[int] + sp_local_stage_actual_token_counts: list[int] + + +class AsyncMoeUbatchMetadata(_AsyncMoeUbatchMetadataOptional): + """Async MoE ubatch sidecar metadata carried by the forward context. + + The required fields (``attn_metadata``, ``ubatch_slices``) are always + populated by the attention model runner before the model forward starts. + The optional layout fields describe padded stage inputs and their real + token coverage: + + - ``use_sp_stage_resharding``: set to ``True`` by the attention model + runner when Ascend sequence parallelism is active and the applied split + is TP-aligned. + - ``sp_local_stage_slices``: set by + ``build_async_moe_stage_inputs`` during model forward; its lengths + describe the equal per-rank stage shards, not ranges into the original + full-batch local shard. + - ``stage_actual_token_counts`` and + ``sp_local_stage_actual_token_counts`` distinguish real token rows from + DP/SP padding globally and on the current TP rank. + """ + attn_metadata: object ubatch_slices: UBatchSlices diff --git a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py index f1f677c7..42a02b0f 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py @@ -19,6 +19,11 @@ AFDTransferMetadata, ) from afd_plugin.model_executor.models import AsyncMoeUbatchMetadata +from afd_plugin.model_executor.models.npu.ubatch_sp import ( + build_async_moe_stage_inputs, + restore_async_moe_stage_outputs, + sp_local_actual_token_count, +) from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield if TYPE_CHECKING: @@ -124,35 +129,87 @@ def run_async_moe_ubatch_afd_forward( afd_connector = afd_metadata.connector first_moe_layer = int(model.config.first_k_dense_replace) dense_end_layer = min(model.end_layer, first_moe_layer) - for layer in islice(model.layers, model.start_layer, dense_end_layer): - hidden_states, residual = layer( - positions, - hidden_states, - residual, - llama_4_scaling, + + ( + stage_hidden_states, + stage_residual, + stage_positions, + stage_llama_4_scaling, + sp_local_stage_slices, + ) = build_async_moe_stage_inputs( + hidden_states=hidden_states, + residual=residual, + positions=positions, + llama_4_scaling=llama_4_scaling, + ubatch_slices=ubatch_slices, + use_sp_stage_resharding=bool( + async_moe_ubatch_metadata.get("use_sp_stage_resharding", False), + ), + ) + async_moe_ubatch_metadata["sp_local_stage_slices"] = sp_local_stage_slices + stage_actual_token_counts = async_moe_ubatch_metadata.get( + "stage_actual_token_counts", + [int(ubatch_slice.num_tokens) for ubatch_slice in ubatch_slices], + ) + use_sp_stage_resharding = ( + bool( + async_moe_ubatch_metadata.get("use_sp_stage_resharding", False), ) - if dense_end_layer == model.end_layer: - return hidden_states, residual + and sp_local_stage_slices is not ubatch_slices + ) + if use_sp_stage_resharding: + sp_local_stage_actual_token_counts = [ + sp_local_actual_token_count( + stage_actual_tokens=int(stage_actual_tokens), + stage_input_tokens=int(ubatch_slice.num_tokens), + ) + for stage_actual_tokens, ubatch_slice in zip( + stage_actual_token_counts, + ubatch_slices, + strict=True, + ) + ] + else: + sp_local_stage_actual_token_counts = [ + int(stage_actual_tokens) + for stage_actual_tokens in stage_actual_token_counts + ] + async_moe_ubatch_metadata["sp_local_stage_actual_token_counts"] = ( + sp_local_stage_actual_token_counts + ) - stage_hidden_states = [ - hidden_states[ubatch_slice.token_slice] for ubatch_slice in ubatch_slices - ] - stage_residual = [ - _slice_optional_first_dim(residual, ubatch_slice.token_slice) - for ubatch_slice in ubatch_slices - ] - stage_positions = [ - _slice_positions(positions, ubatch_slice.token_slice) - for ubatch_slice in ubatch_slices - ] - stage_llama_4_scaling = [ - _slice_llama_4_scaling( - llama_4_scaling, - ubatch_slice.token_slice, - num_tokens=int(hidden_states.shape[0]), + # Per-stage Ascend metadata is built before model execution. Run the dense + # prefix under the matching stage context too, so its RoPE/KV inputs never + # mix the full-batch context with buffers prepared by a stage builder. + for stage_idx in range(len(ubatch_slices)): + with _use_async_moe_ubatch_forward_context( + forward_context=forward_context, + parent_afd_metadata=afd_metadata, + async_moe_ubatch_metadata=async_moe_ubatch_metadata, + stage_idx=stage_idx, + ): + for layer in islice( + model.layers, + model.start_layer, + dense_end_layer, + ): + ( + stage_hidden_states[stage_idx], + stage_residual[stage_idx], + ) = layer( + stage_positions[stage_idx], + stage_hidden_states[stage_idx], + stage_residual[stage_idx], + stage_llama_4_scaling[stage_idx], + ) + + if dense_end_layer == model.end_layer: + return _restore_async_moe_stage_state( + stage_hidden_states, + stage_residual, + ubatch_slices, + use_sp_stage_resharding=use_sp_stage_resharding, ) - for ubatch_slice in ubatch_slices - ] moe_start_layer = max(model.start_layer, first_moe_layer) moe_layers = list(islice(model.layers, moe_start_layer, model.end_layer)) @@ -161,7 +218,7 @@ def compute_stage_attention( layer: AFDDeepseekV2DecoderLayer, stage_idx: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: - ubatch_slice = ubatch_slices[stage_idx] + sp_local_stage_slice = sp_local_stage_slices[stage_idx] with _use_async_moe_ubatch_forward_context( forward_context=forward_context, parent_afd_metadata=afd_metadata, @@ -184,7 +241,7 @@ def compute_stage_attention( raise RuntimeError( "async_moe_ubatching requires Attention-side topk payloads", ) - expected_tokens = int(ubatch_slice.num_tokens) + expected_tokens = int(sp_local_stage_slice.num_tokens) if int(stage_hidden_states[stage_idx].shape[0]) != expected_tokens: raise RuntimeError( "async_moe_ubatching stage output token count mismatch: " @@ -200,7 +257,7 @@ def send_stage_attention( topk_ids: torch.Tensor, router_logits: torch.Tensor | None, ) -> None: - expected_tokens = int(ubatch_slices[stage_idx].num_tokens) + expected_tokens = int(sp_local_stage_slices[stage_idx].num_tokens) stage_metadata = AFDTransferMetadata.create_attention_metadata( layer_idx=layer.layer_idx, stage_idx=stage_idx, @@ -279,12 +336,42 @@ def recv_stage_ffn(stage_idx: int) -> None: router_logits, ) recv_stage_ffn(1) - output_hidden_states = torch.cat(stage_hidden_states, dim=0) - return ( - output_hidden_states, - _cat_optional_async_moe_stage_outputs( + return _restore_async_moe_stage_state( + stage_hidden_states, + stage_residual, + ubatch_slices, + use_sp_stage_resharding=use_sp_stage_resharding, + ) + + +def _restore_async_moe_stage_state( + stage_hidden_states: list[torch.Tensor], + stage_residual: list[torch.Tensor | None], + ubatch_slices: UBatchSlices, + *, + use_sp_stage_resharding: bool, +) -> tuple[torch.Tensor, torch.Tensor | None]: + output_hidden_states = restore_async_moe_stage_outputs( + stage_hidden_states, + ubatch_slices, + use_sp_stage_resharding=use_sp_stage_resharding, + ) + resolved_stage_residual = [ + stage_output if stage_output is not None else fallback_output + for stage_output, fallback_output in zip( stage_residual, stage_hidden_states, + strict=True, + ) + ] + return ( + output_hidden_states, + None + if all(stage_output is None for stage_output in stage_residual) + else restore_async_moe_stage_outputs( + resolved_stage_residual, + ubatch_slices, + use_sp_stage_resharding=use_sp_stage_resharding, ), ) @@ -301,28 +388,37 @@ def _use_async_moe_ubatch_forward_context( stage_idx: int, ) -> Iterator[None]: ubatch_slices = async_moe_ubatch_metadata["ubatch_slices"] + sp_local_stage_slices = async_moe_ubatch_metadata.get( + "sp_local_stage_slices", + ubatch_slices, + ) attn_metadata = async_moe_ubatch_metadata["attn_metadata"] stage_afd_metadata = _build_async_moe_stage_afd_metadata( parent_afd_metadata, - ubatch_slices, + sp_local_stage_slices, + async_moe_ubatch_metadata.get( + "sp_local_stage_actual_token_counts", + [int(stage_slice.num_tokens) for stage_slice in sp_local_stage_slices], + ), stage_idx, ) + stage_context_attr_names = ( + "attn_metadata", + "additional_kwargs", + "ubatch_idx", + "num_ubatches", + "num_tokens", + "pad_size", + "padded_length", + "max_tokens_across_dp", + "padded_num_tokens", + "mc2_mask", + "dbo_enabled", + ) saved_attrs = { - "attn_metadata": _read_forward_context_attr( - forward_context, - "attn_metadata", - ), - "additional_kwargs": _read_forward_context_attr( - forward_context, - "additional_kwargs", - ), - "ubatch_idx": _read_forward_context_attr(forward_context, "ubatch_idx"), - "num_ubatches": _read_forward_context_attr( - forward_context, - "num_ubatches", - ), - "num_tokens": _read_forward_context_attr(forward_context, "num_tokens"), + name: _read_forward_context_attr(forward_context, name) + for name in stage_context_attr_names } original_kwargs = ( @@ -332,13 +428,36 @@ def _use_async_moe_ubatch_forward_context( ) stage_kwargs = dict(original_kwargs or {}) stage_kwargs["afd_metadata"] = stage_afd_metadata + stage_input_tokens = int(ubatch_slices[stage_idx].num_tokens) + stage_actual_tokens = int( + async_moe_ubatch_metadata.get( + "stage_actual_token_counts", + [int(ubatch_slice.num_tokens) for ubatch_slice in ubatch_slices], + )[stage_idx], + ) try: forward_context.attn_metadata = attn_metadata[stage_idx] forward_context.additional_kwargs = stage_kwargs forward_context.ubatch_idx = stage_idx forward_context.num_ubatches = len(ubatch_slices) - forward_context.num_tokens = int(ubatch_slices[stage_idx].num_tokens) + # Ascend MLA all-gathers the local stage shards before RoPE and uses + # this value to allocate its global o-projection input. + forward_context.num_tokens = stage_input_tokens + forward_context.pad_size = 0 + forward_context.padded_length = stage_input_tokens + forward_context.max_tokens_across_dp = stage_input_tokens + forward_context.padded_num_tokens = stage_input_tokens + forward_context.dbo_enabled = True + saved_mc2_mask = saved_attrs["mc2_mask"] + if isinstance(saved_mc2_mask, torch.Tensor): + stage_mc2_mask = torch.zeros( + (stage_input_tokens,), + dtype=saved_mc2_mask.dtype, + device=saved_mc2_mask.device, + ) + stage_mc2_mask[:stage_actual_tokens] = True + forward_context.mc2_mask = stage_mc2_mask yield finally: for name, value in saved_attrs.items(): @@ -348,6 +467,7 @@ def _use_async_moe_ubatch_forward_context( def _build_async_moe_stage_afd_metadata( parent_afd_metadata: AFDForwardContextMetadata, ubatch_slices: UBatchSlices, + stage_actual_token_counts: list[int], stage_idx: int, ) -> AFDForwardContextMetadata: ubatch_slice = ubatch_slices[stage_idx] @@ -357,63 +477,12 @@ def _build_async_moe_stage_afd_metadata( stage_metadata.tokens_start_loc = [ubatch_slice.token_slice.start] stage_metadata.requests_start_loc = [ubatch_slice.request_slice.start] stage_metadata.tokens_lens = [ubatch_slice.num_tokens] - if len(parent_afd_metadata.tokens_unpadded_lens) > stage_idx: - unpadded_len = parent_afd_metadata.tokens_unpadded_lens[stage_idx] - else: - unpadded_len = ubatch_slice.num_tokens - stage_metadata.tokens_unpadded_lens = [int(unpadded_len)] + stage_metadata.tokens_unpadded_lens = [ + int(stage_actual_token_counts[stage_idx]), + ] return stage_metadata -def _cat_optional_async_moe_stage_outputs( - stage_outputs: list[torch.Tensor | None], - fallback_outputs: list[torch.Tensor], -) -> torch.Tensor | None: - if all(stage_output is None for stage_output in stage_outputs): - return None - return torch.cat( - [ - stage_output if stage_output is not None else fallback_output - for stage_output, fallback_output in zip( - stage_outputs, - fallback_outputs, - strict=True, - ) - ], - dim=0, - ) - - -def _slice_optional_first_dim( - tensor: torch.Tensor | None, - token_slice: slice, -) -> torch.Tensor | None: - if tensor is None: - return None - return tensor[token_slice] - - -def _slice_positions(positions: torch.Tensor, token_slice: slice) -> torch.Tensor: - if positions.dim() <= 1: - return positions[token_slice] - return positions[..., token_slice] - - -def _slice_llama_4_scaling( - llama_4_scaling: torch.Tensor | None, - token_slice: slice, - *, - num_tokens: int, -) -> torch.Tensor | None: - if llama_4_scaling is None: - return None - if llama_4_scaling.shape[0] == num_tokens: - return llama_4_scaling[token_slice] - if llama_4_scaling.dim() > 1 and llama_4_scaling.shape[1] == num_tokens: - return llama_4_scaling[:, token_slice] - return llama_4_scaling - - def _read_forward_context_attr(forward_context: object, name: str) -> object: try: return getattr(forward_context, name) diff --git a/afd_plugin/model_executor/models/npu/ubatch_sp.py b/afd_plugin/model_executor/models/npu/ubatch_sp.py new file mode 100644 index 00000000..8c0ff7c9 --- /dev/null +++ b/afd_plugin/model_executor/models/npu/ubatch_sp.py @@ -0,0 +1,406 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Sequence-parallel stage layout helpers for async MoE ubatching. + +Ascend sequence parallelism stores the full-batch tensor as one contiguous +rank-major shard per TP rank. Async MoE stages instead need a contiguous +global stage reconstructed by the attention operator's TP all-gather. The +helpers below explicitly transpose between those two layouts: + +* full-batch rank shards -> global full batch -> per-stage rank shards; +* per-stage rank shards -> global stages -> original full-batch rank shards. + +Slicing each full-batch rank shard independently cannot implement this +transpose and silently pairs hidden states with another stage's positions. +""" + +from __future__ import annotations + +from typing import overload + +import torch +from vllm.distributed import tensor_model_parallel_all_gather +from vllm.distributed.parallel_state import get_tp_group +from vllm.v1.worker.ubatch_utils import UBatchSlice, UBatchSlices + + +def build_async_moe_stage_inputs( + *, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + positions: torch.Tensor, + llama_4_scaling: torch.Tensor | None, + ubatch_slices: UBatchSlices, + use_sp_stage_resharding: bool, +) -> tuple[ + list[torch.Tensor], + list[torch.Tensor | None], + list[torch.Tensor], + list[torch.Tensor | None], + UBatchSlices, +]: + """Build stage inputs in the layout expected by Ascend attention.""" + if not use_sp_stage_resharding: + return _build_non_sp_stage_inputs( + hidden_states=hidden_states, + residual=residual, + positions=positions, + llama_4_scaling=llama_4_scaling, + ubatch_slices=ubatch_slices, + ) + + tp_group = get_tp_group() + tp_rank = int(tp_group.rank_in_group) + tp_size = int(tp_group.world_size) + global_num_tokens = sum( + int(ubatch_slice.num_tokens) for ubatch_slice in ubatch_slices + ) + if tp_size <= 1 or int(hidden_states.shape[0]) == global_num_tokens: + return _build_non_sp_stage_inputs( + hidden_states=hidden_states, + residual=residual, + positions=positions, + llama_4_scaling=llama_4_scaling, + ubatch_slices=ubatch_slices, + ) + if global_num_tokens % tp_size != 0: + raise ValueError( + "Async MoE SP full-batch token count must be divisible by TP size; " + f"got num_tokens={global_num_tokens}, tp_size={tp_size}", + ) + expected_local_tokens = global_num_tokens // tp_size + if int(hidden_states.shape[0]) != expected_local_tokens: + raise ValueError( + "Async MoE SP hidden-state layout mismatch: expected " + f"{expected_local_tokens} local rows for {global_num_tokens} " + f"global tokens at TP={tp_size}, got {int(hidden_states.shape[0])}", + ) + + # This transpose is paid once before and once after the complete MoE + # pipeline, not once per layer. A future optimization can replace the + # temporary global tensors with variable-split TP all-to-all after vLLM + # exposes a stable backend-neutral API for it. + global_hidden_states = _gather_sp_sequence_tensor( + hidden_states, + token_dim=0, + expected_global_tokens=global_num_tokens, + ) + global_residual = ( + None + if residual is None + else _gather_sp_sequence_tensor( + residual, + token_dim=0, + expected_global_tokens=global_num_tokens, + ) + ) + global_positions, positions_token_dim = _to_global_sequence_tensor( + positions, + local_num_tokens=expected_local_tokens, + global_num_tokens=global_num_tokens, + ) + if llama_4_scaling is None: + global_llama_4_scaling = None + scaling_token_dim = None + else: + global_llama_4_scaling, scaling_token_dim = _to_global_sequence_tensor( + llama_4_scaling, + local_num_tokens=expected_local_tokens, + global_num_tokens=global_num_tokens, + ) + + stage_hidden_states: list[torch.Tensor] = [] + stage_residual: list[torch.Tensor | None] = [] + stage_positions: list[torch.Tensor] = [] + stage_llama_4_scaling: list[torch.Tensor | None] = [] + sp_local_stage_slices: UBatchSlices = [] + local_stage_start = 0 + + for ubatch_slice in ubatch_slices: + stage_tokens = int(ubatch_slice.num_tokens) + if stage_tokens % tp_size != 0: + raise ValueError( + "Async MoE SP stage token count must be divisible by TP size; " + f"got token_slice={ubatch_slice.token_slice}, tp_size={tp_size}", + ) + local_stage_tokens = stage_tokens // tp_size + global_stage_start = int(ubatch_slice.token_slice.start) + local_global_start = global_stage_start + tp_rank * local_stage_tokens + local_global_stop = local_global_start + local_stage_tokens + local_global_slice = slice(local_global_start, local_global_stop) + + stage_hidden_states.append(global_hidden_states[local_global_slice]) + stage_residual.append( + _slice_optional_first_dim(global_residual, local_global_slice), + ) + stage_positions.append( + _slice_sequence_dim( + global_positions, + positions_token_dim, + local_global_slice, + ), + ) + stage_llama_4_scaling.append( + None + if global_llama_4_scaling is None or scaling_token_dim is None + else _slice_sequence_dim( + global_llama_4_scaling, + scaling_token_dim, + local_global_slice, + ) + ) + + local_stage_stop = local_stage_start + local_stage_tokens + sp_local_stage_slices.append( + UBatchSlice( + ubatch_slice.request_slice, + slice(local_stage_start, local_stage_stop), + ), + ) + local_stage_start = local_stage_stop + + return ( + stage_hidden_states, + stage_residual, + stage_positions, + stage_llama_4_scaling, + sp_local_stage_slices, + ) + + +def restore_async_moe_stage_outputs( + stage_outputs: list[torch.Tensor], + ubatch_slices: UBatchSlices, + *, + use_sp_stage_resharding: bool, +) -> torch.Tensor: + """Restore stage outputs to the model's original full-batch SP layout.""" + if not use_sp_stage_resharding: + return torch.cat(stage_outputs, dim=0) + + tp_group = get_tp_group() + tp_rank = int(tp_group.rank_in_group) + tp_size = int(tp_group.world_size) + global_num_tokens = sum( + int(ubatch_slice.num_tokens) for ubatch_slice in ubatch_slices + ) + if len(stage_outputs) != len(ubatch_slices): + raise ValueError( + "Async MoE stage output count does not match ubatch metadata; " + f"got {len(stage_outputs)} outputs and {len(ubatch_slices)} slices", + ) + + global_stage_outputs: list[torch.Tensor] = [] + for stage_output, ubatch_slice in zip( + stage_outputs, + ubatch_slices, + strict=True, + ): + global_stage_outputs.append( + _gather_sp_sequence_tensor( + stage_output, + token_dim=0, + expected_global_tokens=int(ubatch_slice.num_tokens), + ), + ) + global_output = torch.cat(global_stage_outputs, dim=0) + local_tokens = global_num_tokens // tp_size + local_start = tp_rank * local_tokens + return global_output[local_start : local_start + local_tokens] + + +def sp_local_actual_token_count( + *, + stage_actual_tokens: int, + stage_input_tokens: int, +) -> int: + """Return real token rows in the current rank's stage-local SP shard.""" + tp_group = get_tp_group() + tp_rank = int(tp_group.rank_in_group) + tp_size = int(tp_group.world_size) + local_stage_tokens = int(stage_input_tokens) // tp_size + local_start = tp_rank * local_stage_tokens + local_stop = local_start + local_stage_tokens + return max(0, min(int(stage_actual_tokens), local_stop) - local_start) + + +def _build_non_sp_stage_inputs( + *, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + positions: torch.Tensor, + llama_4_scaling: torch.Tensor | None, + ubatch_slices: UBatchSlices, +) -> tuple[ + list[torch.Tensor], + list[torch.Tensor | None], + list[torch.Tensor], + list[torch.Tensor | None], + UBatchSlices, +]: + num_tokens = int(hidden_states.shape[0]) + return ( + [ + _slice_and_pad_first_dim(hidden_states, ubatch_slice.token_slice) + for ubatch_slice in ubatch_slices + ], + [ + _slice_and_pad_first_dim(residual, ubatch_slice.token_slice) + for ubatch_slice in ubatch_slices + ], + [ + _slice_positions(positions, ubatch_slice.token_slice) + for ubatch_slice in ubatch_slices + ], + [ + _slice_llama_4_scaling( + llama_4_scaling, + ubatch_slice.token_slice, + num_tokens=num_tokens, + ) + for ubatch_slice in ubatch_slices + ], + ubatch_slices, + ) + + +def _gather_sp_sequence_tensor( + tensor: torch.Tensor, + *, + token_dim: int, + expected_global_tokens: int, +) -> torch.Tensor: + gathered = tensor_model_parallel_all_gather(tensor.contiguous(), token_dim) + if int(gathered.shape[token_dim]) != int(expected_global_tokens): + raise RuntimeError( + "Async MoE SP all-gather returned an unexpected token count: " + f"expected {expected_global_tokens}, got " + f"{int(gathered.shape[token_dim])}", + ) + return gathered + + +def _to_global_sequence_tensor( + tensor: torch.Tensor, + *, + local_num_tokens: int, + global_num_tokens: int, +) -> tuple[torch.Tensor, int]: + global_token_dim = _sequence_tensor_token_dim(tensor, global_num_tokens) + if global_token_dim is not None: + return tensor, global_token_dim + + local_token_dim = _sequence_tensor_token_dim(tensor, local_num_tokens) + if local_token_dim is None: + raise ValueError( + "Sequence tensor token dimension must be on axis 0 or 1; " + f"got tensor shape {tuple(tensor.shape)} with " + f"local_num_tokens={local_num_tokens}, " + f"global_num_tokens={global_num_tokens}", + ) + return ( + _gather_sp_sequence_tensor( + tensor, + token_dim=local_token_dim, + expected_global_tokens=global_num_tokens, + ), + local_token_dim, + ) + + +def _sequence_tensor_token_dim(tensor: torch.Tensor, num_tokens: int) -> int | None: + if tensor.dim() > 0 and int(tensor.shape[0]) == int(num_tokens): + return 0 + if tensor.dim() > 1 and int(tensor.shape[1]) == int(num_tokens): + return 1 + return None + + +def _slice_sequence_dim( + tensor: torch.Tensor, + token_dim: int, + token_slice: slice, +) -> torch.Tensor: + if token_dim == 0: + return tensor[token_slice] + return tensor[:, token_slice] + + +@overload +def _slice_optional_first_dim( + tensor: torch.Tensor, + token_slice: slice, +) -> torch.Tensor: ... + + +@overload +def _slice_optional_first_dim( + tensor: None, + token_slice: slice, +) -> None: ... + + +def _slice_optional_first_dim( + tensor: torch.Tensor | None, + token_slice: slice, +) -> torch.Tensor | None: + if tensor is None: + return None + return tensor[token_slice] + + +@overload +def _slice_and_pad_first_dim( + tensor: torch.Tensor, + token_slice: slice, +) -> torch.Tensor: ... + + +@overload +def _slice_and_pad_first_dim( + tensor: None, + token_slice: slice, +) -> None: ... + + +def _slice_and_pad_first_dim( + tensor: torch.Tensor | None, + token_slice: slice, +) -> torch.Tensor | None: + if tensor is None: + return None + stage_tensor = tensor[token_slice] + expected_tokens = int(token_slice.stop) - int(token_slice.start) + missing_tokens = expected_tokens - int(stage_tensor.shape[0]) + if missing_tokens <= 0: + return stage_tensor + pad_shape = (missing_tokens, *tensor.shape[1:]) + return torch.cat([stage_tensor, tensor.new_zeros(pad_shape)], dim=0) + + +def _slice_positions(positions: torch.Tensor, token_slice: slice) -> torch.Tensor: + if positions.dim() <= 1: + return positions[token_slice] + return positions[..., token_slice] + + +def _slice_llama_4_scaling( + llama_4_scaling: torch.Tensor | None, + token_slice: slice, + *, + num_tokens: int, +) -> torch.Tensor | None: + if llama_4_scaling is None: + return None + if int(llama_4_scaling.shape[0]) == num_tokens: + return llama_4_scaling[token_slice] + if llama_4_scaling.dim() > 1 and int(llama_4_scaling.shape[1]) == num_tokens: + return llama_4_scaling[:, token_slice] + return llama_4_scaling + + +__all__ = [ + "build_async_moe_stage_inputs", + "restore_async_moe_stage_outputs", + "sp_local_actual_token_count", +] diff --git a/afd_plugin/v1/worker/npu/attention_model_runner.py b/afd_plugin/v1/worker/npu/attention_model_runner.py index dcc71974..5f6f88ae 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner.py @@ -90,8 +90,9 @@ snapshot_pcp_manager_state, ) from afd_plugin.v1.worker.npu.ubatch_utils import ( + ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP, check_enable_ubatch, - create_request_boundary_ubatch_slices, + create_async_moe_ubatch_slices, maybe_create_ubatch_slices, pad_out_ubatch_slices, split_attn_metadata, @@ -237,6 +238,14 @@ def _build_attention_metadata_with_async_moe_ubatches( kwargs: dict[str, Any], values: dict[str, Any], ) -> Any: + """Build attention metadata for async MoE ubatches. + + PCP compatibility: request-boundary stages remain supported when + sequence parallelism is disabled. Under SP, every stage must have a + TP-divisible token count and must be explicitly re-sharded from the + full-batch rank-major layout. A request-boundary stage cannot guarantee + that invariant, so such steps safely run without async MoE ubatching. + """ full_metadata = super()._build_attention_metadata(*args, **kwargs) self._afd_async_moe_ubatch_metadata = None self._afd_pending_metadata = self._build_afd_metadata( @@ -248,20 +257,50 @@ def _build_attention_metadata_with_async_moe_ubatches( if num_scheduled_tokens_np is None: return full_metadata - ubatch_slices = create_request_boundary_ubatch_slices( + ubatch_slices, split_mode = create_async_moe_ubatch_slices( + self.vllm_config, num_scheduled_tokens_np, + num_tokens=int(values.get("num_tokens", 0)), + num_tokens_padded=values.get("num_tokens_padded"), + num_reqs_padded=values.get("num_reqs_padded"), num_ubatches=self.afd_async_extra_info.async_moe_num_ubatches, + split=self.afd_async_extra_info.async_moe_split, + ) + can_ubatch_all_dp_ranks = torch.tensor( + int(ubatch_slices is not None), + dtype=torch.int32, + device="cpu", ) - if ubatch_slices is None: + if int(self.vllm_config.parallel_config.data_parallel_size) > 1: + dist.all_reduce( + can_ubatch_all_dp_ranks, + op=dist.ReduceOp.MIN, + group=get_dp_group().cpu_group, + ) + if not bool(can_ubatch_all_dp_ranks.item()): + logger.debug( + "AFD NPU async MoE ubatching disabled for this step because " + "at least one DP rank cannot build two valid stages", + ) + return full_metadata + assert ubatch_slices is not None + use_sp_stage_resharding = bool(enable_sp(self.vllm_config)) + if use_sp_stage_resharding and split_mode != ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP: + logger.debug( + "AFD NPU async MoE ubatching disabled for this SP step " + "because split_mode=%s is not TP-aligned", + split_mode, + ) return full_metadata logger.debug( "AFD NPU async MoE ubatch split; num_reqs=%s num_tokens=%s " - "num_scheduled_tokens=%s request_slices=%s token_slices=%s " - "stage_num_tokens=%s", + "num_scheduled_tokens=%s split_mode=%s request_slices=%s " + "token_slices=%s stage_num_tokens=%s", len(num_scheduled_tokens_np), int(values.get("num_tokens", 0)), num_scheduled_tokens_np.tolist(), + split_mode, [ (ubatch_slice.request_slice.start, ubatch_slice.request_slice.stop) for ubatch_slice in ubatch_slices @@ -289,6 +328,18 @@ def _build_attention_metadata_with_async_moe_ubatches( self._afd_async_moe_ubatch_metadata = { "attn_metadata": stage_attn_metadata, "ubatch_slices": ubatch_slices, + "use_sp_stage_resharding": use_sp_stage_resharding, + "stage_actual_token_counts": [ + max( + 0, + min( + int(ubatch_slice.token_slice.stop), + int(values.get("num_tokens", 0)), + ) + - int(ubatch_slice.token_slice.start), + ) + for ubatch_slice in ubatch_slices + ], } return full_metadata diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index c14f67d5..c7798a9d 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -272,6 +272,11 @@ def _ffn_forward( return rank_ffn_output def _ffn_forward_connector_driven(self) -> Any: + # CAM carries the authoritative layer index and token counts, but its + # dispatch-recv metadata has no stage identifier. Stage pairing is + # therefore FIFO in CAM collective-call order; this local value is only + # a placeholder for the per-item AFD context and must not be used to + # infer how Attention split the current scheduler step. stage_idx = 0 rank_ffn_output = None recv_work_item = getattr(self.connector, "recv_ffn_work_item", None) diff --git a/afd_plugin/v1/worker/npu/ubatch_utils.py b/afd_plugin/v1/worker/npu/ubatch_utils.py index 331eb1ca..a1260011 100644 --- a/afd_plugin/v1/worker/npu/ubatch_utils.py +++ b/afd_plugin/v1/worker/npu/ubatch_utils.py @@ -2,8 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """Ascend ubatch helpers owned by the AFD plugin. -Copied from vLLM-Ascend commit cdd212830271249a1cafcb850c210133f21771c5; -kept plugin-owned so AFD retains DBO support independent of upstream changes. +Ubatch token-space planning (split policies, padding, attention-metadata +rebuilds) plus the TP/SP-aware local-token mapping consumed by model-side +stage slicing. Parts of this module mirror the Ascend DBO logic from +vLLM-Ascend commit cdd212830271249a1cafcb850c210133f21771c5, kept +plugin-owned so AFD retains DBO support independent of upstream changes. """ import numpy as np @@ -141,6 +144,129 @@ def create_request_boundary_ubatch_slices( ] +ASYNC_MOE_SPLIT_REQUEST_BOUNDARY = "request_boundary" +ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP = "token_balanced_tp" + + +def enable_token_balanced_async_moe_split(vllm_config: VllmConfig) -> bool: + """Whether async MoE ubatches split on TP-aligned token counts. + + Non-PCP DP+TP/SP topologies can schedule requests with heavily skewed + lengths, where a request-boundary split produces imbalanced stages; a + token-balanced split keeps both pipeline stages close to half the token + workload. PCP topologies keep the request-boundary split because their + metadata is rebuilt per request. + """ + parallel_config = vllm_config.parallel_config + if int(parallel_config.tensor_parallel_size) <= 1: + return False + if int(parallel_config.prefill_context_parallel_size) != 1: + return False + return int(parallel_config.decode_context_parallel_size) == 1 + + +def token_balanced_split_points( + num_tokens: int, + num_ubatches: int, + tp_size: int, +) -> list[int] | None: + """Even token split points aligned down to multiples of ``tp_size``. + + Alignment keeps every stage divisible by the TP size so that, under + sequence parallelism, each rank's local stage shards have equal length. + Returns None when an aligned split is impossible (too few tokens, or + alignment collapses a stage boundary). + """ + if num_ubatches <= 1 or num_tokens < num_ubatches: + return None + split_points: list[int] = [] + for idx in range(1, num_ubatches): + split_point = (num_tokens * idx) // num_ubatches + if tp_size > 1: + split_point = (split_point // tp_size) * tp_size + if split_point <= 0 or split_point >= num_tokens: + return None + if split_points and split_point <= split_points[-1]: + return None + split_points.append(split_point) + return split_points + + +def create_async_moe_ubatch_slices( + vllm_config: VllmConfig, + num_scheduled_tokens_np: np.ndarray, + *, + num_tokens: int, + num_tokens_padded: int | None, + num_reqs_padded: int | None, + num_ubatches: int, + split: str, +) -> tuple[UBatchSlices | None, str]: + """Split async MoE work into ubatches and report the split policy used. + + ``split`` is the connector's ``async_moe_split`` value. ``"token"`` + selects the TP-aligned token-balanced split on non-PCP DP+TP/SP + topologies: the padded token count is split evenly with split points + aligned to the TP size, so both stages — and, under SP, every rank's + local stage shard — carry close to half the workload. When DP padding + makes the balanced split land past this rank's real tokens, the split + point is moved left to the largest TP-aligned boundary that still leaves + real tokens in both stages. If no such boundary exists, the step runs + unbatched; falling back to request boundaries would create stage lengths + that cannot be represented by equal SP shards. + + Any non-token policy selects the request-boundary split. Unsupported token + topologies also retain the request-boundary result because startup + validation rejects that configuration before serving. + + The returned policy string (``"token_balanced_tp"`` or + ``"request_boundary"``) describes the split actually applied to this + batch, so callers can adapt downstream handling without re-deriving the + decision; it is a string (not a boolean) so additional split policies + can be introduced without changing the contract. + + Returns ``(None, policy)`` when the batch is too small to split; callers + should run such steps unbatched. Raises AssertionError if + ``num_ubatches`` is not 2. + """ + assert num_ubatches == 2, "Async MoE ubatching currently supports 2 stages." + if split == "token" and enable_token_balanced_async_moe_split(vllm_config): + split_total = int(num_tokens_padded or num_tokens) + tp_size = int(vllm_config.parallel_config.tensor_parallel_size) + token_split_points = token_balanced_split_points( + split_total, + num_ubatches, + tp_size, + ) + if token_split_points is not None: + largest_real_split = ((int(num_tokens) - 1) // tp_size) * tp_size + token_split_points[-1] = min( + token_split_points[-1], + largest_real_split, + ) + if token_split_points is not None and token_split_points[-1] > 0: + ubatch_slices = create_ubatch_slices( + num_scheduled_tokens_np, + token_split_points, + ) + if num_tokens_padded is not None and num_reqs_padded is not None: + ubatch_slices = pad_out_ubatch_slices( + ubatch_slices, + int(num_tokens_padded), + int(num_reqs_padded), + ) + return ubatch_slices, ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP + return None, ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP + + return ( + create_request_boundary_ubatch_slices( + num_scheduled_tokens_np, + num_ubatches=num_ubatches, + ), + ASYNC_MOE_SPLIT_REQUEST_BOUNDARY, + ) + + def maybe_create_ubatch_slices( should_ubatch: bool, num_scheduled_tokens_per_request: np.ndarray, @@ -191,10 +317,35 @@ def _make_metadata_with_slice( request_slice = ubatch_slice.request_slice token_slice = ubatch_slice.token_slice start_locs = attn_metadata.query_start_loc_cpu + num_parent_actual_tokens = int(attn_metadata.num_actual_tokens) + actual_token_stop = min(int(token_slice.stop), num_parent_actual_tokens) + num_actual_tokens = actual_token_stop - int(token_slice.start) + if num_actual_tokens <= 0: + raise ValueError( + "Async MoE stage must contain at least one real token; " + f"got token_slice={token_slice} and " + f"num_actual_tokens={num_parent_actual_tokens}", + ) + + # DP/SP padding extends the final stage's input token range without + # creating real requests. Exclude repeated padded request entries from + # query metadata while retaining the full input range for positions and + # slot_mapping. + actual_request_stop = int( + np.searchsorted( + start_locs.numpy(), + actual_token_stop, + side="left", + ), + ) + request_slice = slice( + int(request_slice.start), + min(int(request_slice.stop), actual_request_stop), + ) first_req = request_slice.start first_tok = token_slice.start last_req = request_slice.stop - 1 - last_tok = token_slice.stop - 1 + last_tok = actual_token_stop - 1 assert start_locs[first_req] <= first_tok < start_locs[first_req + 1], ( "Token slice start outside of first request" @@ -222,7 +373,7 @@ def _make_metadata_with_slice( ) if splits_last_request: - tokens_skipped = start_locs[last_req + 1] - token_slice.stop + tokens_skipped = start_locs[last_req + 1] - actual_token_stop query_start_loc[-1] -= tokens_skipped query_start_loc_cpu[-1] -= tokens_skipped seq_lens = seq_lens.clone() @@ -242,7 +393,7 @@ def _make_metadata_with_slice( ) num_requests = request_slice.stop - request_slice.start - num_actual_tokens = token_slice.stop - token_slice.start + num_input_tokens = int(token_slice.stop) - int(token_slice.start) max_query_len = int( torch.max(torch.abs(query_start_loc_cpu[1:] - query_start_loc_cpu[:-1])).item() ) @@ -250,7 +401,9 @@ def _make_metadata_with_slice( max_query_len = attn_metadata.max_query_len if len(attn_metadata.actual_seq_lengths_q) > 0: - actual_seq_lengths_q = attn_metadata.actual_seq_lengths_q[token_slice] + actual_seq_lengths_q = attn_metadata.actual_seq_lengths_q[ + int(token_slice.start) : actual_token_stop + ] if max_num_tokens and len(actual_seq_lengths_q) == 0: actual_seq_lengths_q = list( range( @@ -277,7 +430,7 @@ def _make_metadata_with_slice( block_table_tensor=attn_metadata.block_table_tensor[request_slice], slot_mapping=attn_metadata.slot_mapping[token_slice], causal=attn_metadata.causal, - num_input_tokens=num_actual_tokens, + num_input_tokens=num_input_tokens, actual_seq_lengths_q=actual_seq_lengths_q, positions=attn_metadata.positions[token_slice], attn_state=attn_metadata.attn_state, @@ -312,14 +465,19 @@ def split_attn_metadata( __all__ = [ + "ASYNC_MOE_SPLIT_REQUEST_BOUNDARY", + "ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP", "UBatchSlice", "UBatchSlices", "check_enable_ubatch", + "create_async_moe_ubatch_slices", "create_request_boundary_ubatch_slices", "create_ubatch_slices", + "enable_token_balanced_async_moe_split", "is_last_ubatch_empty", "maybe_create_ubatch_slices", "pad_out_ubatch_slices", "slice_query_start_locs", "split_attn_metadata", + "token_balanced_split_points", ] diff --git a/docs/design/module/connector_contracts.md b/docs/design/module/connector_contracts.md index 6b108d8d..aabbd1f4 100644 --- a/docs/design/module/connector_contracts.md +++ b/docs/design/module/connector_contracts.md @@ -91,7 +91,7 @@ than an `AFDConfig` field. Unknown fields fail in the selected connector parser. | --- | --- | | `P2pNcclAFDConnector` | None; the mapping must be empty. | | `CAMP2pAFDConnector` | `core_num`, optional `attn_core_num` / `ffn_core_num`, `compute_gate_on_attention`, and `quant_mode`. Core counts must be positive; the current runtime rejects gate-on-Attention and any nonzero quantization mode. | -| `CAMAsyncAFDConnector` | `dynamicQuant`, `attn_ranks_per_dp`, `async_moe_ubatching`, `async_moe_num_ubatches`, and `async_moe_split`. Runtime validation further limits dynamic quantization and the optional request-boundary pipeline. | +| `CAMAsyncAFDConnector` | `dynamicQuant`, `attn_ranks_per_dp`, `async_moe_ubatching`, `async_moe_num_ubatches`, and `async_moe_split`. Runtime validation further limits dynamic quantization and the optional two-stage pipeline (request-boundary by default; token-balanced on non-PCP DP+TP/SP topologies). | The common `compute_gate_on_attention` field remains on `AFDConfig` and is the model-routing selector. CAMP2P also parses a connector-local field with that diff --git a/docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md b/docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md index 5ec5276b..cd3bc684 100644 --- a/docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md +++ b/docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md @@ -19,7 +19,8 @@ prefill path when all of the following are true: - Attention performs MoE gating before dispatch to FFN ranks; - execution is eager and AFD async-DP is enabled; **`async=true` is required**; - the service is the prefill stage of a prefill/decode-disaggregated deployment; -- optional MoE ubatching is managed by AFD as two request-boundary stages. +- optional MoE ubatching is managed by AFD as two stages (request-boundary by + default; token-balanced is available on non-PCP DP+TP/SP topologies). CAM async currently does not support decode, ACL graph execution, or vLLM native DBO. @@ -65,9 +66,11 @@ For `A = num_attention_ranks` and `F = num_ffn_ranks`: - world size is `A + F`; - each role rank must be unique and within its role's configured rank count. -Attention ranks are normally `DP x PCP`. `attn_ranks_per_dp` is the PCP width -and is also passed to CAM as its Attention TP width. For an Attention process -whose first data-parallel rank is `d`, use: +Attention ranks are `DP x PCP x TP`. `attn_ranks_per_dp` is the number of +Attention ranks in one DP replica (`PCP x TP`) and is also passed to CAM as its +Attention grouping width. It describes the Attention topology only; the FFN +process may use a different local TP size. For an Attention process whose first +data-parallel rank is `d`, use: ```text afd_role_rank = d * attn_ranks_per_dp @@ -87,6 +90,22 @@ FFN EP8 process: afd_role_rank = 0 CAM world ranks: A0..A23 = 0..23, F0..F7 = 24..31 ``` +For an Attention `DP3 x TP2` and FFN `DP2 x TP1 + EP2` deployment: + +```text +num_attention_ranks = 3 * 2 = 6 +num_ffn_ranks = 2 * 1 = 2 +attn_ranks_per_dp = 2 + +Attention role ranks: 0..5 +FFN role ranks: 0..1 +CAM world ranks: A0..A5 = 0..5, F0..F1 = 6..7 +``` + +Enable FlashComm1/SP only in the Attention process for this topology. The FFN +process uses TP1 with FlashComm1/SP disabled and relies on CAM plus expert +parallelism to route work across its two DP ranks. + FFN ranks follow expert parallel placement. The runtime derives experts per rank from the model routed-expert count and `num_ffn_ranks`; use a model/topology in which routed experts divide evenly across FFN ranks. All roles must use the same @@ -148,10 +167,10 @@ spelling used by the recipes. | Field | Type | Default | Meaning and constraint | | --- | --- | --- | --- | | `dynamicQuant` | `int` | `0` | Enables CAM dispatch/combine dynamic-quant metadata. Only `0` and `1` are accepted. With `1`, FFN receives quantized routed activations plus scale tensors and must return output compatible with combine-send. | -| `attn_ranks_per_dp` | `int` | `1` | Positive Attention rank count per DP replica, normally the PCP width. It affects Attention role-rank derivation and CAM TP size. | +| `attn_ranks_per_dp` | `int` | `1` | Positive Attention rank count per DP replica (`PCP x TP`). It must match the Attention distributed layout and is passed to CAM as its Attention grouping width; it is independent of the FFN process's local TP size. | | `async_moe_ubatching` | `bool` | `false` | Enables AFD-managed asynchronous MoE-only ubatching. | | `async_moe_num_ubatches` | `int` | `2` | Number of asynchronous MoE stages. Only `2` is supported. | -| `async_moe_split` | `str` | `"request"` | Stage split policy. The current async connector supports request-boundary splitting only. | +| `async_moe_split` | `str` | `"request"` | Stage split policy. `"request"` (default) splits at request boundaries and is required on PCP topologies. `"token"` splits the padded prefill token workload into two TP-aligned stages of approximately equal size; the Attention process must use a non-PCP DP+TP/SP topology (`tensor_parallel_size > 1`, no prefill/decode context parallel). The FFN process consumes CAM-routed stage payloads and may independently use TP1. | ## Native DBO and async MoE ubatching are different @@ -172,10 +191,19 @@ synchronous connector deployments. ### AFD-managed asynchronous MoE ubatching `async_moe_ubatching` pipelines only the MoE portion of CAM async execution. -Requests are divided at request boundaries into exactly two stages. Each stage -keeps its own pending Attention routing metadata so dispatch and combine remain -paired while Attention and FFN work overlap. It does not enable vLLM native DBO -and does not use the DBO threshold flags. +Requests are divided into exactly two stages. With the default +`async_moe_split="request"` the split happens at request boundaries, which is +required on PCP topologies. With `async_moe_split="token"` the padded prefill +token workload is split into two TP-aligned stages of approximately equal +size, which is useful on non-PCP DP+TP/SP topologies where request lengths +can be skewed. Sequence-parallel deployments must use the token policy: +request boundaries are not guaranteed to form equal TP shards, so a request +split runs safely without MoE stage pipelining for that step. If any DP rank +has too few real tokens for two TP-aligned stages, all DP ranks make the same +per-step unbatched decision to keep CAM collective call counts consistent. +Each stage keeps its own pending Attention routing metadata so dispatch and +combine remain paired while Attention and FFN work overlap. It does not enable +vLLM native DBO and does not use the DBO threshold flags. When `async_moe_ubatching=true`, all roles must set: @@ -230,8 +258,13 @@ available: `async_dispatch_send`, `async_dispatch_recv`, - Eager execution only; ACL graph mode is unsupported. - Prefill stage only in a prefill/decode-disaggregated deployment. - vLLM native DBO/ubatching is unsupported. -- AFD-managed MoE ubatching supports exactly two request-boundary stages. -- Decode context parallel metadata is unsupported with async MoE ubatching. +- AFD-managed MoE ubatching supports exactly two stages; request-boundary + splitting is the default, and token-balanced splitting is available on + non-PCP Attention DP+TP/SP topologies. Sequence parallelism requires the + token policy for stage pipelining. The FFN topology is independent and may + use TP1. +- Decode context parallel metadata is unsupported on Attention with async MoE + ubatching. - Routed experts should divide evenly across FFN ranks. - Other Ascend hardware, full unmodified DeepSeek-V3.2, different model families, CAM/CANN/container versions, cross-version combinations, and diff --git a/recipe/npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md b/recipe/npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md index a906d057..e9b67eff 100644 --- a/recipe/npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md +++ b/recipe/npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md @@ -69,13 +69,13 @@ for both sides. | `dynamicQuant` | Enables dynamic quantization metadata for CAM dispatch/combine. | | `async_moe_ubatching` | Enables AFD-managed MoE ubatching instead of vLLM native DBO. | | `async_moe_num_ubatches` | Number of async MoE stages. The current CAM async setup uses `2`. | -| `async_moe_split` | Split policy for async MoE ubatches. This recipe uses request-level splitting. | +| `async_moe_split` | Split policy for async MoE ubatches. `"request"` (default) splits at request boundaries; `"token"` splits the padded token workload into two TP-aligned stages of approximately equal size on non-PCP DP+TP/SP topologies. This recipe uses request-level splitting. | | `attn_ranks_per_dp` | Number of attention ranks per DP replica. With `PCP8`, this value is `8`. | Do not add `--enable-dbo`, `--dbo-decode-token-threshold`, or `--dbo-prefill-token-threshold` to these commands. Those flags enable vLLM native DBO, which CAM async rejects. `async_moe_ubatching` is AFD-managed, -MoE-only request-boundary staging and is not vLLM native DBO. +MoE-only two-stage staging and is not vLLM native DBO. ## Experiment Configuration diff --git a/tests/unit/model_executor/models/test_forward_context.py b/tests/unit/model_executor/models/test_forward_context.py index a7d4e776..8bb185b1 100644 --- a/tests/unit/model_executor/models/test_forward_context.py +++ b/tests/unit/model_executor/models/test_forward_context.py @@ -191,7 +191,18 @@ def test_deepseek_async_moe_ubatching_runs_attention_inside_stage_context(): assert "dense_end_layer = min(model.end_layer, first_moe_layer)" in ( async_ubatch_forward ) - assert "stage_hidden_states = [" in async_ubatch_forward + assert "stage_hidden_states," in async_ubatch_forward + assert "build_async_moe_stage_inputs(" in async_ubatch_forward + assert async_ubatch_forward.index( + "build_async_moe_stage_inputs(", + ) < async_ubatch_forward.index( + "for stage_idx in range(len(ubatch_slices)):", + ) + assert ( + "stage_hidden_states[stage_idx],\n" + " stage_residual[stage_idx],\n" + " ) = layer(" + ) in async_ubatch_forward assert ( "moe_layers = list(islice(model.layers, moe_start_layer, model.end_layer))" in async_ubatch_forward @@ -203,7 +214,7 @@ def test_deepseek_async_moe_ubatching_runs_attention_inside_stage_context(): async_ubatch_forward ) assert "def flush_pending_ffn_outputs()" not in async_ubatch_forward - assert "torch.cat(stage_hidden_states, dim=0)" in async_ubatch_forward + assert "restore_async_moe_stage_outputs(" in async_ubatch_forward assert "_run_async_moe_ubatch_layer(" not in executor_source assert "_recv_async_moe_ubatch_outputs(" not in executor_source assert "forward_context.attn_metadata = attn_metadata[stage_idx]" in executor_source @@ -233,6 +244,81 @@ def test_deepseek_async_moe_ubatching_runs_attention_inside_stage_context(): ) < (async_ubatch_forward.rindex("recv_stage_ffn(1)")) +def test_async_moe_stage_context_uses_global_input_and_local_transfer_counts(): + torch = pytest.importorskip("torch") + from vllm.v1.worker.ubatch_utils import UBatchSlice + + from afd_plugin.connectors import AFDForwardContextMetadata + from afd_plugin.model_executor.models.npu.deepseek_v2_async_cam_forward import ( + _use_async_moe_ubatch_forward_context, + ) + + parent_afd_metadata = AFDForwardContextMetadata( + tokens_start_loc=[0], + requests_start_loc=[0], + stage_idx=0, + connector=object(), + tokens_lens=[56], + num_stages=1, + transaction_id=1, + tokens_unpadded_lens=[40], + ) + parent_mask = torch.ones(112, dtype=torch.bool) + forward_context = SimpleNamespace( + attn_metadata="parent-attention", + additional_kwargs={"afd_metadata": parent_afd_metadata, "keep": "value"}, + ubatch_idx=0, + num_ubatches=1, + num_tokens=112, + pad_size=72, + padded_length=112, + max_tokens_across_dp=112, + padded_num_tokens=112, + mc2_mask=parent_mask, + dbo_enabled=False, + ) + global_slices = [ + UBatchSlice(slice(0, 1), slice(0, 38)), + UBatchSlice(slice(0, 1), slice(38, 112)), + ] + local_slices = [ + UBatchSlice(slice(0, 1), slice(0, 19)), + UBatchSlice(slice(0, 1), slice(19, 56)), + ] + sidecar = { + "attn_metadata": ["stage-0-attention", "stage-1-attention"], + "ubatch_slices": global_slices, + "sp_local_stage_slices": local_slices, + "stage_actual_token_counts": [38, 2], + "sp_local_stage_actual_token_counts": [19, 2], + } + + with _use_async_moe_ubatch_forward_context( + forward_context=forward_context, + parent_afd_metadata=parent_afd_metadata, + async_moe_ubatch_metadata=sidecar, + stage_idx=1, + ): + assert forward_context.attn_metadata == "stage-1-attention" + assert forward_context.num_tokens == 74 + assert forward_context.pad_size == 0 + assert forward_context.padded_length == 74 + assert forward_context.padded_num_tokens == 74 + assert forward_context.dbo_enabled is True + assert forward_context.mc2_mask[:2].all() + assert not forward_context.mc2_mask[2:].any() + stage_afd_metadata = forward_context.additional_kwargs["afd_metadata"] + assert stage_afd_metadata.tokens_lens == [37] + assert stage_afd_metadata.tokens_unpadded_lens == [2] + assert forward_context.additional_kwargs["keep"] == "value" + + assert forward_context.attn_metadata == "parent-attention" + assert forward_context.num_tokens == 112 + assert forward_context.pad_size == 72 + assert forward_context.mc2_mask is parent_mask + assert forward_context.dbo_enabled is False + + def test_deepseek_afd_ffn_path_reuses_ascend_moe_mlp_after_attention_gate(): source = Path("afd_plugin/model_executor/models/deepseek_v2.py").read_text() gate_source = Path( diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index ad323fd2..9952300c 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -173,6 +173,7 @@ def _parallel_config(**overrides): values = { "data_parallel_size": 1, "data_parallel_rank": 0, + "tensor_parallel_size": 1, "enable_dbo": False, "use_ubatching": False, "num_ubatches": 1, @@ -475,13 +476,20 @@ def test_npu_attention_metadata_positional_args_and_padded_slices(): assert normalized[-1].token_slice == slice(4, 8) -def test_npu_request_boundary_ubatch_slices_balance_tokens(monkeypatch): - np = pytest.importorskip("numpy") - fake_torch = ModuleType("torch") - fake_torch.Tensor = object +@contextmanager +def _vllm_stubbed_modules(monkeypatch, *, stub_torch=True): + """Stub torch/vLLM/vLLM-Ascend so plugin ubatch modules import on CPU.""" + if stub_torch: + fake_torch = ModuleType("torch") + fake_torch.Tensor = object + monkeypatch.setitem(sys.modules, "torch", fake_torch) fake_vllm = ModuleType("vllm") fake_vllm_config = ModuleType("vllm.config") fake_vllm_config.VllmConfig = object + fake_vllm_forward_context = ModuleType("vllm.forward_context") + fake_vllm_forward_context.DPMetadata = object + fake_vllm_forward_context.ForwardContext = object + fake_vllm_forward_context.get_forward_context = lambda: None fake_vllm_v1 = ModuleType("vllm.v1") fake_vllm_worker = ModuleType("vllm.v1.worker") fake_vllm_ubatch_utils = ModuleType("vllm.v1.worker.ubatch_utils") @@ -501,16 +509,26 @@ def is_empty(self): fake_vllm_ubatch_utils.UBatchSlice = UBatchSlice fake_vllm_ubatch_utils.UBatchSlices = list fake_vllm_ubatch_utils.check_ubatch_thresholds = lambda *_args, **_kwargs: False + fake_distributed = ModuleType("vllm.distributed") + fake_distributed.tensor_model_parallel_all_gather = lambda tensor, _dim: tensor + fake_parallel_state = ModuleType("vllm.distributed.parallel_state") + fake_parallel_state.get_tp_group = lambda: SimpleNamespace( + world_size=1, + rank_in_group=0, + ) + fake_distributed.parallel_state = fake_parallel_state fake_vllm_ascend = ModuleType("vllm_ascend") + fake_ascend_utils = ModuleType("vllm_ascend.utils") + fake_ascend_utils.enable_sp = lambda *_args, **_kwargs: False fake_forward_context = ModuleType("vllm_ascend.ascend_forward_context") fake_forward_context.MoECommType = type("MoECommType", (), {}) fake_attention = ModuleType("vllm_ascend.attention") fake_attention_utils = ModuleType("vllm_ascend.attention.utils") fake_attention_utils.AscendCommonAttentionMetadata = object - monkeypatch.setitem(sys.modules, "torch", fake_torch) monkeypatch.setitem(sys.modules, "vllm", fake_vllm) monkeypatch.setitem(sys.modules, "vllm.config", fake_vllm_config) + monkeypatch.setitem(sys.modules, "vllm.forward_context", fake_vllm_forward_context) monkeypatch.setitem(sys.modules, "vllm.v1", fake_vllm_v1) monkeypatch.setitem(sys.modules, "vllm.v1.worker", fake_vllm_worker) monkeypatch.setitem( @@ -518,7 +536,14 @@ def is_empty(self): "vllm.v1.worker.ubatch_utils", fake_vllm_ubatch_utils, ) + monkeypatch.setitem(sys.modules, "vllm.distributed", fake_distributed) + monkeypatch.setitem( + sys.modules, + "vllm.distributed.parallel_state", + fake_parallel_state, + ) monkeypatch.setitem(sys.modules, "vllm_ascend", fake_vllm_ascend) + monkeypatch.setitem(sys.modules, "vllm_ascend.utils", fake_ascend_utils) monkeypatch.setitem( sys.modules, "vllm_ascend.ascend_forward_context", @@ -530,11 +555,47 @@ def is_empty(self): "vllm_ascend.attention.utils", fake_attention_utils, ) + yield - module_name = "afd_plugin.v1.worker.npu.ubatch_utils" - original_module = sys.modules.pop(module_name, None) + +@contextmanager +def _fresh_import(*module_names): + """Re-import plugin modules so they bind to the stubbed runtime.""" + originals = {name: sys.modules.pop(name, None) for name in module_names} try: - ubatch_utils = importlib.import_module(module_name) + yield [importlib.import_module(name) for name in module_names] + finally: + for name in module_names: + sys.modules.pop(name, None) + for name, module in originals.items(): + if module is not None: + sys.modules[name] = module + + +@contextmanager +def _ubatch_utils_with_stubbed_runtime(monkeypatch): + with ( + _vllm_stubbed_modules(monkeypatch), + _fresh_import("afd_plugin.v1.worker.npu.ubatch_utils") as (module,), + ): + yield module + + +@contextmanager +def _ubatch_sp_with_stubbed_runtime(monkeypatch, *, real_torch): + with ( + _vllm_stubbed_modules(monkeypatch, stub_torch=not real_torch), + _fresh_import( + "afd_plugin.v1.worker.npu.ubatch_utils", + "afd_plugin.model_executor.models.npu.ubatch_sp", + ) as modules, + ): + yield modules[1] + + +def test_npu_request_boundary_ubatch_slices_balance_tokens(monkeypatch): + np = pytest.importorskip("numpy") + with _ubatch_utils_with_stubbed_runtime(monkeypatch) as ubatch_utils: slices = ubatch_utils.create_request_boundary_ubatch_slices( np.array([2, 3, 5, 7], dtype=np.int32), ) @@ -558,10 +619,816 @@ def is_empty(self): ) is None ) - finally: - sys.modules.pop(module_name, None) - if original_module is not None: - sys.modules[module_name] = original_module + + +def test_npu_token_balanced_split_points_align_to_tp_size(monkeypatch): + with _ubatch_utils_with_stubbed_runtime(monkeypatch) as ubatch_utils: + split_points = ubatch_utils.token_balanced_split_points + + # 100 padded tokens over 2 stages at tp=8: 50 aligned down to 48. + assert split_points(100, 2, 8) == [48] + # Already aligned split stays put. + assert split_points(96, 2, 8) == [48] + # Odd totals split as evenly as alignment allows (12/13). + assert split_points(25, 2, 2) == [12] + # A single stage never splits. + assert split_points(100, 1, 8) is None + # Fewer tokens than stages cannot split. + assert split_points(1, 2, 8) is None + # Alignment collapsing the split point disables the split. + assert split_points(4, 2, 8) is None + + +def test_npu_create_async_moe_ubatch_slices_topology_selection(monkeypatch): + np = pytest.importorskip("numpy") + with _ubatch_utils_with_stubbed_runtime(monkeypatch) as ubatch_utils: + + def _config(tp_size=1, pcp_size=1, dcp_size=1): + return SimpleNamespace( + parallel_config=SimpleNamespace( + tensor_parallel_size=tp_size, + prefill_context_parallel_size=pcp_size, + decode_context_parallel_size=dcp_size, + ), + ) + + def _create( + config, + scheduled, + split, + num_tokens_padded=None, + num_reqs_padded=None, + ): + return ubatch_utils.create_async_moe_ubatch_slices( + config, + scheduled, + num_tokens=int(scheduled.sum()), + num_tokens_padded=num_tokens_padded, + num_reqs_padded=num_reqs_padded, + num_ubatches=2, + split=split, + ) + + # Skewed lengths with async_moe_split='token' on a capable topology: + # token split balances 56/56 where the request-boundary split would + # produce 100/12. + scheduled = np.array([100, 4, 4, 4], dtype=np.int32) + slices, split_mode = _create( + _config(tp_size=2), + scheduled, + "token", + num_tokens_padded=112, + num_reqs_padded=4, + ) + assert split_mode == ubatch_utils.ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP + assert [int(stage.num_tokens) for stage in slices] == [56, 56] + # Complete coverage, stable order, no overlap; the straddling long + # request appears in both stages' request ranges. + assert slices[0].token_slice == slice(0, 56) + assert slices[1].token_slice == slice(56, 112) + assert slices[0].request_slice == slice(0, 1) + assert slices[1].request_slice == slice(0, 4) + + # async_moe_split='request' keeps request boundaries even on a + # capable topology (100/12 here). + slices, split_mode = _create(_config(tp_size=2), scheduled, "request") + assert split_mode == ubatch_utils.ASYNC_MOE_SPLIT_REQUEST_BOUNDARY + assert [int(stage.num_tokens) for stage in slices] == [100, 12] + + # async_moe_split='token' on unsupported topologies falls back to + # request boundaries (startup validation rejects this combination + # before serving). + for config in ( + _config(tp_size=1), + _config(tp_size=2, pcp_size=2), + _config(tp_size=2, dcp_size=2), + ): + slices, split_mode = _create(config, scheduled, "token") + assert split_mode == ubatch_utils.ASYNC_MOE_SPLIT_REQUEST_BOUNDARY + assert [int(stage.num_tokens) for stage in slices] == [100, 12] + + +def test_npu_create_async_moe_ubatch_slices_padded_batches(monkeypatch): + np = pytest.importorskip("numpy") + with _ubatch_utils_with_stubbed_runtime(monkeypatch) as ubatch_utils: + config = SimpleNamespace( + parallel_config=SimpleNamespace( + tensor_parallel_size=2, + prefill_context_parallel_size=1, + decode_context_parallel_size=1, + ), + ) + + # Normal padding: split points come from the padded total (26 -> + # 13 -> aligned 12) and the last stage extends over the pad rows. + scheduled = np.array([6, 6, 6, 6], dtype=np.int32) + slices, split_mode = ubatch_utils.create_async_moe_ubatch_slices( + config, + scheduled, + num_tokens=24, + num_tokens_padded=26, + num_reqs_padded=5, + num_ubatches=2, + split="token", + ) + assert split_mode == ubatch_utils.ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP + assert slices[0].request_slice == slice(0, 2) + assert slices[0].token_slice == slice(0, 12) + assert slices[1].request_slice == slice(2, 5) + assert slices[1].token_slice == slice(12, 26) + + # Heavy padding (4 real tokens padded to a 64-token capture size) + # has no TP=8-aligned boundary that leaves real tokens in both + # stages. It runs unbatched instead of falling back to unsafe + # request-boundary SP stages. + config_tp8 = SimpleNamespace( + parallel_config=SimpleNamespace( + tensor_parallel_size=8, + prefill_context_parallel_size=1, + decode_context_parallel_size=1, + ), + ) + scheduled = np.array([2, 2], dtype=np.int32) + slices, split_mode = ubatch_utils.create_async_moe_ubatch_slices( + config_tp8, + scheduled, + num_tokens=4, + num_tokens_padded=64, + num_reqs_padded=2, + num_ubatches=2, + split="token", + ) + assert split_mode == ubatch_utils.ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP + assert slices is None + + # Reproduce the failing light-DP-rank shape: the padded midpoint is + # 56, but only 40 rows are real. Move the boundary to 38 so both + # stages remain TP aligned and contain real work. + scheduled = np.array([20, 20], dtype=np.int32) + slices, split_mode = ubatch_utils.create_async_moe_ubatch_slices( + config, + scheduled, + num_tokens=40, + num_tokens_padded=112, + num_reqs_padded=2, + num_ubatches=2, + split="token", + ) + assert split_mode == ubatch_utils.ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP + assert [stage.token_slice for stage in slices] == [ + slice(0, 38), + slice(38, 112), + ] + assert [int(stage.num_tokens) for stage in slices] == [38, 74] + + # A single long request can straddle both token-balanced stages. The + # shared request range is intentional: stage 1's attention metadata + # retains the KV prefix produced by stage 0. + scheduled = np.array([40], dtype=np.int32) + slices, split_mode = ubatch_utils.create_async_moe_ubatch_slices( + config, + scheduled, + num_tokens=40, + num_tokens_padded=112, + num_reqs_padded=1, + num_ubatches=2, + split="token", + ) + assert split_mode == ubatch_utils.ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP + assert [stage.request_slice for stage in slices] == [ + slice(0, 1), + slice(0, 1), + ] + assert [stage.token_slice for stage in slices] == [ + slice(0, 38), + slice(38, 112), + ] + + # One real token cannot form two non-empty TP-aligned stages. + scheduled = np.array([1], dtype=np.int32) + slices, split_mode = ubatch_utils.create_async_moe_ubatch_slices( + config, + scheduled, + num_tokens=1, + num_tokens_padded=112, + num_reqs_padded=1, + num_ubatches=2, + split="token", + ) + assert split_mode == ubatch_utils.ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP + assert slices is None + + # Odd token counts still cover every token exactly once. + scheduled = np.array([13, 12], dtype=np.int32) + slices, split_mode = ubatch_utils.create_async_moe_ubatch_slices( + config, + scheduled, + num_tokens=25, + num_tokens_padded=26, + num_reqs_padded=2, + num_ubatches=2, + split="token", + ) + assert split_mode == ubatch_utils.ASYNC_MOE_SPLIT_TOKEN_BALANCED_TP + assert slices[0].token_slice == slice(0, 12) + assert slices[1].token_slice == slice(12, 26) + assert [int(stage.num_tokens) for stage in slices] == [12, 14] + + +def test_npu_split_async_moe_metadata_separates_real_and_padded_tokens(monkeypatch): + torch = pytest.importorskip("torch") + with _ubatch_sp_with_stubbed_runtime(monkeypatch, real_torch=True): + ubatch_utils = sys.modules["afd_plugin.v1.worker.npu.ubatch_utils"] + monkeypatch.setattr( + ubatch_utils, + "AscendCommonAttentionMetadata", + lambda **kwargs: SimpleNamespace(**kwargs), + ) + parent = SimpleNamespace( + query_start_loc=torch.tensor([0, 20, 40], dtype=torch.int32), + query_start_loc_cpu=torch.tensor([0, 20, 40], dtype=torch.int32), + seq_lens=torch.tensor([20, 20], dtype=torch.int32), + seq_lens_cpu=torch.tensor([20, 20], dtype=torch.int32), + num_computed_tokens_cpu=torch.tensor([0, 0], dtype=torch.int32), + num_reqs=2, + num_actual_tokens=40, + max_query_len=20, + max_seq_len=20, + block_table_tensor=torch.zeros((2, 1), dtype=torch.int32), + slot_mapping=torch.arange(112), + causal=True, + num_input_tokens=112, + actual_seq_lengths_q=list(range(1, 41)), + positions=torch.arange(112), + attn_state=object(), + graph_pad_size=0, + decode_token_per_req=1, + kvcomp_metadata=None, + encoder_seq_lens=None, + encoder_seq_lens_cpu=None, + logits_indices_padded=None, + num_logits_indices=0, + ) + stage = ubatch_utils.UBatchSlice(slice(1, 2), slice(38, 112)) + + metadata = ubatch_utils.split_attn_metadata([stage], parent)[0] + + assert metadata.num_actual_tokens == 2 + assert metadata.num_input_tokens == 74 + assert metadata.num_reqs == 1 + assert metadata.query_start_loc_cpu.tolist() == [0, 2] + assert metadata.positions.tolist() == list(range(38, 112)) + assert metadata.slot_mapping.tolist() == list(range(38, 112)) + assert metadata.actual_seq_lengths_q == [39, 40] + + +def test_ubatch_sp_stage_inputs_reshard_global_stage_ranges(monkeypatch): + torch = pytest.importorskip("torch") + with _ubatch_sp_with_stubbed_runtime(monkeypatch, real_torch=True) as ubatch_sp: + ubatch_utils = sys.modules["afd_plugin.v1.worker.npu.ubatch_utils"] + ubatch_slices = [ + ubatch_utils.UBatchSlice(slice(0, 2), slice(0, 10)), + ubatch_utils.UBatchSlice(slice(2, 4), slice(10, 16)), + ] + + global_hidden_states = torch.arange(32, dtype=torch.float32).reshape(16, 2) + global_residual = global_hidden_states + 100 + positions = torch.arange(16) + + for tp_rank, expected_stage_positions in ( + (0, [[0, 1, 2, 3, 4], [10, 11, 12]]), + (1, [[5, 6, 7, 8, 9], [13, 14, 15]]), + ): + monkeypatch.setattr( + ubatch_sp, + "get_tp_group", + lambda rank=tp_rank: SimpleNamespace( + world_size=2, + rank_in_group=rank, + ), + ) + local_slice = slice(tp_rank * 8, (tp_rank + 1) * 8) + hidden_states = global_hidden_states[local_slice] + residual = global_residual[local_slice] + + def _all_gather(tensor, token_dim): + assert token_dim == 0 + if float(tensor[0, 0]) >= 100: + return global_residual + return global_hidden_states + + monkeypatch.setattr( + ubatch_sp, + "tensor_model_parallel_all_gather", + _all_gather, + ) + ( + stage_hidden_states, + stage_residual, + stage_positions, + stage_llama_4_scaling, + sp_local_stage_slices, + ) = ubatch_sp.build_async_moe_stage_inputs( + hidden_states=hidden_states, + residual=residual, + positions=positions, + llama_4_scaling=None, + ubatch_slices=ubatch_slices, + use_sp_stage_resharding=True, + ) + + assert [ + tensor_slice.token_slice for tensor_slice in sp_local_stage_slices + ] == [ + slice(0, 5), + slice(5, 8), + ] + assert [int(stage.shape[0]) for stage in stage_hidden_states] == [5, 3] + assert [int(stage.shape[0]) for stage in stage_residual] == [5, 3] + assert [stage.tolist() for stage in stage_positions] == ( + expected_stage_positions + ) + assert stage_llama_4_scaling == [None, None] + expected_hidden_rows = ( + [global_hidden_states[:5], global_hidden_states[10:13]] + if tp_rank == 0 + else [global_hidden_states[5:10], global_hidden_states[13:16]] + ) + assert all( + torch.equal(actual, expected) + for actual, expected in zip( + stage_hidden_states, + expected_hidden_rows, + strict=True, + ) + ) + + # use_sp_stage_resharding=False keeps the global slicing path. + full_hidden_states = torch.arange(32, dtype=torch.float32).reshape(16, 2) + ( + stage_hidden_states, + stage_residual, + stage_positions, + _, + sp_local_stage_slices, + ) = ubatch_sp.build_async_moe_stage_inputs( + hidden_states=full_hidden_states, + residual=None, + positions=positions, + llama_4_scaling=None, + ubatch_slices=ubatch_slices, + use_sp_stage_resharding=False, + ) + assert sp_local_stage_slices is ubatch_slices + assert [int(stage.shape[0]) for stage in stage_hidden_states] == [10, 6] + assert stage_residual == [None, None] + assert stage_positions[0].tolist() == list(range(10)) + assert stage_positions[1].tolist() == list(range(10, 16)) + + +def test_ubatch_sp_stage_inputs_2d_positions_and_llama4_scaling(monkeypatch): + torch = pytest.importorskip("torch") + with _ubatch_sp_with_stubbed_runtime(monkeypatch, real_torch=True) as ubatch_sp: + ubatch_utils = sys.modules["afd_plugin.v1.worker.npu.ubatch_utils"] + ubatch_slices = [ + ubatch_utils.UBatchSlice(slice(0, 2), slice(0, 10)), + ubatch_utils.UBatchSlice(slice(2, 4), slice(10, 16)), + ] + + global_hidden_states = torch.arange(32, dtype=torch.float32).reshape(16, 2) + global_residual = global_hidden_states + 100 + # 2D positions with tokens on dim=1. + positions = torch.arange(16, dtype=torch.float32).reshape(1, 16) + # 2D llama_4_scaling with tokens on dim=1. + llama_4_scaling = torch.ones(2, 16, dtype=torch.float32) + + for tp_rank, expected_stage_positions, expected_stage_scaling in ( + ( + 0, + [[[0, 1, 2, 3, 4]], [[10, 11, 12]]], + [[[1.0] * 5] * 2, [[1.0] * 3] * 2], + ), + ( + 1, + [[[5, 6, 7, 8, 9]], [[13, 14, 15]]], + [[[1.0] * 5] * 2, [[1.0] * 3] * 2], + ), + ): + monkeypatch.setattr( + ubatch_sp, + "get_tp_group", + lambda rank=tp_rank: SimpleNamespace( + world_size=2, + rank_in_group=rank, + ), + ) + local_slice = slice(tp_rank * 8, (tp_rank + 1) * 8) + hidden_states = global_hidden_states[local_slice] + residual = global_residual[local_slice] + + def _all_gather(tensor, token_dim): + assert token_dim == 0 + if float(tensor[0, 0]) >= 100: + return global_residual + return global_hidden_states + + monkeypatch.setattr( + ubatch_sp, + "tensor_model_parallel_all_gather", + _all_gather, + ) + ( + stage_hidden_states, + stage_residual, + stage_positions, + stage_llama_4_scaling, + sp_local_stage_slices, + ) = ubatch_sp.build_async_moe_stage_inputs( + hidden_states=hidden_states, + residual=residual, + positions=positions, + llama_4_scaling=llama_4_scaling, + ubatch_slices=ubatch_slices, + use_sp_stage_resharding=True, + ) + + assert [int(stage.shape[0]) for stage in stage_hidden_states] == [5, 3] + assert [int(stage.shape[0]) for stage in stage_residual] == [5, 3] + assert [stage.tolist() for stage in stage_positions] == ( + expected_stage_positions + ) + assert [stage.tolist() for stage in stage_llama_4_scaling] == ( + expected_stage_scaling + ) + + +def test_ubatch_sp_stage_outputs_restore_full_batch_rank_layout(monkeypatch): + torch = pytest.importorskip("torch") + with _ubatch_sp_with_stubbed_runtime(monkeypatch, real_torch=True) as ubatch_sp: + ubatch_utils = sys.modules["afd_plugin.v1.worker.npu.ubatch_utils"] + ubatch_slices = [ + ubatch_utils.UBatchSlice(slice(0, 2), slice(0, 10)), + ubatch_utils.UBatchSlice(slice(2, 4), slice(10, 16)), + ] + global_stages = [ + torch.arange(10, dtype=torch.float32).reshape(10, 1), + torch.arange(10, 16, dtype=torch.float32).reshape(6, 1), + ] + + for tp_rank, expected in ( + (0, list(range(8))), + (1, list(range(8, 16))), + ): + monkeypatch.setattr( + ubatch_sp, + "get_tp_group", + lambda rank=tp_rank: SimpleNamespace( + world_size=2, + rank_in_group=rank, + ), + ) + stage_outputs = [ + stage[ + tp_rank * (int(stage.shape[0]) // 2) : (tp_rank + 1) + * (int(stage.shape[0]) // 2) + ] + for stage in global_stages + ] + + def _all_gather(tensor, token_dim): + assert token_dim == 0 + return ( + global_stages[0] if int(tensor.shape[0]) == 5 else global_stages[1] + ) + + monkeypatch.setattr( + ubatch_sp, + "tensor_model_parallel_all_gather", + _all_gather, + ) + restored = ubatch_sp.restore_async_moe_stage_outputs( + stage_outputs, + ubatch_slices, + use_sp_stage_resharding=True, + ) + assert restored[:, 0].tolist() == expected + assert ubatch_sp.sp_local_actual_token_count( + stage_actual_tokens=2, + stage_input_tokens=74, + ) == (2 if tp_rank == 0 else 0) + + +def test_ubatch_sp_stage_inputs_rejects_high_dim_sequence_tensor(monkeypatch): + torch = pytest.importorskip("torch") + with _ubatch_sp_with_stubbed_runtime(monkeypatch, real_torch=True) as ubatch_sp: + ubatch_utils = sys.modules["afd_plugin.v1.worker.npu.ubatch_utils"] + ubatch_slices = [ + ubatch_utils.UBatchSlice(slice(0, 2), slice(0, 10)), + ubatch_utils.UBatchSlice(slice(2, 4), slice(10, 16)), + ] + + hidden_states = torch.arange(16, dtype=torch.float32).reshape(8, 2) + # 3D positions: token dimension is on axis 2, which is unsupported. + positions = torch.zeros(2, 2, 16, dtype=torch.float32) + + monkeypatch.setattr( + ubatch_sp, + "get_tp_group", + lambda: SimpleNamespace(world_size=2, rank_in_group=0), + ) + monkeypatch.setattr( + ubatch_sp, + "tensor_model_parallel_all_gather", + lambda tensor, _token_dim: torch.cat([tensor, tensor], dim=0), + ) + with pytest.raises(ValueError, match="token dimension must be on axis 0 or 1"): + ubatch_sp.build_async_moe_stage_inputs( + hidden_states=hidden_states, + residual=None, + positions=positions, + llama_4_scaling=None, + ubatch_slices=ubatch_slices, + use_sp_stage_resharding=True, + ) + + +def test_npu_async_moe_dense_prefix_runs_inside_each_stage_context(monkeypatch): + _require_npu_runtime() + from vllm.v1.worker.ubatch_utils import UBatchSlice + + from afd_plugin.model_executor.models.npu import deepseek_v2_async_cam_forward + + calls = [] + forward_context = SimpleNamespace(stage_idx=None) + ubatch_slices = [ + UBatchSlice(slice(0, 1), slice(0, 4)), + UBatchSlice(slice(1, 2), slice(4, 8)), + ] + + class DenseLayer: + def __init__(self, layer_idx): + self.layer_idx = layer_idx + + def __call__( + self, + positions, + hidden_states, + residual, + llama_4_scaling, + ): + calls.append( + ( + forward_context.stage_idx, + self.layer_idx, + positions, + hidden_states, + llama_4_scaling, + ), + ) + return ( + f"{hidden_states}:dense{self.layer_idx}", + f"residual:{forward_context.stage_idx}:{self.layer_idx}", + ) + + @contextmanager + def stage_context(*, stage_idx, **_kwargs): + previous_stage_idx = forward_context.stage_idx + forward_context.stage_idx = stage_idx + try: + yield + finally: + forward_context.stage_idx = previous_stage_idx + + monkeypatch.setattr( + deepseek_v2_async_cam_forward, + "get_forward_context", + lambda: forward_context, + ) + monkeypatch.setattr( + deepseek_v2_async_cam_forward, + "build_async_moe_stage_inputs", + lambda **_kwargs: ( + ["hidden0", "hidden1"], + [None, None], + ["positions0", "positions1"], + ["scaling0", "scaling1"], + ubatch_slices, + ), + ) + monkeypatch.setattr( + deepseek_v2_async_cam_forward, + "_use_async_moe_ubatch_forward_context", + stage_context, + ) + monkeypatch.setattr( + deepseek_v2_async_cam_forward, + "restore_async_moe_stage_outputs", + lambda stage_outputs, *_args, **_kwargs: tuple(stage_outputs), + ) + + model = SimpleNamespace( + config=SimpleNamespace(first_k_dense_replace=2), + start_layer=0, + end_layer=2, + layers=[DenseLayer(0), DenseLayer(1)], + ) + sidecar = { + "attn_metadata": ["attention0", "attention1"], + "ubatch_slices": ubatch_slices, + "stage_actual_token_counts": [4, 4], + } + output, residual = deepseek_v2_async_cam_forward.run_async_moe_ubatch_afd_forward( + model=model, + hidden_states="full-hidden", + residual=None, + positions="full-positions", + afd_metadata=SimpleNamespace(connector=object()), + async_moe_ubatch_metadata=sidecar, + llama_4_scaling="full-scaling", + ) + + assert calls == [ + (0, 0, "positions0", "hidden0", "scaling0"), + (0, 1, "positions0", "hidden0:dense0", "scaling0"), + (1, 0, "positions1", "hidden1", "scaling1"), + (1, 1, "positions1", "hidden1:dense0", "scaling1"), + ] + assert output == ( + "hidden0:dense0:dense1", + "hidden1:dense0:dense1", + ) + assert residual == ( + "residual:0:1", + "residual:1:1", + ) + assert forward_context.stage_idx is None + + +def test_npu_attention_runner_builds_async_moe_ubatch_metadata(monkeypatch): + _require_npu_runtime() + import numpy as np + from vllm_ascend.worker.model_runner_v1 import NPUModelRunner + + from afd_plugin.connectors.npu.async_cam import AFDAsyncExtraInfo + from afd_plugin.v1.worker.npu import attention_model_runner + from afd_plugin.v1.worker.npu.attention_model_runner import ( + AFDNPUAttentionModelRunner, + ) + + runner = _new_attention_runner() + runner.vllm_config = _vllm_config( + role="attention", + connector="CAMAsyncAFDConnector", + async_dp=True, + tensor_parallel_size=2, + extra_config={ + "async_moe_ubatching": True, + "async_moe_split": "token", + }, + ) + runner.connector = _AsyncRecordingConnector() + runner._is_warmup = False + runner._afd_is_graph_capturing = False + runner._afd_pending_metadata = None + runner._afd_transaction_counter = 0 + runner.afd_async_extra_info = AFDAsyncExtraInfo( + async_moe_ubatching=True, + async_moe_split="token", + ) + + full_metadata = SimpleNamespace(name="full_metadata") + stage_attn_metadata = SimpleNamespace(name="stage_attn_metadata") + monkeypatch.setattr( + NPUModelRunner, + "_build_attention_metadata", + lambda self, *args, **kwargs: full_metadata, + ) + monkeypatch.setattr( + AFDNPUAttentionModelRunner, + "_build_attention_metadata_with_ubatches", + lambda self, *args, **kwargs: (stage_attn_metadata, None), + ) + monkeypatch.setattr( + attention_model_runner, + "enable_sp", + lambda *_args, **_kwargs: True, + ) + + values = { + "num_tokens": 16, + "num_tokens_padded": 16, + "num_reqs_padded": 4, + "num_scheduled_tokens_np": np.array([4, 4, 4, 4], dtype=np.int32), + } + result = runner._build_attention_metadata_with_async_moe_ubatches( + (), + {}, + values, + ) + + assert result is full_metadata + metadata = runner._afd_async_moe_ubatch_metadata + assert metadata["attn_metadata"] is stage_attn_metadata + assert metadata["use_sp_stage_resharding"] is True + assert metadata["stage_actual_token_counts"] == [8, 8] + assert len(metadata["ubatch_slices"]) == 2 + assert metadata["ubatch_slices"][0].token_slice == slice(0, 8) + assert metadata["ubatch_slices"][1].token_slice == slice(8, 16) + assert metadata["ubatch_slices"][0].request_slice == slice(0, 2) + assert metadata["ubatch_slices"][1].request_slice == slice(2, 4) + + # One incapable DP rank disables stage splitting for every DP rank so + # CAM participants keep an identical collective call count. + runner.vllm_config.parallel_config.data_parallel_size = 2 + monkeypatch.setattr( + attention_model_runner, + "get_dp_group", + lambda: SimpleNamespace(cpu_group=object()), + ) + + def _reject_ubatching_on_another_dp_rank(flag, **_kwargs): + flag.fill_(0) + + monkeypatch.setattr( + attention_model_runner.dist, + "all_reduce", + _reject_ubatching_on_another_dp_rank, + ) + result = runner._build_attention_metadata_with_async_moe_ubatches( + (), + {}, + values, + ) + assert result is full_metadata + assert runner._afd_async_moe_ubatch_metadata is None + + +def test_npu_attention_runner_builds_async_moe_ubatch_metadata_request_split( + monkeypatch, +): + _require_npu_runtime() + import numpy as np + from vllm_ascend.worker.model_runner_v1 import NPUModelRunner + + from afd_plugin.connectors.npu.async_cam import AFDAsyncExtraInfo + from afd_plugin.v1.worker.npu import attention_model_runner + from afd_plugin.v1.worker.npu.attention_model_runner import ( + AFDNPUAttentionModelRunner, + ) + + runner = _new_attention_runner() + runner.vllm_config = _vllm_config( + role="attention", + connector="CAMAsyncAFDConnector", + async_dp=True, + tensor_parallel_size=2, + extra_config={ + "async_moe_ubatching": True, + "async_moe_split": "request", + }, + ) + runner.connector = _AsyncRecordingConnector() + runner._is_warmup = False + runner._afd_is_graph_capturing = False + runner._afd_pending_metadata = None + runner._afd_transaction_counter = 0 + runner.afd_async_extra_info = AFDAsyncExtraInfo( + async_moe_ubatching=True, + async_moe_split="request", + ) + + full_metadata = SimpleNamespace(name="full_metadata") + stage_attn_metadata = SimpleNamespace(name="stage_attn_metadata") + monkeypatch.setattr( + NPUModelRunner, + "_build_attention_metadata", + lambda self, *args, **kwargs: full_metadata, + ) + monkeypatch.setattr( + AFDNPUAttentionModelRunner, + "_build_attention_metadata_with_ubatches", + lambda self, *args, **kwargs: (stage_attn_metadata, None), + ) + monkeypatch.setattr( + attention_model_runner, + "enable_sp", + lambda *_args, **_kwargs: True, + ) + + values = { + "num_tokens": 16, + "num_tokens_padded": 16, + "num_reqs_padded": 4, + "num_scheduled_tokens_np": np.array([4, 4, 4, 4], dtype=np.int32), + } + result = runner._build_attention_metadata_with_async_moe_ubatches( + (), + {}, + values, + ) + + assert result is full_metadata + assert runner._afd_async_moe_ubatch_metadata is None def test_npu_create_ascend_forward_context_marks_current_ubatch(monkeypatch): @@ -1181,6 +2048,55 @@ def test_npu_async_feature_validation_allows_dynamic_quant_zero_or_one(): ) +def test_npu_async_moe_ubatching_validation_suggests_token_split(caplog): + with caplog.at_level(logging.WARNING): + fail_if_unsupported_npu_afd_features( + _vllm_config( + connector="CAMAsyncAFDConnector", + async_dp=True, + compute_gate_on_attention=True, + tensor_parallel_size=2, + extra_config={"async_moe_ubatching": True}, + ), + ) + assert "async_moe_split='token'" in caplog.text + + +def test_npu_async_moe_token_split_allows_ffn_tp1_ep_topology(): + config = _vllm_config( + role="ffn", + connector="CAMAsyncAFDConnector", + async_dp=True, + compute_gate_on_attention=True, + data_parallel_size=2, + tensor_parallel_size=1, + enable_expert_parallel=True, + extra_config={ + "async_moe_ubatching": True, + "async_moe_split": "token", + }, + ) + fail_if_unsupported_npu_afd_features(config) + assert npu_afd_num_ubatches(config) == 1 + + with pytest.raises(RuntimeError, match="must be 'request' or 'token'"): + fail_if_unsupported_npu_afd_features( + _vllm_config( + role="ffn", + connector="CAMAsyncAFDConnector", + async_dp=True, + compute_gate_on_attention=True, + data_parallel_size=2, + tensor_parallel_size=1, + enable_expert_parallel=True, + extra_config={ + "async_moe_ubatching": True, + "async_moe_split": "layer", + }, + ), + ) + + def test_npu_async_moe_ubatching_validation_requires_supported_shape(): fail_if_unsupported_npu_afd_features( _vllm_config( @@ -1215,7 +2131,7 @@ def test_npu_async_moe_ubatching_validation_requires_supported_shape(): ), ) - with pytest.raises(RuntimeError, match="request-boundary"): + with pytest.raises(RuntimeError, match=r"non-PCP Attention DP\+TP/SP"): fail_if_unsupported_npu_afd_features( _vllm_config( connector="CAMAsyncAFDConnector", @@ -1228,6 +2144,33 @@ def test_npu_async_moe_ubatching_validation_requires_supported_shape(): ), ) + with pytest.raises(RuntimeError, match="must be 'request' or 'token'"): + fail_if_unsupported_npu_afd_features( + _vllm_config( + connector="CAMAsyncAFDConnector", + async_dp=True, + compute_gate_on_attention=True, + extra_config={ + "async_moe_ubatching": True, + "async_moe_split": "layer", + }, + ), + ) + + # async_moe_split="token" is accepted on non-PCP DP+TP/SP topologies. + fail_if_unsupported_npu_afd_features( + _vllm_config( + connector="CAMAsyncAFDConnector", + async_dp=True, + compute_gate_on_attention=True, + tensor_parallel_size=2, + extra_config={ + "async_moe_ubatching": True, + "async_moe_split": "token", + }, + ), + ) + fail_if_unsupported_npu_afd_features( _vllm_config( connector="CAMAsyncAFDConnector",