[torch.compile] Canonicalize functionalized split slices for fusion pa… (#57299)
Signed-off-by: Laith Sakka <lsakka@meta.com>
Laith Sakka committed
Sep 28, 2026 at 17:21 UTC
b39f760cc705d945ab6eb56bd8dd15c931890725
2 files changed
+397
-34
tests/compile/passes/test_split_coalescing.py
+214
-6
@@ -1,6 +1,8 @@
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
@@ -29,6 +31,33 @@ class SplitCoalescingModel(torch.nn.Module):
31
return q + 1, k + 2, v + 3
32
33
34
+class SliceCoalescingModel(torch.nn.Module):
35
+ """Model with a complete QKV partition represented as contiguous slices."""
36
+
37
+ def __init__(self, q_size: int, kv_size: int) -> None:
38
+ super().__init__()
39
+ self.q_size = q_size
40
+ self.kv_size = kv_size
41
+
42
+ def forward(self, qkv: torch.Tensor):
43
+ q_end = self.q_size
44
+ k_end = q_end + self.kv_size
45
+ v_end = k_end + self.kv_size
46
+ q = torch.ops.aten.slice.Tensor(qkv, -1, 0, q_end)
47
+ k = torch.ops.aten.slice.Tensor(qkv, -1, q_end, k_end)
48
+ v = torch.ops.aten.slice.Tensor(qkv, -1, k_end, v_end)
49
+ return q + 1, k + 2, v + 3
50
+
51
+
52
+def make_vllm_config() -> VllmConfig:
53
+ return VllmConfig(
54
+ compilation_config=CompilationConfig(
55
+ mode=CompilationMode.VLLM_COMPILE,
56
+ pass_config=PassConfig(),
57
+ )
58
+ )
59
+
60
+
61
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
62
def test_split_coalescing(dtype):
63
torch.set_default_device(DEVICE_TYPE)
@@ -37,12 +66,7 @@ def test_split_coalescing(dtype):
66
67
q_size, kv_size = 2048, 512
68
40
- vllm_config = VllmConfig(
41
- compilation_config=CompilationConfig(
42
- mode=CompilationMode.VLLM_COMPILE,
43
- pass_config=PassConfig(),
44
- )
45
- )
69
+ vllm_config = make_vllm_config()
70
with vllm.config.set_current_vllm_config(vllm_config):
71
coalesce_pass = SplitCoalescingPass(vllm_config)
72
backend = TestBackend(coalesce_pass)
@@ -63,3 +87,187 @@ def test_split_coalescing(dtype):
87
torch.testing.assert_close(eager, compiled, atol=ATOL, rtol=RTOL)
88
89
assert backend.op_count(torch.ops.aten.split_with_sizes.default) == 1
90
+
91
+
92
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
93
+def test_slice_partition_canonicalization(dtype):
94
+ torch.set_default_device(DEVICE_TYPE)
95
+ torch.set_default_dtype(dtype)
96
+ torch.manual_seed(0)
97
+
98
+ q_size, kv_size = 2048, 512
99
+ vllm_config = make_vllm_config()
100
+ with vllm.config.set_current_vllm_config(vllm_config):
101
+ coalesce_pass = SplitCoalescingPass(vllm_config)
102
+ backend = TestBackend(coalesce_pass)
103
+ model = SliceCoalescingModel(q_size, kv_size)
104
+
105
+ qkv = torch.randn(5, q_size + 2 * kv_size)
106
+ torch._dynamo.mark_dynamic(qkv, 0)
107
+ result_eager = model(qkv)
108
+ result_compiled = torch.compile(model, backend=backend, fullgraph=True)(qkv)
109
+
110
+ for eager, compiled in zip(result_eager, result_compiled):
111
+ torch.testing.assert_close(eager, compiled, atol=2e-3, rtol=2e-3)
112
+
113
+ assert backend.op_count(torch.ops.aten.slice.Tensor, before=True) == 3
114
+ assert backend.op_count(torch.ops.aten.slice.Tensor) == 0
115
+ assert backend.op_count(torch.ops.aten.split_with_sizes.default) == 1
116
+
117
+ split = next(
118
+ node
119
+ for node in backend.final_graph.nodes
120
+ if node.target == torch.ops.aten.split_with_sizes.default
121
+ )
122
+ assert list(split.args[1]) == [q_size, kv_size, kv_size]
123
+ assert split.args[2] == -1
124
+
125
+ getitems = [
126
+ node
127
+ for node in split.users
128
+ if node.op == "call_function" and node.target == operator.getitem
129
+ ]
130
+ assert sorted(node.args[1] for node in getitems) == [0, 1, 2]
131
+ assert isinstance(split.meta["val"], list)
132
+ for getitem in getitems:
133
+ index = getitem.args[1]
134
+ assert getitem.meta["val"] is split.meta["val"][index]
135
+
136
+
137
+def make_slice_graph(
138
+ ranges: list[tuple[int, int, int]],
139
+) -> torch.fx.Graph:
140
+ graph = torch.fx.Graph()
141
+ source = graph.placeholder("qkv")
142
+ source.meta["val"] = torch.empty((2, 8), device="meta")
143
+
144
+ slices = []
145
+ for start, end, step in ranges:
146
+ node = graph.call_function(
147
+ torch.ops.aten.slice.Tensor,
148
+ args=(source, -1, start, end, step),
149
+ )
150
+ node.meta["val"] = torch.empty(
151
+ (2, len(range(start, min(end, 8), step))), device="meta"
152
+ )
153
+ slices.append(node)
154
+ graph.output(tuple(slices))
155
+ return graph
156
+
157
+
158
+def test_slice_partition_accepts_safe_group():
159
+ graph = make_slice_graph([(0, 4, 1), (4, 8, 1)])
160
+ vllm_config = make_vllm_config()
161
+ with vllm.config.set_current_vllm_config(vllm_config):
162
+ SplitCoalescingPass(vllm_config)(graph)
163
+
164
+ graph.lint()
165
+ assert sum(node.target == torch.ops.aten.slice.Tensor for node in graph.nodes) == 0
166
+ assert (
167
+ sum(
168
+ node.target == torch.ops.aten.split_with_sizes.default
169
+ for node in graph.nodes
170
+ )
171
+ == 1
172
+ )
173
+
174
+
175
+@pytest.mark.parametrize(
176
+ "ranges",
177
+ [
178
+ [(0, 3, 1), (4, 8, 1)], # gap
179
+ [(0, 5, 1), (4, 8, 1)], # overlap
180
+ [(1, 4, 1), (4, 8, 1)], # missing prefix
181
+ [(0, 4, 1), (4, 7, 1)], # missing suffix
182
+ [(0, 8, 1)], # single full-range slice
183
+ [(0, 8, 2)], # non-unit step
184
+ ],
185
+)
186
+def test_slice_partition_rejects_unsafe_groups(ranges):
187
+ graph = make_slice_graph(ranges)
188
+ vllm_config = make_vllm_config()
189
+ with vllm.config.set_current_vllm_config(vllm_config):
190
+ SplitCoalescingPass(vllm_config)(graph)
191
+
192
+ slice_count = sum(
193
+ node.target == torch.ops.aten.slice.Tensor for node in graph.nodes
194
+ )
195
+ split_count = sum(
196
+ node.target == torch.ops.aten.split_with_sizes.default for node in graph.nodes
197
+ )
198
+ assert slice_count == len(ranges)
199
+ assert split_count == 0
200
+
201
+
202
+def test_nested_slice_partitions_use_live_source():
203
+ """Use the rewired source when a slice partition consumes another slice."""
204
+ graph = torch.fx.Graph()
205
+ source = graph.placeholder("input")
206
+ source.meta["val"] = torch.empty((2, 8), device="meta")
207
+
208
+ outer_left = graph.call_function(
209
+ torch.ops.aten.slice.Tensor, args=(source, -1, 0, 4)
210
+ )
211
+ outer_left.meta["val"] = torch.empty((2, 4), device="meta")
212
+ outer_right = graph.call_function(
213
+ torch.ops.aten.slice.Tensor, args=(source, -1, 4, 8)
214
+ )
215
+ outer_right.meta["val"] = torch.empty((2, 4), device="meta")
216
+ inner_left = graph.call_function(
217
+ torch.ops.aten.slice.Tensor, args=(outer_left, -1, 0, 2)
218
+ )
219
+ inner_left.meta["val"] = torch.empty((2, 2), device="meta")
220
+ inner_right = graph.call_function(
221
+ torch.ops.aten.slice.Tensor, args=(outer_left, -1, 2, 4)
222
+ )
223
+ inner_right.meta["val"] = torch.empty((2, 2), device="meta")
224
+ graph.output((inner_left, inner_right, outer_right))
225
+
226
+ vllm_config = make_vllm_config()
227
+ with vllm.config.set_current_vllm_config(vllm_config):
228
+ SplitCoalescingPass(vllm_config)(graph)
229
+
230
+ graph.lint()
231
+ assert sum(node.target == torch.ops.aten.slice.Tensor for node in graph.nodes) == 0
232
+ splits = [
233
+ node
234
+ for node in graph.nodes
235
+ if node.target == torch.ops.aten.split_with_sizes.default
236
+ ]
237
+ assert len(splits) == 2
238
+
239
+ outer_split = next(node for node in splits if list(node.args[1]) == [4, 4])
240
+ inner_split = next(node for node in splits if list(node.args[1]) == [2, 2])
241
+ assert outer_split.args[0] is source
242
+ assert inner_split.args[0].target == operator.getitem
243
+ assert inner_split.args[0].args == (outer_split, 0)
244
+
245
+
246
+def test_split_coalescing_preserves_different_dimensions():
247
+ """Keep equal-sized splits on different dimensions as separate nodes."""
248
+ graph = torch.fx.Graph()
249
+ source = graph.placeholder("input")
250
+ source.meta["val"] = torch.empty((8, 8), device="meta")
251
+
252
+ outputs = []
253
+ for dim in (0, 1):
254
+ split = graph.call_function(
255
+ torch.ops.aten.split_with_sizes.default,
256
+ args=(source, [4, 4], dim),
257
+ )
258
+ for index in range(2):
259
+ outputs.append(graph.call_function(operator.getitem, args=(split, index)))
260
+ graph.output(tuple(outputs))
261
+
262
+ vllm_config = make_vllm_config()
263
+ with vllm.config.set_current_vllm_config(vllm_config):
264
+ SplitCoalescingPass(vllm_config)(graph)
265
+
266
+ splits = [
267
+ node
268
+ for node in graph.nodes
269
+ if node.target == torch.ops.aten.split_with_sizes.default
270
+ ]
271
+ assert len(splits) == 2
272
+ assert all(split.args[0] is source for split in splits)
273
+ assert {split.args[2] for split in splits} == {0, 1}
vllm/compilation/passes/utility/split_coalescing.py
+183
-28
@@ -1,21 +1,24 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Coalesce duplicate ``split_with_sizes`` nodes that operate on the same
4
-input tensor with the same split sizes.
3
+"""Canonicalize tensor partitions to a single ``split_with_sizes`` node.
4
6
-On certain hardware/dtype combinations (e.g. B200 + FP8) the Inductor
7
-graph may contain multiple ``split_with_sizes`` calls on the same tensor
8
-that CSE fails to merge. This pass detects and replaces the duplicates
9
-so that downstream pattern-matching passes (e.g. QK-Norm+RoPE fusion)
10
-see a single split node with all users attached.
5
+Functionalization may replay the individual outputs of a multi-output split as
6
+contiguous ``slice`` nodes. This pass converts complete slice partitions back
7
+to ``split_with_sizes`` and coalesces equivalent split nodes so downstream
8
+pattern-matching passes see one canonical split with all users attached.
9
10
See Also:
13
- - vLLM #33295 (original issue)
11
+ - vLLM #33295 (original duplicate-split issue)
12
- PyTorch #174472 (upstream CSE gap)
13
+ - PyTorch #194037 (multi-output views replayed as individual slices)
14
15
"""
16
17
import operator
18
+from collections import defaultdict
19
+from copy import copy
20
+from dataclasses import dataclass
21
+from typing import Any
22
23
import torch
24
from torch import fx
@@ -28,35 +31,177 @@ from ..vllm_inductor_pass import VllmInductorPass
31
logger = init_logger(__name__)
32
33
31
-class SplitCoalescingPass(VllmInductorPass):
32
- """Replace duplicate ``split_with_sizes`` nodes with a single canonical
33
- node when they share the same input tensor and split sizes."""
34
+@dataclass(frozen=True)
35
+class _SliceInfo:
36
+ node: fx.Node
37
+ dim: int
38
+ start: int
39
+ end: int
40
35
- @VllmInductorPass.time_and_log
36
- def __call__(self, graph: fx.Graph) -> None:
37
- count = 0
41
39
- # Map from input tensor node -> list of split nodes seen so far.
40
- split_nodes: dict[fx.Node, list[fx.Node]] = {}
42
+class SplitCoalescingPass(VllmInductorPass):
43
+ """Canonicalize tensor partitions to one shared ``split_with_sizes``.
44
+
45
+ A complete contiguous slice partition is rewritten from::
46
+
47
+ q = slice(qkv, 0, q_size)
48
+ k = slice(qkv, q_size, q_size + kv_size)
49
+ v = slice(qkv, q_size + kv_size, q_size + 2 * kv_size)
50
+
51
+ to::
52
+
53
+ qkv_split = split_with_sizes(qkv, [q_size, kv_size, kv_size])
54
+ q = qkv_split[0]
55
+ k = qkv_split[1]
56
+ v = qkv_split[2]
57
+
58
+ Equivalent splits are also coalesced from::
59
+
60
+ q = split_with_sizes(qkv, sizes)[0]
61
+ k = split_with_sizes(qkv, sizes)[1]
62
+
63
+ to::
64
+
65
+ qkv_split = split_with_sizes(qkv, sizes)
66
+ q = qkv_split[0]
67
+ k = qkv_split[1]
68
+ """
69
+
70
+ @staticmethod
71
+ def _get_arg(node: fx.Node, index: int, name: str, default: Any) -> Any:
72
+ """Read an FX argument by position, keyword, or schema default."""
73
+ if len(node.args) > index:
74
+ return node.args[index]
75
+ return node.kwargs.get(name, default)
76
+
77
+ @staticmethod
78
+ def _normalize_dim(dim: int, rank: int) -> int:
79
+ return dim + rank if dim < 0 else dim
80
+
81
+ @staticmethod
82
+ def _tensor_val(node: fx.Node) -> torch.Tensor | None:
83
+ value = node.meta.get("val")
84
+ return value if isinstance(value, torch.Tensor) else None
85
+
86
+ def _slice_info(self, node: fx.Node) -> _SliceInfo | None:
87
+ source = node.args[0]
88
+ source_val = self._tensor_val(source)
89
+ slice_val = self._tensor_val(node)
90
+ # in case someone forgot to set meta !
91
+ if source_val is None or slice_val is None:
92
+ return None
93
+
94
+ dim = self._normalize_dim(self._get_arg(node, 1, "dim", 0), source_val.dim())
95
+ extent = source_val.shape[dim]
96
+
97
+ # do not support dynamic slices.
98
+ if not isinstance(extent, int):
99
+ return None
100
+
101
+ start = self._get_arg(node, 2, "start", None)
102
+ end = self._get_arg(node, 3, "end", None)
103
+ step = self._get_arg(node, 4, "step", 1)
104
+ start = 0 if start is None else start
105
+ end = extent if end is None else end
106
+ if not isinstance(start, int) or not isinstance(end, int) or step != 1:
107
+ return None
108
+
109
+ if start < 0 or end < 0:
110
+ return None
111
+
112
+ # aten.slice clamps an oversized end to the selected dimension.
113
+ end = min(end, extent)
114
+ if start >= end:
115
+ return None
116
+
117
+ return _SliceInfo(node=node, dim=dim, start=start, end=end)
118
+
119
+ def _canonicalize_slices(self, graph: fx.Graph) -> int:
120
+ groups: dict[tuple[fx.Node, int], list[_SliceInfo]] = defaultdict(list)
121
+
122
+ for node in graph.find_nodes(
123
+ op="call_function", target=torch.ops.aten.slice.Tensor
124
+ ):
125
+ slice_info = self._slice_info(node)
126
+ if slice_info is None:
127
+ continue
128
+ groups[(node.args[0], slice_info.dim)].append(slice_info)
129
42
- for node in graph.nodes:
43
- if not is_func(node, torch.ops.aten.split_with_sizes.default):
130
+ count = 0
131
+ for slice_group in groups.values():
132
+ if len(slice_group) < 2:
133
continue
45
- if not all(is_func(user, operator.getitem) for user in node.users):
134
+
135
+ source = slice_group[0].node.args[0]
136
+ dim = slice_group[0].dim
137
+ source_val = self._tensor_val(source)
138
+ assert source_val is not None
139
+ extent = source_val.shape[dim]
140
+ assert isinstance(extent, int)
141
+
142
+ ordered = sorted(slice_group, key=lambda info: info.start)
143
+ # Ensure they are contiguous partitions.
144
+ if (
145
+ ordered[0].start != 0
146
+ or ordered[-1].end != extent
147
+ or any(
148
+ left.end != right.start for left, right in zip(ordered, ordered[1:])
149
+ )
150
+ ):
151
continue
152
48
- arg_node, split_sizes = node.args[:2]
153
+ split_sizes = [info.end - info.start for info in ordered]
154
+ # Use -1 for the last dimension; otherwise use the positive dimension index.
155
+ split_dim = -1 if dim == source_val.dim() - 1 else dim
156
+ first_slice = slice_group[0].node
157
+
158
+ with graph.inserting_before(first_slice):
159
+ split = graph.call_function(
160
+ torch.ops.aten.split_with_sizes.default,
161
+ args=(source, split_sizes, split_dim),
162
+ )
163
+ split.meta["val"] = [info.node.meta["val"] for info in ordered]
164
+
165
+ replacements: list[tuple[fx.Node, fx.Node]] = []
166
+ for index, info in enumerate(ordered):
167
+ getitem = graph.call_function(
168
+ operator.getitem,
169
+ args=(split, index),
170
+ type_expr=info.node.type,
171
+ )
172
+ getitem.meta = copy(info.node.meta)
173
+ replacements.append((info.node, getitem))
174
+
175
+ for slice_node, getitem in replacements:
176
+ slice_node.replace_all_uses_with(getitem)
177
+ graph.erase_node(slice_node)
178
+
179
+ count += 1
180
+
181
+ return count
182
+
183
+ def _coalesce_splits(self, graph: fx.Graph) -> int:
184
+ count = 0
185
+ split_nodes: dict[tuple[fx.Node, int], list[fx.Node]] = defaultdict(list)
186
50
- if arg_node not in split_nodes:
51
- split_nodes[arg_node] = [node]
187
+ for node in graph.find_nodes(
188
+ op="call_function", target=torch.ops.aten.split_with_sizes.default
189
+ ):
190
+ if not all(is_func(user, operator.getitem) for user in node.users):
191
continue
192
54
- # Find existing node with same split_sizes
193
+ source, split_sizes = node.args[:2]
194
+ source_val = self._tensor_val(source)
195
+ if source_val is None:
196
+ continue
197
+ raw_dim = self._get_arg(node, 2, "dim", 0)
198
+ dim = self._normalize_dim(raw_dim, source_val.dim())
199
+ candidates = split_nodes[(source, dim)]
200
canonical = next(
201
(
57
- n
58
- for n in split_nodes[arg_node]
59
- if list(n.args[1]) == list(split_sizes)
202
+ candidate
203
+ for candidate in candidates
204
+ if list(candidate.args[1]) == list(split_sizes)
205
),
206
None,
207
)
@@ -65,6 +210,16 @@ class SplitCoalescingPass(VllmInductorPass):
210
graph.erase_node(node)
211
count += 1
212
else:
68
- split_nodes[arg_node].append(node)
213
+ candidates.append(node)
214
70
- logger.debug("Coalesced %d duplicate split_with_sizes nodes", count)
215
+ return count
216
+
217
+ @VllmInductorPass.time_and_log
218
+ def __call__(self, graph: fx.Graph) -> None:
219
+ canonicalized = self._canonicalize_slices(graph)
220
+ coalesced = self._coalesce_splits(graph)
221
+ logger.debug(
222
+ "Canonicalized %d slice partitions and coalesced %d split nodes",
223
+ canonicalized,
224
+ coalesced,
225
+ )