From efa2e424f659dd6bbc07982bc155f265e32532cd Mon Sep 17 00:00:00 2001 From: Shangdi Yu Date: Thu, 16 Jul 2026 12:57:45 -0700 Subject: [PATCH] [Helion] Fix degenerate scale_ub in kernel input generators (#48868) Signed-off-by: Shangdi Yu --- .../helion/ops/dynamic_per_token_scaled_fp8_quant.py | 8 +++++++- .../helion/ops/rms_norm_dynamic_per_token_quant.py | 11 ++++++++++- vllm/kernels/helion/ops/rms_norm_per_block_quant.py | 11 ++++++++++- .../helion/ops/silu_and_mul_per_block_quant.py | 9 ++++++++- 4 files changed, 35 insertions(+), 4 deletions(-) diff --git a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py index 45bd8f6fcf3..75f240c2aac 100644 --- a/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py +++ b/vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py @@ -41,7 +41,13 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype) result = torch.empty(input.shape, device=input.device, dtype=out_dtype) scale = torch.empty((num_tokens, 1), device=input.device, dtype=scale_dtype) - scale_ub = torch.mean(input).to(scale_dtype) + # scale_ub clamps the per-token amax of |input|. Use a non-degenerate + # upper bound (midway between the mean and max of |input|) so clamping is + # partially active and the baseline comparison is meaningful. + # torch.mean(input) ~= 0 for the zero-mean input would collapse every + # scale to the floor and saturate the output. + input_abs = input.to(torch.float32).abs() + scale_ub = (0.5 * (input_abs.mean() + input_abs.amax())).to(scale_dtype) config_key = CaseKey({"hidden_size": hidden_size, "num_tokens": num_tokens}) inputs[config_key] = (result, input, scale, scale_ub) diff --git a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py index 3e02169db3b..8bf352cc512 100644 --- a/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py +++ b/vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py @@ -48,7 +48,6 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype) result = torch.empty(input.shape, device=input.device, dtype=out_dtype) scale = torch.empty((num_tokens, 1), device=input.device, dtype=scale_dtype) - scale_ub = torch.mean(input).to(scale_dtype) residual = torch.randn_like(input) weight = torch.normal( mean=1.0, @@ -58,6 +57,16 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: device=input.device, ) epsilon = 1e-6 + # scale_ub clamps the per-token amax of the RMS-normed, weighted output. + # Use a non-degenerate upper bound (midway between the mean and max of + # that magnitude) so clamping is partially active and the baseline + # comparison is meaningful. torch.mean(input) ~= 0 for the zero-mean + # input would collapse every scale to the floor and saturate the output. + # Mirrors the reference normalization in baseline() below. + x = input.to(torch.float32) + residual.to(torch.float32) + rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon) + x_norm_abs = ((x * rms).to(input.dtype) * weight).abs().to(torch.float32) + scale_ub = (0.5 * (x_norm_abs.mean() + x_norm_abs.amax())).to(scale_dtype) config_key = CaseKey({"hidden_size": hidden_size, "num_tokens": num_tokens}) inputs[config_key] = (result, input, weight, scale, epsilon, scale_ub, residual) diff --git a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py index da7651a1fa5..fe995b2c626 100644 --- a/vllm/kernels/helion/ops/rms_norm_per_block_quant.py +++ b/vllm/kernels/helion/ops/rms_norm_per_block_quant.py @@ -55,7 +55,6 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: device=input.device, dtype=scale_dtype, ) - scale_ub = torch.mean(input).to(scale_dtype) residual = torch.randn_like(input) weight = torch.normal( mean=1.0, @@ -65,6 +64,16 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: device=input.device, ) epsilon = 1e-6 + # scale_ub clamps the per-group amax of the RMS-normed, weighted output. + # Use a non-degenerate upper bound (midway between the mean and max of + # that magnitude) so clamping is partially active and the baseline + # comparison is meaningful. torch.mean(input) ~= 0 for the zero-mean + # input would collapse every scale to the floor and saturate the output. + # Mirrors the reference normalization in baseline() below. + x = input.to(torch.float32) + residual.to(torch.float32) + rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon) + x_norm_abs = ((x * rms).to(input.dtype) * weight).abs().to(torch.float32) + scale_ub = (0.5 * (x_norm_abs.mean() + x_norm_abs.amax())).to(scale_dtype) config_key = CaseKey( { diff --git a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py index 06b7f10af2f..31386943e7c 100644 --- a/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py +++ b/vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py @@ -60,7 +60,14 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]: device=input.device, dtype=scale_dtype, ) - scale_ub = torch.mean(input).to(scale_dtype) + # scale_ub clamps the per-group amax of the SiLU-and-mul activation. Use + # a non-degenerate upper bound (midway between the mean and max of the + # activation magnitude) so clamping is partially active and the baseline + # comparison is meaningful. torch.mean(input) ~= 0 for the zero-mean + # input would collapse every scale to the floor and saturate the output. + # Mirrors tests/kernels/helion/test_silu_and_mul_per_block_quant.py. + act_abs = SiluAndMul.forward_native(input.to(torch.float32)).abs() + scale_ub = (0.5 * (act_abs.mean() + act_abs.amax())).to(scale_dtype) config_key = CaseKey( {