@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ea2c148fa7
Fork 0

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, 09:30 UTC

[compile][graph_partition]Add tensor size handling (#36038)

Signed-off-by: Xiao Fu <xiaofu@meta.com>

Xiao committed Mar 19, 2026 at 19:55 UTC ea2c148fa7f3e579c68de15b19435079f3b00bef
2 files changed +351
tests/compile/test_graph_partition.py
+295
@@ -5,6 +5,8 @@ import operator
5
6 import pytest
7 import torch
8 +import torch._dynamo
9 +import torch.fx as fx
10 from torch.fx.experimental.proxy_tensor import make_fx
11
12 from vllm.compilation.backends import _is_empty_allocation_node, split_graph
@@ -327,3 +329,296 @@ def test_builtin_empty_only_partition_is_merged():
329 output_original = gm(x)
330 output_split = split_gm(x)
331 assert torch.allclose(output_original, output_split), "Output mismatch after split"
332 +
333 +
334 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
335 +def test_sym_size_whole_shape_boundary():
336 + """
337 + Test that using x.size() (whole shape) across a split boundary can be
338 + compiled by standalone_compile.
339 +
340 + The dynamo graph looks like:
341 + shape = x.size()
342 + y = sigmoid(x) # split point
343 + z = y.clone().view(shape)
344 +
345 + Which splits into:
346 + subgraph0(x) -> shape # returns torch.Size — problematic
347 + subgraph1(x) -> y # sigmoid
348 + subgraph2(y, shape) -> z # view
349 +
350 + Two approaches to fix the torch.Size crossing:
351 +
352 + Approach 1 — move sym_size to consumer (memory implication: x passed to
353 + subgraph2 just for .size()):
354 + subgraph0(x) -> # empty
355 + subgraph1(x) -> y
356 + subgraph2(y, x) -> z # computes shape locally from x
357 +
358 + Approach 2 — decompose shape into individual int/SymInt values:
359 + subgraph0(x) -> s0, val # returns individual scalars, not Size
360 + subgraph1(x) -> y
361 + subgraph2(y, s0, val) -> z # reconstructs view args from scalars
362 + """
363 + from torch._inductor import standalone_compile
364 +
365 + captured_graph = None
366 +
367 + def capturing_backend(gm: fx.GraphModule, example_inputs: list) -> fx.GraphModule:
368 + nonlocal captured_graph
369 + captured_graph = gm
370 + return gm
371 +
372 + def model_fn(x: torch.Tensor) -> torch.Tensor:
373 + shape = x.size()
374 + x = torch.ops.aten.sigmoid.default(x)
375 + x = x.clone().view(shape)
376 + return x
377 +
378 + x = torch.randn(4, 8)
379 + torch._dynamo.mark_dynamic(x, 0)
380 + compiled_fn = torch.compile(model_fn, backend=capturing_backend)
381 + compiled_fn(x)
382 +
383 + split_gm, split_items = split_graph(captured_graph, ["aten::sigmoid"])
384 + assert len(split_items) == 3
385 +
386 + submod_0 = split_gm.submod_0
387 + example_input = torch.randn(4, 8)
388 + compiled = standalone_compile(
389 + submod_0, [example_input, 4], dynamic_shapes="from_example_inputs"
390 + )
391 + assert compiled is not None
392 +
393 +
394 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
395 +def test_symint_crosses_split_boundary():
396 + """
397 + Test that SymInt placeholders from torch.compile + mark_dynamic
398 + cross split boundaries safely via split_module's natural threading.
399 +
400 + SymInt values are threaded through subgraphs by split_module and
401 + handled correctly by inductor — no special replacement is needed.
402 + """
403 + captured_graph = None
404 +
405 + def capturing_backend(gm: fx.GraphModule, example_inputs: list) -> fx.GraphModule:
406 + nonlocal captured_graph
407 + captured_graph = gm
408 + return gm
409 +
410 + def model_fn(x: torch.Tensor) -> torch.Tensor:
411 + batch_size = x.shape[0]
412 + hidden_size = x.shape[1]
413 + x = torch.ops.aten.sigmoid.default(x)
414 + x = x.clone().view(batch_size, hidden_size)
415 + x = torch.ops.aten.sigmoid.default(x)
416 + x = x.clone().view(batch_size, hidden_size)
417 + x = torch.ops.aten.sigmoid.default(x)
418 + x = x.clone().view(batch_size, hidden_size)
419 + return x
420 +
421 + x = torch.randn(4, 8)
422 + torch._dynamo.mark_dynamic(x, 0)
423 +
424 + compiled_fn = torch.compile(model_fn, backend=capturing_backend)
425 + compiled_fn(x)
426 +
427 + assert captured_graph is not None, "Graph should be captured by backend"
428 +
429 + # SymInt placeholders should exist in the captured graph
430 + symint_placeholders = [
431 + node
432 + for node in captured_graph.graph.nodes
433 + if node.op == "placeholder"
434 + and isinstance(node.meta.get("example_value"), torch.SymInt)
435 + ]
436 + assert len(symint_placeholders) > 0, (
437 + "Captured graph should have SymInt placeholders from mark_dynamic."
438 + )
439 +
440 + # split_graph should handle SymInt placeholders without error
441 + split_gm, split_items = split_graph(captured_graph, ["aten::sigmoid"])
442 +
443 + # Should have 3 splitting subgraphs (3 sigmoids)
444 + splitting_subgraphs = [item for item in split_items if item.is_splitting_graph]
445 + assert len(splitting_subgraphs) == 3, (
446 + f"Expected 3 splitting subgraphs (3 sigmoids), got {len(splitting_subgraphs)}"
447 + )
448 + assert len(split_items) >= 6, (
449 + f"Expected at least 6 total subgraphs, got {len(split_items)}"
450 + )
451 +
452 +
453 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
454 +def test_shape_boundary_standalone_compile():
455 + """
456 + Repro for the original production bug:
457 +
458 + AssertionError: out_spec mismatch
459 + TreeSpec(tuple, None, [*, *, TreeSpec(Size, None, [*, *]), *])
460 + vs
461 + TreeSpec(tuple, None, [*, *, *, *])
462 +
463 + A subgraph outputs torch.Size (e.g. torch.Size([s72, 2048])) as one of
464 + its values when shape info crosses a split boundary. aot_autograd / inductor
465 + expect all submodule outputs to be flat tensors or scalars, not torch.Size.
466 +
467 + With the fix, x.size() is decomposed into individual sym_size.int calls
468 + so only scalar SymInts cross the boundary — not the torch.Size.
469 + """
470 + from torch._inductor import standalone_compile
471 +
472 + captured_graph = None
473 +
474 + def capturing_backend(gm: fx.GraphModule, example_inputs: list) -> fx.GraphModule:
475 + nonlocal captured_graph
476 + captured_graph = gm
477 + return gm
478 +
479 + def model_fn(x: torch.Tensor) -> torch.Tensor:
480 + shape = x.size()
481 + x = torch.ops.aten.sigmoid.default(x)
482 + x = x.clone().view(shape)
483 + return x
484 +
485 + x = torch.randn(4, 8)
486 + torch._dynamo.mark_dynamic(x, 0)
487 + torch.compile(model_fn, backend=capturing_backend)(x)
488 +
489 + split_gm, split_items = split_graph(captured_graph, ["aten::sigmoid"])
490 + assert len(split_items) == 3
491 +
492 + # Verify that the consumer subgraph only has a placeholder for the dynamic
493 + # dim (SymInt) — the static dim (8) should be inlined as a literal, not
494 + # threaded as a placeholder.
495 + consumer = split_items[-1] # valid since len == 3: [producer, sigmoid, consumer]
496 + symint_placeholders = [
497 + n
498 + for n in consumer.graph.graph.nodes
499 + if n.op == "placeholder"
500 + and isinstance(n.meta.get("example_value"), torch.SymInt)
501 + ]
502 + static_int_placeholders = [
503 + n
504 + for n in consumer.graph.graph.nodes
505 + if n.op == "placeholder"
506 + and isinstance(n.meta.get("example_value"), int)
507 + and not isinstance(n.meta.get("example_value"), torch.SymInt)
508 + ]
509 + assert len(symint_placeholders) >= 1, (
510 + "Consumer should have a SymInt placeholder for the dynamic dim."
511 + )
512 + assert len(static_int_placeholders) == 0, (
513 + "Static dims should be inlined as literals, not threaded as placeholders."
514 + )
515 +
516 + submod_0 = split_gm.submod_0
517 +
518 + standalone_compile(
519 + submod_0, [torch.randn(4, 8), 4], dynamic_shapes="from_example_inputs"
520 + )
521 +
522 +
523 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
524 +def test_size_used_in_multiple_consumer_subgraphs():
525 + """
526 + Validates that x.size() (whole shape) used by multiple downstream subgraphs
527 + does not cause torch.Size to cross split boundaries.
528 +
529 + Model:
530 + shape = x.size() # whole shape — must not cross as torch.Size
531 + z1 = sigmoid(x) # split point 1
532 + y1 = y.view(shape) # consumer 1 uses shape
533 + z2 = sigmoid(z1) # split point 2
534 + y2 = y.view(shape) # consumer 2 uses shape again
535 +
536 + Without the fix, torch.Size crosses the boundary as a submodule output,
537 + which aot_autograd / standalone_compile rejects.
538 + """
539 + captured_graph = None
540 + captured_inputs = None
541 +
542 + def capturing_backend(gm: fx.GraphModule, example_inputs: list) -> fx.GraphModule:
543 + nonlocal captured_graph, captured_inputs
544 + captured_graph = gm
545 + captured_inputs = example_inputs
546 + return gm
547 +
548 + def model_fn(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
549 + shape = x.size()
550 + z1 = torch.ops.aten.sigmoid.default(x)
551 + y1 = y.view(shape)
552 + z2 = torch.ops.aten.sigmoid.default(z1)
553 + y2 = y.view(shape)
554 + return z2 + y1 + y2
555 +
556 + x = torch.randn(4, 8)
557 + y = torch.randn(4, 8) # same shape as x so view(shape) doesn't specialize dim 0
558 + torch._dynamo.mark_dynamic(x, 0)
559 + torch._dynamo.mark_dynamic(y, 0)
560 + torch.compile(model_fn, backend=capturing_backend)(x, y)
561 +
562 + split_gm, split_items = split_graph(captured_graph, ["aten::sigmoid"])
563 +
564 + splitting_items = [item for item in split_items if item.is_splitting_graph]
565 + assert len(splitting_items) == 2
566 +
567 + # Verify functional correctness — fails without the fix because torch.Size
568 + # would cross a split boundary as a submodule output
569 + output_original = model_fn(x, y)
570 + output_split = split_gm(*captured_inputs)
571 + if isinstance(output_split, tuple):
572 + output_split = next(o for o in output_split if isinstance(o, torch.Tensor))
573 + assert torch.allclose(output_original, output_split), "Output mismatch after split"
574 +
575 +
576 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
577 +def test_sym_size_metadata_propagated():
578 + """
579 + Validates that new sym_size.int nodes created by the pre-pass have
580 + example_value metadata set. Without it, placeholder metadata in consumer
581 + subgraphs would be None, breaking any code that dynamically builds
582 + example inputs from metadata (e.g. standalone_compile per-submodule).
583 + """
584 + from torch._inductor import standalone_compile
585 +
586 + captured_graph = None
587 +
588 + def capturing_backend(gm: fx.GraphModule, example_inputs: list) -> fx.GraphModule:
589 + nonlocal captured_graph
590 + captured_graph = gm
591 + return gm
592 +
593 + def model_fn(x: torch.Tensor) -> torch.Tensor:
594 + shape = x.size()
595 + x = torch.ops.aten.sigmoid.default(x)
596 + x = x.clone().view(shape)
597 + return x
598 +
599 + x = torch.randn(4, 8)
600 + torch._dynamo.mark_dynamic(x, 0)
601 + torch.compile(model_fn, backend=capturing_backend)(x)
602 +
603 + split_gm, split_items = split_graph(captured_graph, ["aten::sigmoid"])
604 +
605 + # For each submodule, build example inputs purely from placeholder metadata.
606 + # This fails if example_value is None on any placeholder (i.e. metadata
607 + # was not propagated to the sym_size.int nodes we created).
608 + for item in split_items:
609 + submod = item.graph
610 + example_inputs = []
611 + for n in submod.graph.nodes:
612 + if n.op != "placeholder":
613 + continue
614 + ev = n.meta.get("example_value")
615 + assert ev is not None, (
616 + f"Placeholder '{n.name}' in {item.submod_name} has no "
617 + "example_value metadata. sym_size.int nodes must propagate "
618 + "metadata so consumer subgraphs can be introspected."
619 + )
620 + if isinstance(ev, torch.Tensor):
621 + example_inputs.append(torch.randn(*(int(d) for d in ev.shape)))
622 + else:
623 + example_inputs.append(int(ev))
624 + standalone_compile(submod, example_inputs, dynamic_shapes="from_example_inputs")
vllm/compilation/backends.py
+56
@@ -473,9 +473,65 @@ def _merge_empty_only_subgraphs(
473 prev_non_splitting_subgraph_id = subgraph_id
474
475
476 +def _decompose_size_nodes(graph: fx.GraphModule) -> None:
477 + """Decompose x.size() into per-dim sym_size.int calls.
478 +
479 + torch.Size objects cannot cross split boundaries because aot_autograd
480 + cannot handle them as submodule outputs. This replaces each size() call
481 + with individual sym_size.int(x, dim) nodes:
482 + - Dynamic dims (SymInt) → new sym_size.int node
483 + - Static dims (plain int) → inlined as literal constant
484 + """
485 + # Dynamo captures x.size()/x.shape as call_method target="size".
486 + size_nodes = list(graph.graph.find_nodes(op="call_method", target="size"))
487 +
488 + for node in size_nodes:
489 + tensor_node = node.args[0]
490 + ev = tensor_node.meta.get("example_value")
491 + assert ev is not None, (
492 + f"Tensor node '{tensor_node.name}' has no example_value metadata. "
493 + f"Cannot decompose size node '{node.name}'."
494 + )
495 +
496 + # Build per-dim replacements: sym_size.int node or literal int.
497 + dims: list[fx.Node | int] = []
498 + with graph.graph.inserting_after(tensor_node):
499 + for i in range(ev.dim()):
500 + dim_val = ev.shape[i]
501 + if isinstance(dim_val, torch.SymInt):
502 + dn = graph.graph.call_function(
503 + torch.ops.aten.sym_size.int, args=(tensor_node, i)
504 + )
505 + dn.meta["example_value"] = dim_val
506 + dims.append(dn)
507 + elif isinstance(dim_val, int):
508 + dims.append(dim_val)
509 + else:
510 + raise AssertionError(
511 + f"dim_val is either torch.SymInt or int, "
512 + f"got {type(dim_val)} for dim {i} of "
513 + f"'{node.name}'"
514 + )
515 +
516 + # Replace size node in each user's args.
517 + # Dynamo always passes size as a direct arg: view(clone, size)
518 + # → view(clone, d0, d1, ...)
519 + for user in list(node.users):
520 + new_args = []
521 + for arg in user.args:
522 + if arg is node:
523 + new_args.extend(dims)
524 + else:
525 + new_args.append(arg)
526 + user.args = tuple(new_args)
527 + graph.graph.erase_node(node)
528 +
529 +
530 def split_graph(
531 graph: fx.GraphModule, splitting_ops: list[str]
532 ) -> tuple[fx.GraphModule, list[SplitItem]]:
533 + _decompose_size_nodes(graph)
534 +
535 # split graph by ops
536 subgraph_id = 0
537 node_to_subgraph_id: dict[fx.Node, int] = {}