[CPU] bypass scoring_func argument in topk for cpu device (#31110)

This commit is contained in:
Ma Mingfei
2026-07-14 21:48:22 +08:00
committed by GitHub
parent bbb5702a3c
commit 31548781e0
2 changed files with 27 additions and 1 deletions
+22
View File
@@ -700,7 +700,25 @@ def fused_topk_cpu(
renormalize: bool,
correction_bias: torch.Tensor = None,
scoring_func: str = "softmax",
routed_scaling_factor: Optional[float] = None,
apply_routed_scaling_factor_on_output: Optional[bool] = False,
num_fused_shared_experts: int = 0,
packed_out: Optional[torch.Tensor] = None,
num_token_non_padded: Optional[torch.Tensor] = None,
):
if num_fused_shared_experts != 0:
raise ValueError(
f"num_fused_shared_experts must be 0 for CPU fused topk, got: {num_fused_shared_experts}"
)
if apply_routed_scaling_factor_on_output:
raise ValueError(
"apply_routed_scaling_factor_on_output is not supported for CPU fused topk"
)
if packed_out is not None:
raise ValueError("packed_out is not supported for CPU fused topk")
if num_token_non_padded is not None:
raise ValueError("num_token_non_padded is not supported for CPU fused topk")
# TODO: add c++ kernel for cpu
# The topk_softmax_cpu kernel only handles vanilla softmax scoring with no
# correction bias. Fall back to the torch-native impl for the rest
@@ -963,8 +981,12 @@ def grouped_topk_cpu(
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",
):
assert not apply_routed_scaling_factor_on_output
if scoring_func != "softmax":
raise ValueError(f"Unsupported scoring function: {scoring_func}")
return torch.ops.sgl_kernel.grouped_topk_cpu(
hidden_states,
gating_output,
@@ -7,7 +7,11 @@ from sglang.test.server_fixtures.spec_eagle_fixture import Eagle3Base
# Estimated: 2 sequential 8B server launches + one 4-prompt greedy method
# (CUDA sibling: 360); tune from CI TIMINGS once it has run there.
register_cpu_ci(est_time=480, suite="base-b-test-cpu")
register_cpu_ci(
est_time=480,
suite="base-b-test-cpu",
disabled="EAGLE3 numerical parity mismatches on CPU intel_amx",
)
class TestEagle3ParityCPU(SpecParityKit, Eagle3Base):