[Doc] Add developer guide for CustomOp (#30886)
Signed-off-by: shen-shanshan <467638484@qq.com>
This commit is contained in:
@@ -103,8 +103,11 @@ def is_mamba(config: Plamo2Config, i: int) -> bool:
|
||||
# Adapted from:
|
||||
# vllm.model_executor.layers.mamba.mamba_mixer2.MambaMixer2
|
||||
# transformers.models.mamba.modeling_mamba.MambaMixer
|
||||
@CustomOp.register(name="plamo2_mamba_mixer")
|
||||
# --8<-- [start:plamo2_mamba_mixer]
|
||||
@CustomOp.register("plamo2_mamba_mixer")
|
||||
class Plamo2MambaMixer(MambaBase, CustomOp):
|
||||
# --8<-- [end:plamo2_mamba_mixer]
|
||||
|
||||
def __init__(self, vllm_config: VllmConfig, *, prefix: str = "", **kwargs) -> None:
|
||||
super().__init__()
|
||||
self.config = vllm_config.model_config.hf_config
|
||||
|
||||
@@ -37,10 +37,13 @@ if TYPE_CHECKING:
|
||||
from vllm.config import VllmConfig
|
||||
|
||||
|
||||
# --8<-- [start:transformers_fused_moe]
|
||||
@CustomOp.register("transformers_fused_moe")
|
||||
class TransformersFusedMoE(FusedMoE):
|
||||
"""Custom FusedMoE for the Transformers modeling backend."""
|
||||
|
||||
# --8<-- [end:transformers_fused_moe]
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self._topk_ids: torch.Tensor = None
|
||||
|
||||
Reference in New Issue
Block a user