forked from Karylab-cklius/vllm
Yongye ZhuandClaude Opus 4.8
06324dd3df
Unify the Triton indexer impl onto the same persistent top-k buffer as the MSA impl, and thread the buffer through the AMD model too. - minimax_m3_index_decode gains an out= param (writes out[:, :total_q]); the merge kernel already writes via strides, so a buffer view works. - MiniMaxM3IndexerTritonImpl.forward writes decode ([:, :nd]) and prefill ([:, nd:]) into the shared topk_indices_buffer and returns views into it (no fresh per-step top-k allocations), matching the MSA impl. - amd/model.py: allocate the model-level topk_indices_buffer and thread it model -> decoder layer -> sparse attention -> indexer (mirrors nvidia). AMD keeps its eager break, so this is purely allocation reuse there. test_msa_indexer_impl_matches_triton now gives each impl its own buffer and asserts both decode/prefill outputs are views into it. 42/42 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>