[Kernel] Add enable_sm120_or_later for SM121 (DGX Spark) CUTLASS support (#33517)

Signed-off-by: code4me2 <velvetmoon222999@gmail.com>
This commit is contained in:
Vel
2026-02-06 20:28:01 -08:00
committed by GitHub
parent 18e8545297
commit bc32444b23
2 changed files with 13 additions and 1 deletions

View File

@@ -152,3 +152,14 @@ struct enable_sm120_only : Kernel {
#endif
}
};
// SM12x family includes SM120 (RTX 5090) and SM121 (DGX Spark GB10)
template <typename Kernel>
struct enable_sm120_family : Kernel {
template <typename... Args>
CUTLASS_DEVICE void operator()(Args&&... args) {
#if defined __CUDA_ARCH__ && (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)
Kernel::operator()(std::forward<Args>(args)...);
#endif
}
};

View File

@@ -103,7 +103,8 @@ struct cutlass_3x_gemm_fp8_blockwise {
MainloopScheduler
>::CollectiveOp;
using KernelType = enable_sm120_only<cutlass::gemm::kernel::GemmUniversal<
// SM12x family to support both SM120 (RTX 5090) and SM121 (DGX Spark)
using KernelType = enable_sm120_family<cutlass::gemm::kernel::GemmUniversal<
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue>>;
struct GemmKernel : public KernelType {};