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 6, 14:30 UTC
[Custom Ops] Add functional + out variant for scaled_fp4_quant (#34389)
Signed-off-by: tianrengao <terrygao87@gmail.com>
Terry Gao committed
Mar 16, 2026 at 15:51 UTC
3e6a1e1686958dcd7eff1438bc5418b8d56daa30
12 files changed
+213
-44
csrc/ops.h
+8
-4
@@ -295,10 +295,14 @@ void cutlass_scaled_sparse_mm(torch::Tensor& out, torch::Tensor const& a,
295
296
std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a);
297
298
-void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
299
- torch::Tensor& output_scale,
300
- torch::Tensor const& input_scale,
301
- bool is_sf_swizzled_layout);
298
+std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
299
+ torch::Tensor const& input, torch::Tensor const& input_scale,
300
+ bool is_sf_swizzled_layout);
301
+
302
+void scaled_fp4_quant_out(torch::Tensor const& input,
303
+ torch::Tensor const& input_scale,
304
+ bool is_sf_swizzled_layout, torch::Tensor& output,
305
+ torch::Tensor& output_scale);
306
307
void scaled_fp4_experts_quant(
308
torch::Tensor& output, torch::Tensor& output_scale,
csrc/quantization/fp4/nvfp4_quant_entry.cu
+34
-3
@@ -16,6 +16,8 @@
16
17
#include <torch/all.h>
18
19
+#include "nvfp4_utils.cuh"
20
+
21
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
22
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
23
void scaled_fp4_quant_sm1xxa(torch::Tensor const& output,
@@ -51,9 +53,10 @@ void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
53
torch::Tensor const& output_scale_offset_by_experts);
54
#endif
55
54
-void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
55
- torch::Tensor& output_sf, torch::Tensor const& input_sf,
56
- bool is_sf_swizzled_layout) {
56
+void scaled_fp4_quant_out(torch::Tensor const& input,
57
+ torch::Tensor const& input_sf,
58
+ bool is_sf_swizzled_layout, torch::Tensor& output,
59
+ torch::Tensor& output_sf) {
60
#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
61
(defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
62
return scaled_fp4_quant_sm1xxa(output, input, output_sf, input_sf,
@@ -62,6 +65,34 @@ void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
65
TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
66
}
67
68
+std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
69
+ torch::Tensor const& input, torch::Tensor const& input_sf,
70
+ bool is_sf_swizzled_layout) {
71
+ int64_t n = input.size(-1);
72
+ int64_t m = input.numel() / n;
73
+ auto device = input.device();
74
+
75
+ // Two fp4 values packed into a uint8
76
+ auto output = torch::empty(
77
+ {m, n / 2}, torch::TensorOptions().device(device).dtype(torch::kUInt8));
78
+
79
+ torch::Tensor output_sf;
80
+ if (is_sf_swizzled_layout) {
81
+ auto [sf_m, sf_n] = vllm::computeSwizzledSFShape(m, n);
82
+ output_sf = torch::empty(
83
+ {sf_m, sf_n},
84
+ torch::TensorOptions().device(device).dtype(torch::kInt32));
85
+ } else {
86
+ output_sf = torch::empty(
87
+ {m, n / CVT_FP4_SF_VEC_SIZE},
88
+ torch::TensorOptions().device(device).dtype(torch::kUInt8));
89
+ }
90
+
91
+ scaled_fp4_quant_out(input, input_sf, is_sf_swizzled_layout, output,
92
+ output_sf);
93
+ return {output, output_sf};
94
+}
95
+
96
void scaled_fp4_experts_quant(
97
torch::Tensor& output, torch::Tensor& output_scale,
98
torch::Tensor const& input, torch::Tensor const& input_global_scale,
csrc/quantization/fp4/nvfp4_utils.cuh
+13
@@ -18,6 +18,7 @@
18
19
#include <cuda_runtime.h>
20
#include <cuda_fp8.h>
21
+#include <utility>
22
23
#include "../../cuda_vec_utils.cuh"
24
@@ -54,6 +55,18 @@ inline int computeEffectiveRows(int m) {
55
return round_up(m, ROW_TILE);
56
}
57
58
+// Compute the shape of the swizzled SF output tensor.
59
+// Returns (rounded_m, rounded_n / 4) where:
60
+// rounded_m = round_up(m, 128)
61
+// rounded_n = round_up(n / CVT_FP4_SF_VEC_SIZE, 4)
62
+inline std::pair<int64_t, int64_t> computeSwizzledSFShape(int64_t m,
63
+ int64_t n) {
64
+ int64_t rounded_m = round_up(m, static_cast<int64_t>(128));
65
+ int64_t scale_n = n / CVT_FP4_SF_VEC_SIZE;
66
+ int64_t rounded_n = round_up(scale_n, static_cast<int64_t>(4));
67
+ return {rounded_m, rounded_n / 4};
68
+}
69
+
70
// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
71
inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) {
72
uint32_t val;
csrc/torch_bindings.cpp
+15
-4
@@ -564,10 +564,21 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
564
565
// Compute NVFP4 block quantized tensor.
566
ops.def(
567
- "scaled_fp4_quant(Tensor! output, Tensor input,"
568
- " Tensor! output_scale, Tensor input_scale, bool "
569
- "is_sf_swizzled_layout) -> ()");
570
- ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);
567
+ "scaled_fp4_quant(Tensor input,"
568
+ " Tensor input_scale, bool "
569
+ "is_sf_swizzled_layout) -> (Tensor, Tensor)");
570
+ ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant_func);
571
+
572
+ // Out variant
573
+ // TODO: Add {at::Tag::out_variant} tag and update all call sites
574
+ // to use the functional variant once vLLM upgrades PyTorch.
575
+ // See pytorch/pytorch#176117.
576
+ ops.def(
577
+ "scaled_fp4_quant.out(Tensor input,"
578
+ " Tensor input_scale, bool "
579
+ "is_sf_swizzled_layout, *, Tensor(a!) output, Tensor(b!) output_scale) "
580
+ "-> ()");
581
+ ops.impl("scaled_fp4_quant.out", torch::kCUDA, &scaled_fp4_quant_out);
582
583
// Compute NVFP4 experts quantization.
584
ops.def(
tests/compile/passes/distributed/test_fusion_all_reduce.py
+1
-1
@@ -179,7 +179,7 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module):
179
def ops_in_model_before(self):
180
return [
181
torch.ops.vllm.all_reduce.default,
182
- torch.ops._C.scaled_fp4_quant.default,
182
+ torch.ops._C.scaled_fp4_quant.out,
183
]
184
185
tests/kernels/quantization/test_nvfp4_quant.py
+46
@@ -159,6 +159,52 @@ def test_quantize_to_fp4(
159
torch.testing.assert_close(scale_ans, scale_ref)
160
161
162
+@pytest.mark.parametrize(
163
+ "shape",
164
+ [(32, 4096), (128, 4096), (1, 64), (127, 1024), (256, 16384)],
165
+)
166
+@pytest.mark.parametrize("is_sf_swizzled_layout", [True, False])
167
+@torch.inference_mode()
168
+def test_python_util_matches_cpp_allocation(
169
+ shape: tuple[int, int],
170
+ is_sf_swizzled_layout: bool,
171
+) -> None:
172
+ """
173
+ Verify that the Python utility (create_fp4_output_tensors) allocates
174
+ tensors with the same shapes and dtypes as the C++ functional variant
175
+ (scaled_fp4_quant_func).
176
+ """
177
+ from vllm._custom_ops import create_fp4_output_tensors
178
+
179
+ torch.set_default_device("cuda:0")
180
+ m, n = shape
181
+ input_tensor = torch.randn((m, n), dtype=torch.bfloat16)
182
+ input_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda:0")
183
+
184
+ # C++ functional variant allocates internally
185
+ cpp_out, cpp_scale = torch.ops._C.scaled_fp4_quant(
186
+ input_tensor, input_scale, is_sf_swizzled_layout
187
+ )
188
+
189
+ # Python utility
190
+ py_out, py_scale = create_fp4_output_tensors(
191
+ m, n, torch.device("cuda:0"), is_sf_swizzled_layout
192
+ )
193
+
194
+ assert py_out.shape == cpp_out.shape, (
195
+ f"Output shape mismatch: Python {py_out.shape} vs C++ {cpp_out.shape}"
196
+ )
197
+ assert py_out.dtype == cpp_out.dtype, (
198
+ f"Output dtype mismatch: Python {py_out.dtype} vs C++ {cpp_out.dtype}"
199
+ )
200
+ assert py_scale.shape == cpp_scale.shape, (
201
+ f"Scale shape mismatch: Python {py_scale.shape} vs C++ {cpp_scale.shape}"
202
+ )
203
+ assert py_scale.dtype == cpp_scale.dtype, (
204
+ f"Scale dtype mismatch: Python {py_scale.dtype} vs C++ {cpp_scale.dtype}"
205
+ )
206
+
207
+
208
@pytest.mark.parametrize("pad_shape", PAD_SHAPES)
209
@torch.inference_mode()
210
def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None:
vllm/_custom_ops.py
+85
-21
@@ -29,6 +29,81 @@ else:
29
from torch.library import impl_abstract as register_fake
30
31
32
+# scaled_fp4_quant functional + out variant for torch.compile buffer management
33
+
34
+
35
+def create_fp4_scale_tensor(
36
+ m: int,
37
+ n: int,
38
+ device: torch.device,
39
+ is_sf_swizzled_layout: bool,
40
+) -> torch.Tensor:
41
+ """
42
+ Allocate the output scale tensor for scaled_fp4_quant.
43
+
44
+ When is_sf_swizzled_layout=True, we use rounded values to store the
45
+ swizzled scales. Due to the requirement of the Tensor Core, the minimum
46
+ tile is 128x4 for the scales. So, we first pad the scales to multiples
47
+ of 128 (rows) and 4 (cols). Then, the scales (in float8_e4m3fn) are
48
+ packed into an int32 for every 4 values. More:
49
+ https://docs.nvidia.com/cuda/parallel-thread-execution/
50
+ #tcgen05-mma-scale-factor-b-layout-4x
51
+ """
52
+ from vllm.utils.math_utils import round_up
53
+
54
+ block_size = 16
55
+ if is_sf_swizzled_layout:
56
+ rounded_m = round_up(m, 128)
57
+ scale_n = n // block_size
58
+ rounded_n = round_up(scale_n, 4)
59
+ return torch.empty(
60
+ (rounded_m, rounded_n // 4), device=device, dtype=torch.int32
61
+ )
62
+ else:
63
+ return torch.empty((m, n // block_size), device=device, dtype=torch.uint8)
64
+
65
+
66
+def create_fp4_output_tensors(
67
+ m: int,
68
+ n: int,
69
+ device: torch.device,
70
+ is_sf_swizzled_layout: bool,
71
+) -> tuple[torch.Tensor, torch.Tensor]:
72
+ """
73
+ Allocate both output tensors for scaled_fp4_quant:
74
+ (quantized_output, output_scale).
75
+
76
+ Must match the C++ scaled_fp4_quant_func allocation exactly.
77
+ """
78
+ output = torch.empty((m, n // 2), device=device, dtype=torch.uint8)
79
+ output_scale = create_fp4_scale_tensor(m, n, device, is_sf_swizzled_layout)
80
+ return output, output_scale
81
+
82
+
83
+if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "scaled_fp4_quant"):
84
+
85
+ @register_fake("_C::scaled_fp4_quant")
86
+ def _scaled_fp4_quant_fake(
87
+ input: torch.Tensor,
88
+ input_scale: torch.Tensor,
89
+ is_sf_swizzled_layout: bool,
90
+ ) -> tuple[torch.Tensor, torch.Tensor]:
91
+ n = input.shape[-1]
92
+ m = input.numel() // n
93
+ return create_fp4_output_tensors(m, n, input.device, is_sf_swizzled_layout)
94
+
95
+ @register_fake("_C::scaled_fp4_quant.out")
96
+ def _scaled_fp4_quant_out_fake(
97
+ input: torch.Tensor,
98
+ input_scale: torch.Tensor,
99
+ is_sf_swizzled_layout: bool,
100
+ *,
101
+ output: torch.Tensor,
102
+ output_scale: torch.Tensor,
103
+ ) -> None:
104
+ return None
105
+
106
+
107
# page attention ops
108
def paged_attention_v1(
109
out: torch.Tensor,
@@ -1644,7 +1719,6 @@ def scaled_fp4_quant(
1719
input = input.reshape(other_dims, input.shape[-1])
1720
m, n = input.shape
1721
block_size = 16
1647
- device = input.device
1722
1723
assert n % block_size == 0, f"last dim has to be multiple of 16, but got {n}."
1724
assert input.dtype in (torch.float16, torch.bfloat16), (
@@ -1658,26 +1732,16 @@ def scaled_fp4_quant(
1732
input, input_global_scale
1733
)
1734
else:
1661
- # Two fp4 values will be packed into an uint8.
1662
- output = torch.empty((m, n // 2), device=device, dtype=torch.uint8)
1663
- if is_sf_swizzled_layout:
1664
- # We use the rounded values to store the swizzled values. Due to the
1665
- # requirement of the Tensor Core, the minimum tile is 128x4 for the scales.
1666
- # So, we first pad the scales to multiples of 128 and 4. Then, the scales
1667
- # (in float8_e4m3fn) are packed into an int32 for every 4 values. More:
1668
- # https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-scale-factor-b-layout-4x
1669
- round_up = lambda x, y: (x + y - 1) // y * y
1670
- rounded_m = round_up(m, 128)
1671
- scale_n = n // block_size
1672
- rounded_n = round_up(scale_n, 4)
1673
- output_scale = torch.empty(
1674
- (rounded_m, rounded_n // 4), device=device, dtype=torch.int32
1675
- )
1676
- else:
1677
- output_scale = torch.empty((m, n // 16), device=device, dtype=torch.uint8)
1678
-
1679
- torch.ops._C.scaled_fp4_quant(
1680
- output, input, output_scale, input_global_scale, is_sf_swizzled_layout
1735
+ # Pre-allocate and call .out variant (same behavior as old in-place API)
1736
+ output, output_scale = create_fp4_output_tensors(
1737
+ m, n, input.device, is_sf_swizzled_layout
1738
+ )
1739
+ torch.ops._C.scaled_fp4_quant.out(
1740
+ input,
1741
+ input_global_scale,
1742
+ is_sf_swizzled_layout,
1743
+ output=output,
1744
+ output_scale=output_scale,
1745
)
1746
1747
output_scale = output_scale.view(torch.float8_e4m3fn)
vllm/compilation/passes/fusion/act_quant_fusion.py
+2
-2
@@ -148,11 +148,11 @@ class SiluMulNvfp4QuantPattern(ActivationQuantPattern):
148
result_silu_mul = self.silu_and_mul_matcher(input)
149
at = auto_functionalized(
150
self.QUANT_OP,
151
- output=result,
151
input=result_silu_mul,
153
- output_scale=output_scale,
152
input_scale=scale,
153
is_sf_swizzled_layout=True,
154
+ output=result,
155
+ output_scale=output_scale,
156
)
157
return at[1], at[2]
158
vllm/compilation/passes/fusion/allreduce_rms_fusion.py
+5
-5
@@ -47,7 +47,7 @@ if find_spec("flashinfer"):
47
pass
48
49
if hasattr(torch.ops._C, "scaled_fp4_quant"):
50
- STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.default
50
+ STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.out
51
52
# Max size of the input tensor per world size per device capability
53
# to use flashinfer fused allreduce
@@ -562,11 +562,11 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):
562
rms = self.rmsnorm_matcher(all_reduce, weight)
563
quant_out_tuple = auto_functionalized(
564
STATIC_FP4_QUANT_OP,
565
- output=quant_result,
565
input=rms,
567
- output_scale=output_scale,
566
input_scale=input_global_scale,
567
is_sf_swizzled_layout=True,
568
+ output=quant_result,
569
+ output_scale=output_scale,
570
)
571
572
# quant_out, allreduce_output, output_scale
@@ -660,11 +660,11 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):
660
rms, residual = self.rmsnorm_matcher(allreduce_output, weight, residual)
661
quant_out_tuple = auto_functionalized(
662
STATIC_FP4_QUANT_OP,
663
- output=quant_result,
663
input=rms,
665
- output_scale=output_scale,
664
input_scale=input_global_scale,
665
is_sf_swizzled_layout=True,
666
+ output=quant_result,
667
+ output_scale=output_scale,
668
)
669
670
# quant_out, allreduce_output, output_scale
vllm/compilation/passes/fusion/attn_quant_fusion.py
+2
-2
@@ -250,11 +250,11 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
250
)
251
at2 = auto_functionalized(
252
self.QUANT_OP,
253
- output=output_quant,
253
input=attn_out_view,
255
- output_scale=output_scale,
254
input_scale=input_scale,
255
is_sf_swizzled_layout=True,
256
+ output=output_quant,
257
+ output_scale=output_scale,
258
)
259
output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)
260
return at2[1], output_scale_view
vllm/compilation/passes/fusion/matcher_utils.py
+1
-1
@@ -38,7 +38,7 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {
38
}
39
40
if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
41
- QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.default # noqa: E501
41
+ QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out # noqa: E501
42
43
if current_platform.is_cuda():
44
QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501
vllm/compilation/passes/fusion/rms_quant_fusion.py
+1
-1
@@ -63,7 +63,7 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {
63
kFp8DynamicTokenSym: torch.ops._C.dynamic_per_token_scaled_fp8_quant.default, # noqa: E501
64
}
65
if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
66
- QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.default
66
+ QUANT_OPS[kNvfp4Dynamic] = torch.ops._C.scaled_fp4_quant.out
67
if current_platform.is_cuda():
68
QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501
69
QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501