xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-09-06 06:16:59 +08:00
Code Issues Packages Projects Releases Wiki Activity

basic working test

Browse Source 
Signed-off-by: Bill Nell <bnell@redhat.com>
This commit is contained in:
Bill Nell 2025-05-21 23:20:54 +00:00
parent 77f95b99a6
commit 9cfebf51ba
2 changed files with 54 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tests/kernels/moe/test_batched_moe.py
View File

@ -198,9 +198,11 @@ def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
#B_scale = torch.ones((N, K), dtype=torch.float32, device=tensors.A.device) #B_scale = torch.ones((N, K), dtype=torch.float32, device=tensors.A.device)
A_scale = torch.ones(1, dtype=torch.float32, device=tensors.A.device) A_scale = torch.ones(1, dtype=torch.float32, device=tensors.A.device)
B_scale = torch.ones(1, dtype=torch.float32, device=tensors.B.device) B_scale = torch.ones(1, dtype=torch.float32, device=tensors.B.device)
quant_block_shape = [1, 1]
else: else:
A_scale = None A_scale = None
B_scale = None B_scale = None
quant_block_shape = None
invoke_moe_batched_triton_kernel( invoke_moe_batched_triton_kernel(
tensors.A, tensors.A,
@ -220,7 +222,9 @@ def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
"BLOCK_SIZE_M": block_shape[0], "BLOCK_SIZE_M": block_shape[0],
"BLOCK_SIZE_N": block_shape[1], "BLOCK_SIZE_N": block_shape[1],
"BLOCK_SIZE_K": block_shape[2], "BLOCK_SIZE_K": block_shape[2],
}) },
block_shape=quant_block_shape,
)
ref_output = ref_output.to(dtype=out_dtype) ref_output = ref_output.to(dtype=out_dtype)
ref_output = ref_impl(tensors.A.to(dtype=out_dtype), ref_output = ref_impl(tensors.A.to(dtype=out_dtype),
@ -246,5 +250,4 @@ def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
}[test_output.dtype] }[test_output.dtype]
torch.testing.assert_close(ref_output, ref_output2, atol=atol, rtol=rtol) torch.testing.assert_close(ref_output, ref_output2, atol=atol, rtol=rtol)
if not use_fp8_w8a8:
torch.testing.assert_close(test_output, ref_output2, atol=atol, rtol=rtol) torch.testing.assert_close(test_output, ref_output2, atol=atol, rtol=rtol)

8
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
View File

@ -45,8 +45,8 @@ def moe_mmk(
BLOCK_K: tl.constexpr, BLOCK_K: tl.constexpr,
compute_type: tl.constexpr, compute_type: tl.constexpr,
use_w8a8: tl.constexpr, use_w8a8: tl.constexpr,
use_w8a16: tl.constexpr): use_w8a16: tl.constexpr
):
offs_k = tl.arange(0, BLOCK_K) offs_k = tl.arange(0, BLOCK_K)
if use_w8a16: if use_w8a16:
@ -324,8 +324,8 @@ def invoke_moe_batched_triton_kernel(
use_int8_w8a16: bool, use_int8_w8a16: bool,
use_int4_w4a16: bool, use_int4_w4a16: bool,
config: dict[str, int], config: dict[str, int],
block_shape: Optional[list[int]] = None): block_shape: Optional[list[int]] = None
):
assert not use_int4_w4a16 assert not use_int4_w4a16
max_num_tokens = A.size(1) max_num_tokens = A.size(1)
K = A.size(2) K = A.size(2)