@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 08d954f036
Fork 0

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

[Doc] Add developer guide for CustomOp (#30886)

Signed-off-by: shen-shanshan <467638484@qq.com>

Shanshan Shen committed Jan 10, 2026 at 00:21 UTC 08d954f03659cb08148b77cd2e0d33b77f6bd6ef
24 files changed +441 -5
docs/design/custom_op.md new
+318
@@ -0,0 +1,318 @@
1 +# CustomOp
2 +
3 +`CustomOp` is an abstract class used for dispatching the forward method of various operations to the appropriate backend. It also offers a mechanism for both vLLM and OOT (Out-Of-Tree) plugins to register their custom operations.
4 +
5 +This document will introduce how CustomOp works in vLLM and how to implement a new `CustomOp`.
6 +
7 +## How CustomOp Works in vLLM
8 +
9 +`CustomOp` manages two dictionaries of all custom ops (i.e., op classes, indexed by registered name) in its class, for vLLM and OOT plugins respectively.
10 +
11 +??? code
12 +
13 + ```python
14 + class CustomOp(nn.Module):
15 +
16 + op_registry: dict[str, type["CustomOp"]] = {}
17 + op_registry_oot: dict[str, type["CustomOp"]] = {}
18 + ```
19 +
20 +We can use `@CustomOp.register("op_name")` to register an op class to the `CustomOp` system. After this, the `op_name` and its class will be added into the `op_registry` dictionary. In addition, We can also register an OOT op by `@CustomOp.register_oot("op_name")`. We will introduce this mechanism in detail later.
21 +
22 +When a `CustomOp` is called (i.e., call its `forward()` method), if it is enabled (i.e., with `--compilation_config.custom_ops '["+op_name"]'`), it will automatically dispatch the forward method to the appropriate backend according to `current_platform`. Otherwise (i.e., it is disabled), it will only call the `forward_native()` method to use PyTorch-native implementation of this forward method.
23 +
24 +- **CPU platform:** dispatch to `forward_cpu()`.
25 +- **CUDA platform:** dispatch to `forward_cuda()`.
26 +- **ROCm platform:** dispatch to `forward_hip()`. If `forward_hip()` is not implemented, it will use `forward_cuda()` as a fallback.
27 +- **XPU platform:** dispatch to `forward_xpu()`.
28 +- **TPU platform:** dispatch to `forward_tpu()`.
29 +- **OOT platform:** dispatch to `forward_oot()`. This will only be called on OOT platforms.
30 +- **Default:** dispatch to `forward_native()` as a final fallback for all platforms.
31 +
32 +!!! note
33 + Note that the dispatching logic might not be absolute because of class inheritance. Derived class might override the behavior.
34 +
35 +Furthermore, vLLM decides whether to enable or disable a `CustomOp` based on `compilation_config.custom_ops`. To be specific, if a `CustomOp` is not registered in `compilation_config.custom_ops` (i.e., uses the default config), it will be enabled if `compilation_config.custom_ops` contains `all`, or will be disabled if it contains `none`.
36 +
37 +!!! note
38 + Note that `all` and `none` cannot coexist in `compilation_config.custom_ops`.
39 +
40 +By default, if `compilation_config.backend == "inductor"` and `compilation_config.mode != CompilationMode.NONE`, a `none` will be appended into `compilation_config.custom_ops`, otherwise a `all` will be appended. In other words, this means `CustomOp` will be disabled in some platforms (i.e., those use `inductor` as dafault backend for `torch.compile`) when running with torch compile mode. In this case, Inductor generates (fused) Triton kernels for those disabled custom ops.
41 +
42 +!!! note
43 + For multi-modal models, vLLM has enforece enabled some custom ops to use device-specific deep-optimized kernels for better performance in ViT part, such as `MMEncoderAttention` and `ApplyRotaryEmb`. We can also pass a `enforce_enable=True` param to the `__init__()` method of the `CustomOp` to enforce enable itself at object-level.
44 +
45 + Note that this `enforce_enable` mechanism will be removed after we adding a separate `compilation_config` for multi-modal part.
46 +
47 +## How to Customise Your Configuration for CustomOp
48 +
49 +vLLM also offers fine-grained control over which custom ops to enable or disable for users, by manually passing a `--compilation_config.custom_ops '["..."]'` when launching a server.
50 +
51 +For example:
52 +
53 +- Use `--compilation_config.custom_ops '["all"]'` to enable all custom ops.
54 +- Use `--compilation_config.custom_ops '["none"]'` to disable all custom ops.
55 +- Use `--compilation_config.custom_ops '["all,-op1"]'` to enable all custom ops except op1 (i.e., prefixed with a `-` means "disable").
56 +- Use `--compilation_config.custom_ops '["none,+op1,+op2"]'` to only enable op1 and op2 (i.e., prefixed with a `+` means "enable").
57 +
58 +## Types of Supported CustomOp in vLLM
59 +
60 +**1. Attention:**
61 +
62 +```python
63 +--8<-- "vllm/attention/layers/mm_encoder_attention.py:mm_encoder_attn"
64 +
65 +--8<-- "vllm/model_executor/layers/mla.py:multi_head_latent_attention"
66 +```
67 +
68 +**2. Activation:**
69 +
70 +```python
71 +--8<-- "vllm/model_executor/layers/activation.py:silu_and_mul"
72 +
73 +--8<-- "vllm/model_executor/layers/activation.py:mul_and_silu"
74 +
75 +--8<-- "vllm/model_executor/layers/activation.py:gelu_new"
76 +
77 +--8<-- "vllm/model_executor/layers/activation.py:gelu_fast"
78 +
79 +--8<-- "vllm/model_executor/layers/activation.py:quick_gelu"
80 +
81 +--8<-- "vllm/model_executor/layers/activation.py:gelu_and_mul"
82 +
83 +--8<-- "vllm/model_executor/layers/activation.py:gelu_and_mul_sparse"
84 +
85 +--8<-- "vllm/model_executor/layers/activation.py:relu2"
86 +
87 +--8<-- "vllm/model_executor/layers/activation.py:xielu"
88 +
89 +--8<-- "vllm/model_executor/layers/activation.py:swigluoai_and_mul"
90 +
91 +--8<-- "vllm/model_executor/layers/activation.py:fatrelu_and_mul"
92 +```
93 +
94 +**3. MM-Conv:**
95 +
96 +```python
97 +--8<-- "vllm/model_executor/layers/conv.py:conv2d"
98 +
99 +--8<-- "vllm/model_executor/layers/conv.py:conv3d"
100 +```
101 +
102 +**4. Embedding:**
103 +
104 +```python
105 +--8<-- "vllm/model_executor/layers/vocab_parallel_embedding.py:vocab_parallel_embedding"
106 +
107 +--8<-- "vllm/model_executor/layers/vocab_parallel_embedding.py:parallel_lm_head"
108 +```
109 +
110 +**5. Linear:**
111 +
112 +```python
113 +--8<-- "vllm/model_executor/layers/linear.py:row_parallel_linear"
114 +
115 +--8<-- "vllm/model_executor/layers/linear.py:column_parallel_linear"
116 +
117 +--8<-- "vllm/model_executor/layers/linear.py:replicated_linear"
118 +```
119 +
120 +**6. Logits Processor:**
121 +
122 +```python
123 +--8<-- "vllm/model_executor/layers/logits_processor.py:logits_processor"
124 +```
125 +
126 +**7. Mamba:**
127 +
128 +```python
129 +--8<-- "vllm/model_executor/layers/mamba/mamba_mixer.py:mamba_mixer"
130 +
131 +--8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mamba_mixer2"
132 +
133 +--8<-- "vllm/model_executor/layers/mamba/mamba_mixer2.py:mixer2_gated_rms_norm"
134 +
135 +--8<-- "vllm/model_executor/models/plamo2.py:plamo2_mamba_mixer"
136 +
137 +--8<-- "vllm/model_executor/layers/mamba/short_conv.py:short_conv"
138 +```
139 +
140 +**8. MoE:**
141 +
142 +```python
143 +--8<-- "vllm/model_executor/layers/fused_moe/layer.py:fused_moe"
144 +
145 +--8<-- "vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py:modular_fused_moe"
146 +
147 +--8<-- "vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py:unquantized_fused_moe"
148 +
149 +--8<-- "vllm/model_executor/models/transformers/moe.py:transformers_fused_moe"
150 +
151 +--8<-- "vllm/model_executor/layers/fused_moe/fused_moe.py:grouped_topk"
152 +```
153 +
154 +**9. Norm:**
155 +
156 +```python
157 +--8<-- "vllm/model_executor/layers/layernorm.py:rms_norm"
158 +
159 +--8<-- "vllm/model_executor/layers/layernorm.py:rms_norm_gated"
160 +
161 +--8<-- "vllm/model_executor/layers/layernorm.py:gemma_rms_norm"
162 +```
163 +
164 +**10. Quantization:**
165 +
166 +```python
167 +--8<-- "vllm/model_executor/layers/quantization/input_quant_fp8.py:quant_fp8"
168 +```
169 +
170 +**11. Rope:**
171 +
172 +```python
173 +--8<-- "vllm/model_executor/layers/rotary_embedding/base.py:rotary_embedding"
174 +
175 +--8<-- "vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py:dual_chunk_rotary_embedding"
176 +
177 +--8<-- "vllm/model_executor/layers/rotary_embedding/common.py:apply_rotary_emb"
178 +```
179 +
180 +## Guidelines for Implementing a New CustomOp
181 +
182 +### Implement a New CustomOp in vLLM
183 +
184 +This part is a tutorial of how to implement a New `CustomOp` in vLLM.
185 +
186 +Steps:
187 +
188 +1. Implement a new op class, which extends from `CustomOp` base class.
189 +2. Add the `@CustomOp.register("op_name")` decorator on this op class to register it into `CustomOp` system.
190 +3. Implement different `forward_xxx()` method according to your needs.
191 +
192 +Taking `MMEncoderAttention` as an example:
193 +
194 +??? code
195 +
196 + ```python
197 + @CustomOp.register("mm_encoder_attn")
198 + class MMEncoderAttention(CustomOp):
199 +
200 + def __init__(
201 + self,
202 + num_heads: int,
203 + head_size: int,
204 + scale: float | None = None,
205 + num_kv_heads: int | None = None,
206 + prefix: str = "",
207 + multimodal_config: MultiModalConfig | None = None,
208 + ) -> None:
209 + super().__init__()
210 + # Init...
211 +
212 + def forward_native(
213 + self,
214 + query: torch.Tensor,
215 + key: torch.Tensor,
216 + value: torch.Tensor,
217 + cu_seqlens: torch.Tensor | None = None,
218 + max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
219 + ) -> torch.Tensor:
220 + # Call TORCH_SDPA implementation...
221 +
222 + def forward_cuda(
223 + self,
224 + query: torch.Tensor,
225 + key: torch.Tensor,
226 + value: torch.Tensor,
227 + cu_seqlens: torch.Tensor | None = None,
228 + max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
229 + ) -> torch.Tensor:
230 + # Call FA or TORCH_SDPA implementation...
231 +
232 + def forward_cpu(
233 + self,
234 + query: torch.Tensor,
235 + key: torch.Tensor,
236 + value: torch.Tensor,
237 + cu_seqlens: torch.Tensor | None = None,
238 + max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
239 + ) -> torch.Tensor:
240 + # Call TORCH_SDPA implementation...
241 +
242 + def forward_xpu(
243 + self,
244 + query: torch.Tensor,
245 + key: torch.Tensor,
246 + value: torch.Tensor,
247 + cu_seqlens: torch.Tensor | None = None,
248 + max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
249 + ) -> torch.Tensor:
250 + # Call FA implementation...
251 +
252 + def forward_tpu(
253 + self,
254 + query: torch.Tensor,
255 + key: torch.Tensor,
256 + value: torch.Tensor,
257 + cu_seqlens: torch.Tensor | None = None,
258 + max_seqlen: torch.Tensor | None = None, # Only used for Flash Attention
259 + ) -> torch.Tensor:
260 + # Call PALLAS implementation...
261 + ```
262 +
263 +### Register a New CustomOp in OOT Device Plugins
264 +
265 +Currently, thanks to [vLLM's hardware-plugin mechanism](./plugin_system.md), there are various OOT device plugins emerging out to enable vLLM seamlessly runs on different hardwares. You can also find more details about this mechanism at [Introducing vLLM Hardware Plugin, Best Practice from Ascend NPU](https://blog.vllm.ai/2025/05/12/hardware-plugin.html).
266 +
267 +- **Official device plugins:** [vllm-ascend](https://github.com/vllm-project/vllm-ascend) (for Huawei Ascend NPU), [vllm-spyre](https://github.com/vllm-project/vllm-spyre)
268 +(for Spyre), [vllm-gaudi](https://github.com/vllm-project/vllm-gaudi) (for Intel Gaudi), [vllm-neuron](https://github.com/vllm-project/vllm-neuron) (for AWS Neuron), [vllm-meta](https://github.com/vllm-project/vllm-metal) (for Apple Silicon), etc.
269 +- **Non-official device plugins:** [vllm-metax](https://github.com/MetaX-MACA/vLLM-metax) (for MetaX GPU), [vllm-kunlun](https://github.com/baidu/vLLM-Kunlun) (for Baidu Kunlun XPU), etc.
270 +
271 +In this case, `CustomOp` can enable these hardware manufacturers to seamlessly replace vLLM's operations with their deep-optimized kernels for specific devices at runtime, by just registering an OOT `CustomOp` and implementing the `forward_oot()` method.
272 +
273 +Now, this part will show you how to register an OOT `CustomOp` for a device plugin.
274 +
275 +Taking `MMEncoderAttention` as an example:
276 +
277 +1. Implement a `CustomMMEncoderAttention` class which extends from `MMEncoderAttention` and implement its `forward_oot()` method.
278 +2. Register your `CustomMMEncoderAttention` into vLLM to replace `MMEncoderAttention`.
279 +
280 +??? code
281 +
282 + ```python
283 + from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
284 + from vllm.model_executor.custom_op import CustomOp
285 +
286 +
287 + @CustomOp.register_oot("MMEncoderAttention")
288 + class CustomMMEncoderAttention(MMEncoderAttention):
289 +
290 + def __init__(...):
291 + super().__init__(...)
292 +
293 + def forward_oot(...):
294 + # Call optimized device-specific kernels.
295 + ...
296 + ```
297 +
298 +In this case, a new item `{"MMEncoderAttention": CustomMMEncoderAttention}` will be added into `op_registry_oot`. When initializing a `MMEncoderAttention` op object, if the class name (i.e., `MMEncoderAttention`) is contained in the keys of `op_registry_oot`, vLLM will replace it with our registered class (i.e., `CustomMMEncoderAttention`) and instantiate it.
299 +
300 +After that, when this `MMEncoderAttention` op is called, your `forward_oot()` will be called if it is enabled. Thus, you will get expected performance on your hardwares without directly modify vLLM.
301 +
302 +In addition, you can also register all your `CustomOp` at one place for better management.
303 +
304 +??? code
305 +
306 + ```python
307 + from vllm.model_executor.custom_op import CustomOp
308 +
309 +
310 + REGISTERED_CUSTOM_OPS = {
311 + "CustomOP1": YourCustomOp1,
312 + "CustomOP2": YourCustomOp2,
313 + "CustomOP3": YourCustomOp3,
314 + }
315 +
316 + for op_name, op_cls in REGISTERED_CUSTOM_OPS.items():
317 + CustomOp.register_oot(_decorated_op_cls=op_cls, name=op_name)
318 + ```
vllm/attention/layers/mm_encoder_attention.py
+3
@@ -18,10 +18,13 @@ from vllm.model_executor.models.vision import get_vit_attn_backend
18 logger = init_logger(__name__)
19
20
21 +# --8<-- [start:mm_encoder_attn]
22 @CustomOp.register("mm_encoder_attn")
23 class MMEncoderAttention(CustomOp):
24 """Multi-headed attention without any cache, used for multimodal encoder."""
25
26 + # --8<-- [end:mm_encoder_attn]
27 +
28 def __init__(
29 self,
30 num_heads: int,
vllm/config/compilation.py
+2 -1
@@ -404,7 +404,8 @@ class CompilationConfig:
404 - 'none,+op1,+op2' to enable only op1 and op2
405
406 By default, all custom ops are enabled when running without Inductor and
407 - disabled when running with Inductor: mode>=VLLM_COMPILE and backend="inductor".
407 + disabled when running with Inductor: mode>CompilationMode.NONE and
408 + backend="inductor".
409 Inductor generates (fused) Triton kernels for disabled custom ops."""
410 splitting_ops: list[str] | None = None
411 """A list of ops to exclude from cudagraphs, used in piecewise compilation.
vllm/model_executor/custom_op.py
+6 -3
@@ -87,9 +87,12 @@ class CustomOp(nn.Module):
87 # specific backend. Currently, we do not support dynamic dispatching.
88 compilation_config = get_cached_compilation_config()
89
90 - # CustomOp object can be enforce enabled, e.g., enable device-specific
91 - # kernels in ViT models when enabling graph mode. By default, it will
92 - # follow the compilation_config to determine whether enable itself.
90 + # NOTE(shen-shanshan): CustomOp object can be enforce enabled, e.g.,
91 + # enable device-specific kernels in ViT models when enabling graph
92 + # mode. By default, it will follow the compilation_config to determine
93 + # whether enable itself.
94 + # This enforce_enable mechanism will be removed after we adding a
95 + # separate compilation_config for multi-modal part.
96 enabled = self._enforce_enable or self.enabled()
97 if enabled:
98 compilation_config.enabled_custom_ops.update([self.__class__.name])
vllm/model_executor/layers/activation.py
+33
@@ -22,6 +22,7 @@ from vllm.utils.collection_utils import LazyDict
22 logger = init_logger(__name__)
23
24
25 +# --8<-- [start:fatrelu_and_mul]
26 @CustomOp.register("fatrelu_and_mul")
27 class FatreluAndMul(CustomOp):
28 """An activation function for FATReLU.
@@ -35,6 +36,8 @@ class FatreluAndMul(CustomOp):
36 return: (num_tokens, d) or (batch_size, seq_len, d)
37 """
38
39 + # --8<-- [end:fatrelu_and_mul]
40 +
41 def __init__(self, threshold: float = 0.0):
42 super().__init__()
43 self.threshold = threshold
@@ -58,6 +61,7 @@ class FatreluAndMul(CustomOp):
61 return out
62
63
64 +# --8<-- [start:silu_and_mul]
65 @CustomOp.register("silu_and_mul")
66 class SiluAndMul(CustomOp):
67 """An activation function for SwiGLU.
@@ -69,6 +73,8 @@ class SiluAndMul(CustomOp):
73 return: (num_tokens, d) or (batch_size, seq_len, d)
74 """
75
76 + # --8<-- [end:silu_and_mul]
77 +
78 def __init__(self):
79 super().__init__()
80 if current_platform.is_cuda_alike():
@@ -101,6 +107,7 @@ class SiluAndMul(CustomOp):
107 return out
108
109
110 +# --8<-- [start:mul_and_silu]
111 @CustomOp.register("mul_and_silu")
112 class MulAndSilu(CustomOp):
113 """An activation function for SwiGLU.
@@ -112,6 +119,8 @@ class MulAndSilu(CustomOp):
119 return: (num_tokens, d) or (batch_size, seq_len, d)
120 """
121
122 + # --8<-- [end:mul_and_silu]
123 +
124 def __init__(self):
125 super().__init__()
126 if current_platform.is_cuda_alike():
@@ -139,6 +148,7 @@ class MulAndSilu(CustomOp):
148 # def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
149
150
151 +# --8<-- [start:gelu_and_mul_sparse]
152 @CustomOp.register("gelu_and_mul_sparse")
153 class GeluAndMulSparse(CustomOp):
154 """An activation function for GeluAndMulSparse.
@@ -153,6 +163,8 @@ class GeluAndMulSparse(CustomOp):
163 return: (num_tokens, d) or (batch_size, seq_len, d)
164 """
165
166 + # --8<-- [end:gelu_and_mul_sparse]
167 +
168 def __init__(self, activation_sparsity: float, approximate: str = "none"):
169 super().__init__()
170 # Gelu.
@@ -195,6 +207,7 @@ class GeluAndMulSparse(CustomOp):
207 return self.forward_native(x)
208
209
210 +# --8<-- [start:gelu_and_mul]
211 @CustomOp.register("gelu_and_mul")
212 class GeluAndMul(CustomOp):
213 """An activation function for GeGLU.
@@ -206,6 +219,8 @@ class GeluAndMul(CustomOp):
219 return: (batch_size, seq_len, d) or (num_tokens, d)
220 """
221
222 + # --8<-- [end:gelu_and_mul]
223 +
224 def __init__(self, approximate: str = "none"):
225 super().__init__()
226 self.approximate = approximate
@@ -257,9 +272,12 @@ class GeluAndMul(CustomOp):
272 return f"approximate={repr(self.approximate)}"
273
274
275 +# --8<-- [start:swigluoai_and_mul]
276 @CustomOp.register("swigluoai_and_mul")
277 class SwigluOAIAndMul(CustomOp):
278 # https://github.com/huggingface/transformers/blob/v4.55.0/src/transformers/models/gpt_oss/modeling_gpt_oss.py#L106-L110
279 + # --8<-- [end:swigluoai_and_mul]
280 +
281 def __init__(self, alpha: float = 1.702, limit: float = 7.0):
282 super().__init__()
283 self.alpha = alpha
@@ -286,8 +304,11 @@ class SwigluOAIAndMul(CustomOp):
304 return f"alpha={repr(self.alpha)}, limit={repr(self.limit)}"
305
306
307 +# --8<-- [start:gelu_new]
308 @CustomOp.register("gelu_new")
309 class NewGELU(CustomOp):
310 + # --8<-- [end:gelu_new]
311 +
312 def __init__(self):
313 super().__init__()
314 if current_platform.is_cuda_alike() or current_platform.is_cpu():
@@ -311,8 +332,11 @@ class NewGELU(CustomOp):
332 return self.op(x)
333
334
335 +# --8<-- [start:gelu_fast]
336 @CustomOp.register("gelu_fast")
337 class FastGELU(CustomOp):
338 + # --8<-- [end:gelu_fast]
339 +
340 def __init__(self):
341 super().__init__()
342 if current_platform.is_cuda_alike() or current_platform.is_cpu():
@@ -335,9 +359,12 @@ class FastGELU(CustomOp):
359 return self.op(x)
360
361
362 +# --8<-- [start:quick_gelu]
363 @CustomOp.register("quick_gelu")
364 class QuickGELU(CustomOp):
365 # https://github.com/huggingface/transformers/blob/main/src/transformers/activations.py#L90
366 + # --8<-- [end:quick_gelu]
367 +
368 def __init__(self):
369 super().__init__()
370 if current_platform.is_cuda_alike() or current_platform.is_cpu():
@@ -365,12 +392,15 @@ class QuickGELU(CustomOp):
392 # def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
393
394
395 +# --8<-- [start:relu2]
396 @CustomOp.register("relu2")
397 class ReLUSquaredActivation(CustomOp):
398 """
399 Applies the relu^2 activation introduced in https://arxiv.org/abs/2109.08668v2
400 """
401
402 + # --8<-- [end:relu2]
403 +
404 def forward_native(self, x: torch.Tensor) -> torch.Tensor:
405 """PyTorch-native implementation equivalent to forward()."""
406 return torch.square(F.relu(x))
@@ -380,6 +410,7 @@ class ReLUSquaredActivation(CustomOp):
410 return self.forward_native(x)
411
412
413 +# --8<-- [start:xielu]
414 @CustomOp.register("xielu")
415 class XIELU(CustomOp):
416 """
@@ -388,6 +419,8 @@ class XIELU(CustomOp):
419 Otherwise, we emit a single warning and use xIELU Python
420 """
421
422 + # --8<-- [end:xielu]
423 +
424 def __init__(
425 self,
426 alpha_p_init: float = 0.8,
vllm/model_executor/layers/conv.py
+6
@@ -105,10 +105,13 @@ class ConvLayerBase(CustomOp):
105 return s
106
107
108 +# --8<-- [start:conv2d]
109 @CustomOp.register("conv2d")
110 class Conv2dLayer(ConvLayerBase):
111 """Conv layer with Conv2d."""
112
113 + # --8<-- [end:conv2d]
114 +
115 num_dim = 2
116
117 def _forward_mulmat(self, x: torch.Tensor) -> torch.Tensor:
@@ -204,10 +207,13 @@ class CausalConv2dLayer(Conv2dLayer):
207 return x
208
209
210 +# --8<-- [start:conv3d]
211 @CustomOp.register("conv3d")
212 class Conv3dLayer(ConvLayerBase):
213 """Conv layer with Conv3d."""
214
215 + # --8<-- [end:conv3d]
216 +
217 num_dim = 3
218
219 def _forward_mulmat(self, x: torch.Tensor) -> torch.Tensor:
vllm/model_executor/layers/fused_moe/fused_moe.py
+3
@@ -1459,10 +1459,13 @@ def grouped_topk(
1459 return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
1460
1461
1462 +# --8<-- [start:grouped_topk]
1463 @CustomOp.register("grouped_topk")
1464 class GroupedTopk(CustomOp):
1465 """GroupedTopk used by the Deepseek-V2 and Deepseek-V3 model."""
1466
1467 + # --8<-- [end:grouped_topk]
1468 +
1469 def __init__(
1470 self,
1471 topk: int,
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
+3
@@ -21,8 +21,11 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
21 logger = init_logger(__name__)
22
23
24 +# --8<-- [start:modular_fused_moe]
25 @CustomOp.register("modular_fused_moe")
26 class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):
27 + # --8<-- [end:modular_fused_moe]
28 +
29 def __init__(
30 self, old_quant_method: FusedMoEMethodBase, experts: FusedMoEModularKernel
31 ):
vllm/model_executor/layers/fused_moe/layer.py
+3
@@ -302,6 +302,7 @@ class FusedMoERouterImpl(FusedMoERouter):
302 return self.layer._select_experts(hidden_states, router_logits)
303
304
305 +# --8<-- [start:fused_moe]
306 @CustomOp.register("fused_moe")
307 class FusedMoE(CustomOp):
308 """FusedMoE layer for MoE models.
@@ -326,6 +327,8 @@ class FusedMoE(CustomOp):
327 router_logits_dtype: Data type for router logits buffers.
328 """
329
330 + # --8<-- [end:fused_moe]
331 +
332 def __init__(
333 self,
334 num_experts: int, # Global number of experts
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+3
@@ -52,10 +52,13 @@ else:
52 logger = init_logger(__name__)
53
54
55 +# --8<-- [start:unquantized_fused_moe]
56 @CustomOp.register("unquantized_fused_moe")
57 class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
58 """MoE method without quantization."""
59
60 + # --8<-- [end:unquantized_fused_moe]
61 +
62 def __init__(self, moe: FusedMoEConfig):
63 super().__init__(moe)
64
vllm/model_executor/layers/layernorm.py
+9
@@ -88,6 +88,7 @@ def dispatch_rocm_rmsnorm_func(
88 return rms_norm
89
90
91 +# --8<-- [start:rms_norm]
92 @CustomOp.register("rms_norm")
93 class RMSNorm(CustomOp):
94 """Root mean square normalization.
@@ -96,6 +97,8 @@ class RMSNorm(CustomOp):
97 Refer to https://arxiv.org/abs/1910.07467
98 """
99
100 + # --8<-- [end:rms_norm]
101 +
102 def __init__(
103 self,
104 hidden_size: int,
@@ -253,6 +256,7 @@ class RMSNorm(CustomOp):
256 return s
257
258
259 +# --8<-- [start:gemma_rms_norm]
260 @CustomOp.register("gemma_rms_norm")
261 class GemmaRMSNorm(CustomOp):
262 """RMS normalization for Gemma.
@@ -262,6 +266,8 @@ class GemmaRMSNorm(CustomOp):
266 2. (x * w).to(orig_dtype) instead of x.to(orig_dtype) * w.
267 """
268
269 + # --8<-- [end:gemma_rms_norm]
270 +
271 def __init__(
272 self,
273 hidden_size: int,
@@ -321,6 +327,7 @@ class GemmaRMSNorm(CustomOp):
327 return self.forward_native(x, residual)
328
329
330 +# --8<-- [start:rms_norm_gated]
331 @CustomOp.register("rms_norm_gated")
332 class RMSNormGated(CustomOp):
333 """RMS Normalization with optional gating.
@@ -331,6 +338,8 @@ class RMSNormGated(CustomOp):
338 - Optional gating with SiLU activation
339 """
340
341 + # --8<-- [end:rms_norm_gated]
342 +
343 def __init__(
344 self,
345 hidden_size: int,
vllm/model_executor/layers/linear.py
+9
@@ -304,6 +304,7 @@ class LinearBase(CustomOp):
304 param.tp_size = self.tp_size
305
306
307 +# --8<-- [start:replicated_linear]
308 @CustomOp.register("replicated_linear")
309 class ReplicatedLinear(LinearBase):
310 """Replicated linear layer.
@@ -321,6 +322,8 @@ class ReplicatedLinear(LinearBase):
322 disable_tp: Take no effect for replicated linear layers.
323 """
324
325 + # --8<-- [end:replicated_linear]
326 +
327 def __init__(
328 self,
329 input_size: int,
@@ -421,6 +424,7 @@ class ReplicatedLinear(LinearBase):
424 return s
425
426
427 +# --8<-- [start:column_parallel_linear]
428 @CustomOp.register("column_parallel_linear")
429 class ColumnParallelLinear(LinearBase):
430 """Linear layer with column parallelism.
@@ -448,6 +452,8 @@ class ColumnParallelLinear(LinearBase):
452 disable_tp: If true, weights matrix won't be sharded through tp rank.
453 """
454
455 + # --8<-- [end:column_parallel_linear]
456 +
457 def __init__(
458 self,
459 input_size: int,
@@ -1289,6 +1295,7 @@ class QKVParallelLinear(ColumnParallelLinear):
1295 param_data.copy_(loaded_weight)
1296
1297
1298 +# --8<-- [start:row_parallel_linear]
1299 @CustomOp.register("row_parallel_linear")
1300 class RowParallelLinear(LinearBase):
1301 """Linear layer with row parallelism.
@@ -1323,6 +1330,8 @@ class RowParallelLinear(LinearBase):
1330 disable_tp: If true, weights matrix won't be sharded through tp rank.
1331 """
1332
1333 + # --8<-- [end:row_parallel_linear]
1334 +
1335 def __init__(
1336 self,
1337 input_size: int,
vllm/model_executor/layers/logits_processor.py
+3
@@ -13,6 +13,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmb
13 from vllm.platforms import current_platform
14
15
16 +# --8<-- [start:logits_processor]
17 @CustomOp.register("logits_processor")
18 class LogitsProcessor(CustomOp):
19 """Process logits and apply logits processors from sampling metadata.
@@ -23,6 +24,8 @@ class LogitsProcessor(CustomOp):
24 3. Apply logits processors (if any).
25 """
26
27 + # --8<-- [end:logits_processor]
28 +
29 def __init__(
30 self,
31 vocab_size: int,
vllm/model_executor/layers/mamba/mamba_mixer.py
+3
@@ -39,6 +39,7 @@ from vllm.v1.attention.backends.mamba1_attn import Mamba1AttentionMetadata
39
40
41 # Adapted from transformers.models.mamba.modeling_mamba.MambaMixer
42 +# --8<-- [start:mamba_mixer]
43 @CustomOp.register("mamba_mixer")
44 class MambaMixer(MambaBase, CustomOp):
45 """
@@ -51,6 +52,8 @@ class MambaMixer(MambaBase, CustomOp):
52 **selective** state spaces)
53 """
54
55 + # --8<-- [end:mamba_mixer]
56 +
57 def __init__(
58 self,
59 hidden_size: int,
vllm/model_executor/layers/mamba/mamba_mixer2.py
+6
@@ -49,8 +49,11 @@ from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
49
50
51 # Adapted from transformers.models.mamba2.modeling_mamba2.MambaRMSNormGated
52 +# --8<-- [start:mixer2_gated_rms_norm]
53 @CustomOp.register("mixer2_gated_rms_norm")
54 class Mixer2RMSNormGated(CustomOp):
55 + # --8<-- [end:mixer2_gated_rms_norm]
56 +
57 def __init__(
58 self,
59 full_hidden_size: int,
@@ -214,6 +217,7 @@ def mamba_v2_sharded_weight_loader(
217
218
219 # Adapted from transformers.models.mamba.modeling_mamba.MambaMixer
220 +# --8<-- [start:mamba_mixer2]
221 @CustomOp.register("mamba_mixer2")
222 class MambaMixer2(MambaBase, CustomOp):
223 """
@@ -226,6 +230,8 @@ class MambaMixer2(MambaBase, CustomOp):
230 **selective** state spaces)
231 """
232
233 + # --8<-- [end:mamba_mixer2]
234 +
235 def __init__(
236 self,
237 hidden_size: int,
vllm/model_executor/layers/mamba/short_conv.py
+3
@@ -27,8 +27,11 @@ from vllm.utils.torch_utils import direct_register_custom_op
27 from vllm.v1.attention.backends.short_conv_attn import ShortConvAttentionMetadata
28
29
30 +# --8<-- [start:short_conv]
31 @CustomOp.register("short_conv")
32 class ShortConv(MambaBase, CustomOp):
33 + # --8<-- [end:short_conv]
34 +
35 def __init__(
36 self,
37 config,
vllm/model_executor/layers/mla.py
+3
@@ -29,6 +29,7 @@ class MLAModules:
29 indexer_rotary_emb: torch.nn.Module | None = None
30
31
32 +# --8<-- [start:multi_head_latent_attention]
33 @CustomOp.register("multi_head_latent_attention")
34 class MultiHeadLatentAttentionWrapper(CustomOp):
35 """MLA layer registered as CustomOp to allow OOT backends to add
@@ -47,6 +48,8 @@ class MultiHeadLatentAttentionWrapper(CustomOp):
48 3. Return the output tensor.
49 """
50
51 + # --8<-- [end:multi_head_latent_attention]
52 +
53 def __init__(
54 self,
55 hidden_size: int,
vllm/model_executor/layers/quantization/input_quant_fp8.py
+3
@@ -18,6 +18,7 @@ _FP8_MIN, _FP8_MAX = get_fp8_min_max()
18 _FP8_MIN_SCALING_FACTOR = 1.0 / (_FP8_MAX * 512.0)
19
20
21 +# --8<-- [start:quant_fp8]
22 @CustomOp.register("quant_fp8")
23 class QuantFP8(CustomOp):
24 """
@@ -25,6 +26,8 @@ class QuantFP8(CustomOp):
26 This CustomOp supports both static and dynamic quantization.
27 """
28
29 + # --8<-- [end:quant_fp8]
30 +
31 def __init__(
32 self,
33 static: bool,
vllm/model_executor/layers/rotary_embedding/base.py
+3
@@ -10,10 +10,13 @@ from vllm.model_executor.custom_op import CustomOp
10 from .common import ApplyRotaryEmb
11
12
13 +# --8<-- [start:rotary_embedding]
14 @CustomOp.register("rotary_embedding")
15 class RotaryEmbeddingBase(CustomOp):
16 """Original rotary positional embedding."""
17
18 + # --8<-- [end:rotary_embedding]
19 +
20 def __init__(
21 self,
22 head_size: int,
vllm/model_executor/layers/rotary_embedding/common.py
+3
@@ -118,8 +118,11 @@ direct_register_custom_op(
118 )
119
120
121 +# --8<-- [start:apply_rotary_emb]
122 @CustomOp.register("apply_rotary_emb")
123 class ApplyRotaryEmb(CustomOp):
124 + # --8<-- [end:apply_rotary_emb]
125 +
126 def __init__(
127 self,
128 enforce_enable: bool = False,
vllm/model_executor/layers/rotary_embedding/dual_chunk_rope.py
+3
@@ -9,10 +9,13 @@ from vllm.model_executor.custom_op import CustomOp
9 from .common import rotate_gptj, rotate_neox
10
11
12 +# --8<-- [start:dual_chunk_rotary_embedding]
13 @CustomOp.register("dual_chunk_rotary_embedding")
14 class DualChunkRotaryEmbedding(CustomOp):
15 """Rotary positional embedding for Dual Chunk Attention."""
16
17 + # --8<-- [end:dual_chunk_rotary_embedding]
18 +
19 def __init__(
20 self,
21 head_size: int,
vllm/model_executor/layers/vocab_parallel_embedding.py
+6
@@ -181,6 +181,7 @@ def get_masked_input_and_mask(
181 return input_, ~vocab_mask
182
183
184 +# --8<-- [start:vocab_parallel_embedding]
185 @CustomOp.register("vocab_parallel_embedding")
186 class VocabParallelEmbedding(CustomOp):
187 """Embedding parallelized in the vocabulary dimension.
@@ -221,6 +222,8 @@ class VocabParallelEmbedding(CustomOp):
222 prefix: full name of the layer in the state dict
223 """ # noqa: E501
224
225 + # --8<-- [end:vocab_parallel_embedding]
226 +
227 def __init__(
228 self,
229 num_embeddings: int,
@@ -492,6 +495,7 @@ class VocabParallelEmbedding(CustomOp):
495 return s
496
497
498 +# --8<-- [start:parallel_lm_head]
499 @CustomOp.register("parallel_lm_head")
500 class ParallelLMHead(VocabParallelEmbedding):
501 """Parallelized LM head.
@@ -509,6 +513,8 @@ class ParallelLMHead(VocabParallelEmbedding):
513 padding_size: padding size for the vocabulary.
514 """
515
516 + # --8<-- [end:parallel_lm_head]
517 +
518 def __init__(
519 self,
520 num_embeddings: int,
vllm/model_executor/models/plamo2.py
+4 -1
@@ -103,8 +103,11 @@ def is_mamba(config: Plamo2Config, i: int) -> bool:
103 # Adapted from:
104 # vllm.model_executor.layers.mamba.mamba_mixer2.MambaMixer2
105 # transformers.models.mamba.modeling_mamba.MambaMixer
106 -@CustomOp.register(name="plamo2_mamba_mixer")
106 +# --8<-- [start:plamo2_mamba_mixer]
107 +@CustomOp.register("plamo2_mamba_mixer")
108 class Plamo2MambaMixer(MambaBase, CustomOp):
109 + # --8<-- [end:plamo2_mamba_mixer]
110 +
111 def __init__(self, vllm_config: VllmConfig, *, prefix: str = "", **kwargs) -> None:
112 super().__init__()
113 self.config = vllm_config.model_config.hf_config
vllm/model_executor/models/transformers/moe.py
+3
@@ -37,10 +37,13 @@ if TYPE_CHECKING:
37 from vllm.config import VllmConfig
38
39
40 +# --8<-- [start:transformers_fused_moe]
41 @CustomOp.register("transformers_fused_moe")
42 class TransformersFusedMoE(FusedMoE):
43 """Custom FusedMoE for the Transformers modeling backend."""
44
45 + # --8<-- [end:transformers_fused_moe]
46 +
47 def __init__(self, *args, **kwargs):
48 super().__init__(*args, **kwargs)
49 self._topk_ids: torch.Tensor = None