forked from Karylab-cklius/vllm
Yongye ZhuandClaude Opus 4.8
d2fbaf73c1
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>