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

[quantization] AWQ Marlin doesn't work when dtype is bfloat16 (#11494)

Browse Source 
Signed-off-by: Kai-Hsun Chen <khchen@x.ai>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
This commit is contained in:
Kai-Hsun Chen
2025-10-26 00:49:45 -07:00
committed by GitHub
parent 8491c794ad
commit 6371f7af27
4 changed files with 42 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
sgl-kernel/python/sgl_kernel/fused_moe.py
View File

@@ -80,6 +80,12 @@ def fused_marlin_moe(
assert w1.is_contiguous(), "Expert weights1 must be contiguous"
assert w2.is_contiguous(), "Expert weights2 must be contiguous"
assert hidden_states.dtype in [torch.float16, torch.bfloat16]
assert (
hidden_states.dtype == w1_scale.dtype
), f"moe_wna16_marlin_gemm assumes hidden_states.dtype ({hidden_states.dtype}) == w1_scale.dtype ({w1_scale.dtype})"
assert (
hidden_states.dtype == w2_scale.dtype
), f"moe_wna16_marlin_gemm assumes hidden_states.dtype ({hidden_states.dtype}) == w2_scale.dtype ({w2_scale.dtype})"
assert num_bits in [4, 8]
M, K = hidden_states.shape