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

[Intel XPU] Enable (biased) grouped topk for xpu (#31126)

Browse Source
This commit is contained in:
gaopengff
2026-07-20 09:35:08 +08:00
committed by GitHub
parent 5325cee7ea
commit bab1dd0d12
2 changed files with 231 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/topk.py
+90 -4
View File
@@ -1006,6 +1006,59 @@ def grouped_topk_cpu(
)
def grouped_topk_xpu(
hidden_states: torch.Tensor,
gating_output: torch.Tensor,
topk: int,
renormalize: bool,
num_expert_group: Optional[int] = None,
topk_group: Optional[int] = None,
num_fused_shared_experts: int = 0,
routed_scaling_factor: Optional[float] = None,
apply_routed_scaling_factor_on_output: Optional[bool] = False,
scoring_func: str = "softmax",
):
num_experts = gating_output.shape[1]
experts_per_group = (
num_experts // num_expert_group if num_expert_group else num_experts
)
# moe_fused_gate kernel ensures that num_experts/num_expert_group does not exceed MAX_VPT=32 now.
if experts_per_group <= 32 and is_power_of_two(num_experts):
from sgl_kernel import moe_fused_gate
return moe_fused_gate(
gating_output.to(torch.float32),
None, # without bias
num_expert_group,
topk_group,
topk,
renormalize=renormalize,
scoring_func=scoring_func,
num_fused_shared_experts=num_fused_shared_experts,
routed_scaling_factor=(
routed_scaling_factor if routed_scaling_factor is not None else 1.0
),
apply_routed_scaling_factor_on_output=bool(
apply_routed_scaling_factor_on_output
),
)
# use default implementation
return grouped_topk_gpu(
hidden_states,
gating_output,
topk,
renormalize,
num_expert_group,
topk_group,
num_fused_shared_experts,
routed_scaling_factor,
apply_routed_scaling_factor_on_output,
scoring_func,
)
@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu)
def kimi_k2_biased_topk_impl(
hidden_states: torch.Tensor,
@@ -1539,8 +1592,15 @@ def biased_grouped_topk_gpu(
and num_experts <= 256
and topk <= 8
):
if not apply_routed_scaling_factor_on_output:
scaling = 1.0
scale = (
routed_scaling_factor
if (
apply_routed_scaling_factor_on_output
and routed_scaling_factor is not None
)
else 1.0
)
num_tokens = gating_output.shape[0]
@@ -1560,8 +1620,34 @@ def biased_grouped_topk_gpu(
gating_output,
renormalize,
correction_bias,
scale,
)
return topk_values, topk_indices
elif (
_is_xpu
# moe_fused_gate kernel ensures that num_experts/num_expert_group does not exceed MAX_VPT=32 now.
and experts_per_group <= 32
and is_power_of_two(num_experts)
):
from sgl_kernel import moe_fused_gate
return moe_fused_gate(
gating_output.to(torch.float32),
correction_bias.to(torch.float32),
num_expert_group,
topk_group,
topk,
renormalize=renormalize,
scoring_func="sigmoid",
num_fused_shared_experts=num_fused_shared_experts,
routed_scaling_factor=(
routed_scaling_factor if routed_scaling_factor is not None else 1.0
),
apply_routed_scaling_factor_on_output=bool(
apply_routed_scaling_factor_on_output
),
)
return topk_values * scaling, topk_indices
else:
return biased_grouped_topk_impl(
@@ -1613,7 +1699,7 @@ if _is_cpu and _is_cpu_amx_available:
fused_topk = fused_topk_cpu
else:
biased_grouped_topk = biased_grouped_topk_gpu
grouped_topk = grouped_topk_gpu
grouped_topk = grouped_topk_xpu if _is_xpu else grouped_topk_gpu
fused_topk_native = fused_topk_torch_native
test/registered/xpu/test_topk.py
+141
View File
@@ -8,12 +8,31 @@ from sglang.srt.layers.moe.topk import (
from sglang.srt.layers.moe.topk import (
biased_grouped_topk_impl as native_biased_grouped_topk,
)
from sglang.srt.layers.moe.topk import grouped_topk_gpu as native_grouped_topk
from sglang.srt.layers.moe.topk import (
grouped_topk_xpu,
)
from sglang.test.ci.ci_register import register_xpu_ci
from sglang.test.test_utils import CustomTestCase
register_xpu_ci(est_time=5, suite="stage-b-test-1-gpu-xpu")
def _scatter_by_expert(
weights: torch.Tensor, indices: torch.Tensor, num_columns: int
) -> torch.Tensor:
"""Scatter (weight, id) pairs into a dense ``[M, num_columns]`` tensor.
Makes the comparison independent of the per-row slot order, so the test does
not depend on how ties between equal scores are broken.
"""
dense = torch.zeros(
(weights.shape[0], num_columns), dtype=torch.float32, device=weights.device
)
dense.scatter_(1, indices.long(), weights.float())
return dense
# Nemotron-3 uses biased_grouped_topk
class TestBiasedGroupedTopK(CustomTestCase):
def _run_single_test(
@@ -96,6 +115,128 @@ class TestBiasedGroupedTopK(CustomTestCase):
routed_scaling_factor,
)
def test_biased_grouped_topk(self):
# DeepSeek-V3 style grouped routing shape
E_num = 256
num_expert_group = 8
topk_value = 8
topk_group = 4
gating_dtype = torch.bfloat16
bias_dtype = torch.float32
renormalize = True
routed_scaling_factor = 2.5
torch.manual_seed(1024)
device = torch.device("xpu")
bs = [1, 2, 4, 8]
seq_len = 1024
num_tokens = [b * seq_len for b in bs]
num_fused_shared_experts_list = [0, 1]
for M in num_tokens:
for num_fused_shared_experts in num_fused_shared_experts_list:
topk_routed = topk_value - num_fused_shared_experts
hidden_states = torch.randn(M, 100, dtype=torch.bfloat16, device=device)
gating_output = torch.randn(M, E_num, dtype=gating_dtype, device=device)
correction_bias = torch.randn(E_num, dtype=bias_dtype, device=device)
ref_topk_weights, ref_topk_ids = native_biased_grouped_topk(
hidden_states.float(),
gating_output.float(),
correction_bias,
topk_value,
renormalize,
num_expert_group,
topk_group,
num_fused_shared_experts,
routed_scaling_factor=routed_scaling_factor,
)
# fused version
topk_weights, topk_ids = biased_grouped_topk_gpu(
hidden_states,
gating_output,
correction_bias,
topk_value,
renormalize,
num_expert_group,
topk_group,
num_fused_shared_experts,
routed_scaling_factor,
)
torch.testing.assert_close(
_scatter_by_expert(
topk_weights[:, :topk_routed], topk_ids[:, :topk_routed], E_num
),
_scatter_by_expert(
ref_topk_weights[:, :topk_routed],
ref_topk_ids[:, :topk_routed],
E_num,
),
)
def test_grouped_topk(self):
# DeepSeek-V3 style grouped routing shape
E_num = 256
num_expert_group = 8
topk_value = 8
topk_group = 4
gating_dtype = torch.bfloat16
renormalize = True
routed_scaling_factor = 2.5
torch.manual_seed(1024)
device = torch.device("xpu")
bs = [1]
seq_len = 1024
num_tokens = [b * seq_len for b in bs]
num_fused_shared_experts_list = [0, 1]
for M in num_tokens:
for num_fused_shared_experts in num_fused_shared_experts_list:
topk_routed = topk_value - num_fused_shared_experts
hidden_states = torch.randn(M, 100, dtype=torch.bfloat16, device=device)
gating_output = torch.randn(M, E_num, dtype=gating_dtype, device=device)
ref_topk_weights, ref_topk_ids = native_grouped_topk(
hidden_states.float(),
gating_output.float(),
topk_value,
renormalize,
num_expert_group,
topk_group,
num_fused_shared_experts,
routed_scaling_factor=routed_scaling_factor,
)
# fused version
topk_weights, topk_ids = grouped_topk_xpu(
hidden_states,
gating_output,
topk_value,
renormalize,
num_expert_group,
topk_group,
num_fused_shared_experts,
routed_scaling_factor,
)
torch.testing.assert_close(
_scatter_by_expert(
topk_weights[:, :topk_routed], topk_ids[:, :topk_routed], E_num
),
_scatter_by_expert(
ref_topk_weights[:, :topk_routed],
ref_topk_ids[:, :topk_routed],
E_num,
),
)
if __name__ == "__main__":
unittest.main()