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 7, 11:30 UTC
[Core] Move test utility to test file (#35672)
Signed-off-by: Turner Jabbour <doubleujabbour@gmail.com>
Turner Jabbour committed
Mar 2, 2026 at 08:56 UTC
4034c3d32e30d01639459edd3ab486f56993876d
4 files changed
+13
-24
tests/kernels/moe/test_gpt_oss_triton_kernels.py
+2
-1
@@ -26,9 +26,10 @@ from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_co
26
from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
27
triton_kernel_moe_forward,
28
)
29
-from vllm.model_executor.layers.utils import shuffle_weight
29
from vllm.utils.math_utils import round_up
30
31
+from .utils import shuffle_weight
32
+
33
34
def deshuffle(w: torch.Tensor):
35
first = w[..., ::2]
tests/kernels/moe/test_modular_oai_triton_moe.py
+1
-2
@@ -33,11 +33,10 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularK
33
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
34
MoEPrepareAndFinalizeNoEP,
35
)
36
-from vllm.model_executor.layers.utils import shuffle_weight
36
from vllm.platforms import current_platform
37
from vllm.utils.torch_utils import set_random_seed
38
40
-from .utils import make_dummy_moe_config
39
+from .utils import make_dummy_moe_config, shuffle_weight
40
41
MNK = [
42
(1, 512, 384),
tests/kernels/moe/utils.py
+10
@@ -33,6 +33,16 @@ from vllm.utils.deep_gemm import per_block_cast_to_fp8
33
from vllm.utils.math_utils import round_up
34
35
36
+def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
37
+ """Fold weights to adjacent locations for Triton MoE / SwiGLU kernel layout."""
38
+ shape = w.shape
39
+ n = shape[-1]
40
+ first = w[..., : n // 2]
41
+ second = w[..., n // 2 :]
42
+ stacked = torch.stack((first, second), dim=-1)
43
+ return stacked.reshape(shape)
44
+
45
+
46
def make_dummy_moe_config(
47
num_experts: int = 1,
48
experts_per_token: int = 1,
vllm/model_executor/layers/utils.py
-21
@@ -31,27 +31,6 @@ def is_layer_moe_router_gate(prefix: str) -> bool:
31
return prefix.rsplit(".", 1)[-1] in MOE_LAYER_ROUTER_GATE_SUFFIXES
32
33
34
-def shuffle_weight(w: torch.Tensor) -> torch.Tensor:
35
- # Shuffle weight along the last dimension so that
36
- # we folded the weights to adjance location
37
- # Example:
38
- # input:
39
- # [[1, 2, 3, 4, 5, 6],
40
- # [7, 8, 9, 10, 11, 12]]
41
- # output:
42
- # [[1, 4, 2, 5, 3, 6],
43
- # [7, 10, 8, 11, 9, 12]]
44
- # This will be used together with triton swiglu kernel
45
- shape = w.shape
46
- N = shape[-1]
47
- first = w[..., : N // 2]
48
- second = w[..., N // 2 :]
49
-
50
- stacked = torch.stack((first, second), dim=-1)
51
- w_shuffled = stacked.reshape(shape)
52
- return w_shuffled
53
-
54
-
34
def get_token_bin_counts_and_mask(
35
tokens: torch.Tensor,
36
vocab_size: int,