4fa7ce46f3 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/compile/distributed/test_fusions_e2e.py

+1 -1

@@ -20,7 +20,7 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer

20

21 from ...utils import flat_product, multi_gpu_test

22

23 -is_blackwell = lambda: current_platform.is_device_capability(100)

23 +is_blackwell = lambda: current_platform.is_device_capability_family(100)

24 """Are we running on Blackwell, a lot of tests depend on it"""

25

26

tests/kernels/attention/test_cutlass_mla_decode.py

+2 -2

@@ -32,8 +32,8 @@ def cal_diff(

32

33

34 CUTLASS_MLA_UNSUPPORTED_REASON = (

35 - "Cutlass MLA Requires compute capability of 10 or above."

36 - if not current_platform.is_device_capability(100)

35 + "Cutlass MLA Requires compute capability of 100 or above."

36 + if not current_platform.is_device_capability_family(100)

37 else "Cutlass MLA is supported"

38 )

39

tests/kernels/attention/test_flashinfer_trtllm_attention.py

+2 -2

@@ -11,7 +11,7 @@ from tests.kernels.quantization.nvfp4_utils import (

11 from vllm.platforms import current_platform

12 from vllm.utils.math_utils import round_up

13

14 -if not current_platform.is_device_capability(100):

14 +if not current_platform.is_device_capability_family(100):

15 pytest.skip(

16 "This TRTLLM kernel requires NVIDIA Blackwell.", allow_module_level=True

17 )

@@ -443,7 +443,7 @@ def test_flashinfer_trtllm_prefill_with_baseline(

443 output_trtllm = output_trtllm.reshape(-1, query.shape[1], query.shape[2])

444

445 if q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP4_DTYPE:

446 - rtol, atol = 1e-1, 2e-1

446 + rtol, atol = 3e-1, 4e-1

447 elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == FP8_DTYPE:

448 rtol, atol = 4e-2, 6e-2

449 elif q_quant_dtype == FP8_DTYPE and o_quant_dtype == dtype:

tests/kernels/moe/test_ocp_mx_moe.py

+2 -2

@@ -17,7 +17,7 @@ QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse(

17 ) >= version.parse("0.8.99")

18

19 TRTLLM_GEN_MXFP4_AVAILABLE = (

20 - current_platform.is_cuda() and current_platform.is_device_capability(100)

20 + current_platform.is_cuda() and current_platform.is_device_capability_family(100)

21 )

22

23 HOPPER_MXFP4_BF16_AVAILABLE = (

@@ -799,7 +799,7 @@ def test_flashinfer_cutlass_mxfp4_fused_moe(

799 @pytest.mark.skipif(

800 not (

801 current_platform.is_cuda()

802 - and current_platform.is_device_capability(100)

802 + and current_platform.is_device_capability_family(100)

803 and has_flashinfer()

804 ),

805 reason="NVIDIA GPU sm100 and flashinfer are required for this test",

tests/quantization/test_blackwell_moe.py

+2 -2

@@ -10,9 +10,9 @@ import pytest

10 from tests.utils import RemoteOpenAIServer

11 from vllm.platforms import current_platform

12

13 -if not current_platform.is_device_capability(100):

13 +if not current_platform.is_device_capability_family(100):

14 pytest.skip(

15 - "This test only runs on Blackwell GPUs (SM100).", allow_module_level=True

15 + "This test only runs on Blackwell GPUs (SM10x).", allow_module_level=True

16 )

17

18

vllm/model_executor/layers/batch_invariant.py

+1 -1

@@ -936,7 +936,7 @@ def enable_batch_invariant_mode():

936 # Batch invariant matmuls are no longer needed after cublas overrides

937 if not is_torch_equal_or_newer("2.10.0.dev"):

938 if (

939 - current_platform.is_device_capability(100)

939 + current_platform.is_device_capability_family(100)

940 or current_platform.is_device_capability(80)

941 or current_platform.is_device_capability(89)

942 ):

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

+4 -1

@@ -287,7 +287,10 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

287 """

288 DeepGemm supports packed ue8m0 activation scales format in devices == sm100

289 """

290 - return is_deep_gemm_e8m0_used() and current_platform.is_device_capability(100)

290 + return (

291 + is_deep_gemm_e8m0_used()

292 + and current_platform.is_device_capability_family(100)

293 + )

294

295 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:

296 # Let PrepareAndFinalize::finalize() decide the impl.

vllm/model_executor/layers/quantization/fp8.py

+3 -3

@@ -137,7 +137,7 @@ def get_fp8_moe_backend(

137 if (

138 current_platform.is_cuda()

139 and (

140 - current_platform.is_device_capability(100)

140 + current_platform.is_device_capability_family(100)

141 or current_platform.is_device_capability(90)

142 )

143 and envs.VLLM_USE_FLASHINFER_MOE_FP8

@@ -148,7 +148,7 @@ def get_fp8_moe_backend(

148 logger.info_once("Using FlashInfer FP8 MoE TRTLLM backend for SM100")

149 return Fp8MoeBackend.FLASHINFER_TRTLLM

150 else:

151 - if block_quant and current_platform.is_device_capability(100):

151 + if block_quant and current_platform.is_device_capability_family(100):

152 raise ValueError(

153 "FlashInfer FP8 MoE throughput backend does not "

154 "support block quantization. Please use "

@@ -193,7 +193,7 @@ def get_fp8_moe_backend(

193 # CUTLASS BlockScaled GroupedGemm on SM100 with block-quantized weights

194 if (

195 current_platform.is_cuda()

196 - and current_platform.is_device_capability(100)

196 + and current_platform.is_device_capability_family(100)

197 and block_quant

198 ):

199 logger.info_once(

vllm/model_executor/layers/quantization/mxfp4.py

+4 -4

@@ -118,19 +118,19 @@ def get_mxfp4_backend(with_lora_support: bool) -> Mxfp4Backend:

118 logger.info_once("Using FlashInfer MXFP4 BF16 backend for SM90")

119 return Mxfp4Backend.SM90_FI_MXFP4_BF16

120 elif (

121 - current_platform.is_device_capability(100)

121 + current_platform.is_device_capability_family(100)

122 and has_flashinfer()

123 and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS

124 ):

125 logger.info_once("Using FlashInfer MXFP4 MXFP8 CUTLASS backend for SM100")

126 return Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS

127 elif (

128 - current_platform.is_device_capability(100)

128 + current_platform.is_device_capability_family(100)

129 and has_flashinfer()

130 and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8

131 ):

132 return Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM

133 - elif current_platform.is_device_capability(100) and has_flashinfer():

133 + elif current_platform.is_device_capability_family(100) and has_flashinfer():

134 logger.info_once(

135 "Using FlashInfer MXFP4 BF16 backend for SM100, "

136 "For faster performance on SM100, consider setting "

@@ -139,7 +139,7 @@ def get_mxfp4_backend(with_lora_support: bool) -> Mxfp4Backend:

139 )

140 return Mxfp4Backend.SM100_FI_MXFP4_BF16

141 elif (

142 - current_platform.is_device_capability(100)

142 + current_platform.is_device_capability_family(100)

143 or current_platform.is_device_capability(90)

144 ) and not has_flashinfer():

145 logger.warning_once(

vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py

+1 -1

@@ -50,7 +50,7 @@ def is_flashinfer_fp4_cutedsl_moe_available() -> bool:

50 envs.VLLM_USE_FLASHINFER_MOE_FP4

51 and has_flashinfer_cutedsl_grouped_gemm_nt_masked()

52 and current_platform.is_cuda()

53 - and current_platform.is_device_capability(100)

53 + and current_platform.is_device_capability_family(100)

54 )

55

56

vllm/model_executor/layers/quantization/utils/flashinfer_utils.py

+1 -1

@@ -290,7 +290,7 @@ def get_flashinfer_moe_backend() -> FlashinferMoeBackend:

290 if flashinfer_moe_backend in backend_map:

291 if (

292 flashinfer_moe_backend == "latency"

293 - and not current_platform.has_device_capability(100)

293 + and not current_platform.is_device_capability_family(100)

294 ):

295 logger.info_once(

296 "Flashinfer TRTLLM MOE backend is only supported on "

vllm/model_executor/layers/quantization/utils/fp8_utils.py

+1 -1

@@ -247,7 +247,7 @@ class W8A8BlockFp8LinearOp:

247 self.act_quant_group_shape = act_quant_group_shape

248 self.is_deep_gemm_supported = is_deep_gemm_supported()

249 self.is_hopper = current_platform.is_device_capability(90)

250 - self.is_blackwell = current_platform.is_device_capability(100)

250 + self.is_blackwell = current_platform.is_device_capability_family(100)

251 self.use_deep_gemm_e8m0 = is_deep_gemm_e8m0_used()

252

253 # Get the correct blockscale mul and input quant operations.

vllm/model_executor/layers/quantization/utils/mxfp4_utils.py

+1 -1

@@ -63,7 +63,7 @@ def _swizzle_mxfp4(quant_tensor, scale, num_warps):

63 "split_k": 1,

64 }

65 opt_flags.update_opt_flags_constraints(constraints)

66 - elif current_platform.is_device_capability(100):

66 + elif current_platform.is_device_capability_family(100):

67 constraints = {

68 "is_persistent": True,

69 "epilogue_subtile": 1,

vllm/model_executor/models/config.py

+1 -1

@@ -363,7 +363,7 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):

363 else:

364 kernel_block_alignment_size = 16

365 if (

366 - current_platform.is_device_capability(100)

366 + current_platform.is_device_capability_family(100)

367 and model_config.get_head_size() == 256

368 and (

369 attention_config.backend is None

vllm/platforms/cuda.py

+1 -1

@@ -182,7 +182,7 @@ class CudaPlatformBase(Platform):

182

183 if vllm_config.attention_config.backend is None:

184 # Default case

185 - if cls.is_device_capability(100) and not use_sparse:

185 + if cls.is_device_capability_family(100) and not use_sparse:

186 # Blackwell => Force CutlassMLA (unless sparse, i.e. DSv3.2).

187 use_cutlass_mla = True

188 # Set the backend in AttentionConfig so it's used during

vllm/platforms/interface.py

+15

@@ -301,6 +301,21 @@ class Platform:

301

302 return current_capability.to_int() == capability

303

304 + @classmethod

305 + def is_device_capability_family(

306 + cls,

307 + capability: int,

308 + device_id: int = 0,

309 + ) -> bool:

310 + """

311 + Returns True if the device capability is any <major>.x.

312 + Mirrors CUDA 13 'family' architecture semantics (e.g. 10.x, 11.x, 12.x).

313 + """

314 + current_capability = cls.get_device_capability(device_id=device_id)

315 + if current_capability is None:

316 + return False

317 + return (current_capability.to_int() // 10) == (capability // 10)

318 +

319 @classmethod

320 def get_device_name(cls, device_id: int = 0) -> str:

321 """Get the name of a device."""

vllm/utils/deep_gemm.py

+2 -2

@@ -38,7 +38,7 @@ class DeepGemmQuantScaleFMT(Enum):

38 return DeepGemmQuantScaleFMT.FLOAT32

39 return (

40 DeepGemmQuantScaleFMT.UE8M0

41 - if current_platform.is_device_capability(100)

41 + if current_platform.is_device_capability_family(100)

42 else DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0

43 )

44

@@ -50,7 +50,7 @@ def is_deep_gemm_supported() -> bool:

50 """

51 is_supported_arch = current_platform.is_cuda() and (

52 current_platform.is_device_capability(90)

53 - or current_platform.is_device_capability(100)

53 + or current_platform.is_device_capability_family(100)

54 )

55 return envs.VLLM_USE_DEEP_GEMM and has_deep_gemm() and is_supported_arch

56

vllm/utils/flashinfer.py

+3 -1

@@ -264,7 +264,9 @@ def supports_trtllm_attention() -> bool:

264 return False

265

266 # Requires SM100 and NVIDIA artifactory to be accessible to download cubins

267 - return current_platform.is_device_capability(100) and has_nvidia_artifactory()

267 + return (

268 + current_platform.is_device_capability_family(100) and has_nvidia_artifactory()

269 + )

270

271

272 def force_use_trtllm_attention() -> bool | None:

vllm/v1/attention/backends/flashinfer.py

+1 -1

@@ -564,7 +564,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

564 )

565 self.paged_kv_last_page_len_np = self.paged_kv_last_page_len_cpu.numpy()

566

567 - if self.head_dim == 256 and current_platform.is_device_capability(100):

567 + if self.head_dim == 256 and current_platform.is_device_capability_family(100):

568 # https://github.com/flashinfer-ai/flashinfer/issues/1993 reports that

569 # head size 256 and block size 16 is not supported on blackwell.

570 assert kv_cache_spec.block_size != 16, (

vllm/v1/attention/backends/mla/common.py

+3 -3

@@ -446,7 +446,7 @@ def use_flashinfer_prefill() -> bool:

446 and flashinfer_available

447 and not vllm_config.attention_config.use_cudnn_prefill

448 and not vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill

449 - and current_platform.is_device_capability(100)

449 + and current_platform.is_device_capability_family(100)

450 )

451

452

@@ -457,7 +457,7 @@ def use_cudnn_prefill() -> bool:

457 return (

458 flashinfer_available

459 and vllm_config.attention_config.use_cudnn_prefill

460 - and current_platform.is_device_capability(100)

460 + and current_platform.is_device_capability_family(100)

461 and has_nvidia_artifactory()

462 )

463

@@ -470,7 +470,7 @@ def use_trtllm_ragged_deepseek_prefill() -> bool:

470 return (

471 flashinfer_available

472 and vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill

473 - and current_platform.is_device_capability(100)

473 + and current_platform.is_device_capability_family(100)

474 )

475

476

vllm/v1/attention/backends/mla/flashmla_sparse.py

+2 -2

@@ -420,7 +420,7 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

420 max_num_sm_parts = int(

421 max((sm_count // 2) / h_k // (cdiv(h_q // h_k, 2 * 64) * s_q), 1)

422 )

423 - if current_platform.is_device_capability(100):

423 + if current_platform.is_device_capability_family(100):

424 max_num_sm_parts *= 2

425 self.tile_scheduler_metadata_buffer = torch.empty(

426 # TileSchedulerMetaDataSize = 8

@@ -719,7 +719,7 @@ class FlashMLASparseImpl(MLACommonBaseImpl[FlashMLASparseMetadata]):

719 self.softmax_scale = scale

720 assert indexer is not None

721 self.topk_indices_buffer = indexer.topk_indices_buffer

722 - self.padding = 128 if current_platform.is_device_capability(100) else 64

722 + self.padding = 128 if current_platform.is_device_capability_family(100) else 64

723

724 if kv_cache_dtype == "fp8_ds_mla":

725 # Reserve workspace during initialization

[Feature] Add SM103 (Blackwell Ultra) Support to vLLM (#30484)