[Attention] Overlap sparse MLA indexer on a native CUDA side stream

[LucasWilkinson]

Summary

Runs the sparse-MLA (DSA) indexer on a shared high-priority CUDA stream so it
overlaps the main stream's path to attention. The indexer is independent of the
q projection, kv layernorm and rope that precede attention; only
topk_indices_buffer couples it to attention, so a single join immediately
before the attention call is sufficient.

Rebased on current main (PyTorch 2.13) and uses PyTorch's native compiled
stream operations — no side-stream begin/end custom ops, and no per-graph
stream mappings.

Mechanics

• Indexer owns the stream as a class attribute and hands it to
  MultiHeadLatentAttentionWrapper via MLAModules. Class-level because the
  fork, the join and the compiled artifacts must reference the same object.
• Indexer.forward forks (wait_stream + with stream:); the wrapper joins
  with current_stream().wait_stream(side) just before self.mla_attn(...).
  Both are plain stream ops, so torch.compile traces them into native stream
  nodes.
• The indexer is issued right after the q_a/kv_a projections rather than after
  rope, so there is main-stream work left to overlap with.
• Fork and join land in different compiled pieces (the indexer is a
  splitting op), so a piece carrying side-stream work cannot be captured as an
  independent piecewise cudagraph — capture would end with unjoined work. Those
  pieces skip that capture and run inside the full graph instead.
• vLLM invokes compiled pieces without dynamo's registry-repopulating bytecode,
  so stream external-object indices are resolved by a small resolver: dynamo's
  reserved index means "current stream", any other index means the side stream.
  It is installed only once a layer registers a side stream.
• Compile-time Triton autotuning is disabled only for graphs containing a
  non-default stream; PyTorch 2.13's compile-time autotune wrapper does not
  initialize the user-stream handle.

Results

GLM-5.2-NVFP4, TP=4, B200, torch.compile + FULL_AND_PIECEWISE CUDA graphs
(not eager, not piecewise-only).

Kernel overlap

Measured per steady-state decode iteration as
union(indexer) + union(other) − union(indexer ∪ other), at ~9.6k-token
context, batch 4. Stream IDs are unreliable for cudagraph-replayed kernels, so
this is time overlap, not stream attribution.

config	main	this PR
TP=4	3.5%	71.0%
TP=4 + DCP=4	1.8%	30.2%

TP=4, no DCP — the main-stream row gaps open exactly where the indexer runs
on main, and runs straight through it with the side stream:

TP=4 + DCP=4 — lower, because the DCP-only tail of the indexer chain
(pack_dcp_topk_candidates → top-k all-gather → stable merge) overlaps poorly:
the top-k all-gather is 0.330 ms at 2.7%.

Per kernel class, TP=4 without DCP:

kernel class	main	this PR
k quant + cache	4.9%	98.0%
top-k select/merge	0.7%	74.3%
q rope + quant	5.2%	64.1%
logits (MQA)	6.9%	45.7%

End-to-end

vllm bench latency, 10 iterations at 4k / 5 at 32k, warmup excluded.

config	main	this PR	delta
32k in / 128 out / batch 8, TP=4	11.6437 s	11.5420 s	+0.87%
32k in / 128 out / batch 8, TP=4 + DCP=4	20.1777 s	19.9829 s	+0.97%
4k in / 256 out / batch 16, TP=4	5.7858 s	5.9437 s	−2.73%
4k in / 256 out / batch 16, TP=4 + DCP=4	9.7436 s	9.9133 s	−1.74%

This regresses short-context decode. Reporting it because it is real and
reproducible — the distributions do not overlap (main's slowest 4k run is faster
than this PR's fastest), sd ≤ 0.03 s.

Two contributing mechanisms are measured; a third is suspected:

1. Concurrency widens the indexer chain. Same kernels, run alongside the
   main stream, take longer: 4.340 → 5.608 ms (+29%) without DCP, 0.975 → 1.014
   ms with. On a saturated device concurrency partitions SMs rather than adding
   throughput. You get 71% of a bigger number.
2. The prize is small at short context. At 4k/batch-16 the indexer is ~4% of
   decode, so perfect hiding saves ≤4%; a 2–3% slowdown of the critical path
   erases it. At 32k the indexer's share grows — it scans the whole KV while
   sparse attention stays fixed at topk_tokens.
3. Suspected, not yet isolated: pieces carrying side-stream work skip
   piecewise cudagraph capture. Decode uses full graphs so it is unaffected, but
   prefill runs piecewise and now runs uncaptured. That cost is roughly fixed
   per prefill while the decode benefit scales with context, which would explain
   the sign flipping with the prefill:decode ratio rather than with DCP.

Accuracy

GSM8K, 300 questions, 5-shot, 0% invalid:

config	accuracy
main, TP=4	0.9367
this PR, TP=4	0.9500
this PR, TP=4 + DCP=4	0.9433

A 4-question spread on 300 is inside one standard error; no regression.

Not a duplicate

This updates the existing side-stream PR rather than opening another.
Related but distinct: #48196 is the DCP output-merge optimization; #45964
(merged) is DCP query replication.

Prior art worth naming: an earlier unlanded branch of mine
(5012b439) did the same overlap using dcp_indexer_stream_begin/restore
custom ops that call torch.cuda.set_stream at runtime, scoped to DCP > 1, and
placed the join after the DCP query all-gather. This PR replaces the custom ops
with native traced streams. I re-tested that later join placement here and it
made no measurable difference (30.7% → 30.2% overlap, e2e slightly worse), so it
is not carried over.

Testing

pytest tests/compile/test_side_stream.py           # 4 passed
pytest tests/compile/test_aot_compile.py           # 28 passed, 1 skipped
pre-commit run --files <changed>                   # passed, incl. manual-stage mypy

test_aot_compile.py includes a new AOT save/reload round-trip of a graph that
uses a side stream, covering the index resolver.

Caveats

• Profiles are at ~9.6k tokens / batch 4; the benches at 32k and 4k. Different
  operating points — the overlap percentages do not transfer directly to the
  end-to-end numbers.
• Single model (GLM-5.2-NVFP4) and single machine (4×B200).
• The short-context regression should be understood before merge; if mechanism 3
  above is confirmed, it is fixable rather than inherent.

AI assistance

AI assistance (Claude, and Codex for earlier revisions) was used to implement,
profile and test this change. The submitter has reviewed every changed line and
is responsible for the results reported here.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LucasWilkinson.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LucasWilkinson.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @LucasWilkinson.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork