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 8, 17:45 UTC
[Bugfix] Fix QK Norm+RoPE fusion pattern matching on B200+FP8 (#33967)
Signed-off-by: Ikenna <ikennachifo@gmail.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Ikenna committed
Feb 6, 2026 at 21:27 UTC
906077181b21128576018604562e55fbc7f70c34
5 files changed
+154
-6
tests/compile/fusions_e2e/models.py
+1
-3
@@ -101,9 +101,7 @@ qwen3_a3b_fp8 = ModelFusionInfo(
101
model_name="Qwen/Qwen3-30B-A3B-FP8",
102
matches=lambda n_layers: Matches(
103
rms_quant_fusion=n_layers,
104
- # TODO broken on Blackwell:
105
- # https://github.com/vllm-project/vllm/issues/33295
106
- norm_rope_fusion=0 if is_blackwell() else n_layers,
104
+ norm_rope_fusion=n_layers,
105
attn_quant_fusion=0, # attn + group quant not supported
106
ar_rms_fusion=n_layers * 2 + 1,
107
sequence_parallel=n_layers * 2 + 1,
tests/compile/passes/test_qk_norm_rope_fusion.py
+19
-3
@@ -16,6 +16,7 @@ from vllm.compilation.passes.fusion.qk_norm_rope_fusion import (
16
)
17
from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass
18
from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
19
+from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
20
from vllm.config import (
21
CompilationConfig,
22
CompilationMode,
@@ -45,6 +46,7 @@ class QKNormRoPETestModel(torch.nn.Module):
46
is_neox: bool,
47
vllm_config: VllmConfig,
48
dtype: torch.dtype,
49
+ test_scattered_split: bool = False,
50
prefix: str = "model.layers.0.self_attn.attn",
51
) -> None:
52
super().__init__()
@@ -78,11 +80,17 @@ class QKNormRoPETestModel(torch.nn.Module):
80
is_neox_style=is_neox,
81
dtype=self.dtype,
82
)
83
+ self.test_scattered_split = test_scattered_split
84
self.enable_rms_norm_custom_op = self.q_norm.enabled()
85
self.enable_rope_custom_op = self.rotary_emb.enabled()
86
87
def forward(self, qkv: torch.Tensor, positions: torch.Tensor):
85
- q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
88
+ if self.test_scattered_split:
89
+ q, _, _ = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
90
+ _, k, _ = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
91
+ _, _, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
92
+ else:
93
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
94
q_by_head = q.view(*q.shape[:-1], q.shape[-1] // self.head_dim, self.head_dim)
95
q_by_head = self.q_norm(q_by_head)
96
q = q_by_head.view(q.shape)
@@ -112,6 +120,7 @@ class QKNormRoPETestModel(torch.nn.Module):
120
return [FUSED_QK_ROPE_OP]
121
122
123
+@pytest.mark.parametrize("scattered_split", [True, False])
124
@pytest.mark.parametrize("eps", [1e-5, 1e-6])
125
@pytest.mark.parametrize("is_neox", [True, False])
126
@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
@@ -122,7 +131,12 @@ class QKNormRoPETestModel(torch.nn.Module):
131
reason="Only test on cuda and rocm platform",
132
)
133
def test_qk_norm_rope_fusion(
125
- eps, is_neox, enable_rms_norm_custom_op, enable_rope_custom_op, dtype
134
+ eps,
135
+ is_neox,
136
+ enable_rms_norm_custom_op,
137
+ enable_rope_custom_op,
138
+ dtype,
139
+ scattered_split,
140
):
141
if not hasattr(torch.ops._C, "fused_qk_norm_rope"):
142
pytest.skip("fused_qk_norm_rope custom op not available")
@@ -161,13 +175,15 @@ def test_qk_norm_rope_fusion(
175
is_neox=is_neox,
176
vllm_config=vllm_config,
177
dtype=dtype,
178
+ test_scattered_split=scattered_split,
179
)
180
181
noop_pass = NoOpEliminationPass(vllm_config)
182
+ coalesce_pass = SplitCoalescingPass(vllm_config)
183
fusion_pass = QKNormRoPEFusionPass(vllm_config)
184
cleanup_pass = PostCleanupPass(vllm_config)
185
170
- backend = TestBackend(noop_pass, fusion_pass, cleanup_pass)
186
+ backend = TestBackend(noop_pass, coalesce_pass, fusion_pass, cleanup_pass)
187
backend_baseline = TestBackend(noop_pass, cleanup_pass)
188
189
qkv = torch.randn(T, model.q_size + 2 * model.kv_size)
tests/compile/passes/test_split_coalescing.py
new
+62
@@ -0,0 +1,62 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+import vllm
8
+from tests.compile.backend import TestBackend
9
+from vllm.compilation.passes.utility.split_coalescing import SplitCoalescingPass
10
+from vllm.config import CompilationConfig, CompilationMode, PassConfig, VllmConfig
11
+
12
+
13
+class SplitCoalescingModel(torch.nn.Module):
14
+ """Model with 3 separate split_with_sizes calls on the same input,
15
+ simulating the B200+FP8 graph where CSE fails to merge them."""
16
+
17
+ def __init__(self, q_size: int, kv_size: int) -> None:
18
+ super().__init__()
19
+ self.q_size = q_size
20
+ self.kv_size = kv_size
21
+
22
+ def forward(self, qkv: torch.Tensor):
23
+ q, _, _ = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
24
+ _, k, _ = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
25
+ _, _, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
26
+ return q + 1, k + 2, v + 3
27
+
28
+
29
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
30
+def test_split_coalescing(dtype):
31
+ torch.set_default_device("cuda")
32
+ torch.set_default_dtype(dtype)
33
+ torch.manual_seed(0)
34
+
35
+ q_size, kv_size = 2048, 512
36
+
37
+ vllm_config = VllmConfig(
38
+ compilation_config=CompilationConfig(
39
+ mode=CompilationMode.VLLM_COMPILE,
40
+ pass_config=PassConfig(),
41
+ )
42
+ )
43
+ with vllm.config.set_current_vllm_config(vllm_config):
44
+ coalesce_pass = SplitCoalescingPass(vllm_config)
45
+ backend = TestBackend(coalesce_pass)
46
+
47
+ model = SplitCoalescingModel(q_size, kv_size)
48
+
49
+ T = 5
50
+ qkv = torch.randn(T, q_size + 2 * kv_size)
51
+ torch._dynamo.mark_dynamic(qkv, 0)
52
+
53
+ result_eager = model(qkv)
54
+
55
+ model_compiled = torch.compile(model, backend=backend)
56
+ result_compiled = model_compiled(qkv)
57
+
58
+ ATOL, RTOL = (2e-3, 2e-3)
59
+ for eager, compiled in zip(result_eager, result_compiled):
60
+ torch.testing.assert_close(eager, compiled, atol=ATOL, rtol=RTOL)
61
+
62
+ assert backend.op_count(torch.ops.aten.split_with_sizes.default) == 1
vllm/compilation/passes/pass_manager.py
+2
@@ -29,6 +29,7 @@ if current_platform.is_cuda_alike():
29
from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass
30
from .fusion.rms_quant_fusion import RMSNormQuantFusionPass
31
from .fusion.sequence_parallelism import SequenceParallelismPass
32
+ from .utility.split_coalescing import SplitCoalescingPass
33
34
if current_platform.is_cuda():
35
from .fusion.allreduce_rms_fusion import AllReduceFusionPass
@@ -139,6 +140,7 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
140
self.passes += [AttnFusionPass(config)]
141
142
if self.pass_config.enable_qk_norm_rope_fusion:
143
+ self.passes += [SplitCoalescingPass(config)]
144
self.passes += [QKNormRoPEFusionPass(config)]
145
146
# needs a functional graph
vllm/compilation/passes/utility/split_coalescing.py
new
+70
@@ -0,0 +1,70 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Coalesce duplicate ``split_with_sizes`` nodes that operate on the same
5
+input tensor with the same split sizes.
6
+
7
+On certain hardware/dtype combinations (e.g. B200 + FP8) the Inductor
8
+graph may contain multiple ``split_with_sizes`` calls on the same tensor
9
+that CSE fails to merge. This pass detects and replaces the duplicates
10
+so that downstream pattern-matching passes (e.g. QK-Norm+RoPE fusion)
11
+see a single split node with all users attached.
12
+
13
+See also:
14
+ - vLLM #33295 (original issue)
15
+ - PyTorch #174472 (upstream CSE gap)
16
+"""
17
+
18
+import operator
19
+
20
+import torch
21
+from torch import fx
22
+
23
+from vllm.logger import init_logger
24
+
25
+from ..fx_utils import is_func
26
+from ..vllm_inductor_pass import VllmInductorPass
27
+
28
+logger = init_logger(__name__)
29
+
30
+
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
+
35
+ @VllmInductorPass.time_and_log
36
+ def __call__(self, graph: fx.Graph) -> None:
37
+ count = 0
38
+
39
+ # Map from input tensor node -> list of split nodes seen so far.
40
+ split_nodes: dict[fx.Node, list[fx.Node]] = {}
41
+
42
+ for node in graph.nodes:
43
+ if not is_func(node, torch.ops.aten.split_with_sizes.default):
44
+ continue
45
+ if not all(is_func(user, operator.getitem) for user in node.users):
46
+ continue
47
+
48
+ arg_node, split_sizes = node.args[:2]
49
+
50
+ if arg_node not in split_nodes:
51
+ split_nodes[arg_node] = [node]
52
+ continue
53
+
54
+ # Find existing node with same split_sizes
55
+ canonical = next(
56
+ (
57
+ n
58
+ for n in split_nodes[arg_node]
59
+ if list(n.args[1]) == list(split_sizes)
60
+ ),
61
+ None,
62
+ )
63
+ if canonical is not None:
64
+ node.replace_all_uses_with(canonical)
65
+ graph.erase_node(node)
66
+ count += 1
67
+ else:
68
+ split_nodes[arg_node].append(node)
69
+
70
+ logger.debug("Coalesced %d duplicate split_with_sizes nodes", count)