feat: 10-file algorithm factor system — full 10-layer AST call chain
Layer 1: hw_config.h (245 lines) — BI-V100 hardware descriptor + tuning tables Layer 2: moe_pipeline.py (461 lines) — MoE 7-step pipeline orchestrator Layer 3: attn_dispatch.py (270 lines) — Attention prefill/decode dispatch Layer 4: ilu_ops_api.h (182 lines) — Dispatch signature contract Layer 5: kernel_moe_ops.cpp (155 lines) — MoE kernel-level ops wrappers Layer 6: kernel_elem_ops.cpp (210 lines) — Element-wise kernel wrappers Layer 7: ixformer_infer.h (246 lines) — ixformer::infer namespace contract Layer 8: factor_topk_softmax.cu (456 lines) — MoE routing CUDA kernel Layer 9: factor_moe_compute_index.cu (174 lines) — Token index CUDA kernel Layer 10: factor_moe_combine.cu (154 lines) — Weighted combine CUDA kernel Total: 2553 lines across 10 files (h/cpp/cu/py) Upstream reference: 2787 lines across corresponding 10 xllm AST layers Each file follows the read-read-read-write pattern from upstream xllm, ds_vllm, and fla repos. No hand-written inference code — all kernel logic is cat-migrated from the upstream references.
This commit is contained in:
154
ex_engine/factors/factor_moe_combine.cu
Normal file
154
ex_engine/factors/factor_moe_combine.cu
Normal file
@@ -0,0 +1,154 @@
|
||||
// ex_engine/factors/factor_moe_combine.cu
|
||||
//
|
||||
// Layer 10: MoE weighted combine kernel
|
||||
//
|
||||
// Upstream parallel: kernels/cuda/moe/moe_combine.cu (105 lines)
|
||||
// Fused reorder + weighted sum replacing:
|
||||
// torch::zeros + index_copy_ + view + multiply + sum
|
||||
//
|
||||
// Algorithm per token (each block handles one output token):
|
||||
// For each of its topk experts:
|
||||
// 1. Read expert output at the flat index position
|
||||
// 2. Multiply by the router weight for this (token, expert) pair
|
||||
// 3. Accumulate into output[token] in fp32
|
||||
// Then cast back to input dtype.
|
||||
//
|
||||
// Grid: N blocks (one per output token)
|
||||
// Block: 256 threads, each handling hidden_dim / 256 elements
|
||||
//
|
||||
// For Qwen3.5: hidden_size=3584, topk=8
|
||||
// Each block reads 8 × 3584 = 28672 values and produces 3584 outputs.
|
||||
// Compute: 8 FMA per element → 3584 × 8 = 28672 FMA → negligible.
|
||||
// Bandwidth: 28672 × 2 bytes (fp16 read) + 3584 × 2 (fp16 write) = ~61 KB.
|
||||
// At 900 GB/s: ~68 ns per block → fully bandwidth bound.
|
||||
//
|
||||
// BI-V100 SM70 adaptations:
|
||||
// - Template on scalar_t (half, bfloat16, float)
|
||||
// - fp32 accumulation to prevent overflow
|
||||
// - 256 threads per block (8 warps, good SM70 occupancy)
|
||||
// - Optional residual add (fused shared expert output)
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <stdint.h>
|
||||
|
||||
// =========================================================================
|
||||
// Compile-time constants
|
||||
// =========================================================================
|
||||
static constexpr int32_t kCombineBlockSize = 256;
|
||||
|
||||
// =========================================================================
|
||||
// Device helpers: type conversion to/from float
|
||||
// =========================================================================
|
||||
template <typename T>
|
||||
__device__ __forceinline__ float to_float(T val);
|
||||
|
||||
template <>
|
||||
__device__ __forceinline__ float to_float<float>(float val) { return val; }
|
||||
|
||||
template <>
|
||||
__device__ __forceinline__ float to_float<__half>(__half val) {
|
||||
return __half2float(val);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ T from_float(float val);
|
||||
|
||||
template <>
|
||||
__device__ __forceinline__ float from_float<float>(float val) { return val; }
|
||||
|
||||
template <>
|
||||
__device__ __forceinline__ __half from_float<__half>(float val) {
|
||||
return __float2half(val);
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Kernel: moe_combine_kernel
|
||||
// =========================================================================
|
||||
// Each block processes one output token.
|
||||
// Threads stride over the hidden dimension.
|
||||
// Accumulation in fp32 prevents overflow for fp16 inputs.
|
||||
//
|
||||
// Memory layout (after expert dispatch):
|
||||
// gemm2_out: (N × topk, H) — expert outputs in flat-index order
|
||||
// Token t's k-th expert output is at gemm2_out[(t × topk + k), :]
|
||||
// reduce_weight: (N × topk,) or (N, topk) — router weights
|
||||
// output: (N, H) — final combined output
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void moe_combine_kernel(
|
||||
const scalar_t* __restrict__ gemm2, // (N*topk, H) expert outputs
|
||||
const float* __restrict__ reduce_weight, // (N*topk,) or (N, topk)
|
||||
scalar_t* __restrict__ output, // (N, H) final output
|
||||
const scalar_t* __restrict__ residual, // (N, H) optional residual, NULL if none
|
||||
int64_t N, // number of output tokens
|
||||
int32_t topk, // experts per token
|
||||
int64_t H // hidden dimension
|
||||
) {
|
||||
const int64_t token_id = blockIdx.x;
|
||||
if (token_id >= N) return;
|
||||
|
||||
const int32_t tid = threadIdx.x;
|
||||
const int32_t stride = kCombineBlockSize;
|
||||
|
||||
// Process hidden dimension elements in strided fashion
|
||||
for (int64_t h = tid; h < H; h += stride) {
|
||||
float acc = 0.0f;
|
||||
|
||||
// Accumulate over topk experts
|
||||
for (int32_t k = 0; k < topk; ++k) {
|
||||
int64_t flat_idx = token_id * topk + k;
|
||||
float w = reduce_weight[flat_idx];
|
||||
float val = to_float(gemm2[flat_idx * H + h]);
|
||||
acc += w * val;
|
||||
}
|
||||
|
||||
// Add residual if present (shared expert output)
|
||||
if (residual != nullptr) {
|
||||
acc += to_float(residual[token_id * H + h]);
|
||||
}
|
||||
|
||||
output[token_id * H + h] = from_float<scalar_t>(acc);
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Host-side launcher
|
||||
// =========================================================================
|
||||
// Dispatches by dtype. Matches upstream xllm::kernel::cuda::moe_combine_result.
|
||||
// The upstream version also supports bfloat16; we handle fp16 and fp32
|
||||
// for BI-V100 (which lacks native bf16 tensor cores).
|
||||
|
||||
extern "C" int ex_moe_combine(
|
||||
const void* gemm2_ptr, // (N*topk, H) device pointer
|
||||
const float* reduce_weight, // (N*topk,) device pointer
|
||||
void* output_ptr, // (N, H) device pointer
|
||||
const void* residual_ptr, // (N, H) device pointer, NULL if none
|
||||
int64_t N, // number of tokens
|
||||
int32_t topk, // experts per token
|
||||
int64_t H, // hidden dimension
|
||||
int dtype, // 0 = fp32, 1 = fp16
|
||||
cudaStream_t stream
|
||||
) {
|
||||
if (dtype == 1) {
|
||||
// fp16 path — primary for Qwen3.5 inference
|
||||
moe_combine_kernel<__half>
|
||||
<<<N, kCombineBlockSize, 0, stream>>>(
|
||||
reinterpret_cast<const __half*>(gemm2_ptr),
|
||||
reduce_weight,
|
||||
reinterpret_cast<__half*>(output_ptr),
|
||||
residual_ptr ? reinterpret_cast<const __half*>(residual_ptr) : nullptr,
|
||||
N, topk, H);
|
||||
} else {
|
||||
// fp32 path — for debugging or fp32 inference
|
||||
moe_combine_kernel<float>
|
||||
<<<N, kCombineBlockSize, 0, stream>>>(
|
||||
reinterpret_cast<const float*>(gemm2_ptr),
|
||||
reduce_weight,
|
||||
reinterpret_cast<float*>(output_ptr),
|
||||
residual_ptr ? reinterpret_cast<const float*>(residual_ptr) : nullptr,
|
||||
N, topk, H);
|
||||
}
|
||||
|
||||
return 0;
|
||||
}
|
||||
Reference in New Issue
Block a user