| 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 | |
| 9 | import vllm |
| 10 | from tests.compile.backend import TestBackend |
| 11 | from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass |
| 12 | from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig |
| 13 | from vllm.platforms import current_platform |
| 14 | |
| 15 | DEVICE_TYPE = current_platform.device_type |
| 16 | |
| 17 | |
| 18 | class SplitCoalescingModel(torch.nn.Module): |
| 19 | """Model with 3 separate split_with_sizes calls on the same input, |
| 20 | simulating the B200+FP8 graph where CSE fails to merge them.""" |
| 21 | |
| 22 | def __init__(self, q_size: int, kv_size: int) -> None: |
| 23 | super().__init__() |
| 24 | self.q_size = q_size |
| 25 | self.kv_size = kv_size |
| 26 | |
| 27 | def forward(self, qkv: torch.Tensor): |
| 28 | q, _, _ = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) |
| 29 | _, k, _ = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) |
| 30 | _, _, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) |
| 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) |
| 64 | torch.set_default_dtype(dtype) |
| 65 | torch.manual_seed(0) |
| 66 | |
| 67 | q_size, kv_size = 2048, 512 |
| 68 | |
| 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) |
| 73 | |
| 74 | model = SplitCoalescingModel(q_size, kv_size) |
| 75 | |
| 76 | T = 5 |
| 77 | qkv = torch.randn(T, q_size + 2 * kv_size) |
| 78 | torch._dynamo.mark_dynamic(qkv, 0) |
| 79 | |
| 80 | result_eager = model(qkv) |
| 81 | |
| 82 | model_compiled = torch.compile(model, backend=backend) |
| 83 | result_compiled = model_compiled(qkv) |
| 84 | |
| 85 | ATOL, RTOL = (2e-3, 2e-3) |
| 86 | for eager, compiled in zip(result_eager, result_compiled): |
| 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} |