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

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

[Misc] DP : Add ExpertTokensMetadata (#20332)

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

Varun Sundar Rabindranath committed Jul 9, 2025 at 20:33 UTC 805d62ca88e4fb534ab58aa884612228c77cc08d

12 files changed +117 -79

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

+4 -2

@@ -260,8 +260,11 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

260 a2_scale: Optional[torch.Tensor],

261 workspace13: torch.Tensor,

262 workspace2: torch.Tensor,

263 - expert_num_tokens: Optional[torch.Tensor],

263 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

264 ):

265 + assert expert_tokens_meta is not None

266 + expert_num_tokens = expert_tokens_meta.expert_num_tokens

267 +

268 import deep_gemm as dg

269 assert hidden_states.ndim == 3

270 assert self.block_shape is not None

@@ -287,7 +290,6 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

290 masked_m=expert_num_tokens,

291 expected_m=expected_m)

292

290 - assert expert_num_tokens is not None

293 a2q, a2q_scale = silu_mul_fp8_quant_deep_gemm(workspace1,

294 expert_num_tokens)

295

vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py

+2 -2

@@ -129,7 +129,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

129 a2_scale: Optional[torch.Tensor],

130 workspace13: torch.Tensor,

131 workspace2: torch.Tensor,

132 - expert_num_tokens: Optional[torch.Tensor],

132 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

133 ):

134 experts = (self.batched_deep_gemm_experts

135 if self.allow_deep_gemm else self.batched_triton_experts)

@@ -137,4 +137,4 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

137 experts.apply(output, hidden_states, w1, w2, topk_ids, activation,

138 global_num_experts, expert_map, w1_scale, w2_scale,

139 w1_zp, w2_zp, a1q_scale, a2_scale, workspace13,

140 - workspace2, expert_num_tokens)

140 + workspace2, expert_tokens_meta)

vllm/model_executor/layers/fused_moe/cutlass_moe.py

+7 -1

@@ -303,11 +303,17 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):

303 a2_scale: Optional[torch.Tensor],

304 workspace13: torch.Tensor,

305 workspace2: torch.Tensor,

306 - expert_num_tokens: Optional[torch.Tensor],

306 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

307 ):

308 assert w1_zp is None, "w1_zp is not supported in CUTLASS MoE"

309 assert w2_zp is None, "w2_zp is not supported in CUTLASS MoE"

310 +

311 + expert_num_tokens = None

312 + if expert_tokens_meta is not None:

313 + expert_num_tokens = expert_tokens_meta.expert_num_tokens

314 +

315 activation_callable = lambda o, i: self.activation(activation, o, i)

316 +

317 in_dtype = hidden_states.dtype

318 run_cutlass_moe_fp8(

319 output, hidden_states, w1, w2, topk_ids, activation_callable,

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+1 -1

@@ -119,7 +119,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

119 a2_scale: Optional[torch.Tensor],

120 workspace13: torch.Tensor,

121 workspace2: torch.Tensor,

122 - expert_num_tokens: Optional[torch.Tensor],

122 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

123 ):

124 import deep_gemm as dg

125 assert self.block_shape is not None

vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py

+17 -9

@@ -62,8 +62,9 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

63 has_scales = token_scales is not None

65 - (num_tokens_per_rank, num_tokens_per_rdma_rank, expert_num_tokens,

66 - is_token_in_rank, event) = self.buffer.get_dispatch_layout(

65 + (num_tokens_per_rank, num_tokens_per_rdma_rank,

66 + dispatch_expert_num_tokens, is_token_in_rank,

67 + event) = self.buffer.get_dispatch_layout(

68 topk_idx=rank_topk_ids,

69 num_experts=num_experts,

70 previous_event=None,

@@ -83,7 +84,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

84 num_tokens_per_rank=num_tokens_per_rank,

85 num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,

86 is_token_in_rank=is_token_in_rank,

86 - num_tokens_per_expert=expert_num_tokens,

87 + num_tokens_per_expert=dispatch_expert_num_tokens,

88 topk_idx=rank_topk_ids,

89 topk_weights=rank_topk_weights,

90 # expert_alignment rounds the number of tokens per expert

@@ -115,7 +116,13 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

116 num_experts - 1 if self.rank_expert_offset == 0 else 0,

117 expert_topk_ids + self.rank_expert_offset)

118

118 - return (expert_x, expert_x_scale, expert_num_tokens, expert_topk_ids,

119 + # Makes a GPU-CPU copy.

120 + # TODO (varun): Maybe it is better to re-compute the expert_num_tokens

121 + # on GPU.

122 + expert_tokens_meta = mk.ExpertTokensMetadata.make_from_list(

123 + expert_num_tokens_per_expert_list, device=expert_x.device)

124 +

125 + return (expert_x, expert_x_scale, expert_tokens_meta, expert_topk_ids,

126 expert_topk_weights)

127

128 def prepare(

@@ -129,8 +136,9 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

136 expert_map: Optional[torch.Tensor],

137 apply_router_weight_on_input: bool,

138 quant_config: FusedMoEQuantConfig,

132 - ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],

133 - Optional[torch.Tensor], Optional[torch.Tensor]]:

139 + ) -> tuple[torch.Tensor, Optional[torch.Tensor],

140 + Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],

141 + Optional[torch.Tensor]]:

142

143 if apply_router_weight_on_input:

144 topk = topk_ids.size(1)

@@ -149,7 +157,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

157 )

158 if a1q_scale is not None and a1q_scale.numel() == 1:

159 a1q_scale = a1q_scale.view(1, 1)

152 - (expert_x, expert_x_scale, expert_num_tokens, expert_topk_ids,

160 + (expert_x, expert_x_scale, expert_tokens_meta, expert_topk_ids,

161 expert_topk_weights) = self._do_dispatch(

162 tokens=a1q,

163 token_scales=a1q_scale,

@@ -159,7 +167,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

167 else:

168 # DeepEP kernels only support dispatching per-token-quant

169 # quantization. dispatch in bfloat16.

162 - (expert_x, _, expert_num_tokens, expert_topk_ids,

170 + (expert_x, _, expert_tokens_meta, expert_topk_ids,

171 expert_topk_weights) = self._do_dispatch(

172 tokens=a1,

173 token_scales=None,

@@ -176,7 +184,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

184 per_act_token_quant=False,

185 block_shape=quant_config.block_shape)

186

179 - return (expert_x, expert_x_scale, expert_num_tokens, expert_topk_ids,

187 + return (expert_x, expert_x_scale, expert_tokens_meta, expert_topk_ids,

188 expert_topk_weights)

189

190 def _apply_weights_and_reduce(self, num_tokens: int,

vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py

+7 -3

@@ -119,8 +119,9 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

119 expert_map: Optional[torch.Tensor],

120 apply_router_weight_on_input: bool,

121 quant_config: FusedMoEQuantConfig,

122 - ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],

123 - Optional[torch.Tensor], Optional[torch.Tensor]]:

122 + ) -> tuple[torch.Tensor, Optional[torch.Tensor],

123 + Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],

124 + Optional[torch.Tensor]]:

125

126 hidden_size = a1.size(1)

127 assert hidden_size in self.SUPPORTED_HIDDEN_SIZES, \

@@ -158,7 +159,10 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

159 expert_x, a1_scale, a2_scale, a1.dtype, quant_config.quant_dtype,

160 quant_config.per_act_token_quant, quant_config.block_shape)

161

161 - return (expert_x, expert_x_scale, expert_num_tokens, None, None)

162 + expert_tokens_meta = mk.ExpertTokensMetadata(

163 + expert_num_tokens=expert_num_tokens, expert_num_tokens_cpu=None)

164 +

165 + return (expert_x, expert_x_scale, expert_tokens_meta, None, None)

166

167 def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,

168 topk_weights: torch.Tensor, topk_ids: torch.Tensor,

vllm/model_executor/layers/fused_moe/fused_batched_moe.py

+32 -44

@@ -505,8 +505,9 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

505 expert_map: Optional[torch.Tensor],

506 apply_router_weight_on_input: bool,

507 quant_config: FusedMoEQuantConfig,

508 - ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],

509 - Optional[torch.Tensor], Optional[torch.Tensor]]:

508 + ) -> tuple[torch.Tensor, Optional[torch.Tensor],

509 + Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],

510 + Optional[torch.Tensor]]:

511 assert a1.dim() == 2

512 assert topk_ids.dim() == 2

513 assert topk_ids.size(0) == a1.size(0)

@@ -587,7 +588,10 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

588

589 assert b_a1_scale is None or b_a1_scale.ndim == 3

590

590 - return b_a1, b_a1_scale, tokens_per_expert, None, None

591 + expert_tokens_meta = mk.ExpertTokensMetadata(

592 + expert_num_tokens=tokens_per_expert, expert_num_tokens_cpu=None)

593 +

594 + return b_a1, b_a1_scale, expert_tokens_meta, None, None

595

596 def finalize(

597 self,

@@ -694,28 +698,19 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

698 else:

699 return t.to(f32) * group_broadcast(scale, t.shape)

700

697 - def apply(

698 - self,

699 - output: torch.Tensor,

700 - hidden_states: torch.Tensor,

701 - w1: torch.Tensor,

702 - w2: torch.Tensor,

703 - topk_ids: torch.Tensor,

704 - activation: str,

705 - global_num_experts: int,

706 - expert_map: Optional[torch.Tensor],

707 - w1_scale: Optional[torch.Tensor],

708 - w2_scale: Optional[torch.Tensor],

709 - w1_zp: Optional[torch.Tensor],

710 - w2_zp: Optional[torch.Tensor],

711 - a1q_scale: Optional[torch.Tensor],

712 - a2_scale: Optional[torch.Tensor],

713 - workspace13: torch.Tensor,

714 - workspace2: torch.Tensor,

715 - expert_num_tokens: Optional[torch.Tensor],

716 - ):

701 + def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,

702 + w1: torch.Tensor, w2: torch.Tensor, topk_ids: torch.Tensor,

703 + activation: str, global_num_experts: int,

704 + expert_map: Optional[torch.Tensor],

705 + w1_scale: Optional[torch.Tensor],

706 + w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],

707 + w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],

708 + a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,

709 + workspace2: torch.Tensor,

710 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata]):

711 assert hidden_states.dim() == 3

718 - assert expert_num_tokens is not None

712 + assert expert_tokens_meta is not None

713 + expert_num_tokens = expert_tokens_meta.expert_num_tokens

714

715 num_local_experts = w1.size(0)

716 assert num_local_experts == w1.size(0), (

@@ -902,26 +897,16 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

897 output = (num_experts, max_num_tokens * num_dp, K)

898 return (workspace13, workspace2, output, a.dtype)

899

905 - def apply(

906 - self,

907 - output: torch.Tensor,

908 - hidden_states: torch.Tensor,

909 - w1: torch.Tensor,

910 - w2: torch.Tensor,

911 - topk_ids: torch.Tensor,

912 - activation: str,

913 - global_num_experts: int,

914 - expert_map: Optional[torch.Tensor],

915 - w1_scale: Optional[torch.Tensor],

916 - w2_scale: Optional[torch.Tensor],

917 - w1_zp: Optional[torch.Tensor],

918 - w2_zp: Optional[torch.Tensor],

919 - a1q_scale: Optional[torch.Tensor],

920 - a2_scale: Optional[torch.Tensor],

921 - workspace13: torch.Tensor,

922 - workspace2: torch.Tensor,

923 - expert_num_tokens: Optional[torch.Tensor],

924 - ):

900 + def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,

901 + w1: torch.Tensor, w2: torch.Tensor, topk_ids: torch.Tensor,

902 + activation: str, global_num_experts: int,

903 + expert_map: Optional[torch.Tensor],

904 + w1_scale: Optional[torch.Tensor],

905 + w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],

906 + w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],

907 + a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,

908 + workspace2: torch.Tensor,

909 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata]):

910 # Check constraints.

911 if self.use_int4_w4a16:

912 assert hidden_states.size(-1) // 2 == w1.size(2), (

@@ -938,6 +923,9 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

923 assert hidden_states.dtype in [

924 torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn

925 ]

926 + assert expert_tokens_meta is not None

927 +

928 + expert_num_tokens = expert_tokens_meta.expert_num_tokens

929

930 E, max_num_tokens, N, K, top_k_num = mk._moe_problem_size(

931 hidden_states, w1, w2, topk_ids)

vllm/model_executor/layers/fused_moe/fused_moe.py

+1 -1

@@ -1630,7 +1630,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1630 a2_scale: Optional[torch.Tensor],

1631 workspace13: torch.Tensor,

1632 workspace2: torch.Tensor,

1633 - expert_num_tokens: Optional[torch.Tensor],

1633 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

1634 ):

1635 # Check constraints.

1636 if self.use_int4_w4a16:

vllm/model_executor/layers/fused_moe/modular_kernel.py

+34 -9

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 from abc import ABC, abstractmethod

4 +from dataclasses import dataclass

5 from enum import Enum

6 from math import prod

7 from typing import Optional, final

@@ -95,6 +96,26 @@ class FusedMoEActivationFormat(Enum):

96 BatchedExperts = "batched_experts",

99 +@dataclass

100 +class ExpertTokensMetadata:

101 + """

102 + Metadata regarding expert-token routing.

103 + """

104 + expert_num_tokens: torch.Tensor

105 + expert_num_tokens_cpu: Optional[torch.Tensor]

106 +

107 + @staticmethod

108 + def make_from_list(expert_num_tokens_list: list[int],

109 + device: str) -> "ExpertTokensMetadata":

110 + expert_num_tokens_cpu = torch.tensor(expert_num_tokens_list,

111 + device="cpu",

112 + dtype=torch.int32)

113 + return ExpertTokensMetadata(

114 + expert_num_tokens=expert_num_tokens_cpu.to(device,

115 + non_blocking=True),

116 + expert_num_tokens_cpu=expert_num_tokens_cpu)

117 +

118 +

119 # TODO: pass FusedMoEParallelConfig in as ctor parameter?

120 class FusedMoEPrepareAndFinalize(ABC):

121 """

@@ -114,8 +135,9 @@ class FusedMoEPrepareAndFinalize(ABC):

135 expert_map: Optional[torch.Tensor],

136 apply_router_weight_on_input: bool,

137 quant_config: FusedMoEQuantConfig,

117 - ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],

118 - Optional[torch.Tensor], Optional[torch.Tensor]]:

138 + ) -> tuple[torch.Tensor, Optional[torch.Tensor],

139 + Optional[ExpertTokensMetadata], Optional[torch.Tensor],

140 + Optional[torch.Tensor]]:

141 """

142 Perform any quantization (and/or) dispatching needed

143 for this kernel.

@@ -134,7 +156,8 @@ class FusedMoEPrepareAndFinalize(ABC):

156 Returns a tuple of:

157 - quantized + dispatched a.

158 - quantized + dispatched a1_scales.

137 - - Optional tensor as big as number of local experts that contains the

159 + - Optional ExpertTokensMetadata containing gpu/cpu tensors

160 + as big as the number of local experts with the information about the

161 number of tokens assigned to each local expert.

162 - Optional dispatched expert topk IDs

163 - Optional dispatched expert topk weight

@@ -318,7 +341,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

341 a2_scale: Optional[torch.Tensor],

342 workspace13: torch.Tensor,

343 workspace2: torch.Tensor,

321 - expert_num_tokens: Optional[torch.Tensor],

344 + expert_tokens_meta: Optional[ExpertTokensMetadata],

345 ):

346 """

347 This function computes the intermediate result of a Mixture of Experts

@@ -351,8 +374,10 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

374 must be large enough to hold output of either MoE gemm.

375 - workspace2 (torch.Tensor): A scratch tensor used for the activation

376 function.

354 - - expert_num_tokens: An optional tensor containing the number of tokens

355 - assigned to each expert when using batched experts format input.

377 + - expert_tokens_meta (Optional[ExpertTokensMetadata]) - An optional

378 + ExpertTokensMetadata object containing gpu/cpu tensors

379 + as big as the number of local experts with the information about the

380 + number of tokens assigned to each local expert.

381 """

382 raise NotImplementedError

383

@@ -458,7 +483,7 @@ class FusedMoEModularKernel(torch.nn.Module):

483 if global_num_experts == -1:

484 global_num_experts = local_num_experts

485

461 - (a1q, a1q_scale, expert_num_tokens, _expert_topk_ids,

486 + (a1q, a1q_scale, expert_tokens_meta, _expert_topk_ids,

487 _expert_topk_weights) = self.prepare_finalize.prepare(

488 a1,

489 a1_scale,

@@ -542,7 +567,7 @@ class FusedMoEModularKernel(torch.nn.Module):

567 a2_scale=a2_scale,

568 workspace13=workspace13,

569 workspace2=workspace2,

545 - expert_num_tokens=expert_num_tokens,

570 + expert_tokens_meta=expert_tokens_meta,

571 )

572 else:

573 # The leading output dimension may not be equal to M, so

@@ -589,7 +614,7 @@ class FusedMoEModularKernel(torch.nn.Module):

614 a2_scale=curr_a2_scale,

615 workspace13=workspace13,

616 workspace2=workspace2,

592 - expert_num_tokens=expert_num_tokens,

617 + expert_tokens_meta=expert_tokens_meta,

618 )

619

620 self.prepare_finalize.finalize(output, fused_out, topk_weights,

vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py

+7 -3

@@ -94,8 +94,9 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

94 expert_map: Optional[torch.Tensor],

95 apply_router_weight_on_input: bool,

96 quant_config: FusedMoEQuantConfig,

97 - ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],

98 - Optional[torch.Tensor], Optional[torch.Tensor]]:

97 + ) -> tuple[torch.Tensor, Optional[torch.Tensor],

98 + Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],

99 + Optional[torch.Tensor]]:

100 num_tokens = a1.size(0) # M

101 hidden_dim = a1.size(-1) # K

102

@@ -200,7 +201,10 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

201 expert_x_scale = expert_x_scale[:, :, :orig_a_scale_block_shape]

202 assert expert_x_scale.ndim == 3

203

203 - return expert_x, expert_x_scale, expert_num_tokens, None, None

204 + expert_tokens_meta = mk.ExpertTokensMetadata(

205 + expert_num_tokens=expert_num_tokens, expert_num_tokens_cpu=None)

206 +

207 + return expert_x, expert_x_scale, expert_tokens_meta, None, None

208

209 def finalize(

210 self,

vllm/model_executor/layers/fused_moe/prepare_finalize.py

+3 -2

@@ -38,8 +38,9 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):

38 expert_map: Optional[torch.Tensor],

39 apply_router_weight_on_input: bool,

40 quant_config: FusedMoEQuantConfig,

41 - ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],

42 - Optional[torch.Tensor], Optional[torch.Tensor]]:

41 + ) -> tuple[torch.Tensor, Optional[torch.Tensor],

42 + Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],

43 + Optional[torch.Tensor]]:

45 if apply_router_weight_on_input:

46 topk = topk_ids.size(1)

vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py

+2 -2

@@ -110,7 +110,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

110 a2_scale: Optional[torch.Tensor],

111 workspace13: torch.Tensor,

112 workspace2: torch.Tensor,

113 - expert_num_tokens: Optional[torch.Tensor],

113 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

114 ):

115 use_deep_gemm = (self.allow_deep_gemm

116 and _valid_deep_gemm(hidden_states, w1, w2))

@@ -135,5 +135,5 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

135 a2_scale,

136 workspace13,

137 workspace2,

138 - expert_num_tokens,

138 + expert_tokens_meta,

139 )