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 6, 20:30 UTC
[Kernel][Helion][13/N] Force static_shapes=False in helion register (#36677)
Signed-off-by: Yanan Cao <gmagogsfm@gmail.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Yanan Cao committed
Mar 11, 2026 at 22:35 UTC
584a3f56deb35f3b25bdb079fabc88f453bdfe1e
2 files changed
+9
-66
tests/kernels/helion/test_register.py
+4
-57
@@ -134,14 +134,14 @@ class TestValidateHelionSettings:
134
validate_helion_settings(settings, "test_kernel")
135
136
def test_warns_on_static_shapes_true(self):
137
- """Test that static_shapes=True emits a warning."""
137
+ """Test that static_shapes=True emits a warning about being overridden."""
138
settings = helion.Settings()
139
settings.static_shapes = True
140
141
with patch("vllm.kernels.helion.register.logger") as mock_logger:
142
validate_helion_settings(settings, "test_kernel")
143
mock_logger.warning.assert_called_once()
144
- assert "static_shapes=True" in mock_logger.warning.call_args[0][0]
144
+ assert "overridden to False" in mock_logger.warning.call_args[0][0]
145
146
147
def create_configured_kernel_with_configs(
@@ -259,7 +259,6 @@ class TestConfiguredHelionKernel:
259
260
settings = helion.Settings()
261
settings.print_output_code = True
262
- # Note: helion.Settings() defaults static_shapes to True
262
263
mock_config_manager = Mock(spec=ConfigManager)
264
mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
@@ -288,46 +287,8 @@ class TestConfiguredHelionKernel:
287
call_kwargs = mock_kernel.call_args[1]
288
assert "print_output_code" in call_kwargs
289
assert call_kwargs["print_output_code"] is True
291
- # helion.Settings() defaults to static_shapes=True, so it should remain True
292
- assert call_kwargs["static_shapes"] is True
293
-
294
- def test_create_decorated_kernel_preserves_static_shapes_true(
295
- self, sample_kernel, sample_configs
296
- ):
297
- """Test that explicit static_shapes=True is preserved."""
298
-
299
- def default_picker(args, config_keys):
300
- return "default"
301
-
302
- settings = helion.Settings()
303
- settings.static_shapes = True
304
-
305
- mock_config_manager = Mock(spec=ConfigManager)
306
- mock_config_manager.get_platform_configs = Mock(return_value=sample_configs)
307
-
308
- with (
309
- patch("vllm.kernels.helion.register.helion.kernel") as mock_kernel,
310
- patch(
311
- "vllm.kernels.helion.config_manager.ConfigManager.get_instance",
312
- return_value=mock_config_manager,
313
- ),
314
- patch(
315
- "vllm.kernels.helion.utils.get_canonical_gpu_name",
316
- return_value="nvidia_h200",
317
- ),
318
- ):
319
- mock_decorated = Mock()
320
- mock_kernel.return_value = Mock(return_value=mock_decorated)
321
-
322
- ConfiguredHelionKernel(
323
- op_name="test_kernel",
324
- config_picker=default_picker,
325
- raw_kernel_func=sample_kernel,
326
- helion_settings=settings,
327
- )
328
-
329
- call_kwargs = mock_kernel.call_args[1]
330
- assert call_kwargs["static_shapes"] is True
290
+ # static_shapes is always forced to False by vLLM
291
+ assert call_kwargs["static_shapes"] is False
292
293
def test_key_and_config_selector_use_same_logic(
294
self, sample_kernel, sample_configs
@@ -761,20 +722,6 @@ class TestKernelRegistry:
722
def test_kernel(x):
723
return x
724
764
- def test_register_kernel_warns_with_static_shapes_true(self):
765
- """Test register_kernel warns when static_shapes=True."""
766
- mock_settings = Mock()
767
- mock_settings.to_dict.return_value = {"static_shapes": True}
768
-
769
- with patch("vllm.kernels.helion.register.logger") as mock_logger:
770
-
771
- @register_kernel("test", helion_settings=mock_settings)
772
- def test_kernel(x):
773
- return x
774
-
775
- mock_logger.warning.assert_called_once()
776
- assert "static_shapes=True" in mock_logger.warning.call_args[0][0]
777
-
725
def test_register_kernel_no_warning_with_static_shapes_false(self):
726
"""Test register_kernel doesn't warn with static_shapes=False."""
727
mock_settings = Mock()
vllm/kernels/helion/register.py
+5
-9
@@ -98,13 +98,11 @@ def validate_helion_settings(
98
f"@{op_name}.register_config_picker instead."
99
)
100
101
- # Warn if static_shapes is explicitly set to True since most vLLM ops need
102
- # dynamic shapes for variable batch sizes and sequence lengths
101
if settings_dict.get("static_shapes") is True:
102
logger.warning(
105
- "Kernel '%s' has static_shapes=True in helion_settings. "
106
- "Most vLLM ops require dynamic shapes for variable batch sizes "
107
- "and sequence lengths. Consider removing this setting.",
103
+ "Kernel '%s' has static_shapes=True in helion_settings, "
104
+ "which will be overridden to False. vLLM requires dynamic "
105
+ "shapes for variable batch sizes and sequence lengths.",
106
op_name,
107
)
108
@@ -118,10 +116,8 @@ def create_helion_decorated_kernel(
116
if helion_settings:
117
kernel_kwargs.update(helion_settings.to_dict())
118
121
- # Set static_shapes=False by default if user didn't explicitly set it
122
- # This is needed for dynamic batch sizes and sequence lengths in vLLM
123
- if kernel_kwargs.get("static_shapes") is not True:
124
- kernel_kwargs["static_shapes"] = False
119
+ # vLLM requires dynamic shapes for variable batch sizes and sequence lengths
120
+ kernel_kwargs["static_shapes"] = False
121
122
if extra_kwargs:
123
kernel_kwargs.update(extra_kwargs)