[ROCm][Perf] Tune MiniMax-M3 decode top-k for short contexts - #55235
Conversation
Use a single 128-entry CTA for the gfx950 BF16 decode selector when the graph capacity is at most 128 sparse blocks. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: fai <fangzhouai@gmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. 📝 SummarySummary by CodeRabbit
WalkthroughThe AMD decode top-k policy now selects block size, warp count, and stage count. The fused kernel receives these values. Tests cover short-context boundaries, parameterized sequence sizes, ordering, and graph replay. ChangesAMD decode top-k dispatch
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to This change specializes the gfx950 short-context BF16 decode top-k launch configuration while retaining existing policies elsewhere; no concrete merge-blocking risk is identified. Sequence Diagram(s)sequenceDiagram
participant Decode as minimax_m3_index_decode
participant Policy as _decode_topk_launch_policy
participant Kernel as _decode_topk_fused_kernel
Decode->>Policy: Request top-k launch configuration
Policy-->>Decode: Return block size, warps, stages, chunk count, and flags
Decode->>Kernel: Launch with selected configuration
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
AndreasKaratzas
added
the
ready
|
/ci run |
|
✅ @Fangzhou-Ai, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #87142 for commit |
|
/ci run |
|
✅ Triggered Buildkite CI #87276 for commit |
|
/ci run |
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
✅ Triggered Buildkite CI #87371 for commit |
|
LGTM |
|
/ci run |
|
✅ Triggered Buildkite CI #88455 for commit |
|
hi @zhou9402 can we merge it now? |
3 participants
Summary
topk == 16and the graph capacity is at most 128 sparse blocks.Kernel performance
Measured on gfx950 at the production 8k1k graph capacity of 74 blocks, with query length 4, one local index head, top-k 16, and the recorded c1/c2/c4/c6/c8/c12/c16/c24/c32/c64 input rows:
Each timing sample used a 64-node HIP graph. Every case had five seconds of same-workload graph warmup, 16 full-graph pre-replays before each sample, and 31 balanced HIP-event samples.
A separate exact boundary sweep over capacities 65, 74, and 128; query lengths 1, 4, and 8; and local index-head counts 1, 2, and 4 measured 32.78-35.32% lower selector time and 26.19-29.21% lower selector-plus-table time.
Precision
This changes launch geometry only. The BF16 index cache, FP32 score output, score arithmetic, and deterministic score-descending/index-ascending ordering are unchanged. Tests matched an independent CPU total-order reference exactly, including NaN, infinities, signed zero, and tied boundary scores. Top-k IDs, sparse tables, context lengths, completion-counter reset, eager execution, and graph replay all matched exactly.
This PR preserves the current vLLM BF16 precision contract; it does not switch the indexer to the separate AITER FP8 path.
Validation
No model evaluation was rerun because the change is output-exact and does not alter model arithmetic or precision.
Duplicate-work check
This is a narrow launch-policy follow-up to merged #54682. The required open-PR searches found no PR tuning this fused BF16 decode selector to B128/C1/W2/S1. #53448 targets prefill plus an optional FP8 index cache, #52664 integrates an external AITER indexer path, and #52059 targets prefill scoring; none makes this decode launch-policy change.
OpenAI Codex assisted with analysis, implementation, benchmarking, validation, and drafting this description. The human submitter must review every changed line and be able to reproduce and defend the change.