[Helion] Fix degenerate scale_ub in kernel input generators (#48868)
Signed-off-by: Shangdi Yu <shangdiy@meta.com>
Shangdi Yu committed
Jul 16, 2026 at 12:57 UTC
efa2e424f659dd6bbc07982bc155f265e32532cd
4 files changed
+35
-4
vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py
+7
-1
@@ -41,7 +41,13 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]:
41
input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype)
42
result = torch.empty(input.shape, device=input.device, dtype=out_dtype)
43
scale = torch.empty((num_tokens, 1), device=input.device, dtype=scale_dtype)
44
- scale_ub = torch.mean(input).to(scale_dtype)
44
+ # scale_ub clamps the per-token amax of |input|. Use a non-degenerate
45
+ # upper bound (midway between the mean and max of |input|) so clamping is
46
+ # partially active and the baseline comparison is meaningful.
47
+ # torch.mean(input) ~= 0 for the zero-mean input would collapse every
48
+ # scale to the floor and saturate the output.
49
+ input_abs = input.to(torch.float32).abs()
50
+ scale_ub = (0.5 * (input_abs.mean() + input_abs.amax())).to(scale_dtype)
51
52
config_key = CaseKey({"hidden_size": hidden_size, "num_tokens": num_tokens})
53
inputs[config_key] = (result, input, scale, scale_ub)
vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py
+10
-1
@@ -48,7 +48,6 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]:
48
input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype)
49
result = torch.empty(input.shape, device=input.device, dtype=out_dtype)
50
scale = torch.empty((num_tokens, 1), device=input.device, dtype=scale_dtype)
51
- scale_ub = torch.mean(input).to(scale_dtype)
51
residual = torch.randn_like(input)
52
weight = torch.normal(
53
mean=1.0,
@@ -58,6 +57,16 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]:
57
device=input.device,
58
)
59
epsilon = 1e-6
60
+ # scale_ub clamps the per-token amax of the RMS-normed, weighted output.
61
+ # Use a non-degenerate upper bound (midway between the mean and max of
62
+ # that magnitude) so clamping is partially active and the baseline
63
+ # comparison is meaningful. torch.mean(input) ~= 0 for the zero-mean
64
+ # input would collapse every scale to the floor and saturate the output.
65
+ # Mirrors the reference normalization in baseline() below.
66
+ x = input.to(torch.float32) + residual.to(torch.float32)
67
+ rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon)
68
+ x_norm_abs = ((x * rms).to(input.dtype) * weight).abs().to(torch.float32)
69
+ scale_ub = (0.5 * (x_norm_abs.mean() + x_norm_abs.amax())).to(scale_dtype)
70
71
config_key = CaseKey({"hidden_size": hidden_size, "num_tokens": num_tokens})
72
inputs[config_key] = (result, input, weight, scale, epsilon, scale_ub, residual)
vllm/kernels/helion/ops/rms_norm_per_block_quant.py
+10
-1
@@ -55,7 +55,6 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]:
55
device=input.device,
56
dtype=scale_dtype,
57
)
58
- scale_ub = torch.mean(input).to(scale_dtype)
58
residual = torch.randn_like(input)
59
weight = torch.normal(
60
mean=1.0,
@@ -65,6 +64,16 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]:
64
device=input.device,
65
)
66
epsilon = 1e-6
67
+ # scale_ub clamps the per-group amax of the RMS-normed, weighted output.
68
+ # Use a non-degenerate upper bound (midway between the mean and max of
69
+ # that magnitude) so clamping is partially active and the baseline
70
+ # comparison is meaningful. torch.mean(input) ~= 0 for the zero-mean
71
+ # input would collapse every scale to the floor and saturate the output.
72
+ # Mirrors the reference normalization in baseline() below.
73
+ x = input.to(torch.float32) + residual.to(torch.float32)
74
+ rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon)
75
+ x_norm_abs = ((x * rms).to(input.dtype) * weight).abs().to(torch.float32)
76
+ scale_ub = (0.5 * (x_norm_abs.mean() + x_norm_abs.amax())).to(scale_dtype)
77
78
config_key = CaseKey(
79
{
vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py
+8
-1
@@ -60,7 +60,14 @@ def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]:
60
device=input.device,
61
dtype=scale_dtype,
62
)
63
- scale_ub = torch.mean(input).to(scale_dtype)
63
+ # scale_ub clamps the per-group amax of the SiLU-and-mul activation. Use
64
+ # a non-degenerate upper bound (midway between the mean and max of the
65
+ # activation magnitude) so clamping is partially active and the baseline
66
+ # comparison is meaningful. torch.mean(input) ~= 0 for the zero-mean
67
+ # input would collapse every scale to the floor and saturate the output.
68
+ # Mirrors tests/kernels/helion/test_silu_and_mul_per_block_quant.py.
69
+ act_abs = SiluAndMul.forward_native(input.to(torch.float32)).abs()
70
+ scale_ub = (0.5 * (act_abs.mean() + act_abs.amax())).to(scale_dtype)
71
72
config_key = CaseKey(
73
{