Revert "use registry mixin"

This reverts commit 6dd0320.

Author: kylesayrs

src/llmcompressor/modeling/deepseek_v3.py

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

src/llmcompressor/modeling/llama4.py

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

src/llmcompressor/modeling/moe_context.py

@@ -14,20 +14,22 @@
 
 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, RegistryMixin):
+class MoECalibrationModule(ABC, torch.nn.Module):
     """
     Abstract base class for MoE calibration modules.
 
@@ -60,6 +62,32 @@ def restore(self) -> 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,
@@ -99,10 +127,9 @@ 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_class_names:
+        if class_name in MOE_CALIBRATION_MODULES:
             modules_to_replace.append((name, module, class_name))
 
     # Step 2: Replace modules with progress bar
@@ -111,8 +138,8 @@ def moe_calibration_context(
         for name, module, class_name in tqdm(
             modules_to_replace, desc="Replacing MoE modules for calibration"
         ):
-            replacement = MoECalibrationModule.load_from_registry(
-                class_name,
+            calibration_cls = MOE_CALIBRATION_MODULES[class_name]
+            replacement = calibration_cls(
                 module,
                 model.config,
                 calibrate_all_experts=calibrate_all_experts,

src/llmcompressor/modeling/qwen3_moe.py

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

src/llmcompressor/modeling/qwen3_vl_moe.py

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