[Bugfix][Spec Decode] Support heterogeneous QK fusion geometry (#48671)

Signed-off-by: aoshen02 <aoshen@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>

tests/compile/passes/test_qk_norm_rope_fusion.py

@@ -212,3 +212,73 @@ def test_qk_norm_rope_fusion(
 
         backend.check_before_ops(model.ops_in_model_before())
         backend.check_after_ops(model.ops_in_model_after())
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="Only test on CUDA and ROCm platforms",
+)
+def test_qk_norm_rope_fusion_with_heterogeneous_draft_model():
+    """The fusion pass should support target and draft attention geometries."""
+    if not hasattr(torch.ops._C, "fused_qk_norm_rope"):
+        pytest.skip("fused_qk_norm_rope custom op not available")
+
+    dtype = torch.bfloat16
+    torch.set_default_device(current_platform.device_type)
+    torch.set_default_dtype(dtype)
+    torch.manual_seed(0)
+
+    vllm_config = VllmConfig(
+        model_config=ModelConfig(dtype=dtype),
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.VLLM_COMPILE,
+            custom_ops=["+rms_norm", "+rotary_embedding"],
+            pass_config=PassConfig(
+                enable_qk_norm_rope_fusion=True,
+                eliminate_noops=True,
+            ),
+        ),
+    )
+
+    with (
+        set_current_vllm_config(vllm_config),
+        vllm_config.kernel_config.ir_op_priority.set_priority(),
+    ):
+        target_model = QKNormRoPETestModel(
+            num_heads=32,
+            num_kv_heads=8,
+            head_dim=128,
+            eps=1e-6,
+            is_neox=True,
+            vllm_config=vllm_config,
+            dtype=dtype,
+            prefix="model.layers.0.self_attn.attn",
+        )
+        draft_model = QKNormRoPETestModel(
+            num_heads=16,
+            num_kv_heads=8,
+            head_dim=128,
+            eps=1e-6,
+            is_neox=True,
+            vllm_config=vllm_config,
+            dtype=dtype,
+            prefix="draft_model.layers.0.self_attn.attn",
+        )
+
+        fusion_pass = QKNormRoPEFusionPass(vllm_config)
+        backend = TestBackend(
+            NoOpEliminationPass(vllm_config),
+            SplitCoalescingPass(vllm_config),
+            fusion_pass,
+            PostCleanupPass(vllm_config),
+        )
+
+        for model in (target_model, draft_model):
+            num_tokens = 5
+            qkv = torch.randn(num_tokens, model.q_size + 2 * model.kv_size)
+            positions = torch.arange(num_tokens, dtype=torch.long, device=qkv.device)
+            torch._dynamo.mark_dynamic(qkv, 0)
+            torch._dynamo.mark_dynamic(positions, 0)
+            torch.compile(model, backend=backend)(qkv, positions)
+
+            assert fusion_pass.matched_count == 1

vllm/compilation/passes/fusion/qk_norm_rope_fusion.py

@@ -123,7 +123,11 @@ class QkNormRopePattern:
             cos_sin_cache: torch.Tensor,
         ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
             # split qkv -> q,k,v
-            q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+            try:
+                q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
+            except ValueError as e:
+                # register_replacement treats RuntimeError as a trace mismatch.
+                raise RuntimeError from e
 
             # Q path: view -> RMS -> view back to q.shape
             q_by_head = q.view(
@@ -194,6 +198,7 @@ class QKNormRoPEFusionPass(VllmPatternMatcherPass):
         self.patterns: PatternMatcherPass = PatternMatcherPass(
             pass_name="qk_norm_rope_fusion_pass"
         )
+        self._attention_geometries: tuple[tuple[int, int, int], ...] = ()
 
         dtype = config.model_config.dtype
         if dtype not in (torch.bfloat16, torch.float16):
@@ -202,7 +207,6 @@ class QKNormRoPEFusionPass(VllmPatternMatcherPass):
             )
             return
 
-        # use one attn layer to get meta (such as head_dim) for QkNormRopePattern
         attn_layers: dict[str, Attention] = get_layers_from_vllm_config(
             config, Attention
         )
@@ -211,28 +215,31 @@ class QKNormRoPEFusionPass(VllmPatternMatcherPass):
                 "QK Norm+RoPE fusion enabled, but no Attention layers were discovered."
             )
             return
-        layer = next(iter(attn_layers.values()))
 
-        for epsilon in [1e-5, 1e-6]:
-            for neox in [True, False]:
-                if RotaryEmbedding.enabled():
-                    for rope_flashinfer in [False, True]:
+        self._attention_geometries = tuple(
+            sorted(
+                {
+                    (layer.head_size, layer.num_heads, layer.num_kv_heads)
+                    for layer in attn_layers.values()
+                }
+            )
+        )
+
+        rope_flashinfer_options = (
+            [False, True] if RotaryEmbedding.enabled() else [False]
+        )
+        for head_dim, num_heads, num_kv_heads in self._attention_geometries:
+            for epsilon in [1e-5, 1e-6]:
+                for neox in [True, False]:
+                    for rope_flashinfer in rope_flashinfer_options:
                         QkNormRopePattern(
-                            head_dim=layer.head_size,
-                            num_heads=layer.num_heads,
-                            num_kv_heads=layer.num_kv_heads,
+                            head_dim=head_dim,
+                            num_heads=num_heads,
+                            num_kv_heads=num_kv_heads,
                             eps=epsilon,
                             is_neox=neox,
                             rope_flashinfer=rope_flashinfer,
                         ).register(self.patterns)
-                else:
-                    QkNormRopePattern(
-                        head_dim=layer.head_size,
-                        num_heads=layer.num_heads,
-                        num_kv_heads=layer.num_kv_heads,
-                        eps=epsilon,
-                        is_neox=neox,
-                    ).register(self.patterns)
 
         self.dump_patterns(config, self.patterns)
 
@@ -242,4 +249,6 @@ class QKNormRoPEFusionPass(VllmPatternMatcherPass):
         logger.debug("Fused QK Norm+RoPE on %s sites", self.matched_count)
 
     def uuid(self) -> str:
-        return VllmInductorPass.hash_source(self, QkNormRopePattern)
+        return VllmInductorPass.hash_source(
+            self, QkNormRopePattern, repr(self._attention_geometries)
+        )