[LoRA] Cleanup VocabParallelEmbedding

vllm/lora/layers/__init__.py

@@ -20,7 +20,7 @@
     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

@@ -20,7 +20,7 @@
 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 @@ def _get_logits(
             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

@@ -21,34 +21,13 @@ 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,
         max_loras: int,
         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 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         # 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

@@ -145,7 +145,7 @@ def __init__(
             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 @@ def _update_base_metadata(
         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,17 +183,16 @@ def _update_base_metadata(
             lora_index_to_id,
             max_loras,
             vocab_size,
-            extra_vocab_size,
             self.device,
         )
         self._token_lora_indices[: base_indices.shape[0]].copy_(base_indices)
         self._sampler_indices[: sampler_indices.shape[0]].copy_(sampler_indices)
         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 @@ def sampler_indices_padded(self) -> torch.Tensor:
         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

@@ -44,7 +44,7 @@ def __init__(
     ):
         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

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