From 8c78ba9ca386016168cde995e3d75b900cbd193a Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Sun, 27 Sep 2026 19:18:31 -0500 Subject: [PATCH 01/10] perf(cuda): the PQ2_0 tensor-core launches take their tiles past their own from a counter Each block of mmvq_pq2_mma (one an SM) owned every gridDim.x-th tile, so a launch ended when its slowest block did: the blocks the DRAM served later finished a tile or so after the rest, and the launch's last microseconds ran on a few SMs. Now a block's first tiles, as many as its ring holds before the dependency wait, are still its own; past them its producer takes the next tile by an atomic add on its stream's counter (ggml_cuda_pq2_tile_counters, one int a stream, made zeroed before a graph evaluation), after its own dependency wait, when every launch before it on the stream has ended. Every block's last ticket is past the tiles, so a launch takes n_tiles - dyn_base + gridDim.x of them, and the block that takes the last sets the counter back to 0 for the next launch (CUDA graph replays included). The producer tells the consumers each slot's tile (or the end) in shared memory beside the slot's barriers. The arithmetic and its order are unchanged. GGML_CUDA_PQ2_MMA_TILES_LEGACY=1 gives every tile to its block again. RTX 5080, Ternary Bonsai 2 27B, q4_0 K/V, f16 state, one build with and without the switch: - llama-bench with graphs on, legs N-L-L-N: tg128 at depth 0 +1.46 % (106.18 / 105.74 / 106.04 / 106.07 against 103.38 / 104.83 / 105.01 / 104.69 tok/s), at 16,384 +1.34 %, tg64 at 245,760 +1.47 % (68.19 / 68.26 against 67.92 / 66.55); pp4 with -rs 3 (the MTP verify), the median of 30 samples a leg, +0.62 % at 0 and +0.81 % at 16,384. - llama-server as served (MTP drafts of 3, one slot, 32,768 context), 12 greedy agentic requests a leg: the same text in every pair, 231.53 -> 234.18 tok/s against the undisturbed legacy leg (same-text geomean +1.8 %, se 0.5 %). - nsys, graphs off, 16 tokens at 16,384: 169,594 -> 166,576 us. - test-backend-ops MUL_MAT (pq2_0) 99/99, MUL_MAT_GROUP 24/24, MUL_MAT_VEC_FUSION 1010/1010, both ways; without the last ticket's reset the counter build fails 29, 7 and 26 of them and passes them all under the switch. llama-server at one slot and at 4 slots with --kv-unified, 128 greedy tokens with top-5 log-probabilities: bit-identical to the published engine-32e695e build both ways. --- ggml/src/ggml-cuda/common.cuh | 26 ++++++ ggml/src/ggml-cuda/ggml-cuda.cu | 5 + ggml/src/ggml-cuda/mmvq-pq2-mma.cu | 136 ++++++++++++++++++++-------- ggml/src/ggml-cuda/mmvq-pq2-mma.cuh | 10 +- ggml/src/ggml-cuda/mmvq.cu | 4 +- 5 files changed, 135 insertions(+), 46 deletions(-) diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 9d5d04c50336..fb618f03e891 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1512,6 +1512,28 @@ 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. +struct ggml_cuda_pq2_tile_counters { + int * ptr = nullptr; // GGML_CUDA_MAX_STREAMS ints + + void ensure() { + if (ptr == nullptr) { + CUDA_CHECK(cudaMalloc(&ptr, GGML_CUDA_MAX_STREAMS * sizeof(int))); + CUDA_CHECK(cudaMemset(ptr, 0, GGML_CUDA_MAX_STREAMS * sizeof(int))); + } + } + + void release() { + if (ptr != nullptr) { + CUDA_CHECK(cudaFree(ptr)); + ptr = nullptr; + } + } +}; + // Owned by the backend context that evaluates the graph: registrations are keyed by node pointer, so they only mean // something for the evaluation that made them. Cleared at the start of every graph evaluation/capture. struct ggml_cuda_gdn_gather_context { @@ -1792,9 +1814,13 @@ struct ggml_backend_cuda_context { ggml_cuda_ssm_conv_update_context ssm_conv_update_context; ggml_cuda_fattn_kv_live_context fattn_kv_live_context; ggml_cuda_pq2_prefetch pq2_next; // for the node being dispatched + ggml_cuda_pq2_tile_counters pq2_tile_counters; ~ggml_backend_cuda_context(); + // the current stream's PQ2_0 tile counter, or nullptr before the first graph evaluation made them + int * pq2_tile_counter() { return pq2_tile_counters.ptr != nullptr ? pq2_tile_counters.ptr + curr_stream_no : nullptr; } + cudaStream_t stream(int device, int stream) { if (streams[device][stream] == nullptr) { ggml_cuda_set_device(device); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 22e8a327f667..60bdd8e41f65 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -738,6 +738,10 @@ ggml_backend_cuda_context::~ggml_backend_cuda_context() { ggml_cuda_set_device(device); fattn_kv_live_context.release(); } + if (pq2_tile_counters.ptr != nullptr) { + ggml_cuda_set_device(device); + pq2_tile_counters.release(); + } for (int i = 0; i < GGML_CUDA_MAX_DEVICES; ++i) { for (int j = 0; j < GGML_CUDA_MAX_STREAMS; ++j) { if (streams[i][j] != nullptr) { @@ -5494,6 +5498,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; ggml_cuda_set_device(cuda_ctx->device); + cuda_ctx->pq2_tile_counters.ensure(); // before a capture can begin bool use_cuda_graph = false; bool cuda_graph_update_required = false; diff --git a/ggml/src/ggml-cuda/mmvq-pq2-mma.cu b/ggml/src/ggml-cuda/mmvq-pq2-mma.cu index 3f01d65c5f1c..08d9efd8fa94 100644 --- a/ggml/src/ggml-cuda/mmvq-pq2-mma.cu +++ b/ggml/src/ggml-cuda/mmvq-pq2-mma.cu @@ -8,8 +8,11 @@ #include #include -// Weights: a tile is 16 rows, and each block (one per SM) owns whole tiles, every gridDim.x-th from its index, so -// nothing is summed across blocks: no workspace, no arrival counters, no fixup at the end of a launch. The block's +// Weights: a tile is 16 rows, and each block (one per SM) computes whole tiles, so nothing is summed across blocks: no +// workspace, no fixup at the end of a launch. A block's first tiles, the ones its ring holds before the dependency wait, +// are its own (blockIdx.x, then every gridDim.x-th); the rest go to whichever block asks next, by an atomic add on its +// stream's counter (ggml_cuda_pq2_tile_counters, common.cuh), so the blocks the DRAM serves faster take more and all end +// within a tile of each other, where owned tiles left a launch's last microseconds to the slowest few. The block's // producer warp streams its tiles through a ring of shared-memory slots with 2D TMA. A slot is one box of a tile, 16 // rows x M*1,024 weights (M*272 bytes of each row: the whole row when it fits, else the row in NKB boxes), and with a // gate matrix the gate's box of the same rows beside it, both under one barrier. The blocks walk the matrix from its @@ -79,8 +82,13 @@ static constexpr __host__ __device__ int pq2_mma_red_bytes(int nmat) { return nmat * PQ2_MMA_NW * 32 * 4 * (int) sizeof(float); } +// a slot: its boxes, its full and empty barriers, and the tile it holds +static constexpr size_t pq2_mma_slot_bytes(int m, int nmat) { + return (size_t) nmat * pq2_mma_box_bytes(m) + 2 * sizeof(uint64_t) + sizeof(int); +} + static constexpr size_t pq2_mma_smem_bytes(int m, int nmat, int nslots) { - return (size_t) nslots * nmat * pq2_mma_box_bytes(m) + pq2_mma_red_bytes(nmat) + 2 * (size_t) nslots * sizeof(uint64_t); + return (size_t) nslots * pq2_mma_slot_bytes(m, nmat) + pq2_mma_red_bytes(nmat); } #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_HOPPER && !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) @@ -171,12 +179,13 @@ static __device__ __forceinline__ float pq2_mma_dot(const int c) { // tile's map, output, rows and output stride come from its matrix in grp, not from tmap, dst, nrows and stride_col_dst. // M: the box's width in 1,024-weight units, and each consumer warp's share of it in PQ2_0 blocks; nb: PQ2_0 blocks per // row; nkb: boxes per row (the last one may reach past the row: TMA fills that part with zeros, and it is skipped). +// tile_ctr: the stream's tile counter, 0 when the launch starts and when it ends; nullptr, every tile its block's own. template static __device__ __forceinline__ void mmvq_pq2_mma_body( const CUtensorMap * tmap, const CUtensorMap * tmap_gate, const pq2_mma_group * grp, const block_q8_1 * y, const float * x_bias, float * dst, const int nrows, const int ncols, const int nb, const int n_tiles, const int nkb, const int nslots, const int evict_first, const int stride_col_y, const int stride_col_dst, - const ggml_cuda_pq2_prefetch & next) { + const ggml_cuda_pq2_prefetch & next, int * tile_ctr) { static_assert(nmat == 1 || (nmat == 2 && !has_bias), "a gated product fuses no bias"); static_assert(!grouped || (nmat == 1 && !has_bias), "a group's matrices fuse nothing"); #ifdef PQ2_MMA_AVAILABLE @@ -189,14 +198,11 @@ static __device__ __forceinline__ void mmvq_pq2_mma_body( float * red = (float *) (smem + (size_t) nslots * nmat * box_bytes); uint64_t * full = (uint64_t *) ((char *) red + pq2_mma_red_bytes(nmat)); uint64_t * empty = full + nslots; + int * held = (int *) (empty + nslots); // the tile in each slot, -1: the block has no more const int warp = threadIdx.x / 32; const int lane = threadIdx.x % 32; - // this block's tiles: blockIdx.x, then every gridDim.x-th (gridDim.x <= n_tiles) - const int n_mine = (n_tiles - (int) blockIdx.x + (int) gridDim.x - 1) / (int) gridDim.x; - const int n_boxes = n_mine * nkb; - ggml_cuda_pdl_lc(); if (threadIdx.x == 0) { @@ -227,12 +233,38 @@ static __device__ __forceinline__ void mmvq_pq2_mma_body( asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t) tmap_gate) : "memory"); } } - for (int i = 0; i < n_boxes; ++i) { - const int s = i % nslots; - if (i >= nslots) { - pq2_mbar_wait(&empty[s], (uint32_t) ((i / nslots - 1) & 1)); + // the block's own tiles first, as many as the ring holds before the consumers' dependency wait; then, with a + // counter, a ticket a tile, taken after this thread's own wait, when every launch before this one on the + // stream has ended and left the counter at 0. 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 takes the last sets the counter back to 0 + const int n_own = tile_ctr != nullptr ? (nslots + nkb - 1) / nkb : n_tiles; + const int dyn_base = tile_ctr != nullptr ? min(n_own * (int) gridDim.x, n_tiles) : n_tiles; + int i = 0; // the block's box sequence + for (int k = 0;; ++k) { + int tile = -1; + if (k < n_own) { + const int t = (int) blockIdx.x + k * (int) gridDim.x; + tile = t < dyn_base ? t : -1; + } else if (dyn_base < n_tiles) { + if (k == n_own) { + ggml_cuda_pdl_sync(); + } + const int t = atomicAdd(tile_ctr, 1); + if (t == n_tiles - dyn_base + (int) gridDim.x - 1) { + atomicExch(tile_ctr, 0); + } + tile = dyn_base + t < n_tiles ? dyn_base + t : -1; + } + if (tile < 0) { + // the end: a slot with no boxes, its barrier completed by a plain arrival + const int s = i % nslots; + if (i >= nslots) { + pq2_mbar_wait(&empty[s], (uint32_t) ((i / nslots - 1) & 1)); + } + held[s] = -1; + pq2_mbar_arrive(&full[s]); + break; } - const int tile = (int) blockIdx.x + (i / nkb) * (int) gridDim.x; const CUtensorMap * map = tmap; int row0 = tile * 16; // the tile's first row in its matrix if constexpr (grouped) { @@ -240,12 +272,19 @@ static __device__ __forceinline__ void mmvq_pq2_mma_body( map = &grp->tmap[g]; row0 = (tile - pq2_mma_group_begin(grp, g)) * 16; } - const int c0 = (i % nkb) * (row_bytes / 8); // 8-byte elements - char * slot = ring + (size_t) s * nmat * box_bytes; - pq2_mbar_arrive_expect_tx(&full[s], nmat * box_bytes); - pq2_tma_load_2d(slot, map, c0, row0, &full[s], policy); - if constexpr (nmat == 2) { - pq2_tma_load_2d(slot + box_bytes, tmap_gate, c0, row0, &full[s], policy); + for (int kb = 0; kb < nkb; ++kb, ++i) { + const int s = i % nslots; + if (i >= nslots) { + pq2_mbar_wait(&empty[s], (uint32_t) ((i / nslots - 1) & 1)); + } + held[s] = tile; + const int c0 = kb * (row_bytes / 8); // 8-byte elements + char * slot = ring + (size_t) s * nmat * box_bytes; + pq2_mbar_arrive_expect_tx(&full[s], nmat * box_bytes); + pq2_tma_load_2d(slot, map, c0, row0, &full[s], policy); + if constexpr (nmat == 2) { + pq2_tma_load_2d(slot + box_bytes, tmap_gate, c0, row0, &full[s], policy); + } } } } @@ -295,7 +334,16 @@ static __device__ __forceinline__ void mmvq_pq2_mma_body( const int r_out = j_out & 3; int i = 0; // the block's box sequence - for (int j = 0, tile = (int) blockIdx.x; j < n_mine; ++j, tile += (int) gridDim.x) { + for (int j = 0;; ++j) { + if (nkb > 1) { + load_b(0); // the first box's fragments, under the wait for it + } + // the tile the producer put in the next slot, or the end + pq2_mbar_wait(&full[i % nslots], (uint32_t) ((i / nslots) & 1)); + const int tile = held[i % nslots]; + if (tile < 0) { + break; + } // the tile's matrix: its output, rows and output stride, and the tile's first row in it float * dst_t = dst; int nrows_t = nrows; @@ -320,11 +368,11 @@ static __device__ __forceinline__ void mmvq_pq2_mma_body( float acc[nmat][4] = {{0.0f}}; for (int kb = 0; kb < nkb; ++kb, ++i) { - if (nkb > 1) { + if (nkb > 1 && kb > 0) { load_b(kb); // before the wait, under the TMA's latency } const int s = i % nslots; - pq2_mbar_wait(&full[s], (uint32_t) ((i / nslots) & 1)); + pq2_mbar_wait(&full[s], (uint32_t) ((i / nslots) & 1)); // the first box's: complete already const char * sw = ring + (size_t) s * nmat * box_bytes; const int valid = nb - (kb*8 + warp)*M; // this warp's blocks inside the row @@ -422,7 +470,7 @@ static __device__ __forceinline__ void mmvq_pq2_mma_body( } #else GGML_UNUSED_VARS(tmap, tmap_gate, grp, y, x_bias, dst, nrows, ncols, nb, n_tiles, nkb, nslots, evict_first, - stride_col_y, stride_col_dst, next); + stride_col_y, stride_col_dst, next, tile_ctr); NO_DEVICE_CODE; #endif // PQ2_MMA_AVAILABLE } @@ -433,18 +481,19 @@ static __global__ void mmvq_pq2_mma( const __grid_constant__ CUtensorMap tmap, const __grid_constant__ CUtensorMap tmap_gate, const block_q8_1 * y, const float * x_bias, float * dst, const int nrows, const int ncols, const int nb, const int n_tiles, const int nkb, const int nslots, const int evict_first, const int stride_col_y, const int stride_col_dst, - const ggml_cuda_pq2_prefetch next) { + const ggml_cuda_pq2_prefetch next, int * tile_ctr) { mmvq_pq2_mma_body(&tmap, &tmap_gate, nullptr, y, x_bias, dst, nrows, ncols, nb, n_tiles, - nkb, nslots, evict_first, stride_col_y, stride_col_dst, next); + nkb, nslots, evict_first, stride_col_y, stride_col_dst, next, tile_ctr); } template __launch_bounds__((PQ2_MMA_NW + 1)*32, 1) static __global__ void mmvq_pq2_mma_group( const __grid_constant__ pq2_mma_group grp, const block_q8_1 * y, const int ncols, const int nb, const int n_tiles, - const int nkb, const int nslots, const int evict_first, const int stride_col_y, const ggml_cuda_pq2_prefetch next) { + const int nkb, const int nslots, const int evict_first, const int stride_col_y, const ggml_cuda_pq2_prefetch next, + int * tile_ctr) { mmvq_pq2_mma_body(nullptr, nullptr, &grp, y, nullptr, nullptr, 0, ncols, nb, n_tiles, nkb, nslots, - evict_first, stride_col_y, 0, next); + evict_first, stride_col_y, 0, next, tile_ctr); } // --------------------------------------------------------------------------------------------------------------------- @@ -480,6 +529,13 @@ static bool pq2_mma_legacy() { return legacy; } +// the stream's tile counter, or nullptr: GGML_CUDA_PQ2_MMA_TILES_LEGACY=1 gives every tile to its block (blockIdx.x, then +// every gridDim.x-th) +static int * pq2_mma_tile_ctr(int * tile_ctr) { + static const bool legacy = ggml_env_switch("GGML_CUDA_PQ2_MMA_TILES_LEGACY"); + return legacy ? nullptr : tile_ctr; +} + // a row's boxes: m (1,024-weight units per box, at most 7) and nkb (boxes per row), the fewest boxes whose slots fit twice // in the shared memory; nslots, as many as fit up to the configured count struct pq2_mma_plan { @@ -495,7 +551,7 @@ static pq2_mma_plan pq2_mma_make_plan(int64_t ncols_x, int nmat) { if (m > PQ2_MMA_MAX_M) { continue; } - const size_t per_slot = (size_t) nmat * pq2_mma_box_bytes(m) + 2 * sizeof(uint64_t); + const size_t per_slot = pq2_mma_slot_bytes(m, nmat); const int fit = (int) ((PQ2_MMA_SMEM_MAX - pq2_mma_red_bytes(nmat)) / per_slot); if (fit >= 2) { return { m, nkb, std::min(fit, pq2_mma_get_config().nslots) }; @@ -543,7 +599,7 @@ static CUtensorMap pq2_mma_tensor_map(const void * vx, int64_t row_bytes, int64_ template static void pq2_mma_launch(const pq2_mma_plan & p, const CUtensorMap & tmap, const CUtensorMap & tmap_gate, const block_q8_1 * y, const float * x_bias, float * dst, int nrows, int ncols, int nb, int n_tiles, - int stride_col_y, int stride_col_dst, const ggml_cuda_pq2_prefetch & next, cudaStream_t stream) { + int stride_col_y, int stride_col_dst, const ggml_cuda_pq2_prefetch & next, int * tile_ctr, cudaStream_t stream) { static_assert(pq2_mma_smem_bytes(M, nmat, 2) <= PQ2_MMA_SMEM_MAX || nmat == 2, "two slots of a plain box fit"); const int nsm = ggml_cuda_info().devices[ggml_cuda_get_device()].nsm; const int nblocks = std::min(nsm, n_tiles); @@ -551,16 +607,16 @@ static void pq2_mma_launch(const pq2_mma_plan & p, const CUtensorMap & tmap, con pq2_mma_smem_bytes(M, nmat, p.nslots), stream); CUDA_SET_SHARED_MEMORY_LIMIT((mmvq_pq2_mma), PQ2_MMA_SMEM_MAX); // every plan's size, once ggml_cuda_kernel_launch(mmvq_pq2_mma, params, tmap, tmap_gate, y, x_bias, dst, nrows, ncols, nb, - n_tiles, p.nkb, p.nslots, (int) pq2_mma_get_config().evict_first, stride_col_y, stride_col_dst, next); + n_tiles, p.nkb, p.nslots, (int) pq2_mma_get_config().evict_first, stride_col_y, stride_col_dst, next, tile_ctr); } template static void pq2_mma_launch_m(const pq2_mma_plan & p, const CUtensorMap & tmap, const CUtensorMap & tmap_gate, const block_q8_1 * y, const float * x_bias, float * dst, int nrows, int ncols, int nb, int n_tiles, - int stride_col_y, int stride_col_dst, const ggml_cuda_pq2_prefetch & next, cudaStream_t stream) { + int stride_col_y, int stride_col_dst, const ggml_cuda_pq2_prefetch & next, int * tile_ctr, cudaStream_t stream) { switch (p.m) { #define PQ2_MMA_CASE(M) case M: pq2_mma_launch(p, tmap, tmap_gate, y, x_bias, dst, nrows, ncols, nb, \ - n_tiles, stride_col_y, stride_col_dst, next, stream); break; + n_tiles, stride_col_y, stride_col_dst, next, tile_ctr, stream); break; PQ2_MMA_CASE(1) PQ2_MMA_CASE(2) PQ2_MMA_CASE(3) @@ -576,7 +632,7 @@ static void pq2_mma_launch_m(const pq2_mma_plan & p, const CUtensorMap & tmap, c void ggml_cuda_mmvq_pq2_mma(const void * vx, const void * vgate, const void * vy, const float * x_bias, float * dst, int64_t ncols_x, int64_t nrows_x, int64_t ncols_dst, int64_t stride_row_x, int64_t stride_col_y, int64_t stride_col_dst, const ggml_cuda_pq2_prefetch & next, - cudaStream_t stream) { + int * tile_ctr, cudaStream_t stream) { const int nmat = vgate != nullptr ? 2 : 1; const pq2_mma_plan p = pq2_mma_make_plan(ncols_x, nmat); GGML_ASSERT(p.m > 0 && "ggml_cuda_mmvq_pq2_mma_usable holds a plan"); @@ -592,32 +648,32 @@ void ggml_cuda_mmvq_pq2_mma(const void * vx, const void * vgate, const void * vy if (nmat == 2) { GGML_ASSERT(x_bias == nullptr); pq2_mma_launch_m<2, false>(p, tmap, tmap_gate, y, nullptr, dst, (int) nrows_x, (int) ncols_dst, nb, n_tiles, - (int) stride_col_y, (int) stride_col_dst, next, stream); + (int) stride_col_y, (int) stride_col_dst, next, pq2_mma_tile_ctr(tile_ctr), stream); } else if (x_bias != nullptr) { pq2_mma_launch_m<1, true>(p, tmap, tmap_gate, y, x_bias, dst, (int) nrows_x, (int) ncols_dst, nb, n_tiles, - (int) stride_col_y, (int) stride_col_dst, next, stream); + (int) stride_col_y, (int) stride_col_dst, next, pq2_mma_tile_ctr(tile_ctr), stream); } else { pq2_mma_launch_m<1, false>(p, tmap, tmap_gate, y, nullptr, dst, (int) nrows_x, (int) ncols_dst, nb, n_tiles, - (int) stride_col_y, (int) stride_col_dst, next, stream); + (int) stride_col_y, (int) stride_col_dst, next, pq2_mma_tile_ctr(tile_ctr), stream); } } template static void pq2_mma_launch_group(const pq2_mma_plan & p, const pq2_mma_group & grp, const block_q8_1 * y, int ncols, - int nb, int n_tiles, int stride_col_y, const ggml_cuda_pq2_prefetch & next, cudaStream_t stream) { + int nb, int n_tiles, int stride_col_y, const ggml_cuda_pq2_prefetch & next, int * tile_ctr, cudaStream_t stream) { const int nsm = ggml_cuda_info().devices[ggml_cuda_get_device()].nsm; const int nblocks = std::min(nsm, n_tiles); const ggml_cuda_kernel_launch_params params(dim3(nblocks), dim3((PQ2_MMA_NW + 1)*32), pq2_mma_smem_bytes(M, 1, p.nslots), stream); CUDA_SET_SHARED_MEMORY_LIMIT((mmvq_pq2_mma_group), PQ2_MMA_SMEM_MAX); ggml_cuda_kernel_launch(mmvq_pq2_mma_group, params, grp, y, ncols, nb, n_tiles, p.nkb, p.nslots, - (int) pq2_mma_get_config().evict_first, stride_col_y, next); + (int) pq2_mma_get_config().evict_first, stride_col_y, next, tile_ctr); } void ggml_cuda_mmvq_pq2_mma_group(int n, const void * const * vx, float * const * dst, const int64_t * nrows_x, const int64_t * stride_row_x, const int64_t * stride_col_dst, const void * vy, int64_t ncols_x, int64_t ncols_dst, int64_t stride_col_y, - const ggml_cuda_pq2_prefetch & next, cudaStream_t stream) { + const ggml_cuda_pq2_prefetch & next, int * tile_ctr, cudaStream_t stream) { GGML_ASSERT(n >= 2 && n <= PQ2_MMA_MAX_GROUP); const pq2_mma_plan p = pq2_mma_make_plan(ncols_x, 1); GGML_ASSERT(p.m > 0 && "ggml_cuda_mmvq_pq2_mma_usable holds a plan"); @@ -641,7 +697,7 @@ void ggml_cuda_mmvq_pq2_mma_group(int n, const void * const * vx, float * const const block_q8_1 * y = (const block_q8_1 *) vy; switch (p.m) { #define PQ2_MMA_CASE(M) case M: pq2_mma_launch_group(p, grp, y, (int) ncols_dst, nb, (int) n_tiles, (int) stride_col_y, \ - next, stream); break; + next, pq2_mma_tile_ctr(tile_ctr), stream); break; PQ2_MMA_CASE(1) PQ2_MMA_CASE(2) PQ2_MMA_CASE(3) diff --git a/ggml/src/ggml-cuda/mmvq-pq2-mma.cuh b/ggml/src/ggml-cuda/mmvq-pq2-mma.cuh index d20e65c7cffa..b0a29cee283e 100644 --- a/ggml/src/ggml-cuda/mmvq-pq2-mma.cuh +++ b/ggml/src/ggml-cuda/mmvq-pq2-mma.cuh @@ -10,15 +10,17 @@ // Whether the kernel serves this matmul, and then it always does: GeForce / RTX PRO Blackwell (sm_120, not yet measured // on GB10's sm_121) built for it, 1-8 columns, K a multiple of 1024 (every row a whole number of 16-byte TMA units), the -// weights 16-byte aligned. It holds no state between launches. GGML_CUDA_PQ2_MMA_LEGACY=1 turns it off. +// weights 16-byte aligned. Its only state between launches is its stream's tile counter, 0 then (tile_ctr below). +// GGML_CUDA_PQ2_MMA_LEGACY=1 turns it off. bool ggml_cuda_mmvq_pq2_mma_usable(int cc, const void * vx, const void * vgate, int64_t ncols_x, int64_t nrows_x, int64_t stride_row_x, int64_t ncols_dst); -// next: the head of the launch after this one, which this one prefetches into L2 (ggml_cuda_pq2_prefetch, common.cuh) +// next: the head of the launch after this one, which this one prefetches into L2 (ggml_cuda_pq2_prefetch, common.cuh). +// tile_ctr: the stream's tile counter (ggml_cuda_pq2_tile_counters, common.cuh), or nullptr: every tile to its own block. void ggml_cuda_mmvq_pq2_mma(const void * vx, const void * vgate, const void * vy, const float * x_bias, float * dst, int64_t ncols_x, int64_t nrows_x, int64_t ncols_dst, int64_t stride_row_x, int64_t stride_col_y, int64_t stride_col_dst, const ggml_cuda_pq2_prefetch & next, - cudaStream_t stream); + int * tile_ctr, cudaStream_t stream); #define PQ2_MMA_MAX_GROUP 4 @@ -28,4 +30,4 @@ void ggml_cuda_mmvq_pq2_mma(const void * vx, const void * vgate, const void * vy void ggml_cuda_mmvq_pq2_mma_group(int n, const void * const * vx, float * const * dst, const int64_t * nrows_x, const int64_t * stride_row_x, const int64_t * stride_col_dst, const void * vy, int64_t ncols_x, int64_t ncols_dst, int64_t stride_col_y, - const ggml_cuda_pq2_prefetch & next, cudaStream_t stream); + const ggml_cuda_pq2_prefetch & next, int * tile_ctr, cudaStream_t stream); diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index e328115024fe..0992cbf23870 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -1480,7 +1480,7 @@ void ggml_cuda_mul_mat_vec_q( 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, stream); + dst_d, ne00, ne01, ne1, s01, s11, s1, ctx.pq2_next, ctx.pq2_tile_counter(), stream); return; } @@ -1528,7 +1528,7 @@ void ggml_cuda_mul_mat_vec_q_pq2_group(ggml_backend_cuda_context & ctx, ggml_ten } ggml_cuda_mmvq_pq2_mma_group(n, vx, dst, nrows, stride_row, stride_col_dst, src1_q8_1.get(), ne10, ne11, - ne10_padded / QK8_1, ctx.pq2_next, stream); + ne10_padded / QK8_1, ctx.pq2_next, ctx.pq2_tile_counter(), stream); } void ggml_cuda_op_mul_mat_vec_q( From a28b76874fcd6b54a18e10b2e10757da1fef06b5 Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Sun, 27 Sep 2026 19:18:43 -0500 Subject: [PATCH 02/10] docs(torad): the row for the PQ2_0 launches' tiles from a counter (8c78ba9ca) --- TORAD.md | 1 + 1 file changed, 1 insertion(+) diff --git a/TORAD.md b/TORAD.md index 9bab1a18ec7d..d37785d14dbf 100644 --- a/TORAD.md +++ b/TORAD.md @@ -99,6 +99,7 @@ which pins a commit of this branch as a submodule. | `32e695ecf` | the live-tile split launched every block the SMs hold, where the plain stream-k split rounds down to a multiple of the output tiles. 01f4f0fda's 3 blocks an SM made a decode token's split 510 blocks on an RTX 5090 (127.5 per KV head) and 210 on an RTX 5070 Ti (52.5): each head's blocks started half a block's steps from its neighbour's. A cell's K and V rows hold the heads side by side (q4_0, head 256: 144 bytes each) and the L2 fetches 64-byte units, so neighbouring heads share one, fetched once only when they read the row together: 12 units a row where 9 hold it. ncu on the 5070 Ti at 131,072 cells: DRAM read 202-207 MB against 151.0 MB of K and V; rounded down to a multiple of the Q tile's output tiles (208 blocks), 152-153 MB. test-backend-ops perf, a token, the same library with and without the switch: 5090 -15.7 / -7.8 / -17.0 / -20.9 %, 5070 Ti -2.1 / -18.3 / -21.4 / -23.1 % at 16,384 / 65,536 / 131,072 / 245,760 cells; on the 5090 131,072 and 245,760 cells take 115.72 and 196.44 us, where 4104c47 (340 blocks) took 123.69 and 204.66 and 4b61c54 145.50 and 255.48. A 4-row verify (2 blocks an SM: 340 and 140 blocks) does not move; the 5080's 252 is a multiple already. The served decode on the 5090 at 245K tokens (the gate's pool, 4 x 294,912 with `--kv-unified`, its four questions, 3 legs each): plain 98.75 -> 107.92 tok/s (+9.3 %), drafted 205.43 -> 208.90 (+1.7 %), the texts first differing 35-132 tokens in: the split sums in another order. Held to engine-10's G1 bars at that shape with top-5 log-probabilities, 384 tokens: plain 4 first differences, each at a tie, |dlogprob| p99 0.095 and max 0.128; drafted 4, each at a tie, p99 0.070, max 0.078. The gate's one-slot legs prefill through the live-tile path (the mask-prefix hint covers at most 16 rows), so their bits move too; their decode takes the mask-prefix path, whose split was rounded already, and its rate does not move (101.5 against 100.7 tok/s at 245K). FLASH_ATTN_EXT 3,220/3,220 on the 5070 Ti and the 5080 | `GGML_CUDA_FATTN_LIVE_BLOCKS_ALIGN_LEGACY=1` | | `e75e318e5` | `llama-bench` puts an NVTX range `gen` (domain `llama-bench`, a registered string) around each repetition's generation, so `nsys profile -c nvtx -p gen@llama-bench --capture-range-end=stop` records the tokens and none of the depth's prefill before them; a build without the NVTX headers compiles it out. rig's roofline probe traced the whole run, and at a depth of 245,760 on an RTX 5080 (Ternary Bonsai 2 27B, graphs off, `GGML_CUDA_NVTX=1`) the prefill's ~1.19M kernels kept nsys past its 30-minute limit; started at the range the capture takes 179 s, a 2.3 MB report of the 16 tokens' 18,816 kernels. At 16,384 it reads what the full capture reads: kernels 9,553.6 against 9,559.8 us a token, the same 23 groups, the largest group apart by 6.3 us | (a profiler that does not ask for the range ignores it) | | `fb03d294b` | `rms_norm_fwht_cuda` and `fwht_cuda_block` request their weights (the norm's `w` and the signs, which no kernel writes) before the PDL dependency wait, as the PQ2_0 matmul does its own. The matmul after a rotation starts under it (PDL), and its ring fill and the L2 prefetch of its head (`16ad036c4`) take DRAM for the microseconds the rotation runs; read after the wait, the rotation's weights (DRAM misses each token) queued behind them. RTX 5080, Ternary Bonsai 2 27B, one build both ways: nsys over 15 tokens (graphs off, the capture started at `e75e318e5`'s range), the FFN-input norm after an out-projection ends 6.1-6.2 -> 4.7-4.8 us past it at 16,384 and 5.5-6.1 -> 4.5-4.8 at 245,760, the layer-input norm 2.1 -> 1.7 us; llama-bench with graphs on, legs N-L-L-N: tg128 at 16,384 +0.81 % (105.48 / 104.21 / 104.40 / 104.82 tok/s), tg64 at 245,760 +1.63 % (69.35 / 67.93 / 67.48 / 68.27), all four pairs faster; `test-backend-ops MUL_MAT_HADAMARD` 42/42 both ways; llama-server at one slot and at 4 slots with `--kv-unified`, 128 greedy tokens with top-5 log-probabilities: bit-identical to the published engine-32e695e both ways | `GGML_CUDA_FWHT_PREWAIT_LEGACY=1` | +| `8c78ba9ca` | The PQ2_0 tensor-core launches (`mmvq-pq2-mma.cu`) keep a block's first tiles, as many as its ring holds before the dependency wait, and take the rest from a per-stream counter after the wait (`ggml_cuda_pq2_tile_counters`): each block's last ticket is past the tiles, and the block that takes the launch's last sets the counter back to 0. Owned tiles left a launch's last microseconds to the blocks the DRAM served later; now the blocks end within a tile of each other. RTX 5080, Ternary Bonsai 2 27B, one build both ways, graphs on, legs N-L-L-N: tg128 +1.46 % at 0 and +1.34 % at 16,384, tg64 +1.47 % at 245,760, every pair faster; pp4 `-rs 3` (the MTP verify), medians of 30 samples, +0.62 % and +0.81 %; llama-server with MTP drafts of 3, 12 greedy agentic requests: the same text, +1.14 % (same-text geomean +1.8 %) against the undisturbed legacy leg; `test-backend-ops` MUL_MAT (pq2_0) 99/99, MUL_MAT_GROUP 24/24, MUL_MAT_VEC_FUSION 1010/1010 both ways, and without the counter's reset 29, 7 and 26 of them fail; llama-server at one slot and at 4 with `--kv-unified`: bit-identical to the published engine-32e695e both ways | `GGML_CUDA_PQ2_MMA_TILES_LEGACY=1` | 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 From 833d067e0bb3b25220f7c47137c9314e3b15e549 Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Sun, 27 Sep 2026 23:14:55 -0500 Subject: [PATCH 03/10] feat(tensor): glm5next splits FFN/MoE, mirrors attention and caches Stage 1 of GLM-5.3 tensor parallelism: the split callback mirrors every KDA and DSA weight, all ssm and indexer tensors and all four caches, and keeps the existing splits for dense experts, routed experts, shared experts (as DeepSeek4) and the head, with one all-reduce per layer past the expert adds. glm5next leaves the -sm tensor refusal list. Mirrored mapping also dodges the ssm_out suffix assert and the ssm_d_state granularity divide for this arch. test-llama-archs -a glm5next, local 5080 + 5070 Ti: Meta row OK at 2.82e-06 NMSE vs CPU (1e-4 bar). Mutation proof: ffn_down_exps mirrored while up stays split aborts the Meta row in the matmul split-state handler, so the green row really splits. Per-card weights from the GGUF tensor table: 66.15 GiB plus mirrored latent and indexer caches. --- src/llama-arch.cpp | 1 - src/llama-model.cpp | 41 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index e0a21debf4e8..6ce183fe11ee 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1104,7 +1104,6 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_MISTRAL4: case LLM_ARCH_KIMI_LINEAR: - case LLM_ARCH_GLM5NEXT: case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_KIMI_K3: case LLM_ARCH_QWEN3TTS: diff --git a/src/llama-model.cpp b/src/llama-model.cpp index abfa4c8c0a06..d72b28fa7d36 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -389,6 +389,17 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_attn_q_b_weight ("blk\\.\\d*\\.attn_q_b\\.weight"); static const std::regex pattern_attn_gate_weight("blk\\.\\d*\\.attn_gate.weight"); + // glm5next low-rank attention, KDA state and indexer tensors (all mirrored in stage 1) + static const std::regex pattern_glm_q_a ("blk\\.\\d*\\.attn_q_a\\.weight"); + static const std::regex pattern_glm_q_a_norm ("blk\\.\\d*\\.attn_q_a_norm\\.weight"); + static const std::regex pattern_glm_kv_a_mqa ("blk\\.\\d*\\.attn_kv_a_mqa\\.weight"); + static const std::regex pattern_glm_kv_a_norm ("blk\\.\\d*\\.attn_kv_a_norm\\.weight"); + static const std::regex pattern_glm_k_b ("blk\\.\\d*\\.attn_k_b\\.weight"); + static const std::regex pattern_glm_v_b ("blk\\.\\d*\\.attn_v_b\\.weight"); + static const std::regex pattern_glm_ssm ("blk\\.\\d*\\.ssm_(f_a|f_b|g_a|g_b|norm)\\.weight"); + static const std::regex pattern_glm_ssm_conv ("blk\\.\\d*\\.ssm_conv1d_[qkv]\\.weight"); + static const std::regex pattern_glm_indexer ("blk\\.\\d*\\.indexer.*\\.(weight|bias)"); + static const std::regex pattern_ssm_dt ("blk\\.\\d*\\.ssm_dt.bias"); static const std::regex pattern_ssm_a ("blk\\.\\d*\\.ssm_a"); static const std::regex pattern_ssm_alpha ("blk\\.\\d*\\.ssm_alpha.weight"); @@ -492,6 +503,36 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } } + const bool is_glm5next = ud->model->arch == LLM_ARCH_GLM5NEXT; + if (is_glm5next) { + // Stage 1: split only the FFN/MoE; every KDA and DSA weight and every cache stays + // mirrored, so the meta backend only ever splits dense experts, shared experts and + // the head, with one all-reduce per layer past the expert adds. + if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_kv_weight) || + std::regex_match(tensor_name, pattern_attn_out_weight) || + std::regex_match(tensor_name, pattern_glm_q_a) || std::regex_match(tensor_name, pattern_glm_q_a_norm) || + std::regex_match(tensor_name, pattern_attn_q_b_weight) || + std::regex_match(tensor_name, pattern_glm_kv_a_mqa) || std::regex_match(tensor_name, pattern_glm_kv_a_norm) || + std::regex_match(tensor_name, pattern_glm_k_b) || std::regex_match(tensor_name, pattern_glm_v_b) || + std::regex_match(tensor_name, pattern_ssm_dt) || std::regex_match(tensor_name, pattern_ssm_a) || + std::regex_match(tensor_name, pattern_ssm_alpha) || std::regex_match(tensor_name, pattern_ssm_beta) || + std::regex_match(tensor_name, pattern_ssm_beta_alpha) || + std::regex_match(tensor_name, pattern_glm_ssm) || std::regex_match(tensor_name, pattern_glm_ssm_conv) || + std::regex_match(tensor_name, pattern_glm_indexer) || + std::regex_match(tensor_name, pattern_kv_cache) || + std::regex_match(tensor_name, pattern_r_cache) || std::regex_match(tensor_name, pattern_s_cache)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + // Shared experts split the way DeepSeek4's do. + if (std::regex_match(tensor_name, pattern_ffn_up_shexp_weight) || + std::regex_match(tensor_name, pattern_ffn_gate_shexp_weight)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "ffn_down_shexp.weight"); + } + if (std::regex_match(tensor_name, pattern_ffn_down_shexp_weight)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_0, "ffn_down_shexp.weight"); + } + } + // standard attention if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_kv_weight)) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight"); From 39a37bebbcc4804777a18f78e1d8e76f9ef4a7cb Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Sun, 27 Sep 2026 23:15:45 -0500 Subject: [PATCH 04/10] docs(torad): the row for glm5next tensor-split stage 1 (833d067e0) --- TORAD.md | 1 + 1 file changed, 1 insertion(+) diff --git a/TORAD.md b/TORAD.md index d37785d14dbf..3b9edc81b09d 100644 --- a/TORAD.md +++ b/TORAD.md @@ -100,6 +100,7 @@ which pins a commit of this branch as a submodule. | `e75e318e5` | `llama-bench` puts an NVTX range `gen` (domain `llama-bench`, a registered string) around each repetition's generation, so `nsys profile -c nvtx -p gen@llama-bench --capture-range-end=stop` records the tokens and none of the depth's prefill before them; a build without the NVTX headers compiles it out. rig's roofline probe traced the whole run, and at a depth of 245,760 on an RTX 5080 (Ternary Bonsai 2 27B, graphs off, `GGML_CUDA_NVTX=1`) the prefill's ~1.19M kernels kept nsys past its 30-minute limit; started at the range the capture takes 179 s, a 2.3 MB report of the 16 tokens' 18,816 kernels. At 16,384 it reads what the full capture reads: kernels 9,553.6 against 9,559.8 us a token, the same 23 groups, the largest group apart by 6.3 us | (a profiler that does not ask for the range ignores it) | | `fb03d294b` | `rms_norm_fwht_cuda` and `fwht_cuda_block` request their weights (the norm's `w` and the signs, which no kernel writes) before the PDL dependency wait, as the PQ2_0 matmul does its own. The matmul after a rotation starts under it (PDL), and its ring fill and the L2 prefetch of its head (`16ad036c4`) take DRAM for the microseconds the rotation runs; read after the wait, the rotation's weights (DRAM misses each token) queued behind them. RTX 5080, Ternary Bonsai 2 27B, one build both ways: nsys over 15 tokens (graphs off, the capture started at `e75e318e5`'s range), the FFN-input norm after an out-projection ends 6.1-6.2 -> 4.7-4.8 us past it at 16,384 and 5.5-6.1 -> 4.5-4.8 at 245,760, the layer-input norm 2.1 -> 1.7 us; llama-bench with graphs on, legs N-L-L-N: tg128 at 16,384 +0.81 % (105.48 / 104.21 / 104.40 / 104.82 tok/s), tg64 at 245,760 +1.63 % (69.35 / 67.93 / 67.48 / 68.27), all four pairs faster; `test-backend-ops MUL_MAT_HADAMARD` 42/42 both ways; llama-server at one slot and at 4 slots with `--kv-unified`, 128 greedy tokens with top-5 log-probabilities: bit-identical to the published engine-32e695e both ways | `GGML_CUDA_FWHT_PREWAIT_LEGACY=1` | | `8c78ba9ca` | The PQ2_0 tensor-core launches (`mmvq-pq2-mma.cu`) keep a block's first tiles, as many as its ring holds before the dependency wait, and take the rest from a per-stream counter after the wait (`ggml_cuda_pq2_tile_counters`): each block's last ticket is past the tiles, and the block that takes the launch's last sets the counter back to 0. Owned tiles left a launch's last microseconds to the blocks the DRAM served later; now the blocks end within a tile of each other. RTX 5080, Ternary Bonsai 2 27B, one build both ways, graphs on, legs N-L-L-N: tg128 +1.46 % at 0 and +1.34 % at 16,384, tg64 +1.47 % at 245,760, every pair faster; pp4 `-rs 3` (the MTP verify), medians of 30 samples, +0.62 % and +0.81 %; llama-server with MTP drafts of 3, 12 greedy agentic requests: the same text, +1.14 % (same-text geomean +1.8 %) against the undisturbed legacy leg; `test-backend-ops` MUL_MAT (pq2_0) 99/99, MUL_MAT_GROUP 24/24, MUL_MAT_VEC_FUSION 1010/1010 both ways, and without the counter's reset 29, 7 and 26 of them fail; llama-server at one slot and at 4 with `--kv-unified`: bit-identical to the published engine-32e695e both ways | `GGML_CUDA_PQ2_MMA_TILES_LEGACY=1` | +| `833d067e0` | glm5next leaves the `-sm tensor` refusal list and splits dense experts, routed experts, shared experts (as DeepSeek4) and the head, while every KDA and DSA weight, all ssm and indexer tensors and all four caches stay mirrored: one all-reduce per layer past the expert adds. `test-llama-archs -a glm5next` on 5080 + 5070 Ti: Meta row OK at 2.82e-06 NMSE vs CPU (1e-4 bar). The ffn_down_exps-mirrored mutant aborts the Meta row in the matmul split-state handler, so the green row really splits. Per-card weights from the GGUF tensor table: 66.15 GiB plus mirrored latent and indexer caches. | — | 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 From 881c823c57f171da038d9fce49e8449e5903bebe Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Sun, 27 Sep 2026 23:37:26 -0500 Subject: [PATCH 05/10] perf(cuda): top-k by radix select for a large k, the GLM-5.3 indexer's 512 pools ggml_cuda_op_top_k took a k of 512 over a few thousand columns (the DSA indexer's pools, one launch a layer) through the tiled path, whose one block a tile runs k block-wide max reductions one after another: 194 us a call at 3,520 columns. topk_radix finds the k-th largest key in four 8-bit histogram passes (warp-aggregated shared atomics, one per digit a warp), then writes in column order every column above it and the lowest columns equal to it, so the set and its tie-break are the tiled path's. A row of up to 16,384 columns keeps its keys in registers; a wider one is cut into 8,192- column tiles whose blocks select their k in parallel (a key among the row's k best has fewer than k better keys, so its tile keeps it), and one block a row selects among the tiles' candidates. One wide row stays with CUB's device-wide top-k, which spreads it over the card. Taken for k >= 64 past 1,024 columns on NVIDIA from Volta; GGML_CUDA_TOPK_RADIX_LEGACY=1 keeps the old dispatch. RTX 5080, test-backend-ops perf, legs new-legacy-new-legacy, k 512: 3,520 columns 194.5 -> 6.1 us (1 and 3 rows), 8,192 x 3 199.3 -> 12.4, 32,768 x 3 58.3 -> 17.1, 109,020 x 3 66.3 -> 21.8, and a 512-row ubatch 13.8-18x at every width; one wide row unchanged (CUB both ways). TOP_K 523/523 (new cases at the indexer's widths with and without ties). Two mutants fail it: the last equal key never written (every radix case), and stage 2 writing a candidate's position for its column (the 19 two-stage cases, GLM's 109,020 x 3 among them, and no other). The tests also gain GLM-5.3-Flash perf shapes: its routed experts (288, 8 used, IQ3_XXS / IQ4_XS / Q8_0) at a decode, an MTP verify and a ubatch, and its dense Q8_0 projections. Seen on the legacy dispatch, not fixed here: one of five legacy perf runs aborted at ggml-cuda.cu:3058 (cudaGraphExecUpdate returning neither success nor an update failure) on the graph after CUB's per-row top-k at 151,936 x 16, k 64. The radix path does not take CUB for more than one row. --- ggml/src/ggml-cuda/top-k.cu | 263 ++++++++++++++++++++++++++++++++++++ tests/test-backend-ops.cpp | 35 +++++ 2 files changed, 298 insertions(+) diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index a75c6d540e0c..e73a9e352974 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -143,6 +143,264 @@ static bool ggml_cuda_top_k_tiled(ggml_cuda_pool & pool, const float * src, int return true; } +// Radix select for a large k: one block a row finds the k-th largest key in four passes of 8 bits (a histogram of +// the next digit among the keys that share the digits found so far), then one pass writes, in column order, every +// column above that key and the lowest columns equal to it. Its time does not grow with k, where the tiled path runs +// k block reductions one after another (a DSA indexer's 512 pools of thousands, one launch a layer). +#define TOPK_RADIX_BLOCK 1024 +#define TOPK_RADIX_MIN_K 64 +#define TOPK_RADIX_TILE 8192 // RTX 5080, k 512: 3 rows of 109,020 in 21.7 us (4,096: 26.3; 16,384: 27.3) + +static __device__ __forceinline__ uint32_t topk_radix_key(const float x) { + const uint32_t b = __float_as_uint(x); + return (b & 0x80000000u) ? ~b : (b | 0x80000000u); +} + +static __device__ __forceinline__ uint32_t topk_radix_warp_incl_scan(uint32_t v, const int lane) { +#pragma unroll + for (int o = 1; o < WARP_SIZE; o <<= 1) { + const uint32_t t = __shfl_up_sync(0xFFFFFFFFu, v, o); + if (lane >= o) { + v += t; + } + } + return v; +} + +// One block a segment of seg_len columns, nseg segments a row: a whole row (nseg 1), or a tile of a wide row whose k +// best, keys and columns, a second launch selects among (CAND: src_key/src_col are that first launch's out_key/out_col). +// A tile narrower than k pads its candidates with key 0, below every float's key, so they are never taken. +// ITEMS > 0: the segment's keys stay in registers (seg_len <= ITEMS*BLOCK); 0: every pass reads the segment again. +template +static __global__ void __launch_bounds__(TOPK_RADIX_BLOCK) topk_radix( + const float * src, const uint32_t * src_key, const int * src_col, int * out_col, uint32_t * out_key, + const int64_t row_stride, const int seg_len, const int nseg, const int ncols, const int k) { + constexpr int BLOCK = TOPK_RADIX_BLOCK; + constexpr int NWARPS = BLOCK / WARP_SIZE; + static_assert(NWARPS <= WARP_SIZE, "one warp scans the warps' sums"); + + __shared__ uint32_t hist[256]; + __shared__ uint32_t warp_sums[NWARPS]; + __shared__ uint32_t s_digit; + __shared__ uint32_t s_krem; + + const int tid = threadIdx.x; + const int lane = tid % WARP_SIZE; + const int warp = tid / WARP_SIZE; + + const int row = blockIdx.x / nseg; + const int col0 = (blockIdx.x % nseg) * seg_len; + const int n = min(seg_len, ncols - col0); + const int kn = min(k, n); + const int64_t base = row*row_stride + col0; + int * out = out_col + (size_t) blockIdx.x * k; + uint32_t * out_k = out_key ? out_key + (size_t) blockIdx.x * k : nullptr; + + auto load_key = [&](const int col) -> uint32_t { + if (col >= n) { + return 0; + } + if constexpr (CAND) { + return src_key[base + col]; + } else { + return topk_radix_key(src[base + col]); + } + }; + + const int n_iter = ITEMS > 0 ? ITEMS : (n + BLOCK - 1) / BLOCK; + + uint32_t keys[ITEMS > 0 ? ITEMS : 1]; + if constexpr (ITEMS > 0) { +#pragma unroll + for (int i = 0; i < ITEMS; ++i) { + keys[i] = load_key(i*BLOCK + tid); + } + } + + uint32_t prefix = 0; + uint32_t mask = 0; + uint32_t krem = kn; // how many of the keys matching prefix/mask are still to be taken + +#pragma unroll 1 + for (int shift = 24; shift >= 0; shift -= 8) { + for (int b = tid; b < 256; b += BLOCK) { + hist[b] = 0; + } + __syncthreads(); + +#pragma unroll + for (int i = 0; i < n_iter; ++i) { + const int col = i*BLOCK + tid; + uint32_t key; + if constexpr (ITEMS > 0) { + key = keys[i]; + } else { + key = load_key(col); + } + const bool in = col < n && (key & mask) == prefix; + const uint32_t digit = (key >> shift) & 0xFFu; +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + // keys close in value share their high digits: one atomic a digit a warp, not one a key + const uint32_t active = __ballot_sync(0xFFFFFFFFu, in); + if (in) { + const uint32_t peers = __match_any_sync(active, digit); + if (lane == __ffs(peers) - 1) { + atomicAdd(&hist[digit], (uint32_t) __popc(peers)); + } + } +#else + if (in) { + atomicAdd(&hist[digit], 1u); + } +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + } + __syncthreads(); + + // the digit d with (keys above d) < krem <= (keys at or above d), scanning the bins from the top + const int b = 255 - tid; + const uint32_t h = tid < 256 ? hist[b] : 0; + const uint32_t incl = topk_radix_warp_incl_scan(h, lane); + if (tid < 256 && lane == WARP_SIZE - 1) { + warp_sums[warp] = incl; + } + __syncthreads(); + if (tid < 256) { + uint32_t ge = incl; + for (int w = 0; w < warp; ++w) { + ge += warp_sums[w]; + } + const uint32_t gt = ge - h; + if (gt < krem && ge >= krem) { + s_digit = b; + s_krem = krem - gt; + } + } + __syncthreads(); + + prefix |= s_digit << shift; + mask |= 0xFFu << shift; + krem = s_krem; + } + + // prefix is now the kn-th largest key; krem of the keys equal to it are taken, the lowest columns first + const uint32_t n_gt = kn - krem; + uint32_t base_gt = 0; + uint32_t base_eq = 0; + + auto emit = [&](const uint32_t pos, const int col, const uint32_t key) { + if constexpr (CAND) { + out[pos] = src_col[base + col]; + } else { + out[pos] = col0 + col; + } + if (out_k) { + out_k[pos] = key; + } + }; + if (out_k) { + for (int pos = kn + tid; pos < k; pos += BLOCK) { + out[pos] = -1; + out_k[pos] = 0; + } + } + +#pragma unroll + for (int i = 0; i < n_iter; ++i) { + const int col = i*BLOCK + tid; + uint32_t key; + if constexpr (ITEMS > 0) { + key = keys[i]; + } else { + key = load_key(col); + } + const bool gt = col < n && key > prefix; + const bool eq = col < n && key == prefix; + + // a block's counts fit in 16 bits: the column's rank among the chunk's gt keys low, among its eq keys high + const uint32_t v = (uint32_t) gt | ((uint32_t) eq << 16); + const uint32_t incl = topk_radix_warp_incl_scan(v, lane); + if (lane == WARP_SIZE - 1) { + warp_sums[warp] = incl; + } + __syncthreads(); + if (warp == 0) { + const uint32_t w = topk_radix_warp_incl_scan(lane < NWARPS ? warp_sums[lane] : 0, lane); + if (lane < NWARPS) { + warp_sums[lane] = w; + } + } + __syncthreads(); + + const uint32_t excl = (warp > 0 ? warp_sums[warp - 1] : 0) + incl - v; + const uint32_t total = warp_sums[NWARPS - 1]; + if (gt) { + emit(base_gt + (excl & 0xFFFFu), col, key); + } + if (eq) { + const uint32_t r = base_eq + (excl >> 16); + if (r < krem) { + emit(n_gt + r, col, key); + } + } + base_gt += total & 0xFFFFu; + base_eq += total >> 16; + if (base_gt == n_gt && base_eq >= krem) { + break; // uniform: every thread read the same total + } + __syncthreads(); // warp_sums is written again by the next chunk + } +} + +template +static void topk_radix_launch(const int nblocks, const float * src, const uint32_t * src_key, const int * src_col, + int * out_col, uint32_t * out_key, const int64_t row_stride, const int seg_len, + const int nseg, const int ncols, const int k, cudaStream_t stream) { + if (seg_len <= 4*TOPK_RADIX_BLOCK) { + topk_radix<4, CAND><<>>(src, src_key, src_col, out_col, out_key, row_stride, seg_len, nseg, ncols, k); + } else if (seg_len <= 8*TOPK_RADIX_BLOCK) { + topk_radix<8, CAND><<>>(src, src_key, src_col, out_col, out_key, row_stride, seg_len, nseg, ncols, k); + } else if (seg_len <= 16*TOPK_RADIX_BLOCK) { + topk_radix<16, CAND><<>>(src, src_key, src_col, out_col, out_key, row_stride, seg_len, nseg, ncols, k); + } else { + topk_radix<0, CAND><<>>(src, src_key, src_col, out_col, out_key, row_stride, seg_len, nseg, ncols, k); + } +} + +static bool ggml_cuda_top_k_radix(ggml_cuda_pool & pool, const float * src, int * dst, const int ncols, const int nrows, + const int k, const int cc, cudaStream_t stream) { + // Narrow rows stay with the bitonic sort, and a small k with the tiled path, whose blocks split a row. + // GGML_CUDA_TOPK_RADIX_LEGACY=1: never taken. + static const bool legacy = ggml_env_switch("GGML_CUDA_TOPK_RADIX_LEGACY"); + if (legacy || ncols <= TOPK_CAND || k < TOPK_RADIX_MIN_K || !GGML_CUDA_CC_IS_NVIDIA(cc) || cc < GGML_CUDA_CC_VOLTA) { + return false; + } + if (ncols <= 16*TOPK_RADIX_BLOCK) { + topk_radix_launch(nrows, src, nullptr, nullptr, dst, nullptr, ncols, ncols, 1, ncols, k, stream); + return true; + } +#ifdef CUB_TOP_K_AVAILABLE + // one wide row: CUB's device-wide top-k spreads it over the card (RTX 5080, k 512: 12.7 us at 109,020 columns + // against the two stages' 22.4; at 3 rows its launch a row costs 68.4 against 21.7) + if (nrows == 1) { + return false; + } +#endif // CUB_TOP_K_AVAILABLE + // A row too wide for one block's registers is cut into tiles whose blocks select their k in parallel: a key among + // the row's k best has fewer than k better keys (by value, then lower column), so its tile keeps it too. + const int tile = TOPK_RADIX_TILE; + if (2*k > tile) { + topk_radix_launch(nrows, src, nullptr, nullptr, dst, nullptr, ncols, ncols, 1, ncols, k, stream); + return true; + } + const int ntiles = (ncols + tile - 1) / tile; + const int ncand = ntiles * k; + ggml_cuda_pool_alloc cand_key(pool, (size_t) nrows * ncand); + ggml_cuda_pool_alloc cand_col(pool, (size_t) nrows * ncand); + topk_radix_launch(nrows * ntiles, src, nullptr, nullptr, cand_col.get(), cand_key.get(), ncols, tile, ntiles, ncols, k, stream); + topk_radix_launch(nrows, nullptr, cand_key.get(), cand_col.get(), dst, nullptr, ncand, ncand, 1, ncand, k, stream); + return true; +} + void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const float * src0_d = (const float *) src0->data; @@ -159,6 +417,11 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int64_t k = dst->ne[0]; ggml_cuda_pool & pool = ctx.pool(); + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + if (ggml_cuda_top_k_radix(pool, src0_d, dst_d, ncols, nrows, k, cc, stream)) { + return; + } + if (ggml_cuda_top_k_tiled(pool, src0_d, dst_d, ncols, nrows, k, stream)) { return; } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index d9406f6abe64..042b1a289209 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10979,6 +10979,11 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {202048, nrows, 1, 1}, k, true)); } } + // the GLM-5.3 DSA indexer: 512 of a context's pools, at the radix path's three row widths + for (int64_t cols : {3520, 16384, 109020}) { + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, 3, 1, 1}, 512)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, 3, 1, 1}, 512, true)); + } for (int k : {1, 2, 3, 7, 15}) { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {16, 10, 10, 10}, k)); @@ -11913,6 +11918,22 @@ static std::vector> make_test_cases_perf() { } } + // GLM-5.3-Flash (glm5next): 288 routed experts, 8 used, gate/up 4096 -> 2048 and down 2048 -> 4096, IQ3_XXS as the + // public GGUF serves them (Q8_0 in its MTP layer, IQ4_XS the requantization candidate); a decode token, an MTP verify + // of 2 drafts and a ubatch. Then its dense Q8_0 projections: KDA q/k/v and output, MLA output. + for (int bs : {1, 3, 512}) { + for (ggml_type type_a : {GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ4_XS, GGML_TYPE_Q8_0}) { + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 288, 8, false, 2048, bs, 4096)); + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 288, 8, false, 4096, bs, 2048)); + test_cases.emplace_back(new test_mul_mat_id_fusion(type_a, GGML_TYPE_F32, 288, 8, false, 2048, bs, 4096, 1)); + } + } + for (int bs : {1, 3}) { + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 8192, bs, 4096, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 4096, bs, 8192, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 4096, bs, 16384, {1, 1}, {1, 1})); + } + for (int K : {3, 5}) { for (int IC : {256, 2560}) { for (int IW_IH : {32, 64, 256}) { @@ -12057,6 +12078,20 @@ static std::vector> make_test_cases_perf() { } } } + // the GLM-5.3 DSA indexer's 512 pools: a decode, an MTP verify and a ubatch, from 14k cells of context to 436k + for (int64_t cols : {3520, 8192, 32768, 109020}) { + for (int64_t nrows : {1, 3, 512}) { + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, 512)); + } + } + // around the radix path's k threshold + for (auto k : {32, 64, 128}) { + for (auto cols : {8192, 151936}) { + for (auto nrows : {1, 16}) { + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, k)); + } + } + } for (auto nrows : {1, 4, 8, 16}) { for (auto cols : {128, 1024, 4096, 8192, 16384, 32768, 65536, 131072, 200000, 2000000}) { From 3e76a0e77d17984667737f61b2f46bc8ffdde921 Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Sun, 27 Sep 2026 23:37:42 -0500 Subject: [PATCH 06/10] docs(torad): the row for the radix top-k (881c823c5) --- TORAD.md | 1 + 1 file changed, 1 insertion(+) diff --git a/TORAD.md b/TORAD.md index 3b9edc81b09d..2cc0eb422724 100644 --- a/TORAD.md +++ b/TORAD.md @@ -101,6 +101,7 @@ which pins a commit of this branch as a submodule. | `fb03d294b` | `rms_norm_fwht_cuda` and `fwht_cuda_block` request their weights (the norm's `w` and the signs, which no kernel writes) before the PDL dependency wait, as the PQ2_0 matmul does its own. The matmul after a rotation starts under it (PDL), and its ring fill and the L2 prefetch of its head (`16ad036c4`) take DRAM for the microseconds the rotation runs; read after the wait, the rotation's weights (DRAM misses each token) queued behind them. RTX 5080, Ternary Bonsai 2 27B, one build both ways: nsys over 15 tokens (graphs off, the capture started at `e75e318e5`'s range), the FFN-input norm after an out-projection ends 6.1-6.2 -> 4.7-4.8 us past it at 16,384 and 5.5-6.1 -> 4.5-4.8 at 245,760, the layer-input norm 2.1 -> 1.7 us; llama-bench with graphs on, legs N-L-L-N: tg128 at 16,384 +0.81 % (105.48 / 104.21 / 104.40 / 104.82 tok/s), tg64 at 245,760 +1.63 % (69.35 / 67.93 / 67.48 / 68.27), all four pairs faster; `test-backend-ops MUL_MAT_HADAMARD` 42/42 both ways; llama-server at one slot and at 4 slots with `--kv-unified`, 128 greedy tokens with top-5 log-probabilities: bit-identical to the published engine-32e695e both ways | `GGML_CUDA_FWHT_PREWAIT_LEGACY=1` | | `8c78ba9ca` | The PQ2_0 tensor-core launches (`mmvq-pq2-mma.cu`) keep a block's first tiles, as many as its ring holds before the dependency wait, and take the rest from a per-stream counter after the wait (`ggml_cuda_pq2_tile_counters`): each block's last ticket is past the tiles, and the block that takes the launch's last sets the counter back to 0. Owned tiles left a launch's last microseconds to the blocks the DRAM served later; now the blocks end within a tile of each other. RTX 5080, Ternary Bonsai 2 27B, one build both ways, graphs on, legs N-L-L-N: tg128 +1.46 % at 0 and +1.34 % at 16,384, tg64 +1.47 % at 245,760, every pair faster; pp4 `-rs 3` (the MTP verify), medians of 30 samples, +0.62 % and +0.81 %; llama-server with MTP drafts of 3, 12 greedy agentic requests: the same text, +1.14 % (same-text geomean +1.8 %) against the undisturbed legacy leg; `test-backend-ops` MUL_MAT (pq2_0) 99/99, MUL_MAT_GROUP 24/24, MUL_MAT_VEC_FUSION 1010/1010 both ways, and without the counter's reset 29, 7 and 26 of them fail; llama-server at one slot and at 4 with `--kv-unified`: bit-identical to the published engine-32e695e both ways | `GGML_CUDA_PQ2_MMA_TILES_LEGACY=1` | | `833d067e0` | glm5next leaves the `-sm tensor` refusal list and splits dense experts, routed experts, shared experts (as DeepSeek4) and the head, while every KDA and DSA weight, all ssm and indexer tensors and all four caches stay mirrored: one all-reduce per layer past the expert adds. `test-llama-archs -a glm5next` on 5080 + 5070 Ti: Meta row OK at 2.82e-06 NMSE vs CPU (1e-4 bar). The ffn_down_exps-mirrored mutant aborts the Meta row in the matmul split-state handler, so the green row really splits. Per-card weights from the GGUF tensor table: 66.15 GiB plus mirrored latent and indexer caches. | — | +| `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` | 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 From 33eb70bc01960c3eee53bd89267d8e2abd16aaff Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Sun, 27 Sep 2026 23:40:32 -0500 Subject: [PATCH 07/10] perf(graph): a sparse attention's top-k mask is not tagged as a prefix of the cells build_attn_mha tagged every one-sequence causal mask as a prefix, the hint under which the CUDA flash attention leaves out its range scan and live-tile skipping and applies the mask over the whole range (fattn-common.cuh, mask_prefix). The DSA layers (GLM-5.3's MLA with the lightning indexer, and the kpool path) pass a mask that keeps the indexer's top-k cells of the whole cache, so a decode read every cell's K and V to keep 2,048 of them. build_attn_mha takes mask_is_prefix, true by default; both sparse callers pass false, and the kernel finds their live tiles by scanning the mask. The dense layers keep the hint. LLAMA_ATTN_SPARSE_MASK_PREFIX_LEGACY=1 tags the sparse masks again. The result is the same attention by a different split of the cells (stream-k over the live steps), so its bits can move at float rounding. No other path changes: the no-hint route is what every masked decode without the hint already takes, and test-backend-ops' FLASH_ATTN_EXT covers the hint both ways at 576/512. Checked here only that the tiny glm5next model decodes at 6,000 cells both ways; the gain sits where the cache is long against the top-k (a 436k-cell GLM-5.3 decode), measured on the served model. --- src/llama-graph.cpp | 15 ++++++++++----- src/llama-graph.h | 3 ++- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index b3f8e93bbb49..805625477cf7 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2874,7 +2874,8 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * sinks, ggml_tensor * v_mla, float kq_scale, - int il) const { + int il, + bool mask_is_prefix) const { const bool v_trans = v->nb[1] > v->nb[2]; // split the batch into streams if needed @@ -2912,8 +2913,12 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_flash_attn_ext_add_sinks(cur, sinks); ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); - // one sequence, causal attention and no window: every row's mask is a prefix of the cells - ggml_flash_attn_ext_set_mask_prefix(cur, cparams.n_seq_max == 1 && cparams.causal_attn && il >= 0 && !hparams.is_swa(il)); + // one sequence, causal attention and no window: every row's mask is a prefix of the cells, unless the caller's + // mask selects cells (sparse attention's top-k), whose live cells the kernel then finds by scanning the mask. + // LLAMA_ATTN_SPARSE_MASK_PREFIX_LEGACY=1 tags the sparse masks too: the kernel then walks every cell + static const bool sparse_prefix_legacy = ggml_env_switch("LLAMA_ATTN_SPARSE_MASK_PREFIX_LEGACY"); + ggml_flash_attn_ext_set_mask_prefix(cur, (mask_is_prefix || sparse_prefix_legacy) && cparams.n_seq_max == 1 && + cparams.causal_attn && il >= 0 && !hparams.is_swa(il)); if (v_mla) { #if 0 @@ -3342,7 +3347,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, kq_scale, il, /*mask_is_prefix =*/ false); cb(cur, "kqv_out", il); if (wo) { @@ -4112,7 +4117,7 @@ ggml_tensor * llm_graph_context::build_attn_sparse( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); - ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, mask_top_k, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, mask_top_k, sinks, v_mla, kq_scale, il, /*mask_is_prefix =*/ false); cb(cur, "kqv_out", il); if (wo) { diff --git a/src/llama-graph.h b/src/llama-graph.h index e868929a58cb..18115fa1ae5a 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1274,7 +1274,8 @@ struct llm_graph_context { ggml_tensor * sinks, // [n_head_q] ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] float kq_scale, - int il) const; + int il, + bool mask_is_prefix = true) const; // false: kq_mask selects cells (sparse attention's top-k) llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const; From 7a7a0a238829d2596027d72940acb9550ace11eb Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Sun, 27 Sep 2026 23:40:42 -0500 Subject: [PATCH 08/10] docs(torad): the row for the sparse masks untagged as prefixes (33eb70bc0) --- TORAD.md | 1 + 1 file changed, 1 insertion(+) diff --git a/TORAD.md b/TORAD.md index 2cc0eb422724..67e551e50401 100644 --- a/TORAD.md +++ b/TORAD.md @@ -102,6 +102,7 @@ which pins a commit of this branch as a submodule. | `8c78ba9ca` | The PQ2_0 tensor-core launches (`mmvq-pq2-mma.cu`) keep a block's first tiles, as many as its ring holds before the dependency wait, and take the rest from a per-stream counter after the wait (`ggml_cuda_pq2_tile_counters`): each block's last ticket is past the tiles, and the block that takes the launch's last sets the counter back to 0. Owned tiles left a launch's last microseconds to the blocks the DRAM served later; now the blocks end within a tile of each other. RTX 5080, Ternary Bonsai 2 27B, one build both ways, graphs on, legs N-L-L-N: tg128 +1.46 % at 0 and +1.34 % at 16,384, tg64 +1.47 % at 245,760, every pair faster; pp4 `-rs 3` (the MTP verify), medians of 30 samples, +0.62 % and +0.81 %; llama-server with MTP drafts of 3, 12 greedy agentic requests: the same text, +1.14 % (same-text geomean +1.8 %) against the undisturbed legacy leg; `test-backend-ops` MUL_MAT (pq2_0) 99/99, MUL_MAT_GROUP 24/24, MUL_MAT_VEC_FUSION 1010/1010 both ways, and without the counter's reset 29, 7 and 26 of them fail; llama-server at one slot and at 4 with `--kv-unified`: bit-identical to the published engine-32e695e both ways | `GGML_CUDA_PQ2_MMA_TILES_LEGACY=1` | | `833d067e0` | glm5next leaves the `-sm tensor` refusal list and splits dense experts, routed experts, shared experts (as DeepSeek4) and the head, while every KDA and DSA weight, all ssm and indexer tensors and all four caches stay mirrored: one all-reduce per layer past the expert adds. `test-llama-archs -a glm5next` on 5080 + 5070 Ti: Meta row OK at 2.82e-06 NMSE vs CPU (1e-4 bar). The ffn_down_exps-mirrored mutant aborts the Meta row in the matmul split-state handler, so the green row really splits. Per-card weights from the GGUF tensor table: 66.15 GiB plus mirrored latent and indexer caches. | — | | `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` | 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 From 78e0fadf74fc1d75a29afaf5e2f5325cbeae6d90 Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Sun, 27 Sep 2026 23:56:47 -0500 Subject: [PATCH 09/10] perf(cuda): top-k takes a narrow row by radix at any k, not CUB one row after another A row of up to 1,024 columns has no tiled path, so with CUB's top-k available ggml_cuda_op_top_k looped over the rows, four launches a row: a 512-token ubatch through a DSA indexer under 4,096 cells (its pools at most 1,024) cost 2,048 launches a layer (seen under nsys on the tiny glm5next model's depth prefill). topk_radix takes those rows at any k, one block a row; a wide row with a small k keeps the tiled path. RTX 5080, test-backend-ops perf, legs new-legacy-new-legacy: a narrow row 3.3-3.9 us against 8.2-12.7 (1 row) and 129-171 (16 rows), every k from 1 to 400. TOP_K 523/523, its narrow cases (ties included) now through the radix kernel that the two mutants of 881c823c5 fail. The GLM-5.3 routed-expert perf cases build 32 experts of the model's 288: the case quantizes every expert on the CPU at setup, and all 288 took over ten minutes; 32 are still past a 5080's L2. --- ggml/src/ggml-cuda/top-k.cu | 7 ++++--- tests/test-backend-ops.cpp | 11 ++++++----- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index e73a9e352974..779ba8c5d3e6 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -368,10 +368,11 @@ static void topk_radix_launch(const int nblocks, const float * src, const uint32 static bool ggml_cuda_top_k_radix(ggml_cuda_pool & pool, const float * src, int * dst, const int ncols, const int nrows, const int k, const int cc, cudaStream_t stream) { - // Narrow rows stay with the bitonic sort, and a small k with the tiled path, whose blocks split a row. - // GGML_CUDA_TOPK_RADIX_LEGACY=1: never taken. + // A wide row with a small k stays with the tiled path, whose blocks split the row. A narrow row (no tiling) at any k + // is taken here: the other route is CUB's top-k one row after another, four launches a row (a 512-token ubatch of a + // DSA indexer under 4,096 cells: 2,048 launches a layer). GGML_CUDA_TOPK_RADIX_LEGACY=1: never taken. static const bool legacy = ggml_env_switch("GGML_CUDA_TOPK_RADIX_LEGACY"); - if (legacy || ncols <= TOPK_CAND || k < TOPK_RADIX_MIN_K || !GGML_CUDA_CC_IS_NVIDIA(cc) || cc < GGML_CUDA_CC_VOLTA) { + if (legacy || (ncols > TOPK_CAND && k < TOPK_RADIX_MIN_K) || !GGML_CUDA_CC_IS_NVIDIA(cc) || cc < GGML_CUDA_CC_VOLTA) { return false; } if (ncols <= 16*TOPK_RADIX_BLOCK) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 042b1a289209..95963b8d233d 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -11918,14 +11918,15 @@ static std::vector> make_test_cases_perf() { } } - // GLM-5.3-Flash (glm5next): 288 routed experts, 8 used, gate/up 4096 -> 2048 and down 2048 -> 4096, IQ3_XXS as the + // GLM-5.3-Flash (glm5next): routed experts, 8 used, gate/up 4096 -> 2048 and down 2048 -> 4096, IQ3_XXS as the // public GGUF serves them (Q8_0 in its MTP layer, IQ4_XS the requantization candidate); a decode token, an MTP verify - // of 2 drafts and a ubatch. Then its dense Q8_0 projections: KDA q/k/v and output, MLA output. + // of 2 drafts and a ubatch. 32 experts of its 288: the case quantizes every expert on the CPU at setup (all 288 took + // over 10 minutes), and 32 are still past a 5080's L2. Then its dense Q8_0 projections: KDA q/k/v and output, MLA output. for (int bs : {1, 3, 512}) { for (ggml_type type_a : {GGML_TYPE_IQ3_XXS, GGML_TYPE_IQ4_XS, GGML_TYPE_Q8_0}) { - test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 288, 8, false, 2048, bs, 4096)); - test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 288, 8, false, 4096, bs, 2048)); - test_cases.emplace_back(new test_mul_mat_id_fusion(type_a, GGML_TYPE_F32, 288, 8, false, 2048, bs, 4096, 1)); + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 32, 8, false, 2048, bs, 4096)); + test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 32, 8, false, 4096, bs, 2048)); + test_cases.emplace_back(new test_mul_mat_id_fusion(type_a, GGML_TYPE_F32, 32, 8, false, 2048, bs, 4096, 1)); } } for (int bs : {1, 3}) { From 41d8963c488238d9c66a63b8cf8146b3dc3cb148 Mon Sep 17 00:00:00 2001 From: Marcos Damasceno Date: Sun, 27 Sep 2026 23:56:47 -0500 Subject: [PATCH 10/10] docs(torad): the row for the narrow top-k rows by radix (78e0fadf7) --- TORAD.md | 1 + 1 file changed, 1 insertion(+) diff --git a/TORAD.md b/TORAD.md index 67e551e50401..05f1d8f59223 100644 --- a/TORAD.md +++ b/TORAD.md @@ -103,6 +103,7 @@ which pins a commit of this branch as a submodule. | `833d067e0` | glm5next leaves the `-sm tensor` refusal list and splits dense experts, routed experts, shared experts (as DeepSeek4) and the head, while every KDA and DSA weight, all ssm and indexer tensors and all four caches stay mirrored: one all-reduce per layer past the expert adds. `test-llama-archs -a glm5next` on 5080 + 5070 Ti: Meta row OK at 2.82e-06 NMSE vs CPU (1e-4 bar). The ffn_down_exps-mirrored mutant aborts the Meta row in the matmul split-state handler, so the green row really splits. Per-card weights from the GGUF tensor table: 66.15 GiB plus mirrored latent and indexer caches. | — | | `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` | 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