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 48242ec334 [CUDA][TIRx] Add collector-qualified tcgen05.mma
block_scale forms and bind tcgen05.ld.red redval as output (#20266)
48242ec334 is described below
commit 48242ec33403f2b6e4fac6e763ca7a683fb9d5df
Author: Bohan Hou <[email protected]>
AuthorDate: Thu Sep 3 21:20:55 2026 -0400
[CUDA][TIRx] Add collector-qualified tcgen05.mma block_scale forms and bind
tcgen05.ld.red redval as output (#20266)
## Summary
Follow-up to #20261 for two `tcgen05` forms in the `T.ptx` table
(`python/tvm/backend/cuda/ptx/table.py`), per the CUDA 13.4 / PTX ISA
9.4 manual.
- **`tcgen05.mma` block-scaled with explicit block size**
(`tcgen05_mma_block_scale_block_{ss,ts}`): the entries move from the
base table into `_PTX_94_ENTRIES` and gain optional `collector_a` /
`collector_b` slots, 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 for `kind::mxf8f6f4` and `kind::mxf4`,
block16/block32 for `kind::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`**: `redval` is 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 from
`test_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`: `redval` is rendered
with an output constraint.
## Downstream
The FP4 FlashAttention-4 port in mlc-ai/tirx-kernels uses
`tcgen05.ld.red...max.f32` and depends on the `redval` binding fix.
Co-authored-by: Bohan Hou <[email protected]>
---
python/tvm/backend/cuda/ptx/table.py | 73 +++++++-----
tests/python/tirx/codegen/test_ptx_dialect.py | 163 +++++++++++++++++++++++++-
2 files changed, 199 insertions(+), 37 deletions(-)
diff --git a/python/tvm/backend/cuda/ptx/table.py
b/python/tvm/backend/cuda/ptx/table.py
index 9359781caf..1aeb7dda50 100644
--- a/python/tvm/backend/cuda/ptx/table.py
+++ b/python/tvm/backend/cuda/ptx/table.py
@@ -3050,8 +3050,7 @@ def _check_tcgen05_mma_block_scale(m):
def _check_tcgen05_mma_block_scale_block(m):
- """Valid block sizes per kind: mxf8f6f4/mxf4 use block32, while
- mxf4nvf4 supports block16 and block32."""
+ """Validate documented block sizes and require collector A before
collector B."""
valid = {
"kind::mxf8f6f4": ("block32",),
"kind::mxf4": ("block32",),
@@ -3059,6 +3058,9 @@ def _check_tcgen05_mma_block_scale_block(m):
}[m["kind"]]
if m["block_size"] not in valid:
return f"{m['kind']} supports {'/'.join(valid)}"
+ collector_b = m.get("collector_b", "")
+ if collector_b and not m.get("collector_a", ""):
+ return "collector B requires collector A"
return None
@@ -5347,6 +5349,42 @@ _PTX_94_ENTRIES = [
for sparse in (False, True)
for form in ("ss", "ts")
],
+ # PTX ISA 9.4, 9.7.18.10.10.1 syntax forms 2 and 4 have the same operand
+ # shape for each A location, so each entry owns its no-collector and
+ # collector-qualified variants. Table 68 permits block32 for mxf8f6f4,
+ # block32 for mxf4, and block16/block32 for mxf4nvf4. Form 4 requires
+ # collector A and makes collector B optional; collector B requires sm_107f.
+ *[
+ InstructionEntry(
+ name=f"tcgen05_mma_block_scale_block_{form}",
+ mnemonic="tcgen05",
+ slots=(
+ ModifierSlot("action", ("mma",)),
+ ModifierSlot("cta_group", ("cta_group::1", "cta_group::2")),
+ ModifierSlot("kind", ("kind::mxf8f6f4", "kind::mxf4",
"kind::mxf4nvf4")),
+ ModifierSlot("block_scale", ("block_scale",)),
+ ModifierSlot("block_size", ("block16", "block32")),
+ ModifierSlot("collector_a", _TCGEN05_COLLECTOR_A,
optional=True),
+ ModifierSlot("collector_b", _TCGEN05_COLLECTOR_B,
optional=True),
+ ),
+ check=_check_tcgen05_mma_block_scale_block,
+ cert_arch="sm_107f",
+ operands=(
+ OperandSlot("d_tmem", kind="addr", space="tmem"),
+ *(
+ (OperandSlot("a_desc", dtype="u64"),)
+ if form == "ss"
+ else (OperandSlot("a_tmem", kind="addr", space="tmem"),)
+ ),
+ OperandSlot("b_desc", dtype="u64"),
+ OperandSlot("idesc", dtype="u32"),
+ OperandSlot("sfa_tmem", kind="addr", space="tmem"),
+ OperandSlot("sfb_tmem", kind="addr", space="tmem"),
+ OperandSlot("enable_input_d", dtype="pred"),
+ ),
+ )
+ for form in ("ss", "ts")
+ ],
# PTX ISA 9.4, 9.7.10.28.5.2/3: the address override composes with the
# im2col load modes and report mechanism. These siblings preserve the
# trailing im2col operand while placing global_address inside the address.
@@ -11317,7 +11355,7 @@ _ENTRIES = [
cert_arch="sm_103a",
operands=(
OperandSlot("r", rw="w", lanes=_tcgen05_ldst_lanes),
- OperandSlot("redval"),
+ OperandSlot("redval", rw="w"),
OperandSlot("taddr", kind="addr", space="tmem"),
*((OperandSlot("imm_half_splitoff", kind="imm"),) if split
else ()),
),
@@ -11485,35 +11523,6 @@ _ENTRIES = [
)
for form in ("ss", "ts")
],
- *[
- InstructionEntry( # block-scaled with an explicit scale block size
- name=f"tcgen05_mma_block_scale_block_{form}",
- mnemonic="tcgen05",
- slots=(
- ModifierSlot("action", ("mma",)),
- ModifierSlot("cta_group", ("cta_group::1", "cta_group::2")),
- ModifierSlot("kind", ("kind::mxf8f6f4", "kind::mxf4",
"kind::mxf4nvf4")),
- ModifierSlot("block_scale", ("block_scale",)),
- ModifierSlot("block_size", ("block16", "block32")),
- ),
- cert_arch="sm_100a",
- check=_check_tcgen05_mma_block_scale_block,
- operands=(
- OperandSlot("d_tmem", kind="addr", space="tmem"),
- *(
- (OperandSlot("a_desc", dtype="u64"),)
- if form == "ss"
- else (OperandSlot("a_tmem", kind="addr", space="tmem"),)
- ),
- OperandSlot("b_desc", dtype="u64"),
- OperandSlot("idesc", dtype="u32"),
- OperandSlot("sfa_tmem", kind="addr", space="tmem"),
- OperandSlot("sfb_tmem", kind="addr", space="tmem"),
- OperandSlot("enable_input_d", dtype="pred"),
- ),
- )
- for form in ("ss", "ts")
- ],
*[
InstructionEntry( # weight-stationary: no mask vector, a zero-column
desc
name=f"tcgen05_mma_ws_{form}",
diff --git a/tests/python/tirx/codegen/test_ptx_dialect.py
b/tests/python/tirx/codegen/test_ptx_dialect.py
index 3f0ff6ab58..2304824933 100644
--- a/tests/python/tirx/codegen/test_ptx_dialect.py
+++ b/tests/python/tirx/codegen/test_ptx_dialect.py
@@ -2354,6 +2354,40 @@ def test_ptx_tcgen05_mma_block_size_form():
assert "tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32" in src
_assert_ptxas_ok(src, arch="sm_100a")
+ @T.prim_func
+ def sm107_collector_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::mxf4nvf4.block_scale.block16"
+ ".collector::a::discard.collector::b::fill"
+ ](tmem, desc, desc, idesc, tmem, tmem, T.ptx.pred(flag))
+ T.ptx[
+ "tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32"
+ ".collector::a::fill.collector::b::lastuse"
+ ](tmem, tmem, desc, idesc, tmem, tmem, T.ptx.pred(flag))
+ A[tx] = A[tx]
+
+ collector_src = _cuda_source(sm107_collector_kernel)
+ ss_collector_opcode = (
+ "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16"
+ ".collector::a::discard.collector::b::fill"
+ )
+ ts_collector_opcode = (
+ "tcgen05.mma.cta_group::1.kind::mxf4.block_scale.block32"
+ ".collector::a::fill.collector::b::lastuse"
+ )
+ assert ss_collector_opcode in collector_src
+ assert ts_collector_opcode in collector_src
+ _assert_ptxas_ok(collector_src, arch="sm_107f")
+
with pytest.raises((ValueError, tvm.error.DiagnosticError),
match="mxf4.*block32"):
@T.prim_func
@@ -2368,6 +2402,77 @@ def test_ptx_tcgen05_mma_block_size_form():
)
+def test_ptx_tcgen05_mma_block_size_collector_legality():
+ from tvm.backend.cuda.ptx.table import TABLE, tokens_for
+
+ kind_blocks = (
+ ("kind::mxf8f6f4", "block32"),
+ ("kind::mxf4", "block32"),
+ ("kind::mxf4nvf4", "block16"),
+ ("kind::mxf4nvf4", "block32"),
+ )
+ for form in ("ss", "ts"):
+ entry = TABLE[f"tcgen05_mma_block_scale_block_{form}"]
+ for kind, block_size in kind_blocks:
+ required = {
+ "action": "mma",
+ "cta_group": "cta_group::1",
+ "kind": kind,
+ "block_scale": "block_scale",
+ "block_size": block_size,
+ }
+ tokens_for(entry, **required, collector_a="collector::a::fill")
+ tokens_for(
+ entry,
+ **required,
+ collector_a="collector::a::fill",
+ collector_b="collector::b::lastuse",
+ )
+ with pytest.raises(ValueError, match="collector B requires collector
A"):
+ tokens_for(
+ entry,
+ action="mma",
+ cta_group="cta_group::1",
+ kind="kind::mxf4",
+ block_scale="block_scale",
+ block_size="block32",
+ collector_b="collector::b::fill",
+ )
+
+
+@requires_nvcc
+def test_ptx_tcgen05_mma_block_size_no_b_certifies_at_sm100f():
+ """No-collector and collector-A-only forms retain their documented lower
floor."""
+ from tvm.backend.cuda.ptx.render import render_variant
+ from tvm.backend.cuda.ptx.table import TABLE, tokens_for
+
+ kind_blocks = (
+ ("kind::mxf8f6f4", "block32"),
+ ("kind::mxf4", "block32"),
+ ("kind::mxf4nvf4", "block16"),
+ ("kind::mxf4nvf4", "block32"),
+ )
+ sources = []
+ for form in ("ss", "ts"):
+ entry = TABLE[f"tcgen05_mma_block_scale_block_{form}"]
+ for kind, block_size in kind_blocks:
+ for collector_a in ("", "collector::a::fill"):
+ kwargs = {
+ "action": "mma",
+ "cta_group": "cta_group::1",
+ "kind": kind,
+ "block_scale": "block_scale",
+ "block_size": block_size,
+ }
+ if collector_a:
+ kwargs["collector_a"] = collector_a
+ tokens = tokens_for(entry, **kwargs)
+ _, helper, helper_source = render_variant(entry, tokens)
+ sources.append(_certification_kernel(helper, helper_source,
len(sources)))
+
+ _assert_ptxas_ok("\n".join((_CERT_PRELUDE, *sources)), arch="sm_100f")
+
+
def test_ptx_pred_operand_rejects_untagged_integer():
"""An untagged integer at a `.pred` position is refused, by name.
@@ -2856,11 +2961,11 @@ def test_ptx_94_cp_bulk_semantic_negative_grids():
)
-def test_ptx_94_family_specific_arch_floors_and_delta():
- """SM107 family-only forms certify at 107f and remain in the 9.4 delta."""
+def test_ptx_94_sm107_arch_floors_and_delta():
+ """Entries owning SM107 variants certify at 107f and remain in the 9.4
delta."""
from tvm.backend.cuda.ptx.table import _PTX_94_ENTRIES, TABLE
- family_specific = {
+ sm107_entries = {
"add_mixed_vec_up",
"sub_mixed_vec_up",
"add_mixed_vec_down_f16",
@@ -2873,8 +2978,10 @@ def test_ptx_94_family_specific_arch_floors_and_delta():
"mul_mixed_vec_bf16_f16",
"mul_mixed_vec_f16_bf16",
"set_packed",
+ "tcgen05_mma_block_scale_block_ss",
+ "tcgen05_mma_block_scale_block_ts",
}
- assert {TABLE[name].cert_arch for name in family_specific} == {"sm_107f"}
+ assert {TABLE[name].cert_arch for name in sm107_entries} == {"sm_107f"}
delta_names = {entry.name for entry in _PTX_94_ENTRIES}
noftz_siblings = {
@@ -3990,7 +4097,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 == 761703 # update when the table grows or a ptxas gap
narrows it
+ assert total == 762023 # update when the table grows or a ptxas gap
narrows it
def test_ptx_no_instruction_registered_twice():
@@ -4391,5 +4498,51 @@ def test_ptx_ld_st_gpu_roundtrip():
tvm.testing.run_with_gpu_lock(run_and_check)
+def test_ptx_tcgen05_ld_red_binds_redval_as_output():
+ """ISA 9.7.18.8.3: `tcgen05.ld.red... r, redval, [taddr]` writes the
reduction result into
+ `redval`. The helper must bind it with an output constraint; an input
binding compiles but
+ the kernel never observes the hardware max (measured on GB300: the probe's
redval stayed at
+ its initial value until the binding was fixed)."""
+ from tvm.backend.cuda.ptx.render import render_variant
+ from tvm.backend.cuda.ptx.table import TABLE, tokens_for
+
+ for name, modifiers, binding in (
+ (
+ "tcgen05_ld_red",
+ dict(
+ action="ld",
+ red="red",
+ sync="sync",
+ aligned="aligned",
+ shape="32x32b",
+ num="x2",
+ redop="max",
+ type="f32",
+ ),
+ '"=f"(__redval)',
+ ),
+ (
+ "tcgen05_ld_red_split",
+ dict(
+ action="ld",
+ red="red",
+ sync="sync",
+ aligned="aligned",
+ shape="16x32bx2",
+ num="x2",
+ redop="min",
+ type="s32",
+ ),
+ '"=r"(__redval)',
+ ),
+ ):
+ entry = TABLE[name]
+ assert next(s for s in entry.operands if s.name == "redval").rw == "w"
+ imms = ("0",) if name.endswith("_split") else ()
+ _, _, source = render_variant(entry, tokens_for(entry, **modifiers),
imms=imms)
+ assert binding in source, source
+ assert '"r"(__redval)' not in source and '"f"(__redval)' not in source
+
+
if __name__ == "__main__":
tvm.testing.main()
