From 1374eca25b2360b240aebe5cd79d236c56996d09 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:36:08 +0200 Subject: [PATCH 01/13] [GPU] Run MXFP4 routed experts on DeepEP M2N with vLLM's Marlin repack MXFP4 checkpoints cannot use the M2N DeepGEMM executor, which reads BF16 or block-FP8 FusedMoE weights. Add a "marlin" expert backend that reads the parameters vLLM's Mxfp4MoEMethod leaves after the Marlin repack and runs fused_marlin_moe on the DeepEP-received rows (top_k=1 per row, the expert derived on the device from the prefix sum, graph safe). - mxfp4_experts.py: Mxfp4MarlinExpertWeights, the executor and the expected tensor specs for the model-info handshake. - The routed FFN pins --moe-backend marlin (auto -> marlin) when the checkpoint's experts are MXFP4 and keeps triton otherwise. - The handshake accepts mxfp4_marlin_w4a16 experts with the BF16 payload and checks that every FFN rank agrees on the weight format. - has_mxfp4_experts also recognizes compressed-tensors W4A8 MXFP4 experts (for example RedHatAI/GLM-5.3-MXFP4xFP8), resolved before load with the same scheme lookup and predicate vLLM uses to pick its MoE method. Ported from the DeepSeek-V4.1-Flash work (connector-generic parts only). Tests: backend pin from a real compressed-tensors GLM-5.3 config (MXFP4 experts pin marlin, FP8 experts keep triton), handshake and model-info checks, and tests/gpu Marlin parity against a dense MXFP4 reference (eager and CUDA graph), 6/6 on H100. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../connectors/gpu/deepep_m2n/connector.py | 94 +++++- .../connectors/gpu/deepep_m2n/experts.py | 59 +++- .../gpu/deepep_m2n/mxfp4_experts.py | 315 ++++++++++++++++++ .../worker/ffn_routed_experts_model_runner.py | 39 ++- tests/gpu/test_deepep_m2n_mxfp4_experts.py | 256 ++++++++++++++ .../connectors/test_deepep_m2n_connector.py | 111 ++++++ tests/unit/fakes/glm_moe_dsa.py | 119 +++++++ .../test_ffn_routed_experts_model_runner.py | 83 ++++- 8 files changed, 1052 insertions(+), 24 deletions(-) create mode 100644 afd_plugin/connectors/gpu/deepep_m2n/mxfp4_experts.py create mode 100644 tests/gpu/test_deepep_m2n_mxfp4_experts.py create mode 100644 tests/unit/fakes/glm_moe_dsa.py 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/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/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/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/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 From 720386af49d1d600483db1e8fd9fee027f4bd4cc Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:36:20 +0200 Subject: [PATCH 02/13] [GPU] GLM-5.x over DeepEP M2N on vLLM's NVIDIA DSA model vLLM 0.32 serves GlmMoeDsaForCausalLM on CUDA with vllm.models.deepseek_v32.nvidia (DSA backbone, fused embed/RMSNorm, MLP output reduced in the next layer's fused all-reduce + RMSNorm, breakable CUDA graphs), while AFD routed it to the DeepSeek V2-derived wrapper around the generic deepseek_v2 module that vLLM now uses only on XPU/CPU. Add a CUDA wrapper over the native DSA model that keeps its forward and loader and makes only construction and checkpoint loading role-aware: - Attention: DSA attention with the indexer top-k buffer, the dense MLPs, and per MoE layer the existing routed-experts shell (gate, native router, shared experts, dispatch/combine). - FFN: the routed experts only (native DeepseekV2MoE without shared experts, external routing). - One-pass checkpoint filter in front of the native loader. - Fail closed outside the first contract: DeepEP M2N only, Attention TP 1 (the native layer would all-reduce complete connector rows TP times), no PP/SP, no DBO, no MTP. GlmMoeDsaForCausalLM resolves to the new wrapper on CUDA and keeps the DeepSeek V2-derived wrapper on NPU; the M2N connector admits it. The idea of admitting GLM-5.3 on M2N follows an earlier teammate branch. Tests: construction on GLM-5.3's 78-layer schedule per role, the fail-closed table, role filtering of real GLM-5.3 FP8 and MXFP4xFP8 tensor names, registration and the M2N admission check. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- afd_plugin/__init__.py | 4 + .../model_executor/models/deepseek_v32.py | 324 ++++++++++++++++++ afd_plugin/validation.py | 1 + tests/unit/config/test_validation.py | 2 + .../models/test_deepseek_v32_construction.py | 208 +++++++++++ .../models/test_deepseek_v32_weight_policy.py | 110 ++++++ tests/unit/package/test_package.py | 6 +- 7 files changed, 654 insertions(+), 1 deletion(-) create mode 100644 afd_plugin/model_executor/models/deepseek_v32.py create mode 100644 tests/unit/model_executor/models/test_deepseek_v32_construction.py create mode 100644 tests/unit/model_executor/models/test_deepseek_v32_weight_policy.py diff --git a/afd_plugin/__init__.py b/afd_plugin/__init__.py index eeff9c467..bd2708568 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" ), } 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..b6fcbb98b --- /dev/null +++ b/afd_plugin/model_executor/models/deepseek_v32.py @@ -0,0 +1,324 @@ +# 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 first version is limited to the DeepEP M2N connector with Attention +TP = 1, no MTP and no DBO. 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.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" + + +def require_supported_deployment( + vllm_config: VllmConfig, + afd_config: AFDConfig, +) -> None: + """Fail closed outside the first GLM-5.x CUDA contract. + + 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.enable_dbo or parallel_config.use_ubatching: + raise RuntimeError("AFD GLM-5.x does not support DBO/ubatching") + if vllm_config.speculative_config is not None: + raise RuntimeError("AFD GLM-5.x does not support MTP/speculative decoding") + + +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 + + 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/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/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/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..7a6a93d0d --- /dev/null +++ b/tests/unit/model_executor/models/test_deepseek_v32_construction.py @@ -0,0 +1,208 @@ +# 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.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, + **parallel_overrides, + } + return SimpleNamespace( + model_config=SimpleNamespace(hf_config=glm53_hf_config()), + 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): + 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(), 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_supported_m2n_deployment_is_accepted(): + 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({}, {"enable_dbo": True}, None, "DBO", id="dbo"), + pytest.param({}, {"use_ubatching": True}, None, "DBO", id="ubatching"), + pytest.param({}, {}, SimpleNamespace(method="mtp"), "MTP", id="mtp"), + ], +) +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) + ) 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..ebccb6261 100644 --- a/tests/unit/package/test_package.py +++ b/tests/unit/package/test_package.py @@ -96,7 +96,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" From e5fcc64de391942665bd987a8305cdfae8efa714 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:46:03 +0200 Subject: [PATCH 03/13] [Doc][GPU] GLM-5.3 M2N recipes, native control, E2E case and docs - recipe/gpu/DeepEPM2NAFDConnector/glm_5_3: MXFP4xFP8 4A4F eager and graph on one 8-GPU node (Marlin experts, BF16 payload) and a block-128 FP8 4A4F,4A4F two-pod graph recipe (DeepGEMM, FP8 payload). - recipe/gpu/baseline/glm_5_3/native_dp8.sh: the native vLLM DP8 + EP control with the same checkpoint, KV cache dtype and Marlin MoE kernel. - tests/e2e/models/glm_5_3: native baseline and M2N 4A4F eager/graph GSM8K cases; Deployment.connector_extra_config accepts string values (expert_backend). - Connector guide, support matrix and model-integration design document list GLM-5.x, the marlin expert backend and the MXFP4 backend pin. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- docs/design/module/model_integration.md | 35 ++++++- docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md | 15 ++- docs/support.md | 3 +- recipe/gpu/DeepEPM2NAFDConnector/README.md | 4 +- .../DeepEPM2NAFDConnector/glm_5_3/README.md | 53 +++++++++++ .../4a4f_eager_dp4tp1.sh | 21 +++++ .../4a4f_graph_dp4tp1.sh | 22 +++++ .../multipod_4a4f_4a4f_fp8_graph.sh | 24 +++++ recipe/gpu/baseline/glm_5_3/native_dp8.sh | 42 +++++++++ tests/e2e/case.py | 2 +- tests/e2e/models/glm_5_3/cases.py | 93 +++++++++++++++++++ tests/e2e/models/glm_5_3/test_glm_5_3.py | 59 ++++++++++++ 12 files changed, 362 insertions(+), 11 deletions(-) create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md create mode 100755 recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_eager_dp4tp1.sh create mode 100755 recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_graph_dp4tp1.sh create mode 100755 recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/multipod_4a4f_4a4f_fp8_graph.sh create mode 100755 recipe/gpu/baseline/glm_5_3/native_dp8.sh create mode 100644 tests/e2e/models/glm_5_3/cases.py create mode 100644 tests/e2e/models/glm_5_3/test_glm_5_3.py diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 7af50ca4a..0163b2670 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,30 @@ 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, +sequence-parallel MoE, DBO/ubatching and MTP. 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..06cfff8c1 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -27,8 +27,14 @@ 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. +- 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. Attention TP > 1, MTP and DBO are rejected. -DeepSeek-V3, GLM-DSA and Qwen3.5/Qwen3.6 MoE are rejected with this connector. +DeepSeek-V3 and Qwen3.5/Qwen3.6 MoE are rejected with this connector. ## How it works @@ -139,7 +145,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 +192,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 +230,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/support.md b/docs/support.md index d45f44791..31eb4ef99 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; rejects P2P, Attention TP > 1, PP, SP MoE, DBO and MTP. 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). [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..df818aeb3 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md @@ -0,0 +1,53 @@ +# 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 +└── 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 | +| `--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. + +## Limitations + +- Attention TP1 only: the native DSA layer reduces the MLP output inside the + next layer's RMSNorm, so Attention TP > 1 is rejected. +- No MTP, no DBO, no P2P connector. 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/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) From 9488051f856f12f529c99a923281b89e9e512cf6 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Fri, 13 Nov 2026 15:15:30 +0200 Subject: [PATCH 04/13] [Doc] GLM-5.3 M2N GPU evidence and the afdEP NCCL requirement GLM-5.3 MXFP4xFP8 on one 8 x H200 node, vLLM 0.32.0, GSM8K-200 (flexible/strict): native DP8 + EP eager with Marlin 0.925/0.925, AFD DeepEP M2N 4A4F eager 0.930/0.930, graph (FULL_DECODE_ONLY, breakable CUDA graphs as native) 0.935/0.935. The stock vllm-openai v0.32.0 image's nvidia-nccl-cu13 2.29.7 lacks the GIN request API that afdEP-comm v1.1.0 compiles against; the recipe README asks for the 2.30.7 wheel. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- docs/support.md | 2 +- recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/docs/support.md b/docs/support.md index 31eb4ef99..a4e402d95 100644 --- a/docs/support.md +++ b/docs/support.md @@ -27,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), 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). [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 ([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/glm_5_3/README.md b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md index df818aeb3..e012fc309 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md @@ -45,6 +45,10 @@ The native vLLM control on the same node is 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 From 34fb33446a44391d7d234ab25aa8aef1baf40e6d Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Fri, 13 Nov 2026 17:06:12 +0200 Subject: [PATCH 05/13] [GPU] GLM-5.x MTP on Attention and decode-only DBO with per-micro-batch top-k MTP: the native DeepseekV32MTP draft runs on the Attention ranks with its own (EP-local) experts; the FFN serves only the target's routed experts. The wrapper accepts method "mtp" on the Attention role only. DBO: GLM-5.x runs the lightning indexer on one layer in four and later layers reuse its top-k from one model-level buffer and per-indexer SparseMLAIndexGroup workspaces. The routed-experts connector yields between decoder layers, so the other micro-batch overwrote that state (GSM8K-200 0.685 with garbled output). gpu_deepseek_v32_dbo gives each micro-batch its own buffer and group copies (fresh side stream/events), resolving every holder (model, DSA attention, indexer, indexer op, sparse MLA impl incl. its private buffer) by dbo_current_ubatch_id(), and fails at startup if any object still holds the shared state. It also leases the micro-batch compute SM budget to the DSA indexer planner (as the V4 DBO path does, vLLM #60732); without it the AFD world deadlocked. Remaining limits, fail closed: - micro-batched prefill still corrupts a fraction of requests (0.90 with garbled output vs 0.93): DBO requires dbo_prefill_token_threshold above max_num_batched_tokens; - with CUDA graphs a padding-only micro-batch trips vLLM 0.32's sparse FP8 MLA metadata builder (assert decode_query_len > 0): DBO requires eager; - vLLM ModelRunnerV2 rejects DBO with speculative decoding: no MTP + DBO. Tests: per-micro-batch resolution and isolation over vLLM's SharedTopkIndicesBuffer mixin and real index groups, missed-holder detection, the DBO install on construction and the updated deployment table. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../compat/patches/gpu_deepseek_v32_dbo.py | 232 ++++++++++++++++++ .../model_executor/models/deepseek_v32.py | 64 ++++- .../v1/worker/attention_model_runner_v2.py | 9 + .../patches/test_gpu_deepseek_v32_dbo.py | 142 +++++++++++ .../models/test_deepseek_v32_construction.py | 68 ++++- 5 files changed, 500 insertions(+), 15 deletions(-) create mode 100644 afd_plugin/compat/patches/gpu_deepseek_v32_dbo.py create mode 100644 tests/unit/compat/patches/test_gpu_deepseek_v32_dbo.py 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/model_executor/models/deepseek_v32.py b/afd_plugin/model_executor/models/deepseek_v32.py index b6fcbb98b..ab608a612 100644 --- a/afd_plugin/model_executor/models/deepseek_v32.py +++ b/afd_plugin/model_executor/models/deepseek_v32.py @@ -15,8 +15,11 @@ - FFN builds only each MoE layer's routed experts (native ``DeepseekV2MoE`` without shared experts); the connector runs them on the dispatched rows. -The first version is limited to the DeepEP M2N connector with Attention -TP = 1, no MTP and no DBO. Other deployments fail at construction. +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``). MTP together with DBO (rejected by vLLM +ModelRunnerV2) and other deployments fail at construction. """ from collections.abc import Iterable @@ -28,6 +31,9 @@ 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, @@ -39,14 +45,17 @@ _ATTENTION_ROLE = "attention" _FFN_ROLE = "ffn" +MTP_METHOD = "mtp" def require_supported_deployment( vllm_config: VllmConfig, afd_config: AFDConfig, ) -> None: - """Fail closed outside the first GLM-5.x CUDA contract. + """Fail closed outside the GLM-5.x CUDA contract. + MTP is accepted on the Attention role only, not together with DBO. DBO + is accepted eager and with decode-only micro-batching. 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 @@ -70,10 +79,46 @@ def require_supported_deployment( or parallel_config.use_sequence_parallel_moe ): raise RuntimeError("AFD GLM-5.x does not support PP or sequence-parallel MoE") - if parallel_config.enable_dbo or parallel_config.use_ubatching: - raise RuntimeError("AFD GLM-5.x does not support DBO/ubatching") - if vllm_config.speculative_config is not None: - raise RuntimeError("AFD GLM-5.x does not support MTP/speculative decoding") + 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" + ) + if parallel_config.use_ubatching and not vllm_config.model_config.enforce_eager: + # With CUDA graphs the padded split can leave a micro-batch holding + # only padding (one request with zero query tokens); vLLM 0.32's sparse + # FP8 MLA metadata builder asserts on it + # (FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode). + raise RuntimeError("AFD GLM-5.x DBO requires --enforce-eager") + 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}" + ) + if parallel_config.use_ubatching: + # vLLM 0.32 ModelRunnerV2, which AFD requires, rejects DBO together + # with speculative decoding at config validation. + raise RuntimeError( + "AFD GLM-5.x does not support MTP together with DBO: vLLM " + "ModelRunnerV2 does not support DBO with speculative decoding" + ) class AFDDeepseekV32DecoderLayer(native.DeepseekV32DecoderLayer): @@ -256,6 +301,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) # ### 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: diff --git a/afd_plugin/v1/worker/attention_model_runner_v2.py b/afd_plugin/v1/worker/attention_model_runner_v2.py index 3ff495685..76d0c4540 100644 --- a/afd_plugin/v1/worker/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/attention_model_runner_v2.py @@ -126,6 +126,15 @@ def _use_afd_ubatch_preparation( ) model_preparation = _use_v4_ubatch_preparation(runner) + 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) with model_preparation: state = vars(ubatch_runner) had_override = "prepare" in state 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/model_executor/models/test_deepseek_v32_construction.py b/tests/unit/model_executor/models/test_deepseek_v32_construction.py index 7a6a93d0d..52999d52c 100644 --- a/tests/unit/model_executor/models/test_deepseek_v32_construction.py +++ b/tests/unit/model_executor/models/test_deepseek_v32_construction.py @@ -19,6 +19,7 @@ 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 @@ -56,10 +57,13 @@ def _vllm_config(**parallel_overrides): "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()), + model_config=SimpleNamespace(hf_config=glm53_hf_config(), enforce_eager=True), parallel_config=SimpleNamespace(**parallel), quant_config=None, scheduler_config=SimpleNamespace( @@ -108,10 +112,12 @@ def __init__(self, *args, **kwargs): ) monkeypatch.setattr(adapter, "current_platform", SimpleNamespace(device_type="cpu")) - def build(role: str): + 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(), prefix="model") + model = adapter.AFDDeepseekV32Model( + vllm_config=_vllm_config(**parallel_overrides), prefix="model" + ) return model, calls return build @@ -155,8 +161,35 @@ def test_ffn_builds_only_routed_experts_without_shared_experts(build_model): assert model.get_experts_layer_indices() == GLM53_MOE_LAYERS -def test_supported_m2n_deployment_is_accepted(): - adapter.require_supported_deployment(_vllm_config(), M2N_AFD_CONFIG) +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"), + ], +) +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( @@ -191,9 +224,20 @@ def test_supported_m2n_deployment_is_accepted(): "sequence-parallel", id="sp-moe", ), - pytest.param({}, {"enable_dbo": True}, None, "DBO", id="dbo"), - pytest.param({}, {"use_ubatching": True}, None, "DBO", id="ubatching"), - pytest.param({}, {}, SimpleNamespace(method="mtp"), "MTP", id="mtp"), + pytest.param({"role": "ffn"}, {}, MTP, "Attention role only", id="mtp-on-ffn"), + pytest.param( + {}, {"use_ubatching": True}, MTP, "MTP together with DBO", id="mtp-dbo" + ), + 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( @@ -206,3 +250,11 @@ def test_unsupported_deployments_fail_at_construction( adapter.require_supported_deployment( vllm_config, replace(M2N_AFD_CONFIG, **afd_changes) ) + + +def test_dbo_with_cuda_graphs_fails_at_construction(): + vllm_config = _vllm_config(use_ubatching=True) + vllm_config.model_config.enforce_eager = False + + with pytest.raises(RuntimeError, match="enforce-eager"): + adapter.require_supported_deployment(vllm_config, M2N_AFD_CONFIG) From c02df4c1c4cd6efdc46c29659139dd22cfe71b04 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Fri, 13 Nov 2026 17:06:15 +0200 Subject: [PATCH 06/13] [Doc][GPU] GLM-5.3 MTP and DBO recipes, contract and GPU evidence - 4a4f_graph_mtp_dp4tp1.sh: MTP (3 speculative tokens by default) on the Attention role only. - 4a4f_dbo_dp4tp1.sh: eager DBO, decode steps only. - Connector guide, support matrix, model-integration design and recipe README describe the MTP/DBO contract and its fail-closed limits. GLM-5.3 MXFP4xFP8 on one 8 x H200 node, vLLM 0.32.0, GSM8K-200: MTP native graph 0.93 vs M2N 4A4F graph 0.92 (mean acceptance length 3.4-3.5 vs 3.4); DBO eager decode-only 0.93 (concurrency 32) and 0.91 (64) vs 0.93 without DBO. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- docs/design/module/model_integration.md | 12 ++++++-- docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md | 6 +++- docs/support.md | 4 +-- .../DeepEPM2NAFDConnector/glm_5_3/README.md | 7 ++++- .../4a4f_dbo_dp4tp1.sh | 28 +++++++++++++++++++ .../4a4f_graph_mtp_dp4tp1.sh | 27 ++++++++++++++++++ 6 files changed, 78 insertions(+), 6 deletions(-) create mode 100755 recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_dbo_dp4tp1.sh create mode 100755 recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_graph_mtp_dp4tp1.sh diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 0163b2670..ac54f7675 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -155,8 +155,16 @@ 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, -sequence-parallel MoE, DBO/ubatching and MTP. MXFP4 routed experts, including +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 eager, decode-only micro-batching; micro-batched +prefill, DBO with CUDA graphs, and MTP together with DBO (rejected by +ModelRunnerV2) fail at construction. 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. diff --git a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md index 06cfff8c1..94739a1ff 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -32,7 +32,11 @@ Supported models: 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. Attention TP > 1, MTP and DBO are rejected. + DeepGEMM. MTP runs natively on the Attention ranks (its draft experts + stay local). DBO runs eager with decode-only micro-batching, each + micro-batch keeping its own DSA top-k state. Attention TP > 1, MTP + together with DBO, micro-batched prefill and DBO with CUDA graphs are + rejected. DeepSeek-V3 and Qwen3.5/Qwen3.6 MoE are rejected with this connector. diff --git a/docs/support.md b/docs/support.md index a4e402d95..b4e65c134 100644 --- a/docs/support.md +++ b/docs/support.md @@ -9,7 +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; rejects P2P, Attention TP > 1, PP, SP MoE, DBO and MTP. Ascend keeps the DeepSeek V2-derived path. | +| 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 with decode-only micro-batching. Rejects P2P, Attention TP > 1, PP, SP MoE, MTP together with DBO, micro-batched prefill and DBO with CUDA graphs. 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. | @@ -27,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), 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 ([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) | +| 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/glm_5_3/README.md b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md index e012fc309..61fe203df 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md @@ -54,4 +54,9 @@ The native vLLM control on the same node is - Attention TP1 only: the native DSA layer reduces the MLP output inside the next layer's RMSNorm, so Attention TP > 1 is rejected. -- No MTP, no DBO, no P2P connector. +- 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`) is eager and micro-batches decode steps only; + micro-batched prefill and DBO with CUDA graphs are rejected, as is MTP + together with DBO (vLLM ModelRunnerV2). +- 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..8c0c6d85d --- /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 eager with DBO (two +# micro-batches) on one 8-GPU node: FFN DP4 x TP1 (Marlin MXFP4 experts) and +# Attention DP4 x TP1. GLM DBO is eager-only: with CUDA graphs a padding-only +# micro-batch trips vLLM 0.32's sparse FP8 MLA metadata builder. +# 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 +# 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_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 "$@" From ff445773204f30931efd62b2346a12e88fb8a8e1 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Fri, 13 Nov 2026 18:06:50 +0200 Subject: [PATCH 07/13] [GPU] Backport vLLM #61464 (MRV2 DBO with spec decode); GLM follow-ups report - config_dbo_spec_decode: for AFD configs, drop the MRV2 "DBO with speculative decoding" rejection and raise the DBO thresholds to one request per microbatch (#61464). - gpu_mrv2_dbo_spec_decode (worker): request-boundary split of uniform multi-token (verification) batches from the DP-agreed uniform count, and a draft-prefill pre-step that rebuilds full-batch metadata on the draft builders after a microbatched target forward (eager). - GLM-5.x accepts MTP together with eager decode-only DBO. - docs/reports/glm-5.3-afd-followups.md: open work (prefill DBO, graph DBO root cause, MTP+DBO GPU validation) and suggested upstream changes. Unit tests for the split, the gate and the pre-step pass (15); MTP+DBO is not GPU-validated yet. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- afd_plugin/__init__.py | 1 + .../compat/patches/config_dbo_spec_decode.py | 129 +++++++++ .../patches/gpu_mrv2_dbo_spec_decode.py | 249 ++++++++++++++++++ .../model_executor/models/deepseek_v32.py | 16 +- .../v1/worker/attention_model_runner_v2.py | 14 +- docs/reports/glm-5.3-afd-followups.md | 62 +++++ .../patches/test_config_dbo_spec_decode.py | 66 +++++ .../patches/test_gpu_mrv2_dbo_spec_decode.py | 154 +++++++++++ .../models/test_deepseek_v32_construction.py | 4 +- tests/unit/package/test_package.py | 1 + 10 files changed, 681 insertions(+), 15 deletions(-) create mode 100644 afd_plugin/compat/patches/config_dbo_spec_decode.py create mode 100644 afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py create mode 100644 docs/reports/glm-5.3-afd-followups.md create mode 100644 tests/unit/compat/patches/test_config_dbo_spec_decode.py create mode 100644 tests/unit/compat/patches/test_gpu_mrv2_dbo_spec_decode.py diff --git a/afd_plugin/__init__.py b/afd_plugin/__init__.py index bd2708568..2c6dabd5f 100644 --- a/afd_plugin/__init__.py +++ b/afd_plugin/__init__.py @@ -181,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_mrv2_dbo_spec_decode.py b/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py new file mode 100644 index 000000000..832708ebe --- /dev/null +++ b/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py @@ -0,0 +1,249 @@ +# 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. a draft-prefill pre-step that rebuilds full-batch metadata on the draft's + own builders after a microbatched target forward (eager only; FULL draft + prefill graphs are captured on the target's builders at this pin). + +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 contextmanager +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" + + +# 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 +) -> 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.""" + 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, + CUDAGraphMode.NONE, + block_tables, + target_slot_mappings, + speculator.attn_groups, + speculator.kv_cache_config, + ) + return attn_metadata, build_slot_mappings_by_layer( + target_slot_mappings, speculator.kv_cache_config + ) + + +# 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: rebuild draft-prefill metadata after a microbatched step. + if ( + attn_metadata is None + and self.vllm_config.parallel_config.use_ubatching + and not (dummy_run and skip_attn_for_dummy_run) + ): + attn_metadata, slot_mappings = _build_drafter_attn_metadata(self, input_batch) + # ### PATCH END: rebuild draft-prefill metadata after a microbatched step. + 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 pre-step once (worker process).""" + if hasattr(TargetDependentARSpeculator, _ORIGINAL_PROPOSE_ATTR): + return + setattr( + TargetDependentARSpeculator, + _ORIGINAL_PROPOSE_ATTR, + TargetDependentARSpeculator.propose, + ) + TargetDependentARSpeculator.propose = propose + + +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/model_executor/models/deepseek_v32.py b/afd_plugin/model_executor/models/deepseek_v32.py index ab608a612..53f4f9636 100644 --- a/afd_plugin/model_executor/models/deepseek_v32.py +++ b/afd_plugin/model_executor/models/deepseek_v32.py @@ -18,8 +18,9 @@ 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``). MTP together with DBO (rejected by vLLM -ModelRunnerV2) and other deployments fail at construction. +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 @@ -54,8 +55,8 @@ def require_supported_deployment( ) -> None: """Fail closed outside the GLM-5.x CUDA contract. - MTP is accepted on the Attention role only, not together with DBO. DBO - is accepted eager and with decode-only micro-batching. + MTP is accepted on the Attention role only. DBO is accepted eager and + with decode-only micro-batching, 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 @@ -112,13 +113,6 @@ def require_supported_deployment( f"AFD GLM-5.x supports only {MTP_METHOD!r} speculative decoding, " f"got {speculative_config.method!r}" ) - if parallel_config.use_ubatching: - # vLLM 0.32 ModelRunnerV2, which AFD requires, rejects DBO together - # with speculative decoding at config validation. - raise RuntimeError( - "AFD GLM-5.x does not support MTP together with DBO: vLLM " - "ModelRunnerV2 does not support DBO with speculative decoding" - ) class AFDDeepseekV32DecoderLayer(native.DeepseekV32DecoderLayer): diff --git a/afd_plugin/v1/worker/attention_model_runner_v2.py b/afd_plugin/v1/worker/attention_model_runner_v2.py index 76d0c4540..c59f7dddf 100644 --- a/afd_plugin/v1/worker/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/attention_model_runner_v2.py @@ -135,7 +135,19 @@ def _use_afd_ubatch_preparation( ) model_preparation = use_dsa_indexer_sm_budget(ubatch_runner) - with model_preparation: + 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/docs/reports/glm-5.3-afd-followups.md b/docs/reports/glm-5.3-afd-followups.md new file mode 100644 index 000000000..9c489541b --- /dev/null +++ b/docs/reports/glm-5.3-afd-followups.md @@ -0,0 +1,62 @@ +# 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.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) | + +## 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. Leading hypothesis: a prefill request straddling the + token-midpoint split. Diagnostics: D-B (prefix caching off) was running + when work stopped; next D-A (request-aligned split for all micro-batched + batches; needs a DP-agreed split, e.g. under Attention DP1 / `1A4F`), then + a per-layer DBO vs non-DBO hidden-state comparison. Until then the adapter + requires `--dbo-prefill-token-threshold` > `--max-num-batched-tokens`. +2. **Graph DBO.** Root cause found: 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`; a rank with + exactly half the padded tokens gets a padding-only second micro-batch whose + FULL metadata has only zero-length requests, and + `FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode` asserts + `decode_query_len > 0`. Fix: `<=` (copy of the native function in + `gpu_mrv2_ubatching.py`, plus AFD's DP1 mirror + `num_ubatches * num_tokens < ubatch_desc.num_tokens`). Not implemented yet; + then validate graph DBO and lift the eager-only rule. +3. **MTP + DBO.** 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 + metadata rebuild). Unit tests pass; GPU validation (MTP+DBO eager vs MTP + 0.92) not run yet. FULL draft-prefill graphs still capture on the target's + builders (needs #61464's capture-group switch with graph DBO). +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` | `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 | Prefill DBO corruption on GLM (under investigation) | Request-aligned split for prefill (RFC #50738 open question 2) | +| #61464 test coverage | E2E uses ngram only | Add MTP/EAGLE + DBO GSM8K, and a DSA model | +| 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 | 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_mrv2_dbo_spec_decode.py b/tests/unit/compat/patches/test_gpu_mrv2_dbo_spec_decode.py new file mode 100644 index 000000000..78f6a658a --- /dev/null +++ b/tests/unit/compat/patches/test_gpu_mrv2_dbo_spec_decode.py @@ -0,0 +1,154 @@ +# 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 + + +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): + prepared.append((cg_mode, block_tables, slot_mappings, groups)) + 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", + 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( + "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/model_executor/models/test_deepseek_v32_construction.py b/tests/unit/model_executor/models/test_deepseek_v32_construction.py index 52999d52c..c9a7f6084 100644 --- a/tests/unit/model_executor/models/test_deepseek_v32_construction.py +++ b/tests/unit/model_executor/models/test_deepseek_v32_construction.py @@ -183,6 +183,7 @@ def test_attention_dbo_gives_each_micro_batch_its_own_topk_buffer( 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): @@ -225,9 +226,6 @@ def test_supported_deployments_are_accepted(parallel_changes, speculative): id="sp-moe", ), pytest.param({"role": "ffn"}, {}, MTP, "Attention role only", id="mtp-on-ffn"), - pytest.param( - {}, {"use_ubatching": True}, MTP, "MTP together with DBO", id="mtp-dbo" - ), pytest.param( {}, {}, SimpleNamespace(method="eagle"), "supports only", id="eagle" ), diff --git a/tests/unit/package/test_package.py b/tests/unit/package/test_package.py index ebccb6261..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", ): From 0dcdff9d6fa2e9549370c2a2dd78d16e5d34868e Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Fri, 13 Nov 2026 18:10:31 +0200 Subject: [PATCH 08/13] [Doc] GLM-5.3 follow-ups: prefill-DBO prefix-caching diagnostic result Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- docs/reports/glm-5.3-afd-followups.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/reports/glm-5.3-afd-followups.md b/docs/reports/glm-5.3-afd-followups.md index 9c489541b..b04648003 100644 --- a/docs/reports/glm-5.3-afd-followups.md +++ b/docs/reports/glm-5.3-afd-followups.md @@ -20,8 +20,9 @@ Status of branch `glm-5.3-afd` (vLLM 0.32.0 pin). GPU evidence: one workspaces (per micro-batch), MLA prefill backend and chunked-context workspace (per builder), CuTe input alignment, the Marlin FFN executor, stream ordering. Leading hypothesis: a prefill request straddling the - token-midpoint split. Diagnostics: D-B (prefix caching off) was running - when work stopped; next D-A (request-aligned split for all micro-batched + token-midpoint split. 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 but do not explain all of it. Next: D-A (request-aligned split for all micro-batched batches; needs a DP-agreed split, e.g. under Attention DP1 / `1A4F`), then a per-layer DBO vs non-DBO hidden-state comparison. Until then the adapter requires `--dbo-prefill-token-threshold` > `--max-num-batched-tokens`. From 58da66ff938ce874253738e7511cb65bb9600879 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Fri, 13 Nov 2026 21:50:21 +0200 Subject: [PATCH 09/13] [GPU] GLM-5.x graph DBO: reject padding-only trailing micro-batch; MTP+DBO eager vLLM 0.32's sync_cudagraph_and_dp_padding falls back from a microbatched FULL graph only when 2 * min(tokens) < padded. With an even padded count a rank holding exactly half the tokens replays a graph whose second micro-batch is padding only, and the sparse FP8 MLA builder asserts decode_query_len > 0. A marked copy in gpu_mrv2_ubatching uses the exact midpoint-split condition, and AFD's DP1 dispatch uses the same check. GLM-5.3 MXFP4xFP8 4A4F GSM8K-200 on 8 x H200: graph DBO 0.93, eager MTP + DBO 0.94 (MTP alone 0.92). The eager-only DBO rule is lifted; MTP with graph DBO hit an illegal memory access (draft-prefill graphs still capture on the target's builders) and is rejected. Also: 4a4f_dbo_dp4tp1.sh takes MODE (graph default), new 4a4f_mtp_dbo_dp4tp1.sh, docs and the follow-ups report (prefill-DBO request-aligned split diagnostic removes the garbled answers), and the ubatch-prepare signature test fake gains speculative_config. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../compat/patches/gpu_mrv2_ubatching.py | 213 +++++++++++++++++- .../model_executor/models/deepseek_v32.py | 16 +- docs/design/module/model_integration.md | 12 +- docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md | 8 +- docs/reports/glm-5.3-afd-followups.md | 59 +++-- docs/support.md | 2 +- .../DeepEPM2NAFDConnector/glm_5_3/README.md | 14 +- .../4a4f_dbo_dp4tp1.sh | 8 +- .../4a4f_mtp_dbo_dp4tp1.sh | 32 +++ .../compat/patches/test_gpu_mrv2_ubatching.py | 85 ++++++- tests/unit/compat/patches/test_signatures.py | 4 +- .../models/test_deepseek_v32_construction.py | 14 +- 12 files changed, 413 insertions(+), 54 deletions(-) create mode 100755 recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_mtp_dbo_dp4tp1.sh 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/model_executor/models/deepseek_v32.py b/afd_plugin/model_executor/models/deepseek_v32.py index 53f4f9636..02a178e5d 100644 --- a/afd_plugin/model_executor/models/deepseek_v32.py +++ b/afd_plugin/model_executor/models/deepseek_v32.py @@ -55,8 +55,9 @@ def require_supported_deployment( ) -> None: """Fail closed outside the GLM-5.x CUDA contract. - MTP is accepted on the Attention role only. DBO is accepted eager and - with decode-only micro-batching, with or without MTP. + MTP is accepted on the Attention role only. DBO is accepted with + decode-only micro-batching, eager or with CUDA graphs; MTP with DBO is + accepted eager only. 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 @@ -92,12 +93,6 @@ def require_supported_deployment( "AFD GLM-5.x DBO micro-batches decode steps only: set " "--dbo-prefill-token-threshold above --max-num-batched-tokens" ) - if parallel_config.use_ubatching and not vllm_config.model_config.enforce_eager: - # With CUDA graphs the padded split can leave a micro-batch holding - # only padding (one request with zero query tokens); vLLM 0.32's sparse - # FP8 MLA metadata builder asserts on it - # (FlashMLASparseMetadataBuilder._build_fp8_separate_prefill_decode). - raise RuntimeError("AFD GLM-5.x DBO requires --enforce-eager") speculative_config = vllm_config.speculative_config if speculative_config is None: return @@ -113,6 +108,11 @@ def require_supported_deployment( f"AFD GLM-5.x supports only {MTP_METHOD!r} speculative decoding, " f"got {speculative_config.method!r}" ) + if parallel_config.use_ubatching and not vllm_config.model_config.enforce_eager: + # MTP + graph DBO hit an illegal memory access on GLM-5.3: the FULL + # draft-prefill graphs still capture on the target's attention + # builders (vLLM #61464's capture-group switch is not backported). + raise RuntimeError("AFD GLM-5.x MTP with DBO requires --enforce-eager") class AFDDeepseekV32DecoderLayer(native.DeepseekV32DecoderLayer): diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index ac54f7675..3ac8bf248 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -162,9 +162,15 @@ native draft model and its local experts. Under DBO, 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 eager, decode-only micro-batching; micro-batched -prefill, DBO with CUDA graphs, and MTP together with DBO (rejected by -ModelRunnerV2) fail at construction. MXFP4 routed experts, including +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 (request-aligned +split of verification steps, draft-prefill metadata rebuild) and is eager +only: its FULL draft-prefill graphs would still capture on the target's +builders. 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. diff --git a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md index 94739a1ff..0582e5f3d 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -33,10 +33,10 @@ Supported models: (`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 with decode-only micro-batching, each - micro-batch keeping its own DSA top-k state. Attention TP > 1, MTP - together with DBO, micro-batched prefill and DBO with CUDA graphs are - rejected. + stay local). DBO runs eager or with CUDA graphs (MTP with DBO eager + only), with decode-only micro-batching, each micro-batch keeping its own + DSA top-k state. Attention TP > 1, micro-batched prefill and MTP with + graph DBO are rejected. DeepSeek-V3 and Qwen3.5/Qwen3.6 MoE are rejected with this connector. diff --git a/docs/reports/glm-5.3-afd-followups.md b/docs/reports/glm-5.3-afd-followups.md index b04648003..c30545052 100644 --- a/docs/reports/glm-5.3-afd-followups.md +++ b/docs/reports/glm-5.3-afd-followups.md @@ -11,6 +11,8 @@ Status of branch `glm-5.3-afd` (vLLM 0.32.0 pin). GPU evidence: one | AFD DeepEP M2N 4A4F eager / graph | 0.930 / 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) | ## Open work @@ -19,29 +21,39 @@ Status of branch `glm-5.3-afd` (vLLM 0.32.0 pin). GPU evidence: one 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. Leading hypothesis: a prefill request straddling the - token-midpoint split. D-B (prefix caching off, prefill DBO on): 1 garbled + 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 but do not explain all of it. Next: D-A (request-aligned split for all micro-batched - batches; needs a DP-agreed split, e.g. under Attention DP1 / `1A4F`), then - a per-layer DBO vs non-DBO hidden-state comparison. Until then the adapter - requires `--dbo-prefill-token-threshold` > `--max-num-batched-tokens`. -2. **Graph DBO.** Root cause found: 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`; a rank with - exactly half the padded tokens gets a padding-only second micro-batch whose - FULL metadata has only zero-length requests, and + 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. Remaining before lifting the + rule: the accuracy gap to non-DBO eager (see the flip comparison below), and + steps with a single request (no boundary; they still straddle, and all DP + ranks must micro-batch alike). 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`. Fix: `<=` (copy of the native function in - `gpu_mrv2_ubatching.py`, plus AFD's DP1 mirror - `num_ubatches * num_tokens < ubatch_desc.num_tokens`). Not implemented yet; - then validate graph DBO and lift the eager-only rule. -3. **MTP + DBO.** 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 - metadata rebuild). Unit tests pass; GPU validation (MTP+DBO eager vs MTP - 0.92) not run yet. FULL draft-prefill graphs still capture on the target's - builders (needs #61464's capture-group switch with graph DBO). + `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: eager validated, graph rejected.** 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 metadata rebuild). Eager MTP + DBO 0.94 (MTP alone 0.92). + Graph MTP + DBO hit `illegal memory access` about ten minutes into GSM8K: + the FULL draft-prefill graphs still capture on the target's builders + (#61464's capture-group switch is not backported), so the adapter rejects + MTP with DBO unless `--enforce-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 @@ -56,8 +68,9 @@ Status of branch `glm-5.3-afd` (vLLM 0.32.0 pin). GPU evidence: one | --- | --- | --- | | 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` | `2 * min(tokens) <= padded` in `sync_cudagraph_and_dp_padding` | +| 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 | Prefill DBO corruption on GLM (under investigation) | Request-aligned split for prefill (RFC #50738 open question 2) | +| 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 | Land #61464's draft capture-group switch together with the split | | 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 | diff --git a/docs/support.md b/docs/support.md index b4e65c134..70c8132c0 100644 --- a/docs/support.md +++ b/docs/support.md @@ -9,7 +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 with decode-only micro-batching. Rejects P2P, Attention TP > 1, PP, SP MoE, MTP together with DBO, micro-batched prefill and DBO with CUDA graphs. Ascend keeps the DeepSeek V2-derived path. | +| 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 (MTP with DBO eager only). Rejects P2P, Attention TP > 1, PP, SP MoE, micro-batched prefill and MTP with graph DBO. 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. | diff --git a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md index 61fe203df..98c9cd67d 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md @@ -9,6 +9,9 @@ the model. 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, eager └── multipod_4a4f_4a4f_fp8_graph.sh # block-128 FP8, graph, two 8-GPU pods ``` @@ -26,7 +29,7 @@ The native vLLM control on the same node is | Setting | All recipes | | --- | --- | -| `--max-num-seqs` | 64 per Attention DP rank | +| `--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) | @@ -56,7 +59,10 @@ The native vLLM control on the same node is 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`) is eager and micro-batches decode steps only; - micro-batched prefill and DBO with CUDA graphs are rejected, as is MTP - together with DBO (vLLM ModelRunnerV2). +- 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. MTP with DBO is eager only; verification steps split on + request boundaries and the decode threshold is raised to one request per + micro-batch (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 index 8c0c6d85d..4f46f6447 100755 --- 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 @@ -2,10 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project # -# GLM-5.3 MXFP4xFP8 on DeepEPM2NAFDConnector, 4A4F eager with DBO (two +# 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. GLM DBO is eager-only: with CUDA graphs a padding-only -# micro-batch trips vLLM 0.32's sparse FP8 MLA metadata builder. +# 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 @@ -13,7 +12,8 @@ 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 +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 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..34616b37d --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/prefill_decode_colocation/4a4f_mtp_dbo_dp4tp1.sh @@ -0,0 +1,32 @@ +#!/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). Eager only: MTP with graph DBO is +# rejected (the draft-prefill graphs capture on the target's builders). +# 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 +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/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/model_executor/models/test_deepseek_v32_construction.py b/tests/unit/model_executor/models/test_deepseek_v32_construction.py index c9a7f6084..db3a938c3 100644 --- a/tests/unit/model_executor/models/test_deepseek_v32_construction.py +++ b/tests/unit/model_executor/models/test_deepseek_v32_construction.py @@ -250,9 +250,19 @@ def test_unsupported_deployments_fail_at_construction( ) -def test_dbo_with_cuda_graphs_fails_at_construction(): +def test_dbo_with_cuda_graphs_is_accepted(): + # The padding-only trailing microbatch that the sparse FP8 MLA builder + # rejects is now kept eager by the padding sync (gpu_mrv2_ubatching). vllm_config = _vllm_config(use_ubatching=True) vllm_config.model_config.enforce_eager = False - with pytest.raises(RuntimeError, match="enforce-eager"): + adapter.require_supported_deployment(vllm_config, M2N_AFD_CONFIG) + + +def test_mtp_with_graph_dbo_fails_at_construction(): + vllm_config = _vllm_config(use_ubatching=True) + vllm_config.model_config.enforce_eager = False + vllm_config.speculative_config = MTP + + with pytest.raises(RuntimeError, match="MTP with DBO requires --enforce-eager"): adapter.require_supported_deployment(vllm_config, M2N_AFD_CONFIG) From db217a3ee78d357aacd5f525054644564ea87cb0 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Fri, 13 Nov 2026 21:55:16 +0200 Subject: [PATCH 10/13] [Doc] GLM-5.3 follow-ups: prefill-DBO request-aligned split flip comparison Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- docs/reports/glm-5.3-afd-followups.md | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/docs/reports/glm-5.3-afd-followups.md b/docs/reports/glm-5.3-afd-followups.md index c30545052..2c5e319cb 100644 --- a/docs/reports/glm-5.3-afd-followups.md +++ b/docs/reports/glm-5.3-afd-followups.md @@ -8,7 +8,7 @@ Status of branch `glm-5.3-afd` (vLLM 0.32.0 pin). GPU evidence: one | Cell | GSM8K-200 | | --- | --- | | Native DP8 + EP eager (Marlin) | 0.925 | -| AFD DeepEP M2N 4A4F eager / graph | 0.930 / 0.935 | +| 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) | @@ -31,10 +31,14 @@ Status of branch `glm-5.3-afd` (vLLM 0.32.0 pin). GPU evidence: one 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. Remaining before lifting the - rule: the accuracy gap to non-DBO eager (see the flip comparison below), and - steps with a single request (no boundary; they still straddle, and all DP - ranks must micro-batch alike). Until then the adapter requires + 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 From 096f08b6d31d1c0477f42e0c5d203a5657f26122 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Fri, 13 Nov 2026 22:13:13 +0200 Subject: [PATCH 11/13] [GPU] MTP + graph DBO: run draft prefill eager after a microbatched step At vLLM 0.32 the FULL draft-prefill graphs are captured on the target's attention builders and replay the full-batch metadata the target forward last wrote there. A microbatched target forward writes only per-microbatch metadata, so the replay read stale sequence lengths and block tables (illegal memory access with GLM-5.3 MTP + graph DBO). After a microbatched step the speculator now dispatches only draft prefill eager, on the metadata the #61464 backport rebuilds on the draft builders; the target and draft decode keep their graphs. The adapter accepts MTP with graph DBO. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../patches/gpu_mrv2_dbo_spec_decode.py | 78 +++++++++++++------ .../model_executor/models/deepseek_v32.py | 8 +- .../4a4f_mtp_dbo_dp4tp1.sh | 7 +- .../patches/test_gpu_mrv2_dbo_spec_decode.py | 38 +++++++++ .../models/test_deepseek_v32_construction.py | 16 ++-- 5 files changed, 104 insertions(+), 43 deletions(-) diff --git a/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py b/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py index 832708ebe..4825c81eb 100644 --- a/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py +++ b/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py @@ -29,7 +29,7 @@ """ from collections.abc import Callable, Iterator -from contextlib import contextmanager +from contextlib import AbstractContextManager, contextmanager, nullcontext from typing import Any import torch @@ -40,6 +40,9 @@ 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 import ( + speculator as speculator_module, +) from vllm.v1.worker.gpu.spec_decode.target_dependent_ar.speculator import ( TargetDependentARSpeculator, ) @@ -165,6 +168,34 @@ def _build_drafter_attn_metadata( ) +@contextmanager +def _eager_draft_prefill(speculator: TargetDependentARSpeculator) -> Iterator[None]: + """Run this step's draft prefill eager; draft decode keeps its graphs. + + At this pin the FULL draft-prefill graphs are captured on the target's + attention builders and replay whatever full-batch metadata the target + forward last wrote there. A microbatched target forward writes none, so a + replay would read stale sequence lengths and block tables (illegal memory + access with MTP + graph DBO). Every DP rank microbatched this step, so all + of them take this eager path and agree without a new collective. + """ + native_dispatch = speculator_module.dispatch_cg_and_sync_dp + prefill_manager = speculator.prefill_cudagraph_manager + + def dispatch_cg_and_sync_dp( + cudagraph_manager: Any, *args: Any, **kwargs: Any + ) -> tuple[BatchExecutionDescriptor, DPSyncState | None]: + if cudagraph_manager is prefill_manager: + kwargs["need_eager"] = True + return native_dispatch(cudagraph_manager, *args, **kwargs) + + speculator_module.dispatch_cg_and_sync_dp = dispatch_cg_and_sync_dp + try: + yield + finally: + speculator_module.dispatch_cg_and_sync_dp = native_dispatch + + # Upstream source: vLLM v0.32.0 (5df63a7203cdb9328d824a27b682e77c7b9928be), # vllm/v1/worker/gpu/spec_decode/target_dependent_ar/speculator.py:: # TargetDependentARSpeculator.propose. @@ -199,33 +230,36 @@ def propose( TargetDependentARSpeculator, _ORIGINAL_PROPOSE_ATTR ) # ### PATCH START: rebuild draft-prefill metadata after a microbatched step. + draft_prefill_scope: AbstractContextManager[None] = nullcontext() if ( attn_metadata is None and self.vllm_config.parallel_config.use_ubatching and not (dummy_run and skip_attn_for_dummy_run) ): attn_metadata, slot_mappings = _build_drafter_attn_metadata(self, input_batch) - # ### PATCH END: rebuild draft-prefill metadata after a microbatched step. - 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, - ) + draft_prefill_scope = _eager_draft_prefill(self) + with draft_prefill_scope: + # ### PATCH END: rebuild draft-prefill metadata after a microbatched step. + 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: diff --git a/afd_plugin/model_executor/models/deepseek_v32.py b/afd_plugin/model_executor/models/deepseek_v32.py index 02a178e5d..b5bb40f3e 100644 --- a/afd_plugin/model_executor/models/deepseek_v32.py +++ b/afd_plugin/model_executor/models/deepseek_v32.py @@ -56,8 +56,7 @@ def require_supported_deployment( """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; MTP with DBO is - accepted eager only. + 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 @@ -108,11 +107,6 @@ def require_supported_deployment( f"AFD GLM-5.x supports only {MTP_METHOD!r} speculative decoding, " f"got {speculative_config.method!r}" ) - if parallel_config.use_ubatching and not vllm_config.model_config.enforce_eager: - # MTP + graph DBO hit an illegal memory access on GLM-5.3: the FULL - # draft-prefill graphs still capture on the target's attention - # builders (vLLM #61464's capture-group switch is not backported). - raise RuntimeError("AFD GLM-5.x MTP with DBO requires --enforce-eager") class AFDDeepseekV32DecoderLayer(native.DeepseekV32DecoderLayer): 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 index 34616b37d..75319efac 100755 --- 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 @@ -5,8 +5,7 @@ # 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). Eager only: MTP with graph DBO is -# rejected (the draft-prefill graphs capture on the target's builders). +# (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 @@ -14,7 +13,9 @@ 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 +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 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 index 78f6a658a..7adda1ce4 100644 --- a/tests/unit/compat/patches/test_gpu_mrv2_dbo_spec_decode.py +++ b/tests/unit/compat/patches/test_gpu_mrv2_dbo_spec_decode.py @@ -80,6 +80,7 @@ def prepare_attn(input_batch, cg_mode, block_tables, slot_mappings, groups, kv): model_state=SimpleNamespace(prepare_attn=prepare_attn), attn_groups=["draft groups"], kv_cache_config="kv", + prefill_cudagraph_manager="prefill graphs", prepared=prepared, ) @@ -130,6 +131,43 @@ def test_microbatched_step_rebuilds_draft_prefill_metadata(native_propose): assert slot_mappings["draft_layer"] is target_slots +def test_microbatched_step_runs_only_draft_prefill_eager(monkeypatch): + # FULL draft-prefill graphs replay the target builders' full-batch + # metadata, which a microbatched target forward never writes. + speculator = _speculator(use_ubatching=True) + dispatched: list[tuple[str, bool]] = [] + + def fake_native_dispatch(manager, *args, need_eager=False, **kwargs): + dispatched.append((manager, need_eager)) + return "descriptor", None + + def original(self, *args, **kwargs): + # Native propose dispatches draft prefill, then draft decode. + dispatch = backport.speculator_module.dispatch_cg_and_sync_dp + dispatch("prefill graphs", 1, 12, 4, 1, 0, need_eager=False) + dispatch("decode graphs", 1, 3, 1, 1, 0, need_eager=False) + 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} + ) + monkeypatch.setattr( + backport.speculator_module, "dispatch_cg_and_sync_dp", fake_native_dispatch + ) + + assert _propose(speculator, None, None) == "draft tokens" + + assert dispatched == [("prefill graphs", True), ("decode graphs", False)] + # The scope restores the module binding it replaced. + assert backport.speculator_module.dispatch_cg_and_sync_dp is fake_native_dispatch + + @pytest.mark.parametrize( "use_ubatching, attn_metadata, kwargs", [ diff --git a/tests/unit/model_executor/models/test_deepseek_v32_construction.py b/tests/unit/model_executor/models/test_deepseek_v32_construction.py index db3a938c3..2275c8ac7 100644 --- a/tests/unit/model_executor/models/test_deepseek_v32_construction.py +++ b/tests/unit/model_executor/models/test_deepseek_v32_construction.py @@ -250,19 +250,13 @@ def test_unsupported_deployments_fail_at_construction( ) -def test_dbo_with_cuda_graphs_is_accepted(): +@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 now kept eager by the padding sync (gpu_mrv2_ubatching). + # 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) - - -def test_mtp_with_graph_dbo_fails_at_construction(): - vllm_config = _vllm_config(use_ubatching=True) - vllm_config.model_config.enforce_eager = False - vllm_config.speculative_config = MTP - - with pytest.raises(RuntimeError, match="MTP with DBO requires --enforce-eager"): - adapter.require_supported_deployment(vllm_config, M2N_AFD_CONFIG) From ea3c71de5680ff77faad11c07b4e75b841af9542 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Fri, 13 Nov 2026 22:24:24 +0200 Subject: [PATCH 12/13] [GPU] MTP + graph DBO: draft prefill graphs on the draft builders (#61464) Complete the #61464 backport for CUDA graphs: under DBO the FULL draft-prefill graphs are captured on the draft's own attention builders (scoped swap of target_attn_groups around TargetDependentARSpeculator .capture), and each draft-prefill replay first refreshes those builders with this batch's full-batch metadata (scoped wrap of the prefill manager's run_fullgraph within propose). This replaces the eager draft prefill fallback, which confirmed the cause. GLM-5.3 MXFP4xFP8 4A4F graph MTP + DBO GSM8K-200 on 8 x H200: 0.95, no errors, acceptance 3.30-3.69 (eager-draft-prefill fallback 0.92, eager MTP + DBO 0.94, MTP alone 0.92). Docs and the follow-ups report updated. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- .../patches/gpu_mrv2_dbo_spec_decode.py | 114 ++++++++++++------ docs/design/module/model_integration.md | 9 +- docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md | 7 +- docs/reports/glm-5.3-afd-followups.md | 25 ++-- docs/support.md | 2 +- .../DeepEPM2NAFDConnector/glm_5_3/README.md | 9 +- .../patches/test_gpu_mrv2_dbo_spec_decode.py | 79 +++++++++--- 7 files changed, 167 insertions(+), 78 deletions(-) diff --git a/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py b/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py index 4825c81eb..c00206112 100644 --- a/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py +++ b/afd_plugin/compat/patches/gpu_mrv2_dbo_spec_decode.py @@ -19,9 +19,12 @@ 1. a request-boundary split of uniform multi-token batches, fed the uniform token count the target dispatch agreed across DP ranks; -2. a draft-prefill pre-step that rebuilds full-batch metadata on the draft's - own builders after a microbatched target forward (eager only; FULL draft - prefill graphs are captured on the target's builders at this pin). +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``. @@ -40,15 +43,13 @@ 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 import ( - speculator as speculator_module, -) 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), @@ -143,11 +144,15 @@ def split(input_batch: InputBatch, num_ubatches: int) -> UBatchSlices: def _build_drafter_attn_metadata( - speculator: TargetDependentARSpeculator, input_batch: InputBatch + 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.""" + 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 @@ -157,11 +162,12 @@ def _build_drafter_attn_metadata( ] attn_metadata = speculator.model_state.prepare_attn( input_batch, - CUDAGraphMode.NONE, + 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 @@ -169,31 +175,57 @@ def _build_drafter_attn_metadata( @contextmanager -def _eager_draft_prefill(speculator: TargetDependentARSpeculator) -> Iterator[None]: - """Run this step's draft prefill eager; draft decode keeps its graphs. - - At this pin the FULL draft-prefill graphs are captured on the target's - attention builders and replay whatever full-batch metadata the target - forward last wrote there. A microbatched target forward writes none, so a - replay would read stale sequence lengths and block tables (illegal memory - access with MTP + graph DBO). Every DP rank microbatched this step, so all - of them take this eager path and agree without a new collective. +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``. """ - native_dispatch = speculator_module.dispatch_cg_and_sync_dp - prefill_manager = speculator.prefill_cudagraph_manager + manager = speculator.prefill_cudagraph_manager + native_run_fullgraph = manager.run_fullgraph - def dispatch_cg_and_sync_dp( - cudagraph_manager: Any, *args: Any, **kwargs: Any - ) -> tuple[BatchExecutionDescriptor, DPSyncState | None]: - if cudagraph_manager is prefill_manager: - kwargs["need_eager"] = True - return native_dispatch(cudagraph_manager, *args, **kwargs) + 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) - speculator_module.dispatch_cg_and_sync_dp = dispatch_cg_and_sync_dp + manager.run_fullgraph = run_fullgraph try: yield finally: - speculator_module.dispatch_cg_and_sync_dp = native_dispatch + 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), @@ -229,17 +261,19 @@ def propose( original: Callable[..., torch.Tensor] = getattr( TargetDependentARSpeculator, _ORIGINAL_PROPOSE_ATTR ) - # ### PATCH START: rebuild draft-prefill metadata after a microbatched step. + # ### PATCH START: draft-prefill metadata on the draft builders under DBO. draft_prefill_scope: AbstractContextManager[None] = nullcontext() - if ( - attn_metadata is None - and self.vllm_config.parallel_config.use_ubatching - and not (dummy_run and skip_attn_for_dummy_run) + if self.vllm_config.parallel_config.use_ubatching and not ( + dummy_run and skip_attn_for_dummy_run ): - attn_metadata, slot_mappings = _build_drafter_attn_metadata(self, input_batch) - draft_prefill_scope = _eager_draft_prefill(self) + 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: rebuild draft-prefill metadata after a microbatched step. + # ### PATCH END: draft-prefill metadata on the draft builders under DBO. return original( self, input_batch, @@ -263,7 +297,7 @@ def propose( def apply_mrv2_dbo_spec_decode_backport() -> None: - """Install the draft-prefill pre-step once (worker process).""" + """Install the draft-prefill capture and pre-step once (worker process).""" if hasattr(TargetDependentARSpeculator, _ORIGINAL_PROPOSE_ATTR): return setattr( @@ -271,7 +305,13 @@ def apply_mrv2_dbo_spec_decode_backport() -> None: _ORIGINAL_PROPOSE_ATTR, TargetDependentARSpeculator.propose, ) + setattr( + TargetDependentARSpeculator, + _ORIGINAL_CAPTURE_ATTR, + TargetDependentARSpeculator.capture, + ) TargetDependentARSpeculator.propose = propose + TargetDependentARSpeculator.capture = capture apply_mrv2_dbo_spec_decode_backport() diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 3ac8bf248..2e54745d5 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -167,10 +167,11 @@ 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 (request-aligned -split of verification steps, draft-prefill metadata rebuild) and is eager -only: its FULL draft-prefill graphs would still capture on the target's -builders. MXFP4 routed experts, including +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. diff --git a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md index 0582e5f3d..aa10b2ae0 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -33,10 +33,9 @@ Supported models: (`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 (MTP with DBO eager - only), with decode-only micro-batching, each micro-batch keeping its own - DSA top-k state. Attention TP > 1, micro-batched prefill and MTP with - graph DBO are rejected. + 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. diff --git a/docs/reports/glm-5.3-afd-followups.md b/docs/reports/glm-5.3-afd-followups.md index 2c5e319cb..0a08b963a 100644 --- a/docs/reports/glm-5.3-afd-followups.md +++ b/docs/reports/glm-5.3-afd-followups.md @@ -13,6 +13,7 @@ Status of branch `glm-5.3-afd` (vLLM 0.32.0 pin). GPU evidence: one | 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 @@ -50,14 +51,20 @@ Status of branch `glm-5.3-afd` (vLLM 0.32.0 pin). GPU evidence: one 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: eager validated, graph rejected.** 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 metadata rebuild). Eager MTP + DBO 0.94 (MTP alone 0.92). - Graph MTP + DBO hit `illegal memory access` about ten minutes into GSM8K: - the FULL draft-prefill graphs still capture on the target's builders - (#61464's capture-group switch is not backported), so the adapter rejects - MTP with DBO unless `--enforce-eager`. +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 @@ -76,5 +83,5 @@ Status of branch `glm-5.3-afd` (vLLM 0.32.0 pin). GPU evidence: one | 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 | Land #61464's draft capture-group switch together with the split | +| 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 | diff --git a/docs/support.md b/docs/support.md index 70c8132c0..913bd1887 100644 --- a/docs/support.md +++ b/docs/support.md @@ -9,7 +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 (MTP with DBO eager only). Rejects P2P, Attention TP > 1, PP, SP MoE, micro-batched prefill and MTP with graph DBO. Ascend keeps the DeepSeek V2-derived path. | +| 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. | diff --git a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md index 98c9cd67d..262c1f9f6 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/glm_5_3/README.md @@ -11,7 +11,7 @@ prefill_decode_colocation/ ├── 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, 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 ``` @@ -62,7 +62,8 @@ The native vLLM control on the same node is - 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. MTP with DBO is eager only; verification steps split on - request boundaries and the decode threshold is raised to one request per - micro-batch (a backport of vLLM #61464). + 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/tests/unit/compat/patches/test_gpu_mrv2_dbo_spec_decode.py b/tests/unit/compat/patches/test_gpu_mrv2_dbo_spec_decode.py index 7adda1ce4..511c197e0 100644 --- a/tests/unit/compat/patches/test_gpu_mrv2_dbo_spec_decode.py +++ b/tests/unit/compat/patches/test_gpu_mrv2_dbo_spec_decode.py @@ -61,12 +61,26 @@ def test_split_keeps_the_midpoint_when_it_cannot_align( 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): + 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( @@ -80,7 +94,7 @@ def prepare_attn(input_batch, cg_mode, block_tables, slot_mappings, groups, kv): model_state=SimpleNamespace(prepare_attn=prepare_attn), attn_groups=["draft groups"], kv_cache_config="kv", - prefill_cudagraph_manager="prefill graphs", + prefill_cudagraph_manager=_GraphManager(prepared), prepared=prepared, ) @@ -131,21 +145,20 @@ def test_microbatched_step_rebuilds_draft_prefill_metadata(native_propose): assert slot_mappings["draft_layer"] is target_slots -def test_microbatched_step_runs_only_draft_prefill_eager(monkeypatch): - # FULL draft-prefill graphs replay the target builders' full-batch - # metadata, which a microbatched target forward never writes. +@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) - dispatched: list[tuple[str, bool]] = [] - - def fake_native_dispatch(manager, *args, need_eager=False, **kwargs): - dispatched.append((manager, need_eager)) - return "descriptor", None + manager = speculator.prefill_cudagraph_manager def original(self, *args, **kwargs): - # Native propose dispatches draft prefill, then draft decode. - dispatch = backport.speculator_module.dispatch_cg_and_sync_dp - dispatch("prefill graphs", 1, 12, 4, 1, 0, need_eager=False) - dispatch("decode graphs", 1, 3, 1, 1, 0, need_eager=False) + self.prefill_cudagraph_manager.run_fullgraph( + SimpleNamespace(cg_mode=backport.CUDAGraphMode.FULL) + ) return "draft tokens" monkeypatch.setattr( @@ -157,15 +170,43 @@ def original(self, *args, **kwargs): 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.speculator_module, "dispatch_cg_and_sync_dp", fake_native_dispatch + backport.TargetDependentARSpeculator, + backport._ORIGINAL_CAPTURE_ATTR, + original, + raising=False, ) - assert _propose(speculator, None, None) == "draft tokens" + backport.capture(speculator) - assert dispatched == [("prefill graphs", True), ("decode graphs", False)] - # The scope restores the module binding it replaced. - assert backport.speculator_module.dispatch_cg_and_sync_dp is fake_native_dispatch + assert captured_with == [["draft groups"] if use_ubatching else ["target groups"]] + assert speculator.target_attn_groups == ["target groups"] @pytest.mark.parametrize( From cb0281ac69589449051527dd61c18368d5a71538 Mon Sep 17 00:00:00 2001 From: ronenkat <16743404+ronenkat@users.noreply.github.com> Date: Sat, 14 Nov 2026 07:28:01 +0200 Subject: [PATCH 13/13] [Doc] GLM-5.3 follow-ups: upstream vLLM status after v0.32.0 Map each open item and upstream suggestion to the vLLM issues and PRs found on main and the v0.33.0 release branch, with their status and coverage. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ronenkat <16743404+ronenkat@users.noreply.github.com> --- docs/reports/glm-5.3-afd-followups.md | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/docs/reports/glm-5.3-afd-followups.md b/docs/reports/glm-5.3-afd-followups.md index 0a08b963a..565698ccb 100644 --- a/docs/reports/glm-5.3-afd-followups.md +++ b/docs/reports/glm-5.3-afd-followups.md @@ -85,3 +85,28 @@ Status of branch `glm-5.3-afd` (vLLM 0.32.0 pin). GPU evidence: one | #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. |