@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a4d5d663e2

Fork 0

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 10, 10:15 UTC

Add unpermute-aware fused MoE path and small-batch fallback (#29354)

Signed-off-by: Runkai Tao <rt572@physics.rutgers.edu> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

Runkai Tao committed Jan 9, 2026 at 14:58 UTC a4d5d663e295e3b90a515c79df74048bc715e804

2 files changed +176 -21

tests/kernels/moe/test_moe.py

+114

@@ -64,8 +64,10 @@ from vllm.utils.torch_utils import set_random_seed

64 from vllm.v1.worker.workspace import init_workspace_manager

66 NUM_EXPERTS = [8, 64, 192]

67 +NUM_EXPERTS_LARGE = [128, 256]

68 EP_SIZE = [1, 4]

69 TOP_KS = [2, 6]

70 +TOP_KS_SMALL = [1, 2]

72 MOE_MARLIN_QUANT_TEST_CONFIGS = [

73 # AWQ-INT4

@@ -133,6 +135,13 @@ FUSED_MOE_MNK_FACTORS = [

135 (40000, 1024, 1024),

136 ]

137

138 +FUSED_MOE_MNK_FACTORS_SMALL_M = [

139 + (1, 128, 128),

140 + (1, 2048, 128),

141 + (2, 2048, 128),

142 + (2, 2048, 511),

143 +]

144 +

145 FUSED_MOE_WN16_MNK_FACTORS = [

146 (1, 128, 128),

147 (1, 1024, 1024),

@@ -330,6 +339,111 @@ def test_fused_moe(

339 )

340

341

342 +@pytest.mark.parametrize("m,n,k", FUSED_MOE_MNK_FACTORS_SMALL_M)

343 +@pytest.mark.parametrize("e", NUM_EXPERTS_LARGE)

344 +@pytest.mark.parametrize("topk", TOP_KS_SMALL)

345 +@pytest.mark.parametrize("dtype", [torch.bfloat16])

346 +@pytest.mark.parametrize("padding", [True, False])

347 +@pytest.mark.parametrize("chunk_size", [8192])

348 +def test_naive_block_assignment_moe(

349 + m: int,

350 + n: int,

351 + k: int,

352 + e: int,

353 + topk: int,

354 + dtype: torch.dtype,

355 + padding: bool,

356 + chunk_size: int,

357 + monkeypatch,

358 + workspace_init,

359 +):

360 + current_platform.seed_everything(7)

361 +

362 + monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(chunk_size))

363 +

364 + #

365 + # Setup test data

366 + #

367 +

368 + #

369 + # Setup test data

370 + #

371 +

372 + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10

373 + w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10

374 + w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10

375 +

376 + score = torch.randn((m, e), device="cuda", dtype=dtype)

377 +

378 + e_map = None

379 +

380 + #

381 + # Setup test functions

382 + #

383 + quant_config = FUSED_MOE_UNQUANTIZED_CONFIG

384 +

385 + m_fused_moe_fn = modular_triton_fused_moe(quant_config)

386 +

387 + def m_fused_moe(

388 + a: torch.Tensor,

389 + w1: torch.Tensor,

390 + w2: torch.Tensor,

391 + score: torch.Tensor,

392 + topk: int,

393 + global_num_experts: int = -1,

394 + expert_map: torch.Tensor | None = None,

395 + ) -> torch.Tensor:

396 + topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)

397 + return m_fused_moe_fn(

398 + a,

399 + w1,

400 + w2,

401 + topk_weights,

402 + topk_ids,

403 + global_num_experts=global_num_experts,

404 + expert_map=expert_map,

405 + )

406 +

407 + fused_moe_fn = functools.partial(fused_moe, renormalize=False)

408 +

409 + #

410 + # Run tests

411 + #

412 + runner = functools.partial(

413 + run_moe_test,

414 + a=a,

415 + w1=w1,

416 + w2=w2,

417 + score=score,

418 + topk=topk,

419 + global_num_experts=e,

420 + expert_map=e_map,

421 + padding=padding,

422 + )

423 +

424 + # Note: for now use_compile will error out if the problem size is

425 + # large enough to trigger chunking. I'm leaving the flag and

426 + # setup code in case we are able to revisit this later.

427 + use_compile = False

428 +

429 + use_cudagraph = n >= 1024 and k >= 1024 and current_platform.is_cuda_alike()

430 +

431 + with set_current_vllm_config(vllm_config):

432 + baseline_output = runner(torch_moe, iterative_moe)

433 + runner(

434 + baseline_output,

435 + fused_moe_fn,

436 + use_compile=use_compile,

437 + use_cudagraph=use_cudagraph,

438 + )

439 + runner(

440 + baseline_output,

441 + m_fused_moe,

442 + use_compile=use_compile,

443 + use_cudagraph=use_cudagraph,

444 + )

445 +

446 +

447 @pytest.mark.parametrize("m,n,k", FUSED_MOE_WN16_MNK_FACTORS)

448 @pytest.mark.parametrize("e", NUM_EXPERTS)

449 @pytest.mark.parametrize("topk", TOP_KS)

vllm/model_executor/layers/fused_moe/fused_moe.py

+62 -21

@@ -351,6 +351,7 @@ def fused_moe_kernel(

351 # Block size for block-wise quantization

352 group_n: tl.constexpr,

353 group_k: tl.constexpr,

354 + naive_block_assignment: tl.constexpr,

355 # Meta-parameters

356 BLOCK_SIZE_M: tl.constexpr,

357 BLOCK_SIZE_N: tl.constexpr,

@@ -386,6 +387,9 @@ def fused_moe_kernel(

387 - expert_ids: A tensor containing the indices of the expert for each

388 block. It determines which expert matrix from B should be used for

389 each block in A.

390 + - naive_block_assignment: A boolean flag indicating whether to use naive

391 + token wise block assignment. If True, each block corresponds to a

392 + single token.

393 This kernel performs the multiplication of a token by its corresponding

394 expert matrix as determined by `expert_ids`. The sorting of

395 `sorted_token_ids` by expert index and padding ensures divisibility by

@@ -411,11 +415,20 @@ def fused_moe_kernel(

415 # and accumulate

416 # `a_ptrs` is a block of [BLOCK_SIZE_M, BLOCK_SIZE_K] pointers

417 # `b_ptrs` is a block of [BLOCK_SIZE_K, BLOCK_SIZE_N] pointers

418 + offs = tl.arange(0, BLOCK_SIZE_M).to(tl.int64)

419 num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)

420 if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:

421 return

417 - offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)

418 - offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)

422 + if not naive_block_assignment:

423 + offs_token_id = pid_m * BLOCK_SIZE_M + offs

424 + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)

425 + else:

426 + offs_token = tl.where(

427 + offs == 0,

428 + pid_m, # first element = pid_m

429 + num_valid_tokens, # remaining elements = constant

430 + )

431 +

432 token_mask = offs_token < num_valid_tokens

433

434 off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)

@@ -557,7 +570,7 @@ def invoke_fused_moe_wna16_cuda_kernel(

570 B_scale: torch.Tensor | None,

571 B_zp: torch.Tensor | None,

572 topk_weights: torch.Tensor | None,

560 - sorted_token_ids: torch.Tensor,

573 + sorted_token_ids: torch.Tensor | None,

574 expert_ids: torch.Tensor,

575 num_tokens_post_padded: torch.Tensor,

576 mul_routed_weight: bool,

@@ -705,7 +718,7 @@ def invoke_fused_moe_triton_kernel(

718 A_scale: torch.Tensor | None,

719 B_scale: torch.Tensor | None,

720 topk_weights: torch.Tensor | None,

708 - sorted_token_ids: torch.Tensor,

721 + sorted_token_ids: torch.Tensor | None,

722 expert_ids: torch.Tensor,

723 num_tokens_post_padded: torch.Tensor,

724 mul_routed_weight: bool,

@@ -722,7 +735,7 @@ def invoke_fused_moe_triton_kernel(

735 ):

736 assert topk_weights is not None or not mul_routed_weight

737 assert topk_weights is None or topk_weights.stride(1) == 1

725 - assert sorted_token_ids.stride(0) == 1

738 + assert sorted_token_ids is None or sorted_token_ids.stride(0) == 1

739

740 if use_fp8_w8a8 or use_int8_w8a8:

741 assert B_scale is not None

@@ -741,14 +754,18 @@ def invoke_fused_moe_triton_kernel(

754

755 M = A.size(0)

756 num_tokens = M * top_k

744 -

745 - EM = sorted_token_ids.size(0)

746 - if A.size(0) < config["BLOCK_SIZE_M"]:

747 - # optimize for small batch_size.

748 - # We assume that top_ids of each token is unique,

749 - # so num_valid_experts <= batch_size <= BLOCK_SIZE_M,

750 - # and we can skip some invalid blocks.

751 - EM = min(sorted_token_ids.size(0), A.size(0) * top_k * config["BLOCK_SIZE_M"])

757 + if sorted_token_ids is not None:

758 + EM = sorted_token_ids.size(0)

759 + if A.size(0) < config["BLOCK_SIZE_M"]:

760 + # optimize for small batch_size.

761 + # We assume that top_ids of each token is unique,

762 + # so num_valid_experts <= batch_size <= BLOCK_SIZE_M,

763 + # and we can skip some invalid blocks.

764 + EM = min(

765 + sorted_token_ids.size(0), A.size(0) * top_k * config["BLOCK_SIZE_M"]

766 + )

767 + else:

768 + EM = num_tokens * config["BLOCK_SIZE_M"]

769 grid = lambda META: (

770 triton.cdiv(EM, META["BLOCK_SIZE_M"])

771 * triton.cdiv(B.size(1), META["BLOCK_SIZE_N"]),

@@ -798,6 +815,7 @@ def invoke_fused_moe_triton_kernel(

815 use_int8_w8a8=use_int8_w8a8,

816 use_int8_w8a16=use_int8_w8a16,

817 per_channel_quant=per_channel_quant,

818 + naive_block_assignment=(sorted_token_ids is None),

819 HAS_BIAS=HAS_BIAS,

820 BLOCK_SIZE_K=BLOCK_SIZE_K,

821 **config,

@@ -812,7 +830,7 @@ def dispatch_fused_moe_kernel(

830 B_scale: torch.Tensor | None,

831 B_zp: torch.Tensor | None,

832 topk_weights: torch.Tensor | None,

815 - sorted_token_ids: torch.Tensor,

833 + sorted_token_ids: torch.Tensor | None,

834 expert_ids: torch.Tensor,

835 num_tokens_post_padded: torch.Tensor,

836 mul_routed_weight: bool,

@@ -829,7 +847,7 @@ def dispatch_fused_moe_kernel(

847 ) -> None:

848 assert topk_weights is not None or not mul_routed_weight

849 assert topk_weights is None or topk_weights.stride(1) == 1

832 - assert sorted_token_ids.stride(0) == 1

850 + assert sorted_token_ids is None or sorted_token_ids.stride(0) == 1

851

852 M = A.size(0)

853 num_tokens = M * top_k

@@ -2165,14 +2183,37 @@ def fused_experts_impl(

2183 block_shape=block_shape,

2184 )

2185

2168 - sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(

2169 - curr_topk_ids,

2170 - config["BLOCK_SIZE_M"],

2171 - global_num_experts,

2172 - expert_map,

2173 - ignore_invalid_experts=True,

2186 + # SPARSITY_FACTOR is a heuristic margin ensuring tokens_in_chunk * top_k

2187 + # activates only a small fraction of total experts

2188 + SPARSITY_FACTOR = 4

2189 + # block quantized code path is not implemented yet.

2190 + naive_block_assignment = (

2191 + expert_map is None

2192 + and tokens_in_chunk * top_k_num * SPARSITY_FACTOR <= global_num_experts

2193 + and not (

2194 + (use_int8_w8a16 or use_int4_w4a16)

2195 + and block_shape is not None

2196 + and block_shape[1] > 0

2197 + )

2198 )

2199

2200 + if not naive_block_assignment:

2201 + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(

2202 + curr_topk_ids,

2203 + config["BLOCK_SIZE_M"],

2204 + global_num_experts,

2205 + expert_map,

2206 + ignore_invalid_experts=True,

2207 + )

2208 + else:

2209 + max_num_tokens_padded = topk_ids.numel() * config["BLOCK_SIZE_M"]

2210 + expert_ids = curr_topk_ids.view(-1)

2211 + num_tokens_post_padded = torch.empty(

2212 + (1), dtype=torch.int32, device=topk_ids.device

2213 + )

2214 + num_tokens_post_padded.fill_(max_num_tokens_padded)

2215 + sorted_token_ids = None

2216 +

2217 dispatch_fused_moe_kernel(

2218 qcurr_hidden_states,

2219 w1,