[CUDA][TIRx] Add collector-qualified tcgen05.mma block_scale forms and bind tcgen05.ld.red redval as output - #20266
Merged
spectrometerHBH merged 1 commit intoSep 4, 2026
Conversation
tqchen
approved these changes
Sep 4, 2026
This was referenced Sep 4, 2026
spectrometerHBH
added a commit
that referenced
this pull request
Sep 4, 2026
…#20270) ## Summary `test_ptx_tcgen05_mma_block_size_form` assembles the collector-qualified `tcgen05.mma ... block_scale` forms at `sm_107f`. Those are PTX ISA 9.4 forms (#20266) and need the CUDA 13.4 ptxas; on a CUDA 13.2 host the test fails with `Unknown modifier '.collector::b::fill'` instead of skipping. This moves that half into `test_ptx_tcgen05_mma_block_size_collector_form`, guarded by `env.has_nvcc_version(13, 4)`, the same gate the whole-table certification tests already use. The `sm_100a` block32 certification and the `mxf4` block16 legality check stay in the original test and keep running on older toolchains. ## Testing On a CUDA 13.2 host (B200): ``` tests/python/tirx/codegen/test_ptx_dialect.py -k tcgen05_mma_block_size 3 passed, 1 skipped (collector form: need nvcc >= 13.4) ``` The full `tests/python/tirx/` suite on the same host previously had this as its only failure (2749 passed, 321 skipped).
spectrometerHBH
added a commit
that referenced
this pull request
Sep 5, 2026
….async wait_group immediates (#20271) ## Summary Two `T.ptx` table follow-ups to #20266 (`python/tvm/backend/cuda/ptx/table.py`). - **Collector-A-only SM107 block-scale MMA.** On the SM107 `tcgen05.mma` block-scale entries `collector_b` becomes optional, so the activation-stationary FP8 form with collector A alone is accepted alongside the A+B form. - **Open `cp.async` wait-group immediates.** `cp.async.wait_group N` and `cp.async.bulk.wait_group[.read] N` no longer restrict `N` to `0..7`. PTX ISA 9.7.10.28.3.3 / 9.7.10.28.6.2 define `N` only as an integer constant with no value domain, and CUDA 13.4 ptxas at sm_107a accepts values beyond 7 (measured: 8, 9, 16, 255; the bulk `.read` form also 2147483647 and -1). Each call-site constant still becomes its own helper; enumeration and full-table certification sample the open operand at 0 and certify the instruction shape. ## Tests (`tests/python/tirx/codegen/test_ptx_dialect.py`) - New `test_ptx_tcgen05_mma_block_scale_collector_a_without_block_size`: SM107 activation-stationary FP8 accepts collector A without `.block*`. - Wait-group tests updated for the open immediate. ## Downstream Needed by the Rubin masked grouped GEMM and block-scaled fused-MoE SwiGLU ports in mlc-ai/tirx-kernels.
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
2 participants
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
Follow-up to #20261 for two
tcgen05forms in theT.ptxtable (python/tvm/backend/cuda/ptx/table.py), per the CUDA 13.4 / PTX ISA 9.4 manual.tcgen05.mmablock-scaled with explicit block size (tcgen05_mma_block_scale_block_{ss,ts}): the entries move from the base table into_PTX_94_ENTRIESand gain optionalcollector_a/collector_bslots, so each entry owns both the no-collector (ISA 9.7.18.10.10.1 syntax form 2) and collector-qualified (form 4) variants, which share the same operand shape per A location. Table 68 block sizes are unchanged (block32 forkind::mxf8f6f4andkind::mxf4, block16/block32 forkind::mxf4nvf4). The check additionally rejects collector B without collector A. The entries certify at sm_107f because collector B is SM107-only; the no-collector and collector-A-only forms keep their sm_100f floor.tcgen05.ld.red:redvalis now bound as an output operand (rw="w"). ISA 9.7.18.8.3 writes the reduction result into it; with an input binding the kernel compiles but never observes the hardware max (measured on GB300).Tests (
tests/python/tirx/codegen/test_ptx_dialect.py)test_ptx_tcgen05_mma_block_size_collector_legality: collector A/B legality and ptxas certification of the collector forms at sm_107f.test_ptx_tcgen05_mma_block_size_no_b_certifies_at_sm100f: no-collector and collector-A-only forms still certify at sm_100f.test_ptx_94_sm107_arch_floors_and_delta(renamed fromtest_ptx_94_family_specific_arch_floors_and_delta): entries owning SM107 variants certify at sm_107f and remain in the 9.4 delta.test_ptx_tcgen05_ld_red_binds_redval_as_output:redvalis rendered with an output constraint.Downstream
The FP4 FlashAttention-4 port in mlc-ai/tirx-kernels uses
tcgen05.ld.red...max.f32and depends on theredvalbinding fix.