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 15, 11:00 UTC
[TPU] Add kernel test for moe_pallas (#17496)
Signed-off-by: Michael Goin <mgoin64@gmail.com>
Michael Goin committed
May 6, 2025 at 20:59 UTC
e50a1f1a9cc520d59485f5eb5613bf2d766c5102
4 files changed
+96
-3
.buildkite/scripts/hardware_ci/run-tpu-v1-test.sh
+3
-1
@@ -47,7 +47,9 @@ docker run --privileged --net host --shm-size=16G -it \
47
&& echo TEST_10 \
48
&& pytest -s -v /workspace/vllm/tests/v1/tpu/test_pallas.py \
49
&& echo TEST_11 \
50
- && pytest -s -v /workspace/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py" \
50
+ && pytest -s -v /workspace/vllm/tests/v1/entrypoints/llm/test_struct_output_generate.py \
51
+ && echo TEST_12 \
52
+ && pytest -s -v /workspace/vllm/tests/tpu/test_moe_pallas.py" \
53
54
55
# TODO: This test fails because it uses RANDOM_SEED sampling
tests/tpu/test_moe_pallas.py
new
+87
@@ -0,0 +1,87 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""Tests for the Pallas MOE implementation.
3
+
4
+Run `pytest tests/kernels/moe/test_moe_pallas.py`.
5
+"""
6
+import pytest
7
+import torch
8
+
9
+# yapf conflicts with isort for this block
10
+# yapf: disable
11
+from vllm.model_executor.layers.fused_moe.moe_pallas import (
12
+ fused_moe as pallas_moe)
13
+from vllm.model_executor.layers.fused_moe.moe_torch_iterative import (
14
+ fused_moe as torch_moe)
15
+# yapf: enable
16
+from vllm.platforms import current_platform
17
+
18
+if not current_platform.is_tpu():
19
+ pytest.skip("This test needs a TPU.", allow_module_level=True)
20
+
21
+NUM_EXPERTS = [8, 64]
22
+EP_SIZE = [1]
23
+TOP_KS = [2, 6]
24
+
25
+
26
+# The Pallas GMM kernel requires num_tokens * topk to be a multiple of 16
27
+@pytest.mark.parametrize("m", [8, 16, 64, 2048])
28
+@pytest.mark.parametrize("n", [128, 1024, 2048])
29
+@pytest.mark.parametrize("k", [128, 511, 1024])
30
+@pytest.mark.parametrize("e", NUM_EXPERTS)
31
+@pytest.mark.parametrize("topk", TOP_KS)
32
+@pytest.mark.parametrize("ep_size", EP_SIZE)
33
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
34
+def test_pallas_moe(
35
+ m: int,
36
+ n: int,
37
+ k: int,
38
+ e: int,
39
+ topk: int,
40
+ ep_size: int,
41
+ dtype: torch.dtype,
42
+):
43
+ import torch_xla.core.xla_model as xm
44
+ with torch.device(xm.xla_device()):
45
+ a = torch.randn((m, k), dtype=dtype) / 10
46
+ w1 = torch.randn((e, 2 * n, k), dtype=dtype) / 10
47
+ w2 = torch.randn((e, k, n), dtype=dtype) / 10
48
+
49
+ score = torch.randn((m, e), dtype=dtype)
50
+
51
+ # TODO: Support ep
52
+ if ep_size > 1:
53
+ pytest.skip("No support for ep_size > 1 yet")
54
+ else:
55
+ e_map = None
56
+
57
+ # Run both implementations
58
+ torch_output = torch_moe(
59
+ hidden_states=a,
60
+ w1=w1,
61
+ w2=w2,
62
+ gating_output=score,
63
+ topk=topk,
64
+ global_num_experts=e,
65
+ expert_map=e_map,
66
+ renormalize=False,
67
+ )
68
+
69
+ pallas_output = pallas_moe(
70
+ hidden_states=a,
71
+ w1=w1,
72
+ w2=w2,
73
+ gating_output=score,
74
+ topk=topk,
75
+ global_num_experts=e,
76
+ expert_map=e_map,
77
+ renormalize=False,
78
+ )
79
+ xm.mark_step()
80
+
81
+ # Compare outputs
82
+ torch.testing.assert_close(
83
+ pallas_output.cpu(),
84
+ torch_output.cpu(),
85
+ atol=2e-2,
86
+ rtol=0,
87
+ )
vllm/attention/backends/pallas.py
+2
-1
@@ -123,7 +123,8 @@ class PallasAttentionBackendImpl(AttentionImpl):
123
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
124
self.logits_soft_cap = logits_soft_cap
125
if head_size % 128 != 0:
126
- raise NotImplementedError("Head size must be a multiple of 128.")
126
+ raise NotImplementedError(
127
+ f"Head size must be a multiple of 128, found {head_size}.")
128
if alibi_slopes is not None:
129
raise NotImplementedError("Alibi slopes is not supported.")
130
if sliding_window is not None:
vllm/model_executor/layers/fused_moe/moe_pallas.py
+4
-1
@@ -11,7 +11,9 @@ def fused_moe(
11
w2: torch.Tensor,
12
gating_output: torch.Tensor,
13
topk: int,
14
- renormalize: bool,
14
+ global_num_experts: int,
15
+ expert_map: torch.Tensor = None,
16
+ renormalize: bool = False,
17
) -> torch.Tensor:
18
"""
19
Args:
@@ -20,6 +22,7 @@ def fused_moe(
22
w2: [num_experts, hidden_size, intermediate_size]
23
gating_output: [*, num_experts]
24
"""
25
+ assert expert_map is None, "expert_map is not supported for pallas MoE."
26
orig_shape = hidden_states.shape
27
hidden_size = hidden_states.shape[-1]
28
num_tokens = hidden_states.shape[:-1].numel()