[fix] fix DeepGEMM blackwell input quant & ut & fix style and log (#7247)

This commit is contained in:
JieXin Liang
2025-06-17 02:45:54 +08:00
committed by GitHub
parent e30ef368ab
commit 5ca07eed90
10 changed files with 285 additions and 31 deletions

View File

@@ -290,11 +290,12 @@ def sglang_per_token_group_quant_fp8(
x_s_mn, x_s_k = x_q_mn, x_q_k // 128
aligned_mn = align(x_s_mn, 4)
aligned_k = align(x_s_k, 4)
x_s = torch.empty(
# TODO(FIXME): Fix cuda kernel and recover here to empty.
x_s = torch.zeros(
(aligned_k // 4, aligned_mn),
device=x.device,
dtype=torch.int,
).permute(-1, -2)[:x_s_mn, :]
).transpose(0, 1)[:x_s_mn, :]
elif column_major_scales:
if scale_tma_aligned:
# TODO extract "align" function
@@ -768,7 +769,7 @@ def prepare_block_fp8_matmul_inputs(
if As.dtype == torch.float:
assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1]
elif Bs.dtype == torch.int:
elif As.dtype == torch.int:
assert (
triton.cdiv(triton.cdiv(A.shape[-1], block_k), 4) == As.shape[-1]
), f"{A.shape=} {As.shape=} {block_size=}"