[TRTLLM-8637][feat] Optimize the routing kernel for DeepseekV3 (MoE CUTLASS backend); Add support for KimiK2 and Qwen-next (MoE TRTLLM backend) (NVIDIA#7761)

Signed-off-by: Christina Zhang <[email protected]>

Author: ChristinaZ

cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh

@@ -51,7 +51,7 @@ struct TopKRedType
     static __host__ __device__ inline TypeCmp makeCmpVal(T val, int32_t idx = 0)
     {
         auto valueBits = cub::Traits<T>::TwiddleIn(reinterpret_cast<typename cub::Traits<T>::UnsignedBits&>(val));
-        TypeCmp compactTmp = reinterpret_cast<TypeCmp&>(valueBits);
+        TypeCmp compactTmp = valueBits;
         compactTmp = (compactTmp << kMoveBits) | (0xFFFF & (kMaxIdx - idx));
         // Use 65535 minus idx to give higher priority to elements with smaller indices.
         return compactTmp;
@@ -162,9 +162,28 @@ struct Sort<4, RedType>
     }
 };
 
+template <int K, typename Type>
+__forceinline__ __device__ void reduceTopK(cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
+    int32_t (&outIdx)[K], Type value, int32_t idx, Type const minValue, int actualK = K)
+{
+    static_assert(K > 0, "Top K must have K > 0");
+    static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+    using RedType = TopKRedType<Type>;
+    RedType topK{value, idx};
+    typename RedType::TypeCmp packedMax{};
+#pragma unroll
+    for (int kk = 0; kk < actualK; ++kk) //@todo: check if actualK is correct
+    {
+        topK = kk > 0 && packedMax == topK.compValIdx ? RedType{minValue, idx} : topK;
+        // get the next largest value
+        packedMax = topK.reduce(warp);
+        RedType::unpack(out[kk], outIdx[kk], packedMax);
+    }
+};
+
 template <int K, typename Type, int N, bool IsSorted = false>
-__device__ void reduceTopK(cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K], int32_t (&outIdx)[K],
-    Type (&value)[N], int32_t (&idx)[N], Type minValue)
+__device__ void reduceTopKFunc(cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K], int32_t (&outIdx)[K],
+    Type (&value)[N], int32_t (&idx)[N], Type minValue, int actualK = K)
 {
     static_assert(K > 0, "Top K must have K > 0");
     static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
@@ -184,7 +203,7 @@ __device__ void reduceTopK(cg::thread_block_tile<kWARP_SIZE> const& warp, Type (
     }
     typename RedType::TypeCmp packedMax{};
 #pragma unroll
-    for (int kk = 0; kk < K; ++kk)
+    for (int kk = 0; kk < actualK; ++kk)
     {
         bool update = kk > 0 && packedMax == topK[0].compValIdx;
 #pragma unroll
@@ -198,6 +217,67 @@ __device__ void reduceTopK(cg::thread_block_tile<kWARP_SIZE> const& warp, Type (
     }
 };
 
+template <int K, typename Type, int N>
+__forceinline__ __device__ void reduceTopK(cg::thread_block_tile<kWARP_SIZE> const& warp, Type (&out)[K],
+    int32_t (&outIdx)[K], Type (&value)[N], int32_t (&idx)[N], Type const minValue, int actualK = K)
+{
+    static_assert(K > 0, "Top K must have K > 0");
+    static_assert(K < kWARP_SIZE, "Top K must have K < kWARP_SIZE");
+    static_assert(N > 0, "Top K must have N > 0");
+    static_assert(N <= 16, "Only support candidates number less than or equal to 16*32=512");
+    static_assert(
+        N <= 4 || N % 4 == 0, "Only support candidates number is a multiple of 4*32=128 or less than or equal to 4");
+    using RedType = TopKRedType<Type>;
+
+    if constexpr (N <= 4)
+    {
+        reduceTopKFunc<K, Type, N>(warp, out, outIdx, value, idx, minValue, actualK);
+    }
+    else
+    {
+
+        constexpr int numLoops = N / 4;
+        constexpr int numResults = (numLoops * K - 1) / kWARP_SIZE + 1;
+
+        Type topKBufferValue[numResults];
+        int32_t topKBufferIdx[numResults];
+        int32_t laneIdx = threadIdx.x % kWARP_SIZE;
+
+        for (int ii = 0; ii < numResults; ++ii)
+        {
+            topKBufferValue[ii] = minValue;
+            topKBufferIdx[ii] = ii * kWARP_SIZE - 1; //@todo: check if this is correct
+        }
+        for (int loop = 0; loop < numLoops; ++loop)
+        {
+            int start = loop * 4;
+            Type topKValue[K];
+            int32_t topKIdx[K];
+            Type inValue[4];
+            int32_t inIdx[4];
+            for (int i = 0; i < 4; ++i)
+            {
+                inValue[i] = value[start + i];
+                inIdx[i] = idx[start + i];
+            }
+            reduceTopKFunc<K, Type, 4>(warp, topKValue, topKIdx, inValue, inIdx, minValue, actualK);
+            int inOffset = laneIdx % K;
+            if (laneIdx >= loop * K && laneIdx < (loop + 1) * K)
+            {
+                topKBufferValue[0] = topKValue[inOffset];
+                topKBufferIdx[0] = topKIdx[inOffset];
+            }
+            if (loop == numLoops - 1 && (laneIdx < (numLoops * K - kWARP_SIZE)))
+            {
+                topKBufferValue[1] = topKValue[inOffset];
+                topKBufferIdx[1] = topKIdx[inOffset];
+            }
+        }
+
+        reduceTopKFunc<K, Type, numResults>(warp, out, outIdx, topKBufferValue, topKBufferIdx, minValue, actualK);
+    }
+};
+
 #undef TOPK_SWAP
 
 } // namespace reduce_topk