diff --git a/tests/kernels/test_fused_mtp_input_rmsnorm.py b/tests/kernels/test_fused_mtp_input_rmsnorm.py new file mode 100644 index 000000000000..5173e1d59f96 --- /dev/null +++ b/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) diff --git a/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py b/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py index 874b4f607955..28a28e837ea7 100644 --- a/vllm/models/deepseek_v4/common/ops/fused_mtp_input_rmsnorm.py +++ b/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,