@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 72b1c2ae2c

Fork 0

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 13, 15:45 UTC

[Bugfix] Use latency MOE backend as default for Flashinfer and other misc fixes (#27439)

Signed-off-by: Pavani Majety <pmajety@nvidia.com>

Pavani Majety committed Nov 7, 2025 at 04:18 UTC 72b1c2ae2c2d2782ef1850ca1b96a21f797f7af6

7 files changed +47 -12

csrc/quantization/fp4/nvfp4_quant_kernels.cu

+20 -2

@@ -31,6 +31,13 @@

32 namespace vllm {

34 +template <typename Int>

35 +__host__ __device__ inline Int round_up(Int x, Int y) {

36 + static_assert(std::is_integral_v<Int>,

37 + "round_up argument must be integral type");

38 + return (x + y - 1) / y * y;

39 +}

40 +

41 // Use UE4M3 by default.

42 template <class Type, bool UE8M0_SF = false>

43 __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))

@@ -42,10 +49,21 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))

49 static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,

50 "Vec size is not matched.");

52 + int sf_m = round_up<int>(numRows, 128);

53 + int sf_n_unpadded = numCols / CVT_FP4_SF_VEC_SIZE;

54 + int sf_n_int = round_up<int>(sf_n_unpadded, 4) / 4;

55 + for (int row = numRows + blockIdx.x; row < sf_m; row += gridDim.x) {

56 + // Each thread writes 4 uint32_t elements.

57 + for (int col = sf_n_unpadded + threadIdx.x * 4; col < sf_n_int;

58 + col += blockDim.x * 4) {

59 + SFout[row * sf_n_int + col] = 0x00;

60 + }

61 + }

62 +

63 // Get the global scaling factor, which will be applied to the SF.

64 // Note SFScale is the same as next GEMM's alpha, which is

65 // (448.f / (Alpha_A / 6.f)).

48 - float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[0];

66 + float const global_scale = SFScale == nullptr ? 1.0f : SFScale[0];

68 // Input tensor row/col loops.

69 for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {

@@ -64,7 +82,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))

82 rowIdx, colIdx, numCols, SFout);

84 out_pos =

67 - cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);

85 + cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, global_scale, sf_out);

86 }

87 }

88 }

tests/kernels/quantization/test_nvfp4_quant.py

-2

@@ -168,9 +168,7 @@ def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None:

168 out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)

169

170 out, out_scale = ops.scaled_fp4_quant(x, global_scale)

171 -

171 scale_ans = recover_swizzled_scales(out_scale, m, n)

172 out_ans = cast_from_fp4(out, m, n)

174 -

173 torch.testing.assert_close(out_ans, out_ref)

174 torch.testing.assert_close(scale_ans, scale_ref)

vllm/_custom_ops.py

+1 -1

@@ -1385,7 +1385,7 @@ def scaled_fp4_quant(

1385 rounded_m = round_up(m, 128)

1386 scale_n = n // block_size

1387 rounded_n = round_up(scale_n, 4)

1388 - output_scale = torch.zeros(

1388 + output_scale = torch.empty(

1389 (rounded_m, rounded_n // 4), device=device, dtype=torch.int32

1390 )

1391

vllm/envs.py

+3 -3

@@ -155,7 +155,7 @@ if TYPE_CHECKING:

155 VLLM_USE_FLASHINFER_MOE_FP16: bool = False

156 VLLM_USE_FLASHINFER_MOE_FP8: bool = False

157 VLLM_USE_FLASHINFER_MOE_FP4: bool = False

158 - VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency"] = "throughput"

158 + VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency"] = "latency"

159 VLLM_XGRAMMAR_CACHE_MB: int = 0

160 VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256

161 VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False

@@ -1218,7 +1218,7 @@ environment_variables: dict[str, Callable[[], Any]] = {

1218 # - "latency":

1219 # Uses TensorRT-LLM kernels optimized for low-latency inference.

1220 "VLLM_FLASHINFER_MOE_BACKEND": env_with_choices(

1221 - "VLLM_FLASHINFER_MOE_BACKEND", "throughput", ["throughput", "latency"]

1221 + "VLLM_FLASHINFER_MOE_BACKEND", "latency", ["throughput", "latency"]

1222 ),

1223 # Control the maximum number of tokens per expert supported by the

1224 # NVFP4 MoE CUTLASS Kernel. This value is used to create a buffer for

@@ -1325,7 +1325,7 @@ environment_variables: dict[str, Callable[[], Any]] = {

1325 "VLLM_NVFP4_GEMM_BACKEND": env_with_choices(

1326 "VLLM_NVFP4_GEMM_BACKEND",

1327 None,

1328 - ["flashinfer-cudnn", "flashinfer-trtllm", "flashinfer-cutlass"],

1328 + ["flashinfer-cudnn", "flashinfer-trtllm", "flashinfer-cutlass", "cutlass"],

1329 ),

1330 # Controls garbage collection during CUDA graph capture.

1331 # If set to 0 (default), enables GC freezing to speed up capture time.

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py

@@ -50,6 +50,9 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):

50 elif envs.VLLM_NVFP4_GEMM_BACKEND.startswith("flashinfer-"):

51 self.backend = envs.VLLM_NVFP4_GEMM_BACKEND

52 assert has_flashinfer(), f"FlashInfer is required for {self.backend}"

53 + elif envs.VLLM_NVFP4_GEMM_BACKEND == "cutlass":

54 + self.backend = "cutlass"

55 + assert cutlass_fp4_supported(), f"Cutlass is required for {self.backend}"

57 if self.backend == "none":

58 raise ValueError(

vllm/model_executor/layers/quantization/fp8.py

@@ -138,6 +138,13 @@ def get_fp8_moe_backend(block_quant: bool) -> Fp8MoeBackend:

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

139 return Fp8MoeBackend.FLASHINFER_TRTLLM

140 else:

141 + if block_quant:

142 + raise ValueError(

143 + "FlashInfer FP8 MoE throughput backend does not "

144 + "support block quantization. Please use "

145 + "VLLM_FLASHINFER_MOE_BACKEND=latency "

146 + "instead."

147 + )

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

149 return Fp8MoeBackend.FLASHINFER_CUTLASS

150

vllm/model_executor/layers/quantization/modelopt.py

+13 -4

@@ -221,7 +221,10 @@ class ModelOptFp8Config(QuantizationConfig):

221 def get_quant_method(

222 self, layer: torch.nn.Module, prefix: str

223 ) -> Optional["QuantizeMethodBase"]:

224 - from vllm.attention.layer import Attention # Avoid circular import

224 + from vllm.attention.layer import ( # Avoid circular import

225 + Attention,

226 + MLAAttention,

227 + )

228

229 if isinstance(layer, LinearBase):

230 if self.is_layer_excluded(prefix):

@@ -230,7 +233,7 @@ class ModelOptFp8Config(QuantizationConfig):

233 if "vision_tower" in prefix or "vision_model" in prefix:

234 return UnquantizedLinearMethod()

235 return ModelOptFp8LinearMethod(self)

233 - elif isinstance(layer, Attention):

236 + elif isinstance(layer, (Attention, MLAAttention)):

237 return ModelOptFp8KVCacheMethod(self)

238 elif isinstance(layer, FusedMoE):

239 return ModelOptFp8MoEMethod(self, layer)

@@ -888,7 +891,10 @@ class ModelOptNvFp4Config(QuantizationConfig):

891 def get_quant_method(

892 self, layer: torch.nn.Module, prefix: str

893 ) -> Optional["QuantizeMethodBase"]:

891 - from vllm.attention.layer import Attention # Avoid circular import

894 + from vllm.attention.layer import ( # Avoid circular import

895 + Attention,

896 + MLAAttention,

897 + )

898

899 skip_layer = self.is_layer_excluded(prefix)

900 if isinstance(layer, LinearBase):

@@ -898,7 +904,7 @@ class ModelOptNvFp4Config(QuantizationConfig):

904 if "vision_tower" in prefix or "vision_model" in prefix:

905 return UnquantizedLinearMethod()

906 return ModelOptNvFp4LinearMethod(self)

901 - elif isinstance(layer, Attention):

907 + elif isinstance(layer, (Attention, MLAAttention)):

908 return ModelOptFp8KVCacheMethod(self)

909 elif isinstance(layer, FusedMoE):

910 if skip_layer:

@@ -941,6 +947,9 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

947 elif envs.VLLM_NVFP4_GEMM_BACKEND.startswith("flashinfer-"):

948 self.backend = envs.VLLM_NVFP4_GEMM_BACKEND

949 assert has_flashinfer(), f"FlashInfer is required for {self.backend}"

950 + elif envs.VLLM_NVFP4_GEMM_BACKEND == "cutlass":

951 + self.backend = "cutlass"

952 + assert cutlass_fp4_supported(), f"Cutlass is required for {self.backend}"

953

954 if self.backend == "none":

955 raise ValueError(