[fix] fix DeepGEMM blackwell input quant & ut & fix style and log (#7247)
This commit is contained in:
@@ -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=}"
|
||||
|
||||
Reference in New Issue
Block a user