iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[LoRA] Cleanup VocabParallelEmbedding (#53843)

Browse Source 
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
This commit is contained in:
Jee Jee Li
2026-08-26 16:19:35 +08:00
committed by GitHub
parent e376d45e82
commit 75aa189ba0
6 changed files with 19 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/lora/layers/__init__.py
+1 -1
View File
@@ -20,7 +20,7 @@ from vllm.lora.layers.row_parallel_linear import (
RowParallelLinearWithShardedLoRA,
)
from vllm.lora.layers.utils import LoRAMapping, LoRAMappingType
from vllm.lora.layers.vocal_parallel_embedding import VocabParallelEmbeddingWithLoRA
from vllm.lora.layers.vocab_parallel_embedding import VocabParallelEmbeddingWithLoRA
__all__ = [
"BaseLayerWithLoRA",
vllm/lora/layers/logits_processor.py
+8 -8
View File
@@ -20,7 +20,7 @@ from .base import BaseLayerWithLoRA
class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
"""
LoRA wrapper for LogitsProcessor, with extra logic to handle the
application of the LoRA adapter and added LoRA vocabulary.
application of the LoRA adapter.
Args:
base_layer: LogitsProcessor layer
@@ -183,19 +183,19 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA):
return None
if self.sharded_to_full_mapping_gpu is not None:
# Reindex full logits tensor to ensure 1:1 mapping between
# index and token_id
# Reindex the gathered logits so that index == token_id. Each TP
# shard is padded to `num_embeddings_per_partition`, so the gather
# interleaves the shards' padding with real vocab entries.
# Example for:
# org_vocab_size = 4
# added_vocab_size = 2
# org_vocab_size = 6
# pad_to_size = 8
# tp_size = 2
# indices: [0, 1, 2, 3, 4, 5, 6, 7]
# token_id: [0, 1, 4, -1, 2, 3, 5, -1]
# indices: [0, 1, 2, 3, 4, 5, 6, 7]
# token_id: [0, 1, 2, -1, 3, 4, 5, -1]
# Therefore, the mapping is expected to be:
# [0, 1, 4, 6, 2, 3, 5, 7] so that when we reindex,
# [0, 1, 2, 4, 5, 6, 3, 7] so that when we reindex,
# we get:
# indices: [0, 1, 2, 3, 4, 5, 6, 7]
# token_id: [0, 1, 2, 3, 4, 5, -1, -1]
vllm/lora/layers/vocal_parallel_embedding.py → vllm/lora/layers/vocab_parallel_embedding.py
+2 -23
View File
@@ -21,8 +21,6 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
def __init__(self, base_layer: VocabParallelEmbedding) -> None:
super().__init__()
self.base_layer = base_layer
self.embeddings_slice: tuple[int, int] | None
self.embeddings_weights: torch.Tensor | None
def create_lora_weights(
self,
@@ -30,25 +28,6 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
lora_config: LoRAConfig,
model_config: PretrainedConfig | None = None,
) -> None:
if self.base_layer.num_added_embeddings_per_partition > 0:
# We can start adding lora weights
self.embeddings_weights = self.base_layer.weight.data[
self.base_layer.num_org_embeddings_per_partition : self.base_layer.num_org_embeddings_per_partition # noqa: E501
+ self.base_layer.num_added_embeddings_per_partition
]
self.embeddings_slice = (
self.base_layer.shard_indices.added_vocab_start_index
- self.base_layer.org_vocab_size,
self.base_layer.shard_indices.added_vocab_end_index
- self.base_layer.org_vocab_size,
)
self.base_layer.weight.data[
self.base_layer.num_org_embeddings_per_partition :
].fill_(0)
else:
self.embeddings_slice = None
self.embeddings_weights = None
self.lora_a_stacked = torch.zeros(
(
max_loras,
@@ -100,10 +79,10 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
# NB: Don't use torch.narrow here. torch.narrow triggers some
# Dynamic Shape specialization in torch.compile
num_tokens = x.shape[0]
indices_1 = self.punica_wrapper._embeddings_indices[1][:num_tokens]
embeddings_indices = self.punica_wrapper._embeddings_indices[:num_tokens]
full_lora_a_embeddings = F.embedding(
x + indices_1,
x + embeddings_indices,
self.lora_a_stacked_2d,
)
full_output = self.base_layer.forward(x)
vllm/lora/punica_wrapper/punica_base.py
+4 -18
View File
@@ -145,7 +145,7 @@ class PunicaWrapperBase(PunicaWrapperABC):
max_num_batched_tokens, dtype=torch.long, device=device
)
self._embeddings_indices = torch.empty(
2, max_num_batched_tokens, dtype=torch.long, device=device
max_num_batched_tokens, dtype=torch.long, device=device
)
# 4 is the number of indices tensors.
@@ -172,10 +172,6 @@ class PunicaWrapperBase(PunicaWrapperABC):
max_loras: int,
vocab_size: int,
):
# NOTE We have remove lora extra vocab support for now. So we set
# extra_vocab_size always to 0, and extra_vocab_size will be removed.
extra_vocab_size = 0
(
base_indices,
sampler_indices,
@@ -187,7 +183,6 @@ class PunicaWrapperBase(PunicaWrapperABC):
lora_index_to_id,
max_loras,
vocab_size,
extra_vocab_size,
self.device,
)
self._token_lora_indices[: base_indices.shape[0]].copy_(base_indices)
@@ -195,9 +190,9 @@ class PunicaWrapperBase(PunicaWrapperABC):
self._sampler_indices_padded[: sampler_indices_padded.shape[0]].copy_(
sampler_indices_padded
)
self._embeddings_indices[
: embeddings_indices.shape[0], : embeddings_indices.shape[1]
].copy_(embeddings_indices)
self._embeddings_indices[: embeddings_indices.shape[0]].copy_(
embeddings_indices
)
self.indices_len[:] = indices_len
@@ -272,15 +267,6 @@ class PunicaWrapperBase(PunicaWrapperABC):
indices_padded_len = self.indices_len[2]
return self._sampler_indices_padded[:indices_padded_len]
@property
def embeddings_indices(self) -> torch.Tensor:
"""
This property provides access to the indices used for lora embeddings,
specifically for VocabParallelEmbeddingWithLoRA.
"""
embeddings_indices_len = self.indices_len[3]
return self._embeddings_indices[:, :embeddings_indices_len]
def update_metadata(
self,
mapping: "LoRAMapping",
vllm/lora/punica_wrapper/punica_xpu.py
+1 -1
View File
@@ -44,7 +44,7 @@ class PunicaWrapperXPU(PunicaWrapperBase):
):
PunicaWrapperBase.__init__(self, max_num_batched_tokens, max_batches, device)
torch._dynamo.mark_dynamic(self._token_lora_indices, 0)
torch._dynamo.mark_dynamic(self._embeddings_indices, 1)
torch._dynamo.mark_dynamic(self._embeddings_indices, 0)
torch._dynamo.mark_dynamic(self._sampler_indices_padded, 0)
self.lora_config = kwargs["lora_config"]
vllm/lora/punica_wrapper/utils.py
+3 -15
View File
@@ -56,7 +56,6 @@ def convert_mapping(
lora_index_to_id: list[int | None],
max_loras: int,
vocab_size: int,
extra_vocab_size: int,
device: torch.device,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[int]]:
"""Converts LoRAMapping to index tensors.
@@ -66,7 +65,6 @@ def convert_mapping(
lora_index_to_id: List mapping LoRA ids to LoRA indices.
max_loras: Maximum number of LoRAs.
vocab_size: Model vocab size.
extra_vocab_size: Extra vocab size each LoRA can have.
Returns:
A tuple of tensors:
@@ -80,10 +78,8 @@ def convert_mapping(
requests to LoRA indices for sampler with padding.
Same as sampler_indices, but -1 is replaced with
max_loras.
embeddings_indices: Tensor of shape [2, batch_size] mapping
requests to embedding indices. First row is for embeddings
added by the LoRAs, second row is for the LoRA.lora_a
embeddings.
embeddings_indices: Tensor of shape [batch_size] mapping
requests to the row offset of their LoRA.lora_a embeddings.
indices_len: List of lengths of the above tensors. It contains
(base_indices, sampler_indices, sampler_indices_padded,
embeddings_indices).
@@ -124,15 +120,7 @@ def convert_mapping(
prompt_mapping_tensor = async_tensor_h2d(
prompt_mapping, dtype=torch.long, device=device
)
embeddings_indices = torch.stack(
[
indices[2] * extra_vocab_size,
indices[2] * (vocab_size + extra_vocab_size),
]
)
embeddings_indices = torch.where(
embeddings_indices == -1, max_loras - 1, embeddings_indices
)
embeddings_indices = indices[2] * vocab_size
base_indices = indices[1]
sampler_indices = prompt_mapping_tensor
sampler_indices_padded = sampler_indices.clone()