[Performance]: Compact pseudo-KV/topological sparse attention benchmark evidence and backend gap

[teerthsharma]

Proposal to improve performance

I would like to propose a performance/RFC discussion for a compact pseudo-KV carrier for sparse decode attention in the PyTorch / THOP path.

This follows up on the earlier community PoC in #10305, but changes the shape of the proposal based on the feedback there: instead of only predicting sparse real-token indices, the new direction separates two carriers:

• SparsePrediction: existing real-token sparse indices, still useful when every selected row is an existing KV-cache row.
• CompactPseudoKVCarrier: compact synthesized K/V rows plus witness rows, built during prefill/warmup and consumed during decode.

The motivation is that a topology-compressed prefix cannot be faithfully represented by SparsePrediction alone once residual/pseudo rows are synthesized. A typed compact THOP attention overload avoids the overhead of a generic **kwargs carrier path and keeps the runtime call shape kernel-friendly.

Current local evidence from real distilgpt2 Q/K/V captures at sequence length 1024:

Typed imported THOP attention compact pseudo-KV route
- rows: 6/6 pass
- min speedup vs dense SDPA: 2.728x
- max relative L2: 0.02356
- min cosine similarity: 0.999731
- min token/KV reduction: 0.398
- route: imported tensorrt_llm.bindings.internal.thop attention overload
- caveat: fused=False, production=False

SparsePrediction-style real-token witness baseline comparison
- rows: 6/6 pass
- min compact speedup over real-token sparse baseline: 1.859x
- min sparse-baseline speedup over dense: 1.468x
- min compact speedup over dense: 2.729x
- max relative L2: 0.02356
- min cosine similarity: 0.999731
- min token/KV reduction: 0.398
- caveat: baseline is a SparsePrediction-compatible real-token witness replay, not a production sparse backend benchmark

Prefill / in-flight-batching amortization
- rows: 6/6 pass at decode horizon 4096
- min cached decode speedup: 1.989x
- min amortized speedup at 4096 decode tokens: 1.512x
- max break-even decode length: 1660 tokens
- max prefill topology build cost: 149.708 ms
- caveat: topology build is charged once as prefill/warmup work; production scheduling integration is not done

The current proposal is not to merge the scratch prototype as-is. The issue is to discuss whether this carrier shape is worth pursuing in TensorRT-LLM:

1. Add a production compact pseudo-KV carrier alongside SparsePrediction.
2. Build compact topology rows during prefill/warmup or a conservative stability phase.
3. Consume the cached compact rows during decode with a typed/fused THOP or FMHA/CUTLASS backend path.
4. Keep fallback to existing dense/sparse paths when stability, accuracy, or reduction gates fail.

Report of performance regression

This is not a regression report.

Misc discussion on performance

Why this is different from #10305:

• POC/aether sparse attention #10305 proposed AETHER sparse attention as a PoC plugin / sparse kernel direction.
• The old discussion correctly focused on the need for real accuracy and performance data, not only synthetic or mock evidence.
• The new experiments use real captured Q/K/V tensors from distilgpt2 and explicitly track speed, relative L2, cosine similarity, and token/KV reduction.
• The current evidence also shows one negative result: a generic thop.attention(**kwargs) compact ABI is too slow (0/6 pass, min speedup 0.756x). A typed compact overload is required for speed.

Open caveats before a PR:

• The current strongest route is imported/source-built scratch THOP binding evidence, not the final production TensorRT-LLM bindings target.
• The fused production backend is not complete yet (fused=False, production=False in the local audit).
• The result should be treated as a performance/RFC discussion and benchmark target, not as a ready PR claim.
• Model-level quality evaluation is still needed beyond attention replay metrics.

Local artifact names from my scratch workspace:

• compact_bindings_shim_typed_attention_abi_distilgpt2_seq1024.md
• sparse_prediction_baseline_vs_typed_compact_distilgpt2_seq1024.md
• prefill_amortized_pseudokv_distilgpt2_seq1024_long_decode.md
• pr_readiness_audit.md

my current environment (if you think it is necessary)

Benchmark container: nvcr.io/nvidia/tensorrt-llm/devel:1.2.0rc6
PyTorch in container: 2.9.0a0+145a3a7
Benchmark model/captures: real distilgpt2 Q/K/V captures, seq=1024, selected layers/heads
Local status: scratch implementation and artifacts under .donotcommit; not yet a production/fused TensorRT-LLM backend

Before submitting a new issue...

• I searched for related work and am linking the previous PoC PR: POC/aether sparse attention #10305.
• I checked the TensorRT-LLM sparse attention carrier direction and separated real-token SparsePrediction from compact pseudo-KV rows.
• I am opening this as a performance/RFC discussion, not claiming the production backend is already complete.