A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 10, 19:00 UTC
CustomOp: test forward dispatch for grouped_topk (#31530)
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Xinyu Chen committed
Jan 2, 2026 at 23:04 UTC
08f425bad168a8e6ce258b6bf309339473f9d5fe
1 file changed
+13
-1
tests/kernels/moe/test_grouped_topk.py
+13
-1
@@ -8,6 +8,12 @@ Run `pytest tests/kernels/moe/test_grouped_topk.py`.
8
import pytest
9
import torch
10
11
+from vllm.config import (
12
+ CompilationConfig,
13
+ VllmConfig,
14
+ get_cached_compilation_config,
15
+ set_current_vllm_config,
16
+)
17
from vllm.model_executor.layers.fused_moe.fused_moe import (
18
GroupedTopk,
19
fused_grouped_topk,
@@ -41,6 +47,11 @@ def test_grouped_topk(
47
routed_scaling_factor: float,
48
dtype: torch.dtype,
49
):
50
+ vllm_config = VllmConfig(
51
+ compilation_config=CompilationConfig(custom_ops=["all", "+grouped_topk"])
52
+ )
53
+ get_cached_compilation_config.cache_clear()
54
+
55
current_platform.seed_everything(0)
56
hidden_states = torch.randn((n_token, n_hidden), dtype=dtype, device="cuda")
57
gating_output = torch.randn((n_token, n_expert), dtype=dtype, device="cuda")
@@ -48,7 +59,7 @@ def test_grouped_topk(
59
(n_expert,), dtype=torch.float32, device="cuda"
60
)
61
51
- with monkeypatch.context() as m:
62
+ with set_current_vllm_config(vllm_config), monkeypatch.context() as m:
63
m.setenv("VLLM_USE_FUSED_MOE_GROUPED_TOPK", "0")
64
grouped_topk = GroupedTopk(
65
topk=topk,
@@ -58,6 +69,7 @@ def test_grouped_topk(
69
scoring_func=scoring_func,
70
routed_scaling_factor=routed_scaling_factor,
71
)
72
+ assert grouped_topk._forward_method.__name__ == "forward_cuda"
73
baseline_topk_weights, baseline_topk_ids = grouped_topk(
74
hidden_states=hidden_states,
75
gating_output=gating_output,