diff --git a/benchmarks/bench_cake_vs_flashkda_cp.py b/benchmarks/bench_cake_vs_flashkda_cp.py new file mode 100644 index 00000000..4c7a436b --- /dev/null +++ b/benchmarks/bench_cake_vs_flashkda_cp.py @@ -0,0 +1,199 @@ +#!/usr/bin/env python3 +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Compare FlashInfer CAKE with cuLA FlashKDA auto intracard CP.""" + +from __future__ import annotations + +import argparse +import gc +import json +import statistics +from pathlib import Path + +import flashinfer.kda_prefill as flashinfer_kda_prefill +import torch +from flashinfer import recurrent_kda + +from cula.kda.flashkda import cula_kda_prefill +from cula.ops.kda.cp_mode import CPMode +from cula.ops.kda.sm90.cp.plan import plan_prefill + +D = 128 + + +def _parse_ints(value: str) -> list[int]: + return [int(item) for item in value.split(",") if item] + + +def _middle_half_mean(samples: list[float]) -> float: + ordered = sorted(samples) + kept = ordered[len(ordered) // 4 : 3 * len(ordered) // 4] + return statistics.fmean(kept or ordered) + + +def _time_round(fn, *, warmup: int, samples: int) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + starts = [torch.cuda.Event(enable_timing=True) for _ in range(samples)] + ends = [torch.cuda.Event(enable_timing=True) for _ in range(samples)] + for start, end in zip(starts, ends, strict=True): + start.record() + fn() + end.record() + torch.cuda.synchronize() + return _middle_half_mean([start.elapsed_time(end) for start, end in zip(starts, ends, strict=True)]) + + +def _relative_rms(actual: torch.Tensor, reference: torch.Tensor) -> float: + delta = actual.float() - reference.float() + return float((delta.square().mean().sqrt() / reference.float().square().mean().sqrt().clamp_min(1e-8)).item()) + + +def _cake_route(device: torch.device, heads: int, length: int) -> str: + return flashinfer_kda_prefill._select_flash_kda_bf16_route( + compute_capability=torch.cuda.get_device_capability(device), + sm_count=torch.cuda.get_device_properties(device).multi_processor_count, + fixed_layout=True, + num_sequences=1, + num_heads=heads, + uniform_sequences=True, + max_sequence_length=length, + use_initial_state=False, + store_final_state=False, + ) + + +@torch.inference_mode() +def _benchmark_shape(*, heads: int, length: int, device: torch.device, warmup: int, samples: int, rounds: int): + generator = torch.Generator(device=device).manual_seed(20260829 + heads * 10007 + length) + shape = (1, length, heads, D) + q = torch.randn(shape, dtype=torch.bfloat16, device=device, generator=generator) + k = torch.randn(shape, dtype=torch.bfloat16, device=device, generator=generator) + v = torch.randn(shape, dtype=torch.bfloat16, device=device, generator=generator) + g = (0.1 * torch.randn(shape, dtype=torch.bfloat16, device=device, generator=generator)).contiguous() + beta = torch.randn((1, length, heads), dtype=torch.bfloat16, device=device, generator=generator) + a_log = 0.1 * torch.randn(heads, dtype=torch.float32, device=device, generator=generator) + dt_bias = 0.1 * torch.randn((heads, D), dtype=torch.float32, device=device, generator=generator) + cake_out = torch.empty_like(v) + cula_out = torch.empty_like(v) + common = { + "q": q, + "k": k, + "v": v, + "g": g, + "beta": beta, + "A_log": a_log, + "dt_bias": dt_bias, + "scale": D**-0.5, + "initial_state": None, + "output_final_state": False, + "use_qk_l2norm_in_kernel": True, + "use_gate_in_kernel": True, + "lower_bound": -5.0, + } + + def run_cake(): + return recurrent_kda(**common, output=cake_out, beta_is_logit=True, backend="cake") + + def run_cula(): + return cula_kda_prefill( + **common, + out=cula_out, + use_beta_sigmoid_in_kernel=True, + safe_gate=True, + use_intracard_cp="auto", + ) + + run_cake() + run_cula() + torch.cuda.synchronize() + plan = plan_prefill([length // 16], heads, device, CPMode.AUTO) + accuracy = { + "relative_rms": _relative_rms(cake_out, cula_out), + "max_abs": float((cake_out.float() - cula_out.float()).abs().max().item()), + } + + cake_rounds = [] + cula_rounds = [] + for round_index in range(rounds): + paths = (("cake", run_cake), ("cula", run_cula)) + if round_index % 2: + paths = tuple(reversed(paths)) + for name, fn in paths: + elapsed = _time_round(fn, warmup=warmup, samples=samples) + (cake_rounds if name == "cake" else cula_rounds).append(elapsed) + + cake_ms = statistics.median(cake_rounds) + cula_ms = statistics.median(cula_rounds) + return { + "heads": heads, + "length": length, + "cake_route": _cake_route(device, heads, length), + "cake_ms": cake_ms, + "flashkda_cp_ms": cula_ms, + "winner": "cake" if cake_ms < cula_ms else "flashkda_cp", + "winner_speedup": max(cake_ms, cula_ms) / min(cake_ms, cula_ms), + "flashkda_cp_active": not plan.trivial, + "flashkda_segments": plan.n_seg_total, + "accuracy": accuracy, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--heads", default="2,4,8") + parser.add_argument("--lengths", default="16384,32768,65536,131072") + parser.add_argument("--warmup", type=int, default=10) + parser.add_argument("--samples", type=int, default=40) + parser.add_argument("--rounds", type=int, default=3) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + device = torch.device(args.device) + torch.cuda.set_device(device) + results = [] + for heads in _parse_ints(args.heads): + for length in _parse_ints(args.lengths): + result = _benchmark_shape( + heads=heads, + length=length, + device=device, + warmup=args.warmup, + samples=args.samples, + rounds=args.rounds, + ) + results.append(result) + print(json.dumps(result, sort_keys=True), flush=True) + gc.collect() + torch.cuda.empty_cache() + + properties = torch.cuda.get_device_properties(device) + report = { + "environment": { + "gpu": properties.name, + "compute_capability": list(torch.cuda.get_device_capability(device)), + "sm_count": properties.multi_processor_count, + "torch": torch.__version__, + "cuda": torch.version.cuda, + }, + "settings": { + "batch": 1, + "head_dim": D, + "dtype": "bfloat16", + "warmup": args.warmup, + "samples": args.samples, + "rounds": args.rounds, + "output_final_state": False, + "preallocated_output": True, + }, + "results": results, + } + args.output.write_text(json.dumps(report, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/cula/kda/_flashkda_arch.py b/cula/kda/_flashkda_arch.py new file mode 100644 index 00000000..a7718c29 --- /dev/null +++ b/cula/kda/_flashkda_arch.py @@ -0,0 +1,31 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""Architecture checks for the SM90-derived FlashKDA implementation.""" + +from __future__ import annotations + +import torch + +_SUPPORTED_COMPUTE_CAPABILITIES = frozenset({(9, 0), (10, 0)}) + + +def is_flashkda_supported(device: torch.device) -> bool: + """Return whether FlashKDA is supported on *device*. + + The CuTeDSL kernels are implemented with the SM90 FlashKDA pipeline, which + is also executable on SM100 Blackwell GPUs. SM103 is intentionally not + enabled until it has equivalent hardware validation. + """ + return device.type == "cuda" and torch.cuda.get_device_capability(device) in _SUPPORTED_COMPUTE_CAPABILITIES + + +def assert_flashkda_supported(device: torch.device) -> None: + """Raise unless *device* is an SM90 or SM100 CUDA device.""" + if device.type != "cuda": + raise RuntimeError(f"FlashKDA requires a CUDA device, got {device}.") + major, minor = torch.cuda.get_device_capability(device) + if (major, minor) not in _SUPPORTED_COMPUTE_CAPABILITIES: + raise RuntimeError( + f"FlashKDA requires an SM90 (Hopper) or SM100 (Blackwell) device, got compute capability sm_{major}{minor}." + ) diff --git a/cula/kda/backends/flashkda.py b/cula/kda/backends/flashkda.py index c0fcde16..fe903bbf 100644 --- a/cula/kda/backends/flashkda.py +++ b/cula/kda/backends/flashkda.py @@ -8,10 +8,7 @@ import torch from cula.backends import BaseBackend - - -def _is_sm90(device: torch.device) -> bool: - return device.type == "cuda" and torch.cuda.get_device_capability(device) == (9, 0) +from cula.kda._flashkda_arch import is_flashkda_supported class FlashKDABackend(BaseBackend): @@ -42,8 +39,8 @@ def kda_prefill_verifier( chunk_indices=None, **kwargs, ): - if not _is_sm90(q.device): - return False, "requires an SM90 (Hopper) device" + if not is_flashkda_supported(q.device): + return False, "requires an SM90 (Hopper) or SM100 (Blackwell) device" if v.shape[2] != q.shape[2]: return False, f"no GVA support (HV={v.shape[2]} != H={q.shape[2]})" if any(not t.is_contiguous() for t in (q, k, v, g)): diff --git a/cula/kda/flashkda.py b/cula/kda/flashkda.py index 245d4373..a00b7bb3 100644 --- a/cula/kda/flashkda.py +++ b/cula/kda/flashkda.py @@ -1,17 +1,17 @@ # Copyright 2025-2026 Ant Group Co., Ltd. # SPDX-License-Identifier: Apache-2.0 -"""SM90 KDA prefill wrapper for the two-kernel K1+K2 CuTeDSL path""" +"""SM90/SM100 KDA prefill wrapper for the two-kernel K1+K2 CuTeDSL path.""" from typing import Literal import torch from torch.amp import custom_bwd, custom_fwd +from cula.kda._flashkda_arch import assert_flashkda_supported from cula.ops.kda.cp_mode import CPMode from cula.ops.kda.sm90.cp.plan import plan_prefill from cula.ops.kda.sm90.fwd import _seq_tiles_from_problem, _validate_inputs, _validate_launch_options, flash_kda_fwd -from cula.utils import assert_hopper def _beta_logits_bf16(beta: torch.Tensor) -> torch.Tensor: @@ -196,10 +196,11 @@ def cula_kda_prefill( **kwargs, ): r""" - Hopper (SM90) KDA forward prefill using CuTeDSL two-kernel pipeline. + SM90/SM100 KDA forward prefill using the SM90-derived CuTeDSL two-kernel + pipeline. Gate preprocessing (A_log, dt_bias, lower_bound) and L2-norm are handled - internally by the K1 kernel. This SM90 CuTeDSL path supports only the safe + internally by the K1 kernel. This CuTeDSL path supports only the safe in-kernel gate mode: ``use_gate_in_kernel=True`` and ``safe_gate=True``. ``use_qk_l2norm_in_kernel`` is accepted for API compatibility; CuTeDSL always applies L2-norm internally. @@ -251,7 +252,7 @@ def cula_kda_prefill( chunk_indices (torch.IntTensor): Accepted for API compatibility; unused by CuTeDSL. use_intracard_cp (Literal["auto"] | bool): - Whether to use the SM90 intracard-CP path when profitable. ``True`` + Whether to use the FlashKDA intracard-CP path when profitable. ``True`` requires CP support and raises on rejection, ``"auto"`` falls back to the serial K1+K2 path, and ``False`` disables CP. out (Optional[torch.Tensor]): @@ -269,14 +270,14 @@ def cula_kda_prefill( final_state (torch.Tensor): Final state of shape `[N, H, V, K]` if `output_final_state=True` else `None`. """ - assert_hopper(q.device) + assert_flashkda_supported(q.device) if not use_gate_in_kernel: raise NotImplementedError( - "SM90 CuTeDSL KDA prefill only supports use_gate_in_kernel=True. " + "FlashKDA CuTeDSL prefill only supports use_gate_in_kernel=True. " "Passing preprocessed gates would otherwise fall back to the slow reference path." ) if not safe_gate: - raise NotImplementedError("SM90 CuTeDSL KDA prefill only supports safe_gate=True.") + raise NotImplementedError("FlashKDA CuTeDSL prefill only supports safe_gate=True.") num_qk_heads, head_dim = q.shape[2], q.shape[3] num_kv_heads = v.shape[2] A_log = kwargs.pop("A_log", None) @@ -292,7 +293,7 @@ def cula_kda_prefill( raise TypeError("beta must be in bfloat16 or float32.") if num_kv_heads != num_qk_heads: raise NotImplementedError( - "SM90 CuTeDSL KDA prefill does not support grouped-value attention yet " + "FlashKDA CuTeDSL prefill does not support grouped-value attention yet " f"(num_kv_heads={num_kv_heads} != num_qk_heads={num_qk_heads}); native GVA is a follow-up change." ) diff --git a/tests/test_kda_flashkda_arch.py b/tests/test_kda_flashkda_arch.py new file mode 100644 index 00000000..b37b8df1 --- /dev/null +++ b/tests/test_kda_flashkda_arch.py @@ -0,0 +1,35 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from cula.kda._flashkda_arch import assert_flashkda_supported, is_flashkda_supported + + +@pytest.mark.parametrize( + ("capability", "supported"), + [ + ((9, 0), True), + ((10, 0), True), + ((8, 0), False), + ((10, 3), False), + ], +) +def test_flashkda_supported_compute_capabilities(monkeypatch, capability, supported): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: capability) + device = torch.device("cuda") + + assert is_flashkda_supported(device) is supported + if supported: + assert_flashkda_supported(device) + else: + with pytest.raises(RuntimeError, match="SM90.*SM100"): + assert_flashkda_supported(device) + + +def test_flashkda_rejects_cpu(): + device = torch.device("cpu") + assert not is_flashkda_supported(device) + with pytest.raises(RuntimeError, match="requires a CUDA device"): + assert_flashkda_supported(device) diff --git a/tests/test_kda_flashkda_sm100.py b/tests/test_kda_flashkda_sm100.py new file mode 100644 index 00000000..82807e65 --- /dev/null +++ b/tests/test_kda_flashkda_sm100.py @@ -0,0 +1,73 @@ +# Copyright 2025-2026 Ant Group Co., Ltd. +# SPDX-License-Identifier: Apache-2.0 + +"""GB200 validation for the SM90-derived FlashKDA compatibility path.""" + +import pytest +import torch +from fla.ops.kda.chunk import chunk_kda as fla_chunk_kda +from fla.utils import assert_close + +from cula.kda import kda_prefill + +pytestmark = [ + pytest.mark.sm100_only, + pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA"), +] + + +@pytest.mark.kda_fast +def test_flashkda_auto_cp_matches_fla(monkeypatch): + from cula.ops.kda.sm90.cp import driver + + B, T, H, D = 1, 4096, 2, 128 + device = torch.device("cuda") + torch.manual_seed(42) + q = torch.rand(B, T, H, D, dtype=torch.bfloat16, device=device) + k = torch.rand(B, T, H, D, dtype=torch.bfloat16, device=device) + v = torch.rand(B, T, H, D, dtype=torch.bfloat16, device=device) + g = torch.randn(B, T, H, D, dtype=torch.bfloat16, device=device) + beta = torch.randn(B, T, H, dtype=torch.float32, device=device).sigmoid().to(torch.bfloat16) + A_log = torch.randn(H, dtype=torch.float32, device=device) + dt_bias = torch.randn(H * D, dtype=torch.float32, device=device) + + original_run_cp = driver._run_cp + entered_cp = False + + def record_run_cp(*args, **kwargs): + nonlocal entered_cp + entered_cp = True + return original_run_cp(*args, **kwargs) + + monkeypatch.setattr(driver, "_run_cp", record_run_cp) + + with torch.inference_mode(): + ref_o, ref_ht = fla_chunk_kda( + q, + k, + v, + g, + beta, + A_log=A_log, + dt_bias=dt_bias, + output_final_state=True, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + safe_gate=True, + lower_bound=-5.0, + ) + actual_o, actual_ht_vk = kda_prefill( + q, + k, + v, + g, + beta, + A_log=A_log, + dt_bias=dt_bias, + output_final_state=True, + use_intracard_cp="auto", + ) + + assert entered_cp, "expected the SM100 public backend call to enter FlashKDA intracard CP" + assert_close("o", ref_o, actual_o, 0.005) + assert_close("ht", ref_ht, actual_ht_vk.transpose(-2, -1), 0.005)