diff --git a/README.md b/README.md
index 57d7fc2..ca90731 100644
--- a/README.md
+++ b/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
 
diff --git a/benchmarks/bench_sm103_dispatch.py b/benchmarks/bench_sm103_dispatch.py
new file mode 100644
index 0000000..fb168ba
--- /dev/null
+++ b/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()
diff --git a/csrc/flash_kda.cpp b/csrc/flash_kda.cpp
index 81f5483..2e577bc 100644
--- a/csrc/flash_kda.cpp
+++ b/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,7 +263,15 @@ 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);
@@ -214,6 +279,7 @@ void fwd(
 
     #undef DISPATCH_STATE
     #undef LAUNCH
+    #undef LAUNCH_IMPL
 }
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
diff --git a/csrc/fwd.h b/csrc/fwd.h
index deb82b6..ca367b0 100644
--- a/csrc/fwd.h
+++ b/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,
diff --git a/csrc/k2_dispatch.h b/csrc/k2_dispatch.h
new file mode 100644
index 0000000..8950b81
--- /dev/null
+++ b/csrc/k2_dispatch.h
@@ -0,0 +1,114 @@
+#pragma once
+
+#include <cstdint>
+#include <string_view>
+
+namespace flash_kda {
+
+constexpr int64_t kK2ChunkSize = 16;
+constexpr int64_t kV1AAutoMinChunks = 128;
+
+enum class K2Mode { Baseline, V1A, V1AE, Auto, Invalid };
+enum class K2Implementation { Baseline, V1A, V1AE, Unsupported };
+
+struct K2DispatchConfig {
+    bool v1a_compiled = false;
+    bool v1ae_compiled = false;
+    int compute_major = 0;
+    int compute_minor = 0;
+    bool is_varlen = false;
+    bool has_state_in = false;
+    bool has_state_out = false;
+    bool state_fp32 = false;
+    int64_t total_tokens = 0;
+    int64_t sequences = 0;
+    int64_t heads = 0;
+};
+
+constexpr K2Mode parse_k2_mode(std::string_view value) {
+    return value == "baseline" ? K2Mode::Baseline
+         : value == "v1a" ? K2Mode::V1A
+         : value == "v1ae" ? K2Mode::V1AE
+         : value == "auto" ? K2Mode::Auto
+         : K2Mode::Invalid;
+}
+
+constexpr bool v1a_supported(K2DispatchConfig const& config) {
+    return config.v1a_compiled && config.compute_major == 10 &&
+           config.compute_minor == 3 && !config.is_varlen &&
+           config.has_state_in && config.has_state_out &&
+           !config.state_fp32 && config.sequences > 0;
+}
+
+constexpr int64_t chunks_per_sequence(K2DispatchConfig const& config) {
+    if (config.is_varlen || config.sequences <= 0) return 0;
+    const int64_t sequence_length = config.total_tokens / config.sequences;
+    return (sequence_length + kK2ChunkSize - 1) / kK2ChunkSize;
+}
+
+constexpr bool v1ae_supported(K2DispatchConfig const& config) {
+    return config.v1ae_compiled && config.compute_major == 10 &&
+           config.compute_minor == 3 && !config.is_varlen &&
+           config.has_state_in && config.has_state_out &&
+           !config.state_fp32 && config.sequences > 0;
+}
+
+// Empirical B300 H=64 policy, not a universal grid-size rule.
+constexpr K2Implementation h64_auto_candidate(int64_t batch, int64_t chunks) {
+    using Impl = K2Implementation;
+    switch (batch) {
+    case 1:
+        return chunks < 80 ? Impl::Baseline : chunks < 160 ? Impl::V1AE : Impl::V1A;
+    case 2:
+        return chunks < 32 ? Impl::Baseline : chunks < 96 ? Impl::V1AE : Impl::V1A;
+    case 3:
+        // Below 24 is a conservative, unmeasured lower fallback.
+        return chunks < 24 ? Impl::Baseline : chunks < 320 ? Impl::V1AE : Impl::V1A;
+    case 4:
+        return chunks < 32 ? Impl::Baseline : chunks < 256 ? Impl::V1AE : Impl::V1A;
+    case 5:
+    case 6:
+    case 7:
+        // V1aE wins through 320 inclusive. Neither <24 nor >320 was mapped;
+        // >320 -> V1a is a conservative fallback, NOT a measured crossover.
+        return chunks < 24 ? Impl::Baseline : chunks <= 320 ? Impl::V1AE : Impl::V1A;
+    case 8:
+        return chunks < 24 || chunks >= 320 ? Impl::Baseline : Impl::V1AE;
+    default:
+        // B>8 is unvalidated; do not extrapolate B8 behavior.
+        return Impl::Baseline;
+    }
+}
+
+constexpr K2Implementation select_auto_k2(K2DispatchConfig const& config) {
+    const auto chunks = chunks_per_sequence(config);
+    if (config.heads != 64) {
+        return v1a_supported(config) && chunks >= kV1AAutoMinChunks
+            ? K2Implementation::V1A : K2Implementation::Baseline;
+    }
+    const auto candidate = h64_auto_candidate(config.sequences, chunks);
+    if (candidate == K2Implementation::V1AE && v1ae_supported(config)) return candidate;
+    if (candidate == K2Implementation::V1A && v1a_supported(config)) return candidate;
+    // Unsupported/uncompiled winners fall back to baseline.
+    return K2Implementation::Baseline;
+}
+
+constexpr K2Implementation select_k2_implementation(
+    K2Mode mode, K2DispatchConfig const& config) {
+    switch (mode) {
+    case K2Mode::Baseline:
+        return K2Implementation::Baseline;
+    case K2Mode::V1A:
+        return v1a_supported(config)
+            ? K2Implementation::V1A : K2Implementation::Unsupported;
+    case K2Mode::V1AE:
+        return v1ae_supported(config)
+            ? K2Implementation::V1AE : K2Implementation::Unsupported;
+    case K2Mode::Auto:
+        return select_auto_k2(config);
+    default:
+        return K2Implementation::Unsupported;
+    }
+}
+
+}  // namespace flash_kda
diff --git a/csrc/smxx/fwd_kernel2.cuh b/csrc/smxx/fwd_kernel2.cuh
index 26f73fe..1b56ba8 100644
--- a/csrc/smxx/fwd_kernel2.cuh
+++ b/csrc/smxx/fwd_kernel2.cuh
@@ -112,6 +112,48 @@ struct SharedStorageK2 {
 };
 
 // ==================== Kernel 2: Recurrence ====================
+// V1a owns the recurrent state in the 128 compute threads. This storage keeps
+// only the production pipelines; no full-state SMEM allocation or FP32 alias.
+template <class Layouts, int InputStages, int OutputStages>
+struct SharedStorageK2V1A {
+    using BF16 = cutlass::bfloat16_t;
+    using VOLayout = typename Layouts::VOLayout;
+    using BetaSmemLayout = typename Layouts::BetaSmemLayout;
+    using GTotalLayout = typename Layouts::GTotalLayout;
+    using LMLayout = typename Layouts::LMLayout;
+    using MMALayout = typename Layouts::MMALayout;
+    struct InputStorage {
+        alignas(128) cute::ArrayEngine<BF16, cute::cosize_v<VOLayout>> v;
+        alignas(128) cute::ArrayEngine<BF16, cute::cosize_v<BetaSmemLayout>> beta;
+        alignas(128) cute::ArrayEngine<BF16, cute::cosize_v<MMALayout>> k_decayed;
+        alignas(128) cute::ArrayEngine<BF16, cute::cosize_v<MMALayout>> q_decayed;
+        alignas(128) cute::ArrayEngine<BF16, cute::cosize_v<MMALayout>> k_restored;
+        alignas(128) cute::ArrayEngine<float, cute::cosize_v<GTotalLayout>> g_total;
+        alignas(128) cute::ArrayEngine<BF16, cute::cosize_v<LMLayout>> INV;
+        alignas(128) cute::ArrayEngine<BF16, cute::cosize_v<LMLayout>> Mqk;
+    };
+    struct OutputStorage { alignas(128) cute::ArrayEngine<BF16, cute::cosize_v<VOLayout>> out; };
+    InputStorage input[InputStages];
+    OutputStorage output[OutputStages];
+    typename cutlass::PipelineTmaAsync<InputStages>::SharedStorage load_pipeline;
+    typename cutlass::PipelineAsync<OutputStages>::SharedStorage store_pipeline;
+};
+
+// V1aE final-state egress only: the recurrent state remains in registers.
+template <class Layouts, int InputStages, int OutputStages>
+struct SharedStorageK2V1AE : SharedStorageK2V1A<Layouts, InputStages, OutputStages> {
+    using BF16 = cutlass::bfloat16_t;
+    using StateSmemLayout = typename Layouts::StateSmemLayout;
+    alignas(128) cute::ArrayEngine<BF16, cute::cosize_v<StateSmemLayout>> final_state;
+};
+
+template <class Layouts, int InputStages, int OutputStages, bool UseV1A,
+          int V1AEgress = 0>
+using SelectedSharedStorageK2 =
+    cute::conditional_t<(V1AEgress != 0), SharedStorageK2V1AE<Layouts, InputStages, OutputStages>,
+      cute::conditional_t<UseV1A, SharedStorageK2V1A<Layouts, InputStages, OutputStages>,
+                          SharedStorageK2<Layouts, InputStages, OutputStages>>>;
+
 template <
     class TmaLoadV,
     class TmaLoadBeta,
@@ -128,7 +170,9 @@ template <
     bool HasStateIn = true,
     bool HasStateOut = true,
     bool StateFP32 = false,
-    bool IsVarlen = true
+    bool IsVarlen = true,
+    bool UseV1A = false,
+    int V1AEgress = 0
 >
 __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
     CUTE_GRID_CONSTANT TmaLoadV const tma_load_v,
@@ -142,6 +186,8 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
     CUTE_GRID_CONSTANT TmaLoadState const tma_load_initial_state,
     CUTE_GRID_CONSTANT TmaStoreState const tma_store_final_state,
     CUTE_GRID_CONSTANT TmaStoreOut const tma_store_out,
+    void const* initial_state_raw_ptr,
+    void* final_state_raw_ptr,
     cutlass::bfloat16_t* out_raw_ptr,
     int T_total,
     int H,
@@ -149,6 +195,11 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
     int64_t const* cu_seqlens,
     int total_tiles
 ) {
+    constexpr bool UseRegisterState = UseV1A || V1AEgress != 0;
+    static_assert(!UseRegisterState || (D == 128 && CHUNK == 16 && HasStateIn && HasStateOut &&
+                              !StateFP32 && !IsVarlen),
+                  "V1a supports fixed-length BF16 state input/output only");
+    static_assert(V1AEgress == 0 || V1AEgress == 1, "unknown V1a egress mode");
     using BF16 = cutlass::bfloat16_t;
     using FP16 = cutlass::half_t;
     using Layouts = K2Layouts<D, CHUNK>;
@@ -182,7 +233,8 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
 
     // --- shared memory
     extern __shared__ __align__(128) unsigned char shared_mem[];
-    using SharedStorageT = SharedStorageK2<Layouts, InputStages, OutputStages>;
+    using SharedStorageT = SelectedSharedStorageK2<Layouts, InputStages, OutputStages,
+                                                   UseV1A, V1AEgress>;
     SharedStorageT& shared_storage = *reinterpret_cast<SharedStorageT*>(shared_mem);
 
     // --- warp specialization
@@ -238,7 +290,9 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
 
     // --- Load initial state
 #ifndef TMA_DISABLE_ALL
-    if constexpr (HasStateIn && !StateFP32) {
+    if constexpr (UseRegisterState) {
+        // V1a loads state directly into compute-thread registers below.
+    } else if constexpr (HasStateIn && !StateFP32) {
         // BF16 state: TMA load directly into state_acc
         if (warp_role == WarpRole::LOAD_QKG && lane_predicate) {
             using BarrierType = cutlass::arch::ClusterTransactionBarrier::ValueType;
@@ -431,6 +485,41 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
 #endif
         int compute_tid = threadIdx.x;
 
+        // Persistent V1a state is stored exactly in the Phase-1 MMA B
+        // distribution: two 16-column blocks per warp, eight key blocks.
+        auto state_mma = make_tiled_mma(
+            MMA_Atom<SM80_16x8x16_F32BF16BF16F32_TN>{},
+            Layout<Shape<_1,_1>>{}, Tile<_16,_16,_16>{});
+        auto state_thr = state_mma.get_slice(compute_tid % 32);
+        auto state_identity = make_identity_tensor(make_shape(Int<16>{}, Int<16>{}));
+        auto state_ref = make_tensor(make_gmem_ptr(static_cast<BF16*>(nullptr)),
+                                     make_layout(make_shape(Int<16>{}, Int<16>{}), LayoutRight{}));
+        using PersistentBFrag = decltype(state_thr.partition_fragment_B(state_ref));
+        PersistentBFrag state_regs[2][8];
+        if constexpr (UseRegisterState) {
+            auto state_ptr = static_cast<BF16 const*>(initial_state_raw_ptr) +
+                             size_t(seq_idx * H + head_idx) * D * D;
+            const int state_warp = compute_tid / 32;
+#pragma unroll
+            for (int bi = 0; bi < 2; ++bi) {
+#pragma unroll
+                for (int kb = 0; kb < 8; ++kb) {
+                    auto coords = state_thr.partition_B(state_identity);
+#pragma unroll
+                    for (int i = 0; i < size(state_regs[bi][kb]); ++i) {
+                        auto coord = coords(i);
+                        // CuTe's B identity coordinates are exposed as (N,K)
+                        // for this TN atom. Match the baseline C->MOVM_T->B
+                        // distribution; treating them as (K,N) transposes
+                        // every 16x16 state tile.
+                        int key = kb * 16 + int(get<1>(coord));
+                        int value = (state_warp * 2 + bi) * 16 + int(get<0>(coord));
+                        state_regs[bi][kb](i) = state_ptr[value * D + key];
+                    }
+                }
+            }
+        }
+
         for (int t = 0; t < t_tiles; ++t) {
 #ifndef TMA_DISABLE_ALL
             store_pipeline.producer_acquire(out_write);
@@ -454,8 +543,11 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
             Tensor INV = make_tensor(make_smem_ptr(shared_storage.input[load_stage].INV.begin()), LMLayout{});
             Tensor Mqk = make_tensor(make_smem_ptr(shared_storage.input[load_stage].Mqk.begin()), LMLayout{});
 
-            Tensor s_acc = make_tensor(make_smem_ptr(shared_storage.state_acc.begin()), StateSmemLayout{});
-            Tensor s_acc_T = make_tensor(make_smem_ptr(shared_storage.state_acc.begin()), TransposedStateSmemLayout{});
+            BF16* canonical_state_ptr;
+            if constexpr (UseRegisterState) canonical_state_ptr = shared_storage.input[load_stage].v.begin();
+            else canonical_state_ptr = shared_storage.state_acc.begin();
+            Tensor s_acc = make_tensor(make_smem_ptr(canonical_state_ptr), StateSmemLayout{});
+            Tensor s_acc_T = make_tensor(make_smem_ptr(canonical_state_ptr), TransposedStateSmemLayout{});
 
             // Fused MMA: v_sub, v_beta, U=INV@v, out=q@s, out+=Mqk@U, s_acc_update
             // Each warp handles TWO 16x16 column blocks (N=128 / 4 warps = 32 = 2 x 16)
@@ -537,34 +629,43 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
                 local_tile(k_decayed, make_shape(Int<16>{}, Int<16>{}), make_coord(0, 0))), tCrAi_k_view);
             copy(smem_tiled_copy_A, smem_thr_copy_A.partition_S(
                 local_tile(q_decayed, make_shape(Int<16>{}, Int<16>{}), make_coord(0, 0))), tCrAi_q_view);
-            copy(smem_tiled_copy_B, smem_thr_copy_B.partition_S(
-                local_tile(s_acc, make_shape(Int<16>{}, Int<16>{}), make_coord(warp_id * 2, 0))), tCrBi_view);
+            if constexpr (!UseRegisterState) {
+                copy(smem_tiled_copy_B, smem_thr_copy_B.partition_S(
+                    local_tile(s_acc, make_shape(Int<16>{}, Int<16>{}), make_coord(warp_id * 2, 0))), tCrBi_view);
+            }
 
             #pragma unroll
             for (int k = 0; k < K_BLOCKS; ++k) {
                 cute::transform(tCrAi_k, tCrA_k, cute::identity{});
                 cute::transform(tCrAi_q, tCrA_q, cute::identity{});
-                cute::transform(tCrBi, tCrB, cute::identity{});
-
-                copy(smem_tiled_copy_B, smem_thr_copy_B.partition_S(
-                    local_tile(s_acc, make_shape(Int<16>{}, Int<16>{}), make_coord(warp_id * 2 + 1, k))), tCrBi_view);
-
-                gemm(thr_mma, tCrA_k(_,_,Int<0>{}), tCrB(_,_,Int<0>{}), u_acc[0]);
-                gemm(thr_mma, tCrA_q(_,_,Int<0>{}), tCrB(_,_,Int<0>{}), out_acc[0]);
-
-                cute::transform(tCrBi, tCrB, cute::identity{});
+                if constexpr (UseRegisterState) {
+                    gemm(thr_mma, tCrA_k(_,_,Int<0>{}), state_regs[0][k](_,_,Int<0>{}), u_acc[0]);
+                    gemm(thr_mma, tCrA_q(_,_,Int<0>{}), state_regs[0][k](_,_,Int<0>{}), out_acc[0]);
+                    gemm(thr_mma, tCrA_k(_,_,Int<0>{}), state_regs[1][k](_,_,Int<0>{}), u_acc[1]);
+                    gemm(thr_mma, tCrA_q(_,_,Int<0>{}), state_regs[1][k](_,_,Int<0>{}), out_acc[1]);
+                } else {
+                    cute::transform(tCrBi, tCrB, cute::identity{});
+                    copy(smem_tiled_copy_B, smem_thr_copy_B.partition_S(
+                        local_tile(s_acc, make_shape(Int<16>{}, Int<16>{}), make_coord(warp_id * 2 + 1, k))), tCrBi_view);
+                    gemm(thr_mma, tCrA_k(_,_,Int<0>{}), tCrB(_,_,Int<0>{}), u_acc[0]);
+                    gemm(thr_mma, tCrA_q(_,_,Int<0>{}), tCrB(_,_,Int<0>{}), out_acc[0]);
+                    cute::transform(tCrBi, tCrB, cute::identity{});
+                }
 
                 if (k + 1 < K_BLOCKS) {
                     copy(smem_tiled_copy_A, smem_thr_copy_A.partition_S(
                         local_tile(k_decayed, make_shape(Int<16>{}, Int<16>{}), make_coord(0, k + 1))), tCrAi_k_view);
                     copy(smem_tiled_copy_A, smem_thr_copy_A.partition_S(
                         local_tile(q_decayed, make_shape(Int<16>{}, Int<16>{}), make_coord(0, k + 1))), tCrAi_q_view);
-                    copy(smem_tiled_copy_B, smem_thr_copy_B.partition_S(
-                        local_tile(s_acc, make_shape(Int<16>{}, Int<16>{}), make_coord(warp_id * 2, k + 1))), tCrBi_view);
+                    if constexpr (!UseRegisterState) {
+                        copy(smem_tiled_copy_B, smem_thr_copy_B.partition_S(
+                            local_tile(s_acc, make_shape(Int<16>{}, Int<16>{}), make_coord(warp_id * 2, k + 1))), tCrBi_view);
+                    }
+                }
+                if constexpr (!UseRegisterState) {
+                    gemm(thr_mma, tCrA_k(_,_,Int<0>{}), tCrB(_,_,Int<0>{}), u_acc[1]);
+                    gemm(thr_mma, tCrA_q(_,_,Int<0>{}), tCrB(_,_,Int<0>{}), out_acc[1]);
                 }
-
-                gemm(thr_mma, tCrA_k(_,_,Int<0>{}), tCrB(_,_,Int<0>{}), u_acc[1]);
-                gemm(thr_mma, tCrA_q(_,_,Int<0>{}), tCrB(_,_,Int<0>{}), out_acc[1]);
             }
 
             // ======== Phase 2: Cast out (keep in regs), load v/INV/beta ========
@@ -657,6 +758,43 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
             }
 
             // ======== Phase 6: s_acc update ========
+            if constexpr (UseRegisterState) {
+                constexpr int S_M_BLOCKS = 8;
+                Tensor tCrAi_kr = make_fragment_like<BF16>(thr_mma.partition_fragment_A(A_ref));
+                auto tCrAi_kr_view = smem_thr_copy_A_T.retile_D(tCrAi_kr);
+                AFragT a_kr;
+#pragma unroll
+                for (int m = 0; m < S_M_BLOCKS; ++m) {
+                    Tensor kr_block = local_tile(k_restored_t, make_shape(Int<16>{}, Int<16>{}), make_coord(m, 0));
+                    copy(smem_tiled_copy_A_T, smem_thr_copy_A_T.partition_S(kr_block), tCrAi_kr_view);
+                    cute::transform(tCrAi_kr, a_kr, cute::identity{});
+#pragma unroll
+                    for (int bi = 0; bi < 2; ++bi) {
+                        clear(u_acc[bi]);
+                        gemm(thr_mma, a_kr(_,_,Int<0>{}), tCrB_u_arr[bi](_,_,Int<0>{}), u_acc[bi]);
+                        SFragT state_c;
+                        auto* b = reinterpret_cast<uint32_t*>(&state_regs[bi][m](0));
+                        auto* c = reinterpret_cast<uint32_t*>(&state_c(0));
+#pragma unroll
+                        for (int word = 0; word < 4; ++word) SM75_U32x1_MOVM_T::copy(b[word], c[word]);
+                        float g0 = g_total(m * 16 + group_id);
+                        float g1 = g_total(m * 16 + group_id + 8);
+#pragma unroll
+                        for (int a = 0; a < 2; ++a) {
+#pragma unroll
+                            for (int d = 0; d < 2; ++d) {
+                                auto c0 = make_coord(make_coord(a, 0), 0, d);
+                                auto c1 = make_coord(make_coord(a, 1), 0, d);
+                                state_c(c0) = BF16(bf16_to_f32(state_c(c0)) * g0 + u_acc[bi](c0));
+                                state_c(c1) = BF16(bf16_to_f32(state_c(c1)) * g1 + u_acc[bi](c1));
+                            }
+                        }
+#pragma unroll
+                        for (int word = 0; word < 4; ++word) SM75_U32x1_MOVM_T::copy(c[word], b[word]);
+                    }
+                }
+            }
+            else {
             // s_acc[D, D] = s_acc * g_total + k_restored_t[D, 16] @ U[16, D]
             // Each warp handles columns [warp_id*32, (warp_id+1)*32] = 2 x 16x16 blocks
             // U is already in tCrB_u_arr[0..1] as B operands (from Phase 4 MOVM_T)
@@ -730,6 +868,7 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
                 }
             }
             }
+            }
             compute_barrier.arrive_and_wait();
 
 #ifndef TMA_DISABLE_ALL
@@ -740,6 +879,43 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
             ++out_write;
 #endif
         }
+        if constexpr (UseRegisterState && V1AEgress == 0) {
+            auto state_ptr = static_cast<BF16*>(final_state_raw_ptr) +
+                             size_t(seq_idx * H + head_idx) * D * D;
+            const int state_warp = compute_tid / 32;
+#pragma unroll
+            for (int bi = 0; bi < 2; ++bi) {
+#pragma unroll
+                for (int kb = 0; kb < 8; ++kb) {
+                    auto coords = state_thr.partition_B(state_identity);
+#pragma unroll
+                    for (int i = 0; i < size(state_regs[bi][kb]); ++i) {
+                        auto coord = coords(i);
+                        int key = kb * 16 + int(get<1>(coord));
+                        int value = (state_warp * 2 + bi) * 16 + int(get<0>(coord));
+                        state_ptr[value * D + key] = state_regs[bi][kb](i);
+                    }
+                }
+            }
+        } else if constexpr (V1AEgress == 1) {
+            Tensor s_final = make_tensor(
+                make_smem_ptr(shared_storage.final_state.begin()), StateSmemLayout{});
+            const int state_warp = compute_tid / 32;
+#pragma unroll
+            for (int bi = 0; bi < 2; ++bi) {
+#pragma unroll
+                for (int kb = 0; kb < 8; ++kb) {
+                    auto coords = state_thr.partition_B(state_identity);
+#pragma unroll
+                    for (int i = 0; i < size(state_regs[bi][kb]); ++i) {
+                        auto coord = coords(i);
+                        int key = kb * 16 + int(get<1>(coord));
+                        int value = (state_warp * 2 + bi) * 16 + int(get<0>(coord));
+                        s_final(value, key) = state_regs[bi][kb](i);
+                    }
+                }
+            }
+        }
     }
 
 #ifndef TMA_DISABLE_ALL
@@ -783,7 +959,7 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
             ++out_read;
         }
 
-        if constexpr (HasStateOut && !StateFP32) {
+        if constexpr (HasStateOut && !StateFP32 && !UseRegisterState) {
             // BF16 state: TMA store directly from state_acc
             Tensor g_final = tma_store_final_state.get_tma_tensor(make_shape(N * H, D, D));
             auto state_off = g_final.layout()(seq_idx * H + head_idx, 0, 0);
@@ -801,6 +977,26 @@ __global__ void __launch_bounds__(NumThreads) _flash_kda_fwd_recurrence(
         }
     }
 
+    if constexpr (V1AEgress == 1) {
+        // The STORE warp has drained the normal output pipeline before this
+        // one-time recurrence epilogue rendezvous.
+        cutlass::arch::fence_view_async_shared();
+        __syncthreads();
+        if (warp_role == WarpRole::STORE && lane_predicate) {
+            Tensor g_final = tma_store_final_state.get_tma_tensor(make_shape(N * H, D, D));
+            auto state_off = g_final.layout()(seq_idx * H + head_idx, 0, 0);
+            Tensor g_final_tile = make_tensor(g_final.data() + state_off,
+                make_layout(make_shape(Int<1>{}, Int<D>{}, Int<D>{}), stride(g_final.layout())));
+            Tensor s_final = make_tensor(
+                make_smem_ptr(shared_storage.final_state.begin()), TMAStateSmemLayout{});
+            auto cta_tma_store_state = tma_store_final_state.get_slice(Int<0>{});
+            cute::copy(tma_store_final_state,
+                       cta_tma_store_state.partition_S(s_final),
+                       cta_tma_store_state.partition_D(g_final_tile));
+            tma_store_arrive();
+        }
+    }
+
     if constexpr (HasStateOut && StateFP32) {
         // FP32 state: all threads sync, convert bf16->fp32, then STORE warp does TMA
         using FP32StateSmemLayout = typename Layouts::FP32StateSmemLayout;
diff --git a/csrc/smxx/fwd_launch.cu b/csrc/smxx/fwd_launch.cu
index 91a67a7..8c105d5 100644
--- a/csrc/smxx/fwd_launch.cu
+++ b/csrc/smxx/fwd_launch.cu
@@ -3,7 +3,8 @@
 #include "fwd_kernel2.cuh"
 
 // ==================== launch_fwd ====================
-template <int D, bool HasStateIn, bool HasStateOut, bool StateFP32, bool IsVarlen>
+template <int D, bool HasStateIn, bool HasStateOut, bool StateFP32, bool IsVarlen,
+          bool UseV1A, int V1AEgress>
 void launch_fwd(
     cutlass::bfloat16_t const* q_ptr,
     cutlass::bfloat16_t const* k_ptr,
@@ -184,7 +185,8 @@ void launch_fwd(
 #if BLOCK_LEVEL_K2 >= 0
     {
         constexpr int kK2Threads = 32 * 2 + 128;
-        using SharedStorageK2T = SharedStorageK2<K2L, kInputStages, kOutputStages>;
+        using SharedStorageK2T = SelectedSharedStorageK2<K2L, kInputStages, kOutputStages,
+                                                        UseV1A, V1AEgress>;
         int smem_size_k2 = sizeof(SharedStorageK2T);
 
         auto kernel2 = _flash_kda_fwd_recurrence<
@@ -195,11 +197,18 @@ void launch_fwd(
             decltype(tma_store_final_state),
             decltype(tma_store_out),
             CHUNK, D, kInputStages, kOutputStages, kK2Threads,
-            HasStateIn, HasStateOut, StateFP32, IsVarlen
+            HasStateIn, HasStateOut, StateFP32, IsVarlen, UseV1A, V1AEgress
         >;
 
         cudaFuncSetAttribute(kernel2, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size_k2);
 
+        if constexpr (UseV1A || V1AEgress != 0) {
+            // Preserve the validated register-state shared-memory carveout
+            // preference for these specializations only.
+            cudaFuncSetAttribute(kernel2, cudaFuncAttributePreferredSharedMemoryCarveout,
+                                 cudaSharedmemCarveoutMaxShared);
+        }
+
         dim3 grid_k2(N, H);
         dim3 block_k2(kK2Threads);
 
@@ -210,6 +219,7 @@ void launch_fwd(
             tma_load_initial_state,
             tma_store_final_state,
             tma_store_out,
+            initial_state_ptr, final_state_ptr,
             out_ptr, T_total, H, N, cu_seqlens_ptr, total_tiles
         );
     }
@@ -217,8 +227,8 @@ void launch_fwd(
 }
 
 // Explicit instantiations
-#define INSTANTIATE_LAUNCH_FWD(D, HI, HO, FP32, VL) \
-    template void launch_fwd<D, HI, HO, FP32, VL>( \
+#define INSTANTIATE_LAUNCH_FWD(D, HI, HO, FP32, VL, V1A, EGRESS) \
+    template void launch_fwd<D, HI, HO, FP32, VL, V1A, EGRESS>( \
         cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, \
         cutlass::bfloat16_t const*, cutlass::bfloat16_t const*, \
         cutlass::bfloat16_t const*, void const*, float, void*, \
@@ -226,13 +236,20 @@ void launch_fwd(
         int64_t const*, float const*, float const*, float, cudaStream_t);
 
 #define INSTANTIATE_STATE_VARIANTS(VL) \
-    INSTANTIATE_LAUNCH_FWD(128, true,  true,  false, VL) \
-    INSTANTIATE_LAUNCH_FWD(128, true,  true,  true,  VL) \
-    INSTANTIATE_LAUNCH_FWD(128, false, false, false, VL) \
-    INSTANTIATE_LAUNCH_FWD(128, false, true,  false, VL) \
-    INSTANTIATE_LAUNCH_FWD(128, true,  false, false, VL) \
-    INSTANTIATE_LAUNCH_FWD(128, false, true,  true,  VL) \
-    INSTANTIATE_LAUNCH_FWD(128, true,  false, true,  VL)
+    INSTANTIATE_LAUNCH_FWD(128, true,  true,  false, VL, false, 0) \
+    INSTANTIATE_LAUNCH_FWD(128, true,  true,  true,  VL, false, 0) \
+    INSTANTIATE_LAUNCH_FWD(128, false, false, false, VL, false, 0) \
+    INSTANTIATE_LAUNCH_FWD(128, false, true,  false, VL, false, 0) \
+    INSTANTIATE_LAUNCH_FWD(128, true,  false, false, VL, false, 0) \
+    INSTANTIATE_LAUNCH_FWD(128, false, true,  true,  VL, false, 0) \
+    INSTANTIATE_LAUNCH_FWD(128, true,  false, true,  VL, false, 0)
 
 INSTANTIATE_STATE_VARIANTS(true)   // varlen
 INSTANTIATE_STATE_VARIANTS(false)  // non-varlen
+
+#if defined(FLASH_KDA_ENABLE_V1A)
+INSTANTIATE_LAUNCH_FWD(128, true, true, false, false, true, 0)
+#endif
+#if defined(FLASH_KDA_ENABLE_V1AE)
+INSTANTIATE_LAUNCH_FWD(128, true, true, false, false, false, 1)
+#endif
diff --git a/setup.py b/setup.py
index 76e44ed..9c600f5 100644
--- a/setup.py
+++ b/setup.py
@@ -66,7 +66,9 @@ def get_arch_flags():
             os.path.join(this_dir, 'csrc'),
         ],
         extra_compile_args={
-            'cxx': ['-O3', '-Wno-psabi'],
+            'cxx': ['-O3', '-Wno-psabi',
+                    *(['-DFLASH_KDA_ENABLE_V1A=1'] if is_flag_set('FLASH_KDA_ENABLE_V1A') else []),
+                    *(['-DFLASH_KDA_ENABLE_V1AE=1'] if is_flag_set('FLASH_KDA_ENABLE_V1AE') else [])],
             'nvcc': [
                 '-O3',
                 '-U__CUDA_NO_HALF_OPERATORS__',
@@ -80,6 +82,8 @@ def get_arch_flags():
                 '-lineinfo',
                 *get_nvcc_thread_args(),
                 *get_arch_flags(),
+                *(['-DFLASH_KDA_ENABLE_V1A=1'] if is_flag_set('FLASH_KDA_ENABLE_V1A') else []),
+                *(['-DFLASH_KDA_ENABLE_V1AE=1'] if is_flag_set('FLASH_KDA_ENABLE_V1AE') else []),
             ],
         },
     )
diff --git a/tests/test_k2_dispatch.cpp b/tests/test_k2_dispatch.cpp
new file mode 100644
index 0000000..cb90de1
--- /dev/null
+++ b/tests/test_k2_dispatch.cpp
@@ -0,0 +1,117 @@
+#include "k2_dispatch.h"
+
+#include <cstdio>
+#include <cstdlib>
+#include <initializer_list>
+
+using flash_kda::K2DispatchConfig;
+using flash_kda::K2Implementation;
+using flash_kda::K2Mode;
+using flash_kda::parse_k2_mode;
+using flash_kda::select_k2_implementation;
+using Impl = K2Implementation;
+
+static int checks = 0;
+
+static void check(bool passed, const char* expression, int line) {
+    ++checks;
+    if (!passed) {
+        std::fprintf(stderr, "line %d: %s\n", line, expression);
+        std::exit(EXIT_FAILURE);
+    }
+}
+
+// Keep checks active even when the compiler defines NDEBUG.
+#define CHECK(expression) check((expression), #expression, __LINE__)
+
+static K2DispatchConfig supported(int64_t batch, int64_t chunks, int64_t heads = 64) {
+    K2DispatchConfig config;
+    config.v1a_compiled = config.v1ae_compiled = true;
+    config.compute_major = 10;
+    config.compute_minor = 3;
+    config.has_state_in = config.has_state_out = true;
+    config.sequences = batch;
+    config.total_tokens = batch * chunks * 16;
+    config.heads = heads;
+    return config;
+}
+
+int main() {
+    CHECK(parse_k2_mode("baseline") == K2Mode::Baseline);
+    CHECK(parse_k2_mode("v1a") == K2Mode::V1A);
+    CHECK(parse_k2_mode("v1ae") == K2Mode::V1AE);
+    CHECK(parse_k2_mode("auto") == K2Mode::Auto);
+    for (auto value : {"unknown", "", "V1A"}) {
+        CHECK(parse_k2_mode(value) == K2Mode::Invalid);
+        CHECK(select_k2_implementation(parse_k2_mode(value), supported(1, 160)) == Impl::Unsupported);
+    }
+
+    struct Boundary { int batch, chunks; Impl expected; };
+    const Boundary boundaries[] = {
+        {1, 79, Impl::Baseline}, {1, 80, Impl::V1AE}, {1, 159, Impl::V1AE}, {1, 160, Impl::V1A},
+        {2, 31, Impl::Baseline}, {2, 32, Impl::V1AE}, {2, 95, Impl::V1AE}, {2, 96, Impl::V1A},
+        {3, 23, Impl::Baseline}, {3, 24, Impl::V1AE}, {3, 319, Impl::V1AE}, {3, 320, Impl::V1A},
+        {4, 31, Impl::Baseline}, {4, 32, Impl::V1AE}, {4, 255, Impl::V1AE}, {4, 256, Impl::V1A},
+        {5, 23, Impl::Baseline}, {5, 24, Impl::V1AE}, {5, 320, Impl::V1AE}, {5, 321, Impl::V1A},
+        {6, 23, Impl::Baseline}, {6, 24, Impl::V1AE}, {6, 320, Impl::V1AE}, {6, 321, Impl::V1A},
+        {7, 23, Impl::Baseline}, {7, 24, Impl::V1AE}, {7, 320, Impl::V1AE}, {7, 321, Impl::V1A},
+        {8, 23, Impl::Baseline}, {8, 24, Impl::V1AE}, {8, 319, Impl::V1AE}, {8, 320, Impl::Baseline},
+        {9, 24, Impl::Baseline}, {9, 512, Impl::Baseline}, {16, 512, Impl::Baseline},
+    };
+    for (const auto& test : boundaries) {
+        auto config = supported(test.batch, test.chunks);
+        CHECK(flash_kda::chunks_per_sequence(config) == test.chunks);
+        CHECK(select_k2_implementation(K2Mode::Auto, config) == test.expected);
+        // A partial final tile counts too, including at each AUTO boundary.
+        for (int tail : {1, 15}) {
+            auto partial = config;
+            partial.total_tokens = test.batch * ((test.chunks - 1) * 16 + tail);
+            CHECK(flash_kda::chunks_per_sequence(partial) == test.chunks);
+            CHECK(select_k2_implementation(K2Mode::Auto, partial) == test.expected);
+        }
+        // Explicit requests ignore the empirical AUTO thresholds.
+        CHECK(select_k2_implementation(K2Mode::V1A, config) == Impl::V1A);
+        CHECK(select_k2_implementation(K2Mode::V1AE, config) == Impl::V1AE);
+        CHECK(select_k2_implementation(K2Mode::Baseline, config) == Impl::Baseline);
+        for (bool v1a : {false, true}) for (bool v1ae : {false, true}) {
+            config.v1a_compiled = v1a;
+            config.v1ae_compiled = v1ae;
+            auto expected = test.expected;
+            if ((expected == Impl::V1A && !v1a) || (expected == Impl::V1AE && !v1ae))
+                expected = Impl::Baseline;
+            CHECK(select_k2_implementation(K2Mode::Auto, config) == expected);
+            CHECK(select_k2_implementation(K2Mode::V1A, config) ==
+                  (v1a ? Impl::V1A : Impl::Unsupported));
+            CHECK(select_k2_implementation(K2Mode::V1AE, config) ==
+                  (v1ae ? Impl::V1AE : Impl::Unsupported));
+        }
+    }
+
+    for (int heads : {1, 32, 65, 128}) for (int chunks : {127, 128}) {
+        auto config = supported(4, chunks, heads);
+        CHECK(select_k2_implementation(K2Mode::Auto, config) ==
+              (chunks == 128 ? Impl::V1A : Impl::Baseline));
+        config.v1a_compiled = false;
+        CHECK(select_k2_implementation(K2Mode::Auto, config) == Impl::Baseline);
+    }
+
+    // Exercise unsupported runtime configurations at both kinds of AUTO winner.
+    for (int chunks : {128, 256}) for (int condition = 0; condition < 10; ++condition) {
+        auto config = supported(4, chunks);
+        if (condition == 0) config.is_varlen = true;
+        if (condition == 1) config.state_fp32 = true;
+        if (condition == 2) config.has_state_in = false;
+        if (condition == 3) config.has_state_out = false;
+        if (condition == 4) config.compute_minor = 0; // SM100
+        if (condition == 5) { config.compute_major = 12; config.compute_minor = 0; }
+        if (condition == 6) { config.compute_major = 9; config.compute_minor = 0; }
+        if (condition == 7) { config.compute_major = 8; config.compute_minor = 9; }
+        if (condition == 8) config.sequences = 0;
+        if (condition == 9) config.sequences = -1;
+        CHECK(select_k2_implementation(K2Mode::V1A, config) == Impl::Unsupported);
+        CHECK(select_k2_implementation(K2Mode::V1AE, config) == Impl::Unsupported);
+        CHECK(select_k2_implementation(K2Mode::Auto, config) == Impl::Baseline);
+        CHECK(select_k2_implementation(K2Mode::Baseline, config) == Impl::Baseline);
+    }
+    std::printf("K2 dispatch tests passed (%d checks)\n", checks);
+}
diff --git a/tests/test_k2_dispatch.py b/tests/test_k2_dispatch.py
new file mode 100644
index 0000000..bf66ecd
--- /dev/null
+++ b/tests/test_k2_dispatch.py
@@ -0,0 +1,37 @@
+"""Host-only dispatch test: python tests/test_k2_dispatch.py (no torch/CUDA)."""
+
+import os
+from pathlib import Path
+import shutil
+import subprocess
+import tempfile
+import unittest
+
+
+class TestK2Dispatch(unittest.TestCase):
+    def test_dispatch(self):
+        compiler = shutil.which(os.environ.get("CXX", "c++"))
+        if compiler is None:
+            self.skipTest("requires a C++17 GCC/Clang compiler; set CXX to its executable")
+        root = Path(__file__).resolve().parents[1]
+        env = os.environ.copy()
+        env["PATH"] = str(Path(compiler).parent) + os.pathsep + env.get("PATH", "")
+        with tempfile.TemporaryDirectory(prefix="flashkda-dispatch-") as tmp:
+            executable = Path(tmp) / ("dispatch.exe" if os.name == "nt" else "dispatch")
+            build = subprocess.run(
+                [compiler, "-std=c++17", "-O2", "-Wall", "-Wextra", "-Werror",
+                 "-I", str(root / "csrc"), str(root / "tests/test_k2_dispatch.cpp"),
+                 "-o", str(executable)],
+                capture_output=True, text=True, env=env, timeout=60,
+            )
+            self.assertEqual(build.returncode, 0, build.stdout + build.stderr)
+            run = subprocess.run(
+                [str(executable)], capture_output=True, text=True, env=env, timeout=10,
+            )
+            self.assertEqual(run.returncode, 0, run.stdout + run.stderr)
+            self.assertIn("K2 dispatch tests passed", run.stdout)
+            print(run.stdout.strip())
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/test_sm103_k2.py b/tests/test_sm103_k2.py
new file mode 100644
index 0000000..b2e8fb4
--- /dev/null
+++ b/tests/test_sm103_k2.py
@@ -0,0 +1,52 @@
+"""Opt-in SM103 parity checks; build V1a and V1aE before enabling these tests."""
+
+import os
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+
+pytestmark = pytest.mark.skipif(
+    os.environ.get("FLASH_KDA_TEST_SM103") != "1"
+    or not torch.cuda.is_available()
+    or torch.cuda.get_device_capability() != (10, 3),
+    reason="requires FLASH_KDA_TEST_SM103=1 and an SM103 GPU",
+)
+
+
+@pytest.mark.parametrize(
+    "B,T,H", [(1, 16, 64), (1, 17, 64), (4, 2048, 64), (8, 1024, 64), (1, 2048, 32)]
+)
+@torch.inference_mode()
+def test_sm103_k2_exact(B, T, H, monkeypatch):
+    import flash_kda
+
+    torch.manual_seed(42)
+    D = 128
+    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()
+
+    for mode in ("baseline", "v1a", "v1ae", "auto"):
+        monkeypatch.setenv("FLASH_KDA_K2_IMPL", mode)
+        out = torch.full_like(q, float("nan"))
+        final = torch.full_like(initial, float("nan"))
+        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)
+        torch.cuda.synchronize()
+        assert torch.equal(initial, initial_copy), f"{mode}: initial_state changed"
+        if mode == "baseline":
+            assert torch.isfinite(out).all() and torch.isfinite(final).all()
+            baseline_out, baseline_final = out, final
+        else:
+            assert torch.equal(out, baseline_out), f"{mode}: output mismatch"
+            assert torch.equal(final, baseline_final), f"{mode}: final_state mismatch"