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..98466d43aac --- /dev/null +++ b/modelopt/torch/quantization/plugins/vllm_mega_moe.py @@ -0,0 +1,173 @@ +# 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) + 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. + + 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..160fadf13bf --- /dev/null +++ b/tests/gpu_vllm/torch/quantization/test_vllm_mega_moe.py @@ -0,0 +1,204 @@ +# 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.to(DEVICE) # like vLLM, which loads the model on the GPU + + +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) + + +@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) + )