[Bugfix] Support CuTe DSL 4.8.0 block-scale API (#59480)

Signed-off-by: hjjq <50634613+hjjq@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Shang Wang <samshang.wang@mail.utoronto.ca>
This commit is contained in:
authored and GitHub committed 2026-10-01 02:30:06 +00:00
1 parent 766b4ae54c
commit bc505fc823
1 file changed
+6 -1
+6 -1
View File
@@ -148,6 +148,11 @@ def mma_mxfp8(
loc=None,
ip=None,
) -> None:
# CuTe DSL 4.8 renamed the block-scale enum and keyword.
if hasattr(nvvm, "Tcgen05MMABlockScale"):
scale_kwargs = {"block_scale": nvvm.Tcgen05MMABlockScale.DEFAULT}
else:
scale_kwargs = {"scale_vec_size": nvvm.Tcgen05MMAScaleVecSize.X1}
with cute.arch.elect_one():
nvvm.tcgen05_mma_block_scale(
nvvm.Tcgen05MMAKind.MXF8F6F4,
@@ -159,9 +164,9 @@ def mma_mxfp8(
Boolean(enable_input_d).ir_value(loc=loc, ip=ip),
_make_tmem_llvm_ptr(sfa_tmem, loc=loc, ip=ip),
_make_tmem_llvm_ptr(sfb_tmem, loc=loc, ip=ip),
scale_vec_size=nvvm.Tcgen05MMAScaleVecSize.X1,
loc=loc,
ip=ip,
**scale_kwargs,
)