Post-quant fix via Dockerfile patch to process_weights_after_loading

Forward pre-hook approach didn't work — torch.compile and model
wrappers bypass hooks. Instead, patch vLLM's utils.py to call
model._post_quant_fix() at the end of process_weights_after_loading.
This guarantees the fix runs AFTER quant methods set up their attrs.

Dockerfile now patches:
  model_loader/utils.py → calls model._post_quant_fix() if it exists

DeepseekV4ForCausalLM._post_quant_fix() dequantizes attention
NVFP4 weights to BF16 and replaces quant_method.

Dockerfile

@@ -38,6 +38,26 @@ COPY vllm/nvfp4_cutedsl.py ${VLLM_MODELS_DIR}/nvfp4_cutedsl.py
 RUN sed -i 's/"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),/"DeepseekV32ForCausalLM": ("deepseek_v2", "DeepseekV3ForCausalLM"),\n    "DeepseekV4ForCausalLM": ("deepseek_v4", "DeepseekV4ForCausalLM"),/' \
     ${VLLM_MODELS_DIR}/registry.py
 
+# Patch process_weights_after_loading to call model._post_quant_fix() after quant setup
+ARG VLLM_LOADER_DIR=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/model_loader
+RUN python3 -c "
+import re
+path = '${VLLM_LOADER_DIR}/utils.py'.replace('\$', '')
+with open(path) as f:
+    src = f.read()
+# Add _post_quant_fix() call at end of process_weights_after_loading
+old = '    if model_config.quantization == \"torchao\":'
+new = '''    # Custom: allow models to run post-quant-init fixes
+    if hasattr(model, '_post_quant_fix'):
+        model._post_quant_fix()
+
+    if model_config.quantization == \"torchao\":'''
+src = src.replace(old, new, 1)
+with open(path, 'w') as f:
+    f.write(src)
+print('Patched process_weights_after_loading')
+"
+
 # Verify
 RUN python3 -c "import torch; print(f'PyTorch {torch.__version__} OK')" && \
     python3 -c "import vllm; print('vLLM OK')" && \

vllm/patches/deepseek_v4.py

@@ -2344,8 +2344,6 @@ class DeepseekV4ForCausalLM(nn.Module):
         self.model._convert_nvfp4_post_load()
         print("  Warming up tilelang kernels...", flush=True)
         self._warmup_tilelang()
-        print("  Registering post-quant-init fix...", flush=True)
-        self._register_post_quant_fix()
         print("  NVFP4 model ready ✓", flush=True)
 
         return loaded_params
@@ -2409,65 +2407,43 @@ class DeepseekV4ForCausalLM(nn.Module):
         del residual, fn, hc_scale, hc_base, x, post_mix, comb_mix
         torch.cuda.empty_cache()
 
+    def _post_quant_fix(self) -> None:
+        """Called by vLLM's process_weights_after_loading AFTER quant methods
+        have set up their attributes. Dequantizes attention NVFP4 weights to BF16
+        because FlashInferCutlassNvFp4LinearKernel uses broken input_global_scale_inv."""
+        from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+        
+        E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16)
+        fixed = 0
+        for layer_idx, layer in enumerate(self.model.layers):
+            attn = layer.attn
+            for proj_name in ["fused_wqa_wkv", "wq_b", "wo_b"]:
+                if not hasattr(attn, proj_name):
+                    continue
+                mod = getattr(attn, proj_name)
+                if not hasattr(mod, "weight") or mod.weight.dtype not in (torch.uint8, torch.int8):
+                    continue
+                # Dequantize to BF16
+                self.model._dequant_nvfp4_to_bf16(mod, E2M1_LUT)
+                # Replace quant method
+                mod.quant_method = UnquantizedLinearMethod()
+                # Clean up NVFP4 attributes
+                for attr in ("weight_scale", "weight_scale_2", "input_scale",
+                              "input_global_scale", "input_global_scale_inv",
+                              "weight_global_scale", "alpha",
+                              "weight_scale_inv"):
+                    if hasattr(mod, attr):
+                        try:
+                            delattr(mod, attr)
+                        except (AttributeError, TypeError):
+                            pass
+                fixed += 1
+        print(f"  [CLAWMINE] Post-quant fix: {fixed} attention projections → BF16 ✓", flush=True)
+
     def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
         return self.model.get_expert_mapping()
 
     def _register_post_quant_fix(self) -> None:
-        """Register a one-shot forward pre-hook that fixes attention NVFP4
-        activation scales AFTER process_weights_after_loading has run.
-        
-        process_weights_after_loading (called by vLLM's model loader AFTER
-        load_weights returns) sets up FlashInferCutlassNvFp4LinearKernel
-        with broken input_global_scale_inv from the checkpoint. This hook
-        runs on the first forward call (guaranteed after all init) and:
-        1. Dequantizes attention NVFP4 weights to BF16
-        2. Replaces quant_method with UnquantizedLinearMethod
-        3. Removes itself (one-shot)
-        """
-        import os
-        # Only needed for NVFP4 quantized models
-        quant_method_name = type(getattr(self.model.layers[0].attn.fused_wqa_wkv, 'quant_method', None)).__name__
-        if 'NvFp4' not in quant_method_name and 'nvfp4' not in quant_method_name.lower():
-            print("    No NVFP4 attention fix needed (quant_method={quant_method_name})", flush=True)
-            return
-
-        def _fix_attn_nvfp4(module, args):
-            """One-shot hook: dequant attention NVFP4 to BF16 after quant init."""
-            print("  [CLAWMINE] Running post-quant-init BF16 fix...", flush=True)
-            from vllm.model_executor.layers.linear import UnquantizedLinearMethod
-            
-            E2M1_LUT = torch.tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.bfloat16)
-            fixed = 0
-            for layer_idx, layer in enumerate(module.layers):
-                attn = layer.attn
-                for proj_name in ["fused_wqa_wkv", "wq_b", "wo_b"]:
-                    if not hasattr(attn, proj_name):
-                        continue
-                    mod = getattr(attn, proj_name)
-                    if not hasattr(mod, "weight") or mod.weight.dtype not in (torch.uint8, torch.int8):
-                        continue
-                    # Dequantize to BF16
-                    module._dequant_nvfp4_to_bf16(mod, E2M1_LUT)
-                    # Replace quant method (quant method already ran, won't overwrite again)
-                    mod.quant_method = UnquantizedLinearMethod()
-                    # Clean up NVFP4 attributes that might confuse forward
-                    for attr in ("weight_scale", "weight_scale_2", "input_scale",
-                                  "input_global_scale", "input_global_scale_inv",
-                                  "weight_global_scale", "alpha",
-                                  "weight_scale_inv"):
-                        if hasattr(mod, attr):
-                            try:
-                                delattr(mod, attr)
-                            except (AttributeError, TypeError):
-                                pass
-                    fixed += 1
-            
-            print(f"  [CLAWMINE] Fixed {fixed} attention projections → BF16", flush=True)
-            # Remove this hook (one-shot)
-            module._post_quant_fix_handle.remove()
-            del module._post_quant_fix_handle
-            print("  [CLAWMINE] Post-quant-init fix done ✓", flush=True)
-        
-        handle = self.model.register_forward_pre_hook(_fix_attn_nvfp4)
-        self._post_quant_fix_handle = handle
+        """No-op — we use _post_quant_fix() called from process_weights_after_loading."""
+        pass