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 10, 10:15 UTC
fix lora moe sharding when rank < max_lora_rank (#31994)
Signed-off-by: gnovack <gnovack@amazon.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
gnovack committed
Jan 8, 2026 at 22:43 UTC
bde38c11df0ea066a740efe9b77fff5418be45df
2 files changed
+6
-8
tests/lora/test_gptoss_tp.py
-1
@@ -95,7 +95,6 @@ def test_gpt_oss_lora_tp2(gptoss20b_lora_files, fully_sharded_loras):
95
max_model_len=1024,
96
enable_lora=True,
97
max_loras=2,
98
- max_lora_rank=8,
98
max_num_seqs=2,
99
max_num_batched_tokens=2048,
100
tensor_parallel_size=2,
vllm/lora/layers/fused_moe.py
+6
-7
@@ -428,9 +428,9 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
428
current_lora_rank = w13_lora_a.shape[1]
429
assert current_lora_rank % self.tp_size == 0
430
# Based on S-LoRA, we slice W13/W1/W3 A along the rank dim.
431
- sliced_rank = current_lora_rank // self.tp_size
432
- start_idx = self.tp_rank * sliced_rank
433
- end_idx = (self.tp_rank + 1) * sliced_rank
431
+ shard_size = self.w13_lora_a_stacked[0].shape[2]
432
+ start_idx = self.tp_rank * shard_size
433
+ end_idx = (self.tp_rank + 1) * shard_size
434
return w13_lora_a[:, start_idx:end_idx, :]
435
436
def _slice_w13_b(self, w13_lora_b: torch.Tensor):
@@ -465,11 +465,10 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
465
return w2_lora_b
466
# Based on S-LoRA, we slice W2 B along the hidden_size dim.
467
# w2_lora_b shape (num_experts,output_size,rank)
468
- current_lora_size = w2_lora_b.shape[1]
468
+ shard_size = self.w2_lora_b_stacked[0].shape[2]
469
+ start_idx = self.tp_rank * shard_size
470
+ end_idx = (self.tp_rank + 1) * shard_size
471
470
- sliced_size = current_lora_size // self.tp_size
471
- start_idx = self.tp_rank * sliced_size
472
- end_idx = (self.tp_rank + 1) * sliced_size
472
return w2_lora_b[:, start_idx:end_idx, :]
473
474
def reset_lora(self, index: int):