b4a2f3ac36 · tduncan/Nvidia-Optimization-Stack-VLLM

CMakeLists.txt

+1 -1

@@ -309,7 +309,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")

309 SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library")

310

311 # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building.

312 - set(CUTLASS_REVISION "v4.2.1")

312 + set(CUTLASS_REVISION "v4.4.2")

313

314 # Use the specified CUTLASS source directory for compilation if VLLM_CUTLASS_SRC_DIR is provided

315 if (DEFINED ENV{VLLM_CUTLASS_SRC_DIR})

csrc/quantization/fp4/nvfp4_quant_entry.cu

+29

@@ -16,6 +16,7 @@

16

17 #include <torch/all.h>

18

19 +#include "cutlass_extensions/common.hpp"

20 #include "nvfp4_utils.cuh"

21

22 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \

@@ -53,12 +54,27 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(

54 torch::Tensor const& output_scale_offset_by_experts);

55 #endif

56

57 +static bool nvfp4_quant_sm_supported() {

58 + const int32_t sm = get_sm_version_num();

59 +#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100

60 + if (sm >= 100 && sm < 120) return true;

61 +#endif

62 +#if defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120

63 + if (sm >= 120 && sm < 130) return true;

64 +#endif

65 + return false;

66 +}

67 +

68 void scaled_fp4_quant_out(torch::Tensor const& input,

69 torch::Tensor const& input_sf,

70 bool is_sf_swizzled_layout, torch::Tensor& output,

71 torch::Tensor& output_sf) {

72 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \

73 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)

74 + TORCH_CHECK(nvfp4_quant_sm_supported(),

75 + "No compiled nvfp4 quantization kernel for SM ",

76 + get_sm_version_num(),

77 + ". Recompile with the appropriate CUDA arch.");

78 return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,

79 is_sf_swizzled_layout);

80 #endif

@@ -100,6 +116,10 @@ void scaled_fp4_experts_quant(

116 torch::Tensor const& output_scale_offset_by_experts) {

117 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \

118 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)

119 + TORCH_CHECK(nvfp4_quant_sm_supported(),

120 + "No compiled nvfp4 experts quantization kernel for SM ",

121 + get_sm_version_num(),

122 + ". Recompile with the appropriate CUDA arch.");

123 return scaled_fp4_experts_quant_sm1xxa(

124 output, output_scale, input, input_global_scale, input_offset_by_experts,

125 output_scale_offset_by_experts);

@@ -112,6 +132,10 @@ void silu_and_mul_nvfp4_quant(torch::Tensor& output, torch::Tensor& output_sf,

132 torch::Tensor& input, torch::Tensor& input_sf) {

133 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \

134 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)

135 + TORCH_CHECK(nvfp4_quant_sm_supported(),

136 + "No compiled silu_and_mul nvfp4 quantization kernel for SM ",

137 + get_sm_version_num(),

138 + ". Recompile with the appropriate CUDA arch.");

139 return silu_and_mul_nvfp4_quant_sm1xxa(output, output_sf, input, input_sf);

140 #endif

141 TORCH_CHECK_NOT_IMPLEMENTED(

@@ -125,6 +149,11 @@ void silu_and_mul_scaled_fp4_experts_quant(

149 torch::Tensor const& output_scale_offset_by_experts) {

150 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \

151 (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)

152 + TORCH_CHECK(nvfp4_quant_sm_supported(),

153 + "No compiled silu_and_mul nvfp4 experts quantization kernel "

154 + "for SM ",

155 + get_sm_version_num(),

156 + ". Recompile with the appropriate CUDA arch.");

157 return silu_and_mul_scaled_fp4_experts_quant_sm1xxa(

158 output, output_scale, input, input_global_scale, input_offset_by_experts,

159 output_scale_offset_by_experts);

csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu

+13 -1

@@ -63,5 +63,17 @@ void cutlass_scaled_fp4_mm(torch::Tensor& D, const torch::Tensor& A,

63 bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability) {

64 int runtimeVersion;

65 cudaRuntimeGetVersion(&runtimeVersion);

66 - return cuda_device_capability >= 100 && runtimeVersion >= 12080;

66 + if (runtimeVersion < 12080) return false;

67 + // Only report support when the SM-specific kernel was actually compiled in,

68 + // so the Python-side backend selector does not choose CUTLASS and then hit

69 + // TORCH_CHECK_NOT_IMPLEMENTED (or worse, fall through to Marlin).

70 +#if defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100

71 + if (cuda_device_capability >= 100 && cuda_device_capability < 120)

72 + return true;

73 +#endif

74 +#if defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120

75 + if (cuda_device_capability >= 120 && cuda_device_capability < 130)

76 + return true;

77 +#endif

78 + return false;

79 }

csrc/quantization/machete/machete_mainloop.cuh

+1

@@ -154,6 +154,7 @@ struct MacheteCollectiveMma {

154 struct DispatchPolicy {

155 constexpr static int Stages = PipelineStages;

156 using ClusterShape = ClusterShape_MNK;

157 + using ArchTag = arch::Sm90;

158 using Schedule = KernelScheduleType;

159 };

160

docker/Dockerfile

+4 -1

@@ -590,7 +590,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \

590 # Install FlashInfer JIT cache (requires CUDA-version-specific index URL)

591 # https://docs.flashinfer.ai/installation.html

592 # From versions.json: .flashinfer.version

593 -ARG FLASHINFER_VERSION=0.6.6

593 +# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798)

594 +# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on

595 +# SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738

596 +ARG FLASHINFER_VERSION=0.6.7

597 RUN --mount=type=cache,target=/root/.cache/uv \

598 uv pip install --system flashinfer-jit-cache==${FLASHINFER_VERSION} \

599 --extra-index-url https://flashinfer.ai/whl/cu$(echo $CUDA_VERSION | cut -d. -f1,2 | tr -d '.') \

docker/Dockerfile.nightly_torch

+5 -2

@@ -217,13 +217,16 @@ RUN pip install setuptools==75.6.0 packaging==23.2 ninja==1.11.1.3 build==1.2.2.

217

218

219 # build flashinfer for torch nightly from source around 10 mins

220 -# release version: v0.6.6

220 +# release version: v0.6.7

221 +# 0.6.7: CUTLASS 4.4.2 bump, fixes TMA grouped GEMM on SM12x (flashinfer#2798)

222 +# TODO: bump to 0.6.8 when released for NVFP4/MXFP4 group GEMMs on

223 +# SM120/SM121 (RTX 50 / DGX Spark) via flashinfer#2738

224 # todo(elainewy): cache flashinfer build result for faster build

225 ENV CCACHE_DIR=/root/.cache/ccache

226 RUN --mount=type=cache,target=/root/.cache/ccache \

227 --mount=type=cache,target=/root/.cache/uv \

228 echo "git clone flashinfer..." \

226 - && git clone --depth 1 --branch v0.6.6 --recursive https://github.com/flashinfer-ai/flashinfer.git \

229 + && git clone --depth 1 --branch v0.6.7 --recursive https://github.com/flashinfer-ai/flashinfer.git \

230 && cd flashinfer \

231 && git submodule update --init --recursive \

232 && echo "finish git clone flashinfer..." \

docker/versions.json

+1 -1

@@ -68,7 +68,7 @@

68 "default": "true"

69 },

70 "FLASHINFER_VERSION": {

71 - "default": "0.6.6"

71 + "default": "0.6.7"

72 },

73 "GDRCOPY_CUDA_VERSION": {

74 "default": "12.8"

requirements/cuda.txt

+2 -2

@@ -9,8 +9,8 @@ torchaudio==2.10.0

9 # These must be updated alongside torch

10 torchvision==0.25.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version

11 # FlashInfer should be updated together with the Dockerfile

12 -flashinfer-python==0.6.6

13 -flashinfer-cubin==0.6.6

12 +flashinfer-python==0.6.7

13 +flashinfer-cubin==0.6.7

14 # Cap nvidia-cudnn-frontend (transitive dep of flashinfer) due to

15 # breaking changes in 1.19.0

16 nvidia-cudnn-frontend>=1.13.0,<1.19.0

tests/kernels/moe/test_unquantized_backend_selection.py

-4

@@ -57,7 +57,6 @@ def test_select_default_backend_by_platform(

57 moe_config = make_dummy_moe_config()

58 selected_backend = select_unquantized_moe_backend(

59 moe_config=moe_config,

60 - use_ep=False,

60 use_dp=False,

61 )

62

@@ -90,7 +89,6 @@ def test_select_rocm_aiter_backend(mock_aiter_enabled, mock_has_flashinfer):

89 moe_config = make_dummy_moe_config()

90 selected_backend = select_unquantized_moe_backend(

91 moe_config=moe_config,

93 - use_ep=False,

92 use_dp=False,

93 )

94

@@ -129,7 +127,6 @@ def test_select_cuda_flashinfer_trtllm_backend(

127

128 selected_backend = select_unquantized_moe_backend(

129 moe_config=moe_config,

132 - use_ep=True,

130 use_dp=False,

131 )

132

@@ -171,7 +168,6 @@ def test_select_cuda_flashinfer_cutlass_backend(

168

169 selected_backend = select_unquantized_moe_backend(

170 moe_config=moe_config,

174 - use_ep=True, # CUTLASS requires EP

171 use_dp=False, # CUTLASS doesn't support DP

172 )

173

vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py

+5

@@ -406,6 +406,11 @@ class TrtLlmFp8ExpertsMonolithic(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsMonolit

406 if self.routing_method_type == RoutingMethodType.DeepSeekV3:

407 router_logits = router_logits.to(torch.float32)

408

409 + # Currently FI requires bfloat16 routing bias.

410 + # https://github.com/flashinfer-ai/flashinfer/issues/2909

411 + if e_score_correction_bias is not None:

412 + e_score_correction_bias = e_score_correction_bias.to(torch.bfloat16)

413 +

414 out = flashinfer.fused_moe.trtllm_fp8_per_tensor_scale_moe(

415 routing_logits=router_logits,

416 routing_bias=e_score_correction_bias,

vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py

+8

@@ -5,6 +5,7 @@ import flashinfer

5 import torch

6

7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

8 +from vllm.logger import init_logger

9 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

10 from vllm.model_executor.layers.fused_moe.config import (

11 FusedMoEConfig,

@@ -27,6 +28,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

28 from vllm.platforms import current_platform

29 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe

30

31 +logger = init_logger(__name__)

32 +

33

34 class TrtLlmNvFp4ExpertsBase:

35 """

@@ -315,6 +318,11 @@ class TrtLlmNvFp4ExpertsMonolithic(

318 else router_logits

319 )

320

321 + # Currently FI requires bfloat16 routing bias.

322 + # https://github.com/flashinfer-ai/flashinfer/issues/2909

323 + if e_score_correction_bias is not None:

324 + e_score_correction_bias = e_score_correction_bias.to(torch.bfloat16)

325 +

326 # Invoke kernel.

327 return flashinfer.fused_moe.trtllm_fp4_block_scale_moe(

328 routing_logits=router_logits,

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

+1 -1

@@ -361,7 +361,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):

361 fc1_expert_weights = w1

362 fc2_expert_weights = w2

363 else:

364 - quant_scales = None

364 + quant_scales = []

365 a1q_scale = None

366 fc1_expert_weights = w1

367 fc2_expert_weights = w2

vllm/model_executor/layers/fused_moe/oracle/unquantized.py

+2 -4

@@ -70,7 +70,6 @@ def map_unquantized_backend(runner_backend: MoEBackend) -> UnquantizedMoeBackend

70

71 def select_unquantized_moe_backend(

72 moe_config: FusedMoEConfig,

73 - use_ep: bool,

73 use_dp: bool,

74 ) -> UnquantizedMoeBackend:

75 """

@@ -96,7 +95,6 @@ def select_unquantized_moe_backend(

95 # FlashInfer CUTLASS MoE is only supported on Hopper and later GPUS

96 flashinfer_cutlass_available = (

97 has_flashinfer_cutlass_fused_moe()

99 - and use_ep

98 and (not use_dp)

99 and current_platform.has_device_capability(90)

100 )

@@ -161,9 +159,9 @@ def select_unquantized_moe_backend(

159 "to enable it for better performance.",

160 scope="local",

161 )

164 - elif use_ep and (not use_dp):

162 + elif not use_dp and flashinfer_cutlass_available:

163 logger.info_once(

166 - "FlashInfer MoE is available for EP"

164 + "FlashInfer CUTLASS MoE is available"

165 " but not enabled, consider setting"

166 " VLLM_USE_FLASHINFER_MOE_FP16=1 to enable it.",

167 scope="local",

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

-1

@@ -61,7 +61,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

61 super().__init__(moe)

62 self.unquantized_backend = select_unquantized_moe_backend(

63 moe_config=self.moe,

64 - use_ep=self.moe.moe_parallel_config.use_ep,

64 use_dp=self.moe.moe_parallel_config.dp_size > 1,

65 )

66

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

+14 -2

@@ -55,8 +55,16 @@ def select_nvfp4_linear_backend() -> NvFp4LinearBackend:

55 elif envs.VLLM_USE_NVFP4_CT_EMULATIONS:

56 backend = NvFp4LinearBackend.EMULATION

57 elif envs.VLLM_NVFP4_GEMM_BACKEND is None:

58 - # Auto-select best available backend

59 - if current_platform.has_device_capability(100) and has_flashinfer():

58 + # Auto-select best available backend.

59 + # cutlass_fp4_supported() checks that the vLLM NVFP4 kernels (both

60 + # quantization and GEMM) were compiled for the current SM version.

61 + # FlashInfer backends still rely on the vLLM quantization kernels,

62 + # so we gate them on the same check.

63 + if (

64 + cutlass_fp4_supported()

65 + and current_platform.has_device_capability(100)

66 + and has_flashinfer()

67 + ):

68 backend = NvFp4LinearBackend.FLASHINFER_CUTLASS

69 elif cutlass_fp4_supported():

70 backend = NvFp4LinearBackend.VLLM_CUTLASS

@@ -72,6 +80,10 @@ def select_nvfp4_linear_backend() -> NvFp4LinearBackend:

80 NvFp4LinearBackend.FLASHINFER_CUDNN,

81 ):

82 assert has_flashinfer(), f"FlashInfer is required for {backend}"

83 + assert cutlass_fp4_supported(), (

84 + f"{backend} requires vLLM NVFP4 quantization kernels compiled "

85 + f"for the current GPU (SM {current_platform.get_device_capability()})"

86 + )

87 elif backend == NvFp4LinearBackend.VLLM_CUTLASS:

88 assert cutlass_fp4_supported(), f"Cutlass is required for {backend}"

89 elif backend == NvFp4LinearBackend.MARLIN:

[NVIDIA] Bugfix NVFP4 DGX Spark and RTX50 (#38423)