minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Optimize MiniMax-M2.7 on CPU (#31956)

Browse Source
This commit is contained in:
Xinguo Zhu
2026-08-13 15:04:16 +08:00
committed by GitHub
parent 889c2f31aa
commit 3f6ef01322
8 changed files with 687 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/cpu/test_norm.py
+81
View File
@@ -234,6 +234,87 @@ class TestFusedRMSNormGated:
torch.testing.assert_close(ref_out, out, atol=atol, rtol=rtol)
class TestFusedQKRMSNorm:
@pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS)
@pytest.mark.parametrize(
"batch_size,q_size,k_size,v_size",
[(1, 256, 64, 64), (17, 512, 128, 128)],
)
def test_fused_qk_rmsnorm(
self, batch_size: int, q_size: int, k_size: int, v_size: int, dtype
):
"""Q and K split views must be normalized over their distinct full widths."""
qkv = torch.randn([batch_size, q_size + k_size + v_size], dtype=dtype)
q, k, _ = qkv.split([q_size, k_size, v_size], dim=-1)
q_weight = torch.randn(q_size, dtype=dtype)
k_weight = torch.randn(k_size, dtype=dtype)
q_out, k_out = torch.ops.sgl_kernel.fused_qk_rmsnorm_cpu(
q, k, q_weight, k_weight, eps
)
ref_q_out = TestNorm()._forward_native(q, q_weight, eps)
ref_k_out = TestNorm()._forward_native(k, k_weight, eps)
atol = rtol = precision[dtype]
torch.testing.assert_close(q_out, ref_q_out, atol=atol, rtol=rtol)
torch.testing.assert_close(k_out, ref_k_out, atol=atol, rtol=rtol)
@pytest.mark.parametrize("dtype", DTYPES, ids=DTYPE_IDS)
@pytest.mark.parametrize(
"batch_size,q_size,k_size,tp_world_size",
[(1, 256, 64, 2), (17, 512, 128, 4)],
)
def test_fused_qk_rmsnorm_tp(
self,
batch_size: int,
q_size: int,
k_size: int,
tp_world_size: int,
dtype,
):
q = torch.randn([batch_size, q_size], dtype=dtype)
k = torch.randn([batch_size, k_size], dtype=dtype)
q_weight = torch.randn(q_size, dtype=dtype)
k_weight = torch.randn(k_size, dtype=dtype)
q_shards = q.chunk(tp_world_size, dim=-1)
k_shards = k.chunk(tp_world_size, dim=-1)
q_weight_shards = q_weight.chunk(tp_world_size)
k_weight_shards = k_weight.chunk(tp_world_size)
local_sum_sq = [
torch.ops.sgl_kernel.fused_qk_rmsnorm_sumsq_cpu(q_shard, k_shard)
for q_shard, k_shard in zip(q_shards, k_shards)
]
global_sum_sq = torch.stack(local_sum_sq).sum(dim=0)
shard_outputs = [
torch.ops.sgl_kernel.fused_qk_rmsnorm_apply_from_stats_cpu(
q_shard,
k_shard,
q_weight_shard,
k_weight_shard,
global_sum_sq,
tp_world_size,
eps,
)
for q_shard, k_shard, q_weight_shard, k_weight_shard in zip(
q_shards, k_shards, q_weight_shards, k_weight_shards
)
]
q_out = torch.cat([output[0] for output in shard_outputs], dim=-1)
k_out = torch.cat([output[1] for output in shard_outputs], dim=-1)
ref_q_out = TestNorm()._forward_native(q, q_weight, eps)
ref_k_out = TestNorm()._forward_native(k, k_weight, eps)
atol = rtol = precision[dtype]
torch.testing.assert_close(q_out, ref_q_out, atol=atol, rtol=rtol)
torch.testing.assert_close(k_out, ref_k_out, atol=atol, rtol=rtol)
assert global_sum_sq.shape == (batch_size, 2)
assert global_sum_sq.dtype == torch.float32
class TestLayerNorm:
def _forward_native(
test/registered/cpu/test_topk.py
+157
View File
@@ -1,3 +1,4 @@
import itertools
import unittest
import torch
@@ -207,6 +208,89 @@ class TestTopK(CustomTestCase):
self._run_single_test(123, 256, 4, renormalize, torch.bfloat16)
self._run_single_test(123, 160, 6, renormalize, torch.bfloat16)
def test_topk_softmax_mixed_input_dtypes(self):
torch.manual_seed(0)
hidden_states = torch.randn((17, 16), dtype=torch.bfloat16)
gating_output = torch.randn((17, 128), dtype=torch.float32)
correction_bias = torch.randn(128, dtype=torch.float32)
topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu(
hidden_states=hidden_states,
gating_output=gating_output,
topk=8,
renormalize=True,
correction_bias=correction_bias,
)
scores = torch.softmax(gating_output, dim=-1)
expected_ids = torch.topk(
scores + correction_bias.unsqueeze(0), k=8, dim=-1
).indices
expected_weights = scores.gather(1, topk_ids.to(torch.int64))
expected_weights /= expected_weights.sum(dim=-1, keepdim=True)
self.assertEqual(
torch.sort(topk_ids.to(torch.int64), dim=-1).values.tolist(),
torch.sort(expected_ids, dim=-1).values.tolist(),
)
torch.testing.assert_close(topk_weights, expected_weights)
def test_topk_softmax_with_correction_bias(self):
"""Bias must affect expert selection without becoming a routing weight."""
for num_tokens, num_experts, topk, with_bias, renormalize in itertools.product(
[1, 17, 128],
[16, 128, 384, 512],
[1, 2, 4, 8],
[False, True],
[False, True],
):
torch.manual_seed(0)
hidden_states = torch.randn((num_tokens, 16), dtype=torch.bfloat16)
gating_output = torch.randn((num_tokens, num_experts), dtype=torch.bfloat16)
correction_bias = torch.randn(num_experts) if with_bias else None
topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu(
hidden_states=hidden_states,
gating_output=gating_output,
topk=topk,
renormalize=renormalize,
correction_bias=correction_bias,
)
scores = torch.softmax(gating_output.float(), dim=-1)
scores_for_choice = scores
if correction_bias is not None:
scores_for_choice = scores_for_choice + correction_bias.unsqueeze(0)
expected_choice_scores = torch.topk(
scores_for_choice, k=topk, dim=-1, sorted=True
).values
selected_choice_scores = torch.sort(
scores_for_choice.gather(1, topk_ids.to(torch.int64)),
dim=-1,
descending=True,
).values
expected_weights = scores.gather(1, topk_ids.to(torch.int64))
if renormalize:
expected_weights = expected_weights / expected_weights.sum(
dim=-1, keepdim=True
)
self.assertEqual(topk_ids.dtype, torch.int32)
self.assertEqual(topk_weights.dtype, torch.float32)
self.assertTrue(torch.all((topk_ids >= 0) & (topk_ids < num_experts)))
sorted_ids = torch.sort(topk_ids, dim=-1).values
self.assertTrue(torch.all(sorted_ids[:, 1:] != sorted_ids[:, :-1]))
torch.testing.assert_close(
selected_choice_scores,
expected_choice_scores,
atol=1e-4,
rtol=1e-4,
)
torch.testing.assert_close(
topk_weights, expected_weights, atol=1e-4, rtol=1e-4
)
class TestCustomTopK(CustomTestCase):
def _run_single_test(
@@ -251,6 +335,79 @@ class TestCustomTopK(CustomTestCase):
123, 32, 1, False, torch.bfloat16, native_custom_f, fused_custom_f
)
def test_topk_sigmoid_with_correction_bias(self):
"""Biased scores must select experts while returned weights stay unbiased."""
for num_tokens, num_experts, topk, with_bias, renormalize in itertools.product(
[1, 17, 128],
[16, 128, 256, 384, 512],
[1, 2, 4, 8],
[False, True],
[False, True],
):
torch.manual_seed(0)
hidden_states = torch.randn((num_tokens, 16), dtype=torch.bfloat16)
gating_output = torch.randn((num_tokens, num_experts), dtype=torch.bfloat16)
correction_bias = torch.randn(num_experts) if with_bias else None
topk_weights, topk_ids = torch.ops.sgl_kernel.topk_sigmoid_cpu(
hidden_states=hidden_states,
gating_output=gating_output,
topk=topk,
renormalize=renormalize,
correction_bias=correction_bias,
)
scores = torch.sigmoid(gating_output.float())
scores_for_choice = scores
if correction_bias is not None:
scores_for_choice = scores_for_choice + correction_bias.unsqueeze(0)
expected_choice_scores = torch.topk(
scores_for_choice, k=topk, dim=-1
).values
selected_choice_scores = torch.sort(
scores_for_choice.gather(1, topk_ids.to(torch.int64)),
dim=-1,
descending=True,
).values
expected_weights = scores.gather(1, topk_ids.to(torch.int64))
if renormalize:
expected_weights /= expected_weights.sum(dim=-1, keepdim=True)
self.assertEqual(topk_ids.dtype, torch.int32)
self.assertEqual(topk_weights.dtype, torch.float32)
self.assertTrue(torch.equal(selected_choice_scores, expected_choice_scores))
torch.testing.assert_close(
topk_weights, expected_weights, atol=1e-4, rtol=1e-4
)
def test_topk_sigmoid_mixed_input_dtypes(self):
torch.manual_seed(0)
hidden_states = torch.randn((17, 16), dtype=torch.bfloat16)
gating_output = torch.randn((17, 256), dtype=torch.float32)
topk_weights, topk_ids = torch.ops.sgl_kernel.topk_sigmoid_cpu(
hidden_states=hidden_states,
gating_output=gating_output,
topk=8,
renormalize=True,
correction_bias=None,
)
scores = torch.sigmoid(gating_output)
expected_ids = torch.topk(scores, k=8, dim=-1).indices
expected_weights = scores.gather(1, topk_ids.to(torch.int64))
expected_weights /= expected_weights.sum(dim=-1, keepdim=True)
self.assertTrue(
torch.equal(
torch.sort(topk_ids.to(torch.int64), dim=-1).values,
torch.sort(expected_ids, dim=-1).values,
)
)
torch.testing.assert_close(topk_weights, expected_weights, atol=1e-5, rtol=1e-5)
if __name__ == "__main__":
unittest.main()