From 0b1f2ef280012441d83b02c989909fb896efed70 Mon Sep 17 00:00:00 2001 From: Enwei Zhu <21126786+syuoni@users.noreply.github.com> Date: Wed, 29 Jul 2026 04:08:03 +0000 Subject: [PATCH] fix(moe): pad trtllm-gen route map by one element to avoid OOB read Routed batched-GEMM kernels read one int32 past the end of ptrRouteMap from the last batch-dim CTA on every launch. The value is never consumed, so results stay correct, but the read faults with an illegal address whenever the allocation happens to end at a mapped-region boundary -- which is why it surfaces as flaky MoE autotune and inference crashes. The real fix is kernel-side and is being handled separately. Pad the allocation by one element so the already-shipped cubins stay in bounds; the +1 can be dropped once regenerated cubins land. Co-Authored-By: Claude Opus 5 Signed-off-by: Enwei Zhu <21126786+syuoni@users.noreply.github.com> --- csrc/trtllm_fused_moe_kernel_launcher.cu | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/csrc/trtllm_fused_moe_kernel_launcher.cu b/csrc/trtllm_fused_moe_kernel_launcher.cu index 02cc8a3ab79..27556772fd8 100644 --- a/csrc/trtllm_fused_moe_kernel_launcher.cu +++ b/csrc/trtllm_fused_moe_kernel_launcher.cu @@ -450,8 +450,10 @@ class FusedMoeLauncher { expanded_idx_to_permuted_idx = alloc_tensor({args->num_tokens * totalExpertsPerToken}, dl_int32, hidden_states.device()); + // WAR: the routed batched-GEMM kernels read one int32 past the end of the route map. + // TODO: drop the +1 once the fixed kernel cubins land. permuted_idx_to_token_idx = - alloc_tensor({max_num_padded_tokens}, dl_int32, hidden_states.device()); + alloc_tensor({max_num_padded_tokens + 1}, dl_int32, hidden_states.device()); if (gemm1_bias_type == batchedGemm::gemm::BiasType::Mn) { permuted_idx_to_expanded_idx = @@ -1940,8 +1942,10 @@ class FP4BlockScaleLauncher : public FusedMoeLauncher { total_num_padded_tokens = alloc_tensor({1}, dl_int32, hidden_states.device()); expanded_idx_to_permuted_idx = alloc_tensor({args->num_tokens * args->top_k}, dl_int32, hidden_states.device()); + // WAR: the routed batched-GEMM kernels read one int32 past the end of the route map. + // TODO: drop the +1 once the fixed kernel cubins land. permuted_idx_to_token_idx = - alloc_tensor({max_num_padded_tokens}, dl_int32, hidden_states.device()); + alloc_tensor({max_num_padded_tokens + 1}, dl_int32, hidden_states.device()); int64_t const size_of_expert_count_histogram = std::max(args->num_experts * 2, 256 * 2); expert_count_histogram =