From db80634308615d2a8b1cb6f887cccc4aa87a380b Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Mon, 12 Oct 2026 13:28:08 -0700 Subject: [PATCH 1/2] Add MoE dispatch/combine fake quantization for vLLM's DeepGEMM Mega-MoE Give vLLM's DeepGEMM Mega-MoE experts (--moe-backend deep_gemm_mega_moe; DeepSeek-V4/V4.1 and Kimi-K3) the same dispatch_quantizer and combine_quantizer as RoutedExperts. The kernel fuses dispatch, expert GEMMs and combine, so: - dispatch_quantizer fake-quantizes the experts' input tokens, which only feed the kernel's FP8 staging; - combine_quantizer fake-quantizes the router-weighted per-(token, top-k slot) rows the kernel leaves at the end of its symmetric buffer, then sums them again in the kernel's order (ascending slot, FP32 from zero, one BF16 rounding; a fused shared expert's slot last and unquantized) and replaces the kernel output. On the first batch with routed rows, each buffer checks that its rows reproduce the kernel output bit for bit, so a changed DeepGEMM layout raises. Serving needs a calibrated or constant combine scale (each rank re-quantizes rows from all expert ranks), and quantizing the dispatch needs the shared expert unfused. No host syncs after that check, so it runs under CUDA graphs. Signed-off-by: Shiyang Chen --- CHANGELOG.rst | 2 +- examples/vllm_serve/README.md | 13 +- .../torch/quantization/plugins/__init__.py | 1 + .../quantization/plugins/vllm_mega_moe.py | 163 +++++++++++++++++ .../torch/quantization/test_vllm_mega_moe.py | 171 ++++++++++++++++++ 5 files changed, 345 insertions(+), 5 deletions(-) create mode 100644 modelopt/torch/quantization/plugins/vllm_mega_moe.py create mode 100644 tests/gpu_vllm/torch/quantization/test_vllm_mega_moe.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index df983a8804e..6819f566649 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -36,7 +36,7 @@ Changelog - Add experimental dynamic FP8 fake quantization of GatedDeltaNet chunk-boundary states and WY activations for training through the standard ``quant_cfg`` interface. The fused GDN path requires ``fla-core==0.5.1`` and chunk size 64; state emulation requires SM89 or newer. - Add fake quantization of the sparse-attention indexer key cache and query for DeepSeek-V4 (vLLM and Megatron-Core) and GLM-5.3-Flash (vLLM) through the new ``indexer_k_quantizer`` and ``indexer_q_quantizer``. Enable them by importing the ``configs/ptq/units/indexer_k_nvfp4`` and ``configs/ptq/units/indexer_q_nvfp4`` units (NVFP4 with the global scale fixed to 1) into a recipe. - Add the ``configs/ptq/units/kv_nvfp4_mla`` recipe unit for fake quantization of the MLA KV cache (DeepSeek-V3, GLM-5.3-Flash, ...) in vLLM fake-quant serving: NVFP4 for the latent and FP8 for the RoPE key. See `examples/vllm_serve/README.md `_ for an example recipe. -- Add fake quantization of MoE expert-parallel communication to vLLM fake-quant serving through the new ``dispatch_quantizer`` and ``combine_quantizer`` on vLLM's routed experts, for MoE layers built on vLLM's fused-MoE layer and any all-to-all backend. Enable them by importing the ``configs/ptq/units/moe_dispatch_nvfp4`` and ``configs/ptq/units/moe_combine_nvfp4`` units into a recipe; see ``examples/vllm_serve/README.md``. +- Add fake quantization of MoE expert-parallel communication to vLLM fake-quant serving through the new ``dispatch_quantizer`` and ``combine_quantizer`` on vLLM's routed experts, for MoE layers built on vLLM's fused-MoE layer and any all-to-all backend, and for the DeepGEMM Mega-MoE kernel of DeepSeek-V4 and Kimi-K3 (``--moe-backend deep_gemm_mega_moe``). Enable them by importing the ``configs/ptq/units/moe_dispatch_nvfp4`` and ``configs/ptq/units/moe_combine_nvfp4`` units into a recipe; see ``examples/vllm_serve/README.md``. - vLLM fake-quant serving now runs on pre-quantized checkpoints such as FP8 when the recipe leaves those layers unquantized (for example a KV-cache-only recipe), and on MLA models with an FP8 KV cache; both previously failed during quantization. - ``KV_QUANT_CFG`` presets in vLLM fake-quant serving now quantize the MLA KV cache; on vLLM 0.16 and later they silently quantized nothing for MLA models. - vLLM fake-quant serving of NVFP4 and cast-mode KV-cache configs no longer needs ``--enforce-eager``. ``examples/vllm_serve/vllm_serve_fakequant.py`` now sets ``VLLM_DISABLE_COMPILE_CACHE=1`` for fake-quant serves by default, because a cached torch.compile graph of the same model without the fake quant would otherwise be reused. diff --git a/examples/vllm_serve/README.md b/examples/vllm_serve/README.md index 98bfe45c3ce..3e821367e2e 100644 --- a/examples/vllm_serve/README.md +++ b/examples/vllm_serve/README.md @@ -368,10 +368,15 @@ Notes: the dispatched tokens again, so the expert inputs carry both roundings. - Without data parallelism (`--enable-expert-parallel` with tensor parallelism only), vLLM runs no all-to-all; the quantizers then emulate one over the same expert partition. -- The Mega-MoE backend of DeepSeek-V4 and Kimi-K3 (`--moe-backend deep_gemm_mega_moe`) fuses - dispatch, expert GEMMs and combine into one kernel, and the FlashInfer MoE-EP backends - (`flashinfer_moe_ep_*`) combine inside their kernel. A recipe that enables a quantizer such a - backend cannot reach raises an error; use the default MoE backend instead. +- The DeepGEMM Mega-MoE backend of DeepSeek-V4 and Kimi-K3 (`--moe-backend deep_gemm_mega_moe`) + fuses dispatch, expert GEMMs and combine into one kernel. The dispatch quantizer sees its input + tokens. The combine quantizer sees the per-(token, expert) outputs, already weighted by the + router, which the kernel leaves in its buffer; they are summed again after quantization. There, + the combine needs a calibrated (`algorithm: max`) or constant global scale, and quantizing the + dispatch needs `VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION=1` (the fused shared expert would + read the quantized tokens too). +- The FlashInfer MoE-EP backends (`flashinfer_moe_ep_*`) combine inside their kernel, so a recipe + enabling the combine quantizer raises an error there. ## Serve a model with sparse attention in vLLM diff --git a/modelopt/torch/quantization/plugins/__init__.py b/modelopt/torch/quantization/plugins/__init__.py index 7234ca81f76..c1c50be74f4 100644 --- a/modelopt/torch/quantization/plugins/__init__.py +++ b/modelopt/torch/quantization/plugins/__init__.py @@ -72,6 +72,7 @@ with import_plugin("vllm"): from .vllm import * from .vllm_indexer import * + from .vllm_mega_moe import * with import_plugin("trl"): from .trl import * diff --git a/modelopt/torch/quantization/plugins/vllm_mega_moe.py b/modelopt/torch/quantization/plugins/vllm_mega_moe.py new file mode 100644 index 00000000000..c2cc11aff2a --- /dev/null +++ b/modelopt/torch/quantization/plugins/vllm_mega_moe.py @@ -0,0 +1,163 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Fake quantization of the MoE expert-parallel communication of DeepGEMM's Mega-MoE kernel in vLLM. + +Covers ``--moe-backend deep_gemm_mega_moe`` of DeepSeek-V4/V4.1 and Kimi-K3. The kernel fuses +dispatch, expert GEMMs and combine, so the quantizers of vLLM's ``RoutedExperts`` cannot reach it. +``dispatch_quantizer`` fake-quantizes the tokens before the kernel's FP8 staging, which only these +tokens feed. ``combine_quantizer`` fake-quantizes the per-(token, top-k slot) expert outputs the +kernel sends back, which stay in its symmetric buffer: after the kernel, they are quantized and +summed again in the kernel's order, replacing its output. +""" + +import importlib + +import torch + +from ..nn import ( + QuantLinearConvBase, + QuantModule, + QuantModuleRegistry, + SequentialQuantizer, + TensorQuantizer, +) +from .vllm import create_parallel_state + +__all__ = [] + +# NotImplementedError: some vLLM model packages reject unsupported platforms at import. +try: + from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts +except (ImportError, NotImplementedError): + DeepseekV4MegaMoEExperts = None +try: + from vllm.models.kimi_k3.nvidia.model import KimiK3MegaMoEExperts +except (ImportError, NotImplementedError): + KimiK3MegaMoEExperts = None + +# Symmetric buffers whose combine-row layout was checked against the kernel output. +_checked_buffers: set[int] = set() + + +def _combine_slots(symm_buffer, num_tokens: int) -> torch.Tensor: + """The kernel's BF16 combine rows as ``[num_topk (+1 fused shared expert), num_tokens, hidden]``. + + They are the last region of the symmetric buffer (DeepGEMM ``layout/mega_moe.cuh``), so their + offset is the configuration's buffer size minus their size. The buffer can be larger when it is + shared with a bigger configuration. + """ + buffer = symm_buffer + num_slots = buffer.num_topk + (1 if buffer.num_shared_experts else 0) + max_tokens = buffer.num_max_tokens_per_rank + num_bytes = num_slots * max_tokens * buffer.hidden * 2 + deep_gemm_c = importlib.import_module(type(buffer).__module__)._C # the build that sized it + size = deep_gemm_c.get_symm_buffer_size_for_mega_moe( + buffer.group.size(), + buffer.num_experts, + max_tokens, + buffer.num_topk, + buffer.hidden, + buffer.intermediate_hidden, + buffer.mma_type, + buffer.activation, + buffer.num_shared_experts, + )[0] + if buffer.mma_type != "fp8xfp4" or size < num_bytes or size > buffer.buffer.numel(): + raise RuntimeError(f"Unsupported DeepGEMM Mega-MoE buffer layout ({buffer.mma_type}).") + rows = buffer.buffer[size - num_bytes : size].view(torch.bfloat16) + return rows.view(num_slots, max_tokens, buffer.hidden)[:, :num_tokens] + + +def _ordered_sum(slots, valid, num_topk, quantize=None) -> torch.Tensor: + """Sum the combine rows like the kernel does. + + Valid top-k slots in ascending order, then the fused shared expert's slot, summed in FP32 from + zero and rounded to BF16 once. A Python loop keeps the order fixed. + """ + total = torch.zeros(slots.shape[1:], dtype=torch.float32, device=slots.device) + for k in range(num_topk): + contribution = torch.where(valid[k], slots[k], 0) + total += (quantize(contribution) if quantize else contribution).float() + if slots.shape[0] > num_topk: # the fused shared expert's local output, not communicated + total += slots[num_topk].float() + return total.to(torch.bfloat16) + + +class _QuantVLLMMegaMoEExperts(QuantModule): + """``dispatch_quantizer`` and ``combine_quantizer`` for vLLM's DeepGEMM Mega-MoE experts.""" + + def _setup(self): + self.dispatch_quantizer = TensorQuantizer(QuantLinearConvBase.default_quant_desc_input) + self.combine_quantizer = TensorQuantizer(QuantLinearConvBase.default_quant_desc_input) + self.dispatch_quantizer.disable() + self.combine_quantizer.disable() + self.parallel_state = create_parallel_state() + + def forward(self, hidden_states, *args, **kwargs): + if self.dispatch_quantizer.is_enabled and self.has_fused_shared_experts: + raise RuntimeError( + f"{self.prefix}: the fused shared expert reads the dispatched tokens too; set" + " VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION=1 to quantize the dispatch." + ) + output = super().forward(self.dispatch_quantizer(hidden_states), *args, **kwargs) + if self.combine_quantizer.is_enabled: + self._check_combine_scale() + symm_buffer = self.get_symm_buffer() + slots = _combine_slots(symm_buffer, output.shape[0]) + valid = (symm_buffer.topk_idx[: output.shape[0]] >= 0).T.unsqueeze(-1) + num_topk = symm_buffer.num_topk + checked = id(symm_buffer) in _checked_buffers + if not checked and not torch.cuda.is_current_stream_capturing() and valid.any(): + # Once per buffer, on a batch with routed rows (warmup batches are all padding): + # the rows must reproduce the kernel's output bit for bit. + if not torch.equal(_ordered_sum(slots, valid, num_topk), output): + raise RuntimeError( + f"{self.prefix}: the DeepGEMM Mega-MoE combine rows do not reproduce the" + " kernel output: its buffer layout changed, or DG_COMM_KERNEL_DEBUG=1 zeroed it." + ) + _checked_buffers.add(id(symm_buffer)) + output.copy_(_ordered_sum(slots, valid, num_topk, self.combine_quantizer)) + return output + + def _check_combine_scale(self) -> None: + """Require a global scale fixed across ranks while serving. + + Each rank re-quantizes the rows it received from all expert ranks, so a scale computed per + call would cover other rows than the sending rank's. Calibration collects the amax instead. + """ + quantizer = self.combine_quantizer + members = quantizer if isinstance(quantizer, SequentialQuantizer) else (quantizer,) + if not all( + q._if_calib + or q.is_mx_format + or q._use_constant_amax + or q._constant_amax is not None + or q.amax is not None + for q in members + if q.is_enabled + ): + raise RuntimeError( + f"{self.prefix}: Mega-MoE combine_quantizer needs a calibrated (algorithm: max) or" + " constant (constant_amax) global scale." + ) + + +for _experts_cls, _name in ( + (DeepseekV4MegaMoEExperts, "vllm_DeepseekV4MegaMoEExperts"), + (KimiK3MegaMoEExperts, "vllm_KimiK3MegaMoEExperts"), # overrides forward, so registered too +): + if _experts_cls is not None: + QuantModuleRegistry.register({_experts_cls: _name})(_QuantVLLMMegaMoEExperts) diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_mega_moe.py b/tests/gpu_vllm/torch/quantization/test_vllm_mega_moe.py new file mode 100644 index 00000000000..45a2c1da93f --- /dev/null +++ b/tests/gpu_vllm/torch/quantization/test_vllm_mega_moe.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Dispatch/combine fake quantization of vLLM's DeepGEMM Mega-MoE experts, with stand-ins for the +kernel and its symmetric buffer.""" + +import sys +from types import ModuleType, SimpleNamespace + +import pytest +import torch + +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.quantization.plugins import vllm_mega_moe +from modelopt.torch.utils.distributed import ParallelState + +# The first fake-quant call may build ModelOpt's CUDA extension. +pytestmark = pytest.mark.timeout(300) + +NUM_TOPK, MAX_TOKENS, HIDDEN, PREFIX_BYTES = 3, 8, 16, 64 +DEVICE = "cuda" + + +class _StandInSymmBuffer: + """DeepGEMM's symmetric buffer: inputs first, the BF16 combine rows ``[slots, T, H]`` last.""" + + def __init__(self, num_shared_experts): + self.group = SimpleNamespace(size=lambda: 2) + self.num_experts, self.intermediate_hidden = 8, 32 + self.num_topk, self.num_max_tokens_per_rank, self.hidden = NUM_TOPK, MAX_TOKENS, HIDDEN + self.num_shared_experts = num_shared_experts + self.mma_type, self.activation = "fp8xfp4", "swiglu" + num_slots = NUM_TOPK + (1 if num_shared_experts else 0) + self.size = PREFIX_BYTES + num_slots * MAX_TOKENS * HIDDEN * 2 + self.buffer = torch.zeros( + self.size + 32, dtype=torch.int8, device=DEVICE + ) # reused buffers can be larger + self.topk_idx = torch.full((MAX_TOKENS, NUM_TOPK), -1, dtype=torch.int64, device=DEVICE) + + +_deep_gemm_mega = ModuleType("_stand_in_deep_gemm_mega") +_deep_gemm_mega._C = SimpleNamespace( + get_symm_buffer_size_for_mega_moe=lambda world, experts, tokens, topk, hidden, inter, mma, act, shared: ( + (PREFIX_BYTES + (topk + (1 if shared else 0)) * tokens * hidden * 2), + None, + ) +) +_StandInSymmBuffer.__module__ = _deep_gemm_mega.__name__ + + +class _NativeMegaMoEExperts(torch.nn.Module): + """Writes per-slot rows like the kernel's L2 epilogue and sums them like its combine.""" + + def forward(self, hidden_states, topk_weights, topk_ids, *, activation_clamp=None): + self.dispatched = hidden_states + num_tokens = hidden_states.shape[0] + slots = vllm_mega_moe._combine_slots(self.symm_buffer, num_tokens) + self.symm_buffer.topk_idx[:num_tokens] = topk_ids + for k in range(slots.shape[0]): + slots[k].copy_(hidden_states * (k + 1) + self.corruption) + valid = (topk_ids >= 0).T.unsqueeze(-1) + return vllm_mega_moe._ordered_sum(slots, valid, NUM_TOPK) + self.corruption + + def get_symm_buffer(self): + return self.symm_buffer + + +class _TestQuantMegaMoEExperts(vllm_mega_moe._QuantVLLMMegaMoEExperts, _NativeMegaMoEExperts): + pass + + +@pytest.fixture(autouse=True) +def _stand_in_deep_gemm(monkeypatch): + monkeypatch.setitem(sys.modules, _deep_gemm_mega.__name__, _deep_gemm_mega) + monkeypatch.setattr(vllm_mega_moe, "create_parallel_state", lambda: ParallelState(None)) + monkeypatch.setattr(vllm_mega_moe, "_checked_buffers", set()) + + +def _mega_moe(num_shared_experts=0, constant_amax=1.0, corruption=0.0): + moe = object.__new__(_TestQuantMegaMoEExperts) + torch.nn.Module.__init__(moe) + moe.prefix = "model.layers.0.ffn.experts" + moe.symm_buffer = _StandInSymmBuffer(num_shared_experts) + moe.has_fused_shared_experts = bool(num_shared_experts) + moe.corruption = corruption + moe._setup() + for quantizer in (moe.dispatch_quantizer, moe.combine_quantizer): + quantizer.set_from_attribute_config({"enable": True, "constant_amax": constant_amax}) + return moe + + +def _forward(moe, num_tokens=5): + hidden_states = torch.randn(num_tokens, HIDDEN, dtype=torch.bfloat16, device=DEVICE) + topk_ids = torch.randint(0, 8, (num_tokens, NUM_TOPK), device=DEVICE) + topk_ids[0, 1] = topk_ids[3, 0] = -1 # unused slots + topk_ids[4] = -1 # a padding token + return hidden_states, topk_ids, moe(hidden_states, None, topk_ids, activation_clamp=None) + + +@pytest.mark.parametrize("num_shared_experts", [0, 1], ids=("routed_only", "fused_shared")) +def test_mega_moe_quantizes_dispatch_and_combine(num_shared_experts): + moe = _mega_moe(num_shared_experts) + if num_shared_experts: + moe.dispatch_quantizer.disable() # the fused shared expert reads the dispatched tokens + hidden_states, topk_ids, output = _forward(moe) + + dispatched = moe.dispatch_quantizer(hidden_states) + assert torch.equal(moe.dispatched, dispatched) + valid = (topk_ids >= 0).T.unsqueeze(-1) + total = torch.zeros(hidden_states.shape, device=DEVICE) + for k in range(NUM_TOPK): # each communicated row is quantized; the shared one is not + total += moe.combine_quantizer(torch.where(valid[k], dispatched * (k + 1), 0)).float() + if num_shared_experts: + total += (dispatched * (NUM_TOPK + 1)).float() + assert torch.equal(output, total.to(torch.bfloat16)) + assert not torch.equal(output, vllm_mega_moe._ordered_sum(*_slots_and_valid(moe, topk_ids))) + if not num_shared_experts: + assert not output[4].any() # padding tokens get no routed rows + + +def _slots_and_valid(moe, topk_ids): + slots = vllm_mega_moe._combine_slots(moe.symm_buffer, topk_ids.shape[0]) + return slots, (topk_ids >= 0).T.unsqueeze(-1), NUM_TOPK + + +@pytest.mark.parametrize( + ("num_shared_experts", "constant_amax", "corruption", "match"), + [ + (1, 1.0, 0.0, "fused shared expert reads the dispatched tokens"), + (0, None, 0.0, "needs a calibrated .* or constant"), + (0, 1.0, 1.0, "do not reproduce the kernel output"), + ], + ids=("fused_shared_dispatch", "per_call_combine_scale", "changed_buffer_layout"), +) +def test_mega_moe_rejects_unsupported_setups(num_shared_experts, constant_amax, corruption, match): + moe = _mega_moe(num_shared_experts, constant_amax, corruption) + if constant_amax is None: + moe.dispatch_quantizer.disable() + with pytest.raises(RuntimeError, match=match): + _forward(moe) + + +def test_mega_moe_experts_are_registered(): + for cls in (vllm_mega_moe.DeepseekV4MegaMoEExperts, vllm_mega_moe.KimiK3MegaMoEExperts): + if cls is None: + pytest.skip("this vLLM has no DeepGEMM Mega-MoE experts") + assert cls in vllm_mega_moe.QuantModuleRegistry + assert isinstance(_mega_moe().combine_quantizer, TensorQuantizer) + + +def test_mega_moe_layout_check_waits_for_routed_rows(): + moe = _mega_moe(corruption=1.0) # the rows no longer reproduce the kernel output + moe.dispatch_quantizer.disable() + hidden_states = torch.randn(4, HIDDEN, dtype=torch.bfloat16, device=DEVICE) + padding = torch.full((4, NUM_TOPK), -1, device=DEVICE) # like vLLM's warmup batches + moe(hidden_states, None, padding, activation_clamp=None) # nothing to compare yet + assert not vllm_mega_moe._checked_buffers + with pytest.raises(RuntimeError, match="do not reproduce the kernel output"): + _forward(moe) From 32796ffbf97f860abe888da21fd280ac394f6719 Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Mon, 12 Oct 2026 14:52:16 -0700 Subject: [PATCH 2/2] Check the Mega-MoE combine layout during CUDA graph replay A CUDA graph captured before any eager forward with routed rows would never check the combine rows, since replays skip the Python check. Such a capture now records the check in the graph: every replay compares the rows with the kernel output on the device and fails with a device-side assert on a mismatch. Eager forwards still check each buffer once and raise a RuntimeError. Signed-off-by: Shiyang Chen --- .../quantization/plugins/vllm_mega_moe.py | 34 +++++++++++------- .../torch/quantization/test_vllm_mega_moe.py | 35 ++++++++++++++++++- 2 files changed, 56 insertions(+), 13 deletions(-) diff --git a/modelopt/torch/quantization/plugins/vllm_mega_moe.py b/modelopt/torch/quantization/plugins/vllm_mega_moe.py index c2cc11aff2a..98466d43aac 100644 --- a/modelopt/torch/quantization/plugins/vllm_mega_moe.py +++ b/modelopt/torch/quantization/plugins/vllm_mega_moe.py @@ -118,20 +118,30 @@ def forward(self, hidden_states, *args, **kwargs): symm_buffer = self.get_symm_buffer() slots = _combine_slots(symm_buffer, output.shape[0]) valid = (symm_buffer.topk_idx[: output.shape[0]] >= 0).T.unsqueeze(-1) - num_topk = symm_buffer.num_topk - checked = id(symm_buffer) in _checked_buffers - if not checked and not torch.cuda.is_current_stream_capturing() and valid.any(): - # Once per buffer, on a batch with routed rows (warmup batches are all padding): - # the rows must reproduce the kernel's output bit for bit. - if not torch.equal(_ordered_sum(slots, valid, num_topk), output): - raise RuntimeError( - f"{self.prefix}: the DeepGEMM Mega-MoE combine rows do not reproduce the" - " kernel output: its buffer layout changed, or DG_COMM_KERNEL_DEBUG=1 zeroed it." - ) - _checked_buffers.add(id(symm_buffer)) - output.copy_(_ordered_sum(slots, valid, num_topk, self.combine_quantizer)) + if id(symm_buffer) not in _checked_buffers: + self._check_combine_rows(symm_buffer, slots, valid, output) + output.copy_(_ordered_sum(slots, valid, symm_buffer.num_topk, self.combine_quantizer)) return output + def _check_combine_rows(self, symm_buffer, slots, valid, output) -> None: + """Check that the combine rows reproduce the kernel output bit for bit. + + Eager forwards check each buffer once, on a batch with routed rows (warmup batches are all + padding). A CUDA graph captured before that check checks every replay on the device instead, + failing with a device-side assert. + """ + message = ( + f"{self.prefix}: the DeepGEMM Mega-MoE combine rows do not reproduce the kernel output:" + " its buffer layout changed, or DG_COMM_KERNEL_DEBUG=1 zeroed it." + ) + reproduced = (_ordered_sum(slots, valid, symm_buffer.num_topk) == output).all() + if torch.cuda.is_current_stream_capturing(): + torch._assert_async(reproduced, message) + elif valid.any(): + if not reproduced: + raise RuntimeError(message) + _checked_buffers.add(id(symm_buffer)) + def _check_combine_scale(self) -> None: """Require a global scale fixed across ranks while serving. diff --git a/tests/gpu_vllm/torch/quantization/test_vllm_mega_moe.py b/tests/gpu_vllm/torch/quantization/test_vllm_mega_moe.py index 45a2c1da93f..160fadf13bf 100644 --- a/tests/gpu_vllm/torch/quantization/test_vllm_mega_moe.py +++ b/tests/gpu_vllm/torch/quantization/test_vllm_mega_moe.py @@ -98,7 +98,7 @@ def _mega_moe(num_shared_experts=0, constant_amax=1.0, corruption=0.0): moe._setup() for quantizer in (moe.dispatch_quantizer, moe.combine_quantizer): quantizer.set_from_attribute_config({"enable": True, "constant_amax": constant_amax}) - return moe + return moe.to(DEVICE) # like vLLM, which loads the model on the GPU def _forward(moe, num_tokens=5): @@ -169,3 +169,36 @@ def test_mega_moe_layout_check_waits_for_routed_rows(): assert not vllm_mega_moe._checked_buffers with pytest.raises(RuntimeError, match="do not reproduce the kernel output"): _forward(moe) + + +@pytest.mark.parametrize("corruption", [0.0, 1.0], ids=("same_layout", "changed_buffer_layout")) +def test_mega_moe_graph_checks_layout_on_replay(monkeypatch, corruption): + """A CUDA graph captured after padding-only warmup checks the rows on every replay.""" + reproduced = torch.full((), bool(corruption), device=DEVICE) # wrong until the replay checks + # Record the device-side check instead of letting it kill the CUDA context. + monkeypatch.setattr( + torch, "_assert_async", lambda condition, message: reproduced.copy_(condition) + ) + moe = _mega_moe(corruption=corruption) + moe.dispatch_quantizer.disable() + hidden_states = torch.randn(5, HIDDEN, dtype=torch.bfloat16, device=DEVICE) + topk_ids = torch.full((5, NUM_TOPK), -1, device=DEVICE) # like vLLM's warmup batches + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + moe(hidden_states, None, topk_ids, activation_clamp=None) # eager warmup checks nothing + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = moe(hidden_states, None, topk_ids, activation_clamp=None) + assert not vllm_mega_moe._checked_buffers + + topk_ids.copy_(torch.randint(0, 8, topk_ids.shape, device=DEVICE)) + graph.replay() + assert reproduced.item() == (not corruption) + if not corruption: + slots, valid, num_topk = _slots_and_valid(moe, topk_ids) + assert valid.all() + assert torch.equal( + output, vllm_mega_moe._ordered_sum(slots, valid, num_topk, moe.combine_quantizer) + )