@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 180 lines 5.18 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Stride-aware FP8 quantization with head_dim padding for ViT attention.
4
5 Reads directly from non-contiguous QKV views using 3D strides and pads
6 head_dim to a multiple of 16 for cuDNN compatibility.
7 """
8
9 import torch
10
11 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
12 from vllm.model_executor.layers.quantization.utils.quant_utils import (
13 get_fp8_min_max,
14 )
15 from vllm.platforms import current_platform
16 from vllm.triton_utils import HAS_TRITON, tl, triton
17 from vllm.utils.math_utils import round_up
18
19 _FP8_MIN, _FP8_MAX = get_fp8_min_max()
20
21
22 @triton.jit
23 def _quantize_pad_fp8_kernel(
24 x_ptr,
25 y_ptr,
26 scale_ptr,
27 stride_xs,
28 stride_xh,
29 stride_xd,
30 stride_ys,
31 stride_yh,
32 stride_yd,
33 num_heads,
34 n_rows,
35 n_cols,
36 n_cols_padded,
37 fp8_min: tl.constexpr,
38 fp8_max: tl.constexpr,
39 SKIP_SCALE: tl.constexpr,
40 BLOCK_M: tl.constexpr,
41 BLOCK_N: tl.constexpr,
42 ):
43 pid_m = tl.program_id(0)
44 pid_n = tl.program_id(1)
45
46 offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
47 offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
48 mask_m = offs_m < n_rows
49 mask_out = mask_m[:, None] & (offs_n[None, :] < n_cols_padded)
50 mask_in = mask_m[:, None] & (offs_n[None, :] < n_cols)
51
52 # Decompose flattened row into (token, head) for 3D stride indexing.
53 s = offs_m // num_heads
54 h = offs_m % num_heads
55
56 x_ptrs = (
57 x_ptr
58 + s[:, None] * stride_xs
59 + h[:, None] * stride_xh
60 + offs_n[None, :] * stride_xd
61 )
62 x = tl.load(x_ptrs, mask=mask_in, other=0.0).to(tl.float32)
63 if SKIP_SCALE:
64 x_q = x
65 else:
66 scale = tl.load(scale_ptr)
67 x_q = x / scale
68 x_q = tl.clamp(x_q, fp8_min, fp8_max).to(y_ptr.dtype.element_ty)
69
70 y_ptrs = (
71 y_ptr
72 + s[:, None] * stride_ys
73 + h[:, None] * stride_yh
74 + offs_n[None, :] * stride_yd
75 )
76 tl.store(y_ptrs, x_q, mask=mask_out)
77
78
79 def _get_fp8_pad_quant_config(padded_head_dim: int) -> tuple[int, int, int]:
80 block_n = triton.next_power_of_2(padded_head_dim)
81 block_n = max(16, min(block_n, 128))
82 block_m = 16
83 num_warps = 4
84 return block_m, block_n, num_warps
85
86
87 def quantize_fp8_pad_head_dim_triton(
88 tensor: torch.Tensor,
89 scale: torch.Tensor,
90 skip_scale: bool = False,
91 block_m: int | None = None,
92 block_n: int | None = None,
93 num_warps: int | None = None,
94 ) -> torch.Tensor:
95 """Quantize a 3D/4D tensor to FP8, padding head_dim to a multiple of 16.
96
97 Reads directly from the input using its 3D strides, so non-contiguous
98 views (e.g. Q/K/V slices from an interleaved QKV buffer) are handled
99 without an extra copy. Output is always a fresh contiguous tensor
100 with shape (S, H, padded_D).
101 """
102 if not HAS_TRITON:
103 raise RuntimeError("Triton is required to quantize with head_dim padding.")
104
105 original_shape = tensor.shape
106 if tensor.dim() == 4:
107 tensor = tensor.view(-1, tensor.shape[-2], tensor.shape[-1])
108 assert tensor.dim() == 3, f"Expected 3D input (S, H, D), got {tensor.dim()}D"
109 S, H, D = tensor.shape
110 padded_head_dim = round_up(D, 16)
111 out_dtype = current_platform.fp8_dtype()
112 output = torch.empty(
113 (S, H, padded_head_dim),
114 device=tensor.device,
115 dtype=out_dtype,
116 )
117
118 scale_1d = scale.reshape(-1)
119 n_rows = S * H
120
121 if block_m is None or block_n is None or num_warps is None:
122 block_m, block_n, num_warps = _get_fp8_pad_quant_config(padded_head_dim)
123
124 grid = (
125 triton.cdiv(n_rows, block_m),
126 triton.cdiv(padded_head_dim, block_n),
127 )
128
129 _quantize_pad_fp8_kernel[grid](
130 tensor,
131 output,
132 scale_1d,
133 tensor.stride(0),
134 tensor.stride(1),
135 tensor.stride(2),
136 output.stride(0),
137 output.stride(1),
138 output.stride(2),
139 H,
140 n_rows,
141 D,
142 padded_head_dim,
143 _FP8_MIN,
144 _FP8_MAX,
145 SKIP_SCALE=skip_scale,
146 BLOCK_M=block_m,
147 BLOCK_N=block_n,
148 num_warps=num_warps,
149 )
150
151 return output.view((*original_shape[:-1], padded_head_dim))
152
153
154 def quantize_fp8_maybe_pad_head_dim(
155 tensor: torch.Tensor,
156 scale: torch.Tensor,
157 fp8_quant: QuantFP8,
158 skip_scale: bool = False,
159 ) -> torch.Tensor:
160 """Quantize a 3D/4D tensor to FP8, padding head_dim to a multiple of 16
161 only when needed.
162
163 Accepts (S, H, D) or (B, S, H, D) input. Uses ``fp8_quant`` (a
164 :class:`QuantFP8` CustomOp) when head_dim is already aligned to 16
165 (no padding); otherwise falls back to a stride-aware Triton kernel
166 that pads head_dim to a multiple of 16.
167 """
168 head_dim = tensor.shape[-1]
169 if head_dim % 16 != 0:
170 return quantize_fp8_pad_head_dim_triton(tensor, scale, skip_scale=skip_scale)
171
172 if skip_scale:
173 return tensor.to(current_platform.fp8_dtype())
174
175 # QuantFP8 expects 2D: flatten all dims except (H, D).
176 orig_shape = tensor.shape
177 total_tokens = tensor.numel() // (orig_shape[-1] * orig_shape[-2])
178 tensor_2d = tensor.reshape(total_tokens, -1)
179 fp8_tensor, _ = fp8_quant(tensor_2d, scale=scale)
180 return fp8_tensor.reshape(orig_shape)