| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | 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 ( |
| 13 | _decompose_size_nodes, |
| 14 | _is_empty_allocation_node, |
| 15 | split_graph, |
| 16 | ) |
| 17 | from vllm.compilation.passes.fx_utils import find_op_nodes |
| 18 | from vllm.platforms import current_platform |
| 19 | |
| 20 | # This import automatically registers `torch.ops.silly.attention` |
| 21 | from . import silly_attention # noqa: F401 |
| 22 | |
| 23 | DEVICE_TYPE = current_platform.device_type |
| 24 | |
| 25 | |
| 26 | def test_getitem_moved_to_producer_subgraph(): |
| 27 | """Test that getitem operations are moved to the same subgraph as their input, |
| 28 | preventing tuple inputs to submodules. |
| 29 | """ |
| 30 | |
| 31 | def model_fn(x: torch.Tensor) -> torch.Tensor: |
| 32 | # torch.split returns a tuple, creating real getitem operations |
| 33 | # Should become first submodule that produces tuple |
| 34 | chunks = torch.split(x, x.shape[0] // 2, dim=0) |
| 35 | |
| 36 | # Following ops should become second submodule that consumes tuple |
| 37 | result_0 = torch.relu(chunks[0]) |
| 38 | result_1 = torch.relu(chunks[1]) |
| 39 | return torch.cat([result_0, result_1], dim=0) |
| 40 | |
| 41 | x = torch.randn(4, 3) |
| 42 | gm = make_fx(model_fn)(x) |
| 43 | |
| 44 | has_getitem = any( |
| 45 | node.op == "call_function" and node.target == operator.getitem |
| 46 | for node in gm.graph.nodes |
| 47 | ) |
| 48 | assert has_getitem, "Test setup failed: graph should contain getitem operations" |
| 49 | |
| 50 | # Split on tuple producer aten::split |
| 51 | split_ops = ["aten::split.Tensor"] |
| 52 | split_gm, split_items = split_graph(gm, split_ops) |
| 53 | assert len(split_items) == 2, "Graph should be split into 2 submodules" |
| 54 | |
| 55 | for split_item in split_items: |
| 56 | submodule = split_item.graph |
| 57 | |
| 58 | getitem_on_placeholder = [] |
| 59 | for node in submodule.graph.nodes: |
| 60 | if ( |
| 61 | node.op == "call_function" |
| 62 | and node.target == operator.getitem |
| 63 | and node.args[0].op == "placeholder" |
| 64 | ): |
| 65 | getitem_on_placeholder.append(node) |
| 66 | |
| 67 | assert len(getitem_on_placeholder) == 0, ( |
| 68 | f"Submodule {split_item.submod_name} has getitem operations on " |
| 69 | f"placeholder nodes: {[n.name for n in getitem_on_placeholder]}. " |
| 70 | "This means tuple inputs were not properly eliminated." |
| 71 | ) |
| 72 | |
| 73 | new_x = torch.randn(4, 3) |
| 74 | output_original = gm(new_x) |
| 75 | output_split = split_gm(new_x) |
| 76 | |
| 77 | assert torch.allclose(output_original, output_split), "Output mismatch" |
| 78 | |
| 79 | |
| 80 | def test_no_tuple_inputs_with_multiple_consumers(): |
| 81 | """Test that when a tuple is consumed by multiple split operations, |
| 82 | getitem operations are properly moved to avoid tuple inputs. |
| 83 | """ |
| 84 | |
| 85 | def model_fn(x: torch.Tensor) -> torch.Tensor: |
| 86 | # torch.split returns a tuple, creating real getitem operations |
| 87 | # Should become first submodule that produces tuple |
| 88 | chunks = torch.split(x, x.shape[0] // 2, dim=0) |
| 89 | |
| 90 | # These should become second submodule consuming tuple |
| 91 | result_1 = torch.relu(chunks[0]) |
| 92 | result_2 = torch.relu(chunks[1]) |
| 93 | |
| 94 | # Artificial graph splitting point to create another |
| 95 | # independent submodule that consumes tuple later |
| 96 | # This would become the third submodule |
| 97 | result_1 = torch.sigmoid(result_1) |
| 98 | |
| 99 | # Fourth submodule that consumes tuple |
| 100 | result = torch.cat([chunks[0], chunks[1], result_1, result_2]) |
| 101 | return result |
| 102 | |
| 103 | x = torch.randn(4, 3) |
| 104 | gm = make_fx(model_fn)(x) |
| 105 | |
| 106 | has_getitem = any( |
| 107 | node.op == "call_function" and node.target == operator.getitem |
| 108 | for node in gm.graph.nodes |
| 109 | ) |
| 110 | assert has_getitem, "Test setup failed: graph should contain getitem operations" |
| 111 | |
| 112 | split_ops = ["aten::split.Tensor", "aten::sigmoid"] |
| 113 | split_gm, split_items = split_graph(gm, split_ops) |
| 114 | assert len(split_items) == 4, "Graph should be split into 4 submodules" |
| 115 | |
| 116 | for split_item in split_items: |
| 117 | submodule = split_item.graph |
| 118 | |
| 119 | for node in submodule.graph.nodes: |
| 120 | if ( |
| 121 | node.op == "call_function" |
| 122 | and node.target == operator.getitem |
| 123 | and node.args[0].op == "placeholder" |
| 124 | ): |
| 125 | pytest.fail( |
| 126 | f"Submodule {split_item.submod_name} has getitem on " |
| 127 | f"placeholder {node.args[0].name}, indicating it receives " |
| 128 | "a tuple input" |
| 129 | ) |
| 130 | |
| 131 | new_x = torch.randn(4, 3) |
| 132 | output_original = gm(new_x) |
| 133 | output_split = split_gm(new_x) |
| 134 | |
| 135 | assert torch.allclose(output_original, output_split), "Output mismatch after split" |
| 136 | |
| 137 | |
| 138 | def test_consecutive_ops_in_split(): |
| 139 | """Test that consecutive splitting operations are grouped into the same subgraph.""" |
| 140 | |
| 141 | def model_fn(x: torch.Tensor) -> torch.Tensor: |
| 142 | """Define a simple model where consecutive operations create opportunities |
| 143 | for splitting subgraphs. |
| 144 | """ |
| 145 | # Apply silly attention followed by consecutive operations |
| 146 | intermediate = torch.relu(x) |
| 147 | attn_inout = torch.sqrt(intermediate) |
| 148 | torch.ops.silly.attention(intermediate, intermediate, attn_inout, attn_inout) |
| 149 | final_result = torch.sigmoid(attn_inout) |
| 150 | return final_result |
| 151 | |
| 152 | torch.set_default_device(DEVICE_TYPE) |
| 153 | |
| 154 | # Create the traced FX graph for the model |
| 155 | x = torch.randn(8, 4) |
| 156 | |
| 157 | gm = make_fx(model_fn)(x) |
| 158 | |
| 159 | # Assert presence of the expected operations in the setup |
| 160 | assert ( |
| 161 | len(list(find_op_nodes(torch.ops.aten.relu, gm.graph))) == 1 |
| 162 | and len(list(find_op_nodes(torch.ops.aten.sqrt, gm.graph))) == 1 |
| 163 | ), "Test setup failed: Expected sqrt and relu operations in the graph." |
| 164 | |
| 165 | # Configure split operations to test |
| 166 | splitting_ops = ["silly::attention", "aten::sqrt"] |
| 167 | split_gm, split_items = split_graph(gm, splitting_ops) |
| 168 | |
| 169 | # Validate the number of partitions |
| 170 | assert len(split_items) == 3, ( |
| 171 | "Consecutive splitting operations were not grouped correctly." |
| 172 | ) |
| 173 | |
| 174 | # Validate that correctness is preserved |
| 175 | new_x = torch.randn(8, 4) |
| 176 | output_original = gm(new_x) |
| 177 | output_split = split_gm(new_x) |
| 178 | assert torch.allclose(output_original, output_split), ( |
| 179 | "Output mismatch after splitting." |
| 180 | ) |
| 181 | |
| 182 | # Check the splitting item has 2 nodes exactly (relu and attn) |
| 183 | splitting_items = list(s for s in split_items if s.is_splitting_graph) |
| 184 | assert len(splitting_items) == 1, "Expecting a single splitting graph" |
| 185 | print(splitting_items[0].graph.graph) |
| 186 | splitting_gm = splitting_items[0].graph |
| 187 | assert len(splitting_gm.graph.nodes) == 4, "Expecting 4 nodes in splitting graph" |
| 188 | assert [node.op for node in splitting_gm.graph.nodes] == ["placeholder"] + 2 * [ |
| 189 | "call_function" |
| 190 | ] + ["output"] |
| 191 | |
| 192 | |
| 193 | def _get_empty_nodes(split_item): |
| 194 | return [ |
| 195 | node for node in split_item.graph.graph.nodes if _is_empty_allocation_node(node) |
| 196 | ] |
| 197 | |
| 198 | |
| 199 | def _subgraphs_with_empty_nodes(split_items, *, is_splitting_graph): |
| 200 | return [ |
| 201 | split_item |
| 202 | for split_item in split_items |
| 203 | if split_item.is_splitting_graph == is_splitting_graph |
| 204 | and _get_empty_nodes(split_item) |
| 205 | ] |
| 206 | |
| 207 | |
| 208 | def test_empty_only_partition_stays_separate_after_splitting_predecessor(): |
| 209 | """Empty-only subgraphs should not be merged when the only predecessor is |
| 210 | a splitting-op subgraph. |
| 211 | """ |
| 212 | |
| 213 | def model_fn(x: torch.Tensor) -> torch.Tensor: |
| 214 | y = torch.sin(x) |
| 215 | out = torch.empty_like(y) |
| 216 | torch.ops.aten.cos.out(y, out=out) |
| 217 | return out |
| 218 | |
| 219 | x = torch.randn(4, 3) |
| 220 | gm = make_fx(model_fn)(x) |
| 221 | |
| 222 | split_ops = ["aten::sin", "aten::cos.out"] |
| 223 | split_gm, split_items = split_graph(gm, split_ops) |
| 224 | |
| 225 | # Graph partitioning for this pattern is: |
| 226 | # [sin], [empty_like], [cos.out]. |
| 227 | assert len(split_items) == 3, ( |
| 228 | "Empty-only partition should not merge into splitting-op subgraph" |
| 229 | ) |
| 230 | |
| 231 | splitting_with_empty = _subgraphs_with_empty_nodes( |
| 232 | split_items, is_splitting_graph=True |
| 233 | ) |
| 234 | assert len(splitting_with_empty) == 0, ( |
| 235 | "Splitting-op subgraphs should not contain empty allocation nodes: " |
| 236 | f"{[item.submod_name for item in splitting_with_empty]}" |
| 237 | ) |
| 238 | |
| 239 | output_original = gm(x) |
| 240 | output_split = split_gm(x) |
| 241 | assert torch.allclose(output_original, output_split), "Output mismatch after split" |
| 242 | |
| 243 | |
| 244 | def test_empty_only_partition_is_merged(): |
| 245 | """Empty-only subgraphs should still be merged when a non-splitting predecessor |
| 246 | exists. The merged empty node must remain outside splitting-op subgraphs. |
| 247 | """ |
| 248 | |
| 249 | def model_fn(x: torch.Tensor) -> torch.Tensor: |
| 250 | base = x + 1 |
| 251 | y = torch.sin(base) |
| 252 | out = torch.empty_like(base) |
| 253 | torch.ops.aten.cos.out(base, out=out) |
| 254 | return out + y |
| 255 | |
| 256 | x = torch.randn(4, 3) |
| 257 | gm = make_fx(model_fn)(x) |
| 258 | split_gm, split_items = split_graph(gm, ["aten::sin", "aten::cos.out"]) |
| 259 | |
| 260 | # Partitioning should be: |
| 261 | # [add, empty_like], [sin], [cos.out], [add]. |
| 262 | assert len(split_items) == 4, ( |
| 263 | "Empty-only partition should be merged into non-splitting predecessor" |
| 264 | ) |
| 265 | |
| 266 | splitting_with_empty = _subgraphs_with_empty_nodes( |
| 267 | split_items, is_splitting_graph=True |
| 268 | ) |
| 269 | assert len(splitting_with_empty) == 0, ( |
| 270 | "Splitting-op subgraphs should not contain empty allocation nodes: " |
| 271 | f"{[item.submod_name for item in splitting_with_empty]}" |
| 272 | ) |
| 273 | |
| 274 | non_splitting_with_empty = _subgraphs_with_empty_nodes( |
| 275 | split_items, is_splitting_graph=False |
| 276 | ) |
| 277 | assert len(non_splitting_with_empty) == 1, ( |
| 278 | "Exactly one non-splitting subgraph should contain the merged empty node" |
| 279 | ) |
| 280 | assert len(_get_empty_nodes(non_splitting_with_empty[0])) == 1, ( |
| 281 | "Expected exactly one empty allocation node in merged subgraph" |
| 282 | ) |
| 283 | |
| 284 | output_original = gm(x) |
| 285 | output_split = split_gm(x) |
| 286 | assert torch.allclose(output_original, output_split), "Output mismatch after split" |
| 287 | |
| 288 | |
| 289 | def test_builtin_empty_only_partition_is_merged(): |
| 290 | """In Dynamo graphs, torch.empty/empty_like may appear as builtin call targets |
| 291 | (not aten OpOverload). Ensure empty-only partitions are still merged. |
| 292 | """ |
| 293 | |
| 294 | def model_fn(x: torch.Tensor) -> torch.Tensor: |
| 295 | hidden = x + 1 |
| 296 | out1 = torch.empty_like(hidden) |
| 297 | torch.ops.silly.attention(hidden, hidden, hidden, out1) |
| 298 | out2 = torch.empty_like(hidden) |
| 299 | torch.ops.silly.attention(out1, out1, hidden, out2) |
| 300 | return out2 + hidden |
| 301 | |
| 302 | gm = torch.fx.symbolic_trace(model_fn) |
| 303 | split_gm, split_items = split_graph(gm, ["silly::attention"]) |
| 304 | |
| 305 | # Without empty-only merge, this graph would split into: |
| 306 | # [add, empty_like], [attention], [empty_like], [attention], [add]. |
| 307 | assert len(split_items) == 4, "Builtin empty-only partition should be merged" |
| 308 | |
| 309 | splitting_with_empty = _subgraphs_with_empty_nodes( |
| 310 | split_items, is_splitting_graph=True |
| 311 | ) |
| 312 | assert len(splitting_with_empty) == 0, ( |
| 313 | "Splitting-op subgraphs should not contain empty allocation nodes: " |
| 314 | f"{[item.submod_name for item in splitting_with_empty]}" |
| 315 | ) |
| 316 | |
| 317 | non_splitting_with_empty = _subgraphs_with_empty_nodes( |
| 318 | split_items, is_splitting_graph=False |
| 319 | ) |
| 320 | assert len(non_splitting_with_empty) == 1, ( |
| 321 | "Exactly one non-splitting subgraph should contain merged empty nodes" |
| 322 | ) |
| 323 | assert len(_get_empty_nodes(non_splitting_with_empty[0])) == 2, ( |
| 324 | "Expected two builtin empty_like nodes in merged non-splitting subgraph" |
| 325 | ) |
| 326 | |
| 327 | x = torch.randn(2, 3, device=DEVICE_TYPE) |
| 328 | output_original = gm(x) |
| 329 | output_split = split_gm(x) |
| 330 | assert torch.allclose(output_original, output_split), "Output mismatch after split" |
| 331 | |
| 332 | |
| 333 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 334 | def test_sym_size_whole_shape_boundary(): |
| 335 | """Test that using x.size() (whole shape) across a split boundary can be |
| 336 | compiled by standalone_compile. |
| 337 | |
| 338 | The dynamo graph looks like: |
| 339 | shape = x.size() |
| 340 | y = sigmoid(x) # split point |
| 341 | z = y.clone().view(shape) |
| 342 | |
| 343 | Which splits into: |
| 344 | subgraph0(x) -> shape # returns torch.Size — problematic |
| 345 | subgraph1(x) -> y # sigmoid |
| 346 | subgraph2(y, shape) -> z # view |
| 347 | |
| 348 | Two approaches to fix the torch.Size crossing: |
| 349 | |
| 350 | Approach 1 — move sym_size to consumer (memory implication: x passed to |
| 351 | subgraph2 just for .size()): |
| 352 | subgraph0(x) -> # empty |
| 353 | subgraph1(x) -> y |
| 354 | subgraph2(y, x) -> z # computes shape locally from x |
| 355 | |
| 356 | Approach 2 — decompose shape into individual int/SymInt values: |
| 357 | subgraph0(x) -> s0, val # returns individual scalars, not Size |
| 358 | subgraph1(x) -> y |
| 359 | subgraph2(y, s0, val) -> z # reconstructs view args from scalars |
| 360 | """ |
| 361 | from torch._inductor import standalone_compile |
| 362 | |
| 363 | captured_graph = None |
| 364 | |
| 365 | def capturing_backend(gm: fx.GraphModule, example_inputs: list) -> fx.GraphModule: |
| 366 | nonlocal captured_graph |
| 367 | captured_graph = gm |
| 368 | return gm |
| 369 | |
| 370 | def model_fn(x: torch.Tensor) -> torch.Tensor: |
| 371 | shape = x.size() |
| 372 | x = torch.ops.aten.sigmoid.default(x) |
| 373 | x = x.clone().view(shape) |
| 374 | return x |
| 375 | |
| 376 | x = torch.randn(4, 8) |
| 377 | torch._dynamo.mark_dynamic(x, 0) |
| 378 | compiled_fn = torch.compile(model_fn, backend=capturing_backend) |
| 379 | compiled_fn(x) |
| 380 | |
| 381 | split_gm, split_items = split_graph(captured_graph, ["aten::sigmoid"]) |
| 382 | assert len(split_items) == 3 |
| 383 | |
| 384 | submod_0 = split_gm.submod_0 |
| 385 | example_input = torch.randn(4, 8) |
| 386 | compiled = standalone_compile( |
| 387 | submod_0, [example_input, 4], dynamic_shapes="from_example_inputs" |
| 388 | ) |
| 389 | assert compiled is not None |
| 390 | |
| 391 | |
| 392 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 393 | def test_symint_crosses_split_boundary(): |
| 394 | """Test that SymInt placeholders from torch.compile + mark_dynamic |
| 395 | cross split boundaries safely via split_module's natural threading. |
| 396 | |
| 397 | SymInt values are threaded through subgraphs by split_module and |
| 398 | handled correctly by inductor — no special replacement is needed. |
| 399 | """ |
| 400 | captured_graph = None |
| 401 | |
| 402 | def capturing_backend(gm: fx.GraphModule, example_inputs: list) -> fx.GraphModule: |
| 403 | nonlocal captured_graph |
| 404 | captured_graph = gm |
| 405 | return gm |
| 406 | |
| 407 | def model_fn(x: torch.Tensor) -> torch.Tensor: |
| 408 | batch_size = x.shape[0] |
| 409 | hidden_size = x.shape[1] |
| 410 | x = torch.ops.aten.sigmoid.default(x) |
| 411 | x = x.clone().view(batch_size, hidden_size) |
| 412 | x = torch.ops.aten.sigmoid.default(x) |
| 413 | x = x.clone().view(batch_size, hidden_size) |
| 414 | x = torch.ops.aten.sigmoid.default(x) |
| 415 | x = x.clone().view(batch_size, hidden_size) |
| 416 | return x |
| 417 | |
| 418 | x = torch.randn(4, 8) |
| 419 | torch._dynamo.mark_dynamic(x, 0) |
| 420 | |
| 421 | compiled_fn = torch.compile(model_fn, backend=capturing_backend) |
| 422 | compiled_fn(x) |
| 423 | |
| 424 | assert captured_graph is not None, "Graph should be captured by backend" |
| 425 | |
| 426 | # SymInt placeholders should exist in the captured graph |
| 427 | symint_placeholders = [ |
| 428 | node |
| 429 | for node in captured_graph.graph.nodes |
| 430 | if node.op == "placeholder" |
| 431 | and isinstance(node.meta.get("example_value"), torch.SymInt) |
| 432 | ] |
| 433 | assert len(symint_placeholders) > 0, ( |
| 434 | "Captured graph should have SymInt placeholders from mark_dynamic." |
| 435 | ) |
| 436 | |
| 437 | # split_graph should handle SymInt placeholders without error |
| 438 | split_gm, split_items = split_graph(captured_graph, ["aten::sigmoid"]) |
| 439 | |
| 440 | # Should have 3 splitting subgraphs (3 sigmoids) |
| 441 | splitting_subgraphs = [item for item in split_items if item.is_splitting_graph] |
| 442 | assert len(splitting_subgraphs) == 3, ( |
| 443 | f"Expected 3 splitting subgraphs (3 sigmoids), got {len(splitting_subgraphs)}" |
| 444 | ) |
| 445 | assert len(split_items) >= 6, ( |
| 446 | f"Expected at least 6 total subgraphs, got {len(split_items)}" |
| 447 | ) |
| 448 | |
| 449 | |
| 450 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 451 | def test_shape_boundary_standalone_compile(): |
| 452 | """Repro for the original production bug: |
| 453 | |
| 454 | AssertionError: out_spec mismatch |
| 455 | TreeSpec(tuple, None, [*, *, TreeSpec(Size, None, [*, *]), *]) |
| 456 | vs |
| 457 | TreeSpec(tuple, None, [*, *, *, *]) |
| 458 | |
| 459 | A subgraph outputs torch.Size (e.g. torch.Size([s72, 2048])) as one of |
| 460 | its values when shape info crosses a split boundary. aot_autograd / inductor |
| 461 | expect all submodule outputs to be flat tensors or scalars, not torch.Size. |
| 462 | |
| 463 | With the fix, x.size() is decomposed into individual sym_size.int calls |
| 464 | so only scalar SymInts cross the boundary — not the torch.Size. |
| 465 | """ |
| 466 | from torch._inductor import standalone_compile |
| 467 | |
| 468 | captured_graph = None |
| 469 | |
| 470 | def capturing_backend(gm: fx.GraphModule, example_inputs: list) -> fx.GraphModule: |
| 471 | nonlocal captured_graph |
| 472 | captured_graph = gm |
| 473 | return gm |
| 474 | |
| 475 | def model_fn(x: torch.Tensor) -> torch.Tensor: |
| 476 | shape = x.size() |
| 477 | x = torch.ops.aten.sigmoid.default(x) |
| 478 | x = x.clone().view(shape) |
| 479 | return x |
| 480 | |
| 481 | x = torch.randn(4, 8) |
| 482 | torch._dynamo.mark_dynamic(x, 0) |
| 483 | torch.compile(model_fn, backend=capturing_backend)(x) |
| 484 | |
| 485 | split_gm, split_items = split_graph(captured_graph, ["aten::sigmoid"]) |
| 486 | assert len(split_items) == 3 |
| 487 | |
| 488 | # Verify that the consumer subgraph only has a placeholder for the dynamic |
| 489 | # dim (SymInt) — the static dim (8) should be inlined as a literal, not |
| 490 | # threaded as a placeholder. |
| 491 | consumer = split_items[-1] # valid since len == 3: [producer, sigmoid, consumer] |
| 492 | symint_placeholders = [ |
| 493 | n |
| 494 | for n in consumer.graph.graph.nodes |
| 495 | if n.op == "placeholder" |
| 496 | and isinstance(n.meta.get("example_value"), torch.SymInt) |
| 497 | ] |
| 498 | static_int_placeholders = [ |
| 499 | n |
| 500 | for n in consumer.graph.graph.nodes |
| 501 | if n.op == "placeholder" |
| 502 | and isinstance(n.meta.get("example_value"), int) |
| 503 | and not isinstance(n.meta.get("example_value"), torch.SymInt) |
| 504 | ] |
| 505 | assert len(symint_placeholders) >= 1, ( |
| 506 | "Consumer should have a SymInt placeholder for the dynamic dim." |
| 507 | ) |
| 508 | assert len(static_int_placeholders) == 0, ( |
| 509 | "Static dims should be inlined as literals, not threaded as placeholders." |
| 510 | ) |
| 511 | |
| 512 | submod_0 = split_gm.submod_0 |
| 513 | |
| 514 | standalone_compile( |
| 515 | submod_0, [torch.randn(4, 8), 4], dynamic_shapes="from_example_inputs" |
| 516 | ) |
| 517 | |
| 518 | |
| 519 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 520 | def test_size_used_in_multiple_consumer_subgraphs(): |
| 521 | """Validates that x.size() (whole shape) used by multiple downstream subgraphs |
| 522 | does not cause torch.Size to cross split boundaries. |
| 523 | |
| 524 | Model: |
| 525 | shape = x.size() # whole shape — must not cross as torch.Size |
| 526 | z1 = sigmoid(x) # split point 1 |
| 527 | y1 = y.view(shape) # consumer 1 uses shape |
| 528 | z2 = sigmoid(z1) # split point 2 |
| 529 | y2 = y.view(shape) # consumer 2 uses shape again |
| 530 | |
| 531 | Without the fix, torch.Size crosses the boundary as a submodule output, |
| 532 | which aot_autograd / standalone_compile rejects. |
| 533 | """ |
| 534 | captured_graph = None |
| 535 | captured_inputs = None |
| 536 | |
| 537 | def capturing_backend(gm: fx.GraphModule, example_inputs: list) -> fx.GraphModule: |
| 538 | nonlocal captured_graph, captured_inputs |
| 539 | captured_graph = gm |
| 540 | captured_inputs = example_inputs |
| 541 | return gm |
| 542 | |
| 543 | def model_fn(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: |
| 544 | shape = x.size() |
| 545 | z1 = torch.ops.aten.sigmoid.default(x) |
| 546 | y1 = y.view(shape) |
| 547 | z2 = torch.ops.aten.sigmoid.default(z1) |
| 548 | y2 = y.view(shape) |
| 549 | return z2 + y1 + y2 |
| 550 | |
| 551 | x = torch.randn(4, 8) |
| 552 | y = torch.randn(4, 8) # same shape as x so view(shape) doesn't specialize dim 0 |
| 553 | torch._dynamo.mark_dynamic(x, 0) |
| 554 | torch._dynamo.mark_dynamic(y, 0) |
| 555 | torch.compile(model_fn, backend=capturing_backend)(x, y) |
| 556 | assert captured_graph is not None, "Graph should be captured by backend" |
| 557 | assert captured_inputs is not None, "Example inputs should be captured by backend" |
| 558 | |
| 559 | split_gm, split_items = split_graph(captured_graph, ["aten::sigmoid"]) |
| 560 | |
| 561 | splitting_items = [item for item in split_items if item.is_splitting_graph] |
| 562 | assert len(splitting_items) == 2 |
| 563 | |
| 564 | # Verify functional correctness — fails without the fix because torch.Size |
| 565 | # would cross a split boundary as a submodule output |
| 566 | output_original = model_fn(x, y) |
| 567 | output_split = split_gm(*captured_inputs) |
| 568 | if isinstance(output_split, tuple): |
| 569 | output_split = next(o for o in output_split if isinstance(o, torch.Tensor)) |
| 570 | assert torch.allclose(output_original, output_split), "Output mismatch after split" |
| 571 | |
| 572 | |
| 573 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") |
| 574 | def test_sym_size_metadata_propagated(): |
| 575 | """Validates that new sym_size.int nodes created by the pre-pass have |
| 576 | example_value metadata set. Without it, placeholder metadata in consumer |
| 577 | subgraphs would be None, breaking any code that dynamically builds |
| 578 | example inputs from metadata (e.g. standalone_compile per-submodule). |
| 579 | """ |
| 580 | from torch._inductor import standalone_compile |
| 581 | |
| 582 | captured_graph = None |
| 583 | |
| 584 | def capturing_backend(gm: fx.GraphModule, example_inputs: list) -> fx.GraphModule: |
| 585 | nonlocal captured_graph |
| 586 | captured_graph = gm |
| 587 | return gm |
| 588 | |
| 589 | def model_fn(x: torch.Tensor) -> torch.Tensor: |
| 590 | shape = x.size() |
| 591 | x = torch.ops.aten.sigmoid.default(x) |
| 592 | x = x.clone().view(shape) |
| 593 | return x |
| 594 | |
| 595 | x = torch.randn(4, 8) |
| 596 | torch._dynamo.mark_dynamic(x, 0) |
| 597 | torch.compile(model_fn, backend=capturing_backend)(x) |
| 598 | |
| 599 | split_gm, split_items = split_graph(captured_graph, ["aten::sigmoid"]) |
| 600 | |
| 601 | # For each submodule, build example inputs purely from placeholder metadata. |
| 602 | # This fails if example_value is None on any placeholder (i.e. metadata |
| 603 | # was not propagated to the sym_size.int nodes we created). |
| 604 | for item in split_items: |
| 605 | submod = item.graph |
| 606 | example_inputs = [] |
| 607 | for n in submod.graph.nodes: |
| 608 | if n.op != "placeholder": |
| 609 | continue |
| 610 | ev = n.meta.get("example_value") |
| 611 | assert ev is not None, ( |
| 612 | f"Placeholder '{n.name}' in {item.submod_name} has no " |
| 613 | "example_value metadata. sym_size.int nodes must propagate " |
| 614 | "metadata so consumer subgraphs can be introspected." |
| 615 | ) |
| 616 | if isinstance(ev, torch.Tensor): |
| 617 | example_inputs.append(torch.randn(*(int(d) for d in ev.shape))) |
| 618 | else: |
| 619 | example_inputs.append(int(ev)) |
| 620 | standalone_compile(submod, example_inputs, dynamic_shapes="from_example_inputs") |
| 621 | |
| 622 | |
| 623 | def test_decompose_size_with_getitem_user(): |
| 624 | """Regression test: _decompose_size_nodes must handle getitem users of size() |
| 625 | correctly. |
| 626 | |
| 627 | When a graph contains x.shape[i], it can appear as: |
| 628 | |
| 629 | %size = call_method[target="size"](args = (%x,)) |
| 630 | %getitem = call_function[target=operator.getitem](args = (%size, 1)) |
| 631 | |
| 632 | The old code spliced *all* per-dim values into every user's args |
| 633 | unconditionally, turning the 2-arg getitem into a malformed 3-arg node: |
| 634 | |
| 635 | %getitem(args = (%sym_size_int, 5120, 1)) # TypeError at runtime |
| 636 | |
| 637 | The fix detects getitem users and replaces them with dims[idx] directly. |
| 638 | """ |
| 639 | # Build a graph manually to guarantee the size() + getitem pattern. |
| 640 | # |
| 641 | # Graph: |
| 642 | # %x = placeholder |
| 643 | # %size = x.size() |
| 644 | # %dim1 = getitem(%size, 1) <-- the getitem branch we're testing |
| 645 | # %relu = relu(%x) |
| 646 | # %view = view(%relu, -1, %dim1) |
| 647 | # return %view |
| 648 | graph = fx.Graph() |
| 649 | x = graph.placeholder("x") |
| 650 | size_node = graph.call_method("size", args=(x,)) |
| 651 | getitem_node = graph.call_function(operator.getitem, args=(size_node, 1)) |
| 652 | relu_node = graph.call_function(torch.ops.aten.relu.default, args=(x,)) |
| 653 | view_node = graph.call_function( |
| 654 | torch.ops.aten.view.default, args=(relu_node, [-1, getitem_node]) |
| 655 | ) |
| 656 | graph.output(view_node) |
| 657 | |
| 658 | # Attach example_value metadata so _decompose_size_nodes can inspect dims. |
| 659 | # dim 0 is dynamic (SymInt), dim 1 is static (8). |
| 660 | from torch._dynamo.source import LocalSource |
| 661 | from torch._subclasses.fake_tensor import FakeTensorMode |
| 662 | from torch.fx.experimental.symbolic_shapes import ShapeEnv |
| 663 | |
| 664 | shape_env = ShapeEnv() |
| 665 | src = LocalSource("batch_size") |
| 666 | sym_batch = shape_env.create_symintnode(shape_env.create_symbol(4, src), hint=4) |
| 667 | fake_mode = FakeTensorMode(shape_env=shape_env) |
| 668 | with fake_mode: |
| 669 | fake_x = torch.empty_strided((sym_batch, 8), (8, 1)) |
| 670 | x.meta["example_value"] = fake_x |
| 671 | |
| 672 | gm = fx.GraphModule(torch.nn.Module(), graph) |
| 673 | |
| 674 | # Run decomposition — this would produce a 3-arg getitem without the fix |
| 675 | _decompose_size_nodes(gm) |
| 676 | |
| 677 | # Verify no size() nodes remain |
| 678 | remaining_size_nodes = list(gm.graph.find_nodes(op="call_method", target="size")) |
| 679 | assert len(remaining_size_nodes) == 0, ( |
| 680 | f"size() nodes should be fully decomposed, found {len(remaining_size_nodes)}" |
| 681 | ) |
| 682 | |
| 683 | # Verify no malformed getitem nodes (3+ args) |
| 684 | for node in gm.graph.nodes: |
| 685 | if node.op == "call_function" and node.target is operator.getitem: |
| 686 | assert len(node.args) == 2, ( |
| 687 | f"getitem node '{node.name}' has {len(node.args)} args " |
| 688 | f"(expected 2): {node.args}" |
| 689 | ) |
| 690 | |
| 691 | |
| 692 | def test_decompose_size_leaves_scalar_size_with_dim(): |
| 693 | """Regression test: _decompose_size_nodes must leave x.size(dim) alone. |
| 694 | |
| 695 | x.size() returns a torch.Size tuple that can't cross split boundaries and |
| 696 | must be decomposed. x.size(dim), however, already returns a scalar |
| 697 | SymInt/int that crosses fine, so the pass must not touch it. |
| 698 | |
| 699 | The punica LoRA path traces token_lora_mapping[:x.size(0)] under a dynamic |
| 700 | batch dim, so the size(0) node ends up nested inside a slice object: |
| 701 | |
| 702 | %size = call_method[target="size"](args = (%x, 0)) |
| 703 | %slice = call_function[target=getitem]( |
| 704 | args = (%mapping, slice(None, %size, None))) |
| 705 | |
| 706 | The old pass tried to decompose this scalar node too and then erase it, but |
| 707 | the slice still referenced it, raising "Tried to erase Node size but it |
| 708 | still had N users". The fix skips size calls that carry a dim argument. |
| 709 | """ |
| 710 | from torch._dynamo.source import LocalSource |
| 711 | from torch._subclasses.fake_tensor import FakeTensorMode |
| 712 | from torch.fx.experimental.symbolic_shapes import ShapeEnv |
| 713 | |
| 714 | # Build graph: |
| 715 | # %x = placeholder |
| 716 | # %mapping = placeholder |
| 717 | # %size = x.size(0) # scalar, with a dim arg |
| 718 | # %sliced = mapping[slice(None, %size, None)] # size node inside a slice |
| 719 | graph = fx.Graph() |
| 720 | x = graph.placeholder("x") |
| 721 | mapping = graph.placeholder("token_lora_mapping") |
| 722 | size_node = graph.call_method("size", args=(x, 0)) |
| 723 | sliced_node = graph.call_function( |
| 724 | operator.getitem, |
| 725 | args=(mapping, slice(None, size_node, None)), |
| 726 | ) |
| 727 | graph.output((sliced_node,)) |
| 728 | |
| 729 | # dim 0 dynamic (SymInt) — the realistic Unsloth + LoRA case. Without the |
| 730 | # skip, the pass would build per-dim replacements and then crash trying to |
| 731 | # erase the still-referenced size node. |
| 732 | shape_env = ShapeEnv() |
| 733 | src = LocalSource("tokens") |
| 734 | sym_tokens = shape_env.create_symintnode(shape_env.create_symbol(4, src), hint=4) |
| 735 | fake_mode = FakeTensorMode(shape_env=shape_env) |
| 736 | with fake_mode: |
| 737 | fake_x = torch.empty_strided((sym_tokens, 8), (8, 1)) |
| 738 | x.meta["example_value"] = fake_x |
| 739 | |
| 740 | gm = fx.GraphModule(torch.nn.Module(), graph) |
| 741 | |
| 742 | # Must not raise "Tried to erase Node ... still had N users". |
| 743 | _decompose_size_nodes(gm) |
| 744 | |
| 745 | # The scalar x.size(0) node is left in place, untouched. |
| 746 | remaining = list(gm.graph.find_nodes(op="call_method", target="size")) |
| 747 | assert len(remaining) == 1, ( |
| 748 | f"x.size(0) should be left untouched, found {len(remaining)} size nodes" |
| 749 | ) |
| 750 | assert remaining[0].args == (x, 0), ( |
| 751 | f"size node args changed: {remaining[0].args} (expected (x, 0))" |
| 752 | ) |