Conversation
d2b3032 to
82e6db6
Compare
… top-k
The indexer scores every 128-token block against the current query and keeps the
top 16.
Score kernel: one workgroup owned the whole causal scan for a query block, so the
launch waited on its longest program. Split it over the KV axis in fixed
block-aligned chunks, capped by the deepest possible scan, and widen the query
tile from 64 to 128 rows (gfx950; other archs keep 64 until swept).
Top-k kernel: sorted one query token per program, so its fixed compare-exchange
pass count bought a single row. Carry several rows per program instead. The tile
is chosen host-side from the launch shape and clamped by TOPK_TILE_AREA, since
the tile and the autotuned BLOCK_SIZE_K are both live in registers and cannot be
picked independently. Autotune configs are pruned up front, so a topk with no
legal BLOCK_SIZE_K raises an actionable error instead of a Triton compile assert.
Adds an opt-in fp8 (e4m3) indexer side cache behind
--attention-config '{"indexer_kv_dtype": "fp8"}', matching NVIDIA's M3 recipe,
gated to gfx950 and stored unscaled with saturation at the e4m3 max like the
fused CUDA writer. Default stays bf16.
Kernel time on MI350X, bf16 / fp8 index cache: prefill score 1.15-1.77x /
1.06-2.10x, prefill top-k 0.94-3.41x at topk=16 (below 1.0x only at 1k context
batch 1), decode flat by construction / up to 1.08x.
Selected blocks match the base tree on every bf16 shape except two rows in ~10^6,
where two blocks hold equal fp32 scores and the networks break the tie
differently; both are valid top-k sets.
Adds tests for the ROCm indexer kernels, which had none: prefill and decode top-k
against a dense fp32 reference in bf16 and fp8, tail rows for query lengths that
do not divide the tile, both query tile widths, the full servable topk range plus
the error outside it, and e4m3 saturation on cache insert.
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
82e6db6 to
88265b8
Compare
akii96
changed the title
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| @@ -496,10 +500,10 @@ def select_indexer_impl_cls( | |||
| indexer_kv_dtype, | |||
| ) | |||
| return MiniMaxM3IndexerMSAImpl | |||
| if indexer_kv_dtype != "bf16": | |||
| if indexer_kv_dtype not in SUPPORTED_INDEXER_KV_DTYPES: | |||
There was a problem hiding this comment.
why not append fp8 to the SUPPORTED_INDEXER_KV_DTYPES if fmha_sm100 support it?
There was a problem hiding this comment.
fp8 already works on that path. use_msa accepts bf16 and fp8, and returns the MSA impl before this check is reached, so adding it here would not change SM100 behaviour.
The tuple only covers the Triton fallback, which is bf16 everywhere plus fp8 on gfx950 after this PR. The old name sounded like a global list of what is supported, so I renamed it to TRITON_INDEXER_KV_DTYPES 😄
| @@ -306,6 +302,8 @@ def minimax_m3_insert_index_cache( | |||
| raise ValueError("MiniMax-M3 index cache requires contiguous head dimension") | |||
|
|
|||
| head_dim = index_k.shape[1] | |||
| fp8_out = index_cache.element_size() == 1 | |||
| fp8_max = torch.finfo(index_cache.dtype).max if fp8_out else 0.0 | |||
There was a problem hiding this comment.
Use this from vllm.model_executor.layers.quantization.utils.quant_utils import ( get_fp8_min_max, ) instead of torch.finfo(index_cache.dtype).max because it is not suitable for gfx942
There was a problem hiding this comment.
Good point on gfx942, though I don't think get_fp8_min_max() fits here. It picks its range from the platform, but this cache dtype is pinned to e4m3fn so the Triton and fused CUDA writers agree on 448. On gfx942 the helper returns 224 (deliberately narrower than e4m3fnuz's own 240) while the buffer is still 448, so values in between would get clamped when they shouldn't. It matches on gfx950 so nothing is broken today, but it would hurt whoever enables gfx942.
So I've taken the bounds from index_cache.dtype and passed the min through instead of assuming it mirrors the max. Happy to use the helper if you'd also want the cache dtype switched to current_platform.fp8_dtype(), but that needs the fused writer to accept fnuz too.
Clamp the index-cache write to both bounds of the buffer's own dtype rather than assuming symmetry around its max. The bounds deliberately come from index_cache.dtype and not from current_platform: the cache dtype is pinned to e4m3fn to match the fused CUDA writer, so a platform-derived range would clamp to fnuz limits against an e4m3fn buffer. Rename SUPPORTED_INDEXER_KV_DTYPES to TRITON_INDEXER_KV_DTYPES. The tuple only describes what the Triton fallback accepts, since the fmha_sm100 path returns before that check, and the old name read like a global capability list. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
ChuanLi1101
added
the
verified
| # Query rows per program in the prefill scorer. 128 is swept on gfx950; other | ||
| # archs keep the original 64 until they are (cf. _SPARSE_ATTN_SUB_K in | ||
| # sparse_attn.py). Sets grid dim 0, so it cannot be an autotune key. | ||
| PREFILL_BLOCK_SIZE_Q = 128 if on_gfx950() else 64 |
There was a problem hiding this comment.
@akii96 I ran some perf test on the prefill score kernel using an MI300X and it looks like Q=64 causes a perf regression, depending on context length. 256 wins for some context sizes, but 128 is within 1-2% on all but the largest shape, where Q=256 is 5% better.
What do you think about expanding this check to also set Q=128 for MI300X?
Please double check me 🙏
PREFILL_BLOCK_SIZE_Q sweep — prefill score kernel, MI300X (gfx942), do_bench median of 5 reps, ms:
| ctx / batch | Q=64 (current on gfx942) | Q=128 (gfx950 value) | Q=256 | best → vs Q=64 |
|---|---|---|---|---|
| 2048 / 8 | 0.045 | 0.040 | 0.041 | Q=128 → 1.13× |
| 8192 / 1 | 0.052 | 0.046 | 0.055 | Q=128 → 1.13× |
| 16384 / 1 | 0.139 | 0.130 | 0.129 | Q=256 → 1.07× |
| 16384 / 4 | 0.464 | 0.438 | 0.442 | Q=128 → 1.06× |
| 32768 / 1 | 0.459 | 0.389 | 0.390 | Q=128 → 1.18× |
| 32768 / 4 | 1.782 | 1.461 | 1.439 | Q=256 → 1.24× |
| 65536 / 1 | 1.738 | 1.438 | 1.359 | Q=256 → 1.28× |
There was a problem hiding this comment.
Thanks so much again for catching this! Q=64 was indeed the wrong default for gfx942.
Fixed in f89676934, using the Q=128 you suggested.
my gfx942 sweep (MI325X, best over chunk counts, ms):
| shape | Q=64 | Q=128 | Q=256 |
|---|---|---|---|
| 1k / 1 | 0.0110 | 0.0121 | 0.0147 |
| 2k / 8 | 0.0323 | 0.0276 | 0.0287 |
| 8k / 1 | 0.0438 | 0.0400 | 0.0420 |
| 8k / 8 | 0.2281 | 0.1945 | 0.2232 |
| 16k / 4 | 0.4111 | 0.3436 | 0.3532 |
| 32k / 1 | 0.4045 | 0.3221 | 0.3252 |
| 32k / 4 | 1.5241 | 1.2425 | 1.2283 |
| 64k / 1 | 1.4884 | 1.2144 | 1.1568 |
| 8k / 1, 56k prefix | 0.3491 | 0.2723 | 0.2731 |
| 8k / 8, 56k prefix | 2.7362 | 2.1484 | 2.0819 |
| 8k / 1, 120k prefix | 0.7086 | 0.5294 | 0.5600 |
I got to the same conclusion as the you reached, 128 is best or within 2% on 8 of 11. Q=256 is 3–5% ahead on the two longest shapes, and Q=64 is 10% ahead at 1k/1. I kept the single value rather than branching on shape.
The same commit also retunes the split-K heuristic, which was swept on gfx950 and had left 2k ctx / batch 8 at 0.77× of base on gfx942. Prefill score on MI325X now spans 1.20–2.89× (mean 1.80×), against 0.77–2.15× (mean 1.47×) for the PR as submitted. End to end at TP=4: +2.4% throughput at 128k input, +1.3% at 64k, flat at 8k.
There was a problem hiding this comment.
Really happy to hear that, great work!!
|
I'm extremely excited to see this push the benchmark forward on InferenceX. Once this lands, I'll make sure that the relevant benchmark jobs pick up Also, wanted to make sure the ROCm folks are tracking this related PR: #47665 - enables the same index-K cache for all non-SM100, which would include gfx942 |
| """ | ||
| topk = 16 | ||
| t = _build([q_len], [0], 1, torch.bfloat16, seed=3) | ||
| assert _topk_query_tile(q_len, 1, 1, topk) == TOPK_QUERY_TILE |
There was a problem hiding this comment.
Last comment from me here - this assertion fails on MI300X (_min_topk_programs() = 304 on MI300X vs 256 on gfx950)
> assert _topk_query_tile(q_len, 1, 1, topk) == TOPK_QUERY_TILE
E assert 8 == 32
E + where 8 = _topk_query_tile(8225, 1, 1, 16)
Given that this test is run in the AMD pipeline on a MI300X box, this will cause a non-blocking test failure (all tests in that pipeline are optional), but blocking once it's moved to test_areas.
There was a problem hiding this comment.
Also fixed now in f89676934: the test now derives its query length from _min_topk_programs() * TOPK_QUERY_TILE instead of hard-coding one. That's 8192 on a 256-CU part, so gfx950 coverage is unchanged, and 9728 on 304 CUs where it now passes.
Good point on the test_areas move! Deriving the threshold rather than bumping the constant means it won't start blocking on whatever CU count the pipeline lands on next.
|
hey @wjabbour Thanks for catching both issues! I reproduced the gfx942 performance regression that you mentioned and I am working on restoring the performance and fixing the architecture-dependent test. Wasn't aware of #47665. I’ll review it and determine the best path forward to avoid duplicating the index-cache work. Atleast I believe the the prefill scorer and top-k optimizations in this PR are unique at first glance |
The split-K scorer already ran on gfx942 but its constants were swept on gfx950, and one shape lost: at 2k context batch 8 it reached only 0.92x of the base tree. The split was capped by max_block, the deepest scan in the batch, while the mean scan across the query blocks is about half that. The heuristic therefore cut a scan only 8 blocks deep into 16 chunks and launched 4096 programs to issue 2176 dots: one dot per program, 1920 programs that could only early-return, and 8.5x the query-tile reads of the unsplit kernel. Widen the query tile to 128 rows across MI3xx, which is best or within 2% of best on 9 of the 11 shapes swept. Bound the split by a floor of 8 KV blocks per chunk, so a chunk always carries enough dots to pay for re-reading its query tile, and let an occupancy floor override that floor when the unsplit grid cannot fill the device even once, where there is no reuse to protect. The second term is what keeps a 1k prefill at 8 chunks instead of dropping it to 1 and losing 19%. The target grid is unchanged at 4096, and every shape in this PR's published gfx950 table keeps its chunk count except 4k/1, which drops from 32 to 8. Shapes outside that table do move: across a 24-shape sweep at 256 CUs, 9 change, all toward fewer chunks (for example 2k/8 from 16 to 2 and 16k/1 from 32 to 16). Prefill score on MI350X was re-measured across that range at 0.98-1.02x. Prefill score against the base tree on MI325X now spans 1.20x to 2.89x, mean 1.79x, where the PR alone spanned 0.92x to 2.15x. End to end on MiniMax-M3-FP8-dynamic at TP=4, total throughput is +2.4% at 128k input and +1.3% at 64k, and flat at 8k where the indexer is a small share of the step. MI350X is unaffected: prefill score 0.98-1.02x, top-k worst 0.983x, decode 1.00-1.02x. The wide query-tile test took its query length from a hard-coded 8192, which selects the narrow tile on any part with more compute units than the one it was written on and so failed on gfx942's 304. Derive that length from the compute unit count and keep a tail case that does not divide the tile. Pin the split heuristic's own invariants too, since the kernel divides by its result and shifts by its bit length. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
|
On #47665, I don't think it reaches ROCm. It patches gfx942 needs one more piece on top of my PR! Currently I pinned the cache to OCP e4m3fn, which isn't native there, so the fused writer's conversion faults and the Triton path is slower than bf16. Native fnuz measures 1.46–1.92× vs bf16 on the decode scorer. I'll follow up with a separate PR shortly after! |
|
This pull request has merge conflicts that must be resolved before it can be |
4 participants
Motivation
MiniMax-M3 uses sparse attention. Before attention runs, a small indexer scores every 128 token block against the current query and keeps the best 16, and attention only looks at those. Each call is cheap, but the indexer runs for every token on almost every layer.
Two of its prefill kernels do more work than they need to.
The scorer gives one workgroup the entire causal scan for a query block, so the longest sequence in the batch decides when the launch finishes and the rest sit idle.
The top-k sorts one query token per program. The sorting network costs the same number of passes whether it carries one row or thirty two, so nearly all of that work goes unused.
The decode scorer is different: it is memory bound, re-reading the whole index cache on every launch. That is what the optional fp8 cache addresses.
Proposed Fix
Split the prefill scorer over the KV axis. Each query block's causal scan is cut into fixed, block aligned chunks, so no single program owns the whole scan. The scorer also takes 128 query rows per program instead of 64, sharing the same scan across twice as many rows.
Carry several query rows per top-k program. The sorting passes are then shared across a tile instead of one row. A wide tile divides the grid though, so it would starve short prefills: the width is picked on the host from the launch shape, and capped so the score tile and the sort stay small enough to keep registers free. Configs that cannot serve the requested top-k are rejected before launch with a clear error rather than a compile assertion.
The tile also carries only the running best 16 between iterations rather than the full sorted width. Keeping the full sort is bit identical to the baseline but costs 4 to 13 percent at long context, which is why the narrower selection network is here.
Optionally store the index cache as fp8 e4m3. This halves the bytes the decode scorer re-reads on every token. The score only decides which blocks attention looks at and the dot product still accumulates in fp32, so narrower keys are tolerable. It is the same configuration NVIDIA's M3 recipe already uses on B200.
The fp8 cache is opt-in behind
--attention-config '{"indexer_kv_dtype": "fp8"}', on gfx950 only. The default stays bf16 and both prefill changes apply either way.Kernel Benchmarking and E2E Results
Environment:
vllm/vllm-openai-rocm:nightly-6f7df92a8e6cdc74a725b8f10b4d0b48ba2b37efon 4x MI350X (gfx950),amd/MiniMax-M3-MXFP4at TP=4. Baseline is that same image with only the four touched files reverted.End to end benchmarking results comparing the baseline against this PR
The bf16 gain tracks input length only, where the prefill kernels live. The fp8 gain tracks concurrency too, which is the decode scorer reading half as many bytes per token.
GSM8K, full 1319 at 8 shot with chat template,
Reproduction
Serve
For the fp8 index cache, add
--attention-config '{"indexer_kv_dtype": "fp8"}'.GSM8K
pip install "lm_eval[api]" lm_eval --model local-chat-completions \ --model_args model=minimax-m3,base_url=http://localhost:8000/v1/chat/completions,num_concurrent=32,tokenized_requests=False,max_gen_toks=16384,timeout=3600 \ --tasks gsm8k --num_fewshot 8 \ --apply_chat_template --fewshot_as_multiturn \ --batch_size 32