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 12, 23:30 UTC
Relax Transformers modeling backend MoE experts check (#28952)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Nov 19, 2025 at 14:50 UTC
4f5299f7174ffb10bdc640b47d3494083fc39c48
2 files changed
+11
-2
docs/models/supported_models.md
+3
-1
@@ -79,7 +79,9 @@ To make your model compatible with the Transformers modeling backend, it needs:
79
1. Add `is_causal = False` to `MyAttention`.
80
- If your model is mixture-of-experts (MoE):
81
1. Your sparse MoE block must have an attribute called `experts`.
82
- 2. The class of `experts` (`MyExperts`) must inherit from `nn.ModuleList`.
82
+ 2. The class of `experts` (`MyExperts`) must either:
83
+ - Inherit from `nn.ModuleList` (naive).
84
+ - Or contain all 3D `nn.Parameters` (packed).
85
3. `MyExperts.forward` must accept `hidden_states`, `top_k_index`, `top_k_weights`.
86
2. `MyAttention` must use `ALL_ATTENTION_FUNCTIONS` to call attention.
87
3. `MyModel` must contain `_supports_attention_backend = True`.
vllm/model_executor/models/transformers/moe.py
+8
-1
@@ -256,7 +256,14 @@ class MoEMixin(MixtureOfExperts):
256
def _recursive_replace(module: nn.Module, prefix: str):
257
for child_name, child_module in module.named_children():
258
qual_name = maybe_prefix(prefix, child_name)
259
- if child_name == "experts" and isinstance(child_module, nn.ModuleList):
259
+ # Naive implementations will have experts as ModuleList
260
+ is_modulelist = isinstance(child_module, nn.ModuleList)
261
+ # Packed implementations will have experts as 3D tensors of shapes like:
262
+ # gate_up_proj = (num_experts, 2 * intermediate_size, hidden_size)
263
+ # down_proj = (num_experts, intermediate_size, hidden_size)
264
+ params = list(child_module.parameters())
265
+ is_3d = len(params) > 0 and all(p.ndim == 3 for p in params)
266
+ if child_name == "experts" and (is_modulelist or is_3d):
267
# Alias for readability
268
mlp = module
269
experts = child_module