jinhongyii opened a new issue, #20328:
URL: https://github.com/apache/tvm/issues/20328
### Expected behavior
`tcgen05.mma.ws` accepts `collector::b{N}::{fill,use,lastuse,discard}` to
keep the B operand resident across issues. A Blackwell (sm_100a) bf16 kernel
should be able to request it through the table-driven PTX dialect.
### Actual behavior
The token list exists (`_TCGEN05_WS_COLLECTOR_B`,
`python/tvm/backend/cuda/ptx/table.py:3250`) but is only wired into the
`kind::ti16` families, all of which are `cert_arch="sm_107f"`. The two entries
a Blackwell bf16 kernel uses carry no `collector_b` slot at all, so the
modifier is unreachable on that architecture:
```
entry cert_arch collector_b
tcgen05_mma_ws_sp_ti16_ss sm_107f yes
tcgen05_mma_ws_sp_ti16_ss_mask sm_107f yes
tcgen05_mma_ws_sp_ti16_ts sm_107f yes
tcgen05_mma_ws_sp_ti16_ts_mask sm_107f yes
tcgen05_mma_ws_ss sm_100a NO
tcgen05_mma_ws_ti16_ss sm_107f yes
tcgen05_mma_ws_ti16_ss_mask sm_107f yes
tcgen05_mma_ws_ti16_ts sm_107f yes
tcgen05_mma_ws_ti16_ts_mask sm_107f yes
tcgen05_mma_ws_ts sm_100a NO
```
### Environment
- apache/tvm at `3e50db91c8`; `table.py` last touched by #20271.
- Target `{"kind": "cuda", "arch": "sm_100a"}`, NVIDIA GB200 (152 SMs), CUDA
13.2.
### Steps to reproduce
```python
from tvm.backend.cuda.ptx.table import TABLE
ws = {n: e for n, e in TABLE.items() if "mma_ws" in n}
for name, entry in sorted(ws.items()):
slots = {s.name for s in entry.slots}
print(f"{name:34s} {str(entry.cert_arch):10s} {'yes' if 'collector_b' in
slots else 'NO'}")
sm100 = {n for n, e in ws.items() if str(e.cert_arch) == "sm_100a"}
have = {n for n, e in ws.items() if "collector_b" in {s.name for s in
e.slots}}
assert sm100 & have, "no sm_100a weight-stationary entry exposes collector_b"
```
The assertion fails today.
### Why it matters
A weight-stationary B-collector variant of a bf16 linear-attention kernel
was numerically correct on all six of its workloads on a GB200 once the
instruction was emitted, so the instruction itself appears usable on sm_100a —
only the table entry is missing. (That evidence comes from a kernel-evolution
run, not from an independent check in this issue.)
Reaching the modifier today requires mutating `TABLE` at runtime from the
kernel source, e.g.
```python
from tvm.backend.cuda.codegen.registry import register_codegen
from tvm.backend.cuda.ptx.engine import _make_codegen
from tvm.backend.cuda.ptx.table import ModifierSlot, _TCGEN05_WS_COLLECTOR_B
entry = replace(base, slots=base.slots + (ModifierSlot("collector_b",
_TCGEN05_WS_COLLECTOR_B, optional=False),))
namespace._table[name] = entry
register_codegen(f"ptx.{entry.name}")(_make_codegen(entry))
```
which reaches into private engine internals.
### A second, separable problem
A table entry can be replaced *without* its codegen closure being
re-registered, and nothing detects the mismatch. Doing so emits CUDA in which
the trailing marker token lands in an operand position, and the first sign of
trouble is an nvcc type error in generated code:
```
tvm_kernels.cu(2069): error: argument of type "const char *" is incompatible
with
parameter of type "uint64_t"
tvm_builtin_ptx_tcgen05_mma_ws_ss_ws_cta_group__1_kind__f16_collector__b0__fill(
..., (uint)135365776, (uint)0, "mma");
```
A registration-time consistency check would surface this as a table error
instead. This is worth fixing independently of whether the collector-B slot is
added.
### Suggested fix
Add the `collector_b` slot to the sm_100a weight-stationary entries; the
existing `collector B requires collector A` validation (`table.py:3053-3063`)
applies unchanged.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
