perf(deepseek-v4): capture DSA projections and cache writes into the prefill graph

[dongjiyingdjy]

Summary

Narrow the DeepSeek-V4 attention break to the request-dependent DSA core and move the surrounding token-shaped work into the captured prefill-graph segments, removing kernel launches from the eager critical path.

• E1 — phase-1 Q/KV/QR projections and the compressor-score GEMMs run in the pre-break segment, overlapped on an aux stream.
• E2 — the attention output projection joins the post-break segment.
• E3 — the CSA indexer query packing (wq_b / weights_proj / mxfp4) runs pre-break.
• E4a — per-forward memoization of compressed/state slot mappings, shared across the CSA indexer, the 30 CSA layers, and the C4 compressor.
• E4b/E4c — the CSA indexer and main C4 compressor state-save + cache-insert writes are captured in the pre-break segment, fed by persistent staging buffers refreshed by prepare/finish_prefill_graph_replay model hooks, with a -1 fail-open contract and a single ctx.dsa_pregraph_writes arm flag so the eager break skips exactly what the captured segment wrote (exactly-once across capture / replay / eager / mixed / decode / MTP-draft).

Structure is unchanged: 61 eager attention breaks / 62 graph segments / 123 callables.

Verification

Verified on top of current origin/main (rebuilt matching tokenspeed-kernel objs + tokenspeed-scheduler 0.1.2 exposing flat_block_tables_arrays):

• GSM8K (V4-Flash, TP4, 5-shot first-50, greedy): strict 0.98 / flexible 0.98 (gate ≥ 0.92), 0 fatal server lines.
• Prefill breakable graph engaged on all TP ranks (87 segments, buckets 16–8192), no eager-prefill fallback.
• CPU unit tests: 65 passed / 5 skipped — pin the control-flow invariants (arm/disarm matrix, skip gates, hook seam, staging bind). Numerical byte-equality is covered by the GPU gates.

🤖 Generated with Claude Code