[BUG]: check_dtype_support rejects family-conditional (sm_XXXa) gpu_code targets

从 cuda/tile/_passes/check_dtype_support.py 中的 check_dtype_support() 开始，然后跟踪 compile_tile() 和 export_kernel() 如何传递 gpu_code。使用 sm_121 和 sm_121a 运行提供的 export_smoke.py 复现。完成的标准是 family-conditional targets 不再引发所示的 ValueError，同时现有的 base-target 路径仍能继续编译。

Version

1.4.0 (installed via pip; reproduced identically against the 1.5.0 wheel's source, see below)

CUDA Toolkit Version

13.3 (V13.3.73)

Which installation method(s) does this occur on?

Pip

Describe the bug.

export_kernel(..., gpu_code=<family-conditional target>, output_format="cubin") — e.g. gpu_code="sm_121a" — raises ValueError: invalid literal for int() with base 10: '121a' instead of compiling. The failure is in check_dtype_support():

sm_number = int(sm_arch.removeprefix("sm_"))

This doesn't strip a trailing architecture-conditional a suffix (the CUDA convention for family-specific targets like sm_90a, sm_100a, sm_120a, sm_121a that unlock extra instructions beyond the base architecture). Any gpu_code ending in a fails the same way.

This is not a CUDA/ptxas limitation — nvcc -arch=sm_121a -cubin compiles a valid cubin fine on the same toolchain (see "Other" below) — so sm_121a is a real, supported architecture-conditional target that cuTile's own arch-string parsing simply doesn't handle.

Expected: export_kernel/compile_tile should accept sm_XXXa family-conditional targets the same way nvcc-based CUDA C++ kernels can, so cuTile kernels can opt into family-specific instructions on Hopper/Blackwell-class chips.

I confirmed the same int(sm_arch.removeprefix("sm_")) pattern is still present, unchanged, in the 1.5.0 wheel's cuda/tile/_passes/check_dtype_support.py (downloaded and inspected, not installed) — so this isn't fixed by upgrading from 1.4.0.

Minimum reproducible example

# export_smoke.py (a standalone cuTile kernel export script)
import cuda.tile as ct
from cuda.tile.compilation import (
    ArrayConstraint, CallingConvention, KernelSignature, export_kernel,
)

@ct.kernel
def cutile_smoke_add(lhs, rhs, output):
    block = ct.bid(0)
    lhs_tile = ct.load(lhs, block, 256)
    rhs_tile = ct.load(rhs, block, 256)
    ct.store(output, block, lhs_tile + rhs_tile)

array = ArrayConstraint(
    ct.float32, ndim=1, index_dtype=ct.int32,
    stride_lower_bound_incl=(None,), alias_groups=(), may_alias_internally=False,
    stride_constant=(1,), stride_divisible_by=(1,), shape_divisible_by=(256,),
    base_addr_divisible_by=16,
)
sig = KernelSignature(
    parameters=[array, array, array],
    calling_convention=CallingConvention.cutile_python_v1(),
).with_symbol("cutile_smoke_add")

export_kernel(
    kernel=cutile_smoke_add, signatures=[sig],
    output_file="out.cubin", gpu_code="sm_121a", output_format="cubin",
)

Run: python3 export_smoke.py → raises the traceback below. Swapping gpu_code="sm_121a" for gpu_code="sm_121" (base, no a) compiles and runs correctly, confirming the base-target path works and this is specific to the a suffix.

Relevant log output

Traceback (most recent call last):
  File "export_smoke.py", line 76, in <module>
    raise SystemExit(main())
  File "export_smoke.py", line 65, in main
    export_kernel(
        kernel=cutile_smoke_add, signatures=[...],
        ...
    )
  File ".../cuda/tile/compilation/_export.py", line 57, in export_kernel
    res = compile_tile(kernel._annotated_function, signatures, sm_arch=gpu_code, ...)
  File ".../cuda/tile/_compile.py", line 85, in wrapper
    return func(*args, **kwargs)
  File ".../cuda/tile/_compile.py", line 359, in compile_tile
    bytecode_buf = _get_bytecode(ir_keeper, compiler_options, anonymize_debug_info=False)
  File ".../cuda/tile/_compile.py", line 299, in _get_bytecode
    func_body = ir_keeper.get_final_ir(i)
  File ".../cuda/tile/_compile.py", line 283, in get_final_ir
    check_dtype_support(func_body, self.sm_arch, self.bytecode_version)
  File ".../cuda/tile/_passes/check_dtype_support.py", line 127, in check_dtype_support
    sm_number = int(sm_arch.removeprefix("sm_"))
ValueError: invalid literal for int() with base 10: '121a'

Full env printout

GPU: NVIDIA GB10 (DGX Spark), compute capability 12.1
Platform: aarch64
CUDA Toolkit: 13.3 (nvcc: Cuda compilation tools, release 13.3, V13.3.73, build cuda_13.3.r13.3/compiler.38244171_0)
Python: 3.14
cuda-tile: 1.4.0 (pip)

Other/Misc.

Confirmed sm_121a is a real, ptxas-recognized target independent of cuTile:

$ nvcc -arch=sm_121a -cubin -o probe.cubin probe.cu
$ echo $?
0

We hit this trying to add GB10 (sm_121/sm_121a) coverage to an internal cuTile-embedded-kernel smoke test in rapidsai/cuvs — happy to link the eventual cuvs-side fix here once opened, since it's currently limited to targeting the base sm_121 (which does work) rather than the family-specific sm_121a.