diff --git a/exllamav3/exllamav3_ext/bindings.cpp b/exllamav3/exllamav3_ext/bindings.cpp index 4c59860f..d1158821 100644 --- a/exllamav3/exllamav3_ext/bindings.cpp +++ b/exllamav3/exllamav3_ext/bindings.cpp @@ -208,6 +208,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) m.def("blocksparse_mlp_routing", &blocksparse_mlp_routing, "blocksparse_mlp_routing"); m.def("exl3_moe_max_concurrency", &exl3_moe_max_concurrency, "exl3_moe_max_concurrency"); m.def("exl3_moe", &exl3_moe, "exl3_moe"); + m.attr("EXL3_MOE_ADDITIVE_ABI_VERSION") = EXL3_MOE_ADDITIVE_ABI_VERSION; + m.def("exl3_moe_additive", &exl3_moe_additive, "exl3_moe_additive"); + m.def( + "exl3_moe_additive_fused", + &exl3_moe_additive_fused, + "exl3_moe_additive_fused" + ); m.def("bighead_attn", &bighead_attn, "bighead_attn"); m.def("bighead_attn_paged", &bighead_attn_paged, "bighead_attn_paged"); diff --git a/exllamav3/exllamav3_ext/quant/exl3_gemm_inner.cuh b/exllamav3/exllamav3_ext/quant/exl3_gemm_inner.cuh index d2b40f21..e1792338 100644 --- a/exllamav3/exllamav3_ext/quant/exl3_gemm_inner.cuh +++ b/exllamav3/exllamav3_ext/quant/exl3_gemm_inner.cuh @@ -30,7 +30,9 @@ void exl3_gemm_kernel_inner const int size_k, const int size_n, int* __restrict__ locks, - const half* post_scale + const half* post_scale, + const float output_scale = 1.0f, + const bool accumulate = false ) { const int TILEBLOCKS_M = TILESIZE_M / 16; @@ -580,6 +582,17 @@ void exl3_gemm_kernel_inner } #endif + if (output_scale != 1.0f) + { + #pragma unroll + for (int n = 0; n < FRAGS_N_PER_WARP; ++n) + { + #pragma unroll + for (int j = 0; j < 4; ++j) + frag_c[n][j] *= output_scale; + } + } + // First reduce all partial sums along k for the current slice threadblock_reduce(); @@ -594,6 +607,14 @@ void exl3_gemm_kernel_inner bool first = lock_i == 0; bool last = lock_i + lock_d == tiles_k; + // An additive pass starts from the preceding packed stage's result. + // Only the first k-slice reads it; later slices consume this pass's + // partial sum through the existing global reduction protocol. + if (!sub_k && first && accumulate) + { + read_sum_gl(); + } + // Second and subsequent threadblocks in column read back the intermediate sum from global memory if (!sub_k && !first) { diff --git a/exllamav3/exllamav3_ext/quant/exl3_moe.cu b/exllamav3/exllamav3_ext/quant/exl3_moe.cu index f1520c77..41decc06 100644 --- a/exllamav3/exllamav3_ext/quant/exl3_moe.cu +++ b/exllamav3/exllamav3_ext/quant/exl3_moe.cu @@ -1,4 +1,5 @@ #include +#include #include "exl3_gemm.cuh" #include @@ -9,8 +10,139 @@ namespace cg = cooperative_groups; #include "../util.cuh" #include "comp_units/exl3_moe_instances.cuh" #include "exl3_devctx.cuh" +#include +#include #include +// The fused route histogram/scan/stable-pack pipeline is adapted from +// @brandonmmusic-max's draft PR #246 (commit 704aefd), with sentinel routing +// added for expert maps and the additive execution path introduced here. +template +__device__ __forceinline__ int64_t exl3_route_expert +( + id_t route_id, + const int64_t* __restrict__ expert_map, + int64_t expert_map_size, + int num_buckets +) +{ + int64_t id = static_cast(route_id); + if (id < 0 || id >= expert_map_size) return num_buckets - 1; + int64_t expert = expert_map[id]; + if (expert < 0 || expert >= num_buckets) return num_buckets - 1; + return expert; +} + +template +__global__ void exl3_route_histogram_kernel +( + const id_t* __restrict__ topk_ids, + const int64_t* __restrict__ expert_map, + int64_t* __restrict__ expert_count, + int64_t num_routes, + int64_t expert_map_size, + int num_buckets +) +{ + for (int64_t r = blockIdx.x * blockDim.x + threadIdx.x; + r < num_routes; r += (int64_t) blockDim.x * gridDim.x) + { + int64_t e = exl3_route_expert( + topk_ids[r], expert_map, expert_map_size, num_buckets + ); + atomicAdd(reinterpret_cast(expert_count + e), 1ULL); + } +} + +__global__ void exl3_route_scan_kernel +( + const int64_t* __restrict__ expert_count, + int64_t* __restrict__ expert_offsets, + int num_buckets +) +{ + if (blockIdx.x || threadIdx.x) return; + int64_t sum = 0; + for (int e = 0; e < num_buckets; ++e) + { + expert_offsets[e] = sum; + sum += expert_count[e]; + } +} + +template +__device__ __forceinline__ half exl3_route_to_half(weight_t v) +{ + return __float2half(static_cast(v)); +} + +template <> +__device__ __forceinline__ half exl3_route_to_half(half v) +{ + return v; +} + +template <> +__device__ __forceinline__ half exl3_route_to_half<__nv_bfloat16> +( + __nv_bfloat16 v +) +{ + return __float2half(__bfloat162float(v)); +} + +template +__global__ void exl3_route_pack_stable_kernel +( + const id_t* __restrict__ topk_ids, + const weight_t* __restrict__ topk_weights, + const int64_t* __restrict__ expert_map, + const int64_t* __restrict__ expert_offsets, + int64_t* __restrict__ token_sorted, + half* __restrict__ weight_sorted, + int64_t num_routes, + int topk, + int64_t expert_map_size, + int num_buckets +) +{ + // One block per local expert, including the sentinel. Each block scans + // routes in source order to preserve stable expert-grouped ordering. + __shared__ int flags[256]; + __shared__ int running; + if (threadIdx.x == 0) running = 0; + __syncthreads(); + const int64_t expert = blockIdx.x; + for (int64_t base = 0; base < num_routes; base += blockDim.x) + { + int64_t r = base + threadIdx.x; + int flag = 0; + if (r < num_routes) + flag = exl3_route_expert( + topk_ids[r], expert_map, expert_map_size, num_buckets + ) == expert; + flags[threadIdx.x] = flag; + __syncthreads(); + for (int stride = 1; stride < blockDim.x; stride <<= 1) + { + int v = threadIdx.x >= stride ? flags[threadIdx.x - stride] : 0; + __syncthreads(); + flags[threadIdx.x] += v; + __syncthreads(); + } + if (flag) + { + int64_t dst = expert_offsets[expert] + running + + flags[threadIdx.x] - 1; + token_sorted[dst] = r / topk; + weight_sorted[dst] = exl3_route_to_half(topk_weights[r]); + } + __syncthreads(); + if (threadIdx.x == 0) running += flags[blockDim.x - 1]; + __syncthreads(); + } +} + int exl3_moe_max_concurrency(int device) { int num_sms = DevCtx::instance().get_num_sms(device); @@ -33,6 +165,39 @@ fp_exl3_moe_kernel exl3_moe_kernel_instances[] = exl3_moe_kernel_k8_n128_cb1(), exl3_moe_kernel_k8_n256_cb1(), exl3_moe_kernel_k8_n128_cb2(), exl3_moe_kernel_k8_n256_cb2() }; +static void check_cuda_contiguous_same_device +( + const at::Tensor& tensor, + const at::Tensor& reference, + const char* name +) +{ + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK( + tensor.device() == reference.device(), + name, " must be on the same CUDA device as hidden_state" + ); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +static void check_non_overlapping +( + const at::Tensor& first, + const at::Tensor& second, + const char* first_name, + const char* second_name +) +{ + const auto first_begin = reinterpret_cast(first.data_ptr()); + const auto second_begin = reinterpret_cast(second.data_ptr()); + const auto first_end = first_begin + first.nbytes(); + const auto second_end = second_begin + second.nbytes(); + TORCH_CHECK( + first_end <= second_begin || second_end <= first_begin, + first_name, " and ", second_name, " must not overlap" + ); +} + /* Fused mixture-of-experts MLP operation for EXL3 weights @@ -91,12 +256,14 @@ inputs: bool, codebook flags num_active: - number of experts with 0 < token count <= max_tokens_per_expert, i.e. the number of experts this kernel - will process. Used to size the launch: fewer, wider expert groups when few experts are active. Pass -1 if - unknown (defaults to MOE_SMS_PER_EXPERT-wide groups at max concurrency) + launch-size hint. For exl3_moe this is the number of experts with + 0 < token count <= max_tokens_per_expert. Additive entry points tile + oversized route spans, so they count every nonempty expert. Pass -1 + when unknown. For additive calls, 0 with nonempty routes is treated as + unknown rather than dropping work. */ -void exl3_moe +static void exl3_moe_impl ( const at::Tensor& hidden_state, const at::Tensor& output_state, @@ -133,33 +300,92 @@ void exl3_moe const bool down_mul1, const float act_limit, - const int num_active + const int num_active, + const at::Tensor& residual_gate_ptrs_trellis, + const at::Tensor& residual_up_ptrs_trellis, + const at::Tensor& residual_down_ptrs_trellis, + const at::Tensor& residual_gate_scales, + const at::Tensor& residual_up_scales, + const at::Tensor& residual_down_scales, + const at::Tensor& residual_gate_k, + const at::Tensor& residual_up_k, + const at::Tensor& residual_down_k, + const int max_residual_bits, + const bool tile_overflow, + const bool validate_only ) { - const at::cuda::OptionalCUDAGuard device_guard(hidden_state.device()); - cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); - - // Nothing for the fused kernel to do - if (num_active == 0) return; + TORCH_CHECK(hidden_state.is_cuda(), "hidden_state must be a CUDA tensor"); + const bool residual_gate_defined = residual_gate_ptrs_trellis.defined(); + const bool any_residual_defined = + residual_gate_defined || + residual_up_ptrs_trellis.defined() || + residual_down_ptrs_trellis.defined() || + residual_gate_scales.defined() || residual_up_scales.defined() || + residual_down_scales.defined() || residual_gate_k.defined() || + residual_up_k.defined() || residual_down_k.defined(); + const bool all_residual_defined = + residual_gate_defined && + residual_up_ptrs_trellis.defined() && + residual_down_ptrs_trellis.defined() && + residual_gate_scales.defined() && residual_up_scales.defined() && + residual_down_scales.defined() && residual_gate_k.defined() && + residual_up_k.defined() && residual_down_k.defined(); + TORCH_CHECK( + any_residual_defined == all_residual_defined, + "Residual pointer, scale, and K tensors must be all defined or all omitted" + ); + const int num_residual_stages = residual_gate_ptrs_trellis.defined() + ? residual_gate_ptrs_trellis.size(0) + : 0; + TORCH_CHECK( + (num_residual_stages == 0 && max_residual_bits == 0) || + (num_residual_stages > 0 && + max_residual_bits >= 1 && max_residual_bits <= 8), + "max_residual_bits must be zero without residuals or 1..8 with residuals" + ); - // Validate args + // Validate every tensor before taking a raw data_ptr. This is especially + // important for the fused entry point, which mutates routing workspaces + // before launching the MoE kernel. TORCH_CHECK_DTYPE(hidden_state, kHalf); TORCH_CHECK_DIM(hidden_state, 2); + TORCH_CHECK(hidden_state.is_contiguous(), "hidden_state must be contiguous"); size_t bsz = hidden_state.size(0); size_t hidden_dim = hidden_state.size(1); + check_cuda_contiguous_same_device(output_state, hidden_state, "output_state"); TORCH_CHECK_DTYPE(output_state, kFloat); TORCH_CHECK_SHAPES_FULL(output_state, hidden_state); + check_cuda_contiguous_same_device(expert_count, hidden_state, "expert_count"); TORCH_CHECK_DTYPE(expert_count, kLong); TORCH_CHECK_DIM(expert_count, 1); + TORCH_CHECK( + expert_count.size(0) >= 2, + "expert_count must contain at least one expert and one sentinel bucket" + ); size_t num_experts = expert_count.size(0) - 1; + check_cuda_contiguous_same_device(token_sorted, hidden_state, "token_sorted"); + check_cuda_contiguous_same_device(weight_sorted, hidden_state, "weight_sorted"); TORCH_CHECK_DTYPE(token_sorted, kLong); + TORCH_CHECK_DTYPE(weight_sorted, kHalf); TORCH_CHECK_DIM(token_sorted, 1); + TORCH_CHECK_DIM(weight_sorted, 1); TORCH_CHECK_SHAPES_FULL(token_sorted, weight_sorted); - size_t num_experts_per_tok = token_sorted.size(0) / bsz; + TORCH_CHECK( + bsz > 0 || token_sorted.numel() == 0, + "token_sorted must be empty when hidden_state has no rows" + ); + TORCH_CHECK( + bsz == 0 || token_sorted.size(0) % bsz == 0, + "token_sorted length must be divisible by hidden-state rows" + ); + size_t num_experts_per_tok = bsz ? token_sorted.size(0) / bsz : 0; + check_cuda_contiguous_same_device(temp_state_g, hidden_state, "temp_state_g"); + check_cuda_contiguous_same_device(temp_state_u, hidden_state, "temp_state_u"); TORCH_CHECK_DTYPE(temp_state_g, kHalf); TORCH_CHECK_DTYPE(temp_state_u, kHalf); TORCH_CHECK_DIM(temp_state_g, 3); @@ -167,7 +393,15 @@ void exl3_moe TORCH_CHECK_SHAPES_FULL(temp_state_g, temp_state_u); size_t max_tokens_per_expert = temp_state_g.size(1); size_t concurrency = temp_state_g.size(0); + TORCH_CHECK(max_tokens_per_expert > 0, "MoE temp token capacity must be positive"); + TORCH_CHECK(concurrency > 0, "MoE temp concurrency must be positive"); + check_cuda_contiguous_same_device( + temp_intermediate_g, hidden_state, "temp_intermediate_g" + ); + check_cuda_contiguous_same_device( + temp_intermediate_u, hidden_state, "temp_intermediate_u" + ); TORCH_CHECK_DTYPE(temp_intermediate_g, kHalf); TORCH_CHECK_DTYPE(temp_intermediate_u, kHalf); TORCH_CHECK_DIM(temp_intermediate_g, 3); @@ -175,6 +409,24 @@ void exl3_moe TORCH_CHECK_SHAPES_FULL(temp_intermediate_g, temp_intermediate_u); TORCH_CHECK_SHAPES(temp_intermediate_g, 1, temp_state_g, 1, 1); size_t intermediate_dim = temp_intermediate_g.size(2); + TORCH_CHECK( + hidden_dim % 128 == 0 && intermediate_dim % 128 == 0, + "MoE hidden and intermediate dimensions must be multiples of 128" + ); + + TORCH_CHECK( + K_gate >= 1 && K_gate <= 8 && + K_up >= 1 && K_up <= 8 && + K_down >= 1 && K_down <= 8, + "MoE gate/up/down bitrates must be in 1..8" + ); + TORCH_CHECK( + num_active >= -1 && num_active <= static_cast(num_experts), + "num_active must be -1 or in 0..num_experts" + ); + const int effective_num_active = + tile_overflow && num_active == 0 && token_sorted.numel() > 0 + ? -1 : num_active; // TORCH_CHECK(!(gate_mcg && gate_mul1), "Specified both mcg and mul1 (gate)"); // TORCH_CHECK(!(up_mcg && up_mul1), "Specified both mcg and mul1 (up)"); @@ -188,7 +440,35 @@ void exl3_moe int K = 0; if (K_gate == K_up && K_up == K_down) K = K_gate; + // Residual dispatch lives only in the runtime-K (K=0) instances. This + // keeps the common equal-K legacy kernels free of the additional 1..8 + // residual GEMM specializations and their instruction footprint. + if (num_residual_stages > 0) K = 0; + check_cuda_contiguous_same_device( + gate_ptrs_trellis, hidden_state, "gate_ptrs_trellis" + ); + check_cuda_contiguous_same_device(gate_ptrs_suh, hidden_state, "gate_ptrs_suh"); + check_cuda_contiguous_same_device(gate_ptrs_svh, hidden_state, "gate_ptrs_svh"); + check_cuda_contiguous_same_device( + up_ptrs_trellis, hidden_state, "up_ptrs_trellis" + ); + check_cuda_contiguous_same_device(up_ptrs_suh, hidden_state, "up_ptrs_suh"); + check_cuda_contiguous_same_device(up_ptrs_svh, hidden_state, "up_ptrs_svh"); + check_cuda_contiguous_same_device( + down_ptrs_trellis, hidden_state, "down_ptrs_trellis" + ); + check_cuda_contiguous_same_device(down_ptrs_suh, hidden_state, "down_ptrs_suh"); + check_cuda_contiguous_same_device(down_ptrs_svh, hidden_state, "down_ptrs_svh"); + TORCH_CHECK_DTYPE(gate_ptrs_trellis, kLong); + TORCH_CHECK_DTYPE(gate_ptrs_suh, kLong); + TORCH_CHECK_DTYPE(gate_ptrs_svh, kLong); + TORCH_CHECK_DTYPE(up_ptrs_trellis, kLong); + TORCH_CHECK_DTYPE(up_ptrs_suh, kLong); + TORCH_CHECK_DTYPE(up_ptrs_svh, kLong); + TORCH_CHECK_DTYPE(down_ptrs_trellis, kLong); + TORCH_CHECK_DTYPE(down_ptrs_suh, kLong); + TORCH_CHECK_DTYPE(down_ptrs_svh, kLong); TORCH_CHECK_DIM(gate_ptrs_trellis, 1); TORCH_CHECK(gate_ptrs_trellis.size(0) == num_experts, "Number of gate tensors doesn't match num_experts"); TORCH_CHECK_SHAPES_FULL(gate_ptrs_trellis, gate_ptrs_suh); @@ -199,6 +479,84 @@ void exl3_moe TORCH_CHECK_SHAPES_FULL(gate_ptrs_trellis, down_ptrs_trellis); TORCH_CHECK_SHAPES_FULL(gate_ptrs_trellis, down_ptrs_suh); TORCH_CHECK_SHAPES_FULL(gate_ptrs_trellis, down_ptrs_svh); + if (num_residual_stages > 0) + { + // Additive residual trellises use the MCG codebook and reuse the base + // projection's suh/svh. K metadata is graph-resident and must contain + // values in 1..max_residual_bits; callers construct and validate it + // before graph capture. + check_cuda_contiguous_same_device( + residual_gate_ptrs_trellis, hidden_state, + "residual_gate_ptrs_trellis" + ); + check_cuda_contiguous_same_device( + residual_up_ptrs_trellis, hidden_state, + "residual_up_ptrs_trellis" + ); + check_cuda_contiguous_same_device( + residual_down_ptrs_trellis, hidden_state, + "residual_down_ptrs_trellis" + ); + check_cuda_contiguous_same_device( + residual_gate_scales, hidden_state, "residual_gate_scales" + ); + check_cuda_contiguous_same_device( + residual_up_scales, hidden_state, "residual_up_scales" + ); + check_cuda_contiguous_same_device( + residual_down_scales, hidden_state, "residual_down_scales" + ); + check_cuda_contiguous_same_device( + residual_gate_k, hidden_state, "residual_gate_k" + ); + check_cuda_contiguous_same_device( + residual_up_k, hidden_state, "residual_up_k" + ); + check_cuda_contiguous_same_device( + residual_down_k, hidden_state, "residual_down_k" + ); + TORCH_CHECK_DTYPE(residual_gate_ptrs_trellis, kLong); + TORCH_CHECK_DTYPE(residual_up_ptrs_trellis, kLong); + TORCH_CHECK_DTYPE(residual_down_ptrs_trellis, kLong); + TORCH_CHECK_DIM(residual_gate_ptrs_trellis, 2); + TORCH_CHECK_SHAPES_FULL( + residual_gate_ptrs_trellis, residual_up_ptrs_trellis + ); + TORCH_CHECK_SHAPES_FULL( + residual_gate_ptrs_trellis, residual_down_ptrs_trellis + ); + TORCH_CHECK( + residual_gate_ptrs_trellis.size(1) == num_experts, + "Residual pointer tables must have shape (num_stages, num_experts)" + ); + TORCH_CHECK_DTYPE(residual_gate_scales, kFloat); + TORCH_CHECK_DTYPE(residual_up_scales, kFloat); + TORCH_CHECK_DTYPE(residual_down_scales, kFloat); + TORCH_CHECK_SHAPES_FULL( + residual_gate_ptrs_trellis, residual_gate_scales + ); + TORCH_CHECK_SHAPES_FULL( + residual_gate_ptrs_trellis, residual_up_scales + ); + TORCH_CHECK_SHAPES_FULL( + residual_gate_ptrs_trellis, residual_down_scales + ); + TORCH_CHECK_DTYPE(residual_gate_k, kInt); + TORCH_CHECK_DTYPE(residual_up_k, kInt); + TORCH_CHECK_DTYPE(residual_down_k, kInt); + TORCH_CHECK_DIM(residual_gate_k, 1); + TORCH_CHECK_SHAPES_FULL(residual_gate_k, residual_up_k); + TORCH_CHECK_SHAPES_FULL(residual_gate_k, residual_down_k); + TORCH_CHECK( + residual_gate_k.size(0) == num_residual_stages, + "Residual K tensors must have shape (num_stages,)" + ); + } + + if (validate_only || effective_num_active == 0 || bsz == 0) return; + + const at::cuda::OptionalCUDAGuard device_guard(hidden_state.device()); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); // Device properties int device; @@ -214,9 +572,9 @@ void exl3_moe TORCH_CHECK(concurrency * MOE_SMS_PER_EXPERT <= num_sms, "Concurrency too high for device num_sms"); int num_groups = MIN((int) concurrency, MOE_MAX_GROUPS); int group_size = MOE_SMS_PER_EXPERT; - if (num_active > 0) + if (effective_num_active > 0) { - num_groups = MIN(num_groups, num_active); + num_groups = MIN(num_groups, effective_num_active); group_size = MIN(num_sms / num_groups, MOE_MAX_SMS_PER_EXPERT); } dim3 grid_dim(group_size, 1, num_groups); @@ -248,6 +606,33 @@ void exl3_moe void* _down_ptrs_trellis = down_ptrs_trellis.data_ptr(); void* _down_ptrs_suh = down_ptrs_suh.data_ptr(); void* _down_ptrs_svh = down_ptrs_svh.data_ptr(); + void* _residual_gate_ptrs_trellis = num_residual_stages + ? residual_gate_ptrs_trellis.data_ptr() + : nullptr; + void* _residual_up_ptrs_trellis = num_residual_stages + ? residual_up_ptrs_trellis.data_ptr() + : nullptr; + void* _residual_down_ptrs_trellis = num_residual_stages + ? residual_down_ptrs_trellis.data_ptr() + : nullptr; + void* _residual_gate_scales = num_residual_stages + ? residual_gate_scales.data_ptr() + : nullptr; + void* _residual_up_scales = num_residual_stages + ? residual_up_scales.data_ptr() + : nullptr; + void* _residual_down_scales = num_residual_stages + ? residual_down_scales.data_ptr() + : nullptr; + void* _residual_gate_k = num_residual_stages + ? residual_gate_k.data_ptr() + : nullptr; + void* _residual_up_k = num_residual_stages + ? residual_up_k.data_ptr() + : nullptr; + void* _residual_down_k = num_residual_stages + ? residual_down_k.data_ptr() + : nullptr; void* _expert_count = expert_count.data_ptr(); void* _token_sorted = token_sorted.data_ptr(); @@ -270,6 +655,16 @@ void exl3_moe &_down_ptrs_trellis, &_down_ptrs_suh, &_down_ptrs_svh, + &_residual_gate_ptrs_trellis, + &_residual_up_ptrs_trellis, + &_residual_down_ptrs_trellis, + &_residual_gate_scales, + &_residual_up_scales, + &_residual_down_scales, + &_residual_gate_k, + &_residual_up_k, + &_residual_down_k, + (void*) &num_residual_stages, &_expert_count, &_token_sorted, &_weight_sorted, @@ -284,6 +679,7 @@ void exl3_moe (void*) &K_gate, (void*) &K_up, (void*) &K_down, + (void*) &tile_overflow, (void*) &locks }; @@ -299,3 +695,364 @@ void exl3_moe cuda_check(cudaPeekAtLastError()); } + +void exl3_moe +( + const at::Tensor& hidden_state, + const at::Tensor& output_state, + const at::Tensor& expert_count, + const at::Tensor& token_sorted, + const at::Tensor& weight_sorted, + const at::Tensor& temp_state_g, + const at::Tensor& temp_state_u, + const at::Tensor& temp_intermediate_g, + const at::Tensor& temp_intermediate_u, + const int act_function, + const int K_gate, + const int K_up, + const int K_down, + const at::Tensor& gate_ptrs_trellis, + const at::Tensor& gate_ptrs_suh, + const at::Tensor& gate_ptrs_svh, + const at::Tensor& up_ptrs_trellis, + const at::Tensor& up_ptrs_suh, + const at::Tensor& up_ptrs_svh, + const at::Tensor& down_ptrs_trellis, + const at::Tensor& down_ptrs_suh, + const at::Tensor& down_ptrs_svh, + const bool gate_mcg, + const bool gate_mul1, + const bool up_mcg, + const bool up_mul1, + const bool down_mcg, + const bool down_mul1, + const float act_limit, + const int num_active +) +{ + const at::Tensor empty; + exl3_moe_impl + ( + hidden_state, output_state, expert_count, token_sorted, weight_sorted, + temp_state_g, temp_state_u, temp_intermediate_g, temp_intermediate_u, + act_function, K_gate, K_up, K_down, + gate_ptrs_trellis, gate_ptrs_suh, gate_ptrs_svh, + up_ptrs_trellis, up_ptrs_suh, up_ptrs_svh, + down_ptrs_trellis, down_ptrs_suh, down_ptrs_svh, + gate_mcg, gate_mul1, up_mcg, up_mul1, down_mcg, down_mul1, + act_limit, num_active, + empty, empty, empty, empty, empty, empty, empty, empty, empty, 0, + false, false + ); +} + +void exl3_moe_additive +( + const at::Tensor& hidden_state, + const at::Tensor& output_state, + const at::Tensor& expert_count, + const at::Tensor& token_sorted, + const at::Tensor& weight_sorted, + const at::Tensor& temp_state_g, + const at::Tensor& temp_state_u, + const at::Tensor& temp_intermediate_g, + const at::Tensor& temp_intermediate_u, + const int act_function, + const int K_gate, + const int K_up, + const int K_down, + const at::Tensor& gate_ptrs_trellis, + const at::Tensor& gate_ptrs_suh, + const at::Tensor& gate_ptrs_svh, + const at::Tensor& up_ptrs_trellis, + const at::Tensor& up_ptrs_suh, + const at::Tensor& up_ptrs_svh, + const at::Tensor& down_ptrs_trellis, + const at::Tensor& down_ptrs_suh, + const at::Tensor& down_ptrs_svh, + const at::Tensor& residual_gate_ptrs_trellis, + const at::Tensor& residual_up_ptrs_trellis, + const at::Tensor& residual_down_ptrs_trellis, + const at::Tensor& residual_gate_scales, + const at::Tensor& residual_up_scales, + const at::Tensor& residual_down_scales, + const at::Tensor& residual_gate_k, + const at::Tensor& residual_up_k, + const at::Tensor& residual_down_k, + const int max_residual_bits, + const bool gate_mcg, + const bool gate_mul1, + const bool up_mcg, + const bool up_mul1, + const bool down_mcg, + const bool down_mul1, + const float act_limit, + const int num_active +) +{ + exl3_moe_impl + ( + hidden_state, output_state, expert_count, token_sorted, weight_sorted, + temp_state_g, temp_state_u, temp_intermediate_g, temp_intermediate_u, + act_function, K_gate, K_up, K_down, + gate_ptrs_trellis, gate_ptrs_suh, gate_ptrs_svh, + up_ptrs_trellis, up_ptrs_suh, up_ptrs_svh, + down_ptrs_trellis, down_ptrs_suh, down_ptrs_svh, + gate_mcg, gate_mul1, up_mcg, up_mul1, down_mcg, down_mul1, + act_limit, num_active, + residual_gate_ptrs_trellis, + residual_up_ptrs_trellis, + residual_down_ptrs_trellis, + residual_gate_scales, + residual_up_scales, + residual_down_scales, + residual_gate_k, + residual_up_k, + residual_down_k, + max_residual_bits, + true, false + ); +} + +void exl3_moe_additive_fused +( + const at::Tensor& hidden_state, + const at::Tensor& output_state, + const at::Tensor& topk_ids, + const at::Tensor& topk_weights, + const at::Tensor& expert_map, + const at::Tensor& expert_count, + const at::Tensor& expert_offsets, + const at::Tensor& token_sorted, + const at::Tensor& weight_sorted, + const at::Tensor& temp_state_g, + const at::Tensor& temp_state_u, + const at::Tensor& temp_intermediate_g, + const at::Tensor& temp_intermediate_u, + const int act_function, + const int K_gate, + const int K_up, + const int K_down, + const at::Tensor& gate_ptrs_trellis, + const at::Tensor& gate_ptrs_suh, + const at::Tensor& gate_ptrs_svh, + const at::Tensor& up_ptrs_trellis, + const at::Tensor& up_ptrs_suh, + const at::Tensor& up_ptrs_svh, + const at::Tensor& down_ptrs_trellis, + const at::Tensor& down_ptrs_suh, + const at::Tensor& down_ptrs_svh, + const at::Tensor& residual_gate_ptrs_trellis, + const at::Tensor& residual_up_ptrs_trellis, + const at::Tensor& residual_down_ptrs_trellis, + const at::Tensor& residual_gate_scales, + const at::Tensor& residual_up_scales, + const at::Tensor& residual_down_scales, + const at::Tensor& residual_gate_k, + const at::Tensor& residual_up_k, + const at::Tensor& residual_down_k, + const int max_residual_bits, + const bool gate_mcg, + const bool gate_mul1, + const bool up_mcg, + const bool up_mul1, + const bool down_mcg, + const bool down_mul1, + const float act_limit, + const int num_active +) +{ + TORCH_CHECK(hidden_state.is_cuda(), "hidden_state must be a CUDA tensor"); + TORCH_CHECK(hidden_state.is_contiguous(), "hidden_state must be contiguous"); + TORCH_CHECK_DIM(hidden_state, 2); + check_cuda_contiguous_same_device(topk_ids, hidden_state, "topk_ids"); + check_cuda_contiguous_same_device(topk_weights, hidden_state, "topk_weights"); + check_cuda_contiguous_same_device(expert_map, hidden_state, "expert_map"); + check_cuda_contiguous_same_device(expert_count, hidden_state, "expert_count"); + check_cuda_contiguous_same_device(expert_offsets, hidden_state, "expert_offsets"); + check_cuda_contiguous_same_device(token_sorted, hidden_state, "token_sorted"); + check_cuda_contiguous_same_device(weight_sorted, hidden_state, "weight_sorted"); + TORCH_CHECK( + topk_ids.scalar_type() == at::kLong || + topk_ids.scalar_type() == at::kInt, + "topk_ids must be int32 or int64" + ); + TORCH_CHECK( + topk_weights.scalar_type() == at::kFloat || + topk_weights.scalar_type() == at::kHalf || + topk_weights.scalar_type() == at::kBFloat16, + "topk_weights must be float, half, or bfloat16" + ); + TORCH_CHECK_DIM(topk_ids, 2); + TORCH_CHECK_DIM(topk_weights, 2); + TORCH_CHECK_SHAPES_FULL(topk_ids, topk_weights); + TORCH_CHECK_DTYPE(expert_map, kLong); + TORCH_CHECK_DIM(expert_map, 1); + TORCH_CHECK_DTYPE(expert_count, kLong); + TORCH_CHECK_DTYPE(expert_offsets, kLong); + TORCH_CHECK_DIM(expert_count, 1); + TORCH_CHECK_DIM(expert_offsets, 1); + TORCH_CHECK_SHAPES_FULL(expert_count, expert_offsets); + check_non_overlapping( + expert_count, expert_offsets, "expert_count", "expert_offsets" + ); + TORCH_CHECK_DTYPE(token_sorted, kLong); + TORCH_CHECK_DTYPE(weight_sorted, kHalf); + TORCH_CHECK_DIM(token_sorted, 1); + TORCH_CHECK_DIM(weight_sorted, 1); + TORCH_CHECK_SHAPES_FULL(token_sorted, weight_sorted); + TORCH_CHECK( + hidden_state.size(0) == topk_ids.size(0), + "route rows must equal hidden-state rows" + ); + const int64_t num_routes = topk_ids.numel(); + const int num_buckets = expert_count.numel(); + TORCH_CHECK( + num_buckets == gate_ptrs_trellis.numel() + 1, + "expert_count must include one sentinel bucket" + ); + TORCH_CHECK( + token_sorted.numel() >= num_routes, + "route workspace is too small" + ); + + const at::Tensor routed_tokens = token_sorted.narrow(0, 0, num_routes); + const at::Tensor routed_weights = weight_sorted.narrow(0, 0, num_routes); + const int launch_num_active = num_routes == 0 ? 0 : num_active; + exl3_moe_impl + ( + hidden_state, output_state, expert_count, routed_tokens, routed_weights, + temp_state_g, temp_state_u, temp_intermediate_g, temp_intermediate_u, + act_function, K_gate, K_up, K_down, + gate_ptrs_trellis, gate_ptrs_suh, gate_ptrs_svh, + up_ptrs_trellis, up_ptrs_suh, up_ptrs_svh, + down_ptrs_trellis, down_ptrs_suh, down_ptrs_svh, + gate_mcg, gate_mul1, up_mcg, up_mul1, down_mcg, down_mul1, + act_limit, launch_num_active, + residual_gate_ptrs_trellis, + residual_up_ptrs_trellis, + residual_down_ptrs_trellis, + residual_gate_scales, + residual_up_scales, + residual_down_scales, + residual_gate_k, + residual_up_k, + residual_down_k, + max_residual_bits, + true, true + ); + if (num_routes == 0) return; + + const at::cuda::OptionalCUDAGuard device_guard(hidden_state.device()); + cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream(); + + cuda_check(cudaMemsetAsync( + expert_count.data_ptr(), 0, + expert_count.numel() * expert_count.element_size(), stream + )); + const int threads = 256; + const int blocks = std::min( + 1024, (num_routes + threads - 1) / threads + ); + #define LAUNCH_HIST(ID_T, PTR) \ + exl3_route_histogram_kernel<<>>( \ + PTR, expert_map.data_ptr(), \ + expert_count.data_ptr(), num_routes, \ + expert_map.numel(), num_buckets) + if (topk_ids.scalar_type() == at::kInt) + LAUNCH_HIST(int32_t, topk_ids.data_ptr()); + else + LAUNCH_HIST(int64_t, topk_ids.data_ptr()); + #undef LAUNCH_HIST + + exl3_route_scan_kernel<<<1, 1, 0, stream>>> + ( + expert_count.data_ptr(), + expert_offsets.data_ptr(), + num_buckets + ); + + #define LAUNCH_PACK(ID_T, ID_PTR, W_T, W_PTR) \ + exl3_route_pack_stable_kernel \ + <<>>( \ + ID_PTR, W_PTR, expert_map.data_ptr(), \ + expert_offsets.data_ptr(), \ + token_sorted.data_ptr(), \ + reinterpret_cast(weight_sorted.data_ptr()), \ + num_routes, topk_ids.size(1), expert_map.numel(), num_buckets) + #define DISPATCH_WEIGHT(ID_T, ID_PTR) \ + if (topk_weights.scalar_type() == at::kFloat) \ + LAUNCH_PACK( \ + ID_T, ID_PTR, float, topk_weights.data_ptr() \ + ); \ + else if (topk_weights.scalar_type() == at::kHalf) \ + LAUNCH_PACK( \ + ID_T, ID_PTR, half, \ + reinterpret_cast(topk_weights.data_ptr()) \ + ); \ + else if (topk_weights.scalar_type() == at::kBFloat16) \ + LAUNCH_PACK( \ + ID_T, ID_PTR, __nv_bfloat16, \ + reinterpret_cast( \ + topk_weights.data_ptr() \ + ) \ + ); \ + else TORCH_CHECK( \ + false, "topk_weights must be float, half, or bfloat16" \ + ) + if (topk_ids.scalar_type() == at::kInt) + { + DISPATCH_WEIGHT(int32_t, topk_ids.data_ptr()); + } + else + { + DISPATCH_WEIGHT(int64_t, topk_ids.data_ptr()); + } + #undef DISPATCH_WEIGHT + #undef LAUNCH_PACK + cuda_check(cudaPeekAtLastError()); + + exl3_moe_additive + ( + hidden_state, + output_state, + expert_count, + routed_tokens, + routed_weights, + temp_state_g, + temp_state_u, + temp_intermediate_g, + temp_intermediate_u, + act_function, + K_gate, + K_up, + K_down, + gate_ptrs_trellis, + gate_ptrs_suh, + gate_ptrs_svh, + up_ptrs_trellis, + up_ptrs_suh, + up_ptrs_svh, + down_ptrs_trellis, + down_ptrs_suh, + down_ptrs_svh, + residual_gate_ptrs_trellis, + residual_up_ptrs_trellis, + residual_down_ptrs_trellis, + residual_gate_scales, + residual_up_scales, + residual_down_scales, + residual_gate_k, + residual_up_k, + residual_down_k, + max_residual_bits, + gate_mcg, + gate_mul1, + up_mcg, + up_mul1, + down_mcg, + down_mul1, + act_limit, + launch_num_active + ); +} diff --git a/exllamav3/exllamav3_ext/quant/exl3_moe.cuh b/exllamav3/exllamav3_ext/quant/exl3_moe.cuh index 24e227c7..25723c6f 100644 --- a/exllamav3/exllamav3_ext/quant/exl3_moe.cuh +++ b/exllamav3/exllamav3_ext/quant/exl3_moe.cuh @@ -3,6 +3,14 @@ #include #include "../graph.cuh" +// Exported by the Python extension as an integer module attribute with the +// same name. Consumers must require an exact supported value instead of +// inferring compatibility from symbol presence or callable arity. Increment +// whenever the additive entry points' signatures, pointer-table layout, +// routing workspace semantics, residual encoding, or overflow behavior change +// incompatibly. +constexpr int EXL3_MOE_ADDITIVE_ABI_VERSION = 1; + int exl3_moe_max_concurrency(int device); void exl3_moe @@ -45,3 +53,109 @@ void exl3_moe const int num_active ); +// Additive residual contract: +// - all tensor arguments are contiguous CUDA tensors on hidden_state.device(); +// fused routing workspaces must be disjoint (expert_count and expert_offsets +// are checked explicitly); +// - residual trellises use the MCG codebook and reuse the corresponding base +// projection's suh/svh vectors; +// - residual K metadata is int32, one value in 1..max_residual_bits (<= 8) per +// stage, and must be validated before graph capture; +// - a zero residual scale denotes a sparse/missing projection and skips its +// GEMM; +// - int64 pointer tables do not retain their pointees. Callers must keep every +// base/residual trellis and base suh/svh allocation alive and unmoved on the +// same device through asynchronous completion and the lifetime of any graph; +// - additive kernels tile oversized expert route spans internally, so +// num_active counts all nonempty experts and is only a launch-size hint. +void exl3_moe_additive +( + const at::Tensor& hidden_state, + const at::Tensor& output_state, + const at::Tensor& expert_count, + const at::Tensor& token_sorted, + const at::Tensor& weight_sorted, + const at::Tensor& temp_state_g, + const at::Tensor& temp_state_u, + const at::Tensor& temp_intermediate_g, + const at::Tensor& temp_intermediate_u, + const int act_function, + const int K_gate, + const int K_up, + const int K_down, + const at::Tensor& gate_ptrs_trellis, + const at::Tensor& gate_ptrs_suh, + const at::Tensor& gate_ptrs_svh, + const at::Tensor& up_ptrs_trellis, + const at::Tensor& up_ptrs_suh, + const at::Tensor& up_ptrs_svh, + const at::Tensor& down_ptrs_trellis, + const at::Tensor& down_ptrs_suh, + const at::Tensor& down_ptrs_svh, + const at::Tensor& residual_gate_ptrs_trellis, + const at::Tensor& residual_up_ptrs_trellis, + const at::Tensor& residual_down_ptrs_trellis, + const at::Tensor& residual_gate_scales, + const at::Tensor& residual_up_scales, + const at::Tensor& residual_down_scales, + const at::Tensor& residual_gate_k, + const at::Tensor& residual_up_k, + const at::Tensor& residual_down_k, + const int max_residual_bits, + const bool gate_mcg, + const bool gate_mul1, + const bool up_mcg, + const bool up_mul1, + const bool down_mcg, + const bool down_mul1, + const float act_limit, + const int num_active +); + +void exl3_moe_additive_fused +( + const at::Tensor& hidden_state, + const at::Tensor& output_state, + const at::Tensor& topk_ids, + const at::Tensor& topk_weights, + const at::Tensor& expert_map, + const at::Tensor& expert_count, + const at::Tensor& expert_offsets, + const at::Tensor& token_sorted, + const at::Tensor& weight_sorted, + const at::Tensor& temp_state_g, + const at::Tensor& temp_state_u, + const at::Tensor& temp_intermediate_g, + const at::Tensor& temp_intermediate_u, + const int act_function, + const int K_gate, + const int K_up, + const int K_down, + const at::Tensor& gate_ptrs_trellis, + const at::Tensor& gate_ptrs_suh, + const at::Tensor& gate_ptrs_svh, + const at::Tensor& up_ptrs_trellis, + const at::Tensor& up_ptrs_suh, + const at::Tensor& up_ptrs_svh, + const at::Tensor& down_ptrs_trellis, + const at::Tensor& down_ptrs_suh, + const at::Tensor& down_ptrs_svh, + const at::Tensor& residual_gate_ptrs_trellis, + const at::Tensor& residual_up_ptrs_trellis, + const at::Tensor& residual_down_ptrs_trellis, + const at::Tensor& residual_gate_scales, + const at::Tensor& residual_up_scales, + const at::Tensor& residual_down_scales, + const at::Tensor& residual_gate_k, + const at::Tensor& residual_up_k, + const at::Tensor& residual_down_k, + const int max_residual_bits, + const bool gate_mcg, + const bool gate_mul1, + const bool up_mcg, + const bool up_mul1, + const bool down_mcg, + const bool down_mul1, + const float act_limit, + const int num_active +); diff --git a/exllamav3/exllamav3_ext/quant/exl3_moe_common.cuh b/exllamav3/exllamav3_ext/quant/exl3_moe_common.cuh index 2109b270..fff78c4e 100644 --- a/exllamav3/exllamav3_ext/quant/exl3_moe_common.cuh +++ b/exllamav3/exllamav3_ext/quant/exl3_moe_common.cuh @@ -40,6 +40,17 @@ const half** __restrict__ down_suh, \ const half** __restrict__ down_svh, \ \ + const uint16_t** __restrict__ residual_gate_trellis, \ + const uint16_t** __restrict__ residual_up_trellis, \ + const uint16_t** __restrict__ residual_down_trellis, \ + const float* __restrict__ residual_gate_scales, \ + const float* __restrict__ residual_up_scales, \ + const float* __restrict__ residual_down_scales, \ + const int* __restrict__ residual_gate_k, \ + const int* __restrict__ residual_up_k, \ + const int* __restrict__ residual_down_k, \ + const int num_residual_stages, \ + \ const int64_t* __restrict__ expert_count, \ const int64_t* __restrict__ token_sorted, \ const half* __restrict__ weight_sorted, \ @@ -55,5 +66,6 @@ const int K_gate, \ const int K_up, \ const int K_down, \ + const bool tile_overflow, \ \ int* __restrict__ locks diff --git a/exllamav3/exllamav3_ext/quant/exl3_moe_kernel.cuh b/exllamav3/exllamav3_ext/quant/exl3_moe_kernel.cuh index fbe0a82e..55b1089a 100644 --- a/exllamav3/exllamav3_ext/quant/exl3_moe_kernel.cuh +++ b/exllamav3/exllamav3_ext/quant/exl3_moe_kernel.cuh @@ -61,9 +61,11 @@ void exl3_moe_kernel(EXL3_MOE_KERNEL_ARGS) end += expert_count[expert_idx]; int token_count = end - start; - // Skip if no tokens or too many tokens for fused kernel (batch is handled by reconstruct path outside kernel) + // The legacy entry point leaves oversized experts to its reconstruct + // fallback. Additive callers have no such fallback, so they process an + // expert's route span in bounded workspace-sized tiles below. if (token_count == 0) continue; - if (token_count > max_tokens_per_expert) continue; + if (!tile_overflow && token_count > max_tokens_per_expert) continue; // Skip if expert is claimed by a different group if (expert_idx_assign++ != ticket) continue; @@ -79,190 +81,299 @@ void exl3_moe_kernel(EXL3_MOE_KERNEL_ARGS) const half* exp_down_suh = down_suh[expert_idx]; const half* exp_down_svh = down_svh[expert_idx]; - // Gather + input hadamard for g, u. Non-gated mode skips the g staging (and the g GEMM - // below); the activation synthesizes the gate lane from u - const bool gated = act_function != MOE_ACT_RELU2_NOGATE; - auto had_gather_gu_in = [&]() + const int expert_end = end; + for (int chunk_start = start; chunk_start < expert_end; + chunk_start += max_tokens_per_expert) { - const int warps_per_token = hidden_dim / 128; - const int total_warps = token_count * warps_per_token; - const int64_t* top_x = token_sorted + start; - for (int warp_idx = warp_idx0; warp_idx < total_warps; warp_idx += warps_per_group) + start = chunk_start; + token_count = MIN(expert_end - chunk_start, max_tokens_per_expert); + + // Gather + input hadamard for g, u. Non-gated mode skips the g staging (and the g GEMM + // below); the activation synthesizes the gate lane from u + const bool gated = act_function != MOE_ACT_RELU2_NOGATE; + auto had_gather_gu_in = [&]() { - int token_idx = top_x[warp_idx / warps_per_token]; - int token_off = warp_idx % warps_per_token; - const half* in_ptr = hidden_state + token_idx * hidden_dim + token_off * 128; - if (gated) + const int warps_per_token = hidden_dim / 128; + const int total_warps = token_count * warps_per_token; + const int64_t* top_x = token_sorted + start; + for (int warp_idx = warp_idx0; warp_idx < total_warps; warp_idx += warps_per_group) + { + int token_idx = top_x[warp_idx / warps_per_token]; + int token_off = warp_idx % warps_per_token; + const half* in_ptr = hidden_state + token_idx * hidden_dim + token_off * 128; + if (gated) + had_hf_r_128_inner + ( + in_ptr, + temp_state_g + 128 * warp_idx, + exp_gate_suh + 128 * token_off, + 0.088388347648f + ); had_hf_r_128_inner ( in_ptr, - temp_state_g + 128 * warp_idx, - exp_gate_suh + 128 * token_off, + temp_state_u + 128 * warp_idx, + exp_up_suh + 128 * token_off, 0.088388347648f ); - had_hf_r_128_inner - ( - in_ptr, - temp_state_u + 128 * warp_idx, - exp_up_suh + 128 * token_off, - 0.088388347648f - ); - } - group_barrier(group_idx, group_size, barrier_counters_sense); - }; + } + group_barrier(group_idx, group_size, barrier_counters_sense); + }; - had_gather_gu_in(); + had_gather_gu_in(); - // g, u GEMM - auto gemm_up = [&](const half* in_addr, half* out_addr, const uint16_t* trellis, const int K) - { - int size_m = token_count; - while (size_m > 0) + // g, u GEMM + auto gemm_up = [&]( + const half* in_addr, + half* out_addr, + const uint16_t* trellis, + const int K, + const float output_scale, + const bool accumulate + ) { - #define ARGS \ - in_addr, \ - trellis, \ - out_addr, \ - MIN(size_m, 16), \ - hidden_dim, \ - intermediate_dim, \ - locks, \ - nullptr - #define SHAPE_ARGS \ - MOE_TILESIZE_M, \ - MOE_TILESIZE_K, \ - MOE_TILESIZE_N, \ - MOE_SH_STAGES, \ - MOE_FRAG_STAGES - if constexpr (t_bits) - exl3_gemm_kernel_inner(ARGS); - else switch(K) + int size_m = token_count; + while (size_m > 0) { - case 1: exl3_gemm_kernel_inner<1, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 2: exl3_gemm_kernel_inner<2, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 3: exl3_gemm_kernel_inner<3, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 4: exl3_gemm_kernel_inner<4, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 5: exl3_gemm_kernel_inner<5, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 6: exl3_gemm_kernel_inner<6, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 7: exl3_gemm_kernel_inner<7, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 8: exl3_gemm_kernel_inner<8, false, cb, SHAPE_ARGS, false>(ARGS); break; - }; - #undef ARGS - #undef SHAPE_ARGS - - in_addr += 16 * hidden_dim; - out_addr += 16 * intermediate_dim; - size_m -= 16; - } - }; + #define ARGS \ + in_addr, \ + trellis, \ + out_addr, \ + MIN(size_m, 16), \ + hidden_dim, \ + intermediate_dim, \ + locks, \ + nullptr, \ + output_scale, \ + accumulate + #define SHAPE_ARGS \ + MOE_TILESIZE_M, \ + MOE_TILESIZE_K, \ + MOE_TILESIZE_N, \ + MOE_SH_STAGES, \ + MOE_FRAG_STAGES + if constexpr (t_bits == 0) + { + if (accumulate) switch(K) + { + case 1: exl3_gemm_kernel_inner<1, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 2: exl3_gemm_kernel_inner<2, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 3: exl3_gemm_kernel_inner<3, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 4: exl3_gemm_kernel_inner<4, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 5: exl3_gemm_kernel_inner<5, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 6: exl3_gemm_kernel_inner<6, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 7: exl3_gemm_kernel_inner<7, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 8: exl3_gemm_kernel_inner<8, false, 1, SHAPE_ARGS, false>(ARGS); break; + default: __trap(); + } + else switch(K) + { + case 1: exl3_gemm_kernel_inner<1, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 2: exl3_gemm_kernel_inner<2, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 3: exl3_gemm_kernel_inner<3, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 4: exl3_gemm_kernel_inner<4, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 5: exl3_gemm_kernel_inner<5, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 6: exl3_gemm_kernel_inner<6, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 7: exl3_gemm_kernel_inner<7, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 8: exl3_gemm_kernel_inner<8, false, cb, SHAPE_ARGS, false>(ARGS); break; + default: __trap(); + } + } + else + exl3_gemm_kernel_inner(ARGS); + #undef ARGS + #undef SHAPE_ARGS - if (gated) - gemm_up(temp_state_g, temp_intermediate_g, exp_gate_trellis, K_gate); - gemm_up(temp_state_u, temp_intermediate_u, exp_up_trellis, K_up); - group_barrier(group_idx, group_size, barrier_counters_sense); + in_addr += 16 * hidden_dim; + out_addr += 16 * intermediate_dim; + size_m -= 16; + } + }; - // Output hadamard for g, u + activation+gate + input hadamard for d - auto had_guad = [&]() - { - const int warps_per_token = intermediate_dim / 128; - const int total_warps = token_count * warps_per_token; - for (int warp_idx = warp_idx0; warp_idx < total_warps; warp_idx += warps_per_group) - { - int token_off = warp_idx % warps_per_token; - had_hf_r_128_guad_inner - ( - temp_intermediate_g + 128 * warp_idx, - temp_intermediate_u + 128 * warp_idx, - temp_intermediate_g + 128 * warp_idx, - exp_gate_svh + 128 * token_off, - exp_up_svh + 128 * token_off, - exp_down_suh + 128 * token_off, - 0.088388347648f, - act_limit, - act_function + if (gated) + gemm_up( + temp_state_g, temp_intermediate_g, exp_gate_trellis, K_gate, + 1.0f, false ); - } + gemm_up( + temp_state_u, temp_intermediate_u, exp_up_trellis, K_up, 1.0f, false + ); group_barrier(group_idx, group_size, barrier_counters_sense); - }; + if constexpr (t_bits == 0) + for (int stage = 0; stage < num_residual_stages; ++stage) + { + const int residual_idx = stage * num_experts + expert_idx; + const float gate_scale = residual_gate_scales[residual_idx]; + const float up_scale = residual_up_scales[residual_idx]; + if (gated && gate_scale != 0.0f) + gemm_up( + temp_state_g, + temp_intermediate_g, + residual_gate_trellis[residual_idx], + residual_gate_k[stage], + gate_scale, + true + ); + if (up_scale != 0.0f) + gemm_up( + temp_state_u, + temp_intermediate_u, + residual_up_trellis[residual_idx], + residual_up_k[stage], + up_scale, + true + ); + group_barrier(group_idx, group_size, barrier_counters_sense); + } + + // Output hadamard for g, u + activation+gate + input hadamard for d + auto had_guad = [&]() + { + const int warps_per_token = intermediate_dim / 128; + const int total_warps = token_count * warps_per_token; + for (int warp_idx = warp_idx0; warp_idx < total_warps; warp_idx += warps_per_group) + { + int token_off = warp_idx % warps_per_token; + had_hf_r_128_guad_inner + ( + temp_intermediate_g + 128 * warp_idx, + temp_intermediate_u + 128 * warp_idx, + temp_intermediate_g + 128 * warp_idx, + exp_gate_svh + 128 * token_off, + exp_up_svh + 128 * token_off, + exp_down_suh + 128 * token_off, + 0.088388347648f, + act_limit, + act_function + ); + } + group_barrier(group_idx, group_size, barrier_counters_sense); + }; - had_guad(); + had_guad(); - // d GEMM - auto gemm_down = [&](const half* in_addr, half* out_addr, const uint16_t* trellis, const int K) - { - int size_m = token_count; - while (size_m > 0) + // d GEMM + auto gemm_down = [&]( + const half* in_addr, + half* out_addr, + const uint16_t* trellis, + const int K, + const float output_scale, + const bool accumulate + ) { - #define ARGS \ - in_addr, \ - trellis, \ - out_addr, \ - MIN(size_m, 16), \ - intermediate_dim, \ - hidden_dim, \ - locks, \ - nullptr - #define SHAPE_ARGS \ - MOE_TILESIZE_M, \ - MOE_TILESIZE_K, \ - MOE_TILESIZE_N, \ - MOE_SH_STAGES, \ - MOE_FRAG_STAGES - if constexpr (t_bits) - exl3_gemm_kernel_inner(ARGS); - else switch(K) + int size_m = token_count; + while (size_m > 0) { - case 1: exl3_gemm_kernel_inner<1, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 2: exl3_gemm_kernel_inner<2, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 3: exl3_gemm_kernel_inner<3, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 4: exl3_gemm_kernel_inner<4, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 5: exl3_gemm_kernel_inner<5, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 6: exl3_gemm_kernel_inner<6, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 7: exl3_gemm_kernel_inner<7, false, cb, SHAPE_ARGS, false>(ARGS); break; - case 8: exl3_gemm_kernel_inner<8, false, cb, SHAPE_ARGS, false>(ARGS); break; - }; - #undef ARGS - #undef SHAPE_ARGS + #define ARGS \ + in_addr, \ + trellis, \ + out_addr, \ + MIN(size_m, 16), \ + intermediate_dim, \ + hidden_dim, \ + locks, \ + nullptr, \ + output_scale, \ + accumulate + #define SHAPE_ARGS \ + MOE_TILESIZE_M, \ + MOE_TILESIZE_K, \ + MOE_TILESIZE_N, \ + MOE_SH_STAGES, \ + MOE_FRAG_STAGES + if constexpr (t_bits == 0) + { + if (accumulate) switch(K) + { + case 1: exl3_gemm_kernel_inner<1, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 2: exl3_gemm_kernel_inner<2, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 3: exl3_gemm_kernel_inner<3, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 4: exl3_gemm_kernel_inner<4, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 5: exl3_gemm_kernel_inner<5, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 6: exl3_gemm_kernel_inner<6, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 7: exl3_gemm_kernel_inner<7, false, 1, SHAPE_ARGS, false>(ARGS); break; + case 8: exl3_gemm_kernel_inner<8, false, 1, SHAPE_ARGS, false>(ARGS); break; + default: __trap(); + } + else switch(K) + { + case 1: exl3_gemm_kernel_inner<1, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 2: exl3_gemm_kernel_inner<2, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 3: exl3_gemm_kernel_inner<3, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 4: exl3_gemm_kernel_inner<4, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 5: exl3_gemm_kernel_inner<5, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 6: exl3_gemm_kernel_inner<6, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 7: exl3_gemm_kernel_inner<7, false, cb, SHAPE_ARGS, false>(ARGS); break; + case 8: exl3_gemm_kernel_inner<8, false, cb, SHAPE_ARGS, false>(ARGS); break; + default: __trap(); + } + } + else + exl3_gemm_kernel_inner(ARGS); + #undef ARGS + #undef SHAPE_ARGS - in_addr += 16 * intermediate_dim; - out_addr += 16 * hidden_dim; - size_m -= 16; - } - }; + in_addr += 16 * intermediate_dim; + out_addr += 16 * hidden_dim; + size_m -= 16; + } + }; - gemm_down(temp_intermediate_g, temp_state_g, exp_down_trellis, K_down); - group_barrier(group_idx, group_size, barrier_counters_sense); + gemm_down( + temp_intermediate_g, temp_state_g, exp_down_trellis, K_down, + 1.0f, false + ); + group_barrier(group_idx, group_size, barrier_counters_sense); + if constexpr (t_bits == 0) + for (int stage = 0; stage < num_residual_stages; ++stage) + { + const int residual_idx = stage * num_experts + expert_idx; + const float down_scale = residual_down_scales[residual_idx]; + if (down_scale != 0.0f) + gemm_down( + temp_intermediate_g, + temp_state_g, + residual_down_trellis[residual_idx], + residual_down_k[stage], + down_scale, + true + ); + group_barrier(group_idx, group_size, barrier_counters_sense); + } - // Output hadamard for d + scatter add - auto had_d_out = [&]() - { - const int warps_per_token = hidden_dim / 128; - const int total_warps = token_count * warps_per_token; - const int64_t* top_x = token_sorted + start; - const half* weights = weight_sorted + start; - for (int warp_idx = warp_idx0; warp_idx < total_warps; warp_idx += warps_per_group) + // Output hadamard for d + scatter add + auto had_d_out = [&]() { - int token_idx = top_x[warp_idx / warps_per_token]; - half weight = weights[warp_idx / warps_per_token]; - int token_off = warp_idx % warps_per_token; - float* out_ptr = output_state + token_idx * hidden_dim + token_off * 128; - had_hf_r_128_d_inner - ( - temp_state_g + 128 * warp_idx, - out_ptr, - exp_down_svh + 128 * token_off, - 0.088388347648f * __half2float(weight) - ); - } - }; + const int warps_per_token = hidden_dim / 128; + const int total_warps = token_count * warps_per_token; + const int64_t* top_x = token_sorted + start; + const half* weights = weight_sorted + start; + for (int warp_idx = warp_idx0; warp_idx < total_warps; warp_idx += warps_per_group) + { + int token_idx = top_x[warp_idx / warps_per_token]; + half weight = weights[warp_idx / warps_per_token]; + int token_off = warp_idx % warps_per_token; + float* out_ptr = output_state + token_idx * hidden_dim + token_off * 128; + had_hf_r_128_d_inner + ( + temp_state_g + 128 * warp_idx, + out_ptr, + exp_down_svh + 128 * token_off, + 0.088388347648f * __half2float(weight) + ); + } + }; - had_d_out(); + had_d_out(); - // Draw the next ticket and publish it to the group through the end-of-expert barrier, which also protects - // the temp buffers for reuse. Grabbed tickets continue from num_groups since 0..num_groups-1 are implicit - if (block_idx == 0 && threadIdx.x == 0) - sched[2 + group_idx] = num_groups + atomicAdd(&sched[0], 1); - group_barrier(group_idx, group_size, barrier_counters_sense); + // Every tile ends with a barrier that protects the temp buffers. The + // last tile also draws and publishes the group's next expert ticket. + if (chunk_start + token_count == expert_end && + block_idx == 0 && threadIdx.x == 0) + sched[2 + group_idx] = num_groups + atomicAdd(&sched[0], 1); + group_barrier(group_idx, group_size, barrier_counters_sense); + } ticket = sched[2 + group_idx]; } diff --git a/tests/test_exl3_moe_additive.py b/tests/test_exl3_moe_additive.py new file mode 100644 index 00000000..69b1889d --- /dev/null +++ b/tests/test_exl3_moe_additive.py @@ -0,0 +1,169 @@ +import pytest +import torch + +from exllamav3.ext import exllamav3_ext as ext + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason = "CUDA is required" +) + + +def test_additive_moe_abi_version(): + version = ext.EXL3_MOE_ADDITIVE_ABI_VERSION + assert type(version) is int + assert version == 1 + + +def _metadata(device: torch.device): + dim = 128 + trellis = torch.zeros((dim // 16, dim // 16, 16), + dtype = torch.int16, device = device) + suh = torch.ones(dim, dtype = torch.float16, device = device) + svh = torch.ones(dim, dtype = torch.float16, device = device) + pointer = lambda tensor: torch.tensor( + [tensor.data_ptr()], dtype = torch.int64, device = device + ) + + base = ( + pointer(trellis), pointer(suh), pointer(svh), + pointer(trellis), pointer(suh), pointer(svh), + pointer(trellis), pointer(suh), pointer(svh), + ) + # A second sparse stage has a null pointer. A zero scale must prevent the + # kernel from dereferencing it. + residual_ptrs = torch.tensor( + [[trellis.data_ptr()], [0]], dtype = torch.int64, device = device + ) + residual_scales = torch.tensor( + [[0.125], [0.0]], dtype = torch.float32, device = device + ) + residual_k = torch.tensor([1, 1], dtype = torch.int32, device = device) + residual = ( + residual_ptrs, residual_ptrs, residual_ptrs, + residual_scales, residual_scales, residual_scales, + residual_k, residual_k, residual_k, + ) + # Keep every pointee alive for the asynchronous launches below. + retained = (trellis, suh, svh) + return base, residual, retained + + +def _run(rows: int, capacity: int, id_dtype: torch.dtype, + weight_dtype: torch.dtype, base, residual, topk: int = 1): + device = torch.device("cuda", torch.cuda.current_device()) + dim = 128 + concurrency = ext.exl3_moe_max_concurrency(device.index) + hidden = ( + torch.arange(rows * dim, dtype = torch.float32, device = device) + .reshape(rows, dim).remainder(17).sub_(8).mul_(1e-3).half() + ) + output = torch.zeros_like(hidden, dtype = torch.float32) + topk_ids = torch.zeros((rows, topk), dtype = id_dtype, device = device) + topk_weights = torch.linspace( + 0.25, 0.75, rows * topk, dtype = torch.float32, device = device + ).to(weight_dtype).reshape(rows, topk) + expert_map = torch.zeros(1, dtype = torch.int64, device = device) + expert_count = torch.empty(2, dtype = torch.int64, device = device) + expert_offsets = torch.empty_like(expert_count) + route_count = rows * topk + token_sorted = torch.empty(route_count, dtype = torch.int64, device = device) + weight_sorted = torch.empty( + route_count, dtype = torch.float16, device = device + ) + temp_state_g = torch.empty( + (concurrency, capacity, dim), dtype = torch.float16, device = device + ) + temp_state_u = torch.empty_like(temp_state_g) + temp_intermediate_g = torch.empty_like(temp_state_g) + temp_intermediate_u = torch.empty_like(temp_state_g) + + ext.exl3_moe_additive_fused( + hidden, output, topk_ids, topk_weights, expert_map, + expert_count, expert_offsets, token_sorted, weight_sorted, + temp_state_g, temp_state_u, temp_intermediate_g, + temp_intermediate_u, 0, 1, 1, 1, *base, *residual, 1, + True, False, True, False, True, False, 0.0, 1, + ) + torch.cuda.synchronize() + return ( + output, expert_count, expert_offsets, token_sorted, weight_sorted, + topk_weights, + ) + + +@pytest.mark.parametrize("rows", [3, 5]) +@pytest.mark.parametrize("id_dtype", [torch.int32, torch.int64]) +@pytest.mark.parametrize( + "weight_dtype", [torch.float16, torch.bfloat16, torch.float32] +) +@torch.inference_mode() +def test_additive_fused_tiles_overflow_and_skips_sparse_null_stage( + rows, id_dtype, weight_dtype +): + device = torch.device("cuda", torch.cuda.current_device()) + base, residual, retained = _metadata(device) + + tiled = _run(rows, 2, id_dtype, weight_dtype, base, residual) + reference = _run(rows, rows, id_dtype, weight_dtype, base, residual) + + torch.testing.assert_close(tiled[0], reference[0], rtol = 1e-3, atol = 1e-3) + assert tiled[0].abs().max().item() > 0 + assert tiled[1].tolist() == [rows, 0] + assert tiled[2].tolist() == [0, rows] + assert tiled[3].tolist() == list(range(rows)) + torch.testing.assert_close( + tiled[4], tiled[5].reshape(-1).half(), rtol = 0, atol = 0 + ) + assert retained + + +@torch.inference_mode() +def test_additive_fused_topk_routes_do_not_require_topk_times_workspace(): + device = torch.device("cuda", torch.cuda.current_device()) + rows = 5 + base, residual, retained = _metadata(device) + + tiled = _run( + rows, rows, torch.int64, torch.float32, base, residual, topk = 2 + ) + reference = _run( + rows, 2 * rows, torch.int64, torch.float32, base, residual, topk = 2 + ) + + torch.testing.assert_close(tiled[0], reference[0], rtol = 1e-3, atol = 1e-3) + assert tiled[0].abs().max().item() > 0 + assert tiled[1].tolist() == [2 * rows, 0] + assert retained + + +@torch.inference_mode() +def test_legacy_moe_preserves_oversized_expert_fallback_contract(): + device = torch.device("cuda", torch.cuda.current_device()) + rows = 3 + capacity = 2 + dim = 128 + concurrency = ext.exl3_moe_max_concurrency(device.index) + base, _residual, retained = _metadata(device) + hidden = torch.ones((rows, dim), dtype = torch.float16, device = device) + output = torch.zeros_like(hidden, dtype = torch.float32) + expert_count = torch.tensor([rows, 0], dtype = torch.int64, device = device) + token_sorted = torch.arange(rows, dtype = torch.int64, device = device) + weight_sorted = torch.ones(rows, dtype = torch.float16, device = device) + temp_state_g = torch.empty( + (concurrency, capacity, dim), dtype = torch.float16, device = device + ) + temp_state_u = torch.empty_like(temp_state_g) + temp_intermediate_g = torch.empty_like(temp_state_g) + temp_intermediate_u = torch.empty_like(temp_state_g) + + ext.exl3_moe( + hidden, output, expert_count, token_sorted, weight_sorted, + temp_state_g, temp_state_u, temp_intermediate_g, + temp_intermediate_u, 0, 1, 1, 1, *base, + True, False, True, False, True, False, 0.0, -1, + ) + torch.cuda.synchronize() + + assert torch.count_nonzero(output).item() == 0 + assert retained