@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 279 lines 8.47 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import torch
5
6 from vllm.triton_utils import LOG2E, LOGE2, tl, triton
7
8
9 @triton.jit
10 def _bf16_mla_sparse_kernel(
11 q_buffer,
12 k_buffer,
13 v_buffer,
14 indices_ptr,
15 out_ptr,
16 softmax_lse_ptr,
17 max_logits_ptr,
18 seq_q,
19 seq_kv,
20 h_q,
21 dim_qk,
22 dim_v,
23 stride_q_token,
24 stride_q_head,
25 stride_k_token,
26 stride_k_head,
27 stride_v_token,
28 stride_v_head,
29 stride_out_token,
30 stride_out_head,
31 stride_lse,
32 stride_indices_token,
33 stride_indices_head,
34 sm_scale,
35 kv_group_num: tl.constexpr,
36 index_topk: tl.constexpr,
37 BLOCK_H: tl.constexpr, # block size for num heads
38 BLOCK_M: tl.constexpr, # block size for num tokens
39 BLOCK_N: tl.constexpr, # block size for indices
40 BLOCK_DV: tl.constexpr, # block size for dim_v
41 BLOCK_DMODEL: tl.constexpr, # block size for dim_nope
42 BLOCK_DPE: tl.constexpr, # block size for positional embedding
43 LOGE2: tl.constexpr,
44 ):
45 cur_q = tl.program_id(0)
46 cur_head_id = tl.program_id(1)
47 cur_kv_head_id = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H)
48
49 VALID_BLOCK_H: tl.constexpr = BLOCK_H if kv_group_num > BLOCK_H else kv_group_num
50 cur_head = cur_head_id * VALID_BLOCK_H + tl.arange(0, BLOCK_H)
51 mask_h = cur_head < (cur_head_id + 1) * VALID_BLOCK_H
52 mask_h = mask_h & (cur_head < h_q)
53
54 offs_d = tl.arange(0, BLOCK_DMODEL)
55 offs_dv = tl.arange(0, BLOCK_DV)
56
57 off_q = cur_q * stride_q_token + cur_head[:, None] * stride_q_head + offs_d[None, :]
58 mask_dmodel = offs_d < BLOCK_DMODEL
59 q = tl.load(
60 q_buffer + off_q, mask=(mask_h[:, None]) & (mask_dmodel[None, :]), other=0.0
61 )
62
63 if BLOCK_DPE > 0:
64 offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
65 off_qpe = (
66 cur_q * stride_q_token
67 + cur_head[:, None] * stride_q_head
68 + offs_dpe[None, :]
69 )
70 # assume dim_qk == BLOCK_DMODEL + BLOCK_DPE
71 mask_dpe = offs_dpe < dim_qk
72 qpe = tl.load(
73 q_buffer + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0
74 )
75
76 # Use a finite sentinel rather than -inf: if every key in a chunk is
77 # masked (e.g. leading -1 padding in `indices`) an -inf running max gives
78 # re_scale = exp2(-inf - -inf) = NaN, permanently poisoning acc / e_sum.
79 e_max = tl.zeros([BLOCK_H], dtype=tl.float32) - 1.0e30
80 e_sum = tl.zeros([BLOCK_H], dtype=tl.float32)
81 acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32)
82
83 for start_indice in range(0, index_topk, BLOCK_N):
84 offs_indice = start_indice + tl.arange(0, BLOCK_N)
85 mask_indice = offs_indice < index_topk
86 indices = tl.load(
87 indices_ptr
88 + (
89 cur_q * stride_indices_token
90 + cur_kv_head_id * stride_indices_head
91 + offs_indice
92 ),
93 mask=mask_indice,
94 other=-1,
95 )
96
97 mask_kv = (indices >= 0) & (indices < seq_kv)
98 mask_kv_d = mask_dmodel
99 offs_k = (
100 indices[None, :] * stride_k_token
101 + cur_kv_head_id * stride_k_head
102 + offs_d[:, None]
103 )
104
105 # q_nope @ k_nope
106 k = tl.load(
107 k_buffer + offs_k, mask=(mask_kv[None, :]) & (mask_kv_d[:, None]), other=0.0
108 )
109 qk = tl.dot(q, k.to(q.dtype))
110
111 if BLOCK_DPE > 0:
112 # q_rope @ k_rope
113 offs_kpe = (
114 indices[None, :] * stride_k_token
115 + cur_kv_head_id * stride_k_head
116 + offs_dpe[:, None]
117 )
118 mask_k_dpe = offs_dpe < dim_qk
119 kpe = tl.load(
120 k_buffer + offs_kpe,
121 mask=(mask_kv[None, :]) & (mask_k_dpe[:, None]),
122 other=0.0,
123 )
124 qk += tl.dot(qpe, kpe.to(q.dtype))
125
126 # apply scaling
127 qk *= sm_scale
128 qk = tl.where((mask_h[:, None]) & (mask_kv[None, :]), qk, -1.0e30)
129
130 # load v
131 mask_v_d = offs_dv < dim_v
132 offs_v = (
133 indices[:, None] * stride_v_token
134 + cur_kv_head_id * stride_v_head
135 + offs_dv[None, :]
136 )
137 v = tl.load(
138 v_buffer + offs_v, mask=(mask_kv[:, None]) & (mask_v_d[None, :]), other=0.0
139 )
140
141 # online softmax
142 n_e_max = tl.maximum(tl.max(qk, 1), e_max)
143 re_scale = tl.exp2(e_max - n_e_max)
144 p = tl.exp2(qk - n_e_max[:, None])
145 acc *= re_scale[:, None]
146
147 # score @ v
148 acc += tl.dot(p.to(v.dtype), v)
149
150 # update global sum and max
151 e_sum = e_sum * re_scale + tl.sum(p, 1)
152 e_max = n_e_max
153
154 # rescaling
155 acc /= e_sum[:, None]
156
157 max_logits = e_max * LOGE2
158 # calculate lse
159 lse = max_logits + tl.log2(e_sum) * LOGE2
160
161 # write output
162 offs_o = (
163 cur_q * stride_out_token
164 + cur_head[:, None] * stride_out_head
165 + offs_dv[None, :]
166 )
167 mask_out_d = offs_dv < dim_v
168 tl.store(
169 out_ptr + offs_o,
170 acc.to(tl.bfloat16),
171 mask=(mask_h[:, None]) & (mask_out_d[None, :]),
172 )
173
174 offs_lse = cur_q * stride_lse + cur_head
175 tl.store(softmax_lse_ptr + offs_lse, lse, mask=mask_h)
176 tl.store(max_logits_ptr + offs_lse, max_logits, mask=mask_h)
177
178
179 # reference implementation of bf16 sparse prefill kernel
180 def triton_bf16_mla_sparse_interface(
181 q: torch.Tensor, # [num_tokens, num_heads_q, dim_qk]
182 kv: torch.Tensor, # [num_tokens, num_heads_kv, dim_qk]
183 indices: torch.Tensor, # [num_tokens, num_heads_kv, topk]
184 sm_scale: float,
185 d_v: int = 512,
186 block_dpe: int = 64,
187 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
188 """Out : [num_tokens, num_heads_q, d_v]
189 max_logits : [num_tokens, num_heads_q]
190 lse : logsumexp, [num_tokens, num_heads_q]
191
192 Args:
193 q: Query tensor, [num_tokens, num_heads_q, dim_qk].
194 kv: Latent KV cache tensor.
195 indices: Sparse KV indices selected for each query token.
196 sm_scale: Softmax scale applied to the QK product.
197 d_v: Size of the value dimension.
198 block_dpe: Size of positional embedding portion of dim_qk.
199 Set to 0 when q/kv contain only the nope latent (e.g. DSv4
200 prefill where RoPE is not split out).
201
202 """
203 num_tokens, num_heads_q, dim_qk = q.shape
204 _, num_heads_kv, _ = kv.shape
205 assert dim_qk == kv.shape[2], "q and kv have different head dimensions"
206
207 # for deepseek v3.2, index topk should be 2048
208 _, _, index_topk = indices.shape
209
210 BLOCK_H = 16
211 BLOCK_DPE = block_dpe
212 BLOCK_DMODEL = dim_qk - BLOCK_DPE
213 BLOCK_M = 32
214 BLOCK_N = 16
215 BLOCK_DV = 512
216 assert d_v == BLOCK_DV, "only support d_v = 512"
217
218 assert dim_qk == BLOCK_DMODEL + BLOCK_DPE, (
219 "dim_qk does not match BLOCK_DMODEL + BLOCK_DPE"
220 )
221 assert num_heads_kv == 1, "only support kv head = 1 for now"
222 assert index_topk % BLOCK_N == 0, "index_topk must be multiple of BLOCK_N"
223
224 sm_scale *= LOG2E
225
226 kv_group_num = num_heads_q // num_heads_kv
227 grid = (
228 num_tokens,
229 triton.cdiv(num_heads_q, min(BLOCK_H, kv_group_num)),
230 )
231
232 out = torch.zeros((num_tokens, num_heads_q, d_v), dtype=q.dtype, device=q.device)
233 softmax_lse = torch.zeros(
234 (num_tokens, num_heads_q), dtype=torch.float32, device=q.device
235 )
236 max_logits = torch.zeros(
237 (num_tokens, num_heads_q), dtype=torch.float32, device=q.device
238 )
239
240 k = kv
241 v = kv[..., :d_v]
242
243 _bf16_mla_sparse_kernel[grid](
244 q_buffer=q,
245 k_buffer=k,
246 v_buffer=v,
247 indices_ptr=indices,
248 out_ptr=out,
249 softmax_lse_ptr=softmax_lse,
250 max_logits_ptr=max_logits,
251 seq_q=num_tokens,
252 seq_kv=kv.shape[0],
253 h_q=num_heads_q,
254 dim_qk=dim_qk,
255 dim_v=d_v,
256 stride_q_token=q.stride(0),
257 stride_q_head=q.stride(1),
258 stride_k_token=k.stride(0),
259 stride_k_head=k.stride(1),
260 stride_v_token=v.stride(0),
261 stride_v_head=v.stride(1),
262 stride_out_token=out.stride(0),
263 stride_out_head=out.stride(1),
264 stride_lse=softmax_lse.stride(0),
265 stride_indices_token=indices.stride(0),
266 stride_indices_head=indices.stride(1),
267 sm_scale=sm_scale,
268 kv_group_num=kv_group_num,
269 index_topk=index_topk,
270 BLOCK_H=BLOCK_H,
271 BLOCK_M=BLOCK_M,
272 BLOCK_N=BLOCK_N,
273 BLOCK_DV=BLOCK_DV,
274 BLOCK_DMODEL=BLOCK_DMODEL,
275 BLOCK_DPE=BLOCK_DPE,
276 LOGE2=LOGE2,
277 )
278
279 return out, max_logits, softmax_lse