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

[CPU] Add FP8 Bmm support (#9744)

Browse Source 
Co-authored-by: Fan Yin <1106310035@qq.com>
This commit is contained in:
blzheng
2026-03-19 13:19:48 +08:00
committed by GitHub
parent c2b01bd2fc
commit cd22aa27a9
14 changed files with 584 additions and 83 deletions
Download Patch File Download Diff File Expand all files Collapse all files

95
test/srt/cpu/test_bmm.py Normal file
View File

@@ -0,0 +1,95 @@
import itertools
import unittest
# TODO: use interface in cpu.py
import torch
import torch.nn as nn
from utils import precision
from sglang.srt.layers.quantization.fp8_utils import input_to_float8
from sglang.test.test_utils import CustomTestCase
torch.manual_seed(1234)
class Mod(nn.Module):
def __init__(self, input_channel, output_channel, has_bias):
super(Mod, self).__init__()
self.linear = torch.nn.Linear(input_channel, output_channel, has_bias)
def forward(self, x):
return self.linear(x)
class TestBmm(CustomTestCase):
M = [1, 2, 11, 111]
N = [128 + 32, 512]
K = [512 + 32, 128 + 32]
B = [1, 16, 17]
chunk = [True, False]
def _get_bmm_inputs(self, B, M, N, K, chunk, dtype):
if chunk:
mat1 = (
torch.randn(M, B, K + 64, dtype=dtype).narrow(2, 0, K).transpose_(0, 1)
)
mat2 = torch.randn(B, N, K, dtype=dtype).transpose_(1, 2)
mat3 = (
torch.randn(M, B, N + 64, dtype=dtype).narrow(2, 0, N).transpose_(0, 1)
)
else:
mat1 = torch.randn(M, B, K, dtype=dtype).transpose_(0, 1)
mat2 = torch.randn(B, N, K, dtype=dtype).transpose_(1, 2)
mat3 = torch.randn(M, B, N, dtype=dtype).transpose_(0, 1)
return mat1, mat2, mat3
def _bf16_bmm(self, B, M, N, K, chunk, dtype=torch.bfloat16):
mat1, mat2, mat3 = self._get_bmm_inputs(B, M, N, K, chunk, dtype)
ref = torch.bmm(mat1, mat2)
mat2_t = mat2.transpose_(1, 2)
mat3.zero_()
torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, mat2, False, None)
atol = rtol = precision[ref.dtype]
torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol)
packed_B = torch.ops.sgl_kernel.convert_weight_packed(mat2_t)
mat3.zero_()
torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, packed_B, True, None)
torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol)
def _fp8_bmm(self, B, M, N, K, chunk, dtype=torch.bfloat16):
mat1, mat2, mat3 = self._get_bmm_inputs(B, M, N, K, chunk, dtype)
mat2_q, mat2_s = input_to_float8(mat2)
ref = torch.bmm(mat1, mat2_q.to(torch.bfloat16)) * mat2_s
mat2_q_t = mat2_q.transpose_(1, 2).contiguous()
mat3.zero_()
atol = rtol = precision[ref.dtype]
torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, mat2_q_t, False, mat2_s)
torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol)
packed_B_q = torch.ops.sgl_kernel.convert_weight_packed(mat2_q_t)
mat3.zero_()
torch.ops.sgl_kernel.bmm_cpu(mat3, mat1, packed_B_q, True, mat2_s)
torch.testing.assert_close(ref, mat3, atol=atol, rtol=rtol)
def test_bmm(self):
for params in itertools.product(
self.B,
self.M,
self.N,
self.K,
self.chunk,
):
with self.subTest(
B=params[0],
M=params[1],
N=params[2],
K=params[3],
chunk=params[4],
):
self._bf16_bmm(*params)
self._fp8_bmm(*params)
if __name__ == "__main__":
unittest.main()

17
test/srt/cpu/test_qkv_proj_with_rope.py
View File

@@ -8,6 +8,7 @@ from utils import (
precision,
)
from sglang.srt.layers.quantization.fp8_utils import input_to_float8
from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb
from sglang.test.test_utils import CustomTestCase
@@ -186,6 +187,7 @@ class TestQKVProjWithROPE(CustomTestCase):
None,
None,
None,
None,
True,
None,
)
@@ -203,6 +205,7 @@ class TestQKVProjWithROPE(CustomTestCase):
False,
None,
None,
None,
True,
None,
q_lora_rank,
@@ -274,6 +277,7 @@ class TestQKVProjWithROPE(CustomTestCase):
w1_s,
w2_s,
w3_s,
None,
True,
None,
)
@@ -294,6 +298,7 @@ class TestQKVProjWithROPE(CustomTestCase):
False,
fused_weight_s,
w2_s,
None,
True,
None,
q_lora_rank,
@@ -320,6 +325,7 @@ class TestQKVProjWithROPE(CustomTestCase):
torch.randn(num_heads * qk_head_dim, q_lora_rank, dtype=dtype) * 0.1
)
w_kc = torch.randn(num_heads, kv_lora_rank, qk_nope_head_dim, dtype=dtype) * 0.1
w_kc_q, w_kc_s = input_to_float8(w_kc)
kv_a_proj_weight = (
torch.randn(kv_lora_rank + qk_rope_head_dim, hidden_size, dtype=dtype) * 0.1
)
@@ -350,13 +356,14 @@ class TestQKVProjWithROPE(CustomTestCase):
) = convert_weight(
kv_a_proj_weight, [scale_block_size_N, scale_block_size_K], torch.bfloat16
)
w_kc_dq = w_kc_q.to(torch.bfloat16) * w_kc_s
q_ref, k_ref, v_ref = native_torch(
q_input,
hidden_states,
q_a_proj_weight_dq,
norm_weight1,
q_b_proj_weight_dq,
w_kc.transpose(1, 2),
w_kc_dq.transpose(1, 2),
kv_a_proj_with_mqa_weight_dq,
norm_weight2,
pos,
@@ -367,13 +374,13 @@ class TestQKVProjWithROPE(CustomTestCase):
fp8_kv_a_proj_with_mqa_weight_packed = convert_weight_packed(
fp8_kv_a_proj_with_mqa_weight
)
w_kc = convert_weight_packed(w_kc)
w_kc_q = convert_weight_packed(w_kc_q)
q_out, k_out, v_out = qkv_proj_with_rope(
hidden_states,
fp8_q_a_proj_weight_packed,
fp8_q_b_proj_weight_packed,
fp8_kv_a_proj_with_mqa_weight_packed,
w_kc,
w_kc_q,
norm_weight1,
norm_weight2,
pos,
@@ -384,6 +391,7 @@ class TestQKVProjWithROPE(CustomTestCase):
q_a_proj_weight_scale_inv.float(),
q_b_proj_weight_scale_inv.float(),
kv_a_proj_with_mqa_weight_scale_inv.float(),
w_kc_s,
True,
[scale_block_size_N, scale_block_size_K],
)
@@ -399,7 +407,7 @@ class TestQKVProjWithROPE(CustomTestCase):
hidden_states,
fused_weight_packed,
fp8_q_b_proj_weight_packed,
w_kc,
w_kc_q,
norm_weight1,
norm_weight2,
pos,
@@ -409,6 +417,7 @@ class TestQKVProjWithROPE(CustomTestCase):
True,
fused_weight_s.float(),
q_b_proj_weight_scale_inv.float(),
w_kc_s,
True,
[scale_block_size_N, scale_block_size_K],
q_lora_rank,

1
test/srt/run_suite.py
View File

@@ -49,6 +49,7 @@ suite_xeon = {
"per-commit-cpu": [
TestFile("cpu/test_activation.py"),
TestFile("cpu/test_binding.py"),
TestFile("cpu/test_bmm.py"),
TestFile("cpu/test_causal_conv1d.py"),
TestFile("cpu/test_cpu_graph.py"),
TestFile("cpu/test_decode.py"),