A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 17:45 UTC
[Kernel] Add enable_sm120_or_later for SM121 (DGX Spark) CUTLASS support (#33517)
Signed-off-by: code4me2 <velvetmoon222999@gmail.com>
Vel committed
Feb 6, 2026 at 20:28 UTC
bc32444b238d2ec3726f599cf3fc67dbaf51a6c6
2 files changed
+13
-1
csrc/cutlass_extensions/common.hpp
+11
@@ -152,3 +152,14 @@ struct enable_sm120_only : Kernel {
152
#endif
153
}
154
};
155
+
156
+// SM12x family includes SM120 (RTX 5090) and SM121 (DGX Spark GB10)
157
+template <typename Kernel>
158
+struct enable_sm120_family : Kernel {
159
+ template <typename... Args>
160
+ CUTLASS_DEVICE void operator()(Args&&... args) {
161
+#if defined __CUDA_ARCH__ && (__CUDA_ARCH__ >= 1200 && __CUDA_ARCH__ < 1300)
162
+ Kernel::operator()(std::forward<Args>(args)...);
163
+#endif
164
+ }
165
+};
csrc/quantization/w8a8/cutlass/c3x/scaled_mm_blockwise_sm120_fp8_dispatch.cuh
+2
-1
@@ -103,7 +103,8 @@ struct cutlass_3x_gemm_fp8_blockwise {
103
MainloopScheduler
104
>::CollectiveOp;
105
106
- using KernelType = enable_sm120_only<cutlass::gemm::kernel::GemmUniversal<
106
+ // SM12x family to support both SM120 (RTX 5090) and SM121 (DGX Spark)
107
+ using KernelType = enable_sm120_family<cutlass::gemm::kernel::GemmUniversal<
108
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue>>;
109
110
struct GemmKernel : public KernelType {};