[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>
chengchengpei committed
Sep 11, 2026 at 12:07 UTC
0c1e89ceb92b95e02b5a7aa07781340aacfcf011
3 files changed
+239
-43
tests/kernels/test_bf16_skinny_gemm.py
+101
-1
@@ -269,6 +269,70 @@ def test_kda_overlap_configs_match_measured_table() -> None:
269
}
270
271
272
+@pytest.mark.parametrize(
273
+ "compute_capability,supported",
274
+ [((9, 0), False), ((10, 0), True), ((10, 3), True), ((12, 0), True)],
275
+)
276
+def test_kda_mixed_precision_bf16_fma_capability(
277
+ monkeypatch: pytest.MonkeyPatch,
278
+ compute_capability: tuple[int, int],
279
+ supported: bool,
280
+) -> None:
281
+ from vllm.models.kimi_k3.nvidia.ops.cute_dsl import kda_skinny_gemm
282
+
283
+ capability = compute_capability[0] * 10 + compute_capability[1]
284
+ monkeypatch.setattr(
285
+ kda_skinny_gemm.current_platform,
286
+ "has_device_capability",
287
+ lambda target, device_id=0: capability >= target,
288
+ )
289
+
290
+ assert kda_skinny_gemm._has_mixed_precision_bf16_fma() is supported
291
+
292
+
293
+@pytest.mark.parametrize("num_tokens", [2, 8, 14])
294
+def test_kda_skinny_gemms_sm90_cuda_graph(num_tokens: int) -> None:
295
+ """Hopper must compile and capture the pre-SM100 FMA specialization."""
296
+ _require_capability_and_cute((9, 0))
297
+ from vllm.models.kimi_k3.nvidia.ops.cute_dsl.kda_skinny_gemm import (
298
+ KdaSkinnyGemm,
299
+ )
300
+
301
+ torch.manual_seed(43 + num_tokens)
302
+ hidden_states = torch.randn(num_tokens, 7168, dtype=torch.bfloat16, device="cuda")
303
+ f_ab_weight = torch.randn(144, 7168, dtype=torch.bfloat16, device="cuda")
304
+ f_b_weight = torch.randn(1536, 128, dtype=torch.bfloat16, device="cuda")
305
+ gemm = KdaSkinnyGemm()
306
+
307
+ gemm.run_k(gemm.run_n(hidden_states, f_ab_weight), f_b_weight)
308
+ torch.accelerator.synchronize()
309
+
310
+ capture_stream = torch.cuda.Stream()
311
+ capture_stream.wait_stream(torch.cuda.current_stream())
312
+ graph = torch.cuda.CUDAGraph()
313
+ with torch.cuda.graph(graph, stream=capture_stream):
314
+ projected_fab = gemm.run_n(hidden_states, f_ab_weight)
315
+ projected_fb = gemm.run_k(projected_fab, f_b_weight)
316
+ torch.cuda.current_stream().wait_stream(capture_stream)
317
+ graph.replay()
318
+ torch.accelerator.synchronize()
319
+
320
+ expected_fab = torch.nn.functional.linear(
321
+ hidden_states.float(), f_ab_weight.float()
322
+ )
323
+ expected_fb = torch.nn.functional.linear(
324
+ expected_fab[:, :128].to(torch.bfloat16).float(), f_b_weight.float()
325
+ )
326
+ for actual, expected in (
327
+ (projected_fab, expected_fab),
328
+ (projected_fb, expected_fb),
329
+ ):
330
+ cosine = torch.nn.functional.cosine_similarity(
331
+ actual.float().flatten(), expected.flatten(), dim=0
332
+ ).item()
333
+ assert cosine > 0.999
334
+
335
+
336
@pytest.mark.parametrize("tp_size", [1, 2, 4, 16])
337
def test_kda_projection_overlap_is_tp8_only(tp_size: int) -> None:
338
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:
346
assert kda._projection_overlap_max_tokens == 0
347
348
349
+@pytest.mark.parametrize(
350
+ ("capability", "expected"),
351
+ [
352
+ ((9, 0), None),
353
+ ((10, 0), "cute-dsl"),
354
+ ((10, 3), "cute-dsl"),
355
+ ((10, 7), None),
356
+ ((12, 0), None),
357
+ ],
358
+)
359
+def test_kda_qkvg_flashinfer_backend(
360
+ monkeypatch: pytest.MonkeyPatch,
361
+ capability: tuple[int, int],
362
+ expected: str | None,
363
+) -> None:
364
+ monkeypatch.setattr(
365
+ current_platform,
366
+ "is_device_capability",
367
+ lambda candidate: candidate == capability,
368
+ )
369
+
370
+ assert k3_gemm._kda_qkvg_flashinfer_backend() == expected
371
+
372
+
373
+@pytest.mark.parametrize("backend", ["cute-dsl", None])
374
def test_kda_qkvg_autotune_enables_full_overlap(
375
monkeypatch: pytest.MonkeyPatch,
376
+ backend: str | None,
377
) -> None:
378
flashinfer_gemm = pytest.importorskip("flashinfer.gemm")
379
@@ -304,10 +394,20 @@ def test_kda_qkvg_autotune_enables_full_overlap(
394
return torch.empty(a.shape[0], b.shape[1], dtype=a.dtype)
395
396
monkeypatch.setattr(flashinfer_gemm, "mm_bf16", fake_mm_bf16)
397
+ monkeypatch.setattr(
398
+ k3_gemm,
399
+ "_kda_qkvg_flashinfer_backend",
400
+ lambda: backend,
401
+ )
402
403
k3_gemm.autotune_kda_qkvg(kda)
404
310
- assert calls == [(torch.Size([14, 8]), torch.Size([8, 6144]), True, "cute-dsl")]
405
+ expected_calls = []
406
+ if backend is not None:
407
+ expected_calls.append(
408
+ (torch.Size([14, 8]), torch.Size([8, 6144]), True, "cute-dsl")
409
+ )
410
+ assert calls == expected_calls
411
assert (
412
kda._projection_overlap_max_tokens == k3_gemm.KDA_PROJECTION_OVERLAP_MAX_TOKENS
413
)
vllm/models/kimi_k3/nvidia/low_latency_gemm.py
+47
-28
@@ -65,6 +65,20 @@ _KDA_PACKED_SIZE = 6288
65
_KDA_TP_SIZE = 8
66
67
68
+def _kda_qkvg_flashinfer_backend() -> Literal["cute-dsl"] | None:
69
+ """Return the tested FlashInfer backend for the low-M QKVG branch.
70
+
71
+ FlashInfer's ``cute-dsl`` BF16 GEMM supports SM100 and SM103, but rejects
72
+ SM90 during backend validation. Other capabilities use ``torch.mm`` for
73
+ QKVG while retaining the concurrent F_A/beta and F_B branch.
74
+ """
75
+ if current_platform.is_device_capability(
76
+ (10, 0)
77
+ ) or current_platform.is_device_capability((10, 3)):
78
+ return "cute-dsl"
79
+ return None
80
+
81
+
82
@dataclass(frozen=True, slots=True)
83
class ProjectionSpec:
84
n: int
@@ -749,14 +763,16 @@ def run_kda_projection_overlap(
763
config,
764
None,
765
)
752
- if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS:
766
+ if num_tokens <= KDA_PROJECTION_OVERLAP_MAX_TOKENS and (
767
+ backend := _kda_qkvg_flashinfer_backend()
768
+ ):
769
from flashinfer.gemm import mm_bf16
770
771
return mm_bf16(
772
hidden_states,
773
qkvg_weight.t(),
774
pdl=True,
759
- backend="cute-dsl",
775
+ backend=backend,
776
)
777
return torch.mm(hidden_states, qkvg_weight.t())
778
@@ -807,11 +823,10 @@ def run_kda_projection_overlap(
823
824
825
def autotune_kda_qkvg(model: nn.Module) -> None:
810
- """Autotune the FlashInfer QKVG GEMM before CUDA graph capture."""
811
- from flashinfer.gemm import mm_bf16
812
-
826
+ """Autotune the supported QKVG GEMM before CUDA graph capture."""
827
from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
828
829
+ backend = _kda_qkvg_flashinfer_backend()
830
children: list[KimiK3DeltaAttention] = []
831
weights_by_shape: dict[
832
tuple[int, int, int, torch.dtype, torch.device], torch.Tensor
@@ -822,29 +837,33 @@ def autotune_kda_qkvg(model: nn.Module) -> None:
837
if child._projection_overlap_max_tokens <= 0:
838
continue
839
children.append(child)
825
- qkvg_weight = child.in_proj_qkvgfab.weight[:_KDA_QKVG_SIZE]
826
- shape = (
827
- KDA_PROJECTION_OVERLAP_MAX_TOKENS,
828
- qkvg_weight.shape[0],
829
- qkvg_weight.shape[1],
830
- qkvg_weight.dtype,
831
- qkvg_weight.device,
832
- )
833
- weights_by_shape.setdefault(shape, qkvg_weight)
834
-
835
- for shape, qkvg_weight in weights_by_shape.items():
836
- num_tokens = shape[0]
837
- hidden_states = torch.empty(
838
- (num_tokens, qkvg_weight.shape[1]),
839
- dtype=qkvg_weight.dtype,
840
- device=qkvg_weight.device,
841
- )
842
- mm_bf16(
843
- hidden_states,
844
- qkvg_weight.t(),
845
- pdl=True,
846
- backend="cute-dsl",
847
- )
840
+ if backend is not None:
841
+ qkvg_weight = child.in_proj_qkvgfab.weight[:_KDA_QKVG_SIZE]
842
+ shape = (
843
+ KDA_PROJECTION_OVERLAP_MAX_TOKENS,
844
+ qkvg_weight.shape[0],
845
+ qkvg_weight.shape[1],
846
+ qkvg_weight.dtype,
847
+ qkvg_weight.device,
848
+ )
849
+ weights_by_shape.setdefault(shape, qkvg_weight)
850
+
851
+ if backend is not None:
852
+ from flashinfer.gemm import mm_bf16
853
+
854
+ for shape, qkvg_weight in weights_by_shape.items():
855
+ num_tokens = shape[0]
856
+ hidden_states = torch.empty(
857
+ (num_tokens, qkvg_weight.shape[1]),
858
+ dtype=qkvg_weight.dtype,
859
+ device=qkvg_weight.device,
860
+ )
861
+ mm_bf16(
862
+ hidden_states,
863
+ qkvg_weight.t(),
864
+ pdl=True,
865
+ backend=backend,
866
+ )
867
for child in children:
868
child._projection_overlap_max_tokens = KDA_PROJECTION_OVERLAP_MAX_TOKENS
869
vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py
+91
-14
@@ -18,6 +18,8 @@ from cutlass._mlir.dialects import llvm
18
from cutlass.cute.runtime import make_fake_stream, make_fake_tensor
19
from cutlass.cutlass_dsl import T, dsl_user_op
20
21
+from vllm.platforms import current_platform
22
+
23
_FAB_N = 144
24
_FAB_K = 7168
25
_FAB_VECTOR_WIDTH = 16
@@ -91,14 +93,61 @@ def _fma_f32_bf16(
93
return Float32(result)
94
95
96
+@dsl_user_op
97
+def _fma_f32_bf16_portable(
98
+ a: BFloat16,
99
+ b: BFloat16,
100
+ acc: Float32,
101
+ *,
102
+ loc=None,
103
+ ip=None,
104
+) -> Float32:
105
+ """BF16 multiply with FP32 accumulation for pre-SM100 GPUs.
106
+
107
+ PTX ``fma.f32.bf16`` requires SM100 or newer. Converting the operands to
108
+ FP32 first keeps this kernel available on Hopper, where an unconditional
109
+ mixed-precision FMA fails libNVVM compilation for ``sm_90a``.
110
+ """
111
+ a_bits = llvm.bitcast(T.i16(), a.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
112
+ b_bits = llvm.bitcast(T.i16(), b.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
113
+ result = llvm.inline_asm(
114
+ T.f32(),
115
+ [a_bits, b_bits, acc.ir_value(loc=loc, ip=ip)],
116
+ "{\n\t"
117
+ ".reg .f32 a_f32, b_f32;\n\t"
118
+ "cvt.f32.bf16 a_f32, $1;\n\t"
119
+ "cvt.f32.bf16 b_f32, $2;\n\t"
120
+ "fma.rn.f32 $0, a_f32, b_f32, $3;\n\t"
121
+ "}",
122
+ "=f,h,h,f",
123
+ has_side_effects=False,
124
+ is_align_stack=False,
125
+ asm_dialect=llvm.AsmDialect.AD_ATT,
126
+ loc=loc,
127
+ ip=ip,
128
+ )
129
+ return Float32(result)
130
+
131
+
132
+def _has_mixed_precision_bf16_fma() -> bool:
133
+ """Whether PTX ``fma.f32.bf16`` is supported by the current GPU."""
134
+ return current_platform.has_device_capability(100)
135
+
136
+
137
class _KdaSkinnyNGemm:
95
- def __init__(self, num_rows: int, split_n: int) -> None:
138
+ def __init__(
139
+ self,
140
+ num_rows: int,
141
+ split_n: int,
142
+ use_mixed_precision_fma: bool,
143
+ ) -> None:
144
if _FAB_N % split_n:
145
raise ValueError("split_n must divide the FAB output size")
146
if num_rows * split_n > _FAB_BLOCK_SIZE:
147
raise ValueError("the FAB output tile must fit in one block")
148
self.num_rows = num_rows
149
self.split_n = split_n
150
+ self.use_mixed_precision_fma = use_mixed_precision_fma
151
152
@cute.jit
153
def __call__(
@@ -191,11 +240,18 @@ class _KdaSkinnyNGemm:
240
for mi in cutlass.range_constexpr(num_rows):
241
for ni in cutlass.range_constexpr(split_n):
242
for vi in cutlass.range_constexpr(_FAB_VECTOR_WIDTH):
194
- acc[mi, ni] = _fma_f32_bf16(
195
- a_regs[mi, vi],
196
- weight_regs[ni, vi],
197
- acc[mi, ni],
198
- )
243
+ if cutlass.const_expr(self.use_mixed_precision_fma):
244
+ acc[mi, ni] = _fma_f32_bf16(
245
+ a_regs[mi, vi],
246
+ weight_regs[ni, vi],
247
+ acc[mi, ni],
248
+ )
249
+ else:
250
+ acc[mi, ni] = _fma_f32_bf16_portable(
251
+ a_regs[mi, vi],
252
+ weight_regs[ni, vi],
253
+ acc[mi, ni],
254
+ )
255
acc[mi, ni] = cute.arch.warp_reduction_sum(acc[mi, ni])
256
257
partials_layout = cute.make_layout(
@@ -232,13 +288,19 @@ class _KdaSkinnyNGemm:
288
289
290
class _KdaSkinnyKGemm:
235
- def __init__(self, num_rows: int, split_n: int) -> None:
291
+ def __init__(
292
+ self,
293
+ num_rows: int,
294
+ split_n: int,
295
+ use_mixed_precision_fma: bool,
296
+ ) -> None:
297
if _FB_N % split_n:
298
raise ValueError("split_n must divide the FB output size")
299
if split_n * _FB_K_LANES > 1024:
300
raise ValueError("the FB output tile exceeds the block-size limit")
301
self.num_rows = num_rows
302
self.split_n = split_n
303
+ self.use_mixed_precision_fma = use_mixed_precision_fma
304
305
@cute.jit
306
def __call__(
@@ -314,11 +376,18 @@ class _KdaSkinnyKGemm:
376
for mi in cutlass.range_constexpr(num_rows):
377
acc = Float32(0.0)
378
for vi in cutlass.range_constexpr(_FB_VECTOR_WIDTH):
317
- acc = _fma_f32_bf16(
318
- input_regs[mi, vi],
319
- weight_regs[vi],
320
- acc,
321
- )
379
+ if cutlass.const_expr(self.use_mixed_precision_fma):
380
+ acc = _fma_f32_bf16(
381
+ input_regs[mi, vi],
382
+ weight_regs[vi],
383
+ acc,
384
+ )
385
+ else:
386
+ acc = _fma_f32_bf16_portable(
387
+ input_regs[mi, vi],
388
+ weight_regs[vi],
389
+ acc,
390
+ )
391
for offset in (8, 4, 2, 1):
392
acc += cute.arch.shuffle_sync_bfly(acc, offset)
393
output[mi, n_idx] = acc.to(BFloat16)
@@ -341,7 +410,11 @@ class KdaSkinnyGemm:
410
return
411
split_n = SKINNY_N_SPLIT_N[num_rows]
412
self._compiled_n[key] = cute.compile(
344
- _KdaSkinnyNGemm(num_rows, split_n),
413
+ _KdaSkinnyNGemm(
414
+ num_rows,
415
+ split_n,
416
+ use_mixed_precision_fma=_has_mixed_precision_bf16_fma(),
417
+ ),
418
make_fake_tensor(
419
BFloat16,
420
(num_rows, _FAB_K),
@@ -371,7 +444,11 @@ class KdaSkinnyGemm:
444
return
445
split_n = SKINNY_K_SPLIT_N[num_rows]
446
self._compiled_k[key] = cute.compile(
374
- _KdaSkinnyKGemm(num_rows, split_n),
447
+ _KdaSkinnyKGemm(
448
+ num_rows,
449
+ split_n,
450
+ use_mixed_precision_fma=_has_mixed_precision_bf16_fma(),
451
+ ),
452
make_fake_tensor(
453
BFloat16,
454
(num_rows, _FB_K),