diff --git a/afd_plugin/__init__.py b/afd_plugin/__init__.py index eeff9c467..2c6dabd5f 100644 --- a/afd_plugin/__init__.py +++ b/afd_plugin/__init__.py @@ -116,8 +116,12 @@ def get_spawn_context(method: str | None = None): "afd_plugin.model_executor.models.npu.deepseek_v4_dspark:" "AFDDSparkDeepseekV4ForCausalLM" ), + # CUDA follows native vLLM, which serves GLM-5.x with the NVIDIA DSA model; + # NPU keeps the DeepSeek V2-derived wrapper. "GlmMoeDsaForCausalLM": ( "afd_plugin.model_executor.models.deepseek_v2:AFDGlmMoeDsaForCausalLM" + if importlib.util.find_spec("torch_npu") is not None + else "afd_plugin.model_executor.models.deepseek_v32:AFDGlmMoeDsaForCausalLM" ), } @@ -177,6 +181,7 @@ def register_afd() -> None: try: import afd_plugin.compat.patches.async_dp_engine # noqa: F401 import afd_plugin.compat.patches.async_dp_forward_context # noqa: F401 + import afd_plugin.compat.patches.config_dbo_spec_decode # noqa: F401 import afd_plugin.compat.patches.config_validation # noqa: F401 import afd_plugin.compat.patches.engine_core # noqa: F401 except Exception: diff --git a/afd_plugin/compat/patches/config_dbo_spec_decode.py b/afd_plugin/compat/patches/config_dbo_spec_decode.py new file mode 100644 index 000000000..1b0911733 --- /dev/null +++ b/afd_plugin/compat/patches/config_dbo_spec_decode.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Config gate of the vLLM #61464 backport (MRV2 DBO with speculative decoding). + +Upstream pin: 5df63a7203cdb9328d824a27b682e77c7b9928be. Upstream fix: +https://github.com/vllm-project/vllm/pull/61464 (merged after 0.32.0, merge +commit cdd2750cd5e7d90849282a16bc46c21bbc58bcbd). + +vLLM 0.32 rejects DBO together with speculative decoding on ModelRunnerV2. +For AFD configs, whose Attention runner installs the request-boundary split +and the draft-prefill rebuild (``gpu_mrv2_dbo_spec_decode``), drop that entry +and raise the DBO thresholds to one request per microbatch, as #61464 does. +Non-AFD configs keep the native gate. Removal: delete when the pinned vLLM +contains #61464. +""" + +import vllm.config.vllm as config_module +from vllm import envs +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.logger import init_logger + +from afd_plugin.config import parse_optional_afd_config + +logger = init_logger(__name__) + +_ORIGINAL_DBO_UNSUPPORTED_ATTR = "_afd_original_get_dbo_unsupported_features" +SPEC_DECODE_DBO_FEATURE = "dual batch overlap with speculative decoding" + + +def raise_dbo_thresholds_for_spec_decode(vllm_config: VllmConfig) -> None: + """#61464's ``VllmConfig._raise_dbo_thresholds_for_spec_decode``. + + A uniform decode batch below one request per microbatch cannot be split + on a request boundary, and an idle DP rank's one-request dummy batch must + stay below the thresholds too, so raise them to the number of microbatches + times the decode query length rather than reject the config. + """ + parallel_config = vllm_config.parallel_config + min_ubatch_tokens = ( + parallel_config.num_ubatches * vllm_config.uniform_decode_query_len + ) + for name in ("dbo_decode_token_threshold", "dbo_prefill_token_threshold"): + threshold = getattr(parallel_config, name) + if threshold < min_ubatch_tokens: + logger.warning_once( + "Raising %s from %d to %d, the number of microbatches times the " + "decode query length.", + name, + threshold, + min_ubatch_tokens, + ) + setattr(parallel_config, name, min_ubatch_tokens) + + +# Upstream source: vLLM v0.32.0 (5df63a7203cdb9328d824a27b682e77c7b9928be), +# vllm/config/vllm.py::VllmConfig._get_dbo_unsupported_features. +# Patch reason: the V2 runner rejects DBO with speculative decoding. +# Patch functionality: for AFD configs, drop the speculative-decoding entry and +# raise the DBO thresholds (#61464 raises them from __post_init__; this method +# runs in the same validation); every other restriction is kept. +# Signature: matches upstream; no added parameters. +def _get_dbo_unsupported_features(self) -> list[str]: + """Collect what the V2 model runner cannot combine with DBO. + + The V2 runner microbatches a plain decoder forward pass. Anything that + slices or replays the batch differently (drafting, adapters, pipeline + stages, context parallelism, encoders) is not handled yet. + """ + # TODO: DBO with model runner V2 is under development. + # It should be enabled with explicit VLLM_USE_V2_MODEL_RUNNER environ. + # Remove it when stable. + if envs.VLLM_USE_V2_MODEL_RUNNER is None: + return ["dual batch overlap"] + + unsupported: list[str] = [] + model_config = self.model_config + parallel_config = self.parallel_config + + if self.lora_config is not None: + unsupported.append("dual batch overlap with LoRA") + if self.speculative_config is not None: + # ### PATCH START: AFD backports MRV2 DBO with speculative decoding. + if parse_optional_afd_config(self) is None: + unsupported.append(SPEC_DECODE_DBO_FEATURE) + else: + raise_dbo_thresholds_for_spec_decode(self) + # ### PATCH END: AFD backports MRV2 DBO with speculative decoding. + if parallel_config.pipeline_parallel_size > 1: + unsupported.append("dual batch overlap with pipeline parallelism") + if ( + parallel_config.decode_context_parallel_size > 1 + or parallel_config.prefill_context_parallel_size > 1 + ): + unsupported.append("dual batch overlap with context parallelism") + if model_config is not None and ( + model_config.is_multimodal_model or model_config.is_encoder_decoder + ): + unsupported.append("dual batch overlap with multimodal models") + if model_config is not None and model_config.is_hybrid: + unsupported.append("dual batch overlap with hybrid models") + if self.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE: + # DBO captures FULL graphs only. + unsupported.append("dual batch overlap with PIECEWISE CUDA graphs") + if self.is_mm_encoder_only: + unsupported.append("dual batch overlap with encoder only models") + + return unsupported + + +def apply_config_dbo_spec_decode_patch() -> None: + """Install the gate patch once.""" + vllm_config_cls = config_module.VllmConfig + if hasattr(vllm_config_cls, _ORIGINAL_DBO_UNSUPPORTED_ATTR): + return + setattr( + vllm_config_cls, + _ORIGINAL_DBO_UNSUPPORTED_ATTR, + vllm_config_cls._get_dbo_unsupported_features, + ) + vllm_config_cls._get_dbo_unsupported_features = _get_dbo_unsupported_features + + +apply_config_dbo_spec_decode_patch() + +__all__ = [ + "SPEC_DECODE_DBO_FEATURE", + "apply_config_dbo_spec_decode_patch", + "raise_dbo_thresholds_for_spec_decode", +] diff --git a/afd_plugin/compat/patches/gpu_deepseek_v32_dbo.py b/afd_plugin/compat/patches/gpu_deepseek_v32_dbo.py new file mode 100644 index 000000000..1ae649792 --- /dev/null +++ b/afd_plugin/compat/patches/gpu_deepseek_v32_dbo.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Per-micro-batch DSA top-k indices for CUDA GLM-5.x DBO on vLLM 0.32.0. + +Upstream pin: 5df63a7203cdb9328d824a27b682e77c7b9928be. + +GLM-5.x runs its lightning indexer on one decoder layer in four +(``index_topk_freq``); the following layers reuse that layer's top-k. The +native NVIDIA DSA model keeps the result in one model-level +``topk_indices_buffer`` (rows ``[0, num_tokens)`` of the current batch) and in +one ``SparseMLAIndexGroup`` per indexer layer (physical indices, valid counts, +a side stream and two events). Every DBO micro-batch writes those same rows, +and the routed-experts connector yields between consecutive decoder layers, so +the other micro-batch would overwrite the top-k a reusing layer still has to +read. vLLM 0.32 keeps no per-micro-batch copy, and native DBO has the same gap. + +``install_per_ubatch_topk_indices`` gives each micro-batch its own buffer and +group copies. Each object that holds them (model, DSA attention, indexer, +indexer op, sparse MLA impl, including the impl's private +``_topk_indices_buffer``) gets a per-instance subclass whose attribute +resolves by ``dbo_current_ubatch_id()``; micro-batch 0, and every step that +does not micro-batch, keeps the native objects. Native forward code is not +copied. + +Removal tracking: no upstream issue yet. Remove when the pinned vLLM keeps +per-micro-batch sparse-index state. +""" + +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import replace +from types import MethodType + +import torch +from torch import nn +from vllm.config import CUDAGraphMode +from vllm.v1.attention.backends.mla.index_group import ( + HiSparseMLAIndexGroup, + SparseMLAIndexGroup, +) +from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadataBuilder +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.ubatch_utils import UBatchRunner, UBatchState +from vllm.v1.worker.ubatching import dbo_current_ubatch_id + +TOPK_INDICES_ATTRIBUTE = "topk_indices_buffer" +# Sparse MLA impls keep the buffer privately behind a native +# ``topk_indices_buffer`` property (``SharedTopkIndicesBuffer``); a layer +# without an indexer reads this one. +TOPK_INDICES_ATTRIBUTES = (TOPK_INDICES_ATTRIBUTE, "_topk_indices_buffer") +INDEX_GROUP_ATTRIBUTE = "index_group" +_PER_UBATCH_PREFIX = "_afd_per_ubatch_" +_PER_UBATCH_CLASS_CACHE: dict[tuple[type, frozenset[str]], type] = {} + + +def _per_ubatch_property(name: str) -> property: + slot_name = f"{_PER_UBATCH_PREFIX}{name}" + + def get(self: object) -> object: + return vars(self)[slot_name][dbo_current_ubatch_id()] + + def set_(self: object, value: object) -> None: + raise AttributeError(f"{name} is bound per micro-batch by AFD") + + return property(get, set_) + + +def _per_ubatch_class(cls: type, names: frozenset[str]) -> type: + key = (cls, names) + subclass = _PER_UBATCH_CLASS_CACHE.get(key) + if subclass is None: + # Patch reason: native DSA objects keep the top-k state as plain + # attributes bound at construction, shared by every micro-batch. + # Patch functionality: read those attributes from the current + # micro-batch's copy; all native methods are inherited unchanged. + namespace = {name: _per_ubatch_property(name) for name in names} + subclass = type(f"AFDPerUBatch{cls.__name__}", (cls,), namespace) + _PER_UBATCH_CLASS_CACHE[key] = subclass + return subclass + + +def _bind(obj: object, values: dict[str, tuple[object, ...]]) -> None: + """Make ``obj``'s attributes resolve per micro-batch to ``values``.""" + state = vars(obj) + for name, per_ubatch in values.items(): + # Keep the native object for micro-batch 0 and non-micro-batched steps. + if state.get(name) is not per_ubatch[0]: + raise RuntimeError( + f"{type(obj).__name__}.{name} is not the shared DSA top-k state" + ) + state[f"{_PER_UBATCH_PREFIX}{name}"] = per_ubatch + del state[name] + obj.__class__ = _per_ubatch_class(type(obj), frozenset(values)) + + +def _copy_group( + group: SparseMLAIndexGroup, logical: torch.Tensor +) -> SparseMLAIndexGroup: + if isinstance(group, HiSparseMLAIndexGroup): + raise RuntimeError("AFD GLM-5.x DBO does not support HiSparse index groups") + device = logical.device + return replace( + group, + logical_topk_indices=logical, + physical_topk_indices=torch.empty_like(group.physical_topk_indices), + valid_topk_counts=torch.empty_like(group.valid_topk_counts), + side_stream=torch.Stream(device=device), + logical_topk_ready=torch.Event(), + physical_topk_ready=torch.Event(), + ) + + +def _holders(model: nn.Module) -> Iterator[object]: + """Every object that may hold the shared top-k buffer or index group.""" + yield model + for module in model.modules(): + yield module + impl = vars(module).get("impl") + if impl is not None: + yield impl + + +def install_per_ubatch_topk_indices(model: nn.Module, num_ubatches: int) -> None: + """Give each DBO micro-batch its own DSA top-k buffer and index groups. + + ``model`` is the Attention-role DSA model whose ``topk_indices_buffer`` + every DSA layer shares. Raises when an object still holds the shared state + after binding, so a missed holder fails at startup instead of mixing two + micro-batches' indices. + """ + shared = vars(model)[TOPK_INDICES_ATTRIBUTE] + buffers = (shared,) + tuple( + torch.empty_like(shared) for _ in range(1, num_ubatches) + ) + group_copies: dict[int, tuple[SparseMLAIndexGroup, ...]] = {} + holders = list(_holders(model)) + for holder in holders: + state = vars(holder) + values: dict[str, tuple[object, ...]] = {} + for name in TOPK_INDICES_ATTRIBUTES: + if state.get(name) is shared: + values[name] = buffers + group = state.get(INDEX_GROUP_ATTRIBUTE) + if isinstance(group, SparseMLAIndexGroup): + if group.logical_topk_indices is not shared: + raise RuntimeError("DSA index group does not use the model top-k") + copies = group_copies.get(id(group)) + if copies is None: + copies = (group,) + tuple( + _copy_group(group, buffer) for buffer in buffers[1:] + ) + group_copies[id(group)] = copies + values[INDEX_GROUP_ATTRIBUTE] = copies + if values: + _bind(holder, values) + + for holder in holders: + leftovers = [ + name + for name, value in vars(holder).items() + if not name.startswith(_PER_UBATCH_PREFIX) + and (value is shared or isinstance(value, SparseMLAIndexGroup)) + ] + if leftovers: + raise RuntimeError( + f"{type(holder).__name__} still shares DSA top-k state: {leftovers}" + ) + + +@contextmanager +def use_dsa_indexer_sm_budget(ubatch_runner: UBatchRunner) -> Iterator[None]: + """Plan the DSA indexer for the micro-batch compute SM budget under DBO. + + The native indexer metadata builder plans its kernels for every physical + SM. Under DBO the routed-experts connector's dispatch/combine kernels of + the other micro-batch need SMs at the same time, and an indexer planned + for the whole GPU can keep them from running: one Attention rank never + dispatches and the AFD world deadlocks (observed on GLM-5.3 4A4F DBO). + The CUDA DeepSeek-V4 DBO path leases the same budget + (``gpu_deepseek_v4_dbo``); upstream fix https://github.com/vllm-project/vllm/pull/60732. + """ + state = vars(ubatch_runner) + had_override = "prepare" in state + previous_override = state.get("prepare") + original_prepare = ubatch_runner.prepare + + # Patch reason: native preparation plans the DSA indexer for every + # physical SM while DBO reserves SMs for communication. + # Patch functionality: lease the UBatchRunner compute SM budget to the + # native indexer builders during preparation and restore it on exit. + # Signature: matches UBatchRunner.prepare in + # vllm/v1/worker/gpu/ubatch_utils.py at vLLM + # 5df63a7203cdb9328d824a27b682e77c7b9928be. Delegation keeps native slicing + # and metadata construction. Remove when the pin contains #60732. + def prepare( + self: UBatchRunner, + input_batch: InputBatch, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, + for_capture: bool = False, + ) -> UBatchState: + # ### PATCH START: lease the compute SM budget to the DSA indexer. + budgets: list[tuple[DeepseekV32IndexerMetadataBuilder, int]] = [] + try: + for groups in self.attn_groups: + for group in groups: + for stage in range(self.num_ubatches): + builder = group.get_metadata_builder(stage) + if isinstance(builder, DeepseekV32IndexerMetadataBuilder): + budgets.append((builder, builder.num_sms)) + builder.num_sms = self.sm_control.compute_sms + return original_prepare( + input_batch, block_tables, slot_mappings, cg_mode, for_capture + ) + finally: + # Reverse order also restores correctly if groups share a builder. + for builder, num_sms in reversed(budgets): + builder.num_sms = num_sms + # ### PATCH END: lease the compute SM budget to the DSA indexer. + + ubatch_runner.prepare = MethodType(prepare, ubatch_runner) + try: + yield + finally: + if had_override: + ubatch_runner.prepare = previous_override + else: + del ubatch_runner.prepare + + +__all__ = ["install_per_ubatch_topk_indices", "use_dsa_indexer_sm_budget"] diff --git a/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py b/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py new file mode 100644 index 000000000..c00206112 --- /dev/null +++ b/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py @@ -0,0 +1,323 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Backport of vLLM #61464 (MRV2 DBO with speculative decoding) to vLLM 0.32.0. + +Upstream pin: 5df63a7203cdb9328d824a27b682e77c7b9928be. Upstream fix: +https://github.com/vllm-project/vllm/pull/61464 (merged after 0.32.0, merge +commit cdd2750cd5e7d90849282a16bc46c21bbc58bcbd). + +vLLM 0.32's ModelRunnerV2 rejects DBO with speculative decoding because it +only microbatches a plain decoder forward pass: + +* a midpoint split cuts a multi-token verification request in two whenever + the request count is odd, so neither microbatch is a uniform decode; +* draft prefill reuses the target's full-batch attention metadata and slot + mappings, which a microbatched step never builds (``propose`` receives + ``None``). + +This worker-side module backports, for the AFD Attention runner: + +1. a request-boundary split of uniform multi-token batches, fed the uniform + token count the target dispatch agreed across DP ranks; +2. draft prefill on the draft's own builders under DBO: its FULL graphs are + captured there, each replay first refreshes them with this batch's + full-batch metadata, and an eager draft prefill after a microbatched + target forward gets that metadata rebuilt (the pin captures on, and + reuses, the target's builders, whose full-batch buffers a microbatched + forward never writes). + +The config gate lives in ``config_dbo_spec_decode``. + +Removal: delete when the pinned vLLM contains #61464. +""" + +from collections.abc import Callable, Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from typing import Any + +import torch +from vllm.config import CUDAGraphMode +from vllm.v1.worker.gpu import model_runner as model_runner_module +from vllm.v1.worker.gpu import ubatch_utils +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu.dp_utils import DPSyncState +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import ( + TargetDependentARSpeculator, +) +from vllm.v1.worker.ubatch_utils import UBatchSlice, UBatchSlices + +_ORIGINAL_PROPOSE_ATTR = "_afd_original_target_dependent_ar_propose" +_ORIGINAL_CAPTURE_ATTR = "_afd_original_target_dependent_ar_capture" + + +# Upstream source: vLLM v0.32.0 (5df63a7203cdb9328d824a27b682e77c7b9928be), +# vllm/v1/worker/gpu/ubatch_utils.py::create_ubatch_slices. +# Patch reason: a midpoint split cuts a multi-token verification request in two. +# Patch functionality: split uniform multi-token batches on a request +# boundary, as #61464 does. +# Signature: adds ``uniform_token_count`` (the DP-agreed per-request query +# length, or None), as upstream #61464 does. +def create_ubatch_slices( + input_batch: InputBatch, + num_ubatches: int, + uniform_token_count: int | None = None, +) -> UBatchSlices: + """Split a DP-padded batch into the slices its microbatches run on. + + Splitting at the midpoint of the DP-padded token count leaves the trailing + microbatch anywhere from full to entirely inside the padding. Clamping its + request slice onto the last request leaves it holding that request with + zero query tokens -- the shape a straddling request already takes -- so it + stays well-formed with no work to do, like a dummy run. + """ + # ### PATCH START: request-boundary split of uniform multi-token batches. + split_point = None + if uniform_token_count is not None and uniform_token_count > 1: + num_tokens = input_batch.num_tokens_after_padding + reqs_per_ubatch = num_tokens // (uniform_token_count * num_ubatches) + if reqs_per_ubatch > 0: + split_point = reqs_per_ubatch * uniform_token_count + # ### PATCH END: request-boundary split of uniform multi-token batches. + _, ubatch_slices = ubatch_utils.maybe_create_ubatch_slices( + True, + input_batch.num_scheduled_tokens, + input_batch.num_tokens_after_padding, + input_batch.num_reqs_after_padding, + num_ubatches, + split_point=split_point, + ) + assert ubatch_slices is not None + return [ + UBatchSlice( + slice( + min(s.request_slice.start, input_batch.num_reqs - 1), + min(s.request_slice.stop, input_batch.num_reqs_after_padding), + ), + s.token_slice, + ) + for s in ubatch_slices + ] + + +def _agreed_uniform_token_count( + batch_desc: BatchExecutionDescriptor, dp_sync_state: DPSyncState | None +) -> int | None: + """#61464: FULL replays the split its graph captured with; eager steps use + the uniform token count agreed across DP ranks.""" + if batch_desc.cg_mode == CUDAGraphMode.FULL or dp_sync_state is None: + return batch_desc.uniform_token_count + return dp_sync_state.uniform_token_count + + +@contextmanager +def use_request_aligned_ubatch_split() -> Iterator[None]: + """Split one AFD ``execute_model`` call's microbatches on request boundaries. + + The pinned ``UBatchRunner.prepare`` does not receive the dispatch result, + so record the uniform token count the target dispatch agreed and give it + to the split. Enter it inside the scope that binds AFD's + ``dispatch_cg_and_sync_dp``. + """ + agreed: list[int | None] = [None] + original_dispatch = model_runner_module.dispatch_cg_and_sync_dp + original_split = ubatch_utils.create_ubatch_slices + + def dispatch_cg_and_sync_dp( + *args: Any, **kwargs: Any + ) -> tuple[BatchExecutionDescriptor, DPSyncState | None]: + batch_desc, dp_sync_state = original_dispatch(*args, **kwargs) + agreed[0] = _agreed_uniform_token_count(batch_desc, dp_sync_state) + return batch_desc, dp_sync_state + + def split(input_batch: InputBatch, num_ubatches: int) -> UBatchSlices: + return create_ubatch_slices(input_batch, num_ubatches, agreed[0]) + + model_runner_module.dispatch_cg_and_sync_dp = dispatch_cg_and_sync_dp + ubatch_utils.create_ubatch_slices = split + try: + yield + finally: + ubatch_utils.create_ubatch_slices = original_split + model_runner_module.dispatch_cg_and_sync_dp = original_dispatch + + +def _build_drafter_attn_metadata( + speculator: TargetDependentARSpeculator, + input_batch: InputBatch, + cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, + for_capture: bool = False, +) -> tuple[dict[str, Any], dict[str, torch.Tensor]]: + """#61464's ``DraftModelSpeculator._build_drafter_attn_metadata`` for the + rebuild branch: full-batch draft-prefill metadata on the draft's own + builders, from the block tables and slot mappings the target prepared. + With ``FULL`` it refreshes the buffers the draft-prefill graphs captured.""" + block_tables = tuple( + block_table[: input_batch.num_reqs_after_padding] + for block_table in speculator.block_tables.input_block_tables + ) + target_slot_mappings = speculator.block_tables.slot_mappings[ + :, : input_batch.num_tokens_after_padding + ] + attn_metadata = speculator.model_state.prepare_attn( + input_batch, + cg_mode, + block_tables, + target_slot_mappings, + speculator.attn_groups, + speculator.kv_cache_config, + for_capture=for_capture, + ) + return attn_metadata, build_slot_mappings_by_layer( + target_slot_mappings, speculator.kv_cache_config + ) + + +@contextmanager +def _refresh_draft_prefill_graph( + speculator: TargetDependentARSpeculator, input_batch: InputBatch, dummy_run: bool +) -> Iterator[None]: + """Refresh the draft builders' buffers before a FULL draft-prefill replay. + + Under DBO the draft-prefill graphs are captured on the draft's own + builders (see ``capture``), so every replay first writes this batch's + full-batch metadata there, as #61464 does before ``run_fullgraph``. + """ + manager = speculator.prefill_cudagraph_manager + native_run_fullgraph = manager.run_fullgraph + + def run_fullgraph(desc: BatchExecutionDescriptor) -> Any: + _build_drafter_attn_metadata( + speculator, input_batch, desc.cg_mode, for_capture=dummy_run + ) + return native_run_fullgraph(desc) + + manager.run_fullgraph = run_fullgraph + try: + yield + finally: + del manager.run_fullgraph + + +# Upstream source: vLLM v0.32.0 (5df63a7203cdb9328d824a27b682e77c7b9928be), +# vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py:: +# TargetDependentARSpeculator.capture. +# Patch reason: draft prefill captures its FULL graphs on the target's +# attention builders, whose full-batch buffers a microbatched target forward +# never refreshes (illegal memory access with MTP + graph DBO). +# Patch functionality: under DBO, capture them on the draft's own builders, +# as #61464 does (``can_reuse_target_attn_metadata`` is False with DBO). +# Expansion exception: the upstream method is about 50 lines and passes the +# groups by attribute; the backport swaps that attribute around the call. +# Signature: matches upstream. +def capture(self) -> None: + original: Callable[[Any], None] = getattr( + TargetDependentARSpeculator, _ORIGINAL_CAPTURE_ATTR + ) + # ### PATCH START: capture draft prefill on the draft builders under DBO. + if not self.vllm_config.parallel_config.use_ubatching: + original(self) + return + target_attn_groups = self.target_attn_groups + self.target_attn_groups = self.attn_groups + try: + original(self) + finally: + self.target_attn_groups = target_attn_groups + # ### PATCH END: capture draft prefill on the draft builders under DBO. + + +# Upstream source: vLLM v0.32.0 (5df63a7203cdb9328d824a27b682e77c7b9928be), +# vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py:: +# TargetDependentARSpeculator.propose. +# Patch reason: after a microbatched target forward the runner passes no +# full-batch attention metadata (``None``) to draft prefill. +# Patch functionality: rebuild it on the draft's builders before delegating, +# as #61464 does inside propose. +# Expansion exception: the upstream method is about 130 lines; the backport +# only adds a pre-step, so it delegates to the original instead of copying it. +# Signature: matches upstream; no added parameters. +def propose( + self, + input_batch: InputBatch, + attn_metadata: dict[str, Any], + slot_mappings: dict[str, torch.Tensor], + last_hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + last_sampled: torch.Tensor, + next_prefill_tokens: torch.Tensor, + temperature: torch.Tensor, + seeds: torch.Tensor, + dp_sync_state: DPSyncState | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + is_profile: bool = False, + num_speculative_tokens: int | None = None, +) -> torch.Tensor: + original: Callable[..., torch.Tensor] = getattr( + TargetDependentARSpeculator, _ORIGINAL_PROPOSE_ATTR + ) + # ### PATCH START: draft-prefill metadata on the draft builders under DBO. + draft_prefill_scope: AbstractContextManager[None] = nullcontext() + if self.vllm_config.parallel_config.use_ubatching and not ( + dummy_run and skip_attn_for_dummy_run + ): + if attn_metadata is None: + # Eager draft prefill after a microbatched target forward. + attn_metadata, slot_mappings = _build_drafter_attn_metadata( + self, input_batch + ) + draft_prefill_scope = _refresh_draft_prefill_graph(self, input_batch, dummy_run) + with draft_prefill_scope: + # ### PATCH END: draft-prefill metadata on the draft builders under DBO. + return original( + self, + input_batch, + attn_metadata, + slot_mappings, + last_hidden_states, + aux_hidden_states, + num_sampled, + num_rejected, + last_sampled, + next_prefill_tokens, + temperature, + seeds, + dp_sync_state=dp_sync_state, + dummy_run=dummy_run, + skip_attn_for_dummy_run=skip_attn_for_dummy_run, + mm_inputs=mm_inputs, + is_profile=is_profile, + num_speculative_tokens=num_speculative_tokens, + ) + + +def apply_mrv2_dbo_spec_decode_backport() -> None: + """Install the draft-prefill capture and pre-step once (worker process).""" + if hasattr(TargetDependentARSpeculator, _ORIGINAL_PROPOSE_ATTR): + return + setattr( + TargetDependentARSpeculator, + _ORIGINAL_PROPOSE_ATTR, + TargetDependentARSpeculator.propose, + ) + setattr( + TargetDependentARSpeculator, + _ORIGINAL_CAPTURE_ATTR, + TargetDependentARSpeculator.capture, + ) + TargetDependentARSpeculator.propose = propose + TargetDependentARSpeculator.capture = capture + + +apply_mrv2_dbo_spec_decode_backport() + +__all__ = [ + "apply_mrv2_dbo_spec_decode_backport", + "create_ubatch_slices", + "use_request_aligned_ubatch_split", +] diff --git a/afd_plugin/compat/patches/gpu_mrv2_ubatching.py b/afd_plugin/compat/patches/gpu_mrv2_ubatching.py index c264fa0e7..433eb4ee5 100644 --- a/afd_plugin/compat/patches/gpu_mrv2_ubatching.py +++ b/afd_plugin/compat/patches/gpu_mrv2_ubatching.py @@ -17,15 +17,18 @@ from typing import Any import torch +import torch.distributed as dist from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig +from vllm.distributed.parallel_state import get_dp_group from vllm.forward_context import DPMetadata, ForwardContext, create_forward_context from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.dp_utils import should_skip_dp_coordination from vllm.v1.worker.gpu import model_runner from vllm.v1.worker.gpu.cudagraph_utils import ( BatchExecutionDescriptor, CudaGraphManager, ) -from vllm.v1.worker.gpu.dp_utils import DPSyncState, sync_cudagraph_and_dp_padding +from vllm.v1.worker.gpu.dp_utils import DPSyncState from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.ubatch_utils import UBatchRunner, logger from vllm.v1.worker.ubatch_utils import ( @@ -218,8 +221,10 @@ def dispatch_cg_and_sync_dp( num_active_loras=num_active_loras, num_ubatches=num_ubatches, ) - if num_ubatches * num_tokens < ubatch_desc.num_tokens: - # Match the native DP>1 guard against an empty trailing stage. + if _trailing_ubatch_is_padding( + num_tokens, ubatch_desc.num_tokens, num_ubatches + ): + # Same guard as the DP>1 copy below: no padding-only stage. ubatch_desc = None if ubatch_desc is None: ubatch_desc = BatchExecutionDescriptor( @@ -268,6 +273,208 @@ def dispatch_cg_and_sync_dp( ) +def _trailing_ubatch_is_padding( + min_num_tokens: int, padded_num_tokens: int, num_ubatches: int +) -> bool: + """Whether the last midpoint-split microbatch holds only padding. + + `maybe_create_ubatch_slices` starts the last microbatch at + `(num_ubatches - 1) * (padded // num_ubatches)`. A rank whose real tokens + end at or before that point replays a FULL graph over a microbatch of + zero-length requests, which the sparse MLA FP8 builder rejects. A + request-aligned split never starts the last microbatch later, so the check + also covers it. + """ + last_start = (num_ubatches - 1) * (padded_num_tokens // num_ubatches) + return min_num_tokens <= last_start + + +# Patch reason: the native "empty second microbatch" guard is off by one. It +# falls back to eager only when 2 * min(tokens) < padded tokens, so a rank with +# exactly half the padded tokens replays a FULL microbatched graph whose second +# microbatch is padding only, and FlashMLASparse's FP8 decode builder asserts +# `decode_query_len > 0` (GLM-5.x graph DBO). +# Patch functionality: identical DP agreement and descriptor selection; the +# microbatched graph is rejected whenever the trailing microbatch is padding. +# Signature: matches sync_cudagraph_and_dp_padding exactly at the pinned commit. +# Upstream: vllm/v1/worker/gpu/dp_utils.py, sync_cudagraph_and_dp_padding; +# 5df63a7203cdb9328d824a27b682e77c7b9928be (vLLM 0.32.0). +# Removal plan: remove when the native guard rejects a padding-only microbatch. +def sync_cudagraph_and_dp_padding( + cudagraph_manager: CudaGraphManager | None, + desired_batch_desc: BatchExecutionDescriptor, + num_tokens: int, + num_reqs: int, + uniform_token_count: int | None, + dp_size: int, + dp_rank: int, + max_query_len: int | None = None, + num_active_loras: int = 0, + parallel_config: ParallelConfig | None = None, + allow_ubatching: bool = False, + uniform_decode: bool = False, +) -> tuple[BatchExecutionDescriptor, DPSyncState | None]: + """Coordinates the batch descriptor and DP padding across all ranks. + + `parallel_config` is only needed to decide whether to microbatch, so callers + that never do (`allow_ubatching=False`) can leave it out. + + Returns (synced_batch_desc, sync). `sync` is None when no rank has work. + """ + assert dp_size > 1, "DP size must be greater than 1" + group = get_dp_group().cpu_group + tensor = torch.zeros(6, dp_size, dtype=torch.int32, device="cpu") + tensor[0][dp_rank] = num_tokens + tensor[1][dp_rank] = desired_batch_desc.cg_mode.value + tensor[2][dp_rank] = uniform_token_count or 0 # (0 means None) + tensor[3][dp_rank] = max_query_len or -1 # (-1 means None) + tensor[4][dp_rank] = int(allow_ubatching) + tensor[5][dp_rank] = num_reqs + if should_skip_dp_coordination(): + tensor[:] = tensor[:, dp_rank, None].clone() + else: + dist.all_reduce(tensor, group=group) + + num_tokens_across_dp = tensor[0] + cg_mode_across_dp = tensor[1] + uniform_token_counts_across_dp = tensor[2] + max_query_lens_across_dp = tensor[3] + allow_ubatching_across_dp = tensor[4] + num_reqs_across_dp = tensor[5] + + # If ranks disagree on the uniform token count, or its 0 (means None) set to None + synced_uniform_token_count: int | None = int(uniform_token_counts_across_dp[0]) + if synced_uniform_token_count == 0 or not torch.all( + uniform_token_counts_across_dp == synced_uniform_token_count + ): + synced_uniform_token_count = None + + if torch.all(num_tokens_across_dp == 0).item(): + synced_desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, num_tokens=0, num_reqs=0 + ) + return synced_desc, None + + if torch.all(allow_ubatching_across_dp == 1): + # This rank voted too, so its caller passed a config. + assert parallel_config is not None + # A uniform decode only if every rank runs the same uniform query + # length, so every rank reaches the same answer here. + uniform_decode_across_dp = uniform_decode and torch.all( + uniform_token_counts_across_dp == (uniform_token_count or 0) + ) + if check_ubatch_thresholds( + parallel_config, + # Thresholds only grow with the token count, so holding the smallest + # rank to them holds every rank to them. + int(num_tokens_across_dp.min()), + uniform_decode=uniform_decode_across_dp, + ): + # Expert all-to-all requires every rank to split and pad equally. + # Empty microbatches run as dummy batches. + ubatch_num_tokens = int(num_tokens_across_dp.max()) + num_ubatches = get_num_ubatches(parallel_config) + ubatch_desc = None + if cudagraph_manager is not None: + # Match a FULL microbatched CUDA graph, falling back to an eager + # descriptor if no graph matches. + ubatch_desc = cudagraph_manager.dispatch( + num_reqs, + ubatch_num_tokens, + synced_uniform_token_count, + num_active_loras=num_active_loras, + num_ubatches=num_ubatches, + ) + # ### PATCH START: reject a padding-only trailing microbatch. + if _trailing_ubatch_is_padding( + int(num_tokens_across_dp.min()), + ubatch_desc.num_tokens, + num_ubatches, + ): + # ### PATCH END: reject a padding-only trailing microbatch. + # If one rank has an empty second microbatch, run without + # CUDA graphs. + ubatch_desc = None + if ubatch_desc is None: + ubatch_desc = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=ubatch_num_tokens, + num_reqs=num_reqs, + num_ubatches=num_ubatches, + ) + # Refresh the token count to include CUDA graph padding. + ubatch_num_tokens = ubatch_desc.num_tokens + num_reqs = int(num_reqs_across_dp.max()) + if ubatch_desc.cg_mode == CUDAGraphMode.FULL: + assert ubatch_desc.num_reqs is not None + num_reqs = ubatch_desc.num_reqs + return ubatch_desc, DPSyncState( + num_tokens_across_dp=torch.full_like( + num_tokens_across_dp, ubatch_num_tokens + ), + uniform_token_count=synced_uniform_token_count, + eager=ubatch_desc.cg_mode == CUDAGraphMode.NONE, + num_reqs=num_reqs, + ) + + synced_cg_mode = CUDAGraphMode(int(cg_mode_across_dp.min())) + + # If any rank wants to run eager, all ranks run eager + if synced_cg_mode == CUDAGraphMode.NONE: + return ( + BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=num_tokens, + num_reqs=num_reqs, + num_active_loras=desired_batch_desc.num_active_loras, + ), + DPSyncState( + num_tokens_across_dp=num_tokens_across_dp, + uniform_token_count=synced_uniform_token_count, + eager=True, + num_reqs=int(num_reqs_across_dp.max()), + ), + ) + + assert cudagraph_manager is not None, ( + "cudagraph_manager should only be None during profile run, " + "where synced_cg_mode must be NONE across all DP ranks" + ) + synced_num_tokens = int(num_tokens_across_dp.max()) + + # Varlen decode graphs are selected by the query-length bound, so ranks must agree + # on it or they pad to different token counts below. + synced_max_query_len: int | None = None + if bool(torch.all(max_query_lens_across_dp != -1)): + synced_max_query_len = int(max_query_lens_across_dp.max()) + + # Dispatch for the final synced values, use num_reqs instead of synced_num_reqs + # so we don't perform request padding for PIECEWISE graphs. + # num_active_loras is per-rank and doesn't need cross-rank agreement. + synced_desc = cudagraph_manager.dispatch( + num_reqs, + synced_num_tokens, + synced_uniform_token_count, + num_active_loras=num_active_loras, + max_query_len=synced_max_query_len, + ) + + # Update num_tokens_across_dp to reflect padded size. + num_tokens_across_dp[:] = synced_desc.num_tokens + + return synced_desc, DPSyncState( + num_tokens_across_dp=num_tokens_across_dp, + uniform_token_count=synced_uniform_token_count, + eager=False, + num_reqs=( + synced_desc.num_reqs + if synced_desc.cg_mode == CUDAGraphMode.FULL + and synced_desc.num_reqs is not None + else int(num_reqs_across_dp.max()) + ), + ) + + @contextmanager def use_gpu_mrv2_ubatch_factory() -> Iterator[None]: """Scope the DP1 factory to AFD's native KV initialization call.""" diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index 4eb11dca3..e7511cc40 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -104,6 +104,14 @@ TransportLayout, plan_transport_layout, ) +from afd_plugin.connectors.gpu.deepep_m2n.mxfp4_experts import ( + MARLIN_EXPERT_BACKEND, + MARLIN_TILE_SIZE, + MXFP4_MARLIN_WEIGHTS, + Mxfp4MarlinExpertWeights, + has_mxfp4_experts, + marlin_mxfp4_expected_tensors, +) from afd_plugin.connectors.gpu.deepep_m2n.network import ( NetworkConfig, NetworkPlan, @@ -141,12 +149,13 @@ FP8_PAYLOAD = "fp8_block128_fp32" PAYLOAD_FORMATS = (BF16_PAYLOAD, FP8_PAYLOAD) FP8_GROUP_SIZE = 128 -# Expert weight formats the FFN must load for each payload format. +# Expert weight formats the FFN may load for each payload format. MXFP4 +# experts run as W4A16 Marlin on BF16 rows. BF16_WEIGHTS = "bf16" FP8_BLOCK128_WEIGHTS = "fp8_block128" -PAYLOAD_WEIGHT_FORMATS = { - BF16_PAYLOAD: BF16_WEIGHTS, - FP8_PAYLOAD: FP8_BLOCK128_WEIGHTS, +PAYLOAD_ALLOWED_WEIGHT_FORMATS: Mapping[str, frozenset[str]] = { + BF16_PAYLOAD: frozenset({BF16_WEIGHTS, MXFP4_MARLIN_WEIGHTS}), + FP8_PAYLOAD: frozenset({FP8_BLOCK128_WEIGHTS}), } # vLLM's all2all backend on the FFN. The connector already moved the rows, so # vLLM's own MoE all-to-all never runs; this backend allocates no buffers and @@ -184,7 +193,9 @@ class DeepEPM2NConfig(ConnectorExtraInfo): uses vLLM's micro-batch count: two under DBO, otherwise one (``resolve_num_lanes``). payload_format: Dispatch wire format. - expert_backend: ``deepgemm`` (fast) or ``fused_experts`` (reference). + expert_backend: ``deepgemm`` (fast), ``fused_experts`` (reference) or + ``marlin`` (MXFP4 experts after vLLM's Marlin repack; required for + MXFP4 checkpoints and rejected for others, checked on the FFN). network: Fabric settings (GIN, hybrid mode, traffic class); see ``network.NetworkConfig``. NICs are NCCL's own per-GPU choice. bucket_granularity: Dispatch bucket is rounded up to this multiple, @@ -466,6 +477,28 @@ def expert_weight_format(quant_config: QuantizationConfig | None) -> str: return f"unsupported ({quant_config.get_name()})" +def ffn_expert_weight_format(vllm_config: VllmConfig) -> str: + """Expert weight format of an FFN deployment, MXFP4 included. + + MXFP4 experts are readable only after vLLM's Marlin repack, so they count + as ``mxfp4_marlin_w4a16`` with ``moe_backend=marlin`` and as unsupported + with any other backend. + """ + if has_mxfp4_experts(vllm_config): + moe_backend = str(vllm_config.kernel_config.moe_backend) + if moe_backend == MARLIN_EXPERT_BACKEND: + return MXFP4_MARLIN_WEIGHTS + return f"unsupported (mxfp4 with moe_backend={moe_backend})" + return expert_weight_format(vllm_config.quant_config) + + +def local_weight_format(weights: LocalExpertWeights) -> str: + """Format name of loaded expert weights, as ``ffn_expert_weight_format``.""" + if isinstance(weights, Mxfp4MarlinExpertWeights): + return weights.weight_format + return FP8_BLOCK128_WEIGHTS if weights.is_fp8 else BF16_WEIGHTS + + def ffn_deployment_info(vllm_config: VllmConfig) -> dict[str, Any]: """FFN-side vLLM settings that decide which experts each rank loads.""" parallel_config = vllm_config.parallel_config @@ -477,7 +510,7 @@ def ffn_deployment_info(vllm_config: VllmConfig) -> dict[str, Any]: "expert_placement_strategy": parallel_config.expert_placement_strategy, "eplb": parallel_config.enable_eplb, "all2all_backend": parallel_config.all2all_backend, - "weight_format": expert_weight_format(vllm_config.quant_config), + "weight_format": ffn_expert_weight_format(vllm_config), } @@ -554,20 +587,23 @@ def _check_ffn_deployment( """ if deployment is None: raise RuntimeError(f"FFN world rank {world_rank} sent no deployment info") - expected = { + allowed_formats = PAYLOAD_ALLOWED_WEIGHT_FORMATS[payload_format] + expected: dict[str, Any] = { "expert_parallel": True, "ep_size": layout.ffn_size, "tensor_parallel_size": 1, "expert_placement_strategy": "linear", "eplb": False, "all2all_backend": IDLE_ALL2ALL_BACKEND, - "weight_format": PAYLOAD_WEIGHT_FORMATS[payload_format], } mismatches = { key: deployment[key] for key, value in expected.items() if deployment[key] != value } + if deployment["weight_format"] not in allowed_formats: + mismatches["weight_format"] = deployment["weight_format"] + expected["weight_format"] = f"one of {sorted(allowed_formats)}" if mismatches: required = {key: expected[key] for key in mismatches} raise RuntimeError( @@ -609,6 +645,13 @@ def validate_handshake( _check_ffn_deployment( world_rank, info["ffn_deployment"], layout, info["payload_format"] ) + weight_formats = { + info["ffn_deployment"]["weight_format"] for info in infos[: layout.ffn_size] + } + if len(weight_formats) != 1: + raise RuntimeError( + f"FFN ranks disagree on the expert weight format: {sorted(weight_formats)}" + ) alignments = {info["expert_alignment"] for info in infos[: layout.ffn_size]} if len(alignments) != 1 or None in alignments: raise RuntimeError(f"FFN ranks disagree on expert_alignment: {alignments}") @@ -692,6 +735,7 @@ def local_model_info( experts = { layer_idx: { "ids": list(layer_experts.expert_ids), + "weight_format": local_weight_format(layer_experts.weights), "w13": _describe_tensor(layer_experts.weights.w13), "w2": _describe_tensor(layer_experts.weights.w2), "w13_scale": _describe_tensor(layer_experts.weights.w13_scale), @@ -707,8 +751,13 @@ def _expected_expert_tensors( hidden_size: int, intermediate_size: int, payload_format: str, + weight_format: str, ) -> dict[str, list[Any] | None]: """dtype and shape of every ``LocalExpertWeights`` tensor.""" + if weight_format == MXFP4_MARLIN_WEIGHTS: + return dict( + marlin_mxfp4_expected_tensors(num_experts, hidden_size, intermediate_size) + ) if payload_format == BF16_PAYLOAD: weight_dtype = str(torch.bfloat16) w13_scale = w2_scale = None @@ -775,9 +824,26 @@ def validate_model_infos( f"[{expected_ids[0]}, {expected_ids[-1] + 1}) to it: use " "linear expert placement with EP size = num_ffn_ranks" ) - intermediate_size = described["w2"][1][2] + weight_format = described["weight_format"] + if weight_format not in PAYLOAD_ALLOWED_WEIGHT_FORMATS[payload_format]: + raise RuntimeError( + f"FFN world rank {world_rank} layer {layer_idx} expert weights " + f"{weight_format!r} do not match payload_format=" + f"{payload_format!r}" + ) + # w2 is [E, H, I] in FusedMoE layout and [E, I / 16, 2 H] after + # the Marlin repack. + intermediate_size = ( + described["w2"][1][1] * MARLIN_TILE_SIZE + if weight_format == MXFP4_MARLIN_WEIGHTS + else described["w2"][1][2] + ) expected = _expected_expert_tensors( - layout.experts_per_rank, hidden_size, intermediate_size, payload_format + layout.experts_per_rank, + hidden_size, + intermediate_size, + payload_format, + weight_format, ) wrong = { name: described[name] @@ -961,6 +1027,14 @@ def __init__( self.payload_dtype = FP8_DTYPE if self.is_fp8_payload else self.dtype if self.is_fp8_payload and self.hidden_size % FP8_GROUP_SIZE != 0: raise ValueError(f"FP8 payload needs hidden_size % {FP8_GROUP_SIZE} == 0") + if self.is_ffn and has_mxfp4_experts(vllm_config) != ( + self.config.expert_backend == MARLIN_EXPERT_BACKEND + ): + raise ValueError( + f"expert_backend={self.config.expert_backend!r} does not match the " + "checkpoint's experts: MXFP4 experts need expert_backend " + f"{MARLIN_EXPERT_BACKEND!r}, and only they can use it" + ) self.world_rank = self.layout.world_rank(afd_config.role, role_rank) # The FFN runs without TP and learns the Attention TP size and the # resulting capacity in the handshake (``_handshake``). diff --git a/afd_plugin/connectors/gpu/deepep_m2n/experts.py b/afd_plugin/connectors/gpu/deepep_m2n/experts.py index 2b5fe0042..cf5b94308 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/experts.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/experts.py @@ -26,6 +26,9 @@ - ``fused_experts``: vLLM's Triton ``fused_experts`` with ``top_k=1`` on the valid rows (FP8 inputs and weights are dequantized to BF16 first). Reads the prefix sum on the host: a correctness reference, not a fast path. +- ``marlin``: BF16 rows with MXFP4 weights after vLLM's Marlin repack + (``mxfp4_experts.py``), with the checkpoint's SwiGLU clamp. Requires + ``expert_alignment`` 1; no host reads. """ from __future__ import annotations @@ -42,10 +45,17 @@ dequantize_fp8_rows, silu_mul_fp8_quant, ) +from afd_plugin.connectors.gpu.deepep_m2n.mxfp4_experts import ( + MARLIN_EXPERT_BACKEND, + Mxfp4MarlinExpertWeights, + marlin_workspace, + run_marlin_experts, + sentinel_expert_map, +) from afd_plugin.connectors.metadata import FP8_DTYPE, LocalExpertWeights -ExpertBackend = Literal["deepgemm", "fused_experts"] -EXPERT_BACKENDS: tuple[ExpertBackend, ...] = ("deepgemm", "fused_experts") +ExpertBackend = Literal["deepgemm", "fused_experts", "marlin"] +EXPERT_BACKENDS: tuple[ExpertBackend, ...] = ("deepgemm", "fused_experts", "marlin") def align_up(value: int, alignment: int) -> int: @@ -75,6 +85,16 @@ def __init__( raise ValueError(f"expert backend must be one of {EXPERT_BACKENDS}") self.backend = backend self.expert_alignment = expert_alignment + self._marlin_workspace: torch.Tensor | None = None + # Sentinel expert maps by local expert count (one per deployment). + self._marlin_expert_maps: dict[int, torch.Tensor] = {} + if backend == MARLIN_EXPERT_BACKEND: + if expert_alignment != 1: + raise ValueError("the marlin expert backend needs expert_alignment 1") + # Allocated here, eagerly, so no CUDA-graph capture allocates it. + self._marlin_workspace = marlin_workspace( + torch.device("cuda", torch.cuda.current_device()) + ) if backend == "deepgemm": deepgemm.set_contiguous_m_alignment(expert_alignment) if num_sms: @@ -117,6 +137,13 @@ def __call__( f"payload dtype {recv_x.dtype} does not match the expert weights " f"({weights.w13.dtype}) or the presence of scales" ) + is_mxfp4_marlin = isinstance(weights, Mxfp4MarlinExpertWeights) + if is_mxfp4_marlin != (self.backend == MARLIN_EXPERT_BACKEND): + raise ValueError( + f"expert backend {self.backend!r} cannot run " + f"{type(weights).__name__}; MXFP4 Marlin weights need the " + f"{MARLIN_EXPERT_BACKEND!r} backend and only it" + ) if out is None: out = torch.empty(recv_x.shape, dtype=torch.bfloat16, device=recv_x.device) elif out.shape != recv_x.shape or out.dtype != torch.bfloat16: @@ -126,6 +153,10 @@ def __call__( ) if recv_x.shape[0] == 0: return out + if isinstance(weights, Mxfp4MarlinExpertWeights): + return self._run_marlin( + recv_x, recv_topk_weights, recv_count_psum, weights, out + ) if self.backend == "fused_experts": return self._run_fused_experts( recv_x, recv_scale, recv_topk_weights, recv_count_psum, weights, out @@ -138,6 +169,30 @@ def __call__( recv_x, recv_topk_weights, recv_count_psum, weights, out ) + def _run_marlin( + self, + recv_x: torch.Tensor, + recv_topk_weights: torch.Tensor, + recv_count_psum: torch.Tensor, + weights: Mxfp4MarlinExpertWeights, + out: torch.Tensor, + ) -> torch.Tensor: + assert self._marlin_workspace is not None + num_experts = weights.num_experts + expert_map = self._marlin_expert_maps.get(num_experts) + if expert_map is None: + expert_map = sentinel_expert_map(num_experts, recv_x.device) + self._marlin_expert_maps[num_experts] = expert_map + return run_marlin_experts( + recv_x, + recv_topk_weights, + recv_count_psum, + weights, + out, + workspace=self._marlin_workspace, + expert_map=expert_map, + ) + def _run_deepgemm_bf16( self, recv_x: torch.Tensor, diff --git a/afd_plugin/connectors/gpu/deepep_m2n/mxfp4_experts.py b/afd_plugin/connectors/gpu/deepep_m2n/mxfp4_experts.py new file mode 100644 index 000000000..1d15c420d --- /dev/null +++ b/afd_plugin/connectors/gpu/deepep_m2n/mxfp4_experts.py @@ -0,0 +1,315 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""MXFP4 routed experts on DeepEP-received BF16 rows, with vLLM's Marlin. + +The FFN loads MXFP4 experts (for example GLM-5.3 MXFP4xFP8, whose +compressed-tensors W4A8 experts use a ``Mxfp4MoEMethod`` subclass) through +vLLM's ``Mxfp4MoEMethod`` with ``--moe-backend marlin``, which repacks them for the +W4A16 Marlin MoE kernel at load. ``mxfp4_marlin_weights_from_vllm`` reads +those parameters, plus the activation config (SwiGLU clamp) and quant type of +the ``MarlinExperts`` object vLLM built for the layer, so the executor runs +exactly what native vLLM runs. + +``run_marlin_experts`` treats each received row as one token with +``top_k = 1``: DeepEP packs each local expert's rows contiguously +(``expert_alignment = 1``), so a row's expert is its position in the +per-expert prefix sum, computed on the device. Rows past the last expert's +end (unused tail of the received buffer) get a sentinel expert that +``expert_map`` sends to -1; Marlin skips them and ``moe_sum`` writes zeros. +The top-k weight multiplies the second GEMM's output inside Marlin. The +routed scale is already part of that weight (applied on Attention). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import torch +from torch import nn +from vllm.config import set_current_vllm_config +from vllm.model_executor.layers.fused_moe.activation import ( + ApplyMoEActivationConfig, + MoEActivation, +) +from vllm.model_executor.layers.fused_moe.experts.marlin_moe import ( + MarlinExpertsBase, + fused_marlin_moe, +) +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 + CompressedTensorsConfig, +) +from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4Config, Mxfp4MoEMethod +from vllm.model_executor.layers.quantization.utils.marlin_utils import ( + MARLIN_MAX_BLOCKS_PER_SM, + marlin_make_workspace_new, +) + +from afd_plugin.connectors.metadata import FFNLocalExperts, LocalExpertWeights + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + +MARLIN_EXPERT_BACKEND = "marlin" +MXFP4_MARLIN_WEIGHTS = "mxfp4_marlin_w4a16" +# vLLM's DeepSeek-V4/V4.1 quantization method; its experts are MXFP4 when the +# checkpoint's expert_dtype is "fp4". +DEEPSEEK_V4_FP8_QUANTIZATION = "deepseek_v4_fp8" +MXFP4_EXPERT_DTYPE = "fp4" +# Checkpoint projection names of one routed expert (DeepSeek and GLM layout). +CHECKPOINT_EXPERT_PROJECTIONS = ("gate_proj", "up_proj", "down_proj") +MXFP4_GROUP_SIZE = 32 +# Marlin packs 4-bit weights into int32 words over 16-row tiles. +MARLIN_TILE_SIZE = 16 +MARLIN_4BIT_PACK_FACTOR = 2 +_MARLIN_MXFP4_BACKENDS = frozenset( + {Mxfp4MoeBackend.MARLIN, Mxfp4MoeBackend.BATCHED_MARLIN} +) + + +def _compressed_tensors_experts_are_w4a8_mxfp4( + quant_config: CompressedTensorsConfig, vllm_config: VllmConfig +) -> bool: + """Whether a compressed-tensors checkpoint quantizes its routed experts + as MXFP4 weights with dynamic FP8 activations. + + vLLM resolves the MoE scheme from the unfused checkpoint names of one + expert (``CompressedTensorsMoEMethod.get_moe_method``) and loads W4A8 + MXFP4 experts with ``CompressedTensorsW4A8Mxfp4MoEMethod``, a subclass of + ``Mxfp4MoEMethod``. Before any layer exists this resolves the same names + against a ``Linear``, the module type the checkpoint's projections had. + """ + # The first routed-expert layer: a layer-index regex target (for example + # an MTP-only group) must not decide the decoder experts' scheme. + first_moe_layer = getattr( + vllm_config.model_config.hf_text_config, "first_k_dense_replace", 0 + ) + probe = nn.Linear(1, 1, bias=False, device="meta") + schemes = [ + quant_config.get_scheme_dict( + probe, f"model.layers.{first_moe_layer}.mlp.experts.0.{projection}" + ) + for projection in CHECKPOINT_EXPERT_PROJECTIONS + ] + return all( + scheme is not None + # The exact predicate vLLM uses to pick the W4A8 MXFP4 MoE method. + and quant_config._is_mxfp4_w4a8_fp8( + scheme["weights"], scheme["input_activations"] + ) + for scheme in schemes + ) + + +def has_mxfp4_experts(vllm_config: VllmConfig) -> bool: + """Whether vLLM will load this checkpoint's routed experts as MXFP4. + + Decided from the quantization config before any MoE layer exists. The + DeepSeek-V4 config resolves ``expert_dtype`` from the current vLLM + config, so it is read under this one. A compressed-tensors checkpoint + counts when its experts are W4A8 MXFP4 (for example GLM-5.3 MXFP4xFP8), + which vLLM loads through ``Mxfp4MoEMethod``. + """ + quant_config = vllm_config.quant_config + if isinstance(quant_config, Mxfp4Config): + return True + if isinstance(quant_config, CompressedTensorsConfig): + return _compressed_tensors_experts_are_w4a8_mxfp4(quant_config, vllm_config) + if quant_config is None or quant_config.get_name() != DEEPSEEK_V4_FP8_QUANTIZATION: + return False + with set_current_vllm_config(vllm_config): + return quant_config.expert_dtype == MXFP4_EXPERT_DTYPE # type: ignore[attr-defined] + + +@dataclass(frozen=True, kw_only=True) +class Mxfp4MarlinExpertWeights(LocalExpertWeights): + """This rank's MXFP4 experts as vLLM's Marlin MXFP4 method left them. + + Attributes: + w13: ``[E_local, H / 16, 2 * I * 2]`` int32 Marlin-packed gate and up. + w2: ``[E_local, I / 16, H * 2]`` int32 Marlin-packed down projection. + w13_scale: ``[E_local, H / 32, 2 * I]`` e8m0 Marlin-permuted scales. + w2_scale: ``[E_local, I / 32, H]`` e8m0 Marlin-permuted scales. + quant_type_id: Marlin weight scalar type (``float4_e2m1f``). + activation_config: The SwiGLU clamp vLLM's ``MarlinExperts`` applies. + input_dtype: Marlin activation quantization (``None`` = BF16, W4A16). + """ + + w13_scale: torch.Tensor + w2_scale: torch.Tensor + quant_type_id: int + activation_config: ApplyMoEActivationConfig + input_dtype: torch.dtype | None = None + weight_format: str = field(default=MXFP4_MARLIN_WEIGHTS) + + def __post_init__(self) -> None: + # Replaces the BF16/FP8 rule of LocalExpertWeights: Marlin MXFP4 has + # packed int32 weights and e8m0 scales. + tensors = (self.w13, self.w2, self.w13_scale, self.w2_scale) + if len({tensor.device for tensor in tensors}) != 1: + raise ValueError("MXFP4 Marlin expert tensors must share one device") + if len({int(tensor.shape[0]) for tensor in tensors}) != 1: + raise ValueError("MXFP4 Marlin expert tensors disagree on the expert count") + + @property + def is_fp8(self) -> bool: + return False + + +def marlin_mxfp4_expected_tensors( + num_experts: int, hidden_size: int, intermediate_size: int +) -> dict[str, list[object]]: + """dtype and shape of every ``Mxfp4MarlinExpertWeights`` tensor. + + Mirrors ``prepare_moe_mxfp4_layer_for_marlin`` (vLLM v0.32.0) for BF16 + activations: ``gptq_marlin_repack`` gives ``[K / 16, N * 2]`` int32 per + expert, and the permuted scales keep ``[K / 32, N]``. + """ + return { + "w13": [ + str(torch.int32), + [ + num_experts, + hidden_size // MARLIN_TILE_SIZE, + 2 * intermediate_size * MARLIN_4BIT_PACK_FACTOR, + ], + ], + "w2": [ + str(torch.int32), + [ + num_experts, + intermediate_size // MARLIN_TILE_SIZE, + hidden_size * MARLIN_4BIT_PACK_FACTOR, + ], + ], + "w13_scale": [ + str(torch.float8_e8m0fnu), + [num_experts, hidden_size // MXFP4_GROUP_SIZE, 2 * intermediate_size], + ], + "w2_scale": [ + str(torch.float8_e8m0fnu), + [num_experts, intermediate_size // MXFP4_GROUP_SIZE, hidden_size], + ], + } + + +def mxfp4_marlin_weights_from_vllm( + routed_experts: RoutedExperts, +) -> FFNLocalExperts: + """This rank's experts of one MXFP4 vLLM MoE layer, after the Marlin repack. + + Reads the parameters ``prepare_moe_mxfp4_layer_for_marlin`` left on the + layer and the ``MarlinExperts`` object vLLM built for it (scales, clamp, + quant type, activation dtype), so the executor matches native. + """ + method = routed_experts.quant_method + if not isinstance(method, Mxfp4MoEMethod): + raise RuntimeError( + f"expected vLLM's MXFP4 MoE method, got {type(method).__name__}" + ) + if method.mxfp4_backend not in _MARLIN_MXFP4_BACKENDS: + raise RuntimeError( + f"vLLM processed the MXFP4 experts for {method.mxfp4_backend.name}, " + "whose layout the M2N executor cannot read; set --moe-backend marlin" + ) + assert method.moe_kernel is not None + experts = method.moe_kernel.fused_experts + if not isinstance(experts, MarlinExpertsBase): + raise RuntimeError( + f"expected vLLM's Marlin MoE experts, got {type(experts).__name__}" + ) + if experts.w1_bias is not None or experts.w2_bias is not None: + raise RuntimeError("MXFP4 experts with biases are not supported on M2N") + assert experts.w1_scale is not None and experts.w2_scale is not None + return FFNLocalExperts( + expert_ids=tuple(routed_experts.expert_map_manager.get_local_expert_ids()), + weights=Mxfp4MarlinExpertWeights( + w13=routed_experts.w13_weight.data, + w2=routed_experts.w2_weight.data, + w13_scale=experts.w1_scale, + w2_scale=experts.w2_scale, + quant_type_id=experts.quant_type_id, + activation_config=experts.activation_config, + input_dtype=experts.input_dtype, + ), + ) + + +def marlin_workspace(device: torch.device) -> torch.Tensor: + """Marlin's lock buffer, allocated once before any CUDA-graph capture. + + The kernel resets its locks after each call, so calls ordered on one + stream (every expert call runs on the caller's stream) can share it. + """ + return marlin_make_workspace_new(device, MARLIN_MAX_BLOCKS_PER_SM) + + +def sentinel_expert_map(num_local_experts: int, device: torch.device) -> torch.Tensor: + """``[0, 1, ..., E_local - 1, -1]``: the extra id ``E_local`` maps nowhere. + + Built with device ops only, so it is safe inside a CUDA-graph capture. + """ + expert_map = torch.arange(num_local_experts + 1, dtype=torch.int32, device=device) + expert_map[num_local_experts:].fill_(-1) + return expert_map + + +def run_marlin_experts( + recv_x: torch.Tensor, + recv_topk_weights: torch.Tensor, + recv_count_psum: torch.Tensor, + weights: Mxfp4MarlinExpertWeights, + out: torch.Tensor, + *, + workspace: torch.Tensor, + expert_map: torch.Tensor, +) -> torch.Tensor: + """``weight * expert(row)`` for every received row, without host reads. + + Requires ``expert_alignment = 1`` (contiguous expert groups); rows past + the last group's end come out as zeros. + """ + num_rows = recv_x.shape[0] + num_local_experts = weights.num_experts + rows = torch.arange(num_rows, dtype=recv_count_psum.dtype, device=recv_x.device) + # Row r belongs to the first expert whose group ends after r. + row_experts = torch.searchsorted(recv_count_psum, rows, right=True) + topk_ids = torch.where( + rows < recv_count_psum[-1], row_experts, num_local_experts + ).to(torch.int32) + return fused_marlin_moe( + hidden_states=recv_x, + w1=weights.w13, + w2=weights.w2, + bias1=None, + bias2=None, + w1_scale=weights.w13_scale, + w2_scale=weights.w2_scale, + topk_weights=recv_topk_weights[:num_rows].float().view(num_rows, 1), + topk_ids=topk_ids.view(num_rows, 1), + quant_type_id=weights.quant_type_id, + # Weight the second GEMM's output, as native top-k reduction does. + apply_router_weight_on_input=False, + global_num_experts=num_local_experts + 1, + activation=MoEActivation.SILU, + expert_map=expert_map, + workspace=workspace, + output=out, + input_dtype=weights.input_dtype, + activation_config=weights.activation_config, + ) + + +__all__ = [ + "MARLIN_EXPERT_BACKEND", + "MXFP4_MARLIN_WEIGHTS", + "Mxfp4MarlinExpertWeights", + "has_mxfp4_experts", + "marlin_mxfp4_expected_tensors", + "marlin_workspace", + "mxfp4_marlin_weights_from_vllm", + "run_marlin_experts", + "sentinel_expert_map", +] diff --git a/afd_plugin/model_executor/models/deepseek_v32.py b/afd_plugin/model_executor/models/deepseek_v32.py new file mode 100644 index 000000000..b5bb40f3e --- /dev/null +++ b/afd_plugin/model_executor/models/deepseek_v32.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CUDA AFD wrapper for vLLM's NVIDIA DSA model (GLM-5.x). + +vLLM v0.32.0 serves ``GlmMoeDsaForCausalLM`` on CUDA with +``vllm.models.deepseek_v32.nvidia.model.DeepseekV32ForCausalLM``: DeepSeek +V2/V3 orchestration with a DSA backbone (MLA + lightning indexer), fused +embedding/RMSNorm, and an MLP output that stays un-reduced until the next +layer's fused all-reduce + RMSNorm. This wrapper keeps that native forward +and loader and only makes construction and checkpoint loading role-aware: + +- Attention builds the DSA attention, the dense MLPs of the first + ``first_k_dense_replace`` layers, and, per MoE layer, the routed-experts + connector shell (gate, native router, shared experts, dispatch/combine). +- FFN builds only each MoE layer's routed experts (native ``DeepseekV2MoE`` + without shared experts); the connector runs them on the dispatched rows. + +The wrapper is limited to the DeepEP M2N connector with Attention TP = 1. +MTP runs natively on the Attention ranks (its draft experts are local there); +DBO (eager, decode steps only) gives each micro-batch its own DSA top-k +state (``gpu_deepseek_v32_dbo``) and runs together with MTP through the +vLLM #61464 backport (``gpu_mrv2_dbo_spec_decode``). Other deployments fail +at construction. +""" + +from collections.abc import Iterable +from copy import copy +from typing import TYPE_CHECKING + +import torch +from vllm.config import VllmConfig +from vllm.models.deepseek_v32.nvidia import model as native +from vllm.platforms import current_platform + +from afd_plugin.compat.patches.gpu_deepseek_v32_dbo import ( + install_per_ubatch_topk_indices, +) +from afd_plugin.config import DEEPEP_M2N_CONNECTORS, AFDConfig, parse_afd_config +from afd_plugin.model_executor.models.deepseek_v2 import ( + AFDDeepseekV2RoutedExpertsMoE, + _iter_role_weights, +) + +if TYPE_CHECKING: + from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts + +_ATTENTION_ROLE = "attention" +_FFN_ROLE = "ffn" +MTP_METHOD = "mtp" + + +def require_supported_deployment( + vllm_config: VllmConfig, + afd_config: AFDConfig, +) -> None: + """Fail closed outside the GLM-5.x CUDA contract. + + MTP is accepted on the Attention role only. DBO is accepted with + decode-only micro-batching, eager or with CUDA graphs, with or without MTP. + Attention TP > 1 is rejected because the native decoder layer leaves the + MLP output un-reduced and all-reduces it in the next layer's RMSNorm, + while the connector returns complete rows on every TP rank; the sum would + count the FFN output TP times. + """ + if afd_config.connector not in DEEPEP_M2N_CONNECTORS: + raise RuntimeError( + f"AFD GLM-5.x on CUDA supports only {DEEPEP_M2N_CONNECTORS}, " + f"got {afd_config.connector!r}" + ) + if not afd_config.compute_gate_on_attention: + raise RuntimeError("AFD GLM-5.x requires compute_gate_on_attention=true") + parallel_config = vllm_config.parallel_config + if parallel_config.tensor_parallel_size != 1: + raise RuntimeError( + "AFD GLM-5.x does not support tensor_parallel_size > 1: the native " + "DSA layer all-reduces the MLP output in the next layer's RMSNorm" + ) + if ( + parallel_config.pipeline_parallel_size != 1 + or parallel_config.use_sequence_parallel_moe + ): + raise RuntimeError("AFD GLM-5.x does not support PP or sequence-parallel MoE") + if ( + parallel_config.use_ubatching + and parallel_config.dbo_prefill_token_threshold + <= vllm_config.scheduler_config.max_num_batched_tokens + ): + # Micro-batched prefill still corrupts a fraction of requests on + # GLM-5.3 (GSM8K-200 0.90 vs 0.93) after the per-micro-batch top-k fix; + # micro-batched decode matches the non-DBO result. + raise RuntimeError( + "AFD GLM-5.x DBO micro-batches decode steps only: set " + "--dbo-prefill-token-threshold above --max-num-batched-tokens" + ) + speculative_config = vllm_config.speculative_config + if speculative_config is None: + return + # The MTP draft runs natively on the Attention ranks, with its own experts; + # the FFN serves the target model's routed experts only. + if afd_config.role != _ATTENTION_ROLE: + raise RuntimeError( + "AFD GLM-5.x takes the MTP speculative config on the Attention role " + "only; the FFN role serves no draft model" + ) + if speculative_config.method != MTP_METHOD: + raise RuntimeError( + f"AFD GLM-5.x supports only {MTP_METHOD!r} speculative decoding, " + f"got {speculative_config.method!r}" + ) + + +class AFDDeepseekV32DecoderLayer(native.DeepseekV32DecoderLayer): + """DSA decoder layer that constructs only the active role's stages.""" + + # Patch reason: the native layer constructs both the DSA attention and the + # dense MLP or MoE. + # Patch functionality: construct only the stages owned by the active AFD + # role; the native forward is kept. + # Signature: matches upstream; no added parameters. + # Upstream: vLLM v0.32.0, vllm/models/deepseek_v32/nvidia/model.py + # Commit: 5df63a7203cdb9328d824a27b682e77c7b9928be + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + config=None, + topk_indices_buffer: torch.Tensor | None = None, + index_group_builder: native.SparseMLAIndexGroupBuilder | None = None, + ) -> None: + # ### PATCH START: require an explicit AFD role before allocation. + afd_role = parse_afd_config(vllm_config, validate=False).role + torch.nn.Module.__init__(self) + # ### PATCH END + + if config is None: + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + + self.hidden_size = config.hidden_size + moe_layer_freq = getattr(config, "moe_layer_freq", 1) + layer_idx = int(prefix.split(sep=".")[-1]) + self.layer_idx = layer_idx + self.use_mha = False + self.use_sequence_parallel = ( + parallel_config.use_sequence_parallel_moe + and parallel_config.pipeline_parallel_size == 1 + ) + + # ### PATCH START: construct only the stages owned by this AFD role. + self.is_moe_layer = ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % moe_layer_freq == 0 + ) + if afd_role == _ATTENTION_ROLE: + self.self_attn = native.DeepseekV32Attention( + vllm_config=vllm_config, + config=config, + prefix=f"{prefix}.self_attn", + topk_indices_buffer=topk_indices_buffer, + index_group_builder=index_group_builder, + ) + if self.is_moe_layer: + # Gate, native router and shared experts; the routed experts + # run on the FFN role behind the connector. + self.mlp = AFDDeepseekV2RoutedExpertsMoE( + config=config, + vllm_config=vllm_config, + layer_idx=layer_idx, + prefix=f"{prefix}.mlp", + ) + else: + self.mlp = native.DeepseekV2MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + reduce_results=False, + is_sequence_parallel=self.use_sequence_parallel, + ) + elif afd_role == _FFN_ROLE: + self.self_attn = native.PPMissingLayer() + if self.is_moe_layer: + # Attention runs the shared experts. HF validates this public + # field as int, but native MoE omits shared construction only + # for None; apply it to a private copy of the config. + moe_config = copy(config) + object.__setattr__(moe_config, "n_shared_experts", None) + self.mlp = native.DeepseekV2MoE( + config=moe_config, + parallel_config=parallel_config, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.mlp", + apply_routed_scale_to_output=False, + ) + # Attention routes; the gate weights stay loadable but the + # runner never computes its own routing (gate is None means + # external routing in the v0.32 runner contract). + self.mlp.experts.gate = None + else: + self.mlp = native.PPMissingLayer() + else: + raise RuntimeError(f"Unknown AFD role {afd_role!r}") + # ### PATCH END + + self.input_layernorm = native.RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.post_attention_layernorm = native.RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + + +class AFDDeepseekV32Model(native.DeepseekV32Model): + """DSA model whose layers and indexer buffer follow the AFD role.""" + + fall_back_to_pt_during_load = False + + # Patch reason: the native model always constructs native decoder layers + # and the indexer top-k buffer. + # Patch functionality: construct role-aware AFD layers and allocate the + # indexer buffer only on Attention, which owns the DSA attention. + # Signature: matches upstream; no added parameters. + # Upstream: vLLM v0.32.0, vllm/models/deepseek_v32/nvidia/model.py + # Commit: 5df63a7203cdb9328d824a27b682e77c7b9928be + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + torch.nn.Module.__init__(self) + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + + self.device = current_platform.device_type + + self.vocab_size = config.vocab_size + parallel_config = vllm_config.parallel_config + self.use_sequence_parallel = ( + parallel_config.use_sequence_parallel_moe + and parallel_config.pipeline_parallel_size == 1 + ) + self.is_v32 = True + # ### PATCH START: allocate the indexer buffer only on Attention. + afd_role = parse_afd_config(vllm_config, validate=False).role + if afd_role == _ATTENTION_ROLE: + # On the model, not a local: the MTP proposer shares it with the + # draft. + self.topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=self.device, + ) + index_group_builder = native.SparseMLAIndexGroupBuilder( + self.topk_indices_buffer, + native.get_sparse_mla_index_group_max_rows(vllm_config), + ) + else: + self.topk_indices_buffer = None + index_group_builder = None + # ### PATCH END + + if native.get_pp_group().is_first_rank: + self.embed_tokens = native.make_input_embedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + tie_word_embeddings=getattr(config, "tie_word_embeddings", False), + ) + else: + self.embed_tokens = native.PPMissingLayer() + # The fused embed+norm gather needs the full table on-rank. + self.replicated_embed = native.has_full_vocab_on_rank(self.embed_tokens) + + # ### PATCH START: use the role-aware decoder layer. + self.start_layer, self.end_layer, self.layers = native.make_layers( + config.num_hidden_layers, + lambda prefix: AFDDeepseekV32DecoderLayer( + vllm_config=vllm_config, + prefix=prefix, + topk_indices_buffer=self.topk_indices_buffer, + index_group_builder=index_group_builder, + ), + prefix=f"{prefix}.layers", + ) + # ### PATCH END + + # ### PATCH START: per-micro-batch DSA top-k state under DBO. + if afd_role == _ATTENTION_ROLE and parallel_config.use_ubatching: + install_per_ubatch_topk_indices(self, parallel_config.num_ubatches) + # ### PATCH END + + if native.get_pp_group().is_last_rank: + self.norm = native.RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + else: + self.norm = native.PPMissingLayer() + self.make_empty_intermediate_tensors = ( + native.make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + ) + + self.aux_hidden_state_layers = tuple[int, ...]() + self.num_redundant_experts = parallel_config.eplb_config.num_redundant_experts + + def get_experts_layer_indices(self) -> tuple[int, ...]: + """Decoder layers whose routed experts run behind the connector.""" + return tuple( + layer_idx + for layer_idx, layer in enumerate(self.layers) + if isinstance(layer, AFDDeepseekV32DecoderLayer) and layer.is_moe_layer + ) + + def get_routed_experts(self, layer_idx: int) -> "RoutedExperts": + """FFN only: the vLLM module holding one MoE layer's local experts. + + Attention folds ``routed_scaling_factor`` into the top-k weights it + dispatches; the FFN never applies it. + """ + return self.layers[layer_idx].mlp.experts.routed_experts + + +class AFDGlmMoeDsaForCausalLM(native.DeepseekV32ForCausalLM): + """GLM-5.x (``glm_moe_dsa``) causal LM over AFD on CUDA.""" + + model_cls = AFDDeepseekV32Model + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + self.afd_config = parse_afd_config(vllm_config, validate=False) + require_supported_deployment(vllm_config, self.afd_config) + super().__init__(vllm_config=vllm_config, prefix=prefix) + + def get_experts_layer_indices(self) -> tuple[int, ...]: + return self.model.get_experts_layer_indices() + + def get_routed_experts(self, layer_idx: int) -> "RoutedExperts": + return self.model.get_routed_experts(layer_idx) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + # The routed-experts connector runs the shared experts and the dense + # MLPs on Attention; FFN keeps only the routed experts and the gate. + return super().load_weights( + _iter_role_weights( + weights, + role=self.afd_config.role, + config=self.config, + compute_gate_on_attention=True, + attention_shared_experts=True, + ) + ) + + +__all__ = [ + "AFDDeepseekV32DecoderLayer", + "AFDDeepseekV32Model", + "AFDGlmMoeDsaForCausalLM", + "require_supported_deployment", +] diff --git a/afd_plugin/v1/worker/attention_model_runner_v2.py b/afd_plugin/v1/worker/attention_model_runner_v2.py index 3ff495685..c59f7dddf 100644 --- a/afd_plugin/v1/worker/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/attention_model_runner_v2.py @@ -126,7 +126,28 @@ def _use_afd_ubatch_preparation( ) model_preparation = _use_v4_ubatch_preparation(runner) - with model_preparation: + elif ( + runner.vllm_config.model_config.hf_config.model_type == "glm_moe_dsa" + and runner.device.type == "cuda" + ): + from afd_plugin.compat.patches.gpu_deepseek_v32_dbo import ( + use_dsa_indexer_sm_budget, + ) + + model_preparation = use_dsa_indexer_sm_budget(ubatch_runner) + spec_decode_split: AbstractContextManager[None] = nullcontext() + if ( + runner.vllm_config.speculative_config is not None + and runner.device.type == "cuda" + ): + # Backport of vLLM #61464: verification batches split on request + # boundaries, and draft prefill rebuilds its full-batch metadata. + from afd_plugin.compat.patches.gpu_mrv2_dbo_spec_decode import ( + use_request_aligned_ubatch_split, + ) + + spec_decode_split = use_request_aligned_ubatch_split() + with model_preparation, spec_decode_split: state = vars(ubatch_runner) had_override = "prepare" in state previous_override = state.get("prepare") diff --git a/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py index 55c26f41d..dac8c3a55 100644 --- a/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py +++ b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py @@ -9,6 +9,7 @@ import torch from vllm.config import VllmConfig from vllm.forward_context import DPMetadata +from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod from afd_plugin.connectors import ( AFDConnectorBase, @@ -18,6 +19,11 @@ FFNLocalExperts, LocalExpertWeights, ) +from afd_plugin.connectors.gpu.deepep_m2n.mxfp4_experts import ( + MARLIN_EXPERT_BACKEND, + has_mxfp4_experts, + mxfp4_marlin_weights_from_vllm, +) from afd_plugin.connectors.metadata import FP8_DTYPE from afd_plugin.v1.worker.ffn_model_runner import GPUFFNModelRunner @@ -31,26 +37,36 @@ # swaps w13 to [up; gate]; TRT-LLM and Marlin also shuffle), which yields wrong # outputs with no error. ROUTED_MOE_BACKEND = "triton" +# MXFP4 experts: the connector's ``marlin`` executor reads vLLM's Marlin +# repack, so these checkpoints pin Marlin instead. +MXFP4_ROUTED_MOE_BACKEND = MARLIN_EXPERT_BACKEND _LAYOUT_PRESERVING_BACKENDS = frozenset({"TRITON"}) def select_routed_moe_backend(vllm_config: VllmConfig) -> None: - """Pin the FFN's vLLM MoE backend to one that keeps the expert layout. + """Pin the FFN's vLLM MoE backend to one whose layout the executor reads. vLLM's MoE forward never runs on a routed-experts FFN; its backend only - decides how the expert weights are processed at load. ``auto`` picks - FlashInfer CUTLASS on Hopper with EP > 1 for block FP8, which reorders - w13, so ``auto`` becomes ``triton`` and any other explicit choice fails. + decides how the expert weights are processed at load. For BF16 and block + FP8, ``auto`` picks FlashInfer CUTLASS on Hopper with EP > 1, which + reorders w13, so ``auto`` becomes ``triton``. For MXFP4 experts ``auto`` + may pick Humming, so it becomes ``marlin``. Any other explicit choice + fails. """ kernel_config = vllm_config.kernel_config + required = ( + MXFP4_ROUTED_MOE_BACKEND + if has_mxfp4_experts(vllm_config) + else ROUTED_MOE_BACKEND + ) backend = str(kernel_config.moe_backend) if backend == "auto": - kernel_config.moe_backend = ROUTED_MOE_BACKEND - elif backend != ROUTED_MOE_BACKEND: + kernel_config.moe_backend = required + elif backend != required: raise RuntimeError( - f"routed-experts AFD FFN needs --moe-backend {ROUTED_MOE_BACKEND} (or " - f"auto); moe_backend={backend!r} reorders the expert weights the " - "connector reads" + f"routed-experts AFD FFN needs --moe-backend {required} (or auto) " + f"for this checkpoint's experts; moe_backend={backend!r} lays out " + "the expert weights differently from what the connector reads" ) @@ -75,8 +91,11 @@ def local_experts_from_vllm(routed_experts: RoutedExperts) -> FFNLocalExperts: Read after ``process_weights_after_loading``, which must have kept the FusedMoE layout (``select_routed_moe_backend``); checked here because a - reordered layout keeps dtypes and shapes. + reordered layout keeps dtypes and shapes. MXFP4 experts come back as + vLLM's Marlin repack (``Mxfp4MarlinExpertWeights``). """ + if isinstance(routed_experts.quant_method, Mxfp4MoEMethod): + return mxfp4_marlin_weights_from_vllm(routed_experts) _check_expert_layout(routed_experts) w13 = routed_experts.w13_weight.data w2 = routed_experts.w2_weight.data diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index b2ef9b0eb..936f21a86 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -53,6 +53,7 @@ ATTENTION_ROUTING_MODEL_ARCHITECTURES: Final[tuple[str, ...]] = ( "Qwen3MoeForCausalLM", "DeepseekV2ForCausalLM", + "GlmMoeDsaForCausalLM", ) diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 7af50ca4a..2e54745d5 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -24,6 +24,7 @@ upstream_refs: - "GPU: vLLM 0.32.0 5df63a7203cdb9328d824a27b682e77c7b9928be" - "NPU target: vLLM 0.32.0 5df63a7203cdb9328d824a27b682e77c7b9928be" - "vLLM vllm.model_executor.models.deepseek_v2" + - "vLLM vllm.models.deepseek_v32.nvidia.model" - "vLLM vllm.model_executor.models.qwen3_moe" - "vLLM vllm.forward_context.ForwardContext" - "vLLM vllm.model_executor.model_loader" @@ -61,6 +62,7 @@ make a backend-specific worker class the shared model API. | --- | --- | --- | | Registration map | [`afd_plugin/__init__.py`](../../../afd_plugin/__init__.py) | [`test_package.py`](../../../tests/unit/package/test_package.py) | | Role-aware model and weight loading | [`deepseek_v2.py`](../../../afd_plugin/model_executor/models/deepseek_v2.py) | [`test_forward_context.py`](../../../tests/unit/model_executor/models/test_forward_context.py), model and accuracy E2E suites | +| GLM-5.x CUDA DSA role boundary | [`deepseek_v32.py`](../../../afd_plugin/model_executor/models/deepseek_v32.py) | [`test_deepseek_v32_construction.py`](../../../tests/unit/model_executor/models/test_deepseek_v32_construction.py), [`test_deepseek_v32_weight_policy.py`](../../../tests/unit/model_executor/models/test_deepseek_v32_weight_policy.py), [GLM-5.3 E2E](../../../tests/e2e/models/glm_5_3/test_glm_5_3.py) | | DeepSeek V4 CUDA role boundary | [`deepseek_v4.py`](../../../afd_plugin/model_executor/models/deepseek_v4.py) | [`test_deepseek_v4_construction.py`](../../../tests/unit/model_executor/models/test_deepseek_v4_construction.py), [`test_deepseek_v4_proxy.py`](../../../tests/unit/model_executor/models/test_deepseek_v4_proxy.py), [`test_deepseek_v4_weight_policy.py`](../../../tests/unit/model_executor/models/test_deepseek_v4_weight_policy.py) | | DeepSeek V4 Ascend Hash-id boundary | [`npu/deepseek_v4.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v4.py), [`npu/deepseek_v4_attention_gate.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py) | [`test_deepseek_v4_hash_ids.py`](../../../tests/unit/model_executor/test_deepseek_v4_hash_ids.py), [NPU model role tests](../../../tests/unit/model_executor/models/npu/test_deepseek_v4.py), ids-mode cases in [`test_camp2p_token_ids.py`](../../../tests/unit/connectors/test_camp2p_token_ids.py) | | Qwen3 MoE role-aware model and weight loading | [`qwen3_moe.py`](../../../afd_plugin/model_executor/models/qwen3_moe.py) | [`test_qwen3_moe_construction.py`](../../../tests/unit/model_executor/models/test_qwen3_moe_construction.py), [`test_qwen3_moe_weight_policy.py`](../../../tests/unit/model_executor/models/test_qwen3_moe_weight_policy.py) | @@ -86,7 +88,7 @@ registers lazy AFD wrapper paths under `AFD`-prefixed aliases. | `DeepseekV32ForCausalLM` | `AFDDeepseekV32ForCausalLM` | `AFDDeepseekV3ForCausalLM` | | `DeepseekV4ForCausalLM` | `AFDDeepseekV4ForCausalLM` | `AFDDeepseekV4ForCausalLM` | | `DSparkDraftModel` | Unsupported | `AFDDSparkDeepseekV4ForCausalLM` (Async CAM: eager/native decode Graph A, eager/FULL F) | -| `GlmMoeDsaForCausalLM` | `AFDGlmMoeDsaForCausalLM` | `AFDGlmMoeDsaForCausalLM` | +| `GlmMoeDsaForCausalLM` | `AFDGlmMoeDsaForCausalLM` | `AFDGlmMoeDsaForCausalLM` (CUDA: NVIDIA DSA wrapper; NPU: DeepSeek V2-derived) | | `Qwen3MoeForCausalLM` | `AFDQwen3MoeForCausalLM` | `AFDQwen3MoeForCausalLM` | | `Qwen3_5MoeForConditionalGeneration` | `AFDQwen3_5MoeForConditionalGeneration` | `AFDQwen3_5MoeForConditionalGeneration` | @@ -94,9 +96,10 @@ Only AFD workers switch their worker-local model configuration to the matching alias before constructing the AFD model runner. Non-AFD workers keep the checkpoint architecture and resolve to vLLM's native model class. -The DeepSeek, DeepSeek V2/V3/V3.2, and GLM aliases share the DeepSeek -V2-derived implementation. DeepSeek V4, Qwen3 MoE, and Qwen3.5/3.6 each have a -separate wrapper around their matching native architecture. These aliases +The DeepSeek and DeepSeek V2/V3/V3.2 aliases, and the GLM alias on NPU, share +the DeepSeek V2-derived implementation. GLM-5.x on CUDA, DeepSeek V4, Qwen3 +MoE, and Qwen3.5/3.6 each have a separate wrapper around their matching native +architecture. These aliases express known compatible architecture families; they do not make any wrapper a generic MoE model API. @@ -133,6 +136,45 @@ The full AFD model remains decorated with vLLM's compile support. Backend-only helpers are imported inside the NPU path so CUDA model import does not require vLLM-Ascend. +### GLM-5.x CUDA boundary + +vLLM serves `GlmMoeDsaForCausalLM` on CUDA with its NVIDIA DSA model +(`vllm.models.deepseek_v32.nvidia.model`): DeepSeek V2/V3 orchestration with +MLA plus the lightning indexer, fused embedding/RMSNorm, and an MLP output +left un-reduced until the next layer's fused all-reduce + RMSNorm. +[`deepseek_v32.py`](../../../afd_plugin/model_executor/models/deepseek_v32.py) +keeps that native forward and loader and copies only the decoder layer and +model constructors. Attention builds the DSA attention and the indexer top-k +buffer, the dense MLPs of the first `first_k_dense_replace` layers, and, per +MoE layer, the routed-experts shell `AFDDeepseekV2RoutedExpertsMoE` (gate, +native router, shared experts, dispatch/combine). FFN builds only the routed +experts (native `DeepseekV2MoE` without shared experts, external routing). A +one-pass checkpoint filter sends `self_attn` (indexer included), dense MLP and +shared-expert paths to Attention, routed experts to FFN, and the gate and +common paths to both. + +The wrapper requires `DeepEPM2NAFDConnector` with +`compute_gate_on_attention=true` and rejects Attention TP > 1 (the native +layer would all-reduce the connector's complete rows TP times), PP and +sequence-parallel MoE. MTP runs natively on the Attention ranks with the +native draft model and its local experts. Under DBO, +[`gpu_deepseek_v32_dbo.py`](../../../afd_plugin/compat/patches/gpu_deepseek_v32_dbo.py) +gives each micro-batch its own DSA top-k buffer and index groups (the +layers that reuse an indexer layer's top-k would otherwise read the other +micro-batch's) and leases the micro-batch compute SM budget to the indexer +planner. DBO is limited to decode-only micro-batching, so micro-batched +prefill fails at construction. With CUDA graphs, +[`gpu_mrv2_ubatching.py`](../../../afd_plugin/compat/patches/gpu_mrv2_ubatching.py) +falls back to an eager micro-batched step whenever a rank's trailing +micro-batch would be padding only, which the sparse FP8 MLA builder rejects. +MTP together with DBO uses a scoped backport of vLLM #61464: verification +steps split on request boundaries, and under DBO draft prefill runs on the +draft's own builders (FULL graphs captured there and refreshed before each +replay; eager steps rebuild the metadata), because a micro-batched target +forward never writes the target's full-batch buffers. MXFP4 routed experts, including +compressed-tensors W4A8 MXFP4, run on the connector's `marlin` expert backend; +block-128 FP8 experts run on DeepGEMM with the FP8 payload. + ### DeepSeek V4 CUDA boundary `AFDDeepseekV4ForCausalLM` wraps vLLM's NVIDIA DeepSeek V4 implementation and diff --git a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md index 76a2333bb..aa10b2ae0 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -27,8 +27,17 @@ Supported models: Attention runs the dense layers and the shared experts; the shared experts overlap the dispatch and the FFN. With Attention TP > 1 each TP rank runs them, with replicated weights, on the token slice it dispatches. - -DeepSeek-V3, GLM-DSA and Qwen3.5/Qwen3.6 MoE are rejected with this connector. +- GLM-5.x (`GlmMoeDsaForCausalLM`, validated with GLM-5.3 MXFP4xFP8) on vLLM's + NVIDIA DSA model. Attention runs the DSA attention, the three dense layers + and the shared expert; the FFN runs the routed experts. MXFP4 experts + (`RedHatAI/GLM-5.3-MXFP4xFP8`) need `expert_backend: "marlin"`; the + block-128 FP8 checkpoint (`zai-org/GLM-5.3`) uses the FP8 payload and + DeepGEMM. MTP runs natively on the Attention ranks (its draft experts + stay local). DBO runs eager or with CUDA graphs, with or without MTP, + with decode-only micro-batching, each micro-batch keeping its own DSA + top-k state. Attention TP > 1 and micro-batched prefill are rejected. + +DeepSeek-V3 and Qwen3.5/Qwen3.6 MoE are rejected with this connector. ## How it works @@ -139,7 +148,8 @@ The connector needs three components in addition to the plugin and vLLM: `cusparse.h` and other CUDA headers shipped by the `nvidia-*` wheels. Add that wheel `include` directory to `CPLUS_INCLUDE_PATH` for the build. - The `fused_experts` expert backend (vLLM's Triton kernels) runs without - DeepGEMM; it is a reference backend. + DeepGEMM; it is a reference backend. The `marlin` backend (MXFP4 + experts) also runs without DeepGEMM. 3. A GPU with compute capability 9.0 or newer on every rank. @@ -185,7 +195,7 @@ Unknown keys are rejected. All keys are optional. | Key | Type | Default | Meaning | | --- | --- | --- | --- | | `payload_format` | `str` | `"bf16"` | Dispatch wire format: `"bf16"` or `"fp8_block128_fp32"` (FP8 e4m3 rows with one fp32 scale per 128 channels). The FFN must load matching expert weights: BF16 weights for `bf16`, block-128 FP8 weights (for example an FP8 checkpoint with `weight_block_size=[128, 128]`) for `fp8_block128_fp32`. | -| `expert_backend` | `str` | `"deepgemm"` | `"deepgemm"` (DeepGEMM grouped GEMMs) or `"fused_experts"` (vLLM `fused_experts`, a reference path). | +| `expert_backend` | `str` | `"deepgemm"` | `"deepgemm"` (DeepGEMM grouped GEMMs), `"fused_experts"` (vLLM `fused_experts`, a reference path) or `"marlin"` (MXFP4 experts after vLLM's Marlin repack, W4A16; required for MXFP4 checkpoints, including compressed-tensors W4A8 MXFP4 experts such as GLM-5.3 MXFP4xFP8, and rejected for others; BF16 payload only). | | `capacity_tokens_per_rank` | `int \| null` | `null` | Largest dispatch one Attention rank issues per call (with Attention TP, its slice); sizes the transport buffers. `null` uses the larger of `max_num_batched_tokens` and the largest CUDA-graph capture size, divided by the Attention TP size. | | `num_lanes` | `int \| null` | `null` | In-flight micro-batches, one transport buffer each, `1` or `2`. `null` uses vLLM's micro-batch count: 2 with DBO, otherwise 1. | | `bucket_granularity` | `int` | `1` | Dispatch buckets are rounded up to this multiple, then to the next rung of the bucket ladder derived from the CUDA-graph capture sizes. Every step uses one of a few kernel sets compiled at startup. | @@ -223,7 +233,7 @@ that sets `capacity_tokens_per_rank` explicitly must agree with Attention. | `--tensor-parallel-size` | FFN | Must be `1`. Attention may use TP > 1. | | Expert placement | FFN | `linear` (vLLM's default `--expert-placement-strategy`). | | `--all2all-backend` | FFN | `allgather_reducescatter` (vLLM's default). The connector moves the rows itself, so vLLM's MoE all-to-all never runs; this backend allocates no buffers and keeps vLLM's own DeepEP out of the process. | -| `--moe-backend` | FFN | `triton` or `auto`. `auto` is pinned to `triton`; any other backend is rejected because it reorders the expert weights that the connector reads. | +| `--moe-backend` | FFN | `triton` or `auto` for BF16 and FP8 experts; `auto` is pinned to `triton`. `marlin` or `auto` for MXFP4 experts; `auto` is pinned to `marlin`. Any other backend is rejected because it lays out the expert weights differently from what the connector reads. | | `--enable-eplb` | Both | Not supported. | | `--enable-dbo` | Both | Optional, but must match on both roles. DBO uses exactly two micro-batches. | | CUDA graphs | Both | Eager or `FULL_DECODE_ONLY`. | diff --git a/docs/reports/glm-5.3-afd-followups.md b/docs/reports/glm-5.3-afd-followups.md new file mode 100644 index 000000000..565698ccb --- /dev/null +++ b/docs/reports/glm-5.3-afd-followups.md @@ -0,0 +1,112 @@ +# GLM-5.3 on AFD: follow-ups and upstream suggestions + +Status of branch `glm-5.3-afd` (vLLM 0.32.0 pin). GPU evidence: one +8 x H200 node, `RedHatAI/GLM-5.3-MXFP4xFP8`, GSM8K-200. + +## Validated + +| Cell | GSM8K-200 | +| --- | --- | +| Native DP8 + EP eager (Marlin) | 0.925 | +| AFD DeepEP M2N 4A4F eager / graph | 0.930 (0.915 on a second run) / 0.935 | +| Native / AFD graph + MTP (3 tokens) | 0.93 / 0.92 (acceptance 3.4-3.5 / 3.4) | +| AFD eager DBO, decode-only micro-batching | 0.93 (conc 32), 0.91 (conc 64) | +| AFD graph DBO, decode-only micro-batching | 0.93 (conc 64) | +| AFD eager MTP (3 tokens) + DBO, decode-only | 0.94 (conc 32, acceptance 3.35-3.73) | +| AFD graph MTP (3 tokens) + DBO, decode-only | 0.95 (conc 32, acceptance 3.30-3.69) | + +## Open work + +1. **Micro-batched prefill corruption (DBO).** With per-micro-batch DSA + top-k state, prefill DBO still scores 0.90 with garbled answers (0.685 + without the top-k fix). Ruled out: top-k/index-group state, vLLM + workspaces (per micro-batch), MLA prefill backend and chunked-context + workspace (per builder), CuTe input alignment, the Marlin FFN executor, + stream ordering. D-B (prefix caching off, prefill DBO on): 1 garbled + answer of 200 (doc 138) instead of several, so prefix-cache hits make the + corruption much more frequent. **D-A (request-aligned split of every + micro-batched step, prefix caching on, 4A4F eager): 0 garbled answers; + every micro-batched step had a request boundary to split on (401 aligned, + 0 straddled), GSM8K-200 0.895.** The garbling is caused by a prefill + request straddling the token-midpoint split (its second half runs in the + other micro-batch as a chunked prefill over context written by the first + one in the same layer); the kernel-level cause is still open. Under M2N + eager the Attention DP ranks need not agree on per-micro-batch token + counts, so the aligned split ran at DP4. Flip comparison against a + non-DBO eager run on the same node (0.915): D-A lost 8 and gained 4 + questions, all coherent reasoning slips (graph DBO: lost 6, gained 9; eager + MTP + DBO: lost 3, gained 8), so the remaining gap is within run-to-run + spread (0.915-0.930 for non-DBO eager). Fix direction: a request-aligned + split for micro-batched prefill. Remaining before lifting the rule: steps + with a single request (no boundary; they would still straddle, and all + Attention ranks must micro-batch alike), and the kernel-level cause. Until then the adapter requires + `--dbo-prefill-token-threshold` > `--max-num-batched-tokens`. +2. **Graph DBO: done.** Off-by-one in the "empty second microbatch" guard: + `sync_cudagraph_and_dp_padding` (vllm/v1/worker/gpu/dp_utils.py) falls + back only when `2 * min(tokens) < padded`; with an even padded count, a + rank with exactly half the tokens replays a FULL graph whose second + micro-batch holds only zero-length requests, and + `FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode` asserts + `decode_query_len > 0`. `gpu_mrv2_ubatching.py` now carries a marked copy + whose guard is exact for the midpoint split + (`min <= (n - 1) * (padded // n)`), used by AFD's DP1 path too. Graph DBO + GSM8K-200 0.93; the eager-only rule is lifted. +3. **MTP + DBO: done, eager and graph.** vLLM #61464 backported as scoped + patches (`config_dbo_spec_decode`, `gpu_mrv2_dbo_spec_decode`: gate, + threshold raise, request-boundary split of verification batches, draft + prefill on the draft's own builders). Eager MTP + DBO 0.94 (MTP alone + 0.92). Graph MTP + DBO first hit `illegal memory access` about ten minutes + into GSM8K: the pinned speculator captures the FULL draft-prefill graphs on + the target's attention builders and replays whatever full-batch metadata + the target last wrote there, and a micro-batched target forward writes + only per-micro-batch metadata (micro-batch 0 shares the full-batch slot). + Confirmed by running only that draft prefill eager after micro-batched + steps (0.92, no errors), then fixed as #61464 does: under DBO the draft + prefill graphs are captured on the draft's builders and each replay first + refreshes them with this batch's metadata (graph MTP + DBO 0.95, acceptance + 3.30-3.69, lost 4 / gained 11 against non-DBO eager). +4. Attention TP > 1 (connector rows must be returned as a partial for the + native fused all-reduce + RMSNorm). +5. FP8 checkpoint on two pods (`multipod_4a4f_4a4f_fp8_graph.sh`) not + GPU-validated; E2E case not run through the harness; no performance + numbers; DeepSeek-V3.2 still on the legacy wrapper. +6. Image: vllm-openai v0.32.0 ships NCCL 2.29.7; afdEP-comm v1.1.0 needs + 2.30.7 (GIN request API). + +## Suggested upstream (vLLM) changes + +| Issue | Evidence | Suggested fix | +| --- | --- | --- | +| DSA top-k state shared across DBO micro-batches | GLM index_topk_freq=4 layers reuse the model-level `topk_indices_buffer` and `SparseMLAIndexGroup` workspaces; DBO yields between layers; GSM8K 0.685 without per-micro-batch copies | Per-micro-batch top-k buffers and index groups (or key them by micro-batch) in `deepseek_v32` and the sparse MLA backends | +| Indexer SM budget under DBO | GLM DBO deadlocked until the indexer planner got the micro-batch compute SMs | Land #60732 for DSA (V3.2/GLM), as the V4 path needs | +| Empty-second-microbatch guard off by one | Padding-only FULL micro-batch -> `assert decode_query_len > 0`; fixed in AFD by a marked copy, graph DBO then 0.93 | `min(tokens) <= (n - 1) * (padded // n)` (for two micro-batches: `2 * min(tokens) <= padded`) in `sync_cudagraph_and_dp_padding` | +| Sparse FP8 builder assumes a non-empty first request | Same assert | Return empty decode metadata when the micro-batch has no real tokens | +| Prefill straddling the token-midpoint split | GLM prefill DBO garbles answers with the midpoint split; a request-aligned split removes them (0 garbled of 200) | Request-aligned split for prefill (RFC #50738 open question 2), and find why a straddling sparse-MLA prefill reads wrong context | +| #61464 test coverage | E2E uses ngram only | Add MTP/EAGLE + DBO GSM8K, and a DSA model | +| Graph DBO with a target-dependent drafter | GLM MTP + graph DBO: illegal memory access until draft prefill captured on, and refreshed through, the draft builders (then 0.95) | #61464 already does this; extend its E2E to graph DBO with MTP | +| Non-uniform-shaped draft buffer sharing | `eagle/utils.py` re-points draft top-k buffers to the target's | Keep in mind if drafting is ever micro-batched | + +## Upstream status after v0.32.0 (checked 2026-11-14) + +Checked against vLLM `main` at `449576cd` (2026-11-14) and the +`releases/v0.33.0` branch (v0.33.0rc3; cut from `main` at `4135590f` on +2026-10-26). "main only" means merged after the v0.33 branch cut, so not in +v0.33.0. All open PRs below are unmerged; none changes the AFD patches yet. + +| Item | vLLM issue / PR | Status | Coverage | +| --- | --- | --- | --- | +| DSA top-k state shared across DBO micro-batches | [#60734](https://github.com/vllm-project/vllm/pull/60734) Keep DSA shared-indexer top-k state per DBO microbatch | Open, `needs-rebase`, last update 2026-10-20 | Full: per-micro-batch top-k and index-group state for `index_topk_freq > 1` (GLM-5.x), GLM-5.3 GSM8K 0.61 -> 0.92 under DBO. Main still has one model-level `topk_indices_buffer`. | +| (same class) | [#56476](https://github.com/vllm-project/vllm/pull/56476) Isolate sparse index state across microbatches (DSv4.1) | Open, `needs-rebase` | DeepSeek-V4.1 only; #60734 would reuse its helper if it lands first. | +| (same class, PP) | [#59961](https://github.com/vllm-project/vllm/pull/59961) Carry the shared indexer top-k across pipeline stages ([#63170](https://github.com/vllm-project/vllm/issues/63170)) | Merged 2026-11-01, main only | Not DBO: stale shared top-k at a PP stage start. Confirms the shared top-k hazard upstream. | +| Indexer SM budget under DBO | [#60732](https://github.com/vllm-project/vllm/pull/60732) Re-plan DSA indexer decode schedule for DBO's reduced DeepGEMM SM count | Open, `needs-rebase`, last update 2026-10-14 | Full. | +| Empty-second-micro-batch guard off by one | none | Not fixed: `sync_cudagraph_and_dp_padding` on main still uses `2 * min(tokens) < padded` | Related: [#63823](https://github.com/vllm-project/vllm/pull/63823) (open draft) removes the guard and replays an all-padding FULL region with zero-length rows. For DSA that reaches the sparse FP8 assert below unless the builder is fixed too. [#63556](https://github.com/vllm-project/vllm/pull/63556) (open) stages MLA decode rows instead. | +| Sparse FP8 builder assumes a non-empty first request | [#66734](https://github.com/vllm-project/vllm/pull/66734) DBO and async EP fixes (padding-only DBO microbatch) | Open | Partial: an eager micro-batch wholly in DP padding skips attention (`cg_mode != FULL` only), found on MiniMax-M3's `assert decode_query_len > 0`. `FlashMLASparseMetadataBuilder` still asserts; its `num_tokens == 0` early return predates v0.32 (PCP). | +| Prefill straddling the token-midpoint split | RFC [#50738](https://github.com/vllm-project/vllm/issues/50738) | Closed 2026-11-09 by its author after the stale notice, no linked change | None. Main splits on a request boundary only for uniform spec-decode verification batches (#61464); prefill still splits at the token midpoint. | +| (related) | [#63706](https://github.com/vllm-project/vllm/pull/63706) `--dbo-scope {all,prefill}` | Open | Opposite scope (prefill-only DBO); would expose the straddle on GLM. | +| (related, ruled out for AFD) | [#61584](https://github.com/vllm-project/vllm/pull/61584) Reserve peer microbatch workspaces after profiling ([#61487](https://github.com/vllm-project/vllm/issues/61487)) | Open | Workspace sizing for micro-batches other than 0; AFD already ruled out workspaces for the prefill corruption. | +| MTP + DBO (#61464 backport) | [#61464](https://github.com/vllm-project/vllm/pull/61464) Support DBO with speculative decoding and DCP | Merged 2026-10-28, main only | Full; AFD carries it as scoped patches until the pin moves past it. | +| #61464 test coverage | #61464 tests | Merged as is | Gap remains: E2E `tests/v1/distributed/test_dbo.py` uses ngram only; no MTP/EAGLE, graph DBO or DSA model. | +| Graph DBO with a target-dependent drafter | #61464 | Merged 2026-10-28, main only | Full (draft prefill graphs on the draft builders). | +| Draft buffer sharing if drafting is micro-batched | RFC [#65119](https://github.com/vllm-project/vllm/issues/65119) DBO for the MTP/EAGLE drafter forward | Open RFC | Would micro-batch the drafter; `eagle/utils.py` on main still re-points draft top-k buffers to the target's, so the hazard applies if it lands. | +| Attention TP > 1 (open work 4) | [#60477](https://github.com/vllm-project/vllm/pull/60477) DSA: all-reduce partial hidden_states before PP send; [#62947](https://github.com/vllm-project/vllm/pull/62947) honor the decoder's deferred MoE all-reduce on the naive all2all path | #60477 merged 2026-10-29 (main only); #62947 open | Not an AFD fix. Both handle the same deferred reduction (the MLP output is reduced in the next layer's fused all-reduce + RMSNorm) at a PP boundary and on the all-gather/reduce-scatter path; #60477 is the pattern for returning connector rows as a partial. | +| NCCL in the image (open work 6) | [#58768](https://github.com/vllm-project/vllm/pull/58768) add vllm-openai-efa target | On `releases/v0.33.0` | `docker/Dockerfile` `NCCL_VERSION` goes 2.30.7 (already the v0.32.0 value) -> 2.31.2. The v0.32.0 image used here still reported 2.29.7, so check the published v0.33.0 image before dropping the NCCL 2.30.7 wheel step. | diff --git a/docs/support.md b/docs/support.md index d45f44791..913bd1887 100644 --- a/docs/support.md +++ b/docs/support.md @@ -9,6 +9,7 @@ source revision. See [installation](installation.md) for the pinned runtime stac | Model family | Execution scope | | --- | --- | | DeepSeek V2 / V3 / V3.2 | CUDA P2P and synchronous Ascend paths; DeepEP M2N admits V2 only. Each role loads its required components. | +| GLM-5.x (`GlmMoeDsaForCausalLM`) | CUDA DeepEP M2N on vLLM's NVIDIA DSA model, Attention TP1, MXFP4 (Marlin) or block-128 FP8 experts; MTP on the Attention ranks; DBO eager or graph with decode-only micro-batching, with or without MTP. Rejects P2P, Attention TP > 1, PP, SP MoE and micro-batched prefill. Ascend keeps the DeepSeek V2-derived path. | | DeepSeek V4 Flash | CUDA P2P or Ascend CAMP2P / Async CAM; backend-specific limits are in the matrix below. | | Qwen3 MoE | CUDA; P2P gates on FFN, DeepEP M2N gates on Attention. Rejects sequence-parallel MoE, EPLB, PP, speculative decoding, LoRA and NPU. | | Qwen3.5 / Qwen3.6 MoE | CUDA text-only with `--language-model-only`; evidence covers Qwen3.6-35B-A3B native DP4/TP1/EP4 and synchronous AFD 2A1F eager/graph. GDN DBO, Attention-side routing, PP, async and multi-node execution are unsupported; NPU/multimodal are unsupported or unverified, quantization is unverified, and no performance claim is made. | @@ -26,7 +27,7 @@ qualify every model/topology combination on the current target. | Platform/path | Execution | Ubatching | Routing/quantization limits | Evidence | | --- | --- | --- | --- | --- | | CUDA V2 + `P2pNcclAFDConnector` | Eager or `FULL_DECODE_ONLY` native V2 CUDA Graph | Exactly two native microbatches; Attention DP1 or DP > 1 | Model-specific gate placement; PP/CP, elastic EP, EPLB, SP MoE, and compile SP rejected; role ranks equal DP x TP | Source-specific GPU validation is linked below; coverage is limited to the recorded model/topology combinations. DeepSeek-V2-Lite has known DBO shutdown failures. Qwen3.5/3.6 GDN DBO is unsupported. | -| CUDA V2 + `DeepEPM2NAFDConnector` | Eager or `FULL_DECODE_ONLY` native V2 CUDA Graph; SM90 or newer | Exactly two native microbatches, enabled on both roles; Attention DP1 or DP > 1, TP1 or TP2 | Qwen3 MoE (all layers MoE, no shared expert) or DeepSeek-V2; `compute_gate_on_attention=true`; FFN TP1 with EP over all FFN ranks, linear placement, `allgather_reducescatter` all2all and the Triton MoE weight layout; BF16 or block-128 FP8 payload; EPLB, PP/CP and SP MoE rejected | Historical vLLM 0.30 Qwen3-30B-A3B 2A2F GSM8K E2E for eager, graph, graph + DBO, and Attention TP2 eager/graph on H100 and H200, and for eager, graph and graph + DBO on GB200; two-node 2A2F eager, graph and graph + DBO with Attention and FFN on separate nodes. DeepSeek-V2-Lite: [E2E definitions](../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py). [GPU connector device tests](../tests/gpu/README.md) | +| CUDA V2 + `DeepEPM2NAFDConnector` | Eager or `FULL_DECODE_ONLY` native V2 CUDA Graph; SM90 or newer | Exactly two native microbatches, enabled on both roles; Attention DP1 or DP > 1, TP1 or TP2 | Qwen3 MoE (all layers MoE, no shared expert), DeepSeek-V2 or GLM-5.x; `compute_gate_on_attention=true`; FFN TP1 with EP over all FFN ranks, linear placement, `allgather_reducescatter` all2all and the Triton MoE weight layout (Marlin for MXFP4 experts); BF16 or block-128 FP8 payload; EPLB, PP/CP and SP MoE rejected | Historical vLLM 0.30 Qwen3-30B-A3B 2A2F GSM8K E2E for eager, graph, graph + DBO, and Attention TP2 eager/graph on H100 and H200, and for eager, graph and graph + DBO on GB200; two-node 2A2F eager, graph and graph + DBO with Attention and FFN on separate nodes. DeepSeek-V2-Lite: [E2E definitions](../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py). GLM-5.3 MXFP4xFP8 on one 8 x H200 node (vLLM 0.32.0, afdEP-comm v1.1.0 on NCCL 2.30.7): GSM8K-200 native DP8 + EP eager 0.925, M2N 4A4F eager 0.930 and graph 0.935; with MTP (3 tokens) native graph 0.93 vs M2N graph 0.92, mean acceptance length 3.4-3.5 vs 3.4; DBO eager with decode-only micro-batching 0.93 (concurrency 32) and 0.91 (64), where micro-batched prefill dropped to 0.90 with garbled outputs and to 0.685 without per-micro-batch top-k ([E2E definitions](../tests/e2e/models/glm_5_3/test_glm_5_3.py)); the FP8 two-pod recipe is not GPU-validated. [GPU connector device tests](../tests/gpu/README.md) | | Ascend V2 + `CAMP2pAFDConnector` | Eager, `FULL`, or `FULL_DECODE_ONLY` native V2 ACL Graph | Eager or graph native DBO with two ubatches and Attention DP > 1; plain MLA graph DBO requires `FULL_DECODE_ONLY` | `compute_gate_on_attention` may be `false`, or `true` for DeepSeek-V4 at TP=1 without DBO/ubatching; PP/CP, elastic EP and EPLB rejected; model-specific SP policy; role ranks equal DP x TP | Historical 0.30 V2-Lite 2A2F FULL 300-question comparison and 2A1F functional checks; current 0.32 Lite evidence is in the [validation record](https://github.com/vllm-project/afd-plugin/pull/507) | | Ascend V2 + `CAMAsyncAFDConnector` | Eager Attention; layered FFN eager or FULL | Two native DBO microbatches on Attention, including DP1 and with FFN FULL; FFN stays outside the native DBO schedule | Experimental DeepSeek-V4 W4A8/W8A8 on A3/910C; `async_dp=true`, `compute_gate_on_attention=true`, and `dynamicQuant=1`; PP/CP unsupported; plugin-owned async CAM ops required | Historical DSV4 layered comparisons and limitations are in the [validation record](https://github.com/vllm-project/afd-plugin/pull/507); they predate removal of the legacy path | diff --git a/recipe/gpu/DeepEPM2NAFDConnector/README.md b/recipe/gpu/DeepEPM2NAFDConnector/README.md index 7ca658db1..0ef7dbc39 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/README.md @@ -16,11 +16,13 @@ Recipe names encode their layout, DP/TP shape, and mode. DeepSeek-V2-Lite and Qwen3-30B-A3B provide `2a2f_{eager,graph,graph_dbo}_dp2tp1.sh` plus `2a2f_eager_dp1tp2.sh`; Qwen3-235B-A22B-FP8 provides `4a2f` eager, graph and DBO recipes plus `multipod_4a1f_4a1f_dbo.sh`. Each model README lists its -exact filenames and defaults. +exact filenames and defaults. GLM-5.3 provides MXFP4 `4a4f` eager and graph +recipes on one 8-GPU node and an FP8 two-pod graph recipe. ```text . ├── deepseek_v2_lite/ # BF16 payload +├── glm_5_3/ # MXFP4 experts on Marlin (BF16 payload) or FP8 payload ├── qwen3_30b_a3b/ # BF16 payload └── qwen3_235b_a22b_fp8/ # FP8 payload (fp8_block128_fp32) ``` diff --git a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md new file mode 100644 index 000000000..262c1f9f6 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md @@ -0,0 +1,69 @@ +# GLM-5.3 with DeepEPM2NAFDConnector + +GLM-5.3 (`GlmMoeDsaForCausalLM`) recipes. The scripts, layouts, settings and +multi-pod runs are described once in the +[connector recipes README](../README.md); this page lists what is specific to +the model. + +```text +prefill_decode_colocation/ +├── 4a4f_eager_dp4tp1.sh # MXFP4xFP8, eager, one 8-GPU node +├── 4a4f_graph_dp4tp1.sh # MXFP4xFP8, graph, one 8-GPU node +├── 4a4f_graph_mtp_dp4tp1.sh # MXFP4xFP8, graph, MTP on Attention +├── 4a4f_dbo_dp4tp1.sh # MXFP4xFP8, DBO, MODE=graph (default) or eager +├── 4a4f_mtp_dbo_dp4tp1.sh # MXFP4xFP8, MTP + DBO, MODE=graph (default) or eager +└── multipod_4a4f_4a4f_fp8_graph.sh # block-128 FP8, graph, two 8-GPU pods +``` + +The native vLLM control on the same node is +[`baseline/glm_5_3/native_dp8.sh`](../../baseline/glm_5_3/native_dp8.sh). + +## Checkpoints + +| Checkpoint | Size | Routed experts on the FFN | Recipe | +| --- | --- | --- | --- | +| `RedHatAI/GLM-5.3-MXFP4xFP8` | 415 GB | MXFP4, `expert_backend: "marlin"` (W4A16 Marlin), BF16 payload; about 96 GB per FFN rank at 4 ranks | `4a4f_*` | +| `zai-org/GLM-5.3` | 756 GB | block-128 FP8, DeepGEMM, FP8 payload; about 91 GB per FFN rank at 8 ranks | `multipod_4a4f_4a4f_fp8_graph.sh` | + +## Defaults + +| Setting | All recipes | +| --- | --- | +| `--max-num-seqs` | 64 per Attention DP rank (128 in `4a4f_dbo_dp4tp1.sh`) | +| `--max-num-batched-tokens` | 4096 | +| `--max-model-len` | 8192 | +| `--kv-cache-dtype` | `fp8` (vLLM's `fp8_ds_mla` sparse MLA cache on Hopper) | +| `CUDAGRAPH_CAPTURE_SIZES` (graph) | 1,2,4,8,16,24,32,48,64 | +| `VLLM_ENGINE_READY_TIMEOUT_S` | 3600 | + +## Notes + +- Attention owns the DSA attention, the indexer, the three dense layers and + the shared expert; Attention ranks hold no routed experts + (`--enable-ep-weight-filter`). +- The FFN pins `--moe-backend marlin` for MXFP4 experts (vLLM loads the + compressed-tensors W4A8 experts with its MXFP4 method and repacks them for + Marlin). Marlin runs BF16 activations, as native vLLM does on Hopper, where + no MXFP4 MoE kernel takes the checkpoint's FP8 activations. Run the native + control with the same `--moe-backend marlin`. +- The FP8 recipe needs DeepGEMM (see the connector guide). The MXFP4 recipes + do not; leave vLLM's bundled DeepGEMM in place for the DSA indexer. +- `vllm/vllm-openai:v0.32.0` ships `nvidia-nccl-cu13` 2.29.7, which lacks the + NCCL GIN request API afdEP-comm's kernels compile against. Install + `nvidia-nccl-cu13==2.30.7` and export `AFDEP_NCCL_ROOT` and + `LD_LIBRARY_PATH` as the connector recipes README shows. + +## Limitations + +- Attention TP1 only: the native DSA layer reduces the MLP output inside the + next layer's RMSNorm, so Attention TP > 1 is rejected. +- MTP (`4a4f_graph_mtp_dp4tp1.sh`) runs the native draft model on the + Attention ranks; only the Attention role takes the speculative config. +- DBO (`4a4f_dbo_dp4tp1.sh`, `4a4f_mtp_dbo_dp4tp1.sh`) micro-batches decode + steps only, eager or with CUDA graphs: `--dbo-prefill-token-threshold` must + stay above `--max-num-batched-tokens`, because micro-batched GLM prefill is + not yet correct. With MTP, verification steps split on request boundaries, + the decode threshold is raised to one request per micro-batch, and draft + prefill runs on the draft's own attention builders (a backport of vLLM + #61464). +- No P2P connector. diff --git a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_dbo_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_dbo_dp4tp1.sh new file mode 100755 index 000000000..4f46f6447 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_dbo_dp4tp1.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# GLM-5.3 MXFP4xFP8 on DeepEPM2NAFDConnector, 4A4F with DBO (two +# micro-batches) on one 8-GPU node: FFN DP4 x TP1 (Marlin MXFP4 experts) and +# Attention DP4 x TP1. MODE is graph (default) or eager. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables + +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-4A4F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/GLM-5.3-MXFP4xFP8} +CONNECTOR_CONFIG='"expert_backend": "marlin"' # MXFP4 experts, BF16 payload +MODE=${MODE:-graph} +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-1,2,4,8,16,24,32,48,64,96,128} +# DBO splits decode steps only: the prefill threshold stays above the step's +# token budget, because micro-batched GLM prefill is not yet correct. +MAX_NUM_BATCHED_TOKENS=4096 +SERVE_ARGS=(--enable-ep-weight-filter --kv-cache-dtype fp8 --api-server-count 1 + --max-num-seqs 128 --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" --max-model-len 8192 + --enable-dbo --dbo-decode-token-threshold 2 + --dbo-prefill-token-threshold $((MAX_NUM_BATCHED_TOKENS + 1))) +# A cold load of the 415 GB checkpoint outlasts vLLM's 600 s front-end wait. +export VLLM_ENGINE_READY_TIMEOUT_S=${VLLM_ENGINE_READY_TIMEOUT_S:-3600} + +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_eager_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_eager_dp4tp1.sh new file mode 100755 index 000000000..83e9474ae --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_eager_dp4tp1.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# GLM-5.3 MXFP4xFP8 on DeepEPM2NAFDConnector, 4A4F eager mode on one 8-GPU +# node: FFN DP4 x TP1 (64 MXFP4 experts per rank, Marlin) and Attention DP4 x TP1. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables + +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-4A4F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/GLM-5.3-MXFP4xFP8} +CONNECTOR_CONFIG='"expert_backend": "marlin"' # MXFP4 experts, BF16 payload +MODE=eager +SERVE_ARGS=(--enable-ep-weight-filter --kv-cache-dtype fp8 --api-server-count 1 + --max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 8192) +# A cold load of the 415 GB checkpoint outlasts vLLM's 600 s front-end wait. +export VLLM_ENGINE_READY_TIMEOUT_S=${VLLM_ENGINE_READY_TIMEOUT_S:-3600} + +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_graph_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_graph_dp4tp1.sh new file mode 100755 index 000000000..8ef752de4 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_graph_dp4tp1.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# GLM-5.3 MXFP4xFP8 on DeepEPM2NAFDConnector, 4A4F graph mode on one 8-GPU +# node: FFN DP4 x TP1 (64 MXFP4 experts per rank, Marlin) and Attention DP4 x TP1. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables + +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-4A4F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/GLM-5.3-MXFP4xFP8} +CONNECTOR_CONFIG='"expert_backend": "marlin"' # MXFP4 experts, BF16 payload +MODE=graph +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-1,2,4,8,16,24,32,48,64} +SERVE_ARGS=(--enable-ep-weight-filter --kv-cache-dtype fp8 --api-server-count 1 + --max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 8192) +# A cold load of the 415 GB checkpoint outlasts vLLM's 600 s front-end wait. +export VLLM_ENGINE_READY_TIMEOUT_S=${VLLM_ENGINE_READY_TIMEOUT_S:-3600} + +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_graph_mtp_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_graph_mtp_dp4tp1.sh new file mode 100755 index 000000000..1eae837f9 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_graph_mtp_dp4tp1.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# GLM-5.3 MXFP4xFP8 on DeepEPM2NAFDConnector, 4A4F graph mode with MTP on one +# 8-GPU node: FFN DP4 x TP1 (Marlin MXFP4 experts) and Attention DP4 x TP1. +# The MTP draft runs natively on the Attention ranks, so only the Attention +# role takes the speculative config (NUM_SPECULATIVE_TOKENS, default 3). +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables + +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-4A4F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/GLM-5.3-MXFP4xFP8} +CONNECTOR_CONFIG='"expert_backend": "marlin"' # MXFP4 experts, BF16 payload +MODE=${MODE:-graph} +# Decode steps carry 1 + NUM_SPECULATIVE_TOKENS tokens per request. +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-4,8,16,32,64,96,128,192,256} +SPECULATIVE_CONFIG="{\"method\":\"mtp\",\"num_speculative_tokens\":${NUM_SPECULATIVE_TOKENS:-3}}" +ATTN_ARGS="--speculative-config=${SPECULATIVE_CONFIG} ${ATTN_ARGS:-}" +SERVE_ARGS=(--enable-ep-weight-filter --kv-cache-dtype fp8 --api-server-count 1 + --max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 8192) +# A cold load of the 415 GB checkpoint outlasts vLLM's 600 s front-end wait. +export VLLM_ENGINE_READY_TIMEOUT_S=${VLLM_ENGINE_READY_TIMEOUT_S:-3600} + +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_mtp_dbo_dp4tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_mtp_dbo_dp4tp1.sh new file mode 100755 index 000000000..75319efac --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_mtp_dbo_dp4tp1.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# GLM-5.3 MXFP4xFP8 on DeepEPM2NAFDConnector, 4A4F with MTP and DBO (two +# micro-batches) on one 8-GPU node: FFN DP4 x TP1 (Marlin MXFP4 experts) and +# Attention DP4 x TP1. The MTP draft runs on the Attention ranks only +# (NUM_SPECULATIVE_TOKENS, default 3). MODE is graph (default) or eager. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables + +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-4A4F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/GLM-5.3-MXFP4xFP8} +CONNECTOR_CONFIG='"expert_backend": "marlin"' # MXFP4 experts, BF16 payload +MODE=${MODE:-graph} +# Decode steps carry 1 + NUM_SPECULATIVE_TOKENS tokens per request. +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-4,8,16,32,64,96,128,192,256} +SPECULATIVE_CONFIG="{\"method\":\"mtp\",\"num_speculative_tokens\":${NUM_SPECULATIVE_TOKENS:-3}}" +ATTN_ARGS="--speculative-config=${SPECULATIVE_CONFIG} ${ATTN_ARGS:-}" +# DBO splits decode (verification) steps only; the AFD gate raises the decode +# threshold to one request per micro-batch. The prefill threshold stays above +# the step's token budget, because micro-batched GLM prefill is not yet correct. +MAX_NUM_BATCHED_TOKENS=4096 +SERVE_ARGS=(--enable-ep-weight-filter --kv-cache-dtype fp8 --api-server-count 1 + --max-num-seqs 64 --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" --max-model-len 8192 + --enable-dbo --dbo-decode-token-threshold 2 + --dbo-prefill-token-threshold $((MAX_NUM_BATCHED_TOKENS + 1))) +# A cold load of the 415 GB checkpoint outlasts vLLM's 600 s front-end wait. +export VLLM_ENGINE_READY_TIMEOUT_S=${VLLM_ENGINE_READY_TIMEOUT_S:-3600} + +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/multipod_4a4f_4a4f_fp8_graph.sh b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/multipod_4a4f_4a4f_fp8_graph.sh new file mode 100755 index 000000000..576b28b05 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/multipod_4a4f_4a4f_fp8_graph.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# GLM-5.3 (block-128 FP8) on DeepEPM2NAFDConnector, graph mode on two 8-GPU +# pods: FFN DP8 x TP1 (32 FP8 experts per rank, DeepGEMM) and Attention DP8 x TP1, +# with the FP8 payload. Run with POD=AFD_0 and POD=AFD_1. +# Placement: one pod per comma-separated LAYOUT entry; FFN on each pod's first GPUs. +# Settings and running in a pod: recipe/gpu/common/README.md +# shellcheck disable=SC1090,SC2034 # the sourced launcher reads these variables + +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +LAYOUT=${LAYOUT:-4A4F,4A4F} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/GLM-5.3} +CONNECTOR_CONFIG='"payload_format": "fp8_block128_fp32"' # FP8 payload +MODE=graph +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-1,2,4,8,16,24,32,48,64} +SERVE_ARGS=(--enable-ep-weight-filter --kv-cache-dtype fp8 --api-server-count 1 + --max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 8192) +# A cold load of the 756 GB checkpoint outlasts vLLM's 600 s front-end wait. +export VLLM_ENGINE_READY_TIMEOUT_S=${VLLM_ENGINE_READY_TIMEOUT_S:-3600} + +source "${AFD_RECIPE_LIB:-$(dirname "${BASH_SOURCE[0]}")/../../../common/afd_colocation.sh}" +afd_colocation_serve "$@" diff --git a/recipe/gpu/baseline/glm_5_3/native_dp8.sh b/recipe/gpu/baseline/glm_5_3/native_dp8.sh new file mode 100755 index 000000000..de8837dbd --- /dev/null +++ b/recipe/gpu/baseline/glm_5_3/native_dp8.sh @@ -0,0 +1,42 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# GLM-5.3 native vLLM baseline (no AFD) on one 8-GPU node: DP8 with expert +# parallelism. It is the control for +# recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_*.sh: +# same checkpoint, KV cache dtype, batch limits and, for MXFP4 experts, the +# same Marlin kernel the M2N FFN runs. +# +# Every knob is an environment variable: +# +# MODEL_PATH weights dir or HF id (/path/model_weights/GLM-5.3-MXFP4xFP8) +# MODE graph | eager (graph: FULL_DECODE_ONLY CUDA graphs) +# MOE_BACKEND vLLM --moe-backend (marlin; use auto for FP8 experts) +# CUDAGRAPH_CAPTURE_SIZES decode graph sizes (1,2,4,8,16,24,32,48,64) +# HOST / PORT HTTP bind address and port (127.0.0.1 / 18000) +# EXTRA_ARGS extra `vllm serve` arguments (empty) + +set -euo pipefail + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/GLM-5.3-MXFP4xFP8} +MODE=${MODE:-eager} +MOE_BACKEND=${MOE_BACKEND:-marlin} +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-1,2,4,8,16,24,32,48,64} +HOST=${HOST:-127.0.0.1} +PORT=${PORT:-18000} +read -ra EXTRA <<<"${EXTRA_ARGS:-}" + +case $MODE in + graph) MODE_ARGS=(--max-cudagraph-capture-size "${CUDAGRAPH_CAPTURE_SIZES##*,}" + --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}") ;; + eager) MODE_ARGS=(--enforce-eager) ;; + *) echo "MODE must be graph or eager, got '$MODE'" >&2; exit 1 ;; +esac + +export VLLM_USE_V2_MODEL_RUNNER=1 +export VLLM_ENGINE_READY_TIMEOUT_S=${VLLM_ENGINE_READY_TIMEOUT_S:-3600} +exec vllm serve "$MODEL_PATH" --data-parallel-size 8 --enable-expert-parallel \ + --moe-backend "$MOE_BACKEND" --kv-cache-dtype fp8 --api-server-count 1 \ + --max-num-seqs 64 --max-num-batched-tokens 4096 --max-model-len 8192 \ + --trust-remote-code "${MODE_ARGS[@]}" --host "$HOST" --port "$PORT" "${EXTRA[@]}" diff --git a/tests/e2e/case.py b/tests/e2e/case.py index 13fc4ef59..e04c16594 100644 --- a/tests/e2e/case.py +++ b/tests/e2e/case.py @@ -80,7 +80,7 @@ class Deployment: max_num_seqs: int | None = None max_num_batched_tokens: int | None = None compute_gate_on_attention: bool = False - connector_extra_config: tuple[tuple[str, int], ...] = () + connector_extra_config: tuple[tuple[str, int | str], ...] = () pin_attention_address: bool = False compilation_config: tuple[tuple[str, str | tuple[int, ...]], ...] = () shutdown_timeout_s: int = DEFAULT_SHUTDOWN_TIMEOUT_S diff --git a/tests/e2e/models/glm_5_3/cases.py b/tests/e2e/models/glm_5_3/cases.py new file mode 100644 index 000000000..0c2d83878 --- /dev/null +++ b/tests/e2e/models/glm_5_3/cases.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""GLM-5.3 MXFP4xFP8 CUDA cases on one 8-GPU node; IDs only at import. + +The native baseline runs DP8 with expert parallelism; the AFD cases run +DeepEP M2N 4A4F, the FFN executing the MXFP4 experts with vLLM's Marlin +repack. Both use the Marlin MoE kernel, so the baseline is the control. +""" + +from __future__ import annotations + +from dataclasses import replace + +from tests.e2e.case import ( + Case, + CaseEnvironment, + Gsm8kSettings, + Gsm8kTask, + Mrv2Check, + Phase, +) +from tests.e2e.presets import DECODE_GRAPH, EAGER + +SCENARIOS = ( + "baseline-eager", + "afd-deepep-m2n-eager-4a4f", + "afd-deepep-m2n-graph-4a4f", +) +GPU_COUNT = 8 +GSM8K_SAMPLES = 50 +GSM8K_THRESHOLD = 0.85 +CAPTURE_SIZE = 16 +MAX_MODEL_LEN = 4096 +COMMON_VLLM_ARGS = ( + f"--max-model-len={MAX_MODEL_LEN}", + "--kv-cache-dtype=fp8", + "--moe-backend=marlin", + "--enable-ep-weight-filter", + "--api-server-count=1", +) +# The FFN reads vLLM's Marlin repack of the MXFP4 experts. +M2N_EXTRA_CONFIG = (("expert_backend", "marlin"),) + + +def build_case(case_id: str, environment: CaseEnvironment) -> Case: + if environment.device_backend != "gpu": + raise ValueError("GLM-5.3 E2E requires the gpu backend") + profiles = { + "baseline-eager": (True, False), + "afd-deepep-m2n-eager-4a4f": (False, False), + "afd-deepep-m2n-graph-4a4f": (False, True), + } + baseline, graph = profiles[case_id] + execution = DECODE_GRAPH if graph else EAGER + half = GPU_COUNT // 2 + deployment = execution.deployment( + backend="gpu", + attention=(GPU_COUNT, 1) if baseline else (half, 1), + ffn=None if baseline else (half, 1), + connector=None if baseline else "DeepEPM2NAFDConnector", + compute_gate_on_attention=not baseline, + capture_size=CAPTURE_SIZE, + ) + deployment = ( + # The native baseline keeps vLLM's default prefill chunk. + replace(deployment, max_num_batched_tokens=None) + if baseline + else replace(deployment, connector_extra_config=M2N_EXTRA_CONFIG) + ) + return Case( + id=case_id, + model="GLM-5.3", + served_model_name_prefix="glm-5-3-afd", + deployment=deployment, + common_vllm_args=COMMON_VLLM_ARGS, + phases=( + Phase( + "evaluation", + Gsm8kTask( + Gsm8kSettings(sample_limit=GSM8K_SAMPLES, threshold=GSM8K_THRESHOLD) + ), + (Mrv2Check(("attention", "ffn"), graph=True, dbo=False),) + if graph + else (), + ), + ), + ) + + +if __name__ == "__main__": + from tests.e2e.runner import case_main + + raise SystemExit(case_main(build_case, SCENARIOS)) diff --git a/tests/e2e/models/glm_5_3/test_glm_5_3.py b/tests/e2e/models/glm_5_3/test_glm_5_3.py new file mode 100644 index 000000000..a270a1b03 --- /dev/null +++ b/tests/e2e/models/glm_5_3/test_glm_5_3.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CUDA GLM-5.3 MXFP4xFP8 E2E scenarios (one 8 x H200-class node).""" + +from __future__ import annotations + +import os +from collections.abc import Iterator +from pathlib import Path + +import pytest + +from tests.conftest import ( + download_dataset, + download_model, + preserve_environment_variable, + run_case_entry, +) +from tests.e2e.environment import model_case_command +from tests.e2e.models.glm_5_3.cases import GPU_COUNT, SCENARIOS + +GSM8K_DATASET_ID = "openai/gsm8k" +GSM8K_DATASET_CONFIG = "main" +GLM53_REPO_ID = "RedHatAI/GLM-5.3-MXFP4xFP8" +# Loading the 415 GB checkpoint from network storage can take most of an hour. +STARTUP_TIMEOUT_S = 3600 + + +def prepare_e2e_assets() -> None: + """Ensure GSM8K and the GLM-5.3 MXFP4xFP8 checkpoint are available.""" + download_dataset(GSM8K_DATASET_ID, GSM8K_DATASET_CONFIG) + if os.environ.get("AFD_E2E_BACKEND") != "gpu": + raise RuntimeError("GLM-5.3 E2E supports only the 'gpu' backend") + existing = os.environ.get("AFD_GPU_E2E_MODEL") + if existing: + print(f"[e2e] Using existing model path {Path(existing).expanduser()}") + return + os.environ["AFD_GPU_E2E_MODEL"] = str(download_model(GLM53_REPO_ID)) + + +@pytest.fixture(scope="module", autouse=True) +def _prepare_e2e_assets() -> Iterator[None]: + with preserve_environment_variable("AFD_GPU_E2E_MODEL"): + prepare_e2e_assets() + yield + + +@pytest.mark.e2e +@pytest.mark.parametrize("scenario", SCENARIOS, ids=SCENARIOS) +def test_glm_5_3(scenario: str, tmp_path: Path) -> None: + command = model_case_command( + "tests.e2e.models.glm_5_3.cases", + scenario, + tmp_path / scenario, + backend="gpu", + default_devices=tuple(str(device) for device in range(GPU_COUNT)), + ) + command += ["--startup-timeout", str(STARTUP_TIMEOUT_S)] + run_case_entry(command, tmp_path / scenario) diff --git a/tests/gpu/test_deepep_m2n_mxfp4_experts.py b/tests/gpu/test_deepep_m2n_mxfp4_experts.py new file mode 100644 index 000000000..15bc2f4e1 --- /dev/null +++ b/tests/gpu/test_deepep_m2n_mxfp4_experts.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""GPU parity of the DeepEP M2N ``marlin`` executor for MXFP4 experts. + +Random MXFP4 experts go through vLLM's own Marlin repack +(``prepare_moe_mxfp4_layer_for_marlin``, as ``Mxfp4MoEMethod`` does at load); +the executor's output on DeepEP-shaped received rows (contiguous expert +groups, unused tail rows) is compared with an independent dense FP32 +reference built from the unpacked MXFP4 values, with the SwiGLU clamp. +One GPU, no transport. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402 + ApplyMoEActivationConfig, +) +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import ( # noqa: E402 + prepare_moe_mxfp4_layer_for_marlin, +) +from vllm.scalar_type import scalar_types # noqa: E402 + +from afd_plugin.connectors.gpu.deepep_m2n.experts import DeepEPM2NExperts # noqa: E402 +from afd_plugin.connectors.gpu.deepep_m2n.mxfp4_experts import ( # noqa: E402 + MXFP4_GROUP_SIZE, + Mxfp4MarlinExpertWeights, + marlin_mxfp4_expected_tensors, +) + +pytestmark = pytest.mark.gpu + +HIDDEN = 1024 +INTERMEDIATE = 512 +NUM_LOCAL_EXPERTS = 4 +NUM_ROWS = 64 +# DeepSeek-V4.1's swiglu_limit. +CLAMP_LIMIT = 10.0 +# E8M0 exponents around 2^-4 keep activations O(1) and let x * 8 hit the clamp. +SCALE_EXPONENT_LOW = 121 +SCALE_EXPONENT_HIGH = 125 +E8M0_BIAS = 127 +E2M1_VALUES = (0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0) +# BF16 activations and intermediate, FP32 accumulation: relative L2 error. +MAX_REL_ERR = 2e-2 +SEED = 1234 + + +def _skip_without_gpu() -> None: + if not torch.cuda.is_available(): + pytest.skip("needs a CUDA GPU") + + +def _unpack_mxfp4(packed: Any, scales: Any) -> Any: + """FP32 values of ``[..., N, K / 2]`` packed E2M1 (low nibble first) with + ``[..., N, K / 32]`` E8M0 scales.""" + lut = torch.tensor(E2M1_VALUES, device=packed.device) + low, high = packed & 0x0F, packed >> 4 + nibbles = torch.stack((low, high), dim=-1).flatten(-2).long() + values = lut[nibbles & 0x07] * torch.where(nibbles >= 8, -1.0, 1.0) + scale = torch.exp2(scales.float() - E8M0_BIAS) + return values * scale.repeat_interleave(MXFP4_GROUP_SIZE, dim=-1) + + +class Mxfp4Experts: + """Random MXFP4 experts in checkpoint form, Marlin form and FP32 form.""" + + def __init__(self, gen: Any) -> None: + def packed(*shape: int) -> Any: + return torch.randint(0, 256, shape, dtype=torch.uint8, generator=gen) + + def scales(*shape: int) -> Any: + return torch.randint( + SCALE_EXPONENT_LOW, + SCALE_EXPONENT_HIGH, + shape, + dtype=torch.uint8, + generator=gen, + ) + + e, i, h = NUM_LOCAL_EXPERTS, INTERMEDIATE, HIDDEN + w13 = packed(e, 2 * i, h // 2).cuda() + w2 = packed(e, h, i // 2).cuda() + w13_scale = scales(e, 2 * i, h // MXFP4_GROUP_SIZE).cuda() + w2_scale = scales(e, h, i // MXFP4_GROUP_SIZE).cuda() + self.w13_fp32 = _unpack_mxfp4(w13, w13_scale) + self.w2_fp32 = _unpack_mxfp4(w2, w2_scale) + layer = SimpleNamespace(params_dtype=torch.bfloat16) + m_w13, m_w2, m_w13_scale, m_w2_scale, _, _ = prepare_moe_mxfp4_layer_for_marlin( + layer, w13, w2, w13_scale, w2_scale, None, None + ) + self.marlin = Mxfp4MarlinExpertWeights( + w13=m_w13, + w2=m_w2, + w13_scale=m_w13_scale, + w2_scale=m_w2_scale, + quant_type_id=scalar_types.float4_e2m1f.id, + activation_config=ApplyMoEActivationConfig(clamp_limit=CLAMP_LIMIT), + ) + + def reference( + self, + x: Any, + weights: Any, + psum: list[int], + clamp: float | None = CLAMP_LIMIT, + ) -> Any: + """``weight * w2 @ swiglu(w13 @ x)`` per row; rows past the groups are 0.""" + out = torch.zeros(x.shape, dtype=torch.float32, device=x.device) + start = 0 + for expert, end in enumerate(psum): + rows = x[start:end].float() + gate, up = (rows @ self.w13_fp32[expert].T).chunk(2, dim=-1) + if clamp is not None: + gate = gate.clamp(max=clamp) + up = up.clamp(-clamp, clamp) + hidden = torch.nn.functional.silu(gate) * up + out[start:end] = (hidden @ self.w2_fp32[expert].T) * weights[ + start:end, None + ] + start = end + return out + + +@pytest.fixture(scope="module") +def experts() -> Mxfp4Experts: + _skip_without_gpu() + return Mxfp4Experts(torch.Generator().manual_seed(SEED)) + + +def _rows(seed: int, x_scale: float = 1.0) -> tuple[Any, Any]: + gen = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randn(NUM_ROWS, HIDDEN, generator=gen, device="cuda") * x_scale + weights = torch.rand(NUM_ROWS, generator=gen, device="cuda") + return x.to(torch.bfloat16), weights + + +def _rel_err(actual: Any, expected: Any) -> float: + return float((actual.float() - expected).norm() / expected.norm().clamp_min(1e-6)) + + +def test_marlin_weights_have_the_shapes_the_handshake_expects(experts): + """The model-info check's expected shapes match vLLM's real Marlin repack.""" + expected = marlin_mxfp4_expected_tensors(NUM_LOCAL_EXPERTS, HIDDEN, INTERMEDIATE) + actual = { + name: [str(tensor.dtype), list(tensor.shape)] + for name, tensor in ( + ("w13", experts.marlin.w13), + ("w2", experts.marlin.w2), + ("w13_scale", experts.marlin.w13_scale), + ("w2_scale", experts.marlin.w2_scale), + ) + } + + assert actual == expected + + +@pytest.mark.parametrize( + "psum", + [ + [10, 10, 40, 52], # empty expert 1, unused tail rows 52..63 + [0, 0, 64, 64], # every row on expert 2, no tail + [16, 32, 48, 64], + ], + ids=["empty-expert-and-tail", "one-hot-expert", "balanced"], +) +def test_marlin_executor_matches_dense_mxfp4_reference(experts, psum): + """Each valid row is its expert's clamped SwiGLU MLP times its weight, and + unused tail rows come back as zeros. + + Why: the executor derives each row's expert from DeepEP's prefix sum on + the device; an off-by-one or a wrong sentinel would route rows to the + wrong expert or leak garbage into rows combine never reads (but a later + lane reuse might). + """ + x, weights = _rows(seed=sum(psum)) + executor = DeepEPM2NExperts("marlin", expert_alignment=1) + out = torch.empty_like(x) + + executor( + x, + None, + weights, + torch.tensor(psum, dtype=torch.int32, device="cuda"), + experts.marlin, + out=out, + ) + + expected = experts.reference(x, weights, psum) + assert _rel_err(out[: psum[-1]], expected[: psum[-1]]) < MAX_REL_ERR + assert torch.count_nonzero(out[psum[-1] :]) == 0 + + +def test_marlin_executor_applies_the_swiglu_clamp(experts): + """Inputs large enough to exceed the clamp match the clamped reference, + not the unclamped one. + + Why: V4.1 experts were trained with swiglu_limit; dropping the clamp + changes outputs only for large activations, which GSM8K might not catch. + """ + psum = [16, 32, 48, 64] + x, weights = _rows(seed=7, x_scale=8.0) + executor = DeepEPM2NExperts("marlin", expert_alignment=1) + out = torch.empty_like(x) + + executor( + x, + None, + weights, + torch.tensor(psum, dtype=torch.int32, device="cuda"), + experts.marlin, + out=out, + ) + + clamped = experts.reference(x, weights, psum) + unclamped = experts.reference(x, weights, psum, clamp=None) + assert _rel_err(unclamped, clamped) > 10 * MAX_REL_ERR + assert _rel_err(out, clamped) < MAX_REL_ERR + + +def test_marlin_executor_replays_in_a_cuda_graph(experts): + """One capture serves different rows and expert groups on replay. + + Why: the FFN captures each bucket once and replays it with new receive + counts; any host read of the prefix sum would freeze the first step's + routing into the graph. + """ + executor = DeepEPM2NExperts("marlin", expert_alignment=1) + static_x, static_w = _rows(seed=1) + static_psum = torch.tensor([16, 32, 48, 64], dtype=torch.int32, device="cuda") + static_out = torch.empty_like(static_x) + # Warm-up outside the capture (JIT, sentinel map), as the FFN does. + executor(static_x, None, static_w, static_psum, experts.marlin, out=static_out) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + executor(static_x, None, static_w, static_psum, experts.marlin, out=static_out) + + for seed, psum in ((2, [5, 30, 30, 50]), (3, [64, 64, 64, 64])): + x, weights = _rows(seed=seed) + static_x.copy_(x) + static_w.copy_(weights) + static_psum.copy_(torch.tensor(psum, dtype=torch.int32)) + graph.replay() + torch.cuda.synchronize() + + expected = experts.reference(x, weights, psum) + assert _rel_err(static_out[: psum[-1]], expected[: psum[-1]]) < MAX_REL_ERR + assert torch.count_nonzero(static_out[psum[-1] :]) == 0 diff --git a/tests/unit/compat/patches/test_config_dbo_spec_decode.py b/tests/unit/compat/patches/test_config_dbo_spec_decode.py new file mode 100644 index 000000000..8ca9a417b --- /dev/null +++ b/tests/unit/compat/patches/test_config_dbo_spec_decode.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Config gate of the vLLM #61464 backport (DBO with speculative decoding).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("vllm") + +from vllm.config import CUDAGraphMode # noqa: E402 + +from afd_plugin.compat.patches import config_dbo_spec_decode as gate # noqa: E402 + +MTP_DECODE_QUERY_LEN = 4 # one sampled token plus three draft tokens + + +def _vllm_config(*, decode_threshold: int = 2, prefill_threshold: int = 4097): + """The VllmConfig fields the pinned DBO gate reads.""" + return SimpleNamespace( + model_config=SimpleNamespace( + is_multimodal_model=False, is_encoder_decoder=False, is_hybrid=False + ), + parallel_config=SimpleNamespace( + pipeline_parallel_size=1, + decode_context_parallel_size=1, + prefill_context_parallel_size=1, + num_ubatches=2, + dbo_decode_token_threshold=decode_threshold, + dbo_prefill_token_threshold=prefill_threshold, + ), + lora_config=None, + speculative_config=SimpleNamespace(method="mtp"), + compilation_config=SimpleNamespace(cudagraph_mode=CUDAGraphMode.NONE), + is_mm_encoder_only=False, + uniform_decode_query_len=MTP_DECODE_QUERY_LEN, + ) + + +@pytest.fixture(autouse=True) +def v2_runner(monkeypatch): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") + + +def test_afd_config_accepts_dbo_with_speculative_decoding(monkeypatch): + monkeypatch.setattr(gate, "parse_optional_afd_config", lambda config: object()) + config = _vllm_config() + + unsupported = gate._get_dbo_unsupported_features(config) + + assert gate.SPEC_DECODE_DBO_FEATURE not in unsupported + # One request per microbatch: num_ubatches * decode query length. + assert config.parallel_config.dbo_decode_token_threshold == 8 + assert config.parallel_config.dbo_prefill_token_threshold == 4097 + + +def test_native_config_keeps_the_speculative_decoding_gate(monkeypatch): + monkeypatch.setattr(gate, "parse_optional_afd_config", lambda config: None) + config = _vllm_config() + + unsupported = gate._get_dbo_unsupported_features(config) + + assert gate.SPEC_DECODE_DBO_FEATURE in unsupported + assert config.parallel_config.dbo_decode_token_threshold == 2 diff --git a/tests/unit/compat/patches/test_gpu_deepseek_v32_dbo.py b/tests/unit/compat/patches/test_gpu_deepseek_v32_dbo.py new file mode 100644 index 000000000..e7b573316 --- /dev/null +++ b/tests/unit/compat/patches/test_gpu_deepseek_v32_dbo.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Per-micro-batch DSA top-k state for GLM-5.x DBO.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + import torch + from torch import nn +else: + torch = pytest.importorskip("torch") + nn = torch.nn + +pytest.importorskip("vllm") + +from vllm.model_executor.layers.attention.sparse_mla_attention import ( # noqa: E402 + SharedTopkIndicesBuffer, +) +from vllm.v1.attention.backends.mla.index_group import ( # noqa: E402 + SparseMLAIndexGroupBuilder, +) + +from afd_plugin.compat.patches import gpu_deepseek_v32_dbo as patch # noqa: E402 + +NUM_UBATCHES = 2 +MAX_TOKENS = 4 +TOPK = 8 +# GLM-5.x: an indexer layer followed by layers that reuse its top-k. +INDEX_PRODUCING_LAYERS = (True, False, False, True, False) + + +class _Impl(SharedTopkIndicesBuffer): + """Sparse MLA impl stand-in with the native top-k buffer mixin: an indexer + layer reads the buffer through its indexer, a reusing layer privately.""" + + def __init__(self, buffer, builder, indexer) -> None: + self.init_topk_indices_buffer(indexer, buffer) + self.index_group, self.index_group_index = builder.register_layer( + indexer is not None + ) + + +class _Indexer(nn.Module): + def __init__(self, buffer) -> None: + super().__init__() + self.topk_indices_buffer = buffer + self.indexer_op = _IndexerOp(buffer) + + +class _IndexerOp(nn.Module): + def __init__(self, buffer) -> None: + super().__init__() + self.topk_indices_buffer = buffer + + +class _Attention(nn.Module): + def __init__(self, buffer, builder, has_indexer: bool) -> None: + super().__init__() + self.topk_indices_buffer = buffer + self.indexer = _Indexer(buffer) if has_indexer else None + self.impl = _Impl(buffer, builder, self.indexer) + + +class _Model(nn.Module): + def __init__(self) -> None: + super().__init__() + self.topk_indices_buffer = torch.zeros(MAX_TOKENS, TOPK, dtype=torch.int32) + builder = SparseMLAIndexGroupBuilder(self.topk_indices_buffer, MAX_TOKENS) + self.layers = nn.ModuleList( + _Attention(self.topk_indices_buffer, builder, has_indexer) + for has_indexer in INDEX_PRODUCING_LAYERS + ) + + +def _topk_views(model: _Model) -> list[torch.Tensor]: + views = [model.topk_indices_buffer] + for layer in model.layers: + views += [layer.topk_indices_buffer, layer.impl.topk_indices_buffer] + views.append(layer.impl.index_group.logical_topk_indices) + if layer.indexer is not None: + views += [ + layer.indexer.topk_indices_buffer, + layer.indexer.indexer_op.topk_indices_buffer, + ] + return views + + +def test_each_micro_batch_resolves_its_own_topk_state(monkeypatch): + model = _Model() + native_buffer = model.topk_indices_buffer + native_groups = [layer.impl.index_group for layer in model.layers] + + patch.install_per_ubatch_topk_indices(model, NUM_UBATCHES) + + monkeypatch.setattr(patch, "dbo_current_ubatch_id", lambda: 0) + assert all(view is native_buffer for view in _topk_views(model)) + assert [layer.impl.index_group for layer in model.layers] == native_groups + + monkeypatch.setattr(patch, "dbo_current_ubatch_id", lambda: 1) + second = model.topk_indices_buffer + assert second is not native_buffer and second.shape == native_buffer.shape + assert all(view is second for view in _topk_views(model)) + groups = [layer.impl.index_group for layer in model.layers] + # Reusing layers keep sharing their indexer layer's group within a + # micro-batch, and no group state is shared with micro-batch 0. + assert groups[0] is groups[1] is groups[2] + assert groups[3] is groups[4] and groups[3] is not groups[0] + for group, native in zip(groups, native_groups, strict=True): + assert group.physical_topk_indices.data_ptr() != ( + native.physical_topk_indices.data_ptr() + ) + assert group.valid_topk_counts.data_ptr() != ( + native.valid_topk_counts.data_ptr() + ) + assert group.physical_topk_ready is not native.physical_topk_ready + assert group.side_stream is not native.side_stream + + +def test_one_micro_batch_writes_do_not_reach_the_other(monkeypatch): + model = _Model() + patch.install_per_ubatch_topk_indices(model, NUM_UBATCHES) + + monkeypatch.setattr(patch, "dbo_current_ubatch_id", lambda: 0) + model.layers[0].impl.topk_indices_buffer.fill_(7) + monkeypatch.setattr(patch, "dbo_current_ubatch_id", lambda: 1) + model.layers[0].indexer.indexer_op.topk_indices_buffer.fill_(9) + + # A reusing layer of micro-batch 0 still reads micro-batch 0's indices. + monkeypatch.setattr(patch, "dbo_current_ubatch_id", lambda: 0) + assert bool((model.layers[2].impl.topk_indices_buffer == 7).all()) + + +def test_an_unknown_holder_of_the_shared_state_fails_at_install(): + model = _Model() + model.layers[1].cached_topk = model.topk_indices_buffer + + with pytest.raises(RuntimeError, match="still shares DSA top-k state"): + patch.install_per_ubatch_topk_indices(model, NUM_UBATCHES) diff --git a/tests/unit/compat/patches/test_gpu_mrv2_dbo_spec_decode.py b/tests/unit/compat/patches/test_gpu_mrv2_dbo_spec_decode.py new file mode 100644 index 000000000..511c197e0 --- /dev/null +++ b/tests/unit/compat/patches/test_gpu_mrv2_dbo_spec_decode.py @@ -0,0 +1,233 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Worker side of the vLLM #61464 backport: split and draft-prefill rebuild.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +np = pytest.importorskip("numpy") +pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.compat.patches import gpu_mrv2_dbo_spec_decode as backport # noqa: E402 + +NUM_UBATCHES = 2 + + +def _input_batch(query_lens: list[int]) -> SimpleNamespace: + """The InputBatch fields the pinned split reads.""" + scheduled = np.array(query_lens, dtype=np.int32) + return SimpleNamespace( + num_scheduled_tokens=scheduled, + num_tokens_after_padding=int(scheduled.sum()), + num_reqs=len(query_lens), + num_reqs_after_padding=len(query_lens), + ) + + +@pytest.mark.parametrize("num_reqs", [3, 4, 5]) +def test_uniform_verification_batches_split_on_request_boundaries(num_reqs): + # MTP with 3 draft tokens verifies 4 tokens per request; a midpoint split + # cuts a request in two whenever the request count is odd. + query_len = 4 + slices = backport.create_ubatch_slices( + _input_batch([query_len] * num_reqs), NUM_UBATCHES, query_len + ) + + first, second = slices + assert first.token_slice.stop == query_len * (num_reqs // 2) + assert first.request_slice.stop == second.request_slice.start + assert second.token_slice.stop == query_len * num_reqs + + +@pytest.mark.parametrize( + "query_lens, uniform_token_count, midpoint", + [ + pytest.param([4] * 5, None, 10, id="no-agreed-count"), + pytest.param([4] * 5, 1, 10, id="single-token"), + pytest.param([4], 4, 2, id="one-request"), + ], +) +def test_split_keeps_the_midpoint_when_it_cannot_align( + query_lens, uniform_token_count, midpoint +): + first, _ = backport.create_ubatch_slices( + _input_batch(query_lens), NUM_UBATCHES, uniform_token_count + ) + + assert first.token_slice.stop == midpoint + + +class _GraphManager: + """The SpeculatorCudaGraphManager replay the backport wraps.""" + + def __init__(self, events: list[tuple]): + self.events = events + + def run_fullgraph(self, desc): + self.events.append(("replay", desc.cg_mode)) + + +def _speculator(*, use_ubatching: bool) -> SimpleNamespace: + """The TargetDependentARSpeculator state the rebuild reads.""" + prepared: list[tuple] = [] + + def prepare_attn( + input_batch, cg_mode, block_tables, slot_mappings, groups, kv, for_capture + ): + prepared.append((cg_mode, block_tables, slot_mappings, groups)) + if cg_mode == backport.CUDAGraphMode.FULL: + prepared.append(("refresh", for_capture)) + return {"draft_layer": "draft metadata"} + + return SimpleNamespace( + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace(use_ubatching=use_ubatching) + ), + block_tables=SimpleNamespace( + input_block_tables=[np.arange(8), np.arange(8)], + slot_mappings=np.arange(32).reshape(2, 16), + ), + model_state=SimpleNamespace(prepare_attn=prepare_attn), + attn_groups=["draft groups"], + kv_cache_config="kv", + prefill_cudagraph_manager=_GraphManager(prepared), + prepared=prepared, + ) + + +def _propose(speculator, attn_metadata, slot_mappings, **kwargs): + batch = SimpleNamespace(num_reqs_after_padding=3, num_tokens_after_padding=12) + tensors = [None] * 8 + return backport.propose( + speculator, batch, attn_metadata, slot_mappings, *tensors, **kwargs + ) + + +@pytest.fixture +def native_propose(monkeypatch): + received: list[tuple] = [] + + def original(self, input_batch, attn_metadata, slot_mappings, *args, **kwargs): + received.append((attn_metadata, slot_mappings)) + return "draft tokens" + + monkeypatch.setattr( + backport.TargetDependentARSpeculator, + backport._ORIGINAL_PROPOSE_ATTR, + original, + raising=False, + ) + monkeypatch.setattr( + backport, + "build_slot_mappings_by_layer", + lambda slot_mappings, kv: {"draft_layer": slot_mappings}, + ) + return received + + +def test_microbatched_step_rebuilds_draft_prefill_metadata(native_propose): + speculator = _speculator(use_ubatching=True) + + assert _propose(speculator, None, None) == "draft tokens" + + [(attn_metadata, slot_mappings)] = native_propose + assert attn_metadata == {"draft_layer": "draft metadata"} + # Built on the draft's own builders over the target's padded batch. + [(cg_mode, block_tables, target_slots, groups)] = speculator.prepared + assert cg_mode == backport.CUDAGraphMode.NONE + assert groups == ["draft groups"] + assert [len(table) for table in block_tables] == [3, 3] + assert target_slots.shape == (2, 12) + assert slot_mappings["draft_layer"] is target_slots + + +@pytest.mark.parametrize("dummy_run", [False, True], ids=["step", "dummy-run"]) +@pytest.mark.parametrize("target_metadata", [None, {"layer": "target"}]) +def test_full_draft_prefill_replay_refreshes_the_draft_builders_first( + monkeypatch, dummy_run, target_metadata +): + # Under DBO the draft-prefill graphs are captured on the draft builders, + # so each replay must first write this batch's metadata there. + speculator = _speculator(use_ubatching=True) + manager = speculator.prefill_cudagraph_manager + + def original(self, *args, **kwargs): + self.prefill_cudagraph_manager.run_fullgraph( + SimpleNamespace(cg_mode=backport.CUDAGraphMode.FULL) + ) + return "draft tokens" + + monkeypatch.setattr( + backport.TargetDependentARSpeculator, + backport._ORIGINAL_PROPOSE_ATTR, + original, + raising=False, + ) + monkeypatch.setattr( + backport, "build_slot_mappings_by_layer", lambda slots, kv: {"l": slots} + ) + + assert ( + _propose(speculator, target_metadata, None, dummy_run=dummy_run) + == "draft tokens" + ) + + refresh_and_replay = [ + event for event in speculator.prepared if event[0] in ("refresh", "replay") + ] + assert refresh_and_replay == [ + ("refresh", dummy_run), + ("replay", backport.CUDAGraphMode.FULL), + ] + # The replay seam is removed with the step. + assert "run_fullgraph" not in vars(manager) + + +@pytest.mark.parametrize("use_ubatching", [True, False], ids=["dbo", "no-dbo"]) +def test_capture_uses_the_draft_builders_only_under_dbo(monkeypatch, use_ubatching): + speculator = _speculator(use_ubatching=use_ubatching) + speculator.target_attn_groups = ["target groups"] + captured_with: list[list[str]] = [] + + def original(self): + captured_with.append(self.target_attn_groups) + + monkeypatch.setattr( + backport.TargetDependentARSpeculator, + backport._ORIGINAL_CAPTURE_ATTR, + original, + raising=False, + ) + + backport.capture(speculator) + + assert captured_with == [["draft groups"] if use_ubatching else ["target groups"]] + assert speculator.target_attn_groups == ["target groups"] + + +@pytest.mark.parametrize( + "use_ubatching, attn_metadata, kwargs", + [ + pytest.param(True, {"layer": "target"}, {}, id="target-metadata"), + pytest.param(False, None, {}, id="no-dbo"), + pytest.param( + True, + None, + {"dummy_run": True, "skip_attn_for_dummy_run": True}, + id="attention-skipped-dummy", + ), + ], +) +def test_other_steps_keep_the_native_draft_prefill_inputs( + native_propose, use_ubatching, attn_metadata, kwargs +): + speculator = _speculator(use_ubatching=use_ubatching) + + _propose(speculator, attn_metadata, "target slots", **kwargs) + + assert native_propose == [(attn_metadata, "target slots")] + assert speculator.prepared == [] diff --git a/tests/unit/compat/patches/test_gpu_mrv2_ubatching.py b/tests/unit/compat/patches/test_gpu_mrv2_ubatching.py index 7738bb6a4..952fc45a6 100644 --- a/tests/unit/compat/patches/test_gpu_mrv2_ubatching.py +++ b/tests/unit/compat/patches/test_gpu_mrv2_ubatching.py @@ -230,7 +230,9 @@ def test_dp1_thresholds_select_eager_microbatches_without_handshake( @pytest.mark.parametrize( "num_tokens, capture_tokens, expected_mode, expected_tokens", [ - (8, 16, _CUDAGraphMode.FULL, 16), + (9, 16, _CUDAGraphMode.FULL, 16), + # Exactly half: the second microbatch would be padding only. + (8, 16, _CUDAGraphMode.NONE, 8), (3, 16, _CUDAGraphMode.NONE, 3), (8, None, _CUDAGraphMode.NONE, 8), ], @@ -256,6 +258,87 @@ def test_dp1_uses_microbatch_graph_bucket_or_native_eager_padding_fallback( assert desc.num_ubatches == 2 +@pytest.mark.parametrize( + "min_tokens, padded_tokens, num_ubatches, padding_only", + [ + (8, 16, 2, True), + (9, 16, 2, False), + (7, 15, 2, True), + (8, 15, 2, False), + (10, 15, 3, True), + (11, 15, 3, False), + ], +) +def test_trailing_microbatch_padding_matches_the_midpoint_split( + patches, min_tokens, padded_tokens, num_ubatches, padding_only +): + # The last microbatch starts at (n - 1) * (padded // n). + assert ( + patches._trailing_ubatch_is_padding(min_tokens, padded_tokens, num_ubatches) + is padding_only + ) + + +@pytest.mark.parametrize( + "num_tokens, expected_mode, expected_tokens", + [(9, "FULL", 16), (8, "NONE", 8)], +) +def test_dp2_padding_sync_rejects_a_padding_only_trailing_microbatch( + monkeypatch, num_tokens, expected_mode, expected_tokens +): + pytest.importorskip("torch") + pytest.importorskip("vllm") + import inspect + + from vllm.config import CUDAGraphMode + from vllm.v1.worker.gpu import dp_utils + from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor + + from afd_plugin.compat.patches import gpu_mrv2_ubatching as patch + + assert inspect.signature(patch.sync_cudagraph_and_dp_padding) == ( + inspect.signature(dp_utils.sync_cudagraph_and_dp_padding) + ) + + class Manager: + def dispatch(self, num_reqs, num_tokens, uniform_token_count, **kwargs): + assert kwargs["num_ubatches"] == 2 + return BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=16, + num_reqs=16, + uniform_token_count=uniform_token_count, + num_ubatches=2, + ) + + # Every rank reports this rank's counts, so min(tokens) is num_tokens. + monkeypatch.setattr(patch, "should_skip_dp_coordination", lambda: True) + monkeypatch.setattr(patch, "get_dp_group", lambda: SimpleNamespace(cpu_group=None)) + desired = BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, num_tokens=num_tokens, num_reqs=num_tokens + ) + + descriptor, sync = patch.sync_cudagraph_and_dp_padding( + Manager(), + desired, + num_tokens, + num_tokens, + 1, + 2, + 0, + parallel_config=_parallel(dp_size=2), + allow_ubatching=True, + uniform_decode=True, + ) + + assert descriptor.cg_mode == CUDAGraphMode[expected_mode] + assert descriptor.num_tokens == expected_tokens + assert descriptor.num_ubatches == 2 + assert sync is not None + assert sync.eager is (expected_mode == "NONE") + assert sync.num_tokens_across_dp.tolist() == [expected_tokens] * 2 + + def test_profile_dispatch_does_not_select_a_graph(patches): manager = _GraphManager(16) desc, _ = patches.dispatch_cg_and_sync_dp( diff --git a/tests/unit/compat/patches/test_signatures.py b/tests/unit/compat/patches/test_signatures.py index d9ebad4bb..70ef69626 100644 --- a/tests/unit/compat/patches/test_signatures.py +++ b/tests/unit/compat/patches/test_signatures.py @@ -411,7 +411,9 @@ def test_patch_preserves_upstream_signature( scope = adapter._use_afd_ubatch_preparation( SimpleNamespace( ubatch_runner=binding, - vllm_config=SimpleNamespace(model_config=model), + vllm_config=SimpleNamespace( + model_config=model, speculative_config=None + ), ) ) else: diff --git a/tests/unit/config/test_validation.py b/tests/unit/config/test_validation.py index f61a7f7df..7a8d878ed 100644 --- a/tests/unit/config/test_validation.py +++ b/tests/unit/config/test_validation.py @@ -295,6 +295,8 @@ def test_routed_transport_preserves_runner_and_gate_policy( pytest.param("AFDQwen3MoeForCausalLM", True, id="afd-qwen3-moe"), pytest.param("DeepseekV2ForCausalLM", True, id="deepseek-v2"), pytest.param("AFDDeepseekV2ForCausalLM", True, id="afd-deepseek-v2"), + pytest.param("GlmMoeDsaForCausalLM", True, id="glm-moe-dsa"), + pytest.param("AFDGlmMoeDsaForCausalLM", True, id="afd-glm-moe-dsa"), pytest.param("DeepseekV3ForCausalLM", False, id="deepseek-v3"), pytest.param("AFDDeepseekV3ForCausalLM", False, id="afd-deepseek-v3"), pytest.param("Qwen3_5MoeForConditionalGeneration", False, id="qwen3-5-moe"), diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index c151249aa..74f93d221 100644 --- a/tests/unit/connectors/test_deepep_m2n_connector.py +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -13,6 +13,9 @@ torch = pytest.importorskip("torch") pytest.importorskip("vllm") +from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402 + ApplyMoEActivationConfig, +) from vllm.platforms.interface import set_assigned_physical_gpu_ids # noqa: E402 from afd_plugin.config import afd_config_from_mapping # noqa: E402 @@ -66,6 +69,10 @@ RankPlacement, plan_transport_layout, ) +from afd_plugin.connectors.gpu.deepep_m2n.mxfp4_experts import ( # noqa: E402 + Mxfp4MarlinExpertWeights, + marlin_mxfp4_expected_tensors, +) from afd_plugin.connectors.gpu.deepep_m2n.network import ( # noqa: E402 NetworkConfig, RankTopology, @@ -1098,6 +1105,42 @@ def test_handshake_accepts_fp8_block_weights_for_fp8_payload(): validate_handshake(infos, layout, NetworkConfig()) +@pytest.mark.parametrize( + "payload_format, weight_formats, error", + [ + ("bf16", ("mxfp4_marlin_w4a16", "mxfp4_marlin_w4a16"), None), + ( + "fp8_block128_fp32", + ("mxfp4_marlin_w4a16", "mxfp4_marlin_w4a16"), + "is deployed with .*weight_format", + ), + ("bf16", ("mxfp4_marlin_w4a16", "bf16"), "disagree on the expert weight"), + ], + ids=["bf16-payload", "fp8-payload", "mixed-formats"], +) +def test_handshake_feeds_mxfp4_marlin_experts_only_bf16_rows( + payload_format, weight_formats, error +): + """MXFP4 experts (W4A16 Marlin) take BF16 rows; an FP8 payload or FFN + ranks loading different formats fail the handshake on every rank. + + Why: Marlin W4A16 has no FP8 row input, and a mixed F would compute some + tokens with other numerics. + """ + layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) + infos = _handshake( + layout, {(rank, "payload_format"): payload_format for rank in range(4)} + ) + for rank, weight_format in enumerate(weight_formats): + infos[rank]["ffn_deployment"] = _ffn_deployment(layout, weight_format) + + if error is None: + validate_handshake(infos, layout, NetworkConfig()) + return + with pytest.raises(RuntimeError, match=error): + validate_handshake(infos, layout, NetworkConfig()) + + def test_handshake_requires_gate_on_attention(): """Both roles must route on Attention; a mix, or false everywhere, fails. @@ -1225,6 +1268,74 @@ def test_model_check_rejects_what_would_compute_wrong_rows( ) +def _mxfp4_local_experts( + expert_ids: range, *, w2_scale_shape: list[int] | None = None +) -> FFNLocalExperts: + """MXFP4 experts with the shapes vLLM's Marlin repack leaves (no data).""" + expected = marlin_mxfp4_expected_tensors( + len(expert_ids), SMALL_HIDDEN, SMALL_INTERMEDIATE + ) + shapes = {name: spec[1] for name, spec in expected.items()} + if w2_scale_shape is not None: + shapes["w2_scale"] = w2_scale_shape + + def e8m0(shape): + return torch.zeros(shape, dtype=torch.uint8).view(torch.float8_e8m0fnu) + + weights = Mxfp4MarlinExpertWeights( + w13=torch.zeros(shapes["w13"], dtype=torch.int32), + w2=torch.zeros(shapes["w2"], dtype=torch.int32), + w13_scale=e8m0(shapes["w13_scale"]), + w2_scale=e8m0(shapes["w2_scale"]), + quant_type_id=0, + activation_config=ApplyMoEActivationConfig(clamp_limit=10.0), + ) + return FFNLocalExperts(expert_ids=tuple(expert_ids), weights=weights) + + +def _mxfp4_model_infos(**kwargs) -> list: + infos = [ + local_model_info( + MOE_LAYERS, + {layer: _mxfp4_local_experts(experts, **kwargs) for layer in MOE_LAYERS}, + ) + for experts in (range(0, 4), range(4, 8)) + ] + infos.append(local_model_info(MOE_LAYERS, None)) + return infos + + +@pytest.mark.parametrize( + "infos_kwargs, payload_format, message", + [ + ({}, "bf16", None), + ({}, "fp8_block128_fp32", "do not match payload_format"), + ({"w2_scale_shape": [4, 2, SMALL_HIDDEN]}, "bf16", "w2_scale"), + ], + ids=["accepted", "fp8-payload", "scale-shape"], +) +def test_model_check_reads_mxfp4_marlin_expert_shapes( + infos_kwargs, payload_format, message +): + """MXFP4 experts are checked against the Marlin repack's shapes and fed + only BF16 rows; a wrong scale shape fails after loading. + + Why: Marlin-packed tensors keep their expert dimension but not the + FusedMoE shapes, so the BF16/FP8 rules would reject a valid deployment + and miss a wrong one. + """ + infos = _mxfp4_model_infos(**infos_kwargs) + if message is None: + validate_model_infos( + infos, SMALL_LAYOUT, hidden_size=SMALL_HIDDEN, payload_format=payload_format + ) + return + with pytest.raises(RuntimeError, match=message): + validate_model_infos( + infos, SMALL_LAYOUT, hidden_size=SMALL_HIDDEN, payload_format=payload_format + ) + + def test_model_check_requires_ffn_experts_once_anyone_reports(): infos = _model_infos() infos[1] = None diff --git a/tests/unit/fakes/glm_moe_dsa.py b/tests/unit/fakes/glm_moe_dsa.py new file mode 100644 index 000000000..34d1837d0 --- /dev/null +++ b/tests/unit/fakes/glm_moe_dsa.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""GLM-5.3 (``glm_moe_dsa``) config and checkpoint facts for CPU contract tests. + +Layer schedule, expert counts and the quantization config come from the +published ``zai-org/GLM-5.3`` and ``RedHatAI/GLM-5.3-MXFP4xFP8`` checkpoints; +widths are shrunk so the fixtures stay small. +""" + +from copy import deepcopy +from types import SimpleNamespace +from typing import Any + +GLM53_NUM_HIDDEN_LAYERS = 78 +GLM53_FIRST_K_DENSE_REPLACE = 3 +GLM53_N_ROUTED_EXPERTS = 256 +GLM53_INDEX_TOPK = 2048 +# The MTP layer follows the decoder layers in the checkpoint. +GLM53_MTP_LAYER = GLM53_NUM_HIDDEN_LAYERS +GLM53_MOE_LAYERS = tuple(range(GLM53_FIRST_K_DENSE_REPLACE, GLM53_NUM_HIDDEN_LAYERS)) + + +def glm53_hf_config() -> SimpleNamespace: + """GLM-5.3's layer schedule and MoE layout with small widths.""" + return SimpleNamespace( + model_type="glm_moe_dsa", + first_k_dense_replace=GLM53_FIRST_K_DENSE_REPLACE, + hidden_act="silu", + hidden_size=8, + index_topk=GLM53_INDEX_TOPK, + intermediate_size=16, + moe_intermediate_size=8, + moe_layer_freq=1, + n_group=1, + n_routed_experts=GLM53_N_ROUTED_EXPERTS, + n_shared_experts=1, + norm_topk_prob=True, + num_experts_per_tok=8, + num_hidden_layers=GLM53_NUM_HIDDEN_LAYERS, + num_nextn_predict_layers=1, + rms_norm_eps=1e-5, + routed_scaling_factor=2.5, + scoring_func="sigmoid", + tie_word_embeddings=False, + topk_group=1, + topk_method="noaux_tc", + vocab_size=32, + ) + + +_FP8_DYNAMIC_GROUP_ACTIVATIONS: dict[str, Any] = { + "dynamic": True, + "group_size": 128, + "num_bits": 8, + "strategy": "group", + "symmetric": True, + "type": "float", +} + +# ``RedHatAI/GLM-5.3-MXFP4xFP8`` quantization_config: MXFP4 weights with FP8 +# activations for every Linear not matched by the FP8 block group (routed +# experts and dense MLPs); FP8 block-128 for attention, shared experts and MTP. +GLM53_MXFP4XFP8_QUANTIZATION_CONFIG: dict[str, Any] = { + "config_groups": { + "fp8_block": { + "format": "float-quantized", + "input_activations": _FP8_DYNAMIC_GROUP_ACTIVATIONS, + "output_activations": None, + "targets": [ + "re:.*self_attn\\..*", + "re:.*shared_experts\\..*", + "re:.*\\.layers\\.78\\..*", + ], + "weights": { + "block_structure": [128, 128], + "dynamic": False, + "num_bits": 8, + "strategy": "block", + "symmetric": True, + "type": "float", + }, + }, + "mxfp4": { + "format": "mxfp4-pack-quantized", + "input_activations": _FP8_DYNAMIC_GROUP_ACTIVATIONS, + "output_activations": None, + "targets": ["Linear"], + "weights": { + "dynamic": False, + "group_size": 32, + "num_bits": 4, + "scale_dtype": "torch.uint8", + "strategy": "group", + "symmetric": True, + "type": "float", + }, + }, + }, + "format": "mixed-precision", + "ignore": [ + "re:.*mlp.gate$", + "re:.*lm_head", + "re:.*embed_tokens$", + "re:.*eh_proj$", + "re:.*self_attn.indexer.weights_proj$", + "re:.*self_attn.indexer.wk$", + "re:.*\\.head$", + ], + "kv_cache_scheme": None, + "quant_method": "compressed-tensors", + "quantization_status": "compressed", +} + + +def glm53_fp8_experts_quantization_config() -> dict[str, Any]: + """The MXFP4xFP8 config with the routed experts moved to FP8 block.""" + config = deepcopy(GLM53_MXFP4XFP8_QUANTIZATION_CONFIG) + config["config_groups"]["fp8_block"]["targets"].append("re:.*mlp\\.experts\\..*") + return config diff --git a/tests/unit/model_executor/models/test_deepseek_v32_construction.py b/tests/unit/model_executor/models/test_deepseek_v32_construction.py new file mode 100644 index 000000000..2275c8ac7 --- /dev/null +++ b/tests/unit/model_executor/models/test_deepseek_v32_construction.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Role-aware construction of the CUDA GLM-5.x (NVIDIA DSA) AFD wrapper.""" + +from __future__ import annotations + +from dataclasses import replace +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + import torch + from torch import nn +else: + torch = pytest.importorskip("torch") + nn = torch.nn + +pytest.importorskip("vllm") + +from afd_plugin.compat.patches import gpu_deepseek_v32_dbo as dbo_patch # noqa: E402 +from afd_plugin.config import DEEPEP_M2N_CONNECTOR, AFDConfig # noqa: E402 +from afd_plugin.model_executor.models import deepseek_v32 as adapter # noqa: E402 +from tests.unit.fakes.glm_moe_dsa import ( # noqa: E402 + GLM53_FIRST_K_DENSE_REPLACE, + GLM53_INDEX_TOPK, + GLM53_MOE_LAYERS, + GLM53_NUM_HIDDEN_LAYERS, + glm53_hf_config, +) +from tests.unit.fakes.model_layers import stage_type # noqa: E402 + +MAX_NUM_BATCHED_TOKENS = 8 +M2N_AFD_CONFIG = AFDConfig( + role="attention", + connector=DEEPEP_M2N_CONNECTOR, + compute_gate_on_attention=True, +) + + +class _FakeRoutedExpertsMoE(nn.Module): + """Native ``DeepseekV2MoE`` stand-in recording the config it was given.""" + + def __init__(self, *, config, **kwargs) -> None: + super().__init__() + self.n_shared_experts = config.n_shared_experts + self.weight = nn.Parameter(torch.empty(1)) + self.experts = SimpleNamespace(gate=object()) + + +def _vllm_config(**parallel_overrides): + parallel = { + "enable_dbo": False, + "eplb_config": SimpleNamespace(num_redundant_experts=0), + "pipeline_parallel_size": 1, + "tensor_parallel_size": 1, + "use_sequence_parallel_moe": False, + "use_ubatching": False, + "num_ubatches": 2, + # Decode-only DBO: prefill steps never reach the micro-batch threshold. + "dbo_prefill_token_threshold": MAX_NUM_BATCHED_TOKENS + 1, + **parallel_overrides, + } + return SimpleNamespace( + model_config=SimpleNamespace(hf_config=glm53_hf_config(), enforce_eager=True), + parallel_config=SimpleNamespace(**parallel), + quant_config=None, + scheduler_config=SimpleNamespace( + max_num_batched_tokens=MAX_NUM_BATCHED_TOKENS, max_num_seqs=4 + ), + speculative_config=None, + ) + + +@pytest.fixture +def build_model(monkeypatch): + """Build the AFD model for one role over native stage stand-ins.""" + calls: dict[str, list[str]] = { + kind: [] for kind in ("attention", "dense", "norm", "embed", "routed_shell") + } + + def bind(kind): + base = stage_type(kind) + + class _Bound(base): # type: ignore[valid-type,misc] + def __init__(self, *args, **kwargs): + super().__init__(calls, *args, **kwargs) + + return _Bound + + native = adapter.native + monkeypatch.setattr(native, "DeepseekV32Attention", bind("attention")) + monkeypatch.setattr(native, "DeepseekV2MLP", bind("dense")) + monkeypatch.setattr(native, "RMSNorm", bind("norm")) + monkeypatch.setattr(native, "make_input_embedding", bind("embed")) + monkeypatch.setattr(native, "DeepseekV2MoE", _FakeRoutedExpertsMoE) + monkeypatch.setattr(adapter, "AFDDeepseekV2RoutedExpertsMoE", bind("routed_shell")) + monkeypatch.setattr( + native, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + monkeypatch.setattr( + native, + "make_layers", + lambda count, layer_fn, prefix: ( + 0, + count, + nn.ModuleList(layer_fn(f"{prefix}.{idx}") for idx in range(count)), + ), + ) + monkeypatch.setattr(adapter, "current_platform", SimpleNamespace(device_type="cpu")) + + def build(role: str, **parallel_overrides): + afd_config = replace(M2N_AFD_CONFIG, role=role) + monkeypatch.setattr(adapter, "parse_afd_config", lambda *a, **k: afd_config) + model = adapter.AFDDeepseekV32Model( + vllm_config=_vllm_config(**parallel_overrides), prefix="model" + ) + return model, calls + + return build + + +def test_attention_builds_dsa_attention_dense_mlps_and_routed_shells(build_model): + model, calls = build_model("attention") + + assert len(calls["attention"]) == GLM53_NUM_HIDDEN_LAYERS + assert calls["dense"] == [ + f"model.layers.{idx}.mlp" for idx in range(GLM53_FIRST_K_DENSE_REPLACE) + ] + assert calls["routed_shell"] == [ + f"model.layers.{idx}.mlp" for idx in GLM53_MOE_LAYERS + ] + # The Attention role owns the DSA indexer and its shared top-k buffer. + assert tuple(model.topk_indices_buffer.shape) == ( + MAX_NUM_BATCHED_TOKENS, + GLM53_INDEX_TOPK, + ) + assert model.get_experts_layer_indices() == GLM53_MOE_LAYERS + + +def test_ffn_builds_only_routed_experts_without_shared_experts(build_model): + model, calls = build_model("ffn") + + assert calls["attention"] == [] + assert calls["dense"] == [] + assert calls["routed_shell"] == [] + assert model.topk_indices_buffer is None + for idx, layer in enumerate(model.layers): + assert isinstance(layer.self_attn, adapter.native.PPMissingLayer) + if idx < GLM53_FIRST_K_DENSE_REPLACE: + assert isinstance(layer.mlp, adapter.native.PPMissingLayer) + continue + # Attention runs the shared experts; FFN builds the routed ones only + # and routes externally. + assert layer.mlp.n_shared_experts is None + assert layer.mlp.experts.gate is None + assert model.config.n_shared_experts == 1, "the model config is not mutated" + assert model.get_experts_layer_indices() == GLM53_MOE_LAYERS + + +def test_attention_dbo_gives_each_micro_batch_its_own_topk_buffer( + build_model, monkeypatch +): + model, _ = build_model("attention", use_ubatching=True) + + monkeypatch.setattr(dbo_patch, "dbo_current_ubatch_id", lambda: 0) + first = model.topk_indices_buffer + monkeypatch.setattr(dbo_patch, "dbo_current_ubatch_id", lambda: 1) + second = model.topk_indices_buffer + + assert first is not second and first.shape == second.shape + + +MTP = SimpleNamespace(method="mtp") + + +@pytest.mark.parametrize( + "parallel_changes, speculative", + [ + pytest.param({}, None, id="m2n"), + pytest.param({"use_ubatching": True}, None, id="dbo"), + pytest.param({}, MTP, id="mtp-on-attention"), + pytest.param({"use_ubatching": True}, MTP, id="mtp-with-dbo"), + ], +) +def test_supported_deployments_are_accepted(parallel_changes, speculative): + vllm_config = _vllm_config(**parallel_changes) + vllm_config.speculative_config = speculative + + adapter.require_supported_deployment(vllm_config, M2N_AFD_CONFIG) + + +@pytest.mark.parametrize( + "afd_changes, parallel_changes, speculative, match", + [ + pytest.param( + {"connector": "P2pNcclAFDConnector"}, + {}, + None, + "supports only", + id="p2p-connector", + ), + pytest.param( + {"compute_gate_on_attention": False}, + {}, + None, + "compute_gate_on_attention", + id="gate-on-ffn", + ), + pytest.param( + {}, + {"tensor_parallel_size": 2}, + None, + "tensor_parallel_size", + id="attention-tp2", + ), + pytest.param({}, {"pipeline_parallel_size": 2}, None, "PP", id="pp2"), + pytest.param( + {}, + {"use_sequence_parallel_moe": True}, + None, + "sequence-parallel", + id="sp-moe", + ), + pytest.param({"role": "ffn"}, {}, MTP, "Attention role only", id="mtp-on-ffn"), + pytest.param( + {}, {}, SimpleNamespace(method="eagle"), "supports only", id="eagle" + ), + pytest.param( + {}, + {"use_ubatching": True, "dbo_prefill_token_threshold": 8}, + None, + "decode steps only", + id="dbo-prefill", + ), + ], +) +def test_unsupported_deployments_fail_at_construction( + afd_changes, parallel_changes, speculative, match +): + vllm_config = _vllm_config(**parallel_changes) + vllm_config.speculative_config = speculative + + with pytest.raises(RuntimeError, match=match): + adapter.require_supported_deployment( + vllm_config, replace(M2N_AFD_CONFIG, **afd_changes) + ) + + +@pytest.mark.parametrize("speculative", [None, MTP], ids=["dbo", "mtp-with-dbo"]) +def test_dbo_with_cuda_graphs_is_accepted(speculative): + # The padding-only trailing microbatch that the sparse FP8 MLA builder + # rejects is kept eager by the padding sync (gpu_mrv2_ubatching), and draft + # prefill after a microbatched step is rebuilt (gpu_mrv2_dbo_spec_decode). + vllm_config = _vllm_config(use_ubatching=True) + vllm_config.model_config.enforce_eager = False + vllm_config.speculative_config = speculative + + adapter.require_supported_deployment(vllm_config, M2N_AFD_CONFIG) diff --git a/tests/unit/model_executor/models/test_deepseek_v32_weight_policy.py b/tests/unit/model_executor/models/test_deepseek_v32_weight_policy.py new file mode 100644 index 000000000..82d51a811 --- /dev/null +++ b/tests/unit/model_executor/models/test_deepseek_v32_weight_policy.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Checkpoint role filtering of the CUDA GLM-5.x (NVIDIA DSA) AFD wrapper.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + import torch + from torch import nn +else: + torch = pytest.importorskip("torch") + nn = torch.nn + +pytest.importorskip("vllm") + +from afd_plugin.config import DEEPEP_M2N_CONNECTOR, AFDConfig # noqa: E402 +from afd_plugin.model_executor.models import deepseek_v32 as adapter # noqa: E402 +from tests.unit.fakes.glm_moe_dsa import GLM53_MTP_LAYER, glm53_hf_config # noqa: E402 + +ATTENTION = frozenset({"attention"}) +FFN = frozenset({"ffn"}) +BOTH = ATTENTION | FFN + +# Checkpoint module paths of GLM-5.3 and their AFD owner under DeepEP M2N: +# Attention runs the DSA attention (indexer included), the dense MLPs and the +# shared experts; FFN runs the routed experts; both load the gate. +GLM53_MODULE_ROLES = { + "model.embed_tokens": BOTH, + "model.norm": BOTH, + "lm_head": BOTH, + "model.layers.0.input_layernorm": BOTH, + "model.layers.0.self_attn.q_a_proj": ATTENTION, + "model.layers.0.self_attn.indexer.wq_b": ATTENTION, + "model.layers.0.self_attn.indexer.weights_proj": ATTENTION, + "model.layers.0.mlp.gate_proj": ATTENTION, + "model.layers.2.mlp.down_proj": ATTENTION, + "model.layers.3.post_attention_layernorm": BOTH, + "model.layers.3.mlp.gate": BOTH, + "model.layers.3.mlp.shared_experts.down_proj": ATTENTION, + "model.layers.3.mlp.experts.255.up_proj": FFN, + "model.layers.4.self_attn.kv_b_proj": ATTENTION, + f"model.layers.{GLM53_MTP_LAYER}.eh_proj": BOTH, + f"model.layers.{GLM53_MTP_LAYER}.self_attn.o_proj": ATTENTION, + f"model.layers.{GLM53_MTP_LAYER}.mlp.experts.0.gate_proj": FFN, +} +# Parameter suffixes per checkpoint: zai-org FP8 block scales, and the +# compressed-tensors MXFP4xFP8 export (packed MXFP4 experts and dense MLPs). +MXFP4_MODULES = (".experts.", "layers.0.mlp.", "layers.2.mlp.") +CHECKPOINT_SUFFIXES = { + "fp8": lambda module: ("weight", "weight_scale_inv"), + "mxfp4xfp8": lambda module: ( + ("weight_packed", "weight_scale") + if any(part in f"{module}." for part in MXFP4_MODULES) + else ("weight", "weight_scale") + ), +} +EXTRA_PARAMETERS = {"model.layers.3.mlp.gate": ("e_score_correction_bias",)} + + +def _checkpoint(fmt: str) -> list[tuple[str, torch.Tensor]]: + names = [] + for module in GLM53_MODULE_ROLES: + suffixes: tuple[str, ...] = CHECKPOINT_SUFFIXES[fmt](module) + if module.endswith(("norm", "embed_tokens", "lm_head", ".gate", "eh_proj")): + suffixes = ("weight",) + names += [f"{module}.{suffix}" for suffix in suffixes] + names += [f"{module}.{extra}" for extra in EXTRA_PARAMETERS.get(module, ())] + return [(name, torch.empty(1)) for name in names] + + +def _owner(name: str) -> frozenset[str]: + return GLM53_MODULE_ROLES[name.rsplit(".", 1)[0]] + + +@pytest.mark.parametrize("fmt", sorted(CHECKPOINT_SUFFIXES)) +@pytest.mark.parametrize("role", ["attention", "ffn"]) +def test_load_weights_passes_role_owned_tensors_to_native_loader_once( + monkeypatch, role, fmt +): + received: list[list[tuple[str, torch.Tensor]]] = [] + + def native_load_weights(self, weights): + received.append(list(weights)) + return {name for name, _ in received[-1]} + + monkeypatch.setattr( + adapter.native.DeepseekV32ForCausalLM, "load_weights", native_load_weights + ) + model = object.__new__(adapter.AFDGlmMoeDsaForCausalLM) + nn.Module.__init__(model) + model.afd_config = AFDConfig( + role=role, connector=DEEPEP_M2N_CONNECTOR, compute_gate_on_attention=True + ) + model.config = glm53_hf_config() + checkpoint = _checkpoint(fmt) + + loaded = model.load_weights(iter(checkpoint)) + + expected = [(name, tensor) for name, tensor in checkpoint if role in _owner(name)] + assert len(received) == 1, "the native loader runs once over one iterator" + assert [name for name, _ in received[0]] == [name for name, _ in expected] + # Original tensor objects reach the native loader unchanged. + assert all( + got is want for (_, got), (_, want) in zip(received[0], expected, strict=True) + ) + assert loaded == {name for name, _ in expected} diff --git a/tests/unit/package/test_package.py b/tests/unit/package/test_package.py index d025f73c7..41927713f 100644 --- a/tests/unit/package/test_package.py +++ b/tests/unit/package/test_package.py @@ -57,6 +57,7 @@ def register_model(name, path): for name in ( "async_dp_engine", "async_dp_forward_context", + "config_dbo_spec_decode", "config_validation", "engine_core", ): @@ -96,7 +97,11 @@ def test_register_afd_preserves_native_models_and_is_idempotent(registration): "AFDDeepseekV2ForCausalLM": deepseek + "AFDDeepseekV2ForCausalLM", "AFDDeepseekV3ForCausalLM": deepseek + "AFDDeepseekV3ForCausalLM", "AFDDeepseekV32ForCausalLM": deepseek + "AFDDeepseekV3ForCausalLM", - "AFDGlmMoeDsaForCausalLM": deepseek + "AFDGlmMoeDsaForCausalLM", + "AFDGlmMoeDsaForCausalLM": ( + deepseek + "AFDGlmMoeDsaForCausalLM" + if registration.device == "npu" + else "afd_plugin.model_executor.models.deepseek_v32:AFDGlmMoeDsaForCausalLM" + ), "AFDDeepseekV4ForCausalLM": ( "afd_plugin.model_executor.models.npu.deepseek_v4:AFDDeepseekV4ForCausalLM" if registration.device == "npu" diff --git a/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py b/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py index a28a99820..24dbaeb4d 100644 --- a/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py @@ -11,6 +11,10 @@ torch = pytest.importorskip("torch") pytest.importorskip("vllm") +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E402, E501 + CompressedTensorsConfig, +) + from afd_plugin.v1.worker.ffn_model_runner import GPUFFNModelRunner # noqa: E402 from afd_plugin.v1.worker.ffn_routed_experts_model_runner import ( # noqa: E402 GPUFFNRoutedExpertsModelRunner, @@ -20,6 +24,11 @@ from tests.unit.fakes.connector import ( # noqa: E402 FakeRoutedConnector, ) +from tests.unit.fakes.glm_moe_dsa import ( # noqa: E402 + GLM53_MXFP4XFP8_QUANTIZATION_CONFIG, + glm53_fp8_experts_quantization_config, + glm53_hf_config, +) MOE_LAYERS = (1, 3) NUM_LOCAL_EXPERTS = 2 @@ -216,8 +225,18 @@ def test_local_experts_reject_backends_that_reorder_weights(fp8, backend): local_experts_from_vllm(_routed_experts(fp8=fp8, backend=backend)) -def _kernel_config(moe_backend: str) -> SimpleNamespace: - return SimpleNamespace(kernel_config=SimpleNamespace(moe_backend=moe_backend)) +def _kernel_config(moe_backend: str, quant_config=None) -> SimpleNamespace: + return SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend=moe_backend), + quant_config=quant_config, + ) + + +def _deepseek_v4_quant_config(expert_dtype: str) -> SimpleNamespace: + """vLLM's DeepSeek-V4/V4.1 quantization config, reduced to what the pin reads.""" + return SimpleNamespace( + get_name=lambda: "deepseek_v4_fp8", expert_dtype=expert_dtype + ) def test_routed_ffn_pins_auto_moe_backend_to_triton(): @@ -241,3 +260,63 @@ def test_routed_ffn_rejects_other_explicit_moe_backends(configured): """ with pytest.raises(RuntimeError, match="moe-backend triton"): select_routed_moe_backend(_kernel_config(configured)) + + +@pytest.mark.parametrize( + "expert_dtype, configured, pinned", + [ + ("fp4", "auto", "marlin"), + ("fp4", "marlin", "marlin"), + ("fp8", "auto", "triton"), + ], + ids=["mxfp4-auto", "mxfp4-marlin", "fp8-experts-auto"], +) +def test_routed_ffn_pins_marlin_for_mxfp4_experts(expert_dtype, configured, pinned): + """DeepSeek-V4.1 MXFP4 experts pin Marlin; its FP8-expert variant keeps + Triton. + + Why: the M2N ``marlin`` executor reads vLLM's Marlin repack, and ``auto`` + may pick Humming on SM90, whose layout it cannot read. + """ + vllm_config = _kernel_config(configured, _deepseek_v4_quant_config(expert_dtype)) + + select_routed_moe_backend(vllm_config) + + assert vllm_config.kernel_config.moe_backend == pinned + + +def test_routed_ffn_rejects_triton_for_mxfp4_experts(): + with pytest.raises(RuntimeError, match="moe-backend marlin"): + select_routed_moe_backend( + _kernel_config("triton", _deepseek_v4_quant_config("fp4")) + ) + + +@pytest.mark.parametrize( + "quantization_config, configured, pinned", + [ + (GLM53_MXFP4XFP8_QUANTIZATION_CONFIG, "auto", "marlin"), + (GLM53_MXFP4XFP8_QUANTIZATION_CONFIG, "marlin", "marlin"), + (glm53_fp8_experts_quantization_config(), "auto", "triton"), + ], + ids=["mxfp4xfp8-auto", "mxfp4xfp8-marlin", "fp8-experts-auto"], +) +def test_routed_ffn_pins_backend_from_compressed_tensors_expert_scheme( + quantization_config, configured, pinned +): + """GLM-5.3 MXFP4xFP8 routed experts pin Marlin; FP8 experts keep Triton. + + Why: vLLM loads compressed-tensors W4A8 MXFP4 experts with an + ``Mxfp4MoEMethod`` subclass, and the M2N ``marlin`` executor reads only + vLLM's Marlin repack. The FP8 block group targets attention, shared + experts and the MTP layer by name, so only the experts' own scheme decides. + """ + vllm_config = SimpleNamespace( + kernel_config=SimpleNamespace(moe_backend=configured), + model_config=SimpleNamespace(hf_text_config=glm53_hf_config()), + quant_config=CompressedTensorsConfig.from_config(quantization_config), + ) + + select_routed_moe_backend(vllm_config) + + assert vllm_config.kernel_config.moe_backend == pinned