MoE: Add fused additive EXL3 kernels

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");

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)
         {