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 13, 08:15 UTC
[Docs] Add some details about what the MoE block needs for the Transformers backend (#28588)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Nov 13, 2025 at 03:12 UTC
32262834610d56f533c53acb953777ac8f3d7942
1 file changed
+23
-1
docs/models/supported_models.md
+23
-1
@@ -75,7 +75,12 @@ This section details the necessary modifications to make to a Transformers compa
75
To make your model compatible with the Transformers backend, it needs:
76
77
1. `kwargs` passed down through all modules from `MyModel` to `MyAttention`.
78
- 1. If your model is encoder-only, you must also add `is_causal = False` to `MyAttention`.
78
+ - If your model is encoder-only:
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`.
83
+ 3. `MyExperts.forward` must accept `hidden_states`, `top_k_index`, `top_k_weights`.
84
2. `MyAttention` must use `ALL_ATTENTION_FUNCTIONS` to call attention.
85
3. `MyModel` must contain `_supports_attention_backend = True`.
86
@@ -102,6 +107,23 @@ class MyAttention(nn.Module):
107
)
108
...
109
110
+# Only do this for mixture-of-experts models
111
+class MyExperts(nn.ModuleList):
112
+ def forward(self, hidden_states, top_k_index, top_k_weights):
113
+ ...
114
+
115
+# Only do this for mixture-of-experts models
116
+class MySparseMoEBlock(nn.Module):
117
+ def __init__(self, config):
118
+ ...
119
+ self.experts = MyExperts(config)
120
+ ...
121
+
122
+ def forward(self, hidden_states: torch.Tensor):
123
+ ...
124
+ hidden_states = self.experts(hidden_states, top_k_index, top_k_weights)
125
+ ...
126
+
127
class MyModel(PreTrainedModel):
128
_supports_attention_backend = True
129
```