@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5ff0d32580

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, 19:45 UTC

[V1] LoRA - Add triton kernels for V1 (#13096)

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

Varun Sundar Rabindranath committed Mar 10, 2025 at 17:27 UTC 5ff0d32580eb4daadadece5ee39d33f43f230108

11 files changed +1162 -188

benchmarks/kernels/benchmark_lora.py

+141 -24

@@ -23,6 +23,7 @@ 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

25 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

27 from vllm.utils import FlexibleArgumentParser

29 DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())

@@ -171,6 +172,8 @@ class OpType(Enum):

172 SGMV_EXPAND = auto()

173 BGMV_EXPAND = auto()

174 BGMV_EXPAND_SLICE = auto()

175 + V1_SHRINK = auto()

176 + V1_EXPAND = auto()

177

178 @staticmethod

179 def from_str(s: str) -> "OpType":

@@ -184,28 +187,43 @@ class OpType(Enum):

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

194 raise ValueError(f"Unrecognized str {s} to convert to OpType")

195

196 def is_shrink_fn(self) -> bool:

190 - return self in [OpType.SGMV_SHRINK, OpType.BGMV_SHRINK]

197 + return self in [

198 + OpType.SGMV_SHRINK, OpType.BGMV_SHRINK, OpType.V1_SHRINK

199 + ]

200

201 def is_expand_fn(self) -> bool:

193 - return self in [OpType.SGMV_EXPAND, OpType.BGMV_EXPAND]

202 + return self in [

203 + OpType.SGMV_EXPAND, OpType.BGMV_EXPAND, OpType.V1_EXPAND

204 + ]

205

206 def is_prefill_op(self) -> bool:

196 - return self in [OpType.SGMV_SHRINK, OpType.SGMV_EXPAND]

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 [

200 - OpType.BGMV_SHRINK, OpType.BGMV_EXPAND, OpType.BGMV_EXPAND_SLICE

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]

220

221 def num_slices(self) -> list[int]:

207 - if self in [OpType.SGMV_EXPAND, OpType.SGMV_SHRINK]:

208 - # SGMV kernels supports slices

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]

@@ -250,11 +268,13 @@ class OpType(Enum):

268 m, k, n = self.mkn(batch_size, seq_length, hidden_size, lora_rank)

269

270 b_shape = (num_loras, n, k) # col-major

253 - if self == OpType.SGMV_SHRINK:

254 - # SGMV shrink supports num_slices inherently in the kernel

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.

274 return ((m, k), b_shape, (num_slices, m, n))

256 - if self == OpType.SGMV_EXPAND:

257 - # SGMV expand supports num_slices inherently in the kernel

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

278 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,25 +301,30 @@ class OpType(Enum):

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

308 +

309 raise ValueError(f"Unrecognized optype {self}")

310

311 def run_ref_group_gemm(self, output: torch.Tensor, input: torch.Tensor,

312 lora_weights: list[torch.Tensor],

313 **kwargs) -> Callable:

289 - """Each benchmark operation expected the input, lora_weights and outputs

314 + """Each benchmark operation expects the input, lora_weights and outputs

315 in a slightly different format. Refer to self.matmul_shapes().

316 run_ref_group_gemm accounts for those differences in executing a

317 reference group gemm for correctness testing.

318 """

319 w_dtype = lora_weights[0].dtype

320 num_slices = len(lora_weights)

296 - if self == OpType.SGMV_SHRINK:

321 + if self in [OpType.SGMV_SHRINK, OpType.V1_SHRINK]:

322 for slice_idx in range(num_slices):

323 ref_group_gemm(ref_out=output[slice_idx, :],

324 input=input,

325 lora_weights=lora_weights[slice_idx],

326 **kwargs)

302 - if self == OpType.SGMV_EXPAND:

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

@@ -308,19 +333,19 @@ class OpType(Enum):

333 input=input[slice_idx].clone().to(dtype=w_dtype),

334 lora_weights=lora_weights[slice_idx],

335 **kwargs)

311 - if self == OpType.BGMV_SHRINK:

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)

317 - if self == OpType.BGMV_EXPAND:

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)

323 - if self == OpType.BGMV_EXPAND_SLICE:

348 + elif self == OpType.BGMV_EXPAND_SLICE:

349 hidden_size = lora_weights[0].shape[1]

350 for slice_idx in range(num_slices):

351 slice_offset = slice_idx * hidden_size

@@ -329,7 +354,8 @@ class OpType(Enum):

354 input=input[slice_idx].clone().to(dtype=w_dtype),

355 lora_weights=lora_weights[slice_idx],

356 **kwargs)

332 - raise ValueError(f"Unrecognized optype {self}")

357 + else:

358 + raise ValueError(f"Unrecognized optype {self}")

359

360

361 @dataclass

@@ -390,6 +416,8 @@ class BenchmarkTensors:

416 seq_start_loc: torch.Tensor

417 prompt_lora_mapping: torch.Tensor

418 token_lora_mapping: torch.Tensor

419 + # v1 kernel metadata

420 + v1_kernel_meta: Optional[V1KernelMeta] = None

421

422 def io_types(self) -> str:

423 return (f"{dtype_to_str(self.input.dtype)}x"

@@ -432,10 +460,19 @@ class BenchmarkTensors:

460 total_tokens, ctx.batch_size, prompt_lora_indices_tensor,

461 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)

471 +

472 return BenchmarkTensors(input_tensor, lora_weights, output_tensor,

473 seq_len_tensor, seq_start_loc_tensor,

474 prompt_lora_indices_tensor,

438 - token_lora_indices_tensor)

475 + token_lora_indices_tensor, v1_kernel_meta)

476

477 def sanity_check(self) -> None:

478 """

@@ -468,6 +505,13 @@ class BenchmarkTensors:

505 for i in range(len(self.lora_weights_lst)):

506 self.lora_weights_lst[i] = to_device(self.lora_weights_lst[i])

507

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))

514 +

515 def metadata(self) -> tuple[int, int, int]:

516 """

517 Return num_seqs, num_tokens and max_seq_len

@@ -667,6 +711,78 @@ class BenchmarkTensors:

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

716 + self.sanity_check()

717 + self.to_device(self.input.device)

718 +

719 + _, num_tokens, _, num_slices = self.metadata()

720 +

721 + # Sanity check matrix shapes.

722 + i_shape, lw_shape, o_shape = self.input.shape, self.lora_weights_lst[

723 + 0].shape, self.output.shape

724 + # Expected input shape [num_tokens, hidden_size]

725 + assert len(i_shape) == 2

726 + assert i_shape[0] == num_tokens

727 + hidden_size = i_shape[1]

728 + # Expected lora weight shape [num_loras, lora_rank, hidden_size]

729 + assert len(lw_shape) == 3

730 + assert lw_shape[2] == hidden_size

731 + lora_rank = lw_shape[1]

732 + # Expected output shape [num_slices, num_tokens, lora_rank]

733 + assert len(o_shape) == 3

734 + assert o_shape == (num_slices, num_tokens, lora_rank)

735 +

736 + return {

737 + 'inputs': self.input,

738 + 'lora_a_weights': self.lora_weights_lst,

739 + 'output_tensor': self.output,

740 + 'token_lora_mapping': self.v1_kernel_meta.token_lora_mapping,

741 + '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,

746 + 'scaling': 1.0,

747 + }

748 +

749 + def as_v1_expand_kwargs(self, add_inputs: bool) -> dict[str, Any]:

750 + assert self.v1_kernel_meta is not None

751 + self.sanity_check()

752 + self.to_device(self.input.device)

753 +

754 + _, num_tokens, _, num_slices = self.metadata()

755 +

756 + # Sanity check matrix shapes.

757 + i_shape, lw_shape, o_shape = self.input.shape, self.lora_weights_lst[

758 + 0].shape, self.output.shape

759 + # Expected input shape : [num_slices, num_tokens, lora_rank]

760 + assert len(i_shape) == 3

761 + assert i_shape[0] == num_slices

762 + assert i_shape[1] == num_tokens

763 + lora_rank = i_shape[2]

764 + # Expected lora weight shape : [num_lora, hidden_size, lora_rank]

765 + assert len(lw_shape) == 3

766 + assert lw_shape[2] == lora_rank

767 + hidden_size = lw_shape[1]

768 + # Expected output shape : [num_tokens, hidden_size * num_slices]

769 + assert len(o_shape) == 2

770 + assert o_shape == (num_tokens, hidden_size * num_slices)

771 +

772 + return {

773 + 'inputs': self.input,

774 + 'lora_b_weights': self.lora_weights_lst,

775 + 'output_tensor': self.output,

776 + 'token_lora_mapping': self.v1_kernel_meta.token_lora_mapping,

777 + '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,

782 + 'offset_start': 0,

783 + 'add_inputs': add_inputs,

784 + }

785 +

786 def bench_fn_kwargs(self,

787 op_type: OpType,

788 add_inputs: Optional[bool] = None) -> dict[str, Any]:

@@ -685,6 +801,10 @@ class BenchmarkTensors:

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)

808 raise ValueError(f"Unrecognized optype {self}")

809

810 def test_correctness(self, op_type: OpType,

@@ -872,12 +992,9 @@ def run(args: argparse.Namespace, bench_ctxs: list[BenchmarkContext]):

992 timers = []

993 for bench_ctx in bench_ctxs:

994 for seq_len in args.seq_lengths:

875 - bench_ops: list[OpType] = []

876 - if seq_len == 1:

877 - # bench all decode ops

878 - bench_ops = [op for op in args.op_types if op.is_decode_op()]

879 - else:

880 - # bench all prefill ops

995 + 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

1000 seq_len_timers = []

tests/lora/test_layers.py

+50 -18

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

3 +import importlib

4 import random

5 from copy import deepcopy

6 from dataclasses import dataclass

@@ -63,6 +64,36 @@ DEVICES = ([

64 # stages, so we need to verify this. prefill stage(True) or decode stage(False)

65 STAGES = [True, False]

67 +# With the inclusion of V1 tests (look at the run_with_both_engines_lora),

68 +# the tests in this file run twice, once with the V0 engine and then with

69 +# the V1 engine.

70 +# The NUM_RANDOM_SEEDS value was set to 10 before. It is cut to half

71 +# with the inclusion of V1 tests to maintain the CI test times.

72 +NUM_RANDOM_SEEDS = 5

73 +# The VOCAB_PARALLEL_EMBEDDING_TEST_NUM_RANDOM_SEEDS value was set to

74 +# 256 before. It is cut to half with the inclusion of V1 tests to maintain

75 +# the CI test times.

76 +VOCAB_PARALLEL_EMBEDDING_TEST_NUM_RANDOM_SEEDS = 128

77 +

78 +

79 +@pytest.fixture(autouse=True)

80 +def v1(run_with_both_engines_lora):

81 + # Simple autouse wrapper to run both engines for each test

82 + # This can be promoted up to conftest.py to run for every

83 + # test in a package

84 +

85 + # Reload punica_gpu as the kernels used are tied to engine type.

86 + from vllm.lora.punica_wrapper import punica_gpu

87 + importlib.reload(punica_gpu)

88 +

89 + # Release any memory we might be holding on to. CI runs OOMs otherwise.

90 + from vllm.lora.ops.triton_ops.utils import (_LORA_A_PTR_DICT,

91 + _LORA_B_PTR_DICT)

92 + _LORA_B_PTR_DICT.clear()

93 + _LORA_A_PTR_DICT.clear()

94 +

95 + yield

96 +

98 def get_random_id_to_index(num_loras: int,

99 num_slots: int,

@@ -226,7 +257,7 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:

257

258 torch.set_default_device(device)

259 max_loras = 8

229 - punica_wrapper = get_punica_wrapper(8192, 256, device)

260 + punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)

261 assert check_punica_wrapper(punica_wrapper)

262 lora_config = LoRAConfig(max_loras=max_loras,

263 max_lora_rank=8,

@@ -241,7 +272,7 @@ def test_embeddings(dist_init, num_loras, device, vocab_size, stage) -> None:

272

273 return embedding, lora_embedding

274

244 - for i in range(10):

275 + for i in range(NUM_RANDOM_SEEDS):

276 set_random_seed(i)

277

278 id_to_index = get_random_id_to_index(num_loras, max_loras)

@@ -329,7 +360,7 @@ def test_embeddings_with_new_embeddings(dist_init, num_loras, device,

360

361 torch.set_default_device(device)

362 max_loras = 8

332 - punica_wrapper = get_punica_wrapper(8192, 256, device)

363 + punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)

364 assert check_punica_wrapper(punica_wrapper)

365 lora_config = LoRAConfig(max_loras=max_loras,

366 max_lora_rank=8,

@@ -353,7 +384,7 @@ def test_embeddings_with_new_embeddings(dist_init, num_loras, device,

384

385 return expanded_embedding, lora_embedding

386

356 - for i in range(10):

387 + for i in range(NUM_RANDOM_SEEDS):

388 set_random_seed(i)

389

390 id_to_index = get_random_id_to_index(num_loras, max_loras)

@@ -468,7 +499,7 @@ def test_lm_head_logits_processor(dist_init, num_loras, device, vocab_size,

499

500 torch.set_default_device(device)

501 max_loras = 8

471 - punica_wrapper = get_punica_wrapper(8192, 256, device)

502 + punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)

503 assert check_punica_wrapper(punica_wrapper)

504 lora_config = LoRAConfig(max_loras=max_loras,

505 max_lora_rank=8,

@@ -490,7 +521,7 @@ def test_lm_head_logits_processor(dist_init, num_loras, device, vocab_size,

521

522 return linear, logits_processor, lora_logits_processor

523

493 - for i in range(10):

524 + for i in range(NUM_RANDOM_SEEDS):

525 set_random_seed(i)

526

527 id_to_index = get_random_id_to_index(num_loras, max_loras)

@@ -600,10 +631,10 @@ def test_linear_replicated(dist_init, num_loras, device, stage,

631 if current_platform.is_cuda_alike():

632 torch.cuda.set_device(device)

633

634 + max_loras = 8

635 torch.set_default_device(device)

604 - punica_wrapper = get_punica_wrapper(8192, 256, device)

636 + punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)

637 assert check_punica_wrapper(punica_wrapper)

606 - max_loras = 8

638 lora_config = LoRAConfig(max_loras=max_loras,

639 max_lora_rank=8,

640 lora_dtype=torch.float16,

@@ -627,7 +658,7 @@ def test_linear_replicated(dist_init, num_loras, device, stage,

658 assert lora_linear.lora_bias_stacked is None

659 return linear, lora_linear

660

630 - for i in range(10):

661 + for i in range(NUM_RANDOM_SEEDS):

662 set_random_seed(i)

663

664 id_to_index = get_random_id_to_index(num_loras, max_loras)

@@ -716,10 +747,10 @@ def test_linear_parallel(dist_init, num_loras, orientation, fully_shard,

747 if current_platform.is_cuda_alike():

748 torch.cuda.set_device(device)

749

750 + max_loras = 8

751 torch.set_default_device(device)

720 - punica_wrapper = get_punica_wrapper(8192, 256, device)

752 + punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)

753 assert check_punica_wrapper(punica_wrapper)

722 - max_loras = 8

754 lora_config = LoRAConfig(max_loras=max_loras,

755 max_lora_rank=8,

756 fully_sharded_loras=fully_shard,

@@ -753,7 +784,7 @@ def test_linear_parallel(dist_init, num_loras, orientation, fully_shard,

784 assert lora_linear.lora_bias_stacked is None

785 return linear, lora_linear

786

756 - for i in range(10):

787 + for i in range(NUM_RANDOM_SEEDS):

788 set_random_seed(i)

789

790 id_to_index = get_random_id_to_index(num_loras, max_loras)

@@ -842,10 +873,10 @@ def test_column_parallel_packed(dist_init, num_loras, repeats, fully_shard,

873 if current_platform.is_cuda_alike():

874 torch.cuda.set_device(device)

875

876 + max_loras = 8

877 torch.set_default_device(device)

846 - punica_wrapper = get_punica_wrapper(8192, 256, device)

878 + punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)

879 assert check_punica_wrapper(punica_wrapper)

848 - max_loras = 8

880 lora_config = LoRAConfig(max_loras=max_loras,

881 max_lora_rank=8,

882 fully_sharded_loras=fully_shard,

@@ -900,7 +931,7 @@ def test_column_parallel_packed(dist_init, num_loras, repeats, fully_shard,

931 assert lora_linear.lora_bias_stacked is None

932 return linear, lora_linear

933

903 - for i in range(10):

934 + for i in range(NUM_RANDOM_SEEDS):

935 set_random_seed(i)

936

937 id_to_index = get_random_id_to_index(num_loras, max_loras)

@@ -1002,12 +1033,12 @@ def test_rotary_embedding_long_context(dist_init, num_loras, device,

1033 is_neox_style, rotary_dim, head_size,

1034 seq_len) -> None:

1035 dtype = torch.float16

1036 + max_loras = 8

1037 seed = 0

1038 current_platform.seed_everything(seed)

1039 torch.set_default_device(device)

1008 - punica_wrapper = get_punica_wrapper(8192, 256, device)

1040 + punica_wrapper = get_punica_wrapper(8192, 256, device, max_loras=max_loras)

1041 assert check_punica_wrapper(punica_wrapper)

1010 - max_loras = 8

1042 lora_config = LoRAConfig(max_loras=max_loras,

1043 max_lora_rank=8,

1044 long_lora_scaling_factors=scaling_factors,

@@ -1083,7 +1114,8 @@ def test_rotary_embedding_long_context(dist_init, num_loras, device,

1114

1115

1116 @pytest.mark.parametrize("tp_size", [1, 2, 4, 8])

1086 -@pytest.mark.parametrize("seed", list(range(256)))

1117 +@pytest.mark.parametrize(

1118 + "seed", list(range(VOCAB_PARALLEL_EMBEDDING_TEST_NUM_RANDOM_SEEDS)))

1119 def test_vocab_parallel_embedding_indices(tp_size, seed):

1120 random.seed(seed)

1121 vocab_size = random.randint(4000, 64000)

tests/lora/test_punica_ops.py

+125 -80

@@ -5,10 +5,12 @@ import pytest

5 import torch

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)

12 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

14 from vllm.platforms import current_platform

16 from .utils import (PunicaTensors, assert_close, generate_data,

@@ -91,12 +93,12 @@ def sgmv_expand_for_nslices(nslices: int, hidden_size: int,

93 _dict_lock = Lock()

94 -def check_sgmv_shrink(batches: int, num_loras: int, rank: int,

95 - hidden_size: int, nslices: int, dtype: torch.dtype,

96 - device: str, seq_length: int, scaling: float):

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):

99 """

98 - Compare outputs of vllm.sgmv_shrink kernel against a reference

99 - implementation.

100 + Compare outputs of vllm.sgmv_shrink and vllm.v1_shrink kernel against a

101 + reference implementation.

102 """

103 data: PunicaTensors = generate_data_for_nslices(

104 batches,

@@ -111,44 +113,63 @@ def check_sgmv_shrink(batches: int, num_loras: int, rank: int,

113 )

114 max_seq_length, token_nums = data.meta()

115

116 + # Setup metadata information for SGMV and reference kernels

117 + sgmv_meta_args = (data.b_seq_start_loc, data.seq_len_tensor,

118 + data.prompt_lora_mapping, batches, max_seq_length,

119 + token_nums)

120 +

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)

126 +

127 + ref_out_tensor = data.ref_out_tensor

128 + sgmv_out_tensor = data.our_out_tensor

129 + v1_out_tensor = data.our_out_tensor.clone()

130 +

131 # Preventing cache error pointer.

132 with _dict_lock:

133 + # SGMV shrink kernel

134 _LORA_A_PTR_DICT.clear()

135 torch.ops.vllm.sgmv_shrink(

136 data.inputs_tensor,

137 data.lora_weights,

120 - data.our_out_tensor,

121 - data.b_seq_start_loc,

122 - data.seq_len_tensor,

123 - data.prompt_lora_mapping,

124 - batches,

125 - max_seq_length,

126 - token_nums,

138 + sgmv_out_tensor,

139 + *sgmv_meta_args,

140 scaling,

141 )

142

130 - sgmv_shrink_for_nslices(

131 - nslices,

143 + # V1 shrink kernel

144 + _LORA_A_PTR_DICT.clear()

145 + torch.ops.vllm.v1_shrink(

146 data.inputs_tensor,

147 data.lora_weights,

134 - data.ref_out_tensor,

135 - data.b_seq_start_loc,

136 - data.seq_len_tensor,

137 - data.prompt_lora_mapping,

138 - batches,

139 - max_seq_length,

140 - token_nums,

148 + v1_out_tensor,

149 + *v1_meta.meta_args(token_nums=token_nums),

150 scaling,

151 )

143 - assert_close(data.our_out_tensor, data.ref_out_tensor)

152 +

153 + # Reference

154 + sgmv_shrink_for_nslices(

155 + nslices,

156 + data.inputs_tensor,

157 + data.lora_weights,

158 + ref_out_tensor,

159 + *sgmv_meta_args,

160 + scaling,

161 + )

162 +

163 + assert_close(sgmv_out_tensor, ref_out_tensor)

164 + assert_close(v1_out_tensor, ref_out_tensor)

165

166

146 -def check_sgmv_expand(batches: int, num_loras: int, rank: int,

147 - hidden_size: int, nslices: int, dtype: torch.dtype,

148 - device: str, seq_length: int, add_inputs: bool):

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):

170 """

150 - Compare outputs of vllm.sgmv_expand kernel against a reference

151 - implementation.

171 + Compare outputs of vllm.sgmv_expand and vllm.v1_expand kernels against a

172 + reference implementation.

173 """

174 data: PunicaTensors = generate_data_for_nslices(

175 batches,

@@ -164,36 +185,54 @@ def check_sgmv_expand(batches: int, num_loras: int, rank: int,

185

186 max_seq_length, token_nums = data.meta()

187

188 + # Setup metadata information for SGMV and reference kernels

189 + sgmv_meta_args = (data.b_seq_start_loc, data.seq_len_tensor,

190 + data.prompt_lora_mapping, batches, max_seq_length,

191 + token_nums)

192 +

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)

198 +

199 + # Setup output tensors

200 + ref_out_tensor = data.ref_out_tensor

201 + sgmv_out_tensor = data.our_out_tensor

202 + v1_out_tensor = data.our_out_tensor.clone()

203 +

204 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,

172 - data.our_out_tensor,

173 - data.b_seq_start_loc,

174 - data.seq_len_tensor,

175 - data.prompt_lora_mapping,

176 - batches,

177 - max_seq_length,

178 - token_nums,

210 + sgmv_out_tensor,

211 + *sgmv_meta_args,

212 offset_start=0,

213 add_inputs=add_inputs,

214 )

215

216 + # V1 expand kernel

217 + _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)

224 +

225 + # Reference

226 sgmv_expand_for_nslices(nslices,

227 hidden_size,

228 data.inputs_tensor,

229 data.lora_weights,

187 - data.ref_out_tensor,

188 - data.b_seq_start_loc,

189 - data.seq_len_tensor,

190 - data.prompt_lora_mapping,

191 - batches,

192 - max_seq_length,

193 - token_nums,

230 + ref_out_tensor,

231 + *sgmv_meta_args,

232 add_inputs=add_inputs)

233

196 - assert_close(data.our_out_tensor, data.ref_out_tensor)

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,

@@ -439,7 +478,7 @@ SEED = [0]

478 @pytest.mark.parametrize("device", DEVICES)

479 @pytest.mark.parametrize("seed", SEED)

480 @pytest.mark.parametrize("op_type", ["shrink", "expand"])

442 -def test_punica_sgmv(

481 +def test_kernels(

482 batches: int,

483 num_loras: int,

484 rank: int,

@@ -450,29 +489,32 @@ def test_punica_sgmv(

489 seed: int,

490 op_type: str,

491 ):

492 + """

493 + Tests SGMV and V1 kernels.

494 + """

495 torch.set_default_device(device)

496 current_platform.seed_everything(seed)

497

498 if op_type == "shrink":

457 - check_sgmv_shrink(batches=batches,

458 - num_loras=num_loras,

459 - rank=rank,

460 - hidden_size=hidden_size,

461 - nslices=nslices,

462 - dtype=dtype,

463 - device=device,

464 - seq_length=128,

465 - scaling=0.5)

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)

508 else:

467 - check_sgmv_expand(batches=batches,

468 - num_loras=num_loras,

469 - rank=rank,

470 - hidden_size=hidden_size,

471 - nslices=nslices,

472 - dtype=dtype,

473 - device=device,

474 - seq_length=128,

475 - add_inputs=True)

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)

518

519

520 @pytest.mark.parametrize("batches", hs_test_params['batches'])

@@ -484,7 +526,7 @@ def test_punica_sgmv(

526 @pytest.mark.parametrize("device", DEVICES)

527 @pytest.mark.parametrize("seed", SEED)

528 @pytest.mark.parametrize("op_type", ["shrink", "expand"])

487 -def test_punica_sgmv_hidden_size(

529 +def test_kernels_hidden_size(

530 batches: int,

531 num_loras: int,

532 rank: int,

@@ -495,29 +537,32 @@ def test_punica_sgmv_hidden_size(

537 seed: int,

538 op_type: str,

539 ):

540 + """

541 + Tests SGMV and V1 kernels.

542 + """

543 torch.set_default_device(device)

544 current_platform.seed_everything(seed)

545

546 if op_type == "shrink":

502 - check_sgmv_shrink(batches=batches,

503 - num_loras=num_loras,

504 - rank=rank,

505 - hidden_size=hidden_size,

506 - nslices=nslices,

507 - dtype=dtype,

508 - device=device,

509 - seq_length=128,

510 - scaling=0.5)

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:

512 - check_sgmv_expand(batches=batches,

513 - num_loras=num_loras,

514 - rank=rank,

515 - hidden_size=hidden_size,

516 - nslices=nslices,

517 - dtype=dtype,

518 - device=device,

519 - seq_length=128,

520 - add_inputs=True)

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'])

vllm/lora/models.py

+5 -3

@@ -326,9 +326,11 @@ class LoRAModelManager(AdapterModelManager):

326 self.lora_index_to_id: List[Optional[int]] = [None] * self.lora_slots

327 self.vocab_size = vocab_size

328 self.long_lora_context: Optional[LongContextLoRAContext] = None

329 - self.punica_wrapper = get_punica_wrapper(max_num_batched_tokens,

330 - max_batches=self.max_num_seqs,

331 - device=self.device)

329 + self.punica_wrapper = get_punica_wrapper(

330 + max_num_batched_tokens,

331 + max_batches=self.max_num_seqs,

332 + device=self.device,

333 + max_loras=self.lora_config.max_loras)

334 # Scaling factor -> offset to the sin_cos_cache to it.

335 # Used for long context lora.

336 self.scaling_factor_to_offset: Dict[float, int] = {}

vllm/lora/ops/triton_ops/utils.py

+2 -2

@@ -54,7 +54,7 @@ _LORA_A_PTR_DICT: Dict[Tuple[int, ...], Tuple[torch.tensor, ...]] = {}

54 _LORA_B_PTR_DICT: Dict[Tuple[int, ...], Tuple[torch.tensor, ...]] = {}

57 -def _get_lora_a_ptr(lora_a_weights: List[torch.Tensor], device: str):

57 +def _get_lora_a_ptr(lora_a_weights: List[torch.Tensor], device: torch.device):

58 """

59 `_LORA_A_PTR_DICT` collects the required information during `profile_run`,

60 After this, it remains constant and subsequent usage is through LUT.

@@ -100,7 +100,7 @@ def _get_lora_a_ptr(lora_a_weights: List[torch.Tensor], device: str):

100

101

102 def _get_lora_b_ptr(lora_weights: List[torch.Tensor], offset_start: int,

103 - device: str):

103 + device: torch.device):

104 """

105 `_LORA_B_PTR_DICT` collects the required information during `profile_run`,

106 After this, it remains constant and subsequent usage is through LUT.

vllm/lora/ops/triton_ops/v1/__init__.py new

+11

@@ -0,0 +1,11 @@

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/ops/triton_ops/v1/v1_expand.py new

+282

@@ -0,0 +1,282 @@

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.lora.ops.triton_ops.kernel_utils import do_expand_kernel

16 +from vllm.lora.ops.triton_ops.utils import _get_lora_b_ptr

17 +from vllm.utils import direct_register_custom_op

18 +

19 +

20 +@triton.jit

21 +def _v1_expand_kernel(

22 + input_ptr,

23 + lora_ptr,

24 + out_ptr,

25 + M,

26 + N,

27 + K,

28 + token_indices_sorted_by_lora_ids,

29 + num_tokens_per_lora,

30 + lora_token_start_loc,

31 + lora_ids,

32 + slice_start_loc,

33 + input_d0_stride,

34 + input_d1_stride,

35 + input_d2_stride, # 1

36 + ls_d0_ptr,

37 + ls_d1_ptr,

38 + ls_d2_ptr, # 1

39 + output_d0_stride,

40 + output_d1_stride, # 1

41 + output_hs_ptr,

42 + BLOCK_M: tl.constexpr,

43 + BLOCK_N: tl.constexpr,

44 + BLOCK_K: tl.constexpr,

45 + EVEN_K: tl.constexpr,

46 + ADD_INPUTS: tl.constexpr,

47 + CAST_TYPE: tl.constexpr,

48 + SLICE_NUM: tl.constexpr,

49 + SAME_STRIDE: tl.constexpr):

50 +

51 + cta_n_num = tl.cdiv(N, BLOCK_N)

52 + cta_m_num = tl.cdiv(M, BLOCK_M)

53 +

54 + pid_mn = tl.program_id(axis=0)

55 + pid_m = pid_mn % cta_m_num

56 + pid_n = (pid_mn // cta_m_num) % cta_n_num

57 +

58 + slice_id = tl.program_id(axis=1)

59 + lora_idx = tl.program_id(axis=2)

60 +

61 + lora_id = tl.load(lora_ids + lora_idx)

62 + if lora_id == -1:

63 + # Early exit for the no-lora case.

64 + return

65 +

66 + lora_m_size = tl.load(num_tokens_per_lora + lora_idx)

67 +

68 + cta_m_offset = pid_m * BLOCK_M

69 + if cta_m_offset >= lora_m_size:

70 + # Early exit CTA.

71 + return

72 +

73 + # When the output dimensions of each slice are the same,cur_n=N, otherwise

74 + # cur_n=tl.load(output_hs_ptr + slice_id), this situation exists in GQA's

75 + # qkv linear.

76 + curr_N = N if SAME_STRIDE else tl.load(output_hs_ptr + slice_id)

77 + if pid_n * BLOCK_N >= curr_N:

78 + # Early exit CTA.

79 + return

80 +

81 + # num rows this CTA should process.

82 + cta_m_len = min(BLOCK_M, lora_m_size - cta_m_offset)

83 +

84 + # Identify all rows that this CTA should process.

85 + lora_m_indices_start = tl.load(lora_token_start_loc + lora_idx)

86 + cta_lora_seq_indices = (token_indices_sorted_by_lora_ids +

87 + lora_m_indices_start + cta_m_offset)

88 +

89 + # Load all relevant row indices.

90 + offset_m = tl.arange(0, BLOCK_M) % cta_m_len

91 + ram = tl.load(cta_lora_seq_indices + offset_m)

92 +

93 + do_expand_kernel(

94 + pid_n,

95 + lora_id,

96 + slice_id,

97 + input_ptr,

98 + lora_ptr,

99 + out_ptr,

100 + curr_N,

101 + K,

102 + cta_m_len,

103 + ram, # array identifying the rows of Input ptr to operate on

104 + slice_start_loc,

105 + # input ptr strides

106 + input_d0_stride,

107 + input_d1_stride,

108 + input_d2_stride,

109 + # lora ptr strides

110 + ls_d0_ptr,

111 + ls_d1_ptr,

112 + ls_d2_ptr,

113 + # out ptr strides

114 + output_d0_stride,

115 + output_d1_stride,

116 + # constants

117 + BLOCK_M,

118 + BLOCK_N,

119 + BLOCK_K,

120 + SAME_STRIDE,

121 + SLICE_NUM,

122 + EVEN_K,

123 + CAST_TYPE,

124 + ADD_INPUTS)

125 +

126 +

127 +@torch.inference_mode()

128 +def _v1_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]

132 + output_tensor: torch.

133 + Tensor, # shape [num_tokens, hidden_size * num_slices]

134 + token_lora_mapping: torch.Tensor, # shape [num_tokens]

135 + token_indices_sorted_by_lora_ids: torch.Tensor, # shape [num_tokens]

136 + num_tokens_per_lora: torch.Tensor, # shape [max-loras + 1]

137 + lora_token_start_loc: torch.Tensor, # shape [max-loras + 2]

138 + lora_ids: torch.Tensor, # shape [max-loras + 1]

139 + offset_start: int = 0,

140 + add_inputs: bool = False,

141 +) -> None:

142 + """

143 + Args:

144 + inputs (torch.Tensor): input tensor

145 + lora_b_weights (List[torch.Tensor]): lora'b weight

146 + output_tensor (torch.Tensor): output tensor

147 + token_lora_mapping (torch.Tensor): A tensor mapping each input token

148 + to the lora-id related to that token. A value of -1 indicates that

149 + LoRA doesn't apply to that token.

150 + token_indices_sorted_by_lora_ids (torch.Tensor): Row/Token indices from

151 + the A matrix grouped by LoRA IDs.

152 + num_tokens_per_lora (torch.Tensor): num_tokens_per_lora[i] is the number

153 + of tokens that are to be processed by LoRA ID lora_ids[i]

154 + lora_token_start_loc (torch.Tensor): A cumulative sum of

155 + num_tokens_per_lora. lora_token_start_loc[0] is always 0 so that

156 + lora_token_start_loc[i], along with num_tokens_per_lora[i]

157 + identifies the the region in token_indices_sorted_by_lora_ids that

158 + LoRA lora_ids[i] should process.

159 + lora_ids (torch.Tensor): LoRA ids to process.

160 + offset_start (int, optional): Offset start for output_tensor.

161 + Defaults to 0.

162 + add_inputs (bool, optional): Whether to add the input tensor to the

163 + output tensor. Defaults to False.

164 + """

165 + assert inputs.dtype in [torch.float16, torch.bfloat16, torch.float32]

166 + for weight in lora_b_weights:

167 + assert weight.dtype in [torch.float16, torch.bfloat16]

168 +

169 + assert inputs.size(0) == len(lora_b_weights)

170 + assert output_tensor.is_contiguous()

171 +

172 + # metadata sanity check.

173 + assert token_lora_mapping.size(0) == token_indices_sorted_by_lora_ids.size(

174 + 0)

175 + assert lora_ids.size(0) == num_tokens_per_lora.size(0)

176 + assert lora_token_start_loc.size(0) == lora_ids.size(0) + 1

177 +

178 + (slice_start_tensor, lora_ptr_tensor, lora_strides_d0_tensor,

179 + lora_strides_d1_tensor, lora_strides_d2_tensor, hidden_sizes_tensor,

180 + same_stride, MAX_N) = _get_lora_b_ptr(lora_b_weights, offset_start,

181 + inputs.device)

182 +

183 + K = lora_b_weights[0].shape[-1] # K= rank

184 + M = inputs.size(1)

185 + ADD_INPUTS = add_inputs

186 + MAX_LORAS = lora_ids.size(0)

187 + CAST_TYPE = False

188 + NUM_SLICES = len(lora_b_weights)

189 +

190 + # Triton kernel configs.

191 + BLOCK_M = 64

192 + BLOCK_N = 128

193 + BLOCK_K = 16

194 + NUM_WARPS = 4

195 + NUM_CTAS = 1

196 + NUM_STAGES = 2

197 + MAX_NREG = None

198 +

199 + EVEN_K = K % BLOCK_K == 0 # type: ignore

200 +

201 + if inputs.dtype == torch.float32 and lora_b_weights[0].dtype in [

202 + torch.float16,

203 + torch.bfloat16,

204 + ]:

205 + CAST_TYPE = True

206 +

207 + # TODO (varun): This grid formulation maximizes parallelization at the

208 + # cost of wasteful thread block launch when only a few input tokens require

209 + # LoRA. This might not be the best in all cases.

210 + grid = (

211 + triton.cdiv(M, BLOCK_M) * triton.cdiv(MAX_N, BLOCK_N),

212 + NUM_SLICES,

213 + # Each LoRA receives its own set of thread blocks for output

214 + # computation. If some LoRA doesn't have any tokens to process, its

215 + # thread blocks simply exit.

216 + MAX_LORAS,

217 + )

218 +

219 + _v1_expand_kernel[grid](

220 + inputs,

221 + lora_ptr_tensor,

222 + output_tensor,

223 + M,

224 + MAX_N,

225 + K,

226 + token_indices_sorted_by_lora_ids,

227 + num_tokens_per_lora,

228 + lora_token_start_loc,

229 + lora_ids,

230 + slice_start_tensor,

231 + inputs.stride(0),

232 + inputs.stride(1),

233 + inputs.stride(2),

234 + lora_strides_d0_tensor,

235 + lora_strides_d1_tensor,

236 + lora_strides_d2_tensor,

237 + output_tensor.stride(0),

238 + output_tensor.stride(1),

239 + hidden_sizes_tensor,

240 + BLOCK_M,

241 + BLOCK_N,

242 + BLOCK_K,

243 + EVEN_K,

244 + ADD_INPUTS,

245 + CAST_TYPE,

246 + NUM_SLICES,

247 + same_stride,

248 + num_warps=NUM_WARPS,

249 + num_ctas=NUM_CTAS,

250 + num_stages=NUM_STAGES,

251 + maxnreg=MAX_NREG,

252 + )

253 +

254 + return

255 +

256 +

257 +def _v1_expand_fake(

258 + inputs: torch.Tensor,

259 + lora_b_weights: List[torch.Tensor],

260 + output_tensor: torch.Tensor,

261 + token_lora_mapping: torch.Tensor,

262 + token_indices_sorted_by_lora_ids: torch.Tensor,

263 + num_tokens_per_lora: torch.Tensor,

264 + lora_token_start_loc: torch.Tensor,

265 + lora_ids: torch.Tensor,

266 + offset_start: int = 0,

267 + add_inputs: bool = False,

268 +) -> None:

269 + return

270 +

271 +

272 +try:

273 + direct_register_custom_op(

274 + op_name="v1_expand",

275 + op_func=_v1_expand,

276 + mutates_args=["output_tensor"],

277 + fake_impl=_v1_expand_fake,

278 + )

279 + v1_expand = torch.ops.vllm.v1_expand

280 +

281 +except AttributeError:

282 + v1_expand = _v1_expand

vllm/lora/ops/triton_ops/v1/v1_kernel_metadata.py new

+117

@@ -0,0 +1,117 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +"""

3 +V1 LoRA kernels metadata preparation utilities.

4 +"""

5 +

6 +from dataclasses import dataclass

7 +from typing import Tuple, Union

8 +

9 +import torch

10 +

11 +

12 +@dataclass

13 +class V1KernelMeta:

14 + token_lora_mapping: torch.Tensor

15 + token_indices_sorted_by_lora_ids: torch.Tensor

16 + active_lora_ids: torch.Tensor

17 + num_tokens_per_lora: torch.Tensor

18 + lora_token_start_loc: torch.Tensor

19 +

20 + @staticmethod

21 + def make(max_loras: int, max_num_tokens: int,

22 + device: Union[torch.device, str]) -> "V1KernelMeta":

23 +

24 + token_lora_mapping = torch.empty(max_num_tokens,

25 + dtype=torch.int32,

26 + device=device)

27 +

28 + token_indices_sorted_by_lora_ids = torch.empty(max_num_tokens,

29 + dtype=torch.int32,

30 + device=device)

31 +

32 + # +1 because "no-lora" is also a possibility

33 + # example: let max_loras be 3, active_lora_ids of [-1, 0, 2, 1]

34 + # is a possibility.

35 + active_lora_ids = torch.empty(max_loras + 1,

36 + dtype=torch.int32,

37 + device=device)

38 +

39 + # using running example, [3, 10, 5, 2] is a possibility.

40 + num_tokens_per_lora = torch.zeros(max_loras + 1,

41 + dtype=torch.int32,

42 + device=device)

43 +

44 + # +2 for this because, the first index is always 0.

45 + # using running example, lora_token_start_loc

46 + # is [0, 3, 13, 18, 20].

47 + lora_token_start_loc = torch.zeros(max_loras + 2,

48 + dtype=torch.int32,

49 + device=device)

50 + return V1KernelMeta(

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,

54 + num_tokens_per_lora=num_tokens_per_lora,

55 + lora_token_start_loc=lora_token_start_loc)

56 +

57 + def _reset(self):

58 + self.active_lora_ids.fill_(-1)

59 + self.num_tokens_per_lora.fill_(0)

60 + self.lora_token_start_loc.fill_(0)

61 +

62 + def prepare_tensors(self, token_lora_mapping: torch.Tensor) -> None:

63 + """

64 + Prepare kernel metadata tensors for the current forward pass.

65 +

66 + Args:

67 + token_lora_tensor (torch.Tensor): Tensor containing lora indices

68 + for each input token.

69 + """

70 +

71 + self._reset()

72 +

73 + num_tokens = token_lora_mapping.size(0)

74 +

75 + # copy token lora mapping

76 + self.token_lora_mapping[:num_tokens].copy_(token_lora_mapping,

77 + non_blocking=True)

78 +

79 + # token_indices_sorted_by_lora_ids

80 + _, token_indices_sorted_by_lora_ids = torch.sort(token_lora_mapping,

81 + stable=True)

82 + # start gpu transfer

83 + self.token_indices_sorted_by_lora_ids[:num_tokens].copy_(

84 + token_indices_sorted_by_lora_ids, non_blocking=True)

85 +

86 + # active_lora_ids, num_tokens_per_lora

87 + lora_ids, num_tokens_per_lora = torch.unique(token_lora_mapping,

88 + sorted=False,

89 + return_counts=True)

90 + self.active_lora_ids[:lora_ids.size(0)].copy_(lora_ids,

91 + non_blocking=True)

92 + self.num_tokens_per_lora[:num_tokens_per_lora.size(0)].copy_(

93 + num_tokens_per_lora, non_blocking=True)

94 +

95 + # lora_token_start_loc

96 + lora_token_start_loc = torch.cumsum(num_tokens_per_lora, dim=0)

97 + self.lora_token_start_loc[1:1 + lora_token_start_loc.size(0)].copy_(

98 + lora_token_start_loc, non_blocking=True)

99 +

100 + def meta_args(

101 + self, token_nums: int

102 + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor,

103 + torch.Tensor]:

104 + """

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.

109 +

110 + Args:

111 + token_nums (int): Number of input tokens in the current forward

112 + pass.

113 + """

114 + return (self.token_lora_mapping[:token_nums],

115 + self.token_indices_sorted_by_lora_ids[:token_nums],

116 + self.num_tokens_per_lora, self.lora_token_start_loc,

117 + self.active_lora_ids)

vllm/lora/ops/triton_ops/v1/v1_shrink.py new

+236

@@ -0,0 +1,236 @@

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.lora.ops.triton_ops.kernel_utils import do_shrink_kernel

16 +from vllm.lora.ops.triton_ops.utils import _get_lora_a_ptr

17 +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):

30 +

31 + cta_n_num = tl.cdiv(N, BLOCK_N)

32 + cta_m_num = tl.cdiv(M, BLOCK_M)

33 +

34 + pid_sk_m_n = tl.program_id(axis=0)

35 + pid_sk = pid_sk_m_n % SPLIT_K

36 + pid_m = (pid_sk_m_n // SPLIT_K) % cta_m_num

37 + pid_n = pid_sk_m_n // (SPLIT_K * cta_m_num) % cta_n_num

38 +

39 + slice_id = tl.program_id(axis=1)

40 + lora_idx = tl.program_id(axis=2)

41 +

42 + lora_id = tl.load(lora_ids + lora_idx)

43 + if lora_id == -1:

44 + # Early exit for the no-lora case.

45 + return

46 +

47 + lora_m_size = tl.load(num_tokens_per_lora + lora_idx)

48 +

49 + cta_m_offset = pid_m * BLOCK_M

50 + if cta_m_offset >= lora_m_size:

51 + # Early exit CTA.

52 + return

53 +

54 + # num rows this CTA should process.

55 + cta_m_len = min(BLOCK_M, lora_m_size - cta_m_offset)

56 +

57 + # Identify all rows that this CTA should process.

58 + lora_m_indices_start = tl.load(lora_token_start_loc + lora_idx)

59 + cta_lora_seq_indices = (token_indices_sorted_by_lora_ids +

60 + lora_m_indices_start + cta_m_offset)

61 +

62 + # Load all relevant row indices.

63 + offset_m = tl.arange(0, BLOCK_M) % cta_m_len

64 + ram = tl.load(cta_lora_seq_indices + offset_m)

65 +

66 + do_shrink_kernel(

67 + pid_n,

68 + pid_sk,

69 + slice_id,

70 + lora_id,

71 + input_ptr,

72 + lora_ptr,

73 + out_ptr,

74 + N,

75 + K,

76 + cta_m_len,

77 + ram, # array identifying the rows of Input ptr to operate on

78 + # input strides

79 + input_d0_stride,

80 + input_d1_stride,

81 + # lora strides

82 + lora_d0_stride,

83 + lora_d1_stride,

84 + lora_d2_stride,

85 + # output strides

86 + output_d0_stride,

87 + output_d1_stride,

88 + output_d2_stride,

89 + scaling,

90 + BLOCK_M,

91 + BLOCK_N,

92 + BLOCK_K,

93 + EVEN_K,

94 + SPLIT_K,

95 + SLICE_NUM)

96 +

97 +

98 +@torch.inference_mode()

99 +def _v1_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]

103 + output_tensor: torch.Tensor, # shape [num_slices, num_tokens, lora_rank]

104 + token_lora_mapping: torch.Tensor, # shape [num_tokens]

105 + token_indices_sorted_by_lora_ids: torch.Tensor, # shape [num_tokens]

106 + num_tokens_per_lora: torch.Tensor, # shape [max-loras + 1]

107 + lora_token_start_loc: torch.Tensor, # shape [max-loras + 2]

108 + lora_ids: torch.Tensor, # shape [max-loras + 1]

109 + scaling: float,

110 +) -> None:

111 + """

112 + Args:

113 + inputs (torch.Tensor): Input tensor

114 + lora_a_weights (List[torch.Tensor]): LoRA weights

115 + output_tensor (torch.Tensor): output tensor

116 + token_lora_mapping (torch.Tensor): A tensor mapping each input token

117 + to the lora-id related to that token. A value of -1 indicates that

118 + LoRA doesn't apply to that token.

119 + token_indices_sorted_by_lora_ids (torch.Tensor): Row/Token indices from

120 + the A matrix grouped by LoRA IDs.

121 + num_tokens_per_lora (torch.Tensor): num_tokens_per_lora[i] is the number

122 + of tokens that are to be processed by LoRA ID lora_ids[i]

123 + lora_token_start_loc (torch.Tensor): A cumulative sum of

124 + num_tokens_per_lora. lora_token_start_loc[0] is always 0 so that

125 + lora_token_start_loc[i], along with num_tokens_per_lora[i]

126 + identifies the region in token_indices_sorted_by_lora_ids that

127 + LoRA lora_ids[i] should process.

128 + lora_ids (torch.Tensor): LoRA ids to process.

129 + scaling (float): Scaling factor.

130 + """

131 + assert inputs.dtype == lora_a_weights[0].dtype

132 + assert inputs.dtype in [torch.float16, torch.bfloat16]

133 + for weight in lora_a_weights:

134 + assert weight.dtype in [torch.float16, torch.bfloat16]

135 +

136 + assert inputs.size(1) == lora_a_weights[0].size(-1)

137 + assert inputs.is_contiguous()

138 + assert output_tensor.is_contiguous()

139 +

140 + # metadata sanity check

141 + assert token_lora_mapping.size(0) == token_indices_sorted_by_lora_ids.size(

142 + 0)

143 + assert lora_ids.size(0) == num_tokens_per_lora.size(0)

144 + assert lora_token_start_loc.size(0) == lora_ids.size(0) + 1

145 +

146 + (lora_ptr_tensor, lora_strides_d0, lora_strides_d1,

147 + lora_strides_d2) = _get_lora_a_ptr(lora_a_weights, inputs.device)

148 + N, K = lora_a_weights[0].shape[-2:] # K=hidden_size,N=rank

149 + M = inputs.size(0)

150 + NUM_SLICES = len(lora_a_weights)

151 + MAX_LORAS = lora_ids.size(0)

152 +

153 + # Triton kernel configs

154 + BLOCK_M = 32

155 + BLOCK_N = 16

156 + BLOCK_K = 256 if M < 128 else 32

157 + SPLIT_K = 64 if M < 128 else 8

158 + NUM_WARPS = 4

159 + NUM_CTAS = 1

160 + NUM_STAGES = 2

161 + MAX_NREG = None

162 +

163 + EVEN_K = K % (BLOCK_K * SPLIT_K) == 0 # type: ignore

164 +

165 + # TODO (varun): This grid formulation maximizes parallelization at the

166 + # cost of wasteful thread block launch when only few of the input tokens

167 + # require LoRA. This might not be the best in all cases.

168 + grid = (

169 + SPLIT_K * triton.cdiv(M, BLOCK_M) * triton.cdiv(N, BLOCK_N),

170 + NUM_SLICES,

171 + # Each LoRA receives its own set of thread blocks for output

172 + # computation. If some LoRA doesn't have any tokens to process, its

173 + # thread blocks exit early.

174 + MAX_LORAS,

175 + )

176 +

177 + _v1_shrink_kernel[grid](

178 + inputs,

179 + lora_ptr_tensor,

180 + output_tensor,

181 + M,

182 + N,

183 + K,

184 + token_indices_sorted_by_lora_ids,

185 + num_tokens_per_lora,

186 + lora_token_start_loc,

187 + lora_ids,

188 + scaling,

189 + inputs.stride(0),

190 + inputs.stride(1),

191 + lora_strides_d0,

192 + lora_strides_d1,

193 + lora_strides_d2,

194 + output_tensor.stride(0),

195 + output_tensor.stride(1),

196 + output_tensor.stride(2),

197 + BLOCK_M,

198 + BLOCK_N,

199 + BLOCK_K,

200 + EVEN_K,

201 + SPLIT_K,

202 + NUM_SLICES,

203 + num_warps=NUM_WARPS,

204 + num_ctas=NUM_CTAS,

205 + num_stages=NUM_STAGES,

206 + maxnreg=MAX_NREG,

207 + )

208 +

209 + return

210 +

211 +

212 +def _v1_shrink_fake(

213 + inputs: torch.Tensor,

214 + lora_a_weights: List[torch.Tensor],

215 + output_tensor: torch.Tensor,

216 + token_lora_mapping: torch.Tensor,

217 + token_indices_sorted_by_lora_ids: torch.Tensor,

218 + num_tokens_per_lora: torch.Tensor,

219 + lora_token_start_loc: torch.Tensor,

220 + lora_ids: torch.Tensor,

221 + scaling: float,

222 +) -> None:

223 + return

224 +

225 +

226 +try:

227 + direct_register_custom_op(

228 + op_name="v1_shrink",

229 + op_func=_v1_shrink,

230 + mutates_args=["output_tensor"],

231 + fake_impl=_v1_shrink_fake,

232 + )

233 + v1_shrink = torch.ops.vllm.v1_shrink

234 +

235 +except AttributeError:

236 + v1_shrink = _v1_shrink

vllm/lora/punica_wrapper/punica_gpu.py

+190 -58

@@ -6,24 +6,83 @@ Punica: Multi-Tenant LoRA Serving.

6 https://arxiv.org/abs/2310.18547

7 """

9 -from typing import Optional, Tuple, Union, final

9 +from typing import TYPE_CHECKING, List, Optional, Tuple, Union, final

11 import torch

13 +import vllm.envs as env

14 +from vllm.lora.layers import LoRAMapping

15 from vllm.triton_utils import HAS_TRITON

17 if HAS_TRITON:

16 - from vllm.lora.ops.triton_ops import bgmv_expand

17 - from vllm.lora.ops.triton_ops import bgmv_expand_slice

18 - from vllm.lora.ops.triton_ops import bgmv_shrink

19 - from vllm.lora.ops.triton_ops import sgmv_expand

20 - from vllm.lora.ops.triton_ops import sgmv_shrink

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

28 from .punica_base import PunicaWrapperBase

30 +if TYPE_CHECKING:

31 + # avoid circuit import

32 + from vllm.lora.models import LongContextLoRAContext

33 +

34 +

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 +

84 @final

26 -class PunicaWrapperGPU(PunicaWrapperBase):

85 +class PunicaWrapperGPU(PunicaWrapperBase, V1KernelMixin):

86 """

87 PunicaWrapperGPU is designed to manage and provide metadata for the punica

88 kernel. The main function is to maintain the state information for

@@ -35,6 +94,36 @@ class PunicaWrapperGPU(PunicaWrapperBase):

94 PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches,

95 device)

97 + self.max_loras = kwargs['max_loras']

98 +

99 + if env.VLLM_USE_V1:

100 + self._v1_make_metadata(self.max_loras, max_num_batched_tokens,

101 + max_batches, device)

102 +

103 + def update_metadata(

104 + self,

105 + mapping: LoRAMapping,

106 + lora_index_to_id: List[Optional[int]],

107 + max_loras: int,

108 + vocab_size: int,

109 + extra_vocab_size: int,

110 + long_lora_context: Optional["LongContextLoRAContext"] = None,

111 + **kwargs):

112 +

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,

@@ -66,7 +155,7 @@ class PunicaWrapperGPU(PunicaWrapperBase):

155 self,

156 y: torch.Tensor,

157 x: torch.Tensor,

69 - w_t_all: torch.Tensor,

158 + w_t_all: Tuple[torch.Tensor, ...],

159 offset_start: int,

160 add_inputs: bool,

161 ):

@@ -118,14 +207,21 @@ class PunicaWrapperGPU(PunicaWrapperBase):

207

208 x = x.view(-1, x.shape[-1])

209

121 - if self.is_prefill:

122 - # NOTE fused kernel

123 - self._apply_shrink_prefill(y, x, lora_a_stacked, scale)

210 + if env.VLLM_USE_V1:

211 + self._v1_apply_shrink(y, x, lora_a_stacked, scale) # type: ignore

212 else:

125 - # TODO fuse these kernels

126 - for slice_idx in range(len(lora_a_stacked)):

127 - self._apply_shrink_decode(y[slice_idx], x,

128 - lora_a_stacked[slice_idx], scale)

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)

225

226 def add_expand(self,

227 y: torch.Tensor,

@@ -160,25 +256,38 @@ class PunicaWrapperGPU(PunicaWrapperBase):

256 if lora_bias_stacked is not None:

257 self._apply_bias(self.token_lora_indices, y, output_slices,

258 lora_bias_stacked)

163 - if self.is_prefill:

164 - # NOTE fused kernel

165 - self._apply_expand_prefill(y,

166 - x,

167 - lora_b_stacked,

168 - offset_start,

169 - add_inputs=True)

259 +

260 + if env.VLLM_USE_V1:

261 + # TODO (varun): Profile with add_inputs = False. i.e. move the

262 + # addition out of the kernel

263 + self._v1_apply_expand(

264 + y,

265 + x, # type: ignore

266 + lora_b_stacked,

267 + offset_start,

268 + add_inputs=True)

269 else:

171 - # TODO fuse these kernels

172 - for slice_idx in range(len(lora_b_stacked)):

173 - self._apply_expand_decode(

270 +

271 + if self.is_prefill:

272 + # NOTE fused kernel

273 + self._apply_expand_prefill(

274 y,

175 - x[slice_idx],

176 - lora_b_stacked[slice_idx],

275 + x, # type: ignore

276 + lora_b_stacked,

277 offset_start,

178 - output_slices[slice_idx],

179 - add_inputs=add_inputs,

180 - )

181 - offset_start += output_slices[slice_idx]

278 + add_inputs=True)

279 + else:

280 + # TODO fuse these kernels

281 + for slice_idx in range(len(lora_b_stacked)):

282 + self._apply_expand_decode(

283 + y,

284 + x[slice_idx],

285 + lora_b_stacked[slice_idx],

286 + offset_start,

287 + output_slices[slice_idx],

288 + add_inputs=add_inputs,

289 + )

290 + offset_start += output_slices[slice_idx]

291 y = y.view_as(y_org)

292

293 def add_lora_embedding(self,

@@ -200,18 +309,24 @@ class PunicaWrapperGPU(PunicaWrapperBase):

309 add_inputs (bool): Default to True.

310 """

311

203 - if self.is_prefill:

204 - sgmv_expand(

205 - x.unsqueeze(dim=0),

206 - [lora_b_stacked],

207 - y,

208 - *self.prefill_metadata,

209 - offset_start=0,

210 - add_inputs=add_inputs,

211 - )

312 + if env.VLLM_USE_V1:

313 + self._v1_apply_expand(y,

314 + x.unsqueeze(dim=0), (lora_b_stacked, ),

315 + offset_start=0,

316 + add_inputs=add_inputs)

317 else:

213 - bgmv_expand(x, lora_b_stacked, y, self.token_lora_indices,

214 - add_inputs)

318 + if self.is_prefill:

319 + sgmv_expand(

320 + x.unsqueeze(dim=0),

321 + (lora_b_stacked, ),

322 + y,

323 + *self.prefill_metadata,

324 + offset_start=0,

325 + add_inputs=add_inputs,

326 + )

327 + else:

328 + bgmv_expand(x, lora_b_stacked, y, self.token_lora_indices,

329 + add_inputs)

330

331 def add_lora_linear(self,

332 y: torch.Tensor,

@@ -257,19 +372,25 @@ class PunicaWrapperGPU(PunicaWrapperBase):

372 r = lora_b_stacked[0].size(-1)

373 # We set the buffer to be float32 by default ,refer to:

374 # https://github.com/triton-lang/triton/issues/1387

260 - buffer = torch.zeros(

375 + buffer = torch.zeros( # type: ignore

376 (len(output_slices), x.size(0), r),

377 dtype=torch.float32,

378 device=x.device,

379 )

265 - self.add_shrink(buffer, x, lora_a_stacked, scale, **kwargs)

266 - self.add_expand(y,

267 - buffer,

268 - lora_b_stacked,

269 - None,

270 - output_slices,

271 - add_inputs=True,

272 - **kwargs)

380 + self.add_shrink(

381 + buffer, # type: ignore

382 + x,

383 + lora_a_stacked,

384 + scale,

385 + **kwargs)

386 + self.add_expand(

387 + y,

388 + buffer, # type: ignore

389 + lora_b_stacked,

390 + None,

391 + output_slices,

392 + add_inputs=True,

393 + **kwargs)

394

395 def add_lora_logits(self,

396 y: torch.Tensor,

@@ -305,11 +426,22 @@ class PunicaWrapperGPU(PunicaWrapperBase):

426 buffer = torch.zeros((x.size(0), r),

427 dtype=torch.float32,

428 device=x.device)

308 - # LogitsProcessorWithLoRA always using bgmv.

309 - bgmv_shrink(x, lora_a_stacked, buffer, self.sampler_indices, scale)

310 - bgmv_expand(buffer,

311 - lora_b_stacked,

312 - y,

313 - self.sampler_indices,

314 - add_inputs=True)

429 +

430 + if env.VLLM_USE_V1:

431 + v1_shrink(x, [lora_a_stacked], buffer.unsqueeze(dim=0),

432 + *self.prompt_mapping_v1_meta.meta_args(x.size(0)), scale)

433 +

434 + v1_expand(buffer.unsqueeze(dim=0), [lora_b_stacked],

435 + y,

436 + *self.prompt_mapping_v1_meta.meta_args(buffer.size(0)),

437 + add_inputs=True)

438 + else:

439 +

440 + # V0 LogitsProcessorWithLoRA always using bgmv.

441 + bgmv_shrink(x, lora_a_stacked, buffer, self.sampler_indices, scale)

442 + bgmv_expand(buffer,

443 + lora_b_stacked,

444 + y,

445 + self.sampler_indices,

446 + add_inputs=True)

447 y = y.view_as(y_org)

vllm/v1/worker/lora_model_runner_mixin.py

+3 -3

@@ -62,9 +62,9 @@ class LoRAModelRunnerMixin:

62 if not self.lora_manager:

63 raise RuntimeError("LoRA is not enabled.")

65 - # We dont make any distinction between prefills and decodes in the

66 - # scheduler. To that effect, set is_prefill to True so we use the

67 - # sgmv punica kernels always.

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`.

68 lora_mapping = LoRAMapping(token_lora_mapping,

69 prompt_lora_mapping,

70 is_prefill=True)