# LeapQuant Kimi Delta Attention checkpoint flush

CuTe DSL window-boundary re-quantization of the recurrent state for Kimi Delta Attention.
This is the flush operation, with the checkpoint format and one-round subspace
iteration specified in [NVlabs/kda#10](https://github.com/NVlabs/kda/pull/10).
It does not implement the per-token decode operation or a complete model layer.

The final accepted implementation achieved **2.999x overall**
and **3.128x on batches 128, 256 and 512** versus the
request's TileLang baseline on NVIDIA B300 SXM6 AC. All six workloads passed the
original accuracy criterion. These are the final paired KDA measurements from
2026-10-24T12:38:34Z; they are kernel timings for the deployment interface, not timings
of allocation, input packing or the functional `run` wrapper.

## Requirements and supported inputs

- NVIDIA B300, CUDA compute capability **10.3**. This release advertises only the
  tested architecture; other GPU architectures are not validated.
- CUDA 13-compatible driver; tested with PyTorch 2.12.1+cu130,
  `nvidia-cutlass-dsl[cu13]==4.6.0` and `cuda-bindings==13.3.1`.
- Python 3.12 and `kernels==0.17.2` were used for package validation.
- Fixed geometry: 32 heads, value/key dimension 128, window 16, rank 4.
  Dense benchmark batches: 16, 32, 64, 128, 256 and 512.
- `codes`: int8 `[B,32,128,128]`; `s_k`/`s_v`: fp32 `[B,32,128]`;
  `u`/`q`: fp16 `[B,32,4,128]`; `kbuf`/`ubuf`: bf16 `[B,32,16,128]`;
  `w`: fp32 `[B,32,16]`; `q0`: contiguous fp32 `[128,4]` on the same CUDA device.
- `gbuf`: fp16 `[B,32,16,128]` in `(0,1]`; `pcum`: fp32 `[B,32,128]`, cumulative non-positive log-gates.

Checkpoint scales are positive and the request's deployment domain has state
magnitudes below 1e4. See the complete [upstream input contract](https://github.com/NVlabs/kda/tree/48f02e6ba0e98a1841ab6f0a2d99a90b2de1a326/requests/conless-leapquant-kda-flush)
and [reference definition](https://huggingface.co/kernels/humanfia-lab/leapquant/blob/v1/kda/definition.json).

## Load and use

```python
import torch
from kernels import get_kernel

op = get_kernel("humanfia-lab/leapquant", version=1, trust_remote_code=True).kda
B, H, V, K, L, R = 16, 32, 128, 128, 16, 4
codes = torch.zeros(B, H, V, K, dtype=torch.int8, device="cuda")
s_k = torch.full((B, H, K), 0.05, device="cuda")
s_v = torch.full((B, H, V), 0.3, device="cuda")
u = torch.zeros(B, H, R, K, dtype=torch.float16, device="cuda")
q = torch.zeros(B, H, R, V, dtype=torch.float16, device="cuda")
kbuf = (torch.randn(B, H, L, K, device="cuda") * 0.3).bfloat16()
ubuf = (torch.randn(B, H, L, V, device="cuda") * 0.3).bfloat16()
w = torch.rand(B, H, L, device="cuda")
gbuf = torch.ones_like(kbuf, dtype=torch.float16)
pcum = torch.zeros(B, H, K, device="cuda")
q0 = torch.linalg.qr(torch.randn(V, R, device="cuda"))[0].contiguous()

new_codes, new_s_k, new_s_v, new_u, new_q = op.run(
    codes, s_k, s_v, u, q, kbuf, ubuf, gbuf, w, pcum, q0
)
```

`run` returns five fresh checkpoint tensors and leaves the supplied inputs
unchanged. It includes allocation/packing overhead. Use serial calls to `run`;
for independent concurrent streams, use an independent caller-owned `Pool` per
stream and keep calls on each pool ordered.

For serving, use the original caller-owned pool interface:

```python
pool = op.Pool(B + 8, device="cuda")  # allocates buffers and JIT-compiles kernels
sl = slice(1, B + 1)                 # slot 0 is reserved
pool.codes[sl], pool.s_k[sl], pool.s_v[sl] = codes, s_k, s_v
pool.uq[sl, :, :, :K], pool.uq[sl, :, :, K:] = u, q
pool.kbuf[sl], pool.ubuf[sl], pool.w[sl] = kbuf, ubuf, w
pool.gbuf[sl] = gbuf
pool.pcum[sl] = pcum
pool.hcnt[sl] = L
idx = torch.arange(1, B + 1, dtype=torch.int32, device="cuda")
op.flush_inplace(pool, idx, q0)       # checkpoint stays in pool; returns None
```

`idx` is a contiguous CUDA int32 vector of valid, unique slot indices. Index 0
is skipped. Only selected heads whose fill counter equals 16 are flushed;
non-due heads and unselected slots remain unchanged. A due head resets `w = 0`, `gbuf = 1`, `pcum = 0`.
The fill counter remains at `L`: the decode step interprets this as "flushed,
empty" and advances it. Flush does not clear it, because other CTAs may still
be scanning the due flags.
`q0` is the supplied fixed start basis, not the previous window's factors.

Construct the pool and perform warmup before CUDA graph capture. Refill the
window buffers and counters before each new window. Calls use the current CUDA
stream on the pool's device. Keep the pool, index vector and `q0` buffers alive
while a graph using their addresses is alive.

## Accuracy and timing

The original benchmark compares decoded checkpoints with the exactly rebuilt
fp64 state for structured edge-case heads, a second window and unstructured
in-domain random inputs at each batch size. For every head, reconstruction
error may be at most 1.01x the reference's error; nearly exact reference heads
use the original absolute cap `1e-3 * mean(abs(S))`. Outputs must be finite and
int8 codes must remain in [-127,127]. The upstream benchmark and thresholds are
included unchanged.

Deployment timing uses `flush_inplace`, all programs due, eight disjoint slot
sets rotated to avoid L2 reuse, four warmups and 16 measured iterations, and
CUDA kernel time from the PyTorch profiler. Large means batches >=128.

| Batch | Final paired speedup | Accuracy |
|---:|---:|---|
| 16 | 2.669x | pass |
| 32 | 2.811x | pass |
| 64 | 3.170x | pass |
| 128 | 3.219x | pass |
| 256 | 3.094x | pass |
| 512 | 3.073x | pass |
| Geomean, all six | 2.999x | 6/6 pass |
| Geomean, batch >=128 | 3.128x | 3/3 pass |

These measurements are workload- and hardware-specific. The functional `run`
wrapper is not the timed deployment entrypoint.

## Reproduce

Download the combined repository, then enter this component directory:

```bash
hf download humanfia-lab/leapquant --repo-type kernel --revision v1 --local-dir leapquant-kernels
cd leapquant-kernels/kda
python -m pip install torch==2.12.1 --index-url https://download.pytorch.org/whl/cu130
python -m pip install -r ../requirements.txt
python benchmark.py --impl ../build/torch-cuda/_kda.py  # six original workloads
python benchmark.py                                    # original TileLang baseline
```

TileLang and `nvcc` on `PATH` are needed only to benchmark the original baseline.
The CuTe implementation JIT-compiles on first use; do not include compilation in
steady-state timings or perform first-use compilation inside graph capture.


The unchanged kernel passed all six original workloads (three input families
each) and deployment checks on B300 before consolidation. The speedup table
retains the paired optimization-run measurement, not a new baseline comparison.

## Matched precision comparison

The output dtypes are unchanged. The original baseline already uses mixed
precision. This implementation retains the rebuilt state and residual in FP32
TMEM, uses scaled FP16 projection operands, and includes local approximations
such as reciprocal operations. See [precision_comparison.json](precision_comparison.json)
for the 2026-10-26 B300 comparison on identical inputs for both implementations.

Both pass all 18 original cases (six batches and three input families).
The original judge averages the batch dimension for each head index. The report
also applies the same bound independently to each (batch item, head), an extra,
stricter check. See the repository root for those results and the KDA outlier.
Reproduce from the repository root with `python compare_precision.py kda`.

## Files and provenance

The combined runtime is in `../build/torch-cuda/`; this component's unchanged
kernel is `_kda.py`. The files in this directory provide its original
benchmark, baseline, definition, workloads, benchmark and precision results.
Shared dependencies and licenses are at the repository root.

The runtime kernel is byte-for-byte the final accepted submission
`e797d585430d45e1` (SHA-256 `2f635896b70d13ddfd76822d711f75c8c3abb4606f0ccd25dd04b583b95b7ec6`), from a 48-hour KDA experiment
with GPT 6.1 and Claude Opus 5.5. Request sources are pinned to public commit
`48f02e6ba0e98a1841ab6f0a2d99a90b2de1a326`. The release adds a public import wrapper, metadata
and precision comparison tooling; it does not change the kernel algorithm. Apache-2.0;
see `../LICENSE.txt` and `../NOTICE`.

## Available functions

- `run`
- `Pool`
- `flush_inplace`