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, 22:45 UTC
Change the type signature of MixtureOfExperts.expert_weights to MutableSequence[Sequence[Tensor]] (#33573)
Signed-off-by: Sage Moore <sagmoore@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Sage Moore committed
Feb 4, 2026 at 14:02 UTC
ce498a6d61a083b1bbbd1cc92754961b43860ff6
2 files changed
+18
-9
vllm/distributed/eplb/rebalance_execute.py
+9
-7
@@ -6,7 +6,7 @@ The actual execution of the rearrangement.
6
This involves the exchange of expert weights between GPUs.
7
"""
8
9
-from collections.abc import Iterable, Sequence
9
+from collections.abc import Sequence
10
from dataclasses import dataclass
11
12
import numpy as np
@@ -153,7 +153,7 @@ def move_to_buffer(
153
num_local_experts: int,
154
old_indices: np.ndarray,
155
new_indices: np.ndarray,
156
- expert_weights: Iterable[torch.Tensor],
156
+ expert_weights: Sequence[torch.Tensor],
157
expert_weights_buffers: Sequence[torch.Tensor],
158
cuda_stream: torch.cuda.Stream | None,
159
ep_group: ProcessGroup,
@@ -355,7 +355,7 @@ def move_to_buffer(
355
356
357
def move_from_buffer(
358
- expert_weights: Iterable[torch.Tensor],
358
+ expert_weights: Sequence[torch.Tensor],
359
expert_weights_buffers: list[torch.Tensor],
360
is_unchanged: np.ndarray,
361
is_received_locally: np.ndarray,
@@ -436,7 +436,7 @@ def move_from_buffer(
436
async def transfer_layer(
437
old_global_expert_indices: torch.Tensor,
438
new_global_expert_indices: torch.Tensor,
439
- expert_weights: Sequence[Iterable[torch.Tensor]],
439
+ expert_weights: Sequence[Sequence[torch.Tensor]],
440
expert_weights_buffer: Sequence[torch.Tensor],
441
ep_group: ProcessGroup,
442
is_profile: bool = False,
@@ -488,7 +488,8 @@ async def transfer_layer(
488
assert old_global_expert_indices.shape[1] == new_global_expert_indices.shape[1]
489
num_moe_layers, num_physical_experts = old_global_expert_indices.shape
490
assert len(expert_weights) == num_moe_layers
491
- num_local_physical_experts = next(iter(expert_weights[0])).shape[0]
491
+ assert len(expert_weights[0]) >= 1
492
+ num_local_physical_experts = expert_weights[0][0].shape[0]
493
assert new_global_expert_indices.shape == (num_moe_layers, num_physical_experts)
494
assert num_physical_experts == ep_size * num_local_physical_experts
495
@@ -510,7 +511,7 @@ async def transfer_layer(
511
def rearrange_expert_weights_inplace(
512
old_global_expert_indices: torch.Tensor,
513
new_global_expert_indices: torch.Tensor,
513
- expert_weights: Sequence[Iterable[torch.Tensor]],
514
+ expert_weights: Sequence[Sequence[torch.Tensor]],
515
ep_group: ProcessGroup,
516
is_profile: bool = False,
517
rank_mapping: dict[int, int] | None = None,
@@ -553,8 +554,9 @@ def rearrange_expert_weights_inplace(
554
555
num_moe_layers, num_physical_experts = old_global_expert_indices.shape
556
assert len(expert_weights) == num_moe_layers
557
+ assert len(expert_weights[0]) >= 1
558
557
- num_local_physical_experts = next(iter(expert_weights[0])).shape[0]
559
+ num_local_physical_experts = expert_weights[0][0].shape[0]
560
assert new_global_expert_indices.shape == (num_moe_layers, num_physical_experts)
561
562
ep_size = ep_group.size()
vllm/model_executor/models/interfaces.py
+9
-2
@@ -2,7 +2,14 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import asyncio
5
-from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, MutableSequence
5
+from collections.abc import (
6
+ AsyncGenerator,
7
+ Callable,
8
+ Iterable,
9
+ Mapping,
10
+ MutableSequence,
11
+ Sequence,
12
+)
13
from contextlib import ExitStack, contextmanager, nullcontext
14
from typing import (
15
TYPE_CHECKING,
@@ -818,7 +825,7 @@ class MixtureOfExperts(Protocol):
825
Check if the model is a mixture of experts (MoE) model.
826
"""
827
821
- expert_weights: MutableSequence[Iterable[Tensor]]
828
+ expert_weights: MutableSequence[Sequence[Tensor]]
829
"""
830
Expert weights saved in this rank.
831