Skip to content

Commit a7338f0

Browse files
committed
make FusedMoEModularKernel a Leaf
Signed-off-by: Bill Nell <bnell@redhat.com>
1 parent d1b9b99 commit a7338f0

File tree

1 file changed

+6
-3
lines changed

1 file changed

+6
-3
lines changed

vllm/model_executor/layers/fused_moe/modular_kernel.py

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from abc import ABC, abstractmethod
44
from enum import Enum
55
from math import prod
6-
from typing import Optional
6+
from typing import Optional, final
77

88
import torch
99

@@ -202,7 +202,8 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
202202

203203
@property
204204
@abstractmethod
205-
def activation_formats(self) -> tuple[FusedMoEActivationFormat, FusedMoEActivationFormat]:
205+
def activation_formats(
206+
self) -> tuple[FusedMoEActivationFormat, FusedMoEActivationFormat]:
206207
"""
207208
Add comment
208209
"""
@@ -325,6 +326,7 @@ def _chunk_scales(scales: Optional[torch.Tensor], start: int,
325326
return None
326327

327328

329+
@final
328330
class FusedMoEModularKernel(torch.nn.Module):
329331
"""
330332
This class combines a FusedMoEPrepareAndFinalize instance and
@@ -346,7 +348,8 @@ def __init__(
346348
super().__init__()
347349
self.prepare_finalize = prepare_finalize
348350
self.fused_experts = fused_experts
349-
assert prepare_finalize.activation_format == fused_experts.activation_formats[0]
351+
assert prepare_finalize.activation_format == fused_experts.activation_formats[ # noqa: E501
352+
0]
350353

351354
def forward(
352355
self,

0 commit comments

Comments
 (0)