minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[FP8] Fix weight_scale shape to match with x_scale shape for per-tensor quant under torch.compile (#16356)

Browse Source
This commit is contained in:
Binyao Jiang
2026-01-05 13:55:26 -08:00
committed by GitHub
parent 4221b7c573
commit 9a9f996f8d
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/fp8_utils.py
+3 -1
View File
@@ -1041,7 +1041,9 @@ def apply_fp8_linear(
return _process_scaled_mm_output(output, input_2d.shape, output_shape) return _process_scaled_mm_output(output, input_2d.shape, output_shape)
if per_tensor_weights and per_tensor_activations: 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( output = torch._scaled_mm(
qinput, qinput,
weight, weight,