[Bugfix][Spec Decode] Support heterogeneous QK fusion geometry (#48671)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
aoshen02 committed
Jul 15, 2026 at 13:37 UTC
3b39fd284aa3bcfcad57736cdd5ddb9eb4b746d2
2 files changed
+98
-19
tests/compile/passes/test_qk_norm_rope_fusion.py
+70
@@ -212,3 +212,73 @@ def test_qk_norm_rope_fusion(
212
213
backend.check_before_ops(model.ops_in_model_before())
214
backend.check_after_ops(model.ops_in_model_after())
215
+
216
+
217
+@pytest.mark.skipif(
218
+ not current_platform.is_cuda_alike(),
219
+ reason="Only test on CUDA and ROCm platforms",
220
+)
221
+def test_qk_norm_rope_fusion_with_heterogeneous_draft_model():
222
+ """The fusion pass should support target and draft attention geometries."""
223
+ if not hasattr(torch.ops._C, "fused_qk_norm_rope"):
224
+ pytest.skip("fused_qk_norm_rope custom op not available")
225
+
226
+ dtype = torch.bfloat16
227
+ torch.set_default_device(current_platform.device_type)
228
+ torch.set_default_dtype(dtype)
229
+ torch.manual_seed(0)
230
+
231
+ vllm_config = VllmConfig(
232
+ model_config=ModelConfig(dtype=dtype),
233
+ compilation_config=CompilationConfig(
234
+ mode=CompilationMode.VLLM_COMPILE,
235
+ custom_ops=["+rms_norm", "+rotary_embedding"],
236
+ pass_config=PassConfig(
237
+ enable_qk_norm_rope_fusion=True,
238
+ eliminate_noops=True,
239
+ ),
240
+ ),
241
+ )
242
+
243
+ with (
244
+ set_current_vllm_config(vllm_config),
245
+ vllm_config.kernel_config.ir_op_priority.set_priority(),
246
+ ):
247
+ target_model = QKNormRoPETestModel(
248
+ num_heads=32,
249
+ num_kv_heads=8,
250
+ head_dim=128,
251
+ eps=1e-6,
252
+ is_neox=True,
253
+ vllm_config=vllm_config,
254
+ dtype=dtype,
255
+ prefix="model.layers.0.self_attn.attn",
256
+ )
257
+ draft_model = QKNormRoPETestModel(
258
+ num_heads=16,
259
+ num_kv_heads=8,
260
+ head_dim=128,
261
+ eps=1e-6,
262
+ is_neox=True,
263
+ vllm_config=vllm_config,
264
+ dtype=dtype,
265
+ prefix="draft_model.layers.0.self_attn.attn",
266
+ )
267
+
268
+ fusion_pass = QKNormRoPEFusionPass(vllm_config)
269
+ backend = TestBackend(
270
+ NoOpEliminationPass(vllm_config),
271
+ SplitCoalescingPass(vllm_config),
272
+ fusion_pass,
273
+ PostCleanupPass(vllm_config),
274
+ )
275
+
276
+ for model in (target_model, draft_model):
277
+ num_tokens = 5
278
+ qkv = torch.randn(num_tokens, model.q_size + 2 * model.kv_size)
279
+ positions = torch.arange(num_tokens, dtype=torch.long, device=qkv.device)
280
+ torch._dynamo.mark_dynamic(qkv, 0)
281
+ torch._dynamo.mark_dynamic(positions, 0)
282
+ torch.compile(model, backend=backend)(qkv, positions)
283
+
284
+ assert fusion_pass.matched_count == 1
vllm/compilation/passes/fusion/qk_norm_rope_fusion.py
+28
-19
@@ -123,7 +123,11 @@ class QkNormRopePattern:
123
cos_sin_cache: torch.Tensor,
124
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
125
# split qkv -> q,k,v
126
- q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
126
+ try:
127
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
128
+ except ValueError as e:
129
+ # register_replacement treats RuntimeError as a trace mismatch.
130
+ raise RuntimeError from e
131
132
# Q path: view -> RMS -> view back to q.shape
133
q_by_head = q.view(
@@ -194,6 +198,7 @@ class QKNormRoPEFusionPass(VllmPatternMatcherPass):
198
self.patterns: PatternMatcherPass = PatternMatcherPass(
199
pass_name="qk_norm_rope_fusion_pass"
200
)
201
+ self._attention_geometries: tuple[tuple[int, int, int], ...] = ()
202
203
dtype = config.model_config.dtype
204
if dtype not in (torch.bfloat16, torch.float16):
@@ -202,7 +207,6 @@ class QKNormRoPEFusionPass(VllmPatternMatcherPass):
207
)
208
return
209
205
- # use one attn layer to get meta (such as head_dim) for QkNormRopePattern
210
attn_layers: dict[str, Attention] = get_layers_from_vllm_config(
211
config, Attention
212
)
@@ -211,28 +215,31 @@ class QKNormRoPEFusionPass(VllmPatternMatcherPass):
215
"QK Norm+RoPE fusion enabled, but no Attention layers were discovered."
216
)
217
return
214
- layer = next(iter(attn_layers.values()))
218
216
- for epsilon in [1e-5, 1e-6]:
217
- for neox in [True, False]:
218
- if RotaryEmbedding.enabled():
219
- for rope_flashinfer in [False, True]:
219
+ self._attention_geometries = tuple(
220
+ sorted(
221
+ {
222
+ (layer.head_size, layer.num_heads, layer.num_kv_heads)
223
+ for layer in attn_layers.values()
224
+ }
225
+ )
226
+ )
227
+
228
+ rope_flashinfer_options = (
229
+ [False, True] if RotaryEmbedding.enabled() else [False]
230
+ )
231
+ for head_dim, num_heads, num_kv_heads in self._attention_geometries:
232
+ for epsilon in [1e-5, 1e-6]:
233
+ for neox in [True, False]:
234
+ for rope_flashinfer in rope_flashinfer_options:
235
QkNormRopePattern(
221
- head_dim=layer.head_size,
222
- num_heads=layer.num_heads,
223
- num_kv_heads=layer.num_kv_heads,
236
+ head_dim=head_dim,
237
+ num_heads=num_heads,
238
+ num_kv_heads=num_kv_heads,
239
eps=epsilon,
240
is_neox=neox,
241
rope_flashinfer=rope_flashinfer,
242
).register(self.patterns)
228
- else:
229
- QkNormRopePattern(
230
- head_dim=layer.head_size,
231
- num_heads=layer.num_heads,
232
- num_kv_heads=layer.num_kv_heads,
233
- eps=epsilon,
234
- is_neox=neox,
235
- ).register(self.patterns)
243
244
self.dump_patterns(config, self.patterns)
245
@@ -242,4 +249,6 @@ class QKNormRoPEFusionPass(VllmPatternMatcherPass):
249
logger.debug("Fused QK Norm+RoPE on %s sites", self.matched_count)
250
251
def uuid(self) -> str:
245
- return VllmInductorPass.hash_source(self, QkNormRopePattern)
252
+ return VllmInductorPass.hash_source(
253
+ self, QkNormRopePattern, repr(self._attention_geometries)
254
+ )