From 640a12c4151197cede857a7b5041bc97826f13b5 Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Mon, 28 Sep 2026 04:06:34 -0500 Subject: [PATCH 1/3] perf(cuda): a SwiGLU limit's two clamps fuse into the quantized mat-vec's GLU GLM-5.3 clamps before the GLU (a SwiGLU limit of 10, in its routed experts, shared experts and dense FFN): llama-graph puts a CLAMP of the gate to [-inf, 10] and of the up projection to [-10, 10] between the mat-vecs and the SWIGLU, and those two nodes kept the gate/up mat-vec + GLU fusion from matching, so a decode ran the two mat-vecs apart, then two clamps and the GLU: five launches where the fusion takes one. ggml_cuda_can_fuse now takes {MUL_MAT, CLAMP, MUL_MAT, CLAMP, GLU} and its MUL_MAT_ID twin (either mat-vec first: the GLU's first source is the gate's clamp) when the clamps are exactly a limit L (gate [-inf, L], up [-L, L], L finite and above 0), nothing sits between a mat-vec and its clamp, and the GLU is a SwiGLU; mul_mat_vec_q applies the limit in its epilogue as one float, glu_limit. The PQ2_0 tensor-core path takes no limit and stays unfused. GGML_CUDA_GLU_LIMIT_FUSE_LEGACY=1 leaves the clamps as nodes. In place (moe-graph: GLM-5.3-Flash's routed FFN, 8 layers of 32 experts, CUDA graphs and PDL; RTX 5080, legs A-B-B-A): 807/808 us at 1 token against 865/857 with the clamps as nodes. 3 tokens are unchanged (1,956/1,955 against 1,953/1,974): the fused mat-vec takes one column, and a verify's experts run apart as before. test-backend-ops: MUL_MAT_VEC_FUSION builds the limit as llama-graph does, at 0.5 so that most products clamp and a kernel that drops it fails, for Q8_0, IQ3_XXS and Q4_0, with and without ids, at 1 and 3 columns and with a batch. --- ggml/src/ggml-cuda/common.cuh | 2 + ggml/src/ggml-cuda/ggml-cuda.cu | 74 +++++++++++++++++++++++++++++++-- ggml/src/ggml-cuda/mmvq.cu | 12 +++++- tests/test-backend-ops.cpp | 27 +++++++++++- 4 files changed, 107 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index fb618f03e891..923c547a1e16 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1881,6 +1881,7 @@ struct ggml_cuda_mm_fusion_args_host { const ggml_tensor * x_scale = nullptr; const ggml_tensor * gate_scale = nullptr; ggml_glu_op glu_op; + float glu_limit = 0.0f; // > 0: before the GLU, the gate clamped to at most glu_limit and x to +-glu_limit (a SwiGLU limit) }; struct ggml_cuda_mm_fusion_args_device { const void * x_bias = nullptr; @@ -1889,6 +1890,7 @@ struct ggml_cuda_mm_fusion_args_device { const void * x_scale = nullptr; const void * gate_scale = nullptr; ggml_glu_op glu_op; + float glu_limit = 0.0f; }; struct ggml_cuda_kernel_launch_params { diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 60bdd8e41f65..139d13c9d66c 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1701,13 +1701,38 @@ static void ggml_cuda_mul_mat_cublas(ggml_backend_cuda_context & ctx, const ggml } } +// A SwiGLU limit (GLM-5.3, DeepSeek-V4): the gate clamped to [-inf, L] and the up projection to [-L, L] before the GLU, +// two CLAMP nodes the mat-vec epilogue takes as one float. 0 when the pair is not exactly that. +static float ggml_cuda_glu_limit(const ggml_tensor * gate_clamp, const ggml_tensor * up_clamp) { + if (gate_clamp->op != GGML_OP_CLAMP || up_clamp->op != GGML_OP_CLAMP) { + return 0.0f; + } + const float limit = ggml_get_op_params_f32(up_clamp, 1); + if (!(limit > 0.0f) || !std::isfinite(limit) || ggml_get_op_params_f32(up_clamp, 0) != -limit || + ggml_get_op_params_f32(gate_clamp, 0) != -INFINITY || ggml_get_op_params_f32(gate_clamp, 1) != limit) { + return 0.0f; + } + return limit; +} + static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, const ggml_tensor * ffn_gate, const ggml_tensor * glu, const ggml_tensor * ffn_up_bias = nullptr, const ggml_tensor * ffn_gate_bias = nullptr, const ggml_tensor * ffn_up_scale = nullptr, - const ggml_tensor * ffn_gate_scale = nullptr) { + const ggml_tensor * ffn_gate_scale = nullptr, + const ggml_tensor * ffn_up_clamp = nullptr, + const ggml_tensor * ffn_gate_clamp = nullptr) { + const bool has_clamp = ffn_up_clamp != nullptr || ffn_gate_clamp != nullptr; + if (has_clamp) { + // a limit fuses alone: no bias or scale between the mat-vec and its clamp, and only into a SwiGLU + if (!ffn_up_clamp || !ffn_gate_clamp || ffn_up_bias || ffn_gate_bias || ffn_up_scale || ffn_gate_scale || + ffn_up_clamp->src[0] != ffn_up || ffn_gate_clamp->src[0] != ffn_gate || + ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || ggml_cuda_glu_limit(ffn_gate_clamp, ffn_up_clamp) == 0.0f) { + return false; + } + } const bool has_bias = ffn_up_bias != nullptr || ffn_gate_bias != nullptr; const bool has_scale = ffn_up_scale != nullptr || ffn_gate_scale != nullptr; @@ -1730,8 +1755,8 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, const ggml_op expected_bias_op = is_mul_mat ? GGML_OP_ADD : GGML_OP_ADD_ID; const ggml_tensor * ffn_up_bias_src = has_scale ? ffn_up_scale : ffn_up; const ggml_tensor * ffn_gate_bias_src = has_scale ? ffn_gate_scale : ffn_gate; - const ggml_tensor * ffn_up_out = has_bias ? ffn_up_bias : ffn_up_bias_src; - const ggml_tensor * ffn_gate_out = has_bias ? ffn_gate_bias : ffn_gate_bias_src; + const ggml_tensor * ffn_up_out = has_clamp ? ffn_up_clamp : has_bias ? ffn_up_bias : ffn_up_bias_src; + const ggml_tensor * ffn_gate_out = has_clamp ? ffn_gate_clamp : has_bias ? ffn_gate_bias : ffn_gate_bias_src; if (glu->src[0] != ffn_gate_out || glu->src[1] != ffn_up_out) { return false; @@ -3738,6 +3763,26 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, } } + std::initializer_list mul_mat_id_clamp_glu_ops = { GGML_OP_MUL_MAT_ID, GGML_OP_CLAMP, GGML_OP_MUL_MAT_ID, GGML_OP_CLAMP, GGML_OP_GLU }; + std::initializer_list mul_mat_clamp_glu_ops = { GGML_OP_MUL_MAT, GGML_OP_CLAMP, GGML_OP_MUL_MAT, GGML_OP_CLAMP, GGML_OP_GLU }; + + if ((is_equal(mul_mat_id_clamp_glu_ops, ops) || is_equal(mul_mat_clamp_glu_ops, ops)) && + ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 4 })) { + // each mat-vec is followed by its clamp; the GLU's first source names which pair is the gate + const ggml_tensor * glu = cgraph->nodes[node_idx + 4]; + const bool gate_first = glu->src[0] == cgraph->nodes[node_idx + 1]; + const ggml_tensor * ffn_gate = cgraph->nodes[node_idx + (gate_first ? 0 : 2)]; + const ggml_tensor * gate_clamp = cgraph->nodes[node_idx + (gate_first ? 1 : 3)]; + const ggml_tensor * ffn_up = cgraph->nodes[node_idx + (gate_first ? 2 : 0)]; + const ggml_tensor * up_clamp = cgraph->nodes[node_idx + (gate_first ? 3 : 1)]; + + if (ggml_cuda_should_fuse_mul_mat(ffn_up, ffn_gate, glu, nullptr, nullptr, nullptr, nullptr, up_clamp, gate_clamp)) { + int out_nodes[] = { node_idx + 4 }; + return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1, false, + ggml_cuda_mmvq_staged_src1(ffn_up)); + } + } + if ((is_equal(mul_mat_id_glu_ops, ops) || is_equal(mul_mat_glu_ops, ops)) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) { const ggml_tensor * ffn_gate = cgraph->nodes[node_idx]; @@ -4426,7 +4471,9 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return bias == nullptr || ids != nullptr || ggml_cuda_mmvq_fusion_operand_ok(bias, out); }; - // gate + glu + up, with optional scale/bias on both lanes. + // gate + glu + up, with optional scale/bias on both lanes, or a SwiGLU limit's two clamps. + // GGML_CUDA_GLU_LIMIT_FUSE_LEGACY=1: the clamps stay nodes of their own and the mat-vecs run apart. + static const bool glu_limit_legacy = ggml_env_switch("GGML_CUDA_GLU_LIMIT_FUSE_LEGACY"); for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; @@ -4712,6 +4759,25 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph fused_node_count = 3; break; } + } else if (!glu_limit_legacy && ggml_cuda_can_fuse(cgraph, i, { op, GGML_OP_CLAMP, op, GGML_OP_CLAMP, GGML_OP_GLU }, {})) { + // a SwiGLU limit (GLM-5.3's experts, shared experts and dense FFN): the two clamps ride the quantized + // mat-vec's epilogue with the GLU, one launch for five nodes + ggml_tensor * glu = cgraph->nodes[i + 4]; + const bool gate_first = glu->src[0] == cgraph->nodes[i + 1]; + ggml_tensor * gate = cgraph->nodes[i + (gate_first ? 0 : 2)]; + ggml_tensor * up = cgraph->nodes[i + (gate_first ? 2 : 0)]; + + if (ggml_cuda_should_fuse_mul_mat_vec_q(up, /*with_gate =*/ true)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate->src[0]; + fusion_data.glu_op = ggml_get_glu_op(glu); + fusion_data.glu_limit = ggml_cuda_glu_limit(glu->src[0], glu->src[1]); + + ggml_cuda_mul_mat_vec_q(*cuda_ctx, up->src[0], up->src[1], up->src[2], glu, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = 5; + break; + } } } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 0992cbf23870..135018c29fc7 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -609,6 +609,7 @@ static __global__ void mul_mat_vec_q( const float * x_scale = nullptr; const float * gate_scale = nullptr; ggml_glu_op active_glu; + float glu_limit = 0.0f; if constexpr (has_fusion) { use_bias = fusion.x_bias != nullptr; @@ -617,6 +618,7 @@ static __global__ void mul_mat_vec_q( x_bias = (const float *) fusion.x_bias; gate_bias = (const float *) fusion.gate_bias; active_glu = fusion.glu_op; + glu_limit = fusion.glu_limit; if constexpr (type == GGML_TYPE_NVFP4) { use_scale = fusion.x_scale != nullptr; use_gate_scale = fusion.gate_scale != nullptr && use_gate; @@ -819,6 +821,10 @@ static __global__ void mul_mat_vec_q( gate_value *= gate_scales; } gate_value += gate_biases[j]; + if (glu_limit > 0.0f) { + gate_value = fminf(gate_value, glu_limit); + result = fminf(fmaxf(result, -glu_limit), glu_limit); + } switch (active_glu) { case GGML_GLU_OP_SWIGLU: result *= ggml_cuda_op_silu_single(gate_value); @@ -841,7 +847,7 @@ static __global__ void mul_mat_vec_q( } if constexpr (!has_fusion) { - GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, use_scale, use_gate_scale, active_glu, gate_bias, x_bias, x_scale, gate_scale, tmp_gate); + GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, use_scale, use_gate_scale, active_glu, glu_limit, gate_bias, x_bias, x_scale, gate_scale, tmp_gate); } if constexpr (type != GGML_TYPE_NVFP4) { GGML_UNUSED_VARS(use_scale, use_gate_scale, x_scale, gate_scale, x_scales, gate_scales); @@ -1390,6 +1396,7 @@ void ggml_cuda_mul_mat_vec_q( // which also takes a gated FFN (SWIGLU, no biases) and a residual add at columns past MMVQ_MAX_FUSED_NCOLS const bool pq2_mma = src0->type == GGML_TYPE_PQ2_0 && !ids && ne02 == 1 && ne03 == 1 && ne12 == 1 && ne13 == 1 && (fusion == nullptr || (fusion->gate_bias == nullptr && fusion->x_scale == nullptr && fusion->gate_scale == nullptr && + fusion->glu_limit == 0.0f && (fusion->gate == nullptr || (fusion->x_bias == nullptr && fusion->glu_op == GGML_GLU_OP_SWIGLU)))) && ggml_cuda_mmvq_pq2_mma_usable(ggml_cuda_info().devices[ctx.device].cc, src0->data, fusion && fusion->gate ? fusion->gate->data : nullptr, ne00, ne01, nb01 / ts_src0, ne1); @@ -1433,7 +1440,8 @@ void ggml_cuda_mul_mat_vec_q( GGML_ASSERT(ggml_nelements(fusion->gate_scale) == (ids ? src0->ne[2] : 1)); fusion_local.gate_scale = fusion->gate_scale->data; } - fusion_local.glu_op = fusion->glu_op; + fusion_local.glu_op = fusion->glu_op; + fusion_local.glu_limit = fusion->glu_limit; } // If src0 is a temporary compute buffer, clear any potential padding. diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 95963b8d233d..fe443e456b3c 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7233,14 +7233,17 @@ struct test_mul_mat_vec_fusion : public test_case { // PQ2_0 block scales 1e-4..1e-3: the products stay far under the bias, which a kernel that drops it then fails by // ~100 % (at K 5120 and block scales 0.01..1 a dropped bias moves the output by ~0.1 %, under the tolerance) const bool small_scales; + // > 0: a SwiGLU limit (GLM-5.3, DeepSeek-V4): gate clamped to [-inf, limit], up to [-limit, limit], then the GLU. + // Small enough that most products clamp, so a kernel that drops it fails. + const float glu_limit; test_mul_mat_vec_fusion(ggml_type type, ggml_glu_op op, int64_t m, int64_t n, int64_t k, bool use_id = false, int n_mats = 1, int n_used = 1, bool b = false, bool with_bias = false, bool with_gate = true, bool with_lane_scale = false, std::array batch_dims = {4, 2}, bool alias_out = false, - bool full_bias = false, bool small_scales = false) + bool full_bias = false, bool small_scales = false, float glu_limit = 0.0f) : type(type), glu_op(op), m(m), n(n), k(k), use_id(use_id), n_mats(n_mats), n_used(n_used), b(b), with_bias(with_bias), with_gate(with_gate), with_lane_scale(with_lane_scale), batch_dims(batch_dims), alias_out(alias_out), - full_bias(full_bias), small_scales(small_scales) { + full_bias(full_bias), small_scales(small_scales), glu_limit(glu_limit) { if (use_id) { GGML_ASSERT(n_used <= n_mats); } @@ -7257,6 +7260,9 @@ struct test_mul_mat_vec_fusion : public test_case { if (small_scales) { v += "," + VAR_TO_STR(small_scales); } + if (glu_limit > 0.0f) { + v += "," + VAR_TO_STR(glu_limit); + } return v; } @@ -7275,6 +7281,11 @@ struct test_mul_mat_vec_fusion : public test_case { constexpr float alpha = 1.702f; constexpr float limit = 7.0f; out = ggml_swiglu_oai(ctx, ffn_gate, ffn_up, alpha, limit); + } else if (glu_limit > 0.0f) { + // as llama-graph builds it: the GLU's first source is the gate's clamp + ffn_gate = ggml_clamp(ctx, ffn_gate, -INFINITY, glu_limit); + ffn_up = ggml_clamp(ctx, ffn_up, -glu_limit, glu_limit); + out = ggml_glu_split(ctx, ffn_gate, ffn_up, glu_op); } else { out = ggml_glu_split(ctx, ffn_gate, ffn_up, glu_op); } @@ -11384,6 +11395,18 @@ static std::vector> make_test_cases_eval() { } } + // a SwiGLU limit (GLM-5.3's experts, shared experts and dense FFN clamp at 10): its clamps fuse into the mat-vec + for (ggml_type type : { GGML_TYPE_Q8_0, GGML_TYPE_IQ3_XXS, GGML_TYPE_Q4_0 }) { + for (bool use_id : { false, true }) { + for (int64_t m_batch : { 1, 3 }) { + test_cases.emplace_back(new test_mul_mat_vec_fusion(type, GGML_GLU_OP_SWIGLU, m_batch, 32, 256, + use_id, 16, 8, false, false, true, false, {1, 1}, false, false, false, 0.5f)); + } + test_cases.emplace_back(new test_mul_mat_vec_fusion(type, GGML_GLU_OP_SWIGLU, 1, 32, 256, + use_id, 16, 8, false, false, true, false, {4, 2}, false, false, false, 0.5f)); + } + } + // Ternary Bonsai 2 27B's FFN at decode (qwen35: n_embd 5120, n_ff 17408, PQ2_0; its MTP layer Q8_0), also with the // GLU output over src1's bytes as the served graph places it; f16 takes mul_mat_vec_f, which reads src1 in the // kernel that writes the output, so it must not fuse over it From b8d44d4b8744f157c93afab3678faf5cc88d6a8d Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Mon, 28 Sep 2026 04:06:34 -0500 Subject: [PATCH 2/3] perf(cuda): routed experts stream through a ring of bulk copies, a shared expert read once A decode's routed experts are a few thousand rows each (GLM-5.3-Flash: 8 of 288, gate and up 2,048 rows of 4,096 weights at IQ3_XXS, down 4,096 of 2,048), and mul_mat_vec_q reads them a row or two a block: a block's loads are its only bytes in flight, and at an MTP verify every token reads its experts again where the tokens share them. mmvq-moe.cu runs one block an SM: a producer warp lists the distinct experts once, each with every token/slot pair that routes to it, and streams tiles of their rows through a ring of shared-memory slots with cp.async.bulk, while two teams of eight warps take the dot products for all of a tile's pairs. An IQ3_XXS warp takes its lanes' fragments of a tile into registers and frees the slot before the math (vecdotq.cuh: iq3_xxs_frag), so the slots stay in flight, and keeps the grid in shared memory a copy a lane. The last tiles go by tickets on the stream's tile counter (ggml_cuda_pq2_tile_counters), so the blocks end within a tile of each other. It takes what it measured faster at in place (moe-graph: GLM-5.3-Flash's routed FFN on 32 experts, CUDA graphs and PDL; RTX 5080, legs A-B-B-A, ring against mul_mat_vec_q): - IQ3_XXS, 8 layers: 815/814 us against 818/879 at 1 token, 1,823/1,833 against 1,953/1,992 at 3. - Q8_0 (an MTP layer's experts), 4 layers: past 1 token only, 2,270/2,260 against 2,384/2,381 at 3; at 1 token it stays on mul_mat_vec_q (988/991 against 979/977), its dot products holding the slot. - IQ4_XS not at all: 527/525 against 506/505 at 1 token, 1,243/1,240 against 1,199/1,200 at 3. GGML_CUDA_MMVQ_MOE_LEGACY=1 turns it off. Per-tile %globaltimer stamps put its steady state at DRAM's ceiling (a 12.5 KB tile every 1.15 us an SM, ~913 GB/s); what it loses is each launch's start: a team's first q8_1 vector is loaded from global memory behind the stream's queues, 3-6 us, as is a down projection's at each new expert. The IQ3_XXS dot product takes its signs as byte masks from ksigns64, 3 integer ops for 4 weights where __vcmpne4 and __vsub4 are emulated (GGML_CUDA_IQ3_XXS_SIGNS_LEGACY builds the old ones), in mul_mat_vec_q too, and get_vec_dot_q_cuda and get_vdr_mmvq move to vecdotq.cuh for both kernels. test-backend-ops: GLM-5.3-Flash's shapes through the ring, MUL_MAT_ID at 1 and 3 tokens (gate/up and down) and MUL_MAT_VEC_FUSION gated at a SwiGLU limit; MUL_MAT_ID 1066/1066, MUL_MAT_VEC_FUSION 1030/1030. --- ggml/src/ggml-cuda/common.cuh | 9 +- ggml/src/ggml-cuda/mmvq-moe.cu | 632 ++++++++++++++++++++++++++++++++ ggml/src/ggml-cuda/mmvq-moe.cuh | 48 +++ ggml/src/ggml-cuda/mmvq.cu | 96 ++--- ggml/src/ggml-cuda/vecdotq.cuh | 133 ++++++- tests/test-backend-ops.cpp | 13 + 6 files changed, 847 insertions(+), 84 deletions(-) create mode 100644 ggml/src/ggml-cuda/mmvq-moe.cu create mode 100644 ggml/src/ggml-cuda/mmvq-moe.cuh diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 923c547a1e16..cb70a62408b0 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1512,10 +1512,11 @@ struct ggml_cuda_pq2_prefetch { int64_t bytes = 0; }; -// One tile counter a stream for the PQ2_0 tensor-core launches (mmvq-pq2-mma.cu): past its own first tiles a block takes -// the next by an atomic add on it, and the block that takes the launch's last ticket sets it back to 0. A launch takes -// tickets only after its dependency wait, when every launch before it on the stream has ended, so one counter serves -// every launch of a stream, CUDA graph replays included. Made zeroed before a graph evaluation, never inside a capture. +// One tile counter a stream for the ring launches (the PQ2_0 tensor-core ones, mmvq-pq2-mma.cu, and the routed experts', +// mmvq-moe.cu): past its own first tiles a block takes the next by an atomic add on it, and the block that takes the +// launch's last ticket sets it back to 0. A launch takes tickets only after its dependency wait, when every launch before +// it on the stream has ended, so one counter serves every launch of a stream, CUDA graph replays included. Made zeroed +// before a graph evaluation, never inside a capture. struct ggml_cuda_pq2_tile_counters { int * ptr = nullptr; // GGML_CUDA_MAX_STREAMS ints diff --git a/ggml/src/ggml-cuda/mmvq-moe.cu b/ggml/src/ggml-cuda/mmvq-moe.cu new file mode 100644 index 000000000000..22084a31c9eb --- /dev/null +++ b/ggml/src/ggml-cuda/mmvq-moe.cu @@ -0,0 +1,632 @@ +#include "mmvq-moe.cuh" +#include "mmvq.cuh" +#include "unary.cuh" +#include "vecdotq.cuh" + +#include + +// A token's routed experts are a few thousand rows each (GLM-5.3-Flash: 8 experts x 2,048 rows of 4,096 weights, 1,568 +// bytes a row at IQ3_XXS), and the mat-vec kernels read them a row or two a block: a block's loads are its only bytes +// in flight, so an SM holds a few KB of requests and DRAM idles between them (RTX 5080, cold L2: 84.6 % of its +// bandwidth at gate/up, 61.1 % at down, where a 2,048-weight row leaves half of each block's threads without a block). +// Here each SM runs one block, and the block's producer warp streams its tiles of rows through a ring of shared-memory +// slots, a tile a slot, with 1D bulk copies (cp.async.bulk): the ring's slots are the SM's bytes in flight, 50-90 KB, +// whatever the row length. +// +// The work: the launch's tokens route to experts, and an expert that several tokens route to (an MTP verify's drafts +// share many) is one: every block reads ids after the dependency wait and lists the distinct experts in the order they +// first appear, each with the token/slot pairs that route to it, and a tile is 8*RPW rows of one of them. Tile t is +// expert t / ntr's row tile t % ntr. A block's first tiles, as many as its ring holds, are its own (blockIdx.x, then +// every gridDim.x-th); the rest go to whichever block asks next, by a ticket on the stream's counter +// (ggml_cuda_pq2_tile_counters, as mmvq-pq2-mma.cu takes them), so the blocks end within a tile of each other where +// owned tiles left the launch's end to the slowest (2,048 tiles on 84 SMs: 25 against 24, a tile's 1.3 us). Two teams of +// eight consumer warps take the block's tiles in turn (team 0 the even ones). In a team, warp g takes the tile's rows +// [g*RPW, (g+1)*RPW), its 32 lanes along k as mul_mat_vec_q's (vec_dot_q_cuda at the type's vdr), for each of the +// expert's pairs: the rows are read from DRAM once. +// +// A gate (vgate) is the up rows' twin, the same rows of the gate matrix beside them in the slot under one barrier, and +// the GLU (with the SwiGLU limit and the biases) is applied as mul_mat_vec_q applies it. +// +// The bytes in flight are the rate (qgre's bulk-copy ladder on sm_120: 861 B/ns at 16 KB an SM, 877 at 32, 884 at 64), +// and a slot a team holds is not in flight: a team that kept its slot through its dot products held two of four, so the +// ring streamed ~25 KB an SM and 87.5 % of DRAM with the math on against 94 % with it compiled out. So at IQ3_XXS a warp +// takes its lanes' share of its rows out of the slot into registers (iq3_xxs_frag: the 8 grid indices, the signs and +// scale, and d of each block a lane reads, 4 registers each) and releases the slot before the math, which then runs on +// registers while the producer refills it; the plan keeps a lane's fragments to 8 (mmvq_moe_fits), past which the +// kernel spills. The math is bound by shared memory and L1, not by the dot products: the grid is gathered at random, +// so past the ring the shared memory holds it a copy per lane (grid_rep[i*32 + lane]: a warp's 32 gathers in 32 banks, +// whatever the indices), and a lane keeps its q8_1 fragments in registers, loaded once per token vector, not once per +// row. The other types read as mul_mat_vec_q does (vec_dot_q_cuda), from the slot, which their team releases after it. +// +// No pointer carries __restrict__: with PDL a restrict load may compile to ld.global.nc, which the compiler can move +// above the grid dependency wait (upstream #24030). Nothing is read before that wait: the experts come from ids. + +#define MMVQ_MOE_NG 8 // a team's warps: a tile's row groups +#define MMVQ_MOE_NT 2 // teams, each on its own tiles +#define MMVQ_MOE_NW (MMVQ_MOE_NG * MMVQ_MOE_NT) // consumer warps, and one producer warp +#define MMVQ_MOE_MAX_PAIRS 64 // tokens x experts used +#define MMVQ_MOE_MAX_SLOTS 8 +#define MMVQ_MOE_SMEM_MAX (99 * 1024) // the shared memory a block may take on sm_90 - sm_120 +#define MMVQ_MOE_SLOT_TARGET (16 * 1024) // a slot's bytes past one row a warp: more slots, a shorter tail + +static_assert(MMVQ_MOE_MAX_PAIRS == 64, "the producer warp lists the pairs, two a lane"); +static_assert(MMVQ_MOE_MAX_PAIRS <= 64, "an expert's pairs are a 64-bit mask"); + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER && !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +#define MMVQ_MOE_AVAILABLE +#endif + +struct mmvq_moe_dev_args { + const char * vx; + const char * vgate; + const block_q8_1 * y; + const int32_t * ids; + const float * x_bias; + const float * gate_bias; + float * dst; + int64_t stride_channel_x_bytes; + int64_t stride_col_dst; + int64_t stride_channel_dst; + int64_t stride_bias; + int row_bytes; + int box_bytes; // a matrix's part of a slot: 8*RPW rows, padded to 128 bytes + int nrows; + int nb; // quant blocks a row + int ntr; // row tiles an expert + int n_used; + int ntokens; + int ids_stride; + int nchannels_y; + int stride_col_y; + int stride_channel_y; + int glu_op; + float glu_limit; + int nslots; + int * tile_ctr; +}; + +// IQ3_XXS: a lane's k iterations at rpw rows a warp, 4 blocks of a row an iteration, the most its registers hold. The +// plan's rpw bounds a row (MMVQ_MOE_SLOT_TARGET: 8*rpw rows in 16 KB) and so the iterations it needs: rpw 1, 16 blocks +// (K 4,096, the most taken); rpw 2, 1,024-byte rows, 10 blocks; rpw 4, 512 bytes, 5 +static constexpr __host__ __device__ int mmvq_moe_iq3_nit(const int rpw) { + return rpw == 1 ? 4 : rpw == 2 ? 3 : 2; +} + +// Whether a warp's rows fit its lanes' registers: at IQ3_XXS a lane holds rpw*nmat*nit fragments of the slot (4 registers +// each) through the math, and past 8 the kernel spills (a thread has 96 registers, 17 warps an SM). The plan takes only +// what fits, and only what fits is built. +static constexpr bool mmvq_moe_fits(const ggml_type type, const int nmat, const int rpw) { + return type != GGML_TYPE_IQ3_XXS || rpw*nmat*mmvq_moe_iq3_nit(rpw) <= 8; +} + +#ifdef MMVQ_MOE_AVAILABLE +static __device__ __forceinline__ uint32_t mmvq_moe_smem_u32(const void * p) { + return (uint32_t) __cvta_generic_to_shared(p); +} + +static __device__ __forceinline__ void mmvq_moe_mbar_init(uint64_t * bar, uint32_t count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" :: "r"(mmvq_moe_smem_u32(bar)), "r"(count) : "memory"); +} + +static __device__ __forceinline__ void mmvq_moe_mbar_arrive_expect_tx(uint64_t * bar, uint32_t bytes) { + asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" + :: "r"(mmvq_moe_smem_u32(bar)), "r"(bytes) : "memory"); +} + +static __device__ __forceinline__ void mmvq_moe_mbar_arrive(uint64_t * bar) { + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0];" :: "r"(mmvq_moe_smem_u32(bar)) : "memory"); +} + +static __device__ __forceinline__ void mmvq_moe_mbar_wait(uint64_t * bar, uint32_t parity) { + const uint32_t addr = mmvq_moe_smem_u32(bar); + uint32_t done = 0; + do { + asm volatile( + "{\n" + ".reg .pred p;\n" + "mbarrier.try_wait.parity.shared::cta.b64 p, [%1], %2;\n" + "selp.u32 %0, 1, 0, p;\n" + "}\n" + : "=r"(done) : "r"(addr), "r"(parity) : "memory"); + } while (!done); +} + +// Into this block's own shared memory, as mmvq-pq2-mma.cu's TMA loads: .shared::cluster would compile on sm_120 to a +// branch through a driver syscall that costs every resident thread a 14.5 KB stack. .shared::cta needs PTX ISA 8.6. +#if __CUDACC_VER_MAJOR__ > 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 8) +#define MMVQ_MOE_BULK_DST "shared::cta" +#else +#define MMVQ_MOE_BULK_DST "shared::cluster" +#endif + +// bytes (a multiple of 16) from src (16-byte aligned) to dst (16-byte aligned), counted on bar +static __device__ __forceinline__ void mmvq_moe_bulk_load(void * dst, const void * src, uint32_t bytes, uint64_t * bar, + uint64_t policy) { + asm volatile( + "cp.async.bulk." MMVQ_MOE_BULK_DST ".global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;" + :: "r"(mmvq_moe_smem_u32(dst)), "l"((uint64_t) src), "r"(bytes), "r"(mmvq_moe_smem_u32(bar)), "l"(policy) + : "memory"); +} +#endif // MMVQ_MOE_AVAILABLE + +// pair p's q8_1 vector: token p / n_used's, for expert slot p % n_used +static __device__ __forceinline__ const block_q8_1 * mmvq_moe_y(const mmvq_moe_dev_args & a, const int p) { + const int t = p / a.n_used; + const int slot = p % a.n_used; + return a.y + (int64_t) (slot % a.nchannels_y)*a.stride_channel_y + (int64_t) t*a.stride_col_y; +} + +template +__launch_bounds__((MMVQ_MOE_NW + 1)*32, 1) +static __global__ void mmvq_moe(const mmvq_moe_dev_args a) { +#ifdef MMVQ_MOE_AVAILABLE + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = get_vdr_mmvq(type); + constexpr int lanes_per_block = qi / vdr; + constexpr int blocks_per_iter = 32 / lanes_per_block; + constexpr int R = MMVQ_MOE_NG * rpw; + constexpr vec_dot_q_cuda_t vec_dot = get_vec_dot_q_cuda(type); + static_assert(32 % lanes_per_block == 0, "a warp takes whole blocks of a row"); + static_assert(nmat == 1 || nmat == 2, "a matrix, or a gate beside it"); + + extern __shared__ __align__(128) char ring[]; + __shared__ uint64_t full[MMVQ_MOE_MAX_SLOTS]; + __shared__ uint64_t empty[MMVQ_MOE_MAX_SLOTS]; + __shared__ int held[MMVQ_MOE_MAX_SLOTS]; // the tile in each slot, -1: no more + __shared__ int expert[MMVQ_MOE_MAX_PAIRS]; // the distinct experts + __shared__ int pair_e[32]; // pair t*n_used + s's expert, the first 32 + __shared__ int pair_u[MMVQ_MOE_MAX_PAIRS]; // a first pair's index in expert + __shared__ unsigned long long pairs_of[MMVQ_MOE_MAX_PAIRS]; // each distinct expert's pairs, a bit each + __shared__ __align__(16) uint32_t grid_rep[type == GGML_TYPE_IQ3_XXS ? 256*32 : 1]; // iq3xxs_grid, a copy a lane + __shared__ uint64_t ksigns_s[type == GGML_TYPE_IQ3_XXS ? 128 : 1]; + + const int warp = threadIdx.x / 32; + const int lane = threadIdx.x % 32; + + ggml_cuda_pdl_lc(); + + if (threadIdx.x == 0) { + for (int s = 0; s < a.nslots; ++s) { + mmvq_moe_mbar_init(&full[s], 1); + mmvq_moe_mbar_init(&empty[s], MMVQ_MOE_NG); // the warps of the team the slot's tile goes to + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + __syncthreads(); // the barriers, before any arrival or wait + + if (warp == MMVQ_MOE_NW) { + // The producer warp alone lists the distinct experts, in the order of their first pairs, as soon as the + // dependency wait lets it read ids, and starts the loads while the consumers fill their tables (they read the + // lists after a slot's barrier, which lane 0's arrival publishes): pairs lane and lane + 32, a pair past the + // launch's with an expert of its own + ggml_cuda_pdl_sync(); // ids is a previous kernel's result + const int npairs = a.ntokens * a.n_used; + const int p1 = lane + 32; + const int e0 = lane < npairs ? a.ids[lane % a.n_used + (lane / a.n_used)*a.ids_stride] : -1 - lane; + const int e1 = p1 < npairs ? a.ids[p1 % a.n_used + (p1 / a.n_used)*a.ids_stride] : -1 - p1; + pair_e[lane] = e0; + pairs_of[lane] = 0; + pairs_of[p1] = 0; + __syncwarp(); + const unsigned same0 = __match_any_sync(0xFFFFFFFF, e0); + const unsigned same1 = __match_any_sync(0xFFFFFFFF, e1); + const int first0 = __ffs(same0) - 1; + int first1 = 32 + __ffs(same1) - 1; + for (int q = 0; q < 32; ++q) { + if (pair_e[q] == e1) { + first1 = q; + break; + } + } + const unsigned heads0 = __ballot_sync(0xFFFFFFFF, lane < npairs && first0 == lane); + const unsigned heads1 = __ballot_sync(0xFFFFFFFF, p1 < npairs && first1 == p1); + const unsigned below = (1u << lane) - 1; + if (heads0 & (1u << lane)) { + expert[__popc(heads0 & below)] = e0; + pair_u[lane] = __popc(heads0 & below); + } + if (heads1 & (1u << lane)) { + expert[__popc(heads0) + __popc(heads1 & below)] = e1; + pair_u[p1] = __popc(heads0) + __popc(heads1 & below); + } + __syncwarp(); + if (lane < npairs) { + atomicOr(&pairs_of[pair_u[first0]], 1ull << lane); + } + if (p1 < npairs) { + atomicOr(&pairs_of[pair_u[first1]], 1ull << p1); + } + __syncwarp(); + const int n_tiles = (__popc(heads0) + __popc(heads1)) * a.ntr; + + // the producer: tile i of the block's sequence into slot i % nslots, once the consumers have released it. The + // block's own tiles first (blockIdx.x, then every gridDim.x-th: the blocks stream the tiles in waves, a wave a + // few MB of neighbouring rows), all the waves but the last; then, with a counter, a ticket a tile as + // mmvq-pq2-mma.cu takes them: each block's last ticket is past the tiles, so the launch takes + // n_tiles - dyn_base + gridDim.x of them, and the block that uses the last sets the counter back to 0. A ticket + // is asked for once the tile before it is issued and used a sequence later, so its round trip to L2 overlaps the + // wait for a slot: the barrier arrival that issues a tile is a release, which waits for the atomic before it + // (a ticket a tile, each asked for before its tile's arrival, cost GLM's up projection 43.0 against 40.9 us on + // an RTX 5080) + if (lane == 0) { + uint64_t policy; // the rows stream through L2 once + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" : "=l"(policy)); + const int n_own = a.tile_ctr != nullptr ? max(a.nslots, n_tiles / (int) gridDim.x - 1) : n_tiles; + const int dyn_base = a.tile_ctr != nullptr ? min(n_own * (int) gridDim.x, n_tiles) : n_tiles; + const int last = n_tiles - dyn_base + (int) gridDim.x - 1; // the launch's last ticket + int n_end = 0; + int ticket = -1; // the next sequence's, once asked for + for (int i = 0;; ++i) { + const int s = i % a.nslots; + if (i >= a.nslots) { + mmvq_moe_mbar_wait(&empty[s], (uint32_t) ((i / a.nslots - 1) & 1)); + } + int tile = -1; + if (n_end == 0) { + if (i < n_own) { + const int t = (int) blockIdx.x + i * (int) gridDim.x; + tile = t < dyn_base ? t : -1; + } else if (dyn_base < n_tiles) { + if (ticket == last) { + atomicExch(a.tile_ctr, 0); + } + tile = dyn_base + ticket < n_tiles ? dyn_base + ticket : -1; + } + } + if (tile < 0) { + // the end: a slot with no rows for each team (the next NT in the sequence are one a team), its barrier + // completed by a plain arrival + held[s] = -1; + mmvq_moe_mbar_arrive(&full[s]); + if (++n_end == MMVQ_MOE_NT) { + break; + } + continue; + } + held[s] = tile; + const int u = tile / a.ntr; + const int row0 = (tile % a.ntr) * R; + const uint32_t bytes = (uint32_t) (min(R, a.nrows - row0) * a.row_bytes); + const int64_t off = expert[u]*a.stride_channel_x_bytes + (int64_t) row0*a.row_bytes; + char * slot = ring + (size_t) s*nmat*a.box_bytes; + mmvq_moe_mbar_arrive_expect_tx(&full[s], nmat*bytes); + mmvq_moe_bulk_load(slot, a.vx + off, bytes, &full[s], policy); + if constexpr (nmat == 2) { + mmvq_moe_bulk_load(slot + a.box_bytes, a.vgate + off, bytes, &full[s], policy); + } + if (i + 1 >= n_own && dyn_base < n_tiles) { + ticket = atomicAdd(a.tile_ctr, 1); // the next sequence's tile, asked for now this one is issued + } + } + } + return; + } + + // the consumers: the tables depend on nothing, so they are filled under the previous kernel, before the dependency + // wait, and the consumer warps meet on barrier 1 once they are whole + if constexpr (type == GGML_TYPE_IQ3_XXS) { + for (int i = threadIdx.x; i < 256*8; i += MMVQ_MOE_NW*32) { + const uint32_t v = iq3xxs_grid[i / 8]; + *(uint4 *) &grid_rep[4*i] = make_uint4(v, v, v, v); // entry i/8's copies 4*(i%8) to 4*(i%8) + 3 + } + for (int i = threadIdx.x; i < 128; i += MMVQ_MOE_NW*32) { + ksigns_s[i] = ksigns64[i]; + } + asm volatile("bar.sync 1, %0;" :: "n"(MMVQ_MOE_NW*32) : "memory"); + } + ggml_cuda_pdl_sync(); // the tokens are the previous kernels' results, and dst may still be read + + const int team = warp / MMVQ_MOE_NG; // the block's tiles team, team + NT, ... + const int g = warp % MMVQ_MOE_NG; // the warp's rows of a tile, [g*rpw, (g+1)*rpw) + const int kb0 = lane / lanes_per_block; // the lane's first block of a row + const int kqs = vdr * (lane % lanes_per_block); // and its quant ints in each + + // IQ3_XXS: the lane's q8_1 fragments, a k iteration each (its block's 8 ints and scale), of the vector y_key + // (token * nchannels_y + channel), kept across pairs and tiles until the vector changes + constexpr int nit = type == GGML_TYPE_IQ3_XXS ? mmvq_moe_iq3_nit(rpw) : 1; + int yu[nit][8]; + float yd[nit]; + int y_key = -1; + [[maybe_unused]] const auto grid = [&](const int i) { return grid_rep[i*32 + lane]; }; + [[maybe_unused]] const auto ksigns = [&](const int i) { return ksigns_s[i]; }; + + for (int i = team;; i += MMVQ_MOE_NT) { + const int s = i % a.nslots; + mmvq_moe_mbar_wait(&full[s], (uint32_t) ((i / a.nslots) & 1)); + const int tile = held[s]; + if (tile < 0) { + break; + } + const int u = tile / a.ntr; + const int e = expert[u]; + const int row0 = (tile % a.ntr) * R + g*rpw; // this warp's first row + const char * box = ring + (size_t) s*nmat*a.box_bytes; + + // IQ3_XXS: the lane's fragments of its rows out of the slot, and the slot released before the math, so the + // producer refills it while the warps compute: a slot is held for these loads, not for the dot products, and + // the ring's slots are nearly all in flight at once + [[maybe_unused]] iq3_xxs_frag wf[rpw][nmat][nit]; + if constexpr (type == GGML_TYPE_IQ3_XXS) { +#pragma unroll + for (int it = 0; it < nit; ++it) { + const int kb = kb0 + it*blocks_per_iter; + if (kb < a.nb) { +#pragma unroll + for (int r = 0; r < rpw; ++r) { +#pragma unroll + for (int m = 0; m < nmat; ++m) { + wf[r][m][it] = iq3_xxs_frag_load( + (const block_iq3_xxs *) (box + m*a.box_bytes) + (g*rpw + r)*a.nb + kb, kqs); + } + } + } + } + __syncwarp(); // every lane has read this slot: release it + if (lane == 0) { + mmvq_moe_mbar_arrive(&empty[s]); + } + } + + // every pair that routes to the expert, each with its own tokens (a down projection's differ by slot) + for (unsigned long long pairs = pairs_of[u]; pairs != 0; pairs &= pairs - 1) { + const int p = __ffsll(pairs) - 1; + const int t = p / a.n_used; + const int slot = p % a.n_used; + const block_q8_1 * y = mmvq_moe_y(a, p); + + float acc[rpw][nmat] = {{0.0f}}; + if constexpr (type == GGML_TYPE_IQ3_XXS) { + const int key = t*a.nchannels_y + slot % a.nchannels_y; + if (key != y_key) { + y_key = key; +#pragma unroll + for (int it = 0; it < nit; ++it) { + const int kb = kb0 + it*blocks_per_iter; + if (kb < a.nb) { + const block_q8_1 * yb = y + kb*(qk/QK8_1) + kqs/2; +#pragma unroll + for (int l = 0; l < 8; ++l) { + yu[it][l] = get_int_b4(yb->qs, l); + } + yd[it] = __low2float(yb->ds); + } + } + } +#pragma unroll + for (int it = 0; it < nit; ++it) { + const int kb = kb0 + it*blocks_per_iter; + if (kb < a.nb) { +#pragma unroll + for (int r = 0; r < rpw; ++r) { +#pragma unroll + for (int m = 0; m < nmat; ++m) { + acc[r][m] += vec_dot_iq3_xxs_frag(wf[r][m][it], yu[it], yd[it], grid, ksigns); + } + } + } + } + } else { +#pragma unroll 4 + for (int kb = kb0; kb < a.nb; kb += blocks_per_iter) { + const int kby = kb * (qk/QK8_1); +#pragma unroll + for (int r = 0; r < rpw; ++r) { +#pragma unroll + for (int m = 0; m < nmat; ++m) { + acc[r][m] += vec_dot(box + m*a.box_bytes, &y[kby], (g*rpw + r)*a.nb + kb, kqs); + } + } + } + } +#pragma unroll + for (int r = 0; r < rpw; ++r) { +#pragma unroll + for (int m = 0; m < nmat; ++m) { + acc[r][m] = warp_reduce_sum<32>(acc[r][m]); + } + } + +#pragma unroll + for (int r = 0; r < rpw; ++r) { + const float * v = acc[r]; + const int row = row0 + r; + if (lane == 0 && row < a.nrows) { + float result = v[0]; + if (a.x_bias != nullptr) { + result += a.x_bias[e*a.stride_bias + row]; + } + if constexpr (nmat == 2) { + float gate_value = v[1]; + if (a.gate_bias != nullptr) { + gate_value += a.gate_bias[e*a.stride_bias + row]; + } + if (a.glu_limit > 0.0f) { + gate_value = fminf(gate_value, a.glu_limit); + result = fminf(fmaxf(result, -a.glu_limit), a.glu_limit); + } + switch (a.glu_op) { + case GGML_GLU_OP_SWIGLU: + result *= ggml_cuda_op_silu_single(gate_value); + break; + case GGML_GLU_OP_GEGLU: + result *= ggml_cuda_op_gelu_single(gate_value); + break; + case GGML_GLU_OP_SWIGLU_OAI: + result = ggml_cuda_op_swiglu_oai_single(gate_value, result); + break; + default: + result = result * gate_value; + break; + } + } + a.dst[t*a.stride_col_dst + slot*a.stride_channel_dst + row] = result; + } + } + } + + if constexpr (type != GGML_TYPE_IQ3_XXS) { + __syncwarp(); // every lane has read this slot: release it + if (lane == 0) { + mmvq_moe_mbar_arrive(&empty[s]); + } + } + } +#else + GGML_UNUSED(a); + NO_DEVICE_CODE; +#endif // MMVQ_MOE_AVAILABLE +} + +// --------------------------------------------------------------------------------------------------------------------- +// host + +static bool mmvq_moe_legacy() { + static const bool legacy = ggml_env_switch("GGML_CUDA_MMVQ_MOE_LEGACY"); + return legacy; +} + +// The types, at the token counts, where the ring measured faster than mul_mat_vec_q in place, or level with it (moe-graph: +// GLM-5.3-Flash's routed FFN on 32 experts, CUDA graphs and PDL, RTX 5080). IQ3_XXS: 815 against 818 us for 8 layers at +// 1 token, 1,823 against 1,953 at 3. Q8_0: 989 against 978 us for 4 layers at 1 token, 2,265 against 2,382 at 3; its dot +// products read the slot, so a team holds the slot through them, and only an MTP verify's shared experts make up for +// that. IQ4_XS was slower at both, 527 against 505 and 1,241 against 1,199, so it keeps mul_mat_vec_q. +static bool mmvq_moe_takes(ggml_type type, int64_t ntokens) { + return type == GGML_TYPE_IQ3_XXS || (type == GGML_TYPE_Q8_0 && ntokens > 1); +} + +static int64_t mmvq_moe_row_bytes(ggml_type type, int64_t ncols_x) { + return ncols_x / ggml_blck_size(type) * (int64_t) ggml_type_size(type); +} + +// the static shared memory a block takes past the ring, with room to spare: barriers, experts, pairs and sums (1,960 +// bytes), and at IQ3_XXS its tables (33,792) +static int mmvq_moe_meta_bytes(ggml_type type) { + return type == GGML_TYPE_IQ3_XXS ? 36 * 1024 : 3 * 1024; +} + +static int mmvq_moe_ring_max(ggml_type type) { + return MMVQ_MOE_SMEM_MAX - mmvq_moe_meta_bytes(type); +} + +// rpw: rows a row group, the tile 8*rpw rows; the most whose slot stays under MMVQ_MOE_SLOT_TARGET (one row a group in +// any case), and as many slots as fit, a whole number a team: sequence i's slot is i % nslots and its team i % NT, so a +// team always uses the same slots, and waits on each slot's full barrier one phase after the last it consumed (a team +// waiting on a slot another team has not consumed yet could see the parity of the phase before and pass early) +struct mmvq_moe_plan { + int rpw = 0; + int nslots = 0; + int box_bytes = 0; +}; + +static mmvq_moe_plan mmvq_moe_make_plan(ggml_type type, int64_t row_bytes, int nmat) { + for (int rpw : { 4, 2, 1 }) { + const int64_t box = GGML_PAD(MMVQ_MOE_NG * rpw * row_bytes, 128); + const int64_t slot = nmat * box; + if (rpw > 1 && (slot > MMVQ_MOE_SLOT_TARGET || !mmvq_moe_fits(type, nmat, rpw))) { + continue; + } + const int fit = (int) std::min(mmvq_moe_ring_max(type) / slot, MMVQ_MOE_MAX_SLOTS) / MMVQ_MOE_NT * + MMVQ_MOE_NT; + if (fit < MMVQ_MOE_NT) { + return {}; + } + return { rpw, fit, (int) box }; + } + return {}; +} + +bool ggml_cuda_mmvq_moe_usable(int cc, ggml_type type, const void * vx, const void * vgate, int64_t ncols_x, + int64_t nrows_x, int64_t stride_row_x, int64_t stride_channel_x, int64_t n_used, + int64_t ntokens) { + if (mmvq_moe_legacy() || !GGML_CUDA_CC_IS_NVIDIA(cc) || cc < GGML_CUDA_CC_HOPPER || + ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_HOPPER || !mmvq_moe_takes(type, ntokens)) { + return false; + } + const int64_t ts = ggml_type_size(type); + const int64_t row_bytes = mmvq_moe_row_bytes(type, ncols_x); + const mmvq_moe_plan p = mmvq_moe_make_plan(type, row_bytes, vgate != nullptr ? 2 : 1); + // a tile is 8*rpw dense rows: 16-byte aligned wherever it starts when the rows are an even number of bytes, and the + // last tile of an expert when the expert is a whole number of 16-byte units. IQ3_XXS: a lane's k iterations (4 + // blocks of a row each) fit its fragments + return ncols_x % ggml_blck_size(type) == 0 && stride_row_x*ts == row_bytes && row_bytes % 2 == 0 && + (nrows_x*row_bytes) % 16 == 0 && (stride_channel_x*ts) % 16 == 0 && + (uintptr_t) vx % 16 == 0 && (uintptr_t) vgate % 16 == 0 && + ntokens >= 1 && ntokens <= MMVQ_MAX_BATCH_SIZE && n_used >= 1 && ntokens*n_used <= MMVQ_MOE_MAX_PAIRS && + nrows_x < (1 << 30) && row_bytes < (1 << 20) && p.rpw > 0 && + (type != GGML_TYPE_IQ3_XXS || (ncols_x/QK_K + 3)/4 <= mmvq_moe_iq3_nit(p.rpw)); +} + +template +static void mmvq_moe_launch(const mmvq_moe_dev_args & a, const int nblocks, cudaStream_t stream) { + if constexpr (mmvq_moe_fits(type, nmat, rpw)) { + const size_t smem = (size_t) a.nslots * nmat * a.box_bytes; + CUDA_SET_SHARED_MEMORY_LIMIT((mmvq_moe), mmvq_moe_ring_max(type)); // every plan's size, once + const ggml_cuda_kernel_launch_params params(dim3(nblocks), dim3((MMVQ_MOE_NW + 1)*32), smem, stream); + ggml_cuda_kernel_launch(mmvq_moe, params, a); + } else { + GGML_ABORT("%s: no plan takes %d matrices at %d rows a warp for %s", __func__, nmat, rpw, ggml_type_name(type)); + } +} + +template +static void mmvq_moe_launch_type(const mmvq_moe_dev_args & a, const int nmat, const int rpw, const int nblocks, + cudaStream_t stream) { + switch (nmat*8 + rpw) { + case 1*8 + 1: mmvq_moe_launch(a, nblocks, stream); break; + case 1*8 + 2: mmvq_moe_launch(a, nblocks, stream); break; + case 1*8 + 4: mmvq_moe_launch(a, nblocks, stream); break; + case 2*8 + 1: mmvq_moe_launch(a, nblocks, stream); break; + case 2*8 + 2: mmvq_moe_launch(a, nblocks, stream); break; + case 2*8 + 4: mmvq_moe_launch(a, nblocks, stream); break; + default: GGML_ABORT("%s: no instance for %d matrices at %d rows a warp", __func__, nmat, rpw); + } +} + +void ggml_cuda_mmvq_moe(const ggml_cuda_mmvq_moe_args & args, cudaStream_t stream) { + const int nmat = args.vgate != nullptr ? 2 : 1; + const int64_t row_bytes = mmvq_moe_row_bytes(args.type, args.ncols_x); + const mmvq_moe_plan p = mmvq_moe_make_plan(args.type, row_bytes, nmat); + GGML_ASSERT(p.rpw > 0 && "ggml_cuda_mmvq_moe_usable holds a plan"); + + const int R = MMVQ_MOE_NG * p.rpw; + const int ntr = (int) ((args.nrows_x + R - 1) / R); + + mmvq_moe_dev_args a; + a.vx = (const char *) args.vx; + a.vgate = (const char *) args.vgate; + a.y = (const block_q8_1 *) args.y; + a.ids = args.ids; + a.x_bias = args.x_bias; + a.gate_bias = args.gate_bias; + a.dst = args.dst; + a.stride_channel_x_bytes = args.stride_channel_x_bytes; + a.stride_col_dst = args.stride_col_dst; + a.stride_channel_dst = args.stride_channel_dst; + a.stride_bias = args.stride_bias; + a.row_bytes = (int) row_bytes; + a.box_bytes = p.box_bytes; + a.nrows = (int) args.nrows_x; + a.nb = (int) (args.ncols_x / ggml_blck_size(args.type)); + a.ntr = ntr; + a.n_used = (int) args.n_used; + a.ntokens = (int) args.ntokens; + a.ids_stride = (int) args.ids_stride; + a.nchannels_y = (int) args.nchannels_y; + a.stride_col_y = (int) args.stride_col_y; + a.stride_channel_y = (int) args.stride_channel_y; + a.glu_op = (int) args.glu_op; + a.glu_limit = args.glu_limit; + a.nslots = p.nslots; + a.tile_ctr = args.tile_ctr; + + // one block an SM, never more than the tiles of the most distinct experts the pairs can name + const int nsm = ggml_cuda_info().devices[ggml_cuda_get_device()].nsm; + const int nblocks = (int) std::min(nsm, args.ntokens*args.n_used*ntr); + + switch (args.type) { + case GGML_TYPE_IQ3_XXS: mmvq_moe_launch_type(a, nmat, p.rpw, nblocks, stream); break; + case GGML_TYPE_Q8_0: mmvq_moe_launch_type (a, nmat, p.rpw, nblocks, stream); break; + default: GGML_ABORT("%s: no instance for %s", __func__, ggml_type_name(args.type)); + } +} diff --git a/ggml/src/ggml-cuda/mmvq-moe.cuh b/ggml/src/ggml-cuda/mmvq-moe.cuh new file mode 100644 index 000000000000..f9238528f036 --- /dev/null +++ b/ggml/src/ggml-cuda/mmvq-moe.cuh @@ -0,0 +1,48 @@ +#pragma once + +#include "common.cuh" + +// Routed experts (MUL_MAT_ID) at 1-MMVQ_MAX_BATCH_SIZE tokens, each SM streaming its tiles of expert rows through a ring +// of shared-memory slots with 1D bulk copies (mmvq-moe.cu). An expert that several tokens route to is read once for all +// of them, and a gate and its GLU fuse at any token count. +// +// The weights: expert e's row r at vx + e*stride_channel_x_bytes + r*row_bytes, rows dense (row_bytes = ncols_x / +// the type's block size * its block bytes). The tokens: q8_1 (quantize_row_q8_1_cuda's MMVQ layout), token t's column +// for expert slot s at y + (s % nchannels_y)*stride_channel_y + t*stride_col_y blocks. ids[s + t*ids_stride] is the +// expert in slot s of token t, and the result for that pair is dst[t*stride_col_dst + s*stride_channel_dst + r]. +// A bias (x_bias for vx's product, gate_bias for vgate's) is expert e's row r at e*stride_bias + r. +struct ggml_cuda_mmvq_moe_args { + ggml_type type; + const void * vx; + const void * vgate; // nullptr: dst = vx's product (+ x_bias); else the GLU of the two products + const void * y; + const int32_t * ids; + const float * x_bias; + const float * gate_bias; + float * dst; + ggml_glu_op glu_op; + float glu_limit; // > 0: the gate clamped to [-inf, limit] and vx's product to [-limit, limit] before the GLU + int64_t ncols_x; + int64_t nrows_x; + int64_t stride_channel_x_bytes; + int64_t n_used; // experts a token routes to: ids' columns + int64_t ntokens; + int64_t ids_stride; + int64_t nchannels_y; + int64_t stride_col_y; + int64_t stride_channel_y; + int64_t stride_col_dst; + int64_t stride_channel_dst; + int64_t stride_bias; + int * tile_ctr; // the stream's tile counter (ggml_cuda_pq2_tile_counters), or nullptr: every tile its + // block's own +}; + +// Whether ggml_cuda_mmvq_moe takes these weights: an NVIDIA GPU from Hopper on (cp.async.bulk), a type and token count +// it measured faster at (IQ3_XXS; Q8_0 past one token), 16-byte aligned rows, tiles and experts, a ring that fits the +// shared memory, at most MMVQ_MAX_BATCH_SIZE tokens and 64 token/expert pairs. GGML_CUDA_MMVQ_MOE_LEGACY=1: never. +bool ggml_cuda_mmvq_moe_usable(int cc, ggml_type type, const void * vx, const void * vgate, int64_t ncols_x, + int64_t nrows_x, int64_t stride_row_x, int64_t stride_channel_x, int64_t n_used, + int64_t ntokens); + +void ggml_cuda_mmvq_moe(const ggml_cuda_mmvq_moe_args & args, cudaStream_t stream); diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 135018c29fc7..2dfcb38d91f8 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -1,4 +1,5 @@ #include "mmvq.cuh" +#include "mmvq-moe.cuh" #include "mmvq-pq2-mma.cuh" #include "quantize.cuh" #include "unary.cuh" @@ -7,68 +8,6 @@ #include #include -typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs); - -static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) { - switch (type) { - case GGML_TYPE_Q1_0: return vec_dot_q1_0_q8_1; - case GGML_TYPE_Q2_0: return vec_dot_q2_0_q8_1; - case GGML_TYPE_PQ2_0: return vec_dot_pq2_0_q8_1; - case GGML_TYPE_PTQ1_0: return vec_dot_ptq1_0_q8_1; - case GGML_TYPE_Q4_0: return vec_dot_q4_0_q8_1; - case GGML_TYPE_Q4_1: return vec_dot_q4_1_q8_1; - case GGML_TYPE_Q5_0: return vec_dot_q5_0_q8_1; - case GGML_TYPE_Q5_1: return vec_dot_q5_1_q8_1; - case GGML_TYPE_Q8_0: return vec_dot_q8_0_q8_1; - case GGML_TYPE_MXFP4: return vec_dot_mxfp4_q8_1; - case GGML_TYPE_NVFP4: return vec_dot_nvfp4_q8_1; - case GGML_TYPE_Q2_K: return vec_dot_q2_K_q8_1; - case GGML_TYPE_Q3_K: return vec_dot_q3_K_q8_1; - case GGML_TYPE_Q4_K: return vec_dot_q4_K_q8_1; - case GGML_TYPE_Q5_K: return vec_dot_q5_K_q8_1; - case GGML_TYPE_Q6_K: return vec_dot_q6_K_q8_1; - case GGML_TYPE_IQ2_XXS: return vec_dot_iq2_xxs_q8_1; - case GGML_TYPE_IQ2_XS: return vec_dot_iq2_xs_q8_1; - case GGML_TYPE_IQ2_S: return vec_dot_iq2_s_q8_1; - case GGML_TYPE_IQ3_XXS: return vec_dot_iq3_xxs_q8_1; - case GGML_TYPE_IQ1_S: return vec_dot_iq1_s_q8_1; - case GGML_TYPE_IQ1_M: return vec_dot_iq1_m_q8_1; - case GGML_TYPE_IQ4_NL: return vec_dot_iq4_nl_q8_1; - case GGML_TYPE_IQ4_XS: return vec_dot_iq4_xs_q8_1; - case GGML_TYPE_IQ3_S: return vec_dot_iq3_s_q8_1; - default: return nullptr; - } -} - -static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) { - switch (type) { - case GGML_TYPE_Q1_0: return VDR_Q1_0_Q8_1_MMVQ; - case GGML_TYPE_Q2_0: return VDR_Q2_0_Q8_1_MMVQ; - case GGML_TYPE_PQ2_0: return VDR_PQ2_0_Q8_1_MMVQ; - case GGML_TYPE_PTQ1_0: return VDR_PTQ1_0_Q8_1_MMVQ; - case GGML_TYPE_Q4_0: return VDR_Q4_0_Q8_1_MMVQ; - case GGML_TYPE_Q4_1: return VDR_Q4_1_Q8_1_MMVQ; - case GGML_TYPE_Q5_0: return VDR_Q5_0_Q8_1_MMVQ; - case GGML_TYPE_Q5_1: return VDR_Q5_1_Q8_1_MMVQ; - case GGML_TYPE_Q8_0: return VDR_Q8_0_Q8_1_MMVQ; - case GGML_TYPE_MXFP4: return VDR_MXFP4_Q8_1_MMVQ; - case GGML_TYPE_NVFP4: return VDR_NVFP4_Q8_1_MMVQ; - case GGML_TYPE_Q2_K: return VDR_Q2_K_Q8_1_MMVQ; - case GGML_TYPE_Q3_K: return VDR_Q3_K_Q8_1_MMVQ; - case GGML_TYPE_Q4_K: return VDR_Q4_K_Q8_1_MMVQ; - case GGML_TYPE_Q5_K: return VDR_Q5_K_Q8_1_MMVQ; - case GGML_TYPE_Q6_K: return VDR_Q6_K_Q8_1_MMVQ; - case GGML_TYPE_IQ2_XXS: return VDR_IQ2_XXS_Q8_1_MMVQ; - case GGML_TYPE_IQ2_XS: return VDR_IQ2_XS_Q8_1_MMVQ; - case GGML_TYPE_IQ2_S: return VDR_IQ2_S_Q8_1_MMVQ; - case GGML_TYPE_IQ3_XXS: return VDR_IQ3_XXS_Q8_1_MMVQ; - case GGML_TYPE_IQ3_S: return VDR_IQ3_S_Q8_1_MMVQ; - case GGML_TYPE_IQ4_NL: return VDR_IQ4_NL_Q8_1_MMVQ; - case GGML_TYPE_IQ4_XS: return VDR_IQ4_XS_Q8_1_MMVQ; - default: return 1; - } -} - enum mmvq_parameter_table_id { MMVQ_PARAMETERS_GENERIC = 0, MMVQ_PARAMETERS_TURING, @@ -1486,6 +1425,39 @@ void ggml_cuda_mul_mat_vec_q( const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0; + // routed experts at 1-8 tokens, each SM streaming its tiles of expert rows through a ring of bulk copies, and each + // distinct expert read once (mmvq-moe.cu) + if (ids && ne03 == 1 && ne13 == 1 && (fusion == nullptr || (fusion->x_scale == nullptr && fusion->gate_scale == nullptr)) && + ggml_cuda_mmvq_moe_usable(ggml_cuda_info().devices[ctx.device].cc, src0->type, src0->data, + fusion_local.gate, ne00, ne01, s01, s02, ne1, ne2)) { + ggml_cuda_mmvq_moe_args args{}; + args.type = src0->type; + args.vx = src0->data; + args.vgate = fusion_local.gate; + args.y = src1_q8_1.get(); + args.ids = ids_d; + args.x_bias = (const float *) fusion_local.x_bias; + args.gate_bias = (const float *) fusion_local.gate_bias; + args.dst = dst_d; + args.glu_op = fusion_local.glu_op; + args.glu_limit = fusion_local.glu_limit; + args.ncols_x = ne00; + args.nrows_x = ne01; + args.stride_channel_x_bytes = s02 * (int64_t) ts_src0; + args.n_used = ne1; + args.ntokens = ne2; + args.ids_stride = ids_stride; + args.nchannels_y = nchannels_y; + args.stride_col_y = stride_col_y; + args.stride_channel_y = stride_channel_y; + args.stride_col_dst = stride_col_dst; + args.stride_channel_dst = stride_channel_dst; + args.stride_bias = stride_channel_dst; // as mul_mat_vec_q reads an expert's bias + args.tile_ctr = ctx.pq2_tile_counter(); + ggml_cuda_mmvq_moe(args, stream); + return; + } + if (pq2_mma) { ggml_cuda_mmvq_pq2_mma(src0->data, fusion_local.gate, src1_q8_1.get(), (const float *) fusion_local.x_bias, dst_d, ne00, ne01, ne1, s01, s11, s1, ctx.pq2_next, ctx.pq2_tile_counter(), stream); diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index 6c70b6cb5a39..710f290ea5d8 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -1379,39 +1379,72 @@ static __device__ __forceinline__ float vec_dot_iq2_s_q8_1( #define VDR_IQ3_XXS_Q8_1_MMVQ 2 #define VDR_IQ3_XXS_Q8_1_MMQ 2 -static __device__ __forceinline__ float vec_dot_iq3_xxs_q8_1( - const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { +// One lane's 32 weights of an IQ3_XXS block (its quant ints iqs and iqs + 1): the 8 grid indices, the signs and scale, +// and the block's d, all a lane reads of the block, so a caller can take them out of shared memory and free it before +// the math (mmvq-moe.cu) +struct iq3_xxs_frag { + int2 q3; + uint32_t aux32; + float d; +}; + +static __device__ __forceinline__ iq3_xxs_frag iq3_xxs_frag_load(const block_iq3_xxs * bq3, const int iqs) { + return { make_int2(get_int_b2(bq3->qs, iqs), get_int_b2(bq3->qs, iqs+1)), (uint32_t) get_int_b2(bq3->qs, QK_K/16 + iqs/2), + __half2float(bq3->d) }; +} - const block_iq3_xxs * bq3 = (const block_iq3_xxs *) vbq + kbx; +// A fragment against the 8 ints u of the q8_1 block it meets and its scale d8. grid(i) and ksigns(i) read iq3xxs_grid +// and ksigns64 wherever the caller keeps them (mmvq-moe.cu: the grid a copy per lane in shared memory, so a warp's 32 +// gathers meet no bank conflict). +template +static __device__ __forceinline__ float vec_dot_iq3_xxs_frag( + const iq3_xxs_frag & w, const int * u, const float d8, grid_t grid, ksigns_t ksigns) { - const int2 q3_packed = make_int2(get_int_b2(bq3->qs, iqs), get_int_b2(bq3->qs, iqs+1)); - const uint8_t * q3 = (const uint8_t *) &q3_packed; - const uint32_t aux32 = get_int_b2(bq3->qs, QK_K/16 + iqs/2); + const uint8_t * q3 = (const uint8_t *) &w.q3; + const uint32_t aux32 = w.aux32; int sumi = 0; #pragma unroll for (int l0 = 0; l0 < 8; l0 += 2) { - const int2 grid_pos = make_int2(iq3xxs_grid[q3[l0 + 0]], iq3xxs_grid[q3[l0 + 1]]); + const int2 grid_pos = make_int2(grid(q3[l0 + 0]), grid(q3[l0 + 1])); + // the 8 weights' signs as byte masks, 0xFF where negative (the 8th: the parity of the 7 stored), and a negative + // weight as ~g + 1, with no carry into the next byte as every grid byte is 4-62: 3 integer ops for 4 weights, + // where __vcmpne4 and __vsub4 are emulated in several each (built with -DGGML_CUDA_IQ3_XXS_SIGNS_LEGACY) +#ifdef GGML_CUDA_IQ3_XXS_SIGNS_LEGACY + GGML_UNUSED(ksigns); const uint32_t signs = unpack_ksigns(aux32 >> (7*l0/2)); - const int signs0 = __vcmpne4(signs & 0x08040201, 0); const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); - - const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0); - const int signs1 = __vcmpne4(signs & 0x80402010, 0); const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); - - const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1); - - sumi = ggml_cuda_dp4a(grid_l, u0, sumi); - sumi = ggml_cuda_dp4a(grid_h, u1, sumi); +#else + const uint64_t signs = ksigns((aux32 >> (7*l0/2)) & 0x7F); + const int signs0 = (int) (uint32_t) signs; + const int signs1 = (int) (uint32_t) (signs >> 32); + const int grid_l = (grid_pos.x ^ signs0) + (signs0 & 0x01010101); + const int grid_h = (grid_pos.y ^ signs1) + (signs1 & 0x01010101); +#endif // GGML_CUDA_IQ3_XXS_SIGNS_LEGACY + + sumi = ggml_cuda_dp4a(grid_l, u[l0 + 0], sumi); + sumi = ggml_cuda_dp4a(grid_h, u[l0 + 1], sumi); } const int ls = aux32 >> 28; sumi = (ls*sumi + sumi/2)/2; - const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs/2].ds); - return d * sumi; + return w.d * d8 * sumi; +} + +static __device__ __forceinline__ float vec_dot_iq3_xxs_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq3_xxs * bq3 = (const block_iq3_xxs *) vbq + kbx; + int u[8]; +#pragma unroll + for (int l = 0; l < 8; ++l) { + u[l] = get_int_b4(bq8_1[iqs/2].qs, l); + } + return vec_dot_iq3_xxs_frag(iq3_xxs_frag_load(bq3, iqs), u, __low2float(bq8_1[iqs/2].ds), + [](const int i) { return iq3xxs_grid[i]; }, [](const int i) { return ksigns64[i]; }); } #define VDR_IQ3_S_Q8_1_MMVQ 2 @@ -1588,3 +1621,67 @@ static __device__ __forceinline__ float vec_dot_iq4_xs_q8_1( const float d = __half2float(bq4->d) * __low2float(bq8_1[iqs/4].ds); return d * sumi; } + +// The mat-vec dot product of each type against q8_1, and the quant ints a thread takes of a block per call: the +// entry points of the MMVQ kernels (mmvq.cu, mmvq-moe.cu). +typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs); + +static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) { + switch (type) { + case GGML_TYPE_Q1_0: return vec_dot_q1_0_q8_1; + case GGML_TYPE_Q2_0: return vec_dot_q2_0_q8_1; + case GGML_TYPE_PQ2_0: return vec_dot_pq2_0_q8_1; + case GGML_TYPE_PTQ1_0: return vec_dot_ptq1_0_q8_1; + case GGML_TYPE_Q4_0: return vec_dot_q4_0_q8_1; + case GGML_TYPE_Q4_1: return vec_dot_q4_1_q8_1; + case GGML_TYPE_Q5_0: return vec_dot_q5_0_q8_1; + case GGML_TYPE_Q5_1: return vec_dot_q5_1_q8_1; + case GGML_TYPE_Q8_0: return vec_dot_q8_0_q8_1; + case GGML_TYPE_MXFP4: return vec_dot_mxfp4_q8_1; + case GGML_TYPE_NVFP4: return vec_dot_nvfp4_q8_1; + case GGML_TYPE_Q2_K: return vec_dot_q2_K_q8_1; + case GGML_TYPE_Q3_K: return vec_dot_q3_K_q8_1; + case GGML_TYPE_Q4_K: return vec_dot_q4_K_q8_1; + case GGML_TYPE_Q5_K: return vec_dot_q5_K_q8_1; + case GGML_TYPE_Q6_K: return vec_dot_q6_K_q8_1; + case GGML_TYPE_IQ2_XXS: return vec_dot_iq2_xxs_q8_1; + case GGML_TYPE_IQ2_XS: return vec_dot_iq2_xs_q8_1; + case GGML_TYPE_IQ2_S: return vec_dot_iq2_s_q8_1; + case GGML_TYPE_IQ3_XXS: return vec_dot_iq3_xxs_q8_1; + case GGML_TYPE_IQ1_S: return vec_dot_iq1_s_q8_1; + case GGML_TYPE_IQ1_M: return vec_dot_iq1_m_q8_1; + case GGML_TYPE_IQ4_NL: return vec_dot_iq4_nl_q8_1; + case GGML_TYPE_IQ4_XS: return vec_dot_iq4_xs_q8_1; + case GGML_TYPE_IQ3_S: return vec_dot_iq3_s_q8_1; + default: return nullptr; + } +} + +static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) { + switch (type) { + case GGML_TYPE_Q1_0: return VDR_Q1_0_Q8_1_MMVQ; + case GGML_TYPE_Q2_0: return VDR_Q2_0_Q8_1_MMVQ; + case GGML_TYPE_PQ2_0: return VDR_PQ2_0_Q8_1_MMVQ; + case GGML_TYPE_PTQ1_0: return VDR_PTQ1_0_Q8_1_MMVQ; + case GGML_TYPE_Q4_0: return VDR_Q4_0_Q8_1_MMVQ; + case GGML_TYPE_Q4_1: return VDR_Q4_1_Q8_1_MMVQ; + case GGML_TYPE_Q5_0: return VDR_Q5_0_Q8_1_MMVQ; + case GGML_TYPE_Q5_1: return VDR_Q5_1_Q8_1_MMVQ; + case GGML_TYPE_Q8_0: return VDR_Q8_0_Q8_1_MMVQ; + case GGML_TYPE_MXFP4: return VDR_MXFP4_Q8_1_MMVQ; + case GGML_TYPE_NVFP4: return VDR_NVFP4_Q8_1_MMVQ; + case GGML_TYPE_Q2_K: return VDR_Q2_K_Q8_1_MMVQ; + case GGML_TYPE_Q3_K: return VDR_Q3_K_Q8_1_MMVQ; + case GGML_TYPE_Q4_K: return VDR_Q4_K_Q8_1_MMVQ; + case GGML_TYPE_Q5_K: return VDR_Q5_K_Q8_1_MMVQ; + case GGML_TYPE_Q6_K: return VDR_Q6_K_Q8_1_MMVQ; + case GGML_TYPE_IQ2_XXS: return VDR_IQ2_XXS_Q8_1_MMVQ; + case GGML_TYPE_IQ2_XS: return VDR_IQ2_XS_Q8_1_MMVQ; + case GGML_TYPE_IQ2_S: return VDR_IQ2_S_Q8_1_MMVQ; + case GGML_TYPE_IQ3_XXS: return VDR_IQ3_XXS_Q8_1_MMVQ; + case GGML_TYPE_IQ3_S: return VDR_IQ3_S_Q8_1_MMVQ; + case GGML_TYPE_IQ4_NL: return VDR_IQ4_NL_Q8_1_MMVQ; + case GGML_TYPE_IQ4_XS: return VDR_IQ4_XS_Q8_1_MMVQ; + default: return 1; + } +} diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index fe443e456b3c..f6053dc91869 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10625,6 +10625,14 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_MXFP4, GGML_TYPE_F32, 32, 2, false, 2880, 32, 2880)); test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q4_0, GGML_TYPE_F32, 32, 2, false, 2880, 32, 2880)); + // GLM-5.3-Flash's routed experts (IQ3_XXS, 8 used; 32 experts here, of its 288) at decode and an MTP verify, where + // the CUDA backend streams them through mmvq-moe.cu's ring: gate/up rows of 4,096 weights with a token's vector + // shared by its experts, down rows of 2,048 with a vector each, and three tokens that share experts + for (int n : { 1, 3 }) { + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_IQ3_XXS, GGML_TYPE_F32, 32, 8, true, 2048, n, 4096)); + test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_IQ3_XXS, GGML_TYPE_F32, 32, 8, false, 4096, n, 2048)); + } + for (ggml_type type_a : all_types) { test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 4, 2, false, 64, 16, 3*ggml_blck_size(type_a))); } @@ -11406,6 +11414,11 @@ static std::vector> make_test_cases_eval() { use_id, 16, 8, false, false, true, false, {4, 2}, false, false, false, 0.5f)); } } + // and at GLM-5.3-Flash's routed gate/up (IQ3_XXS, 4,096 -> 2,048, 8 of 32 experts), a gated ring tile in mmvq-moe.cu + for (int64_t m_batch : { 1, 3 }) { + test_cases.emplace_back(new test_mul_mat_vec_fusion(GGML_TYPE_IQ3_XXS, GGML_GLU_OP_SWIGLU, m_batch, 2048, 4096, + true, 32, 8, false, false, true, false, {1, 1}, false, false, false, 0.5f)); + } // Ternary Bonsai 2 27B's FFN at decode (qwen35: n_embd 5120, n_ff 17408, PQ2_0; its MTP layer Q8_0), also with the // GLU output over src1's bytes as the served graph places it; f16 takes mul_mat_vec_f, which reads src1 in the From 217bbd77ff8a8003de46b3eda713d28d04b2cc37 Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Mon, 28 Sep 2026 04:06:53 -0500 Subject: [PATCH 3/3] docs(torad): the rows for the SwiGLU limit's fusion (640a12c41) and the routed experts' ring (b8d44d4b8) --- TORAD.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/TORAD.md b/TORAD.md index 05f1d8f59223..70ff7de97216 100644 --- a/TORAD.md +++ b/TORAD.md @@ -104,6 +104,8 @@ which pins a commit of this branch as a submodule. | `881c823c5` | `ggml_cuda_op_top_k` with a k of 64 or more past 1,024 columns selects by radix (`topk_radix`): four 8-bit histogram passes find the k-th largest key, then one ordered pass writes every column above it and the lowest columns equal to it, the tiled path's set and tie-break. Up to 16,384 columns a block keeps its row in registers; a wider row is cut into 8,192-column tiles selected in parallel, then one block a row selects among their candidates; one wide row stays with CUB's top-k. The GLM-5.3 DSA indexer (512 pools a layer) paid the tiled path's k serial block reductions: RTX 5080, k 512, 3,520 columns 194.5 -> 6.1 us, 109,020 x 3 66.3 -> 21.8, a 512-row ubatch 13.8-18x. TOP_K 523/523; two mutants fail it (the last equal key unwritten; stage 2 writing a candidate's position). | `GGML_CUDA_TOPK_RADIX_LEGACY` | | `33eb70bc0` | `build_attn_mha` takes `mask_is_prefix` (true by default) and the two sparse-attention callers (the DSA layers' top-k mask, and the kpool path's) pass false: every one-sequence causal mask was tagged as a prefix of the cells, the hint under which the CUDA flash attention skips its range scan and live tiles and applies the mask over the whole range, so a GLM-5.3 DSA decode read every cell's K and V to keep the indexer's 2,048. The dense layers keep the hint. Same attention by another split of the cells, bits can move at float rounding; the gain sits at long caches (436k) and is measured on the served model. | `LLAMA_ATTN_SPARSE_MASK_PREFIX_LEGACY` | | `78e0fadf7` | `topk_radix` also takes a row of up to 1,024 columns at any k: with CUB's top-k available such a row had no path but CUB one row after another, four launches a row (a DSA indexer's ubatch under 4,096 cells: 2,048 launches a layer). RTX 5080: 3.3-3.9 us against 8.2-12.7 (1 row) and 129-171 (16 rows). | `GGML_CUDA_TOPK_RADIX_LEGACY` | +| `640a12c41` | A SwiGLU limit (GLM-5.3's routed experts, shared experts and dense FFN: the gate clamped to `[-inf, 10]`, the up projection to `[-10, 10]`, then the GLU) fuses into the quantized mat-vec's GLU epilogue: `ggml_cuda_can_fuse` takes `{MUL_MAT(_ID), CLAMP, MUL_MAT(_ID), CLAMP, GLU}` when the clamps are exactly a limit and the GLU a SwiGLU, and `mul_mat_vec_q` applies it as one float, where the two nodes had kept the gate/up fusion from matching (five launches for one). In place (GLM-5.3-Flash's routed FFN, 8 layers of 32 experts, RTX 5080): 807 against 861 us at 1 token; a verify's 3 tokens are unchanged, their experts unfused. | `GGML_CUDA_GLU_LIMIT_FUSE_LEGACY` | +| `b8d44d4b8` | Routed experts at 1-8 tokens stream through `mmvq-moe.cu`: one block an SM, a producer warp listing the distinct experts once (every token/slot pair that routes to one reads it once) and streaming tiles of their rows through a ring of shared-memory slots by `cp.async.bulk`, two teams of eight warps on the dot products; an IQ3_XXS warp frees its slot before the math, and the last tiles go by tickets on the stream's tile counter. It takes IQ3_XXS, and Q8_0 past 1 token: in place (GLM-5.3-Flash's routed FFN on 32 experts, CUDA graphs and PDL, RTX 5080) IQ3_XXS 815 against 818 us for 8 layers at 1 token and 1,823 against 1,953 at 3, Q8_0 2,265 against 2,382 for 4 layers at 3; IQ4_XS measured slower and keeps `mul_mat_vec_q`. Its steady state runs at DRAM's ceiling (~913 GB/s); a launch's first q8_1 vector, loaded from global memory behind the stream, costs 3-6 us. | `GGML_CUDA_MMVQ_MOE_LEGACY` | Every switch in the last column is read once per process and parses as an integer: a `*_LEGACY` switch set to `0` is the same as unset (the change stays on), and `=0` turns off `GGML_CUDA_LORA_RANK1_FUSE` and