@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 918 lines 28 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 tl, triton
7 from vllm.utils.torch_utils import aux_stream
8
9 LOW_BLOCK_M = 32
10 LOW_BLOCK_N = 64
11 LOW_NUM_WARPS = 4
12 THROUGHPUT_BLOCK_M = 32
13 THROUGHPUT_BLOCK_N = 128
14 THROUGHPUT_GROUP_M = 2
15 THROUGHPUT_NUM_WARPS = 4
16 SMALL_TOKEN_THRESHOLD = 128
17 SMALL_NUM_WARPS = 2
18 Q_BLOCK_ROWS = 8
19 Q_NUM_WARPS = 2
20 KV_BLOCK_ROWS = 4
21 KV_NUM_WARPS = 2
22
23
24 @triton.jit(do_not_specialize=["rows"])
25 def _rel_proj_low_latency_kernel(
26 qkvr_ptr,
27 rel_proj_ptr,
28 rel_out_ptr,
29 log_scaling_ptr,
30 rows,
31 stride_x_t,
32 R_OFFSET: tl.constexpr,
33 NUM_Q_HEADS: tl.constexpr,
34 REL_EXTENT: tl.constexpr,
35 BLOCK_M: tl.constexpr,
36 BLOCK_N: tl.constexpr,
37 APPLY_LOG_SCALING: tl.constexpr,
38 ):
39 row = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
40 col = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
41 inner = tl.arange(0, 16)
42 token = row // NUM_Q_HEADS
43 head = row % NUM_Q_HEADS
44 relative = tl.load(
45 qkvr_ptr
46 + token[:, None] * stride_x_t
47 + R_OFFSET
48 + head[:, None] * 16
49 + inner[None, :],
50 mask=row[:, None] < rows,
51 other=0.0,
52 )
53 projection = tl.load(
54 rel_proj_ptr + inner[:, None] * REL_EXTENT + col[None, :],
55 mask=col[None, :] < REL_EXTENT,
56 other=0.0,
57 )
58 values = tl.dot(relative, projection, out_dtype=tl.float32).to(
59 rel_out_ptr.dtype.element_ty
60 )
61 values = values.to(tl.float32)
62 if APPLY_LOG_SCALING:
63 values *= tl.load(log_scaling_ptr + token, mask=row < rows, other=1.0)[:, None]
64 tl.store(
65 rel_out_ptr + row[:, None] * REL_EXTENT + col[None, :],
66 values.to(rel_out_ptr.dtype.element_ty),
67 mask=(row[:, None] < rows) & (col[None, :] < REL_EXTENT),
68 )
69
70
71 @triton.jit(do_not_specialize=["rows"])
72 def _rel_proj_throughput_kernel(
73 qkvr_ptr,
74 rel_proj_ptr,
75 rel_out_ptr,
76 log_scaling_ptr,
77 rows,
78 stride_x_t,
79 R_OFFSET: tl.constexpr,
80 NUM_Q_HEADS: tl.constexpr,
81 REL_EXTENT: tl.constexpr,
82 BLOCK_M: tl.constexpr,
83 BLOCK_N: tl.constexpr,
84 GROUP_M: tl.constexpr,
85 APPLY_LOG_SCALING: tl.constexpr,
86 ):
87 row_group = tl.program_id(0)
88 col = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
89 inner = tl.arange(0, 16)
90 projection = tl.load(
91 rel_proj_ptr + inner[:, None] * REL_EXTENT + col[None, :],
92 mask=col[None, :] < REL_EXTENT,
93 other=0.0,
94 )
95 row_offsets = tl.arange(0, BLOCK_M)
96 for group_offset in tl.static_range(GROUP_M):
97 row = (row_group * GROUP_M + group_offset) * BLOCK_M + row_offsets
98 token = row // NUM_Q_HEADS
99 head = row % NUM_Q_HEADS
100 relative = tl.load(
101 qkvr_ptr
102 + token[:, None] * stride_x_t
103 + R_OFFSET
104 + head[:, None] * 16
105 + inner[None, :],
106 mask=row[:, None] < rows,
107 other=0.0,
108 )
109 values = tl.dot(relative, projection, out_dtype=tl.float32).to(
110 rel_out_ptr.dtype.element_ty
111 )
112 values = values.to(tl.float32)
113 if APPLY_LOG_SCALING:
114 values *= tl.load(log_scaling_ptr + token, mask=row < rows, other=1.0)[
115 :, None
116 ]
117 tl.store(
118 rel_out_ptr + row[:, None] * REL_EXTENT + col[None, :],
119 values.to(rel_out_ptr.dtype.element_ty),
120 mask=(row[:, None] < rows) & (col[None, :] < REL_EXTENT),
121 )
122
123
124 def use_rel_proj_throughput(rows: int, rel_extent: int) -> bool:
125 min_rows = 8192 if rel_extent == 512 else 2048
126 return rows >= min_rows
127
128
129 def qkvr_rel_proj(
130 qkvr: torch.Tensor,
131 rel_proj: torch.Tensor,
132 rel_out: torch.Tensor,
133 log_scaling: torch.Tensor | None,
134 *,
135 num_q_heads: int,
136 num_kv_heads: int,
137 head_dim: int,
138 d_rel: int,
139 ) -> None:
140 rows = qkvr.shape[0] * num_q_heads
141 rel_extent = rel_proj.shape[1]
142 assert d_rel == 16 and rel_proj.shape[0] == 16
143 r_offset = num_q_heads * head_dim + 2 * num_kv_heads * head_dim
144 log_scaling_ptr = log_scaling if log_scaling is not None else qkvr
145 common = dict(
146 R_OFFSET=r_offset,
147 NUM_Q_HEADS=num_q_heads,
148 REL_EXTENT=rel_extent,
149 APPLY_LOG_SCALING=log_scaling is not None,
150 )
151
152 if use_rel_proj_throughput(rows, rel_extent):
153 grid = (
154 triton.cdiv(rows, THROUGHPUT_BLOCK_M * THROUGHPUT_GROUP_M),
155 triton.cdiv(rel_extent, THROUGHPUT_BLOCK_N),
156 )
157 _rel_proj_throughput_kernel[grid](
158 qkvr,
159 rel_proj,
160 rel_out,
161 log_scaling_ptr,
162 rows,
163 qkvr.stride(0),
164 BLOCK_M=THROUGHPUT_BLOCK_M,
165 BLOCK_N=THROUGHPUT_BLOCK_N,
166 GROUP_M=THROUGHPUT_GROUP_M,
167 num_warps=THROUGHPUT_NUM_WARPS,
168 **common,
169 )
170 return
171
172 grid = (
173 triton.cdiv(rows, LOW_BLOCK_M),
174 triton.cdiv(rel_extent, LOW_BLOCK_N),
175 )
176 _rel_proj_low_latency_kernel[grid](
177 qkvr,
178 rel_proj,
179 rel_out,
180 log_scaling_ptr,
181 rows,
182 qkvr.stride(0),
183 BLOCK_M=LOW_BLOCK_M,
184 BLOCK_N=LOW_BLOCK_N,
185 num_warps=LOW_NUM_WARPS,
186 **common,
187 )
188
189
190 @triton.jit(do_not_specialize=["tokens", "stride_block_table_req", "max_blocks"])
191 def _qkvr_qkv_kernel(
192 qkvr_ptr,
193 q_norm_weight_ptr,
194 q_out_ptr,
195 rel_proj_ptr,
196 rel_out_ptr,
197 k_weight_ptr,
198 v_weight_ptr,
199 k_norm_weight_ptr,
200 conv_cache_ptr,
201 key_cache_ptr,
202 value_cache_ptr,
203 positions_ptr,
204 seq_idx_ptr,
205 conv_slot_mapping_ptr,
206 conv_block_table_ptr,
207 query_start_ptr,
208 attention_slot_mapping_ptr,
209 log_scaling_ptr,
210 tokens,
211 eps,
212 stride_x_t,
213 stride_cc_block,
214 stride_cc_head,
215 stride_cc_token,
216 stride_cc_dim,
217 stride_kc_block,
218 stride_kc_token,
219 stride_kc_head,
220 stride_vc_block,
221 stride_vc_token,
222 stride_vc_head,
223 stride_block_table_req,
224 max_blocks,
225 conv_block_size,
226 attention_page_size,
227 Q_WIDTH: tl.constexpr,
228 KV_WIDTH: tl.constexpr,
229 NUM_Q_HEADS: tl.constexpr,
230 NUM_KV_HEADS: tl.constexpr,
231 HEAD_DIM: tl.constexpr,
232 WINDOW_SIZE: tl.constexpr,
233 OFF_K: tl.constexpr,
234 OFF_V: tl.constexpr,
235 APPLY_LOG_SCALING: tl.constexpr,
236 D_REL: tl.constexpr,
237 REL_EXTENT: tl.constexpr,
238 REL_EXTENT_PADDED: tl.constexpr,
239 ):
240 block = tl.program_id(0)
241 num_q_rows = tokens * NUM_Q_HEADS
242 dims = tl.arange(0, HEAD_DIM)
243
244 if block < num_q_rows:
245 row = block
246 token = row // NUM_Q_HEADS
247 head = row % NUM_Q_HEADS
248 values = tl.load(
249 qkvr_ptr + token * stride_x_t + head * HEAD_DIM + dims,
250 ).to(tl.float32)
251 weight = tl.load(q_norm_weight_ptr + dims).to(tl.float32)
252 rstd = tl.rsqrt(tl.sum(values * values, axis=0) / HEAD_DIM + eps)
253 normalized = values * rstd * weight
254 if APPLY_LOG_SCALING:
255 normalized = normalized.to(q_out_ptr.dtype.element_ty).to(tl.float32)
256 normalized *= tl.load(log_scaling_ptr + token)
257 tl.store(
258 q_out_ptr + row * HEAD_DIM + dims,
259 normalized.to(q_out_ptr.dtype.element_ty),
260 )
261 rel_cols = tl.arange(0, REL_EXTENT_PADDED)
262 rel_mask = rel_cols < REL_EXTENT
263 projected = tl.zeros([REL_EXTENT_PADDED], dtype=tl.float32)
264 rel_offset = Q_WIDTH + 2 * KV_WIDTH + head * D_REL
265 for rel_dim in tl.static_range(D_REL):
266 rel_value = tl.load(
267 qkvr_ptr + token * stride_x_t + rel_offset + rel_dim
268 ).to(tl.float32)
269 proj = tl.load(
270 rel_proj_ptr + rel_dim * REL_EXTENT + rel_cols,
271 mask=rel_mask,
272 other=0.0,
273 ).to(tl.float32)
274 projected += rel_value * proj
275 projected = projected.to(rel_out_ptr.dtype.element_ty).to(tl.float32)
276 if APPLY_LOG_SCALING:
277 projected *= tl.load(log_scaling_ptr + token)
278 tl.store(
279 rel_out_ptr + row * REL_EXTENT + rel_cols,
280 projected.to(rel_out_ptr.dtype.element_ty),
281 mask=rel_mask,
282 )
283 else:
284 row = block - num_q_rows
285 if row < tokens * NUM_KV_HEADS:
286 token = row // NUM_KV_HEADS
287 head = row % NUM_KV_HEADS
288 position = tl.load(positions_ptr + token)
289 request = tl.load(seq_idx_ptr + token)
290 conv_slot = tl.load(conv_slot_mapping_ptr + token)
291 query_start = tl.load(query_start_ptr + token)
292 attention_slot = tl.load(attention_slot_mapping_ptr + token)
293 valid = conv_slot >= 0
294
295 k_col = Q_WIDTH + head * HEAD_DIM
296 v_col = Q_WIDTH + KV_WIDTH + head * HEAD_DIM
297 k_value = tl.load(qkvr_ptr + token * stride_x_t + k_col + dims)
298 v_value = tl.load(qkvr_ptr + token * stride_x_t + v_col + dims)
299
300 safe_slot = tl.maximum(conv_slot, 0)
301 cache_base = (
302 conv_cache_ptr
303 + (safe_slot // conv_block_size) * stride_cc_block
304 + head * stride_cc_head
305 + (safe_slot % conv_block_size) * stride_cc_token
306 )
307 tl.store(
308 cache_base + (OFF_K + dims) * stride_cc_dim,
309 k_value,
310 mask=valid,
311 )
312 tl.store(
313 cache_base + (OFF_V + dims) * stride_cc_dim,
314 v_value,
315 mask=valid,
316 )
317
318 acc_k = tl.zeros([HEAD_DIM], dtype=tl.float32)
319 acc_v = tl.zeros([HEAD_DIM], dtype=tl.float32)
320 for tap in tl.static_range(WINDOW_SIZE):
321 source_position = position - (WINDOW_SIZE - 1) + tap
322 source_row = token - (WINDOW_SIZE - 1) + tap
323 in_window = valid & (source_position >= 0)
324 intra = in_window & (source_row >= query_start)
325 cached = in_window & (source_row < query_start)
326 safe_row = tl.maximum(source_row, 0)
327 source_k = tl.load(
328 qkvr_ptr + safe_row * stride_x_t + k_col + dims,
329 mask=intra,
330 other=0.0,
331 ).to(tl.float32)
332 source_v = tl.load(
333 qkvr_ptr + safe_row * stride_x_t + v_col + dims,
334 mask=intra,
335 other=0.0,
336 ).to(tl.float32)
337 safe_position = tl.maximum(source_position, 0)
338 logical_block = tl.minimum(
339 safe_position // conv_block_size, max_blocks - 1
340 )
341 physical_block = tl.load(
342 conv_block_table_ptr
343 + request * stride_block_table_req
344 + logical_block,
345 mask=cached,
346 other=0,
347 ).to(tl.int64)
348 tap_base = (
349 conv_cache_ptr
350 + physical_block * stride_cc_block
351 + head * stride_cc_head
352 + (safe_position % conv_block_size) * stride_cc_token
353 )
354 source_k += tl.load(
355 tap_base + (OFF_K + dims) * stride_cc_dim,
356 mask=cached,
357 other=0.0,
358 ).to(tl.float32)
359 source_v += tl.load(
360 tap_base + (OFF_V + dims) * stride_cc_dim,
361 mask=cached,
362 other=0.0,
363 ).to(tl.float32)
364 k_weight = tl.load(
365 k_weight_ptr + (head * HEAD_DIM + dims) * WINDOW_SIZE + tap
366 ).to(tl.float32)
367 v_weight = tl.load(
368 v_weight_ptr + (head * HEAD_DIM + dims) * WINDOW_SIZE + tap
369 ).to(tl.float32)
370 acc_k += source_k * k_weight
371 acc_v += source_v * v_weight
372
373 k_rounded = (acc_k + k_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty)
374 v_rounded = (acc_v + v_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty)
375 k_float = k_rounded.to(tl.float32)
376 k_norm_weight = tl.load(k_norm_weight_ptr + dims).to(tl.float32)
377 rstd = tl.rsqrt(tl.sum(k_float * k_float, axis=0) / HEAD_DIM + eps)
378 k_normalized = (k_float * rstd * k_norm_weight).to(
379 qkvr_ptr.dtype.element_ty
380 )
381
382 safe_attention_slot = tl.maximum(attention_slot, 0)
383 attention_block = safe_attention_slot // attention_page_size
384 attention_offset = safe_attention_slot % attention_page_size
385 tl.store(
386 key_cache_ptr
387 + attention_block * stride_kc_block
388 + attention_offset * stride_kc_token
389 + head * stride_kc_head
390 + dims,
391 k_normalized,
392 mask=attention_slot >= 0,
393 )
394 tl.store(
395 value_cache_ptr
396 + attention_block * stride_vc_block
397 + attention_offset * stride_vc_token
398 + head * stride_vc_head
399 + dims,
400 v_rounded,
401 mask=attention_slot >= 0,
402 )
403
404
405 @triton.jit(do_not_specialize=["num_rows"])
406 def _q_kernel(
407 qkvr_ptr,
408 q_norm_weight_ptr,
409 q_out_ptr,
410 log_scaling_ptr,
411 num_rows,
412 stride_x_t,
413 eps,
414 NUM_Q_HEADS: tl.constexpr,
415 HEAD_DIM: tl.constexpr,
416 BLOCK_ROWS: tl.constexpr,
417 APPLY_LOG_SCALING: tl.constexpr,
418 ):
419 rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
420 dims = tl.arange(0, HEAD_DIM)
421 row_mask = rows < num_rows
422 tokens = rows // NUM_Q_HEADS
423 heads = rows % NUM_Q_HEADS
424 values = tl.load(
425 qkvr_ptr
426 + tokens[:, None] * stride_x_t
427 + heads[:, None] * HEAD_DIM
428 + dims[None, :],
429 mask=row_mask[:, None],
430 other=0.0,
431 ).to(tl.float32)
432 weight = tl.load(q_norm_weight_ptr + dims).to(tl.float32)
433 rstd = tl.rsqrt(tl.sum(values * values, axis=1) / HEAD_DIM + eps)
434 normalized = values * rstd[:, None] * weight[None, :]
435 if APPLY_LOG_SCALING:
436 normalized = normalized.to(q_out_ptr.dtype.element_ty).to(tl.float32)
437 tau = tl.load(log_scaling_ptr + tokens, mask=row_mask, other=1.0)
438 normalized *= tau[:, None]
439 tl.store(
440 q_out_ptr + rows[:, None] * HEAD_DIM + dims[None, :],
441 normalized.to(q_out_ptr.dtype.element_ty),
442 mask=row_mask[:, None],
443 )
444
445
446 @triton.jit(do_not_specialize=["tokens", "stride_block_table_req", "max_blocks"])
447 def _kv_kernel(
448 qkvr_ptr,
449 k_weight_ptr,
450 v_weight_ptr,
451 k_norm_weight_ptr,
452 conv_cache_ptr,
453 key_cache_ptr,
454 value_cache_ptr,
455 positions_ptr,
456 seq_idx_ptr,
457 conv_slot_mapping_ptr,
458 conv_block_table_ptr,
459 query_start_ptr,
460 attention_slot_mapping_ptr,
461 tokens,
462 eps,
463 stride_x_t,
464 stride_cc_block,
465 stride_cc_head,
466 stride_cc_token,
467 stride_cc_dim,
468 stride_kc_block,
469 stride_kc_token,
470 stride_kc_head,
471 stride_vc_block,
472 stride_vc_token,
473 stride_vc_head,
474 stride_block_table_req,
475 max_blocks,
476 conv_block_size,
477 attention_page_size,
478 Q_WIDTH: tl.constexpr,
479 KV_WIDTH: tl.constexpr,
480 NUM_KV_HEADS: tl.constexpr,
481 HEAD_DIM: tl.constexpr,
482 WINDOW_SIZE: tl.constexpr,
483 OFF_K: tl.constexpr,
484 OFF_V: tl.constexpr,
485 BLOCK_ROWS: tl.constexpr,
486 ):
487 token = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
488 dims = tl.arange(0, HEAD_DIM)
489 row_mask = token < tokens
490 head_id = tl.program_id(1)
491 head = tl.full([BLOCK_ROWS], head_id, tl.int64)
492 position = tl.load(positions_ptr + token, mask=row_mask, other=0)
493 request = tl.load(seq_idx_ptr + token, mask=row_mask, other=0)
494 conv_slot = tl.load(conv_slot_mapping_ptr + token, mask=row_mask, other=-1)
495 query_start = tl.load(query_start_ptr + token, mask=row_mask, other=0)
496 attention_slot = tl.load(
497 attention_slot_mapping_ptr + token, mask=row_mask, other=-1
498 )
499 valid = row_mask & (conv_slot >= 0)
500
501 k_col = Q_WIDTH + head * HEAD_DIM
502 v_col = Q_WIDTH + KV_WIDTH + head * HEAD_DIM
503 k_value = tl.load(
504 qkvr_ptr + token[:, None] * stride_x_t + k_col[:, None] + dims[None, :],
505 mask=row_mask[:, None],
506 other=0.0,
507 )
508 v_value = tl.load(
509 qkvr_ptr + token[:, None] * stride_x_t + v_col[:, None] + dims[None, :],
510 mask=row_mask[:, None],
511 other=0.0,
512 )
513 safe_slot = tl.maximum(conv_slot, 0)
514 cache_base = (
515 conv_cache_ptr
516 + (safe_slot // conv_block_size) * stride_cc_block
517 + head * stride_cc_head
518 + (safe_slot % conv_block_size) * stride_cc_token
519 )
520 tl.store(
521 cache_base[:, None] + (OFF_K + dims[None, :]) * stride_cc_dim,
522 k_value,
523 mask=valid[:, None],
524 )
525 tl.store(
526 cache_base[:, None] + (OFF_V + dims[None, :]) * stride_cc_dim,
527 v_value,
528 mask=valid[:, None],
529 )
530
531 acc_k = tl.zeros([BLOCK_ROWS, HEAD_DIM], dtype=tl.float32)
532 acc_v = tl.zeros([BLOCK_ROWS, HEAD_DIM], dtype=tl.float32)
533 for tap in tl.static_range(WINDOW_SIZE):
534 source_position = position - (WINDOW_SIZE - 1) + tap
535 source_row = token - (WINDOW_SIZE - 1) + tap
536 in_window = valid & (source_position >= 0)
537 intra = in_window & (source_row >= query_start)
538 cached = in_window & (source_row < query_start)
539 safe_row = tl.maximum(source_row, 0)
540 source_k = tl.load(
541 qkvr_ptr + safe_row[:, None] * stride_x_t + k_col[:, None] + dims[None, :],
542 mask=intra[:, None],
543 other=0.0,
544 ).to(tl.float32)
545 source_v = tl.load(
546 qkvr_ptr + safe_row[:, None] * stride_x_t + v_col[:, None] + dims[None, :],
547 mask=intra[:, None],
548 other=0.0,
549 ).to(tl.float32)
550 safe_position = tl.maximum(source_position, 0)
551 logical_block = tl.minimum(safe_position // conv_block_size, max_blocks - 1)
552 physical_block = tl.load(
553 conv_block_table_ptr + request * stride_block_table_req + logical_block,
554 mask=cached,
555 other=0,
556 ).to(tl.int64)
557 tap_base = (
558 conv_cache_ptr
559 + physical_block * stride_cc_block
560 + head * stride_cc_head
561 + (safe_position % conv_block_size) * stride_cc_token
562 )
563 source_k += tl.load(
564 tap_base[:, None] + (OFF_K + dims[None, :]) * stride_cc_dim,
565 mask=cached[:, None],
566 other=0.0,
567 ).to(tl.float32)
568 source_v += tl.load(
569 tap_base[:, None] + (OFF_V + dims[None, :]) * stride_cc_dim,
570 mask=cached[:, None],
571 other=0.0,
572 ).to(tl.float32)
573 k_weight = tl.load(
574 k_weight_ptr + (head_id * HEAD_DIM + dims) * WINDOW_SIZE + tap
575 ).to(tl.float32)
576 v_weight = tl.load(
577 v_weight_ptr + (head_id * HEAD_DIM + dims) * WINDOW_SIZE + tap
578 ).to(tl.float32)
579 acc_k += source_k * k_weight[None, :]
580 acc_v += source_v * v_weight[None, :]
581
582 k_rounded = (acc_k + k_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty)
583 v_rounded = (acc_v + v_value.to(tl.float32)).to(qkvr_ptr.dtype.element_ty)
584 k_float = k_rounded.to(tl.float32)
585 k_norm_weight = tl.load(k_norm_weight_ptr + dims).to(tl.float32)
586 rstd = tl.rsqrt(tl.sum(k_float * k_float, axis=1) / HEAD_DIM + eps)
587 k_normalized = (k_float * rstd[:, None] * k_norm_weight[None, :]).to(
588 qkvr_ptr.dtype.element_ty
589 )
590
591 safe_attention_slot = tl.maximum(attention_slot, 0)
592 attention_block = safe_attention_slot // attention_page_size
593 attention_offset = safe_attention_slot % attention_page_size
594 attention_mask = row_mask & (attention_slot >= 0)
595 tl.store(
596 key_cache_ptr
597 + attention_block[:, None] * stride_kc_block
598 + attention_offset[:, None] * stride_kc_token
599 + head[:, None] * stride_kc_head
600 + dims[None, :],
601 k_normalized,
602 mask=attention_mask[:, None],
603 )
604 tl.store(
605 value_cache_ptr
606 + attention_block[:, None] * stride_vc_block
607 + attention_offset[:, None] * stride_vc_token
608 + head[:, None] * stride_vc_head
609 + dims[None, :],
610 v_rounded,
611 mask=attention_mask[:, None],
612 )
613
614
615 def _run_tiled_q(
616 qkvr: torch.Tensor,
617 q_norm_weight: torch.Tensor,
618 q_out: torch.Tensor,
619 positions: torch.Tensor,
620 *,
621 eps: float,
622 num_q_heads: int,
623 head_dim: int,
624 log_scaling: torch.Tensor | None,
625 ) -> None:
626 num_rows = qkvr.shape[0] * num_q_heads
627 _q_kernel[(triton.cdiv(num_rows, Q_BLOCK_ROWS),)](
628 qkvr,
629 q_norm_weight,
630 q_out,
631 log_scaling if log_scaling is not None else positions,
632 num_rows,
633 qkvr.stride(0),
634 eps,
635 NUM_Q_HEADS=num_q_heads,
636 HEAD_DIM=head_dim,
637 BLOCK_ROWS=Q_BLOCK_ROWS,
638 APPLY_LOG_SCALING=log_scaling is not None,
639 num_warps=Q_NUM_WARPS,
640 )
641
642
643 def _run_tiled_kv(
644 qkvr: torch.Tensor,
645 k_weight: torch.Tensor,
646 v_weight: torch.Tensor,
647 k_norm_weight: torch.Tensor,
648 conv_cache: torch.Tensor,
649 key_cache: torch.Tensor,
650 value_cache: torch.Tensor,
651 positions: torch.Tensor,
652 seq_idx: torch.Tensor,
653 conv_slot_mapping: torch.Tensor,
654 conv_block_table: torch.Tensor,
655 query_start: torch.Tensor,
656 attention_slot_mapping: torch.Tensor,
657 *,
658 eps: float,
659 num_q_heads: int,
660 num_kv_heads: int,
661 head_dim: int,
662 off_k: int,
663 off_v: int,
664 conv_block_size: int,
665 ) -> None:
666 tokens = qkvr.shape[0]
667 _kv_kernel[(triton.cdiv(tokens, KV_BLOCK_ROWS), num_kv_heads)](
668 qkvr,
669 k_weight,
670 v_weight,
671 k_norm_weight,
672 conv_cache,
673 key_cache,
674 value_cache,
675 positions,
676 seq_idx,
677 conv_slot_mapping,
678 conv_block_table,
679 query_start,
680 attention_slot_mapping,
681 tokens,
682 eps,
683 qkvr.stride(0),
684 conv_cache.stride(0),
685 conv_cache.stride(1),
686 conv_cache.stride(2),
687 conv_cache.stride(3),
688 key_cache.stride(0),
689 key_cache.stride(1),
690 key_cache.stride(2),
691 value_cache.stride(0),
692 value_cache.stride(1),
693 value_cache.stride(2),
694 conv_block_table.stride(0),
695 conv_block_table.shape[1],
696 conv_block_size,
697 key_cache.shape[1],
698 Q_WIDTH=num_q_heads * head_dim,
699 KV_WIDTH=num_kv_heads * head_dim,
700 NUM_KV_HEADS=num_kv_heads,
701 HEAD_DIM=head_dim,
702 WINDOW_SIZE=k_weight.shape[1],
703 OFF_K=off_k,
704 OFF_V=off_v,
705 BLOCK_ROWS=KV_BLOCK_ROWS,
706 num_warps=KV_NUM_WARPS,
707 )
708
709
710 def _run_fused_small(
711 qkvr: torch.Tensor,
712 q_norm_weight: torch.Tensor,
713 q_out: torch.Tensor,
714 rel_proj: torch.Tensor,
715 rel_out: torch.Tensor,
716 k_weight: torch.Tensor,
717 v_weight: torch.Tensor,
718 k_norm_weight: torch.Tensor,
719 conv_cache: torch.Tensor,
720 key_cache: torch.Tensor,
721 value_cache: torch.Tensor,
722 positions: torch.Tensor,
723 seq_idx: torch.Tensor,
724 conv_slot_mapping: torch.Tensor,
725 conv_block_table: torch.Tensor,
726 query_start: torch.Tensor,
727 attention_slot_mapping: torch.Tensor,
728 *,
729 eps: float,
730 num_q_heads: int,
731 num_kv_heads: int,
732 head_dim: int,
733 off_k: int,
734 off_v: int,
735 conv_block_size: int,
736 log_scaling: torch.Tensor | None,
737 ) -> None:
738 tokens = qkvr.shape[0]
739
740 num_q_rows = tokens * num_q_heads
741 grid = (num_q_rows + tokens * num_kv_heads,)
742 _qkvr_qkv_kernel[grid](
743 qkvr,
744 q_norm_weight,
745 q_out,
746 rel_proj,
747 rel_out,
748 k_weight,
749 v_weight,
750 k_norm_weight,
751 conv_cache,
752 key_cache,
753 value_cache,
754 positions,
755 seq_idx,
756 conv_slot_mapping,
757 conv_block_table,
758 query_start,
759 attention_slot_mapping,
760 log_scaling if log_scaling is not None else positions,
761 tokens,
762 eps,
763 qkvr.stride(0),
764 conv_cache.stride(0),
765 conv_cache.stride(1),
766 conv_cache.stride(2),
767 conv_cache.stride(3),
768 key_cache.stride(0),
769 key_cache.stride(1),
770 key_cache.stride(2),
771 value_cache.stride(0),
772 value_cache.stride(1),
773 value_cache.stride(2),
774 conv_block_table.stride(0),
775 conv_block_table.shape[1],
776 conv_block_size,
777 key_cache.shape[1],
778 Q_WIDTH=num_q_heads * head_dim,
779 KV_WIDTH=num_kv_heads * head_dim,
780 NUM_Q_HEADS=num_q_heads,
781 NUM_KV_HEADS=num_kv_heads,
782 HEAD_DIM=head_dim,
783 WINDOW_SIZE=k_weight.shape[1],
784 OFF_K=off_k,
785 OFF_V=off_v,
786 APPLY_LOG_SCALING=log_scaling is not None,
787 D_REL=16,
788 REL_EXTENT=rel_proj.shape[1],
789 REL_EXTENT_PADDED=triton.next_power_of_2(rel_proj.shape[1]),
790 num_warps=SMALL_NUM_WARPS,
791 )
792
793
794 def fused_qkvr_prep(
795 qkvr: torch.Tensor,
796 k_weight: torch.Tensor,
797 v_weight: torch.Tensor,
798 q_norm_weight: torch.Tensor,
799 k_norm_weight: torch.Tensor,
800 rel_proj: torch.Tensor,
801 eps: float,
802 num_q_heads: int,
803 num_kv_heads: int,
804 head_dim: int,
805 d_rel: int,
806 conv_cache: torch.Tensor,
807 key_cache: torch.Tensor,
808 value_cache: torch.Tensor,
809 positions: torch.Tensor,
810 conv_block_table: torch.Tensor,
811 seq_idx: torch.Tensor,
812 conv_slot_mapping: torch.Tensor,
813 query_start: torch.Tensor,
814 attention_slot_mapping: torch.Tensor,
815 off_k: int,
816 off_v: int,
817 conv_block_size: int,
818 log_scaling: torch.Tensor | None = None,
819 ) -> tuple[torch.Tensor, torch.Tensor]:
820 assert d_rel == 16 and rel_proj.shape[0] == 16
821 assert head_dim == 128
822 assert qkvr.is_contiguous()
823 assert k_weight.stride() == (k_weight.shape[1], 1)
824 assert v_weight.stride() == (v_weight.shape[1], 1)
825 assert rel_proj.stride() == (rel_proj.shape[1], 1)
826 assert conv_cache.stride(3) == 1
827 assert key_cache.stride(3) == 1 and value_cache.stride(3) == 1
828 tokens = qkvr.shape[0]
829 q_out = torch.empty(
830 (tokens, num_q_heads * head_dim), dtype=qkvr.dtype, device=qkvr.device
831 )
832 rel_out = torch.empty(
833 (tokens, num_q_heads, rel_proj.shape[1]),
834 dtype=qkvr.dtype,
835 device=qkvr.device,
836 )
837 if tokens == 0:
838 return q_out, rel_out
839
840 if tokens < SMALL_TOKEN_THRESHOLD:
841 _run_fused_small(
842 qkvr,
843 q_norm_weight,
844 q_out,
845 rel_proj,
846 rel_out,
847 k_weight,
848 v_weight,
849 k_norm_weight,
850 conv_cache,
851 key_cache,
852 value_cache,
853 positions,
854 seq_idx,
855 conv_slot_mapping,
856 conv_block_table,
857 query_start,
858 attention_slot_mapping,
859 eps=eps,
860 num_q_heads=num_q_heads,
861 num_kv_heads=num_kv_heads,
862 head_dim=head_dim,
863 off_k=off_k,
864 off_v=off_v,
865 conv_block_size=conv_block_size,
866 log_scaling=log_scaling,
867 )
868 return q_out, rel_out
869
870 kv_stream = aux_stream()
871 assert kv_stream is not None
872 current_stream = torch.cuda.current_stream()
873 kv_stream.wait_stream(current_stream)
874 with torch.cuda.stream(kv_stream):
875 _run_tiled_kv(
876 qkvr,
877 k_weight,
878 v_weight,
879 k_norm_weight,
880 conv_cache,
881 key_cache,
882 value_cache,
883 positions,
884 seq_idx,
885 conv_slot_mapping,
886 conv_block_table,
887 query_start,
888 attention_slot_mapping,
889 eps=eps,
890 num_q_heads=num_q_heads,
891 num_kv_heads=num_kv_heads,
892 head_dim=head_dim,
893 off_k=off_k,
894 off_v=off_v,
895 conv_block_size=conv_block_size,
896 )
897 _run_tiled_q(
898 qkvr,
899 q_norm_weight,
900 q_out,
901 positions,
902 eps=eps,
903 num_q_heads=num_q_heads,
904 head_dim=head_dim,
905 log_scaling=log_scaling,
906 )
907 qkvr_rel_proj(
908 qkvr,
909 rel_proj,
910 rel_out,
911 log_scaling,
912 num_q_heads=num_q_heads,
913 num_kv_heads=num_kv_heads,
914 head_dim=head_dim,
915 d_rel=d_rel,
916 )
917 current_stream.wait_stream(kv_stream)
918 return q_out, rel_out