Optimize K2 recurrent-state dataflow for SM103/B300

README.md

@@ -27,6 +27,34 @@ FLASH_KDA_CUDA_ARCHS=all pip install -v --no-build-isolation .
 
 Supported values are `auto` (default), `all`, or a comma-separated arch list such as `90a,100a`.
 
+### Optional SM103 K2 register state
+
+On SM103 (B300), build the V1a and V1aE specializations with:
+
+```bash
+FLASH_KDA_CUDA_ARCHS=103a FLASH_KDA_ENABLE_V1A=1 FLASH_KDA_ENABLE_V1AE=1 \
+  pip install -v --no-build-isolation .
+```
+
+`FLASH_KDA_ENABLE_V1A` enables persistent register state; `FLASH_KDA_ENABLE_V1AE`
+enables the same recurrence with a one-time shared-memory/TMA final-state store.
+Both flags default to off and may be enabled independently. The optimization
+changes recurrent-state dataflow, preserving the existing MMA math.
+
+Set `FLASH_KDA_K2_IMPL=auto` at runtime to select by workload, or use `baseline`,
+`v1a`, or `v1ae` explicitly. Unset defaults to `baseline`, even when both
+specializations are built. V1a/V1aE require SM103, fixed lengths, and both BF16
+`initial_state` and `final_state`. Unsupported or uncompiled explicit selections
+raise an error; `auto` falls back to baseline when its preferred path is unavailable.
+
+The H=64, B1–B8 thresholds in [`csrc/k2_dispatch.h`](csrc/k2_dispatch.h) are
+empirical B300 policy. For B5–B7, `>320` chunks selecting V1a is a conservative,
+unvalidated fallback, not a measured crossover. At H=64, B8 with `>=320` chunks
+and all B>8 workloads use baseline. For H!=64, supported calls use V1a at
+`>=128` chunks and baseline below that. A chunk is 16 tokens; the policy counts
+`ceil(T/16)` tiles per sequence, including a partial final tile. These thresholds
+are not claims for other GPUs.
+
 ## Using FlashKDA as an FLA backend
 
 Once installed, FlashKDA is auto-dispatched from `flash-linear-attention`'s `chunk_kda`. See [fla-org/flash-linear-attention#852](https://github.com/fla-org/flash-linear-attention/pull/852) for integration details.
@@ -67,6 +95,18 @@ Once installed, FlashKDA is auto-dispatched from `flash-linear-attention`'s `chu
 
 See [BENCHMARK_H20.md](BENCHMARK_H20.md).
 
+After the SM103 build above, compare all four K2 selections on B1/T8192,
+B4/T2048, and B8/T1024 (H=64, D=128):
+
+```bash
+python benchmarks/bench_sm103_dispatch.py --warmup 30 --iters 200 --seed 42
+```
+
+The script uses identical seeded inputs, checks exact output/final-state equality
+and unchanged initial state, and prints mean/median CUDA-event latency in
+microseconds plus the expected AUTO selection. Timing calls the public
+`flash_kda.fwd` wrapper, including its workspace allocation; checks are untimed.
+
 ## Tests
 
 ```bash
@@ -75,6 +115,24 @@ bash tests/test.sh
 
 - `tests/test_fwd.py` — correctness tests (exact match against the torch reference; compared with `flash-linear-attention`)
 
+Host-only K2 policy tests require Python and a C++17 GCC/Clang compiler, with
+no PyTorch or CUDA dependency (`CXX` may name the compiler executable):
+
+```bash
+python tests/test_k2_dispatch.py
+```
+
+The runner compiles into a temporary directory and also supports pytest.
+After building both specializations, opt into SM103 parity checks with:
+
+```bash
+FLASH_KDA_TEST_SM103=1 python -m pytest tests/test_sm103_k2.py -q
+```
+
+These GPU tests skip by default and on other architectures. Run the existing
+full regression separately under `FLASH_KDA_K2_IMPL=baseline` and `auto` using
+`python -m pytest tests/test_fwd_full.py -q`.
+
 
 ## Kernel API
 

benchmarks/bench_sm103_dispatch.py

@@ -0,0 +1,105 @@
+"""Reproduce the three B300 K2 workloads through the public flash_kda.fwd API."""
+
+import argparse
+import os
+import statistics
+
+
+def bench_fn(fn, warmup, iters):
+    import torch
+
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    starts = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
+    ends = [torch.cuda.Event(enable_timing=True) for _ in range(iters)]
+    for start, end in zip(starts, ends):
+        start.record()
+        fn()
+        end.record()
+    torch.cuda.synchronize()
+    times_us = [start.elapsed_time(end) * 1000 for start, end in zip(starts, ends)]
+    return statistics.mean(times_us), statistics.median(times_us)
+
+
+def run_case(B, T, expected_auto, warmup, iters, seed):
+    import torch
+    import torch.nn.functional as F
+    import flash_kda
+
+    H, D = 64, 128
+    torch.manual_seed(seed)
+    shape = (B, T, H, D)
+    q = F.normalize(torch.randn(shape, dtype=torch.float32, device="cuda"), dim=-1).to(torch.bfloat16)
+    k = F.normalize(torch.randn(shape, dtype=torch.float32, device="cuda"), dim=-1).to(torch.bfloat16)
+    v = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+    g = torch.randn_like(v)
+    beta = torch.randn((B, T, H), dtype=torch.bfloat16, device="cuda")
+    A_log = torch.rand(H, dtype=torch.float32, device="cuda")
+    dt_bias = torch.rand(H, D, dtype=torch.float32, device="cuda")
+    initial = torch.randn((B, H, D, D), dtype=torch.bfloat16, device="cuda")
+    initial_copy = initial.clone()
+    print(f"\nB={B} T={T} H={H} D={D} chunks={(T + 15) // 16} expected_auto={expected_auto}")
+
+    previous = os.environ.get("FLASH_KDA_K2_IMPL")
+    try:
+        for mode in ("baseline", "v1a", "v1ae", "auto"):
+            os.environ["FLASH_KDA_K2_IMPL"] = mode
+            out = torch.full_like(q, float("nan"))
+            final = torch.full_like(initial, float("nan"))
+
+            def run():
+                # Include the public wrapper's workspace allocation in each call.
+                flash_kda.fwd(q, k, v, g, beta, D ** -0.5, out,
+                              A_log=A_log, dt_bias=dt_bias, lower_bound=-5.0,
+                              initial_state=initial, final_state=final)
+
+            run()
+            torch.cuda.synchronize()
+            if mode == "baseline":
+                if not (torch.isfinite(out).all() and torch.isfinite(final).all()):
+                    raise RuntimeError("baseline produced non-finite results")
+                baseline_out, baseline_final = out.clone(), final.clone()
+
+            def check_results():
+                if not torch.equal(initial, initial_copy):
+                    raise RuntimeError(f"{mode}: initial_state changed")
+                if not torch.equal(out, baseline_out):
+                    raise RuntimeError(f"{mode}: output mismatch")
+                if not torch.equal(final, baseline_final):
+                    raise RuntimeError(f"{mode}: final_state mismatch")
+
+            check_results()
+            mean, median = bench_fn(run, warmup, iters)
+            check_results()
+            print(f"  {mode:8s} mean={mean:.3f} us median={median:.3f} us exact=PASS")
+    finally:
+        if previous is None:
+            os.environ.pop("FLASH_KDA_K2_IMPL", None)
+        else:
+            os.environ["FLASH_KDA_K2_IMPL"] = previous
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--warmup", type=int, default=30)
+    parser.add_argument("--iters", type=int, default=200)
+    parser.add_argument("--seed", type=int, default=42)
+    args = parser.parse_args()
+    if args.warmup < 0 or args.iters <= 0:
+        parser.error("--warmup must be nonnegative and --iters must be positive")
+
+    import torch
+
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 3):
+        parser.error("requires an SM103 GPU and FlashKDA built with V1a and V1aE")
+    print(f"GPU={torch.cuda.get_device_name()} CC=10.3 "
+          f"PyTorch={torch.__version__} CUDA={torch.version.cuda}")
+    print(f"warmup={args.warmup} iters={args.iters} seed={args.seed}")
+    with torch.inference_mode():
+        for B, T, expected in ((1, 8192, "v1a"), (4, 2048, "v1ae"), (8, 1024, "v1ae")):
+            run_case(B, T, expected, args.warmup, args.iters, args.seed)
+
+
+if __name__ == "__main__":
+    main()

csrc/flash_kda.cpp

@@ -1,6 +1,9 @@
 #include <torch/extension.h>
 #include <c10/cuda/CUDAStream.h>
 #include "fwd.h"
+#include "k2_dispatch.h"
+#include <cstdlib>
+#include <string>
 
 int64_t get_workspace_size(
     int64_t T_total,
@@ -109,6 +112,14 @@ void fwd(
 
     TORCH_CHECK(D == 128, "currently only supports D == 128");
 
+    // Preserve the public default (baseline). Selection is read per call so
+    // correctness and benchmark tests can interleave implementations.
+    const char* requested_k2 = std::getenv("FLASH_KDA_K2_IMPL");
+    const std::string k2_impl_name = requested_k2 ? requested_k2 : "baseline";
+    const auto k2_mode = flash_kda::parse_k2_mode(k2_impl_name);
+    TORCH_CHECK(k2_mode != flash_kda::K2Mode::Invalid,
+                "FLASH_KDA_K2_IMPL must be auto, baseline, v1a, or v1ae");
+
     // Flatten [B, T, H, D] -> [B*T, H, D] (contiguous, same data pointer)
     auto q_3d = q.reshape({T_total, H, D});
     auto k_3d = k.reshape({T_total, H, D});
@@ -159,6 +170,50 @@ void fwd(
         N_val = B;
     }
 
+    int compute_major = 0, compute_minor = 0;
+    if (k2_mode != flash_kda::K2Mode::Baseline) {
+        TORCH_CHECK(
+            cudaDeviceGetAttribute(&compute_major, cudaDevAttrComputeCapabilityMajor,
+                                   q.get_device()) == cudaSuccess &&
+            cudaDeviceGetAttribute(&compute_minor, cudaDevAttrComputeCapabilityMinor,
+                                   q.get_device()) == cudaSuccess,
+            "Cannot query CUDA device capability for K2 implementation selection");
+    }
+
+    flash_kda::K2DispatchConfig k2_config;
+#if defined(FLASH_KDA_ENABLE_V1A)
+    k2_config.v1a_compiled = true;
+#endif
+#if defined(FLASH_KDA_ENABLE_V1AE)
+    k2_config.v1ae_compiled = true;
+#endif
+    k2_config.compute_major = compute_major;
+    k2_config.compute_minor = compute_minor;
+    k2_config.is_varlen = is_varlen;
+    k2_config.has_state_in = has_state_in;
+    k2_config.has_state_out = has_state_out;
+    k2_config.state_fp32 = state_fp32;
+    k2_config.total_tokens = T_total;
+    k2_config.sequences = N_val;
+    k2_config.heads = H;
+
+    const auto k2_implementation =
+        flash_kda::select_k2_implementation(k2_mode, k2_config);
+    if (k2_implementation == flash_kda::K2Implementation::Unsupported) {
+        if (k2_mode == flash_kda::K2Mode::V1AE) {
+            TORCH_CHECK(k2_config.v1ae_compiled,
+                        "Rebuild with FLASH_KDA_ENABLE_V1AE=1 to enable v1ae");
+            TORCH_CHECK(false,
+                        "v1ae requires SM103 and fixed-length BF16 initial_state and final_state");
+        }
+        TORCH_CHECK(k2_config.v1a_compiled,
+                    "Rebuild with FLASH_KDA_ENABLE_V1A=1 to enable v1a");
+        TORCH_CHECK(false,
+                    "v1a requires SM103 and fixed-length BF16 initial_state and final_state");
+    }
+    const bool use_v1a =
+        k2_implementation == flash_kda::K2Implementation::V1A;
+
     // Validate state shapes: always [N, H, D, D]
     if (has_state_in) {
         auto& is = initial_state.value();
@@ -181,14 +236,16 @@ void fwd(
     }
 
     // Dispatch based on state configuration and varlen
-    #define LAUNCH(HI, HO, FP32, VL) \
-        launch_fwd<128, HI, HO, FP32, VL>( \
+    #define LAUNCH_IMPL(HI, HO, FP32, VL, V1A, EGRESS) \
+        launch_fwd<128, HI, HO, FP32, VL, V1A, EGRESS>( \
             q_ptr, k_ptr, v_ptr, g_ptr, beta_t_ptr, \
             initial_state_raw, scale_f, final_state_raw, out_ptr, \
             workspace_ptr, total_tiles, \
             int(T_total), int(H), int(N_val), cu_seqlens_dev, \
             A_log_ptr, dt_bias_ptr, gate_scale, stream)
 
+    #define LAUNCH(HI, HO, FP32, VL) LAUNCH_IMPL(HI, HO, FP32, VL, false, 0)
+
     #define DISPATCH_STATE(VL) \
         if (!has_state_in && !has_state_out) { \
             LAUNCH(false, false, false, VL); \
@@ -206,14 +263,23 @@ void fwd(
             LAUNCH(true, false, false, VL); \
         }
 
-    if (is_varlen) {
+    if (k2_implementation == flash_kda::K2Implementation::V1AE) {
+#if defined(FLASH_KDA_ENABLE_V1AE)
+        LAUNCH_IMPL(true, true, false, false, false, 1);
+#endif
+    } else if (use_v1a) {
+#if defined(FLASH_KDA_ENABLE_V1A)
+        LAUNCH_IMPL(true, true, false, false, true, 0);
+#endif
+    } else if (is_varlen) {
         DISPATCH_STATE(true);
     } else {
         DISPATCH_STATE(false);
     }
 
     #undef DISPATCH_STATE
     #undef LAUNCH
+    #undef LAUNCH_IMPL
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {

csrc/fwd.h

@@ -3,7 +3,8 @@
 
 #include <cutlass/bfloat16.h>
 
-template <int D, bool HasStateIn = true, bool HasStateOut = true, bool StateFP32 = false, bool IsVarlen = true>
+template <int D, bool HasStateIn = true, bool HasStateOut = true, bool StateFP32 = false,
+          bool IsVarlen = true, bool UseV1A = false, int V1AEgress = 0>
 void launch_fwd(
     cutlass::bfloat16_t const* q_ptr,
     cutlass::bfloat16_t const* k_ptr,