[Kernels] Isolate modular kernel code from FusedMoEMethodBase subclasses. (#27123)

This commit is contained in:
bnellnm
2025-11-04 08:59:45 -05:00
committed by GitHub
parent e4ee658672
commit 938772af03
16 changed files with 271 additions and 311 deletions

View File

@@ -266,14 +266,14 @@ class DeviceCommunicatorBase:
module
for module in model.modules()
# TODO(bnell): Should use isinstance but can't. Maybe search for
# presence of quant_method.init_prepare_finalize?
# presence of quant_method.maybe_init_modular_kernel?
if (
module.__class__.__name__ == "FusedMoE"
or module.__class__.__name__ == "SharedFusedMoE"
)
]
for module in moe_modules:
module.quant_method.init_prepare_finalize(module)
module.maybe_init_modular_kernel()
def dispatch(
self,