perf(text2semantic): pass only active KV cache prefix to SDPA

[quantumxiaol]

Is this PR adding new feature or fix a BUG?

Performance optimization for the native PyTorch autoregressive inference path.

Problem

During cached generation, the physical KV cache is preallocated at
max_seq_len. At every decode step, the current implementation expands the
full-capacity K/V tensors with repeat_interleave and passes the full K/V
sequence to SDPA. For S2-Pro, this means expanding and passing a key sequence
length of 32768 to SDPA even when only a few hundred positions have been
populated.

This describes the tensors supplied to SDPA; it does not assume how a
particular SDPA backend handles the masked, unpopulated tail internally.

Change

This change keeps the physical KV cache at full capacity but limits attention
to the active prefix:

• prompt prefill uses kv_len=T;
• decode step i uses kv_len=T+i+1;
• the causal mask is sliced to the active K length;
• cached K/V are sliced before head expansion and SDPA.

The regular full-sequence forward() path used for training does not use the
inference KV cache and is unchanged. The new slicing is gated by
self.kv_cache is not None and is exercised by cached forward_generate()
inference.

Benchmark results

S2-Pro, FP16, native PyTorch CLI inference, prompt length 234,
max_seq_len=32768. Throughput is the median of three warm samples:

Mode	Full-capacity KV	Active-prefix KV	Speedup	Peak CUDA memory reserved (GB)
Eager	2.73 tok/s	12.57 tok/s	4.60x	17.33 -> 15.16
torch.compile	16.65 tok/s	32.74 tok/s	1.97x	15.72 -> 15.16

CUDA environment: Quadro GV100 32 GB, PyTorch 2.8.0+cu128, Triton 3.4.0.

Generated lengths differed slightly because FP16/kernel differences can change
stochastic sampling trajectories, so throughput is normalized by the number of
generated tokens.

The memory values are decimal GB, matching the existing log calculation
torch.cuda.max_memory_reserved() / 1e9.

The same change was also validated on Apple MPS (M4 Pro 48 GB, FP16,
max_seq_len=4096), where semantic generation improved from 4.08 to
8.31 frames/s (2.04x). The MPS service metric and CUDA CLI tok/s are not
directly comparable; only their within-device A/B ratios are reported.

Cold compile time is excluded from the throughput comparison. The existing
Bandwidth achieved log value is also excluded because it is derived from
model size and tok/s rather than measured hardware DRAM bandwidth.

Correctness and compatibility validation

• Upstream pre-commit hooks pass for both modified files.
• A CPU shape check confirmed that SDPA receives only the active K/V prefix
  while the physical KV cache retains full capacity.
• A CPU float32 comparison confirmed that active-prefix attention matches
  full-cache attention with a masked tail (rtol=1e-5, atol=1e-6).
• CUDA eager and torch.compile generation completed successfully.
• With TORCH_LOGS=recompiles, the compiled active-KV path recompiled once
  when kv_len first changed, not once per token; warm throughput remained
  stable at approximately 32.7 tok/s.
• Four representative CUDA outputs (full/active x eager/compiled) were decoded
  and manually listened to; no audible regression was observed.
• MPS end-to-end generation and manual listening completed successfully.

The benchmark results above were collected from
c4146e7.
The current draft contains only the two production source changes; the
standalone correctness tests used for validation remain available in that
commit.

Is this pull request related to any issue? If yes, please link the issue.

Closes #1310

[Whale-Dolphin]

The active-prefix optimization is effective. I verified the current head on an NVIDIA H200 with S2-Pro FP16 through the official text2semantic CLI: eager warm throughput improved from 9.22 to 19.73 tok/s (2.14x), and torch.compile warm throughput improved from 68.60 to 113.33 tok/s (1.65x).

Before merging, please retain regression tests in the repository for the behavior this PR changes:

1. Compare full-cache attention with active-prefix attention on a small deterministic fixture, with an explicit numerical tolerance and argmax assertion rather than requiring bitwise equality.
2. Cover kv_len boundary validation and the prefill/decode progression, including the first decode position and the maximum valid prefix.
3. Add an end-to-end autoregressive decode regression that exercises the real kv_len call chain and checks output shape, finite logits, and a defined token or distribution-level quality contract.

This is important because my FP16 comparison was not bitwise identical: the first five decode positions kept the same argmax, but the maximum absolute logit difference was 0.078125, and same-seed full generations diverged in 170 of 630 acoustic code elements after autoregressive propagation. This does not by itself indicate a quality regression, but the acceptable numerical and generation behavior needs to be encoded in tests so future SDPA or cache changes cannot silently break it.

[quantumxiaol]

Thanks for the independent H200 validation and the regression-test
requirements.

I added the requested regression coverage in commit
3263548,
under
tests/test_active_kv_attention.py.
Five deterministic CPU tests cover:

1. Active-prefix vs. full-cache attention

   • Compares float32 outputs with rtol=1e-5, atol=1e-6, and a separate
     argmax assertion, without requiring bitwise equality.
   • Verifies that SDPA receives only the active K/V prefix while the physical
     cache retains full capacity. The observed maximum absolute difference was
     5.96e-08.

2. kv_len validation and progression

   • Covers valid boundaries 1 and 8, invalid boundaries 0 and 9,
     prompt prefill, the first and subsequent decode positions, and the maximum
     valid prefix.
   • The observed (K, V, mask) lengths were
     [(3, 3, 3), (4, 4, 4), (8, 8, 8)]; the end-to-end kv_len progression
     was [3, 4, 5].

3. End-to-end autoregressive regression

   • Exercises the real
     generate -> decode_n_tokens -> decode_one_token_ar -> forward_generate
     path.
   • Checks output shape, finite slow/fast logits, probability and argmax
     agreement, and deterministic greedy-token equality.

The end-to-end fixture uses top_k=1, so exact token equality is its explicit
greedy-token contract; it does not assume stochastic FP16 CUDA generations are
bitwise identical.

Local result:

Ran 5 tests
OK

The pre-commit.ci check also passes.