Fix grouped GEMM with CUTLASS DSL 4.7 (#3369)

The Hopper grouped GEMM wrapper selected the current NVVM TMA bindings only when the CUTLASS version prefix was exactly 4.6. A 4.7 wheel therefore fell back to the removed CpAsyncBulkTensorLoadMode/loadMode interface.

Detect the TMALoadMode capability directly, preserving the legacy fallback for older bindings. For example, 4.7 now uses the TMALoadMode/mode path and descriptor address-space cast.

Verification: Python syntax compilation and diff checks.
This commit is contained in:
Jiaxuan Bai
2026-07-08 10:16:39 +08:00
committed by GitHub
parent 4471280efc
commit e37889f489

View File

@@ -121,9 +121,9 @@ def CpAsyncBulkTensorGlobalToSharedClusterOpWrapper(
loc = None,
ip = None,
) -> None:
if cutlass.__version__[0:3] == "4.6":
if hasattr(_nvvm_d, "TMALoadMode"):
from cutlass._mlir.dialects import llvm as _llvm_d
# Handle new nvvm API in 4.6+
# Handle the current NVVM API when available.
tmaDescriptor = _llvm_d.addrspacecast(
_llvm_d.PointerType.get(_CuteAddressSpace.generic), tmaDescriptor, loc=loc, ip=ip
)