@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 400d483e87
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 15, 18:30 UTC

[Kernels] LoRA - Retire SGMV and BGMV Kernels (#14685)

Signed-off-by: Varun Sundar Rabindranath <varun@neuralmagic.com> Co-authored-by: Varun Sundar Rabindranath <varun@neuralmagic.com>

Varun Sundar Rabindranath committed Mar 18, 2025 at 05:47 UTC 400d483e87b71315bbb73edb0da9fd629212ca82
15 files changed +247 -2094
benchmarks/kernels/benchmark_lora.py
+63 -373
@@ -17,13 +17,8 @@ from torch.utils.benchmark import Measurement as TMeasurement
17 from utils import ArgPool, Bench, CudaGraphBenchParams
18 from weight_shapes import WEIGHT_SHAPES
19
20 -from vllm.lora.ops.triton_ops.bgmv_expand import bgmv_expand
21 -from vllm.lora.ops.triton_ops.bgmv_expand_slice import bgmv_expand_slice
22 -from vllm.lora.ops.triton_ops.bgmv_shrink import bgmv_shrink
23 -from vllm.lora.ops.triton_ops.sgmv_expand import sgmv_expand
24 -from vllm.lora.ops.triton_ops.sgmv_shrink import sgmv_shrink
20 +from vllm.lora.ops.triton_ops import LoRAKernelMeta, lora_expand, lora_shrink
21 from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
26 -from vllm.lora.ops.triton_ops.v1 import V1KernelMeta, v1_expand, v1_shrink
22 from vllm.utils import FlexibleArgumentParser
23
24 DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())
@@ -167,69 +162,25 @@ class OpType(Enum):
162 """
163 LoRA Ops to benchmark and its properties.
164 """
170 - SGMV_SHRINK = auto()
171 - BGMV_SHRINK = auto()
172 - SGMV_EXPAND = auto()
173 - BGMV_EXPAND = auto()
174 - BGMV_EXPAND_SLICE = auto()
175 - V1_SHRINK = auto()
176 - V1_EXPAND = auto()
165 + LORA_SHRINK = auto()
166 + LORA_EXPAND = auto()
167
168 @staticmethod
169 def from_str(s: str) -> "OpType":
180 - if s.lower() == 'sgmv_shrink':
181 - return OpType.SGMV_SHRINK
182 - if s.lower() == 'sgmv_expand':
183 - return OpType.SGMV_EXPAND
184 - if s.lower() == 'bgmv_shrink':
185 - return OpType.BGMV_SHRINK
186 - if s.lower() == 'bgmv_expand':
187 - return OpType.BGMV_EXPAND
188 - if s.lower() == "bgmv_expand_slice":
189 - return OpType.BGMV_EXPAND_SLICE
190 - if s.lower() == "v1_shrink":
191 - return OpType.V1_SHRINK
192 - if s.lower() == "v1_expand":
193 - return OpType.V1_EXPAND
170 + if s.lower() == "lora_shrink":
171 + return OpType.LORA_SHRINK
172 + if s.lower() == "lora_expand":
173 + return OpType.LORA_EXPAND
174 raise ValueError(f"Unrecognized str {s} to convert to OpType")
175
176 def is_shrink_fn(self) -> bool:
197 - return self in [
198 - OpType.SGMV_SHRINK, OpType.BGMV_SHRINK, OpType.V1_SHRINK
199 - ]
177 + return self in [OpType.LORA_SHRINK]
178
179 def is_expand_fn(self) -> bool:
202 - return self in [
203 - OpType.SGMV_EXPAND, OpType.BGMV_EXPAND, OpType.V1_EXPAND
204 - ]
205 -
206 - def is_prefill_op(self) -> bool:
207 - return self in [
208 - OpType.SGMV_SHRINK, OpType.SGMV_EXPAND, OpType.V1_SHRINK,
209 - OpType.V1_EXPAND
210 - ]
211 -
212 - def is_decode_op(self) -> bool:
213 - return self in [
214 - OpType.BGMV_SHRINK, OpType.BGMV_EXPAND, OpType.BGMV_EXPAND_SLICE,
215 - OpType.V1_SHRINK, OpType.V1_EXPAND
216 - ]
217 -
218 - def is_expand_slice_fn(self) -> bool:
219 - return self in [OpType.BGMV_EXPAND_SLICE]
180 + return self in [OpType.LORA_EXPAND]
181
182 def num_slices(self) -> list[int]:
222 - if self in [
223 - OpType.SGMV_EXPAND, OpType.SGMV_SHRINK, OpType.V1_SHRINK,
224 - OpType.V1_EXPAND
225 - ]:
226 - # SGMV kernels and v1 kernels supports slices
227 - return [1, 2, 3]
228 - if self in [OpType.BGMV_SHRINK, OpType.BGMV_EXPAND]:
229 - return [1]
230 - if self in [OpType.BGMV_EXPAND_SLICE]:
231 - return [2, 3]
232 - raise ValueError(f"Unrecognized OpType {self}")
183 + return [1, 2, 3]
184
185 def mkn(self, batch_size: int, seq_length: int, hidden_size: int,
186 lora_rank: int) -> tuple[int, int, int]:
@@ -239,7 +190,7 @@ class OpType(Enum):
190 k = hidden_size
191 n = lora_rank
192 else:
242 - assert self.is_expand_fn() or self.is_expand_slice_fn()
193 + assert self.is_expand_fn()
194 m = num_tokens
195 k = lora_rank
196 n = hidden_size
@@ -254,7 +205,7 @@ class OpType(Enum):
205 if self.is_shrink_fn():
206 return op_dtype, op_dtype, torch.float32
207 else:
257 - assert self.is_expand_fn() or self.is_expand_slice_fn()
208 + assert self.is_expand_fn()
209 return torch.float32, op_dtype, op_dtype
210
211 def matmul_shapes(
@@ -268,43 +219,19 @@ class OpType(Enum):
219 m, k, n = self.mkn(batch_size, seq_length, hidden_size, lora_rank)
220
221 b_shape = (num_loras, n, k) # col-major
271 - if self in [OpType.SGMV_SHRINK, OpType.V1_SHRINK]:
272 - # SGMV shrink and V1 shrink kernels support num_slices inherently
273 - # in the kernel.
222 + if self in [OpType.LORA_SHRINK]:
223 + # LoRA shrink kernels support num_slices inherently in the kernel.
224 return ((m, k), b_shape, (num_slices, m, n))
275 - if self in [OpType.SGMV_EXPAND, OpType.V1_EXPAND]:
276 - # SGMV expand and V1 expand kernels support num_slices inherently
277 - # in the kernel
225 + if self in [OpType.LORA_EXPAND]:
226 + # LoRA expand kernels support num_slices inherently in the kernel
227 return ((num_slices, m, k), b_shape, (m, n * num_slices))
279 - if self == OpType.BGMV_SHRINK:
280 - return ((m, k), b_shape, (m, n))
281 - if self == OpType.BGMV_EXPAND:
282 - return ((m, k), b_shape, (m, n))
283 - if self == OpType.BGMV_EXPAND_SLICE:
284 - return ((num_slices, m, k), b_shape, (m, n * num_slices))
285 -
228 raise ValueError(f"Unrecognized op_type {self}")
229
230 def bench_fn(self) -> Callable:
289 -
290 - def emulate_bgmv_expand_slice(kwargs_list: list[dict[str, Any]]):
291 - for x in kwargs_list:
292 - bgmv_expand_slice(**x)
293 -
294 - if self == OpType.SGMV_SHRINK:
295 - return sgmv_shrink
296 - if self == OpType.SGMV_EXPAND:
297 - return sgmv_expand
298 - if self == OpType.BGMV_SHRINK:
299 - return bgmv_shrink
300 - if self == OpType.BGMV_EXPAND:
301 - return bgmv_expand
302 - if self == OpType.BGMV_EXPAND_SLICE:
303 - return emulate_bgmv_expand_slice
304 - if self == OpType.V1_SHRINK:
305 - return v1_shrink
306 - if self == OpType.V1_EXPAND:
307 - return v1_expand
231 + if self == OpType.LORA_SHRINK:
232 + return lora_shrink
233 + if self == OpType.LORA_EXPAND:
234 + return lora_expand
235
236 raise ValueError(f"Unrecognized optype {self}")
237
@@ -318,34 +245,13 @@ class OpType(Enum):
245 """
246 w_dtype = lora_weights[0].dtype
247 num_slices = len(lora_weights)
321 - if self in [OpType.SGMV_SHRINK, OpType.V1_SHRINK]:
248 + if self in [OpType.LORA_SHRINK]:
249 for slice_idx in range(num_slices):
250 ref_group_gemm(ref_out=output[slice_idx, :],
251 input=input,
252 lora_weights=lora_weights[slice_idx],
253 **kwargs)
327 - elif self in [OpType.SGMV_EXPAND, OpType.V1_EXPAND]:
328 - hidden_size = lora_weights[0].shape[1]
329 - for slice_idx in range(num_slices):
330 - slice_offset = slice_idx * hidden_size
331 - ref_group_gemm(
332 - ref_out=output[:, slice_offset:slice_offset + hidden_size],
333 - input=input[slice_idx].clone().to(dtype=w_dtype),
334 - lora_weights=lora_weights[slice_idx],
335 - **kwargs)
336 - elif self == OpType.BGMV_SHRINK:
337 - assert num_slices == 1
338 - ref_group_gemm(ref_out=output,
339 - input=input,
340 - lora_weights=lora_weights[0],
341 - **kwargs)
342 - elif self == OpType.BGMV_EXPAND:
343 - assert num_slices == 1
344 - ref_group_gemm(ref_out=output,
345 - input=input.clone().to(dtype=w_dtype),
346 - lora_weights=lora_weights[0],
347 - **kwargs)
348 - elif self == OpType.BGMV_EXPAND_SLICE:
254 + elif self in [OpType.LORA_EXPAND]:
255 hidden_size = lora_weights[0].shape[1]
256 for slice_idx in range(num_slices):
257 slice_offset = slice_idx * hidden_size
@@ -411,13 +317,11 @@ class BenchmarkTensors:
317 input: torch.Tensor
318 lora_weights_lst: list[torch.Tensor]
319 output: torch.Tensor
414 - # metadata tensors
320 + # LoRA kernel metadata
321 + lora_kernel_meta: LoRAKernelMeta
322 + # Metadata tensors used in testing correctness
323 seq_lens: torch.Tensor
416 - seq_start_loc: torch.Tensor
324 prompt_lora_mapping: torch.Tensor
418 - token_lora_mapping: torch.Tensor
419 - # v1 kernel metadata
420 - v1_kernel_meta: Optional[V1KernelMeta] = None
325
326 def io_types(self) -> str:
327 return (f"{dtype_to_str(self.input.dtype)}x"
@@ -444,35 +348,29 @@ class BenchmarkTensors:
348 assert ctx.num_active_loras <= ctx.num_loras
349 total_tokens = ctx.batch_size * ctx.seq_length
350
351 + # Make metadata tensors involved in correctness testing.
352 # Prepare seq lens tensor
353 seq_len_tensor = torch.randint(ctx.seq_length, ctx.seq_length + 1,
354 (ctx.batch_size, ))
450 - # Prepare seq_start_loc tensor
451 - seq_start_loc_tensor = torch.cumsum(torch.tensor(
452 - [0] + seq_len_tensor[:-1].tolist(), dtype=torch.long),
453 - dim=0)
355 assert total_tokens == seq_len_tensor.sum()
356 # Prepare prompt lora indices tensor
357 prompt_lora_indices_tensor = make_prompt_lora_mapping(
358 ctx.batch_size, ctx.num_active_loras, ctx.sort_by_lora_id, "cpu")
458 - # Prepare token lora indices tensor
359 +
360 + # Make LoRAKernelMeta
361 token_lora_indices_tensor = make_token_lora_mapping(
362 total_tokens, ctx.batch_size, prompt_lora_indices_tensor,
363 seq_len_tensor, "cpu")
462 -
463 - v1_kernel_meta = None
464 - if op_type in [OpType.V1_SHRINK, OpType.V1_EXPAND]:
465 - v1_kernel_meta = V1KernelMeta.make(
466 - max_loras=ctx.num_loras,
467 - max_num_tokens=token_lora_indices_tensor.size(0),
468 - device="cpu")
469 - v1_kernel_meta.prepare_tensors(
470 - token_lora_mapping=token_lora_indices_tensor)
364 + lora_kernel_meta = LoRAKernelMeta.make(
365 + max_loras=ctx.num_loras,
366 + max_num_tokens=token_lora_indices_tensor.size(0),
367 + device="cpu")
368 + lora_kernel_meta.prepare_tensors(
369 + token_lora_mapping=token_lora_indices_tensor)
370
371 return BenchmarkTensors(input_tensor, lora_weights, output_tensor,
473 - seq_len_tensor, seq_start_loc_tensor,
474 - prompt_lora_indices_tensor,
475 - token_lora_indices_tensor, v1_kernel_meta)
372 + lora_kernel_meta, seq_len_tensor,
373 + prompt_lora_indices_tensor)
374
375 def sanity_check(self) -> None:
376 """
@@ -482,9 +380,9 @@ class BenchmarkTensors:
380 # check metadata tensors
381 assert torch.sum(self.seq_lens) == num_tokens
382 num_seqs = self.seq_lens.shape[0]
485 - assert self.seq_start_loc.shape[0] == num_seqs
383 + #assert self.seq_start_loc.shape[0] == num_seqs
384 assert self.prompt_lora_mapping.shape[0] == num_seqs
487 - assert self.token_lora_mapping.shape[0] == num_tokens
385 + assert self.lora_kernel_meta.token_lora_mapping.shape[0] == num_tokens
386
387 def to_device(self, device: str):
388 """
@@ -499,220 +397,27 @@ class BenchmarkTensors:
397 self.input = to_device(self.input)
398 self.output = to_device(self.output)
399 self.seq_lens = to_device(self.seq_lens)
502 - self.seq_start_loc = to_device(self.seq_start_loc)
400 self.prompt_lora_mapping = to_device(self.prompt_lora_mapping)
504 - self.token_lora_mapping = to_device(self.token_lora_mapping)
401 for i in range(len(self.lora_weights_lst)):
402 self.lora_weights_lst[i] = to_device(self.lora_weights_lst[i])
403
508 - # v1 meta
509 - if self.v1_kernel_meta:
510 - for field_name in V1KernelMeta.__dataclass_fields__:
511 - field = getattr(self.v1_kernel_meta, field_name)
512 - assert isinstance(field, torch.Tensor)
513 - setattr(self.v1_kernel_meta, field_name, to_device(field))
404 + # LoRA meta
405 + for field_name in LoRAKernelMeta.__dataclass_fields__:
406 + field = getattr(self.lora_kernel_meta, field_name)
407 + assert isinstance(field, torch.Tensor)
408 + setattr(self.lora_kernel_meta, field_name, to_device(field))
409
410 def metadata(self) -> tuple[int, int, int]:
411 """
412 Return num_seqs, num_tokens and max_seq_len
413 """
414 num_seqs = self.seq_lens.shape[0]
520 - num_tokens = self.token_lora_mapping.shape[0]
415 + num_tokens = self.lora_kernel_meta.token_lora_mapping.shape[0]
416 max_seq_len = torch.max(self.seq_lens).item()
417 num_slices = len(self.lora_weights_lst)
418 return num_seqs, num_tokens, max_seq_len, num_slices
419
525 - def convert_to_sgmv_benchmark_tensors(self):
526 - """
527 - For sgmv punica kernels, when consecutive sequences have the
528 - same LoRA ID, we just merge them together.
529 - This happens in punica.py::compute_metadata
530 - """
531 -
532 - # Collapse seq_lens and seq_start_loc
533 - _, seq_lens = torch.unique_consecutive(self.token_lora_mapping,
534 - return_counts=True)
535 - cum_result = torch.cumsum(seq_lens, dim=0)
536 - seq_start_loc = torch.zeros_like(seq_lens)
537 - seq_start_loc[1:].copy_(cum_result[:-1])
538 -
539 - # Collapse prompt mapping
540 - prompt_lora_mapping = torch.unique_consecutive(
541 - self.prompt_lora_mapping)
542 -
543 - assert torch.sum(seq_lens) == torch.sum(self.seq_lens), \
544 - f"dont match - new {torch.sum(seq_lens)} vs {torch.sum(self.seq_lens)}"
545 -
546 - self.prompt_lora_mapping = prompt_lora_mapping.to(
547 - dtype=self.prompt_lora_mapping.dtype)
548 - self.seq_lens = seq_lens.to(dtype=self.seq_lens.dtype)
549 - self.seq_start_loc = seq_start_loc.to(dtype=self.seq_start_loc.dtype)
550 -
551 - def as_sgmv_shrink_kwargs(self) -> dict[str, Any]:
552 - self.convert_to_sgmv_benchmark_tensors()
553 - self.sanity_check()
554 - self.to_device(self.input.device)
555 -
556 - num_seqs, num_tokens, max_seq_len, num_slices = self.metadata()
557 -
558 - # Sanity check matrix shapes.
559 - i_shape, lw_shape, o_shape = self.input.shape, self.lora_weights_lst[
560 - 0].shape, self.output.shape
561 - # Expected input shape [num_tokens, hidden_size]
562 - assert len(i_shape) == 2
563 - assert i_shape[0] == num_tokens
564 - hidden_size = i_shape[1]
565 - # Expected lora weight shape [num_loras, lora_rank, hidden_size]
566 - assert len(lw_shape) == 3
567 - assert lw_shape[2] == hidden_size
568 - lora_rank = lw_shape[1]
569 - # Expected output shape [num_slices, num_tokens, lora_rank]
570 - assert len(o_shape) == 3
571 - assert o_shape == (num_slices, num_tokens, lora_rank)
572 -
573 - return {
574 - 'inputs': self.input,
575 - 'lora_a_weights': self.lora_weights_lst,
576 - 'output_tensor': self.output,
577 - 'b_seq_start_loc': self.seq_start_loc,
578 - 'seq_len_tensor': self.seq_lens,
579 - 'lora_indices_tensor': self.prompt_lora_mapping,
580 - 'batches': num_seqs,
581 - 'max_seq_length': max_seq_len,
582 - 'token_nums': num_tokens,
583 - 'scaling': 1.0,
584 - }
585 -
586 - def as_sgmv_expand_kwargs(self, add_inputs: bool) -> dict[str, Any]:
587 -
588 - self.convert_to_sgmv_benchmark_tensors()
589 - self.sanity_check()
590 - self.to_device(self.input.device)
591 -
592 - num_seqs, num_tokens, max_seq_len, num_slices = self.metadata()
593 -
594 - # Sanity check matrix shapes.
595 - i_shape, lw_shape, o_shape = self.input.shape, self.lora_weights_lst[
596 - 0].shape, self.output.shape
597 - # Expected input shape : [num_slices, num_tokens, lora_rank]
598 - assert len(i_shape) == 3
599 - assert i_shape[0] == num_slices
600 - assert i_shape[1] == num_tokens
601 - lora_rank = i_shape[2]
602 - # Expected lora weight shape : [num_lora, hidden_size, lora_rank]
603 - assert len(lw_shape) == 3
604 - assert lw_shape[2] == lora_rank
605 - hidden_size = lw_shape[1]
606 - # Expected output shape : [num_tokens, hidden_size * num_slices]
607 - assert len(o_shape) == 2
608 - assert o_shape == (num_tokens, hidden_size * num_slices)
609 -
610 - return {
611 - 'inputs': self.input,
612 - 'lora_b_weights': self.lora_weights_lst,
613 - 'output_tensor': self.output,
614 - 'b_seq_start_loc': self.seq_start_loc,
615 - 'seq_len_tensor': self.seq_lens,
616 - 'lora_indices_tensor': self.prompt_lora_mapping,
617 - 'batches': num_seqs,
618 - 'max_seq_length': max_seq_len,
619 - 'token_nums': num_tokens,
620 - 'offset_start': 0,
621 - 'add_inputs': add_inputs,
622 - }
623 -
624 - def as_bgmv_shrink_kwargs(self) -> dict[str, Any]:
625 - assert len(self.lora_weights_lst) == 1
626 - self.to_device(self.input.device)
627 -
628 - _, num_tokens, _, _ = self.metadata()
629 - # Sanity check shapes
630 - i_shape, lw_shape, o_shape = self.input.shape, self.lora_weights_lst[
631 - 0].shape, self.output.shape
632 - # Expected input shape [num_tokens, hidden_size]
633 - assert len(i_shape) == 2
634 - assert i_shape[0] == num_tokens
635 - hidden_size = i_shape[1]
636 - # Expected lora weight shape [num_loras, lora_rank, hidden_size]
637 - assert len(lw_shape) == 3
638 - assert lw_shape[2] == hidden_size
639 - lora_rank = lw_shape[1]
640 - # Expected output shape [num_tokens, lora_rank]
641 - assert len(o_shape) == 2
642 - assert o_shape == (num_tokens, lora_rank)
643 -
644 - return {
645 - 'inputs': self.input,
646 - 'lora_a_weights': self.lora_weights_lst[0],
647 - 'output_tensor': self.output,
648 - 'lora_indices_tensor': self.token_lora_mapping,
649 - 'scaling': 1.0
650 - }
651 -
652 - def as_bgmv_expand_kwargs(self, add_inputs: bool):
653 - assert len(self.lora_weights_lst) == 1
654 - self.to_device(self.input.device)
655 -
656 - _, num_tokens, _, _ = self.metadata()
657 - # Sanity check shapes
658 - i_shape, lw_shape, o_shape = self.input.shape, self.lora_weights_lst[
659 - 0].shape, self.output.shape
660 - # Expected input shape [num_tokens, lora_rank]
661 - assert len(i_shape) == 2
662 - assert i_shape[0] == num_tokens
663 - lora_rank = i_shape[1]
664 - # Expected lora weight shape [num_loras, hidden_size, lora_rank]
665 - assert len(lw_shape) == 3
666 - assert lw_shape[2] == lora_rank
667 - hidden_size = lw_shape[1]
668 - # Expected output shape [num_tokens, hidden_size]
669 - assert len(o_shape) == 2
670 - assert o_shape == (num_tokens, hidden_size)
671 -
672 - return {
673 - 'inputs': self.input,
674 - 'lora_b_weights': self.lora_weights_lst[0],
675 - 'output_tensor': self.output,
676 - 'lora_indices_tensor': self.token_lora_mapping,
677 - 'add_inputs': add_inputs
678 - }
679 -
680 - def as_bgmv_expand_slice_kwargs(self, add_inputs: bool) -> dict[str, Any]:
681 -
682 - _, num_tokens, _, num_slices = self.metadata()
683 - # Sanity check shapes
684 - i_shape, lw_shape, o_shape = self.input.shape, self.lora_weights_lst[
685 - 0].shape, self.output.shape
686 - # Expected input shape [num_slices, num_tokens, lora_rank]
687 - assert len(i_shape) == 3
688 - assert i_shape[0] == num_slices
689 - assert i_shape[1] == num_tokens
690 - lora_rank = i_shape[2]
691 - # Expected lora weight shape [num_loras, hidden_size, lora_rank]
692 - assert len(lw_shape) == 3
693 - assert lw_shape[2] == lora_rank
694 - hidden_size = lw_shape[1]
695 - # Expected output shape [num_tokens, hidden_size * num_slices]
696 - assert len(o_shape) == 2
697 - assert o_shape == (num_tokens, hidden_size * num_slices)
698 -
699 - self.to_device(self.input.device)
700 -
701 - kwargs_list = []
702 - for i in range(num_slices):
703 - kwargs_list.append({
704 - 'inputs': self.input[i],
705 - 'lora_b_weights': self.lora_weights_lst[i],
706 - 'output_tensor': self.output,
707 - 'lora_indices_tensor': self.token_lora_mapping,
708 - 'slice_offset': i * hidden_size,
709 - 'slice_size': hidden_size,
710 - 'add_inputs': add_inputs,
711 - })
712 - return {'kwargs_list': kwargs_list}
713 -
714 - def as_v1_shrink_kwargs(self) -> dict[str, Any]:
715 - assert self.v1_kernel_meta is not None
420 + def as_lora_shrink_kwargs(self) -> dict[str, Any]:
421 self.sanity_check()
422 self.to_device(self.input.device)
423
@@ -737,17 +442,16 @@ class BenchmarkTensors:
442 'inputs': self.input,
443 'lora_a_weights': self.lora_weights_lst,
444 'output_tensor': self.output,
740 - 'token_lora_mapping': self.v1_kernel_meta.token_lora_mapping,
445 + 'token_lora_mapping': self.lora_kernel_meta.token_lora_mapping,
446 'token_indices_sorted_by_lora_ids':
742 - self.v1_kernel_meta.token_indices_sorted_by_lora_ids,
743 - 'num_tokens_per_lora': self.v1_kernel_meta.num_tokens_per_lora,
744 - 'lora_token_start_loc': self.v1_kernel_meta.lora_token_start_loc,
745 - 'lora_ids': self.v1_kernel_meta.active_lora_ids,
447 + self.lora_kernel_meta.token_indices_sorted_by_lora_ids,
448 + 'num_tokens_per_lora': self.lora_kernel_meta.num_tokens_per_lora,
449 + 'lora_token_start_loc': self.lora_kernel_meta.lora_token_start_loc,
450 + 'lora_ids': self.lora_kernel_meta.active_lora_ids,
451 'scaling': 1.0,
452 }
453
749 - def as_v1_expand_kwargs(self, add_inputs: bool) -> dict[str, Any]:
750 - assert self.v1_kernel_meta is not None
454 + def as_lora_expand_kwargs(self, add_inputs: bool) -> dict[str, Any]:
455 self.sanity_check()
456 self.to_device(self.input.device)
457
@@ -773,12 +477,12 @@ class BenchmarkTensors:
477 'inputs': self.input,
478 'lora_b_weights': self.lora_weights_lst,
479 'output_tensor': self.output,
776 - 'token_lora_mapping': self.v1_kernel_meta.token_lora_mapping,
480 + 'token_lora_mapping': self.lora_kernel_meta.token_lora_mapping,
481 'token_indices_sorted_by_lora_ids':
778 - self.v1_kernel_meta.token_indices_sorted_by_lora_ids,
779 - 'num_tokens_per_lora': self.v1_kernel_meta.num_tokens_per_lora,
780 - 'lora_token_start_loc': self.v1_kernel_meta.lora_token_start_loc,
781 - 'lora_ids': self.v1_kernel_meta.active_lora_ids,
482 + self.lora_kernel_meta.token_indices_sorted_by_lora_ids,
483 + 'num_tokens_per_lora': self.lora_kernel_meta.num_tokens_per_lora,
484 + 'lora_token_start_loc': self.lora_kernel_meta.lora_token_start_loc,
485 + 'lora_ids': self.lora_kernel_meta.active_lora_ids,
486 'offset_start': 0,
487 'add_inputs': add_inputs,
488 }
@@ -791,20 +495,10 @@ class BenchmarkTensors:
495 else:
496 assert add_inputs is not None
497
794 - if op_type == OpType.SGMV_SHRINK:
795 - return self.as_sgmv_shrink_kwargs()
796 - if op_type == OpType.SGMV_EXPAND:
797 - return self.as_sgmv_expand_kwargs(add_inputs)
798 - if op_type == OpType.BGMV_SHRINK:
799 - return self.as_bgmv_shrink_kwargs()
800 - if op_type == OpType.BGMV_EXPAND:
801 - return self.as_bgmv_expand_kwargs(add_inputs)
802 - if op_type == OpType.BGMV_EXPAND_SLICE:
803 - return self.as_bgmv_expand_slice_kwargs(add_inputs)
804 - if op_type == OpType.V1_SHRINK:
805 - return self.as_v1_shrink_kwargs()
806 - if op_type == OpType.V1_EXPAND:
807 - return self.as_v1_expand_kwargs(add_inputs)
498 + if op_type == OpType.LORA_SHRINK:
499 + return self.as_lora_shrink_kwargs()
500 + if op_type == OpType.LORA_EXPAND:
501 + return self.as_lora_expand_kwargs(add_inputs)
502 raise ValueError(f"Unrecognized optype {self}")
503
504 def test_correctness(self, op_type: OpType,
@@ -993,10 +687,6 @@ def run(args: argparse.Namespace, bench_ctxs: list[BenchmarkContext]):
687 for bench_ctx in bench_ctxs:
688 for seq_len in args.seq_lengths:
689 bench_ops: list[OpType] = args.op_types
996 - if seq_len > 1:
997 - # bench only prefill ops
998 - bench_ops = [op for op in args.op_types if op.is_prefill_op()]
999 -
690 seq_len_timers = []
691 for bench_op in bench_ops:
692 for num_slices in bench_op.num_slices():
@@ -1206,13 +896,13 @@ Benchmark LoRA kernels:
896 {use_cuda_graph_recommendation()}
897
898 list_bench example:
1209 - python3 benchmarks/kernels/benchmark_lora.py list_bench --arg-pool-size 32 --batch-sizes 1 16 32 --dtype torch.float16 --hidden-sizes 2048 --lora-ranks 16 --num-loras 1 4 --op-types bgmv_shrink bgmv_expand sgmv_shrink sgmv_expand bgmv_expand_slice --seq-lengths 1 16 --sort-by-lora-id 1 --cuda-graph-nops 32
899 + python3 benchmarks/kernels/benchmark_lora.py list_bench --arg-pool-size 32 --batch-sizes 1 16 32 --dtype torch.float16 --hidden-sizes 2048 --lora-ranks 16 --num-loras 1 4 --op-types lora_shrink lora_expand --seq-lengths 1 16 --sort-by-lora-id 1 --cuda-graph-nops 32
900
901 model_bench example:
1212 - python3 benchmarks/kernels/benchmark_lora.py model_bench --models meta-llama/Llama-3-8b --arg-pool-size 32 --batch-sizes 1 16 32 --dtype torch.float16 --lora-ranks 16 --num-loras 1 4 --op-types bgmv_shrink bgmv_expand sgmv_shrink sgmv_expand bgmv_expand_slice --seq-lengths 1 16 --sort-by-lora-id 1 --cuda-graph-nops 32
902 + python3 benchmarks/kernels/benchmark_lora.py model_bench --models meta-llama/Llama-3-8b --arg-pool-size 32 --batch-sizes 1 16 32 --dtype torch.float16 --lora-ranks 16 --num-loras 1 4 --op-types lora_shrink lora_expand --seq-lengths 1 16 --sort-by-lora-id 1 --cuda-graph-nops 32
903
904 range_bench example:
1215 - python3 benchmarks/kernels/benchmark_lora.py range_bench --arg-pool-size 32 --batch-sizes 1 16 32 --dtype torch.float16 --num-loras 1 4 --op-types bgmv_shrink bgmv_expand sgmv_shrink sgmv_expand bgmv_expand_slice --seq-lengths 1 16 --sort-by-lora-id 1 --cuda-graph-nops 32 --hidden-sizes-start 1024 --hidden-sizes-end 4096 --hidden-sizes-increment 1024 --lora-ranks-start 8 --lora-ranks-end 24 --lora-ranks-increment 8
905 + python3 benchmarks/kernels/benchmark_lora.py range_bench --arg-pool-size 32 --batch-sizes 1 16 32 --dtype torch.float16 --num-loras 1 4 --op-types lora_shrink lora_expand --seq-lengths 1 16 --sort-by-lora-id 1 --cuda-graph-nops 32 --hidden-sizes-start 1024 --hidden-sizes-end 4096 --hidden-sizes-increment 1024 --lora-ranks-start 8 --lora-ranks-end 24 --lora-ranks-increment 8
906 """, # noqa: E501
907 formatter_class=argparse.RawTextHelpFormatter)
908
tests/lora/test_punica_ops.py
+84 -359
@@ -4,18 +4,13 @@ from threading import Lock
4 import pytest
5 import torch
6
7 -import vllm.lora.ops.triton_ops # noqa: F401
8 -import vllm.lora.ops.triton_ops.v1 # noqa: F401
9 -from vllm.lora.ops.torch_ops import (bgmv_expand, bgmv_expand_slice,
10 - bgmv_shrink, sgmv_expand,
11 - sgmv_expand_slice, sgmv_shrink)
7 +import vllm.lora.ops.torch_ops as torch_ops
8 +import vllm.lora.ops.triton_ops as triton_ops
9 +from vllm.lora.ops.triton_ops import LoRAKernelMeta
10 from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
13 -from vllm.lora.ops.triton_ops.v1 import V1KernelMeta
11 from vllm.platforms import current_platform
12
16 -from .utils import (PunicaTensors, assert_close, generate_data,
17 - generate_data_for_expand_nslices,
18 - generate_data_for_nslices)
13 +from .utils import PunicaTensors, assert_close, generate_data_for_nslices
14
15
16 # Utility shrink and expand operations used as reference implementations.
@@ -26,10 +21,10 @@ def sgmv_shrink_for_nslices(
21 prompt_lora_mapping: torch.Tensor, batches: int, max_seq_length: int,
22 num_tokens: int, scaling: float):
23 """
29 - Wrapper around sgmv_shrink that handles any nslices.
24 + Wrapper around torch_ops.sgmv_shrink that handles any nslices.
25 """
26 for index in range(nslices):
32 - sgmv_shrink(
27 + torch_ops.sgmv_shrink(
28 inputs_tensor,
29 lora_weights_lst[index],
30 out_tensor[index],
@@ -53,11 +48,11 @@ def sgmv_expand_for_nslices(nslices: int, hidden_size: int,
48 max_seq_length: int, num_tokens: int,
49 add_inputs: bool) -> None:
50 """
56 - Wrapper around sgmv_expand that handles any nslices.
51 + Wrapper around torch_ops.sgmv_expand that handles any nslices.
52 """
53 if nslices == 1:
54 # Verify the torch's sgmv_expand op
60 - sgmv_expand(
55 + torch_ops.sgmv_expand(
56 inputs_tensor[0],
57 lora_weights_lst[0],
58 out_tensor,
@@ -73,7 +68,7 @@ def sgmv_expand_for_nslices(nslices: int, hidden_size: int,
68 slice_offset = 0
69 for index in range(nslices):
70 lora_weights = lora_weights_lst[index]
76 - sgmv_expand_slice(
71 + torch_ops.sgmv_expand_slice(
72 inputs_tensor[index],
73 lora_weights,
74 out_tensor,
@@ -93,12 +88,13 @@ def sgmv_expand_for_nslices(nslices: int, hidden_size: int,
88 _dict_lock = Lock()
89
90
96 -def check_shrink_kernels(batches: int, num_loras: int, rank: int,
97 - hidden_size: int, nslices: int, dtype: torch.dtype,
98 - device: str, seq_length: int, scaling: float):
91 +def check_lora_shrink_kernel(batches: int, num_loras: int, rank: int,
92 + hidden_size: int, nslices: int,
93 + dtype: torch.dtype, device: str, seq_length: int,
94 + scaling: float):
95 """
100 - Compare outputs of vllm.sgmv_shrink and vllm.v1_shrink kernel against a
101 - reference implementation.
96 + Compare outputs of torch_ops.sgmv_shrink and triton_ops.lora_shrink
97 + kernels.
98 """
99 data: PunicaTensors = generate_data_for_nslices(
100 batches,
@@ -118,35 +114,24 @@ def check_shrink_kernels(batches: int, num_loras: int, rank: int,
114 data.prompt_lora_mapping, batches, max_seq_length,
115 token_nums)
116
121 - # Setup metadata information for the V1 kernel.
122 - v1_meta = V1KernelMeta.make(max_loras=num_loras,
123 - max_num_tokens=token_nums,
124 - device='cuda')
125 - v1_meta.prepare_tensors(data.token_lora_mapping)
117 + # Setup metadata information for the LoRA kernel.
118 + lora_meta = LoRAKernelMeta.make(max_loras=num_loras,
119 + max_num_tokens=token_nums,
120 + device='cuda')
121 + lora_meta.prepare_tensors(data.token_lora_mapping)
122
123 ref_out_tensor = data.ref_out_tensor
128 - sgmv_out_tensor = data.our_out_tensor
129 - v1_out_tensor = data.our_out_tensor.clone()
124 + out_tensor = data.our_out_tensor.clone()
125
126 # Preventing cache error pointer.
127 with _dict_lock:
133 - # SGMV shrink kernel
128 + # lora_shrink kernel
129 _LORA_A_PTR_DICT.clear()
135 - torch.ops.vllm.sgmv_shrink(
130 + triton_ops.lora_shrink(
131 data.inputs_tensor,
132 data.lora_weights,
138 - sgmv_out_tensor,
139 - *sgmv_meta_args,
140 - scaling,
141 - )
142 -
143 - # V1 shrink kernel
144 - _LORA_A_PTR_DICT.clear()
145 - torch.ops.vllm.v1_shrink(
146 - data.inputs_tensor,
147 - data.lora_weights,
148 - v1_out_tensor,
149 - *v1_meta.meta_args(token_nums=token_nums),
133 + out_tensor,
134 + *lora_meta.meta_args(token_nums=token_nums),
135 scaling,
136 )
137
@@ -160,16 +145,16 @@ def check_shrink_kernels(batches: int, num_loras: int, rank: int,
145 scaling,
146 )
147
163 - assert_close(sgmv_out_tensor, ref_out_tensor)
164 - assert_close(v1_out_tensor, ref_out_tensor)
148 + assert_close(out_tensor, ref_out_tensor)
149
150
167 -def check_expand_kernels(batches: int, num_loras: int, rank: int,
168 - hidden_size: int, nslices: int, dtype: torch.dtype,
169 - device: str, seq_length: int, add_inputs: bool):
151 +def check_lora_expand_kernel(batches: int, num_loras: int, rank: int,
152 + hidden_size: int, nslices: int,
153 + dtype: torch.dtype, device: str, seq_length: int,
154 + add_inputs: bool):
155 """
171 - Compare outputs of vllm.sgmv_expand and vllm.v1_expand kernels against a
172 - reference implementation.
156 + Compare outputs of torch_ops.sgmv_expand and triton_ops.lora_expand
157 + kernels.
158 """
159 data: PunicaTensors = generate_data_for_nslices(
160 batches,
@@ -190,37 +175,25 @@ def check_expand_kernels(batches: int, num_loras: int, rank: int,
175 data.prompt_lora_mapping, batches, max_seq_length,
176 token_nums)
177
193 - # Setup metadata information for the V1 kernel.
194 - v1_meta = V1KernelMeta.make(max_loras=num_loras,
195 - max_num_tokens=token_nums,
196 - device='cuda')
197 - v1_meta.prepare_tensors(data.token_lora_mapping)
178 + # Setup metadata information for the LoRA kernel.
179 + lora_meta = LoRAKernelMeta.make(max_loras=num_loras,
180 + max_num_tokens=token_nums,
181 + device='cuda')
182 + lora_meta.prepare_tensors(data.token_lora_mapping)
183
184 # Setup output tensors
185 ref_out_tensor = data.ref_out_tensor
201 - sgmv_out_tensor = data.our_out_tensor
202 - v1_out_tensor = data.our_out_tensor.clone()
186 + out_tensor = data.our_out_tensor.clone()
187
188 with _dict_lock:
205 - # SGMV expand kernel
206 - _LORA_B_PTR_DICT.clear()
207 - torch.ops.vllm.sgmv_expand(
208 - data.inputs_tensor,
209 - data.lora_weights,
210 - sgmv_out_tensor,
211 - *sgmv_meta_args,
212 - offset_start=0,
213 - add_inputs=add_inputs,
214 - )
215 -
216 - # V1 expand kernel
189 + # lora_expand kernel
190 _LORA_B_PTR_DICT.clear()
218 - torch.ops.vllm.v1_expand(data.inputs_tensor,
219 - data.lora_weights,
220 - v1_out_tensor,
221 - *v1_meta.meta_args(token_nums=token_nums),
222 - offset_start=0,
223 - add_inputs=add_inputs)
191 + triton_ops.lora_expand(data.inputs_tensor,
192 + data.lora_weights,
193 + out_tensor,
194 + *lora_meta.meta_args(token_nums=token_nums),
195 + offset_start=0,
196 + add_inputs=add_inputs)
197
198 # Reference
199 sgmv_expand_for_nslices(nslices,
@@ -231,124 +204,7 @@ def check_expand_kernels(batches: int, num_loras: int, rank: int,
204 *sgmv_meta_args,
205 add_inputs=add_inputs)
206
234 - assert_close(sgmv_out_tensor, ref_out_tensor)
235 - assert_close(v1_out_tensor, ref_out_tensor)
236 -
237 -
238 -def check_bgmv_shrink(batches: int, num_loras: int, rank: int,
239 - hidden_size: int, dtype: torch.dtype, device: str,
240 - scaling: float):
241 - """
242 - Compare vllm.bgmv_shrink against a reference implementation.
243 - """
244 - seq_length = 1
245 - data: PunicaTensors = generate_data(
246 - batches,
247 - hidden_size,
248 - num_loras,
249 - rank,
250 - seq_length,
251 - dtype,
252 - "shrink",
253 - device,
254 - )
255 -
256 - torch.ops.vllm.bgmv_shrink(
257 - data.inputs_tensor,
258 - data.lora_weights,
259 - data.our_out_tensor,
260 - data.token_lora_mapping,
261 - scaling,
262 - )
263 -
264 - bgmv_shrink(
265 - data.inputs_tensor,
266 - data.lora_weights,
267 - data.ref_out_tensor,
268 - data.token_lora_mapping,
269 - scaling,
270 - )
271 -
272 - data.ref_out_tensor = data.ref_out_tensor.to(torch.float32)
273 - assert_close(data.our_out_tensor, data.ref_out_tensor)
274 -
275 -
276 -def check_bgmv_expand(batches: int, num_loras: int, rank: int,
277 - hidden_size: int, dtype: torch.dtype, device: str,
278 - add_inputs: bool):
279 - """
280 - Compare vllm.bgmv_expand against a reference implementation.
281 - """
282 - seq_length = 1
283 - data: PunicaTensors = generate_data(
284 - batches,
285 - hidden_size,
286 - num_loras,
287 - rank,
288 - seq_length,
289 - dtype,
290 - "expand",
291 - device,
292 - )
293 -
294 - torch.ops.vllm.bgmv_expand(
295 - data.inputs_tensor,
296 - data.lora_weights,
297 - data.our_out_tensor,
298 - data.token_lora_mapping,
299 - add_inputs=add_inputs,
300 - )
301 - bgmv_expand(
302 - data.inputs_tensor,
303 - data.lora_weights,
304 - data.ref_out_tensor,
305 - data.token_lora_mapping,
306 - add_inputs=add_inputs,
307 - )
308 - assert_close(data.our_out_tensor, data.ref_out_tensor)
309 -
310 -
311 -def check_bgmv_expand_slice(batches: int, num_loras: int, rank: int,
312 - hidden_size: int, nslices: int, dtype: torch.dtype,
313 - device: str, add_inputs: bool):
314 - """
315 - Compare vllm.bgmv_expand_slice against a reference implementation.
316 - """
317 - seq_length = 1
318 - data: PunicaTensors = generate_data_for_expand_nslices(
319 - batches,
320 - hidden_size,
321 - num_loras,
322 - rank,
323 - seq_length,
324 - dtype,
325 - nslices,
326 - device,
327 - )
328 -
329 - slice_offset = 0
330 - for index in range(nslices):
331 - torch.ops.vllm.bgmv_expand_slice(
332 - data.inputs_tensor,
333 - data.lora_weights[index],
334 - data.our_out_tensor,
335 - data.token_lora_mapping,
336 - slice_offset,
337 - slice_size=hidden_size,
338 - add_inputs=add_inputs,
339 - )
340 - bgmv_expand_slice(
341 - data.inputs_tensor,
342 - data.lora_weights[index],
343 - data.ref_out_tensor,
344 - data.token_lora_mapping,
345 - slice_offset,
346 - slice_size=hidden_size,
347 - add_inputs=add_inputs,
348 - )
349 -
350 - slice_offset += hidden_size
351 - assert_close(data.our_out_tensor, data.ref_out_tensor)
207 + assert_close(out_tensor, ref_out_tensor)
208
209
210 # Tests
@@ -490,31 +346,31 @@ def test_kernels(
346 op_type: str,
347 ):
348 """
493 - Tests SGMV and V1 kernels.
349 + Tests LoRA kernels.
350 """
351 torch.set_default_device(device)
352 current_platform.seed_everything(seed)
353
354 if op_type == "shrink":
499 - check_shrink_kernels(batches=batches,
500 - num_loras=num_loras,
501 - rank=rank,
502 - hidden_size=hidden_size,
503 - nslices=nslices,
504 - dtype=dtype,
505 - device=device,
506 - seq_length=128,
507 - scaling=0.5)
355 + check_lora_shrink_kernel(batches=batches,
356 + num_loras=num_loras,
357 + rank=rank,
358 + hidden_size=hidden_size,
359 + nslices=nslices,
360 + dtype=dtype,
361 + device=device,
362 + seq_length=128,
363 + scaling=0.5)
364 else:
509 - check_expand_kernels(batches=batches,
510 - num_loras=num_loras,
511 - rank=rank,
512 - hidden_size=hidden_size,
513 - nslices=nslices,
514 - dtype=dtype,
515 - device=device,
516 - seq_length=128,
517 - add_inputs=True)
365 + check_lora_expand_kernel(batches=batches,
366 + num_loras=num_loras,
367 + rank=rank,
368 + hidden_size=hidden_size,
369 + nslices=nslices,
370 + dtype=dtype,
371 + device=device,
372 + seq_length=128,
373 + add_inputs=True)
374
375
376 @pytest.mark.parametrize("batches", hs_test_params['batches'])
@@ -538,159 +394,28 @@ def test_kernels_hidden_size(
394 op_type: str,
395 ):
396 """
541 - Tests SGMV and V1 kernels.
397 + Tests SGMV and LoRA kernels.
398 """
399 torch.set_default_device(device)
400 current_platform.seed_everything(seed)
401
402 if op_type == "shrink":
547 - check_shrink_kernels(batches=batches,
548 - num_loras=num_loras,
549 - rank=rank,
550 - hidden_size=hidden_size,
551 - nslices=nslices,
552 - dtype=dtype,
553 - device=device,
554 - seq_length=128,
555 - scaling=0.5)
556 - else:
557 - check_expand_kernels(batches=batches,
558 - num_loras=num_loras,
559 - rank=rank,
560 - hidden_size=hidden_size,
561 - nslices=nslices,
562 - dtype=dtype,
563 - device=device,
564 - seq_length=128,
565 - add_inputs=True)
566 -
567 -
568 -@pytest.mark.parametrize("batches", test_params['batches'])
569 -@pytest.mark.parametrize("num_loras", test_params['num_loras'])
570 -@pytest.mark.parametrize("rank", test_params['max_ranks'])
571 -@pytest.mark.parametrize("hidden_size", test_params['hidden_sizes'])
572 -@pytest.mark.parametrize("dtype", DTYPES)
573 -@pytest.mark.parametrize("device", DEVICES)
574 -@pytest.mark.parametrize("seed", SEED)
575 -@pytest.mark.parametrize("op_type", ["shrink", "expand"])
576 -def test_punica_bgmv(
577 - batches: int,
578 - num_loras: int,
579 - rank: int,
580 - hidden_size: int,
581 - dtype: torch.dtype,
582 - device: str,
583 - seed: int,
584 - op_type: str,
585 -):
586 - torch.set_default_device(device)
587 - current_platform.seed_everything(seed)
588 -
589 - if op_type == "shrink":
590 - check_bgmv_shrink(batches=batches,
591 - num_loras=num_loras,
592 - rank=rank,
593 - hidden_size=hidden_size,
594 - dtype=dtype,
595 - device=device,
596 - scaling=0.5)
597 - else:
598 - check_bgmv_expand(batches=batches,
599 - num_loras=num_loras,
600 - rank=rank,
601 - hidden_size=hidden_size,
602 - dtype=dtype,
603 - device=device,
604 - add_inputs=True)
605 -
606 -
607 -@pytest.mark.parametrize("batches", hs_test_params['batches'])
608 -@pytest.mark.parametrize("num_loras", hs_test_params['num_loras'])
609 -@pytest.mark.parametrize("rank", hs_test_params['max_ranks'])
610 -@pytest.mark.parametrize("hidden_size", hs_test_params['hidden_sizes'])
611 -@pytest.mark.parametrize("dtype", DTYPES)
612 -@pytest.mark.parametrize("device", DEVICES)
613 -@pytest.mark.parametrize("seed", SEED)
614 -@pytest.mark.parametrize("op_type", ["shrink", "expand"])
615 -def test_punica_bgmv_hidden_size(
616 - batches: int,
617 - num_loras: int,
618 - rank: int,
619 - hidden_size: int,
620 - dtype: torch.dtype,
621 - device: str,
622 - seed: int,
623 - op_type: str,
624 -):
625 - torch.set_default_device(device)
626 - current_platform.seed_everything(seed)
627 -
628 - if op_type == "shrink":
629 - check_bgmv_shrink(batches=batches,
630 - num_loras=num_loras,
631 - rank=rank,
632 - hidden_size=hidden_size,
633 - dtype=dtype,
634 - device=device,
635 - scaling=0.5)
403 + check_lora_shrink_kernel(batches=batches,
404 + num_loras=num_loras,
405 + rank=rank,
406 + hidden_size=hidden_size,
407 + nslices=nslices,
408 + dtype=dtype,
409 + device=device,
410 + seq_length=128,
411 + scaling=0.5)
412 else:
637 - check_bgmv_expand(batches=batches,
638 - num_loras=num_loras,
639 - rank=rank,
640 - hidden_size=hidden_size,
641 - dtype=dtype,
642 - device=device,
643 - add_inputs=True)
644 -
645 -
646 -@pytest.mark.parametrize("batches", test_params['batches'])
647 -@pytest.mark.parametrize("num_loras", test_params['num_loras'])
648 -@pytest.mark.parametrize("rank", test_params['max_ranks'])
649 -@pytest.mark.parametrize("hidden_size", test_params['hidden_sizes'])
650 -@pytest.mark.parametrize("nslices", [2, 3])
651 -@pytest.mark.parametrize("dtype", DTYPES)
652 -@pytest.mark.parametrize("device", DEVICES)
653 -@pytest.mark.parametrize("seed", SEED)
654 -def test_punica_bgmv_expand_nslices(batches: int, num_loras: int, rank: int,
655 - hidden_size: int, nslices: int,
656 - dtype: torch.dtype, device: str,
657 - seed: int):
658 -
659 - torch.set_default_device(device)
660 - current_platform.seed_everything(seed)
661 -
662 - check_bgmv_expand_slice(batches=batches,
663 - num_loras=num_loras,
664 - rank=rank,
665 - hidden_size=hidden_size,
666 - nslices=nslices,
667 - dtype=dtype,
668 - device=device,
669 - add_inputs=True)
670 -
671 -
672 -@pytest.mark.parametrize("batches", hs_test_params['batches'])
673 -@pytest.mark.parametrize("num_loras", hs_test_params['num_loras'])
674 -@pytest.mark.parametrize("rank", hs_test_params['max_ranks'])
675 -@pytest.mark.parametrize("hidden_size", hs_test_params['hidden_sizes'])
676 -@pytest.mark.parametrize("nslices", [2, 3])
677 -@pytest.mark.parametrize("dtype", DTYPES)
678 -@pytest.mark.parametrize("device", DEVICES)
679 -@pytest.mark.parametrize("seed", SEED)
680 -def test_punica_bgmv_expand_nslices_hidden_size(batches: int, num_loras: int,
681 - rank: int, hidden_size: int,
682 - nslices: int,
683 - dtype: torch.dtype,
684 - device: str, seed: int):
685 -
686 - torch.set_default_device(device)
687 - current_platform.seed_everything(seed)
688 -
689 - check_bgmv_expand_slice(batches=batches,
690 - num_loras=num_loras,
691 - rank=rank,
692 - hidden_size=hidden_size,
693 - nslices=nslices,
694 - dtype=dtype,
695 - device=device,
696 - add_inputs=True)
413 + check_lora_expand_kernel(batches=batches,
414 + num_loras=num_loras,
415 + rank=rank,
416 + hidden_size=hidden_size,
417 + nslices=nslices,
418 + dtype=dtype,
419 + device=device,
420 + seq_length=128,
421 + add_inputs=True)
vllm/lora/ops/triton_ops/__init__.py
+6 -10
@@ -1,15 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 -from vllm.lora.ops.triton_ops.bgmv_expand import bgmv_expand
4 -from vllm.lora.ops.triton_ops.bgmv_expand_slice import bgmv_expand_slice
5 -from vllm.lora.ops.triton_ops.bgmv_shrink import bgmv_shrink
6 -from vllm.lora.ops.triton_ops.sgmv_expand import sgmv_expand
7 -from vllm.lora.ops.triton_ops.sgmv_shrink import sgmv_shrink # noqa: F401
3 +from vllm.lora.ops.triton_ops.lora_expand import lora_expand
4 +from vllm.lora.ops.triton_ops.lora_kernel_metadata import LoRAKernelMeta
5 +from vllm.lora.ops.triton_ops.lora_shrink import lora_shrink
6
7 __all__ = [
10 - "bgmv_expand",
11 - "bgmv_expand_slice",
12 - "bgmv_shrink",
13 - "sgmv_expand",
14 - "sgmv_shrink",
8 + "lora_expand",
9 + "lora_shrink",
10 + "LoRAKernelMeta",
11 ]
vllm/lora/ops/triton_ops/bgmv_expand.py deleted
-188
@@ -1,188 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -"""
3 -Based on:
4 -Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
5 -Punica: Multi-Tenant LoRA Serving.
6 -https://arxiv.org/abs/2310.18547
7 -"""
8 -
9 -import torch
10 -import triton
11 -import triton.language as tl
12 -
13 -from vllm.utils import direct_register_custom_op
14 -
15 -from .utils import get_lora_op_configs
16 -
17 -
18 -@triton.jit
19 -def _bgmv_expand_kernel(
20 - input_ptr,
21 - lora_ptr,
22 - out_ptr,
23 - N,
24 - K,
25 - lora_indices,
26 - xm_stride,
27 - xk_stride,
28 - l0_stride,
29 - lora_k_stride,
30 - lora_n_stride,
31 - cm_stride,
32 - cn_stride,
33 - BLOCK_N: tl.constexpr,
34 - BLOCK_K: tl.constexpr,
35 - SPLIT_N: tl.constexpr,
36 - EVEN_K: tl.constexpr,
37 - ADD_INPUTS: tl.constexpr,
38 - CAST_TYPE: tl.constexpr,
39 -):
40 - """
41 - GroupGEMV, additionally, introducing SPLIT_N can improve large hidden_size's
42 - performance
43 - """
44 - pid_sn = tl.program_id(axis=0)
45 - cur_batch = tl.program_id(axis=1)
46 - lora_index = tl.load(lora_indices + cur_batch)
47 - if lora_index == -1:
48 - return
49 - offset_k = tl.arange(0, BLOCK_K)
50 - offset_n = tl.arange(0, BLOCK_N)
51 - if EVEN_K:
52 - tiled_a = tl.load(input_ptr + cur_batch * xm_stride +
53 - offset_k * xk_stride, ) # [BLOCK_K]
54 - else:
55 - tiled_a = tl.load(
56 - input_ptr + cur_batch * xm_stride + offset_k * xk_stride,
57 - mask=offset_k < K,
58 - other=0,
59 - ) # [BLOCK_K]
60 - # N must be divisible by SPLIT_N
61 - split_n_length = tl.cdiv(N, SPLIT_N)
62 - if CAST_TYPE:
63 - tiled_a = tiled_a.to(lora_ptr.dtype.element_ty)
64 - # sliding to next row-block
65 - b_ptr = (lora_ptr + l0_stride * lora_index +
66 - pid_sn * split_n_length * lora_k_stride)
67 - c_ptr = out_ptr + cur_batch * cm_stride + pid_sn * split_n_length
68 - for n in range(0, split_n_length, BLOCK_N):
69 - current_n = n + offset_n
70 - current_n_c = tl.max_contiguous(current_n, BLOCK_N)
71 - b_ptr_mask = (current_n[:, None] < split_n_length) & (offset_k[None, :]
72 - < K)
73 - c_mask = current_n < split_n_length
74 - tiled_b = tl.load(
75 - b_ptr + current_n_c[:, None] * lora_k_stride +
76 - offset_k[None, :] * lora_n_stride,
77 - mask=b_ptr_mask,
78 - other=0.0,
79 - ) # [BLOCK_N,BLOCK_K]
80 - if ADD_INPUTS:
81 - tiled_out = tl.load(c_ptr + current_n * cn_stride,
82 - mask=c_mask,
83 - other=0.0)
84 - accumulator = tl.sum(tiled_a * tiled_b, 1) + tiled_out
85 - else:
86 - accumulator = tl.sum(tiled_a * tiled_b, 1)
87 -
88 - tl.store(c_ptr + current_n * cn_stride, accumulator, mask=c_mask)
89 -
90 -
91 -@torch.inference_mode()
92 -def _bgmv_expand(
93 - inputs: torch.Tensor,
94 - lora_b_weights: torch.Tensor,
95 - output_tensor: torch.Tensor,
96 - lora_indices_tensor: torch.Tensor,
97 - add_inputs: bool = True,
98 -) -> None:
99 - """
100 - Args:
101 - inputs (torch.Tensor): input tensor
102 - lora_b_weights (torch.Tensor): lora'a weight
103 - output_tensor (torch.Tensor): output tensor
104 - lora_indices_tensor (torch.Tensor): (batch_size,). The LoRA index
105 - corresponding to each batch, An index of -1 means no lora should be
106 - applied.
107 - batches (int): batch size
108 - add_inputs (bool, optional): Defaults to False, adds the final lora
109 - results to the output.
110 - """
111 - assert inputs.dtype in [torch.float16, torch.bfloat16, torch.float32]
112 - assert lora_b_weights.dtype in [
113 - torch.float16,
114 - torch.bfloat16,
115 - ]
116 - assert inputs.size(1) == lora_b_weights.size(-1)
117 -
118 - assert inputs.is_contiguous()
119 - assert output_tensor.is_contiguous()
120 -
121 - if lora_b_weights.ndim == 4: # shape:(lora_num,1,size,rank)
122 - assert lora_b_weights.size(1) == 1
123 - lora_b_weights = lora_b_weights.squeeze(dim=1)
124 - else:
125 - assert lora_b_weights.ndim == 3 # shape:(lora_num,size,rank)
126 - assert lora_b_weights.is_contiguous()
127 -
128 - # TODO tuning this config
129 - N, K = lora_b_weights.shape[-2:] # K= rank,N=hidden_size
130 - BLOCK_K = triton.next_power_of_2(K)
131 - EVEN_K = K % BLOCK_K == 0
132 - ADD_INPUTS = add_inputs
133 - CAST_TYPE = False
134 - if inputs.dtype == torch.float32 and lora_b_weights.dtype in [
135 - torch.float16,
136 - torch.bfloat16,
137 - ]:
138 - CAST_TYPE = True
139 - batches = lora_indices_tensor.size(0)
140 - config = get_lora_op_configs("expand", batches, N)
141 - grid = lambda META: (
142 - META["SPLIT_N"],
143 - batches,
144 - )
145 - _bgmv_expand_kernel[grid](
146 - inputs,
147 - lora_b_weights,
148 - output_tensor,
149 - N,
150 - K,
151 - lora_indices_tensor,
152 - inputs.stride(0),
153 - inputs.stride(1),
154 - lora_b_weights.stride(0),
155 - lora_b_weights.stride(1),
156 - lora_b_weights.stride(2),
157 - output_tensor.stride(0),
158 - output_tensor.stride(1),
159 - BLOCK_K=BLOCK_K,
160 - EVEN_K=EVEN_K,
161 - ADD_INPUTS=ADD_INPUTS,
162 - CAST_TYPE=CAST_TYPE,
163 - **config,
164 - )
165 - return
166 -
167 -
168 -def bgmv_expand_fake(
169 - inputs: torch.Tensor,
170 - lora_b_weights: torch.Tensor,
171 - output_tensor: torch.Tensor,
172 - lora_indices_tensor: torch.Tensor,
173 - add_inputs: bool = True,
174 -) -> None:
175 - return
176 -
177 -
178 -try:
179 - direct_register_custom_op(
180 - op_name="bgmv_expand",
181 - op_func=_bgmv_expand,
182 - mutates_args=["output_tensor"],
183 - fake_impl=bgmv_expand_fake,
184 - )
185 - bgmv_expand = torch.ops.vllm.bgmv_expand
186 -
187 -except AttributeError:
188 - bgmv_expand = _bgmv_expand
vllm/lora/ops/triton_ops/bgmv_expand_slice.py deleted
-207
@@ -1,207 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -"""
3 -Based on:
4 -Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
5 -Punica: Multi-Tenant LoRA Serving.
6 -https://arxiv.org/abs/2310.18547
7 -"""
8 -
9 -import torch
10 -import triton
11 -import triton.language as tl
12 -
13 -from vllm.utils import direct_register_custom_op
14 -
15 -from .utils import get_lora_op_configs
16 -
17 -
18 -@triton.jit
19 -def _bgmv_expand_slice_kernel(
20 - input_ptr,
21 - lora_ptr,
22 - out_ptr,
23 - N,
24 - K,
25 - lora_indices,
26 - xm_stride,
27 - xk_stride,
28 - l0_stride,
29 - lora_k_stride,
30 - lora_n_stride,
31 - cm_stride,
32 - cn_stride,
33 - slice_offset,
34 - BLOCK_N: tl.constexpr,
35 - BLOCK_K: tl.constexpr,
36 - SPLIT_N: tl.constexpr,
37 - EVEN_K: tl.constexpr,
38 - ADD_INPUTS: tl.constexpr,
39 - CAST_TYPE: tl.constexpr,
40 -):
41 - """
42 - GroupGEMV, additionally, introducing SPLIT_N can improve large hidden_size's
43 - performance
44 - """
45 - pid_sn = tl.program_id(axis=0)
46 - cur_batch = tl.program_id(axis=1)
47 - lora_index = tl.load(lora_indices + cur_batch)
48 - if lora_index == -1:
49 - return
50 - offset_k = tl.arange(0, BLOCK_K)
51 - offset_n = tl.arange(0, BLOCK_N)
52 - if EVEN_K:
53 - tiled_a = tl.load(input_ptr + cur_batch * xm_stride +
54 - offset_k * xk_stride, ) # [BLOCK_K]
55 - else:
56 - tiled_a = tl.load(
57 - input_ptr + cur_batch * xm_stride + offset_k * xk_stride,
58 - mask=offset_k < K,
59 - other=0,
60 - ) # [BLOCK_K]
61 - # N must be divisible by SPLIT_N
62 - split_n_length = tl.cdiv(N, SPLIT_N)
63 - if CAST_TYPE:
64 - tiled_a = tiled_a.to(lora_ptr.dtype.element_ty)
65 - # sliding to next row-block
66 - b_ptr = (lora_ptr + l0_stride * lora_index +
67 - pid_sn * split_n_length * lora_k_stride)
68 - c_ptr = (out_ptr + cur_batch * cm_stride + pid_sn * split_n_length +
69 - slice_offset * cn_stride)
70 -
71 - for n in range(0, split_n_length, BLOCK_N):
72 - current_n = n + offset_n
73 - b_ptr_mask = (current_n[:, None] < split_n_length) & (offset_k[None, :]
74 - < K)
75 - c_mask = current_n < split_n_length
76 - tiled_b = tl.load(
77 - b_ptr + current_n[:, None] * lora_k_stride +
78 - offset_k[None, :] * lora_n_stride,
79 - mask=b_ptr_mask,
80 - other=0.0,
81 - ) # [BLOCK_N,BLOCK_K]
82 -
83 - if ADD_INPUTS:
84 - # explicitly pass in other=None to tell triton that masked values
85 - # can be uninitialized. This is OK because the later tl.store
86 - # operation uses the same mask, eliminating the risk of garbage
87 - # values propagating
88 - tiled_out = tl.load(c_ptr + current_n * cn_stride,
89 - mask=c_mask,
90 - other=None)
91 - accumulator = tl.sum(tiled_a * tiled_b, 1) + tiled_out
92 - else:
93 - accumulator = tl.sum(tiled_a * tiled_b, 1)
94 -
95 - tl.store(c_ptr + current_n * cn_stride, accumulator, mask=c_mask)
96 -
97 -
98 -@torch.inference_mode()
99 -def _bgmv_expand_slice(
100 - inputs: torch.Tensor,
101 - lora_b_weights: torch.Tensor,
102 - output_tensor: torch.Tensor,
103 - lora_indices_tensor: torch.Tensor,
104 - slice_offset: int,
105 - slice_size: int,
106 - add_inputs: bool = True,
107 -) -> None:
108 - """
109 - Args:
110 - inputs (torch.Tensor): input tensor
111 - lora_b_weights (torch.Tensor): lora'b weight
112 - output_tensor (torch.Tensor): output tensor
113 - lora_indices_tensor (torch.Tensor): (batch_size,). The LoRA index
114 - corresponding to each batch, An index of -1 means no lora should be
115 - applied.
116 - slice_offset (int): output_tensor's offset
117 - slice_size (int): current output_tensor's size
118 - batches (int): batch size
119 - add_inputs (bool, optional): Defaults to False.
120 - """
121 - assert inputs.dtype in [torch.float16, torch.bfloat16, torch.float32]
122 - assert lora_b_weights.dtype in [
123 - torch.float16,
124 - torch.bfloat16,
125 - ]
126 - assert inputs.size(1) == lora_b_weights.size(-1)
127 -
128 - assert slice_size == lora_b_weights.size(-2)
129 - assert inputs.is_contiguous()
130 - assert output_tensor.is_contiguous()
131 -
132 - if lora_b_weights.ndim == 4: # shape:(lora_num,1,size,rank)
133 - assert lora_b_weights.size(1) == 1
134 - lora_b_weights = lora_b_weights.squeeze(dim=1)
135 - else:
136 - assert lora_b_weights.ndim == 3 # shape:(lora_num,size,rank)
137 -
138 - assert lora_b_weights.is_contiguous()
139 -
140 - # TODO tuning this config
141 -
142 - N, K = lora_b_weights.shape[-2:] # K= rank,N=hidden_size
143 - BLOCK_K = triton.next_power_of_2(K)
144 - EVEN_K = K % BLOCK_K == 0
145 - ADD_INPUTS = add_inputs
146 - CAST_TYPE = False
147 - if inputs.dtype == torch.float32 and lora_b_weights.dtype in [
148 - torch.float16,
149 - torch.bfloat16,
150 - ]:
151 - CAST_TYPE = True
152 -
153 - batches = lora_indices_tensor.size(0)
154 -
155 - config = get_lora_op_configs("expand", batches, N)
156 -
157 - grid = lambda META: (
158 - META["SPLIT_N"],
159 - batches,
160 - )
161 - _bgmv_expand_slice_kernel[grid](
162 - inputs,
163 - lora_b_weights,
164 - output_tensor,
165 - N,
166 - K,
167 - lora_indices_tensor,
168 - inputs.stride(0),
169 - inputs.stride(1),
170 - lora_b_weights.stride(0),
171 - lora_b_weights.stride(1),
172 - lora_b_weights.stride(2),
173 - output_tensor.stride(0),
174 - output_tensor.stride(1),
175 - slice_offset,
176 - BLOCK_K=BLOCK_K,
177 - EVEN_K=EVEN_K,
178 - ADD_INPUTS=ADD_INPUTS,
179 - CAST_TYPE=CAST_TYPE,
180 - **config,
181 - )
182 - return
183 -
184 -
185 -def bgmv_expand_slice_fake(
186 - inputs: torch.Tensor,
187 - lora_b_weights: torch.Tensor,
188 - output_tensor: torch.Tensor,
189 - lora_indices_tensor: torch.Tensor,
190 - slice_offset: int,
191 - slice_size: int,
192 - add_inputs: bool = True,
193 -) -> None:
194 - return
195 -
196 -
197 -try:
198 - direct_register_custom_op(
199 - op_name="bgmv_expand_slice",
200 - op_func=_bgmv_expand_slice,
201 - mutates_args=["output_tensor"],
202 - fake_impl=bgmv_expand_slice_fake,
203 - )
204 - bgmv_expand_slice = torch.ops.vllm.bgmv_expand_slice
205 -
206 -except AttributeError:
207 - bgmv_expand_slice = _bgmv_expand_slice
vllm/lora/ops/triton_ops/bgmv_shrink.py deleted
-168
@@ -1,168 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -"""
3 -Based on:
4 -Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
5 -Punica: Multi-Tenant LoRA Serving.
6 -https://arxiv.org/abs/2310.18547
7 -"""
8 -
9 -import torch
10 -import triton
11 -import triton.language as tl
12 -
13 -from vllm.utils import direct_register_custom_op
14 -
15 -from .utils import get_lora_op_configs
16 -
17 -
18 -@triton.jit
19 -def _bgmv_shrink_kernel(
20 - input_ptr,
21 - lora_ptr,
22 - out_ptr,
23 - N,
24 - K,
25 - lora_indices,
26 - scaling,
27 - xm_stride,
28 - xk_stride,
29 - l0_stride,
30 - lora_k_stride,
31 - lora_n_stride,
32 - cm_stride,
33 - cn_stride,
34 - BLOCK_N: tl.constexpr,
35 - BLOCK_K: tl.constexpr,
36 - SPLIT_K: tl.constexpr,
37 -):
38 - """
39 - GroupGEMV, additionally, introducing SPLIT-K can improve large hidden_size's
40 - performance
41 - """
42 - pid_sk = tl.program_id(axis=0)
43 - cur_batch = tl.program_id(axis=1)
44 - lora_index = tl.load(lora_indices + cur_batch)
45 - if lora_index == -1:
46 - return
47 -
48 - offset_n = tl.arange(0, BLOCK_N)
49 - offset_k = tl.arange(0, BLOCK_K) + pid_sk * BLOCK_K
50 - a_ptr = input_ptr + cur_batch * xm_stride
51 - b_ptr = lora_ptr + l0_stride * lora_index
52 - accumulator = tl.zeros((BLOCK_N, ), dtype=tl.float32)
53 - for k in range(0, K, BLOCK_K * SPLIT_K):
54 - current_k = k + offset_k
55 - current_k_c = tl.max_contiguous(current_k, BLOCK_K)
56 - tiled_a = tl.load(
57 - a_ptr + current_k_c,
58 - mask=current_k < K,
59 - other=0.0,
60 - ) # [BLOCK_K]
61 - b_ptr_mask = (offset_n[:, None] < N) & (current_k[None, :] < K)
62 -
63 - tiled_b = tl.load(
64 - b_ptr + offset_n[:, None] * lora_k_stride +
65 - current_k[None, :] * lora_n_stride,
66 - mask=b_ptr_mask,
67 - other=0.0,
68 - ) # [BLOCK_N,BLOCK_K]
69 -
70 - accumulator += tl.sum(tiled_a * tiled_b, 1)
71 - accumulator *= scaling
72 - offset_cn = tl.arange(0, BLOCK_N)
73 - c_ptr = out_ptr + cur_batch * cm_stride + offset_cn * cn_stride
74 - c_mask = offset_cn < N
75 - if SPLIT_K == 1:
76 - tl.store(c_ptr, accumulator, mask=c_mask)
77 - else:
78 - tl.atomic_add(c_ptr, accumulator, mask=c_mask)
79 -
80 -
81 -@torch.inference_mode()
82 -def _bgmv_shrink(
83 - inputs: torch.Tensor,
84 - lora_a_weights: torch.Tensor,
85 - output_tensor: torch.Tensor,
86 - lora_indices_tensor: torch.Tensor,
87 - scaling: float = 1.0,
88 -) -> None:
89 - """
90 - Args:
91 - inputs (torch.Tensor): input tensor
92 - lora_a_weights (torch.Tensor): lora'a weight
93 - output_tensor (torch.Tensor): output tensor
94 - lora_indices_tensor (torch.Tensor): (batch_size,). The LoRA index
95 - corresponding to each batch. An index of -1 means no lora should be
96 - applied.
97 - batches (int): batch size
98 - scaling (float): Scaling factor.
99 - """
100 - assert inputs.dtype == lora_a_weights.dtype
101 - assert inputs.dtype in [torch.float16, torch.bfloat16]
102 - assert lora_a_weights.dtype in [
103 - torch.float16,
104 - torch.bfloat16,
105 - ]
106 - assert inputs.size(1) == lora_a_weights.size(-1)
107 - assert inputs.is_contiguous()
108 -
109 - if lora_a_weights.ndim == 4: # shape:(lora_num,1,rank, size)
110 - assert lora_a_weights.size(1) == 1
111 - lora_a_weights = lora_a_weights.squeeze(dim=1)
112 - else:
113 - assert lora_a_weights.ndim == 3 # shape:(lora_num,rank, size)
114 - assert lora_a_weights.is_contiguous()
115 - assert output_tensor.is_contiguous()
116 - # TODO tuning this config
117 - batches = lora_indices_tensor.size(0)
118 - N, K = lora_a_weights.shape[-2:] # K=hidden_size,N=rank
119 - BLOCK_N = triton.next_power_of_2(N)
120 - # First try to load optimal config from the file
121 - config = get_lora_op_configs("bgmv_shrink", batches, K)
122 -
123 - grid = lambda META: (
124 - META["SPLIT_K"],
125 - batches,
126 - )
127 - _bgmv_shrink_kernel[grid](
128 - inputs,
129 - lora_a_weights,
130 - output_tensor,
131 - N,
132 - K,
133 - lora_indices_tensor,
134 - scaling,
135 - inputs.stride(0),
136 - inputs.stride(1),
137 - lora_a_weights.stride(0),
138 - lora_a_weights.stride(1),
139 - lora_a_weights.stride(2),
140 - output_tensor.stride(0),
141 - output_tensor.stride(1),
142 - BLOCK_N=BLOCK_N,
143 - **config,
144 - )
145 - return
146 -
147 -
148 -def bgmv_shrink_fake(
149 - inputs: torch.Tensor,
150 - lora_a_weights: torch.Tensor,
151 - output_tensor: torch.Tensor,
152 - lora_indices_tensor: torch.Tensor,
153 - scaling: float = 1.0,
154 -) -> None:
155 - return
156 -
157 -
158 -try:
159 - direct_register_custom_op(
160 - op_name="bgmv_shrink",
161 - op_func=_bgmv_shrink,
162 - mutates_args=["output_tensor"],
163 - fake_impl=bgmv_shrink_fake,
164 - )
165 - bgmv_shrink = torch.ops.vllm.bgmv_shrink
166 -
167 -except AttributeError:
168 - bgmv_shrink = _bgmv_shrink
vllm/lora/ops/triton_ops/lora_expand.py renamed
+9 -9
@@ -18,7 +18,7 @@ from vllm.utils import direct_register_custom_op
18
19
20 @triton.jit
21 -def _v1_expand_kernel(
21 +def _lora_expand_kernel(
22 input_ptr,
23 lora_ptr,
24 out_ptr,
@@ -125,7 +125,7 @@ def _v1_expand_kernel(
125
126
127 @torch.inference_mode()
128 -def _v1_expand(
128 +def _lora_expand(
129 inputs: torch.Tensor, # shape [num_slices, num_tokens, lora_rank]
130 lora_b_weights: List[
131 torch.Tensor], # shape [num_lora, hidden_size, lora_rank]
@@ -216,7 +216,7 @@ def _v1_expand(
216 MAX_LORAS,
217 )
218
219 - _v1_expand_kernel[grid](
219 + _lora_expand_kernel[grid](
220 inputs,
221 lora_ptr_tensor,
222 output_tensor,
@@ -254,7 +254,7 @@ def _v1_expand(
254 return
255
256
257 -def _v1_expand_fake(
257 +def _lora_expand_fake(
258 inputs: torch.Tensor,
259 lora_b_weights: List[torch.Tensor],
260 output_tensor: torch.Tensor,
@@ -271,12 +271,12 @@ def _v1_expand_fake(
271
272 try:
273 direct_register_custom_op(
274 - op_name="v1_expand",
275 - op_func=_v1_expand,
274 + op_name="lora_expand",
275 + op_func=_lora_expand,
276 mutates_args=["output_tensor"],
277 - fake_impl=_v1_expand_fake,
277 + fake_impl=_lora_expand_fake,
278 )
279 - v1_expand = torch.ops.vllm.v1_expand
279 + lora_expand = torch.ops.vllm.lora_expand
280
281 except AttributeError:
282 - v1_expand = _v1_expand
282 + lora_expand = _lora_expand
vllm/lora/ops/triton_ops/lora_kernel_metadata.py renamed
+5 -5
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """
3 -V1 LoRA kernels metadata preparation utilities.
3 +LoRA kernels metadata preparation utilities.
4 """
5
6 from dataclasses import dataclass
@@ -10,7 +10,7 @@ import torch
10
11
12 @dataclass
13 -class V1KernelMeta:
13 +class LoRAKernelMeta:
14 token_lora_mapping: torch.Tensor
15 token_indices_sorted_by_lora_ids: torch.Tensor
16 active_lora_ids: torch.Tensor
@@ -19,7 +19,7 @@ class V1KernelMeta:
19
20 @staticmethod
21 def make(max_loras: int, max_num_tokens: int,
22 - device: Union[torch.device, str]) -> "V1KernelMeta":
22 + device: Union[torch.device, str]) -> "LoRAKernelMeta":
23
24 token_lora_mapping = torch.empty(max_num_tokens,
25 dtype=torch.int32,
@@ -47,7 +47,7 @@ class V1KernelMeta:
47 lora_token_start_loc = torch.zeros(max_loras + 2,
48 dtype=torch.int32,
49 device=device)
50 - return V1KernelMeta(
50 + return LoRAKernelMeta(
51 token_lora_mapping=token_lora_mapping,
52 token_indices_sorted_by_lora_ids=token_indices_sorted_by_lora_ids,
53 active_lora_ids=active_lora_ids,
@@ -105,7 +105,7 @@ class V1KernelMeta:
105 This function returns the kernel metadata required for the current
106 forward pass execution of the kernel. The function returns all the
107 metadata required by the kernel, in order, as a tuple, so it can be
108 - unpacked directly during the v1_shrink/v1_expand function call.
108 + unpacked directly during the lora_shrink/lora_expand function call.
109
110 Args:
111 token_nums (int): Number of input tokens in the current forward
vllm/lora/ops/triton_ops/lora_shrink.py renamed
+17 -17
@@ -18,15 +18,15 @@ from vllm.utils import direct_register_custom_op
18
19
20 @triton.jit
21 -def _v1_shrink_kernel(input_ptr, lora_ptr, out_ptr, M, N, K,
22 - token_indices_sorted_by_lora_ids, num_tokens_per_lora,
23 - lora_token_start_loc, lora_ids, scaling, input_d0_stride,
24 - input_d1_stride, lora_d0_stride, lora_d1_stride,
25 - lora_d2_stride, output_d0_stride, output_d1_stride,
26 - output_d2_stride, BLOCK_M: tl.constexpr,
27 - BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
28 - EVEN_K: tl.constexpr, SPLIT_K: tl.constexpr,
29 - SLICE_NUM: tl.constexpr):
21 +def _lora_shrink_kernel(input_ptr, lora_ptr, out_ptr, M, N, K,
22 + token_indices_sorted_by_lora_ids, num_tokens_per_lora,
23 + lora_token_start_loc, lora_ids, scaling,
24 + input_d0_stride, input_d1_stride, lora_d0_stride,
25 + lora_d1_stride, lora_d2_stride, output_d0_stride,
26 + output_d1_stride, output_d2_stride,
27 + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
28 + BLOCK_K: tl.constexpr, EVEN_K: tl.constexpr,
29 + SPLIT_K: tl.constexpr, SLICE_NUM: tl.constexpr):
30
31 cta_n_num = tl.cdiv(N, BLOCK_N)
32 cta_m_num = tl.cdiv(M, BLOCK_M)
@@ -96,7 +96,7 @@ def _v1_shrink_kernel(input_ptr, lora_ptr, out_ptr, M, N, K,
96
97
98 @torch.inference_mode()
99 -def _v1_shrink(
99 +def _lora_shrink(
100 inputs: torch.Tensor, # shape [num_tokens, hidden_size]
101 lora_a_weights: List[
102 torch.Tensor], # shape [num_loras, lora_rank, hidden_size]
@@ -174,7 +174,7 @@ def _v1_shrink(
174 MAX_LORAS,
175 )
176
177 - _v1_shrink_kernel[grid](
177 + _lora_shrink_kernel[grid](
178 inputs,
179 lora_ptr_tensor,
180 output_tensor,
@@ -209,7 +209,7 @@ def _v1_shrink(
209 return
210
211
212 -def _v1_shrink_fake(
212 +def _lora_shrink_fake(
213 inputs: torch.Tensor,
214 lora_a_weights: List[torch.Tensor],
215 output_tensor: torch.Tensor,
@@ -225,12 +225,12 @@ def _v1_shrink_fake(
225
226 try:
227 direct_register_custom_op(
228 - op_name="v1_shrink",
229 - op_func=_v1_shrink,
228 + op_name="lora_shrink",
229 + op_func=_lora_shrink,
230 mutates_args=["output_tensor"],
231 - fake_impl=_v1_shrink_fake,
231 + fake_impl=_lora_shrink_fake,
232 )
233 - v1_shrink = torch.ops.vllm.v1_shrink
233 + lora_shrink = torch.ops.vllm.lora_shrink
234
235 except AttributeError:
236 - v1_shrink = _v1_shrink
236 + lora_shrink = _lora_shrink
vllm/lora/ops/triton_ops/sgmv_expand.py deleted
-249
@@ -1,249 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -"""
3 -Based on:
4 -Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
5 -Punica: Multi-Tenant LoRA Serving.
6 -https://arxiv.org/abs/2310.18547
7 -"""
8 -
9 -from typing import List
10 -
11 -import torch
12 -import triton
13 -import triton.language as tl
14 -
15 -from vllm.utils import direct_register_custom_op
16 -
17 -from .kernel_utils import do_expand_kernel
18 -from .utils import _get_lora_b_ptr
19 -
20 -
21 -@triton.jit
22 -def _sgmv_expand_kernel(
23 - input_ptr,
24 - lora_ptr,
25 - out_ptr,
26 - N,
27 - K,
28 - b_seq_start_loc,
29 - seq_lens,
30 - lora_indices,
31 - slice_start_loc,
32 - input_d0_stride,
33 - input_d1_stride,
34 - input_d2_stride, # 1
35 - ls_d0_ptr,
36 - ls_d1_ptr,
37 - ls_d2_ptr, # 1
38 - output_d0_stride,
39 - output_d1_stride, # 1
40 - output_hs_ptr,
41 - BLOCK_M: tl.constexpr,
42 - BLOCK_N: tl.constexpr,
43 - BLOCK_K: tl.constexpr,
44 - EVEN_K: tl.constexpr,
45 - ADD_INPUTS: tl.constexpr,
46 - CAST_TYPE: tl.constexpr,
47 - SLICE_NUM: tl.constexpr,
48 - SAME_STRIDE: tl.constexpr):
49 - """
50 -
51 - Similar to the 'sgmv_expand' operator, but with an added parameter
52 - 'slice_offset'. The reason for not reusing the 'sgmv_expand' operator
53 - might be that in the future, we could implement a fusion operator to
54 - achieve the current functionality instead of having to call it multiple
55 - times.
56 - """
57 - pid = tl.program_id(axis=0)
58 - cur_batch = tl.program_id(axis=1)
59 - slice_id = tl.program_id(axis=2)
60 - cta_n_num = tl.cdiv(N, BLOCK_N)
61 - # When the output dimensions of each slice are the same,cur_n=N, otherwise
62 - # cur_n=tl.load(output_hs_ptr + slice_id), this situation exists in GQA's
63 - # qkv linear.
64 - curr_N = N if SAME_STRIDE else tl.load(output_hs_ptr + slice_id)
65 - pid_m = pid // cta_n_num
66 - pid_n = pid % cta_n_num
67 -
68 - M = tl.load(seq_lens + cur_batch)
69 - if pid_m * BLOCK_M >= M:
70 - return
71 - if pid_n * BLOCK_N >= curr_N:
72 - return
73 - lora_index = tl.load(lora_indices + cur_batch)
74 - if lora_index == -1:
75 - return
76 -
77 - m_offset = tl.load(b_seq_start_loc + cur_batch)
78 -
79 - cta_m_len = min(BLOCK_M, M - (pid_m * BLOCK_M))
80 - cta_m_offset = m_offset + (pid_m * BLOCK_M)
81 - offset_m = tl.arange(0, BLOCK_M)
82 - ram = cta_m_offset + tl.max_contiguous(
83 - tl.multiple_of(offset_m % cta_m_len, BLOCK_M), BLOCK_M)
84 - do_expand_kernel(
85 - pid_n,
86 - lora_index,
87 - slice_id,
88 - input_ptr,
89 - lora_ptr,
90 - out_ptr,
91 - curr_N,
92 - K,
93 - cta_m_len,
94 - ram, # array identifying the rows of Input ptr to operate on
95 - slice_start_loc,
96 - # input ptr strides
97 - input_d0_stride,
98 - input_d1_stride,
99 - input_d2_stride,
100 - # lora ptr strides
101 - ls_d0_ptr,
102 - ls_d1_ptr,
103 - ls_d2_ptr,
104 - # out ptr strides
105 - output_d0_stride,
106 - output_d1_stride,
107 - # constants
108 - BLOCK_M,
109 - BLOCK_N,
110 - BLOCK_K,
111 - SAME_STRIDE,
112 - SLICE_NUM,
113 - EVEN_K,
114 - CAST_TYPE,
115 - ADD_INPUTS,
116 - )
117 -
118 -
119 -@torch.inference_mode()
120 -def _sgmv_expand(
121 - inputs: torch.Tensor,
122 - lora_b_weights: List[torch.Tensor],
123 - output_tensor: torch.Tensor,
124 - b_seq_start_loc: torch.Tensor,
125 - seq_len_tensor: torch.Tensor,
126 - lora_indices_tensor: torch.Tensor,
127 - batches: int,
128 - max_seq_length: int,
129 - token_nums: int,
130 - offset_start: int = 0,
131 - add_inputs: bool = False,
132 -) -> None:
133 - """
134 - Args:
135 - inputs (torch.Tensor): input tensor
136 - lora_b_weights (List[torch.Tensor]): lora'b weight
137 - output_tensor (torch.Tensor): output tensor
138 - b_seq_start_loc (torch.Tensor): (batch_size,). The cumulative
139 - sequence lengths of the sequences in the batch, used to index
140 - into sequence. E.g., if the sequence length is [4, 6], it is
141 - [0, 4].
142 - seq_len_tensor (torch.Tensor): (batch_size,). Record the sequence
143 - length of the sequences in the batch.
144 - lora_indices_tensor (torch.Tensor): (batch_size,). The LoRA index
145 - corresponding to each batch. An index of -1 means no lora should be
146 - applied.
147 - batches (int): batch size
148 - max_seq_length (int): The max sequence lengths of the sequences in the
149 - batch.
150 - token_nums (int): The token numbers in the batch. Used to verify if the
151 - token numbers in the inputs matches the one in the metadata.
152 - offset_start (int, optional): Offset start for output_tensor.
153 - Defaults to 0.
154 - add_inputs (bool, optional): Whether to add the input tensor to the
155 - output tensor. Defaults to False.
156 - """
157 - assert inputs.dtype in [torch.float16, torch.bfloat16, torch.float32]
158 - for weight in lora_b_weights:
159 - assert weight.dtype in [torch.float16, torch.bfloat16]
160 -
161 - assert inputs.size(1) == token_nums
162 - assert inputs.size(0) == len(lora_b_weights)
163 -
164 - assert b_seq_start_loc.size(0) == batches
165 - assert lora_indices_tensor.size(0) == batches
166 - assert output_tensor.is_contiguous()
167 - (slice_start_tensor, lora_ptr_tensor, lora_strides_d0_tensor,
168 - lora_strides_d1_tensor, lora_strides_d2_tensor, hidden_sizes_tensor,
169 - same_stride, MAX_N) = _get_lora_b_ptr(lora_b_weights, offset_start,
170 - b_seq_start_loc.device)
171 -
172 - # TODO tuning this config
173 - K = lora_b_weights[0].shape[-1] # K= rank
174 -
175 - BLOCK_M = 64
176 - BLOCK_N = 128
177 - BLOCK_K = 16
178 - EVEN_K = K % BLOCK_K == 0
179 - ADD_INPUTS = add_inputs
180 - CAST_TYPE = False
181 -
182 - if inputs.dtype == torch.float32 and lora_b_weights[0].dtype in [
183 - torch.float16,
184 - torch.bfloat16,
185 - ]:
186 - CAST_TYPE = True
187 - grid = (
188 - triton.cdiv(max_seq_length, BLOCK_M) * triton.cdiv(MAX_N, BLOCK_N),
189 - batches,
190 - len(lora_b_weights),
191 - )
192 - _sgmv_expand_kernel[grid](
193 - inputs,
194 - lora_ptr_tensor,
195 - output_tensor,
196 - MAX_N,
197 - K,
198 - b_seq_start_loc,
199 - seq_len_tensor,
200 - lora_indices_tensor,
201 - slice_start_tensor,
202 - inputs.stride(0),
203 - inputs.stride(1),
204 - inputs.stride(2),
205 - lora_strides_d0_tensor,
206 - lora_strides_d1_tensor,
207 - lora_strides_d2_tensor,
208 - output_tensor.stride(0),
209 - output_tensor.stride(1),
210 - hidden_sizes_tensor,
211 - BLOCK_M,
212 - BLOCK_N,
213 - BLOCK_K,
214 - EVEN_K,
215 - ADD_INPUTS,
216 - CAST_TYPE,
217 - len(lora_b_weights),
218 - same_stride,
219 - )
220 - return
221 -
222 -
223 -def _sgmv_expand_fake(
224 - inputs: torch.Tensor,
225 - lora_b_weights: List[torch.Tensor],
226 - output_tensor: torch.Tensor,
227 - b_seq_start_loc: torch.Tensor,
228 - seq_len_tensor: torch.Tensor,
229 - lora_indices_tensor: torch.Tensor,
230 - batches: int,
231 - max_seq_length: int,
232 - token_nums: int,
233 - offset_start: int = 0,
234 - add_inputs: bool = False,
235 -) -> None:
236 - return
237 -
238 -
239 -try:
240 - direct_register_custom_op(
241 - op_name="sgmv_expand",
242 - op_func=_sgmv_expand,
243 - mutates_args=["output_tensor"],
244 - fake_impl=_sgmv_expand_fake,
245 - )
246 - sgmv_expand = torch.ops.vllm.sgmv_expand
247 -
248 -except AttributeError:
249 - sgmv_expand = _sgmv_expand
vllm/lora/ops/triton_ops/sgmv_shrink.py deleted
-224
@@ -1,224 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -"""
3 -Based on:
4 -Chen, L., Ye, Z., Wu, Y., Zhuo, D., Ceze, L., & Krishnamurthy, A. (2023).
5 -Punica: Multi-Tenant LoRA Serving.
6 -https://arxiv.org/abs/2310.18547
7 -"""
8 -
9 -from typing import List
10 -
11 -import torch
12 -import triton
13 -import triton.language as tl
14 -
15 -from vllm.utils import direct_register_custom_op
16 -
17 -from .kernel_utils import do_shrink_kernel
18 -from .utils import _get_lora_a_ptr
19 -
20 -
21 -@triton.jit
22 -def _sgmv_shrink_kernel(
23 - input_ptr,
24 - lora_ptr, #1-3
25 - out_ptr,
26 - N,
27 - K,
28 - b_seq_start_loc,
29 - seq_lens,
30 - lora_indices,
31 - scaling,
32 - input_d0_stride,
33 - input_d1_stride, # 1
34 - lora_d0_stride,
35 - lora_d1_stride,
36 - lora_d2_stride, # 1
37 - output_d0_stride,
38 - output_d1_stride,
39 - output_d2_stride, # 1
40 - BLOCK_M: tl.constexpr,
41 - BLOCK_N: tl.constexpr,
42 - BLOCK_K: tl.constexpr,
43 - EVEN_K: tl.constexpr,
44 - SPLIT_K: tl.constexpr,
45 - SLICE_NUM: tl.constexpr):
46 - """
47 - The sgmv's shrink triton kernel is based on GroupGEMM+SPLIT-K.
48 - The GEMM of Multi-LoRA can be considered as GroupGEMM. Additionally,
49 - introducing SPLIT-K can improve performance
50 - """
51 - pid = tl.program_id(axis=0)
52 - pid_mix = tl.program_id(axis=1)
53 - cur_batch = tl.program_id(axis=2)
54 - cta_n_num = tl.cdiv(N, BLOCK_N)
55 - pid_m = pid // cta_n_num
56 - pid_n = pid % cta_n_num
57 - if SLICE_NUM == 1:
58 - slice_id: tl.constexpr = 0
59 - pid_sk = tl.program_id(axis=1)
60 - else:
61 - pid_mix = tl.program_id(axis=1)
62 - slice_id = pid_mix // SPLIT_K
63 - pid_sk = pid_mix % SPLIT_K
64 -
65 - M = tl.load(seq_lens + cur_batch)
66 - if pid_m * BLOCK_M >= M:
67 - return
68 - lora_index = tl.load(lora_indices + cur_batch)
69 - if lora_index == -1:
70 - return
71 -
72 - m_offset = tl.load(b_seq_start_loc + cur_batch)
73 -
74 - cta_m_len = min(BLOCK_M, M - (pid_m * BLOCK_M))
75 - cta_m_offset = m_offset + (pid_m * BLOCK_M)
76 - offset_m = tl.arange(0, BLOCK_M)
77 - ram = cta_m_offset + tl.max_contiguous(
78 - tl.multiple_of(offset_m % cta_m_len, BLOCK_M), BLOCK_M)
79 -
80 - do_shrink_kernel(
81 - pid_n,
82 - pid_sk,
83 - slice_id,
84 - lora_index,
85 - input_ptr,
86 - lora_ptr,
87 - out_ptr,
88 - N,
89 - K,
90 - cta_m_len,
91 - ram,
92 - # input strides
93 - input_d0_stride,
94 - input_d1_stride,
95 - # lora strides
96 - lora_d0_stride,
97 - lora_d1_stride,
98 - lora_d2_stride,
99 - # output strides
100 - output_d0_stride,
101 - output_d1_stride,
102 - output_d2_stride,
103 - scaling,
104 - BLOCK_M,
105 - BLOCK_N,
106 - BLOCK_K,
107 - EVEN_K,
108 - SPLIT_K,
109 - SLICE_NUM)
110 -
111 -
112 -@torch.inference_mode()
113 -def _sgmv_shrink(
114 - inputs: torch.Tensor,
115 - lora_a_weights: List[torch.Tensor],
116 - output_tensor: torch.Tensor,
117 - b_seq_start_loc: torch.Tensor,
118 - seq_len_tensor: torch.Tensor,
119 - lora_indices_tensor: torch.Tensor,
120 - batches: int,
121 - max_seq_length: int,
122 - token_nums: int,
123 - scaling: float,
124 -) -> None:
125 - """
126 - Args:
127 - inputs (torch.Tensor): input tensor
128 - lora_a_weights (List[torch.Tensor]): lora'a weight
129 - output_tensor (torch.Tensor): output tensor
130 - b_seq_start_loc (torch.Tensor): (batch_size,). The cumulative
131 - sequence lengths of the sequences in the batch, used to index
132 - into sequence. E.g., if the sequence length is [4, 6], it is
133 - [0, 4].
134 - seq_len_tensor (torch.Tensor): (batch_size,). Record the sequence
135 - length of the sequences in the batch.
136 - lora_indices_tensor (torch.Tensor): (batch_size,). The LoRA index
137 - corresponding to each batch. An index of -1 means no lora should be
138 - applied.
139 - batches (int): batch size
140 - max_seq_length (int): The max sequence lengths of the sequences in the
141 - batch.
142 - token_nums (int): The token numbers in the batch. Used to verify if the
143 - token numbers in the inputs matches the one in the metadata.
144 - scaling (float): Scaling factor.
145 - """
146 - assert inputs.dtype == lora_a_weights[0].dtype
147 - assert inputs.dtype in [torch.float16, torch.bfloat16]
148 - for weight in lora_a_weights:
149 - assert weight.dtype in [torch.float16, torch.bfloat16]
150 -
151 - assert inputs.size(0) == token_nums
152 - assert inputs.size(1) == lora_a_weights[0].size(-1)
153 - assert b_seq_start_loc.size(0) == batches
154 - assert lora_indices_tensor.size(0) == batches
155 - assert inputs.is_contiguous()
156 - assert output_tensor.is_contiguous()
157 - (lora_ptr_tensor, lora_strides_d0, lora_strides_d1,
158 - lora_strides_d2) = _get_lora_a_ptr(lora_a_weights, b_seq_start_loc.device)
159 - # TODO tuning this config
160 - N, K = lora_a_weights[0].shape[-2:] # K=hidden_size,N=rank
161 - BLOCK_M = 32
162 - BLOCK_N = 16
163 - BLOCK_K = 32
164 - SPLIT_K = 8
165 - EVEN_K = K % (BLOCK_K * SPLIT_K) == 0
166 - grid = (
167 - triton.cdiv(max_seq_length, BLOCK_M) * triton.cdiv(N, BLOCK_N),
168 - SPLIT_K * len(lora_a_weights),
169 - batches,
170 - )
171 - _sgmv_shrink_kernel[grid](
172 - inputs,
173 - lora_ptr_tensor,
174 - output_tensor,
175 - N,
176 - K,
177 - b_seq_start_loc,
178 - seq_len_tensor,
179 - lora_indices_tensor,
180 - scaling,
181 - inputs.stride(0),
182 - inputs.stride(1),
183 - lora_strides_d0,
184 - lora_strides_d1,
185 - lora_strides_d2,
186 - output_tensor.stride(0),
187 - output_tensor.stride(1),
188 - output_tensor.stride(2),
189 - BLOCK_M,
190 - BLOCK_N,
191 - BLOCK_K,
192 - EVEN_K,
193 - SPLIT_K,
194 - len(lora_a_weights),
195 - )
196 - return
197 -
198 -
199 -def sgmv_shrink_fake(
200 - inputs: torch.Tensor,
201 - lora_a_weights: List[torch.Tensor],
202 - output_tensor: torch.Tensor,
203 - b_seq_start_loc: torch.Tensor,
204 - seq_len_tensor: torch.Tensor,
205 - lora_indices_tensor: torch.Tensor,
206 - batches: int,
207 - max_seq_length: int,
208 - token_nums: int,
209 - scaling: float,
210 -) -> None:
211 - return
212 -
213 -
214 -try:
215 - direct_register_custom_op(
216 - op_name="sgmv_shrink",
217 - op_func=_sgmv_shrink,
218 - mutates_args=["output_tensor"],
219 - fake_impl=sgmv_shrink_fake,
220 - )
221 - sgmv_shrink = torch.ops.vllm.sgmv_shrink
222 -
223 -except AttributeError:
224 - sgmv_shrink = _sgmv_shrink
vllm/lora/ops/triton_ops/utils.py
-46
@@ -1,55 +1,9 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 -import functools
3 from typing import Dict, List, Tuple
4
5 import torch
6
8 -
9 -@functools.lru_cache
10 -def _get_op_configs(op_type: str, batch: int, hidden_size: int):
11 - # TODO: add optimal configurations
12 - return None
13 -
14 -
15 -def _check_divisibility(hidden_size: int):
16 - # The bgmv_expand kernel requires that the hidden_size be divisible by
17 - # the number below.
18 - divisibility = [2, 4, 8, 16, 32, 64]
19 - divisibility.sort(reverse=True)
20 - for div in divisibility:
21 - if hidden_size % div == 0:
22 - return div
23 - # hidden_size is an odd number
24 - return 1
25 -
26 -
27 -def _get_default_config(op_type: str, batch: int, hidden_size: int):
28 - if op_type == "expand":
29 - return {
30 - "BLOCK_N": 256,
31 - "SPLIT_N": _check_divisibility(hidden_size),
32 - "num_warps": 8
33 - }
34 - else:
35 - return {"BLOCK_K": 256, "SPLIT_K": 64, "num_warps": 8}
36 -
37 -
38 -def get_lora_op_configs(op_type: str, batch: int,
39 - hidden_size: int) -> Dict[str, int]:
40 - """Inspired by `fused_moe_kernel`
41 - The return value will be a dictionary mapping an irregular grid of batch
42 - sizes and hidden_size to configurations of the bgmv-related kernel.
43 - NOTE: It currently only supports the default configuration. We plan to
44 - generate optimal configurations for different hardware in the future using
45 - scripts similar to `benchmark_moe.py`.
46 - """
47 - config = _get_op_configs(op_type, batch, hidden_size)
48 - if not config:
49 - config = _get_default_config(op_type, batch, hidden_size)
50 - return config
51 -
52 -
7 _LORA_A_PTR_DICT: Dict[Tuple[int, ...], Tuple[torch.tensor, ...]] = {}
8 _LORA_B_PTR_DICT: Dict[Tuple[int, ...], Tuple[torch.tensor, ...]] = {}
9
vllm/lora/ops/triton_ops/v1/__init__.py deleted
-11
@@ -1,11 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -
3 -from vllm.lora.ops.triton_ops.v1.v1_expand import v1_expand
4 -from vllm.lora.ops.triton_ops.v1.v1_kernel_metadata import V1KernelMeta
5 -from vllm.lora.ops.triton_ops.v1.v1_shrink import v1_shrink
6 -
7 -__all__ = [
8 - "v1_expand",
9 - "v1_shrink",
10 - "V1KernelMeta",
11 -]
\ No newline at end of file
vllm/lora/punica_wrapper/punica_gpu.py
+59 -225
@@ -10,20 +10,12 @@ from typing import TYPE_CHECKING, List, Optional, Tuple, Union, final
10
11 import torch
12
13 -import vllm.envs as env
13 from vllm.lora.layers import LoRAMapping
14 from vllm.triton_utils import HAS_TRITON
15
16 if HAS_TRITON:
18 - if env.VLLM_USE_V1:
19 - from vllm.lora.ops.triton_ops.v1 import (V1KernelMeta, v1_expand,
20 - v1_shrink)
21 - else:
22 - from vllm.lora.ops.triton_ops import bgmv_expand
23 - from vllm.lora.ops.triton_ops import bgmv_expand_slice
24 - from vllm.lora.ops.triton_ops import bgmv_shrink
25 - from vllm.lora.ops.triton_ops import sgmv_expand
26 - from vllm.lora.ops.triton_ops import sgmv_shrink
17 + from vllm.lora.ops.triton_ops import (LoRAKernelMeta, lora_expand,
18 + lora_shrink)
19
20 from .punica_base import PunicaWrapperBase
21
@@ -32,57 +24,8 @@ if TYPE_CHECKING:
24 from vllm.lora.models import LongContextLoRAContext
25
26
35 -class V1KernelMixin:
36 -
37 - def _v1_make_metadata(self, max_loras: int, max_num_batched_tokens: int,
38 - max_batches: int, device: Union[torch.device, str]):
39 - self.token_mapping_v1_meta = V1KernelMeta.make(max_loras,
40 - max_num_batched_tokens,
41 - device=device)
42 - self.prompt_mapping_v1_meta = V1KernelMeta.make(max_loras,
43 - max_batches,
44 - device=device)
45 -
46 - def _v1_prepare_metadata_tensors(self, token_lora_indices: torch.Tensor,
47 - sampler_indices: torch.Tensor):
48 - self.token_mapping_v1_meta.prepare_tensors(token_lora_indices)
49 - self.prompt_mapping_v1_meta.prepare_tensors(sampler_indices)
50 -
51 - def _v1_apply_shrink(
52 - self,
53 - y: torch.Tensor,
54 - x: torch.Tensor,
55 - w_t_all: Tuple[torch.Tensor, ...],
56 - scale: float,
57 - ):
58 - v1_shrink(
59 - x,
60 - w_t_all,
61 - y,
62 - *self.token_mapping_v1_meta.meta_args(x.size(0)),
63 - scale,
64 - )
65 -
66 - def _v1_apply_expand(
67 - self,
68 - y: torch.Tensor,
69 - x: torch.Tensor,
70 - w_t_all: Tuple[torch.Tensor, ...],
71 - offset_start: int,
72 - add_inputs: bool,
73 - ):
74 - v1_expand(
75 - x,
76 - w_t_all,
77 - y,
78 - *self.token_mapping_v1_meta.meta_args(x.size(0)),
79 - offset_start=offset_start,
80 - add_inputs=add_inputs,
81 - )
82 -
83 -
27 @final
85 -class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):
28 +class PunicaWrapperGPU(PunicaWrapperBase):
29 """
30 PunicaWrapperGPU is designed to manage and provide metadata for the punica
31 kernel. The main function is to maintain the state information for
@@ -96,9 +39,12 @@ class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):
39
40 self.max_loras = kwargs['max_loras']
41
99 - if env.VLLM_USE_V1:
100 - self._v1_make_metadata(self.max_loras, max_num_batched_tokens,
101 - max_batches, device)
42 + self.token_mapping_meta = LoRAKernelMeta.make(self.max_loras,
43 + max_num_batched_tokens,
44 + device=device)
45 + self.prompt_mapping_meta = LoRAKernelMeta.make(self.max_loras,
46 + max_batches,
47 + device=device)
48
49 def update_metadata(
50 self,
@@ -110,83 +56,18 @@ class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):
56 long_lora_context: Optional["LongContextLoRAContext"] = None,
57 **kwargs):
58
113 - if env.VLLM_USE_V1:
114 - self.is_prefill = mapping.is_prefill
115 - self._update_base_metadata(mapping, lora_index_to_id, max_loras,
116 - vocab_size, extra_vocab_size,
117 - long_lora_context)
118 - self._v1_prepare_metadata_tensors(self.token_lora_indices,
119 - self.sampler_indices)
120 - else:
121 - # Forward to base class update_metadata
122 - PunicaWrapperBase.update_metadata(self, mapping, lora_index_to_id,
123 - max_loras, vocab_size,
124 - extra_vocab_size,
125 - long_lora_context, **kwargs)
126 -
127 - def _apply_shrink_prefill(
128 - self,
129 - y: torch.Tensor,
130 - x: torch.Tensor,
131 - w_t_all: Tuple[torch.Tensor, ...],
132 - scale: float,
133 - ):
134 - #No LoRA request, so return directly
135 - if self.no_lora:
136 - return
137 - sgmv_shrink(
138 - x,
139 - w_t_all,
140 - y,
141 - *self.prefill_metadata,
142 - scale,
143 - )
59 + self.is_prefill = mapping.is_prefill
60 + self._update_base_metadata(mapping, lora_index_to_id, max_loras,
61 + vocab_size, extra_vocab_size,
62 + long_lora_context)
63
145 - def _apply_shrink_decode(
146 - self,
147 - y: torch.Tensor,
148 - x: torch.Tensor,
149 - w_t_all: torch.Tensor,
150 - scale: float,
151 - ):
152 - bgmv_shrink(x, w_t_all, y, self.token_lora_indices, scale)
153 -
154 - def _apply_expand_prefill(
155 - self,
156 - y: torch.Tensor,
157 - x: torch.Tensor,
158 - w_t_all: Tuple[torch.Tensor, ...],
159 - offset_start: int,
160 - add_inputs: bool,
161 - ):
162 - #No LoRA request, so return directly
163 - if self.no_lora:
164 - return
165 -
166 - sgmv_expand(
167 - x,
168 - w_t_all,
169 - y,
170 - *self.prefill_metadata,
171 - offset_start=offset_start,
172 - add_inputs=add_inputs,
173 - )
64 + # Prepare cuda kernel metadata tensors
65 + self.token_mapping_meta.prepare_tensors(self.token_lora_indices)
66 + self.prompt_mapping_meta.prepare_tensors(self.sampler_indices)
67
175 - def _apply_expand_decode(
176 - self,
177 - y: torch.Tensor,
178 - x: torch.Tensor,
179 - w_t_all: torch.Tensor,
180 - y_offset: Optional[int],
181 - y_slice_size: Optional[int],
182 - add_inputs: bool,
183 - ):
184 - bgmv_expand_slice(x, w_t_all, y, self.token_lora_indices, y_offset,
185 - y_slice_size, add_inputs)
186 -
187 - def add_shrink(self, y: Union[Tuple[torch.Tensor, ...], torch.Tensor],
188 - x: torch.Tensor, lora_a_stacked: Tuple[torch.Tensor, ...],
189 - scale: float, **kwargs):
68 + def add_shrink(self, y: torch.Tensor, x: torch.Tensor,
69 + lora_a_stacked: Tuple[torch.Tensor,
70 + ...], scale: float, **kwargs):
71 """
72 Performs GEMM for multiple slices of lora_a.
73 When `is_prefill is` true, it indicates that it is currently the
@@ -199,33 +80,24 @@ class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):
80 y[i] += (x @ lora_a_stacked[i]) * scale
81
82 Args:
202 - y (Union[Tuple[torch.Tensor, ...], torch.Tensor]): Output tensors
83 + y (torch.Tensor): Output tensors
84 x (torch.Tensor): Input tensor
85 lora_a_stacked (Tuple[torch.Tensor, ...]): lora_a's weights
86 scale (float): Scaling factor for the operation
87 """
88
89 x = x.view(-1, x.shape[-1])
209 -
210 - if env.VLLM_USE_V1:
211 - self._v1_apply_shrink(y, x, lora_a_stacked, scale) # type: ignore
212 - else:
213 - if self.is_prefill:
214 - # NOTE fused kernel
215 - self._apply_shrink_prefill(
216 - y, # type: ignore
217 - x,
218 - lora_a_stacked,
219 - scale)
220 - else:
221 - # TODO fuse these kernels
222 - for slice_idx in range(len(lora_a_stacked)):
223 - self._apply_shrink_decode(y[slice_idx], x,
224 - lora_a_stacked[slice_idx], scale)
90 + lora_shrink(
91 + x,
92 + lora_a_stacked,
93 + y,
94 + *self.token_mapping_meta.meta_args(x.size(0)),
95 + scale,
96 + )
97
98 def add_expand(self,
99 y: torch.Tensor,
228 - x: Union[Tuple[torch.Tensor, ...], torch.Tensor],
100 + x: torch.Tensor,
101 lora_b_stacked: Tuple[torch.Tensor, ...],
102 lora_bias_stacked: Optional[Tuple[torch.Tensor, ...]],
103 output_slices: Tuple[int, ...],
@@ -244,7 +116,7 @@ class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):
116
117 Args:
118 y (torch.Tensor): Output tensor.
247 - x (Union[Tuple[torch.Tensor, ...], torch.Tensor]): Input tensors
119 + x (torch.Tensor): Input tensors
120 lora_b_stacked (Tuple[torch.Tensor, ...]): lora_b's weight
121 lora_bias_stacked (Optional[Tuple[torch.Tensor, ...]]):
122 bias's weight
@@ -259,37 +131,19 @@ class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):
131 self._apply_bias(token_lora_indices, y, output_slices,
132 lora_bias_stacked)
133
262 - if env.VLLM_USE_V1:
263 - # TODO (varun): Profile with add_inputs = False. i.e. move the
264 - # addition out of the kernel
265 - self._v1_apply_expand(
266 - y,
267 - x, # type: ignore
268 - lora_b_stacked,
269 - offset_start,
270 - add_inputs=True)
271 - else:
272 -
273 - if self.is_prefill:
274 - # NOTE fused kernel
275 - self._apply_expand_prefill(
276 - y,
277 - x, # type: ignore
278 - lora_b_stacked,
279 - offset_start,
280 - add_inputs=True)
281 - else:
282 - # TODO fuse these kernels
283 - for slice_idx in range(len(lora_b_stacked)):
284 - self._apply_expand_decode(
285 - y,
286 - x[slice_idx],
287 - lora_b_stacked[slice_idx],
288 - offset_start,
289 - output_slices[slice_idx],
290 - add_inputs=add_inputs,
291 - )
292 - offset_start += output_slices[slice_idx]
134 + assert x.ndim == 3
135 + assert x.size(0) == len(output_slices)
136 + num_tokens = x.size(1) # first dimension is the num slices
137 +
138 + lora_expand(
139 + x,
140 + lora_b_stacked,
141 + y,
142 + *self.token_mapping_meta.meta_args(num_tokens),
143 + offset_start=offset_start,
144 + add_inputs=True,
145 + )
146 +
147 y = y.view_as(y_org)
148
149 def add_lora_embedding(self,
@@ -311,24 +165,14 @@ class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):
165 add_inputs (bool): Default to True.
166 """
167
314 - if env.VLLM_USE_V1:
315 - self._v1_apply_expand(y,
316 - x.unsqueeze(dim=0), (lora_b_stacked, ),
317 - offset_start=0,
318 - add_inputs=add_inputs)
319 - else:
320 - if self.is_prefill:
321 - sgmv_expand(
322 - x.unsqueeze(dim=0),
323 - (lora_b_stacked, ),
324 - y,
325 - *self.prefill_metadata,
326 - offset_start=0,
327 - add_inputs=add_inputs,
328 - )
329 - else:
330 - bgmv_expand(x, lora_b_stacked, y, self.token_lora_indices,
331 - add_inputs)
168 + lora_expand(
169 + x.unsqueeze(dim=0),
170 + (lora_b_stacked, ),
171 + y,
172 + *self.token_mapping_meta.meta_args(x.size(0)),
173 + offset_start=0,
174 + add_inputs=add_inputs,
175 + )
176
177 def add_lora_linear(self,
178 y: torch.Tensor,
@@ -339,7 +183,7 @@ class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):
183 scale: float,
184 output_slices: Tuple[int, ...],
185 *,
342 - buffer: Optional[Tuple[torch.Tensor, ...]] = None,
186 + buffer: Optional[torch.Tensor] = None,
187 **kwargs) -> None:
188 """
189 Applicable to linear-related lora.
@@ -361,7 +205,7 @@ class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):
205 lora_bias_stacked (Optional[Tuple[torch.Tensor, ...]]): lora's bias.
206 scale (float): Scaling factor.
207 output_slices (Tuple[int, ...]): Every slice's size.
364 - buffer (Optional[Tuple[torch.Tensor, ...]]): Defaults to None.
208 + buffer (Optional[torch.Tensor]): Defaults to None.
209 """
210
211 assert len(lora_a_stacked) == len(lora_b_stacked) == len(output_slices)
@@ -431,21 +275,11 @@ class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):
275 dtype=torch.float32,
276 device=x.device)
277
434 - if env.VLLM_USE_V1:
435 - v1_shrink(x, [lora_a_stacked], buffer.unsqueeze(dim=0),
436 - *self.prompt_mapping_v1_meta.meta_args(x.size(0)), scale)
437 -
438 - v1_expand(buffer.unsqueeze(dim=0), [lora_b_stacked],
439 - y,
440 - *self.prompt_mapping_v1_meta.meta_args(buffer.size(0)),
441 - add_inputs=True)
442 - else:
443 -
444 - # V0 LogitsProcessorWithLoRA always using bgmv.
445 - bgmv_shrink(x, lora_a_stacked, buffer, self.sampler_indices, scale)
446 - bgmv_expand(buffer,
447 - lora_b_stacked,
448 - y,
449 - self.sampler_indices,
450 - add_inputs=True)
278 + lora_shrink(x, [lora_a_stacked], buffer.unsqueeze(dim=0),
279 + *self.prompt_mapping_meta.meta_args(x.size(0)), scale)
280 +
281 + lora_expand(buffer.unsqueeze(dim=0), [lora_b_stacked],
282 + y,
283 + *self.prompt_mapping_meta.meta_args(buffer.size(0)),
284 + add_inputs=True)
285 y = y.view_as(y_org)
vllm/v1/worker/lora_model_runner_mixin.py
+4 -3
@@ -62,9 +62,10 @@ class LoRAModelRunnerMixin:
62 if not self.lora_manager:
63 raise RuntimeError("LoRA is not enabled.")
64
65 - # Set is_prefill to True, so we always use the SGMV kernels.
66 - # For cuda platforms, we have specialized triton kernels, and
67 - # the cuda path ignores `is_prefill`.
65 + # Set is_prefill to True, so we always use the SGMV kernels on
66 + # non-cuda platforms.
67 + # On cuda platforms we use the same kernels for prefill and
68 + # decode and this flag is generally ignored.
69 lora_mapping = LoRAMapping(token_lora_mapping,
70 prompt_lora_mapping,
71 is_prefill=True)