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 ff72e2a96c17..5fab150d9fce 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 @@ -495,7 +500,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) @@ -580,6 +585,104 @@ 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", +) +@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") + + 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) if num_tokens else residual + ) + ( + 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, + ) + args = ( + x, + residual, + previous_mix, + post_mix, + res_mix, + fn, + scale, + base, + 2e-5, + 3e-4, + 1.25, + 2e-6, + 10, + norm_weight, + 7e-6, + ) + 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: 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..9c9719e97729 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,25 +403,23 @@ 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, - ) + residual, post_mix, res_mix, x, attn_pre, aux = mhc_shifted_post_pre( + 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 @@ -433,24 +431,22 @@ def forward( if self.use_sequence_parallel: x = sp_reduce_scatter(x) - residual, post_mix, res_mix, x, ffn_pre, _ = ( - mhc_fused_post_pre_delayed_tilelang( - x, - residual, - 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_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, - ) + residual, post_mix, res_mix, x, ffn_pre, _ = mhc_shifted_post_pre( + x, + residual, + 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_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, ) 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..f2c74902616f --- /dev/null +++ b/vllm/models/deepseek_v41/nvidia/ops/mega_mhc.py @@ -0,0 +1,156 @@ +# 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_fused_post_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, + post_layer_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + 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, + pre_mix, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + 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]) + + return mhc_fused_post_pre_delayed_tilelang( + x, + residual, + post_layer_mix, + comb_res_mix, + fn, + hc_scale, + hc_base, + rms_eps, + 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, + ) 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",