fix: Use is_family_of() for SM90 arch guard in warpgroup MmaOp

[blake-snc]

Summary

• Replace hardcoded arch == Arch.sm_90a with arch.is_family_of(Arch.sm_90a) in warpgroup/mma.py's MmaOp.__post_init__

Problem

MmaOp in warpgroup/mma.py uses if not arch == Arch.sm_90a to guard its arch check. While functionally correct today (sm_90a is the only Hopper "a"-suffix arch), this is inconsistent with the is_family_of() pattern used elsewhere in CuTe DSL (see #3082 for the SM12x equivalent fix).

Fix

# Before
if not arch == Arch.sm_90a:

# After
if not arch.is_family_of(Arch.sm_90a):

Validation

$ python3 -c "from cutlass.base_dsl.arch import Arch; print(Arch.sm_90a.is_family_of(Arch.sm_90a))"
True

$ python3 -c "from cutlass.base_dsl.arch import Arch; print(Arch.sm_120a.is_family_of(Arch.sm_90a))"
False

$ python3 -c "from cutlass.base_dsl.arch import Arch; print(Arch.sm_100a.is_family_of(Arch.sm_90a))"
False

sm_90a accepted, non-Hopper arches correctly rejected.

Related

• fix: Use is_family_of() for SM12x arch guard in MmaSM120BlockScaledOp #3082 — Same fix for SM12x arch guard in warp/mma.py

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

[blake-snc]

Ping for review — same is_family_of() pattern as #3082, applied to the SM90 warpgroup MmaOp guard.

[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/warpgroup/mma.py:134:

131:     def __post_init__(self) -> None:
132:         # Verify arch
133:         arch = BaseDSL._get_dsl().get_arch_enum()
134:         if not arch == Arch.sm_90a:
135:             raise OpError(
136:                 self,
137:                 f"expects arch to be {Arch.sm_90a}, but got {arch}",

Pure consistency fix with the is_family_of() pattern used elsewhere in CuTe DSL (see #3082 for the SM12x equivalent). Happy to rebase or address reviewer comments.

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

[github-actions]

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

[blake-snc]

Responding to the inactivity label: this is still relevant at current main. The guard in python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py remains a strict arch == sm_90a equality check, so SM90 family variants still trip it; the two-line change here needs no special hardware to review. Our SM120 hardware is currently allocated to other work, so my own follow-up bandwidth is limited for the time being; I am also reachable at github.com/bledden.