Conversation
_auto_variant returned f"mfma_r2_w{wpb}", so rows_per_block was pinned at 2 and
seven of the nine registered variants -- including every member of the r4 family
-- could never be selected.
r4 amortizes each KV tile load over twice as many query rows and is faster from
seq_len 8 upward. The gap is widest exactly where it costs most: vLLM chunks
indexer prefill to fit VLLM_SPARSE_INDEXER_MAX_LOGITS_MB (512 MB), which caps
seq_len at 1024 when seq_len_kv is 131072, so the existing seq_len >= 2048 branch
cannot fire at long context and every such call took mfma_r2_w4 -- the median of
the nine by speed, with the best 1.65x faster.
Measured on MI325X (gfx942), seq_len_kv 131072, best variant vs the r2 pick:
seq_len 1 r2 25.2 us r4 79.1 us r4 3.1x worse
seq_len 4 r2 24.3 us r4 31.5 us r4 1.3x worse
seq_len 8 r2 34.0 us r4 32.9 us r4 1.03x better
seq_len 16 r2 56.5 us r4 47.5 us r4 1.19x better
seq_len 1024 r2 2564.3 us r4 1520.3 us r4 1.69x better
Below seq_len 8 the host padding of seq_len up to a multiple of RPB dominates --
at seq_len 1 an r4 kernel computes 4 rows to obtain 1 -- so r2 is kept there and
behaviour for those shapes is unchanged.
Logits are bitwise identical across all variants at every shape tested, so this
is purely a blocking/occupancy change.
End to end on 8x MI325X, TP8, GLM-5.2-FP8, 131072 in / 1024 out, concurrency 8:
median TPOT improves 7.22% and output throughput 6.60%. This kernel is 16.8% of
GPU time at that point.
Signed-off-by: Jin Tao <jin.tao@amd.com>
…t a divisor
Refines the previous commit's rule after a 2-D sweep. That rule keyed RPB off
seq_len alone with a crossover measured only at seq_len_kv=131072; sweeping the
other contexts shows the crossover is not a seq_len threshold at all, and that a
second effect was being read as one.
RPB tracks the logits element count. Over seq_len 1..8192 x seq_len_kv
1024..262144 on MI325X, the boundaries land on the same element count at every
context: RPB=1 wins below 2**19 elements (27/27 shapes), RPB=2 at 2**19 (6/6),
RPB=4 from 2**21 up (38/38), with 2**20 a transition band split 3/3. Keying off
seq_len instead put the previous rule on the wrong side at low context: at
seq_len 16, seq_len_kv 1024 it chose RPB=4 and ran 1.26x slower than RPB=1.
RPB must also divide seq_len. When it does not, the launcher pads with four
torch.cat calls; that is a flat ~44 us of host-side overhead, independent of
seq_len_kv, and it is the whole of the "small seq_len" penalty the previous
commit attributed to wasted rows. At seq_len 1, seq_len_kv 131072: RPB=1 23.1 us,
RPB=2 67.8 us, of which the four cats are 44.1 us and pre-padding by hand
recovers all of it (21.9 us). So the penalty is not proportional to the padding
-- 1 wasted row of 2 costs the same as 3 of 4 -- and it applies to every odd
seq_len, which the old rule sent to RPB=2 unconditionally.
Stepping down to a divisor is only right while the kernel is cheap relative to
that fixed cost, so it is gated to the same 2**21 elements: at seq_len 1025,
seq_len_kv 131072 the dividing RPB=1 takes 3880 us against 2601 us for a padded
RPB=2.
Measured on MI325X, no FLYDSL_FP8_MQA_LOGITS_VARIANT set, old pick vs new:
seq_len seq_len_kv old new speedup
1024 131072 r2_w4 2598.7 r4_w4 1623.5 1.60x
1025 131072 r2_w4 2609.3 r4_w4 1594.0 1.64x
512 131072 r2_w4 1190.1 r4_w4 795.3 1.50x
700 50000 r2_w4 661.6 r4_w4 440.3 1.50x
333 12000 r2_w4 118.4 r4_w4 97.1 1.22x
16 131072 r2_w4 58.3 r4_w4 47.4 1.23x
3 131072 r2_w4 66.4 r1_w4 28.1 2.36x
1 131072 r2_w4 69.4 r1_w4 23.2 2.99x
1 1024 r2_w4 66.2 r1_w4 22.2 2.98x
No shape measured regresses; the smallest gain is 1.04x. Against the best of the
nine variants at each shape, pooled over held-out data (non-power-of-two shapes,
a fine seq_len sweep, and head counts 16 and 64), the geometric mean cost falls
from 1.45x to 1.03x and the worst case from 3.17x to 1.41x.
Logits are bitwise identical across all nine variants at all 180 shapes swept
(1620 timings), so this remains purely a blocking/occupancy change.
WPB is deliberately left alone. It is worth a few percent at most here, and
unlike RPB its optimum moves with the head count -- at 64 heads the current
WPB rule costs 1.61x worst case where a fixed WPB=4 costs 1.07x -- so it needs
its own sweep rather than a change fitted to one head count.
Signed-off-by: Jin Tao <jin.tao@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags: |
There was a problem hiding this comment.
Pull request overview
Updates the gfx942 FlyDSL fp8_mqa_logits kernel variant auto-selector so rows_per_block (RPB) is no longer effectively pinned to 2, allowing the faster r4 variants to be selected for large “long-context prefill” shapes (and selecting r1 for small shapes where padding overhead dominates).
Changes:
- Add element-count-based thresholds (
seq_len * seq_len_kv) to choose RPB ∈ {1, 2, 4}, including a divisibility-aware fallback to avoid padding overhead on small/medium shapes. - Keep the existing WPB heuristic unchanged, while returning the full
mfma_r{rpb}_w{wpb}variant tag sor1/r4families are reachable.
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
|
@vpietila-amd @valarLip — would appreciate a look when you have a moment. Short version: on gfx942 |
samremes
left a comment
samremes
left a comment
There was a problem hiding this comment.
Please trim the comments down, they don't need to tell the whole story how it ended up with these. State the logic very briefly.
| _RPB2_MIN_ELEMS = 2**19 | ||
| _RPB4_MIN_ELEMS = 2**21 |
There was a problem hiding this comment.
Could move this inside the _auto_variant as they are very local to the heuristic.
Move RPB element-count thresholds into _auto_variant and shorten the docstring; logic unchanged. Signed-off-by: Jin Tao <jin.tao@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
| """Pick (RPB, WPB) from the problem shape. | ||
|
|
||
| RPB from ``seq_len * seq_len_kv`` thresholds; step down to a divisor of | ||
| ``seq_len`` when padding overhead would dominate. WPB unchanged. | ||
| """ |
| rpb2_min_elems = 2**19 | ||
| rpb4_min_elems = 2**21 | ||
| elems = seq_len * seq_len_kv | ||
| if elems < rpb2_min_elems: | ||
| rpb = 1 | ||
| elif elems < rpb4_min_elems: | ||
| rpb = 2 if seq_len % 2 == 0 else 1 | ||
| else: | ||
| rpb = 4 | ||
| wpb = 2 if (seq_len >= 2048 and seq_len_kv >= 8192) else 4 | ||
| return f"mfma_r2_w{wpb}" | ||
| return f"mfma_r{rpb}_w{wpb}" |
Thank you @samremes Both addressed in d863e4d. The comments are trimmed down to a two-line docstring that just states the rule (RPB from seq_len * seq_len_kv thresholds, stepped down to a divisor of seq_len when padding overhead would dominate; WPB unchanged), and the two element-count constants are now locals inside _auto_variant. The rationale and measurements stay in the commit message. No logic change. |
There was a problem hiding this comment.
🟢 Approval recommended
The change is localized to variant selection logic, matches the available registered variant tags, and does not alter kernel math/output semantics.
Review details
- Files reviewed: 1/1 changed files
- Comments generated: 0 new
- Review effort level: Lite
There was a problem hiding this comment.
🟢 Approval recommended
The change is a small, self-contained update to variant selection logic that only returns already-registered variants and preserves existing WPB behavior.
Review details
- Files reviewed: 1/1 changed files
- Comments generated: 0 new
- Review effort level: Lite
Motivation
On gfx942
_auto_variantreturnsf"mfma_r2_w{wpb}", sorows_per_blockis pinned at 2 andseven of the nine registered variants — including the entire
r4family — can never beselected. At the shape vLLM actually issues for long-context indexer prefill (
seq_len1024,seq_len_kv131072) the pinned pick is the median of the nine by speed, with the best 1.65xfaster. vLLM chunks prefill to fit
VLLM_SPARSE_INDEXER_MAX_LOGITS_MB(512 MB), which capsseq_lenat 1024 at that context, so the selector's ownseq_len >= 2048branch can neverfire there. Since vllm-project/vllm#49544 routes gfx942
fp8_mqa_logitsto this kernel bydefault, every MI300X/MI325X user inherits it.
Technical Details
Sweeping all nine variants over 180 shapes surfaced two separable effects.
RPB tracks the logits element count, not
seq_len. Overseq_len ∈ [1, 8192] × seq_len_kv ∈ [1024, 262144]the boundaries land on the same value ofseq_len * seq_len_kvat all six contexts:r1below 2^19 (27/27 shapes),r2at 2^19 (6/6),r4from 2^21 up (38/38), with 2^20 a transition band split 3/3.RPB must also divide
seq_len. When it does not, the launcher pads with fourtorch.catcalls. That is a flat ~44 us of host-side overhead, independent of
seq_len_kv: atseq_len1,seq_len_kv131072,r1runs 23.1 us andr267.8 us, of which the four cats measure 44.1 us(0.7 us unexplained), and pre-padding the input by hand recovers all of it. So the cost is not
proportional to the wasted rows — one padded row of two costs the same as three of four — and the
pinned
r2walks into it on every oddseq_len, 37 of the 180 shapes swept, at a median 2.97xpenalty.
The step-down to a divisor has to be gated to small shapes, where the kernel is cheap relative to
44 us. Above the threshold the reuse wins even when padding: at
seq_len1025,seq_len_kv131072, the dividing
r1takes 3880 us against 2601 us for a paddedr2.waves_per_blockis deliberately left unchanged. It is worth a few percent at most here, andunlike RPB its optimum moves with the head count — at 64 heads the current rule costs 1.61x worst
case where a fixed
w4costs 1.07x — so it warrants its own sweep rather than a change fitted toone head count.
Test Plan
MI325X (gfx942), all nine variants timed at 180 shapes (1620 timings): a power-of-two
seq_len × seq_len_kvgrid, plus three held-out sets the rule was not fitted on —non-power-of-two shapes, a fine
seq_lensweep through the divisibility cliff, and head counts16 and 64. Selector changes were verified end to end through
flydsl_fp8_mqa_logits(variant=None)with no
FLYDSL_FP8_MQA_LOGITS_VARIANTset, so the selector itself is under test. Every variant'soutput was compared against a reference at every shape.
Test Result
Old pick vs new, no env override:
No shape measured regresses; the smallest gain is 1.04x. Scored against the fastest of the nine
at each shape, pooled over the held-out data, the geometric mean cost falls from 1.45x to 1.03x
and the worst case from 3.17x to 1.41x.
Logits are bitwise identical across all nine variants at all 180 shapes, so this is purely a
blocking/occupancy change with no numerical trade-off.
End to end on 8x MI325X, TP8, GLM-5.2-FP8, 131072 in / 1024 out, concurrency 8: median TPOT
improves 7.22% and output throughput 6.60%. This kernel is 16.8% of GPU time at that point, the
largest single item in the profile.
Submission Checklist