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 5, 03:30 UTC
[BugFix][Graph] fix: handle empty sym_shape_indices in PiecewiseBackend. (#39395)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Apr 15, 2026 at 21:28 UTC
db8d4a4a06ac747894534982b6e52f84eb125fab
2 files changed
+59
-5
tests/compile/test_dynamic_shapes_compilation.py
+44
@@ -222,3 +222,47 @@ def test_model_specialization_with_evaluate_guards(
222
torch.randn(1, 10).cuda(),
223
is_01_specialization=True,
224
)
225
+
226
+
227
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
228
+def test_piecewise_backend_empty_sym_shape_indices():
229
+ """Test that PiecewiseBackend handles empty sym_shape_indices correctly.
230
+
231
+ When all inputs have static shapes (no torch.SymInt), sym_shape_indices
232
+ will be empty. The fix in PiecewiseBackend.__call__ handles this case
233
+ by using the first compiled range_entry.
234
+ """
235
+ gc.collect()
236
+ torch.accelerator.empty_cache()
237
+ torch.accelerator.synchronize()
238
+
239
+ # Use small max_model_len and max_num_batched_tokens to encourage
240
+ # static shape compilation with empty sym_shape_indices
241
+ llm = LLM(
242
+ model="Qwen/Qwen3-0.6B",
243
+ max_model_len=512,
244
+ max_num_batched_tokens=1,
245
+ compilation_config={
246
+ "mode": CompilationMode.VLLM_COMPILE,
247
+ "dynamic_shapes_config": {
248
+ "type": DynamicShapesType.BACKED.value,
249
+ },
250
+ },
251
+ )
252
+
253
+ sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
254
+
255
+ # Generate with static shape inputs
256
+ output = llm.generate("Hello, my name is", sampling_params=sampling_params)
257
+ result = output[0].outputs[0].text
258
+ assert len(result) > 0, "Should generate non-empty output"
259
+
260
+ # Generate again to verify compilation works with empty sym_shape_indices
261
+ output = llm.generate("The capital of France is", sampling_params=sampling_params)
262
+ result = output[0].outputs[0].text
263
+ assert len(result) > 0, "Should generate non-empty output on second run"
264
+
265
+ del llm
266
+ gc.collect()
267
+ torch.accelerator.empty_cache()
268
+ torch.accelerator.synchronize()
vllm/compilation/piecewise_backend.py
+15
-5
@@ -354,12 +354,22 @@ class PiecewiseBackend:
354
return None
355
356
def __call__(self, *args: Any) -> Any:
357
- runtime_shape = args[self.sym_shape_indices[0]]
358
- range_entry = self._find_range_for_shape(runtime_shape)
357
+ if self.sym_shape_indices:
358
+ runtime_shape = args[self.sym_shape_indices[0]]
359
+ range_entry = self._find_range_for_shape(runtime_shape)
360
+ assert range_entry is not None, (
361
+ f"Shape: {runtime_shape} out of considered ranges: "
362
+ f"{self.compile_ranges}"
363
+ )
364
+ else:
365
+ # All inputs have static shapes; use the only compiled range_entry
366
+ compiled_entries = [re for re in self.range_entries.values() if re.compiled]
367
+ assert len(compiled_entries) == 1, (
368
+ f"Expected exactly one compiled range_entry for static shape "
369
+ f"compilation, but found {len(compiled_entries)}"
370
+ )
371
+ range_entry = compiled_entries[0]
372
360
- assert range_entry is not None, (
361
- f"Shape: {runtime_shape} out of considered ranges: {self.compile_ranges}"
362
- )
373
assert range_entry.compiled, (
374
"All ranges should be compiled or loaded up front in "
375
"PiecewiseBackend.__init__. "