| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 8 | from tests.kernels.moe.utils import ( |
| 9 | make_dummy_moe_config, |
| 10 | make_test_quant_config, |
| 11 | make_test_weights, |
| 12 | modular_triton_fused_moe, |
| 13 | ) |
| 14 | from tests.kernels.quant_utils import native_w8a8_block_matmul |
| 15 | from vllm.config import VllmConfig, set_current_vllm_config |
| 16 | from vllm.model_executor.layers.activation import SiluAndMul |
| 17 | from vllm.model_executor.layers.fused_moe import ( |
| 18 | fused_experts, |
| 19 | fused_topk, |
| 20 | ) |
| 21 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 22 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 23 | maybe_make_prepare_finalize, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.fused_moe.config import ( |
| 26 | fp8_w8a8_moe_quant_config, |
| 27 | ) |
| 28 | from vllm.model_executor.layers.fused_moe.experts import triton_moe |
| 29 | from vllm.model_executor.layers.fused_moe.experts.deep_gemm_moe import ( |
| 30 | _valid_deep_gemm_shape, |
| 31 | ) |
| 32 | from vllm.model_executor.layers.fused_moe.experts.triton_deep_gemm_moe import ( |
| 33 | TritonOrDeepGemmExperts, |
| 34 | ) |
| 35 | from vllm.model_executor.layers.quantization.utils.fp8_utils import ( |
| 36 | per_token_group_quant_fp8, |
| 37 | ) |
| 38 | from vllm.platforms import current_platform |
| 39 | from vllm.utils.deep_gemm import ( |
| 40 | get_mk_alignment_for_contiguous_layout, |
| 41 | is_deep_gemm_e8m0_used, |
| 42 | ) |
| 43 | from vllm.utils.import_utils import has_deep_gemm |
| 44 | |
| 45 | dg_available = has_deep_gemm() |
| 46 | |
| 47 | if current_platform.get_device_capability() < (9, 0): |
| 48 | pytest.skip("FP8 Triton requires CUDA 9.0 or higher", allow_module_level=True) |
| 49 | if current_platform.is_fp8_fnuz(): |
| 50 | pytest.skip( |
| 51 | "Tests in this file require float8_e4m3fn and platform does not support", |
| 52 | allow_module_level=True, |
| 53 | ) |
| 54 | |
| 55 | vllm_config = VllmConfig() |
| 56 | |
| 57 | # Test configurations |
| 58 | DTYPES = [torch.bfloat16] # [torch.half, torch.bfloat16, torch.float32] |
| 59 | # Deepseek-V3's intermediate size 18432, so N is 18432*2/8=4608 at TP8 |
| 60 | # and its hidden size is 7168. |
| 61 | MNK_FACTORS = [ |
| 62 | (1, 128, 128), |
| 63 | (1, 128, 7168), |
| 64 | (1, 1024, 7168), |
| 65 | (1, 4608, 128), |
| 66 | (1, 4608, 7168), |
| 67 | (83, 128, 128), |
| 68 | (83, 512, 512), |
| 69 | (83, 4608, 512), |
| 70 | (83, 4608, 7168), |
| 71 | (128, 512, 512), |
| 72 | (128, 1024, 7168), |
| 73 | (128, 4608, 7168), |
| 74 | (2048, 128, 128), |
| 75 | (2048, 1024, 7168), |
| 76 | (2048, 4608, 512), |
| 77 | (2048, 4608, 7168), |
| 78 | (8192, 128, 128), |
| 79 | (8192, 128, 7168), |
| 80 | (8192, 1024, 7168), |
| 81 | (8192, 4608, 7168), |
| 82 | ] |
| 83 | |
| 84 | MNK_FACTORS_DG = [ |
| 85 | (128, 128, 128), |
| 86 | (128, 128, 7168), |
| 87 | (128, 1024, 7168), |
| 88 | (128, 4608, 128), |
| 89 | (128, 4608, 7168), |
| 90 | (192, 512, 512), |
| 91 | (192, 1024, 7168), |
| 92 | (192, 4608, 7168), |
| 93 | (1335, 128, 128), |
| 94 | (1335, 1024, 7168), |
| 95 | (1335, 4608, 512), |
| 96 | (1335, 4608, 7168), |
| 97 | (2048, 128, 128), |
| 98 | (2048, 128, 7168), |
| 99 | (2048, 1024, 7168), |
| 100 | (2048, 4608, 7168), |
| 101 | ] |
| 102 | |
| 103 | BLOCK_SIZE = [[128, 128]] |
| 104 | E = [2, 8, 16] # [128, 256] |
| 105 | TOP_KS = [1, 2, 6] |
| 106 | SEEDS = [0] |
| 107 | |
| 108 | |
| 109 | def torch_w8a8_block_fp8_moe( |
| 110 | a, |
| 111 | w1, |
| 112 | w2, |
| 113 | w1_s, |
| 114 | w2_s, |
| 115 | topk_weight, |
| 116 | topk_ids, |
| 117 | block_shape, |
| 118 | silu_fp32=False, |
| 119 | expert_map=None, |
| 120 | ): |
| 121 | """Fused MoE with block-wise fp8 quantization using native torch. |
| 122 | |
| 123 | silu_fp32=True computes the intermediate SiLU in fp32 and quantizes it |
| 124 | directly, matching the modular Helion silu_and_mul_per_block_quant kernel; |
| 125 | False rounds the SiLU output to bf16 first, matching fused_experts. Each |
| 126 | kernel is checked against the reference variant matching its precision. |
| 127 | """ |
| 128 | B, D = a.shape |
| 129 | topk = topk_ids.size(1) |
| 130 | a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D) |
| 131 | out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device) |
| 132 | |
| 133 | topk_weight = topk_weight.view(-1) |
| 134 | topk_ids = topk_ids.view(-1) |
| 135 | if expert_map is not None: |
| 136 | # w1/w2 hold this rank's experts; slots mapped to -1 contribute nothing. |
| 137 | topk_ids = expert_map[topk_ids] |
| 138 | |
| 139 | _, block_k = block_shape[0], block_shape[1] |
| 140 | # Quantize with the production per-token-group fp8 kernel (same HIP/CUDA op the |
| 141 | # kernels use) so the reference is bit-identical here; removes ~0.06-0.10% fp8 |
| 142 | # boundary-flip divergence that otherwise accumulates over K. Matmul stays fp32. |
| 143 | a_q, a_s = per_token_group_quant_fp8(a, block_k, dtype=current_platform.fp8_dtype()) |
| 144 | a_q = a_q.to(torch.float32) |
| 145 | for i in range(w1.shape[0]): |
| 146 | mask = topk_ids == i |
| 147 | if mask.sum(): |
| 148 | inter_out = native_w8a8_block_matmul( |
| 149 | a_q[mask], w1[i], a_s[mask], w1_s[i], block_shape, output_dtype=a.dtype |
| 150 | ) |
| 151 | act_out = SiluAndMul().forward_native( |
| 152 | inter_out.float() if silu_fp32 else inter_out |
| 153 | ) |
| 154 | act_out_q, act_out_s = per_token_group_quant_fp8( |
| 155 | act_out, block_k, dtype=current_platform.fp8_dtype() |
| 156 | ) |
| 157 | out[mask] = native_w8a8_block_matmul( |
| 158 | act_out_q, w2[i], act_out_s, w2_s[i], block_shape, output_dtype=a.dtype |
| 159 | ) |
| 160 | return ( |
| 161 | out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype) |
| 162 | ).sum(dim=1) |
| 163 | |
| 164 | |
| 165 | # Skip all tests if CUDA is not available |
| 166 | pytest.importorskip("torch.cuda") |
| 167 | |
| 168 | |
| 169 | @pytest.fixture(autouse=True) |
| 170 | def setup_cuda(): |
| 171 | torch.set_default_device("cuda") |
| 172 | |
| 173 | |
| 174 | @pytest.mark.parametrize(("M", "N", "K"), MNK_FACTORS) |
| 175 | @pytest.mark.parametrize("E", E) |
| 176 | @pytest.mark.parametrize("topk", TOP_KS) |
| 177 | @pytest.mark.parametrize("block_size", BLOCK_SIZE) |
| 178 | @pytest.mark.parametrize("dtype", DTYPES) |
| 179 | @pytest.mark.parametrize("seed", SEEDS) |
| 180 | @torch.inference_mode() |
| 181 | def test_w8a8_block_fp8_fused_moe( |
| 182 | M, N, K, E, topk, block_size, dtype, seed, monkeypatch, workspace_init |
| 183 | ): |
| 184 | if topk > E: |
| 185 | pytest.skip(f"Skipping test; topk={topk} > E={E}") |
| 186 | |
| 187 | torch.manual_seed(seed) |
| 188 | |
| 189 | a = torch.randn((M, K), dtype=dtype) / 10 |
| 190 | score = torch.randn((M, E), dtype=dtype) |
| 191 | |
| 192 | w1, w2, quant_config = make_test_quant_config( |
| 193 | E, |
| 194 | N, |
| 195 | K, |
| 196 | dtype, |
| 197 | quant_dtype=torch.float8_e4m3fn, |
| 198 | per_act_token_quant=False, |
| 199 | block_shape=block_size, |
| 200 | ) |
| 201 | |
| 202 | m_fused_moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config) |
| 203 | |
| 204 | topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False) |
| 205 | |
| 206 | # Set the context to avoid lots of warning spam. |
| 207 | with set_current_vllm_config(vllm_config): |
| 208 | ref_out = torch_w8a8_block_fp8_moe( |
| 209 | a, |
| 210 | w1, |
| 211 | w2, |
| 212 | quant_config.w1_scale, |
| 213 | quant_config.w2_scale, |
| 214 | topk_weights, |
| 215 | topk_ids, |
| 216 | block_size, |
| 217 | ) |
| 218 | |
| 219 | out = fused_experts( |
| 220 | a, w1, w2, topk_weights, topk_ids, quant_config=quant_config |
| 221 | ) |
| 222 | |
| 223 | m_out = m_fused_moe.apply( |
| 224 | a, |
| 225 | w1, |
| 226 | w2, |
| 227 | topk_weights, |
| 228 | topk_ids, |
| 229 | activation=MoEActivation.SILU, |
| 230 | apply_router_weight_on_input=False, |
| 231 | expert_map=None, |
| 232 | global_num_experts=w1.shape[0], |
| 233 | ) |
| 234 | |
| 235 | # 0.039 only needed for M >= 8192 |
| 236 | tol = 0.035 if M < 8192 else 0.039 |
| 237 | |
| 238 | # The modular path fuses SiLU+quant in fp32 (silu_and_mul_per_block_quant), |
| 239 | # while fused_experts/the reference round SiLU to bf16 first. On large K/N this |
| 240 | # ~1-ULP gap pushes m_out past the base tol, so validate m_out against an |
| 241 | # fp32-SiLU reference — keeping the tight base tolerance, no widened override. |
| 242 | if current_platform.is_rocm() and K >= 4096 and N >= 1024: |
| 243 | with set_current_vllm_config(vllm_config): |
| 244 | ref_out_m = torch_w8a8_block_fp8_moe( |
| 245 | a, |
| 246 | w1, |
| 247 | w2, |
| 248 | quant_config.w1_scale, |
| 249 | quant_config.w2_scale, |
| 250 | topk_weights, |
| 251 | topk_ids, |
| 252 | block_size, |
| 253 | silu_fp32=True, |
| 254 | ) |
| 255 | else: |
| 256 | ref_out_m = ref_out |
| 257 | |
| 258 | torch.testing.assert_close(out, ref_out, atol=tol, rtol=tol) |
| 259 | torch.testing.assert_close(m_out, ref_out_m, atol=tol, rtol=tol) |
| 260 | |
| 261 | |
| 262 | @pytest.mark.parametrize("clamp_limit", [None, 10.0]) |
| 263 | @torch.inference_mode() |
| 264 | def test_w8a8_block_fp8_clamp_skips_fused_silu_quant( |
| 265 | clamp_limit, monkeypatch, workspace_init, disable_deepgemm_ue8m0 |
| 266 | ): |
| 267 | """A SwiGLU clamp must bypass the fused SiLU+quant kernel, which cannot clamp.""" |
| 268 | torch.manual_seed(0) |
| 269 | M, N, K, E, topk = 16, 256, 512, 8, 2 |
| 270 | block_size = [128, 128] |
| 271 | |
| 272 | a = torch.randn((M, K), dtype=torch.bfloat16) / 10 |
| 273 | score = torch.randn((M, E), dtype=torch.bfloat16) |
| 274 | (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights( |
| 275 | E, N, K, torch.bfloat16, torch.float8_e4m3fn, block_shape=block_size |
| 276 | ) |
| 277 | quant_config = fp8_w8a8_moe_quant_config( |
| 278 | w1_scale=w1_s, |
| 279 | w2_scale=w2_s, |
| 280 | block_shape=block_size, |
| 281 | gemm1_clamp_limit=clamp_limit, |
| 282 | ) |
| 283 | m_fused_moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config) |
| 284 | topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False) |
| 285 | |
| 286 | fused_calls = 0 |
| 287 | fused_op = triton_moe.ops.silu_and_mul_per_block_quant |
| 288 | |
| 289 | def counting_fused_op(*args, **kwargs): |
| 290 | nonlocal fused_calls |
| 291 | fused_calls += 1 |
| 292 | return fused_op(*args, **kwargs) |
| 293 | |
| 294 | monkeypatch.setattr( |
| 295 | triton_moe.ops, "silu_and_mul_per_block_quant", counting_fused_op |
| 296 | ) |
| 297 | |
| 298 | with set_current_vllm_config(vllm_config): |
| 299 | m_fused_moe.apply( |
| 300 | a, |
| 301 | w1, |
| 302 | w2, |
| 303 | topk_weights, |
| 304 | topk_ids, |
| 305 | activation=MoEActivation.SILU, |
| 306 | apply_router_weight_on_input=False, |
| 307 | expert_map=None, |
| 308 | global_num_experts=E, |
| 309 | ) |
| 310 | |
| 311 | assert (fused_calls > 0) == (clamp_limit is None), ( |
| 312 | f"fused fast path ran {fused_calls} times with clamp_limit={clamp_limit}" |
| 313 | ) |
| 314 | |
| 315 | |
| 316 | @pytest.mark.parametrize(("M", "N", "K"), [(83, 512, 512), (2048, 4608, 512)]) |
| 317 | @pytest.mark.parametrize("topk", [2, 6]) |
| 318 | @pytest.mark.parametrize("ep_size", [2, 4]) |
| 319 | @torch.inference_mode() |
| 320 | def test_w8a8_block_fp8_fused_moe_expert_map(M, N, K, topk, ep_size, workspace_init): |
| 321 | """One EP rank: local expert shard plus an expert_map with -1 entries.""" |
| 322 | E, block_size, dtype = 16, [128, 128], torch.bfloat16 |
| 323 | torch.manual_seed(0) |
| 324 | |
| 325 | a = torch.randn((M, K), dtype=dtype) / 10 |
| 326 | score = torch.randn((M, E), dtype=dtype) |
| 327 | w1, w2, quant_config = make_test_quant_config( |
| 328 | E, N, K, dtype, quant_dtype=torch.float8_e4m3fn, block_shape=block_size |
| 329 | ) |
| 330 | |
| 331 | local_e = E // ep_size |
| 332 | e_ids = torch.randperm(E, dtype=torch.int32)[:local_e] |
| 333 | e_map = torch.full((E,), -1, dtype=torch.int32) |
| 334 | e_map[e_ids] = torch.arange(local_e, dtype=torch.int32) |
| 335 | w1, w2 = w1[e_ids].contiguous(), w2[e_ids].contiguous() |
| 336 | w1_s = quant_config.w1_scale[e_ids].contiguous() |
| 337 | w2_s = quant_config.w2_scale[e_ids].contiguous() |
| 338 | quant_config = fp8_w8a8_moe_quant_config( |
| 339 | w1_scale=w1_s, w2_scale=w2_s, block_shape=block_size |
| 340 | ) |
| 341 | |
| 342 | m_fused_moe = modular_triton_fused_moe(make_dummy_moe_config(), quant_config) |
| 343 | topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False) |
| 344 | |
| 345 | with set_current_vllm_config(vllm_config): |
| 346 | ref_out = torch_w8a8_block_fp8_moe( |
| 347 | a, w1, w2, w1_s, w2_s, topk_weights, topk_ids, block_size, expert_map=e_map |
| 348 | ) |
| 349 | m_out = m_fused_moe.apply( |
| 350 | a, |
| 351 | w1, |
| 352 | w2, |
| 353 | topk_weights, |
| 354 | topk_ids, |
| 355 | activation=MoEActivation.SILU, |
| 356 | apply_router_weight_on_input=False, |
| 357 | expert_map=e_map, |
| 358 | global_num_experts=E, |
| 359 | ) |
| 360 | |
| 361 | torch.testing.assert_close(m_out, ref_out, atol=0.035, rtol=0.035) |
| 362 | |
| 363 | |
| 364 | @pytest.mark.parametrize(("M", "N", "K"), MNK_FACTORS_DG) |
| 365 | @pytest.mark.parametrize("E", E) |
| 366 | @pytest.mark.parametrize("topk", TOP_KS) |
| 367 | @pytest.mark.parametrize("seed", SEEDS) |
| 368 | @pytest.mark.skipif(not dg_available, reason="DeepGemm kernels not available.") |
| 369 | @pytest.mark.skipif(is_deep_gemm_e8m0_used(), reason="Not E8M0 scale MOE") |
| 370 | @torch.inference_mode() |
| 371 | def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch): |
| 372 | if topk > E: |
| 373 | pytest.skip(f"Skipping test: topk={topk} > E={E}") |
| 374 | |
| 375 | if not _valid_deep_gemm_shape(M, N, K): |
| 376 | pytest.skip(f"Skipping test: invalid size m={M}, n={N}, k={K}") |
| 377 | |
| 378 | torch.manual_seed(seed) |
| 379 | |
| 380 | block_size = get_mk_alignment_for_contiguous_layout() |
| 381 | dtype = torch.bfloat16 |
| 382 | |
| 383 | a = torch.randn((M, K), dtype=dtype) / 10 |
| 384 | score = torch.randn((M, E), dtype=dtype) |
| 385 | |
| 386 | (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights( |
| 387 | E, |
| 388 | N, |
| 389 | K, |
| 390 | dtype, |
| 391 | torch.float8_e4m3fn, |
| 392 | per_out_ch_quant=False, |
| 393 | block_shape=block_size, |
| 394 | ) |
| 395 | |
| 396 | # Note: for now use_compile will error out if the problem size is |
| 397 | # large enough to trigger chunking. I'm leaving the flag and |
| 398 | # setup code in case we are able to revisit this later. |
| 399 | use_compile = False |
| 400 | |
| 401 | use_cudagraph = N >= 1024 and K >= 1024 and current_platform.is_cuda_alike() |
| 402 | |
| 403 | topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False) |
| 404 | |
| 405 | quant_config = fp8_w8a8_moe_quant_config( |
| 406 | w1_scale=w1_s, |
| 407 | w2_scale=w2_s, |
| 408 | block_shape=block_size, |
| 409 | ) |
| 410 | moe_config = make_dummy_moe_config() |
| 411 | |
| 412 | deep_gemm_experts = mk.FusedMoEKernel( |
| 413 | prepare_finalize=maybe_make_prepare_finalize( |
| 414 | moe=moe_config, |
| 415 | quant_config=quant_config, |
| 416 | allow_new_interface=True, |
| 417 | use_monolithic=False, |
| 418 | ), |
| 419 | fused_experts=TritonOrDeepGemmExperts( |
| 420 | moe_config=moe_config, |
| 421 | quant_config=quant_config, |
| 422 | ), |
| 423 | ) |
| 424 | |
| 425 | def deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids): |
| 426 | return deep_gemm_experts.apply( |
| 427 | hidden_states=a, |
| 428 | w1=w1, |
| 429 | w2=w2, |
| 430 | topk_weights=topk_weights, |
| 431 | topk_ids=topk_ids, |
| 432 | global_num_experts=E, |
| 433 | activation=MoEActivation.SILU, |
| 434 | apply_router_weight_on_input=False, |
| 435 | expert_map=False, |
| 436 | ) |
| 437 | |
| 438 | # Set the context to avoid lots of warning spam. |
| 439 | with set_current_vllm_config(vllm_config): |
| 440 | ref_out = torch_w8a8_block_fp8_moe( |
| 441 | a, w1, w2, w1_s, w2_s, topk_weights, topk_ids, block_size |
| 442 | ) |
| 443 | |
| 444 | if use_compile: |
| 445 | deep_gemm_moe_fp8_fn = torch.compile( |
| 446 | deep_gemm_moe_fp8, backend="inductor", fullgraph=True |
| 447 | ) |
| 448 | torch._dynamo.mark_dynamic(a, 0) |
| 449 | torch._dynamo.mark_dynamic(topk_weights, 0) |
| 450 | torch._dynamo.mark_dynamic(topk_ids, 0) |
| 451 | else: |
| 452 | deep_gemm_moe_fp8_fn = deep_gemm_moe_fp8 |
| 453 | |
| 454 | out = deep_gemm_moe_fp8_fn(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids) |
| 455 | |
| 456 | if use_cudagraph: |
| 457 | out.fill_(0) |
| 458 | stream = torch.cuda.Stream() |
| 459 | graph = torch.cuda.CUDAGraph() |
| 460 | with torch.cuda.graph(graph, stream=stream): |
| 461 | out = deep_gemm_moe_fp8_fn( |
| 462 | a, w1, w2, w1_s, w2_s, topk_weights, topk_ids |
| 463 | ) |
| 464 | torch.accelerator.synchronize() |
| 465 | graph.replay() |
| 466 | torch.accelerator.synchronize() |
| 467 | |
| 468 | torch.testing.assert_close(out, ref_out, atol=0.035, rtol=0.035) |
| 469 | |
| 470 | |
| 471 | def _dequant_block_fp8( |
| 472 | w_fp8: torch.Tensor, w_s: torch.Tensor, block_shape: list[int] |
| 473 | ) -> torch.Tensor: |
| 474 | """Dequantize a [N, K] fp8 weight with [Nb, Kb] block scales to fp32.""" |
| 475 | n, k = w_fp8.shape |
| 476 | bn, bk = block_shape |
| 477 | s = w_s.repeat_interleave(bn, dim=0).repeat_interleave(bk, dim=1) |
| 478 | return w_fp8.to(torch.float32) * s[:n, :k] |
| 479 | |
| 480 | |
| 481 | @pytest.mark.parametrize("tp_rank", [0, 1, 2, 3]) |
| 482 | @torch.inference_mode() |
| 483 | def test_w8a8_block_fp8_fused_moe_refined_block_scales(tp_rank, workspace_init): |
| 484 | """TP-misaligned blockwise FP8 MoE (e.g. Qwen4Exp, intermediate |
| 485 | 640 at TP=4 -> 160 per rank, not divisible by the checkpoint's 128 block). |
| 486 | |
| 487 | The per-rank scale grid is refined from 128 to 32 (lossless: each 32-block |
| 488 | lies inside one global 128-block) so the Triton kernel can consume exact |
| 489 | per-shard scales. Simulates one TP rank's shard and checks: |
| 490 | 1. the refined+sharded scales dequantize exactly like the global scales; |
| 491 | 2. the Triton fused MoE kernel with block_shape=[32, 32] matches the |
| 492 | native-torch blockwise reference on the same shard. |
| 493 | """ |
| 494 | torch.manual_seed(0) |
| 495 | dtype = torch.bfloat16 |
| 496 | M, E, topk = 8, 4, 2 |
| 497 | hidden = 256 |
| 498 | inter_full = 640 |
| 499 | tp_size = 4 |
| 500 | n_shard = inter_full // tp_size # 160 |
| 501 | ckpt_block = [128, 128] |
| 502 | refined_block = [32, 32] |
| 503 | factor = ckpt_block[0] // refined_block[0] # 4 |
| 504 | |
| 505 | a = torch.randn((M, hidden), dtype=dtype) / 10 |
| 506 | score = torch.randn((M, E), dtype=dtype) |
| 507 | |
| 508 | (_, w1_full, w1_s_full, _), (_, w2_full, w2_s_full, _) = make_test_weights( |
| 509 | E, |
| 510 | inter_full, |
| 511 | hidden, |
| 512 | dtype, |
| 513 | torch.float8_e4m3fn, |
| 514 | per_out_ch_quant=False, |
| 515 | block_shape=ckpt_block, |
| 516 | ) |
| 517 | assert w1_s_full is not None |
| 518 | assert w2_s_full is not None |
| 519 | |
| 520 | # TP shard the weights (exact slicing, no scale involvement). |
| 521 | lo, hi = tp_rank * n_shard, (tp_rank + 1) * n_shard |
| 522 | w1 = torch.cat( |
| 523 | [w1_full[:, lo:hi], w1_full[:, inter_full + lo : inter_full + hi]], dim=1 |
| 524 | ) |
| 525 | w2 = w2_full[:, :, lo:hi] |
| 526 | |
| 527 | # Refine the global 128-block scales to 32 blocks, then take the shard's |
| 528 | # slice -- mirrors the Fp8MoEMethod/RoutedExperts loading path. |
| 529 | w1_s32 = w1_s_full.repeat_interleave(factor, dim=-2).repeat_interleave( |
| 530 | factor, dim=-1 |
| 531 | ) |
| 532 | nb = n_shard // refined_block[0] # 5 local 32-blocks per projection |
| 533 | gate_hi = w1_s32.shape[1] // 2 |
| 534 | w1_s = torch.cat( |
| 535 | [ |
| 536 | w1_s32[:, tp_rank * nb : (tp_rank + 1) * nb], |
| 537 | w1_s32[:, gate_hi + tp_rank * nb : gate_hi + (tp_rank + 1) * nb], |
| 538 | ], |
| 539 | dim=1, |
| 540 | ) |
| 541 | w2_s32 = w2_s_full.repeat_interleave(factor, dim=-2).repeat_interleave( |
| 542 | factor, dim=-1 |
| 543 | ) |
| 544 | w2_s = w2_s32[:, :, tp_rank * nb : (tp_rank + 1) * nb] |
| 545 | |
| 546 | # The refined per-shard scales must reproduce the global-scale dequant |
| 547 | # exactly (dequant full-width, then slice to the shard). |
| 548 | for e in range(E): |
| 549 | d32 = _dequant_block_fp8(w1[e, :n_shard], w1_s[e, :nb], refined_block) |
| 550 | d128 = _dequant_block_fp8(w1_full[e], w1_s_full[e], ckpt_block)[lo:hi] |
| 551 | assert torch.equal(d32, d128) |
| 552 | d32 = _dequant_block_fp8(w2[e], w2_s[e], refined_block) |
| 553 | d128 = _dequant_block_fp8(w2_full[e], w2_s_full[e], ckpt_block)[:, lo:hi] |
| 554 | assert torch.equal(d32, d128) |
| 555 | |
| 556 | quant_config = fp8_w8a8_moe_quant_config( |
| 557 | w1_scale=w1_s, |
| 558 | w2_scale=w2_s, |
| 559 | block_shape=refined_block, |
| 560 | ) |
| 561 | |
| 562 | topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False) |
| 563 | |
| 564 | with set_current_vllm_config(vllm_config): |
| 565 | ref_out = torch_w8a8_block_fp8_moe( |
| 566 | a, w1, w2, w1_s, w2_s, topk_weights, topk_ids, refined_block |
| 567 | ) |
| 568 | out = fused_experts( |
| 569 | a, w1, w2, topk_weights, topk_ids, quant_config=quant_config |
| 570 | ) |
| 571 | |
| 572 | torch.testing.assert_close(out, ref_out, atol=0.035, rtol=0.035) |