From 788f284b5361d5192e6fc22856ea97e0e579530c Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Thu, 27 Feb 2025 05:00:00 -0500 Subject: [PATCH] Fix test_block_fp8.py test for MoE (#13915) Signed-off-by: mgoin --- tests/kernels/test_block_fp8.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/kernels/test_block_fp8.py b/tests/kernels/test_block_fp8.py index 20eff1c207239..6206cbd5f76f7 100644 --- a/tests/kernels/test_block_fp8.py +++ b/tests/kernels/test_block_fp8.py @@ -30,8 +30,8 @@ M_moe = [1, 7, 83, 512, 2048] N_moe = [4608] # [128, 4608, 13824] K_moe = [7168] # [256, 7168, 13824] BLOCK_SIZE = [[128, 128]] -E = [256] # [8, 24, 128, 256] -TOP_KS = [1] # [1, 2, 6] +E = [8, 24] # [8, 24, 128, 256] +TOP_KS = [2] # [1, 2, 6] OUT_DTYPES = [torch.bfloat16] # [torch.float32, torch.half, torch.bfloat16] SEEDS = [0]