[Helion] Fix degenerate scale_ub in kernel input generators (#48868)

Signed-off-by: Shangdi Yu <shangdiy@meta.com>

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)

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)

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(
             {

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(
             {