iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: jdebache <jdebache@nvidia.com>
This commit is contained in:
Julien Debache
2026-09-01 15:30:11 -04:00
committed by GitHub
parent aa71f9bcc4
commit ab54f5bd83
6 changed files with 17 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
+2 -4
View File
@@ -20,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
)
from vllm.model_executor.layers.fused_moe.utils import (
fi_moe_largest_bucket,
trtllm_moe_pack_topk_ids_weights,
)
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
activation_to_flashinfer_int,
@@ -199,11 +198,10 @@ class TrtLlmBf16ExpertsModular(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsModular)
import flashinfer
from flashinfer.fused_moe import WeightLayout
# Pack topk ids and weights into format expected by the TRTLLM kernel.
packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
topk_ids = topk_ids.to(dtype=torch.int32)
result = flashinfer.fused_moe.trtllm_bf16_routed_moe(
topk_ids=packed_topk_ids,
topk_ids=(topk_ids, topk_weights),
hidden_states=hidden_states,
gemm1_weights=w1,
gemm2_weights=w2,
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+2 -4
View File
@@ -17,7 +17,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
)
from vllm.model_executor.layers.fused_moe.utils import (
fi_moe_largest_bucket,
trtllm_moe_pack_topk_ids_weights,
)
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
activation_to_flashinfer_int,
@@ -273,8 +272,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
import flashinfer
from flashinfer.fused_moe import Fp8QuantizationType, WeightLayout
# Pack topk ids and weights into format expected by the kernel.
packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
topk_ids = topk_ids.to(dtype=torch.int32)
if a1q_scale is None:
raise RuntimeError(
@@ -295,7 +293,7 @@ class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
hidden_states_scale = prepare_deepseek_fp8_x_sf(hidden_states, a1q_scale)
flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe(
topk_ids=packed_topk_ids,
topk_ids=(topk_ids, topk_weights),
routing_bias=None,
hidden_states=hidden_states,
hidden_states_scale=hidden_states_scale,
vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py
+5 -12
View File
@@ -38,9 +38,6 @@ from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import (
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
TopKWeightAndReduceNoOP,
)
from vllm.model_executor.layers.fused_moe.utils import (
trtllm_moe_pack_topk_ids_weights,
)
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
@@ -206,7 +203,7 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
packed_topk_ids: torch.Tensor,
topk_ids_and_weights: tuple[torch.Tensor, torch.Tensor],
gemm1_lora_delta: torch.Tensor | None,
global_num_experts: int,
a1q_scale: torch.Tensor | None,
@@ -250,10 +247,6 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
intermediate_size = self.intermediate_size_per_partition
K = output.size(1)
# Routing is computed outside the MoE; pack it into the
# (eid<<16)|w.bf16 format the routed API expects.
packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
# ---- Base-model fast path ----
# When no token in the batch selects a LoRA adapter, skip the LoRA machinery
# and run the plain base MoE with do_finalize=True, which writes the finalized
@@ -263,7 +256,7 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
hidden_states=hidden_states,
w1=w1,
w2=w2,
packed_topk_ids=packed_topk_ids,
topk_ids_and_weights=(topk_ids, topk_weights),
gemm1_lora_delta=None, # without LoRA, no delta
global_num_experts=global_num_experts,
a1q_scale=a1q_scale,
@@ -337,7 +330,7 @@ class _TrtLlmLoRAExpertsBase(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
hidden_states=hidden_states,
w1=w1,
w2=w2,
packed_topk_ids=packed_topk_ids,
topk_ids_and_weights=(topk_ids, topk_weights),
gemm1_lora_delta=gemm1_lora_delta,
global_num_experts=global_num_experts,
a1q_scale=a1q_scale,
@@ -518,7 +511,7 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase):
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
packed_topk_ids: torch.Tensor,
topk_ids_and_weights: tuple[torch.Tensor, torch.Tensor],
gemm1_lora_delta: torch.Tensor | None,
global_num_experts: int,
a1q_scale: torch.Tensor | None,
@@ -533,7 +526,7 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase):
# the caller's buffer via output= so it finalizes in place -- no copy.
do_finalize = gemm1_lora_delta is None
ret = flashinfer.fused_moe.trtllm_bf16_routed_moe(
topk_ids=packed_topk_ids,
topk_ids=topk_ids_and_weights,
hidden_states=hidden_states,
gemm1_weights=w1,
gemm2_weights=w2,
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+3 -3
View File
@@ -20,7 +20,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
)
from vllm.model_executor.layers.fused_moe.utils import (
fi_moe_largest_bucket,
trtllm_moe_pack_topk_ids_weights,
)
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
activation_to_flashinfer_int,
@@ -363,9 +362,8 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
) -> None:
from flashinfer import trtllm_fp4_block_scale_routed_moe
packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
trtllm_fp4_block_scale_routed_moe(
topk_ids=packed_tensor,
topk_ids=(topk_ids, topk_weights),
routing_bias=None,
hidden_states=x_quant,
hidden_states_scale=x_scale,
@@ -417,6 +415,8 @@ class TrtLlmMxfp4ExpertsModular(TrtLlmMxfp4ExpertsBase, mk.FusedMoEExpertsModula
expert_tokens_meta: mk.ExpertTokensMetadata | None,
apply_router_weight_on_input: bool,
):
topk_ids = topk_ids.to(dtype=torch.int32)
topk = topk_ids.size(-1)
local_num_experts = w1.size(0)
local_expert_offset = self.moe_config.ep_rank * local_num_experts
vllm/model_executor/layers/fused_moe/experts/trtllm_nvfp4_moe.py
+5 -8
View File
@@ -20,10 +20,7 @@ from vllm.model_executor.layers.fused_moe.moe_output import (
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
TopKWeightAndReduceNoOP,
)
from vllm.model_executor.layers.fused_moe.utils import (
fi_moe_largest_bucket,
trtllm_moe_pack_topk_ids_weights,
)
from vllm.model_executor.layers.fused_moe.utils import fi_moe_largest_bucket
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
activation_to_flashinfer_int,
has_flashinfer_situ_activation,
@@ -369,14 +366,11 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
)
else:
block_scale, per_token_scale = a1q_scale, None
# Pack topk ids and weights into format expected by the kernel.
packed_tensor = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights)
output1_scale_gate_scalar = self.quant_config.g1_alphas
# Invoke kernel.
flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe(
topk_ids=packed_tensor,
topk_ids=(topk_ids, topk_weights),
routing_bias=None,
hidden_states=hidden_states,
hidden_states_scale=block_scale.view(torch.float8_e4m3fn).reshape(
@@ -434,6 +428,9 @@ class TrtLlmNvFp4ExpertsModular(TrtLlmNvFp4ExpertsBase, mk.FusedMoEExpertsModula
# Per-token defers input quant to _invoke_kernel, so a1q_scale is None.
assert a1q_scale is not None or self.per_token_activation
# DeepEP produces int64 indexes.
topk_ids = topk_ids.to(dtype=torch.int32)
M = hidden_states.shape[0]
chunk_size = self._get_chunk_size()
vllm/model_executor/layers/fused_moe/utils.py
-29
View File
@@ -503,35 +503,6 @@ def fi_moe_largest_bucket(moe_config: "FusedMoEConfig") -> int:
return max(moe_config.max_num_tokens * moe_config.dp_size, 8192)
def trtllm_moe_pack_topk_ids_weights(
topk_ids: torch.Tensor,
topk_weights: torch.Tensor,
block_size: int = 1024,
) -> torch.Tensor:
assert topk_ids.shape == topk_weights.shape
assert topk_ids.is_contiguous() and topk_weights.is_contiguous()
original_shape = topk_ids.shape
ids_flat = topk_ids.reshape(-1)
weights_flat = topk_weights.reshape(-1)
n_elements = ids_flat.numel()
output = torch.empty(n_elements, dtype=torch.int32, device=topk_ids.device)
use_gdc = current_platform.is_cuda() and current_platform.has_device_capability(90)
grid = (triton.cdiv(n_elements, block_size),)
_pack_topk_ids_weights_kernel[grid](
ids_flat,
weights_flat,
output,
n_elements,
BLOCK_SIZE=block_size,
USE_GDC=use_gdc,
launch_pdl=use_gdc,
)
return output.reshape(original_shape)
@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
def _swiglu_limit_torch(
output: torch.Tensor,