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 8, 10:00 UTC
[BugFix] Fix DP chunking (#34379)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Bill Nell <bnell@redhat.com>
Lucas Wilkinson committed
Feb 11, 2026 at 23:44 UTC
136b0bfa59377ed2bbd3b3716036a96267cfe80b
1 file changed
+12
-2
vllm/model_executor/layers/fused_moe/runner/default_moe_runner.py
+12
-2
@@ -421,7 +421,7 @@ class DefaultMoERunner(MoERunner):
421
layer: torch.nn.Module,
422
full_hidden_states: torch.Tensor,
423
full_router_logits: torch.Tensor,
424
- shared_input: torch.Tensor | None,
424
+ full_shared_input: torch.Tensor | None,
425
has_separate_shared_experts: bool,
426
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
427
assert self.batched_hidden_states is not None
@@ -449,6 +449,11 @@ class DefaultMoERunner(MoERunner):
449
chunk_size = chunk_end - chunk_start
450
hidden_states = full_hidden_states[chunk_start:chunk_end, :]
451
router_logits = full_router_logits[chunk_start:chunk_end, :]
452
+ shared_input = (
453
+ full_shared_input[chunk_start:chunk_end, :]
454
+ if full_shared_input is not None
455
+ else None
456
+ )
457
458
assert self.batched_hidden_states is not None
459
assert self.batched_router_logits is not None
@@ -476,8 +481,13 @@ class DefaultMoERunner(MoERunner):
481
staged_hidden_states.copy_(hidden_states, non_blocking=True)
482
staged_router_logits.copy_(router_logits, non_blocking=True)
483
484
+ shared_input = (
485
+ shared_input if shared_input is not None else staged_hidden_states
486
+ )
487
+
488
# Matrix multiply.
489
if self.quant_method.is_monolithic:
490
+ assert has_separate_shared_experts or self.shared_experts is None
491
final_hidden_states = self.quant_method.apply_monolithic(
492
layer=layer,
493
x=staged_hidden_states,
@@ -501,7 +511,7 @@ class DefaultMoERunner(MoERunner):
511
assert not isinstance(final_hidden_states, tuple)
512
assert self.shared_experts is not None
513
504
- shared_output = self.shared_experts(staged_hidden_states)
514
+ shared_output = self.shared_experts(shared_input)
515
516
final_hidden_states = (
517
shared_output,