fix: Use is_family_of() for SM12x arch guard in MmaSM120BlockScaledOp

[blake-snc]

Summary

• Replace hardcoded arch == Arch.sm_120a equality check with arch.is_family_of(Arch.sm_120a) in MmaSM120BlockScaledOp.__post_init__
• Add "sm_121a" to admissible_archs list for consistency with the error message

Problem

MmaSM120BlockScaledOp guards its arch check with if not arch == Arch.sm_120a, which rejects sm_121a (DGX Spark / GB10) even though the block-scaled MMA instruction set (mma.sync.aligned.block_scale) is identical across the SM12x family. The error message already references admissible_archs, showing the intent was to support multiple archs, but the guard ignores the list entirely:

admissible_archs = [
    "sm_120a",
]

def __post_init__(self) -> None:
    arch = BaseDSL._get_dsl().get_arch_enum()
    if not arch == Arch.sm_120a:  # ignores admissible_archs
        raise OpError(
            self,
            f"expects arch to be one of {self.admissible_archs}, but got {arch}",
            ...
        )

Fix

Use the existing is_family_of() method which was designed for exactly this purpose:

if not arch.is_family_of(Arch.sm_120a):

This accepts sm_120a, sm_120f, sm_121a, sm_121f — all SM12x family members that share the same block-scaled MMA instructions.

Validation (DGX Spark, SM121a)

Tested on NVIDIA GB10 (sm_121a) by patching the installed nvidia-cutlass-dsl package and running each test in a separate process (the DSL caches arch at init):

Before fix — sm_121a rejected:

$ CUTE_DSL_ARCH=sm_121a python3 -c "from cutlass.cute.nvgpu.warp.mma import MmaMXF4Op; ..."
FAILED: OpError: expects arch to be one of ['sm_120a'], but got Arch.sm_121a

After fix — all tests pass:

$ CUTE_DSL_ARCH=sm_121a python3 -c "...MmaMXF4Op..."
PASS: sm_121a accepted

$ CUTE_DSL_ARCH=sm_121a python3 -c "...MmaMXF4NVF4Op..."
PASS: MmaMXF4NVF4Op on sm_121a accepted

$ CUTE_DSL_ARCH=sm_120a python3 -c "...MmaMXF4Op..."
PASS: sm_120a accepted (no regression)

$ CUTE_DSL_ARCH=sm_90a python3 -c "...MmaMXF4Op..."
PASS: sm_90a correctly rejected (non-sm12x guard still works)

Contributed by Second Nature Computing (https://joinsecondnature.com)

[blake-snc]

Ping for review — small fix to use is_family_of() for the SM12x arch guard in MmaSM120BlockScaledOp so SM121a works correctly.

[johnnynunez]

cc @depaulmillz

[github-actions]

This PR has been labeled inactive-30d due to no recent activity in the past 30 days. Please close this PR if it is no longer required. Otherwise, please respond with a comment indicating any updates. This PR will be labeled inactive-90d if there is no activity in the next 60 days.

[blake-snc]

Still relevant on current main (b46b16d003). The hardcoded equality check is unchanged in python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py:138:

135:     def __post_init__(self) -> None:
136:         # Verify arch
137:         arch = BaseDSL._get_dsl().get_arch_enum()
138:         if not arch == Arch.sm_120a:
139:             raise OpError(
140:                 self,
141:                 f"expects arch to be one of {self.admissible_archs}, but got {arch}",

sm_121a (DGX Spark / GB10) still rejected by this path despite mma.sync.aligned.block_scale being identical across the SM12x family. Happy to rebase or address reviewer comments.

Note: commenting to reset the inactive-30d label. I will revalidate on SM121a hardware this weekend to confirm this PR is still relevant, and will close if it turns out it no longer is.

[blake-snc]

Runtime confirmation on cutlass-dsl 4.4.2 / NVIDIA GB10 (sm_121a) today:

$ CUTE_DSL_ARCH=sm_121a python3 -c "
from cutlass.cute.nvgpu.warp.mma import MmaMXF4Op
from cutlass.cute.typing import Float4E2M1FN, Float8E8M0FNU, Float32
MmaMXF4Op(ab_dtype=Float4E2M1FN, acc_dtype=Float32, sf_type=Float8E8M0FNU)"

OpError: expects arch to be one of ['sm_120a'], but got Arch.sm_121a
Suggestions: Ensure env CUTE_DSL_ARCH matches your GPU architecture

Same line 138 of warp/mma.py, same hardcoded equality check this PR replaces with is_family_of(Arch.sm_120a). Bug is observable on real SM12x hardware today.

[AethoceSora]

depaulmillz commented 3 weeks ago

tcgen05 is only supported on sm_100 and sm_103. I believe you are looking for these MMA atoms which are supported on sm_120 and sm_121 -

cutlass/include/cute/atom/mma_traits_sm120.hpp

Line 121 in 418d38a

struct MMA_Traits<SM120_16x8x32_TN<a_type, b_type, c_type>>

cutlass/python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py

Line 123 in 418d38a

class MmaSM120BlockScaledOp(MmaOp):

As of the current main branch, the admissible architectures for the MmaSM120BlockScaledOp operator are limited to sm_120a. This restriction prevents devices with the sm_121 architecture from utilizing the corresponding MMA atoms.

This PR addresses the issue by replacing the strict architecture check with is_family_of(), allowing proper support for SM12x-family GPUs.

I would appreciate it if you could consider merging this PR :)
@depaulmillz

[depaulmillz]

There will be a fix for this in the next 4.5-wheel release.

[blake-snc]

Closing this — it's been superseded by main. MmaSM120BlockScaledOp.admissible_archs now lists sm_120a/sm_120f/sm_121a/sm_121f with the conventional if arch not in self.admissible_archs guard, matching the pattern in tcgen05/copy.py and tcgen05/mma.py. That already accepts sm_121a (DGX Spark / GB10), which was the motivation here, so the is_family_of() change is no longer needed. Thanks!