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

Create col-major and tma-aligned x_scale for deep_gemm.gemm_fp8_fp8_bf16_nt (#4515)

Browse Source 
Co-authored-by: Zhang Kaihong <zhangkaihong.zkh@alibaba-inc.com>
This commit is contained in:
strgrb
2025-03-19 15:02:43 +08:00
committed by GitHub
parent 90532b7627
commit f9c53cbb42
3 changed files with 28 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/sglang/srt/layers/quantization/fp8_utils.py
View File

@@ -3,6 +3,7 @@ from typing import List, Optional, Tuple
import torch
from sglang.srt.layers.quantization.fp8_kernel import (
_enable_jit_deepgemm,
per_token_group_quant_fp8,
static_quant_fp8,
w8a8_block_fp8_matmul,
@@ -129,9 +130,17 @@ def apply_w8a8_block_fp8_linear(
)
gemm_a8w8_blockscale(q_input, weight, x_scale, weight_scale, output)
else:
q_input, x_scale = per_token_group_quant_fp8(
input_2d, block_size[1], column_major_scales=False
)
if _enable_jit_deepgemm:
q_input, x_scale = per_token_group_quant_fp8(
input_2d,
block_size[1],
column_major_scales=True,
scale_tma_aligned=True,
)
else:
q_input, x_scale = per_token_group_quant_fp8(
input_2d, block_size[1], column_major_scales=False
)
output = w8a8_block_fp8_matmul(
q_input, weight, x_scale, weight_scale, block_size, output_dtype=input.dtype
)