[Bugfix][Kimi-K3] Fix KDA projection overlap on Hopper (#55426)

Signed-off-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com> Co-authored-by: Chengcheng Pei <5881383+chengchengpei@users.noreply.github.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
2026-09-11 19:07:30 +00:00
parent 00e5cda926
commit 0c1e89ceb9
3 changed files with 238 additions and 42 deletions
@@ -269,6 +269,70 @@ def test_kda_overlap_configs_match_measured_table() -> None:
    }


+@pytest.mark.parametrize(
+    "compute_capability,supported",
+    [((9, 0), False), ((10, 0), True), ((10, 3), True), ((12, 0), True)],
+)
+def test_kda_mixed_precision_bf16_fma_capability(
+    monkeypatch: pytest.MonkeyPatch,
+    compute_capability: tuple[int, int],
+    supported: bool,
+) -> None:
+    from vllm.models.kimi_k3.nvidia.ops.cute_dsl import kda_skinny_gemm
+
+    capability = compute_capability[0] * 10 + compute_capability[1]
+    monkeypatch.setattr(
+        kda_skinny_gemm.current_platform,
+        "has_device_capability",
+        lambda target, device_id=0: capability >= target,
+    )
+
+    assert kda_skinny_gemm._has_mixed_precision_bf16_fma() is supported
+
+
+@pytest.mark.parametrize("num_tokens", [2, 8, 14])
+def test_kda_skinny_gemms_sm90_cuda_graph(num_tokens: int) -> None:
+    """Hopper must compile and capture the pre-SM100 FMA specialization."""
+    _require_capability_and_cute((9, 0))
+    from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import (
+        KdaSkinnyGemm,
+    )
+
+    torch.manual_seed(43 + num_tokens)
+    hidden_states = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda")
+    f_ab_weight = torch.randn(144, 7168, dtype=torch.bfloat16, device="cuda")
+    f_b_weight = torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda")
+    gemm = KdaSkinnyGemm()
+
+    gemm.run_k(gemm.run_n(hidden_states, f_ab_weight), f_b_weight)
+    torch.accelerator.synchronize()
+
+    capture_stream = torch.cuda.Stream()
+    capture_stream.wait_stream(torch.cuda.current_stream())
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph, stream=capture_stream):
+        projected_fab = gemm.run_n(hidden_states, f_ab_weight)
+        projected_fb = gemm.run_k(projected_fab, f_b_weight)
+    torch.cuda.current_stream().wait_stream(capture_stream)
+    graph.replay()
+    torch.accelerator.synchronize()
+
+    expected_fab = torch.nn.functional.linear(
+        hidden_states.float(), f_ab_weight.float()
+    )
+    expected_fb = torch.nn.functional.linear(
+        expected_fab[:, :128].to(torch.bfloat16).float(), f_b_weight.float()
+    )
+    for actual, expected in (
+        (projected_fab, expected_fab),
+        (projected_fb, expected_fb),
+    ):
+        cosine = torch.nn.functional.cosine_similarity(
+            actual.float().flatten(), expected.flatten(), dim=0
+        ).item()
+        assert cosine > 0.999
+
+
@pytest.mark.parametrize("tp_size", [1, 2, 4, 16])
 def test_kda_projection_overlap_is_tp8_only(tp_size: int) -> None:
    from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
@@ -282,8 +346,34 @@ def test_kda_projection_overlap_is_tp8_only(tp_size: int) -> None:
    assert kda._projection_overlap_max_tokens == 0


+@pytest.mark.parametrize(
+    ("capability", "expected"),
+    [
+        ((9, 0), None),
+        ((10, 0), "cute-dsl"),
+        ((10, 3), "cute-dsl"),
+        ((10, 7), None),
+        ((12, 0), None),
+    ],
+)
+def test_kda_qkvg_flashinfer_backend(
+    monkeypatch: pytest.MonkeyPatch,
+    capability: tuple[int, int],
+    expected: str | None,
+) -> None:
+    monkeypatch.setattr(
+        current_platform,
+        "is_device_capability",
+        lambda candidate: candidate == capability,
+    )
+
+    assert k3_gemm._kda_qkvg_flashinfer_backend() == expected
+
+
+@pytest.mark.parametrize("backend", ["cute-dsl", None])
 def test_kda_qkvg_autotune_enables_full_overlap(
    monkeypatch: pytest.MonkeyPatch,
+    backend: str | None,
 ) -> None:
    flashinfer_gemm = pytest.importorskip("flashinfer.gemm")

@@ -304,10 +394,20 @@ def test_kda_qkvg_autotune_enables_full_overlap(
        return torch.empty(a.shape[0], b.shape[1], dtype=a.dtype)

    monkeypatch.setattr(flashinfer_gemm, "mm_bf16", fake_mm_bf16)
+    monkeypatch.setattr(
+        k3_gemm,
+        "_kda_qkvg_flashinfer_backend",
+        lambda: backend,
+    )

    k3_gemm.autotune_kda_qkvg(kda)

-    assert calls == [(torch.Size([14, 8]), torch.Size([8, 6144]), True, "cute-dsl")]
+    expected_calls = []
+    if backend is not None:
+        expected_calls.append(
+            (torch.Size([14, 8]), torch.Size([8, 6144]), True, "cute-dsl")
+        )
+    assert calls == expected_calls
    assert (
        kda._projection_overlap_max_tokens == k3_gemm.KDA_PROJECTION_OVERLAP_MAX_TOKENS
    )
@@ -65,6 +65,20 @@ _KDA_PACKED_SIZE = 6288
 _KDA_TP_SIZE = 8


+def _kda_qkvg_flashinfer_backend() -> Literal["cute-dsl"] | None:
+    """Return the tested FlashInfer backend for the low-M QKVG branch.
+
+    FlashInfer's ``cute-dsl`` BF16 GEMM supports SM100 and SM103, but rejects
+    SM90 during backend validation. Other capabilities use ``torch.mm`` for
+    QKVG while retaining the concurrent F_A/beta and F_B branch.
+    """
+    if current_platform.is_device_capability(
+        (10, 0)
+    ) or current_platform.is_device_capability((10, 3)):
+        return "cute-dsl"
+    return None
+
+
@dataclass(frozen=True, slots=True)
 class ProjectionSpec:
    n: int
@@ -749,14 +763,16 @@ def run_kda_projection_overlap(
                config,
                None,
            )
-        if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS:
+        if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS and (
+            backend := _kda_qkvg_flashinfer_backend()
+        ):
            from flashinfer.gemm import mm_bf16

            return mm_bf16(
                hidden_states,
                qkvg_weight.t(),
                pdl=True,
-                backend="cute-dsl",
+                backend=backend,
            )
        return torch.mm(hidden_states, qkvg_weight.t())

@@ -807,11 +823,10 @@ def run_kda_projection_overlap(


 def autotune_kda_qkvg(model: nn.Module) -> None:
-    """Autotune the FlashInfer QKVG GEMM before CUDA graph capture."""
-    from flashinfer.gemm import mm_bf16
-
+    """Autotune the supported QKVG GEMM before CUDA graph capture."""
    from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention

+    backend = _kda_qkvg_flashinfer_backend()
    children: list[KimiK3DeltaAttention] = []
    weights_by_shape: dict[
        tuple[int, int, int, torch.dtype, torch.device], torch.Tensor
@@ -822,29 +837,33 @@ def autotune_kda_qkvg(model: nn.Module) -> None:
        if child._projection_overlap_max_tokens <= 0:
            continue
        children.append(child)
-        qkvg_weight = child.in_proj_qkvgfab.weight[:_KDA_QKVG_SIZE]
-        shape = (
-            KDA_PROJECTION_OVERLAP_MAX_TOKENS,
-            qkvg_weight.shape[0],
-            qkvg_weight.shape[1],
-            qkvg_weight.dtype,
-            qkvg_weight.device,
-        )
-        weights_by_shape.setdefault(shape, qkvg_weight)
+        if backend is not None:
+            qkvg_weight = child.in_proj_qkvgfab.weight[:_KDA_QKVG_SIZE]
+            shape = (
+                KDA_PROJECTION_OVERLAP_MAX_TOKENS,
+                qkvg_weight.shape[0],
+                qkvg_weight.shape[1],
+                qkvg_weight.dtype,
+                qkvg_weight.device,
+            )
+            weights_by_shape.setdefault(shape, qkvg_weight)

-    for shape, qkvg_weight in weights_by_shape.items():
-        num_tokens = shape[0]
-        hidden_states = torch.empty(
-            (num_tokens, qkvg_weight.shape[1]),
-            dtype=qkvg_weight.dtype,
-            device=qkvg_weight.device,
-        )
-        mm_bf16(
-            hidden_states,
-            qkvg_weight.t(),
-            pdl=True,
-            backend="cute-dsl",
-        )
+    if backend is not None:
+        from flashinfer.gemm import mm_bf16
+
+        for shape, qkvg_weight in weights_by_shape.items():
+            num_tokens = shape[0]
+            hidden_states = torch.empty(
+                (num_tokens, qkvg_weight.shape[1]),
+                dtype=qkvg_weight.dtype,
+                device=qkvg_weight.device,
+            )
+            mm_bf16(
+                hidden_states,
+                qkvg_weight.t(),
+                pdl=True,
+                backend=backend,
+            )
    for child in children:
        child._projection_overlap_max_tokens = KDA_PROJECTION_OVERLAP_MAX_TOKENS

@@ -18,6 +18,8 @@ from cutlass._mlir.dialects import llvm
 from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
 from cutlass.cutlass_dsl import T, dsl_user_op

+from vllm.platforms import current_platform
+
 _FAB_N = 144
 _FAB_K = 7168
 _FAB_VECTOR_WIDTH = 16
@@ -91,14 +93,61 @@ def _fma_f32_bf16(
    return Float32(result)


+@dsl_user_op
+def _fma_f32_bf16_portable(
+    a: BFloat16,
+    b: BFloat16,
+    acc: Float32,
+    *,
+    loc=None,
+    ip=None,
+) -> Float32:
+    """BF16 multiply with FP32 accumulation for pre-SM100 GPUs.
+
+    PTX ``fma.f32.bf16`` requires SM100 or newer. Converting the operands to
+    FP32 first keeps this kernel available on Hopper, where an unconditional
+    mixed-precision FMA fails libNVVM compilation for ``sm_90a``.
+    """
+    a_bits = llvm.bitcast(T.i16(), a.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
+    b_bits = llvm.bitcast(T.i16(), b.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
+    result = llvm.inline_asm(
+        T.f32(),
+        [a_bits, b_bits, acc.ir_value(loc=loc, ip=ip)],
+        "{\n\t"
+        ".reg .f32 a_f32, b_f32;\n\t"
+        "cvt.f32.bf16 a_f32, $1;\n\t"
+        "cvt.f32.bf16 b_f32, $2;\n\t"
+        "fma.rn.f32 $0, a_f32, b_f32, $3;\n\t"
+        "}",
+        "=f,h,h,f",
+        has_side_effects=False,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+    return Float32(result)
+
+
+def _has_mixed_precision_bf16_fma() -> bool:
+    """Whether PTX ``fma.f32.bf16`` is supported by the current GPU."""
+    return current_platform.has_device_capability(100)
+
+
 class _KdaSkinnyNGemm:
-    def __init__(self, num_rows: int, split_n: int) -> None:
+    def __init__(
+        self,
+        num_rows: int,
+        split_n: int,
+        use_mixed_precision_fma: bool,
+    ) -> None:
        if _FAB_N % split_n:
            raise ValueError("split_n must divide the FAB output size")
        if num_rows * split_n > _FAB_BLOCK_SIZE:
            raise ValueError("the FAB output tile must fit in one block")
        self.num_rows = num_rows
        self.split_n = split_n
+        self.use_mixed_precision_fma = use_mixed_precision_fma

    @cute.jit
    def __call__(
@@ -191,11 +240,18 @@ class _KdaSkinnyNGemm:
        for mi in cutlass.range_constexpr(num_rows):
            for ni in cutlass.range_constexpr(split_n):
                for vi in cutlass.range_constexpr(_FAB_VECTOR_WIDTH):
-                    acc[mi, ni] = _fma_f32_bf16(
-                        a_regs[mi, vi],
-                        weight_regs[ni, vi],
-                        acc[mi, ni],
-                    )
+                    if cutlass.const_expr(self.use_mixed_precision_fma):
+                        acc[mi, ni] = _fma_f32_bf16(
+                            a_regs[mi, vi],
+                            weight_regs[ni, vi],
+                            acc[mi, ni],
+                        )
+                    else:
+                        acc[mi, ni] = _fma_f32_bf16_portable(
+                            a_regs[mi, vi],
+                            weight_regs[ni, vi],
+                            acc[mi, ni],
+                        )
                acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni])

        partials_layout = cute.make_layout(
@@ -232,13 +288,19 @@ class _KdaSkinnyNGemm:


 class _KdaSkinnyKGemm:
-    def __init__(self, num_rows: int, split_n: int) -> None:
+    def __init__(
+        self,
+        num_rows: int,
+        split_n: int,
+        use_mixed_precision_fma: bool,
+    ) -> None:
        if _FB_N % split_n:
            raise ValueError("split_n must divide the FB output size")
        if split_n * _FB_K_LANES > 1024:
            raise ValueError("the FB output tile exceeds the block-size limit")
        self.num_rows = num_rows
        self.split_n = split_n
+        self.use_mixed_precision_fma = use_mixed_precision_fma

    @cute.jit
    def __call__(
@@ -314,11 +376,18 @@ class _KdaSkinnyKGemm:
        for mi in cutlass.range_constexpr(num_rows):
            acc = Float32(0.0)
            for vi in cutlass.range_constexpr(_FB_VECTOR_WIDTH):
-                acc = _fma_f32_bf16(
-                    input_regs[mi, vi],
-                    weight_regs[vi],
-                    acc,
-                )
+                if cutlass.const_expr(self.use_mixed_precision_fma):
+                    acc = _fma_f32_bf16(
+                        input_regs[mi, vi],
+                        weight_regs[vi],
+                        acc,
+                    )
+                else:
+                    acc = _fma_f32_bf16_portable(
+                        input_regs[mi, vi],
+                        weight_regs[vi],
+                        acc,
+                    )
            for offset in (8, 4, 2, 1):
                acc += cute.arch.shuffle_sync_bfly(acc, offset)
            output[mi, n_idx] = acc.to(BFloat16)
@@ -341,7 +410,11 @@ class KdaSkinnyGemm:
            return
        split_n = SKINNY_N_SPLIT_N[num_rows]
        self._compiled_n[key] = cute.compile(
-            _KdaSkinnyNGemm(num_rows, split_n),
+            _KdaSkinnyNGemm(
+                num_rows,
+                split_n,
+                use_mixed_precision_fma=_has_mixed_precision_bf16_fma(),
+            ),
            make_fake_tensor(
                BFloat16,
                (num_rows, _FAB_K),
@@ -371,7 +444,11 @@ class KdaSkinnyGemm:
            return
        split_n = SKINNY_K_SPLIT_N[num_rows]
        self._compiled_k[key] = cute.compile(
-            _KdaSkinnyKGemm(num_rows, split_n),
+            _KdaSkinnyKGemm(
+                num_rows,
+                split_n,
+                use_mixed_precision_fma=_has_mixed_precision_bf16_fma(),
+            ),
            make_fake_tensor(
                BFloat16,
                (num_rows, _FB_K),