biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: shen-shanshan <467638484@qq.com>
This commit is contained in:
Shanshan Shen
2026-01-10 00:21:11 +08:00
committed by GitHub
parent ac9f9330e6
commit 08d954f036
24 changed files with 441 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/model_executor/layers/mamba/mamba_mixer2.py
View File

@@ -49,8 +49,11 @@ from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata
# Adapted from transformers.models.mamba2.modeling_mamba2.MambaRMSNormGated
# --8<-- [start:mixer2_gated_rms_norm]
@CustomOp.register("mixer2_gated_rms_norm")
class Mixer2RMSNormGated(CustomOp):
# --8<-- [end:mixer2_gated_rms_norm]
def __init__(
self,
full_hidden_size: int,
@@ -214,6 +217,7 @@ def mamba_v2_sharded_weight_loader(
# Adapted from transformers.models.mamba.modeling_mamba.MambaMixer
# --8<-- [start:mamba_mixer2]
@CustomOp.register("mamba_mixer2")
class MambaMixer2(MambaBase, CustomOp):
"""
@@ -226,6 +230,8 @@ class MambaMixer2(MambaBase, CustomOp):
**selective** state spaces)
"""
# --8<-- [end:mamba_mixer2]
def __init__(
self,
hidden_size: int,