[Kernel][Helion] Disable unsafe B200 RMS reduction warp specialization (#50345)
Signed-off-by: Shangdi Yu <shangdiy@meta.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Shangdi Yu committed
Jul 30, 2026 at 13:37 UTC
45b60e39140a5576505d986c37b382b3985a3d9e
2 files changed
+19
-6
tests/kernels/helion/test_rms_norm_per_block_quant.py
+13
@@ -136,6 +136,19 @@ class TestRmsNormPerBlockQuantConfigPicker:
136
{"hidden_size": 4096, "group_size": 128, "num_tokens": 32}
137
)
138
139
+ def test_b200_configs_disable_reduction_warp_specialization(self):
140
+ config_set = ConfigManager.get_instance().load_config_set(
141
+ "rms_norm_per_block_quant"
142
+ )
143
+ configs = config_set.to_dict()["nvidia_b200"].values()
144
+
145
+ # TODO: Remove once the Triton pin includes
146
+ # https://github.com/triton-lang/triton/pull/9716. Tracked by
147
+ # https://github.com/triton-lang/triton/issues/10901.
148
+ assert all(
149
+ config["range_warp_specializes"][1] is not True for config in configs
150
+ )
151
+
152
153
DTYPES = [torch.bfloat16, torch.float]
154
QUANT_DTYPES = [torch.int8, FP8_DTYPE]
vllm/kernels/helion/configs/rms_norm_per_block_quant/nvidia_b200.json
+6
-6
@@ -234,7 +234,7 @@
234
],
235
"range_warp_specializes": [
236
null,
237
- true,
237
+ false,
238
null,
239
null
240
],
@@ -1354,7 +1354,7 @@
1354
],
1355
"range_warp_specializes": [
1356
null,
1357
- true,
1357
+ false,
1358
null,
1359
null
1360
],
@@ -1564,7 +1564,7 @@
1564
],
1565
"range_warp_specializes": [
1566
null,
1567
- true,
1567
+ false,
1568
null,
1569
null
1570
],
@@ -1634,7 +1634,7 @@
1634
],
1635
"range_warp_specializes": [
1636
null,
1637
- true,
1637
+ false,
1638
null,
1639
null
1640
],
@@ -2546,7 +2546,7 @@
2546
],
2547
"range_warp_specializes": [
2548
null,
2549
- true,
2549
+ false,
2550
null,
2551
null
2552
],
@@ -2686,7 +2686,7 @@
2686
],
2687
"range_warp_specializes": [
2688
null,
2689
- true,
2689
+ false,
2690
null,
2691
null
2692
],