Repository navigation
[FlyDSL] [Bugfix] gfx950 A16W16: fix int32 C-offset overflow - #6853
Open
josh-schell-amd wants to merge 12 commits into
Open
josh-schell-amd wants to merge 12 commits into
josh-schell-amd wants to merge 12 commits into
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
josh-schell-amd
changed the title
1 task done
josh-schell-amd
added a commit
to josh-schell-amd/aiter
that referenced
this pull request
Oct 23, 2026
…past buffer limits The gfx950 A16W16 epilogue computed C offsets as int32 `row * n + col`, so any GEMM with M * N >= 2^31 stored to a wrapped address. Tuned BF16 rows are looked up by padded M, and getPaddedM(gl=1) maps every M > 8192 with N > 4096 to the M=8192 row, so a row tuned at M=8192 serves arbitrarily large M. With the MiniMax-M3 ViT rows from ROCm#6235, vLLM's encoder-cache profile run (M=768000, N=6144, K=1280) hits this and faults with "Write access to a read-only page" on MI355X (reproduced standalone via tgemm.mm at M>=360000). - Compute the non-split-K C offsets in int64 in both the full-tile and half-tile-interleaved epilogues. - Reject shapes whose A or B exceeds the 2 GiB buffer-resource range in make_gemm_a16w16_param_and_validate (also prunes them from tuning). - Skip a FlyDSL tuned row in get_GEMM_A16W16_tuned_config when the kernel cannot run the actual shape, falling back like an unrecognized kernel. - Add the large-M MiniMax-M3 ViT shape and a shape-limit check to op_tests/test_flydsl_hgemm.py. Co-Authored-By: Claude <noreply@anthropic.com>
The buffer-resource loads index elements with int32 offsets, so the limit is 2^31 elements (4 GiB of 16-bit data), not 2^31 bytes. The byte bound rejected the tuned Kimi-K3 row M=32768 N=7168 K=35840 (A = 2.19 GiB, tuned with err_ratio 0), which failed `python -m aiter.aot.flydsl.gemm` and would have fallen back to torch at runtime. All 2543 tuned gfx950 FlyDSL rows now pass. Co-Authored-By: Claude <noreply@anthropic.com>
…nning the last call FlyDSL 0.4.0's layout-dynamic `from_dlpack` argument stores each launch's DLTensorAdaptor in a thread-local (`_shared.dladaptor`) that is never cleared, and the fill closure lives as long as the cached CallState. Every tensor argument of the A16W16 kernel therefore keeps its most recent call's tensor alive. After vLLM's MiniMax-M3 encoder profile run (M=768000 ViT projector) that pins A+C = 10.6 GiB, which vLLM counts as persistent memory: 7.4 GiB less KV cache than when the GEMM ran on torch. `from_torch_tensor` reads data_ptr/shape/stride straight from the tensor and retains nothing. `_dynamic_tensor_arg` is shared with the AOT driver, so runtime and AOT keep the same argument type and cache key. Measured on MI355X: torch reserved after `del` + `empty_cache` drops from 10.64 to 0.02 GiB; output matches torch at M=8192/360000/768000. Co-Authored-By: Claude <noreply@anthropic.com>
…B limit A and B are read through one flat buffer descriptor each, so like the A8W8 preshuffle GEMM's _check_preshuffle_flat_buffer_capacity, operands must stay below 4 GiB. State the bound in bytes against the same limit instead of 2^31 elements (equivalent for the 16-bit inputs this kernel takes), and explain why C needs none. Note that flydsl_hgemm_supports_shape only models the gfx950 kernel; gfx1250 rows keep their previous behavior. No behavior change. Co-Authored-By: Claude <noreply@anthropic.com>
Shorten the new comments to the codebase's style and test the A/B bound at its exact boundary instead of narrated example shapes. No behavior change. Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
…riant flydsl_hgemm accepts a row-strided `out` on gfx1250 and returns the caller's tensor, but the gfx950 kernel asserted `out.is_contiguous()` and returned a new view, so op_tests/flydsl_tests/test_gemm_a16w16.py's kernel-config table failed on gfx950 and tgemm.mm(out=) with a strided view would hit the assert. A contiguous `out` keeps C's row stride as N and compiles exactly as before. A row-strided `out` selects a separate variant (`_ldc`) that reads the stride from `out`'s layout, so only strided calls pay for it. Split-K's zero-fill and atomics use the same stride; the host checks it and split-K's 4 GiB atomic range. A 2D `out` is returned as-is. Tests: the config table skips (otype, config) pairs the gfx950 kernel cannot build (a 256x256 fp32 C tile exceeds LDS), and test_flydsl_hgemm.py checks ft, hti and split-K writes into a strided view. Co-Authored-By: Claude <noreply@anthropic.com>
josh-schell-amd
force-pushed
the
fix/flydsl-hgemm-int64-offsets
branch
from
October 23, 2026 21:57
38f564e to
7b73ede
Compare
josh-schell-amd
marked this pull request as ready for review
October 23, 2026 21:57
…st 2^31 - gemm_a16w16 accepts strided A/B views, whose i32 buffer offsets are row * stride + k; reject operands whose real span reaches 4 GiB even when M * K * 2 is below it. The shape-only check stays for tuned_gemm and AOT, which only have M/N/K. - A row-strided `out` must have a row stride below 2^31: the kernel reads it as i32. - test_flydsl_hgemm.py: a meta-device check for the A span, a numeric check of the first and last rows at M = 1677721 (A just under 4 GiB, byte offsets past 2^31), and the extra checks only run on the default sweep. Co-Authored-By: Claude <noreply@anthropic.com>
Conflict: gemm_a16w16_gfx950_utils.py, where ROCm#7067 raised SPLIT_K_SEMAPHORE_MAX_LEN to 2048 next to the new GFX950_FLAT_BUFFER_LIMIT_BYTES; both kept.
Contributor
|
Found this fix independently; I think int64 output offset plus a fallback to torch for exceeding buffer is the right approach. My equivalent fix solved the Minimax issue on MI350. Nice work |
josh-schell-amd
changed the title
READY FOR HUMAN REVIEWPR Self-ReviewRevision Author self-review accepted. Readiness report: Author answers
Show the author's self-review answersBehavior and intentValidation and unverified scopePR-specific questionChanged op testsPR state
Copilot findings
Code and op-test hygiene
Required CI
Shadow mode reports what would block but keeps the status successful. Set |
Conflict: op_tests/test_flydsl_hgemm.py, where ROCm#7042 added test_hgemm_concurrent_copy at the same spot as this PR's check_input_span / check_large_a / check_strided_out; both kept.
- Dense operands (row- or column-major) are covered by the shape check, so the per-call span computation now runs only for strided views, keeping host work on the launch path flat. - test_flydsl_hgemm.py: run the extra checks whenever -s is not given. The default shape list now also holds the concurrent-copy shape, so comparing against _default_mnk() skipped them. Co-Authored-By: Claude <noreply@anthropic.com>
josh-schell-amd
force-pushed
the
fix/flydsl-hgemm-int64-offsets
branch
from
October 29, 2026 18:35
bf9caa7 to
8b6061a
Compare
micah-wil
added a commit
to ROCm/vllm
that referenced
this pull request
Oct 29, 2026
AITER's gfx950 FlyDSL A16W16 kernel addresses C with int32 offsets and faults once M * N >= 2^31 (ROCm/aiter#6853). With --mm-encoder-tp-mode data, the MiniMax-M3 ViT profile run issues an M=768000, N=6144, K=1280 GEMM that AITER 0.1.26 routes to that kernel, crashing engine startup with a memory access fault. Fall back to torch (hipBLASLt) for those shapes. They are compute-bound and hipBLASLt matches the FlyDSL kernel's throughput there. Signed-off-by: Micah Williamson <micah.williamson@amd.com>
2 of 6 tasks
micah-wil
added a commit
to ROCm/vllm
that referenced
this pull request
Oct 29, 2026
AITER's gfx950 FlyDSL A16W16 kernel reads A and B through flat buffer descriptors with i32 offsets, so an operand of 4 GiB or more is read out of bounds and the kernel silently returns wrong rows. For example, Llama-405B's TP8 gate_up GEMM (N=13312, K=16384) maps to a FlyDSL row and returns wrong output at 147456 tokens, while C is still under the 2^31 limit. Route these shapes to torch as well, matching the limit ROCm/aiter#6853 enforces. Drop the on_gfx950() check, which is_tgemm_enabled() already implies. Signed-off-by: Micah Williamson <micah.williamson@amd.com>
xytpai
requested changes
Oct 30, 2026
xytpai
left a comment
xytpai
left a comment
Contributor
There was a problem hiding this comment.
Review result: Request changes
Reviewed commit 8b6061a64836fa3ca7a782ea1211f857c55bac57, including all six changed files and the relevant call paths. I found one correctness issue that should be addressed before merging: the newly accepted strided output views can violate the alignment required by 16-bit split-K packed atomics. Details and requested changes are in the inline comment.
Validation and limitations
- Executed the PR's host-side validation with the launch mocked and confirmed that an odd-column-offset output view reaches dispatch.
- An independent fp16
BufferAtomicPkAddtest on MI308X/gfx942 succeeded with aligned addresses but triggeredHSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATIONwith a one-element offset. - This was a primitive-level reproduction using FlyDSL 0.3.4.1, not an end-to-end gfx950/FlyDSL 0.4.x GEMM reproduction.
- The host-side input-size boundary checks behaved as expected (
M=1677721accepted andM=1677722rejected forN=6144, K=1280with the tested configuration). No other actionable findings were identified within the reviewed scope.
Comment on lines
+1663
to
+1666
| out.stride(1) != 1 | ||
| or out.stride(0) < n | ||
| or out.stride(0) >= 1 << 31 | ||
| or out.stride(0) % param.cshuffle_r2g_vec_size != 0 |
Contributor
There was a problem hiding this comment.
[P1] Validate output pointer alignment for strided split-K outputs
The new validation checks output strides but not the starting address alignment. For example:
parent = torch.empty((16, 5184), dtype=torch.float16, device="cuda")
out = parent[:, 1:5121]
# shape: (16, 5120)
# stride: (5184, 1)
# data_ptr() % 4 == 2
With the existing M=16, N=5120, K=2880, split_k=3 configuration, this view passes validation. However, the fp16/bf16 split-K path uses BufferAtomicPkAdd, which requires a 4-byte-aligned address for each pair of 16-bit elements. An odd column offset violates that requirement even when the row stride is valid, potentially causing a GPU fault.
This case is newly accepted by this PR; previously, the non-contiguous output was rejected by out.is_contiguous().
Please check out.data_ptr() % 4 == 0 for 16-bit split-K outputs before launching, rejecting unsupported alignment or providing a safe alternative. Please also add an odd-column-offset rejection test and an aligned nonzero-offset positive test, checking that elements outside the view remain unchanged. The current check_strided_out() only covers parent[:m, :n].
I confirmed that this view reaches dispatch in the PR's host-side validation with the launch mocked. An independent fp16 packed-atomic test on gfx942 reproduced a GPU memory-aperture fault for the misaligned address; this is not an end-to-end gfx950 GEMM reproduction (see the review summary for environment details).
Author
There was a problem hiding this comment.
ack. I made the check dtype-independent (split_k > 1 and out.data_ptr() % 4 != 0): the fp32 split-K path also uses 4-byte atomics, so it has the same requirement. fp32 tensors are always 4-byte aligned, so in practice this only rejects 16-bit outputs.
- Split-K's buffer atomics write 4-byte units: an fp32 element or a packed 16-bit pair. A 16-bit out that is not 4-byte aligned faults (HSA_STATUS_ERROR_MEMORY_APERTURE_VIOLATION on MI355X). Reject a misaligned out before launch, for contiguous and strided out alike. Non-split-K stores handle any element offset. - test_flydsl_hgemm.py: check_out_column_offset writes an out that starts at a nonzero column and checks nothing outside it changes; check_split_k_out_alignment expects a ValueError for a misaligned split-K out, strided and contiguous. Co-Authored-By: Claude <noreply@anthropic.com>
josh-schell-amd
force-pushed
the
fix/flydsl-hgemm-int64-offsets
branch
from
October 30, 2026 04:48
b9682ad to
14b38ab
Compare
Contributor
|
LGTM. |
xytpai
approved these changes
Oct 30, 2026
This branch has not been deployed
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.
Summary
Fixes a GPU memory fault in the gfx950 FlyDSL A16W16 GEMM when M·N ≥ 2^31 (int32 C offsets), which crashes vLLM
amd/MiniMax-M3-MXFP4at startup on AITER v0.1.26. It also stops the kernel's FlyDSL arguments from keeping the last call's tensors alive, rejects shapes the kernel can't address, and accepts a row-stridedouton gfx950 so this kernel's own config test passes.Motivation
Since AITER v0.1.26, vLLM
amd/MiniMax-M3-MXFP4(MI355X, TP4) crashes at engine startup withMemory access fault ... Write access to a read-only pageduring the encoder-cache profile run.M=8192, N=6144, K=1280row tominimax_m3_mxfp4_bf16_tuned_gemm.csv.getPaddedM(gl=1)maps every M > 8192 (with N > 4096) to that row, so the 768000-row ViT GEMM runs the gfx950 FlyDSL A16W16 kernel.global_row * n + global_col, in int32, which wraps once M·N ≥ 2^31 (row 349,526 at N=6144). The wrap causes:Memory access faultvLLM reports at M=768000.tgemm.mmwith N=6144, K=1280 passes at M=340000 and faults at M=360000.Fixing the crash exposed a second issue. The kernel's layout-dynamic
from_dlpackarguments keep their last call's tensors alive (a FlyDSL 0.4.0 bug, fixed upstream in ROCm/FlyDSL#1321). After the ViT profile run that pins 10.6 GiB, which costs vLLM 7.4 GiB of KV cache.AITER main at
537ac7f6d(2026-10-30) still has the int32 offsets, thefrom_dlpackarguments and the8192,6144,1280FlyDSL row.Changes
SPLIT_K_SEMAPHORE_MAX_LEN).gemm_a16w16also checks the real span of strided A/B views (rows × row stride), which can reach 4 GiB while M·K·2 does not; dense operands skip that check.get_GEMM_A16W16_tuned_configskips a FlyDSL row the kernel can't run at the real M/N/K (flydsl_hgemm_supports_shape) and falls back, as it does for an unrecognized kernel. It reuses the lookup's owngfx, so the check models the architecture the row was selected for._dynamic_tensor_argusesfrom_torch_tensor, which retains nothing and is cheaper per launch. FlyDSL v0.4.1 does not include [Fix] Don't retain the last tensor in layout-dynamic DLPack fills FlyDSL#1321, so this is still needed.out:flydsl_hgemmalready accepts a row-stridedouton gfx1250 and returns it; gfx950 now does too. A stridedoutbuilds a separate variant (_ldc) that reads C's row stride fromout; a contiguousoutcompiles to the same code as before. It's in this PR becauseop_tests/flydsl_tests/test_gemm_a16w16.py, this kernel's own test, fails without it, and it's its own commit (7b73ede). The row stride must be below 2^31. Strided variants aren't in the AOT cache, so the first strided call compiles at runtime.Performance
torch.mm: 4729 µs)gemm_a16w16host time per call, 8x4096x4096Kernel perf vs v0.1.26. v0.1.26's kernel and this PR's kernel are built side by side in one process and timed with CUDA graphs over 20 rounds that alternate which one runs first. Each result is the mean time ratio with a 95% confidence interval. An A/A control (v0.1.26 against an identical copy of itself) shows how large a difference the method reports when nothing changed: up to about ±0.9% on 5 µs kernels.
Shapes: the 31 tuned
ntrows intest_flydsl_hgemm.pythat v0.1.26 can run (M·N < 2^31), plusnn/tn/ttat 4096³ and 98304x6144x1280, for 37 in total.360448x6144x1280, which v0.1.26 can't run (M·N > 2^31), is compared with
torch.mminstead: 4702 µs vs 4729 µs.In the large interleaved kernel (256x256 tiles), the 64-bit offset adds 26 instructions to the epilogue, outside the main loop; the A/B shows no difference on those shapes (4096x4096x8192, 8192³, and the 98304-row layouts). Computing a 64-bit tile base once and 32-bit offsets per store was tried: the compiler then added more per-store address math (for example 965 → 1043 epilogue instructions on 4096x4096x8192), so it isn't used.
The later commits (strided
out, the A/B span check) don't change these kernels: for a contiguousoutthey compile byte-identical to the build measured here (checked on all three kernel paths: standard tiles, large interleaved tiles, and split-K).vLLM
amd/MiniMax-M3-MXFP4(MI355X, TP4, perf-evalminimax_m3-mi355x). On v0.1.26 the engine faults while profiling the encoder cache (CI build 709):With this PR the same step runs that GEMM on FlyDSL (
AITER_LOG_TUNED_CONFIG=1), and the engine starts and completes the workload:MiniMax-M3's decode GEMMs run on this kernel, so serving exercises the changed store path. The columns are different vLLM images.
Testing
Tested on MI355X (gfx950); the kernel and the new dispatch check are gfx950-only. Setup: the vLLM image where the nightly fails (vLLM
df6603f0, AITER v0.1.26), with and without this PR's changes applied, on FlyDSL 0.4.0 and again on 0.4.1 (main's pin).Reproduction. With the image's default tuned configs (which include
minimax_m3_mxfp4_bf16_tuned_gemm.csv), this faults on v0.1.26 and matches torch with this PR:op_tests/test_flydsl_hgemm.pyop_tests/flydsl_tests/test_gemm_a16w16.pyouttest_a16w16_tuned_config_lookup.pypython -m aiter.aot.flydsl.gemmDocumentation
gemm_a16w16: a 2Doutmay be row-strided)Dependencies
Breaking Changes
None for existing correct calls. Behavior changes on gfx950
gemm_a16w16:ValueError; before, it computed wrong results.outis now accepted (it used to fail an assert), and a 2Doutis returned as the same tensor instead of a new view.PR Self-Review
Reviewed revision
14b38abBehavior and intent
gemm_a16w16calls raiseValueError.from_torch_tensor, so FlyDSL 0.4.x no longer keeps the last call's tensors alive (10.6 GiB after MiniMax's profile run).outis accepted on gfx950, as on gfx1250, through a compile-time variant; a contiguousoutcompiles to the same code as before.Validation and unverified scope
On MI355X, in the vLLM image from the failing CI run (AITER v0.1.26 with this PR's changes applied), on FlyDSL 0.4.0 and 0.4.1:
python op_tests/test_flydsl_hgemm.py: all tuned rows and this PR's checks passed.python op_tests/flydsl_tests/test_gemm_a16w16.py: 86/86 passed.python -m unittest op_tests.tuning_tests.test_a16w16_tuned_config_lookup: OK.python -m aiter.aot.flydsl.gemm: 2505 ok, 0 failed.tgemmrepro, the memory check, the kernel A/B against v0.1.26, the host-time measurement and the vLLM perf-eval run, all in Performance and Testing above.Unverified:
test_flydsl_hgemm.py's newtest_hgemm_concurrent_copy([FlyDSL] [Bugfix] Sync both wave groups before the gfx950 A16W16 HTI GEMM reads its first and last stages #7042) fails because the base lacks [FlyDSL] [Bugfix] Sync both wave groups before the gfx950 A16W16 HTI GEMM reads its first and last stages #7042's kernel fix; it is not exercised against main here.op_tests/test_gemm_a16w16.py -t out: needstgemm.mm(out=)from after v0.1.26.Not caused by this PR:
op_tests/test_gemm_a16w16.pyhas two directgemm_a16w16_asmchecks at 64x256x5120 (bf16 output) that also fail on main CI, andop_tests/test_graph_alloc.pyis fixed separately by #7039.PR-specific question
Without this PR, the 360448x6144x1280 row crashes, the strided-
outand A/B span checks fail, and the config table inflydsl_tests/test_gemm_a16w16.pystops at its first gfx950 config. No existing test covered these cases; details per file are in Changed op tests.Changed op tests
op_tests/test_flydsl_hgemm.py: the canonicalflydsl_hgemmtest, extended rather than duplicated. It adds the M·N > 2^31 row, the shape-limit and real-span checks, the M = 1677721 result check, and strided-outchecks for standard tiles, large interleaved tiles and split-K. The extra checks run whenever-sisn't given.op_tests/flydsl_tests/test_gemm_a16w16.py: skips (with a warning) config/output pairs the gfx950 kernel can't build (a 256x256 fp32 C tile exceeds LDS). Without the skip, the config table, the only test that sweeps explicitflydsl_hgemmconfigs against triton, can't run on gfx950.This PR was authored with assistance from Claude Code (AI-assisted). Results were run and checked on MI355X.
🤖 Generated with Claude Code