@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / fdadb6f43a
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 15, 02:15 UTC

[Bugfix] Fused MoE Modular Kernel chunking loop (#20392)

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>

Varun Sundar Rabindranath committed Jul 10, 2025 at 16:31 UTC fdadb6f43a3fbba1e0bedfb904ef920257963c50
4 files changed +404 -107
tests/kernels/moe/test_count_expert_num_tokens.py new
+140
@@ -0,0 +1,140 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Tests compute_expert_num_tokens kernels
5 +"""
6 +
7 +import dataclasses
8 +from typing import Optional
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm.model_executor.layers.fused_moe.utils import count_expert_num_tokens
14 +
15 +
16 +@dataclasses.dataclass
17 +class TestTensors:
18 +
19 + topk_ids: torch.Tensor
20 + expert_map: Optional[torch.Tensor] = None
21 +
22 + def to_device(self, device: str):
23 + self.topk_ids = self.topk_ids.to(device=device)
24 + if self.expert_map is not None:
25 + self.expert_map = self.expert_map.to(device=device)
26 +
27 + @staticmethod
28 + def make(num_tokens: int, num_topk: int, num_experts: int, device: str,
29 + topk_ids_dtype: torch.dtype) -> "TestTensors":
30 +
31 + # make topk ids
32 + topk_ids = torch.empty((num_tokens, num_topk),
33 + device=device,
34 + dtype=torch.int64)
35 + for x in range(num_tokens):
36 + topk_ids[x] = torch.randperm(num_experts)[:num_topk]
37 + topk_ids = topk_ids.to(dtype=torch.int64)
38 + return TestTensors(topk_ids=topk_ids)
39 +
40 + def with_ep_rank(self, ep_rank: int, num_global_experts: int,
41 + num_local_experts: int, device: str):
42 + # make an expert map
43 + expert_map = torch.empty((num_global_experts),
44 + device=device,
45 + dtype=torch.int32)
46 + expert_map.fill_(-1)
47 + s = ep_rank * num_local_experts
48 + e = s + num_local_experts
49 + expert_map[s:e] = torch.tensor(list(range(num_local_experts)),
50 + device=device)
51 +
52 + return TestTensors(topk_ids=self.topk_ids.clone(),
53 + expert_map=expert_map)
54 +
55 +
56 +def ref_impl(tt: TestTensors, expert_num_tokens: torch.Tensor):
57 + # do the reference in cpu
58 + tt.to_device("cpu")
59 + expert_ids, counts = tt.topk_ids.unique(return_counts=True)
60 +
61 + for eid, count in zip(expert_ids, counts):
62 + if eid != -1 and tt.expert_map is not None:
63 + eid = tt.expert_map[eid]
64 +
65 + if eid == -1:
66 + continue
67 +
68 + expert_num_tokens[eid] += count
69 +
70 +
71 +def do_test_compute_expert_num_tokens(num_tokens: int, num_topk: int,
72 + num_experts: int, ep_size: int,
73 + topk_ids_dtype: torch.dtype):
74 +
75 + assert num_topk <= num_experts
76 +
77 + tt = TestTensors.make(num_tokens,
78 + num_topk,
79 + num_experts,
80 + topk_ids_dtype=topk_ids_dtype,
81 + device="cpu")
82 +
83 + num_global_experts = num_experts
84 + assert num_global_experts % ep_size == 0
85 + num_local_experts = num_global_experts // ep_size
86 + for ep_rank in range(ep_size):
87 + tt_rank = tt.with_ep_rank(ep_rank, num_global_experts,
88 + num_local_experts, "cpu")
89 +
90 + ref_expert_num_tokens = torch.zeros((num_local_experts),
91 + device="cpu",
92 + dtype=torch.int32)
93 + ref_impl(tt_rank, ref_expert_num_tokens)
94 + ref_expert_num_tokens = ref_expert_num_tokens.to("cuda")
95 +
96 + tt_rank.to_device("cuda")
97 + # Test with expert_map
98 + triton_expert_num_tokens_w_emap = count_expert_num_tokens(
99 + tt_rank.topk_ids, num_local_experts, tt_rank.expert_map)
100 +
101 + # Test without expert map
102 + topk_ids = tt_rank.expert_map[tt_rank.topk_ids].to(topk_ids_dtype)
103 + triton_expert_num_tokens_wo_emap = count_expert_num_tokens(
104 + topk_ids, num_local_experts, expert_map=None)
105 +
106 + torch.testing.assert_close(ref_expert_num_tokens,
107 + triton_expert_num_tokens_w_emap,
108 + atol=0,
109 + rtol=0)
110 + torch.testing.assert_close(ref_expert_num_tokens,
111 + triton_expert_num_tokens_wo_emap,
112 + atol=0,
113 + rtol=0)
114 +
115 +
116 +@pytest.mark.parametrize(
117 + "num_tokens", [1, 4, 8, 11, 19, 128, 127, 405, 1024, 3333, 6666, 7317])
118 +@pytest.mark.parametrize("num_topk", [2, 6, 8])
119 +@pytest.mark.parametrize("num_experts", [64])
120 +@pytest.mark.parametrize("ep_size", [1, 2, 4])
121 +@pytest.mark.parametrize("topk_ids_dtype", [torch.int64])
122 +def test_compute_expert_num_tokens(num_tokens: int, num_topk: int,
123 + num_experts: int, ep_size: int,
124 + topk_ids_dtype: torch.dtype):
125 + do_test_compute_expert_num_tokens(num_tokens, num_topk, num_experts,
126 + ep_size, topk_ids_dtype)
127 +
128 +
129 +@pytest.mark.parametrize("numel", list(range(1, 8192, 11)))
130 +@pytest.mark.parametrize("num_experts", [32])
131 +@pytest.mark.parametrize("ep_size", [2])
132 +@pytest.mark.parametrize("topk_ids_dtype", [torch.int64])
133 +def test_compute_expert_num_tokens_from_numel(numel: int, num_experts: int,
134 + ep_size: int,
135 + topk_ids_dtype: torch.dtype):
136 + do_test_compute_expert_num_tokens(num_tokens=numel,
137 + num_topk=1,
138 + num_experts=num_experts,
139 + ep_size=ep_size,
140 + topk_ids_dtype=topk_ids_dtype)
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+2 -2
@@ -98,7 +98,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
98 M_sum = round_up(M_sum, block_m)
99 workspace1 = (M_sum, max(N * 2, K))
100 workspace2 = (M_sum, max(N, K))
101 - output = (M * topk, K)
101 + output = (M, topk, K)
102 return (workspace1, workspace2, output, a.dtype)
103
104 def apply(
@@ -172,7 +172,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
172 dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
173 (a2q, a2q_scale), (w2, w2_scale), mm2_out, expert_ids)
174
175 - torch.index_select(mm2_out, 0, inv_perm, out=output)
175 + torch.index_select(mm2_out, 0, inv_perm, out=output.view((-1, K)))
176
177
178 def deep_gemm_moe_fp8(
vllm/model_executor/layers/fused_moe/modular_kernel.py
+190 -105
@@ -10,7 +10,8 @@ import torch
10
11 import vllm.envs as envs
12 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
13 -from vllm.model_executor.layers.fused_moe.utils import _resize_cache
13 +from vllm.model_executor.layers.fused_moe.utils import ( # yapf: disable
14 + _resize_cache, count_expert_num_tokens)
15 from vllm.utils import cdiv
16
17 #
@@ -421,6 +422,177 @@ class FusedMoEModularKernel(torch.nn.Module):
422 f"{fused_experts.__class__.__name__}."
423 f"{fused_experts.activation_formats[0]}")
424
425 + def _do_fused_experts(
426 + self, fused_out: Optional[torch.Tensor], a1: torch.Tensor,
427 + a1q: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor,
428 + topk_ids: torch.Tensor, activation: str, global_num_experts: int,
429 + local_num_experts: int, expert_map: Optional[torch.Tensor],
430 + w1_scale: Optional[torch.Tensor], w2_scale: Optional[torch.Tensor],
431 + w1_zp: Optional[torch.Tensor], w2_zp: Optional[torch.Tensor],
432 + a1q_scale: Optional[torch.Tensor],
433 + a2_scale: Optional[torch.Tensor],
434 + expert_tokens_meta: Optional[ExpertTokensMetadata]
435 + ) -> torch.Tensor:
436 +
437 + _, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
438 +
439 + (workspace13_shape, workspace2_shape, fused_out_shape,
440 + workspace_dtype) = self.fused_experts.workspace_shapes(
441 + a1, a1q, M, N, K, top_k, global_num_experts, local_num_experts)
442 +
443 + # We can reuse the memory between cache1 and cache3 because by the
444 + # time we need cache3, we're done with cache1.
445 + workspace13 = torch.empty(prod(workspace13_shape),
446 + device=a1.device,
447 + dtype=workspace_dtype)
448 + workspace2 = torch.empty(prod(workspace2_shape),
449 + device=a1.device,
450 + dtype=workspace_dtype)
451 +
452 + assert fused_out is None or fused_out.shape == fused_out_shape, (
453 + f"fused_out {fused_out.shape} but expected {fused_out_shape}")
454 + if fused_out is None:
455 + # reuse workspace13 for the output
456 + fused_out = _resize_cache(workspace13, fused_out_shape)
457 +
458 + self.fused_experts.apply(fused_out,
459 + a1q,
460 + w1,
461 + w2,
462 + topk_ids=topk_ids,
463 + activation=activation,
464 + global_num_experts=global_num_experts,
465 + expert_map=expert_map,
466 + w1_scale=w1_scale,
467 + w2_scale=w2_scale,
468 + w1_zp=w1_zp,
469 + w2_zp=w2_zp,
470 + a1q_scale=a1q_scale,
471 + a2_scale=a2_scale,
472 + workspace13=workspace13,
473 + workspace2=workspace2,
474 + expert_tokens_meta=expert_tokens_meta)
475 +
476 + return fused_out
477 +
478 + def _maybe_chunk_fused_experts(
479 + self, a1: torch.Tensor, a1q: torch.Tensor, w1: torch.Tensor,
480 + w2: torch.Tensor, topk_ids: torch.Tensor, activation: str,
481 + global_num_experts: int, local_num_experts: int,
482 + expert_map: Optional[torch.Tensor],
483 + w1_scale: Optional[torch.Tensor], w2_scale: Optional[torch.Tensor],
484 + w1_zp: Optional[torch.Tensor], w2_zp: Optional[torch.Tensor],
485 + a1q_scale: Optional[torch.Tensor],
486 + a2_scale: Optional[torch.Tensor],
487 + expert_tokens_meta: Optional[ExpertTokensMetadata]
488 + ) -> torch.Tensor:
489 +
490 + _, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
491 +
492 + CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
493 + num_chunks = cdiv(M, CHUNK_SIZE)
494 +
495 + if not self.fused_experts.supports_chunking() or num_chunks == 1:
496 + return self._do_fused_experts(
497 + fused_out=None,
498 + a1=a1,
499 + a1q=a1q,
500 + w1=w1,
501 + w2=w2,
502 + topk_ids=topk_ids,
503 + activation=activation,
504 + global_num_experts=global_num_experts,
505 + local_num_experts=local_num_experts,
506 + expert_map=expert_map,
507 + w1_scale=w1_scale,
508 + w2_scale=w2_scale,
509 + w1_zp=w1_zp,
510 + w2_zp=w2_zp,
511 + a1q_scale=a1q_scale,
512 + a2_scale=a2_scale,
513 + expert_tokens_meta=expert_tokens_meta)
514 +
515 + # Chunking required case
516 + assert num_chunks > 1
517 +
518 + # Construct the entire output that can then be processed in chunks.
519 + (_, _, fused_out_shape,
520 + _) = self.fused_experts.workspace_shapes(a1, a1q, M, N, K, top_k,
521 + global_num_experts,
522 + local_num_experts)
523 + fused_out = torch.empty(fused_out_shape,
524 + device=a1q.device,
525 + dtype=a1.dtype)
526 +
527 + def slice_input_tensors(
528 + chunk_idx: int
529 + ) -> tuple[torch.Tensor, Optional[torch.Tensor],
530 + Optional[torch.Tensor], torch.Tensor]:
531 + s = chunk_idx * CHUNK_SIZE
532 + e = min(s + CHUNK_SIZE, M)
533 + return (a1q[s:e], _chunk_scales(a1q_scale, s, e),
534 + _chunk_scales(a2_scale, s, e), topk_ids[s:e])
535 +
536 + def slice_output_tensor(chunk_idx: int) -> torch.Tensor:
537 + assert fused_out.size(0) % M == 0, (
538 + f"fused_out shape {fused_out.shape} vs M {M}")
539 + factor = fused_out.size(0) // M
540 + out_chunk_size = CHUNK_SIZE * factor
541 + s = chunk_idx * out_chunk_size
542 + e = min(s + out_chunk_size, fused_out.size(0))
543 + return fused_out[s:e]
544 +
545 + def slice_expert_tokens_metadata(
546 + full_expert_tokens_meta: ExpertTokensMetadata,
547 + chunk_topk_ids: torch.Tensor, local_num_experts: int,
548 + expert_map: Optional[torch.Tensor]) -> ExpertTokensMetadata:
549 + # The existing expert_num_tokens is for the entire a1q
550 + # input. Chunking forces recomputation of the number
551 + # of tokens assigned to each expert.
552 + c_expert_num_tokens = count_expert_num_tokens(
553 + chunk_topk_ids, local_num_experts, expert_map)
554 +
555 + c_expert_num_tokens_cpu = None
556 + need_expert_num_tokens_cpu = (
557 + full_expert_tokens_meta.expert_num_tokens_cpu is not None)
558 + if need_expert_num_tokens_cpu:
559 + c_expert_num_tokens_cpu = c_expert_num_tokens.to(
560 + "cpu", non_blocking=True)
561 +
562 + return ExpertTokensMetadata(
563 + expert_num_tokens=c_expert_num_tokens,
564 + expert_num_tokens_cpu=c_expert_num_tokens_cpu)
565 +
566 + for chunk_idx in range(num_chunks):
567 + c_a1q, c_a1q_scale, c_a2_scale, c_topk_ids = (
568 + slice_input_tensors(chunk_idx))
569 +
570 + c_expert_tokens_meta = None
571 + if expert_tokens_meta is not None:
572 + c_expert_tokens_meta = slice_expert_tokens_metadata(
573 + expert_tokens_meta, c_topk_ids, local_num_experts,
574 + expert_map)
575 +
576 + self._do_fused_experts(fused_out=slice_output_tensor(chunk_idx),
577 + a1=a1,
578 + a1q=c_a1q,
579 + w1=w1,
580 + w2=w2,
581 + topk_ids=c_topk_ids,
582 + activation=activation,
583 + global_num_experts=global_num_experts,
584 + local_num_experts=local_num_experts,
585 + expert_map=expert_map,
586 + w1_scale=w1_scale,
587 + w2_scale=w2_scale,
588 + w1_zp=w1_zp,
589 + w2_zp=w2_zp,
590 + a1q_scale=c_a1q_scale,
591 + a2_scale=c_a2_scale,
592 + expert_tokens_meta=c_expert_tokens_meta)
593 +
594 + return fused_out
595 +
596 def forward(
597 self,
598 hidden_states: torch.Tensor,
@@ -512,110 +684,23 @@ class FusedMoEModularKernel(torch.nn.Module):
684 # and can never run into the tensor.numel() == 0 case.
685 fused_out = torch.empty_like(a1q).to(dtype=a1.dtype)
686 else:
515 - _, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
516 -
517 - if self.fused_experts.enable_chunking():
518 - CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
519 - num_chunks = cdiv(M, CHUNK_SIZE)
520 - else:
521 - CHUNK_SIZE = M
522 - num_chunks = 1
523 -
524 - if num_chunks == 1:
525 - (workspace13_shape, workspace2_shape, fused_out_shape,
526 - workspace_dtype) = self.fused_experts.workspace_shapes(
527 - a1, a1q, M, N, K, top_k, global_num_experts,
528 - local_num_experts)
529 - else:
530 - # Use the full M to get the final output shape.
531 - _, _, fused_out_shape, _ = (
532 - self.fused_experts.workspace_shapes(
533 - a1, a1q, M, N, K, top_k, global_num_experts,
534 - local_num_experts))
535 - # Use the CHUNK_SIZE to get the workspace shapes.
536 - workspace13_shape, workspace2_shape, _, workspace_dtype = (
537 - self.fused_experts.workspace_shapes(
538 - a1, a1q, CHUNK_SIZE, N, K, top_k, global_num_experts,
539 - local_num_experts))
540 -
541 - # We can reuse the memory between cache1 and cache3 because by the
542 - # time we need cache3, we're done with cache1.
543 - workspace13 = torch.empty(prod(workspace13_shape),
544 - device=a1.device,
545 - dtype=workspace_dtype)
546 - workspace2 = torch.empty(prod(workspace2_shape),
547 - device=a1.device,
548 - dtype=workspace_dtype)
549 -
550 - if num_chunks == 1:
551 - fused_out = _resize_cache(workspace13, fused_out_shape)
552 -
553 - self.fused_experts.apply(
554 - fused_out,
555 - a1q,
556 - w1,
557 - w2,
558 - topk_ids,
559 - activation=activation,
560 - global_num_experts=global_num_experts,
561 - expert_map=expert_map,
562 - w1_scale=w1_scale,
563 - w2_scale=w2_scale,
564 - w1_zp=w1_zp,
565 - w2_zp=w2_zp,
566 - a1q_scale=a1q_scale,
567 - a2_scale=a2_scale,
568 - workspace13=workspace13,
569 - workspace2=workspace2,
570 - expert_tokens_meta=expert_tokens_meta,
571 - )
572 - else:
573 - # The leading output dimension may not be equal to M, so
574 - # we compute output indices separately.
575 - M_out = fused_out_shape[0]
576 - assert M_out >= M
577 - factor = M_out // M
578 - assert factor > 0
579 - OUT_CHUNK_SIZE = CHUNK_SIZE * factor
580 -
581 - fused_out = torch.empty(fused_out_shape,
582 - device=a1q.device,
583 - dtype=workspace_dtype)
584 -
585 - assert cdiv(M_out, OUT_CHUNK_SIZE) == num_chunks, (
586 - f"{cdiv(M_out, OUT_CHUNK_SIZE)} == {num_chunks}")
587 -
588 - for chunk in range(num_chunks):
589 - begin_chunk_idx = chunk * CHUNK_SIZE
590 - end_chunk_idx = min((chunk + 1) * CHUNK_SIZE, M)
591 - begin_out_idx = chunk * OUT_CHUNK_SIZE
592 - end_out_idx = min((chunk + 1) * OUT_CHUNK_SIZE, M_out)
593 - curr_a1q = a1q[begin_chunk_idx:end_chunk_idx]
594 - curr_a1q_scale = _chunk_scales(a1q_scale, begin_chunk_idx,
595 - end_chunk_idx)
596 - curr_a2_scale = _chunk_scales(a2_scale, begin_chunk_idx,
597 - end_chunk_idx)
598 - curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
599 -
600 - self.fused_experts.apply(
601 - fused_out[begin_out_idx:end_out_idx],
602 - curr_a1q,
603 - w1,
604 - w2,
605 - curr_topk_ids,
606 - activation=activation,
607 - global_num_experts=global_num_experts,
608 - expert_map=expert_map,
609 - w1_scale=w1_scale,
610 - w2_scale=w2_scale,
611 - w1_zp=w1_zp,
612 - w2_zp=w2_zp,
613 - a1q_scale=curr_a1q_scale,
614 - a2_scale=curr_a2_scale,
615 - workspace13=workspace13,
616 - workspace2=workspace2,
617 - expert_tokens_meta=expert_tokens_meta,
618 - )
687 + fused_out = self._maybe_chunk_fused_experts(
688 + a1=a1,
689 + a1q=a1q,
690 + w1=w1,
691 + w2=w2,
692 + topk_ids=topk_ids,
693 + activation=activation,
694 + global_num_experts=global_num_experts,
695 + local_num_experts=local_num_experts,
696 + expert_map=expert_map,
697 + w1_scale=w1_scale,
698 + w2_scale=w2_scale,
699 + w1_zp=w1_zp,
700 + w2_zp=w2_zp,
701 + a1q_scale=a1q_scale,
702 + a2_scale=a2_scale,
703 + expert_tokens_meta=expert_tokens_meta)
704
705 self.prepare_finalize.finalize(output, fused_out, topk_weights,
706 topk_ids, apply_router_weight_on_input)
vllm/model_executor/layers/fused_moe/utils.py
+72
@@ -13,9 +13,81 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
13 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
14 quant_dequant_mxfp4)
15 from vllm.platforms import current_platform
16 +from vllm.triton_utils import tl, triton
17 from vllm.utils import cdiv
18
19
20 +@triton.jit
21 +def _count_expert_num_tokens(topk_ids_ptr, expert_num_tokens_ptr, num_experts,
22 + topk_numel, expert_map,
23 + HAS_EXPERT_MAP: tl.constexpr,
24 + BLOCK_SIZE: tl.constexpr):
25 +
26 + curr_expert = tl.program_id(0)
27 +
28 + offsets = tl.arange(0, BLOCK_SIZE)
29 + topk_ids_ptrs = topk_ids_ptr + offsets
30 +
31 + acc = tl.zeros((BLOCK_SIZE, ), dtype=tl.int32)
32 + for x in range(tl.cdiv(topk_numel, BLOCK_SIZE)):
33 + mask = offsets < (topk_numel - x * BLOCK_SIZE)
34 + expert_ids = tl.load(topk_ids_ptrs, mask=mask, other=-1)
35 + if HAS_EXPERT_MAP:
36 + expert_map_ptrs = expert_map + expert_ids
37 + expert_map_mask = expert_ids >= 0
38 + expert_ids = tl.load(expert_map_ptrs,
39 + mask=expert_map_mask,
40 + other=-1)
41 +
42 + has_curr_expert = tl.where(expert_ids == curr_expert, 1, 0)
43 + acc = acc + has_curr_expert
44 + topk_ids_ptrs += BLOCK_SIZE
45 +
46 + if curr_expert < num_experts:
47 + tl.store(expert_num_tokens_ptr + curr_expert, tl.sum(acc))
48 +
49 +
50 +def count_expert_num_tokens(
51 + topk_ids: torch.Tensor, num_local_experts: int,
52 + expert_map: Optional[torch.Tensor]) -> torch.Tensor:
53 + """
54 + Count the number to tokens assigned to each expert.
55 +
56 + Parameters:
57 + - topk_ids (torch.Tensor): Tensor mapping each token to its
58 + list of experts.
59 + - num_local_experts (int): Number of experts in this rank.
60 + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
61 + from the global expert space to the local expert space of the expert
62 + parallel shard.
63 +
64 + Returns:
65 + A tensor of size num_local_experts, where tensor[i] holds the number
66 + of tokens assigned to the ith expert.
67 + """
68 + assert topk_ids.dtype.is_signed, (
69 + "The kernel uses -1 to represent invalid topk_ids")
70 + expert_num_tokens = torch.empty((num_local_experts),
71 + device=topk_ids.device,
72 + dtype=torch.int32)
73 +
74 + grid = num_local_experts
75 + BLOCK_SIZE = min(topk_ids.numel(), 1024)
76 + BLOCK_SIZE = triton.next_power_of_2(BLOCK_SIZE)
77 +
78 + _count_expert_num_tokens[(grid, )](
79 + topk_ids,
80 + expert_num_tokens,
81 + num_local_experts,
82 + topk_ids.numel(),
83 + expert_map,
84 + HAS_EXPERT_MAP=expert_map is not None,
85 + BLOCK_SIZE=BLOCK_SIZE,
86 + )
87 +
88 + return expert_num_tokens
89 +
90 +
91 def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor:
92 """
93 Shrink the given tensor and apply the given view to it. This is