@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 227 lines 7.51 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Fused SwiGLU-OAI activation (split layout) for AMD ROCm via Triton.
4
5 SwiGLU-OAI on a ``[*, 2I]`` split-layout input (gate = first half, up = second
6 half):
7
8 gate = clamp(gate, max=limit)
9 up = clamp(up, -limit, +limit)
10 out = gate * sigmoid(alpha * gate) * (up + beta)
11
12 On ROCm the dense MLP and the native MXFP8 MoE (between its two GEMMs) fell back
13 to a chain of elementwise PyTorch ops with fp32 intermediates: vLLM's shared
14 ``SiluAndMulWithClamp`` blanket-routes ROCm to ``forward_native``, and the MoE
15 applies the activation inline in PyTorch. This Triton kernel collapses that into
16 a single pass producing the ``[*, I]`` output directly, and computes in fp32
17 (rel ~1e-6 vs reference).
18
19 Note: the vectorized ``torch.ops._C.silu_and_mul_with_clamp`` op IS built on
20 ROCm and is ~1.2-2.2x faster in isolation, but the win is launch overhead that
21 HIP graphs already eliminate — measured end-to-end throughput is identical
22 (within noise), so we keep the fp32-accurate Triton kernel.
23 """
24
25 import torch
26
27 from vllm.triton_utils import tl, triton
28
29
30 @triton.jit
31 def _swiglu_oai_kernel(
32 g_ptr,
33 out_ptr,
34 n_inter,
35 stride_gm,
36 stride_gn,
37 stride_om,
38 stride_on,
39 alpha,
40 beta,
41 limit,
42 HAS_LIMIT: tl.constexpr,
43 BLOCK_I: tl.constexpr,
44 ):
45 row = tl.program_id(0)
46 pid_i = tl.program_id(1)
47 cols = pid_i * BLOCK_I + tl.arange(0, BLOCK_I)
48 mask = cols < n_inter
49 gate = tl.load(g_ptr + row * stride_gm + cols * stride_gn, mask=mask, other=0.0).to(
50 tl.float32
51 )
52 up = tl.load(
53 g_ptr + row * stride_gm + (n_inter + cols) * stride_gn,
54 mask=mask,
55 other=0.0,
56 ).to(tl.float32)
57 if HAS_LIMIT:
58 gate = tl.minimum(gate, limit)
59 up = tl.minimum(tl.maximum(up, -limit), limit)
60 out = gate * tl.sigmoid(alpha * gate) * (up + beta)
61 tl.store(
62 out_ptr + row * stride_om + cols * stride_on,
63 out.to(out_ptr.dtype.element_ty),
64 mask=mask,
65 )
66
67
68 @triton.jit
69 def _swiglu_oai_quant_kernel(
70 g_ptr,
71 aq_ptr,
72 as_ptr,
73 M,
74 n_inter,
75 stride_gm,
76 stride_gn,
77 stride_qm,
78 stride_qn,
79 stride_sm,
80 stride_sk,
81 alpha,
82 beta,
83 limit,
84 HAS_LIMIT: tl.constexpr,
85 BLOCK_M: tl.constexpr,
86 FP8_MAX: tl.constexpr,
87 TINY: tl.constexpr,
88 ):
89 """SwiGLU-OAI (split layout) fused with per-32-block MXFP8 (E4M3 + E8M0)
90 quant. Each program handles ``[BLOCK_M, 32]`` of the ``[M, I]`` output (one
91 MX block): it reads the matching gate/up columns from ``g1`` (``[M, 2I]``),
92 computes the SwiGLU in fp32, then derives the block E8M0 scale and emits the
93 FP8 values + scale in a single pass — no bf16 ``act`` round-trip to HBM.
94 """
95 pid_m = tl.program_id(0)
96 pid_b = tl.program_id(1) # which 32-element block along I
97 offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
98 offs_c = pid_b * 32 + tl.arange(0, 32)
99 m_mask = offs_m < M
100 gate = tl.load(
101 g_ptr + offs_m[:, None] * stride_gm + offs_c[None, :] * stride_gn,
102 mask=m_mask[:, None],
103 other=0.0,
104 ).to(tl.float32)
105 up = tl.load(
106 g_ptr + offs_m[:, None] * stride_gm + (n_inter + offs_c)[None, :] * stride_gn,
107 mask=m_mask[:, None],
108 other=0.0,
109 ).to(tl.float32)
110 if HAS_LIMIT:
111 gate = tl.minimum(gate, limit)
112 up = tl.minimum(tl.maximum(up, -limit), limit)
113 act = gate * tl.sigmoid(alpha * gate) * (up + beta) # [BLOCK_M, 32] fp32
114 amax = tl.maximum(tl.max(tl.abs(act), axis=1), TINY) # [BLOCK_M]
115 sb = tl.minimum(tl.maximum(tl.ceil(tl.log2(amax / FP8_MAX)) + 127.0, 0.0), 254.0)
116 # Computing the E8M0 byte-0 divisor (2^-127) can flush to zero on CDNA.
117 # Use its finite reciprocal so all-zero blocks do not produce 0/0.
118 rescale = tl.exp2(127.0 - sb)
119 aq = (act * rescale[:, None]).to(aq_ptr.dtype.element_ty)
120 tl.store(
121 aq_ptr + offs_m[:, None] * stride_qm + offs_c[None, :] * stride_qn,
122 aq,
123 mask=m_mask[:, None],
124 )
125 tl.store(
126 as_ptr + offs_m * stride_sm + pid_b * stride_sk, sb.to(tl.uint8), mask=m_mask
127 )
128
129
130 def swiglu_oai_quantize_mxfp8(
131 gate_up: torch.Tensor,
132 alpha: float,
133 beta: float,
134 limit: float | None,
135 block_m: int = 64,
136 ) -> tuple[torch.Tensor, torch.Tensor]:
137 """SwiGLU-OAI on split-layout ``[M, 2I]`` fused with MXFP8 activation-quant.
138
139 Returns ``(act_q [M, I] float8_e4m3fn, act_scale [M, I//32] uint8 E8M0)``,
140 identical to ``mxfp8_e4m3_quantize(swiglu_oai_split(gate_up))`` but in a
141 single Triton pass (no bf16 intermediate). Used between the two GEMMs of the
142 native MXFP8 MoE. Numerically equivalent to the unfused chain (bit-exact on
143 measured MoE shapes); marginally more accurate (fp32 act, no bf16 round-trip).
144 """
145 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
146 MXFP8_BLOCK_SIZE,
147 MXFP8_SCALE_DTYPE,
148 MXFP8_VALUE_DTYPE,
149 )
150
151 two_i = gate_up.shape[-1]
152 n_inter = two_i // 2
153 assert n_inter % MXFP8_BLOCK_SIZE == 0, (
154 f"fused swiglu+quant needs I % {MXFP8_BLOCK_SIZE} == 0, got I={n_inter}"
155 )
156 g1 = gate_up.reshape(-1, two_i).contiguous()
157 M = g1.shape[0]
158 aq = torch.empty((M, n_inter), dtype=MXFP8_VALUE_DTYPE, device=g1.device)
159 asc = torch.empty(
160 (M, n_inter // MXFP8_BLOCK_SIZE), dtype=MXFP8_SCALE_DTYPE, device=g1.device
161 )
162 grid = (triton.cdiv(M, block_m), n_inter // MXFP8_BLOCK_SIZE)
163 _swiglu_oai_quant_kernel[grid](
164 g1,
165 aq,
166 asc,
167 M,
168 n_inter,
169 g1.stride(0),
170 g1.stride(1),
171 aq.stride(0),
172 aq.stride(1),
173 asc.stride(0),
174 asc.stride(1),
175 float(alpha),
176 float(beta),
177 0.0 if limit is None else float(limit),
178 HAS_LIMIT=limit is not None,
179 BLOCK_M=block_m,
180 FP8_MAX=float(torch.finfo(MXFP8_VALUE_DTYPE).max),
181 TINY=float(torch.finfo(torch.float32).tiny),
182 num_warps=4,
183 )
184 return aq, asc
185
186
187 def swiglu_oai_split(
188 gate_up: torch.Tensor,
189 alpha: float,
190 beta: float,
191 limit: float | None,
192 out_dtype: torch.dtype | None = None,
193 ) -> torch.Tensor:
194 """SwiGLU-OAI on a split-layout ``[*, 2I]`` tensor -> ``[*, I]``."""
195 orig_shape = gate_up.shape
196 two_i = orig_shape[-1]
197 n_inter = two_i // 2
198 x2 = gate_up.reshape(-1, two_i)
199 m = x2.shape[0]
200 dt = out_dtype if out_dtype is not None else gate_up.dtype
201 out = torch.empty((m, n_inter), dtype=dt, device=gate_up.device)
202 # Tile tuned on gfx950. The SwiGLU intermediate is sharded across tensor
203 # parallel ranks (per-rank n_inter = I / tp: dense I=12288, MoE I=3072), and
204 # a 512-wide tile (4 warps, ~2 elems/lane) only helps once the per-rank slice
205 # is large enough to be bandwidth-bound — at TP=1 prefill that is ~1.25-1.35x
206 # faster than 256. For small sharded slices (high TP) the kernel is launch-
207 # bound (~12us) and a wide tile can slightly regress, so fall back to 256.
208 # Decode is launch-bound at every TP. num_warps=8 underfills this tile, so it
209 # is pinned to 4.
210 block_i = 512 if n_inter >= 2048 else 256
211 grid = (m, triton.cdiv(n_inter, block_i))
212 _swiglu_oai_kernel[grid](
213 x2,
214 out,
215 n_inter,
216 x2.stride(0),
217 x2.stride(1),
218 out.stride(0),
219 out.stride(1),
220 float(alpha),
221 float(beta),
222 0.0 if limit is None else float(limit),
223 HAS_LIMIT=limit is not None,
224 BLOCK_I=block_i,
225 num_warps=4,
226 )
227 return out.reshape(*orig_shape[:-1], n_inter)