A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 11:30 UTC
[Bugfix] Fix dtype mismatch in RMSNormGated.forward_native() during torch.compile (#35256)
Signed-off-by: haosdent <haosdent@gmail.com>
haosdent committed
Mar 2, 2026 at 04:14 UTC
6290470843c131681e3e1318ae71070a34f33225
2 files changed
+71
-4
tests/kernels/test_fla_layernorm_guard.py
+63
-1
@@ -74,7 +74,7 @@ def layer_norm_ref(
74
return out.to(dtype)
75
76
77
-DTYPES = [torch.bfloat16, torch.float32]
77
+DTYPES = [torch.float16, torch.bfloat16, torch.float32]
78
# Test various M sizes to ensure rows_per_block logic works correctly
79
NUM_TOKENS = [
80
1,
@@ -380,6 +380,68 @@ def test_multidimensional_input(
380
torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
381
382
383
+@pytest.mark.parametrize("num_tokens", [1, 128, 1024])
384
+@pytest.mark.parametrize("hidden_size", [64, 256, 1024])
385
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
386
+@pytest.mark.parametrize("has_gate", [True, False])
387
+@pytest.mark.parametrize("group_size", [None, 64])
388
+@pytest.mark.parametrize("norm_before_gate", [True, False])
389
+@torch.inference_mode()
390
+def test_rmsnorm_gated_forward_native_dtype(
391
+ default_vllm_config,
392
+ num_tokens: int,
393
+ hidden_size: int,
394
+ dtype: torch.dtype,
395
+ has_gate: bool,
396
+ group_size: int | None,
397
+ norm_before_gate: bool,
398
+):
399
+ """Test that RMSNormGated.forward_native preserves input dtype."""
400
+ if group_size is not None and hidden_size % group_size != 0:
401
+ pytest.skip(
402
+ f"hidden_size {hidden_size} not divisible by group_size {group_size}"
403
+ )
404
+
405
+ from vllm.model_executor.layers.layernorm import RMSNormGated
406
+
407
+ device = torch.device("cuda:0")
408
+ set_random_seed(42)
409
+
410
+ layer = RMSNormGated(
411
+ hidden_size,
412
+ eps=1e-5,
413
+ group_size=group_size,
414
+ norm_before_gate=norm_before_gate,
415
+ device=device,
416
+ dtype=dtype,
417
+ )
418
+
419
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
420
+ z = (
421
+ torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
422
+ if has_gate
423
+ else None
424
+ )
425
+
426
+ out = layer.forward_native(x, z)
427
+
428
+ # Verify dtype preservation
429
+ assert out.dtype == dtype, f"Expected {dtype}, got {out.dtype}"
430
+
431
+ # Verify numerical correctness against reference
432
+ ref_out = rms_norm_ref(
433
+ x,
434
+ layer.weight,
435
+ layer.bias,
436
+ z=z,
437
+ eps=1e-5,
438
+ group_size=group_size,
439
+ norm_before_gate=norm_before_gate,
440
+ upcast=True,
441
+ )
442
+ torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
443
+
444
+
445
if __name__ == "__main__":
446
# Run a quick smoke test
447
test_layer_norm_fwd_basic(128, 1024, torch.float16, 42, False)
vllm/model_executor/layers/layernorm.py
+8
-3
@@ -557,6 +557,11 @@ class RMSNormGated(CustomOp):
557
- norm_before_gate=True: out = norm(x) * silu(z)
558
- norm_before_gate=False: out = norm(x * silu(z))
559
"""
560
+ orig_dtype = x.dtype
561
+ x = x.float()
562
+ weight = self.weight.float()
563
+ z = z.float() if z is not None else None
564
+
565
# Apply gating before normalization if needed
566
if z is not None and not self.norm_before_gate:
567
x = x * F.silu(z)
@@ -566,7 +571,7 @@ class RMSNormGated(CustomOp):
571
# Standard RMS norm across the last dimension
572
variance = x.pow(2).mean(dim=-1, keepdim=True)
573
x_normed = x * torch.rsqrt(variance + self.eps)
569
- out = x_normed * self.weight
574
+ out = x_normed * weight
575
else:
576
# Group RMS norm
577
from einops import rearrange
@@ -574,13 +579,13 @@ class RMSNormGated(CustomOp):
579
x_group = rearrange(x, "... (g d) -> ... g d", d=self.group_size)
580
variance = x_group.pow(2).mean(dim=-1, keepdim=True)
581
x_normed = x_group * torch.rsqrt(variance + self.eps)
577
- out = rearrange(x_normed, "... g d -> ... (g d)") * self.weight
582
+ out = rearrange(x_normed, "... g d -> ... (g d)") * weight
583
584
# Apply gating after normalization if needed
585
if z is not None and self.norm_before_gate:
586
out = out * F.silu(z)
587
583
- return out.to(x.dtype)
588
+ return out.to(orig_dtype)
589
590
def forward_cuda(
591
self, x: torch.Tensor, z: torch.Tensor | None = None