From 8a194b4696d061b48824f7ba7925380227629a06 Mon Sep 17 00:00:00 2001 From: Truong Vu Date: Fri, 9 Oct 2026 09:59:49 -0400 Subject: [PATCH 1/2] [HIP] topk_gating kernels: trap when vec_size does not divide num_experts The selection loop scans only full vectors, so the last num_experts % vec_size experts could never be selected. Dispatch never picks such a vec_size today; make the precondition explicit. --- csrc/include/topk_gating_kernels.cuh | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/csrc/include/topk_gating_kernels.cuh b/csrc/include/topk_gating_kernels.cuh index 4ccd356f8f2..f388ea24cc9 100644 --- a/csrc/include/topk_gating_kernels.cuh +++ b/csrc/include/topk_gating_kernels.cuh @@ -902,6 +902,11 @@ __global__ void topk_gating_kernel( static constexpr int vec_size = opus::vector_traits::size(); using vec_i = opus::vector_t; const int num_experts_vec = num_experts / vec_size; + // The top-k selection below scans whole vectors only, so experts past the + // last full vector would never be picked. topk_gating_launch chooses f32vec + // so that vec_size divides num_experts; stop loudly if a caller does not. + if(num_experts % vec_size != 0) + __builtin_trap(); // Step 1: load + score function // For softmax, bias is NOT added here -- it's added AFTER normalization @@ -1076,6 +1081,11 @@ __global__ void topk_gating_kernel_smem_n( static constexpr int vec_size = opus::vector_traits::size(); using vec_i = opus::vector_t; const int num_experts_vec = num_experts / vec_size; + // The top-k selection below scans whole vectors only, so experts past the + // last full vector would never be picked. topk_gating_launch chooses f32vec + // so that vec_size divides num_experts; stop loudly if a caller does not. + if(num_experts % vec_size != 0) + __builtin_trap(); // Guard: trailing tokens of the last block may be out of range. const bool valid = token_idx < num_tokens; From d54a07b93036c1909dbd7543bced73c55f936b64 Mon Sep 17 00:00:00 2001 From: Truong Vu Date: Tue, 13 Oct 2026 10:09:20 -0400 Subject: [PATCH 2/2] [HIP] topk_gating: shorten the vec_size check comment --- csrc/include/topk_gating_kernels.cuh | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/csrc/include/topk_gating_kernels.cuh b/csrc/include/topk_gating_kernels.cuh index f388ea24cc9..e85e41045d8 100644 --- a/csrc/include/topk_gating_kernels.cuh +++ b/csrc/include/topk_gating_kernels.cuh @@ -902,9 +902,8 @@ __global__ void topk_gating_kernel( static constexpr int vec_size = opus::vector_traits::size(); using vec_i = opus::vector_t; const int num_experts_vec = num_experts / vec_size; - // The top-k selection below scans whole vectors only, so experts past the - // last full vector would never be picked. topk_gating_launch chooses f32vec - // so that vec_size divides num_experts; stop loudly if a caller does not. + // Top-k scans whole vectors only, so a partial last vector would be skipped. + // topk_gating_launch picks f32vec so that vec_size divides num_experts. if(num_experts % vec_size != 0) __builtin_trap(); @@ -1081,9 +1080,6 @@ __global__ void topk_gating_kernel_smem_n( static constexpr int vec_size = opus::vector_traits::size(); using vec_i = opus::vector_t; const int num_experts_vec = num_experts / vec_size; - // The top-k selection below scans whole vectors only, so experts past the - // last full vector would never be picked. topk_gating_launch chooses f32vec - // so that vec_size divides num_experts; stop loudly if a caller does not. if(num_experts % vec_size != 0) __builtin_trap();