[Bugfix][Model Loader] Restore kernel state post model loading via TensorizerLoader.

vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py

@@ -90,6 +90,12 @@ def input_quant_key(self) -> QuantKey | None:
     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         raise NotImplementedError
 
+    def initialize_runtime_state_after_loading(
+        self, layer: torch.nn.Module
+    ) -> None:
+        """Restore transient state for weights already in kernel format."""
+        return
+
     @abstractmethod
     def apply_weights(
         self,

vllm/model_executor/kernels/linear/scaled_mm/marlin.py

@@ -9,6 +9,9 @@
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     process_fp8_weight_block_strategy,
 )
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+    marlin_make_workspace_new,
+)
 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
     apply_fp8_marlin_linear,
     is_fp8_marlin_supported,
@@ -79,6 +82,19 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         )
         del layer.input_scale
 
+    def initialize_runtime_state_after_loading(
+        self, layer: torch.nn.Module
+    ) -> None:
+        if layer.weight.dtype != torch.int32:
+            raise ValueError(
+                "Native FP8 Marlin weights must retain their packed int32 dtype; "
+                f"got {layer.weight.dtype}."
+            )
+        layer.workspace = marlin_make_workspace_new(
+            layer.weight.device,
+            existing=getattr(layer, "workspace", None),
+        )
+
     def apply_weights(
         self,
         layer: torch.nn.Module,

vllm/model_executor/layers/quantization/base_config.py

@@ -84,6 +84,16 @@ def process_weights_after_loading(self, layer: nn.Module) -> None:
         """
         return
 
+    def initialize_runtime_state_after_loading(
+        self, layer: nn.Module
+    ) -> None:
+        """Initialize non-serialized runtime state.
+
+        Unlike process_weights_after_loading, this must not transform
+        registered parameters or buffers.
+        """
+        return
+
 
 def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> bool:
     """Not all quant methods have embedding implemented, so we need to check that

vllm/model_executor/layers/quantization/fp8.py

@@ -415,6 +415,15 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
 
         self.fp8_linear.process_weights_after_loading(layer)
 
+    def initialize_runtime_state_after_loading(
+        self, layer: torch.nn.Module
+    ) -> None:
+        if not self.use_marlin:
+            return
+
+        self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype
+        self.fp8_linear.initialize_runtime_state_after_loading(layer)
+
     def apply(
         self,
         layer: torch.nn.Module,

vllm/model_executor/model_loader/tensorizer.py

@@ -527,7 +527,9 @@ def init_tensorizer_model(
 
 
 def deserialize_tensorizer_model(
-    model: nn.Module, tensorizer_config: TensorizerConfig
+    model: nn.Module,
+    tensorizer_config: TensorizerConfig,
+    preserve_serialized_dtype: bool = False,
 ) -> None:
     tensorizer_args = tensorizer_config._construct_tensorizer_args()
     if not is_valid_deserialization_uri(tensorizer_config.tensorizer_uri):
@@ -547,7 +549,7 @@ def deserialize_tensorizer_model(
         ) as stream,
         TensorDeserializer(
             stream,
-            dtype=tensorizer_config.dtype,
+            dtype=None if preserve_serialized_dtype else tensorizer_config.dtype,
             device=f"{device_type}:{device_index}",
             **tensorizer_args.deserialization_kwargs,
         ) as deserializer,

vllm/model_executor/model_loader/tensorizer_loader.py

@@ -10,6 +10,7 @@
 from vllm.config import ModelConfig, ParallelConfig, VllmConfig
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
 from vllm.model_executor.model_loader.base_loader import BaseModelLoader
 from vllm.model_executor.model_loader.tensorizer import (
     TensorizerConfig,
@@ -22,6 +23,7 @@
 from vllm.model_executor.model_loader.utils import (
     get_model_architecture,
     initialize_model,
+    initialize_runtime_state_after_loading,
 )
 from vllm.utils.torch_utils import set_default_torch_dtype
 
@@ -108,7 +110,16 @@ def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
         for serializing vLLM models."""
         if is_vllm_tensorized(self.tensorizer_config):
             tensorizer_config = self._patch_tensorizer_config(model_config)
-            deserialize_tensorizer_model(model, tensorizer_config)
+            # Marlin's packed FP8 weights include integer tensors. A global
+            # deserialization dtype would cast them to the model dtype.
+            has_fp8_marlin = any(
+                isinstance(method, Fp8LinearMethod) and method.use_marlin
+                for module in model.modules()
+                if (method := getattr(module, "quant_method", None)) is not None
+            )
+            deserialize_tensorizer_model(
+                model, tensorizer_config, preserve_serialized_dtype=has_fp8_marlin
+            )
         else:
             model.load_weights(self._get_weights_iterator())
 
@@ -134,8 +145,13 @@ def load_model(
                     model = init_tensorizer_model(
                         tensorizer_config=tensorizer_config, vllm_config=vllm_config
                     )
+
             self.load_weights(model, model_config)
-            return model
+            initialize_runtime_state_after_loading(
+                model,
+                torch.device(device_config.device),
+            )
+            return model.eval()
         return self._load_model_serialized_cpu(vllm_config=vllm_config, prefix=prefix)
 
     @staticmethod

vllm/model_executor/model_loader/utils.py

@@ -199,6 +199,17 @@ def process_weights_after_loading(
         set_torchao_reload_attrs(model, model_config)
 
 
+def initialize_runtime_state_after_loading(
+    model: nn.Module,
+    target_device: torch.device,
+) -> None:
+    for _, module in model.named_modules():
+        quant_method = getattr(module, "quant_method", None)
+        if isinstance(quant_method, QuantizeMethodBase):
+            with device_loading_context(module, target_device):
+                quant_method.initialize_runtime_state_after_loading(module)
+
+
 @contextmanager
 def device_loading_context(module: torch.nn.Module, target_device: torch.device):
     if target_device.type == "cpu":