From 9a9f996f8de7bc51a007ad3d79dc4b0a03b9a9d4 Mon Sep 17 00:00:00 2001 From: Binyao Jiang Date: Mon, 5 Jan 2026 13:55:26 -0800 Subject: [PATCH] [FP8] Fix weight_scale shape to match with x_scale shape for per-tensor quant under torch.compile (#16356) --- python/sglang/srt/layers/quantization/fp8_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/quantization/fp8_utils.py b/python/sglang/srt/layers/quantization/fp8_utils.py index 6da7ec66c..e7548dbce 100644 --- a/python/sglang/srt/layers/quantization/fp8_utils.py +++ b/python/sglang/srt/layers/quantization/fp8_utils.py @@ -1041,7 +1041,9 @@ def apply_fp8_linear( return _process_scaled_mm_output(output, input_2d.shape, output_shape) if per_tensor_weights and per_tensor_activations: - # Fused GEMM_DQ + # Fused GEMM_DQ; _scaled_mm with torch.compile requires len(weight_scale.shape) == len(x_scale.shape) + if weight_scale.ndim == 0 and x_scale.ndim == 1: + weight_scale = weight_scale.unsqueeze(0) output = torch._scaled_mm( qinput, weight,