Follow-ups on the cudagraph-capturable MSA indexer:
- Top-k: both decode and prefill now write into the single shared, persistent
topk_indices_buffer (decode at [:, :nd], prefill at [:, nd:]) and return views
into it -- no fresh per-step top-k allocations.
- Build the decode plan + flat page table entirely with torch on-GPU: drop numpy
and CpuGpuBuffer; segment offsets/lengths are computed via torch.cumsum into the
persistent int32 buffers, and the request-major page table is scattered into the
buffer via the on-GPU page indptr (the run bounds reads by indptr, so the full
buffer is passed and no host page count is needed).
- No GPU->CPU sync on the decode path: scalars come from host ints
(num_decode_tokens // num_decodes), and seq_lens.cpu() is confined to the eager
prefill branch. The impl forward (fmha OnlyScore + Triton top-k) was already
sync-free.
test_msa_indexer_impl_matches_triton now also asserts both outputs are views into
the persistent buffer. 42/42 in test_minimax_m3.py pass.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
