@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b3ce711b93
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 19:15 UTC

Fp8 lora dense kernel (#35242)

Signed-off-by: Yu Gong <yu3.gong@gmail.com>

yugong333 committed Mar 13, 2026 at 12:05 UTC b3ce711b93c6d960078aea0490c73bcde96adfd8
6 files changed +2439 -1
tests/lora/test_punica_ops_fp8.py new
+999
@@ -0,0 +1,999 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""FP8 accuracy tests for LoRA shrink and expand kernels.
4 +
5 +Tests the FP8 kernels by:
6 +1. Quantizing bf16 inputs/weights to FP8
7 +2. Dequantizing them back to bf16
8 +3. Running the bf16 reference (sgmv_shrink/sgmv_expand) with dequantized values
9 +4. Comparing FP8 kernel output against this dequantized reference
10 +
11 +This isolates kernel correctness from quantization precision loss,
12 +allowing much tighter tolerances than comparing against the original bf16.
13 +"""
14 +
15 +import math
16 +from threading import Lock
17 +
18 +import pytest
19 +import torch
20 +
21 +import vllm.lora.ops.torch_ops as torch_ops
22 +import vllm.lora.ops.triton_ops as triton_ops
23 +from vllm.lora.ops.triton_ops import LoRAKernelMeta
24 +from vllm.lora.ops.triton_ops.lora_expand_fp8_op import (
25 + _EXPAND_LORA_SCALE_PTR_DICT,
26 +)
27 +from vllm.lora.ops.triton_ops.lora_shrink_fp8_op import (
28 + _SHRINK_LORA_SCALE_PTR_DICT,
29 +)
30 +from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
31 +from vllm.utils.torch_utils import set_random_seed
32 +
33 +DEVICES = [f"cuda:{0}"]
34 +SEED = [0]
35 +
36 +_dict_lock = Lock()
37 +
38 +
39 +@pytest.fixture(autouse=True)
40 +def reset_device(reset_default_device):
41 + pass
42 +
43 +
44 +# ============================================================================
45 +# Reference implementations (bf16 baseline)
46 +# ============================================================================
47 +
48 +
49 +def sgmv_shrink_for_nslices(
50 + nslices,
51 + inputs_tensor,
52 + lora_weights_lst,
53 + out_tensor,
54 + b_seq_start_loc,
55 + seq_len_tensor,
56 + prompt_lora_mapping,
57 + batches,
58 + max_seq_length,
59 + num_tokens,
60 + scaling,
61 +):
62 + """Wrapper around torch_ops.sgmv_shrink that handles any nslices."""
63 + for index in range(nslices):
64 + torch_ops.sgmv_shrink(
65 + inputs_tensor,
66 + lora_weights_lst[index],
67 + out_tensor[index],
68 + b_seq_start_loc,
69 + seq_len_tensor,
70 + prompt_lora_mapping,
71 + batches,
72 + max_seq_length,
73 + num_tokens,
74 + scaling,
75 + )
76 +
77 +
78 +def sgmv_expand_for_nslices(
79 + nslices,
80 + hidden_size,
81 + inputs_tensor,
82 + lora_weights_lst,
83 + out_tensor,
84 + b_seq_start_loc,
85 + seq_len_tensor,
86 + prompt_lora_mapping,
87 + batches,
88 + max_seq_length,
89 + num_tokens,
90 + add_inputs,
91 +):
92 + """Wrapper around torch_ops.sgmv_expand that handles any nslices."""
93 + if nslices == 1:
94 + torch_ops.sgmv_expand(
95 + inputs_tensor[0],
96 + lora_weights_lst[0],
97 + out_tensor,
98 + b_seq_start_loc,
99 + seq_len_tensor,
100 + prompt_lora_mapping,
101 + batches,
102 + max_seq_length,
103 + num_tokens,
104 + add_inputs=add_inputs,
105 + )
106 + else:
107 + slice_offset = 0
108 + for index in range(nslices):
109 + torch_ops.sgmv_expand_slice(
110 + inputs_tensor[index],
111 + lora_weights_lst[index],
112 + out_tensor,
113 + b_seq_start_loc,
114 + seq_len_tensor,
115 + prompt_lora_mapping,
116 + batches,
117 + max_seq_length,
118 + num_tokens,
119 + slice_offset,
120 + hidden_size,
121 + add_inputs=add_inputs,
122 + )
123 + slice_offset += hidden_size
124 +
125 +
126 +# ============================================================================
127 +# FP8 Quantization Helpers
128 +# ============================================================================
129 +
130 +FP8_DTYPE = torch.float8_e4m3fn
131 +FP8_MAX = torch.finfo(FP8_DTYPE).max
132 +FP8_MIN = torch.finfo(FP8_DTYPE).min
133 +
134 +
135 +def quantize_to_fp8_per_tensor(
136 + tensor: torch.Tensor,
137 +) -> tuple[torch.Tensor, torch.Tensor]:
138 + """Quantize a tensor to FP8 with per-tensor scaling."""
139 + amax = tensor.abs().float().max().clamp(min=1e-12)
140 + scale = (amax / FP8_MAX).to(torch.float32)
141 + fp8_tensor = (tensor.float() / scale).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
142 + return fp8_tensor, scale.reshape(1)
143 +
144 +
145 +def quantize_to_fp8_per_channel(
146 + tensor: torch.Tensor,
147 + channel_dim: int = 0,
148 +) -> tuple[torch.Tensor, torch.Tensor]:
149 + """Quantize a tensor to FP8 with per-channel scaling.
150 +
151 + For shrink lora_a weights of shape (num_loras, rank, hidden_size):
152 + channel_dim=1 gives per-rank scaling -> scale shape (num_loras, rank)
153 + For expand lora_b weights of shape (num_loras, hidden_size, rank):
154 + channel_dim=1 gives per-hidden scaling -> scale shape (num_loras, hidden_size)
155 + """
156 + # Compute amax along all dims except the leading dims up to channel_dim+1
157 + reduce_dims = list(range(channel_dim + 1, tensor.ndim))
158 + if reduce_dims:
159 + amax = tensor.abs().float().amax(dim=reduce_dims).clamp(min=1e-12)
160 + else:
161 + amax = tensor.abs().float().clamp(min=1e-12)
162 + scale = (amax / FP8_MAX).to(torch.float32)
163 +
164 + # Expand scale for broadcasting
165 + for _ in reduce_dims:
166 + scale = scale.unsqueeze(-1)
167 + fp8_tensor = (tensor.float() / scale).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
168 + scale = scale.squeeze()
169 + if scale.ndim == 0:
170 + scale = scale.unsqueeze(0)
171 + return fp8_tensor, scale
172 +
173 +
174 +def quantize_to_fp8_per_token(
175 + tensor: torch.Tensor,
176 +) -> tuple[torch.Tensor, torch.Tensor]:
177 + """Quantize a 2D tensor to FP8 with per-token (per-row) scaling.
178 +
179 + Input shape: (num_tokens, hidden_size)
180 + Returns: (fp8_tensor, scale) where scale shape is (num_tokens, 1)
181 + """
182 + assert tensor.ndim == 2
183 + amax = tensor.abs().float().amax(dim=1, keepdim=True).clamp(min=1e-12)
184 + scale = (amax / FP8_MAX).to(torch.float32)
185 + fp8_tensor = (tensor.float() / scale).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
186 + return fp8_tensor, scale
187 +
188 +
189 +def quantize_to_fp8_blockwise(
190 + tensor: torch.Tensor,
191 + group_n: int,
192 + group_k: int,
193 +) -> tuple[torch.Tensor, torch.Tensor]:
194 + """Quantize a 2D or 3D tensor to FP8 with block-wise scaling.
195 +
196 + For a 2D tensor (num_tokens, hidden_size):
197 + Blocks of size (1, group_k) ->
198 + scale shape (num_tokens, ceil(hidden_size/group_k))
199 +
200 + For a 3D tensor (num_loras, N, K):
201 + Blocks of size (group_n, group_k) ->
202 + scale shape (num_loras, ceil(N/group_n), ceil(K/group_k))
203 + """
204 + if tensor.ndim == 2:
205 + M, K = tensor.shape
206 + n_blocks_k = math.ceil(K / group_k)
207 + scale = torch.zeros(M, n_blocks_k, dtype=torch.float32, device=tensor.device)
208 + fp8_tensor = torch.zeros_like(tensor, dtype=FP8_DTYPE)
209 + for m in range(M):
210 + for bk in range(n_blocks_k):
211 + k_start = bk * group_k
212 + k_end = min(k_start + group_k, K)
213 + block = tensor[m, k_start:k_end].float()
214 + amax = block.abs().max().clamp(min=1e-12)
215 + s = (amax / FP8_MAX).to(torch.float32)
216 + scale[m, bk] = s
217 + fp8_tensor[m, k_start:k_end] = (
218 + (block / s).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
219 + )
220 + return fp8_tensor, scale
221 + elif tensor.ndim == 3:
222 + L, N, K = tensor.shape
223 + n_blocks_n = math.ceil(N / group_n)
224 + n_blocks_k = math.ceil(K / group_k)
225 + scale = torch.zeros(
226 + L, n_blocks_n, n_blocks_k, dtype=torch.float32, device=tensor.device
227 + )
228 + fp8_tensor = torch.zeros_like(tensor, dtype=FP8_DTYPE)
229 + for li in range(L):
230 + for bn in range(n_blocks_n):
231 + for bk in range(n_blocks_k):
232 + n_start = bn * group_n
233 + n_end = min(n_start + group_n, N)
234 + k_start = bk * group_k
235 + k_end = min(k_start + group_k, K)
236 + block = tensor[li, n_start:n_end, k_start:k_end].float()
237 + amax = block.abs().max().clamp(min=1e-12)
238 + s = (amax / FP8_MAX).to(torch.float32)
239 + scale[li, bn, bk] = s
240 + fp8_tensor[li, n_start:n_end, k_start:k_end] = (
241 + (block / s).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
242 + )
243 + return fp8_tensor, scale
244 + else:
245 + raise ValueError(f"Unsupported tensor ndim: {tensor.ndim}")
246 +
247 +
248 +# ============================================================================
249 +# FP8 Dequantization Helpers
250 +# ============================================================================
251 +
252 +
253 +def dequantize_fp8_per_tensor(
254 + fp8_tensor: torch.Tensor,
255 + scale: torch.Tensor,
256 + output_dtype: torch.dtype = torch.bfloat16,
257 +) -> torch.Tensor:
258 + """Dequantize FP8 tensor with per-tensor scale back to output_dtype."""
259 + return (fp8_tensor.float() * scale.float()).to(output_dtype)
260 +
261 +
262 +def dequantize_fp8_per_channel(
263 + fp8_tensor: torch.Tensor,
264 + scale: torch.Tensor,
265 + channel_dim: int,
266 + output_dtype: torch.dtype = torch.bfloat16,
267 +) -> torch.Tensor:
268 + """Dequantize FP8 tensor with per-channel scale back to output_dtype.
269 +
270 + For 3D tensor (num_loras, N, K) with channel_dim=1:
271 + scale shape is (num_loras, N), broadcast over K.
272 + """
273 + expand_scale = scale.float()
274 + # Add trailing dims for broadcasting
275 + for _ in range(channel_dim + 1, fp8_tensor.ndim):
276 + expand_scale = expand_scale.unsqueeze(-1)
277 + return (fp8_tensor.float() * expand_scale).to(output_dtype)
278 +
279 +
280 +def dequantize_fp8_per_token(
281 + fp8_tensor: torch.Tensor,
282 + scale: torch.Tensor,
283 + output_dtype: torch.dtype = torch.bfloat16,
284 +) -> torch.Tensor:
285 + """Dequantize FP8 2D tensor with per-token scale back to output_dtype.
286 +
287 + fp8_tensor: (num_tokens, hidden_size), scale: (num_tokens, 1)
288 + """
289 + return (fp8_tensor.float() * scale.float()).to(output_dtype)
290 +
291 +
292 +def dequantize_fp8_blockwise(
293 + fp8_tensor: torch.Tensor,
294 + scale: torch.Tensor,
295 + group_n: int,
296 + group_k: int,
297 + output_dtype: torch.dtype = torch.bfloat16,
298 +) -> torch.Tensor:
299 + """Dequantize FP8 tensor with block-wise scale back to output_dtype."""
300 + if fp8_tensor.ndim == 2:
301 + M, K = fp8_tensor.shape
302 + out = torch.zeros(M, K, dtype=output_dtype, device=fp8_tensor.device)
303 + n_blocks_k = math.ceil(K / group_k)
304 + for m in range(M):
305 + for bk in range(n_blocks_k):
306 + k_start = bk * group_k
307 + k_end = min(k_start + group_k, K)
308 + out[m, k_start:k_end] = (
309 + fp8_tensor[m, k_start:k_end].float() * scale[m, bk].float()
310 + ).to(output_dtype)
311 + return out
312 + elif fp8_tensor.ndim == 3:
313 + L, N, K = fp8_tensor.shape
314 + out = torch.zeros(L, N, K, dtype=output_dtype, device=fp8_tensor.device)
315 + n_blocks_n = math.ceil(N / group_n)
316 + n_blocks_k = math.ceil(K / group_k)
317 + for l_idx in range(L):
318 + for bn in range(n_blocks_n):
319 + for bk in range(n_blocks_k):
320 + n_start = bn * group_n
321 + n_end = min(n_start + group_n, N)
322 + k_start = bk * group_k
323 + k_end = min(k_start + group_k, K)
324 + out[l_idx, n_start:n_end, k_start:k_end] = (
325 + fp8_tensor[l_idx, n_start:n_end, k_start:k_end].float()
326 + * scale[l_idx, bn, bk].float()
327 + ).to(output_dtype)
328 + return out
329 + else:
330 + raise ValueError(f"Unsupported tensor ndim: {fp8_tensor.ndim}")
331 +
332 +
333 +# ============================================================================
334 +# FP8 Data Generation
335 +# ============================================================================
336 +
337 +
338 +def generate_fp8_shrink_data(
339 + batches: int,
340 + hidden_size: int,
341 + num_loras: int,
342 + rank: int,
343 + seq_length: int,
344 + nslices: int,
345 + dtype: torch.dtype,
346 + device: str,
347 + quant_mode: str, # "per_tensor", "per_channel", "blockwise"
348 + group_k: int = 128,
349 + group_n: int = 128,
350 +):
351 + """Generate test data for FP8 shrink kernel.
352 +
353 + Shrink: output = input @ lora_a^T * scaling
354 + input: (num_tokens, hidden_size) -> quantized to FP8
355 + lora_a: (num_loras, rank, hidden_size) -> quantized to FP8
356 +
357 + Returns bf16 reference tensors, FP8 quantized tensors with scales,
358 + and dequantized bf16 tensors for accurate reference computation.
359 + """
360 + seq_len_tensor = torch.randint(seq_length, seq_length + 1, (batches,)).to(device)
361 + b_seq_start_loc = torch.cumsum(
362 + torch.tensor([0] + seq_len_tensor[:-1].tolist(), dtype=torch.long),
363 + dim=0,
364 + ).to(device)
365 + total_tokens = seq_len_tensor.sum().item()
366 +
367 + # Generate bf16 reference data
368 + inputs_bf16 = torch.randn(total_tokens, hidden_size, dtype=dtype, device=device)
369 +
370 + lora_a_weights_bf16 = []
371 + for _ in range(nslices):
372 + lora_a_weights_bf16.append(
373 + torch.randn(num_loras, rank, hidden_size, dtype=dtype, device=device)
374 + )
375 +
376 + # Quantize inputs to FP8 and dequantize back for reference
377 + if quant_mode == "blockwise":
378 + inputs_fp8, a_scale = quantize_to_fp8_blockwise(
379 + inputs_bf16, group_n=1, group_k=group_k
380 + )
381 + inputs_dequant = dequantize_fp8_blockwise(
382 + inputs_fp8,
383 + a_scale,
384 + group_n=1,
385 + group_k=group_k,
386 + output_dtype=dtype,
387 + )
388 + elif quant_mode == "per_tensor":
389 + # Per-tensor: kernel loads a single scalar from a_scale_ptr
390 + inputs_fp8, a_scale = quantize_to_fp8_per_tensor(inputs_bf16)
391 + inputs_dequant = dequantize_fp8_per_tensor(
392 + inputs_fp8,
393 + a_scale,
394 + output_dtype=dtype,
395 + )
396 + else:
397 + # per_channel: kernel loads per-token a_scale via ram indexing
398 + inputs_fp8, a_scale = quantize_to_fp8_per_token(inputs_bf16)
399 + inputs_dequant = dequantize_fp8_per_token(
400 + inputs_fp8,
401 + a_scale,
402 + output_dtype=dtype,
403 + )
404 +
405 + # Quantize lora_a weights to FP8 and dequantize back for reference
406 + b_scales = []
407 + lora_a_weights_fp8 = []
408 + lora_a_weights_dequant = []
409 + for w in lora_a_weights_bf16:
410 + if quant_mode == "per_tensor":
411 + w_fp8, w_scale = quantize_to_fp8_per_tensor(w)
412 + w_dequant = dequantize_fp8_per_tensor(w_fp8, w_scale, output_dtype=dtype)
413 + # Scale shape: (1,) -> need (num_loras,) for the kernel
414 + w_scale = w_scale.expand(num_loras).contiguous()
415 + lora_a_weights_fp8.append(w_fp8)
416 + b_scales.append(w_scale)
417 + lora_a_weights_dequant.append(w_dequant)
418 + elif quant_mode == "per_channel":
419 + # Per-channel along rank dim: scale shape (num_loras, rank)
420 + w_fp8, w_scale = quantize_to_fp8_per_channel(w, channel_dim=1)
421 + w_dequant = dequantize_fp8_per_channel(
422 + w_fp8,
423 + w_scale,
424 + channel_dim=1,
425 + output_dtype=dtype,
426 + )
427 + lora_a_weights_fp8.append(w_fp8)
428 + b_scales.append(w_scale)
429 + lora_a_weights_dequant.append(w_dequant)
430 + elif quant_mode == "blockwise":
431 + w_fp8, w_scale = quantize_to_fp8_blockwise(
432 + w, group_n=group_n, group_k=group_k
433 + )
434 + w_dequant = dequantize_fp8_blockwise(
435 + w_fp8,
436 + w_scale,
437 + group_n=group_n,
438 + group_k=group_k,
439 + output_dtype=dtype,
440 + )
441 + lora_a_weights_fp8.append(w_fp8)
442 + b_scales.append(w_scale)
443 + lora_a_weights_dequant.append(w_dequant)
444 +
445 + # Output tensor (float32 for shrink)
446 + out_tensor = torch.zeros(
447 + nslices, total_tokens, rank, dtype=torch.float32, device=device
448 + )
449 + ref_out_tensor = out_tensor.clone()
450 +
451 + # Token-to-lora mapping
452 + lora_indices_tensor = torch.randint(0, max(num_loras - 1, 1), (batches,)).to(device)
453 + token_lora_mapping = torch.zeros(total_tokens, dtype=torch.long, device=device)
454 + current_offset = 0
455 + for b_id in range(batches):
456 + lora_index = lora_indices_tensor[b_id]
457 + sl = seq_len_tensor[b_id].item()
458 + token_lora_mapping[current_offset : current_offset + sl] = lora_index
459 + current_offset += sl
460 +
461 + return {
462 + "inputs_bf16": inputs_bf16,
463 + "inputs_fp8": inputs_fp8,
464 + "inputs_dequant": inputs_dequant,
465 + "lora_a_bf16": lora_a_weights_bf16,
466 + "lora_a_fp8": lora_a_weights_fp8,
467 + "lora_a_dequant": lora_a_weights_dequant,
468 + "a_scale": a_scale,
469 + "b_scales": b_scales,
470 + "out_tensor": out_tensor,
471 + "ref_out_tensor": ref_out_tensor,
472 + "token_lora_mapping": token_lora_mapping,
473 + "seq_len_tensor": seq_len_tensor,
474 + "b_seq_start_loc": b_seq_start_loc,
475 + "lora_indices_tensor": lora_indices_tensor,
476 + "total_tokens": total_tokens,
477 + }
478 +
479 +
480 +def generate_fp8_expand_data(
481 + batches: int,
482 + hidden_size: int,
483 + num_loras: int,
484 + rank: int,
485 + seq_length: int,
486 + nslices: int,
487 + dtype: torch.dtype,
488 + device: str,
489 + quant_mode: str, # "per_tensor", "per_channel", "blockwise"
490 + group_k: int = 128,
491 + group_n: int = 128,
492 +):
493 + """Generate test data for FP8 expand kernel (w8a8).
494 +
495 + Expand: output += input @ lora_b^T
496 + input: (nslices, num_tokens, rank) -> quantized to FP8 (activations)
497 + lora_b: (num_loras, hidden_size, rank) -> quantized to FP8 (weights)
498 +
499 + In w8a8 mode, both activations and weights are FP8.
500 + Returns bf16 reference tensors, FP8 quantized tensors with scales,
501 + and dequantized bf16 tensors for accurate reference computation.
502 + """
503 + seq_len_tensor = torch.randint(seq_length, seq_length + 1, (batches,)).to(device)
504 + b_seq_start_loc = torch.cumsum(
505 + torch.tensor([0] + seq_len_tensor[:-1].tolist(), dtype=torch.long),
506 + dim=0,
507 + ).to(device)
508 + total_tokens = seq_len_tensor.sum().item()
509 +
510 + # Generate bf16 input (shrink output) and quantize to FP8
511 + inputs_bf16 = torch.randn(nslices, total_tokens, rank, dtype=dtype, device=device)
512 +
513 + # Quantize input to FP8 and dequantize back for reference
514 + inputs_2d_all = inputs_bf16.reshape(-1, rank)
515 + if quant_mode == "blockwise":
516 + # For blockwise, the kernel indexes a_scale by token id (0..total_tokens-1)
517 + # shared across slices. Compute shared scale across slices, then quantize.
518 + # First compute per-token-per-block scale across all slices
519 + n_blocks_k = math.ceil(rank / group_k)
520 + a_scale = torch.zeros(
521 + total_tokens, n_blocks_k, dtype=torch.float32, device=device
522 + )
523 + for m in range(total_tokens):
524 + for bk in range(n_blocks_k):
525 + k_start = bk * group_k
526 + k_end = min(k_start + group_k, rank)
527 + # Max across all slices for this token and block
528 + block_amax = torch.tensor(0.0, device=device)
529 + for s in range(nslices):
530 + block = inputs_bf16[s, m, k_start:k_end].float()
531 + block_amax = torch.max(
532 + block_amax, block.abs().max().clamp(min=1e-12)
533 + )
534 + a_scale[m, bk] = (block_amax / FP8_MAX).to(torch.float32)
535 +
536 + # Quantize all slices with the shared scale
537 + inputs_fp8_list = []
538 + inputs_dequant_list = []
539 + for s in range(nslices):
540 + slice_2d = inputs_bf16[s] # (total_tokens, rank)
541 + fp8_slice = torch.zeros_like(slice_2d, dtype=FP8_DTYPE)
542 + dequant_slice = torch.zeros_like(slice_2d)
543 + for m in range(total_tokens):
544 + for bk in range(n_blocks_k):
545 + k_start = bk * group_k
546 + k_end = min(k_start + group_k, rank)
547 + block = slice_2d[m, k_start:k_end].float()
548 + s_val = a_scale[m, bk]
549 + fp8_slice[m, k_start:k_end] = (
550 + (block / s_val).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
551 + )
552 + dequant_slice[m, k_start:k_end] = (
553 + fp8_slice[m, k_start:k_end].float() * s_val.float()
554 + ).to(dtype)
555 + inputs_fp8_list.append(fp8_slice)
556 + inputs_dequant_list.append(dequant_slice)
557 + inputs_fp8 = torch.stack(inputs_fp8_list, dim=0)
558 + inputs_dequant = torch.stack(inputs_dequant_list, dim=0)
559 + elif quant_mode == "per_tensor":
560 + # Per-tensor: kernel loads a single scalar from a_scale_ptr
561 + inputs_fp8_2d, a_scale = quantize_to_fp8_per_tensor(inputs_2d_all)
562 + inputs_dequant_2d = dequantize_fp8_per_tensor(
563 + inputs_fp8_2d,
564 + a_scale,
565 + output_dtype=dtype,
566 + )
567 + inputs_fp8 = inputs_fp8_2d.reshape(nslices, total_tokens, rank)
568 + inputs_dequant = inputs_dequant_2d.reshape(nslices, total_tokens, rank)
569 + else:
570 + # per_channel: kernel loads per-token a_scale via ram indexing.
571 + # The kernel uses the same a_scale for all slices (indexed by token
572 + # id 0..total_tokens-1), so we compute a shared per-token scale
573 + # across all slices, then quantize each slice with that shared scale.
574 + per_slice_views = [inputs_bf16[s] for s in range(nslices)]
575 + # (nslices, total_tokens, rank) -> max across slices per token
576 + stacked = torch.stack(per_slice_views, dim=0) # (nslices, tokens, rank)
577 + amax = stacked.abs().float().amax(dim=(0, 2), keepdim=False).clamp(min=1e-12)
578 + # amax shape: (total_tokens,)
579 + a_scale = (amax / FP8_MAX).to(torch.float32).unsqueeze(1) # (tokens, 1)
580 + # Quantize all slices with the shared scale
581 + inputs_fp8_2d = (
582 + (inputs_2d_all.float() / a_scale.repeat(nslices, 1))
583 + .clamp(FP8_MIN, FP8_MAX)
584 + .to(FP8_DTYPE)
585 + )
586 + inputs_dequant_2d = (
587 + inputs_fp8_2d.float() * a_scale.repeat(nslices, 1).float()
588 + ).to(dtype)
589 + inputs_fp8 = inputs_fp8_2d.reshape(nslices, total_tokens, rank)
590 + inputs_dequant = inputs_dequant_2d.reshape(nslices, total_tokens, rank)
591 +
592 + # Generate bf16 LoRA B weights
593 + lora_b_weights_bf16 = []
594 + for _ in range(nslices):
595 + lora_b_weights_bf16.append(
596 + torch.randn(num_loras, hidden_size, rank, dtype=dtype, device=device)
597 + )
598 +
599 + # Quantize LoRA B weights to FP8 and dequantize back for reference
600 + b_scales = []
601 + lora_b_weights_fp8 = []
602 + lora_b_weights_dequant = []
603 + for w in lora_b_weights_bf16:
604 + if quant_mode == "per_tensor":
605 + w_fp8, w_scale = quantize_to_fp8_per_tensor(w)
606 + w_dequant = dequantize_fp8_per_tensor(w_fp8, w_scale, output_dtype=dtype)
607 + w_scale = w_scale.expand(num_loras).contiguous()
608 + lora_b_weights_fp8.append(w_fp8)
609 + b_scales.append(w_scale)
610 + lora_b_weights_dequant.append(w_dequant)
611 + elif quant_mode == "per_channel":
612 + # Per-channel along hidden_size dim: scale (num_loras, hidden_size)
613 + w_fp8, w_scale = quantize_to_fp8_per_channel(w, channel_dim=1)
614 + w_dequant = dequantize_fp8_per_channel(
615 + w_fp8,
616 + w_scale,
617 + channel_dim=1,
618 + output_dtype=dtype,
619 + )
620 + lora_b_weights_fp8.append(w_fp8)
621 + b_scales.append(w_scale)
622 + lora_b_weights_dequant.append(w_dequant)
623 + elif quant_mode == "blockwise":
624 + w_fp8, w_scale = quantize_to_fp8_blockwise(
625 + w, group_n=group_n, group_k=group_k
626 + )
627 + w_dequant = dequantize_fp8_blockwise(
628 + w_fp8,
629 + w_scale,
630 + group_n=group_n,
631 + group_k=group_k,
632 + output_dtype=dtype,
633 + )
634 + lora_b_weights_fp8.append(w_fp8)
635 + b_scales.append(w_scale)
636 + lora_b_weights_dequant.append(w_dequant)
637 +
638 + # Output tensor (initialized randomly for add_inputs)
639 + out_tensor = torch.randn(
640 + total_tokens, hidden_size * nslices, dtype=dtype, device=device
641 + )
642 + ref_out_tensor = out_tensor.clone()
643 +
644 + # Token-to-lora mapping
645 + lora_indices_tensor = torch.randint(0, max(num_loras - 1, 1), (batches,)).to(device)
646 + token_lora_mapping = torch.zeros(total_tokens, dtype=torch.long, device=device)
647 + current_offset = 0
648 + for b_id in range(batches):
649 + lora_index = lora_indices_tensor[b_id]
650 + sl = seq_len_tensor[b_id].item()
651 + token_lora_mapping[current_offset : current_offset + sl] = lora_index
652 + current_offset += sl
653 +
654 + return {
655 + "inputs_bf16": inputs_bf16,
656 + "inputs_fp8": inputs_fp8,
657 + "inputs_dequant": inputs_dequant,
658 + "a_scale": a_scale,
659 + "lora_b_bf16": lora_b_weights_bf16,
660 + "lora_b_fp8": lora_b_weights_fp8,
661 + "lora_b_dequant": lora_b_weights_dequant,
662 + "b_scales": b_scales,
663 + "out_tensor": out_tensor,
664 + "ref_out_tensor": ref_out_tensor,
665 + "token_lora_mapping": token_lora_mapping,
666 + "seq_len_tensor": seq_len_tensor,
667 + "b_seq_start_loc": b_seq_start_loc,
668 + "lora_indices_tensor": lora_indices_tensor,
669 + "total_tokens": total_tokens,
670 + }
671 +
672 +
673 +# ============================================================================
674 +# FP8 Shrink Kernel Check
675 +# ============================================================================
676 +
677 +
678 +def check_lora_shrink_fp8_kernel(
679 + batches: int,
680 + num_loras: int,
681 + rank: int,
682 + hidden_size: int,
683 + nslices: int,
684 + dtype: torch.dtype,
685 + device: str,
686 + seq_length: int,
687 + scaling: float,
688 + quant_mode: str,
689 + group_k: int = 128,
690 + group_n: int = 128,
691 +):
692 + """Test FP8 shrink kernel against dequantized bf16 reference.
693 +
694 + Instead of comparing FP8 kernel output against the original bf16 reference
695 + (which conflates quantization error with kernel error), we:
696 + 1. Quantize bf16 inputs/weights to FP8
697 + 2. Dequantize them back to bf16
698 + 3. Run the bf16 reference (sgmv_shrink) with the dequantized values
699 + 4. Compare FP8 kernel output against this dequantized reference
700 +
701 + This isolates kernel correctness from quantization precision loss,
702 + allowing much tighter tolerances.
703 + """
704 + data = generate_fp8_shrink_data(
705 + batches,
706 + hidden_size,
707 + num_loras,
708 + rank,
709 + seq_length,
710 + nslices,
711 + dtype,
712 + device,
713 + quant_mode,
714 + group_k,
715 + group_n,
716 + )
717 +
718 + total_tokens = data["total_tokens"]
719 +
720 + # Setup LoRA kernel metadata
721 + lora_meta = LoRAKernelMeta.make(
722 + max_loras=num_loras, max_num_tokens=total_tokens, device=device
723 + )
724 + lora_meta.prepare_tensors(data["token_lora_mapping"])
725 +
726 + out_tensor = data["out_tensor"]
727 +
728 + # Determine quantization params for the kernel
729 + per_channel = quant_mode == "per_channel"
730 + gk = group_k if quant_mode == "blockwise" else 0
731 + gn = group_n if quant_mode == "blockwise" else 0
732 +
733 + with _dict_lock:
734 + _LORA_A_PTR_DICT.clear()
735 + _SHRINK_LORA_SCALE_PTR_DICT.clear()
736 + triton_ops.lora_shrink_fp8(
737 + data["inputs_fp8"],
738 + data["lora_a_fp8"],
739 + out_tensor,
740 + *lora_meta.meta_args(token_nums=total_tokens, specialize_active_lora=False),
741 + scaling,
742 + data["b_scales"],
743 + a_scale=data["a_scale"],
744 + group_k=gk,
745 + group_n=gn,
746 + use_fp8_w8a8=True,
747 + per_channel_quant=per_channel,
748 + )
749 +
750 + # Compute reference using dequantized (round-tripped) tensors.
751 + # This means the reference sees the same quantization error as the kernel,
752 + # so any difference is purely kernel error.
753 + ref_out_tensor = data["ref_out_tensor"]
754 + max_seq_length = data["seq_len_tensor"].max().item()
755 + sgmv_shrink_for_nslices(
756 + nslices,
757 + data["inputs_dequant"],
758 + data["lora_a_dequant"],
759 + ref_out_tensor,
760 + data["b_seq_start_loc"],
761 + data["seq_len_tensor"],
762 + data["lora_indices_tensor"],
763 + batches,
764 + max_seq_length,
765 + total_tokens,
766 + scaling,
767 + )
768 +
769 + # With dequantized reference, we can use much tighter tolerances
770 + # since we're only measuring kernel error, not quantization error.
771 + # Blockwise accumulation order differs from the bf16 reference, so
772 + # allow a slightly larger margin for sporadic rounding outliers.
773 + rtol, atol = 0.1, 0.25
774 + torch.testing.assert_close(
775 + out_tensor.to(dtype), ref_out_tensor.to(dtype), rtol=rtol, atol=atol
776 + )
777 +
778 +
779 +# ============================================================================
780 +# FP8 Expand Kernel Check
781 +# ============================================================================
782 +
783 +
784 +def check_lora_expand_fp8_kernel(
785 + batches: int,
786 + num_loras: int,
787 + rank: int,
788 + hidden_size: int,
789 + nslices: int,
790 + dtype: torch.dtype,
791 + device: str,
792 + seq_length: int,
793 + add_inputs: bool,
794 + quant_mode: str,
795 + group_k: int = 128,
796 + group_n: int = 128,
797 +):
798 + """Test FP8 expand kernel (w8a8) against dequantized bf16 reference.
799 +
800 + Instead of comparing FP8 kernel output against the original bf16 reference
801 + (which conflates quantization error with kernel error), we:
802 + 1. Quantize bf16 inputs/weights to FP8
803 + 2. Dequantize them back to bf16
804 + 3. Run the bf16 reference (sgmv_expand) with the dequantized values
805 + 4. Compare FP8 kernel output against this dequantized reference
806 +
807 + This isolates kernel correctness from quantization precision loss,
808 + allowing much tighter tolerances.
809 + """
810 + data = generate_fp8_expand_data(
811 + batches,
812 + hidden_size,
813 + num_loras,
814 + rank,
815 + seq_length,
816 + nslices,
817 + dtype,
818 + device,
819 + quant_mode,
820 + group_k,
821 + group_n,
822 + )
823 +
824 + total_tokens = data["total_tokens"]
825 +
826 + # Setup LoRA kernel metadata
827 + lora_meta = LoRAKernelMeta.make(
828 + max_loras=num_loras, max_num_tokens=total_tokens, device=device
829 + )
830 + lora_meta.prepare_tensors(data["token_lora_mapping"])
831 +
832 + out_tensor = data["out_tensor"]
833 +
834 + # Determine quantization params for the kernel
835 + per_channel = quant_mode == "per_channel"
836 + gk = group_k if quant_mode == "blockwise" else 0
837 + gn = group_n if quant_mode == "blockwise" else 0
838 +
839 + with _dict_lock:
840 + _LORA_B_PTR_DICT.clear()
841 + _EXPAND_LORA_SCALE_PTR_DICT.clear()
842 + triton_ops.lora_expand_fp8(
843 + data["inputs_fp8"],
844 + data["lora_b_fp8"],
845 + out_tensor,
846 + *lora_meta.meta_args(token_nums=total_tokens, specialize_active_lora=False),
847 + data["b_scales"],
848 + a_scale=data["a_scale"],
849 + offset_start=0,
850 + add_inputs=add_inputs,
851 + group_k=gk,
852 + group_n=gn,
853 + use_fp8_w8a8=True,
854 + per_channel_quant=per_channel,
855 + )
856 +
857 + # Compute reference using dequantized (round-tripped) tensors.
858 + ref_out_tensor = data["ref_out_tensor"]
859 + max_seq_length = data["seq_len_tensor"].max().item()
860 + sgmv_expand_for_nslices(
861 + nslices,
862 + hidden_size,
863 + data["inputs_dequant"],
864 + data["lora_b_dequant"],
865 + ref_out_tensor,
866 + data["b_seq_start_loc"],
867 + data["seq_len_tensor"],
868 + data["lora_indices_tensor"],
869 + batches,
870 + max_seq_length,
871 + total_tokens,
872 + add_inputs=add_inputs,
873 + )
874 +
875 + # With dequantized reference, we can use much tighter tolerances
876 + # since we're only measuring kernel error, not quantization error.
877 + rtol, atol = 0.1, 0.15
878 + torch.testing.assert_close(out_tensor, ref_out_tensor, rtol=rtol, atol=atol)
879 +
880 +
881 +# ============================================================================
882 +# FP8 Test Parameters
883 +# ============================================================================
884 +
885 +fp8_test_params = {
886 + "hidden_sizes": [512, 1024, 2048],
887 + "batches": [1, 4, 16],
888 + "num_loras": [1, 4, 8],
889 + "max_ranks": [8, 16, 32, 64],
890 +}
891 +
892 +
893 +# ============================================================================
894 +# FP8 Shrink Tests
895 +# ============================================================================
896 +
897 +
898 +@pytest.mark.parametrize("batches", fp8_test_params["batches"])
899 +@pytest.mark.parametrize("num_loras", fp8_test_params["num_loras"])
900 +@pytest.mark.parametrize("rank", fp8_test_params["max_ranks"])
901 +@pytest.mark.parametrize("hidden_size", fp8_test_params["hidden_sizes"])
902 +@pytest.mark.parametrize("nslices", [1, 2, 3])
903 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
904 +@pytest.mark.parametrize("device", DEVICES)
905 +@pytest.mark.parametrize("seed", SEED)
906 +@pytest.mark.parametrize("quant_mode", ["per_tensor", "per_channel", "blockwise"])
907 +def test_lora_shrink_fp8(
908 + batches: int,
909 + num_loras: int,
910 + rank: int,
911 + hidden_size: int,
912 + nslices: int,
913 + dtype: torch.dtype,
914 + device: str,
915 + seed: int,
916 + quant_mode: str,
917 +):
918 + """Test FP8 shrink kernel with per-tensor, per-channel, and block-wise
919 + quantization, comparing against the bf16 baseline."""
920 + torch.set_default_device(device)
921 + set_random_seed(seed)
922 +
923 + # For blockwise, group sizes must divide evenly or be handled by the kernel
924 + group_k = 128
925 + group_n = 128
926 +
927 + # Adjust group sizes if they're larger than the dimensions
928 + if quant_mode == "blockwise":
929 + group_k = min(group_k, hidden_size)
930 + group_n = min(group_n, rank)
931 +
932 + check_lora_shrink_fp8_kernel(
933 + batches=batches,
934 + num_loras=num_loras,
935 + rank=rank,
936 + hidden_size=hidden_size,
937 + nslices=nslices,
938 + dtype=dtype,
939 + device=device,
940 + seq_length=128,
941 + scaling=0.5,
942 + quant_mode=quant_mode,
943 + group_k=group_k,
944 + group_n=group_n,
945 + )
946 +
947 +
948 +# ============================================================================
949 +# FP8 Expand Tests
950 +# ============================================================================
951 +
952 +
953 +@pytest.mark.parametrize("batches", fp8_test_params["batches"])
954 +@pytest.mark.parametrize("num_loras", fp8_test_params["num_loras"])
955 +@pytest.mark.parametrize("rank", fp8_test_params["max_ranks"])
956 +@pytest.mark.parametrize("hidden_size", fp8_test_params["hidden_sizes"])
957 +@pytest.mark.parametrize("nslices", [1, 2, 3])
958 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
959 +@pytest.mark.parametrize("device", DEVICES)
960 +@pytest.mark.parametrize("seed", SEED)
961 +@pytest.mark.parametrize("quant_mode", ["per_tensor", "per_channel", "blockwise"])
962 +def test_lora_expand_fp8(
963 + batches: int,
964 + num_loras: int,
965 + rank: int,
966 + hidden_size: int,
967 + nslices: int,
968 + dtype: torch.dtype,
969 + device: str,
970 + seed: int,
971 + quant_mode: str,
972 +):
973 + """Test FP8 expand kernel with per-tensor, per-channel, and block-wise
974 + quantization, comparing against the bf16 baseline."""
975 + torch.set_default_device(device)
976 + set_random_seed(seed)
977 +
978 + group_k = 128
979 + group_n = 128
980 +
981 + # Adjust group sizes if they're larger than the dimensions
982 + if quant_mode == "blockwise":
983 + group_k = min(group_k, rank)
984 + group_n = min(group_n, hidden_size)
985 +
986 + check_lora_expand_fp8_kernel(
987 + batches=batches,
988 + num_loras=num_loras,
989 + rank=rank,
990 + hidden_size=hidden_size,
991 + nslices=nslices,
992 + dtype=dtype,
993 + device=device,
994 + seq_length=128,
995 + add_inputs=True,
996 + quant_mode=quant_mode,
997 + group_k=group_k,
998 + group_n=group_n,
999 + )
vllm/lora/ops/triton_ops/__init__.py
+4
@@ -12,13 +12,17 @@ from vllm.lora.ops.triton_ops.fused_moe_lora_op import (
12 fused_moe_lora_expand,
13 fused_moe_lora_shrink,
14 )
15 +from vllm.lora.ops.triton_ops.lora_expand_fp8_op import lora_expand_fp8
16 from vllm.lora.ops.triton_ops.lora_expand_op import lora_expand
17 from vllm.lora.ops.triton_ops.lora_kernel_metadata import LoRAKernelMeta
18 +from vllm.lora.ops.triton_ops.lora_shrink_fp8_op import lora_shrink_fp8
19 from vllm.lora.ops.triton_ops.lora_shrink_op import lora_shrink
20
21 __all__ = [
22 "lora_expand",
23 + "lora_expand_fp8",
24 "lora_shrink",
25 + "lora_shrink_fp8",
26 "LoRAKernelMeta",
27 "fused_moe_lora",
28 "fused_moe_lora_shrink",
vllm/lora/ops/triton_ops/fp8_kernel_utils.py new
+603
@@ -0,0 +1,603 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Utilities for Punica kernel construction.
5 +"""
6 +
7 +from vllm.triton_utils import tl, triton
8 +
9 +
10 +@triton.jit
11 +def _accumulate_mm(
12 + tiled_a,
13 + tiled_b,
14 + accumulator,
15 + a_scale_ptr,
16 + b_scale_ptr,
17 + a_scale_k_stride,
18 + b_scale_k_stride,
19 + iter_k,
20 + group_k: tl.constexpr,
21 + group_n: tl.constexpr,
22 + use_fp8_w8a8: tl.constexpr,
23 +):
24 + """
25 + Core matrix multiplication and accumulation logic with quantization support.
26 +
27 + Args:
28 + tiled_a (tl.tensor): Loaded tile from A matrix
29 + tiled_b (tl.tensor): Loaded tile from B matrix
30 + accumulator (tl.tensor): Current accumulator value
31 + a_scale_ptr (tl.tensor): Scale pointer for A matrix
32 + b_scale_ptr (tl.tensor): Scale pointer for B matrix
33 + a_scale_k_stride (int): K dimension stride for A's block-wise scales
34 + b_scale_k_stride (int): K dimension stride for B's block-wise scales
35 + iter_k (int): Current iteration's global K offset
36 + group_k: Block size for K dimension in block-wise quantization
37 + group_n: Block size for N dimension in block-wise quantization
38 + use_fp8_w8a8: Whether using FP8 W8A8 quantization
39 + """
40 +
41 + if use_fp8_w8a8:
42 + if group_k > 0 and group_n > 0:
43 + # Block-wise quantization: scales are loaded per block
44 + offs_ks = iter_k // group_k
45 + # a_scale_ptr is (BLOCK_M,) tensor of base pointers per row
46 + # Load scale for current K-group, result shape: (BLOCK_M,)
47 + a_scale = tl.load(a_scale_ptr + offs_ks * a_scale_k_stride)
48 + # b_scale_ptr is (BLOCK_N,) tensor with N-offset pre-baked
49 + # Load scale for current K-group, result shape: (BLOCK_N,)
50 + b_scale = tl.load(b_scale_ptr + offs_ks * b_scale_k_stride)
51 + accumulator += (
52 + tl.dot(tiled_a, tiled_b) * a_scale[:, None] * b_scale[None, :]
53 + )
54 + else:
55 + # Tensor-wise or per-channel: accumulate and scale at end
56 + accumulator = tl.dot(tiled_a, tiled_b, acc=accumulator)
57 + else:
58 + accumulator += tl.dot(tiled_a, tiled_b)
59 + return accumulator
60 +
61 +
62 +@triton.jit
63 +def fp8_mm_k(
64 + a_ptr,
65 + b_ptr,
66 + a_scale_ptr,
67 + b_scale_ptr,
68 + ak_stride,
69 + bk_stride,
70 + a_scale_k_stride,
71 + b_scale_k_stride,
72 + offset_k,
73 + K: tl.constexpr,
74 + BLOCK_M: tl.constexpr,
75 + BLOCK_N: tl.constexpr,
76 + BLOCK_K: tl.constexpr,
77 + EVEN_K: tl.constexpr,
78 + SPLIT_K: tl.constexpr,
79 + group_k: tl.constexpr,
80 + group_n: tl.constexpr,
81 + use_fp8_w8a8: tl.constexpr,
82 + per_channel_quant: tl.constexpr,
83 + CAST_TYPE: tl.constexpr,
84 + b_dtype: tl.constexpr,
85 + USE_GDC: tl.constexpr,
86 + base_k,
87 +):
88 + """
89 + FP8-compatible matrix multiplication kernel with quantization support.
90 + Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of
91 + B (k x n), iterate through the K dimension to compute the partial/complete
92 + matrix block product with proper dequantization.
93 +
94 + Args:
95 + a_ptr (tl.tensor): Array of pointers, identifying rows of A
96 + (FP8 or other dtype)
97 + b_ptr (tl.tensor): Array of pointers, identifying columns of B
98 + (FP8 dtype)
99 + a_scale_ptr (tl.tensor): Scale pointer for A matrix
100 + (per-token or block-wise)
101 + b_scale_ptr (tl.tensor): Scale pointer for B matrix
102 + (per-channel or block-wise)
103 + ak_stride (int): K dimension stride of the A matrix
104 + bk_stride (int): K dimension stride of the B matrix
105 + a_scale_k_stride (int): K dimension stride for A's block-wise scales
106 + b_scale_k_stride (int): K dimension stride for B's block-wise scales
107 + offset_k (int): Base offset along K dimension
108 + K: Length of the K dimension
109 + BLOCK_M: M dimension of the output block m x n
110 + BLOCK_N: N dimension of the output block m x n
111 + BLOCK_K: K dimension atom
112 + EVEN_K: True if the blocks of A and B can be loaded without masking
113 + SPLIT_K: Parameter signifying parallelism in the K dimension
114 + group_k: Block size for K dimension in block-wise quantization
115 + group_n: Block size for N dimension in block-wise quantization
116 + use_fp8_w8a8: Whether using FP8 W8A8 quantization
117 + per_channel_quant: Whether using per-channel quantization
118 + CAST_TYPE: if True, cast the values from the A matrix to the B
119 + matrix dtype.
120 + b_dtype: datatype of the B matrix
121 + USE_GDC: Whether to use PDL. True indicates use.
122 + base_k (int): Base offset along K dimension for current SPLIT_K group
123 + """
124 + accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
125 +
126 + # Step size along K for each iteration
127 + STEP_K = BLOCK_K * SPLIT_K
128 +
129 + # Total number of iterations (compile-time constant)
130 + num_iters = tl.cdiv(K, STEP_K)
131 +
132 + for k in range(num_iters):
133 + # Current iteration's global K offset
134 + iter_k = k * STEP_K + base_k
135 + block_end = iter_k + BLOCK_K
136 +
137 + # Skip iterations that are entirely past the K boundary
138 + if not EVEN_K and iter_k >= K:
139 + pass
140 + elif EVEN_K or block_end <= K:
141 + # No masking needed: either K is evenly divisible (EVEN_K)
142 + # or this block fits entirely within K
143 + tiled_b = tl.load(b_ptr)
144 + if USE_GDC:
145 + tl.extra.cuda.gdc_wait()
146 + tiled_a = tl.load(a_ptr)
147 + if CAST_TYPE:
148 + tiled_a = tiled_a.to(b_dtype)
149 +
150 + accumulator = _accumulate_mm(
151 + tiled_a,
152 + tiled_b,
153 + accumulator,
154 + a_scale_ptr,
155 + b_scale_ptr,
156 + a_scale_k_stride,
157 + b_scale_k_stride,
158 + iter_k,
159 + group_k,
160 + group_n,
161 + use_fp8_w8a8,
162 + )
163 + else:
164 + # Partial block at the tail: mask out-of-bounds elements
165 + k_offsets = tl.arange(0, BLOCK_K)
166 + mask = iter_k + k_offsets < K
167 + tiled_b = tl.load(b_ptr, mask=mask[:, None], other=0.0)
168 + if USE_GDC:
169 + tl.extra.cuda.gdc_wait()
170 + tiled_a = tl.load(a_ptr, mask=mask[None, :], other=0.0)
171 + if CAST_TYPE:
172 + tiled_a = tiled_a.to(b_dtype)
173 +
174 + accumulator = _accumulate_mm(
175 + tiled_a,
176 + tiled_b,
177 + accumulator,
178 + a_scale_ptr,
179 + b_scale_ptr,
180 + a_scale_k_stride,
181 + b_scale_k_stride,
182 + iter_k,
183 + group_k,
184 + group_n,
185 + use_fp8_w8a8,
186 + )
187 +
188 + a_ptr += STEP_K * ak_stride
189 + b_ptr += STEP_K * bk_stride
190 +
191 + return accumulator
192 +
193 +
194 +@triton.jit
195 +def do_shrink_kernel_fp8(
196 + pid_n,
197 + pid_sk,
198 + slice_id,
199 + lora_index,
200 + input_ptr,
201 + lora_ptr,
202 + out_ptr,
203 + a_scale_ptr,
204 + b_scale_ptr,
205 + N,
206 + K,
207 + M_LEN,
208 + ram,
209 + # input strides
210 + input_d0_stride,
211 + input_d1_stride,
212 + # lora strides
213 + lora_d0_stride,
214 + lora_d1_stride,
215 + lora_d2_stride,
216 + # scale strides
217 + a_scale_m_stride,
218 + a_scale_k_stride,
219 + b_scale_l_stride,
220 + b_scale_n_stride,
221 + b_scale_k_stride,
222 + # output strides
223 + output_d0_stride,
224 + output_d1_stride,
225 + output_d2_stride,
226 + scaling,
227 + # block size for block-wise quantization
228 + group_n: tl.constexpr,
229 + group_k: tl.constexpr,
230 + BLOCK_M: tl.constexpr,
231 + BLOCK_N: tl.constexpr,
232 + BLOCK_K: tl.constexpr,
233 + EVEN_K: tl.constexpr,
234 + SPLIT_K: tl.constexpr,
235 + SLICE_NUM: tl.constexpr,
236 + USE_GDC: tl.constexpr,
237 + use_fp8_w8a8: tl.constexpr,
238 + per_channel_quant: tl.constexpr,
239 + launch_pdl: tl.constexpr,
240 +):
241 + """
242 + Given an array of integers that identifies the rows of A, ram,
243 + a lora index that identifies which LoRA to use from lora_ptr, lora_index,
244 + a slice_id that identifies the input/output slice, compute the
245 + matrix product and store in the appropriate output location.
246 + """
247 +
248 + # Identify the lora_ptr from slice_id.
249 + if SLICE_NUM == 1:
250 + cur_lora_ptr = lora_ptr
251 + cur_b_scale_ptr = b_scale_ptr
252 + else:
253 + cur_lora_ptr = (
254 + tl.load(lora_ptr + slice_id).to(tl.pointer_type(tl.float8e4nv))
255 + if b_scale_ptr is not None
256 + else tl.load(lora_ptr + slice_id).to(
257 + tl.pointer_type(input_ptr.dtype.element_ty)
258 + )
259 + )
260 + cur_b_scale_ptr = (
261 + tl.load(b_scale_ptr + slice_id).to(tl.pointer_type(tl.float32))
262 + if b_scale_ptr is not None
263 + else b_scale_ptr
264 + )
265 +
266 + # Identify the column indices of B to process.
267 + offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
268 + rbn = tl.max_contiguous(tl.multiple_of(offset_n % N, BLOCK_N), BLOCK_N)
269 +
270 + # Identify A and B block pointers
271 + offset_k = pid_sk * BLOCK_K + tl.arange(0, BLOCK_K)
272 + a_ptr = (
273 + input_ptr + ram[:, None] * input_d0_stride + offset_k[None, :] * input_d1_stride
274 + )
275 + b_ptr = (
276 + cur_lora_ptr
277 + + lora_d0_stride * lora_index
278 + + rbn[None, :] * lora_d1_stride
279 + + offset_k[:, None] * lora_d2_stride
280 + )
281 +
282 + # Load scales for tensor-wise or per-channel quantization (outside the loop)
283 + # Block-wise scales are loaded inside fp8_mm_k
284 + if use_fp8_w8a8:
285 + if group_k > 0 and group_n > 0:
286 + # Block-wise: compute scale pointers for fp8_mm_k
287 + # a_scale: per-row base pointers, shape (BLOCK_M,)
288 + # Each pointer points to the start of that row's scale data
289 + mm_a_scale_ptr = a_scale_ptr + ram * a_scale_m_stride
290 +
291 + # b_scale: pre-compute N-dimension offset
292 + # We need to bake in the N-group offset since fp8_mm_k doesn't know pid_n
293 + n_offset = pid_n * BLOCK_N
294 + offs_ns = (n_offset + tl.arange(0, BLOCK_N)) // group_n
295 + # Base pointer with lora offset + N-group offset baked in, shape (BLOCK_N,)
296 + mm_b_scale_ptr = (
297 + cur_b_scale_ptr
298 + + lora_index * b_scale_l_stride
299 + + offs_ns * b_scale_n_stride
300 + )
301 + elif per_channel_quant:
302 + # Per-channel for weights, per-token for activations
303 + b_scale_ptrs = (
304 + cur_b_scale_ptr + lora_index * b_scale_l_stride + rbn * b_scale_n_stride
305 + )
306 + b_scale = tl.load(b_scale_ptrs)
307 + # Per-token activation scale
308 + a_scale = tl.load(a_scale_ptr + ram * a_scale_m_stride)[:, None]
309 + # For non-block-wise, pass original pointers (not used in mm loop)
310 + mm_a_scale_ptr = a_scale_ptr
311 + mm_b_scale_ptr = cur_b_scale_ptr
312 + else:
313 + # Tensor-wise quantization
314 + a_scale = tl.load(a_scale_ptr) if a_scale_ptr is not None else 1.0
315 + b_scale = tl.load(cur_b_scale_ptr + lora_index * b_scale_l_stride)
316 + # For non-block-wise, pass original pointers (not used in mm loop)
317 + mm_a_scale_ptr = a_scale_ptr
318 + mm_b_scale_ptr = cur_b_scale_ptr
319 + else:
320 + # Non-quantized path
321 + mm_a_scale_ptr = a_scale_ptr
322 + mm_b_scale_ptr = cur_b_scale_ptr
323 +
324 + # Compute partial/complete block matrix product.
325 + accumulator = fp8_mm_k(
326 + a_ptr,
327 + b_ptr,
328 + mm_a_scale_ptr,
329 + mm_b_scale_ptr,
330 + input_d1_stride,
331 + lora_d2_stride,
332 + a_scale_k_stride,
333 + b_scale_k_stride,
334 + offset_k,
335 + K,
336 + BLOCK_M,
337 + BLOCK_N,
338 + BLOCK_K,
339 + EVEN_K,
340 + SPLIT_K,
341 + group_k,
342 + group_n,
343 + use_fp8_w8a8,
344 + per_channel_quant,
345 + False,
346 + cur_lora_ptr.dtype.element_ty,
347 + USE_GDC,
348 + base_k=pid_sk * BLOCK_K,
349 + )
350 + # GDC launch dependents hints the runtime system to launch dependent kernels.
351 + if USE_GDC:
352 + tl.extra.cuda.gdc_launch_dependents()
353 +
354 + # Apply dequantization scales for tensor-wise/per-channel quantization
355 + if use_fp8_w8a8:
356 + if group_k > 0 and group_n > 0:
357 + # Block-wise: already applied in fp8_mm_k
358 + pass
359 + else:
360 + # Tensor-wise or per-channel: apply scales after accumulation
361 + accumulator = accumulator * a_scale * b_scale
362 +
363 + # Apply LoRA scaling factor
364 + accumulator *= scaling
365 +
366 + # Identify the C output pointers to store the results of the accumulator.
367 + offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
368 + offset_cm = tl.arange(0, BLOCK_M)
369 + cur_out_ptr = out_ptr if SLICE_NUM == 1 else out_ptr + slice_id * output_d0_stride
370 + c_ptr = (
371 + cur_out_ptr
372 + + ram[:, None] * output_d1_stride
373 + + offset_cn[None, :] * output_d2_stride
374 + )
375 + c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < N)
376 +
377 + # Cast accumulator to output dtype
378 + accumulator = accumulator.to(out_ptr.dtype.element_ty)
379 +
380 + # handles write-back with reduction-splitting
381 + if SPLIT_K == 1:
382 + tl.store(c_ptr, accumulator, mask=c_mask)
383 + else:
384 + tl.atomic_add(c_ptr, accumulator, mask=c_mask, sem="relaxed")
385 +
386 +
387 +@triton.jit
388 +def do_expand_kernel_fp8(
389 + pid_n,
390 + lora_index,
391 + slice_id,
392 + input_ptr,
393 + lora_ptr,
394 + out_ptr,
395 + a_scale_ptr,
396 + b_scale_ptr,
397 + N,
398 + K,
399 + M_LEN,
400 + ram, # array identifying the rows of Input ptr to operate on
401 + slice_start_loc,
402 + # input ptr strides
403 + input_d0_stride,
404 + input_d1_stride,
405 + input_d2_stride,
406 + # lora ptr strides
407 + ls_d0_ptr,
408 + ls_d1_ptr,
409 + ls_d2_ptr,
410 + # scale strides
411 + a_scale_m_stride,
412 + a_scale_k_stride,
413 + b_scale_l_stride,
414 + b_scale_n_stride,
415 + b_scale_k_stride,
416 + # out ptr strides
417 + output_d0_stride,
418 + output_d1_stride,
419 + # block size for block-wise quantization
420 + group_n: tl.constexpr,
421 + group_k: tl.constexpr,
422 + # constants
423 + BLOCK_M: tl.constexpr,
424 + BLOCK_N: tl.constexpr,
425 + BLOCK_K: tl.constexpr,
426 + SAME_STRIDE: tl.constexpr,
427 + SLICE_NUM: tl.constexpr,
428 + EVEN_K: tl.constexpr,
429 + CAST_TYPE: tl.constexpr,
430 + ADD_INPUTS: tl.constexpr,
431 + USE_GDC: tl.constexpr,
432 + use_fp8_w8a8: tl.constexpr,
433 + per_channel_quant: tl.constexpr,
434 +):
435 + """
436 + FP8-compatible expand kernel for LoRA.
437 + Given an array of integers that identifies the rows of A, ram,
438 + a lora index that identifies which LoRA to use from lora_ptr, lora_index,
439 + a slice_id that identifies the input/output slice,
440 + compute the matrix product with FP8 quantization support and store in
441 + the appropriate output location.
442 +
443 + For expand kernel, the input (shrink output) may be in FP32/FP16/BF16,
444 + while the LoRA B weights can be in FP8.
445 +
446 + Supports:
447 + - FP8 W8A8 quantization for LoRA B weights
448 + - Block-wise quantization with configurable group_k and group_n
449 + - Per-channel quantization
450 + - Tensor-wise quantization
451 + """
452 +
453 + # ls_d*_ptr can be either an integer or a pointer
454 + if SAME_STRIDE:
455 + cur_lora_d0_stride = ls_d0_ptr
456 + cur_lora_d1_stride = ls_d1_ptr
457 + cur_lora_d2_stride = ls_d2_ptr
458 + else:
459 + cur_lora_d0_stride = tl.load(ls_d0_ptr + slice_id)
460 + cur_lora_d1_stride = tl.load(ls_d1_ptr + slice_id)
461 + cur_lora_d2_stride = tl.load(ls_d2_ptr + slice_id)
462 +
463 + # Identify the input_ptr and lora_ptr from slice_id.
464 + if SLICE_NUM == 1:
465 + cur_input_ptr = input_ptr
466 + if use_fp8_w8a8:
467 + cur_lora_ptr = lora_ptr
468 + cur_b_scale_ptr = b_scale_ptr
469 + else:
470 + cur_lora_ptr = lora_ptr
471 + cur_b_scale_ptr = b_scale_ptr # May be None for non-quantized
472 + else:
473 + cur_input_ptr = input_ptr + slice_id * input_d0_stride
474 + if use_fp8_w8a8:
475 + cur_lora_ptr = tl.load(lora_ptr + slice_id).to(
476 + tl.pointer_type(tl.float8e4nv)
477 + )
478 + cur_b_scale_ptr = tl.load(b_scale_ptr + slice_id).to(
479 + tl.pointer_type(tl.float32)
480 + )
481 + else:
482 + cur_lora_ptr = tl.load(lora_ptr + slice_id).to(
483 + tl.pointer_type(out_ptr.dtype.element_ty)
484 + )
485 + cur_b_scale_ptr = (
486 + tl.load(b_scale_ptr + slice_id).to(tl.pointer_type(tl.float32))
487 + if b_scale_ptr is not None
488 + else None
489 + )
490 +
491 + # Identify the column indices of B to process.
492 + offset_n = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
493 + rbn = tl.max_contiguous(tl.multiple_of(offset_n % N, BLOCK_N), BLOCK_N)
494 +
495 + # Identify A and B block pointers
496 + offset_k = tl.arange(0, BLOCK_K)
497 + a_ptr = (
498 + cur_input_ptr
499 + + ram[:, None] * input_d1_stride
500 + + offset_k[None, :] * input_d2_stride
501 + )
502 + b_ptr = (
503 + cur_lora_ptr
504 + + cur_lora_d0_stride * lora_index
505 + + offset_k[:, None] * cur_lora_d2_stride
506 + + rbn[None, :] * cur_lora_d1_stride
507 + )
508 +
509 + # Setup scale pointers for FP8/INT8 quantization
510 + if use_fp8_w8a8:
511 + if group_k > 0 and group_n > 0:
512 + # Block-wise quantization - compute scale pointers for fp8_mm_k
513 + # a_scale: per-row base pointers, shape (BLOCK_M,)
514 + mm_a_scale_ptr = a_scale_ptr + ram * a_scale_m_stride
515 +
516 + # b_scale: pre-compute N-dimension offset since fp8_mm_k doesn't know pid_n
517 + n_offset = pid_n * BLOCK_N
518 + offs_ns = (n_offset + tl.arange(0, BLOCK_N)) // group_n
519 + # Base pointer with lora offset + N-group offset baked in, shape (BLOCK_N,)
520 + mm_b_scale_ptr = (
521 + cur_b_scale_ptr
522 + + lora_index * b_scale_l_stride
523 + + offs_ns * b_scale_n_stride
524 + )
525 + elif per_channel_quant:
526 + # Per-channel for weights, shape (BLOCK_N,)
527 + b_scale_ptrs = (
528 + cur_b_scale_ptr + lora_index * b_scale_l_stride + rbn * b_scale_n_stride
529 + )
530 + b_scale = tl.load(b_scale_ptrs)
531 + # Per-token activation scale, only if a_scale_ptr provided
532 + a_scale = tl.load(a_scale_ptr + ram * a_scale_m_stride)[:, None]
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 + # Tensor-wise quantization
538 + a_scale = tl.load(a_scale_ptr) if a_scale_ptr is not None else 1.0
539 + b_scale = tl.load(cur_b_scale_ptr + lora_index * b_scale_l_stride)
540 + # For non-block-wise, pass original pointers (not used in mm loop)
541 + mm_a_scale_ptr = a_scale_ptr
542 + mm_b_scale_ptr = cur_b_scale_ptr
543 + else:
544 + # Non-quantized path
545 + mm_a_scale_ptr = a_scale_ptr
546 + mm_b_scale_ptr = cur_b_scale_ptr
547 +
548 + # Compute the block matrix product using fp8_mm_k
549 + # Note: For expand kernel, SPLIT_K=1, so we pass 1 for SPLIT_K
550 + accumulator = fp8_mm_k(
551 + a_ptr,
552 + b_ptr,
553 + mm_a_scale_ptr,
554 + mm_b_scale_ptr,
555 + input_d2_stride, # ak_stride
556 + cur_lora_d2_stride, # bk_stride
557 + a_scale_k_stride,
558 + b_scale_k_stride,
559 + offset_k,
560 + K,
561 + BLOCK_M,
562 + BLOCK_N,
563 + BLOCK_K,
564 + EVEN_K,
565 + 1, # SPLIT_K = 1 for expand kernel
566 + group_k,
567 + group_n,
568 + use_fp8_w8a8,
569 + per_channel_quant,
570 + CAST_TYPE, # CAST_TYPE - cast FP8 B to A's dtype
571 + cur_lora_ptr.dtype.element_ty,
572 + USE_GDC,
573 + base_k=0,
574 + )
575 +
576 + # Apply dequantization scales for non-block-wise quantization
577 + if use_fp8_w8a8:
578 + if group_k > 0 and group_n > 0:
579 + pass # Already applied per block in fp8_mm_k
580 + else:
581 + # Tensor-wise or per-channel: apply scales after accumulation
582 + accumulator = accumulator * a_scale * b_scale
583 +
584 + tiled_c = accumulator.to(out_ptr.dtype.element_ty)
585 + if SLICE_NUM == 1:
586 + cur_slice_start = slice_start_loc
587 + else:
588 + cur_slice_start = tl.load(slice_start_loc + slice_id)
589 +
590 + # Identify the C output pointers to store the results of the accumulator.
591 + offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N + cur_slice_start
592 + offset_cm = tl.arange(0, BLOCK_M)
593 + c_ptr = (
594 + out_ptr
595 + + ram[:, None] * output_d0_stride
596 + + offset_cn[None, :] * output_d1_stride
597 + )
598 + c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < (cur_slice_start + N))
599 +
600 + if ADD_INPUTS:
601 + tiled_out = tl.load(c_ptr, mask=c_mask)
602 + tiled_c += tiled_out
603 + tl.store(c_ptr, tiled_c, mask=c_mask)
vllm/lora/ops/triton_ops/lora_expand_fp8_op.py new
+403
@@ -0,0 +1,403 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Based on:
5 +Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
6 +Punica: Multi-Tenant LoRA Serving.
7 +https://arxiv.org/abs/2310.18547
8 +"""
9 +
10 +import torch
11 +
12 +from vllm.lora.ops.triton_ops.fp8_kernel_utils import do_expand_kernel_fp8
13 +from vllm.lora.ops.triton_ops.utils import (
14 + _get_lora_b_ptr,
15 + get_lora_op_configs,
16 +)
17 +from vllm.triton_utils import tl, triton
18 +from vllm.utils.torch_utils import direct_register_custom_op
19 +
20 +_EXPAND_LORA_SCALE_PTR_DICT: dict[tuple[int, ...], torch.tensor] = {}
21 +
22 +
23 +def _get_expand_lora_scale_ptr(lora_weights: list[torch.Tensor], device: torch.device):
24 + """
25 + `_EXPAND_LORA_SCALE_PTR_DICT` collects the required information during
26 + `profile_run`,
27 + After this, it remains constant and subsequent usage is through LUT.
28 + Refer to:
29 + https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py
30 + """
31 + key = tuple(lora_weight.data_ptr() for lora_weight in lora_weights)
32 +
33 + if (ptr_tensor := _EXPAND_LORA_SCALE_PTR_DICT.get(key)) is not None:
34 + return ptr_tensor
35 +
36 + if len(lora_weights) > 1:
37 + tensor_ptrs = []
38 + for lora_weight in lora_weights:
39 + tensor_ptrs.append(lora_weight.data_ptr())
40 + ptr_tensor = torch.tensor(tensor_ptrs, device=device, dtype=torch.uint64)
41 + else:
42 + # Single slice: return the actual tensor so the kernel can use it
43 + # directly without pointer indirection (matches SLICE_NUM == 1 path).
44 + ptr_tensor = lora_weights[0]
45 +
46 + _EXPAND_LORA_SCALE_PTR_DICT[key] = ptr_tensor
47 + return _EXPAND_LORA_SCALE_PTR_DICT.get(key)
48 +
49 +
50 +@triton.jit
51 +def _lora_expand_kernel_fp8(
52 + input_ptr,
53 + lora_ptr,
54 + out_ptr,
55 + a_scale_ptr,
56 + b_scale_ptr,
57 + M,
58 + N,
59 + K,
60 + token_indices_sorted_by_lora_ids,
61 + num_tokens_per_lora,
62 + lora_token_start_loc,
63 + lora_ids,
64 + slice_start_loc,
65 + input_d0_stride,
66 + input_d1_stride,
67 + input_d2_stride,
68 + ls_d0_ptr,
69 + ls_d1_ptr,
70 + ls_d2_ptr,
71 + a_scale_m_stride,
72 + a_scale_k_stride,
73 + b_scale_l_stride,
74 + b_scale_n_stride,
75 + b_scale_k_stride,
76 + output_d0_stride,
77 + output_d1_stride,
78 + output_hs_ptr,
79 + group_n: tl.constexpr,
80 + group_k: tl.constexpr,
81 + BLOCK_M: tl.constexpr,
82 + BLOCK_N: tl.constexpr,
83 + BLOCK_K: tl.constexpr,
84 + EVEN_K: tl.constexpr,
85 + ADD_INPUTS: tl.constexpr,
86 + CAST_TYPE: tl.constexpr,
87 + SLICE_NUM: tl.constexpr,
88 + SAME_STRIDE: tl.constexpr,
89 + USE_GDC: tl.constexpr,
90 + use_fp8_w8a8: tl.constexpr,
91 + per_channel_quant: tl.constexpr,
92 + launch_pdl: tl.constexpr,
93 +):
94 + """
95 + FP8-compatible expand kernel wrapper.
96 + """
97 + cta_n_num = tl.cdiv(N, BLOCK_N)
98 + cta_m_num = tl.cdiv(M, BLOCK_M)
99 +
100 + pid_mn = tl.program_id(axis=0)
101 + pid_m = pid_mn % cta_m_num
102 + pid_n = (pid_mn // cta_m_num) % cta_n_num
103 +
104 + slice_id = tl.program_id(axis=1)
105 + lora_idx = tl.program_id(axis=2)
106 +
107 + lora_id = tl.load(lora_ids + lora_idx)
108 + if lora_id == -1:
109 + return
110 +
111 + lora_m_size = tl.load(num_tokens_per_lora + lora_idx)
112 +
113 + cta_m_offset = pid_m * BLOCK_M
114 + if cta_m_offset >= lora_m_size:
115 + return
116 +
117 + curr_N = N if SAME_STRIDE else tl.load(output_hs_ptr + slice_id)
118 + if pid_n * BLOCK_N >= curr_N:
119 + return
120 +
121 + cta_m_len = min(BLOCK_M, lora_m_size - cta_m_offset)
122 +
123 + lora_m_indices_start = tl.load(lora_token_start_loc + lora_idx)
124 + cta_lora_seq_indices = (
125 + token_indices_sorted_by_lora_ids + lora_m_indices_start + cta_m_offset
126 + )
127 +
128 + offset_m = tl.arange(0, BLOCK_M) % cta_m_len
129 + ram = tl.load(cta_lora_seq_indices + offset_m)
130 +
131 + do_expand_kernel_fp8(
132 + pid_n,
133 + lora_id,
134 + slice_id,
135 + input_ptr,
136 + lora_ptr,
137 + out_ptr,
138 + a_scale_ptr,
139 + b_scale_ptr,
140 + curr_N,
141 + K,
142 + cta_m_len,
143 + ram,
144 + slice_start_loc,
145 + input_d0_stride,
146 + input_d1_stride,
147 + input_d2_stride,
148 + ls_d0_ptr,
149 + ls_d1_ptr,
150 + ls_d2_ptr,
151 + a_scale_m_stride,
152 + a_scale_k_stride,
153 + b_scale_l_stride,
154 + b_scale_n_stride,
155 + b_scale_k_stride,
156 + output_d0_stride,
157 + output_d1_stride,
158 + group_n,
159 + group_k,
160 + BLOCK_M,
161 + BLOCK_N,
162 + BLOCK_K,
163 + SAME_STRIDE,
164 + SLICE_NUM,
165 + EVEN_K,
166 + CAST_TYPE,
167 + ADD_INPUTS,
168 + USE_GDC,
169 + use_fp8_w8a8,
170 + per_channel_quant,
171 + )
172 +
173 +
174 +@torch.inference_mode()
175 +def _lora_expand_fp8(
176 + inputs: torch.Tensor, # shape [num_slices, num_tokens, lora_rank]
177 + lora_b_weights: list[torch.Tensor], # FP8 [num_lora, hidden_size, lora_rank]
178 + output_tensor: torch.Tensor, # shape [num_tokens, hidden_size * num_slices]
179 + token_lora_mapping: torch.Tensor,
180 + token_indices_sorted_by_lora_ids: torch.Tensor,
181 + num_tokens_per_lora: torch.Tensor,
182 + lora_token_start_loc: torch.Tensor,
183 + lora_ids: torch.Tensor,
184 + no_lora_flag_cpu: torch.Tensor, # shape [1]
185 + num_active_loras: int, # number of active LoRAs (unused here, for API compat)
186 + b_scale: list[torch.Tensor], # LoRA B weight scale per slice
187 + a_scale: torch.Tensor | None = None, # Scale for shrink output (optional)
188 + offset_start: int = 0,
189 + add_inputs: bool = False,
190 + group_k: int = 0,
191 + group_n: int = 0,
192 + use_fp8_w8a8: bool = False,
193 + per_channel_quant: bool = False,
194 +) -> None:
195 + """
196 + FP8-compatible LoRA expand operation.
197 +
198 + Args:
199 + inputs: Input tensor from shrink operation [num_slices, num_tokens, lora_rank]
200 + lora_b_weights: List of FP8 LoRA B weights per slice
201 + output_tensor: Output tensor
202 + a_scale: Optional scale for input (if input is quantized)
203 + b_scale: Weight quantization scales per slice
204 + token_lora_mapping: Token to LoRA ID mapping
205 + token_indices_sorted_by_lora_ids: Sorted token indices
206 + num_tokens_per_lora: Number of tokens per LoRA
207 + lora_token_start_loc: Start location for each LoRA's tokens
208 + lora_ids: LoRA IDs to process
209 + no_lora_flag_cpu (torch.Tensor): A CPU tensor of size 1, that indicates
210 + if there are any requests that require LoRA.
211 + offset_start (int, optional): Offset start for output_tensor.
212 + Defaults to 0.
213 + add_inputs (bool, optional): Whether to add the input tensor to the
214 + output tensor. Defaults to False.
215 + group_k (int, optional): Block size for K in block-wise quantization.
216 + group_n (int, optional): Block size for N in block-wise quantization.
217 + use_fp8_w8a8 (bool, optional): Whether to use FP8 W8A8 quantization.
218 + per_channel_quant (bool, optional): Whether to use per-channel quantization.
219 + """
220 + assert no_lora_flag_cpu.numel() == 1
221 + if no_lora_flag_cpu.item():
222 + # None of the inputs require LoRA.
223 + return
224 +
225 + if use_fp8_w8a8:
226 + assert inputs.dtype in [
227 + torch.float8_e4m3fn,
228 + torch.float8_e5m2,
229 + ]
230 + for weight in lora_b_weights:
231 + assert weight.dtype in [
232 + torch.float8_e5m2,
233 + torch.float8_e4m3fn,
234 + ]
235 + else:
236 + assert inputs.dtype in [torch.float16, torch.bfloat16, torch.float32]
237 + for weight in lora_b_weights:
238 + assert weight.dtype in [torch.float16, torch.bfloat16]
239 + assert inputs.size(0) == len(lora_b_weights)
240 + assert output_tensor.is_contiguous()
241 +
242 + # metadata sanity check.
243 + M = inputs.size(1)
244 + assert token_lora_mapping.size(0) == M
245 + assert token_lora_mapping.size(0) == token_indices_sorted_by_lora_ids.size(0)
246 + assert lora_ids.size(0) == num_tokens_per_lora.size(0)
247 + assert lora_token_start_loc.size(0) == lora_ids.size(0) + 1
248 +
249 + (
250 + slice_start_tensor,
251 + lora_ptr_tensor,
252 + lora_strides_d0_tensor,
253 + lora_strides_d1_tensor,
254 + lora_strides_d2_tensor,
255 + hidden_sizes_tensor,
256 + same_stride,
257 + MAX_N,
258 + ) = _get_lora_b_ptr(lora_b_weights, offset_start, inputs.device)
259 +
260 + # Get scale pointers
261 + if b_scale is not None:
262 + b_scale_ptr_tensor = _get_expand_lora_scale_ptr(b_scale, inputs.device)
263 + else:
264 + b_scale_ptr_tensor = None
265 + K = lora_b_weights[0].shape[-1]
266 + ADD_INPUTS = add_inputs
267 + MAX_LORAS = lora_ids.size(0)
268 +
269 + CAST_TYPE = False
270 + NUM_SLICES = len(lora_b_weights)
271 +
272 + # Triton kernel configs.
273 + kernel_config = get_lora_op_configs(
274 + op_type="expand",
275 + max_loras=MAX_LORAS,
276 + batch=M,
277 + hidden_size=MAX_N,
278 + rank=K,
279 + num_slices=NUM_SLICES,
280 + add_inputs=add_inputs,
281 + )
282 + BLOCK_M = kernel_config["block_m"]
283 + BLOCK_N = kernel_config["block_n"]
284 + BLOCK_K = kernel_config["block_k"]
285 + NUM_WARPS = kernel_config["num_warps"]
286 + NUM_CTAS = kernel_config.get("num_ctas", 1)
287 + NUM_STAGES = kernel_config["num_stages"]
288 +
289 + EVEN_K = K % BLOCK_K == 0
290 +
291 + grid = (
292 + triton.cdiv(M, BLOCK_M) * triton.cdiv(MAX_N, BLOCK_N),
293 + NUM_SLICES,
294 + num_active_loras,
295 + )
296 + # We disable PDL temporarily because LoRA kernels are not launching back-to-back,
297 + # making PDL invalid and affecting the kernel performance.
298 + use_gdc = False # supports_pdl(inputs.device)
299 + # Get scale strides
300 + if a_scale is not None:
301 + a_scale_m_stride = a_scale.stride(0) if a_scale.dim() > 1 else 0
302 + a_scale_k_stride = a_scale.stride(-1) if a_scale.dim() > 1 else 0
303 + else:
304 + a_scale_m_stride = 0
305 + a_scale_k_stride = 0
306 +
307 + if b_scale is not None and b_scale[0].dim() > 0:
308 + b_scale_l_stride = b_scale[0].stride(0) if b_scale[0].dim() > 0 else 0
309 + b_scale_n_stride = (
310 + b_scale[0].stride(-2)
311 + if b_scale[0].dim() > 2
312 + else (b_scale[0].stride(-1) if b_scale[0].dim() > 1 else 1)
313 + )
314 + b_scale_k_stride = b_scale[0].stride(-1) if b_scale[0].dim() > 2 else 0
315 + else:
316 + b_scale_l_stride = 1
317 + b_scale_n_stride = 0
318 + b_scale_k_stride = 0
319 +
320 + _lora_expand_kernel_fp8[grid](
321 + inputs,
322 + lora_ptr_tensor,
323 + output_tensor,
324 + a_scale,
325 + b_scale_ptr_tensor,
326 + M,
327 + MAX_N,
328 + K,
329 + token_indices_sorted_by_lora_ids,
330 + num_tokens_per_lora,
331 + lora_token_start_loc,
332 + lora_ids,
333 + slice_start_tensor,
334 + inputs.stride(0),
335 + inputs.stride(1),
336 + inputs.stride(2),
337 + lora_strides_d0_tensor,
338 + lora_strides_d1_tensor,
339 + lora_strides_d2_tensor,
340 + a_scale_m_stride,
341 + a_scale_k_stride,
342 + b_scale_l_stride,
343 + b_scale_n_stride,
344 + b_scale_k_stride,
345 + output_tensor.stride(0),
346 + output_tensor.stride(1),
347 + hidden_sizes_tensor,
348 + group_n,
349 + group_k,
350 + BLOCK_M,
351 + BLOCK_N,
352 + BLOCK_K,
353 + EVEN_K,
354 + ADD_INPUTS,
355 + CAST_TYPE,
356 + NUM_SLICES,
357 + same_stride,
358 + use_gdc,
359 + use_fp8_w8a8=use_fp8_w8a8,
360 + per_channel_quant=per_channel_quant,
361 + num_warps=NUM_WARPS,
362 + num_ctas=NUM_CTAS,
363 + num_stages=NUM_STAGES,
364 + launch_pdl=use_gdc,
365 + )
366 +
367 + return
368 +
369 +
370 +def _lora_expand_fp8_fake(
371 + inputs: torch.Tensor,
372 + lora_b_weights: list[torch.Tensor],
373 + output_tensor: torch.Tensor,
374 + token_lora_mapping: torch.Tensor,
375 + token_indices_sorted_by_lora_ids: torch.Tensor,
376 + num_tokens_per_lora: torch.Tensor,
377 + lora_token_start_loc: torch.Tensor,
378 + lora_ids: torch.Tensor,
379 + no_lora_flag_cpu: torch.Tensor,
380 + num_active_loras: int,
381 + b_scale: list[torch.Tensor],
382 + a_scale: torch.Tensor | None = None,
383 + offset_start: int = 0,
384 + add_inputs: bool = False,
385 + group_k: int = 0,
386 + group_n: int = 0,
387 + use_fp8_w8a8: bool = False,
388 + per_channel_quant: bool = False,
389 +) -> None:
390 + return
391 +
392 +
393 +try:
394 + direct_register_custom_op(
395 + op_name="lora_expand_fp8",
396 + op_func=_lora_expand_fp8,
397 + mutates_args=["output_tensor"],
398 + fake_impl=_lora_expand_fp8_fake,
399 + )
400 + lora_expand_fp8 = torch.ops.vllm.lora_expand_fp8
401 +
402 +except AttributeError:
403 + lora_expand_fp8 = _lora_expand_fp8
vllm/lora/ops/triton_ops/lora_shrink_fp8_op.py new
+429
@@ -0,0 +1,429 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Based on:
5 +Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
6 +Punica: Multi-Tenant LoRA Serving.
7 +https://arxiv.org/abs/2310.18547
8 +"""
9 +
10 +import torch
11 +
12 +from vllm.lora.ops.triton_ops.fp8_kernel_utils import do_shrink_kernel_fp8
13 +from vllm.lora.ops.triton_ops.utils import _get_lora_a_ptr, get_lora_op_configs
14 +from vllm.triton_utils import tl, triton
15 +from vllm.utils.torch_utils import direct_register_custom_op
16 +
17 +_SHRINK_LORA_SCALE_PTR_DICT: dict[tuple[int, ...], tuple] = {}
18 +
19 +
20 +def _get_shrink_lora_scale_ptr(
21 + lora_scale_weights: list[torch.Tensor], device: torch.device
22 +):
23 + """
24 + `_SHRINK_LORA_SCALE_PTR_DICT` collects the required information during
25 + `profile_run`. After this, it remains constant and subsequent usage is
26 + through LUT.
27 +
28 + Returns a tuple of (scale_ptr_tensor, l_stride, n_stride, k_stride).
29 +
30 + Supports scale tensors of varying dimensionality:
31 + - 1D: (lora_num,) — tensor-wise quantization
32 + - 2D: (lora_num, N) — per-channel quantization
33 + - 3D: (lora_num, N, K) — block-wise quantization
34 + - 4D: (lora_num, 1, N, K) — block-wise with extra dim (squeezed to 3D)
35 +
36 + Refer to:
37 + https://github.com/triton-lang/triton/blob/release/3.1.x/python/tutorials/08-grouped-gemm.py
38 + """
39 + key = tuple(lora_weight.data_ptr() for lora_weight in lora_scale_weights)
40 +
41 + if values := _SHRINK_LORA_SCALE_PTR_DICT.get(key):
42 + return values
43 +
44 + tensor_ptrs = []
45 + scale_l_strides = []
46 + scale_n_strides = []
47 + scale_k_strides = []
48 + for lora_scale_weight in lora_scale_weights:
49 + if lora_scale_weight.ndim == 4: # shape:(lora_num,1,size,rank)
50 + assert lora_scale_weight.size(1) == 1
51 + lora_scale_weight = lora_scale_weight.squeeze(dim=1)
52 + assert 1 <= lora_scale_weight.ndim <= 3
53 + assert lora_scale_weight.is_contiguous()
54 + tensor_ptrs.append(lora_scale_weight.data_ptr())
55 + scale_l_strides.append(
56 + lora_scale_weight.stride(0) if lora_scale_weight.ndim > 0 else 0
57 + )
58 + scale_n_strides.append(
59 + lora_scale_weight.stride(-2)
60 + if lora_scale_weight.ndim > 2
61 + else (lora_scale_weight.stride(-1) if lora_scale_weight.ndim > 1 else 1)
62 + )
63 + scale_k_strides.append(
64 + lora_scale_weight.stride(-1) if lora_scale_weight.ndim > 2 else 0
65 + )
66 + if len(lora_scale_weights) > 1:
67 + scale_ptr_tensor = torch.tensor(tensor_ptrs, device=device, dtype=torch.uint64)
68 + else:
69 + scale_ptr_tensor = lora_scale_weights[0]
70 +
71 + if (
72 + len(set(scale_l_strides)) > 1
73 + or len(set(scale_n_strides)) > 1
74 + or len(set(scale_k_strides)) > 1
75 + ):
76 + raise ValueError("All LoRA scale weights must have the same stride.")
77 +
78 + _SHRINK_LORA_SCALE_PTR_DICT[key] = (
79 + scale_ptr_tensor,
80 + scale_l_strides[0],
81 + scale_n_strides[0],
82 + scale_k_strides[0],
83 + )
84 + return _SHRINK_LORA_SCALE_PTR_DICT.get(key)
85 +
86 +
87 +@triton.jit
88 +def _lora_shrink_kernel_fp8(
89 + input_ptr,
90 + lora_ptr,
91 + out_ptr,
92 + a_scale_ptr,
93 + b_scale_ptr,
94 + M,
95 + N,
96 + K,
97 + token_indices_sorted_by_lora_ids,
98 + num_tokens_per_lora,
99 + lora_token_start_loc,
100 + lora_ids,
101 + scaling,
102 + input_d0_stride,
103 + input_d1_stride,
104 + lora_d0_stride,
105 + lora_d1_stride,
106 + lora_d2_stride,
107 + a_scale_m_stride,
108 + a_scale_k_stride,
109 + b_scale_l_stride,
110 + b_scale_n_stride,
111 + b_scale_k_stride,
112 + output_d0_stride,
113 + output_d1_stride,
114 + output_d2_stride,
115 + group_n: tl.constexpr,
116 + group_k: tl.constexpr,
117 + BLOCK_M: tl.constexpr,
118 + BLOCK_N: tl.constexpr,
119 + BLOCK_K: tl.constexpr,
120 + EVEN_K: tl.constexpr,
121 + SPLIT_K: tl.constexpr,
122 + GROUP_SIZE_M: tl.constexpr,
123 + SLICE_NUM: tl.constexpr,
124 + USE_GDC: tl.constexpr, ## should always be false in shrink kernel
125 + use_fp8_w8a8: tl.constexpr,
126 + per_channel_quant: tl.constexpr,
127 + launch_pdl: tl.constexpr,
128 +):
129 + cta_n_num = tl.cdiv(N, BLOCK_N)
130 + cta_m_num = tl.cdiv(M, BLOCK_M)
131 +
132 + pid_sk_m_n = tl.program_id(axis=0)
133 + pid_sk = pid_sk_m_n % SPLIT_K
134 +
135 + pid_m_n = pid_sk_m_n // SPLIT_K
136 + num_pid_in_group = GROUP_SIZE_M * cta_n_num
137 + group_id = pid_m_n // num_pid_in_group
138 + first_pid_m = group_id * GROUP_SIZE_M
139 +
140 + group_size_m = min(cta_m_num - first_pid_m, GROUP_SIZE_M)
141 +
142 + # Column-major ordering within groups for better cache reuse
143 + pid_m = first_pid_m + ((pid_m_n % num_pid_in_group) % group_size_m)
144 + pid_n = (pid_m_n % num_pid_in_group) // group_size_m
145 +
146 + slice_id = tl.program_id(axis=1)
147 + lora_idx = tl.program_id(axis=2)
148 +
149 + lora_id = tl.load(lora_ids + lora_idx)
150 + if lora_id == -1:
151 + # Early exit for the no-lora case.
152 + return
153 +
154 + lora_m_size = tl.load(num_tokens_per_lora + lora_idx)
155 +
156 + cta_m_offset = pid_m * BLOCK_M
157 + if cta_m_offset >= lora_m_size:
158 + # Early exit CTA.
159 + return
160 +
161 + # num rows this CTA should process.
162 + cta_m_len = min(BLOCK_M, lora_m_size - cta_m_offset)
163 +
164 + # Identify all rows that this CTA should process.
165 + lora_m_indices_start = tl.load(lora_token_start_loc + lora_idx)
166 + cta_lora_seq_indices = (
167 + token_indices_sorted_by_lora_ids + lora_m_indices_start + cta_m_offset
168 + )
169 +
170 + # Load all relevant row indices.
171 + offset_m = tl.arange(0, BLOCK_M) % cta_m_len
172 + ram = tl.load(cta_lora_seq_indices + offset_m)
173 +
174 + do_shrink_kernel_fp8(
175 + pid_n,
176 + pid_sk,
177 + slice_id,
178 + lora_id,
179 + input_ptr,
180 + lora_ptr,
181 + out_ptr,
182 + a_scale_ptr,
183 + b_scale_ptr,
184 + N,
185 + K,
186 + cta_m_len,
187 + ram, # array identifying the rows of Input ptr to operate on
188 + # input strides
189 + input_d0_stride,
190 + input_d1_stride,
191 + # lora strides
192 + lora_d0_stride,
193 + lora_d1_stride,
194 + lora_d2_stride,
195 + # scale strides
196 + a_scale_m_stride,
197 + a_scale_k_stride,
198 + b_scale_l_stride,
199 + b_scale_n_stride,
200 + b_scale_k_stride,
201 + # output strides
202 + output_d0_stride,
203 + output_d1_stride,
204 + output_d2_stride,
205 + scaling,
206 + # block size for block-wise quantization
207 + group_n,
208 + group_k,
209 + BLOCK_M,
210 + BLOCK_N,
211 + BLOCK_K,
212 + EVEN_K,
213 + SPLIT_K,
214 + SLICE_NUM,
215 + USE_GDC,
216 + use_fp8_w8a8,
217 + per_channel_quant,
218 + launch_pdl,
219 + )
220 +
221 +
222 +@torch.inference_mode()
223 +def _lora_shrink_fp8(
224 + inputs: torch.Tensor, # shape [num_tokens, hidden_size] - FP8 or FP16/BF16
225 + lora_a_weights: list[
226 + torch.Tensor
227 + ], # shape [num_loras, lora_rank, hidden_size] - FP8 or FP16/BF16
228 + output_tensor: torch.Tensor, # shape [num_slices, num_tokens, lora_rank]
229 + token_lora_mapping: torch.Tensor, # shape [num_tokens]
230 + token_indices_sorted_by_lora_ids: torch.Tensor, # shape [num_tokens]
231 + num_tokens_per_lora: torch.Tensor, # shape [max-loras + 1]
232 + lora_token_start_loc: torch.Tensor, # shape [max-loras + 2]
233 + lora_ids: torch.Tensor, # shape [max-loras + 1]
234 + no_lora_flag_cpu: torch.Tensor, # shape [1]
235 + num_active_loras: int, # number of active LoRAs (unused here, for API compat)
236 + scaling: float,
237 + b_scale: list[torch.Tensor], # LoRA weight scale per slice
238 + a_scale: torch.Tensor | None = None, # Activation scale - per-token or block-wise
239 + group_k: int = 0, # Block size for K in block-wise quantization (0 = tensor-wise)
240 + group_n: int = 0, # Block size for N in block-wise quantization
241 + use_fp8_w8a8: bool = False,
242 + per_channel_quant: bool = False,
243 +) -> None:
244 + """
245 + Args:
246 + inputs: FP8 or FP16/BF16 input tensor [num_tokens, hidden_size]
247 + lora_a_weights: List of FP8 or FP16/BF16 LoRA A weights per slice
248 + output_tensor: Output tensor (FP16/BF16/FP32)
249 + token_lora_mapping: Token to LoRA ID mapping
250 + token_indices_sorted_by_lora_ids: Sorted token indices
251 + num_tokens_per_lora: Number of tokens per LoRA
252 + lora_token_start_loc: Start location for each LoRA's tokens
253 + lora_ids: LoRA IDs to process
254 + scaling: LoRA scaling factor
255 + a_scale: Activation quantization scales
256 + b_scale: Weight quantization scales per slice
257 + group_k: Block size for K dimension quantization
258 + group_n: Block size for N dimension quantization
259 + use_fp8_w8a8: Whether to use FP8 weights and activations
260 + per_channel_quant: Whether to use per-channel quantization
261 + """
262 + assert no_lora_flag_cpu.numel() == 1
263 + if no_lora_flag_cpu.item():
264 + # None of the inputs require LoRA.
265 + return
266 +
267 + assert inputs.size(1) == lora_a_weights[0].size(-1)
268 + assert inputs.is_contiguous()
269 + assert output_tensor.is_contiguous()
270 +
271 + # metadata sanity check
272 + M = inputs.size(0)
273 + assert token_lora_mapping.size(0) == M
274 + assert token_lora_mapping.size(0) == token_indices_sorted_by_lora_ids.size(0)
275 + assert lora_ids.size(0) == num_tokens_per_lora.size(0)
276 + assert lora_token_start_loc.size(0) == lora_ids.size(0) + 1
277 +
278 + output_tensor.zero_()
279 +
280 + # Get LoRA weight pointers
281 + (lora_ptr_tensor, lora_strides_d0, lora_strides_d1, lora_strides_d2) = (
282 + _get_lora_a_ptr(lora_a_weights, inputs.device)
283 + )
284 +
285 + # Get scale pointers if using FP8
286 + if use_fp8_w8a8:
287 + assert a_scale is not None, "a_scale required for FP8 w8a8"
288 + assert b_scale is not None, "b_scale required for FP8"
289 +
290 + b_scale_ptr_tensor, b_scale_l_stride, b_scale_n_stride, b_scale_k_stride = (
291 + _get_shrink_lora_scale_ptr(b_scale, inputs.device)
292 + )
293 + a_scale_ptr = (
294 + a_scale if a_scale is not None else torch.tensor(1.0, device=inputs.device)
295 + )
296 + else:
297 + b_scale_ptr_tensor = torch.tensor(0, device=inputs.device)
298 + b_scale_l_stride = 0
299 + b_scale_n_stride = 0
300 + b_scale_k_stride = 0
301 + a_scale_ptr = torch.tensor(0, device=inputs.device)
302 +
303 + N, K = lora_a_weights[0].shape[-2:] # K=hidden_size, N=rank
304 + NUM_SLICES = len(lora_a_weights)
305 + MAX_LORAS = lora_ids.size(0)
306 +
307 + # Triton kernel configs
308 + kernel_config = get_lora_op_configs(
309 + "shrink",
310 + max_loras=MAX_LORAS,
311 + batch=M,
312 + hidden_size=K,
313 + rank=N,
314 + num_slices=NUM_SLICES,
315 + )
316 + BLOCK_M = kernel_config["block_m"]
317 + BLOCK_N = kernel_config["block_n"]
318 + BLOCK_K = kernel_config["block_k"]
319 + SPLIT_K = kernel_config["split_k"]
320 + NUM_WARPS = kernel_config["num_warps"]
321 + NUM_STAGES = kernel_config["num_stages"]
322 + NUM_CTAS = kernel_config["num_ctas"]
323 + GROUP_SIZE_M = kernel_config.get("group_size_m", 8)
324 + assert BLOCK_K is not None and SPLIT_K is not None
325 + EVEN_K = K % (BLOCK_K * SPLIT_K) == 0
326 +
327 + # Grid configuration with column-major ordering support
328 + grid = (
329 + SPLIT_K * triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),
330 + NUM_SLICES,
331 + num_active_loras,
332 + )
333 +
334 + # Determine scale strides
335 + if use_fp8_w8a8:
336 + if a_scale is not None and a_scale.ndim == 2:
337 + a_scale_m_stride = a_scale.stride(0)
338 + a_scale_k_stride = a_scale.stride(1)
339 + else:
340 + a_scale_m_stride = 0
341 + a_scale_k_stride = 0
342 + else:
343 + a_scale_m_stride = 0
344 + a_scale_k_stride = 0
345 +
346 + # We disable PDL temporarily because LoRA kernels are not launching back-to-back,
347 + # making PDL invalid and affecting the kernel performance.
348 + use_gdc = False # supports_pdl(inputs.device)
349 + _lora_shrink_kernel_fp8[grid](
350 + inputs,
351 + lora_ptr_tensor,
352 + output_tensor,
353 + a_scale_ptr,
354 + b_scale_ptr_tensor,
355 + M,
356 + N,
357 + K,
358 + token_indices_sorted_by_lora_ids,
359 + num_tokens_per_lora,
360 + lora_token_start_loc,
361 + lora_ids,
362 + scaling,
363 + inputs.stride(0),
364 + inputs.stride(1),
365 + lora_strides_d0,
366 + lora_strides_d1,
367 + lora_strides_d2,
368 + a_scale_m_stride,
369 + a_scale_k_stride,
370 + b_scale_l_stride,
371 + b_scale_n_stride,
372 + b_scale_k_stride,
373 + output_tensor.stride(0),
374 + output_tensor.stride(1),
375 + output_tensor.stride(2),
376 + group_n,
377 + group_k,
378 + BLOCK_M,
379 + BLOCK_N,
380 + BLOCK_K,
381 + EVEN_K,
382 + SPLIT_K,
383 + GROUP_SIZE_M,
384 + NUM_SLICES,
385 + use_gdc,
386 + use_fp8_w8a8,
387 + per_channel_quant,
388 + use_gdc,
389 + num_warps=NUM_WARPS,
390 + num_ctas=NUM_CTAS,
391 + num_stages=NUM_STAGES,
392 + )
393 +
394 + return
395 +
396 +
397 +def _lora_shrink_fp8_fake(
398 + inputs: torch.Tensor,
399 + lora_a_weights: list[torch.Tensor],
400 + output_tensor: torch.Tensor,
401 + token_lora_mapping: torch.Tensor,
402 + token_indices_sorted_by_lora_ids: torch.Tensor,
403 + num_tokens_per_lora: torch.Tensor,
404 + lora_token_start_loc: torch.Tensor,
405 + lora_ids: torch.Tensor,
406 + no_lora_flag_cpu: torch.Tensor,
407 + num_active_loras: int,
408 + scaling: float,
409 + b_scale: list[torch.Tensor], # LoRA weight scale per slice
410 + a_scale: torch.Tensor | None = None, # Activation scale - per-token or block-wise
411 + group_k: int = 0, # Block size for K in block-wise quantization (0 = tensor-wise)
412 + group_n: int = 0, # Block size for N in block-wise quantization
413 + use_fp8_w8a8: bool = False,
414 + per_channel_quant: bool = False,
415 +) -> None:
416 + return
417 +
418 +
419 +try:
420 + direct_register_custom_op(
421 + op_name="lora_shrink_fp8",
422 + op_func=_lora_shrink_fp8,
423 + mutates_args=["output_tensor"],
424 + fake_impl=_lora_shrink_fp8_fake,
425 + )
426 + lora_shrink_fp8 = torch.ops.vllm.lora_shrink_fp8
427 +
428 +except AttributeError:
429 + lora_shrink_fp8 = _lora_shrink_fp8
vllm/lora/ops/triton_ops/utils.py
+1 -1
@@ -252,7 +252,7 @@ def get_lora_op_configs(
252 default = {
253 "block_m": 64,
254 "block_n": 64 if num_slices > 1 else 128,
255 - "block_k": 16,
255 + "block_k": 32,
256 "num_warps": 4,
257 "num_ctas": 1,
258 "num_stages": 2,