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.