Add Qwen3.6 MoE inference support

csrc/moe/moe_sparse.cu

@@ -34,6 +34,51 @@
 
 #define SPARSE_TILE_N 8  // output columns per block (= warps per block)
 
+// Weights FP8 E4M3 [local_experts, N, K], activations BF16 (W8A16).
+// Decode is memory-bound (~2 FLOP/byte), so dequant-in-registers GEMV
+// loses nothing to tensor cores and skips activation quantization.
+__global__ void moe_sparse_gemv_bf16_bf16_kernel(
+    const __nv_bfloat16* __restrict__ x,     // [T, K] or [T*top_k, K]
+    const __nv_bfloat16* __restrict__ w,     // [local_experts, N, K]
+    const int* __restrict__ topk_ids,        // [T, top_k] global expert ids
+    __nv_bfloat16* __restrict__ y,           // [T, top_k, N]
+    int N, int K, int top_k,
+    int expert_start, int local_experts,
+    int x_per_slot
+) {
+    int token = blockIdx.z;
+    int slot = blockIdx.y;
+    int eid = topk_ids[token * top_k + slot];
+    int lid = eid - expert_start;
+    if (lid < 0 || lid >= local_experts) return;
+
+    extern __shared__ float xs[];
+    const __nv_bfloat16* xrow =
+        x + (long long)(x_per_slot ? token * top_k + slot : token) * K;
+    for (int i = threadIdx.x; i < K; i += blockDim.x) {
+        xs[i] = __bfloat162float(xrow[i]);
+    }
+    __syncthreads();
+
+    int n = blockIdx.x * SPARSE_TILE_N + (threadIdx.x >> 5);
+    if (n >= N) return;
+    int lane = threadIdx.x & 31;
+
+    const __nv_bfloat16* wrow = w + ((long long)lid * N + n) * K;
+    float acc = 0.0f;
+    for (int i = lane; i < K; i += 32) {
+        acc += xs[i] * __bfloat162float(wrow[i]);
+    }
+
+    #pragma unroll
+    for (int o = 16; o > 0; o >>= 1) {
+        acc += __shfl_down_sync(0xffffffffu, acc, o);
+    }
+    if (lane == 0) {
+        y[((long long)token * top_k + slot) * N + n] = __float2bfloat16(acc);
+    }
+}
+
 // Weights FP8 E4M3 [local_experts, N, K], activations BF16 (W8A16).
 // Decode is memory-bound (~2 FLOP/byte), so dequant-in-registers GEMV
 // loses nothing to tensor cores and skips activation quantization.
@@ -194,6 +239,23 @@ __global__ void moe_weighted_sum_sparse_bf16_kernel(
 
 extern "C" {
 
+void launch_moe_sparse_gemv_bf16_bf16(
+    const void* x, const void* w, const void* topk_ids, void* y,
+    int num_tokens, int N, int K, int top_k,
+    int expert_start, int local_experts, int x_per_slot,
+    void* stream
+) {
+    dim3 grid((N + SPARSE_TILE_N - 1) / SPARSE_TILE_N, top_k, num_tokens);
+    int block = SPARSE_TILE_N * 32;
+    size_t smem = (size_t)K * sizeof(float);
+    moe_sparse_gemv_bf16_bf16_kernel<<<grid, block, smem, (cudaStream_t)stream>>>(
+        (const __nv_bfloat16*)x, (const __nv_bfloat16*)w, (const int*)topk_ids,
+        (__nv_bfloat16*)y,
+        N, K, top_k, expert_start, local_experts, x_per_slot
+    );
+    CUDA_CHECK_LAST_ERROR();
+}
+
 void launch_moe_sparse_gemv_fp8_bf16(
     const void* x, const void* w, const void* w_scales, const void* bias,
     const void* topk_ids, void* y,