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

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] Added flashinfer fp8 per-tensor gemms (#22895)

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

nvjullin committed Aug 26, 2025 at 21:54 UTC f66673a39d9f364194c249f28098cad8a5584ccb

9 files changed +198 -36

.buildkite/test-pipeline.yaml

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

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

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

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

658 + - pytest -v -s tests/kernels/quantization/test_flashinfer_scaled_mm.py

659 - pytest -v -s tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py

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

661 - pytest -v -s tests/kernels/moe/test_mxfp4_moe.py

tests/compile/test_fusion.py

+7 -8

@@ -15,7 +15,7 @@ from vllm.model_executor.layers.layernorm import RMSNorm

15 from vllm.model_executor.layers.quantization.utils.quant_utils import (

16 GroupShape, QuantKey, ScaleDesc)

17 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

18 - CUTLASS_FP8_SUPPORTED, Fp8LinearOp, maybe_create_device_identity)

18 + Fp8LinearOp, maybe_create_device_identity)

19 from vllm.platforms import current_platform

21 from .backend import TestBackend

@@ -26,9 +26,9 @@ FP8_DTYPE = current_platform.fp8_dtype()

26 class TestModel(torch.nn.Module):

28 def __init__(self, hidden_size: int, eps: float, static: bool,

29 - cutlass_fp8_enabled: bool, *args, **kwargs):

29 + force_fp8_e4m3fnuz: bool, *args, **kwargs):

30 super().__init__(*args, **kwargs)

31 - self.cutlass_fp8_enabled = cutlass_fp8_enabled

31 + self.force_fp8_e4m3fnuz = force_fp8_e4m3fnuz

32 self.norm = [RMSNorm(hidden_size, eps) for _ in range(3)]

33 self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(2)]

34 group_shape = GroupShape.PER_TENSOR if static else GroupShape.PER_TOKEN

@@ -43,7 +43,7 @@ class TestModel(torch.nn.Module):

43 for _ in range(2)

44 ]

45 self.fp8_linear = Fp8LinearOp(

46 - cutlass_fp8_supported=cutlass_fp8_enabled,

46 + force_fp8_e4m3fnuz=force_fp8_e4m3fnuz,

47 act_quant_static=static,

48 act_quant_group_shape=group_shape,

49 )

@@ -81,12 +81,11 @@ class TestModel(torch.nn.Module):

81 @pytest.mark.parametrize("num_tokens", [7, 256, 533, 2048, 2049])

82 @pytest.mark.parametrize("eps", [1e-5, 1e-6])

83 @pytest.mark.parametrize("static", [True, False])

84 -@pytest.mark.parametrize("cutlass_fp8_enabled",

85 - [True, False] if CUTLASS_FP8_SUPPORTED else [False])

84 +@pytest.mark.parametrize("force_fp8_e4m3fnuz", [True, False])

85 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"],

86 reason="Only test on CUDA and ROCm")

87 def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,

89 - cutlass_fp8_enabled):

88 + force_fp8_e4m3fnuz):

89 torch.set_default_device("cuda")

90 torch.set_default_dtype(dtype)

91 torch.manual_seed(1)

@@ -103,7 +102,7 @@ def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,

102 fusion_pass = FusionPass.instance(vllm_config)

103

104 backend = TestBackend(noop_pass, fusion_pass)

106 - model = TestModel(hidden_size, eps, static, cutlass_fp8_enabled)

105 + model = TestModel(hidden_size, eps, static, force_fp8_e4m3fnuz)

106

107 # First dimension dynamic

108 x = torch.rand(num_tokens, hidden_size)

tests/compile/test_sequence_parallelism.py

+1 -2

@@ -104,8 +104,7 @@ class TestQuantModel(torch.nn.Module):

104 # Initialize weights

105 torch.nn.init.normal_(self.gate_proj, std=0.02)

106

107 - self.fp8_linear = Fp8LinearOp(cutlass_fp8_supported=True,

108 - use_per_token_if_dynamic=False)

107 + self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=False)

108

109 self.scale = torch.rand(1, dtype=torch.float32)

110 # Create a weight that is compatible with torch._scaled_mm,

tests/compile/test_silu_mul_quant_fusion.py

+6 -7

@@ -12,7 +12,7 @@ from vllm.model_executor.layers.activation import SiluAndMul

12 from vllm.model_executor.layers.quantization.utils.quant_utils import (

13 GroupShape)

14 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (

15 - CUTLASS_FP8_SUPPORTED, Fp8LinearOp)

15 + Fp8LinearOp)

16 from vllm.platforms import current_platform

18 from .backend import TestBackend

@@ -20,7 +20,7 @@ from .backend import TestBackend

21 class TestModel(torch.nn.Module):

23 - def __init__(self, hidden_size: int, cutlass_fp8_enabled: bool, *args,

23 + def __init__(self, hidden_size: int, force_fp8_e4m3fnuz: bool, *args,

24 **kwargs):

25 super().__init__(*args, **kwargs)

26 self.silu_and_mul = SiluAndMul()

@@ -32,7 +32,7 @@ class TestModel(torch.nn.Module):

32 hidden_size).to(dtype=current_platform.fp8_dtype()).t())

34 self.fp8_linear = Fp8LinearOp(

35 - cutlass_fp8_supported=cutlass_fp8_enabled,

35 + force_fp8_e4m3fnuz=force_fp8_e4m3fnuz,

36 act_quant_static=True,

37 act_quant_group_shape=GroupShape.PER_TENSOR,

38 )

@@ -48,12 +48,11 @@ class TestModel(torch.nn.Module):

49 @pytest.mark.parametrize("num_tokens", [256])

50 @pytest.mark.parametrize("hidden_size", [64])

51 -@pytest.mark.parametrize("cutlass_fp8_enabled",

52 - [True, False] if CUTLASS_FP8_SUPPORTED else [False])

51 +@pytest.mark.parametrize("force_fp8_e4m3fnuz", [True, False])

52 @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"],

53 reason="Only test on CUDA and ROCm")

54 def test_fusion_silu_and_mul_quant(num_tokens, hidden_size,

56 - cutlass_fp8_enabled):

55 + force_fp8_e4m3fnuz):

56 torch.set_default_device("cuda")

57 torch.set_default_dtype(torch.float16)

@@ -64,7 +63,7 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size,

63 fusion_pass = ActivationQuantFusionPass(config)

65 backend = TestBackend(NoOpEliminationPass(config), fusion_pass)

67 - model = TestModel(hidden_size, cutlass_fp8_enabled)

66 + model = TestModel(hidden_size, force_fp8_e4m3fnuz)

68 # First dimension dynamic

69 x = torch.rand(num_tokens, hidden_size * 2)

tests/kernels/quantization/test_flashinfer_scaled_mm.py new

+73

@@ -0,0 +1,73 @@

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

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

3 +import pytest

4 +import torch

5 +

6 +from vllm import _custom_ops as ops

7 +from vllm.platforms import current_platform

8 +from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm

9 +

10 +if not current_platform.has_device_capability(100):

11 + pytest.skip(

12 + reason=

13 + "Flashinfer FP8 gemms requires compute capability of 10.0 or above.",

14 + allow_module_level=True,

15 + )

16 +

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

18 +# m, n, k

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

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

21 +SHAPES.extend(PAD_SHAPES)

22 +

23 +SEEDS = [42]

24 +CUDA_DEVICES = ["cuda:0"]

25 +

26 +

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

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

29 +@pytest.mark.parametrize("use_bias", [True, False])

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

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

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

33 +@torch.inference_mode()

34 +def test_flashinfer_fp8_gemm(

35 + dtype: torch.dtype,

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

37 + use_bias: bool,

38 + seed: int,

39 + device: str,

40 + autotune: bool,

41 +) -> None:

42 + current_platform.seed_everything(seed)

43 + m, n, k = shape

44 + a = torch.randn((m, k), dtype=dtype, device=device)

45 + b = torch.randn((n, k), dtype=dtype, device=device) / k

46 +

47 + a_fp8, a_scale = ops.scaled_fp8_quant(a)

48 + b_fp8, b_scale = ops.scaled_fp8_quant(b)

49 +

50 + expected_out = torch.mm(

51 + a_scale * a_fp8.to(dtype=torch.float32),

52 + b_scale * b_fp8.to(dtype=torch.float32).t(),

53 + ).to(dtype=dtype)

54 +

55 + if use_bias:

56 + bias = torch.randn((n, ), dtype=dtype, device=device)

57 + expected_out = expected_out + bias

58 + else:

59 + bias = None

60 +

61 + import flashinfer

62 +

63 + with flashinfer.autotune(autotune):

64 + out = flashinfer_scaled_fp8_mm(

65 + a_fp8,

66 + b_fp8.t(),

67 + a_scale,

68 + b_scale,

69 + dtype,

70 + bias=bias,

71 + )

72 +

73 + torch.testing.assert_close(out, expected_out, atol=1e-2, rtol=1e-2)

vllm/model_executor/layers/quantization/fp8.py

+3 -2

@@ -223,8 +223,7 @@ class Fp8LinearMethod(LinearMethodBase):

223

224 self.fp8_linear = Fp8LinearOp(

225 act_quant_static=self.act_q_static,

226 - act_quant_group_shape=self.act_q_group_shape,

227 - cutlass_fp8_supported=cutlass_fp8_supported())

226 + act_quant_group_shape=self.act_q_group_shape)

227

228 def create_weights(

229 self,

@@ -376,6 +375,8 @@ class Fp8LinearMethod(LinearMethodBase):

375 # Update the layer with the new values.

376 layer.weight = Parameter(qweight.t(), requires_grad=False)

377 layer.weight_scale = Parameter(weight_scale, requires_grad=False)

378 + # layer.input_scale is None indicates dynamic quant and scale is

379 + # computed from input.

380 layer.input_scale = None

381

382 # If checkpoint is fp8, handle that there are N scales for N

vllm/model_executor/layers/quantization/ptpc_fp8.py

+2 -2

@@ -97,8 +97,8 @@ class PTPCFp8LinearMethod(Fp8LinearMethod):

97 self.quant_config.is_checkpoint_fp8_serialized = False

98 self.fp8_linear = Fp8LinearOp(

99 act_quant_static=False,

100 - cutlass_fp8_supported=False,

101 - act_quant_group_shape=GroupShape.PER_TOKEN)

100 + act_quant_group_shape=GroupShape.PER_TOKEN,

101 + force_fp8_e4m3fnuz=True)

102

103 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

104 layer.weight = torch.nn.Parameter(layer.weight.data,

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

+44 -15

@@ -14,6 +14,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

14 GroupShape)

15 from vllm.platforms import current_platform

16 from vllm.utils import direct_register_custom_op

17 +from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm, has_flashinfer

19 # Input scaling factors are no longer optional in _scaled_mm starting

20 # from pytorch 2.5. Allocating a dummy tensor to pass as input_scale

@@ -157,6 +158,19 @@ def cutlass_w8a8_scaled_mm(*, qinput: torch.Tensor, weight: torch.Tensor,

158 return output.view(*output_shape)

159

160

161 +def flashinfer_w8a8_scaled_mm(*, qinput: torch.Tensor, weight: torch.Tensor,

162 + out_dtype: torch.dtype, scale_a: torch.Tensor,

163 + scale_b: torch.Tensor, bias: torch.Tensor,

164 + output_shape: list, **kwargs) -> torch.Tensor:

165 +

166 + return flashinfer_scaled_fp8_mm(qinput,

167 + weight,

168 + out_dtype=out_dtype,

169 + scale_a=scale_a,

170 + scale_b=scale_b,

171 + bias=bias)

172 +

173 +

174 def rocm_per_tensor_w8a8_scaled_mm_impl(

175 qinput: torch.Tensor, weight: torch.Tensor, out_dtype: torch.dtype,

176 scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.Tensor,

@@ -231,8 +245,8 @@ def torch_per_token_w8a8_scaled_mm(*, qinput: torch.Tensor,

245 out_dtype: torch.dtype,

246 scale_a: torch.Tensor,

247 scale_b: torch.Tensor, bias: torch.Tensor,

234 - input_2d: torch.Tensor,

235 - output_shape: list) -> torch.Tensor:

248 + input_2d: torch.Tensor, output_shape: list,

249 + **kwargs) -> torch.Tensor:

250 # Note: Callers of this function should check USE_ROWWISE_TORCH_SCALED_MM

251 # when using it.

252 # For now it has only been validated on ROCm platform.

@@ -303,16 +317,22 @@ def torch_channelwise_w8a8_scaled_mm(*, qinput: torch.Tensor,

317

318

319 def dispatch_w8a8_scaled_mm(

306 - cutlass_fp8_supported: bool, per_tensor_weights: bool,

320 + preferred_backend: str, per_tensor_weights: bool,

321 per_tensor_activations: bool) -> Callable[..., torch.Tensor]:

322

309 - # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A

310 - if cutlass_fp8_supported:

311 - return cutlass_w8a8_scaled_mm

323 if per_tensor_weights and per_tensor_activations:

313 - if current_platform.is_rocm():

324 + if preferred_backend == "rocm":

325 return rocm_per_tensor_w8a8_scaled_mm

326 + if preferred_backend == "flashinfer":

327 + return flashinfer_w8a8_scaled_mm

328 + if preferred_backend == "cutlass":

329 + return cutlass_w8a8_scaled_mm

330 return torch_per_tensor_w8a8_scaled_mm

331 +

332 + # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A

333 + if preferred_backend == "cutlass" or preferred_backend == "flashinfer":

334 + return cutlass_w8a8_scaled_mm

335 +

336 # If torch.scaled_mm supports per-channel (weights) per-token (inputs)

337 if not per_tensor_weights and not per_tensor_activations \

338 and USE_ROWWISE_TORCH_SCALED_MM:

@@ -334,10 +354,20 @@ class Fp8LinearOp:

354

355 def __init__(self,

356 act_quant_static: bool,

337 - cutlass_fp8_supported: bool = cutlass_fp8_supported(),

357 act_quant_group_shape: GroupShape = GroupShape.PER_TENSOR,

339 - pad_output: Optional[bool] = None):

340 - self.cutlass_fp8_supported = cutlass_fp8_supported

358 + pad_output: Optional[bool] = None,

359 + force_fp8_e4m3fnuz: bool = False):

360 + if current_platform.is_rocm():

361 + self.preferred_backend = "rocm"

362 + elif current_platform.is_cuda(

363 + ) and not force_fp8_e4m3fnuz and cutlass_fp8_supported():

364 + if has_flashinfer() and current_platform.has_device_capability(

365 + 100):

366 + self.preferred_backend = "flashinfer"

367 + else:

368 + self.preferred_backend = "cutlass"

369 + else:

370 + self.preferred_backend = "torch"

371

372 # Note: we pad the input because torch._scaled_mm is more performant

373 # for matrices with batch dimension > 16.

@@ -347,8 +377,7 @@ class Fp8LinearOp:

377 if pad_output is None:

378 config = get_current_vllm_config().compilation_config

379 pad_output = config.level < CompilationLevel.PIECEWISE and \

350 - not cutlass_fp8_supported and \

351 - not current_platform.is_rocm()

380 + self.preferred_backend == "torch"

381

382 self.output_padding = 17 if pad_output else None

383 self.act_quant_static = act_quant_static

@@ -393,9 +422,9 @@ class Fp8LinearOp:

422 per_tensor_activations = (x_scale.numel() == 1)

423

424 # TODO(luka) do this dispatch during init (after ScaledMM refactor)

396 - w8a8_scaled_mm_func = dispatch_w8a8_scaled_mm(

397 - self.cutlass_fp8_supported, per_tensor_weights,

398 - per_tensor_activations)

425 + w8a8_scaled_mm_func = dispatch_w8a8_scaled_mm(self.preferred_backend,

426 + per_tensor_weights,

427 + per_tensor_activations)

428

429 return w8a8_scaled_mm_func(qinput=qinput,

430 weight=weight,

vllm/utils/flashinfer.py

+61

@@ -265,6 +265,37 @@ if has_flashinfer():

265 dtype=dtype,

266 device=A.device)

267

268 + @torch.library.custom_op(

269 + "vllm::bmm_fp8",

270 + mutates_args=[],

271 + device_types="cuda",

272 + )

273 + def bmm_fp8(

274 + A: torch.Tensor,

275 + B: torch.Tensor,

276 + A_scale: torch.Tensor,

277 + B_scale: torch.Tensor,

278 + dtype: torch.dtype,

279 + backend: str,

280 + ) -> torch.Tensor:

281 + from flashinfer import bmm_fp8 as bmm_fp8_

282 + return bmm_fp8_(A, B, A_scale, B_scale, dtype, None, backend)

283 +

284 + @torch.library.register_fake("vllm::bmm_fp8", )

285 + def bmm_fp8_fake(

286 + A: torch.Tensor,

287 + B: torch.Tensor,

288 + A_scale: torch.Tensor,

289 + B_scale: torch.Tensor,

290 + dtype: torch.dtype,

291 + backend: str,

292 + ) -> torch.Tensor:

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

294 + A.shape[1],

295 + B.shape[2],

296 + dtype=dtype,

297 + device=A.device)

298 +

299

300 def flashinfer_scaled_fp4_mm(a: torch.Tensor, b: torch.Tensor,

301 block_scale_a: torch.Tensor,

@@ -293,6 +324,35 @@ def flashinfer_scaled_fp4_mm(a: torch.Tensor, b: torch.Tensor,

324 )

325

326

327 +def flashinfer_scaled_fp8_mm(

328 + a: torch.Tensor,

329 + b: torch.Tensor,

330 + scale_a: torch.Tensor,

331 + scale_b: torch.Tensor,

332 + out_dtype: torch.dtype,

333 + bias: Optional[torch.Tensor] = None) -> torch.Tensor:

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

335 + assert a.shape[1] == b.shape[0]

336 + assert scale_a.numel() == 1 and scale_b.numel() == 1

337 + assert a.dtype == torch.float8_e4m3fn and b.dtype == torch.float8_e4m3fn

338 + assert a.device.type == "cuda" and b.device.type == "cuda"

339 + assert scale_a.dtype == torch.float32 and scale_b.dtype == torch.float32

340 + assert scale_a.device.type == "cuda" and scale_b.device.type == "cuda"

341 +

342 + output = bmm_fp8(

343 + a.unsqueeze(0),

344 + b.unsqueeze(0),

345 + scale_a,

346 + scale_b,

347 + out_dtype,

348 + "auto",

349 + ).view(a.shape[0], b.shape[1])

350 +

351 + if bias is not None:

352 + output = output + bias

353 + return output

354 +

355 +

356 __all__ = [

357 "has_flashinfer",

358 "flashinfer_trtllm_fp8_block_scale_moe",

@@ -307,4 +367,5 @@ __all__ = [

367 "supports_trtllm_attention",

368 "use_trtllm_attention",

369 "flashinfer_scaled_fp4_mm",

370 + "flashinfer_scaled_fp8_mm",

371 ]