Optimize native AR decoding by attending only to the active KV-cache prefix (1.97x compiled speedup on GV100)

[quantumxiaol]

Self Checks

• I have thoroughly reviewed the project documentation (installation, training, inference) but couldn't find any relevant information that meets my needs. English 中文 日本語 Portuguese (Brazil)
• I have searched for existing issues search for existing issues, including closed ones.
• I confirm that I am using English to submit this report (我已阅读并同意 Language Policy).
• [FOR CHINESE USERS] 请务必使用英文提交 Issue，否则会被关闭。谢谢！:）
• Please do not modify this template :) and fill in all the required fields.

1. Is this request related to a challenge you're experiencing? Tell us your story.

I am using the native PyTorch text2semantic/inference.py path on both Apple MPS and NVIDIA CUDA. While profiling S2-Pro autoregressive decoding, I found that each decode step attends over the physical capacity of the KV cache rather than only the prefix populated so far.

At upstream commit e5e2926:

• When input_pos is present, forward_generate() builds an attention mask whose K dimension is self.max_seq_len.
• KVCache.update() returns the full preallocated K/V tensors, which are then expanded with repeat_interleave before SDPA.

Masking the unused tail preserves correctness, but it does not prevent the full-capacity K/V tensors from being expanded and passed to attention.

For S2-Pro, the model configuration has max_seq_len=32768, n_head=32, n_local_heads=8, head_dim=128, and 36 slow-transformer layers. My representative prompt contained 234 tokens and generation ended at roughly 400 context tokens, so only about 0.7%-1.3% of the physical cache was active.

At FP16, after expanding 8 local KV heads to 32 query heads, full-capacity K+V represents about 512 MiB of logical tensor data per layer, or about 18 GiB across 36 layers per semantic frame. This is an estimate of logical tensor work, not simultaneous memory residency or measured GPU bandwidth, but it explains why the work can scale with cache capacity instead of the current context length.

This is especially visible in native eager inference, but torch.compile does not eliminate all of the redundant work either.

2. What is your suggested solution?

Keep the KV cache physically preallocated at max_seq_len, but pass an explicit active length through the AR generation path and slice K/V before head expansion and SDPA:

1. Use kv_len=T for prompt prefill.
2. Use kv_len=T+i+1 for decode step i.
3. Build the attention mask with :kv_len, not :max_seq_len.
4. After updating the physical KV cache, slice k[:, :, :kv_len] and v[:, :, :kv_len].
5. Only then apply repeat_interleave and call SDPA.

The physical cache allocation and in-place update behavior remain unchanged. The optimization only reduces the logical K dimension consumed by attention.

I prepared a minimal implementation based directly on the upstream commit:

• Exact upstream baseline: benchmark/upstream-native-cuda, commit e5e2926
• Active-KV branch: perf/active-kv-native, commit c4146e7
• Minimal diff: e5e2926...c4146e7
• Correctness tests: tests/test_active_kv_attention.py

The implementation passes the Python position as an integer from the generation loop, avoiding .item() or another device synchronization in the production path.

CUDA benchmark

Environment:

• GPU: NVIDIA Quadro GV100, 32 GB
• OS: Ubuntu x86_64
• NVIDIA driver: 525.105.17 (nvidia-smi reports CUDA compatibility 12.0)
• Python: 3.12
• PyTorch: 2.8.0+cu128
• Triton: 3.4.0
• Model: S2-Pro, FP16, native PyTorch inference
• Model max_seq_len: 32768
• Encoded prompt shape: [11, 234]
• Seed: 42
• temperature=0.8, top_p=0.8, top_k=30
• max_new_tokens=1024
• Four samples per mode

I used the same model, reference tokens, target text, sampling parameters, and seed on both branches. The first sample was treated as warm-up; for compiled runs it also contains compilation. The table reports the median of samples 1-3. Generated lengths differ slightly because small FP16/kernel differences can change a stochastic sampling trajectory, so throughput is normalized as generated tokens per second.

Native inference mode	Full-capacity KV	Active-prefix KV	Speedup	Per-token latency	Peak CUDA memory reserved
Eager	2.73 tok/s	12.57 tok/s	4.60x	366.3 -> 79.6 ms	17.33 -> 15.16 GB
torch.compile	16.65 tok/s	32.74 tok/s	1.97x	60.1 -> 30.5 ms	15.72 -> 15.16 GB

Individual measurements:

Eager, full KV:    2.71, 2.73, 2.73, 2.73 tok/s
Eager, active KV: 12.19, 12.57, 12.56, 12.58 tok/s

Compile, full KV:   4.01 [compile], 16.65, 16.64, 16.68 tok/s
Compile, active KV: 3.32 [compile], 32.69, 32.76, 32.74 tok/s

The existing log's Bandwidth achieved value is intentionally not used here because it is derived from model size multiplied by tokens/s; it is not a hardware DRAM-bandwidth measurement.

Cold compile times are also not compared because the runs shared persistent Inductor caches. With TORCH_LOGS=recompiles, the active-KV implementation showed one recompile when kv_len first changed from 235. It did not recompile once per token, and subsequent samples remained stable at approximately 32.7 tok/s. This one-time dynamic-shape compile behavior may still be worth improving for one-shot CLI use, while it is amortized in a long-running service.

3. Additional context or comments

PR #1306 identifies the related full-capacity KV-cache allocation and proposes lowering max_seq_len as a memory workaround. This proposal addresses a different part of the same path: it preserves the full cache capacity while preventing the unused tail from being expanded and passed to attention on every decode step.

Issue #1168 contains a possibly related discussion about slow native inference and torch.compile. Its hardware, prompt, platform, and cold/warm compile conditions are not directly comparable with this benchmark, so I am not using its reported throughput as evidence here. This report isolates the full-capacity KV path with an exact upstream-commit A/B comparison.

I first found this while profiling the MPS backend. On an Apple M4 Pro with 48 GB unified memory, FP16, and max_seq_len=4096, the same active-prefix change improved the measured semantic generation rate from 4.08 to 8.31 frames/s (2.04x) for a 234-token prompt. This MPS result came from a service-level diagnostic path, so its absolute metric should not be compared directly with the CUDA CLI's tok/s. I include it only as secondary evidence that the optimization is backend-independent.

For output validation:

• All four representative CUDA outputs (full/active x eager/compiled) decoded successfully.
• I manually listened to all four WAV files and observed no audible regression.
• The files are not bit-identical, which is expected for stochastic FP16 inference when changing reduction shapes or compiled kernels.
• The unit test provides the stricter deterministic float32 equivalence check.

Current limitations of this benchmark:

• One CUDA architecture was tested (Volta GV100); Ampere/Ada/Hopper results would still be useful.
• The benchmark uses one target/reference pair and four samples per mode.
• No Nsight or nvidia-smi dmon trace was collected during generation, so I am not claiming a measured GPU-occupancy or hardware-bandwidth result.
• Cold compilation was excluded from the speed comparison.

The CUDA comparison is nevertheless narrowly controlled: perf/active-kv-native is based on the exact upstream baseline commit and changes only two inference source files plus focused tests.

I am interested in contributing this optimization and can prepare a clean PR if the approach is acceptable.

4. Can you help us with this feature?

• I am interested in contributing to this feature.

[quantumxiaol]

A draft implementation is now available in #1312.

The PR contains the minimal production changes for limiting cached attention
to the active KV prefix. Initial CUDA and MPS benchmark results, correctness
checks, and audio validation are documented in the PR description.

Additional manual and cross-device testing is planned. Independent testing
and feedback would be very welcome.