Fix block wise fp8 torch compile (#3232)

2025-01-31 19:56:02 +08:00
parent 734daedd8f
commit c02e313914
1 changed files with 7 additions and 0 deletions
--- a/python/sglang/srt/layers/quantization/fp8.py
+++ b/python/sglang/srt/layers/quantization/fp8.py
@@ -290,6 +290,13 @@ class Fp8LinearMethod(LinearMethodBase):
                    weight_scale, requires_grad=False
                )
                layer.input_scale = None
+            else:
+                layer.weight = torch.nn.Parameter(
+                    layer.weight.data, requires_grad=False
+                )
+                layer.weight_scale_inv = torch.nn.Parameter(
+                    layer.weight_scale_inv.data, requires_grad=False
+                )
            return
        layer.weight = torch.nn.Parameter(layer.weight.data, requires_grad=False)
        # If checkpoint not serialized fp8, quantize the weights.