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

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 10, 05:15 UTC

[Kernel][Performance] Enable smaller Scaling Factor tiling for NVFP4 small-batch decoding (#30885)

Signed-off-by: LopezCastroRoberto <roberto.lopez.castro@udc.es> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com> Signed-off-by: LopezCastroRoberto <rocastro@redhat.com>

Roberto L. Castro committed Jan 14, 2026 at 00:22 UTC 8ef50d9a6b91b7800e69a846219069a29a0298a4

9 files changed +177 -32

.buildkite/test-pipeline.yaml

+3 -1

@@ -951,7 +951,7 @@ steps:

951 # Whisper needs spawn method to avoid deadlock

952 - VLLM_WORKER_MULTIPROC_METHOD=spawn python3 examples/offline_inference/audio_language.py --model-type whisper

953

954 -- label: Blackwell Test # 21 min

954 +- label: Blackwell Test # 23 min

955 timeout_in_minutes: 30

956 working_dir: "/vllm-workspace/"

957 gpu: b200

@@ -991,6 +991,8 @@ steps:

991 - pytest -v -s tests/kernels/moe/test_ocp_mx_moe.py

992 - pytest -v -s tests/kernels/moe/test_flashinfer.py

993 - pytest -v -s tests/kernels/moe/test_cutedsl_moe.py

994 + # e2e

995 + - pytest -v -s tests/models/quantization/test_nvfp4.py

996

997 - label: Blackwell Fusion and Compile Tests # 30 min

998 timeout_in_minutes: 40

tests/kernels/quantization/nvfp4_utils.py

+24 -2

@@ -23,8 +23,26 @@ def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size):

23 return out[0:m, 0:k]

26 +def convert_swizzled_8x4_layout_to_linear(

27 + a_sf_swizzled: torch.Tensor, m, k, block_size

28 +):

29 + m_tiles = (m + 8 - 1) // 8

30 + f = block_size * 4

31 + k_tiles = (k + f - 1) // f

32 + tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 8, 4))

33 + tmp = torch.permute(tmp, (0, 1, 3, 2, 4))

34 + out = tmp.reshape(m_tiles * 8, k_tiles * f // block_size)

35 + return out[0:m, 0:k]

36 +

37 +

38 def dequantize_nvfp4_to_dtype(

27 - tensor_fp4, tensor_sf, global_scale, dtype, device, block_size=16

39 + tensor_fp4,

40 + tensor_sf,

41 + global_scale,

42 + dtype,

43 + device,

44 + block_size=16,

45 + is_sf_128x4_layout=True,

46 ):

47 """Dequantize the fp4 tensor back to high precision."""

48 # Two fp4 values are packed into one uint8.

@@ -34,7 +52,11 @@ def dequantize_nvfp4_to_dtype(

52 tensor_f32 = break_fp4_bytes(tensor_fp4, dtype)

53 tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)

54 tensor_sf = tensor_sf.view(torch.float8_e4m3fn)

37 - tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)

55 + if is_sf_128x4_layout:

56 + tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)

57 + else:

58 + tensor_sf = convert_swizzled_8x4_layout_to_linear(tensor_sf, m, k, block_size)

59 +

60 tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale

62 # scale the tensor

tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py

+26 -8

@@ -11,7 +11,9 @@ from nvfp4_utils import (

12 from vllm import _custom_ops as ops

13 from vllm.platforms import current_platform

14 -from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm

14 +from vllm.utils.flashinfer import (

15 + flashinfer_scaled_fp4_mm,

16 +)

17 from vllm.utils.torch_utils import set_random_seed

19 if not current_platform.has_device_capability(100):

@@ -22,8 +24,14 @@ if not current_platform.has_device_capability(100):

25 DTYPES = [torch.float16, torch.bfloat16]

26 # m, n, k

25 -SHAPES = [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)]

26 -PAD_SHAPES = [(150, 128, 64), (128, 128, 96)]

27 +SHAPES = [

28 + (128, 128, 64),

29 + (128, 128, 128),

30 + (256, 128, 64),

31 + (128, 256, 128),

32 + (1, 128, 128),

33 +]

34 +PAD_SHAPES = [(150, 128, 64), (128, 128, 96), (2, 128, 64), (3, 128, 96)]

35 SHAPES.extend(PAD_SHAPES)

37 SEEDS = [42]

@@ -42,12 +50,19 @@ def get_ref_results(

50 dtype,

51 block_size,

52 device,

53 + is_sf_128x4_layout,

54 ):

55 _, m_k = a_fp4.shape

56 _, n_k = b_fp4.shape

57 assert m_k == n_k

58 a_in_dtype = dequantize_nvfp4_to_dtype(

50 - a_fp4, a_sf, a_global_scale, dtype=dtype, device=device, block_size=block_size

59 + a_fp4,

60 + a_sf,

61 + a_global_scale,

62 + dtype=dtype,

63 + device=device,

64 + block_size=block_size,

65 + is_sf_128x4_layout=is_sf_128x4_layout,

66 )

67 b_in_dtype = dequantize_nvfp4_to_dtype(

68 b_fp4, b_sf, b_global_scale, dtype=dtype, device=device, block_size=block_size

@@ -70,7 +85,7 @@ def test_flashinfer_nvfp4_gemm(

85 backend: str,

86 autotune: bool,

87 ) -> None:

73 - if backend == "trtllm" and dtype == torch.float16:

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

89 pytest.skip("Only torch.bfloat16 is supported for TRTLLM FP4 GEMM operations")

91 set_random_seed(seed)

@@ -87,11 +102,14 @@ def test_flashinfer_nvfp4_gemm(

102 (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)

103 ).to(torch.float32)

104 alpha = 1.0 / (a_global_scale * b_global_scale)

105 +

106 # ops.scaled_fp4_quant returns swizzled scales, while weights

107 # from checkpoints are in linear scales.

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

109 # we need to unswizzle for trtllm here.

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

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

111 + is_sf_128x4_layout = not (backend == "trtllm" and m <= 32)

112 +

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

114

115 # get_ref_results unswizzles the scales internally.

@@ -107,14 +125,14 @@ def test_flashinfer_nvfp4_gemm(

125 dtype,

126 block_size,

127 device,

128 + is_sf_128x4_layout,

129 )

130

131 import flashinfer

132

114 - if backend == "trtllm":

133 + if "trtllm" in backend:

134 epilogue_tile_m = 128

135 b_fp4 = flashinfer.shuffle_matrix_a(b_fp4.view(torch.uint8), epilogue_tile_m)

117 -

136 b_scale_interleaved = convert_swizzled_to_linear(

137 b_scale_interleaved, n, k, block_size

138 )

tests/models/quantization/test_nvfp4.py

+26

@@ -14,6 +14,8 @@ from transformers import AutoTokenizer

14 from tests.quantization.utils import is_quant_method_supported

15 from vllm import LLM, SamplingParams

17 +from vllm.platforms import current_platform

18 +

19 os.environ["TOKENIZERS_PARALLELISM"] = "true"

21 MAX_MODEL_LEN = 1024

@@ -83,3 +85,27 @@ def test_models(example_prompts, model_name) -> None:

85 assert expected_str == generated_str, (

86 f"Test{i}:\nExpected: {expected_str!r}\nvLLM: {generated_str!r}"

87 )

88 +

89 +

90 +EAGER = [True, False]

91 +

92 +

93 +@pytest.mark.skipif(

94 + not current_platform.has_device_capability(100),

95 + reason="modelopt_fp4 is not supported on this GPU type.",

96 +)

97 +@pytest.mark.parametrize("model", ["nvidia/Llama-3.1-8B-Instruct-NVFP4"])

98 +@pytest.mark.parametrize("eager", EAGER)

99 +@pytest.mark.parametrize(

100 + "backend",

101 + [

102 + "flashinfer-cudnn",

103 + "flashinfer-trtllm", # the small seq_len ensures trtllm_8x4_layout backend is used

104 + "flashinfer-cutlass",

105 + ],

106 +)

107 +def test_nvfp4(vllm_runner, model, eager, backend, monkeypatch):

108 + monkeypatch.setenv("VLLM_NVFP4_GEMM_BACKEND", backend)

109 + with vllm_runner(model, enforce_eager=eager) as llm:

110 + output = llm.generate_greedy(["1 2 3 4 5"], max_tokens=2)

111 + assert output[0][1] == "1 2 3 4 5 6"

vllm/_custom_ops.py

+30 -16

@@ -9,6 +9,9 @@ import vllm.envs as envs

9 from vllm.logger import init_logger

10 from vllm.platforms import current_platform

11 from vllm.scalar_type import ScalarType

12 +from vllm.utils.flashinfer import (

13 + flashinfer_quant_nvfp4_8x4_sf_layout,

14 +)

16 logger = init_logger(__name__)

@@ -1563,7 +1566,9 @@ def permute_cols(a: torch.Tensor, perm: torch.Tensor) -> torch.Tensor:

1566

1567 # fp4

1568 def scaled_fp4_quant(

1566 - input: torch.Tensor, input_global_scale: torch.Tensor

1569 + input: torch.Tensor,

1570 + input_global_scale: torch.Tensor,

1571 + backend: str = "none",

1572 ) -> tuple[torch.Tensor, torch.Tensor]:

1573 """

1574 Quantize input tensor to FP4 and return quantized tensor and scale.

@@ -1577,6 +1582,7 @@ def scaled_fp4_quant(

1582 Args:

1583 input: The input tensor to be quantized to FP4

1584 input_global_scale: A scalar scaling factor for the entire tensor.

1585 + use_8x4_sf_layout: Whether to use the 8x4 or 128x4 layout for the scaling

1586

1587 Returns:

1588 tuple[torch.Tensor, torch.Tensor]: The output tensor in FP4 but every

@@ -1596,23 +1602,31 @@ def scaled_fp4_quant(

1602 f"input.dtype needs to be fp16 or bf16 but got {input.dtype}."

1603 )

1604

1599 - # Two fp4 values will be packed into an uint8.

1600 - output = torch.empty((m, n // 2), device=device, dtype=torch.uint8)

1605 + use_8x4_sf_layout = True if "trtllm" in backend and m <= 32 else False # noqa: SIM210

1606

1602 - # We use the rounded values to store the swizzled values. Due to the

1603 - # requirement of the Tensor Core, the minimum tile is 128x4 for the scales.

1604 - # So, we first pad the scales to multiples of 128 and 4. Then, the scales

1605 - # (in float8_e4m3fn) are packed into an int32 for every 4 values. More:

1606 - # https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x

1607 - round_up = lambda x, y: (x + y - 1) // y * y

1608 - rounded_m = round_up(m, 128)

1609 - scale_n = n // block_size

1610 - rounded_n = round_up(scale_n, 4)

1611 - output_scale = torch.empty(

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

1613 - )

1607 + if use_8x4_sf_layout:

1608 + output, output_scale = flashinfer_quant_nvfp4_8x4_sf_layout(

1609 + input, input_global_scale

1610 + )

1611 + else:

1612 + # Two fp4 values will be packed into an uint8.

1613 + output = torch.empty((m, n // 2), device=device, dtype=torch.uint8)

1614 +

1615 + # We use the rounded values to store the swizzled values. Due to the

1616 + # requirement of the Tensor Core, the minimum tile is 128x4 for the scales.

1617 + # So, we first pad the scales to multiples of 128 and 4. Then, the scales

1618 + # (in float8_e4m3fn) are packed into an int32 for every 4 values. More:

1619 + # https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x

1620 + round_up = lambda x, y: (x + y - 1) // y * y

1621 + rounded_m = round_up(m, 128)

1622 + scale_n = n // block_size

1623 + rounded_n = round_up(scale_n, 4)

1624 + output_scale = torch.empty(

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

1626 + )

1627 +

1628 + torch.ops._C.scaled_fp4_quant(output, input, output_scale, input_global_scale)

1629

1615 - torch.ops._C.scaled_fp4_quant(output, input, output_scale, input_global_scale)

1630 output_scale = output_scale.view(torch.float8_e4m3fn)

1631 return output, output_scale

1632

vllm/envs.py

+6 -1

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

1444 "VLLM_NVFP4_GEMM_BACKEND": env_with_choices(

1445 "VLLM_NVFP4_GEMM_BACKEND",

1446 None,

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

1447 + [

1448 + "flashinfer-cudnn",

1449 + "flashinfer-trtllm",

1450 + "flashinfer-cutlass",

1451 + "cutlass",

1452 + ],

1453 ),

1454 # Controls garbage collection during CUDA graph capture.

1455 # 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

+7 -2

@@ -23,7 +23,10 @@ from vllm.model_executor.parameter import (

23 ModelWeightParameter,

24 PerTensorScaleParameter,

25 )

26 -from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer

26 +from vllm.utils.flashinfer import (

27 + flashinfer_scaled_fp4_mm,

28 + has_flashinfer,

29 +)

31 logger = init_logger(__name__)

@@ -187,7 +190,9 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):

190 output_shape = [*x.shape[:-1], layer.weight_packed.shape[0]]

191

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

190 - x_fp4, x_blockscale = scaled_fp4_quant(x, layer.input_global_scale)

193 + x_fp4, x_blockscale = scaled_fp4_quant(

194 + x, layer.input_global_scale, self.backend

195 + )

196

197 mm_args = (

198 x_fp4,

vllm/model_executor/layers/quantization/modelopt.py

+1 -1

@@ -1291,7 +1291,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):

1291 output_shape = [x.shape[0], layer.weight.shape[0]]

1292

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

1294 - x_fp4, x_blockscale = scaled_fp4_quant(x, layer.input_scale_inv)

1294 + x_fp4, x_blockscale = scaled_fp4_quant(x, layer.input_scale_inv, self.backend)

1295

1296 # validate dtypes of quantized input, input block scale,

1297 # weight and weight_blockscale

vllm/utils/flashinfer.py

+54 -1

@@ -406,12 +406,21 @@ if has_flashinfer():

406 B_scale: torch.Tensor,

407 g_scale: torch.Tensor,

408 dtype: torch.dtype,

409 + use_8x4_sf_layout: bool,

410 backend: str,

411 ) -> torch.Tensor:

412 from flashinfer import mm_fp4 as flashinfer_mm_fp4_

413

414 return flashinfer_mm_fp4_(

414 - A, B, A_scale, B_scale, g_scale, dtype, block_size=16, backend=backend

415 + A,

416 + B,

417 + A_scale,

418 + B_scale,

419 + g_scale,

420 + dtype,

421 + block_size=16,

422 + use_8x4_sf_layout=use_8x4_sf_layout,

423 + backend=backend,

424 )

425

426 @torch.library.register_fake(

@@ -424,6 +433,7 @@ if has_flashinfer():

433 B_scale: torch.Tensor,

434 g_scale: torch.Tensor,

435 dtype: torch.dtype,

436 + use_8x4_sf_layout: bool,

437 backend: str,

438 ) -> torch.Tensor:

439 return torch.empty(A.shape[0], B.shape[1], dtype=dtype, device=A.device)

@@ -460,6 +470,39 @@ if has_flashinfer():

470 A.shape[0], A.shape[1], B.shape[2], dtype=dtype, device=A.device

471 )

472

473 + @torch.library.custom_op(

474 + "vllm::flashinfer_nvfp4_quantize",

475 + mutates_args=[],

476 + device_types="cuda",

477 + )

478 + def flashinfer_nvfp4_quantize(

479 + a: torch.Tensor, a_global_sf: torch.Tensor

480 + ) -> tuple[torch.Tensor, torch.Tensor]:

481 + from flashinfer import SfLayout

482 + from flashinfer import nvfp4_quantize as nvfp4_quantize_

483 +

484 + return nvfp4_quantize_(

485 + a, a_global_sf, sfLayout=SfLayout.layout_8x4, do_shuffle=False

486 + )

487 +

488 + @torch.library.register_fake(

489 + "vllm::flashinfer_nvfp4_quantize",

490 + )

491 + def flashinfer_nvfp4_quantize_fake(

492 + a: torch.Tensor, a_global_sf: torch.Tensor

493 + ) -> tuple[torch.Tensor, torch.Tensor]:

494 + m, n = a.shape

495 +

496 + round_up = lambda x, y: (x + y - 1) // y * y

497 +

498 + rounded_m = round_up(m, 8)

499 + scale_n = n // 16

500 + rounded_n = round_up(scale_n, 4)

501 +

502 + return torch.empty(m, n // 2, dtype=torch.uint8, device=a.device), torch.empty(

503 + rounded_m, rounded_n, dtype=torch.uint8, device=a.device

504 + )

505 +

506

507 def flashinfer_scaled_fp4_mm(

508 a: torch.Tensor,

@@ -479,6 +522,8 @@ def flashinfer_scaled_fp4_mm(

522 block_scale_a = block_scale_a.view(torch.uint8)

523 block_scale_b = block_scale_b.view(torch.uint8)

524

525 + use_8x4_sf_layout = True if backend == "trtllm" and a.shape[0] <= 32 else False # noqa: SIM210

526 +

527 return flashinfer_mm_fp4(

528 a,

529 b.t(),

@@ -486,6 +531,7 @@ def flashinfer_scaled_fp4_mm(

531 block_scale_b.t(),

532 alpha,

533 out_dtype,

534 + use_8x4_sf_layout=use_8x4_sf_layout,

535 backend=backend,

536 )

537

@@ -520,6 +566,12 @@ def flashinfer_scaled_fp8_mm(

566 return output

567

568

569 +def flashinfer_quant_nvfp4_8x4_sf_layout(

570 + a: torch.Tensor, a_global_sf: torch.Tensor

571 +) -> tuple[torch.Tensor, torch.Tensor]:

572 + return flashinfer_nvfp4_quantize(a, a_global_sf)

573 +

574 +

575 flashinfer_fp8_blockscale_gemm = _lazy_import_wrapper(

576 "flashinfer.gemm", "fp8_blockscale_gemm_sm90"

577 )

@@ -596,6 +648,7 @@ __all__ = [

648 "use_trtllm_attention",

649 "flashinfer_scaled_fp4_mm",

650 "flashinfer_scaled_fp8_mm",

651 + "flashinfer_quant_nvfp4_8x4_sf_layout",

652 "flashinfer_fp8_blockscale_gemm",

653 "should_use_flashinfer_for_blockscale_fp8_gemm",

654 "is_flashinfer_fp8_blockscale_gemm_supported",