Repository navigation
fix(runtime): address each cp.async vector in the destination's own layout - #263
Merged
Merged
Conversation
Contributor
Author
Test planThese ran on a single RTX 6000 Ada Generation (compute capability 8.9) with Python 3.12, PyTorch 2.11 for CUDA 12.9 and nvcc 12.8:
A/B on a single sm_89 GPU
The four new nodes are the new test for each target, the TIR roundtrip of the new fixture and the parser's source location check on it. All pass except the new test's H200 node, which fails on this GPU like every kernel built for sm_90, with The other 21 failures are the same tests on both sides: 19 launch kernels built for sm_90, and 2 need a second GPU. |
zhen8838
added a commit
that referenced
this pull request
Oct 23, 2026
## Why #261 declared the sm_89 E4M3 warp MMA and tested it as one atom whose fragments come straight from global memory. Nothing yet ran a GEMM with it: K slices staged through shared memory, the atom's fragments gathered from there, the product accumulated over K. Writing that kernel found the `cp.async` addressing fixed in #263. ## What One 16 x 8 x 128 E4M3 product, accumulated in F32 on sm_89, from its schedule down to the GPU. The schedule is three new fixtures, `tests/fixtures/schedule/{plain,hir,tir}/sm89_e4m3_gemm.py`: the plain program (each 32-wide K slice loaded into registers, where the SM89 atom reads, and multiplied), the scheduled program (one warp; per K slice, `T.copy_async` into shared memory, `T.copy` into the atom's fragments, `T.tiled_mma`), and the TIR `schedule finalize` writes for it. `tests/schedule/test_fixtures.py` adds their entries to the fixture tables and two `FixtureCase`s that run `tilefoundry check` on both programs against torch with `--atol 0 --rtol 0`: the inputs are integers in [-4, 4], so every partial sum is exact and the results must equal the product. The kernel, `tests/fixtures/tir/e4m3_gemm.py` (run from `tests/cuda/ops/tir/test_mma.py`), implements that finalized TIR with its allocation and order: both staging tiles sit at offset 0 of 512 bytes of dynamic shared memory, and each K step stages and gathers A, then B, then issues the MMA. It adds what finalize leaves to the implementation: the commit, wait and warp syncs between those steps, and the split of each copy over the warp's threads. The test is parametrized over the H200 and sm_89 targets, the sm_89 case under the `sm89` marker, and requires the tile to equal the F32 product exactly. ## Contract None. ## Risk The kernel's H200 case runs only in CI, on an H200; CI on the previous head passed with 2301 passed and 18 skipped. The kernel and the fixtures cover one warp, one 16 x 8 tile and integer inputs, and say nothing about performance. --------- Co-authored-by: Zheng QiHang <zhengqihang0915@qq.com>
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.
Why
ops::copy_asyncwalked the source in vectors but addressed the destination by element index. When a thread's slice holds more than one vector and its contiguous run is not the slice's first mode, later vectors can land at the wrong destination address. An E4M3 GEMM that stages its K slices through shared memory on sm_89, written on top of #261, found it: each thread's slice there is (2, 1, 1, 8) with strides (256, -, -, 1), its second 8-byte vector went to byte 4 instead of byte 256, and in that GEMM this caused a misaligned address error. The GPU tests so far move one vector per thread, so none of them reached it.What
CopyAsyncnow recasts the destination to the source's vector type and copies vectoriof the source to vectoriof the destination. The scalar width it used for the element index is gone.The new fixture
tests/fixtures/tir/split_stage.pyholdsSplitStage, in which each of 128 threads stages two 16-byte runs of f32 that lie 512 elements apart;tests/ops/tir/test_copy_async.pychecks on each CUDA target that the staged tile equals its input exactly, and the TIR roundtrip test picks the fixture up as it does the others.Contract
None. runtime §2.6.2 already states that the entry copies exactly the slice
local_tensorhands each thread; this makes the asynchronous entry do so.Risk
None.