diff --git a/csrc/kernels/legacy/internode_ll.cu b/csrc/kernels/legacy/internode_ll.cu index b6be356a..cc36718a 100644 --- a/csrc/kernels/legacy/internode_ll.cu +++ b/csrc/kernels/legacy/internode_ll.cu @@ -702,7 +702,7 @@ void dispatch(void* packed_recv_x, int num_device_sms, cudaStream_t stream, int phases) { - constexpr int kNumMaxTopK = 11; + constexpr int kNumMaxTopK = 16; // Kimi-K3 routes top-16 const int num_warp_groups = ceil_div(num_experts, num_device_sms); const int num_warps_per_group = 32 / num_warp_groups; EP_HOST_ASSERT(num_warp_groups > 0 and num_warps_per_group > 0); @@ -1880,7 +1880,7 @@ void combine(void* combined_x, ); } - constexpr int kNumMaxTopk = 11; + constexpr int kNumMaxTopk = 16; // Kimi-K3 routes top-16 const int num_warp_groups = ceil_div(num_experts, num_device_sms); const int num_warps_per_group = 32 / num_warp_groups; const int num_recv_per_sm = ceil_div(num_combined_tokens, num_device_sms); diff --git a/csrc/kernels/legacy/launch.cuh b/csrc/kernels/legacy/launch.cuh index 60e30803..36351e52 100644 --- a/csrc/kernels/legacy/launch.cuh +++ b/csrc/kernels/legacy/launch.cuh @@ -118,6 +118,8 @@ case_macro(2560); \ case 3072: \ case_macro(3072); /* for gpt-oss */ \ + case 3584: \ + case_macro(3584); /* for kimi-k3 */ \ case 4096: \ case_macro(4096); \ case 5120: \