@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 279a5f31b3

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 14, 05:30 UTC

[Kernel] Add nvfp4 gemm flashinfer backends (#22346)

Signed-off-by: Julien Lin <jullin@nvidia.com> Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mgoin <mgoin64@gmail.com>

nvjullin committed Aug 15, 2025 at 04:03 UTC 279a5f31b3faa6f40759516efa5c742f637ab8b7

9 files changed +371 -41

.buildkite/test-pipeline.yaml

@@ -669,6 +669,7 @@ steps:

669 - pytest -v -s tests/kernels/quantization/test_cutlass_scaled_mm.py -k 'fp8'

670 - pytest -v -s tests/kernels/quantization/test_nvfp4_quant.py

671 - pytest -v -s tests/kernels/quantization/test_nvfp4_scaled_mm.py

672 + - pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py

673 - pytest -v -s tests/kernels/moe/test_nvfp4_moe.py

674 # Fusion

675 - pytest -v -s tests/compile/test_fusion_all_reduce.py

tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py new

+139

@@ -0,0 +1,139 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import pytest

4 +import torch

5 +from nvfp4_utils import (FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX,

6 + convert_swizzled_to_linear, dequantize_nvfp4_to_dtype)

7 +

8 +from vllm import _custom_ops as ops

9 +from vllm.platforms import current_platform

10 +from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm

11 +

12 +if not current_platform.has_device_capability(100):

13 + pytest.skip(

14 + reason="Nvfp4 Requires compute capability of 10 or above.",

15 + allow_module_level=True,

16 + )

17 +

18 +DTYPES = [torch.float16, torch.bfloat16]

19 +# m, n, k

20 +SHAPES = [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)]

21 +PAD_SHAPES = [(150, 128, 64), (128, 128, 96)]

22 +SHAPES.extend(PAD_SHAPES)

23 +

24 +SEEDS = [42]

25 +CUDA_DEVICES = ["cuda:0"]

26 +

27 +

28 +def get_ref_results(

29 + a_fp4,

30 + b_fp4,

31 + a_sf,

32 + b_sf,

33 + a_global_scale,

34 + b_global_scale,

35 + m,

36 + n,

37 + dtype,

38 + block_size,

39 + device,

40 +):

41 + _, m_k = a_fp4.shape

42 + _, n_k = b_fp4.shape

43 + assert m_k == n_k

44 + a_in_dtype = dequantize_nvfp4_to_dtype(a_fp4,

45 + a_sf,

46 + a_global_scale,

47 + dtype=dtype,

48 + device=device,

49 + block_size=block_size)

50 + b_in_dtype = dequantize_nvfp4_to_dtype(b_fp4,

51 + b_sf,

52 + b_global_scale,

53 + dtype=dtype,

54 + device=device,

55 + block_size=block_size)

56 + return torch.matmul(a_in_dtype, b_in_dtype.t())

57 +

58 +

59 +@pytest.mark.parametrize("dtype", DTYPES)

60 +@pytest.mark.parametrize("shape", SHAPES)

61 +@pytest.mark.parametrize("seed", SEEDS)

62 +@pytest.mark.parametrize("device", CUDA_DEVICES)

63 +@pytest.mark.parametrize("backend", ["cutlass", "trtllm"])

64 +@pytest.mark.parametrize("autotune", [False, True])

65 +@torch.inference_mode()

66 +def test_flashinfer_nvfp4_gemm(

67 + dtype: torch.dtype,

68 + shape: tuple[int, int, int],

69 + seed: int,

70 + device: str,

71 + backend: str,

72 + autotune: bool,

73 +) -> None:

74 + if backend == "trtllm" and dtype == torch.float16:

75 + pytest.skip(

76 + "Only torch.bfloat16 is supported for TRTLLM FP4 GEMM operations")

77 +

78 + current_platform.seed_everything(seed)

79 + m, n, packed_k = shape

80 + k = packed_k * 2

81 + block_size = 16

82 + a_dtype = torch.randn((m, k), dtype=dtype, device=device)

83 + b_dtype = torch.randn((n, k), dtype=dtype, device=device)

84 +

85 + a_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) /

86 + torch.amax(a_dtype.flatten(), dim=-1)).to(torch.float32)

87 + b_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) /

88 + torch.amax(b_dtype.flatten(), dim=-1)).to(torch.float32)

89 + alpha = 1.0 / (a_global_scale * b_global_scale)

90 + # ops.scaled_fp4_quant returns swizzled scales, while weights

91 + # from checkpoints are in linear scales.

92 + # So instead of needing to swizzle for cutlass as in modelopt.py,

93 + # we need to unswizzle for trtllm here.

94 + a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a_dtype, a_global_scale)

95 + b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(b_dtype, b_global_scale)

96 +

97 + # get_ref_results unswizzles the scales internally.

98 + expected_out = get_ref_results(

99 + a_fp4,

100 + b_fp4,

101 + a_scale_interleaved,

102 + b_scale_interleaved,

103 + a_global_scale,

104 + b_global_scale,

105 + m,

106 + n,

107 + dtype,

108 + block_size,

109 + device,

110 + )

111 +

112 + import flashinfer

113 +

114 + if backend == "trtllm":

115 + epilogue_tile_m = 128

116 + b_fp4 = flashinfer.shuffle_matrix_a(b_fp4.view(torch.uint8),

117 + epilogue_tile_m)

118 +

119 + b_scale_interleaved = convert_swizzled_to_linear(

120 + b_scale_interleaved, n, k, block_size)

121 + b_scale_interleaved = (flashinfer.shuffle_matrix_sf_a(

122 + b_scale_interleaved.view(torch.uint8), epilogue_tile_m).reshape(

123 + b_scale_interleaved.shape).view(torch.float8_e4m3fn))

124 +

125 + with flashinfer.autotune(autotune):

126 + out = flashinfer_scaled_fp4_mm(

127 + a_fp4,

128 + b_fp4,

129 + a_scale_interleaved,

130 + b_scale_interleaved,

131 + alpha,

132 + dtype,

133 + backend=backend,

134 + )

135 +

136 + torch.testing.assert_close(out,

137 + expected_out.to(dtype=dtype),

138 + atol=1e-1,

139 + rtol=1e-1)

tests/kernels/quantization/test_nvfp4_scaled_mm.py

@@ -65,9 +65,12 @@ def test_nvfp4_gemm(

65 b_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) /

66 torch.amax(b_dtype.flatten(), dim=-1)).to(torch.float32)

67 alpha = 1. / (a_global_scale * b_global_scale)

68 + # ops.scaled_fp4_quant returns swizzled scales, while weights

69 + # from checkpoints are in linear scales.

70 a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a_dtype, a_global_scale)

71 b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(b_dtype, b_global_scale)

73 + # get_ref_results unswizzles the scales internally.

74 expected_out = get_ref_results(a_fp4, b_fp4, a_scale_interleaved,

75 b_scale_interleaved, a_global_scale,

76 b_global_scale, m, n, dtype, block_size,

vllm/envs.py

@@ -1101,6 +1101,12 @@ environment_variables: dict[str, Callable[[], Any]] = {

1101 "VLLM_USE_TRTLLM_ATTENTION":

1102 lambda: os.getenv("VLLM_USE_TRTLLM_ATTENTION", None),

1103

1104 + # If set to 1, force the use of TRTLLM FP4 GEMM backend in flashinfer.

1105 + # Otherwise, uses the first available of: flashinfer cutlass GEMM,

1106 + # vllm cutlass GEMM, marlin GEMM.

1107 + "VLLM_USE_TRTLLM_FP4_GEMM":

1108 + lambda: bool(int(os.getenv("VLLM_USE_TRTLLM_FP4_GEMM", "0"))),

1109 +

1110 # Controls garbage collection during CUDA graph capture.

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

1112 # If set to 1, allows GC to run during capture.

@@ -1208,6 +1214,7 @@ def compute_hash() -> str:

1214 "VLLM_DP_SIZE",

1215 "VLLM_USE_STANDALONE_COMPILE",

1216 "VLLM_FUSED_MOE_CHUNK_SIZE",

1217 + "VLLM_USE_TRTLLM_FP4_GEMM",

1218 ]

1219 for key in environment_variables_to_hash:

1220 if key in environment_variables:

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

+49 -15

@@ -15,6 +15,7 @@ from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import

15 from vllm.model_executor.parameter import (GroupQuantScaleParameter,

16 ModelWeightParameter,

17 PerTensorScaleParameter)

18 +from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer

20 logger = init_logger(__name__)

@@ -24,6 +25,13 @@ __all__ = ["CompressedTensorsW4A4Fp4"]

25 class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):

27 def __init__(self):

28 + if envs.VLLM_USE_TRTLLM_FP4_GEMM:

29 + assert has_flashinfer(), "TRTLLM FP4 GEMM requires FlashInfer"

30 + self.backend = "flashinfer-trtllm"

31 + elif has_flashinfer():

32 + self.backend = "flashinfer-cutlass"

33 + else:

34 + self.backend = "cutlass"

35 self.group_size = 16

37 @classmethod

@@ -108,16 +116,36 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):

116 layer.weight_global_scale.max().to(torch.float32),

117 requires_grad=False)

118

111 - swizzled_weight_scale = self.swizzle_blockscale(layer.weight_scale)

112 - layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,

113 - requires_grad=False)

114 -

115 - # required by cutlass kernel; need Parameter, not ModelWeightParameter

116 - layer.weight = Parameter(layer.weight_packed.data, requires_grad=False)

117 -

118 - layer.alpha = Parameter(layer.input_global_scale *

119 - layer.weight_global_scale,

120 - requires_grad=False)

119 + if self.backend == "flashinfer-trtllm":

120 + # FlashInfer TRTLLM FP4 GEMM requires a different weight layout.

121 + # FlashInfer provides nvfp4_quantize to quantize + shuffle the

122 + # layout but we use our own quantization so we have to call

123 + # shuffles ourselves.

124 + from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a

125 +

126 + weight = layer.weight_packed.data

127 + weight_scale = layer.weight_scale.data

128 +

129 + epilogue_tile_m = 128

130 + weight = shuffle_matrix_a(weight.view(torch.uint8),

131 + epilogue_tile_m)

132 + weight_scale = (shuffle_matrix_sf_a(weight_scale.view(

133 + torch.uint8), epilogue_tile_m).reshape(

134 + weight_scale.shape).view(torch.float8_e4m3fn))

135 +

136 + layer.weight_scale_swizzled = Parameter(weight_scale,

137 + requires_grad=False)

138 + layer.weight_packed = Parameter(weight, requires_grad=False)

139 + else:

140 + swizzled_weight_scale = self.swizzle_blockscale(layer.weight_scale)

141 + layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,

142 + requires_grad=False)

143 + layer.weight_packed = Parameter(layer.weight_packed.data,

144 + requires_grad=False)

145 +

146 + layer.alpha = Parameter(

147 + 1 / (layer.input_global_scale * layer.weight_global_scale),

148 + requires_grad=False)

149

150 def apply_weights(self,

151 layer: torch.nn.Module,

@@ -128,7 +156,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):

156 out = run_nvfp4_emulations(

157 x=x,

158 input_global_scale=layer.input_global_scale,

131 - weight=layer.weight,

159 + weight=layer.weight_packed,

160 weight_scale_swizzled=layer.weight_scale_swizzled,

161 weight_global_scale=layer.weight_global_scale)

162 if bias is not None:

@@ -136,14 +164,20 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):

164 return out

165

166 output_dtype = x.dtype

139 - output_shape = [x.shape[0], layer.weight.shape[0]]

167 + output_shape = [x.shape[0], layer.weight_packed.shape[0]]

168

169 # quantize BF16 or FP16 to (FP4 and interleaved block scale)

170 x_fp4, x_blockscale = scaled_fp4_quant(x, layer.input_global_scale)

171

144 - out = cutlass_scaled_fp4_mm(x_fp4, layer.weight, x_blockscale,

145 - layer.weight_scale_swizzled,

146 - 1 / layer.alpha, output_dtype)

172 + mm_args = (x_fp4, layer.weight_packed, x_blockscale,

173 + layer.weight_scale_swizzled, layer.alpha, output_dtype)

174 + if self.backend == "flashinfer-trtllm":

175 + out = flashinfer_scaled_fp4_mm(*mm_args, backend="trtllm")

176 + elif self.backend == "flashinfer-cutlass":

177 + out = flashinfer_scaled_fp4_mm(*mm_args, backend="cutlass")

178 + else:

179 + out = cutlass_scaled_fp4_mm(*mm_args)

180 +

181 if bias is not None:

182 out = out + bias

183 return out.view(*output_shape)

vllm/model_executor/layers/quantization/modelopt.py

+61 -23

@@ -38,7 +38,8 @@ from vllm.model_executor.parameter import (ModelWeightParameter,

38 PerTensorScaleParameter)

39 from vllm.scalar_type import scalar_types

40 from vllm.utils import next_power_of_2

41 -from vllm.utils.flashinfer import has_flashinfer_moe

41 +from vllm.utils.flashinfer import (flashinfer_scaled_fp4_mm, has_flashinfer,

42 + has_flashinfer_moe)

44 logger = init_logger(__name__)

@@ -724,16 +725,20 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

725

726 def __init__(self, quant_config: ModelOptNvFp4Config) -> None:

727 self.quant_config = quant_config

727 - self.cutlass_nvfp4_supported = cutlass_fp4_supported()

728 - self.use_marlin = False

728

730 - if not self.cutlass_nvfp4_supported:

731 - if is_fp4_marlin_supported():

732 - self.use_marlin = True

733 - else:

734 - raise ValueError("Current platform does not support NVFP4"

735 - " quantization. Please use Blackwell and"

736 - " above.")

729 + if envs.VLLM_USE_TRTLLM_FP4_GEMM:

730 + assert has_flashinfer(), "TRTLLM FP4 GEMM requires FlashInfer"

731 + self.backend = "flashinfer-trtllm"

732 + elif has_flashinfer():

733 + self.backend = "flashinfer-cutlass"

734 + elif cutlass_fp4_supported():

735 + self.backend = "cutlass"

736 + elif is_fp4_marlin_supported():

737 + self.backend = "marlin"

738 + else:

739 + raise ValueError("Current platform does not support NVFP4"

740 + " quantization. Please use Blackwell and"

741 + " above.")

742

743 def create_weights(

744 self,

@@ -815,17 +820,38 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

820 # block_size = 16;

821 assert (layer.weight_scale.dtype == torch.float8_e4m3fn), (

822 "Weight Block scale must be represented as FP8-E4M3")

818 - swizzled_weight_scale = swizzle_blockscale(layer.weight_scale)

823

820 - layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,

821 - requires_grad=False)

822 - layer.weight = Parameter(layer.weight.data, requires_grad=False)

824 + if self.backend == "flashinfer-trtllm":

825 + # FlashInfer TRTLLM FP4 GEMM requires a different weight layout.

826 + # FlashInfer provides nvfp4_quantize to quantize + shuffle the

827 + # layout but we use our own quantization so we have to call

828 + # shuffles ourselves.

829 + from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a

830 +

831 + weight = layer.weight.data

832 + weight_scale = layer.weight_scale.data

833 +

834 + epilogue_tile_m = 128

835 + weight = shuffle_matrix_a(weight.view(torch.uint8),

836 + epilogue_tile_m)

837 + weight_scale = (shuffle_matrix_sf_a(weight_scale.view(

838 + torch.uint8), epilogue_tile_m).reshape(

839 + weight_scale.shape).view(torch.float8_e4m3fn))

840 +

841 + layer.weight_scale_swizzled = Parameter(weight_scale,

842 + requires_grad=False)

843 + layer.weight = Parameter(weight, requires_grad=False)

844 + else:

845 + swizzled_weight_scale = swizzle_blockscale(layer.weight_scale)

846 + layer.weight_scale_swizzled = Parameter(swizzled_weight_scale,

847 + requires_grad=False)

848 + layer.weight = Parameter(layer.weight.data, requires_grad=False)

849

824 - if self.use_marlin:

825 - prepare_fp4_layer_for_marlin(layer)

826 - del layer.alpha

827 - del layer.input_scale

828 - del layer.weight_scale_swizzled

850 + if self.backend == "marlin":

851 + prepare_fp4_layer_for_marlin(layer)

852 + del layer.alpha

853 + del layer.input_scale

854 + del layer.weight_scale_swizzled

855

856 def apply(

857 self,

@@ -833,7 +859,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

859 x: torch.Tensor,

860 bias: Optional[torch.Tensor] = None,

861 ) -> torch.Tensor:

836 - if self.use_marlin:

862 + if self.backend == "marlin":

863 return apply_fp4_marlin_linear(

864 input=x,

865 weight=layer.weight,

@@ -859,9 +885,21 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

885 assert (layer.weight_scale_swizzled.dtype == torch.float8_e4m3fn)

886 assert (layer.alpha.dtype == torch.float32)

887

862 - out = cutlass_scaled_fp4_mm(x_fp4, layer.weight, x_blockscale,

863 - layer.weight_scale_swizzled, layer.alpha,

864 - output_dtype)

888 + mm_args = (

889 + x_fp4,

890 + layer.weight,

891 + x_blockscale,

892 + layer.weight_scale_swizzled,

893 + layer.alpha,

894 + output_dtype,

895 + )

896 + if self.backend == "flashinfer-trtllm":

897 + out = flashinfer_scaled_fp4_mm(*mm_args, backend="trtllm")

898 + elif self.backend == "flashinfer-cutlass":

899 + out = flashinfer_scaled_fp4_mm(*mm_args, backend="cutlass")

900 + else:

901 + out = cutlass_scaled_fp4_mm(*mm_args)

902 +

903 if bias is not None:

904 out = out + bias

905 return out.view(*output_shape)

vllm/model_executor/warmup/kernel_warmup.py

+38 -1

@@ -5,16 +5,53 @@ Warmup kernels used during model execution.

5 This is useful specifically for JIT'ed kernels as we don't want JIT'ing to

6 happen during model execution.

7 """

8 +from typing import TYPE_CHECKING

9 +

10 import torch

12 import vllm.envs as envs

13 from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup

14 +from vllm.platforms import current_platform

15 from vllm.utils.deep_gemm import is_deep_gemm_supported

16 +from vllm.utils.flashinfer import has_flashinfer

17 +

18 +if TYPE_CHECKING:

19 + from vllm.v1.worker.gpu_model_runner import GPUModelRunner

20 + from vllm.v1.worker.gpu_worker import Worker

15 -def kernel_warmup(model: torch.nn.Module, max_tokens: int):

23 +def kernel_warmup(worker: "Worker"):

24 + # Deep GEMM warmup

25 do_deep_gemm_warmup = (envs.VLLM_USE_DEEP_GEMM

26 and is_deep_gemm_supported()

27 and not envs.VLLM_SKIP_DEEP_GEMM_WARMUP)

28 if do_deep_gemm_warmup:

29 + model = worker.get_model()

30 + max_tokens = worker.scheduler_config.max_num_batched_tokens

31 deep_gemm_warmup(model, max_tokens)

32 +

33 + # FlashInfer autotune for Blackwell (SM 10.0) GPUs

34 + if has_flashinfer() and current_platform.is_device_capability(100):

35 + flashinfer_autotune(worker.model_runner)

36 +

37 +

38 +def flashinfer_autotune(runner: "GPUModelRunner") -> None:

39 + """

40 + Autotune FlashInfer operations.

41 + FlashInfer have many implementations for the same operation,

42 + autotuning runs benchmarks for each implementation and stores

43 + the results. The results are cached transparently and

44 + future calls to FlashInfer will use the best implementation.

45 + Without autotuning, FlashInfer will rely on heuristics, which may

46 + be significantly slower.

47 + """

48 + from vllm.utils.flashinfer import autotune

49 +

50 + with torch.inference_mode(), autotune():

51 + # We skip EPLB here since we don't want to record dummy metrics

52 + # When autotuning with number of tokens m, flashinfer will autotune

53 + # operations for all number of tokens up to m.

54 + # So we only need to run with the max number of tokens.

55 + runner._dummy_run(runner.scheduler_config.max_num_batched_tokens,

56 + skip_eplb=True,

57 + is_profile=True)

vllm/utils/flashinfer.py

+71

@@ -14,6 +14,7 @@ import os

14 from typing import Any, Callable, NoReturn, Optional

16 import requests

17 +import torch

19 import vllm.envs as envs

20 from vllm.logger import init_logger

@@ -193,6 +194,75 @@ def use_trtllm_attention(

194 return use_trtllm

195

196

197 +if has_flashinfer():

198 +

199 + @torch.library.custom_op(

200 + "vllm::flashinfer_mm_fp4",

201 + mutates_args=[],

202 + device_types="cuda",

203 + )

204 + def flashinfer_mm_fp4(

205 + A: torch.Tensor,

206 + B: torch.Tensor,

207 + A_scale: torch.Tensor,

208 + B_scale: torch.Tensor,

209 + g_scale: torch.Tensor,

210 + dtype: torch.dtype,

211 + backend: str,

212 + ) -> torch.Tensor:

213 + from flashinfer import mm_fp4 as flashinfer_mm_fp4_

214 + return flashinfer_mm_fp4_(A,

215 + B,

216 + A_scale,

217 + B_scale,

218 + g_scale,

219 + dtype,

220 + block_size=16,

221 + backend=backend)

222 +

223 + @torch.library.register_fake("vllm::flashinfer_mm_fp4", )

224 + def flashinfer_mm_fp4_fake(

225 + A: torch.Tensor,

226 + B: torch.Tensor,

227 + A_scale: torch.Tensor,

228 + B_scale: torch.Tensor,

229 + g_scale: torch.Tensor,

230 + dtype: torch.dtype,

231 + backend: str,

232 + ) -> torch.Tensor:

233 + return torch.empty(A.shape[0],

234 + B.shape[1],

235 + dtype=dtype,

236 + device=A.device)

237 +

238 +

239 +def flashinfer_scaled_fp4_mm(a: torch.Tensor, b: torch.Tensor,

240 + block_scale_a: torch.Tensor,

241 + block_scale_b: torch.Tensor, alpha: torch.Tensor,

242 + out_dtype: torch.dtype,

243 + backend: str) -> torch.Tensor:

244 + assert a.ndim == 2 and b.ndim == 2

245 + assert block_scale_a.ndim == 2 and block_scale_b.ndim == 2

246 + assert a.stride(-1) == 1 and b.stride(-1) == 1

247 + assert a.shape[1] == b.shape[1]

248 + assert block_scale_a.shape[1] == a.shape[1] // 8

249 + assert block_scale_b.shape[1] == b.shape[1] // 8

250 +

251 + if backend == "cutlass":

252 + block_scale_a = block_scale_a.view(torch.uint8)

253 + block_scale_b = block_scale_b.view(torch.uint8)

254 +

255 + return flashinfer_mm_fp4(

256 + a,

257 + b.t(),

258 + block_scale_a,

259 + block_scale_b.t(),

260 + alpha,

261 + out_dtype,

262 + backend=backend,

263 + )

264 +

265 +

266 __all__ = [

267 "has_flashinfer",

268 "flashinfer_trtllm_fp8_block_scale_moe",

@@ -205,4 +275,5 @@ __all__ = [

275 "has_flashinfer_cutlass_fused_moe",

276 "has_nvidia_artifactory",

277 "use_trtllm_attention",

278 + "flashinfer_scaled_fp4_mm",

279 ]

vllm/v1/worker/gpu_worker.py

+2 -2

@@ -310,6 +310,7 @@ class Worker(WorkerBase):

310 for size in sorted(warmup_sizes, reverse=True):

311 logger.info("Compile and warming up model for size %d", size)

312 self.model_runner._dummy_run(size, skip_eplb=True)

313 +

314 if not self.model_config.enforce_eager:

315 self.model_runner.capture_model()

316

@@ -340,8 +341,7 @@ class Worker(WorkerBase):

341 hidden_states=last_hidden_states)

342

343 # Warmup kernels used during model execution

343 - kernel_warmup(self.get_model(),

344 - max_tokens=self.scheduler_config.max_num_batched_tokens)

344 + kernel_warmup(self)

345

346 # Reset the seed to ensure that the random state is not affected by

347 # the model initialization and profiling.