use registry mixin

Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>

Author: kylesayrs

src/llmcompressor/modeling/deepseek_v3.py

@@ -4,13 +4,10 @@
     DeepseekV3MoE as OriginalDeepseekV3MoE,
 )
 
-from llmcompressor.modeling.moe_context import (
-    MoECalibrationModule,
-    register_moe_calibration,
-)
+from llmcompressor.modeling.moe_context import MoECalibrationModule
 
 
-@register_moe_calibration("DeepseekV3MoE")
+@MoECalibrationModule.register("DeepseekV3MoE")
 class CalibrationDeepseekV3MoE(MoECalibrationModule):
     """
     Calibration version of DeepseekV3MoE that sends all tokens to all experts.

src/llmcompressor/modeling/llama4.py

@@ -11,14 +11,11 @@
     Llama4TextMoe,
 )
 
-from llmcompressor.modeling.moe_context import (
-    MoECalibrationModule,
-    register_moe_calibration,
-)
+from llmcompressor.modeling.moe_context import MoECalibrationModule
 from llmcompressor.utils.dev import skip_weights_initialize
 
 
-@register_moe_calibration("Llama4TextMoe")
+@MoECalibrationModule.register("Llama4TextMoe")
 class SequentialLlama4TextMoe(MoECalibrationModule):
     """
     Calibration version of Llama4TextMoe that unpacks experts for sequential processing.

src/llmcompressor/modeling/moe_context.py

@@ -14,22 +14,20 @@
 
 import contextlib
 from abc import ABC
-from typing import Dict, Type
 
 import torch
+from compressed_tensors.registry import RegistryMixin
 from loguru import logger
 from tqdm import tqdm
 from transformers import PreTrainedModel
 
 __all__ = [
     "MoECalibrationModule",
-    "MOE_CALIBRATION_MODULES",
-    "register_moe_calibration",
     "moe_calibration_context",
 ]
 
 
-class MoECalibrationModule(ABC, torch.nn.Module):
+class MoECalibrationModule(ABC, torch.nn.Module, RegistryMixin):
     """
     Abstract base class for MoE calibration modules.
 
@@ -62,32 +60,6 @@ def restore(self, original: torch.nn.Module) -> torch.nn.Module:
         )
 
 
-# Registry: module class name -> calibration module class
-MOE_CALIBRATION_MODULES: Dict[str, Type[MoECalibrationModule]] = {}
-
-
-def register_moe_calibration(module_class_name: str):
-    """
-    Decorator to register a MoE calibration module.
-
-    Usage:
-        @register_moe_calibration("DeepseekV3MoE")
-        class CalibrationDeepseekV3MoE(MoECalibrationModule):
-            ...
-
-    Args:
-        module_class_name: The class name of the original module to replace
-    """
-
-    def decorator(cls: Type[MoECalibrationModule]) -> Type[MoECalibrationModule]:
-        if not issubclass(cls, MoECalibrationModule):
-            raise TypeError(f"{cls.__name__} must inherit from MoECalibrationModule")
-        MOE_CALIBRATION_MODULES[module_class_name] = cls
-        return cls
-
-    return decorator
-
-
 @contextlib.contextmanager
 def moe_calibration_context(
     model: PreTrainedModel,
@@ -127,9 +99,10 @@ def moe_calibration_context(
     # Step 1: Collect all MoE modules that need replacement
     logger.debug("Entering MoE calibration context")
     modules_to_replace = []
+    moe_class_names = MoECalibrationModule.registered_names()
     for name, module in model.named_modules():
         class_name = module.__class__.__name__
-        if class_name in MOE_CALIBRATION_MODULES:
+        if class_name in moe_class_names:
             modules_to_replace.append((name, module, class_name))
 
     # Step 2: Replace modules with progress bar
@@ -138,8 +111,8 @@ def moe_calibration_context(
         for name, module, class_name in tqdm(
             modules_to_replace, desc="Replacing MoE modules for calibration"
         ):
-            calibration_cls = MOE_CALIBRATION_MODULES[class_name]
-            replacement = calibration_cls(
+            replacement = MoECalibrationModule.load_from_registry(
+                class_name,
                 module,
                 model.config,
                 calibrate_all_experts=calibrate_all_experts,

src/llmcompressor/modeling/qwen3_moe.py

@@ -20,13 +20,10 @@
     Qwen3MoeSparseMoeBlock as OriginalQwen3MoeSparseMoeBlock,
 )
 
-from llmcompressor.modeling.moe_context import (
-    MoECalibrationModule,
-    register_moe_calibration,
-)
+from llmcompressor.modeling.moe_context import MoECalibrationModule
 
 
-@register_moe_calibration("Qwen3MoeSparseMoeBlock")
+@MoECalibrationModule.register("Qwen3MoeSparseMoeBlock")
 class CalibrationQwen3MoeSparseMoeBlock(MoECalibrationModule):
     """
     Calibration version of Qwen3MoeSparseMoeBlock that sends all tokens to all experts.

src/llmcompressor/modeling/qwen3_vl_moe.py

@@ -4,14 +4,11 @@
     Qwen3VLMoeTextSparseMoeBlock as OriginalQwen3VLMoeTextSparseMoeBlock,
 )
 
-from llmcompressor.modeling.moe_context import (
-    MoECalibrationModule,
-    register_moe_calibration,
-)
+from llmcompressor.modeling.moe_context import MoECalibrationModule
 from llmcompressor.utils.dev import skip_weights_initialize
 
 
-@register_moe_calibration("CalibrationQwen3VLMoeTextSparseMoeBlock")
+@MoECalibrationModule.register("CalibrationQwen3VLMoeTextSparseMoeBlock")
 class CalibrateQwen3VLMoeTextSparseMoeBlock(MoECalibrationModule):
     """
     Calibration version of Qwen3VLMoeTextSparseMoeBlock that sends all tokens to all