@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 254 lines 8.56 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Triton input-staging kernel for DeepSeek V4 MegaMoE.
4
5 Quantizes hidden states to fp8 with E8M0 group scales and repacks the
6 routing top-k tensors into the int64/float32 layout that the DeepGEMM
7 MegaMoE kernels consume.
8 """
9
10 import torch
11
12 from vllm.triton_utils import tl, triton
13
14
15 @triton.jit(do_not_specialize=["num_tokens"])
16 def _prepare_megamoe_inputs_kernel(
17 hidden_states,
18 x_fp8,
19 x_sf,
20 shared_x_sf,
21 topk_ids,
22 topk_weights,
23 is_padding,
24 topk_idx_out,
25 topk_weights_out,
26 hidden_stride_m: tl.constexpr,
27 hidden_stride_k: tl.constexpr,
28 x_stride_m: tl.constexpr,
29 x_stride_k: tl.constexpr,
30 x_sf_stride_m: tl.constexpr,
31 x_sf_stride_k: tl.constexpr,
32 shared_x_sf_stride_m: tl.constexpr,
33 shared_x_sf_stride_k: tl.constexpr,
34 topk_ids_stride_m: tl.constexpr,
35 topk_ids_stride_k: tl.constexpr,
36 topk_weights_stride_m: tl.constexpr,
37 topk_weights_stride_k: tl.constexpr,
38 is_padding_stride_m: tl.constexpr,
39 topk_idx_stride_m: tl.constexpr,
40 topk_idx_stride_k: tl.constexpr,
41 topk_weights_out_stride_m: tl.constexpr,
42 topk_weights_out_stride_k: tl.constexpr,
43 num_tokens,
44 hidden_size: tl.constexpr,
45 top_k: tl.constexpr,
46 BLOCK_M: tl.constexpr,
47 BLOCK_K: tl.constexpr,
48 GROUP_K: tl.constexpr,
49 BLOCK_TOPK: tl.constexpr,
50 SHARED_BLOCK_M: tl.constexpr,
51 ) -> None:
52 token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
53 token_mask = token_id < num_tokens
54 k_block_id = tl.program_id(1)
55
56 k_offsets = k_block_id * BLOCK_K + tl.arange(0, BLOCK_K)
57 k_mask = token_mask[:, None] & (k_offsets[None, :] < hidden_size)
58 hidden = tl.load(
59 hidden_states
60 + token_id[:, None] * hidden_stride_m
61 + k_offsets[None, :] * hidden_stride_k,
62 mask=k_mask,
63 other=0.0,
64 ).to(tl.float32)
65
66 num_groups: tl.constexpr = BLOCK_K // GROUP_K
67 hidden_groups = tl.reshape(tl.abs(hidden), [BLOCK_M, num_groups, GROUP_K])
68 amax = tl.max(hidden_groups, axis=2)
69 amax = tl.maximum(amax, 1.0e-4)
70
71 scale = amax / 448.0
72 scale_bits = scale.to(tl.uint32, bitcast=True)
73 scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to(
74 tl.uint32
75 )
76 scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254)
77 rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True)
78
79 hidden_groups = tl.reshape(hidden, [BLOCK_M, num_groups, GROUP_K])
80 scaled = hidden_groups * (1.0 / rounded_scale)[:, :, None]
81 scaled = tl.reshape(scaled, [BLOCK_M, BLOCK_K])
82 fp8 = scaled.to(tl.float8e4nv)
83 tl.store(
84 x_fp8 + token_id[:, None] * x_stride_m + k_offsets[None, :] * x_stride_k,
85 fp8,
86 mask=k_mask,
87 )
88
89 scale_offsets = tl.arange(0, num_groups)
90 packed_scale = tl.sum(scale_exp << (scale_offsets[None, :] * 8), axis=1).to(
91 tl.int32
92 )
93 tl.store(
94 x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k,
95 packed_scale,
96 mask=token_mask,
97 )
98
99 # DeepGEMM's SM100 shared-expert TMA loads require the activation scales
100 # in an MN-major layout whose row permutation depends on the MegaMoE
101 # scheduler's runtime BLOCK_M. Write that view while the packed UE8M0 scale
102 # is already resident, avoiding another kernel and temporary tensor.
103 if shared_x_sf is not None:
104 m_block_id = token_id // SHARED_BLOCK_M
105 m_in_block = token_id % SHARED_BLOCK_M
106 aligned_block_m: tl.constexpr = triton.cdiv(SHARED_BLOCK_M, 128) * 128
107 transposed_m = (
108 (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32
109 )
110 shared_row = m_block_id * aligned_block_m + transposed_m
111 tl.store(
112 shared_x_sf
113 + shared_row * shared_x_sf_stride_m
114 + k_block_id * shared_x_sf_stride_k,
115 packed_scale,
116 mask=token_mask,
117 )
118
119 if k_block_id == 0:
120 topk_offsets = tl.arange(0, BLOCK_TOPK)
121 topk_mask = token_mask[:, None] & (topk_offsets[None, :] < top_k)
122 token_is_padding = tl.full((BLOCK_M,), False, tl.int1)
123 if is_padding is not None:
124 token_is_padding = tl.load(
125 is_padding + token_id * is_padding_stride_m,
126 mask=token_mask,
127 other=True,
128 )
129
130 ids = tl.load(
131 topk_ids
132 + token_id[:, None] * topk_ids_stride_m
133 + topk_offsets[None, :] * topk_ids_stride_k,
134 mask=topk_mask,
135 other=0,
136 ).to(tl.int64)
137 ids = tl.where(token_is_padding[:, None], -1, ids)
138 tl.store(
139 topk_idx_out
140 + token_id[:, None] * topk_idx_stride_m
141 + topk_offsets[None, :] * topk_idx_stride_k,
142 ids,
143 mask=topk_mask,
144 )
145
146 weights = tl.load(
147 topk_weights
148 + token_id[:, None] * topk_weights_stride_m
149 + topk_offsets[None, :] * topk_weights_stride_k,
150 mask=topk_mask,
151 other=0.0,
152 )
153 weights = tl.where(token_is_padding[:, None], 0.0, weights)
154 tl.store(
155 topk_weights_out
156 + token_id[:, None] * topk_weights_out_stride_m
157 + topk_offsets[None, :] * topk_weights_out_stride_k,
158 weights,
159 mask=topk_mask,
160 )
161
162
163 def prepare_megamoe_inputs(
164 hidden_states: torch.Tensor,
165 topk_weights: torch.Tensor,
166 topk_ids: torch.Tensor,
167 x_fp8: torch.Tensor,
168 x_sf: torch.Tensor,
169 topk_idx_out: torch.Tensor,
170 topk_weights_out: torch.Tensor,
171 is_padding: torch.Tensor | None = None,
172 shared_x_sf: torch.Tensor | None = None,
173 shared_block_m: int | None = None,
174 ) -> None:
175 num_tokens, hidden_size = hidden_states.shape
176 if num_tokens == 0:
177 return
178 if hidden_size % 128 != 0:
179 raise ValueError(
180 "DeepSeek V4 MegaMoE input staging requires hidden_size to be "
181 "a multiple of 128."
182 )
183 top_k = topk_ids.shape[1]
184 if topk_weights.shape != topk_ids.shape:
185 raise ValueError(
186 "DeepSeek V4 MegaMoE input staging requires topk_weights and "
187 "topk_ids to have the same shape."
188 )
189 if (shared_x_sf is None) != (shared_block_m is None):
190 raise ValueError(
191 "DeepSeek V4 MegaMoE shared input staging requires both "
192 "shared_x_sf and shared_block_m."
193 )
194 if shared_x_sf is not None:
195 assert shared_block_m is not None
196 if shared_block_m <= 0:
197 raise ValueError("MegaMoE shared_block_m must be positive.")
198 expected_sf_k = hidden_size // 128
199 if shared_x_sf.ndim != 2 or shared_x_sf.shape[1] != expected_sf_k:
200 raise ValueError(
201 "MegaMoE shared_x_sf must have shape "
202 f"(*, {expected_sf_k}), got {tuple(shared_x_sf.shape)}."
203 )
204 aligned_block_m = triton.cdiv(shared_block_m, 128) * 128
205 required_rows = triton.cdiv(num_tokens, shared_block_m) * aligned_block_m
206 if shared_x_sf.shape[0] < required_rows:
207 raise ValueError(
208 "MegaMoE shared_x_sf has insufficient rows: requires "
209 f"{required_rows}, got {shared_x_sf.shape[0]}."
210 )
211
212 block_k = 128
213 # On GB200, eight-row tiles win from 64 tokens; keep smaller batches untiled.
214 block_m = 8 if num_tokens >= 64 else 1
215 grid = (triton.cdiv(num_tokens, block_m), triton.cdiv(hidden_size, block_k))
216 block_topk = triton.next_power_of_2(top_k)
217 padding_stride_m = is_padding.stride(0) if is_padding is not None else 0
218 _prepare_megamoe_inputs_kernel[grid](
219 hidden_states,
220 x_fp8,
221 x_sf,
222 shared_x_sf,
223 topk_ids,
224 topk_weights,
225 is_padding,
226 topk_idx_out,
227 topk_weights_out,
228 hidden_states.stride(0),
229 hidden_states.stride(1),
230 x_fp8.stride(0),
231 x_fp8.stride(1),
232 x_sf.stride(0),
233 x_sf.stride(1),
234 shared_x_sf.stride(0) if shared_x_sf is not None else 0,
235 shared_x_sf.stride(1) if shared_x_sf is not None else 0,
236 topk_ids.stride(0),
237 topk_ids.stride(1),
238 topk_weights.stride(0),
239 topk_weights.stride(1),
240 padding_stride_m,
241 topk_idx_out.stride(0),
242 topk_idx_out.stride(1),
243 topk_weights_out.stride(0),
244 topk_weights_out.stride(1),
245 num_tokens,
246 hidden_size,
247 top_k,
248 BLOCK_M=block_m,
249 BLOCK_K=block_k,
250 GROUP_K=32,
251 BLOCK_TOPK=block_topk,
252 SHARED_BLOCK_M=shared_block_m or 1,
253 num_warps=4,
254 )