From eb08d83ef74897d785a3d14671607a80d896e144 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 14 Sep 2026 17:31:14 +0000 Subject: [PATCH 1/6] [Build] Move DeepGEMM pin to the vLLM fork and update the Mega MoE call convention Repoint the DeepGEMM pin from deepseek-ai/DeepGEMM nv_dev (2.6.1) to vllm-project/DeepGEMM dev at e49b29bb (2.8.0). The fork carries the SM120 port we relied on nv_dev for, plus configurable SwiGLU alpha/beta and the SiTU activation for FP8/FP4 Mega MoE. 2.8.0 renames the Mega MoE activation parameters: the optional situ_beta/situ_linear_beta pair becomes activation_alpha/activation_beta (plain floats, defaulting to 1.0/0.0). For SiTU, activation_alpha is the gate tanh scale and activation_beta the linear/up tanh scale, where 0.0 leaves the up branch untouched. Update the Kimi K3 MegaMoE call site accordingly. The SwiGLU defaults reproduce the previous behavior, so the DeepSeek V4 call sites are unchanged. 2.8.0 also swaps the third-party/fmt submodule for third-party/deep_jit, so update the submodule list, the _C include path, and the header glob that retriggers the _C rebuild. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Yongye Zhu --- cmake/external_projects/deepgemm.cmake | 16 ++++++++++------ docker/Dockerfile | 1 + tools/build_deepgemm_C.py | 2 +- tools/install_deepgemm.sh | 8 +++++--- vllm/models/kimi_k3/nvidia/model.py | 6 ++++-- 5 files changed, 21 insertions(+), 12 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 72f7f5eb66a8..878fbdd0ae2c 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -28,9 +28,11 @@ if(DEEPGEMM_SRC_DIR) message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}") else() # Keep in sync with tools/install_deepgemm.sh - set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/deepseek-ai/DeepGEMM.git") - # Pinned to the tip of the nv_dev branch (SM120 support). - set(_DEEPGEMM_UPSTREAM_TAG "8b1392b978f5a03c828dd1711090d7fb50958b8a") + set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git") + # Pinned to the tip of the fork's dev branch: upstream 2.8.0 plus the SM120 + # port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU + # alpha/beta, and SiTU for FP8/FP4 Mega MoE. + set(_DEEPGEMM_UPSTREAM_TAG "9a86ae2b78991f1c8e6e95945a591e1ca3b19ce7") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) @@ -51,7 +53,7 @@ else() BINARY_DIR "${_deepgemm_bin}" GIT_REPOSITORY "${_DEEPGEMM_UPSTREAM_REPO}" GIT_TAG "${_DEEPGEMM_UPSTREAM_TAG}" - GIT_SUBMODULES "third-party/cutlass" "third-party/fmt" + GIT_SUBMODULES "third-party/cutlass" "third-party/deep_jit" GIT_PROGRESS TRUE ) endif() @@ -119,13 +121,15 @@ if(DEEPGEMM_ARCHS) message(STATUS "DeepGEMM _C will be built for: ${_dg_pythons}") # add_custom_command does no implicit header scanning; glob explicitly so - # header-only edits in DeepGEMM/cutlass/fmt re-trigger the rebuild. + # header-only edits in DeepGEMM/cutlass/deep_jit re-trigger the rebuild. file(GLOB_RECURSE _dg_headers "${deepgemm_SOURCE_DIR}/csrc/*.h" "${deepgemm_SOURCE_DIR}/csrc/*.hpp" "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.h" "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.hpp" - "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh") + "${deepgemm_SOURCE_DIR}/deep_gemm/include/*.cuh" + "${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.h" + "${deepgemm_SOURCE_DIR}/third-party/deep_jit/include/*.hpp") set(_dg_markers) set(_dg_seen_soabis) diff --git a/docker/Dockerfile b/docker/Dockerfile index 9b344552c5ae..b3e32777e7bb 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -131,6 +131,7 @@ RUN dnf install -y --setopt=install_weak_deps=False \ git \ curl \ sudo \ + elfutils-devel \ rdma-core-devel \ && dnf clean all \ && rm -rf /var/cache/dnf diff --git a/tools/build_deepgemm_C.py b/tools/build_deepgemm_C.py index 67a527405e2d..f4438c67552f 100644 --- a/tools/build_deepgemm_C.py +++ b/tools/build_deepgemm_C.py @@ -50,7 +50,7 @@ str(src / "deep_gemm/include"), str(src / "third-party/cutlass/include"), str(src / "third-party/cutlass/tools/util/include"), - str(src / "third-party/fmt/include"), + str(src / "third-party/deep_jit/include"), *cpp_extension.include_paths(device_type="cuda"), ] diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index f6a64d8796c1..d7747cc89d4e 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -6,9 +6,11 @@ set -e # Default values # Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake -DEEPGEMM_GIT_REPO="https://github.com/deepseek-ai/DeepGEMM.git" -# NOTE: This is currently targeting the nv_dev branch tip due to sm120 support -DEEPGEMM_GIT_REF="8b1392b978f5a03c828dd1711090d7fb50958b8a" +DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" +# NOTE: This targets the vLLM fork's dev branch tip, which carries the sm120 +# and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE +# activations +DEEPGEMM_GIT_REF="9a86ae2b78991f1c8e6e95945a591e1ca3b19ce7" WHEEL_DIR="" # Parse command line arguments diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index c4f0092f8d9b..6a26ce78b4c4 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -532,8 +532,10 @@ def forward( symm_buffer, activation_clamp=activation_clamp, activation=self.activation, - situ_beta=self.activation_beta, - situ_linear_beta=self.activation_linear_beta, + # DeepGEMM names the SiTU gate tanh scale `activation_alpha` and the + # linear/up tanh scale `activation_beta`; beta=0 leaves up untouched. + activation_alpha=self.activation_beta or 1.0, + activation_beta=self.activation_linear_beta or 0.0, fast_math=fast_math, ) return y From 28a44a6e18ba9abf4f961cc539dbccd3b007d610 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Mon, 14 Sep 2026 17:47:32 +0000 Subject: [PATCH 2/6] [Bugfix] Resolve the masked grouped FP8 GEMM under DeepGEMM 2.8.0 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit DeepGEMM 2.8.0 removes the legacy `fp8_m_grouped_gemm_nt_masked` Python alias, so the getattr lookup silently resolved to None and every call raised "DeepGEMM backend is unavailable" — breaking BatchedDeepGemmExperts. Look up the canonical `m_grouped_fp8_gemm_nt_masked` first and keep the legacy spelling as a fallback for externally pinned older DeepGEMM. Both names alias the same `m_grouped_fp8_fp4_gemm_nt_masked` binding, whose signature and pybind defaults are unchanged across the two versions. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Yongye Zhu --- vllm/utils/deep_gemm.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index 3fc70c0f5da3..d515d5131e75 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -280,7 +280,11 @@ def _lazy_init() -> None: _fp8_gemm_nt_impl = getattr(_dg, "fp8_gemm_nt", None) _fp8_einsum_impl = getattr(_dg, "fp8_einsum", None) _grouped_impl = getattr(_dg, "m_grouped_fp8_gemm_nt_contiguous", None) - _grouped_masked_impl = getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None) + # DeepGEMM 2.8.0 dropped the legacy `fp8_*` aliases; keep both spellings so + # an externally pinned older DeepGEMM still resolves. + _grouped_masked_impl = getattr( + _dg, "m_grouped_fp8_gemm_nt_masked", None + ) or getattr(_dg, "fp8_m_grouped_gemm_nt_masked", None) _grouped_fp4_impl = getattr(_dg, "m_grouped_fp8_fp4_gemm_nt_contiguous", None) # DeepGEMM exposes fp8_fp4_*_mqa_logits as the canonical symbols that # handle both the FP8 and FP4 Q/K paths via a tuple-typed `q`. From 90005228d2232cd6742b4d079c6b4040df4f3d89 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Tue, 15 Sep 2026 04:46:45 +0000 Subject: [PATCH 3/6] [Bugfix] Repin DeepGEMM for the SM90 paged-MQA zero-context fix The fork commit this was first pinned to, 9a86ae2b, re-ported the SM90 paged-MQA kv_block=32/next_n=4 support onto the 2.8.0 base but dropped the zero-context-length guards that deepseek-ai's nv_dev branch carried. With all context lengths zero the metadata kernel reads prefix_sum[batch_size], one past the end of a shared buffer sized to exactly align(batch_size, 32) ints, and traps with CUDBG_EXCEPTION_WARP_OUT_OF_RANGE_ADDRESS. The DSA indexer warmup runs a dummy batch of max_num_seqs (default 1024, already 32-aligned) rows with all context lengths zero, so every DeepSeek V4 startup on Hopper hit it: nvidia-h100-lm-eval-kv-offload-large failed deterministically on all 4 ranks of both test_gsm8k_offloading_correctness cases. Move the pin to the dev tip, which now carries vllm-project/DeepGEMM#8 restoring the four guards and adding a zero-context regression test. Signed-off-by: Yongye Zhu Co-Authored-By: Claude Opus 5 (1M context) --- cmake/external_projects/deepgemm.cmake | 2 +- tools/install_deepgemm.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 878fbdd0ae2c..2ff69bd8b6b3 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -32,7 +32,7 @@ else() # Pinned to the tip of the fork's dev branch: upstream 2.8.0 plus the SM120 # port, the SM90 paged-MQA kv_block=32/next_n=4 port, configurable SwiGLU # alpha/beta, and SiTU for FP8/FP4 Mega MoE. - set(_DEEPGEMM_UPSTREAM_TAG "9a86ae2b78991f1c8e6e95945a591e1ca3b19ce7") + set(_DEEPGEMM_UPSTREAM_TAG "ad1f1726aa540a76c1d26d6a120effb8de21eaa4") set(_deepgemm_fc_root "${FETCHCONTENT_BASE_DIR}") if(NOT _deepgemm_fc_root) diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh index d7747cc89d4e..db81bf4be310 100755 --- a/tools/install_deepgemm.sh +++ b/tools/install_deepgemm.sh @@ -10,7 +10,7 @@ DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" # NOTE: This targets the vLLM fork's dev branch tip, which carries the sm120 # and sm90 paged-MQA ports plus the SwiGLU alpha/beta and SiTU Mega MoE # activations -DEEPGEMM_GIT_REF="9a86ae2b78991f1c8e6e95945a591e1ca3b19ce7" +DEEPGEMM_GIT_REF="ad1f1726aa540a76c1d26d6a120effb8de21eaa4" WHEEL_DIR="" # Parse command line arguments From dd26723b4c19e4ac04b7d06939653a961e4e9e1e Mon Sep 17 00:00:00 2001 From: Thien Tran Date: Thu, 10 Sep 2026 10:49:42 +0000 Subject: [PATCH 4/6] first commit Co-authored-by: OpenAI Codex Signed-off-by: Thien Tran (cherry picked from commit 74c2daf29d509f75e1dd6fd560886be19ebf25e0) Signed-off-by: Juntian Liu --- tests/kernels/quantization/test_block_fp8.py | 7 +- tests/kernels/test_mhc_kernels.py | 83 +++++++++- vllm/models/deepseek_v41/nvidia/model.py | 71 ++++----- .../deepseek_v41/nvidia/ops/__init__.py | 2 + .../deepseek_v41/nvidia/ops/mega_mhc.py | 144 ++++++++++++++++++ vllm/utils/deep_gemm.py | 15 +- 6 files changed, 279 insertions(+), 43 deletions(-) create mode 100644 vllm/models/deepseek_v41/nvidia/ops/__init__.py create mode 100644 vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py index fd530e7ffb53..d02527aba048 100644 --- a/tests/kernels/quantization/test_block_fp8.py +++ b/tests/kernels/quantization/test_block_fp8.py @@ -27,6 +27,7 @@ from vllm.utils.deep_gemm import ( fp8_gemm_nt, get_tma_aligned_size, + is_deep_gemm_e8m0_used, per_block_cast_to_fp8, should_use_deepgemm_for_fp8_linear, ) @@ -291,7 +292,11 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed): A_fp8, As_fp8 = per_token_group_quant_fp8( A_fp32, block_size[1], column_major_scales=True, tma_aligned_scales=True ) - B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32, block_size=block_size) + B_fp8, Bs_fp8 = per_block_cast_to_fp8( + B_fp32, + block_size=block_size, + use_ue8m0=is_deep_gemm_e8m0_used(), + ) As = As_fp8.to(torch.float32) Bs = Bs_fp8.to(torch.float32) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index f805fe39a57b..e85fb2a7c675 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -11,6 +11,7 @@ import vllm.model_executor.layers.mhc as mhc_layers from vllm.model_executor.kernels.mhc.tilelang import ( _torch_hc_prenorm_gemm, + mhc_post_tilelang, mhc_pre_delayed_tilelang, ) from vllm.model_executor.kernels.mhc.tilelang_kernels import ( @@ -39,6 +40,10 @@ from vllm.models.deepseek_v41.nvidia.model import ( DeepseekV4DecoderLayer as DeepseekV41DecoderLayer, ) +from vllm.models.deepseek_v41.nvidia.ops.mega_mhc import ( + is_mega_mhc_supported, + mhc_shifted_post_pre_deep_gemm, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -474,7 +479,7 @@ def fused_reference(x, residual, post_mix, res_mix, *args, capture_aux=False, ** reference, ) monkeypatch.setattr( - "vllm.models.deepseek_v41.nvidia.model.mhc_fused_post_pre_delayed_tilelang", + "vllm.models.deepseek_v41.nvidia.model.mhc_shifted_post_pre", fused_reference, ) expected = decoder(x, positions, None, **kwargs) @@ -559,6 +564,82 @@ def test_mhc_pre_delayed_custom_op_supports_compile(carried): ) +@pytest.mark.skipif( + not current_platform.is_device_capability_family(100), + reason="DeepGEMM Mega mHC requires SM100-family CUDA", +) +def test_deep_gemm_mega_mhc_correctness(): + num_tokens, hidden_size, hc_mult = 17, 7168, 4 + if not is_mega_mhc_supported(hidden_size, 4): + pytest.skip("DeepGEMM Mega mHC is not available") + + x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16, device=DEVICE) + residual = torch.randn( + num_tokens, hc_mult, hidden_size, dtype=torch.bfloat16, device=DEVICE + ) + previous_mix = torch.rand(num_tokens, hc_mult, device=DEVICE) + post_mix = torch.rand(num_tokens, hc_mult, 1, device=DEVICE) + res_mix = torch.rand(num_tokens, hc_mult, hc_mult, device=DEVICE) + fn = torch.randn(24, hc_mult * hidden_size, device=DEVICE) * 0.01 + scale = torch.randn(3, device=DEVICE) * 0.1 + base = torch.randn(24, device=DEVICE) * 0.1 + norm_weight = torch.empty( + hidden_size, dtype=torch.bfloat16, device=DEVICE + ).uniform_(0.9, 1.1) + + expected_residual = mhc_post_tilelang(x, residual, post_mix, res_mix) + ( + expected_post_mix, + expected_res_mix, + expected_y_bf16, + expected_previous_mix, + ) = mhc_pre_delayed_tilelang( + expected_residual, + fn, + scale, + base, + 2e-5, + 3e-4, + 2e-6, + 1.25, + 10, + pre_mix=previous_mix, + norm_weight=norm_weight, + norm_eps=7e-6, + ) + ( + actual_residual, + actual_post_mix, + actual_res_mix, + actual_y_bf16, + actual_previous_mix, + ) = mhc_shifted_post_pre_deep_gemm( + x, + residual, + previous_mix, + post_mix, + res_mix, + fn, + scale, + base, + 2e-5, + 3e-4, + 1.25, + 2e-6, + 10, + norm_weight, + 7e-6, + ) + torch.accelerator.synchronize() + torch.testing.assert_close(actual_residual, expected_residual, atol=1e-5, rtol=1e-2) + torch.testing.assert_close(actual_post_mix, expected_post_mix, atol=1e-6, rtol=1e-3) + torch.testing.assert_close(actual_res_mix, expected_res_mix, atol=1e-6, rtol=1e-3) + torch.testing.assert_close(actual_y_bf16, expected_y_bf16, atol=1e-5, rtol=1e-2) + torch.testing.assert_close( + actual_previous_mix, expected_previous_mix, atol=1e-6, rtol=1e-3 + ) + + def sinkhorn_normalize_ref(x: torch.Tensor, repeat: int, eps: float) -> torch.Tensor: x = x.softmax(-1) + eps x = x / (x.sum(-2, keepdim=True) + eps) diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index f3530544da7b..7748189df105 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -20,7 +20,6 @@ from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.kernels.mhc.tilelang import ( - mhc_fused_post_pre_delayed_tilelang, mhc_post_tilelang, mhc_pre_delayed_tilelang, ) @@ -79,6 +78,7 @@ from ..common.engram import EngramLayout, NgramHashState from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID, image_sentinel_mask from .engram import Engram, gather_engram_hashes +from .ops.mega_mhc import mhc_shifted_post_pre if typing.TYPE_CHECKING: from vllm.v1.attention.backends.mla.sparse_swa import DeepseekSparseSWAMetadata @@ -403,54 +403,45 @@ def forward( else: # The collapse already reads the post-mapped streams, so the mean # aux consumers want comes out of the same kernel. - residual, post_mix, res_mix, x, attn_pre, aux = ( - mhc_fused_post_pre_delayed_tilelang( - x, - residual, - post_mix, - res_mix, - self.hc_attn_fn, - self.hc_attn_scale, - self.hc_attn_base, - self.rms_norm_eps, - self.hc_eps, - self.hc_eps, - self.hc_post_alpha, - self.hc_sinkhorn_iters, - pre_mix=pre_mix, - norm_weight=self.attn_norm.weight, - norm_eps=self.attn_norm.variance_epsilon, - capture_aux=capture_previous_aux, - ) - ) - if capture_previous_aux: - previous_aux = aux - - if self.use_sequence_parallel: - x = sp_all_gather(x)[: positions.shape[0]] - - x = self.attn(positions, x, None) - if self.use_sequence_parallel: - x = sp_reduce_scatter(x) - - residual, post_mix, res_mix, x, ffn_pre, _ = ( - mhc_fused_post_pre_delayed_tilelang( + residual, post_mix, res_mix, x, attn_pre = mhc_shifted_post_pre( x, residual, + pre_mix, post_mix, res_mix, - self.hc_ffn_fn, - self.hc_ffn_scale, - self.hc_ffn_base, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, self.rms_norm_eps, self.hc_eps, - self.hc_eps, self.hc_post_alpha, self.hc_sinkhorn_iters, - pre_mix=attn_pre, - norm_weight=self.ffn_norm.weight, - norm_eps=self.ffn_norm.variance_epsilon, + self.attn_norm.weight, + self.attn_norm.variance_epsilon, ) + + if self.use_sequence_parallel: + x = sp_all_gather(x)[: positions.shape[0]] + + x = self.attn(positions, x, None) + if self.use_sequence_parallel: + x = sp_reduce_scatter(x) + + residual, post_mix, res_mix, x, ffn_pre = mhc_shifted_post_pre( + x, + residual, + attn_pre, + post_mix, + res_mix, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_post_alpha, + self.hc_sinkhorn_iters, + self.ffn_norm.weight, + self.ffn_norm.variance_epsilon, ) x = self.ffn(x, input_ids) return x, residual, post_mix, res_mix, ffn_pre, previous_aux diff --git a/vllm/models/deepseek_v41/nvidia/ops/__init__.py b/vllm/models/deepseek_v41/nvidia/ops/__init__.py new file mode 100644 index 000000000000..208f01a7cb5e --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py new file mode 100644 index 000000000000..e4c47a639f7e --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import functools + +import torch + +from vllm.model_executor.kernels.mhc.tilelang import ( + mhc_post_tilelang, + mhc_pre_delayed_tilelang, +) +from vllm.platforms import current_platform +from vllm.utils.deep_gemm import ( + _import_deep_gemm, + is_deep_gemm_supported, + mega_mhc, +) + + +@functools.cache +def is_mega_mhc_supported(hidden_size: int, hc_mult: int) -> bool: + """Return whether DeepGEMM can run the DSV4.1 shifted mHC kernel.""" + if not ( + is_deep_gemm_supported() + and current_platform.is_device_capability_family(100) + and hidden_size > 0 + and hidden_size % 1024 == 0 + and hc_mult == 4 + ): + return False + deep_gemm = _import_deep_gemm() + return deep_gemm is not None and callable(getattr(deep_gemm, "mega_mhc", None)) + + +def mhc_shifted_post_pre_deep_gemm( + x: torch.Tensor, + residual: torch.Tensor, + shifted_prev_mix: torch.Tensor, + post_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + mix_scales: torch.Tensor, + mix_bases: torch.Tensor, + hc_norm_eps: float, + hc_pre_eps: float, + hc_post_scale: float, + sinkhorn_eps: float, + num_sinkhorn_iters: int, + rmsnorm_weight: torch.Tensor, + rmsnorm_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Run DSV4.1 shifted post, next pre, and BF16 RMSNorm with Mega mHC.""" + num_tokens, hidden_size = x.shape + hc_mult = residual.shape[1] + new_residual = torch.empty_like(residual) + new_prev_mix = x.new_empty(num_tokens, hc_mult, 1, dtype=torch.float32) + new_post_mix = torch.empty_like(post_mix) + new_comb_res_mix = torch.empty_like(comb_res_mix) + y_bf16 = x.new_empty(num_tokens, hidden_size, dtype=torch.bfloat16) + mega_mhc( + x=x, + residual=residual, + shifted_prev_mix=shifted_prev_mix.unsqueeze(-1), + post_mix=post_mix, + comb_res_mix=comb_res_mix, + fn=fn, + mix_scales=mix_scales, + mix_bases=mix_bases, + hc_mult=hc_mult, + hc_norm_eps=hc_norm_eps, + hc_pre_eps=hc_pre_eps, + hc_post_scale=hc_post_scale, + sinkhorn_eps=sinkhorn_eps, + num_sinkhorn_iters=num_sinkhorn_iters, + rmsnorm_weight=rmsnorm_weight, + rmsnorm_eps=rmsnorm_eps, + rmsnorm_scale=1.0, + new_residual=new_residual, + new_prev_mix=new_prev_mix, + new_post_mix=new_post_mix, + new_comb_res_mix=new_comb_res_mix, + y_bf16=y_bf16, + ) + return ( + new_residual, + new_post_mix, + new_comb_res_mix, + y_bf16, + new_prev_mix.squeeze(-1), + ) + + +def mhc_shifted_post_pre( + x: torch.Tensor, + residual: torch.Tensor, + previous_mix: torch.Tensor, + post_mix: torch.Tensor, + res_mix: torch.Tensor, + fn: torch.Tensor, + scale: torch.Tensor, + base: torch.Tensor, + rms_eps: float, + hc_eps: float, + hc_post_alpha: float, + sinkhorn_iters: int, + norm_weight: torch.Tensor, + norm_eps: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Dispatch the DSV4.1 shifted mHC transition.""" + if is_mega_mhc_supported(x.shape[1], residual.shape[1]): + return mhc_shifted_post_pre_deep_gemm( + x, + residual, + previous_mix, + post_mix, + res_mix, + fn, + scale, + base, + rms_eps, + hc_eps, + hc_post_alpha, + hc_eps, + sinkhorn_iters, + norm_weight, + norm_eps, + ) + + residual = mhc_post_tilelang(x, residual, post_mix, res_mix) + post_mix, res_mix, x, next_mix = mhc_pre_delayed_tilelang( + residual, + fn, + scale, + base, + rms_eps, + hc_eps, + hc_eps, + hc_post_alpha, + sinkhorn_iters, + pre_mix=previous_mix, + norm_weight=norm_weight, + norm_eps=norm_eps, + ) + return residual, post_mix, res_mix, x, next_mix diff --git a/vllm/utils/deep_gemm.py b/vllm/utils/deep_gemm.py index d515d5131e75..5d05c29a84c5 100644 --- a/vllm/utils/deep_gemm.py +++ b/vllm/utils/deep_gemm.py @@ -159,6 +159,7 @@ def _missing(*_: Any, **__: Any) -> NoReturn: _fp8_fp4_paged_mqa_logits_impl: Callable[..., Any] | None = None _get_paged_mqa_logits_metadata_impl: Callable[..., Any] | None = None _tf32_hc_prenorm_gemm_impl: Callable[..., Any] | None = None +_mega_mhc_impl: Callable[..., Any] | None = None _get_mn_major_tma_aligned_tensor_impl: Callable[..., Any] | None = None _get_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = None _get_theoretical_mk_alignment_for_contiguous_layout_impl: Callable[..., Any] | None = ( @@ -231,7 +232,7 @@ def _lazy_init() -> None: global _grouped_impl, _grouped_masked_impl, _grouped_fp4_impl global _fp8_fp4_mqa_logits_impl, _fp8_fp4_paged_mqa_logits_impl global _get_paged_mqa_logits_metadata_impl - global _tf32_hc_prenorm_gemm_impl + global _tf32_hc_prenorm_gemm_impl, _mega_mhc_impl global _get_mn_major_tma_aligned_tensor_impl global _get_mk_alignment_for_contiguous_layout_impl global _get_theoretical_mk_alignment_for_contiguous_layout_impl @@ -251,6 +252,7 @@ def _lazy_init() -> None: or _fp8_fp4_paged_mqa_logits_impl is not None or _get_paged_mqa_logits_metadata_impl is not None or _tf32_hc_prenorm_gemm_impl is not None + or _mega_mhc_impl is not None or _get_mk_alignment_for_contiguous_layout_impl is not None or _transform_sf_into_required_layout_impl is not None or _pack_ue8m0_to_int_impl is not None @@ -294,6 +296,7 @@ def _lazy_init() -> None: _dg, "get_paged_mqa_logits_metadata", None ) _tf32_hc_prenorm_gemm_impl = getattr(_dg, "tf32_hc_prenorm_gemm", None) + _mega_mhc_impl = getattr(_dg, "mega_mhc", None) _get_mn_major_tma_aligned_tensor_impl = getattr( _dg, "get_mn_major_tma_aligned_tensor", None ) @@ -704,6 +707,15 @@ def tf32_hc_prenorm_gemm( ) +def mega_mhc(*args: Any, **kwargs: Any) -> None: + """Run DeepGEMM Mega mHC with caller-owned output tensors.""" + _lazy_init() + if _mega_mhc_impl is None: + _missing() + return + _mega_mhc_impl(*args, **kwargs) + + def _ceil_to_ue8m0(x: torch.Tensor): return torch.pow(2.0, torch.ceil(torch.log2(x.abs()))) @@ -797,6 +809,7 @@ def should_use_deepgemm_for_fp8_linear( "per_block_cast_to_fp8", "is_deep_gemm_e8m0_used", "is_deep_gemm_supported", + "mega_mhc", "get_num_sms", "set_num_sms", "should_use_deepgemm_for_fp8_linear", From 040f3eb3ec7d120a67c9fea851d5f08495ec3ff4 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Tue, 15 Sep 2026 06:14:53 +0000 Subject: [PATCH 5/6] [DSv4.1] Reopen #56255: adapt Mega-mHC to current decoder Rebase the closed #56255 Mega-mHC integration on the #56876 dependency base, retargeting deepseek_v41. Preserve the TileLang fallback and Engram auxiliary-state semantics, and extend correctness and CUDA graph coverage. Co-authored-by: OpenAI Codex Signed-off-by: Juntian Liu --- tests/kernels/test_mhc_kernels.py | 55 +++++++++---- vllm/models/deepseek_v41/nvidia/model.py | 21 +++-- .../deepseek_v41/nvidia/ops/mega_mhc.py | 82 +++++++++++-------- 3 files changed, 98 insertions(+), 60 deletions(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index e85fb2a7c675..3e7f892ca712 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -568,8 +568,11 @@ def test_mhc_pre_delayed_custom_op_supports_compile(carried): not current_platform.is_device_capability_family(100), reason="DeepGEMM Mega mHC requires SM100-family CUDA", ) -def test_deep_gemm_mega_mhc_correctness(): - num_tokens, hidden_size, hc_mult = 17, 7168, 4 +@pytest.mark.parametrize("num_tokens", [0, 1, 17, 128, 1024]) +@pytest.mark.parametrize("hidden_size", [5120, 7168]) +def test_deep_gemm_mega_mhc_correctness(num_tokens, hidden_size): + set_random_seed(0) + hc_mult = 4 if not is_mega_mhc_supported(hidden_size, 4): pytest.skip("DeepGEMM Mega mHC is not available") @@ -587,7 +590,9 @@ def test_deep_gemm_mega_mhc_correctness(): hidden_size, dtype=torch.bfloat16, device=DEVICE ).uniform_(0.9, 1.1) - expected_residual = mhc_post_tilelang(x, residual, post_mix, res_mix) + expected_residual = ( + mhc_post_tilelang(x, residual, post_mix, res_mix) if num_tokens else residual + ) ( expected_post_mix, expected_res_mix, @@ -607,13 +612,7 @@ def test_deep_gemm_mega_mhc_correctness(): norm_weight=norm_weight, norm_eps=7e-6, ) - ( - actual_residual, - actual_post_mix, - actual_res_mix, - actual_y_bf16, - actual_previous_mix, - ) = mhc_shifted_post_pre_deep_gemm( + args = ( x, residual, previous_mix, @@ -630,14 +629,36 @@ def test_deep_gemm_mega_mhc_correctness(): norm_weight, 7e-6, ) - torch.accelerator.synchronize() - torch.testing.assert_close(actual_residual, expected_residual, atol=1e-5, rtol=1e-2) - torch.testing.assert_close(actual_post_mix, expected_post_mix, atol=1e-6, rtol=1e-3) - torch.testing.assert_close(actual_res_mix, expected_res_mix, atol=1e-6, rtol=1e-3) - torch.testing.assert_close(actual_y_bf16, expected_y_bf16, atol=1e-5, rtol=1e-2) - torch.testing.assert_close( - actual_previous_mix, expected_previous_mix, atol=1e-6, rtol=1e-3 + actual = mhc_shifted_post_pre_deep_gemm(*args) + expected = ( + expected_residual, + expected_post_mix, + expected_res_mix, + expected_y_bf16, + expected_previous_mix, ) + for i, (result, ref) in enumerate(zip(actual, expected, strict=True)): + # Match the existing mHC BF16 tolerance: post rounding differences + # can be amplified by cancellation in the carried collapse. + atol, rtol = (1.6e-2, 1e-2) if i in (0, 3) else (1e-6, 1e-3) + torch.testing.assert_close(result, ref, atol=atol, rtol=rtol) + + if num_tokens: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + # DeepGEMM initializes barriers per stream before capture. + mhc_shifted_post_pre_deep_gemm(*args) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = mhc_shifted_post_pre_deep_gemm(*args) + for multiplier in (0.5, -1.0): + x.mul_(multiplier) + graph.replay() + eager = mhc_shifted_post_pre_deep_gemm(*args) + for result, ref in zip(captured, eager, strict=True): + torch.testing.assert_close(result, ref, atol=0, rtol=0) + torch.cuda.current_stream().wait_stream(stream) def sinkhorn_normalize_ref(x: torch.Tensor, repeat: int, eps: float) -> torch.Tensor: diff --git a/vllm/models/deepseek_v41/nvidia/model.py b/vllm/models/deepseek_v41/nvidia/model.py index 7748189df105..9c9719e97729 100644 --- a/vllm/models/deepseek_v41/nvidia/model.py +++ b/vllm/models/deepseek_v41/nvidia/model.py @@ -403,10 +403,9 @@ def forward( else: # The collapse already reads the post-mapped streams, so the mean # aux consumers want comes out of the same kernel. - residual, post_mix, res_mix, x, attn_pre = mhc_shifted_post_pre( + residual, post_mix, res_mix, x, attn_pre, aux = mhc_shifted_post_pre( x, residual, - pre_mix, post_mix, res_mix, self.hc_attn_fn, @@ -414,11 +413,16 @@ def forward( self.hc_attn_base, self.rms_norm_eps, self.hc_eps, + self.hc_eps, self.hc_post_alpha, self.hc_sinkhorn_iters, - self.attn_norm.weight, - self.attn_norm.variance_epsilon, + pre_mix=pre_mix, + norm_weight=self.attn_norm.weight, + norm_eps=self.attn_norm.variance_epsilon, + capture_aux=capture_previous_aux, ) + if capture_previous_aux: + previous_aux = aux if self.use_sequence_parallel: x = sp_all_gather(x)[: positions.shape[0]] @@ -427,10 +431,9 @@ def forward( if self.use_sequence_parallel: x = sp_reduce_scatter(x) - residual, post_mix, res_mix, x, ffn_pre = mhc_shifted_post_pre( + residual, post_mix, res_mix, x, ffn_pre, _ = mhc_shifted_post_pre( x, residual, - attn_pre, post_mix, res_mix, self.hc_ffn_fn, @@ -438,10 +441,12 @@ def forward( self.hc_ffn_base, self.rms_norm_eps, self.hc_eps, + self.hc_eps, self.hc_post_alpha, self.hc_sinkhorn_iters, - self.ffn_norm.weight, - self.ffn_norm.variance_epsilon, + pre_mix=attn_pre, + norm_weight=self.ffn_norm.weight, + norm_eps=self.ffn_norm.variance_epsilon, ) x = self.ffn(x, input_ids) return x, residual, post_mix, res_mix, ffn_pre, previous_aux diff --git a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py index e4c47a639f7e..f2c74902616f 100644 --- a/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py +++ b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py @@ -6,8 +6,7 @@ import torch from vllm.model_executor.kernels.mhc.tilelang import ( - mhc_post_tilelang, - mhc_pre_delayed_tilelang, + mhc_fused_post_pre_delayed_tilelang, ) from vllm.platforms import current_platform from vllm.utils.deep_gemm import ( @@ -93,52 +92,65 @@ def mhc_shifted_post_pre_deep_gemm( def mhc_shifted_post_pre( x: torch.Tensor, residual: torch.Tensor, - previous_mix: torch.Tensor, - post_mix: torch.Tensor, - res_mix: torch.Tensor, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, fn: torch.Tensor, - scale: torch.Tensor, - base: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, rms_eps: float, - hc_eps: float, - hc_post_alpha: float, - sinkhorn_iters: int, - norm_weight: torch.Tensor, - norm_eps: float, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """Dispatch the DSV4.1 shifted mHC transition.""" - if is_mega_mhc_supported(x.shape[1], residual.shape[1]): - return mhc_shifted_post_pre_deep_gemm( + hc_pre_eps: float, + hc_sinkhorn_eps: float, + hc_post_mult_value: float, + sinkhorn_repeat: int, + pre_mix: torch.Tensor | None = None, + norm_weight: torch.Tensor | None = None, + norm_eps: float = 1e-6, + capture_aux: bool = False, +) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor +]: + """Use Mega-mHC for carried mixing, retaining TileLang for aux capture.""" + if ( + pre_mix is not None + and norm_weight is not None + and not capture_aux + and x.shape[0] <= 1 << 20 + and is_mega_mhc_supported(x.shape[1], residual.shape[1]) + ): + outputs = mhc_shifted_post_pre_deep_gemm( x, residual, - previous_mix, - post_mix, - res_mix, + pre_mix, + post_layer_mix, + comb_res_mix, fn, - scale, - base, + hc_scale, + hc_base, rms_eps, - hc_eps, - hc_post_alpha, - hc_eps, - sinkhorn_iters, + hc_pre_eps, + hc_post_mult_value, + hc_sinkhorn_eps, + sinkhorn_repeat, norm_weight, norm_eps, ) + return *outputs, x.new_empty(0, x.shape[1]) - residual = mhc_post_tilelang(x, residual, post_mix, res_mix) - post_mix, res_mix, x, next_mix = mhc_pre_delayed_tilelang( + return mhc_fused_post_pre_delayed_tilelang( + x, residual, + post_layer_mix, + comb_res_mix, fn, - scale, - base, + hc_scale, + hc_base, rms_eps, - hc_eps, - hc_eps, - hc_post_alpha, - sinkhorn_iters, - pre_mix=previous_mix, + hc_pre_eps, + hc_sinkhorn_eps, + hc_post_mult_value, + sinkhorn_repeat, + pre_mix=pre_mix, norm_weight=norm_weight, norm_eps=norm_eps, + capture_aux=capture_aux, ) - return residual, post_mix, res_mix, x, next_mix From e7fcd569dd9edb56dd8676e3111cb5207ad14026 Mon Sep 17 00:00:00 2001 From: Juntian Liu Date: Tue, 15 Sep 2026 07:04:17 +0000 Subject: [PATCH 6/6] clean tests Signed-off-by: Juntian Liu --- tests/kernels/test_mhc_kernels.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 3e7f892ca712..ec7d9b714d05 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -564,6 +564,7 @@ def test_mhc_pre_delayed_custom_op_supports_compile(carried): ) +@pytest.mark.skipif(not HAS_TILELANG_MHC, reason="TileLang MHC support required") @pytest.mark.skipif( not current_platform.is_device_capability_family(100), reason="DeepGEMM Mega mHC requires SM100-family CUDA",