@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2917d6363a
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 4, 11:15 UTC

[NVFP4][Hopper/AMD Instinct] Add Triton kernels for NVFP4 dequantization and QDQ emulation (#40033)

Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Claude <noreply@anthropic.com>

fxmarty-amd committed Apr 30, 2026 at 23:35 UTC 2917d6363ad722bff647168d3a36261254d7ad42
4 files changed +633 -4
tests/kernels/quantization/test_nvfp4_emulation.py new
+308
@@ -0,0 +1,308 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import huggingface_hub
4 +import pytest
5 +import torch
6 +from safetensors import safe_open
7 +
8 +from vllm.model_executor.layers.quantization.utils import (
9 + nvfp4_emulation_utils,
10 +)
11 +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
12 + dequantize_to_dtype,
13 + ref_nvfp4_quant_dequant,
14 +)
15 +from vllm.platforms import current_platform
16 +from vllm.triton_utils import triton
17 +
18 +
19 +@pytest.mark.skipif(
20 + not current_platform.is_cuda_alike(),
21 + reason="Triton NVFP4 kernel requires CUDA.",
22 +)
23 +def test_triton_dequantize_nvfp4(monkeypatch) -> None:
24 + """Test the Triton dequantization kernel against the CPU reference
25 + using real NVFP4 weights from a checkpoint.
26 +
27 + Tests both 2D (attention projection) and 3D (stacked MoE experts).
28 + """
29 + checkpoint_path = huggingface_hub.snapshot_download(
30 + "nvidia/Qwen3-30B-A3B-NVFP4",
31 + allow_patterns=["model-00001-of-00004.safetensors"],
32 + )
33 + shard_path = f"{checkpoint_path}/model-00001-of-00004.safetensors"
34 + block_size = 16
35 +
36 + with safe_open(shard_path, framework="pt", device="cpu") as f:
37 + all_keys = list(f.keys())
38 +
39 + # 2D case: attention projection
40 + tensor_fp4_2d = f.get_tensor("model.layers.9.self_attn.k_proj.weight")
41 + tensor_sf_2d = f.get_tensor("model.layers.9.self_attn.k_proj.weight_scale")
42 + global_scale_2d = f.get_tensor("model.layers.9.self_attn.k_proj.weight_scale_2")
43 +
44 + # 3D case: stack ALL experts for layer 9 up_proj
45 + expert_prefix = "model.layers.9.mlp.experts."
46 + expert_indices = sorted(
47 + int(key.split(".")[5])
48 + for key in all_keys
49 + if key.startswith(expert_prefix) and key.endswith(".up_proj.weight")
50 + )
51 + assert len(expert_indices) > 0
52 +
53 + all_fp4 = []
54 + all_sf = []
55 + all_global_scale = []
56 + for index in expert_indices:
57 + name = f"{expert_prefix}{index}.up_proj"
58 + all_fp4.append(f.get_tensor(f"{name}.weight"))
59 + all_sf.append(f.get_tensor(f"{name}.weight_scale"))
60 + all_global_scale.append(f.get_tensor(f"{name}.weight_scale_2"))
61 +
62 + tensor_fp4_3d = torch.stack(all_fp4)
63 + tensor_sf_3d = torch.stack(all_sf)
64 + global_scale_3d = torch.stack(all_global_scale)
65 +
66 + test_cases = [
67 + ("2D base", tensor_fp4_2d, tensor_sf_2d, global_scale_2d),
68 + (
69 + "2D 2x rows",
70 + tensor_fp4_2d.repeat(2, 1),
71 + tensor_sf_2d.repeat(2, 1),
72 + global_scale_2d,
73 + ),
74 + (
75 + "2D 4x rows",
76 + tensor_fp4_2d.repeat(4, 1),
77 + tensor_sf_2d.repeat(4, 1),
78 + global_scale_2d,
79 + ),
80 + (
81 + "2D 2x cols",
82 + tensor_fp4_2d.repeat(1, 2),
83 + tensor_sf_2d.repeat(1, 2),
84 + global_scale_2d,
85 + ),
86 + ("3D base", tensor_fp4_3d, tensor_sf_3d, global_scale_3d),
87 + (
88 + "3D 2x experts",
89 + tensor_fp4_3d.repeat(2, 1, 1),
90 + tensor_sf_3d.repeat(2, 1, 1),
91 + global_scale_3d.repeat(2),
92 + ),
93 + (
94 + "3D 2x rows",
95 + tensor_fp4_3d.repeat(1, 2, 1),
96 + tensor_sf_3d.repeat(1, 2, 1),
97 + global_scale_3d,
98 + ),
99 + (
100 + "3D 2x cols",
101 + tensor_fp4_3d.repeat(1, 1, 2),
102 + tensor_sf_3d.repeat(1, 1, 2),
103 + global_scale_3d,
104 + ),
105 + ]
106 +
107 + quantiles = [0.5, 0.001, 0.999]
108 +
109 + # Move the E2M1 lookup table to CUDA ahead of time, as would normally
110 + # happen during model loading (process_weights_after_loading). Both the
111 + # Triton and PyTorch reference paths run on CUDA.
112 + nvfp4_emulation_utils.kE2M1ToFloat_handle.val = (
113 + nvfp4_emulation_utils.kE2M1ToFloat_handle.val.cuda()
114 + )
115 +
116 + for label, tensor_fp4, tensor_sf, global_scale in test_cases:
117 + fp4_cuda = tensor_fp4.cuda()
118 + sf_cuda = tensor_sf.cuda()
119 + gs_cuda = global_scale.cuda()
120 +
121 + # Triton path
122 + triton_result = dequantize_to_dtype(
123 + fp4_cuda,
124 + sf_cuda,
125 + gs_cuda,
126 + torch.bfloat16,
127 + block_size,
128 + swizzle=False,
129 + )
130 +
131 + # Reference path (PyTorch ops on CUDA, Triton dispatch disabled)
132 + with monkeypatch.context() as m:
133 + m.setattr(
134 + nvfp4_emulation_utils.current_platform,
135 + "is_cuda_alike",
136 + lambda: False,
137 + )
138 + reference = dequantize_to_dtype(
139 + fp4_cuda,
140 + sf_cuda,
141 + gs_cuda,
142 + torch.bfloat16,
143 + block_size,
144 + swizzle=False,
145 + )
146 +
147 + torch.testing.assert_close(triton_result, reference, atol=0, rtol=0)
148 +
149 + # Benchmark
150 + shape = list(tensor_fp4.shape)
151 +
152 + def _triton_bench(
153 + fp4_cuda=fp4_cuda,
154 + scale_cuda=sf_cuda,
155 + global_scale_cuda=gs_cuda,
156 + block_size=block_size,
157 + ):
158 + return dequantize_to_dtype(
159 + fp4_cuda,
160 + scale_cuda,
161 + global_scale_cuda,
162 + torch.bfloat16,
163 + block_size,
164 + swizzle=False,
165 + )
166 +
167 + triton_ms, triton_min, triton_max = triton.testing.do_bench(
168 + _triton_bench, quantiles=quantiles
169 + )
170 +
171 + def _reference_bench(
172 + fp4_cuda=fp4_cuda,
173 + scale_cuda=sf_cuda,
174 + global_scale_cuda=gs_cuda,
175 + block_size=block_size,
176 + ):
177 + with monkeypatch.context() as m2:
178 + m2.setattr(
179 + nvfp4_emulation_utils.current_platform,
180 + "is_cuda_alike",
181 + lambda: False,
182 + )
183 + dequantize_to_dtype(
184 + fp4_cuda,
185 + scale_cuda,
186 + global_scale_cuda,
187 + torch.bfloat16,
188 + block_size,
189 + swizzle=False,
190 + )
191 +
192 + ref_ms, ref_min, ref_max = triton.testing.do_bench(
193 + _reference_bench, quantiles=quantiles
194 + )
195 +
196 + speedup = ref_ms / triton_ms if triton_ms > 0 else float("inf")
197 + print(f" dequantize {label} {shape}:")
198 + print(
199 + f" triton: median={triton_ms:.3f}ms, "
200 + f"min={triton_min:.3f}ms, max={triton_max:.3f}ms"
201 + )
202 + print(
203 + f" reference: median={ref_ms:.3f}ms, "
204 + f"min={ref_min:.3f}ms, max={ref_max:.3f}ms"
205 + )
206 + print(f" speedup: {speedup:.2f}x")
207 +
208 +
209 +@pytest.mark.skipif(
210 + not current_platform.is_cuda_alike(),
211 + reason="Triton NVFP4 kernel requires CUDA.",
212 +)
213 +@pytest.mark.parametrize(
214 + "m, k",
215 + [
216 + (1, 16),
217 + (1, 4096),
218 + (2, 4096),
219 + (4, 4096),
220 + (8, 4096),
221 + (16, 4096),
222 + (24, 4096),
223 + (32, 4096),
224 + (1, 8192),
225 + (2, 8192),
226 + (4, 8192),
227 + (8, 8192),
228 + (16, 8192),
229 + (24, 8192),
230 + (32, 8192),
231 + (1, 32),
232 + (2, 48),
233 + (7, 64),
234 + (16, 128),
235 + (33, 160),
236 + (128, 256),
237 + (256, 512),
238 + (1024, 1024),
239 + (5120, 2048),
240 + (2048, 4096),
241 + (4096, 7168),
242 + (8192, 8192),
243 + (128, 16384),
244 + ],
245 +)
246 +@pytest.mark.parametrize("global_scale_value", [0.5, 1.0, 0.001])
247 +def test_triton_nvfp4_quant_dequant(
248 + monkeypatch, m: int, k: int, global_scale_value: float
249 +) -> None:
250 + """Test the Triton quant-dequant kernel against the CPU reference."""
251 + block_size = 16
252 + x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
253 + global_scale = torch.tensor(global_scale_value, dtype=torch.float32, device="cuda")
254 +
255 + # Triton path
256 + triton_result = ref_nvfp4_quant_dequant(x, global_scale, block_size)
257 +
258 + # CPU reference path
259 + with monkeypatch.context() as mp:
260 + mp.setattr(
261 + nvfp4_emulation_utils.current_platform,
262 + "is_cuda_alike",
263 + lambda: False,
264 + )
265 + reference = ref_nvfp4_quant_dequant(x.cpu(), global_scale.cpu(), block_size)
266 +
267 + torch.testing.assert_close(triton_result.cpu(), reference, atol=0, rtol=0)
268 +
269 + # Benchmark (both paths on CUDA tensors for fair comparison)
270 + quantiles = [0.5, 0.001, 0.999]
271 +
272 + def _triton_bench(
273 + input_tensor=x, input_global_scale=global_scale, input_block_size=block_size
274 + ):
275 + return ref_nvfp4_quant_dequant(
276 + input_tensor, input_global_scale, input_block_size
277 + )
278 +
279 + triton_ms, triton_min, triton_max = triton.testing.do_bench(
280 + _triton_bench, quantiles=quantiles
281 + )
282 +
283 + def _reference_bench(
284 + input_tensor=x, input_global_scale=global_scale, input_block_size=block_size
285 + ):
286 + with monkeypatch.context() as mp2:
287 + mp2.setattr(
288 + nvfp4_emulation_utils.current_platform,
289 + "is_cuda_alike",
290 + lambda: False,
291 + )
292 + ref_nvfp4_quant_dequant(input_tensor, input_global_scale, input_block_size)
293 +
294 + ref_ms, ref_min, ref_max = triton.testing.do_bench(
295 + _reference_bench, quantiles=quantiles
296 + )
297 +
298 + speedup = ref_ms / triton_ms if triton_ms > 0 else float("inf")
299 + print(f" quant_dequant [{m}x{k}] gs={global_scale_value}:")
300 + print(
301 + f" triton: median={triton_ms:.3f}ms, "
302 + f"min={triton_min:.3f}ms, max={triton_max:.3f}ms"
303 + )
304 + print(
305 + f" reference: median={ref_ms:.3f}ms, "
306 + f"min={ref_min:.3f}ms, max={ref_max:.3f}ms"
307 + )
308 + print(f" speedup: {speedup:.2f}x")
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
+7
@@ -31,6 +31,9 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
31 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
32 prepare_nvfp4_moe_layer_for_marlin,
33 )
34 +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
35 + kE2M1ToFloat_handle,
36 +)
37 from vllm.model_executor.layers.quantization.utils.quant_utils import (
38 QuantKey,
39 )
@@ -376,6 +379,10 @@ def convert_to_nvfp4_moe_kernel_format(
379 is_act_and_mul=is_act_and_mul,
380 )
381 elif nvfp4_backend == NvFp4MoeBackend.EMULATION:
382 + # Move the E2M1 lookup table to the device now, because
383 + # `.to(device)` is not allowed during CUDA graph capture.
384 + kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(layer.weight.device)
385 +
386 if a13_scale is None or a2_scale is None:
387 raise ValueError(
388 "Activation global scales should not be None, got"
vllm/model_executor/layers/fused_moe/utils.py
+2 -1
@@ -298,7 +298,8 @@ def moe_kernel_quantize_input(
298 A, A_scale, is_sf_swizzled_layout=is_fp4_scale_swizzled
299 )
300 else:
301 - return ref_nvfp4_quant_dequant(A, A_scale, block_size=16)
301 + A = ref_nvfp4_quant_dequant(A, A_scale, block_size=16)
302 + return A, None
303 elif quant_dtype == "mxfp4":
304 if not quantization_emulation:
305 raise NotImplementedError(
vllm/model_executor/layers/quantization/utils/nvfp4_emulation_utils.py
+316 -3
@@ -4,7 +4,9 @@ from types import SimpleNamespace
4
5 import torch
6
7 +from vllm.platforms import current_platform
8 from vllm.scalar_type import scalar_types
9 +from vllm.triton_utils import tl, triton
10
11 __all__ = [
12 "break_fp4_bytes",
@@ -20,6 +22,309 @@ kE2M1ToFloat_handle = SimpleNamespace(
22 )
23
24
25 +@triton.jit
26 +def _e2m1_inline(magnitude):
27 + """Inline E2M1 lookup using binary tree - 3 levels instead of 7 sequential.
28 +
29 + Maps 3-bit magnitude to float: [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
30 + Uses bit decomposition for fewer comparisons.
31 + """
32 + # Bit 2 (MSB): separates 0-3 from 4-7
33 + # Bit 1: separates within groups
34 + # Bit 0 (LSB): separates within pairs
35 + b2 = (magnitude >> 2) & 1 # 0 for mag 0-3, 1 for mag 4-7
36 + b1 = (magnitude >> 1) & 1 # middle bit
37 + b0 = magnitude & 1 # LSB
38 +
39 + # For mag 0-3: [0.0, 0.5, 1.0, 1.5]
40 + low_group = tl.where(
41 + b1 == 1, tl.where(b0 == 1, 1.5, 1.0), tl.where(b0 == 1, 0.5, 0.0)
42 + )
43 + # For mag 4-7: [2.0, 3.0, 4.0, 6.0]
44 + high_group = tl.where(
45 + b1 == 1, tl.where(b0 == 1, 6.0, 4.0), tl.where(b0 == 1, 3.0, 2.0)
46 + )
47 + return tl.where(b2 == 1, high_group, low_group)
48 +
49 +
50 +@triton.jit
51 +def _dequantize_nvfp4_kernel(
52 + fp4_ptr,
53 + scale_ptr,
54 + global_scale_ptr,
55 + output_ptr,
56 + rows_per_batch: tl.constexpr,
57 + num_blocks: tl.constexpr,
58 + BLOCK_SIZE: tl.constexpr,
59 + has_batch_global_scale: tl.constexpr,
60 + TILE_BLOCKS: tl.constexpr,
61 +):
62 + """Triton kernel for NVFP4 dequantization (swizzle=False).
63 +
64 + Optimized with 2D tile processing + interleave for coalesced stores.
65 + """
66 + BLOCK_PACKED: tl.constexpr = BLOCK_SIZE // 2
67 +
68 + row_idx = tl.program_id(0)
69 + tile_idx = tl.program_id(1)
70 +
71 + if has_batch_global_scale:
72 + batch_idx = row_idx // rows_per_batch
73 + global_scale = tl.load(global_scale_ptr + batch_idx).to(tl.float32)
74 + else:
75 + global_scale = tl.load(global_scale_ptr).to(tl.float32)
76 +
77 + fp4_row_offset = row_idx * num_blocks * BLOCK_PACKED
78 + scale_row_offset = row_idx * num_blocks
79 + output_row_offset = row_idx * num_blocks * BLOCK_SIZE
80 +
81 + start_block = tile_idx * TILE_BLOCKS
82 +
83 + # Load scales for this tile: [TILE_BLOCKS]
84 + block_offsets = tl.arange(0, TILE_BLOCKS)
85 + block_mask = (start_block + block_offsets) < num_blocks
86 +
87 + raw_scales = tl.load(
88 + scale_ptr + scale_row_offset + start_block + block_offsets,
89 + mask=block_mask,
90 + other=0,
91 + )
92 + scale_f32 = tl.cast(raw_scales, tl.float8e4nv, bitcast=True).to(tl.float32)
93 + scale_values = (scale_f32 * global_scale)[:, None]
94 +
95 + # Load [TILE_BLOCKS, BLOCK_PACKED] packed bytes
96 + packed_offsets = tl.arange(0, BLOCK_PACKED)[None, :]
97 + byte_indices = (
98 + fp4_row_offset
99 + + (start_block + block_offsets[:, None]) * BLOCK_PACKED
100 + + packed_offsets
101 + )
102 + elem_mask = block_mask[:, None]
103 + raw_bytes = tl.load(fp4_ptr + byte_indices, mask=elem_mask, other=0)
104 +
105 + low_nibble = raw_bytes & 0x0F
106 + high_nibble = (raw_bytes >> 4) & 0x0F
107 +
108 + # Binary tree E2M1 decode
109 + low_mag = low_nibble & 0x07
110 + low_val = _e2m1_inline(low_mag)
111 + low_sign = (low_nibble >> 3) & 1
112 + low_result = tl.where(low_sign == 1, -low_val, low_val) * scale_values
113 +
114 + high_mag = high_nibble & 0x07
115 + high_val = _e2m1_inline(high_mag)
116 + high_sign = (high_nibble >> 3) & 1
117 + high_result = tl.where(high_sign == 1, -high_val, high_val) * scale_values
118 +
119 + # Interleave for coalesced contiguous store
120 + result = tl.interleave(low_result, high_result)
121 +
122 + elem_offsets = tl.arange(0, BLOCK_SIZE)[None, :]
123 + out_indices = (
124 + output_row_offset
125 + + (start_block + block_offsets[:, None]) * BLOCK_SIZE
126 + + elem_offsets
127 + )
128 + tl.store(output_ptr + out_indices, result, mask=block_mask[:, None])
129 +
130 +
131 +@triton.jit
132 +def _e2m1_lookup(magnitude):
133 + """Lookup E2M1 float value from 3-bit magnitude."""
134 + result = tl.where(magnitude == 1, 0.5, 0.0)
135 + result = tl.where(magnitude == 2, 1.0, result)
136 + result = tl.where(magnitude == 3, 1.5, result)
137 + result = tl.where(magnitude == 4, 2.0, result)
138 + result = tl.where(magnitude == 5, 3.0, result)
139 + result = tl.where(magnitude == 6, 4.0, result)
140 + result = tl.where(magnitude == 7, 6.0, result)
141 + return result
142 +
143 +
144 +@triton.jit
145 +def _round_to_fp4(x):
146 + """Round float values to the nearest E2M1 representable value.
147 +
148 + Matches the thresholds in the Python ``cast_to_fp4`` exactly.
149 + """
150 + sign = tl.where(x < 0.0, -1.0, 1.0)
151 + abs_x = tl.abs(x)
152 + result = tl.where(abs_x > 5.0, 6.0, 0.0)
153 + result = tl.where((abs_x >= 3.5) & (abs_x <= 5.0), 4.0, result)
154 + result = tl.where((abs_x > 2.5) & (abs_x < 3.5), 3.0, result)
155 + result = tl.where((abs_x >= 1.75) & (abs_x <= 2.5), 2.0, result)
156 + result = tl.where((abs_x > 1.25) & (abs_x < 1.75), 1.5, result)
157 + result = tl.where((abs_x >= 0.75) & (abs_x <= 1.25), 1.0, result)
158 + result = tl.where((abs_x > 0.25) & (abs_x < 0.75), 0.5, result)
159 + return result * sign
160 +
161 +
162 +@triton.jit
163 +def _nvfp4_quant_dequant_kernel(
164 + input_ptr,
165 + output_ptr,
166 + global_scale_ptr,
167 + k: tl.constexpr,
168 + num_blocks: tl.constexpr,
169 + BLOCK_SIZE: tl.constexpr,
170 + FP4_MAX_RECIPROCAL: tl.constexpr,
171 + TILE_BLOCKS: tl.constexpr,
172 +):
173 + """Fused NVFP4 quantize-dequantize kernel.
174 +
175 + Uses a 2D grid (rows x tiles) to parallelize across both rows
176 + and quantization groups within a row. Each program handles
177 + TILE_BLOCKS groups at once using vectorized 2D operations.
178 + """
179 + row_idx = tl.program_id(0)
180 + tile_idx = tl.program_id(1)
181 + global_scale = tl.load(global_scale_ptr).to(tl.float32)
182 + row_offset = row_idx * k
183 +
184 + start_block = tile_idx * TILE_BLOCKS
185 + block_offsets = tl.arange(0, TILE_BLOCKS)
186 + block_mask = (start_block + block_offsets) < num_blocks
187 +
188 + # Load [TILE_BLOCKS, BLOCK_SIZE] elements
189 + indices = (
190 + row_offset
191 + + (start_block + block_offsets[:, None]) * BLOCK_SIZE
192 + + tl.arange(0, BLOCK_SIZE)[None, :]
193 + )
194 + mask_2d = block_mask[:, None]
195 + x = tl.load(input_ptr + indices, mask=mask_2d, other=0.0).to(tl.float32)
196 +
197 + # Per-group scale: [TILE_BLOCKS]
198 + vec_max = tl.max(tl.abs(x), axis=1)
199 + scale = global_scale * (vec_max * FP4_MAX_RECIPROCAL)
200 + scale = tl.clamp(scale, -448.0, 448.0)
201 + scale = scale.to(tl.float8e4nv).to(tl.float32)
202 +
203 + # Safe reciprocal, broadcast to [TILE_BLOCKS, 1]
204 + output_scale = tl.where(scale == 0.0, 0.0, global_scale / scale)[:, None]
205 +
206 + # Quantize: scale, clamp, round to FP4
207 + scaled_x = tl.clamp(x * output_scale, -6.0, 6.0)
208 + fp4_val = _round_to_fp4(scaled_x)
209 +
210 + # Dequantize: fp4_val * (scale / global_scale)
211 + dequant_scale = (scale / global_scale)[:, None]
212 + result = fp4_val * dequant_scale
213 +
214 + tl.store(output_ptr + indices, result, mask=mask_2d)
215 +
216 +
217 +def _triton_nvfp4_quant_dequant(
218 + x: torch.Tensor,
219 + global_scale: torch.Tensor,
220 + block_size: int,
221 +) -> torch.Tensor:
222 + """Triton-accelerated NVFP4 quantize-dequantize."""
223 + x_m, x_k = x.shape
224 +
225 + if not torch.compiler.is_compiling():
226 + assert x_k % block_size == 0, (
227 + f"Weight shape K={x_k} is not divisible by block_size={block_size}"
228 + )
229 +
230 + output_dtype = x.dtype
231 + num_blocks = x_k // block_size
232 +
233 + output = torch.empty(x_m, x_k, dtype=output_dtype, device=x.device)
234 +
235 + tile_blocks = min(64, triton.next_power_of_2(num_blocks))
236 + num_tiles = (num_blocks + tile_blocks - 1) // tile_blocks
237 + grid = (x_m, num_tiles)
238 + _nvfp4_quant_dequant_kernel[grid](
239 + x,
240 + output,
241 + global_scale,
242 + x_k,
243 + num_blocks,
244 + block_size,
245 + FLOAT4_E2M1_MAX_RECIPROCAL,
246 + tile_blocks,
247 + )
248 +
249 + return output
250 +
251 +
252 +def _triton_dequantize_nvfp4(
253 + tensor_fp4: torch.Tensor,
254 + tensor_sf: torch.Tensor,
255 + global_scale: torch.Tensor,
256 + dtype: torch.dtype,
257 + block_size: int = 16,
258 +) -> torch.Tensor:
259 + """Dequantize NVFP4 using Triton (swizzle=False only).
260 +
261 + Supports both 2D and 3D inputs:
262 + - 2D: [m, packed_k] -> [m, k]
263 + - 3D: [dim0, m, packed_k] -> [dim0, m, k]
264 + """
265 + assert tensor_fp4.dtype == torch.uint8
266 +
267 + is_3d = tensor_fp4.ndim == 3
268 + if is_3d:
269 + dim0, m_per_batch, packed_k = tensor_fp4.shape
270 + tensor_fp4_2d = tensor_fp4.reshape(-1, packed_k)
271 + tensor_sf_2d = tensor_sf.reshape(-1, tensor_sf.shape[-1])
272 + total_rows_flat = dim0 * m_per_batch
273 + else:
274 + m_per_batch, packed_k = tensor_fp4.shape
275 + tensor_fp4_2d = tensor_fp4
276 + tensor_sf_2d = tensor_sf
277 + total_rows_flat = m_per_batch
278 +
279 + k = packed_k * 2
280 + num_blocks = k // block_size
281 +
282 + output = torch.empty(total_rows_flat, k, dtype=dtype, device=tensor_fp4.device)
283 +
284 + # View as uint8 so Triton can load raw bytes and bitcast to float8_e4m3fn
285 + scale_raw = tensor_sf_2d.contiguous().view(torch.uint8)
286 +
287 + # Shape-adaptive tile sizing: for large row counts (3D), process
288 + # entire row in one tile. For small row counts (2D), use smaller
289 + # tiles to increase parallelism across CUs.
290 + np2 = triton.next_power_of_2(num_blocks)
291 + if total_rows_flat >= 4096:
292 + # Many rows: maximize work per CTA, one tile per row
293 + tile_blocks = np2
294 + nw = 1
295 + ns = 2
296 + elif total_rows_flat >= 2048:
297 + # Medium-many rows: full row, 2 warps
298 + tile_blocks = np2
299 + nw = 2
300 + ns = 2
301 + else:
302 + # Few rows: use moderate tiles for CU utilization
303 + tile_blocks = min(64, np2)
304 + nw = 4
305 + ns = 2
306 + num_tiles = (num_blocks + tile_blocks - 1) // tile_blocks
307 + grid = (total_rows_flat, num_tiles)
308 + _dequantize_nvfp4_kernel[grid](
309 + tensor_fp4_2d,
310 + scale_raw,
311 + global_scale,
312 + output,
313 + m_per_batch,
314 + num_blocks,
315 + block_size,
316 + is_3d,
317 + tile_blocks,
318 + num_warps=nw,
319 + num_stages=ns,
320 + )
321 +
322 + if is_3d:
323 + output = output.reshape(dim0, m_per_batch, k)
324 +
325 + return output
326 +
327 +
328 def break_fp4_bytes(a, dtype):
329 assert a.dtype == torch.uint8
330 m, n = a.shape
@@ -67,6 +372,11 @@ def dequantize_to_dtype(
372 # Two fp4 values are packed into one uint8.
373 assert tensor_fp4.dtype == torch.uint8
374
375 + if not swizzle and current_platform.is_cuda_alike():
376 + return _triton_dequantize_nvfp4(
377 + tensor_fp4, tensor_sf, global_scale, dtype, block_size
378 + )
379 +
380 # We handle 3D tensors reshaping them to 2D.
381 is_3d = tensor_fp4.ndim == 3
382
@@ -145,12 +455,15 @@ def ref_nvfp4_quant(x, global_scale, block_size):
455
456 def ref_nvfp4_quant_dequant(
457 x: torch.Tensor, global_scale: torch.Tensor, block_size: int
148 -) -> tuple[torch.Tensor, None]:
458 +) -> torch.Tensor:
459 """
460 NVFP4 quantize-dequantize operation.
461
462 `global_scale` is expected to have a single element.
463 """
464 + if current_platform.is_cuda_alike():
465 + return _triton_nvfp4_quant_dequant(x, global_scale, block_size)
466 +
467 x_m, x_k = x.shape
468 output_dtype = x.dtype
469
@@ -162,7 +475,7 @@ def ref_nvfp4_quant_dequant(
475 x_blockscale = x_blockscale.unsqueeze(-1) / global_scale
476 x_dq = (x_fp4 * x_blockscale).reshape(x_m, x_k).to(output_dtype)
477
165 - return x_dq, None
478 + return x_dq
479
480
481 def run_nvfp4_emulations(
@@ -176,7 +489,7 @@ def run_nvfp4_emulations(
489 output_dtype = x.dtype
490 group_size = 16
491
179 - x_dq, _ = ref_nvfp4_quant_dequant(x, input_global_scale, block_size=group_size)
492 + x_dq = ref_nvfp4_quant_dequant(x, input_global_scale, block_size=group_size)
493
494 # dequantize weight
495 w_fp4 = weight.data.view(torch.uint8)