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, 02:15 UTC
[Bugfix] Fix torch.compile x LoRA for PyTorch 2.8 (#20823)
Signed-off-by: rzou <zou3519@gmail.com>
Richard Zou committed
Jul 12, 2025 at 02:06 UTC
a3a5a47e48d3c6610686a489af2bd987062e74df
1 file changed
+8
-6
vllm/lora/layers.py
+8
-6
@@ -240,17 +240,19 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA):
240
def forward(self, x: torch.Tensor) -> torch.Tensor:
241
added_tokens_mask = torch.where(x > self.base_layer.org_vocab_size - 1,
242
1, 0)
243
- embeddings_indices = torch.narrow(
244
- self.punica_wrapper._embeddings_indices, 1, 0, x.size(0))
243
246
- indices = embeddings_indices[1]
244
+ # NB: Don't use torch.narrow here. torch.narrow triggers some
245
+ # Dynamic Shape specialization in torch.compile
246
+ num_tokens = x.shape[0]
247
+ indices_1 = self.punica_wrapper._embeddings_indices[1][:num_tokens]
248
+ indices_0 = self.punica_wrapper._embeddings_indices[0][:num_tokens]
249
+
250
full_lora_a_embeddings = F.embedding(
248
- x + indices,
251
+ x + indices_1,
252
self.lora_a_stacked_2d,
253
)
251
- indices = embeddings_indices[0]
254
full_output = self.base_layer.forward(x +
253
- (indices * added_tokens_mask))
255
+ (indices_0 * added_tokens_mask))
256
257
full_output_org = full_output
258
if full_output.ndim == 3: