06d7713db6
ixformer probe results: ✗ moe_w16a16_group_gemm NOT in ixformer .so ✗ CUTLASS grouped GEMM needs cuda/std (variadic function error on corex) ✓ ixformer_linear EXISTS (fused matmul) ✓ torch.mm works (uses corex cublas) Solution: moe_batched_gemm.cu - C++ loop over experts (eliminates Python overhead) - torch::mm for GEMM (corex cublas, not F.linear Python) - Fused silu_and_mul CUDA kernel (not PyTorch ops) - Weighted scatter-add in C++ - Skips empty experts (no wasted compute) Integration in qwen3_5.py: _USE_XLLM_MOE_GEMM dispatches to moe_experts_forward() Falls back to Python for-loop if not available Build: bash qwen3_6_scripts/build_xllm_kernels.sh
115 lines
4.3 KiB
Plaintext
115 lines
4.3 KiB
Plaintext
// moe_batched_gemm.cu — Fused MoE expert GEMM for BI-V100
|
||
//
|
||
// Replaces the Python for-loop over 256 experts with:
|
||
// 1. Gather tokens by expert (sorted by moe_compute_index)
|
||
// 2. Per-expert GEMM via cublas (torch::mm)
|
||
// 3. Fused silu activation
|
||
// 4. Per-expert down GEMM
|
||
// 5. Weighted scatter-add back to output
|
||
//
|
||
// This eliminates Python loop overhead (~256 iterations) and reduces
|
||
// kernel launch overhead by batching small GEMMs.
|
||
//
|
||
// Dimensions (Qwen3.5-35B-A3B, TP=4):
|
||
// w13: (256, 256, 2048) -> E=256, 2*I=256, H=2048
|
||
// w2: (256, 2048, 128) -> E=256, H=2048, I=128
|
||
|
||
#include <torch/extension.h>
|
||
#include <c10/cuda/CUDAGuard.h>
|
||
#include <ATen/cuda/CUDAContext.h>
|
||
|
||
namespace xllm::kernel::cuda {
|
||
|
||
// Fused SiLU-and-mul kernel (gate_up → act)
|
||
// gate_up: (N, 2*I), output: (N, I)
|
||
__global__ void silu_and_mul_inplace_kernel(
|
||
const __half* __restrict__ gate_up,
|
||
__half* __restrict__ output,
|
||
int64_t N, int64_t I) {
|
||
int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||
if (idx >= N * I) return;
|
||
|
||
int64_t row = idx / I;
|
||
int64_t col = idx % I;
|
||
|
||
float gate = __half2float(gate_up[row * 2 * I + col]);
|
||
float up = __half2float(gate_up[row * 2 * I + I + col]);
|
||
float silu_gate = gate / (1.0f + expf(-gate));
|
||
output[idx] = __float2half(silu_gate * up);
|
||
}
|
||
|
||
// Main function: batched expert forward
|
||
// Called from Python with pre-sorted token indices
|
||
torch::Tensor moe_experts_forward(
|
||
torch::Tensor hidden_states, // (T, H) — all tokens
|
||
torch::Tensor w13, // (E, 2*I, H) — gate+up weights
|
||
torch::Tensor w2, // (E, H, I) — down weights
|
||
torch::Tensor sorted_tok_ids, // (T*topk,) — which token for each slot
|
||
torch::Tensor sorted_weights, // (T*topk,) — routing weight for each slot
|
||
torch::Tensor expert_offsets, // (E+1,) — cumsum of expert_sizes, expert_offsets[0]=0
|
||
int64_t topk) {
|
||
|
||
auto stream = at::cuda::getCurrentCUDAStream();
|
||
int64_t T = hidden_states.size(0);
|
||
int64_t H = hidden_states.size(1);
|
||
int64_t E = w13.size(0);
|
||
int64_t two_I = w13.size(1); // 2*I
|
||
int64_t I = two_I / 2;
|
||
|
||
auto out = torch::zeros({T, H}, hidden_states.options());
|
||
|
||
// Get expert offsets on CPU for loop control
|
||
auto offsets_cpu = expert_offsets.to(torch::kCPU, torch::kInt64);
|
||
auto offsets_ptr = offsets_cpu.data_ptr<int64_t>();
|
||
|
||
for (int64_t eid = 0; eid < E; ++eid) {
|
||
int64_t start = offsets_ptr[eid];
|
||
int64_t end = offsets_ptr[eid + 1];
|
||
int64_t count = end - start;
|
||
if (count == 0) continue;
|
||
|
||
// Gather tokens for this expert
|
||
auto tok_ids = sorted_tok_ids.slice(0, start, end); // (count,)
|
||
auto tokens = hidden_states.index_select(0, tok_ids); // (count, H)
|
||
|
||
// GEMM 1: gate+up projection
|
||
// tokens (count, H) × w13[eid].T (H, 2*I) → (count, 2*I)
|
||
auto gate_up = torch::mm(tokens, w13[eid].t()); // (count, 2*I)
|
||
|
||
// Fused SiLU activation
|
||
auto act = torch::empty({count, I}, hidden_states.options());
|
||
if (hidden_states.dtype() == torch::kFloat16) {
|
||
int64_t total = count * I;
|
||
int block = 256;
|
||
int grid = (total + block - 1) / block;
|
||
silu_and_mul_inplace_kernel<<<grid, block, 0, stream>>>(
|
||
gate_up.data_ptr<at::Half>(),
|
||
act.data_ptr<at::Half>(),
|
||
count, I);
|
||
} else {
|
||
auto chunks = gate_up.chunk(2, /*dim=*/1);
|
||
act = torch::silu(chunks[0]) * chunks[1];
|
||
}
|
||
|
||
// GEMM 2: down projection
|
||
// act (count, I) × w2[eid].T (I, H) → (count, H)
|
||
auto expert_out = torch::mm(act, w2[eid].t()); // (count, H)
|
||
|
||
// Weighted scatter-add
|
||
auto weights = sorted_weights.slice(0, start, end).unsqueeze(1); // (count, 1)
|
||
out.index_add_(0, tok_ids, (expert_out * weights).to(out.dtype()));
|
||
}
|
||
|
||
return out;
|
||
}
|
||
|
||
} // namespace xllm::kernel::cuda
|
||
|
||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||
m.def("moe_experts_forward", &xllm::kernel::cuda::moe_experts_forward,
|
||
"Batched MoE expert forward (gather → GEMM → silu → GEMM → scatter)",
|
||
py::arg("hidden_states"), py::arg("w13"), py::arg("w2"),
|
||
py::arg("sorted_tok_ids"), py::arg("sorted_weights"),
|
||
py::arg("expert_offsets"), py::arg("topk"));
|
||
}
|