[LoRA] Cleanup VocabParallelEmbedding (#53843)
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Jee Jee Li committed
Aug 26, 2026 at 16:19 UTC
75aa189ba06768f2fc7aa9dd255799c7ed672b5e
6 files changed
+19
-66
vllm/lora/layers/__init__.py
+1
-1
@@ -20,7 +20,7 @@ from vllm.lora.layers.row_parallel_linear import (
20
RowParallelLinearWithShardedLoRA,
21
)
22
from vllm.lora.layers.utils import LoRAMapping, LoRAMappingType
23
-from vllm.lora.layers.vocal_parallel_embedding import VocabParallelEmbeddingWithLoRA
23
+from vllm.lora.layers.vocab_parallel_embedding import VocabParallelEmbeddingWithLoRA
24
25
__all__ = [
26
"BaseLayerWithLoRA",
vllm/lora/layers/logits_processor.py
+8
-8
@@ -20,7 +20,7 @@ from .base import BaseLayerWithLoRA
20
class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
21
"""
22
LoRA wrapper for LogitsProcessor, with extra logic to handle the
23
- application of the LoRA adapter and added LoRA vocabulary.
23
+ application of the LoRA adapter.
24
25
Args:
26
base_layer: LogitsProcessor layer
@@ -183,19 +183,19 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
183
return None
184
185
if self.sharded_to_full_mapping_gpu is not None:
186
- # Reindex full logits tensor to ensure 1:1 mapping between
187
- # index and token_id
186
+ # Reindex the gathered logits so that index == token_id. Each TP
187
+ # shard is padded to `num_embeddings_per_partition`, so the gather
188
+ # interleaves the shards' padding with real vocab entries.
189
# Example for:
189
- # org_vocab_size = 4
190
- # added_vocab_size = 2
190
+ # org_vocab_size = 6
191
# pad_to_size = 8
192
# tp_size = 2
193
194
- # indices: [0, 1, 2, 3, 4, 5, 6, 7]
195
- # token_id: [0, 1, 4, -1, 2, 3, 5, -1]
194
+ # indices: [0, 1, 2, 3, 4, 5, 6, 7]
195
+ # token_id: [0, 1, 2, -1, 3, 4, 5, -1]
196
197
# Therefore, the mapping is expected to be:
198
- # [0, 1, 4, 6, 2, 3, 5, 7] so that when we reindex,
198
+ # [0, 1, 2, 4, 5, 6, 3, 7] so that when we reindex,
199
# we get:
200
# indices: [0, 1, 2, 3, 4, 5, 6, 7]
201
# token_id: [0, 1, 2, 3, 4, 5, -1, -1]
vllm/lora/layers/vocab_parallel_embedding.py
renamed
+2
-23
@@ -21,8 +21,6 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
21
def __init__(self, base_layer: VocabParallelEmbedding) -> None:
22
super().__init__()
23
self.base_layer = base_layer
24
- self.embeddings_slice: tuple[int, int] | None
25
- self.embeddings_weights: torch.Tensor | None
24
25
def create_lora_weights(
26
self,
@@ -30,25 +28,6 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
28
lora_config: LoRAConfig,
29
model_config: PretrainedConfig | None = None,
30
) -> None:
33
- if self.base_layer.num_added_embeddings_per_partition > 0:
34
- # We can start adding lora weights
35
- self.embeddings_weights = self.base_layer.weight.data[
36
- self.base_layer.num_org_embeddings_per_partition : self.base_layer.num_org_embeddings_per_partition # noqa: E501
37
- + self.base_layer.num_added_embeddings_per_partition
38
- ]
39
- self.embeddings_slice = (
40
- self.base_layer.shard_indices.added_vocab_start_index
41
- - self.base_layer.org_vocab_size,
42
- self.base_layer.shard_indices.added_vocab_end_index
43
- - self.base_layer.org_vocab_size,
44
- )
45
- self.base_layer.weight.data[
46
- self.base_layer.num_org_embeddings_per_partition :
47
- ].fill_(0)
48
- else:
49
- self.embeddings_slice = None
50
- self.embeddings_weights = None
51
-
31
self.lora_a_stacked = torch.zeros(
32
(
33
max_loras,
@@ -100,10 +79,10 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
79
# NB: Don't use torch.narrow here. torch.narrow triggers some
80
# Dynamic Shape specialization in torch.compile
81
num_tokens = x.shape[0]
103
- indices_1 = self.punica_wrapper._embeddings_indices[1][:num_tokens]
82
+ embeddings_indices = self.punica_wrapper._embeddings_indices[:num_tokens]
83
84
full_lora_a_embeddings = F.embedding(
106
- x + indices_1,
85
+ x + embeddings_indices,
86
self.lora_a_stacked_2d,
87
)
88
full_output = self.base_layer.forward(x)
vllm/lora/punica_wrapper/punica_base.py
+4
-18
@@ -145,7 +145,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
145
max_num_batched_tokens, dtype=torch.long, device=device
146
)
147
self._embeddings_indices = torch.empty(
148
- 2, max_num_batched_tokens, dtype=torch.long, device=device
148
+ max_num_batched_tokens, dtype=torch.long, device=device
149
)
150
151
# 4 is the number of indices tensors.
@@ -172,10 +172,6 @@ class PunicaWrapperBase(PunicaWrapperABC):
172
max_loras: int,
173
vocab_size: int,
174
):
175
- # NOTE We have remove lora extra vocab support for now. So we set
176
- # extra_vocab_size always to 0, and extra_vocab_size will be removed.
177
-
178
- extra_vocab_size = 0
175
(
176
base_indices,
177
sampler_indices,
@@ -187,7 +183,6 @@ class PunicaWrapperBase(PunicaWrapperABC):
183
lora_index_to_id,
184
max_loras,
185
vocab_size,
190
- extra_vocab_size,
186
self.device,
187
)
188
self._token_lora_indices[: base_indices.shape[0]].copy_(base_indices)
@@ -195,9 +190,9 @@ class PunicaWrapperBase(PunicaWrapperABC):
190
self._sampler_indices_padded[: sampler_indices_padded.shape[0]].copy_(
191
sampler_indices_padded
192
)
198
- self._embeddings_indices[
199
- : embeddings_indices.shape[0], : embeddings_indices.shape[1]
200
- ].copy_(embeddings_indices)
193
+ self._embeddings_indices[: embeddings_indices.shape[0]].copy_(
194
+ embeddings_indices
195
+ )
196
197
self.indices_len[:] = indices_len
198
@@ -272,15 +267,6 @@ class PunicaWrapperBase(PunicaWrapperABC):
267
indices_padded_len = self.indices_len[2]
268
return self._sampler_indices_padded[:indices_padded_len]
269
275
- @property
276
- def embeddings_indices(self) -> torch.Tensor:
277
- """
278
- This property provides access to the indices used for lora embeddings,
279
- specifically for VocabParallelEmbeddingWithLoRA.
280
- """
281
- embeddings_indices_len = self.indices_len[3]
282
- return self._embeddings_indices[:, :embeddings_indices_len]
283
-
270
def update_metadata(
271
self,
272
mapping: "LoRAMapping",
vllm/lora/punica_wrapper/punica_xpu.py
+1
-1
@@ -44,7 +44,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
44
):
45
PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches, device)
46
torch._dynamo.mark_dynamic(self._token_lora_indices, 0)
47
- torch._dynamo.mark_dynamic(self._embeddings_indices, 1)
47
+ torch._dynamo.mark_dynamic(self._embeddings_indices, 0)
48
torch._dynamo.mark_dynamic(self._sampler_indices_padded, 0)
49
50
self.lora_config = kwargs["lora_config"]
vllm/lora/punica_wrapper/utils.py
+3
-15
@@ -56,7 +56,6 @@ def convert_mapping(
56
lora_index_to_id: list[int | None],
57
max_loras: int,
58
vocab_size: int,
59
- extra_vocab_size: int,
59
device: torch.device,
60
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[int]]:
61
"""Converts LoRAMapping to index tensors.
@@ -66,7 +65,6 @@ def convert_mapping(
65
lora_index_to_id: List mapping LoRA ids to LoRA indices.
66
max_loras: Maximum number of LoRAs.
67
vocab_size: Model vocab size.
69
- extra_vocab_size: Extra vocab size each LoRA can have.
68
69
Returns:
70
A tuple of tensors:
@@ -80,10 +78,8 @@ def convert_mapping(
78
requests to LoRA indices for sampler with padding.
79
Same as sampler_indices, but -1 is replaced with
80
max_loras.
83
- embeddings_indices: Tensor of shape [2, batch_size] mapping
84
- requests to embedding indices. First row is for embeddings
85
- added by the LoRAs, second row is for the LoRA.lora_a
86
- embeddings.
81
+ embeddings_indices: Tensor of shape [batch_size] mapping
82
+ requests to the row offset of their LoRA.lora_a embeddings.
83
indices_len: List of lengths of the above tensors. It contains
84
(base_indices, sampler_indices, sampler_indices_padded,
85
embeddings_indices).
@@ -124,15 +120,7 @@ def convert_mapping(
120
prompt_mapping_tensor = async_tensor_h2d(
121
prompt_mapping, dtype=torch.long, device=device
122
)
127
- embeddings_indices = torch.stack(
128
- [
129
- indices[2] * extra_vocab_size,
130
- indices[2] * (vocab_size + extra_vocab_size),
131
- ]
132
- )
133
- embeddings_indices = torch.where(
134
- embeddings_indices == -1, max_loras - 1, embeddings_indices
135
- )
123
+ embeddings_indices = indices[2] * vocab_size
124
base_indices = indices[1]
125
sampler_indices = prompt_mapping_tensor
126
sampler_indices_padded = sampler_indices.clone()