diff --git a/tests/kernels/test_bf16_skinny_gemm.py b/tests/kernels/test_bf16_skinny_gemm.py index 98a49c753c..c16b4d44e8 100644 --- a/tests/kernels/test_bf16_skinny_gemm.py +++ b/tests/kernels/test_bf16_skinny_gemm.py @@ -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 ) diff --git a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py index 02be89aa31..e439ec5b66 100644 --- a/vllm/models/kimi_k3/nvidia/low_latency_gemm.py +++ b/vllm/models/kimi_k3/nvidia/low_latency_gemm.py @@ -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 diff --git a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py index da69e2631b..9e8f933300 100644 --- a/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py +++ b/vllm/models/kimi_k3/nvidia/ops/cute_dsl/kda_skinny_gemm.py @@ -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),