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 3, 14:04 UTC
[Bugfix] Fix SM121 (DGX Spark) exclusion from Marlin/CUTLASS FP8 paths (#35568)
Signed-off-by: Blake Ledden <blake@secondnaturecomputing.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Pavani Majety <pmajety@nvidia.com>
Blake Ledden committed
May 15, 2026 at 13:59 UTC
06d020bb6ef70c9549a4e4c29493ebf048c06ee0
8 files changed
+19
-14
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm.cuh
+1
-1
@@ -202,7 +202,7 @@ struct cutlass_3x_gemm_sm120 {
202
sizeof(typename CollectiveEpilogue::SharedStorage))>,
203
KernelSchedule>::CollectiveOp;
204
205
- using GemmKernel = enable_sm120_only<cutlass::gemm::kernel::GemmUniversal<
205
+ using GemmKernel = enable_sm120_family<cutlass::gemm::kernel::GemmUniversal<
206
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
207
};
208
csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh
+1
-1
@@ -72,7 +72,7 @@ struct cutlass_3x_gemm_sm120_custom {
72
sizeof(typename CollectiveEpilogue::SharedStorage))>,
73
KernelSchedule, void>::CollectiveOp;
74
75
- using GemmKernel = enable_sm120_only<cutlass::gemm::kernel::GemmUniversal<
75
+ using GemmKernel = enable_sm120_family<cutlass::gemm::kernel::GemmUniversal<
76
Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue, void>>;
77
};
78
csrc/moe/marlin_moe_wna16/generate_kernels.py
+3
-3
@@ -15,11 +15,11 @@ SUPPORT_SM80 = False
15
for arch in sys.argv[1].split(","):
16
arch = arch[: arch.index(".") + 2].replace(".", "")
17
arch = int(arch)
18
- # only SM89 and SM120 fully support
19
- # mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
18
+ # SM89 and the SM12x family (SM120 RTX 5090, SM121 DGX Spark GB10)
19
+ # fully support mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
20
# SM90 and SM100 can use this PTX, but it’s simulated
21
# with FP16 MMA, so it cannot achieve any acceleration.
22
- if arch in [89, 120]:
22
+ if arch == 89 or arch // 10 == 12:
23
SUPPORT_FP8 = True
24
if arch >= 80:
25
SUPPORT_SM80 = True
csrc/moe/marlin_moe_wna16/ops.cu
+2
-2
@@ -448,8 +448,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* b_bias,
448
"FP8 only support Ada Lovelace or newer GPUs.");
449
TORCH_CHECK(
450
major_capability * 10 + minor_capability == 89 ||
451
- major_capability * 10 + minor_capability == 120,
452
- "Marlin W4A8-FP8 only support SM89 or SM120 device (It is slower than "
451
+ major_capability == 12,
452
+ "Marlin W4A8-FP8 only support SM89 or SM12x device (It is slower than "
453
"Marlin W4A16 on other devices).");
454
}
455
csrc/quantization/marlin/generate_kernels.py
+3
-3
@@ -15,11 +15,11 @@ SUPPORT_SM80 = False
15
for arch in sys.argv[1].split(","):
16
arch = arch[: arch.index(".") + 2].replace(".", "")
17
arch = int(arch)
18
- # only SM89 and SM120 fully support
19
- # mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
18
+ # SM89 and the SM12x family (SM120 RTX 5090, SM121 DGX Spark GB10)
19
+ # fully support mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
20
# SM90 and SM100 can use this PTX, but it’s simulated
21
# with FP16 MMA, so it cannot achieve any acceleration.
22
- if arch in [89, 120]:
22
+ if arch == 89 or arch // 10 == 12:
23
SUPPORT_FP8 = True
24
if arch >= 80:
25
SUPPORT_SM80 = True
tests/kernels/moe/test_moe.py
+5
-1
@@ -746,7 +746,8 @@ def marlin_moe_generate_valid_test_cases():
746
for sub_case in inner_combinations:
747
if (
748
sub_case[0] == scalar_types.float8_e4m3fn
749
- and current_platform.get_device_capability() not in [89, 120]
749
+ and not current_platform.is_device_capability(89)
750
+ and not current_platform.is_device_capability_family(120)
751
):
752
continue
753
@@ -897,6 +898,7 @@ class MarlinMoEWeightData:
898
marlin_moe_generate_valid_test_cases(),
899
)
900
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
901
+@pytest.mark.usefixtures("default_vllm_config")
902
def test_fused_marlin_moe(
903
a_type: ScalarType,
904
b_type: ScalarType,
@@ -1009,6 +1011,7 @@ def test_fused_marlin_moe(
1011
1012
@pytest.mark.flaky(reruns=2)
1013
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
1014
+@pytest.mark.usefixtures("default_vllm_config")
1015
@pytest.mark.parametrize("m", [1, 256])
1016
def test_fused_marlin_moe_with_bias(m):
1017
set_random_seed(0)
@@ -1081,6 +1084,7 @@ def test_fused_marlin_moe_with_bias(m):
1084
1085
@pytest.mark.flaky(reruns=2)
1086
@pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")
1087
+@pytest.mark.usefixtures("default_vllm_config")
1088
@pytest.mark.parametrize("m", [1, 64, 256])
1089
@pytest.mark.parametrize("n,k", [(1024, 1024), (2048, 2048)])
1090
@pytest.mark.parametrize("e,topk", [(8, 2), (64, 4)])
tests/kernels/quantization/test_marlin_gemm.py
+2
-1
@@ -381,7 +381,8 @@ def marlin_generate_valid_test_cases():
381
for sub_case in inner_combinations:
382
if (
383
sub_case[0] == scalar_types.float8_e4m3fn
384
- and current_platform.get_device_capability() not in [89, 120]
384
+ and not current_platform.is_device_capability(89)
385
+ and not current_platform.is_device_capability_family(120)
386
):
387
continue
388
args = sub_case + (size_m, size_n, size_k) + case[4:]
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+2
-2
@@ -480,9 +480,9 @@ def get_marlin_input_dtype(prefix: str | None = None):
480
elif envs.VLLM_MARLIN_INPUT_DTYPE.lower() == "fp8":
481
if not current_platform.is_device_capability(
482
89
483
- ) and not current_platform.is_device_capability(120):
483
+ ) and not current_platform.is_device_capability_family(120):
484
raise ValueError(
485
- "Marlin W4A8-FP8 only support SM89 or SM120 device "
485
+ "Marlin W4A8-FP8 only support SM89 or SM12x device "
486
"(It is slower than Marlin W4A16 on other devices). "
487
"You can consider using W4A8-INT8 instead"
488
"(set VLLM_MARLIN_INPUT_DTYPE=int8)."