From 62b0dde510aa8435ac35bf56e15d1d7a03b63ca6 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 29 Aug 2026 14:44:40 +0800 Subject: [PATCH 1/4] feat(kda): enable FlashKDA on SM100 Signed-off-by: shuyan.ycf --- cula/kda/_flashkda_arch.py | 31 ++++++++++++++ cula/kda/backends/flashkda.py | 9 ++-- cula/kda/flashkda.py | 19 +++++---- tests/test_kda_flashkda_arch.py | 35 +++++++++++++++ tests/test_kda_flashkda_sm100.py | 73 ++++++++++++++++++++++++++++++++ 5 files changed, 152 insertions(+), 15 deletions(-) create mode 100644 cula/kda/_flashkda_arch.py create mode 100644 tests/test_kda_flashkda_arch.py create mode 100644 tests/test_kda_flashkda_sm100.py 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) From 6c13747cde06fe3dcfda4b4505dca22c3f019991 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 29 Aug 2026 15:06:39 +0800 Subject: [PATCH 2/4] bench(kda): compare CAKE with FlashKDA CP Signed-off-by: shuyan.ycf --- benchmarks/bench_cake_vs_flashkda_cp.py | 199 ++++++++++++++++++++++++ 1 file changed, 199 insertions(+) create mode 100644 benchmarks/bench_cake_vs_flashkda_cp.py 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() From 6486b30c165e425d0baf2fd7e2b77d739d553ba2 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 29 Aug 2026 15:11:40 +0800 Subject: [PATCH 3/4] docs(kda): record CAKE versus FlashKDA CP results Signed-off-by: shuyan.ycf --- ...da_sm100_cake_vs_intracard_cp_crossover.md | 55 +++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 docs/kda_sm100_cake_vs_intracard_cp_crossover.md diff --git a/docs/kda_sm100_cake_vs_intracard_cp_crossover.md b/docs/kda_sm100_cake_vs_intracard_cp_crossover.md new file mode 100644 index 00000000..9f466eca --- /dev/null +++ b/docs/kda_sm100_cake_vs_intracard_cp_crossover.md @@ -0,0 +1,55 @@ +# Latest CAKE vs. cuLA FlashKDA Intracard CP on GB200 + +## TL;DR + +For `B=1`, `D=128`, BF16, and `H=2/4/8`, cuLA FlashKDA auto intracard CP +was faster than the latest official CAKE dispatcher at every tested length from +16K through 128K. The measured advantage was **1.74x--4.52x**. There is no +CAKE-to-CP crossover inside this range: FlashKDA CP already wins at 16K. + +## Versions + +| Implementation | Exact version | +|---|---| +| FlashInfer CAKE | [`flashinfer-ai/flashinfer@e425c7b0`](https://github.com/flashinfer-ai/flashinfer/commit/e425c7b029ca90d5d01ff207913b070863d35a5b), latest `main` at measurement time | +| cuLA FlashKDA CP | [`inclusionAI/cuLA@6c13747c`](https://github.com/inclusionAI/cuLA/commit/6c13747cde06fe3dcfda4b4505dca22c3f019991), [PR #124](https://github.com/inclusionAI/cuLA/pull/124) | + +CAKE was called through `recurrent_kda(..., backend="cake")`. cuLA was called +through `cula.kda.flashkda.cula_kda_prefill(..., use_intracard_cp="auto")`; +auto CP was active for all rows. This cuLA path is the SM90-derived FlashKDA +compatibility implementation running on SM100, not the separate SM100-native +modular KDA path. For CAKE kernel background, see +[this CAKE optimization overview](https://zhuanlan.zhihu.com/p/2068499679076259239). + +## Results + +`1K = 1024` tokens. Speedup is `CAKE latency / FlashKDA CP latency`. + +| H | T | CAKE route | CAKE (ms) | FlashKDA CP (ms) | CP speedup | +|---:|---:|---|---:|---:|---:| +| 2 | 16K | small-BH M128 | 0.622 | 0.278 | 2.23x | +| 2 | 32K | small-BH M128 | 1.221 | 0.420 | 2.91x | +| 2 | 64K | BT16 + M64 chain | 1.858 | 0.573 | 3.24x | +| 2 | 128K | BT16 + M64 chain | 3.698 | 0.818 | 4.52x | +| 4 | 16K | small-BH M128 | 0.625 | 0.295 | 2.12x | +| 4 | 32K | small-BH M128 | 1.227 | 0.423 | 2.90x | +| 4 | 64K | BT16 + M64 chain | 1.910 | 0.671 | 2.85x | +| 4 | 128K | BT16 + M64 chain | 3.794 | 1.163 | 3.26x | +| 8 | 16K | small-BH M128 | 0.607 | 0.349 | 1.74x | +| 8 | 32K | small-BH M128 | 1.198 | 0.593 | 2.02x | +| 8 | 64K | BT16 + M64 chain | 2.005 | 1.090 | 1.84x | +| 8 | 128K | BT16 + M64 chain | 3.972 | 2.075 | 1.91x | + +## Method and Scope + +Measurements used an NVIDIA GB200 (SM100, 152 SMs), `B=1`, `H=HV`, +`D=128`, BF16 Q/K/V/G and beta logits, fused Q/K normalization and gate +processing, no initial/final state, and preallocated output buffers. Each point +uses 10 warmups, 40 CUDA-event samples, and three alternating-order rounds; the +reported value is the median of the round-level middle-half means. The harness +is [`bench_cake_vs_flashkda_cp.py`](https://github.com/inclusionAI/cuLA/blob/6c13747cde06fe3dcfda4b4505dca22c3f019991/benchmarks/bench_cake_vs_flashkda_cp.py). + +Across the matrix, output relative RMS difference was `0.492%--0.624%` and +maximum absolute difference was at most `1.10e-3`. Results are specific to the +measured shapes and execution contract; batch size, state output, head +dimension, or another GPU can change the ordering. From 5438898ffb1c34119a3e6b29636b390f31bd38b8 Mon Sep 17 00:00:00 2001 From: "shuyan.ycf" Date: Sat, 29 Aug 2026 15:12:56 +0800 Subject: [PATCH 4/4] revert: keep KDA benchmark notes local Signed-off-by: shuyan.ycf --- ...da_sm100_cake_vs_intracard_cp_crossover.md | 55 ------------------- 1 file changed, 55 deletions(-) delete mode 100644 docs/kda_sm100_cake_vs_intracard_cp_crossover.md diff --git a/docs/kda_sm100_cake_vs_intracard_cp_crossover.md b/docs/kda_sm100_cake_vs_intracard_cp_crossover.md deleted file mode 100644 index 9f466eca..00000000 --- a/docs/kda_sm100_cake_vs_intracard_cp_crossover.md +++ /dev/null @@ -1,55 +0,0 @@ -# Latest CAKE vs. cuLA FlashKDA Intracard CP on GB200 - -## TL;DR - -For `B=1`, `D=128`, BF16, and `H=2/4/8`, cuLA FlashKDA auto intracard CP -was faster than the latest official CAKE dispatcher at every tested length from -16K through 128K. The measured advantage was **1.74x--4.52x**. There is no -CAKE-to-CP crossover inside this range: FlashKDA CP already wins at 16K. - -## Versions - -| Implementation | Exact version | -|---|---| -| FlashInfer CAKE | [`flashinfer-ai/flashinfer@e425c7b0`](https://github.com/flashinfer-ai/flashinfer/commit/e425c7b029ca90d5d01ff207913b070863d35a5b), latest `main` at measurement time | -| cuLA FlashKDA CP | [`inclusionAI/cuLA@6c13747c`](https://github.com/inclusionAI/cuLA/commit/6c13747cde06fe3dcfda4b4505dca22c3f019991), [PR #124](https://github.com/inclusionAI/cuLA/pull/124) | - -CAKE was called through `recurrent_kda(..., backend="cake")`. cuLA was called -through `cula.kda.flashkda.cula_kda_prefill(..., use_intracard_cp="auto")`; -auto CP was active for all rows. This cuLA path is the SM90-derived FlashKDA -compatibility implementation running on SM100, not the separate SM100-native -modular KDA path. For CAKE kernel background, see -[this CAKE optimization overview](https://zhuanlan.zhihu.com/p/2068499679076259239). - -## Results - -`1K = 1024` tokens. Speedup is `CAKE latency / FlashKDA CP latency`. - -| H | T | CAKE route | CAKE (ms) | FlashKDA CP (ms) | CP speedup | -|---:|---:|---|---:|---:|---:| -| 2 | 16K | small-BH M128 | 0.622 | 0.278 | 2.23x | -| 2 | 32K | small-BH M128 | 1.221 | 0.420 | 2.91x | -| 2 | 64K | BT16 + M64 chain | 1.858 | 0.573 | 3.24x | -| 2 | 128K | BT16 + M64 chain | 3.698 | 0.818 | 4.52x | -| 4 | 16K | small-BH M128 | 0.625 | 0.295 | 2.12x | -| 4 | 32K | small-BH M128 | 1.227 | 0.423 | 2.90x | -| 4 | 64K | BT16 + M64 chain | 1.910 | 0.671 | 2.85x | -| 4 | 128K | BT16 + M64 chain | 3.794 | 1.163 | 3.26x | -| 8 | 16K | small-BH M128 | 0.607 | 0.349 | 1.74x | -| 8 | 32K | small-BH M128 | 1.198 | 0.593 | 2.02x | -| 8 | 64K | BT16 + M64 chain | 2.005 | 1.090 | 1.84x | -| 8 | 128K | BT16 + M64 chain | 3.972 | 2.075 | 1.91x | - -## Method and Scope - -Measurements used an NVIDIA GB200 (SM100, 152 SMs), `B=1`, `H=HV`, -`D=128`, BF16 Q/K/V/G and beta logits, fused Q/K normalization and gate -processing, no initial/final state, and preallocated output buffers. Each point -uses 10 warmups, 40 CUDA-event samples, and three alternating-order rounds; the -reported value is the median of the round-level middle-half means. The harness -is [`bench_cake_vs_flashkda_cp.py`](https://github.com/inclusionAI/cuLA/blob/6c13747cde06fe3dcfda4b4505dca22c3f019991/benchmarks/bench_cake_vs_flashkda_cp.py). - -Across the matrix, output relative RMS difference was `0.492%--0.624%` and -maximum absolute difference was at most `1.10e-3`. Results are specific to the -measured shapes and execution contract; batch size, state output, head -dimension, or another GPU can change the ordering.