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

[Intel XPU] Add xpu pass for biased_topk and hash_topk (#33323)

Browse Source 
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
gaopengff
2026-08-24 12:18:22 +08:00
committed by GitHub
co-authored by Ma Mingfei
parent 1daa94a069
commit 56834422a1
3 changed files with 233 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/hash_topk.py
+37 -2
View File
@@ -23,12 +23,13 @@ from sglang.srt.layers.moe.topk import (
)
from sglang.srt.layers.moe.utils import has_per_rank_fused_shared_slots
from sglang.srt.runtime_context import get_exec
from sglang.srt.utils import is_hip, is_npu
from sglang.srt.utils import is_hip, is_npu, is_xpu
logger = logging.getLogger(__name__)
_is_hip = is_hip()
_is_npu = is_npu()
_is_xpu = is_xpu()
class HashTopK(nn.Module):
@@ -177,6 +178,38 @@ class HashTopK(nn.Module):
return topk_weights, topk_ids
def _forward_xpu(
self, router_logits: torch.Tensor, input_ids: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
# The XPU 'hash_topk' kernel currently supports the 'sqrtsoftplus' score func only.
# Other score funcs fall back to the torch implementation; more will be supported in the future.
if self.score_func == "sqrtsoftplus":
from sgl_kernel import hash_topk
num_tokens = router_logits.size(0)
topk_routed = self.tid2eid.size(1)
topk_fused = topk_routed + self.num_fused_shared_experts
topk_ids = torch.empty(
(num_tokens, topk_fused), dtype=torch.int32, device=router_logits.device
)
topk_weights = torch.empty(
(num_tokens, topk_fused),
dtype=torch.float32,
device=router_logits.device,
)
hash_topk(
router_logits,
input_ids,
self.tid2eid,
topk_weights,
topk_ids,
self.routed_scaling_factor,
self.score_func,
)
return topk_weights, topk_ids
else:
return self._forward_torch(router_logits, input_ids)
def forward(
self,
hidden_states: torch.Tensor,
@@ -189,7 +222,9 @@ class HashTopK(nn.Module):
input_ids.shape[0] == hidden_states.shape[0] == router_logits.shape[0]
), f"{input_ids.shape=} {hidden_states.shape=} {router_logits.shape=}"
if envs.SGLANG_OPT_USE_FUSED_HASH_TOPK.get():
if _is_xpu:
topk_weights, topk_ids = self._forward_xpu(router_logits, input_ids)
elif envs.SGLANG_OPT_USE_FUSED_HASH_TOPK.get():
from sglang.kernels.ops.attention.dsv4 import hash_topk
topk_weights, topk_ids = hash_topk(
python/sglang/srt/layers/moe/topk.py
+43 -1
View File
@@ -1291,6 +1291,47 @@ def biased_topk_jit_kernel_impl(
return topk_weights, topk_ids
def biased_topk_xpu(
hidden_states: torch.Tensor,
gating_output: torch.Tensor,
correction_bias: torch.Tensor,
topk: int,
renormalize: bool,
scoring_func: str = "sigmoid",
num_fused_shared_experts: int = 0,
routed_scaling_factor: Optional[float] = None,
num_token_non_padded: Optional[torch.Tensor] = None,
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
apply_routed_scaling_factor_on_output: Optional[bool] = False,
) -> Tuple[torch.Tensor, torch.Tensor]:
assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
num_rows, _ = gating_output.shape
device = gating_output.device
output = torch.empty(num_rows, topk, dtype=torch.float32, device=device)
indices = torch.empty(num_rows, topk, dtype=torch.int32, device=device)
from sgl_kernel import biased_topk
biased_topk(
gating_output,
correction_bias,
output,
indices,
topk,
scoring_func,
num_fused_shared_experts,
renormalize,
routed_scaling_factor=(routed_scaling_factor if routed_scaling_factor else 1.0),
apply_routed_scaling_factor_on_output=bool(
apply_routed_scaling_factor_on_output
),
)
return output, indices
@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu)
def biased_grouped_topk_impl(
hidden_states: torch.Tensor,
@@ -2212,7 +2253,8 @@ def select_experts(
assert not apply_routed_scaling_factor_on_output, "Not implemented"
if scoring_func == "sqrtsoftplus" or scoring_func == "sigmoid":
topk_weights, topk_ids = biased_topk_jit_kernel_impl(
_biased_topk = biased_topk_xpu if _is_xpu else biased_topk_jit_kernel_impl
topk_weights, topk_ids = _biased_topk(
hidden_states=hidden_states,
gating_output=router_logits,
correction_bias=correction_bias,
test/registered/xpu/test_topk.py
+153 -6
View File
@@ -3,6 +3,8 @@ from typing import Optional
import torch
from sglang.srt.layers.moe.hash_topk import HashTopK
torch.use_deterministic_algorithms(True)
from sglang.srt.layers.moe.topk import (
@@ -11,16 +13,26 @@ 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 biased_topk_impl as native_biased_topk
from sglang.srt.layers.moe.topk import (
biased_topk_xpu,
)
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.srt.runtime_context import get_context
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 _set_seed_and_device():
torch.manual_seed(1024)
return torch.device("xpu")
def _scatter_by_expert(
weights: torch.Tensor, indices: torch.Tensor, num_columns: int
) -> torch.Tensor:
@@ -83,8 +95,7 @@ class TestBiasedGroupedTopK(CustomTestCase):
bias_dtype,
routed_scaling_factor,
):
torch.manual_seed(1024)
device = torch.device("xpu")
device = _set_seed_and_device()
# expand gating_output by M, otherwise bfloat16 fall into same value aftering truncating
hidden_states = torch.randn(M, 100, dtype=torch.bfloat16, device=device)
@@ -162,8 +173,7 @@ class TestBiasedGroupedTopK(CustomTestCase):
renormalize = True
routed_scaling_factor = 2.5
torch.manual_seed(1024)
device = torch.device("xpu")
device = _set_seed_and_device()
bs = [1, 2, 4, 8]
seq_len = 1024
@@ -224,8 +234,7 @@ class TestBiasedGroupedTopK(CustomTestCase):
renormalize = True
routed_scaling_factor = 2.5
torch.manual_seed(1024)
device = torch.device("xpu")
device = _set_seed_and_device()
bs = [1]
seq_len = 1024
@@ -271,6 +280,144 @@ class TestBiasedGroupedTopK(CustomTestCase):
seq_len=seq_len,
)
def test_biased_topk(self):
# DeepSeek-V4 style routing shape
E_num_list = [256, 384]
topk_value = 6
gating_dtype = torch.float32
bias_dtype = torch.float32
renormalize = True
scoring_func_list = ["sqrtsoftplus", "sigmoid"]
routed_scaling_factor = 2.5
device = _set_seed_and_device()
bs = [1]
seq_len = 1024
num_tokens = [b * seq_len for b in bs]
num_fused_shared_experts_list = [0, 1]
for E_num in E_num_list:
for M in num_tokens:
for scoring_func in scoring_func_list:
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=gating_dtype, 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_topk(
hidden_states,
gating_output,
correction_bias,
topk_value,
renormalize,
scoring_func,
num_fused_shared_experts,
routed_scaling_factor,
apply_routed_scaling_factor_on_output=True,
)
# fused version
topk_weights, topk_ids = biased_topk_xpu(
hidden_states,
gating_output,
correction_bias,
topk_value,
renormalize,
scoring_func,
num_fused_shared_experts,
routed_scaling_factor,
apply_routed_scaling_factor_on_output=True,
)
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_hash_topk(self):
"""Guard the XPU fused hash-topk path against math/ID drift from torch."""
device = _set_seed_and_device()
E_num_list = [256, 384]
topk = 6
vocab_size = 128
dtype = torch.float32
bs = [1]
seq_len = 1024
num_tokens = [b * seq_len for b in bs]
num_fused_shared_experts_list = [0, 1]
with get_context().override_server_args(enable_waterfill=False):
for E_num in E_num_list:
for M in num_tokens:
for num_fused_shared_experts in num_fused_shared_experts_list:
hidden_states = torch.randn(
M, 1, dtype=torch.float32, device=device
)
router_logits = torch.randn(
M, E_num, dtype=dtype, device=device
)
input_ids = torch.randint(
low=0,
high=vocab_size,
size=(M,),
dtype=torch.int64,
device=device,
)
hash_topk = HashTopK(
topk=topk,
num_experts=E_num,
num_fused_shared_experts=num_fused_shared_experts,
vocab_size=vocab_size,
scoring_func="sqrtsoftplus",
routed_scaling_factor=2.5,
).to(device)
topk_routed = hash_topk.tid2eid.shape[1]
with torch.no_grad():
hash_topk.tid2eid.copy_(
torch.randint(
low=0,
high=E_num,
size=(vocab_size, topk_routed),
dtype=torch.int32,
device=device,
)
)
ref_topk_weights, ref_topk_ids = hash_topk._forward_torch(
router_logits, input_ids
)
output = hash_topk(
hidden_states=hidden_states,
router_logits=router_logits,
input_ids=input_ids,
)
torch.testing.assert_close(output.topk_ids, ref_topk_ids)
torch.testing.assert_close(
output.topk_weights, ref_topk_weights
)
if __name__ == "__main__":
unittest.main()