[Quantization] Fix GPTQ scheme attachment broken by LinearBase.scheme default (#34962)

Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>

python/sglang/srt/layers/linear.py

@@ -155,8 +155,13 @@ class LinearBase(torch.nn.Module):
         quant_config: Quantization configure.
     """
 
-    # Set by quant methods that attach a per-layer scheme (e.g. Quark) inside
-    # get_quant_method(), which runs before create_weights() picks the loader.
+    # Set by quant methods that attach a per-layer scheme, eagerly in
+    # get_quant_method(), which runs before create_weights() picks the loader,
+    # or lazily inside create_weights() itself (GPTQ). The default is what lets
+    # callers probe with `is None`; a hasattr() probe answers "yes" once it
+    # exists. Schemes must stay plain objects -- nn.Module.__setattr__ files a
+    # Module value under self._modules, which this default then shadows on read.
+    # VocabParallelEmbedding and FusedMoE carry the same default.
     scheme = None
 
     def __init__(

python/sglang/srt/layers/moe/fused_moe_triton/layer.py

@@ -228,6 +228,9 @@ class FusedMoE(torch.nn.Module):
     # backend resolution distinguish them from routed experts.
     is_shared_fused_moe = False
 
+    # Attached by quant methods for a quantized MoE layer; see LinearBase.scheme.
+    scheme = None
+
     _skip_aiter_moe_shuffle: bool = False
 
     def __init__(
@@ -1082,7 +1085,7 @@ class FusedMoE(torch.nn.Module):
         # TODO (mgoin): check self.quant_method.quant_config.quant_format
         # against known CompressionFormat enum values that have this quality
         method = self.quant_method
-        if hasattr(self, "scheme"):
+        if self.scheme is not None:
             method = self.scheme
         if method.__class__.__name__ == "KTEPWrapperMethod":
             method = method.gpu_method
@@ -1349,7 +1352,7 @@ class FusedMoE(torch.nn.Module):
         # TODO: check self.quant_method.quant_config.quant_format
         # against known CompressionFormat enum values that have this quality
         method = self.quant_method
-        if hasattr(self, "scheme"):
+        if self.scheme is not None:
             method = self.scheme
         if isinstance(method, Fp8MoEMethod) and (
             get_moe_runner_backend().is_flashinfer_trtllm_routed()

python/sglang/srt/layers/quantization/gptq/gptq.py

@@ -471,7 +471,7 @@ class GPTQLinearMethod(LinearMethodBase):
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
-        if not hasattr(layer, "scheme"):
+        if layer.scheme is None:
             layer.scheme = self.quant_config.get_linear_scheme(layer)
         weight_loader = extra_weight_attrs.get("weight_loader")
         layer.scheme.create_weights(
@@ -511,7 +511,7 @@ class GPTQMoEMethod(FusedMoEMethodBase):
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
-        if not hasattr(layer, "scheme"):
+        if layer.scheme is None:
             layer.scheme = self.quant_config.get_moe_scheme(layer)
         layer.scheme.create_weights(
             layer=layer,
@@ -563,7 +563,7 @@ class GPTQMarlinLinearMethod(LinearMethodBase):
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ) -> None:
-        if not hasattr(layer, "scheme"):
+        if layer.scheme is None:
             layer.scheme = self.quant_config.get_linear_scheme(layer)
         weight_loader = extra_weight_attrs.get("weight_loader")
         layer.scheme.create_weights(
@@ -603,7 +603,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
         params_dtype: torch.dtype,
         **extra_weight_attrs,
     ):
-        if not hasattr(layer, "scheme"):
+        if layer.scheme is None:
             layer.scheme = self.quant_config.get_moe_scheme(layer)
         layer.scheme.create_weights(
             layer=layer,

python/sglang/srt/layers/vocab_parallel_embedding.py

@@ -224,6 +224,10 @@ class VocabParallelEmbedding(torch.nn.Module):
         prefix: full name of the layer in the state dict
     """  # noqa: E501
 
+    # Attached by quant methods for a quantized ParallelLMHead; see
+    # LinearBase.scheme.
+    scheme = None
+
     def __init__(
         self,
         num_embeddings: int,