feat: add packed kda decode

[zhouaihui]

📌 Description

This PR adds a packed-QKV KDA decode entry point, kda_packed_decode, for Ling V3-style decode workloads where the conv output is already laid out as:

mixed_qkv[b] = [q_flat | k_flat | v_flat]

The packed path avoids per-call Q/K/V materialization by slicing mixed_qkv into static-stride views and launching the existing CuTe DSL KDA decode kernel with a packed static-stride specialization. This keeps the kernel body shared with the existing non-packed decode path while removing the caller-side split -> view -> contiguous adapter cost.

The implementation includes:

• public kda_packed_decode export under cula.kda
• packed layout validation and H/HV/K/V inference from mixed_qkv and state
• dense and varlen decode support
• vk and kv state layout support
• state_indices[b] < 0 dummy slot support by skipping state reads, state writes, and output writes
• packed-vs-non-packed tests and a CUDA graph padded-batch smoke test
• a microbenchmark comparing packed, realistic non-packed, and pre-split oracle paths

🔍 Related Issues

Closes #101

🚀 Pull Request Checklist

Thank you for contributing to cuLA! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing.

Added:

tests/test_kda_packed_decode.py

Validated on an H200 (HBM3e) environment:

pre-commit run --all-files
# clang-format: passed
# ruff: passed
# ruff-format: passed

pytest tests/test_kda_packed_decode.py -q
# 89 passed in 217.20s (0:03:37)

⚡ Performance

This PR adds benchmarks/bench_kda_packed_decode.py to measure:

• non-packed (oracle): q/k/v pre-split once outside the timed loop
• non-packed (realistic): per-call split + view + contiguous
• cula_packed: packed static-stride path with no q/k/v materialization

The expected benefit is reduced decode overhead from avoiding:

• Q/K/V contiguous materialization
• intermediate tensor allocation pressure
• extra PyTorch dispatch overhead from the caller-side adapter chain
• CUDA graph nodes caused by Q/K/V materialization copies
• global memory traffic for intermediate Q/K/V tensors

Validated on an H200 (HBM3e) environment:

python benchmarks/bench_kda_packed_decode.py \
  --batch-sizes 1 2 4 8 16 32 64 128 \
  --head-pairs 32:32 64:64 8:16 16:32

Summary: packed decode matches the non-packed reference numerically (out_diff=0, state_diff=0) across all measured configurations. For N >= 2, speedup ranges from 1.05x to 1.48x, averaging about 1.30x, with roughly 12.6-19.6 us saved versus realistic non-packed decode. N=1 is roughly neutral/slightly slower because the avoided materialization cost is very small relative to fixed dispatch overhead.

Column notes:

• oracle ms: non-packed decode with q/k/v pre-split once outside the timed loop.
• non-packed ms: realistic non-packed decode with per-call q/k/v materialization.
• packed ms: kda_packed_decode static-stride packed path.
• save vs non us: (non-packed ms - packed ms) * 1000.
• diff: max numerical diff for both output and state; all measured rows are zero.

K=128, V=128, warmup=30, rep=200.

H=32, HV=32

N	qkv_dim	oracle ms	non-packed ms	packed ms	save vs non us	speedup	diff
1	12288	0.0349	0.0455	0.0465	-0.94	0.980x	0
2	12288	0.0349	0.0682	0.0485	19.64	1.405x	0
4	12288	0.0345	0.0678	0.0488	19.02	1.390x	0
8	12288	0.0363	0.0672	0.0500	17.17	1.343x	0
16	12288	0.0266	0.0592	0.0401	19.07	1.476x	0
32	12288	0.0441	0.0587	0.0437	15.01	1.344x	0
64	12288	0.0884	0.1017	0.0882	13.52	1.153x	0
128	12288	0.1579	0.1728	0.1576	15.21	1.097x	0

H=64, HV=64

N	qkv_dim	oracle ms	non-packed ms	packed ms	save vs non us	speedup	diff
1	24576	0.0346	0.0457	0.0470	-1.31	0.972x	0
2	24576	0.0348	0.0670	0.0488	18.26	1.375x	0
4	24576	0.0361	0.0669	0.0492	17.66	1.359x	0
8	24576	0.0265	0.0585	0.0398	18.79	1.473x	0
16	24576	0.0439	0.0583	0.0443	13.98	1.315x	0
32	24576	0.0882	0.1019	0.0880	13.86	1.157x	0
64	24576	0.1580	0.1735	0.1576	15.84	1.100x	0
128	24576	0.2973	0.3134	0.2977	15.67	1.053x	0

H=8, HV=16

N	qkv_dim	oracle ms	non-packed ms	packed ms	save vs non us	speedup	diff
1	4096	0.0349	0.0457	0.0469	-1.22	0.974x	0
2	4096	0.0347	0.0669	0.0482	18.68	1.387x	0
4	4096	0.0350	0.0668	0.0484	18.45	1.381x	0
8	4096	0.0351	0.0669	0.0491	17.78	1.362x	0
16	4096	0.0366	0.0676	0.0511	16.56	1.324x	0
32	4096	0.0264	0.0600	0.0415	18.50	1.446x	0
64	4096	0.0458	0.0595	0.0457	13.82	1.303x	0
128	4096	0.0813	0.0962	0.0823	13.90	1.169x	0

H=16, HV=32

N	qkv_dim	oracle ms	non-packed ms	packed ms	save vs non us	speedup	diff
1	8192	0.0351	0.0469	0.0480	-1.08	0.977x	0
2	8192	0.0347	0.0669	0.0485	18.34	1.378x	0
4	8192	0.0354	0.0670	0.0494	17.52	1.354x	0
8	8192	0.0365	0.0681	0.0504	17.79	1.353x	0
16	8192	0.0268	0.0592	0.0418	17.35	1.415x	0
32	8192	0.0454	0.0592	0.0454	13.77	1.303x	0
64	8192	0.0812	0.0963	0.0826	13.71	1.166x	0
128	8192	0.1672	0.1826	0.1700	12.58	1.074x	0

[gemini-code-assist]

Code Review

This pull request introduces a packed-QKV variant of the KDA decode kernel (kda_packed_decode) along with associated benchmarks and unit tests. This implementation allows feeding Q/K/V as strided views directly from a packed tensor, avoiding the overhead of separate tensor materialization and contiguous copies. Feedback on the changes points out a potential issue in the general path of kda_packed_decode where the normalized a_kernel and b_kernel are not guaranteed to be contiguous, which could lead to runtime issues since the compiled kernel expects contiguous layouts. Ensuring these tensors are contiguous before passing them to the kernel is recommended.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Copilot]

Pull request overview

This PR adds a packed-QKV KDA decode entry point (kda_packed_decode) that consumes a single mixed_qkv tensor and slices it into strided Q/K/V views to reuse the existing CuTe DSL KDA decode kernel without caller-side Q/K/V materialization.

Changes:

• Added kda_packed_decode (fast path + general path) and a packed static-stride kernel compilation/cache path.
• Exported kda_packed_decode under cula.kda.
• Added thorough packed-vs-non-packed correctness tests and a microbenchmark.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.

File	Description
cula/ops/kda_decode.py	Implements kda_packed_decode, adds packed fast path and packed static-stride kernel compilation/caching.
cula/kda/__init__.py	Exposes kda_packed_decode in the public cula.kda API.
tests/test_kda_packed_decode.py	Adds correctness tests for dense/varlen, KV state layout, dummy indices, padded row stride, and CUDA graph smoke coverage.
benchmarks/bench_kda_packed_decode.py	Adds a benchmark comparing packed vs realistic non-packed vs oracle (pre-split) decode overhead.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[icavan]

LGTM, you need to fix the conflicts

[zheyang0825]

LGTM