@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ab54f5bd83

[Chore] Remove redundant `_pack_topk_ids_weights_kernel` in TrtLLM NvFP4 MoE (#46872)

Signed-off-by: jdebache <jdebache@nvidia.com>

Julien Debache committed Sep 1, 2026 at 21:30 UTC ab54f5bd832d27fbe3bccbfa45af5775473301c3

6 files changed +17 -60

vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py

+2 -4

@@ -20,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

20 )

21 from vllm.model_executor.layers.fused_moe.utils import (

22 fi_moe_largest_bucket,

23 - trtllm_moe_pack_topk_ids_weights,

23 )

24 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

25 activation_to_flashinfer_int,

@@ -199,11 +198,10 @@ class TrtLlmBf16ExpertsModular(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsModular)

198 import flashinfer

199 from flashinfer.fused_moe import WeightLayout

200

202 - # Pack topk ids and weights into format expected by the TRTLLM kernel.

203 - packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)

201 + topk_ids = topk_ids.to(dtype=torch.int32)

202

203 result = flashinfer.fused_moe.trtllm_bf16_routed_moe(

206 - topk_ids=packed_topk_ids,

204 + topk_ids=(topk_ids, topk_weights),

205 hidden_states=hidden_states,

206 gemm1_weights=w1,

207 gemm2_weights=w2,

vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py

+2 -4

@@ -17,7 +17,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

17 )

18 from vllm.model_executor.layers.fused_moe.utils import (

19 fi_moe_largest_bucket,

20 - trtllm_moe_pack_topk_ids_weights,

20 )

21 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

22 activation_to_flashinfer_int,

@@ -273,8 +272,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):

272 import flashinfer

273 from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout

274

276 - # Pack topk ids and weights into format expected by the kernel.

277 - packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)

275 + topk_ids = topk_ids.to(dtype=torch.int32)

276

277 if a1q_scale is None:

278 raise RuntimeError(

@@ -295,7 +293,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):

293 hidden_states_scale = prepare_deepseek_fp8_x_sf(hidden_states, a1q_scale)

294

295 flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe(

298 - topk_ids=packed_topk_ids,

296 + topk_ids=(topk_ids, topk_weights),

297 routing_bias=None,

298 hidden_states=hidden_states,

299 hidden_states_scale=hidden_states_scale,

vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py

+5 -12

@@ -38,9 +38,6 @@ from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import (

38 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

39 TopKWeightAndReduceNoOP,

40 )

41 -from vllm.model_executor.layers.fused_moe.utils import (

42 - trtllm_moe_pack_topk_ids_weights,

43 -)

41 from vllm.platforms import current_platform

42 from vllm.triton_utils import tl, triton

43 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe

@@ -206,7 +203,7 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):

203 hidden_states: torch.Tensor,

204 w1: torch.Tensor,

205 w2: torch.Tensor,

209 - packed_topk_ids: torch.Tensor,

206 + topk_ids_and_weights: tuple[torch.Tensor, torch.Tensor],

207 gemm1_lora_delta: torch.Tensor | None,

208 global_num_experts: int,

209 a1q_scale: torch.Tensor | None,

@@ -250,10 +247,6 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):

247 intermediate_size = self.intermediate_size_per_partition

248 K = output.size(1)

249

253 - # Routing is computed outside the MoE; pack it into the

254 - # (eid<<16)|w.bf16 format the routed API expects.

255 - packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)

256 -

250 # ---- Base-model fast path ----

251 # When no token in the batch selects a LoRA adapter, skip the LoRA machinery

252 # and run the plain base MoE with do_finalize=True, which writes the finalized

@@ -263,7 +256,7 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):

256 hidden_states=hidden_states,

257 w1=w1,

258 w2=w2,

266 - packed_topk_ids=packed_topk_ids,

259 + topk_ids_and_weights=(topk_ids, topk_weights),

260 gemm1_lora_delta=None, # without LoRA, no delta

261 global_num_experts=global_num_experts,

262 a1q_scale=a1q_scale,

@@ -337,7 +330,7 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):

330 hidden_states=hidden_states,

331 w1=w1,

332 w2=w2,

340 - packed_topk_ids=packed_topk_ids,

333 + topk_ids_and_weights=(topk_ids, topk_weights),

334 gemm1_lora_delta=gemm1_lora_delta,

335 global_num_experts=global_num_experts,

336 a1q_scale=a1q_scale,

@@ -518,7 +511,7 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase):

511 hidden_states: torch.Tensor,

512 w1: torch.Tensor,

513 w2: torch.Tensor,

521 - packed_topk_ids: torch.Tensor,

514 + topk_ids_and_weights: tuple[torch.Tensor, torch.Tensor],

515 gemm1_lora_delta: torch.Tensor | None,

516 global_num_experts: int,

517 a1q_scale: torch.Tensor | None,

@@ -533,7 +526,7 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase):

526 # the caller's buffer via output= so it finalizes in place -- no copy.

527 do_finalize = gemm1_lora_delta is None

528 ret = flashinfer.fused_moe.trtllm_bf16_routed_moe(

536 - topk_ids=packed_topk_ids,

529 + topk_ids=topk_ids_and_weights,

530 hidden_states=hidden_states,

531 gemm1_weights=w1,

532 gemm2_weights=w2,

vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py

+3 -3

@@ -20,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

20 )

21 from vllm.model_executor.layers.fused_moe.utils import (

22 fi_moe_largest_bucket,

23 - trtllm_moe_pack_topk_ids_weights,

23 )

24 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

25 activation_to_flashinfer_int,

@@ -363,9 +362,8 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula

362 ) -> None:

363 from flashinfer import trtllm_fp4_block_scale_routed_moe

364

366 - packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)

365 trtllm_fp4_block_scale_routed_moe(

368 - topk_ids=packed_tensor,

366 + topk_ids=(topk_ids, topk_weights),

367 routing_bias=None,

368 hidden_states=x_quant,

369 hidden_states_scale=x_scale,

@@ -417,6 +415,8 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula

415 expert_tokens_meta: mk.ExpertTokensMetadata | None,

416 apply_router_weight_on_input: bool,

417 ):

418 + topk_ids = topk_ids.to(dtype=torch.int32)

419 +

420 topk = topk_ids.size(-1)

421 local_num_experts = w1.size(0)

422 local_expert_offset = self.moe_config.ep_rank * local_num_experts

vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py

+5 -8

@@ -20,10 +20,7 @@ from vllm.model_executor.layers.fused_moe.moe_output import (

20 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

21 TopKWeightAndReduceNoOP,

22 )

23 -from vllm.model_executor.layers.fused_moe.utils import (

24 - fi_moe_largest_bucket,

25 - trtllm_moe_pack_topk_ids_weights,

26 -)

23 +from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket

24 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

25 activation_to_flashinfer_int,

26 has_flashinfer_situ_activation,

@@ -369,14 +366,11 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula

366 )

367 else:

368 block_scale, per_token_scale = a1q_scale, None

372 -

373 - # Pack topk ids and weights into format expected by the kernel.

374 - packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)

369 output1_scale_gate_scalar = self.quant_config.g1_alphas

370

371 # Invoke kernel.

372 flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe(

379 - topk_ids=packed_tensor,

373 + topk_ids=(topk_ids, topk_weights),

374 routing_bias=None,

375 hidden_states=hidden_states,

376 hidden_states_scale=block_scale.view(torch.float8_e4m3fn).reshape(

@@ -434,6 +428,9 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula

428 # Per-token defers input quant to _invoke_kernel, so a1q_scale is None.

429 assert a1q_scale is not None or self.per_token_activation

430

431 + # DeepEP produces int64 indexes.

432 + topk_ids = topk_ids.to(dtype=torch.int32)

433 +

434 M = hidden_states.shape[0]

435 chunk_size = self._get_chunk_size()

436

vllm/model_executor/layers/fused_moe/utils.py

-29

@@ -503,35 +503,6 @@ def fi_moe_largest_bucket(moe_config: "FusedMoEConfig") -> int:

503 return max(moe_config.max_num_tokens * moe_config.dp_size, 8192)

504

505

506 -def trtllm_moe_pack_topk_ids_weights(

507 - topk_ids: torch.Tensor,

508 - topk_weights: torch.Tensor,

509 - block_size: int = 1024,

510 -) -> torch.Tensor:

511 - assert topk_ids.shape == topk_weights.shape

512 - assert topk_ids.is_contiguous() and topk_weights.is_contiguous()

513 -

514 - original_shape = topk_ids.shape

515 - ids_flat = topk_ids.reshape(-1)

516 - weights_flat = topk_weights.reshape(-1)

517 -

518 - n_elements = ids_flat.numel()

519 - output = torch.empty(n_elements, dtype=torch.int32, device=topk_ids.device)

520 -

521 - use_gdc = current_platform.is_cuda() and current_platform.has_device_capability(90)

522 - grid = (triton.cdiv(n_elements, block_size),)

523 - _pack_topk_ids_weights_kernel[grid](

524 - ids_flat,

525 - weights_flat,

526 - output,

527 - n_elements,

528 - BLOCK_SIZE=block_size,

529 - USE_GDC=use_gdc,

530 - launch_pdl=use_gdc,

531 - )

532 - return output.reshape(original_shape)

533 -

534 -

506 @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)

507 def _swiglu_limit_torch(

508 output: torch.Tensor,