Repository navigation
feat(rocke): resolve gfx942 num_sms from the device for up to 14x mid-batch decode - #9739
Merged
Merged
Conversation
AviralGoelAMD
changed the base branch from
users/arai/attention_sink_decode
to
develop
July 22, 2026 23:15
AviralGoelAMD
changed the base branch from
develop
to
users/arai/attention_sink_decode
July 22, 2026 23:16
AviralGoelAMD
marked this pull request as draft
July 22, 2026 23:16
AviralGoelAMD
changed the base branch from
users/arai/attention_sink_decode
to
develop
July 22, 2026 23:17
Contributor
Author
|
waiting for #9391 to be merged |
✅ All Checks Passed — Ready for Review
📖 Need help? See the Policy FAQ for details on every check and how to fix failures. |
AviralGoelAMD
changed the base branch from
develop
to
users/arai/attention_sink_decode
July 22, 2026 23:23
AviralGoelAMD
force-pushed
the
users/avirgoel/ck/attention-num-sms-resolver
branch
from
July 22, 2026 23:28
ca45b15 to
9443b1e
Compare
AviralGoelAMD
changed the base branch from
users/arai/attention_sink_decode
to
develop
July 22, 2026 23:29
AviralGoelAMD
marked this pull request as ready for review
July 22, 2026 23:38
AviralGoelAMD
changed the title
AviralGoelAMD
changed the title
…it-KV segments num_sms (the split-KV device-subscription target driving 2D<->3D routing and the 3D segment count) defaulted to a stale 120, under-subscribing gfx942 (MI300X=304 CUs): mid-batch decode stayed in slow 2D and long-kv decode under-split. Dispatch: resolve num_sms from the live device CU count for gfx942 requests (AttentionRequest.num_sms default 120->0 sentinel), behind an override seam (ROCKE_NUM_SMS env > explicit value > device query > 304 fallback). Arch-guarded so a gfx942 request on another box falls back to the constant. Other archs keep the legacy 120 (future scope). Segments: bound _num_segments so the higher CU-count target cannot over-split already-3D shapes -- clamp to the pre-bump (num_sms=120) per-shape baseline for gfx942 D128 decode at kv<=16384 and for all q>1 (prefill/spec-decode); D128 kv>=32768, D64, and D256 stay uncapped where the finer split is measured to win. Validated on MI300X (fp32-gated, before/after 120->304): decode flips win 2.6-14.6x, q>1 flips 1.4-2.1x, long-kv/D64 re-splits 1.06-1.26x, and the naive-bump regressions (down to 0.50x) are eliminated -- zero regressions across 38 decode + 28 flip/prefill shapes. Adds tests/dispatch/attention/test_num_sms_resolver.py (CPU: resolution order + routing/segment effects).
AviralGoelAMD
force-pushed
the
users/avirgoel/ck/attention-num-sms-resolver
branch
from
July 23, 2026 00:39
9443b1e to
e9cff18
Compare
Pre-commit check failed⛔ pre-commit failed Please run locally:
This repo uses |
Codecov Report✅ All modified and coverable lines are covered by tests. ❌ Your project check has failed because the head coverage (76.84%) is below the target coverage (80.00%). You can increase the head coverage or adjust the target coverage. Additional details and impacted files@@ Coverage Diff @@
## develop #9739 +/- ##
===========================================
- Coverage 64.28% 64.12% -0.16%
===========================================
Files 2764 2765 +1
Lines 450621 450696 +75
Branches 66317 66320 +3
===========================================
- Hits 289656 288987 -669
- Misses 140161 140917 +756
+ Partials 20804 20792 -12
*This pull request uses carry forward flags. Click here to find out more. 🚀 New features to boost your workflow:
|
Contributor
|
Is there any reason that this is flagged rocfft? |
…s accessor Consolidate the two inlined num_sms*4 routing-target sites (select_path 2D<->3D decision and select_3d segment count) into a single _effective_target_ctas accessor on UnifiedAttentionProblem. Byte-identical no-op: same value flows to both call sites, so routing, segment counts, and emitted IR are unchanged. The _num_segments over-split clamp is untouched. Sets up a single source of truth for the device-subscription target ahead of adding a direct target_ctas spec knob.
Add a direct device-subscription target override (target_ctas) to UnifiedAttentionProblem and AttentionRequest. When > 0 it replaces num_sms*4 as the routing/segmentation target via the _effective_target_ctas accessor, letting tuners/benchmarks pin the target without a device CU count. Default 0 preserves num_sms*4 exactly (byte-identical routing, segments, and emitted IR). Remove the ROCKE_NUM_SMS environment pin from _resolve_num_cus (and the now -unused os import); reproducibility comes from an explicit num_sms or the target_ctas spec override. Resolver order is explicit num_sms -> gfx942 live device CU count -> 304/120 fallback. The split-KV segment clamp is unchanged. Tests: drop the env-tier cases, add target_ctas precedence, routing-flip, and _problem threading coverage.
…ency Document that target_ctas, like num_sms, is honoured on the Python dispatch path only. The C++ C-ABI engine (attention_unified_entry.cpp) hardcodes num_sms=120 and has no target_ctas field, so both need the mirror resolver + target_ctas there to take effect in C-ABI/hipDNN production (companion change). The default path (target_ctas=0 -> num_sms*4) stays byte-identical, so nothing breaks today; this only clarifies the production-path dependency.
…cy 120 A gfx942 request that could not verify the build box was gfx942 previously fell back to a hardcoded 304 (the MI300X CU count). That baked the wrong device's count into cross-compiled / off-box kernels (e.g. a gfx942 target that is actually MI300A, 228 CUs -> over-subscription) and re-introduced a hardcoded per-device CU literal. Now the live device CU count is used ONLY when the request targets gfx942 AND the build box IS gfx942. Every other case -- non-gfx942 arches, gfx942 built off-box, or no visible gfx942 device -- returns the legacy 120, matching develop. The _GFX942_NUM_CUS=304 literal is removed (304 now comes only from the live query). An explicit num_sms / target_ctas still supersedes. Shrinks the deviation from develop to the single verified-on-box case, avoids over-subscribing same-arch SKUs, and drops the hardcoded CU count. Byte-identity unaffected (dispatch-only; goldens pin num_segments). Off-box test now asserts 120 (was 304).
Contributor
Author
|
Thank you for nudging me into a better direction via your review comment with a better design I have incorporated that into my PR and simplified things now. |
…ents
Replace the SKU marketing names in the two comments this PR owns with the
arch (gfx942): the num_sms resolver docstring ("MI300X 304 vs MI300A 228" ->
"varies across parts") and the split-KV segment clamp comment ("measured on
MI300X" -> "measured on gfx942"). Comment-only; no logic/IR change.
Scoped to the lines this PR touches; the broader repo-wide SKU->gfx scrub is a
separate follow-up.
yraparti
approved these changes
Aug 4, 2026
This was referenced Oct 17, 2026
Merged
This was referenced Nov 3, 2026
Open
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What was happening before (plain language)
Attention dispatch has a knob called
num_sms— it tells the dispatcher how many compute units (CUs) the GPU has. That number decides two things:num_smswas hard‑coded to a stale120, but an MI300X (gfx942) has 304 CUs. So the dispatcher thought the GPU was less than half its real size, and medium/large‑batch decode that should have been split across the device stayed stuck in the slow 2D kernel, leaving most of the GPU idle. Long‑context decode also under‑split.What we changed
The
num_smsdevice resolver, a directtarget_ctasknob, and a segment-count bound:num_smsfrom the real device — only when verified on-box. For a gfx942 request built on a gfx942 box, read the live CU count (304 MI300X / 228 MI300A) → the reroute win. Everywhere else — non-gfx942 arches, or a gfx942 kernel built off-box / with no visible gfx942 device — keep the legacy 120 (matches develop), so a cross-compile never bakes the wrong device's count into the kernel. (Per review: the earlier hardcoded304off-box fallback was removed — 304 now comes only from the live query — and theROCKE_NUM_SMSenv pin was removed; pin reproducibility via an explicitnum_smsor thetarget_ctasknob.)target_ctasspec knob. A single device-subscription target (concurrent workgroups) that drives 2D↔3D routing and 3D segmentation. Default0→num_sms * 4(byte-identical to before); set>0to pin the target directly, so tuners / benchmarks / hipDNN can select the fast path without a device CU count. Consolidated behind one accessor (_effective_target_ctas) so routing and segmentation share one source of truth.num_smsalone made the 3D kernel over‑split already‑3D shapes; we clamp the segment count to the pre‑bump baseline where it over‑splits (gfx942 D128 decode ≤16384, and all q>1), leaving the cases that genuinely benefit (D128 kv≥32768, D64, D256) uncapped.How it acts as an optimization
num_smschanges only launch geometry, never math. Every result below was fp32‑gated.Benchmarks — GPU: AMD Instinct MI300X (gfx942, 304 CUs), bf16/fp16, before/after =
num_sms120 → 304, fp32‑gatedHeadline: mid‑batch decode flips — Llama / Qwen (develop‑reproducible; child PR #9742 shapes)
These were stuck in the slow 2D kernel before; the resolver reroutes them to 3D:
Mechanism (profiled on MI300X,
llama8b_dec_b80_kv8192, distinct block_table):The 14.6x is pure device-fill, not fewer bytes (~2.6 GiB moved either way): the 2D grid of 512 leaves the 304 CUs almost idle (occ 0.70, HBM 3% utilized); the reroute to 3D fills the device (grid 6400, occ 2.73, HBM 44%).
seg16confirms the segment clamp is active (not the naive 128). rocprofv3, one representative mid-batch D128 shape — mechanism measured, not inferred.Prefill / spec‑decode flips (q>1, causal‑fp32‑gated; Llama/Qwen, develop‑reproducible)
Existing corpus — 38 affected decode shapes (batch 1–2)
Long‑kv / D64 re‑splits (Llama/Qwen, develop‑reproducible):
Neutralized by the clamp (kernel byte‑identical → flat, no regression) — 24 shapes: Llama‑8B/70B/405B & Qwen‑235B, D128, kv4096/8192/16384 (bf16+fp16), all 0.97–1.08× (noise on identical kernels).
gpt‑oss sink flips — REQUIRES #9391 (not on develop), measured on the #9391 base:
Regressions
num_smswithout the clamp) regressed 21 of 38 shapes, worst 0.50× (2× slower) — exactly the over‑split the clamp fixes.0.97–0.99×readings are measurement noise (±~3%): they sit on byte‑identical (clamped) kernels, bf16/fp16 twins disagree run‑to‑run, and a 3× re‑run of the suspect shape read 1.010 / 1.010 / 1.013×.Tests
Adds
tests/dispatch/attention/test_num_sms_resolver.py— 9 CPU tests: resolution order (explicit / gfx942 device / MI300A=228 / other-arch=120), routing scales withnum_sms, segments stay bounded, andtarget_ctasprecedence (overridesnum_sms*4, flips 2D→3D at fixednum_sms, threads through_problem). Re-verified on MI300X: byte-identity GREEN (66 families), installed-golden PASS, and a same-node A/B confirming thetarget_ctaspath is perf-neutral vs pre-change (0 routing changes, 0.11% median timing diff, mid-batch win intact); the knob validated on silicon (2D→3D at fixednum_sms, output matches the 2D default, 14.1×).Scope / notes
ROCKE_NUM_SMSpin was removed per review; pin reproducibility via an explicitnum_smsor thetarget_ctasknob.num_sms=120and has notarget_ctasfield today (routing decided in Python); mirroring the device resolver +target_ctasthere is tracked as a companion follow-up (AICK-1856)._PRE_BUMP_SMS=120) is kept hardcoded here; parametrizing it as a spec knob (reviewer request) is tracked as AICK-1857.Changes are behind the feature flag HIPKERNELPROVIDER_ENABLE_ROCKE which is defaulted to OFF.