@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 596 lines 20 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Utilities for Punica kernel construction."""
4
5 from vllm.triton_utils import tl, triton
6
7
8 @triton.jit
9 def _accumulate_mm(
10 tiled_a,
11 tiled_b,
12 accumulator,
13 a_scale_ptr,
14 b_scale_ptr,
15 a_scale_k_stride,
16 b_scale_k_stride,
17 iter_k,
18 group_k: tl.constexpr,
19 group_n: tl.constexpr,
20 use_fp8_w8a8: tl.constexpr,
21 ):
22 """Core matrix multiplication and accumulation logic with quantization support.
23
24 Args:
25 tiled_a (tl.tensor): Loaded tile from A matrix
26 tiled_b (tl.tensor): Loaded tile from B matrix
27 accumulator (tl.tensor): Current accumulator value
28 a_scale_ptr (tl.tensor): Scale pointer for A matrix
29 b_scale_ptr (tl.tensor): Scale pointer for B matrix
30 a_scale_k_stride (int): K dimension stride for A's block-wise scales
31 b_scale_k_stride (int): K dimension stride for B's block-wise scales
32 iter_k (int): Current iteration's global K offset
33 group_k: Block size for K dimension in block-wise quantization
34 group_n: Block size for N dimension in block-wise quantization
35 use_fp8_w8a8: Whether using FP8 W8A8 quantization
36
37 """
38 if use_fp8_w8a8:
39 if group_k > 0 and group_n > 0:
40 # Block-wise quantization: scales are loaded per block
41 offs_ks = iter_k // group_k
42 # a_scale_ptr is (BLOCK_M,) tensor of base pointers per row
43 # Load scale for current K-group, result shape: (BLOCK_M,)
44 a_scale = tl.load(a_scale_ptr + offs_ks * a_scale_k_stride)
45 # b_scale_ptr is (BLOCK_N,) tensor with N-offset pre-baked
46 # Load scale for current K-group, result shape: (BLOCK_N,)
47 b_scale = tl.load(b_scale_ptr + offs_ks * b_scale_k_stride)
48 accumulator += (
49 tl.dot(tiled_a, tiled_b) * a_scale[:, None] * b_scale[None, :]
50 )
51 else:
52 # Tensor-wise or per-channel: accumulate and scale at end
53 accumulator = tl.dot(tiled_a, tiled_b, acc=accumulator)
54 else:
55 accumulator += tl.dot(tiled_a, tiled_b)
56 return accumulator
57
58
59 @triton.jit
60 def fp8_mm_k(
61 a_ptr,
62 b_ptr,
63 a_scale_ptr,
64 b_scale_ptr,
65 ak_stride,
66 bk_stride,
67 a_scale_k_stride,
68 b_scale_k_stride,
69 offset_k,
70 K: tl.constexpr,
71 BLOCK_M: tl.constexpr,
72 BLOCK_N: tl.constexpr,
73 BLOCK_K: tl.constexpr,
74 EVEN_K: tl.constexpr,
75 SPLIT_K: tl.constexpr,
76 group_k: tl.constexpr,
77 group_n: tl.constexpr,
78 use_fp8_w8a8: tl.constexpr,
79 per_channel_quant: tl.constexpr,
80 CAST_TYPE: tl.constexpr,
81 b_dtype: tl.constexpr,
82 USE_GDC: tl.constexpr,
83 base_k,
84 ):
85 """FP8-compatible matrix multiplication kernel with quantization support.
86 Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of
87 B (k x n), iterate through the K dimension to compute the partial/complete
88 matrix block product with proper dequantization.
89
90 Args:
91 a_ptr (tl.tensor): Array of pointers, identifying rows of A
92 (FP8 or other dtype)
93 b_ptr (tl.tensor): Array of pointers, identifying columns of B
94 (FP8 dtype)
95 a_scale_ptr (tl.tensor): Scale pointer for A matrix
96 (per-token or block-wise)
97 b_scale_ptr (tl.tensor): Scale pointer for B matrix
98 (per-channel or block-wise)
99 ak_stride (int): K dimension stride of the A matrix
100 bk_stride (int): K dimension stride of the B matrix
101 a_scale_k_stride (int): K dimension stride for A's block-wise scales
102 b_scale_k_stride (int): K dimension stride for B's block-wise scales
103 offset_k (int): Base offset along K dimension
104 K: Length of the K dimension
105 BLOCK_M: M dimension of the output block m x n
106 BLOCK_N: N dimension of the output block m x n
107 BLOCK_K: K dimension atom
108 EVEN_K: True if the blocks of A and B can be loaded without masking
109 SPLIT_K: Parameter signifying parallelism in the K dimension
110 group_k: Block size for K dimension in block-wise quantization
111 group_n: Block size for N dimension in block-wise quantization
112 use_fp8_w8a8: Whether using FP8 W8A8 quantization
113 per_channel_quant: Whether using per-channel quantization
114 CAST_TYPE: if True, cast the values from the A matrix to the B
115 matrix dtype.
116 b_dtype: datatype of the B matrix
117 USE_GDC: Whether to use PDL. True indicates use.
118 base_k (int): Base offset along K dimension for current SPLIT_K group
119
120 """
121 accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
122
123 # Step size along K for each iteration
124 STEP_K = BLOCK_K * SPLIT_K
125
126 # Total number of iterations (compile-time constant)
127 num_iters = tl.cdiv(K, STEP_K)
128
129 for k in range(num_iters):
130 # Current iteration's global K offset
131 iter_k = k * STEP_K + base_k
132 block_end = iter_k + BLOCK_K
133
134 # Skip iterations that are entirely past the K boundary
135 if not EVEN_K and iter_k >= K:
136 pass
137 elif EVEN_K or block_end <= K:
138 # No masking needed: either K is evenly divisible (EVEN_K)
139 # or this block fits entirely within K
140 tiled_b = tl.load(b_ptr)
141 if USE_GDC:
142 tl.extra.cuda.gdc_wait()
143 tiled_a = tl.load(a_ptr)
144 if CAST_TYPE:
145 tiled_a = tiled_a.to(b_dtype)
146
147 accumulator = _accumulate_mm(
148 tiled_a,
149 tiled_b,
150 accumulator,
151 a_scale_ptr,
152 b_scale_ptr,
153 a_scale_k_stride,
154 b_scale_k_stride,
155 iter_k,
156 group_k,
157 group_n,
158 use_fp8_w8a8,
159 )
160 else:
161 # Partial block at the tail: mask out-of-bounds elements
162 k_offsets = tl.arange(0, BLOCK_K)
163 mask = iter_k + k_offsets < K
164 tiled_b = tl.load(b_ptr, mask=mask[:, None], other=0.0)
165 if USE_GDC:
166 tl.extra.cuda.gdc_wait()
167 tiled_a = tl.load(a_ptr, mask=mask[None, :], other=0.0)
168 if CAST_TYPE:
169 tiled_a = tiled_a.to(b_dtype)
170
171 accumulator = _accumulate_mm(
172 tiled_a,
173 tiled_b,
174 accumulator,
175 a_scale_ptr,
176 b_scale_ptr,
177 a_scale_k_stride,
178 b_scale_k_stride,
179 iter_k,
180 group_k,
181 group_n,
182 use_fp8_w8a8,
183 )
184
185 a_ptr += STEP_K * ak_stride
186 b_ptr += STEP_K * bk_stride
187
188 return accumulator
189
190
191 @triton.jit
192 def do_shrink_kernel_fp8(
193 pid_n,
194 pid_sk,
195 slice_id,
196 lora_index,
197 input_ptr,
198 lora_ptr,
199 out_ptr,
200 a_scale_ptr,
201 b_scale_ptr,
202 N,
203 K,
204 M_LEN,
205 ram,
206 # input strides
207 input_d0_stride,
208 input_d1_stride,
209 # lora strides
210 lora_d0_stride,
211 lora_d1_stride,
212 lora_d2_stride,
213 # scale strides
214 a_scale_m_stride,
215 a_scale_k_stride,
216 b_scale_l_stride,
217 b_scale_n_stride,
218 b_scale_k_stride,
219 # output strides
220 output_d0_stride,
221 output_d1_stride,
222 output_d2_stride,
223 scaling,
224 # block size for block-wise quantization
225 group_n: tl.constexpr,
226 group_k: tl.constexpr,
227 BLOCK_M: tl.constexpr,
228 BLOCK_N: tl.constexpr,
229 BLOCK_K: tl.constexpr,
230 EVEN_K: tl.constexpr,
231 SPLIT_K: tl.constexpr,
232 SLICE_NUM: tl.constexpr,
233 USE_GDC: tl.constexpr,
234 use_fp8_w8a8: tl.constexpr,
235 per_channel_quant: tl.constexpr,
236 launch_pdl: tl.constexpr,
237 ):
238 """Given an array of integers that identifies the rows of A, ram,
239 a lora index that identifies which LoRA to use from lora_ptr, lora_index,
240 a slice_id that identifies the input/output slice, compute the
241 matrix product and store in the appropriate output location.
242 """
243 # Identify the lora_ptr from slice_id.
244 if SLICE_NUM == 1:
245 cur_lora_ptr = lora_ptr
246 cur_b_scale_ptr = b_scale_ptr
247 else:
248 cur_lora_ptr = (
249 tl.load(lora_ptr + slice_id).to(tl.pointer_type(tl.float8e4nv))
250 if b_scale_ptr is not None
251 else tl.load(lora_ptr + slice_id).to(
252 tl.pointer_type(input_ptr.dtype.element_ty)
253 )
254 )
255 cur_b_scale_ptr = (
256 tl.load(b_scale_ptr + slice_id).to(tl.pointer_type(tl.float32))
257 if b_scale_ptr is not None
258 else b_scale_ptr
259 )
260
261 # Identify the column indices of B to process.
262 offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
263 rbn = tl.max_contiguous(tl.multiple_of(offset_n % N, BLOCK_N), BLOCK_N)
264
265 # Identify A and B block pointers
266 offset_k = pid_sk * BLOCK_K + tl.arange(0, BLOCK_K)
267 a_ptr = (
268 input_ptr + ram[:, None] * input_d0_stride + offset_k[None, :] * input_d1_stride
269 )
270 b_ptr = (
271 cur_lora_ptr
272 + lora_d0_stride * lora_index
273 + rbn[None, :] * lora_d1_stride
274 + offset_k[:, None] * lora_d2_stride
275 )
276
277 # Load scales for tensor-wise or per-channel quantization (outside the loop)
278 # Block-wise scales are loaded inside fp8_mm_k
279 if use_fp8_w8a8:
280 if group_k > 0 and group_n > 0:
281 # Block-wise: compute scale pointers for fp8_mm_k
282 # a_scale: per-row base pointers, shape (BLOCK_M,)
283 # Each pointer points to the start of that row's scale data
284 mm_a_scale_ptr = a_scale_ptr + ram * a_scale_m_stride
285
286 # b_scale: pre-compute N-dimension offset
287 # We need to bake in the N-group offset since fp8_mm_k doesn't know pid_n
288 n_offset = pid_n * BLOCK_N
289 offs_ns = (n_offset + tl.arange(0, BLOCK_N)) // group_n
290 # Base pointer with lora offset + N-group offset baked in, shape (BLOCK_N,)
291 mm_b_scale_ptr = (
292 cur_b_scale_ptr
293 + lora_index * b_scale_l_stride
294 + offs_ns * b_scale_n_stride
295 )
296 elif per_channel_quant:
297 # Per-channel for weights, per-token for activations
298 b_scale_ptrs = (
299 cur_b_scale_ptr + lora_index * b_scale_l_stride + rbn * b_scale_n_stride
300 )
301 b_scale = tl.load(b_scale_ptrs)
302 # Per-token activation scale
303 a_scale = tl.load(a_scale_ptr + ram * a_scale_m_stride)[:, None]
304 # For non-block-wise, pass original pointers (not used in mm loop)
305 mm_a_scale_ptr = a_scale_ptr
306 mm_b_scale_ptr = cur_b_scale_ptr
307 else:
308 # Tensor-wise quantization
309 a_scale = tl.load(a_scale_ptr) if a_scale_ptr is not None else 1.0
310 b_scale = tl.load(cur_b_scale_ptr + lora_index * b_scale_l_stride)
311 # For non-block-wise, pass original pointers (not used in mm loop)
312 mm_a_scale_ptr = a_scale_ptr
313 mm_b_scale_ptr = cur_b_scale_ptr
314 else:
315 # Non-quantized path
316 mm_a_scale_ptr = a_scale_ptr
317 mm_b_scale_ptr = cur_b_scale_ptr
318
319 # Compute partial/complete block matrix product.
320 accumulator = fp8_mm_k(
321 a_ptr,
322 b_ptr,
323 mm_a_scale_ptr,
324 mm_b_scale_ptr,
325 input_d1_stride,
326 lora_d2_stride,
327 a_scale_k_stride,
328 b_scale_k_stride,
329 offset_k,
330 K,
331 BLOCK_M,
332 BLOCK_N,
333 BLOCK_K,
334 EVEN_K,
335 SPLIT_K,
336 group_k,
337 group_n,
338 use_fp8_w8a8,
339 per_channel_quant,
340 False,
341 cur_lora_ptr.dtype.element_ty,
342 USE_GDC,
343 base_k=pid_sk * BLOCK_K,
344 )
345 # GDC launch dependents hints the runtime system to launch dependent kernels.
346 if USE_GDC:
347 tl.extra.cuda.gdc_launch_dependents()
348
349 # Apply dequantization scales for tensor-wise/per-channel quantization
350 if use_fp8_w8a8:
351 if group_k > 0 and group_n > 0:
352 # Block-wise: already applied in fp8_mm_k
353 pass
354 else:
355 # Tensor-wise or per-channel: apply scales after accumulation
356 accumulator = accumulator * a_scale * b_scale
357
358 # Apply LoRA scaling factor
359 accumulator *= scaling
360
361 # Identify the C output pointers to store the results of the accumulator.
362 offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
363 offset_cm = tl.arange(0, BLOCK_M)
364 cur_out_ptr = out_ptr if SLICE_NUM == 1 else out_ptr + slice_id * output_d0_stride
365 c_ptr = (
366 cur_out_ptr
367 + ram[:, None] * output_d1_stride
368 + offset_cn[None, :] * output_d2_stride
369 )
370 c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < N)
371
372 # Cast accumulator to output dtype
373 accumulator = accumulator.to(out_ptr.dtype.element_ty)
374
375 # handles write-back with reduction-splitting
376 if SPLIT_K == 1:
377 tl.store(c_ptr, accumulator, mask=c_mask)
378 else:
379 tl.atomic_add(c_ptr, accumulator, mask=c_mask, sem="relaxed")
380
381
382 @triton.jit
383 def do_expand_kernel_fp8(
384 pid_n,
385 lora_index,
386 slice_id,
387 input_ptr,
388 lora_ptr,
389 out_ptr,
390 a_scale_ptr,
391 b_scale_ptr,
392 N,
393 K,
394 M_LEN,
395 ram, # array identifying the rows of Input ptr to operate on
396 slice_start_loc,
397 # input ptr strides
398 input_d0_stride,
399 input_d1_stride,
400 input_d2_stride,
401 # lora ptr strides
402 ls_d0_ptr,
403 ls_d1_ptr,
404 ls_d2_ptr,
405 # scale strides
406 a_scale_m_stride,
407 a_scale_k_stride,
408 b_scale_l_stride,
409 b_scale_n_stride,
410 b_scale_k_stride,
411 # out ptr strides
412 output_d0_stride,
413 output_d1_stride,
414 # block size for block-wise quantization
415 group_n: tl.constexpr,
416 group_k: tl.constexpr,
417 # constants
418 BLOCK_M: tl.constexpr,
419 BLOCK_N: tl.constexpr,
420 BLOCK_K: tl.constexpr,
421 SAME_STRIDE: tl.constexpr,
422 SLICE_NUM: tl.constexpr,
423 EVEN_K: tl.constexpr,
424 CAST_TYPE: tl.constexpr,
425 ADD_INPUTS: tl.constexpr,
426 USE_GDC: tl.constexpr,
427 use_fp8_w8a8: tl.constexpr,
428 per_channel_quant: tl.constexpr,
429 ):
430 """FP8-compatible expand kernel for LoRA.
431 Given an array of integers that identifies the rows of A, ram,
432 a lora index that identifies which LoRA to use from lora_ptr, lora_index,
433 a slice_id that identifies the input/output slice,
434 compute the matrix product with FP8 quantization support and store in
435 the appropriate output location.
436
437 For expand kernel, the input (shrink output) may be in FP32/FP16/BF16,
438 while the LoRA B weights can be in FP8.
439
440 Supports:
441 - FP8 W8A8 quantization for LoRA B weights
442 - Block-wise quantization with configurable group_k and group_n
443 - Per-channel quantization
444 - Tensor-wise quantization
445 """
446 # ls_d*_ptr can be either an integer or a pointer
447 if SAME_STRIDE:
448 cur_lora_d0_stride = ls_d0_ptr
449 cur_lora_d1_stride = ls_d1_ptr
450 cur_lora_d2_stride = ls_d2_ptr
451 else:
452 cur_lora_d0_stride = tl.load(ls_d0_ptr + slice_id)
453 cur_lora_d1_stride = tl.load(ls_d1_ptr + slice_id)
454 cur_lora_d2_stride = tl.load(ls_d2_ptr + slice_id)
455
456 # Identify the input_ptr and lora_ptr from slice_id.
457 if SLICE_NUM == 1:
458 cur_input_ptr = input_ptr
459 if use_fp8_w8a8:
460 cur_lora_ptr = lora_ptr
461 cur_b_scale_ptr = b_scale_ptr
462 else:
463 cur_lora_ptr = lora_ptr
464 cur_b_scale_ptr = b_scale_ptr # May be None for non-quantized
465 else:
466 cur_input_ptr = input_ptr + slice_id * input_d0_stride
467 if use_fp8_w8a8:
468 cur_lora_ptr = tl.load(lora_ptr + slice_id).to(
469 tl.pointer_type(tl.float8e4nv)
470 )
471 cur_b_scale_ptr = tl.load(b_scale_ptr + slice_id).to(
472 tl.pointer_type(tl.float32)
473 )
474 else:
475 cur_lora_ptr = tl.load(lora_ptr + slice_id).to(
476 tl.pointer_type(out_ptr.dtype.element_ty)
477 )
478 cur_b_scale_ptr = (
479 tl.load(b_scale_ptr + slice_id).to(tl.pointer_type(tl.float32))
480 if b_scale_ptr is not None
481 else None
482 )
483
484 # Identify the column indices of B to process.
485 offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
486 rbn = tl.max_contiguous(tl.multiple_of(offset_n % N, BLOCK_N), BLOCK_N)
487
488 # Identify A and B block pointers
489 offset_k = tl.arange(0, BLOCK_K)
490 a_ptr = (
491 cur_input_ptr
492 + ram[:, None] * input_d1_stride
493 + offset_k[None, :] * input_d2_stride
494 )
495 b_ptr = (
496 cur_lora_ptr
497 + cur_lora_d0_stride * lora_index
498 + offset_k[:, None] * cur_lora_d2_stride
499 + rbn[None, :] * cur_lora_d1_stride
500 )
501
502 # Setup scale pointers for FP8/INT8 quantization
503 if use_fp8_w8a8:
504 if group_k > 0 and group_n > 0:
505 # Block-wise quantization - compute scale pointers for fp8_mm_k
506 # a_scale: per-row base pointers, shape (BLOCK_M,)
507 mm_a_scale_ptr = a_scale_ptr + ram * a_scale_m_stride
508
509 # b_scale: pre-compute N-dimension offset since fp8_mm_k doesn't know pid_n
510 n_offset = pid_n * BLOCK_N
511 offs_ns = (n_offset + tl.arange(0, BLOCK_N)) // group_n
512 # Base pointer with lora offset + N-group offset baked in, shape (BLOCK_N,)
513 mm_b_scale_ptr = (
514 cur_b_scale_ptr
515 + lora_index * b_scale_l_stride
516 + offs_ns * b_scale_n_stride
517 )
518 elif per_channel_quant:
519 # Per-channel for weights, shape (BLOCK_N,)
520 b_scale_ptrs = (
521 cur_b_scale_ptr + lora_index * b_scale_l_stride + rbn * b_scale_n_stride
522 )
523 b_scale = tl.load(b_scale_ptrs)
524 # Per-token activation scale, only if a_scale_ptr provided
525 a_scale = tl.load(a_scale_ptr + ram * a_scale_m_stride)[:, None]
526 # For non-block-wise, pass original pointers (not used in mm loop)
527 mm_a_scale_ptr = a_scale_ptr
528 mm_b_scale_ptr = cur_b_scale_ptr
529 else:
530 # Tensor-wise quantization
531 a_scale = tl.load(a_scale_ptr) if a_scale_ptr is not None else 1.0
532 b_scale = tl.load(cur_b_scale_ptr + lora_index * b_scale_l_stride)
533 # For non-block-wise, pass original pointers (not used in mm loop)
534 mm_a_scale_ptr = a_scale_ptr
535 mm_b_scale_ptr = cur_b_scale_ptr
536 else:
537 # Non-quantized path
538 mm_a_scale_ptr = a_scale_ptr
539 mm_b_scale_ptr = cur_b_scale_ptr
540
541 # Compute the block matrix product using fp8_mm_k
542 # Note: For expand kernel, SPLIT_K=1, so we pass 1 for SPLIT_K
543 accumulator = fp8_mm_k(
544 a_ptr,
545 b_ptr,
546 mm_a_scale_ptr,
547 mm_b_scale_ptr,
548 input_d2_stride, # ak_stride
549 cur_lora_d2_stride, # bk_stride
550 a_scale_k_stride,
551 b_scale_k_stride,
552 offset_k,
553 K,
554 BLOCK_M,
555 BLOCK_N,
556 BLOCK_K,
557 EVEN_K,
558 1, # SPLIT_K = 1 for expand kernel
559 group_k,
560 group_n,
561 use_fp8_w8a8,
562 per_channel_quant,
563 CAST_TYPE, # CAST_TYPE - cast FP8 B to A's dtype
564 cur_lora_ptr.dtype.element_ty,
565 USE_GDC,
566 base_k=0,
567 )
568
569 # Apply dequantization scales for non-block-wise quantization
570 if use_fp8_w8a8:
571 if group_k > 0 and group_n > 0:
572 pass # Already applied per block in fp8_mm_k
573 else:
574 # Tensor-wise or per-channel: apply scales after accumulation
575 accumulator = accumulator * a_scale * b_scale
576
577 tiled_c = accumulator.to(out_ptr.dtype.element_ty)
578 if SLICE_NUM == 1:
579 cur_slice_start = slice_start_loc
580 else:
581 cur_slice_start = tl.load(slice_start_loc + slice_id)
582
583 # Identify the C output pointers to store the results of the accumulator.
584 offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + cur_slice_start
585 offset_cm = tl.arange(0, BLOCK_M)
586 c_ptr = (
587 out_ptr
588 + ram[:, None] * output_d0_stride
589 + offset_cn[None, :] * output_d1_stride
590 )
591 c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < (cur_slice_start + N))
592
593 if ADD_INPUTS:
594 tiled_out = tl.load(c_ptr, mask=c_mask)
595 tiled_c += tiled_out
596 tl.store(c_ptr, tiled_c, mask=c_mask)