[Bugfix][DeepSeek V4] Fix fused MTP RMSNorm dtype mismatch

tests/kernels/test_fused_mtp_input_rmsnorm.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.models.deepseek_v4.common.ops.fused_mtp_input_rmsnorm import (
+    fused_mtp_input_rmsnorm,
+)
+
+
+def _rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor:
+    x_float = x.float()
+    variance = x_float.square().mean(dim=-1, keepdim=True)
+    return (x_float * torch.rsqrt(variance + eps) * weight.float()).to(x.dtype)
+
+
+def test_fused_mtp_input_rmsnorm() -> None:
+    torch.manual_seed(0)
+    device = "cuda"
+    dtype = torch.float16
+    num_tokens = 4
+    hidden_size = 128
+    hc_mult = 2
+    eps = 1e-6
+
+    inputs_embeds = torch.randn(
+        num_tokens, hidden_size, dtype=dtype, device=device
+    )
+    positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
+    previous_hidden_states = torch.randn(
+        num_tokens, hc_mult, hidden_size, dtype=dtype, device=device
+    )
+    enorm_weight = torch.randn(hidden_size, dtype=dtype, device=device)
+    hnorm_weight = torch.randn(hidden_size, dtype=dtype, device=device)
+
+    enorm_output, hnorm_output = fused_mtp_input_rmsnorm(
+        inputs_embeds,
+        positions,
+        previous_hidden_states,
+        enorm_weight,
+        hnorm_weight,
+        eps,
+        hc_mult,
+    )
+
+    masked_inputs = inputs_embeds.clone()
+    masked_inputs[positions == 0] = 0
+    expected_enorm = _rms_norm(masked_inputs, enorm_weight, eps)
+    expected_hnorm = _rms_norm(previous_hidden_states, hnorm_weight, eps)
+
+    torch.testing.assert_close(enorm_output, expected_enorm, rtol=1e-3, atol=1e-3)
+    torch.testing.assert_close(hnorm_output, expected_hnorm, rtol=1e-3, atol=1e-3)

vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py

@@ -70,7 +70,7 @@ def _fused_mtp_input_rmsnorm_kernel(
         keep = pos != 0
         x = tl.load(
             inputs_embeds_ptr + token_idx * HIDDEN + block, mask=mask, other=0.0
-        )
+        ).to(tl.float32)
         x = tl.where(keep, x, 0.0)
         _rmsnorm_row(
             x,
@@ -85,7 +85,9 @@ def _fused_mtp_input_rmsnorm_kernel(
         # hnorm path: load prev_hidden[token, slot, :].
         slot = pid_task - 1
         row_offset = (token_idx * HC_MULT + slot) * HIDDEN
-        x = tl.load(prev_hidden_ptr + row_offset + block, mask=mask, other=0.0)
+        x = tl.load(
+            prev_hidden_ptr + row_offset + block, mask=mask, other=0.0
+        ).to(tl.float32)
         _rmsnorm_row(
             x,
             hnorm_weight_ptr,