@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,502 lines 49 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 # Adapted from https://github.com/sgl-project/sglang/pull/2575
5 import functools
6 import json
7 import os
8 from collections.abc import Callable, Sequence
9 from typing import Any
10
11 import torch
12
13 import vllm.envs as envs
14 from vllm import _custom_ops as ops
15 from vllm.logger import init_logger
16 from vllm.model_executor.layers.quantization.utils.quant_utils import (
17 amax_for_moe_activation_quant,
18 get_fp8_min_max,
19 )
20 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
21 all_close_1d,
22 per_tensor_dequantize,
23 )
24 from vllm.model_executor.parameter import (
25 BlockQuantScaleParameter,
26 ChannelQuantScaleParameter,
27 PerTensorScaleParameter,
28 )
29 from vllm.model_executor.utils import set_weight_attrs
30 from vllm.platforms import current_platform
31 from vllm.triton_utils import tl, triton
32 from vllm.utils.deep_gemm import (
33 get_tma_aligned_size,
34 is_deep_gemm_e8m0_used,
35 transform_sf_into_required_layout,
36 )
37 from vllm.utils.platform_utils import get_device_name_as_file_name
38
39 logger = init_logger(__name__)
40
41 # Pre-fill value for scale parameters whose shards load independently. The
42 # shards are combined with .max(), so an unloaded shard must never win; the
43 # smallest representable float32 guarantees that.
44 FP8_SCALE_SENTINEL = torch.finfo(torch.float32).min
45
46
47 def is_fp8(x: torch.dtype | torch.Tensor) -> bool:
48 if isinstance(x, torch.Tensor):
49 x = x.dtype
50 return x == torch.float8_e4m3fn or x == torch.float8_e4m3fnuz
51
52
53 def get_fp8_block_weight_scale(layer: torch.nn.Module) -> torch.Tensor | None:
54 """Return the block-FP8 weight scale for supported quant methods."""
55 # Local import avoids a circular import: quark_w8a8_fp8 imports this module.
56 from vllm.model_executor.layers.quantization.quark.schemes.quark_w8a8_fp8 import (
57 QuarkW8A8Fp8PerBlock,
58 )
59
60 if isinstance(
61 getattr(layer, "scheme", None),
62 QuarkW8A8Fp8PerBlock,
63 ):
64 return getattr(layer, "weight_scale", None)
65 return getattr(layer, "weight_scale_inv", None)
66
67
68 def input_to_float8(
69 x: torch.Tensor, dtype: torch.dtype | None = None
70 ) -> tuple[torch.Tensor, torch.Tensor]:
71 """This function quantizes input values to float8 values "
72 "with tensor-wise quantization."""
73 dtype = current_platform.fp8_dtype() if dtype is None else dtype
74 finfo = torch.finfo(dtype)
75 min_val, max_val = x.aminmax()
76 amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
77 scale = finfo.max / amax
78 x_scl_sat = (x * scale).clamp(min=finfo.min, max=finfo.max)
79 return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal()
80
81
82 @triton.jit
83 def _per_token_group_quant_fp8(
84 # Pointers to inputs and output
85 y_ptr,
86 y_q_ptr,
87 y_s_ptr,
88 group_size,
89 # Num columns of y
90 y_num_columns,
91 y_row_stride,
92 # Avoid to divide zero
93 eps,
94 # Information for float8
95 fp8_min: tl.constexpr,
96 fp8_max: tl.constexpr,
97 use_ue8m0: tl.constexpr,
98 # Meta-parameters
99 BLOCK: tl.constexpr,
100 ):
101 """A Triton-accelerated function to perform per-token-group
102 quantization on a tensor.
103 This function converts the tensor values into float8 values.
104 """
105 groups_per_row = y_num_columns // group_size
106
107 # Map the program id to the row of X and Y it should compute.
108 g_id = tl.program_id(0)
109 row = g_id // groups_per_row
110 row_g_id = g_id % groups_per_row
111
112 # Ensure offset calculations use int64 to prevent overflow
113 y_ptr_offset = (row.to(tl.int64) * y_row_stride) + (
114 row_g_id.to(tl.int64) * group_size
115 )
116 y_ptr += y_ptr_offset
117
118 y_q_ptr_offset = g_id.to(tl.int64) * group_size
119 y_q_ptr += y_q_ptr_offset
120 y_s_ptr += g_id
121
122 cols = tl.arange(0, BLOCK) # N <= BLOCK
123 mask = cols < group_size
124
125 y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)
126 # Quant
127 # Use multiply-by-reciprocal instead of division to match PyTorch's
128 # tensor/scalar division precision (GPU fast-division for constexpr
129 # divisors can introduce 1-ULP error that flips FP8 quantization at
130 # representable-value boundaries).
131 _absmax = tl.maximum(tl.max(tl.abs(y)), eps)
132 scale_raw = _absmax * (1.0 / fp8_max)
133 y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw
134 y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
135
136 tl.store(y_q_ptr + cols, y_q, mask=mask)
137 tl.store(y_s_ptr, y_s)
138
139
140 @triton.jit
141 def _silu_mul_quant_fp8_packed_kernel(
142 input_ptr,
143 output_q_ptr,
144 output_scale_ptr,
145 M,
146 input_stride_m,
147 output_q_stride_m,
148 output_scale_stride_k,
149 clamp_limit,
150 alpha,
151 beta,
152 N: tl.constexpr,
153 GROUPS_PER_ROW: tl.constexpr,
154 PACKS_PER_ROW: tl.constexpr,
155 fp8_min: tl.constexpr,
156 fp8_max: tl.constexpr,
157 GROUP_SIZE: tl.constexpr,
158 PACKS_PER_CTA: tl.constexpr,
159 BLOCK_M: tl.constexpr,
160 HAS_CLAMP: tl.constexpr,
161 ):
162 GROUPS_PER_PACK: tl.constexpr = 4
163 hidden_size: tl.constexpr = N // 2
164
165 pack_tile = tl.program_id(0)
166 row_start = tl.program_id(1).to(tl.int64) * BLOCK_M
167 row_step = tl.num_programs(1).to(tl.int64) * BLOCK_M
168
169 groups_per_cta: tl.constexpr = PACKS_PER_CTA * GROUPS_PER_PACK
170 elems_per_cta: tl.constexpr = groups_per_cta * GROUP_SIZE
171 col_start = pack_tile * elems_per_cta
172 col_offsets = tl.arange(0, elems_per_cta)
173 row_offsets = tl.arange(0, BLOCK_M)
174 pack_offsets = tl.arange(0, PACKS_PER_CTA)
175
176 col_mask = (col_start + col_offsets) < (GROUPS_PER_ROW * GROUP_SIZE)
177
178 # persistent with grid_m-stride loop
179 while row_start < M:
180 rows = row_start + row_offsets
181 row_mask = rows < M
182 input_row_start = rows[:, None] * input_stride_m
183 output_row_start = rows[:, None] * output_q_stride_m
184
185 gate_flat = tl.load(
186 input_ptr + input_row_start + col_start + col_offsets[None, :],
187 mask=row_mask[:, None] & col_mask[None, :],
188 other=0.0,
189 )
190 up_flat = tl.load(
191 input_ptr
192 + input_row_start
193 + hidden_size
194 + col_start
195 + col_offsets[None, :],
196 mask=row_mask[:, None] & col_mask[None, :],
197 other=0.0,
198 )
199
200 gate = tl.reshape(gate_flat, (BLOCK_M, groups_per_cta, GROUP_SIZE)).to(
201 tl.float32
202 )
203 up = tl.reshape(up_flat, (BLOCK_M, groups_per_cta, GROUP_SIZE)).to(tl.float32)
204
205 if HAS_CLAMP:
206 gate = tl.minimum(gate, clamp_limit)
207 up = tl.clamp(up, -clamp_limit, clamp_limit)
208
209 # Unified gated activation: silu == swigluoai with alpha=1, beta=0.
210 # glu = gate * sigmoid(alpha * gate); y = (up + beta) * glu
211 glu = gate / (1.0 + tl.exp(-gate * alpha))
212 y = (up + beta) * glu
213 # Round through bf16 to match unfused precision path
214 y = y.to(tl.bfloat16).to(tl.float32)
215
216 absmax = tl.max(tl.abs(y), axis=2)
217 scale_raw = tl.maximum(absmax / fp8_max, 1e-10)
218 exponent = tl.ceil(tl.log2(scale_raw))
219 scale = tl.math.exp2(exponent)
220
221 y_q = tl.clamp(y / scale[:, :, None], fp8_min, fp8_max)
222
223 y_q_flat = tl.reshape(y_q, (BLOCK_M, elems_per_cta))
224 tl.store(
225 output_q_ptr + output_row_start + col_start + col_offsets[None, :],
226 y_q_flat.to(output_q_ptr.dtype.element_ty),
227 mask=row_mask[:, None] & col_mask[None, :],
228 )
229
230 scale_byte = tl.clamp(exponent + 127.0, 0.0, 255.0).to(tl.int32)
231 scale_bytes = tl.reshape(scale_byte, (BLOCK_M, PACKS_PER_CTA, GROUPS_PER_PACK))
232 shifts = tl.arange(0, GROUPS_PER_PACK) * 8
233 packed_scale = tl.sum(scale_bytes << shifts[None, None, :], axis=2)
234
235 scale_pack = pack_tile * PACKS_PER_CTA + pack_offsets
236 scale_ptrs = (
237 output_scale_ptr
238 + scale_pack[None, :] * output_scale_stride_k
239 + rows[:, None]
240 )
241 tl.store(
242 scale_ptrs,
243 packed_scale,
244 mask=row_mask[:, None] & (scale_pack[None, :] < PACKS_PER_ROW),
245 )
246
247 row_start += row_step
248
249
250 def silu_mul_quant_fp8_packed_triton(
251 input: torch.Tensor,
252 group_size: int = 128,
253 output_q: torch.Tensor | None = None,
254 clamp_limit: float | None = None,
255 alpha: float = 1.0,
256 beta: float = 0.0,
257 ) -> tuple[torch.Tensor, torch.Tensor]:
258 assert input.dim() == 2
259 assert input.is_contiguous()
260
261 M, N = input.shape
262 hidden_size = N // 2
263
264 assert hidden_size % group_size == 0
265
266 fp8_dtype = torch.float8_e4m3fn
267 finfo = torch.finfo(fp8_dtype)
268 fp8_min, fp8_max = finfo.min, finfo.max
269
270 groups_per_row = hidden_size // group_size
271 groups_per_pack = 4 # pack 4 UE8M0 scales to a single INT32
272 packs_per_row = triton.cdiv(groups_per_row, groups_per_pack)
273
274 if output_q is None:
275 output_q = torch.empty((M, hidden_size), dtype=fp8_dtype, device=input.device)
276
277 aligned_m = triton.cdiv(M, 4) * 4
278 output_scale_packed = torch.empty(
279 (packs_per_row, aligned_m),
280 dtype=torch.int32,
281 device=input.device,
282 ).T[:M, :]
283
284 # Tuned for group_size=32 (MXFP8) and group_size=128 (DeepSeek-V4)
285 num_warps = 4
286 num_stages = 2
287 if group_size < 128:
288 BM = 1
289 packs_per_cta = 8
290 else:
291 BM = 1 if M < 512 else 4
292 packs_per_cta = 2 if M < 512 else 1
293
294 grid_n = triton.cdiv(packs_per_row, packs_per_cta)
295 grid_m = min(triton.cdiv(M, BM), 4096)
296 grid = (grid_n, grid_m)
297
298 has_clamp = clamp_limit is not None
299 _silu_mul_quant_fp8_packed_kernel[grid](
300 input,
301 output_q,
302 output_scale_packed,
303 M,
304 input.stride(0),
305 output_q.stride(0),
306 output_scale_packed.stride(1),
307 clamp_limit if has_clamp else 0.0,
308 alpha,
309 beta,
310 N=N,
311 GROUPS_PER_ROW=groups_per_row,
312 PACKS_PER_ROW=packs_per_row,
313 fp8_min=fp8_min,
314 fp8_max=fp8_max,
315 GROUP_SIZE=group_size,
316 PACKS_PER_CTA=packs_per_cta,
317 BLOCK_M=BM,
318 HAS_CLAMP=has_clamp,
319 num_warps=num_warps,
320 num_stages=num_stages,
321 )
322
323 return output_q, output_scale_packed
324
325
326 @triton.jit
327 def _silu_mul_per_token_group_quant_fp8_colmajor(
328 y_ptr, # [M, N]
329 y_q_ptr, # [M, N // 2]
330 y_s_ptr, # [M, (N // 2) // GROUP_SIZE]
331 M, # num tokens
332 N, # intermediate size
333 # Stride
334 y_s_col_stride: tl.int64,
335 # Information for float8
336 eps,
337 clamp_limit,
338 alpha,
339 beta,
340 fp8_min: tl.constexpr,
341 fp8_max: tl.constexpr,
342 use_ue8m0: tl.constexpr,
343 HAS_CLAMP: tl.constexpr,
344 # Meta-parameters
345 GROUP_SIZE: tl.constexpr,
346 BLOCK_M: tl.constexpr,
347 BLOCK_N: tl.constexpr,
348 ):
349 # TODO(varun) : Add expert_ids so we may early-exit no-op thread blocks.
350 """Each thread block (BLOCK_N) computes [BLOCK_M, GROUP_SIZE] act-mul outputs. Then
351 the thread block quantizes the [BLOCK_M, GROUP_SIZE] block of values and fills
352 the outputs tensors at the right positions.
353 """
354 pid_m = tl.program_id(0)
355 pid_n = tl.program_id(1)
356 N_2 = N // 2
357
358 m_offset = pid_m.to(tl.int64) * BLOCK_M
359 n_offset = pid_n.to(tl.int64) * BLOCK_N
360 if m_offset >= M:
361 return
362
363 offs_n = tl.arange(0, BLOCK_N).to(tl.int64)
364 offs_m = tl.arange(0, BLOCK_M).to(tl.int64)
365
366 base_y_ptr = y_ptr + m_offset * N + n_offset
367
368 act_in_ptrs = base_y_ptr + offs_m[:, None] * N + offs_n[None, :]
369
370 act_in = tl.load(act_in_ptrs)
371 mul_in = tl.load(act_in_ptrs + N_2)
372
373 # silu & mul — match C++ silu_and_mul: clamp in fp32 then store back to the
374 # input dtype, run silu in fp32 then narrow, and do the mul at input
375 # precision so HAS_CLAMP True/False share the same multiplication path.
376 if HAS_CLAMP:
377 act_in = tl.minimum(act_in.to(tl.float32), clamp_limit).to(
378 y_ptr.dtype.element_ty
379 )
380 mul_in = tl.clamp(mul_in.to(tl.float32), -clamp_limit, clamp_limit).to(
381 y_ptr.dtype.element_ty
382 )
383 # Unified gated activation: silu == swigluoai with alpha=1, beta=0.
384 # glu = gate * sigmoid(alpha * gate); y = (up + beta) * glu
385 # Keep glu/up at input precision (narrow before the mul) so the alpha=1,
386 # beta=0 defaults match the C++ silu_and_mul path bit-for-bit.
387 act_in = act_in.to(tl.float32)
388 glu = (act_in / (1.0 + tl.exp(-act_in * alpha))).to(y_ptr.dtype.element_ty)
389 up = (mul_in.to(tl.float32) + beta).to(y_ptr.dtype.element_ty)
390 y = (glu * up).to(tl.float32)
391
392 # quant
393 _absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps)
394 scale_raw = _absmax * (1.0 / fp8_max)
395 y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw
396 y_s = tl.reshape(y_s, (BLOCK_M, 1))
397 y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
398
399 # store y_q
400 base_y_q_ptr = y_q_ptr + m_offset * N_2 + n_offset
401 y_q_ptrs = base_y_q_ptr + offs_m[:, None] * N_2 + offs_n[None, :]
402 tl.store(y_q_ptrs, y_q)
403
404 # store y_s
405 group_id = n_offset // GROUP_SIZE
406 base_y_s_ptr = y_s_ptr + group_id * y_s_col_stride + m_offset
407 y_s_ptrs = base_y_s_ptr + offs_m
408 y_s = tl.reshape(y_s, (BLOCK_M,))
409 tl.store(y_s_ptrs, y_s)
410
411
412 def silu_mul_per_token_group_quant_fp8_colmajor(
413 input: torch.Tensor, # [M, N]
414 output: torch.Tensor | None = None, # [M, N // 2]
415 use_ue8m0: bool | None = None,
416 eps: float = 1e-10,
417 clamp_limit: float | None = None,
418 group_size: int = 128,
419 alpha: float = 1.0,
420 beta: float = 0.0,
421 ):
422 """Gated activation + block-fp8 quant. ``alpha``/``beta`` select the gate
423 (silu: alpha=1, beta=0; swigluoai: alpha, beta from config).
424 """
425 GROUP_SIZE = group_size
426 assert input.ndim == 2
427 if output is not None:
428 assert output.ndim == 2
429 assert input.size(0) % GROUP_SIZE == 0
430 assert input.size(1) % (GROUP_SIZE * 2) == 0
431
432 if use_ue8m0 is None:
433 use_ue8m0 = is_deep_gemm_e8m0_used()
434
435 M, N = input.size()
436 N_2 = N // 2
437
438 fp8_dtype = current_platform.fp8_dtype()
439 if output is None:
440 output = torch.empty((M, N_2), dtype=fp8_dtype, device=input.device)
441
442 output_scales = torch.empty(
443 ((N_2 // GROUP_SIZE), M), dtype=torch.float32, device=input.device
444 ).transpose(0, 1)
445
446 BLOCK_M = 8
447 BLOCK_N = GROUP_SIZE
448 assert M % BLOCK_M == 0
449 assert N_2 % BLOCK_N == 0
450
451 # Using the default value (240.0) from pytorch will cause accuracy
452 # issue on dynamic quantization models. Here use 224.0 for fnuz on ROCm
453 # platforms that use the torch.float8_e4m3fnuz dtype.
454 finfo = torch.finfo(fp8_dtype)
455 fp8_min = -224.0 if current_platform.is_fp8_fnuz() else finfo.min
456 fp8_max = 224.0 if current_platform.is_fp8_fnuz() else finfo.max
457
458 # Force even division so we can avoid edgecases within the kernel.
459 assert M % BLOCK_M == 0
460 assert N_2 % BLOCK_N == 0
461 grid = (M // BLOCK_M, N_2 // BLOCK_N)
462
463 has_clamp = clamp_limit is not None
464 _silu_mul_per_token_group_quant_fp8_colmajor[grid](
465 input,
466 output,
467 output_scales,
468 M,
469 N,
470 output_scales.stride(-1),
471 eps,
472 clamp_limit if has_clamp else 0.0,
473 alpha,
474 beta,
475 fp8_min,
476 fp8_max,
477 use_ue8m0,
478 has_clamp,
479 GROUP_SIZE,
480 BLOCK_M,
481 BLOCK_N,
482 )
483
484 return output, output_scales
485
486
487 @triton.jit
488 def _per_token_group_quant_fp8_colmajor(
489 # Pointers to inputs and output
490 y_ptr,
491 y_q_ptr,
492 y_s_ptr,
493 group_size,
494 # Num columns of y
495 y_num_columns,
496 y_row_stride,
497 # Stride from one column to the next of y_s
498 y_s_col_stride,
499 # Avoid to divide zero
500 eps,
501 # Information for float8
502 fp8_min: tl.constexpr,
503 fp8_max: tl.constexpr,
504 use_ue8m0: tl.constexpr,
505 # Meta-parameters
506 BLOCK: tl.constexpr,
507 ):
508 """A Triton-accelerated function to perform per-token-group
509 quantization on a tensor.
510 This function converts the tensor values into float8 values.
511 """
512 groups_per_row = y_num_columns // group_size
513
514 # Map the program id to the row of X and Y it should compute.
515 g_id = tl.program_id(0)
516 row = g_id // groups_per_row
517 row_g_id = g_id % groups_per_row
518
519 # Ensure offset calculations use int64 to prevent overflow
520 y_ptr_offset = (row.to(tl.int64) * y_row_stride) + (
521 row_g_id.to(tl.int64) * group_size
522 )
523 y_ptr += y_ptr_offset
524
525 y_q_ptr_offset = g_id.to(tl.int64) * group_size
526 y_q_ptr += y_q_ptr_offset
527
528 # Convert g_id the flattened block coordinate to 2D so we can index
529 # into the output y_scales matrix
530 blocks_per_row = y_num_columns // group_size
531 scale_col = g_id % blocks_per_row
532 scale_row = g_id // blocks_per_row
533 # Ensure offset calculation uses int64 for y_s_ptr
534 y_s_ptr_offset = (scale_col.to(tl.int64) * y_s_col_stride) + scale_row.to(tl.int64)
535 y_s_ptr += y_s_ptr_offset
536
537 cols = tl.arange(0, BLOCK) # group_size <= BLOCK
538 mask = cols < group_size
539
540 y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)
541 # Quant
542 _absmax = tl.maximum(tl.max(tl.abs(y)), eps)
543 scale_raw = _absmax * (1.0 / fp8_max)
544 y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw
545 y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
546
547 tl.store(y_q_ptr + cols, y_q, mask=mask)
548 tl.store(y_s_ptr, y_s)
549
550
551 def per_token_group_quant_fp8(
552 x: torch.Tensor,
553 group_size: int,
554 eps: float = 1e-10,
555 dtype: torch.dtype | None = None,
556 column_major_scales: bool = False,
557 tma_aligned_scales: bool = False,
558 out_q: torch.Tensor | None = None,
559 use_ue8m0: bool | None = None,
560 ) -> tuple[torch.Tensor, torch.Tensor]:
561 """Function to perform per-token-group quantization on an input tensor `x`.
562 It converts the tensor values into signed float8 values and returns the
563 quantized tensor along with the scaling factor used for quantization.
564
565 Args:
566 x: The input tensor with ndim >= 2.
567 group_size: The group size used for quantization.
568 eps: The minimum to avoid dividing zero.
569 dtype: The dtype of output tensor. Note that only `torch.float8_e4m3fn`
570 is supported for now.
571 column_major_scales: Outputs scales in column major.
572 tma_aligned_scales: Outputs scales in TMA-aligned layout.
573 out_q: Optional output tensor. If not provided, function will create.
574 use_ue8m0: If True, round each scale down to a power of two (UE8M0
575 format). None defers to the platform default.
576
577 Returns:
578 tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the
579 scaling factor.
580
581 """
582 if use_ue8m0 is None:
583 use_ue8m0 = is_deep_gemm_e8m0_used()
584 dtype = current_platform.fp8_dtype() if dtype is None else dtype
585 assert x.shape[-1] % group_size == 0, (
586 f"the last dimension of `x` {x.shape[-1]} must be divisible "
587 f"by `group_size` {group_size}"
588 )
589 assert x.stride(-1) == 1, "`x` groups must be contiguous"
590
591 fp8_min, fp8_max = get_fp8_min_max()
592
593 assert out_q is None or out_q.shape == x.shape
594 x_q = out_q
595 if x_q is None:
596 x_q = torch.empty(x.shape, device=x.device, dtype=dtype)
597
598 # Allocate the scale tensor in either row- or column-major format.
599 if column_major_scales:
600 if tma_aligned_scales:
601 m = x.shape[-2]
602 sf_k = x.shape[-1] // group_size
603 tma_aligned_m = get_tma_aligned_size(m, 4)
604 shape = x.shape[:-2] + (m, sf_k)
605 stride = (
606 (1, tma_aligned_m)
607 if x.dim() == 2
608 else (tma_aligned_m * sf_k, 1, tma_aligned_m)
609 )
610 x_s = torch.empty_strided(
611 shape, stride, device=x.device, dtype=torch.float32
612 )
613 else:
614 shape = x.shape[:-2] + (x.shape[-1] // group_size, x.shape[-2])
615 x_s = torch.empty(shape, device=x.device, dtype=torch.float32).permute(
616 -1, -2
617 )
618 else:
619 shape = x.shape[:-1] + (x.shape[-1] // group_size,)
620 x_s = torch.empty(shape, device=x.device, dtype=torch.float32)
621
622 # prefer CUDA/XPU kernel if available
623 # TODO(bnell): this causes some fp8 moe test to fail.
624 if (
625 current_platform.is_cuda_alike() or current_platform.is_xpu()
626 ) and x.is_contiguous():
627 torch.ops._C.per_token_group_fp8_quant(
628 x,
629 x_q,
630 x_s,
631 group_size,
632 eps,
633 fp8_min,
634 fp8_max,
635 use_ue8m0,
636 column_major_scales,
637 tma_aligned_scales,
638 )
639 return x_q, x_s
640
641 # TRITON FALLBACK
642 M = x.numel() // group_size
643 N = group_size
644 BLOCK = triton.next_power_of_2(N)
645 # heuristics for number of warps
646 num_warps = min(max(BLOCK // 256, 1), 8)
647 num_stages = 1
648 if column_major_scales:
649 _per_token_group_quant_fp8_colmajor[(M,)](
650 x,
651 x_q,
652 x_s,
653 group_size,
654 x.shape[1],
655 x.stride(0),
656 x_s.stride(1),
657 eps,
658 fp8_min=fp8_min,
659 fp8_max=fp8_max,
660 use_ue8m0=use_ue8m0,
661 BLOCK=BLOCK,
662 num_warps=num_warps,
663 num_stages=num_stages,
664 )
665 else:
666 _per_token_group_quant_fp8[(M,)](
667 x,
668 x_q,
669 x_s,
670 group_size,
671 x.shape[1],
672 x.stride(0),
673 eps,
674 fp8_min=fp8_min,
675 fp8_max=fp8_max,
676 use_ue8m0=use_ue8m0,
677 BLOCK=BLOCK,
678 num_warps=num_warps,
679 num_stages=num_stages,
680 )
681
682 return x_q, x_s
683
684
685 def per_token_group_quant_fp8_packed_for_deepgemm(
686 x: torch.Tensor,
687 group_size: int,
688 eps: float = 1e-10,
689 use_ue8m0: bool | None = None,
690 out_q: torch.Tensor | None = None,
691 ) -> tuple[torch.Tensor, torch.Tensor]:
692 """FP8 per-token-group quantization for DeepGEMM.
693
694 Returns:
695 (x_q, x_s_packed)
696 x_q: FP8 activations, same shape as `x`.
697 x_s_packed: Int32 tensor with logical shape
698 [mn, ceil(num_groups_per_row / 4)], laid out with
699 TMA-aligned stride along the packed-K dimension
700
701 """
702 if use_ue8m0 is None:
703 use_ue8m0 = is_deep_gemm_e8m0_used()
704 # for DeepGEMM UE8M0-packed layout we *require* UE8M0 scales.
705 assert use_ue8m0, (
706 "per_token_group_quant_fp8_packed_for_deepgemm requires UE8M0 scales."
707 )
708
709 dtype = current_platform.fp8_dtype()
710 assert x.shape[-1] % group_size == 0, (
711 f"the last dimension of `x` {x.shape[-1]} must be divisible "
712 f"by `group_size` {group_size}"
713 )
714 assert x.stride(-1) == 1, "`x` groups must be contiguous"
715
716 fp8_min, fp8_max = get_fp8_min_max()
717
718 # compute DeepGEMM-style packed scale tensor shape.
719 hidden_dim = x.shape[-1]
720 mn = x.numel() // hidden_dim
721 num_groups_per_row = hidden_dim // group_size
722 k_num_packed_sf_k = (num_groups_per_row + 3) // 4
723 tma_aligned_mn = ((mn + 3) // 4) * 4
724
725 x_s_packed = torch.empty_strided(
726 (mn, k_num_packed_sf_k),
727 (1, tma_aligned_mn),
728 device=x.device,
729 dtype=torch.int32,
730 )
731
732 # Native kernel (libtorch stable); used with DeepGEMM on CUDA and
733 # available on ROCm for the same packed UE8M0 scale layout.
734 assert current_platform.is_cuda_alike(), (
735 "per_token_group_quant_fp8_packed_for_deepgemm requires a CUDA or ROCm GPU."
736 )
737
738 x_contiguous = x.contiguous()
739 if out_q is not None:
740 x_q_local = out_q
741 else:
742 x_q_local = torch.empty_like(x_contiguous, device=x.device, dtype=dtype)
743
744 torch.ops._C.per_token_group_fp8_quant_packed(
745 x_contiguous,
746 x_q_local,
747 x_s_packed,
748 group_size,
749 eps,
750 fp8_min,
751 fp8_max,
752 )
753
754 # return a tensor with the original logical shape.
755 x_q = x_q_local.view_as(x)
756 return x_q, x_s_packed
757
758
759 @triton.jit
760 def _w8a8_triton_block_scaled_mm(
761 # Pointers to inputs and output
762 A,
763 B,
764 C,
765 As,
766 Bs,
767 # Shape for matmul
768 M,
769 N,
770 K,
771 # Block size for block-wise quantization
772 group_n,
773 group_k,
774 # Stride for inputs and output
775 stride_am,
776 stride_ak,
777 stride_bk,
778 stride_bn,
779 stride_cm,
780 stride_cn,
781 stride_As_m,
782 stride_As_k,
783 stride_Bs_k,
784 stride_Bs_n,
785 # Meta-parameters
786 BLOCK_SIZE_M: tl.constexpr,
787 BLOCK_SIZE_N: tl.constexpr,
788 BLOCK_SIZE_K: tl.constexpr,
789 GROUP_SIZE_M: tl.constexpr,
790 ):
791 """Triton-accelerated function used to perform linear operations (dot
792 product) on input tensors `A` and `B` with block-wise quantization, and
793 store the result in output tensor `C`.
794 """
795 pid = tl.program_id(axis=0)
796 num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
797 num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
798 num_pid_in_group = GROUP_SIZE_M * num_pid_n
799 group_id = pid // num_pid_in_group
800 first_pid_m = group_id * GROUP_SIZE_M
801 group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
802 pid_m = first_pid_m + (pid % group_size_m)
803 pid_n = (pid % num_pid_in_group) // group_size_m
804
805 offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
806 offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
807 offs_k = tl.arange(0, BLOCK_SIZE_K)
808 a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
809 b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
810
811 As_ptrs = As + offs_am * stride_As_m
812 offs_bsn = offs_bn // group_n
813 Bs_ptrs = Bs + offs_bsn * stride_Bs_n
814
815 accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
816 for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
817 a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
818 b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
819
820 k_start = k * BLOCK_SIZE_K
821 offs_ks = k_start // group_k
822 a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
823 b_s = tl.load(Bs_ptrs + offs_ks * stride_Bs_k)
824
825 accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
826 a_ptrs += BLOCK_SIZE_K * stride_ak
827 b_ptrs += BLOCK_SIZE_K * stride_bk
828
829 if C.dtype.element_ty == tl.bfloat16:
830 c = accumulator.to(tl.bfloat16)
831 elif C.dtype.element_ty == tl.float16:
832 c = accumulator.to(tl.float16)
833 else:
834 c = accumulator.to(tl.float32)
835
836 offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
837 offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
838 c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
839 c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
840 tl.store(c_ptrs, c, mask=c_mask)
841
842
843 @functools.lru_cache
844 def get_w8a8_block_fp8_configs(
845 N: int, K: int, block_n: int, block_k: int
846 ) -> dict[int, Any] | None:
847 """Return optimized configurations for the w8a8 block fp8 kernel.
848 The return value will be a dictionary that maps an irregular grid of
849 batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the
850 kernel on a given batch size bs, the closest batch size in the grid should
851 be picked and the associated configuration chosen to invoke the kernel.
852 """
853 # First look up if an optimized configuration is available in the configs
854 # directory
855 device_name = get_device_name_as_file_name()
856 json_file_name = f"N={N},K={K},device_name={device_name},dtype=fp8_w8a8,block_shape=[{block_n},{block_k}].json" # noqa: E501
857
858 config_file_path = os.path.join(
859 os.path.dirname(os.path.realpath(__file__)), "configs", json_file_name
860 )
861 if os.path.exists(config_file_path):
862 with open(config_file_path) as f:
863 logger.info(
864 "Using configuration from %s for W8A8 Block FP8 kernel.",
865 config_file_path,
866 )
867 # If a configuration has been found, return it
868 return {int(key): val for key, val in json.load(f).items()}
869
870 # If no optimized configuration is available, we will use the default
871 # configuration
872 logger.warning(
873 "Using default W8A8 Block FP8 kernel config. Performance might "
874 "be sub-optimal! Config file not found at %s",
875 config_file_path,
876 )
877 return None
878
879
880 def w8a8_triton_block_scaled_mm(
881 A: torch.Tensor,
882 B: torch.Tensor,
883 As: torch.Tensor,
884 Bs: torch.Tensor,
885 block_size: list[int],
886 output_dtype: torch.dtype = torch.float16,
887 ) -> torch.Tensor:
888 """This function performs matrix multiplication with block-wise
889 quantization.
890 It takes two input tensors `A` and `B` with scales `As` and `Bs`.
891 The output is returned in the specified `output_dtype`.
892
893 Args:
894 A: The input tensor, e.g., activation.
895 B: The input tensor, e.g., weight.
896 As: The per-token-group quantization scale for `A`.
897 Bs: The per-block quantization scale for `B`.
898 block_size: The block size for per-block quantization. It should
899 be 2-dim, e.g., [128, 128].
900 output_dtype: The dtype of the returned tensor.
901
902 Returns:
903 torch.Tensor: The result of matmul.
904
905 """
906 _on_gfx1250 = False
907 if current_platform.is_rocm():
908 from vllm.platforms.rocm import on_gfx1250
909
910 _on_gfx1250 = on_gfx1250()
911
912 if _on_gfx1250:
913 # Torch upcast reference: dequantize A,B to fp32 and matmul in fp32.
914 # Avoids the gfx1250 native-fp8 block GEMM NaN bug. Correct but slow.
915 _bn, _bk = block_size[0], block_size[1]
916 _As = (
917 _upcast_e8m0_to_fp32(As)
918 if As.dtype == torch.float8_e8m0fnu
919 else As.to(torch.float32)
920 )
921 _Bs = (
922 _upcast_e8m0_to_fp32(Bs)
923 if Bs.dtype == torch.float8_e8m0fnu
924 else Bs.to(torch.float32)
925 )
926 _K = A.shape[-1]
927 _N = B.shape[0]
928 _Af = A.to(torch.float32).reshape(-1, _K)
929 _Asf = (
930 _As.to(torch.float32)
931 .reshape(-1, _As.shape[-1])
932 .repeat_interleave(_bk, dim=1)[:, :_K]
933 )
934 _Bf = B.to(torch.float32)
935 _Bsf = _Bs.repeat_interleave(_bn, dim=0).repeat_interleave(_bk, dim=1)[:_N, :_K]
936 _out = (_Af * _Asf) @ (_Bf * _Bsf).t()
937 return _out.to(output_dtype).reshape(*A.shape[:-1], _N)
938
939 assert len(block_size) == 2
940 block_n, block_k = block_size[0], block_size[1]
941
942 # Triton cannot currently bind E8M0 scale tensors directly. Checkpoints
943 # with exponent-only UE8M0 scales (e.g. DeepSeek-V4) are decoded to fp32
944 # before launching the kernel.
945 if As.dtype == torch.float8_e8m0fnu:
946 As = _upcast_e8m0_to_fp32(As).contiguous()
947 if Bs.dtype == torch.float8_e8m0fnu:
948 Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
949
950 assert A.shape[-1] == B.shape[-1]
951 assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous()
952 assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
953 M = A.numel() // A.shape[-1]
954
955 assert B.ndim == 2 and Bs.ndim == 2
956 N, K = B.shape
957 assert triton.cdiv(N, block_n) == Bs.shape[0]
958 assert triton.cdiv(K, block_k) == Bs.shape[1]
959
960 C_shape = A.shape[:-1] + (N,)
961 C = A.new_empty(C_shape, dtype=output_dtype)
962
963 configs = get_w8a8_block_fp8_configs(N, K, block_size[0], block_size[1])
964 if configs:
965 # Get the optimal config if there is one
966 config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
967 else:
968 # Default config
969 # Block-wise quant: BLOCK_SIZE_N must be divisible by block_size[0]
970 # BLOCK_SIZE_K must be divisible by block_size[1]
971 config = {
972 "BLOCK_SIZE_M": 64,
973 "BLOCK_SIZE_N": block_size[0],
974 "BLOCK_SIZE_K": block_size[1],
975 "GROUP_SIZE_M": 32,
976 "num_warps": 4,
977 "num_stages": 2,
978 }
979
980 def grid(META):
981 return (
982 triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),
983 )
984
985 _w8a8_triton_block_scaled_mm[grid](
986 A,
987 B,
988 C,
989 As,
990 Bs,
991 M,
992 N,
993 K,
994 block_n,
995 block_k,
996 A.stride(-2),
997 A.stride(-1),
998 B.stride(1),
999 B.stride(0),
1000 C.stride(-2),
1001 C.stride(-1),
1002 As.stride(-2),
1003 As.stride(-1),
1004 Bs.stride(1),
1005 Bs.stride(0),
1006 **config,
1007 )
1008
1009 return C
1010
1011
1012 def requant_weight_ue8m0_inplace(
1013 weight: torch.Tensor,
1014 weight_scale: torch.Tensor,
1015 block_size: Sequence[int] = (128, 128),
1016 ) -> None:
1017 """Re-quantise *weight* so that its per-block scaling factors are in the
1018 UE8M0 (power-of-two) format expected by the new DeepGEMM kernels inplace.
1019
1020 Args:
1021 weight: Block-quantised weight tensor stored in `torch.float8_e4m3fn`.
1022 Expected shape `(..., M, K)`.
1023 weight_scale: Corresponding per-block scale tensor (`torch.float32`)
1024 with shape `(..., M // block_size[0], K // block_size[1])`.
1025 block_size: 2-element iterable `[block_m, block_k]` describing the
1026 block quantisation granularity.
1027
1028 """
1029 if weight.numel() == 0:
1030 return
1031
1032 if weight.dtype != torch.float8_e4m3fn:
1033 raise ValueError(
1034 f"Expected *weight* to be torch.float8_e4m3fn, got {weight.dtype} instead."
1035 )
1036
1037 from vllm.utils.deep_gemm import per_block_cast_to_fp8
1038
1039 block_m, block_k = int(block_size[0]), int(block_size[1])
1040
1041 # Flatten leading dimensions so we can iterate over the last two dims.
1042 leading_shape = weight.shape[:-2]
1043 if len(leading_shape) == 0:
1044 w_view = weight.unsqueeze(0)
1045 s_view = weight_scale.unsqueeze(0)
1046 else:
1047 w_view = weight.reshape(-1, weight.shape[-2], weight.shape[-1])
1048 s_view = weight_scale.reshape(-1, *weight_scale.shape[-2:])
1049
1050 num_mats = w_view.size(0)
1051 for idx in range(num_mats):
1052 w_q = w_view[idx]
1053 s_old = s_view[idx]
1054
1055 # De-quantise with the *old* scaling factors (float32).
1056 m_cur, k_cur = w_q.shape
1057 s_float = s_old.to(torch.float32)
1058 # Expand scales along rows and cols by block size, then crop.
1059 s_exp_r = torch.repeat_interleave(s_float, block_m, dim=0)
1060 s_exp = torch.repeat_interleave(s_exp_r, block_k, dim=1)
1061 s_exp = s_exp[:m_cur, :k_cur]
1062 w_dq = w_q.to(torch.float32) * s_exp
1063 # Re-quantise using power-of-two scaling (UE8M0).
1064 w_requant, s_requant = per_block_cast_to_fp8(
1065 w_dq, [block_m, block_k], use_ue8m0=True
1066 )
1067
1068 # Write back the results in-place.
1069 w_q.copy_(w_requant)
1070 s_old.copy_(s_requant)
1071
1072
1073 def _upcast_e8m0_to_fp32(scale: torch.Tensor) -> torch.Tensor:
1074 """Upcast E8M0 (exponent-only) scale to float32.
1075
1076 E8M0 stores only the 8-bit biased exponent (bias=127). To convert
1077 to float32 we place those 8 bits into the exponent field of an
1078 IEEE-754 float32 (bits 23-30) with sign=0 and mantissa=0.
1079 """
1080 exp_bits = scale.view(torch.uint8).to(torch.int32)
1081 fp32_bits = exp_bits << 23
1082 return fp32_bits.view(torch.float32)
1083
1084
1085 def deepgemm_post_process_weight_scale_block(
1086 ws: torch.Tensor,
1087 mn: int,
1088 k: int,
1089 quant_block_shape: tuple[int, ...],
1090 num_groups: int,
1091 is_sfa: bool = False,
1092 ) -> torch.Tensor:
1093 if ws.dtype in (torch.float8_e8m0fnu, torch.uint8):
1094 ws = _upcast_e8m0_to_fp32(ws)
1095 else:
1096 assert ws.dtype == torch.float32, (
1097 f"Expected tensor scales dtype to be torch.float32 or "
1098 f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead"
1099 )
1100
1101 return transform_sf_into_required_layout(
1102 sf=ws,
1103 mn=mn,
1104 k=k,
1105 recipe=(1, quant_block_shape[0], quant_block_shape[1]),
1106 num_groups=num_groups,
1107 is_sfa=is_sfa,
1108 )
1109
1110
1111 def deepgemm_post_process_fp8_weight_block(
1112 wq: torch.Tensor,
1113 ws: torch.Tensor,
1114 quant_block_shape: tuple[int, ...],
1115 use_e8m0: bool,
1116 is_bmm: bool = False,
1117 bmm_batch_size: int = 0,
1118 ) -> tuple[torch.Tensor, torch.Tensor]:
1119 assert wq.dtype == torch.float8_e4m3fn, (
1120 "Expected quantized tensor dtype "
1121 f"to be torch.float8_e4m3fn, got {wq.dtype} instead."
1122 )
1123
1124 if ws.dtype not in (torch.float8_e8m0fnu, torch.uint8):
1125 assert ws.dtype == torch.float32, (
1126 f"Expected tensor scales dtype to be torch.float32 or "
1127 f"torch.float8_e8m0fnu or torch.uint8, got {ws.dtype} instead"
1128 )
1129 if use_e8m0:
1130 requant_weight_ue8m0_inplace(wq, ws, block_size=quant_block_shape)
1131
1132 if is_bmm:
1133 # Reshape 2D weight/scale to 3D for grouped BMM (einsum):
1134 # wq: (g*r, d) -> (g, r, d)
1135 # Scale dimensions follow quant_block_shape.
1136 g = bmm_batch_size
1137 assert wq.ndim == 2 and ws.ndim == 2
1138 d = wq.size(1)
1139 r = wq.size(0) // g
1140 wq = wq.view(g, r, d)
1141 ws = ws.view(g, r // quant_block_shape[0], d // quant_block_shape[1])
1142 dg_ws = deepgemm_post_process_weight_scale_block(
1143 ws=ws,
1144 mn=r,
1145 k=d,
1146 quant_block_shape=quant_block_shape,
1147 num_groups=g,
1148 )
1149 return wq, dg_ws
1150
1151 original_ndim = wq.ndim
1152 if wq.ndim == 2:
1153 assert ws.ndim == 2
1154 wq = wq.unsqueeze(0)
1155 ws = ws.unsqueeze(0)
1156
1157 dg_ws = deepgemm_post_process_weight_scale_block(
1158 ws=ws,
1159 mn=wq.size(1),
1160 k=wq.size(2),
1161 quant_block_shape=quant_block_shape,
1162 num_groups=wq.size(0),
1163 )
1164
1165 if original_ndim == 2:
1166 wq = wq.squeeze(0)
1167 dg_ws = dg_ws.squeeze(0)
1168
1169 return wq, dg_ws
1170
1171
1172 def prepare_fp8_moe_layer_for_deepgemm(
1173 w13: torch.Tensor,
1174 w2: torch.Tensor,
1175 w13_scale: torch.Tensor,
1176 w2_scale: torch.Tensor,
1177 block_shape: tuple[int],
1178 ):
1179 w13, w13_scale = deepgemm_post_process_fp8_weight_block(
1180 wq=w13,
1181 ws=w13_scale,
1182 quant_block_shape=block_shape,
1183 use_e8m0=is_deep_gemm_e8m0_used(),
1184 )
1185 w2, w2_scale = deepgemm_post_process_fp8_weight_block(
1186 wq=w2,
1187 ws=w2_scale,
1188 quant_block_shape=block_shape,
1189 use_e8m0=is_deep_gemm_e8m0_used(),
1190 )
1191
1192 return w13, w2, w13_scale, w2_scale
1193
1194
1195 def _maybe_pad_fp8_weight(weight: torch.Tensor) -> torch.Tensor:
1196 """Pad the weight tensor. This is an optimization on ROCm platform, which
1197 can benefit from tensors located far enough from one another in memory"""
1198 if (
1199 envs.VLLM_ROCM_FP8_PADDING
1200 and current_platform.is_rocm()
1201 and weight.stride(-1) == 1
1202 and (weight.stride(-2) * weight.element_size()) % 512 == 0
1203 ):
1204 num_pad = 256 // weight.element_size()
1205 import torch.nn.functional as F
1206
1207 weight = F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
1208 torch.accelerator.empty_cache()
1209 return weight
1210
1211
1212 def validate_fp8_block_shape(
1213 layer: torch.nn.Module,
1214 input_size: int,
1215 output_size: int,
1216 input_size_per_partition: int,
1217 output_partition_sizes: list[int],
1218 block_size: list[int],
1219 ) -> None:
1220 """Validate block quantization shapes for tensor parallelism."""
1221 from vllm.distributed import get_tensor_model_parallel_world_size
1222
1223 if getattr(layer, "allow_fp8_block_shape_mismatch", False):
1224 logger.debug(
1225 "Skipping FP8 block shape validation for layer %s due to detected"
1226 " mismatch allowance.",
1227 getattr(layer, "prefix", "<unknown>"),
1228 )
1229 return
1230
1231 tp_size = getattr(layer, "tp_size", get_tensor_model_parallel_world_size())
1232 block_n, block_k = block_size[0], block_size[1]
1233
1234 # Required by row parallel
1235 if (
1236 tp_size > 1
1237 and input_size // input_size_per_partition == tp_size
1238 and input_size_per_partition % block_k != 0
1239 ):
1240 raise ValueError(
1241 f"Weight input_size_per_partition = {input_size_per_partition} "
1242 f"is not divisible by weight quantization block_k = {block_k}."
1243 )
1244
1245 # Required by column parallel or enabling merged weights
1246 is_tp_split = tp_size > 1 and output_size // sum(output_partition_sizes) == tp_size
1247 is_merged_gemm = len(output_partition_sizes) > 1
1248 if is_tp_split or is_merged_gemm:
1249 sizes_to_check = output_partition_sizes
1250 if not is_tp_split and is_merged_gemm:
1251 # In case of merged matrices, we allow the last
1252 # matrix to not be a multiple of block size
1253 sizes_to_check = output_partition_sizes[:-1]
1254 for output_partition_size in sizes_to_check:
1255 if output_partition_size % block_n != 0:
1256 raise ValueError(
1257 f"Weight output_partition_size = "
1258 f"{output_partition_size} is not divisible by "
1259 f"weight quantization block_n = {block_n}."
1260 )
1261
1262
1263 def validate_fp8_block_shape_moe(
1264 intermediate_size_per_partition: int,
1265 block_size: list[int],
1266 ) -> None:
1267 """Validate fused MoE block quantization shapes for tensor parallelism."""
1268 from vllm.distributed import get_tensor_model_parallel_world_size
1269
1270 tp_size = get_tensor_model_parallel_world_size()
1271 block_n, block_k = block_size[0], block_size[1]
1272
1273 # NOTE: To ensure proper alignment of the block-wise quantization
1274 # scales, the output_size of the weights for both the gate and up
1275 # layers must be divisible by block_n.
1276 # Required by column parallel or enabling merged weights
1277 if intermediate_size_per_partition % block_n != 0:
1278 raise ValueError(
1279 f"The output_size of gate's and up's weight = "
1280 f"{intermediate_size_per_partition} is not divisible by "
1281 f"weight quantization block_n = {block_n}."
1282 )
1283 if tp_size > 1 and intermediate_size_per_partition % block_k != 0:
1284 # Required by row parallel
1285 raise ValueError(
1286 f"The input_size of down's weight = "
1287 f"{intermediate_size_per_partition} is not divisible by "
1288 f"weight quantization block_k = {block_k}."
1289 )
1290
1291
1292 def create_fp8_weight_parameter(
1293 output_size_per_partition: int,
1294 input_size_per_partition: int,
1295 weight_loader: Callable | None,
1296 ) -> torch.nn.Parameter:
1297 """Create FP8 weight parameter."""
1298 from vllm.model_executor.parameter import ModelWeightParameter
1299
1300 return ModelWeightParameter(
1301 data=torch.empty(
1302 output_size_per_partition,
1303 input_size_per_partition,
1304 dtype=torch.float8_e4m3fn,
1305 ),
1306 input_dim=1,
1307 output_dim=0,
1308 weight_loader=weight_loader,
1309 )
1310
1311
1312 def create_fp8_scale_parameter(
1313 parameter_type: torch.nn.Parameter,
1314 output_partition_sizes: list[int],
1315 input_size_per_partition: int,
1316 block_size: list[int] | None,
1317 weight_loader: Callable | None,
1318 scale_dtype: torch.dtype | None = None,
1319 ) -> torch.nn.Parameter:
1320 """Create scale parameter based on quantization strategy."""
1321 dtype = scale_dtype if scale_dtype is not None else torch.float32
1322 if parameter_type == ChannelQuantScaleParameter:
1323 scale = parameter_type(
1324 data=torch.empty((sum(output_partition_sizes), 1), dtype=dtype),
1325 output_dim=0,
1326 weight_loader=weight_loader,
1327 )
1328 elif parameter_type == BlockQuantScaleParameter:
1329 assert block_size is not None
1330 block_n, block_k = block_size[0], block_size[1]
1331 output_size_per_partition = sum(output_partition_sizes)
1332 scale = parameter_type(
1333 data=torch.empty(
1334 (output_size_per_partition + block_n - 1) // block_n,
1335 (input_size_per_partition + block_k - 1) // block_k,
1336 dtype=dtype,
1337 ),
1338 input_dim=1,
1339 output_dim=0,
1340 weight_loader=weight_loader,
1341 )
1342 elif parameter_type == PerTensorScaleParameter:
1343 scale = parameter_type(
1344 data=torch.empty(len(output_partition_sizes), dtype=dtype),
1345 weight_loader=weight_loader,
1346 )
1347 else:
1348 raise ValueError(f"Unknown parameter type: {parameter_type}")
1349
1350 if dtype == torch.float32:
1351 scale[:] = torch.finfo(torch.float32).min
1352 set_weight_attrs(scale, {"scale_type": "weight_scale"})
1353 return scale
1354
1355
1356 def create_fp8_input_scale(
1357 output_partition_sizes: list[int], weight_loader: Callable | None
1358 ) -> torch.nn.Parameter:
1359 """Create input scale parameter for static activation quantization."""
1360 from vllm.model_executor.parameter import PerTensorScaleParameter
1361
1362 scale = PerTensorScaleParameter(
1363 data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
1364 weight_loader=weight_loader,
1365 )
1366 scale[:] = torch.finfo(torch.float32).min
1367 return scale
1368
1369
1370 def process_fp8_weight_tensor_strategy(
1371 weight: torch.Tensor,
1372 weight_scale: torch.Tensor,
1373 logical_widths: list[int],
1374 input_scale: torch.Tensor | None = None,
1375 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
1376 """Requantize fused shards to one scale and return ``(K, N)`` weight."""
1377 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
1378 normalize_e4m3fn_to_e4m3fnuz,
1379 requantize_with_max_scale,
1380 )
1381
1382 if current_platform.is_fp8_fnuz() and weight.dtype == torch.float8_e4m3fn:
1383 weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
1384 weight=weight, weight_scale=weight_scale, input_scale=input_scale
1385 )
1386
1387 # Requantize with max scale
1388 weight_scale, weight = requantize_with_max_scale(
1389 weight=weight,
1390 weight_scale=weight_scale,
1391 logical_widths=logical_widths,
1392 )
1393
1394 weight = _maybe_pad_fp8_weight(weight).t()
1395 return weight, weight_scale, input_scale
1396
1397
1398 def process_fp8_weight_channel_strategy(
1399 weight: torch.Tensor,
1400 weight_scale: torch.Tensor,
1401 input_scale: torch.Tensor | None = None,
1402 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
1403 """Normalize FNUZ if needed and return ``(K, N)`` weight."""
1404 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
1405 normalize_e4m3fn_to_e4m3fnuz,
1406 )
1407
1408 if current_platform.is_fp8_fnuz() and weight.dtype == torch.float8_e4m3fn:
1409 weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
1410 weight=weight, weight_scale=weight_scale, input_scale=input_scale
1411 )
1412
1413 return weight.t(), weight_scale, input_scale
1414
1415
1416 def process_fp8_weight_block_strategy(
1417 weight: torch.Tensor,
1418 weight_scale: torch.Tensor,
1419 ) -> tuple[torch.Tensor, torch.Tensor]:
1420 """Normalize FNUZ if needed and return ``(N, K)`` weight (no transpose)."""
1421 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
1422 normalize_e4m3fn_to_e4m3fnuz,
1423 )
1424
1425 if current_platform.is_fp8_fnuz() and weight.dtype == torch.float8_e4m3fn:
1426 if weight_scale.dtype == torch.float8_e8m0fnu:
1427 # UE8M0 scales: e8m0 stores exponent-only values (2^(exp-127)),
1428 # so doubling the dequant scale == incrementing the exponent byte
1429 # by 1. Convert the OCP E4M3 weight bytes to FNUZ in place by
1430 # reinterpreting and patching the NaN sentinel (-128 in int8),
1431 # then double the UE8M0 exponent so the dequantized magnitudes
1432 # match.
1433 weight_as_int8 = weight.view(torch.int8)
1434 ROCM_FP8_NAN_AS_INT = -128
1435 weight_as_int8[weight_as_int8 == ROCM_FP8_NAN_AS_INT] = 0
1436 weight = weight_as_int8.view(torch.float8_e4m3fnuz)
1437 exp_bytes = weight_scale.view(torch.uint8)
1438 weight_scale = (
1439 (exp_bytes.to(torch.int16) + 1)
1440 .clamp(max=254)
1441 .to(torch.uint8)
1442 .view(torch.float8_e8m0fnu)
1443 )
1444 else:
1445 weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
1446 weight=weight, weight_scale=weight_scale
1447 )
1448
1449 weight = _maybe_pad_fp8_weight(weight)
1450 return weight, weight_scale
1451
1452
1453 def process_fp8_weight_tensor_strategy_moe(
1454 weight: torch.Tensor,
1455 weight_scales: torch.Tensor,
1456 shard_size: int,
1457 num_experts: int,
1458 is_act_and_mul: bool = True,
1459 ) -> tuple[torch.Tensor, torch.Tensor]:
1460 """Process moe weights for tensor-wise quantization strategy."""
1461 max_scales = weight_scales.max(dim=1).values
1462
1463 # For w1 case (i.e. not w13): there is already just one scale per expert.
1464 if not is_act_and_mul:
1465 assert weight_scales.shape[1] == 1
1466 # One scale per expert
1467 assert max_scales.shape == (num_experts,)
1468 return weight, max_scales
1469
1470 # For w13 case (common): require single scale for w13 per expert, but
1471 # on disk there is a scale for w1 and w3. Use the max to requantize.
1472 for expert_id in range(num_experts):
1473 start = 0
1474 for shard_id in range(2):
1475 dq_weight = per_tensor_dequantize(
1476 weight[expert_id][start : start + shard_size, :],
1477 weight_scales[expert_id][shard_id],
1478 )
1479 weight[expert_id][start : start + shard_size, :], _ = ops.scaled_fp8_quant(
1480 dq_weight, max_scales[expert_id]
1481 )
1482 start += shard_size
1483 return weight, max_scales
1484
1485
1486 def process_fp8_input_tensor_strategy_moe(
1487 w13_input_scale: torch.Tensor,
1488 w2_input_scale: torch.Tensor,
1489 enable_eplb: bool,
1490 ) -> tuple[torch.Tensor, torch.Tensor]:
1491 """Process moe input scales for tensor-wise quantization strategy."""
1492 if not all_close_1d(w13_input_scale) or not all_close_1d(w2_input_scale):
1493 logger.info_once(
1494 "Found input_scales that are not equal for "
1495 "fp8 MoE layer. Using the maximum across experts "
1496 "for each layer."
1497 )
1498
1499 return (
1500 amax_for_moe_activation_quant(w13_input_scale, enable_eplb),
1501 amax_for_moe_activation_quant(w2_input_scale, enable_eplb),
1502 )