This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 5e49b34304 [CUDA][TIRx] Allow collector-A-only SM107 block-scale MMA
and open cp.async wait_group immediates (#20271)
5e49b34304 is described below
commit 5e49b343042335e32d61fdcba5d3916efb7d07ed
Author: Bohan Hou <[email protected]>
AuthorDate: Fri Sep 4 21:27:52 2026 -0400
[CUDA][TIRx] Allow collector-A-only SM107 block-scale MMA and open cp.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.
---
python/tvm/backend/cuda/ptx/table.py | 18 +++++++++-----
tests/python/tirx/codegen/test_ptx_dialect.py | 36 ++++++++++++++++++++++++++-
2 files changed, 47 insertions(+), 7 deletions(-)
diff --git a/python/tvm/backend/cuda/ptx/table.py
b/python/tvm/backend/cuda/ptx/table.py
index 1aeb7dda50..94e0121f82 100644
--- a/python/tvm/backend/cuda/ptx/table.py
+++ b/python/tvm/backend/cuda/ptx/table.py
@@ -5328,7 +5328,7 @@ _PTX_94_ENTRIES = [
),
ModifierSlot("block_scale", ("block_scale",)),
ModifierSlot("collector_a", _TCGEN05_COLLECTOR_A),
- ModifierSlot("collector_b", _TCGEN05_COLLECTOR_B),
+ ModifierSlot("collector_b", _TCGEN05_COLLECTOR_B,
optional=True),
),
cert_arch="sm_107f",
operands=(
@@ -8516,6 +8516,12 @@ _ENTRIES = [
OperandSlot("b", dtype="u32"),
),
),
+ # PTX ISA 9.7.10.28.3.3 / 9.7.10.28.6.2 defines N only as an integer
+ # constant and declares no value domain. Keep it OPEN: callers may use any
+ # compile-time integer, while certification samples the instruction shape.
+ # MEASURED on CUDA 13.4 ptxas at sm_107a: ordinary and bulk forms accept
+ # values beyond 7 (8, 9, 16, 255), and the bulk `.read` form also accepts
+ # 2147483647 and -1.
*[
InstructionEntry( # cp.async{.bulk}.wait_group{.read} N;
name=f"cp_async{'_bulk' if bulk else ''}_wait_group",
@@ -8527,7 +8533,7 @@ _ENTRIES = [
*((ModifierSlot("read", ("read",), optional=True),) if bulk
else ()),
),
orders_memory=True,
- operands=(OperandSlot("group", kind="imm", choices=tuple(str(n)
for n in range(8))),),
+ operands=(OperandSlot("group", kind="imm"),),
)
for bulk in (False, True)
],
@@ -10179,10 +10185,10 @@ _ENTRIES = [
# and cp.async.wait_all per 9.7.10.28.3.3, cp.async.bulk.commit_group /
# .wait_group per 9.7.10.28.6.1 / 9.7.10.28.6.2.
#
- # The wait_group counts are caller-chosen immediates: the ISA gives N no
- # register form, so each value is its own helper, and the closed `choices`
- # set is what makes every one of them certifiable. 0..7 covers every call
- # site (pipeline depths); widen the tuple if a deeper pipeline appears.
+ # The wait_group counts are caller-chosen OPEN immediates: the ISA gives N
+ # no register form or value domain. Each call-site constant becomes its own
+ # helper; enumeration and full-table certification sample the open operand
+ # at 0 and therefore certify the instruction shape rather than every value.
#
# (The `cp.async` ca/cg copy lines this note once excluded are registered
# in the 9.7.10 group above, ignore-src operand and all.)
diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py
b/tests/python/tirx/codegen/test_ptx_dialect.py
index 4290125f50..835979b6f2 100644
--- a/tests/python/tirx/codegen/test_ptx_dialect.py
+++ b/tests/python/tirx/codegen/test_ptx_dialect.py
@@ -1977,6 +1977,10 @@ def test_ptx_data_movement_dispatch():
T.ptx.applypriority.global_.L2__evict_normal(A.ptr_to([6]))
T.ptx.discard.global_.L2(A.ptr_to([7]))
T.ptx.prefetchu.L1(A.ptr_to([0]))
+ T.ptx.cp.async_.wait_group(255)
+ T.ptx.cp.async_.bulk.wait_group(255)
+ T.ptx.cp.async_.bulk.wait_group.read(8)
+ T.ptx.cp.async_.bulk.wait_group.read(-1)
T.ptx.multimem_ld_reduce.add.u32(v, A.ptr_to([0]))
T.ptx.multimem_red.relaxed.gpu.add.u32(A.ptr_to([0]), v)
smem[tx % 4] = d + p + v
@@ -2000,6 +2004,10 @@ def test_ptx_data_movement_dispatch():
"applypriority.global.L2::evict_normal [%0], 128;",
"discard.global.L2 [%0], 128;",
"prefetchu.L1 [%0];",
+ "cp.async.wait_group 255;",
+ "cp.async.bulk.wait_group 255;",
+ "cp.async.bulk.wait_group.read 8;",
+ "cp.async.bulk.wait_group.read -1;",
"multimem.ld_reduce.add.u32 %0, [%1];",
"multimem.red.relaxed.gpu.add.u32 [%0], %1;",
):
@@ -2411,6 +2419,32 @@ def test_ptx_tcgen05_mma_block_size_collector_form():
_assert_ptxas_ok(collector_src, arch="sm_107f")
+@requires_nvcc
+def test_ptx_tcgen05_mma_block_scale_collector_a_without_block_size():
+ """SM107 activation-stationary FP8 accepts collector A without
`.block*`."""
+
+ @T.prim_func
+ def kernel(a_ptr: T.handle):
+ A = T.match_buffer(a_ptr, (32,), "uint32")
+ T.device_entry()
+ T.cta_id([1])
+ tx = T.thread_id([32])
+ if tx == 0:
+ tmem = T.local_scalar("uint32")
+ desc = T.local_scalar("uint64")
+ idesc = T.local_scalar("uint32")
+ flag = T.local_scalar("uint32")
+
T.ptx["tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.collector::a::discard"](
+ tmem, desc, desc, idesc, tmem, tmem, T.ptx.pred(flag)
+ )
+ A[tx] = A[tx]
+
+ src = _cuda_source(kernel)
+ opcode =
"tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale.collector::a::discard"
+ assert opcode in src
+ _assert_ptxas_ok(src, arch="sm_107a")
+
+
def test_ptx_tcgen05_mma_block_size_collector_legality():
from tvm.backend.cuda.ptx.table import TABLE, tokens_for
@@ -4106,7 +4140,7 @@ def test_ptx_all_variants_render_unique():
_, helper, _ = render_variant(entry, *args,
addr_offsets=addr_offsets)
assert helper not in names, f"address-offset helper name
collision: {helper}"
names.add(helper)
- assert total == 762023 # update when the table grows or a ptxas gap
narrows it
+ assert total == 762050 # update when the table grows or a ptxas gap
narrows it
def test_ptx_no_instruction_registered_twice():
