From c47d9d5895ae5250aab106848b4219feb501c912 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 26 Aug 2026 22:33:09 +0200 Subject: [PATCH 1/7] perf(ds4): tune gfx1151 verify kernels --- .../deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu | 19 +++ .../ggml/src/ggml-cuda/rocmfp2_mix.cu | 117 +++++++++++++----- .../ggml/src/ggml-cuda/rocmfp3_mix.cu | 94 +++++++++----- server/test/bench_rocmfp_mix_gateup_glu.cpp | 102 ++++++++++----- server/test/test_rocmfp_mix_gateup_glu.cpp | 39 +++++- 5 files changed, 281 insertions(+), 90 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu index 6e4db8e8b..6a50ad9c3 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu @@ -4,6 +4,8 @@ #include "mmvf.cuh" #include "convert.cuh" +#include + template static __global__ void mul_mat_vec_f( const T * __restrict__ x, const float * __restrict__ y, const int32_t * __restrict__ ids, const ggml_cuda_mm_fusion_args_device fusion, float * __restrict__ dst, @@ -834,6 +836,23 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 } return ne11 <= 8; } else if (GGML_CUDA_CC_IS_AMD(cc)) { + // DeepSeek V4's hyper-connection projection is a very thin + // [16384, 24] F16 matrix. Speculative verification normally + // reaches it with four columns. gfx1151 is classified as + // RDNA 3.5 below, where the generic crossover stops at three + // columns and sends q=4 to hipBLAS. On Strix Halo that tiny + // GEMM is launch/tiling bound; the row-split Wave32 MMVF path + // is the better fit. Keep this narrowly scoped until wider + // RDNA 3.5 shapes have their own measurements. + static const bool gfx1151_hc_q4_mmvf = [] { + const char * value = std::getenv( + "DFLASH_GFX1151_HC_MMVF_Q4"); + return !value || std::atoi(value) != 0; + }(); + if (gfx1151_hc_q4_mmvf && GGML_CUDA_CC_IS_RDNA3_5(cc) && + src0_ne[0] == 16384 && src0_ne[1] == 24 && ne11 == 4) { + return true; + } if (fp16_mma_hardware_available(cc)) { if (GGML_CUDA_CC_IS_RDNA3(cc)) { return ne11 <= 3; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu index c61f695bc..637f6d27f 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu @@ -8,6 +8,7 @@ // standalone swiglu_ds4 kernel applies, not a re-derivation of the formula. #include "unary.cuh" #include +#include #include #include #include @@ -56,6 +57,7 @@ struct MixEntry { const nv_bfloat16 * codebooks; // n_experts * 2 * 4 const uint8_t * modes; // n_experts int device; // device the side-data lives on; frees must happen in that context + bool gfx1151; // registration-time dispatch tuning; no hot-path device query }; // Dispatch wrappers keep this lock from lookup through kernel enqueue. It is // recursive because MMQ takes the public dispatch lock and then calls the @@ -83,6 +85,17 @@ static int mix_device_of(const void * p) { return dev; } +static bool mix_device_is_gfx1151(int device) { +#if defined(GGML_USE_HIP) + cudaDeviceProp prop{}; + return cudaGetDeviceProperties(&prop, device) == cudaSuccess && + std::strncmp(prop.gcnArchName, "gfx1151", 7) == 0; +#else + (void) device; + return false; +#endif +} + // RAII device switch: restores the previous device even on the early-return error paths. struct MixDeviceGuard { int prev = -1; @@ -187,7 +200,7 @@ void mix_register_impl(const void * base, size_t nb02, int n_experts, int out, i const size_t expert_bytes = (size_t) out * (size_t) (in / MIX_QK) * MIX_BLOCK_BYTES; MixEntry ne{base, nb02, expert_bytes, n_experts, out, in, codebooks, modes, - device}; + device, mix_device_is_gfx1151(device)}; for (auto & e : g_mix_registry) { if (e.base == base) { // update in place — free the old owned buffers first mix_free_entry_device(e); @@ -464,6 +477,22 @@ __device__ __forceinline__ void mix_block_accum( // block reads up to 6 B past the tensor end. That is now REJECTED AT REGISTRATION // (see the in % 128 guard in ggml_cuda_rocmfp2_mix_register_host) rather than // left to chance, so this loop can stay branch-free. +#if defined(__HIP_PLATFORM_AMD__) && defined(__gfx1151__) + // gfx1151 supports unaligned flat loads. Read the exact 10-byte payload as + // one 8-byte code word plus the 2-byte metadata tail. The old aligned-floor + // path fetched two overlapping 8-byte windows for every 10-byte block, so + // adjacent lanes requested 16 bytes for 10 bytes of useful weights. Keeping + // the packed code word in one register also preserves the existing decode + // and accumulation order exactly. + uint64_t codes; + uint16_t meta; + MIX_MEMCPY(&codes, b, sizeof(codes)); + MIX_MEMCPY(&meta, b + MIX_QS, sizeof(meta)); + const uint8_t m0 = (uint8_t) meta; + const uint8_t m1 = (uint8_t) (meta >> 8); +#else + // Portable fallback: load the two aligned 8-byte windows bracketing this + // 10-byte block, then extract the payload in registers. const uintptr_t addr = (uintptr_t) b; const uint8_t * base8 = (const uint8_t *) MIX_ASSUME_ALIGNED( (const void *) (addr & ~(uintptr_t) 7), 8); @@ -474,6 +503,7 @@ __device__ __forceinline__ void mix_block_accum( const uint64_t codes = (sh == 0) ? lo : ((lo >> sh) | (hi << (64 - sh))); const uint8_t m0 = (uint8_t) (hi >> sh); // block byte MIX_QS+0 const uint8_t m1 = (uint8_t) (hi >> (sh + 8)); // block byte MIX_QS+1 +#endif if (mode == 0) { const float s0 = mix_ue4m3(m0), s1 = mix_ue4m3(m1); #pragma unroll @@ -677,21 +707,27 @@ __global__ void mix_matvec_rocmfp2_slice_kernel( } } -// FUSE_GLU folds the SECOND mul_mat_id of a DeepSeek4 gate/up pair plus the SwiGLU into this -// launch. The unfused shape is two matvec launches writing two [n_ff_exp, n_used, ntok] -// intermediates, then a third kernel reading both back to apply the GLU. qtype 107 never paid -// that: ggml_cuda_try_fuse_mul_mat_glu collapses the trio into one mul_mat_vec_q, which is why -// the profile showed 107 at 15050 launches against 106's 30100 plus a 28 ms swiglu_ds4 pass. +// GLU_MODE folds a DeepSeek4 gate/up pair plus SwiGLU into the MoE matvec path: +// +// 0: ordinary one-tensor matvec; +// 1: one-pass gate+up matvec and SwiGLU (best for q <= 2); +// 2: gate matvec that reads a preceding up result from dst and overwrites it +// with SwiGLU(gate, up) (paired with mode 0 for wider verification). // -// TEMPLATED rather than a second kernel on purpose: both instantiations run the SAME +// The two-pass form keeps each kernel at the lower register footprint of a +// one-tensor matvec. On gfx1151 the one-pass kernel wins through q=2 by +// avoiding a launch, but loses once wider verification exposes its occupancy cost. +// Both forms remove the old third, standalone SwiGLU launch. +// +// TEMPLATED rather than copied kernels on purpose: all instantiations run the SAME // accumulation over the SAME fixed block order, so each dot product is bit-identical to what -// the unfused path computes, and the fused result is bit-identical to +// the unfused path computes, and either fused result is bit-identical to // swiglu_ds4(unfused_gate, unfused_up). A copy-pasted kernel would only *probably* stay that way. // // Naming follows the mmvq fusion convention: the PRIMARY tensor is `up` (src0 of the surviving // mul_mat_id) and `gate` arrives as the extra operand, because // ggml_cuda_op_swiglu_ds4_single(gate, up, limit) is not symmetric -- silu() is applied to gate. -template +template __global__ void mix_matvec_rocmfp2_moe_kernel( const uint8_t * __restrict__ data, size_t nb02, const nv_bfloat16 * __restrict__ codebooks, const uint8_t * __restrict__ modes, @@ -700,12 +736,14 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( int64_t ids_s0, int64_t ids_s1, // element strides (int32) over slot, token int64_t src1_s1, int64_t src1_s2, // element strides (float) over ne11, token int64_t dst_s1, int64_t dst_s2, // element strides (float) over slot, token - // FUSE_GLU only. The gate tensor's own registry entry -- separate codebooks and modes, + // GLU_MODE=1 only. The gate tensor's own registry entry -- separate codebooks and modes, // NOT assumed equal to up's. Producers may emit identical codebooks for the two // halves, but the kernel does not rely on that and staging both costs 32 B of LDS. const uint8_t * __restrict__ gdata, size_t gnb02, const nv_bfloat16 * __restrict__ gcodebooks, const uint8_t * __restrict__ gmodes, float glu_limit) { + constexpr bool DUAL_GLU = GLU_MODE == 1; + constexpr bool FINALIZE_GLU = GLU_MODE == 2; const int warps_per_block = blockDim.x / MIX_WARP; const int warp = blockIdx.x * warps_per_block + (threadIdx.x / MIX_WARP); const int row0 = warp * 2; // two output rows per warp @@ -727,12 +765,12 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( const bool bad_expert = expert < 0 || expert >= n_experts; // Both tables in one array so the staging stays a single guarded write per thread and one // barrier. Gate's table occupies [2*MIX_K, 4*MIX_K). - __shared__ float s_lut[FUSE_GLU ? 4 * MIX_K : 2 * MIX_K]; + __shared__ float s_lut[DUAL_GLU ? 4 * MIX_K : 2 * MIX_K]; if (!bad_expert && (int) threadIdx.x < 2 * MIX_K) { s_lut[threadIdx.x] = __bfloat162float(codebooks[(int64_t) expert * 2 * MIX_K + threadIdx.x]); } - if (FUSE_GLU) { + if (DUAL_GLU) { const int gt = (int) threadIdx.x - 2 * MIX_K; if (!bad_expert && gt >= 0 && gt < 2 * MIX_K) { s_lut[2 * MIX_K + gt] = @@ -762,11 +800,11 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( : rowbase0; // Gate shares the shape, the expert and the row indices -- only the bytes and the table // differ -- so it reuses `nb`, `two`, `row0` and the same activation column below. - const uint8_t * gedata = FUSE_GLU ? gdata + (int64_t) expert * gnb02 : nullptr; - const int gmode = FUSE_GLU ? (int) gmodes[expert] : 0; - const uint8_t * growbase0 = FUSE_GLU ? gedata + (int64_t) row0 * nb * MIX_BLOCK_BYTES + const uint8_t * gedata = DUAL_GLU ? gdata + (int64_t) expert * gnb02 : nullptr; + const int gmode = DUAL_GLU ? (int) gmodes[expert] : 0; + const uint8_t * growbase0 = DUAL_GLU ? gedata + (int64_t) row0 * nb * MIX_BLOCK_BYTES : nullptr; - const uint8_t * growbase1 = FUSE_GLU ? (two ? gedata + (int64_t) (row0 + 1) * nb * MIX_BLOCK_BYTES + const uint8_t * growbase1 = DUAL_GLU ? (two ? gedata + (int64_t) (row0 + 1) * nb * MIX_BLOCK_BYTES : growbase0) : nullptr; // src1 is [in, ne11, ntok]; the get_rows-equivalent row for (slot, token) @@ -792,7 +830,7 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( const int b = blk + u * MIX_WARP; mix_block_accum(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, s_lut, acc0); mix_block_accum(rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { + if (DUAL_GLU) { mix_block_accum(growbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); mix_block_accum(growbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); } @@ -801,7 +839,7 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( for (; blk < nb; blk += MIX_WARP) { mix_block_accum(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc0); mix_block_accum(rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { + if (DUAL_GLU) { mix_block_accum(growbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); mix_block_accum(growbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); } @@ -810,18 +848,24 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( for (int off = MIX_WARP/2; off > 0; off >>= 1) { acc0 += mix_warp_shfl_down(acc0, off); acc1 += mix_warp_shfl_down(acc1, off); - if (FUSE_GLU) { + if (DUAL_GLU) { gacc0 += mix_warp_shfl_down(gacc0, off); gacc1 += mix_warp_shfl_down(gacc1, off); } } if (lane == 0) { const int64_t o = (int64_t) token * dst_s2 + (int64_t) slot * dst_s1 + row0; - if (FUSE_GLU) { + if (DUAL_GLU) { // The SAME function the standalone swiglu_ds4 kernel applies, on inputs bit-identical // to the ones it would have read back from the two intermediates. dst[o] = ggml_cuda_op_swiglu_ds4_single(gacc0, acc0, glu_limit); if (two) dst[o + 1] = ggml_cuda_op_swiglu_ds4_single(gacc1, acc1, glu_limit); + } else if (FINALIZE_GLU) { + // The preceding mode-0 launch left the up projection in dst. Same-stream + // launch ordering makes it visible here; each lane-0 owns distinct rows, + // so reading and replacing those values needs no extra synchronization. + dst[o] = ggml_cuda_op_swiglu_ds4_single(acc0, dst[o], glu_limit); + if (two) dst[o + 1] = ggml_cuda_op_swiglu_ds4_single(acc1, dst[o + 1], glu_limit); } else { dst[o] = acc0; if (two) dst[o + 1] = acc1; @@ -889,13 +933,13 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_id( n_expert_used <= 0 || n_tokens <= 0 || ne11 <= 0) { return false; } - const int warps_per_block = 2; // 64 threads (mirror the mmvq path) + const int warps_per_block = e.gfx1151 ? (n_tokens <= 2 ? 8 : 4) : 2; const int threads = warps_per_block * MIX_WARP; // Two output rows per warp (register-blocked activation reuse), so a workgroup // of `warps_per_block` warps covers 2*warps_per_block rows. const int rows_per_block = 2 * warps_per_block; dim3 grid((out + rows_per_block - 1) / rows_per_block, n_expert_used, n_tokens); - mix_matvec_rocmfp2_moe_kernel<<>>( + mix_matvec_rocmfp2_moe_kernel<0><<>>( (const uint8_t *) e.base, e.nb02, e.codebooks, e.modes, src1, ids, dst, in, out, e.n_experts, ne11, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, @@ -931,15 +975,32 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_id_glu( n_expert_used <= 0 || n_tokens <= 0 || ne11 <= 0) { return false; // not a matched pair; the caller's two-launch path is still correct } - const int warps_per_block = 2; + const bool strix_tuned = eu.gfx1151 && eg.gfx1151; + const int warps_per_block = strix_tuned ? (n_tokens <= 2 ? 8 : 4) : 2; const int threads = warps_per_block * MIX_WARP; const int rows_per_block = 2 * warps_per_block; dim3 grid((out + rows_per_block - 1) / rows_per_block, n_expert_used, n_tokens); - mix_matvec_rocmfp2_moe_kernel<<>>( - (const uint8_t *) eu.base, eu.nb02, eu.codebooks, eu.modes, - src1, ids, dst, in, out, eu.n_experts, ne11, - ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, - (const uint8_t *) eg.base, eg.nb02, eg.codebooks, eg.modes, glu_limit); + if (!strix_tuned || n_tokens <= 2) { + mix_matvec_rocmfp2_moe_kernel<1><<>>( + (const uint8_t *) eu.base, eu.nb02, eu.codebooks, eu.modes, + src1, ids, dst, in, out, eu.n_experts, ne11, + ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, + (const uint8_t *) eg.base, eg.nb02, eg.codebooks, eg.modes, glu_limit); + } else { + // A wide verifier makes the dual-tensor kernel's register pressure more + // expensive than its saved launch. Keep both projections at the ordinary + // kernel's occupancy and fold SwiGLU into the second launch in-place. + mix_matvec_rocmfp2_moe_kernel<0><<>>( + (const uint8_t *) eu.base, eu.nb02, eu.codebooks, eu.modes, + src1, ids, dst, in, out, eu.n_experts, ne11, + ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, + nullptr, 0, nullptr, nullptr, 0.0f); + mix_matvec_rocmfp2_moe_kernel<2><<>>( + (const uint8_t *) eg.base, eg.nb02, eg.codebooks, eg.modes, + src1, ids, dst, in, out, eg.n_experts, ne11, + ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, + nullptr, 0, nullptr, nullptr, glu_limit); + } return true; } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu index df10f4c38..d6905334f 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu @@ -276,10 +276,18 @@ static bool mix_lookup_expert_base( } __device__ __forceinline__ float mix_ue4m3(uint8_t e) { +#if defined(__HIP_PLATFORM_AMD__) && defined(__gfx1151__) + int exp = e >> 3, mant = e & 7; + const float normal = ldexpf((float) (8 + mant), exp - 11); + const float sub = (float) mant * 0.0009765625f; // 2^-10 + const float value = (exp == 0) ? sub : normal; + return (e > 0x7E) ? 0.0f : value; +#else if (e > 0x7E) return 0.0f; int exp = e >> 3, mant = e & 7; if (exp == 0) return (float) mant * 0.0009765625f; // 2^-10 return ldexpf((float) (8 + mant), exp - 11); +#endif } __device__ __forceinline__ uint32_t mix_fp3_code(const uint8_t * qs, int i) { @@ -290,6 +298,22 @@ __device__ __forceinline__ uint32_t mix_fp3_code(const uint8_t * qs, int i) { return (v >> shift) & 7u; } +struct MixFp3Words { + uint64_t lo; // code bytes 0..7 + uint32_t hi; // code bytes 8..11 +}; + +// Decode the same 96-bit fp3 stream from three registers. The only code that +// crosses the 64-bit boundary is i=21 (bit 63); the unrolled caller makes all +// three cases compile-time constants. This avoids materializing a 14-byte +// private array in the gfx1151 matvec while preserving every packed bit. +__device__ __forceinline__ uint32_t mix_fp3_code(const MixFp3Words & qs, int i) { + const int bit = 3 * i; + if (bit <= 60) return (uint32_t) (qs.lo >> bit) & 7u; + if (bit == 63) return (uint32_t) ((qs.lo >> 63) | ((uint64_t) qs.hi << 1)) & 7u; + return (qs.hi >> (bit - 64)) & 7u; +} + __device__ __forceinline__ float mix_fp3_fixed(uint32_t code) { uint32_t m = code & 3u; int mag = (m == 3u) ? 4 : (int) m; @@ -373,6 +397,14 @@ void dequantize_rocmfp3_mix_to_fp16_cuda(const void * vx, half * y, int64_t k, c // quantized blocks once avoids the ~10x f16 round-trip of the dequant fallback. #define MIX_WARP 32 #define MIX_UNROLL 4 +#if defined(__HIP_PLATFORM_AMD__) && defined(__gfx1151__) +// The two-row MoE kernel already exposes thousands of independent workgroups. +// Keeping only one strided block live at a time cuts the large q3 decode body +// and is faster on gfx1151 than duplicating it four times for local MLP. +#define MIX_MOE_UNROLL 1 +#else +#define MIX_MOE_UNROLL MIX_UNROLL +#endif // Down-shift warp shuffle confined to a 32-lane logical group. width=MIX_WARP // keeps the reduction self-contained on wave64 (GFX8/9, physical wave = 64) and @@ -409,16 +441,25 @@ __device__ __forceinline__ void mix_block_accum( // decode arithmetic and the fixed j accumulation order are untouched, so // acc is bit-for-bit identical to the per-byte path (the correctness gate // hashes the greedy output; any reassociation flips a token). +#if defined(__HIP_PLATFORM_AMD__) && defined(__gfx1151__) + MixFp3Words qs; + uint16_t meta; + MIX_MEMCPY(&qs.lo, b, sizeof(qs.lo)); + MIX_MEMCPY(&qs.hi, b + sizeof(qs.lo), sizeof(qs.hi)); + MIX_MEMCPY(&meta, b + MIX_QS, sizeof(meta)); + const uint8_t m0 = (uint8_t) meta, m1 = (uint8_t) (meta >> 8); +#else const uint8_t * ba = (const uint8_t *) MIX_ASSUME_ALIGNED(b, 2); - uint8_t buf[MIX_BLOCK_BYTES]; - MIX_MEMCPY(buf, ba, MIX_BLOCK_BYTES); - const uint8_t m0 = buf[MIX_QS + 0], m1 = buf[MIX_QS + 1]; + uint8_t qs[MIX_BLOCK_BYTES]; + MIX_MEMCPY(qs, ba, MIX_BLOCK_BYTES); + const uint8_t m0 = qs[MIX_QS + 0], m1 = qs[MIX_QS + 1]; +#endif if (mode == 0) { const float s0 = mix_ue4m3(m0), s1 = mix_ue4m3(m1); #pragma unroll for (int j = 0; j < MIX_QK; ++j) { const float s = (j < MIX_QK/2) ? s0 : s1; - acc += s * mix_fp3_fixed(mix_fp3_code(buf, j)) * xc[col0 + j]; + acc += s * mix_fp3_fixed(mix_fp3_code(qs, j)) * xc[col0 + j]; } } else { const float s0 = mix_ue4m3(m0 & 0x7F), s1 = mix_ue4m3(m1 & 0x7F); @@ -435,7 +476,7 @@ __device__ __forceinline__ void mix_block_accum( for (int j = 0; j < MIX_QK; ++j) { const float s = (j < MIX_QK/2) ? s0 : s1; const float * bk = (j < MIX_QK/2) ? bk0 : bk1; - acc += s * bk[mix_fp3_code(buf, j)] * xc[col0 + j]; + acc += s * bk[mix_fp3_code(qs, j)] * xc[col0 + j]; } } } @@ -689,32 +730,21 @@ __global__ void mix_matvec_rocmfp3_moe_kernel( float acc0 = 0.0f, acc1 = 0.0f; float gacc0 = 0.0f, gacc1 = 0.0f; // same block order as acc*, so bit-identical per row int blk = lane; - for (; blk + 3 * MIX_WARP < nb; blk += MIX_UNROLL * MIX_WARP) { - const int b0 = blk, b1 = blk + MIX_WARP; - const int b2 = blk + 2 * MIX_WARP, b3 = blk + 3 * MIX_WARP; - mix_block_accum(rowbase0 + (int64_t) b0 * MIX_BLOCK_BYTES, xcol, b0 * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b0 * MIX_BLOCK_BYTES, xcol, b0 * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { - mix_block_accum(growbase0 + (int64_t) b0 * MIX_BLOCK_BYTES, xcol, b0 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) b0 * MIX_BLOCK_BYTES, xcol, b0 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); - } - mix_block_accum(rowbase0 + (int64_t) b1 * MIX_BLOCK_BYTES, xcol, b1 * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b1 * MIX_BLOCK_BYTES, xcol, b1 * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { - mix_block_accum(growbase0 + (int64_t) b1 * MIX_BLOCK_BYTES, xcol, b1 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) b1 * MIX_BLOCK_BYTES, xcol, b1 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); - } - mix_block_accum(rowbase0 + (int64_t) b2 * MIX_BLOCK_BYTES, xcol, b2 * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b2 * MIX_BLOCK_BYTES, xcol, b2 * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { - mix_block_accum(growbase0 + (int64_t) b2 * MIX_BLOCK_BYTES, xcol, b2 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) b2 * MIX_BLOCK_BYTES, xcol, b2 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); - } - mix_block_accum(rowbase0 + (int64_t) b3 * MIX_BLOCK_BYTES, xcol, b3 * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b3 * MIX_BLOCK_BYTES, xcol, b3 * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { - mix_block_accum(growbase0 + (int64_t) b3 * MIX_BLOCK_BYTES, xcol, b3 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) b3 * MIX_BLOCK_BYTES, xcol, b3 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); + for (; blk + (MIX_MOE_UNROLL - 1) * MIX_WARP < nb; + blk += MIX_MOE_UNROLL * MIX_WARP) { + #pragma unroll + for (int u = 0; u < MIX_MOE_UNROLL; ++u) { + const int b = blk + u * MIX_WARP; + mix_block_accum(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, + b * MIX_QK, mode, s_lut, acc0); + mix_block_accum(rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, + b * MIX_QK, mode, s_lut, acc1); + if (FUSE_GLU) { + mix_block_accum(growbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, + b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); + mix_block_accum(growbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, + b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); + } } } for (; blk < nb; blk += MIX_WARP) { @@ -806,7 +836,7 @@ bool ggml_cuda_rocmfp3_mix_mul_mat_id( n_expert_used <= 0 || n_tokens <= 0 || ne11 <= 0) { return false; } - const int warps_per_block = 2; // 64 threads (mirror the mmvq path) + const int warps_per_block = 2; const int threads = warps_per_block * MIX_WARP; // Two output rows per warp (register-blocked activation reuse), so a workgroup // of `warps_per_block` warps covers 2*warps_per_block rows. diff --git a/server/test/bench_rocmfp_mix_gateup_glu.cpp b/server/test/bench_rocmfp_mix_gateup_glu.cpp index b7fdf2a5f..692070f2e 100644 --- a/server/test/bench_rocmfp_mix_gateup_glu.cpp +++ b/server/test/bench_rocmfp_mix_gateup_glu.cpp @@ -1,4 +1,4 @@ -// Microbenchmark: ONE fused gate/up+SwiGLU launch vs TWO unfused matvec launches. +// Microbenchmark: the dynamic fused gate/up+SwiGLU path vs TWO plain matvec launches. // // Tests the claim the whole fusion rests on. Profiling attributed ~102% of the measured 4.6% // adaptive decode penalty to LAUNCH COUNT (30100 vs qtype 107's 15050) rather than to decode @@ -11,9 +11,11 @@ // shape actually costs. The fused arm therefore looks *worse* than reality -- a conservative // comparison, which is the direction to err in. // -// Reports per-iteration wall time over many iterations on one stream, after a warmup, plus the -// relative delta. Single stream and back-to-back launches on purpose: that is the decode -// dependency chain, where launch overhead cannot be hidden. +// The fused wrapper uses one dual-projection launch through q=2 and two +// lower-register launches (with SwiGLU folded into the second) for wider verification. Reports +// per-iteration wall time over many iterations on one stream, after a warmup, +// plus the relative delta. Single stream and back-to-back launches on purpose: +// that is the decode dependency chain, where launch overhead cannot be hidden. #include "ggml-cuda.h" @@ -21,6 +23,7 @@ #include #include +#include #include #include #include @@ -39,9 +42,20 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_id_glu( int64_t ids_s0, int64_t ids_s1, int64_t src1_s1, int64_t src1_s2, int64_t dst_s1, int64_t dst_s2, float glu_limit, hipStream_t stream); +bool ggml_cuda_rocmfp3_mix_mul_mat_id( + const void * vx, const float * src1, const int32_t * ids, float * dst, + int in, int out, int n_expert_used, int n_tokens, int ne11, + int64_t ids_s0, int64_t ids_s1, int64_t src1_s1, int64_t src1_s2, + int64_t dst_s1, int64_t dst_s2, hipStream_t stream); + +bool ggml_cuda_rocmfp3_mix_mul_mat_id_glu( + const void * vx_up, const void * vx_gate, + const float * src1, const int32_t * ids, float * dst, + int in, int out, int n_expert_used, int n_tokens, int ne11, + int64_t ids_s0, int64_t ids_s1, int64_t src1_s1, int64_t src1_s2, + int64_t dst_s1, int64_t dst_s2, float glu_limit, hipStream_t stream); + static constexpr int QK = 32; -static constexpr int BLOCK_BYTES = 10; -static constexpr int K = 4; static uint32_t xs = 0x1234567u; static uint32_t rnd() { xs ^= xs << 13; xs ^= xs >> 17; xs ^= xs << 5; return xs; } @@ -58,26 +72,35 @@ static double median(std::vector v) { return v[v.size() / 2]; } -int main() { +int main(int argc, char ** argv) { int ndev = 0; if (hipGetDeviceCount(&ndev) != hipSuccess || ndev == 0) { std::fprintf(stderr, "SKIP: no HIP device\n"); return 0; } - // DeepSeek-V4-Flash decode geometry: hidden 7168, n_ff_exp 2048, top-k 4 (--ds4-expert-top-k 4 - // is what the measured runs used), one token per step. - const int in = 7168, out = 2048, n_experts = 8, n_used = 4, ntok = 1; + // DeepSeek-V4-Flash verifier geometry: hidden 4096, n_ff_exp 2048, + // model-default top-k 6. The optional first argument selects the verifier + // width so q=1 and q=4 can be compared without rebuilding. + const int in = 4096, out = 2048, n_experts = 8, n_used = 6; + const int ntok = argc > 1 ? std::atoi(argv[1]) : 4; + const bool fp3 = argc > 2 && std::strcmp(argv[2], "q3") == 0; + if (ntok <= 0 || ntok > 16 || (argc > 2 && !fp3 && std::strcmp(argv[2], "q2") != 0)) { + std::fprintf(stderr, "usage: %s [tokens:1..16] [q2|q3]\n", argv[0]); + return 2; + } + const int block_bytes = fp3 ? 14 : 10; + const int levels = fp3 ? 8 : 4; const int nb = in / QK; - const size_t rows_bytes = (size_t) out * nb * BLOCK_BYTES; + const size_t rows_bytes = (size_t) out * nb * block_bytes; std::vector w(rows_bytes * n_experts); for (auto & b : w) b = (uint8_t) rnd(); - for (size_t blk = 0; blk < w.size() / BLOCK_BYTES; ++blk) { - w[blk * BLOCK_BYTES + 8] = 0x30; - w[blk * BLOCK_BYTES + 9] = 0x30; + for (size_t blk = 0; blk < w.size() / block_bytes; ++blk) { + w[blk * block_bytes + block_bytes - 2] = 0x30; + w[blk * block_bytes + block_bytes - 1] = 0x30; } - std::vector books((size_t) n_experts * 2 * K); + std::vector books((size_t) n_experts * 2 * levels); for (size_t i = 0; i < books.size(); ++i) books[i] = f32_to_bf16(-0.5f + 0.2f * (float) (i % 5)); std::vector modes(n_experts, 1); @@ -100,16 +123,25 @@ int main() { for (int i = 0; i < n_used * ntok; ++i) idsh[i] = i % n_experts; HIP_OK(hipMemcpy(d_ids, idsh.data(), sizeof(int32_t) * idsh.size(), hipMemcpyHostToDevice)); - if (!ggml_cuda_rocmfp2_mix_register_host( + const auto register_mix = fp3 ? ggml_cuda_rocmfp3_mix_register_host + : ggml_cuda_rocmfp2_mix_register_host; + const auto unregister_mix = fp3 ? ggml_cuda_rocmfp3_mix_unregister + : ggml_cuda_rocmfp2_mix_unregister; + const auto mul_mat_id = fp3 ? ggml_cuda_rocmfp3_mix_mul_mat_id + : ggml_cuda_rocmfp2_mix_mul_mat_id; + const auto mul_mat_id_glu = fp3 ? ggml_cuda_rocmfp3_mix_mul_mat_id_glu + : ggml_cuda_rocmfp2_mix_mul_mat_id_glu; + + if (!register_mix( d_up, rows_bytes, n_experts, out, in, books.data(), modes.data())) { std::fprintf(stderr, "FAIL: mixed-tensor registration failed\n"); return 1; } - if (!ggml_cuda_rocmfp2_mix_register_host( + if (!register_mix( d_gate, rows_bytes, n_experts, out, in, books.data(), modes.data())) { - ggml_cuda_rocmfp2_mix_unregister(d_up); + unregister_mix(d_up); std::fprintf(stderr, "FAIL: mixed-tensor registration failed\n"); return 1; } @@ -126,17 +158,17 @@ int main() { auto time_unfused = [&]() { hipEvent_t a, b; hipEventCreate(&a); hipEventCreate(&b); for (int i = 0; i < WARM; ++i) { - ggml_cuda_rocmfp2_mix_mul_mat_id(d_up, d_x, d_ids, d_a, in, out, n_used, ntok, 1, + mul_mat_id(d_up, d_x, d_ids, d_a, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, stream); - ggml_cuda_rocmfp2_mix_mul_mat_id(d_gate, d_x, d_ids, d_b, in, out, n_used, ntok, 1, + mul_mat_id(d_gate, d_x, d_ids, d_b, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, stream); } hipStreamSynchronize(stream); hipEventRecord(a, stream); for (int i = 0; i < ITERS; ++i) { - ggml_cuda_rocmfp2_mix_mul_mat_id(d_up, d_x, d_ids, d_a, in, out, n_used, ntok, 1, + mul_mat_id(d_up, d_x, d_ids, d_a, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, stream); - ggml_cuda_rocmfp2_mix_mul_mat_id(d_gate, d_x, d_ids, d_b, in, out, n_used, ntok, 1, + mul_mat_id(d_gate, d_x, d_ids, d_b, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, stream); } hipEventRecord(b, stream); @@ -148,13 +180,13 @@ int main() { auto time_fused = [&]() { hipEvent_t a, b; hipEventCreate(&a); hipEventCreate(&b); for (int i = 0; i < WARM; ++i) { - ggml_cuda_rocmfp2_mix_mul_mat_id_glu(d_up, d_gate, d_x, d_ids, d_a, in, out, n_used, + mul_mat_id_glu(d_up, d_gate, d_x, d_ids, d_a, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, 7.0f, stream); } hipStreamSynchronize(stream); hipEventRecord(a, stream); for (int i = 0; i < ITERS; ++i) { - ggml_cuda_rocmfp2_mix_mul_mat_id_glu(d_up, d_gate, d_x, d_ids, d_a, in, out, n_used, + mul_mat_id_glu(d_up, d_gate, d_x, d_ids, d_a, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, 7.0f, stream); } hipEventRecord(b, stream); @@ -170,20 +202,32 @@ int main() { for (int r = 0; r < REPS; ++r) { u.push_back(time_unfused()); f.push_back(time_fused()); } const double mu = median(u), mf = median(f); - std::fprintf(stderr, "geometry: in=%d out=%d top_k=%d ntok=%d (%d iters x %d interleaved reps)\n", - in, out, n_used, ntok, ITERS, REPS); + std::fprintf(stderr, "geometry: qtype=%s in=%d out=%d top_k=%d ntok=%d " + "(%d iters x %d interleaved reps)\n", + fp3 ? "q3-mix" : "q2-mix", in, out, n_used, ntok, ITERS, REPS); std::fprintf(stderr, " unfused (2 launches, swiglu NOT counted): %8.4f ms/step [", mu); for (double v : u) std::fprintf(stderr, " %.4f", v); - std::fprintf(stderr, " ]\n fused (1 launch, swiglu included ): %8.4f ms/step [", mf); + const int fused_launches = fp3 || ntok <= 2 ? 1 : 2; + std::fprintf(stderr, " ]\n fused (%d launch%s, swiglu included): %8.4f ms/step [", + fused_launches, fused_launches == 1 ? " " : "es", mf); for (double v : f) std::fprintf(stderr, " %.4f", v); std::fprintf(stderr, " ]\n"); std::fprintf(stderr, " fused is %+.2f%% vs unfused (negative = faster)\n", 100.0 * (mf / mu - 1.0)); std::fprintf(stderr, " per-layer saving %.4f ms -> over 43 layers %.3f ms/step\n", mu - mf, (mu - mf) * 43.0); - ggml_cuda_rocmfp2_mix_unregister(d_gate); - ggml_cuda_rocmfp2_mix_unregister(d_up); + std::vector result(yn); + HIP_OK(hipMemcpy(result.data(), d_a, sizeof(float) * yn, hipMemcpyDeviceToHost)); + uint64_t result_hash = 1469598103934665603ull; + const uint8_t * result_bytes = reinterpret_cast(result.data()); + for (size_t i = 0; i < sizeof(float) * yn; ++i) { + result_hash = (result_hash ^ result_bytes[i]) * 1099511628211ull; + } + std::fprintf(stderr, " result fnv1a64: %016llx\n", (unsigned long long) result_hash); + + unregister_mix(d_gate); + unregister_mix(d_up); hipStreamDestroy(stream); HIP_OK(hipFree(d_up)); HIP_OK(hipFree(d_gate)); HIP_OK(hipFree(d_x)); HIP_OK(hipFree(d_a)); HIP_OK(hipFree(d_b)); HIP_OK(hipFree(d_ids)); diff --git a/server/test/test_rocmfp_mix_gateup_glu.cpp b/server/test/test_rocmfp_mix_gateup_glu.cpp index 8c66aa190..a563244ba 100644 --- a/server/test/test_rocmfp_mix_gateup_glu.cpp +++ b/server/test/test_rocmfp_mix_gateup_glu.cpp @@ -122,7 +122,9 @@ TEST_CASE(RocmfpMixGateupGluFixture, fused_gateup_glu) { // in must be a multiple of 128: the wide block load reads 128 weights at a time and would // read past the tensor on the final block (register_host enforces this). - const int in = 256, out = 64, n_experts = 6, n_used = 3, ntok = 2; + // Three tokens exercises the low-register two-pass GLU finalizer. q <= 2 + // uses the one-pass kernel and is checked separately below. + const int in = 256, out = 64, n_experts = 6, n_used = 3, ntok = 3; const int nb = in / QK; const size_t rows_bytes = (size_t) out * nb * BLOCK_BYTES; @@ -224,6 +226,41 @@ TEST_CASE(RocmfpMixGateupGluFixture, fused_gateup_glu) { std::fprintf(stderr, "worst relative deviation from the host reference: %.3e\n", worst); CHECK(worst < 1e-5); + // q <= 2 selects the one-pass dual-projection kernel. Exercise that branch + // separately while the main ntok=3 case above covers the two-pass finalizer. + { + const int one_tok = 1; + const size_t one_yn = (size_t) out * n_used; + CHECK(ggml_cuda_rocmfp2_mix_mul_mat_id(d_up, d_x, d_ids, d_up_out, + in, out, n_used, one_tok, 1, + ids_s0, ids_s1, src1_s1, src1_s2, + dst_s1, dst_s2, nullptr)); + CHECK(ggml_cuda_rocmfp2_mix_mul_mat_id(d_gate, d_x, d_ids, d_gate_out, + in, out, n_used, one_tok, 1, + ids_s0, ids_s1, src1_s1, src1_s2, + dst_s1, dst_s2, nullptr)); + CHECK(ggml_cuda_rocmfp2_mix_mul_mat_id_glu(d_up, d_gate, d_x, d_ids, d_fused, + in, out, n_used, one_tok, 1, + ids_s0, ids_s1, src1_s1, src1_s2, + dst_s1, dst_s2, limit, nullptr)); + HIP_OK(cudaDeviceSynchronize()); + std::vector one_u(one_yn), one_g(one_yn), one_f(one_yn); + HIP_OK(cudaMemcpy(one_u.data(), d_up_out, sizeof(float) * one_yn, + cudaMemcpyDeviceToHost)); + HIP_OK(cudaMemcpy(one_g.data(), d_gate_out, sizeof(float) * one_yn, + cudaMemcpyDeviceToHost)); + HIP_OK(cudaMemcpy(one_f.data(), d_fused, sizeof(float) * one_yn, + cudaMemcpyDeviceToHost)); + double one_worst = 0.0; + for (size_t i = 0; i < one_yn; ++i) { + const float ref = host_swiglu_ds4(one_g[i], one_u[i], limit); + const double denom = std::fmax(1e-6, std::fabs((double) ref)); + one_worst = std::fmax(one_worst, + std::fabs((double) one_f[i] - (double) ref) / denom); + } + CHECK(one_worst < 1e-5); + } + // ---- operand ORDER: swiglu_ds4 applies silu to GATE, so the two are not symmetric ---- float * d_swapped = nullptr; HIP_OK(cudaMalloc(&d_swapped, sizeof(float) * yn)); From 867f573fbad403a95973d6ade9c58b30e320cff2 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 26 Aug 2026 22:33:14 +0200 Subject: [PATCH 2/7] perf(ds4): accelerate long-context verify selection --- server/CMakeLists.txt | 11 ++ .../llama.cpp/ggml/src/ggml-cuda/top-k.cu | 166 +++++++++++++++- .../src/deepseek4/deepseek4_fused_verify.inc | 117 ++++++++--- server/src/deepseek4/deepseek4_graph.cpp | 4 +- server/test/bench_ds4_topk.cpp | 184 ++++++++++++++++++ 5 files changed, 454 insertions(+), 28 deletions(-) create mode 100644 server/test/bench_ds4_topk.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 925d5839d..1fc77c42c 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -946,6 +946,17 @@ if(DFLASH27B_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) target_link_libraries(bench_rocmfp_mix_gateup_glu PRIVATE ggml ggml-base ${DFLASH27B_GGML_BACKEND_TARGET} hip::host) + + # Benchmark, not a test: compares the generic segmented full sort with + # the exact two-stage block-radix TOP_K used by the DS4 long-context + # indexer. It also checks selected-index set parity at every tile count. + add_executable(bench_ds4_topk test/bench_ds4_topk.cpp) + set_source_files_properties(test/bench_ds4_topk.cpp PROPERTIES LANGUAGE HIP) + set_target_properties(bench_ds4_topk PROPERTIES HIP_ARCHITECTURES "${_dflash_archs}") + target_include_directories(bench_ds4_topk PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + target_link_libraries(bench_ds4_topk PRIVATE + ggml ggml-base ${DFLASH27B_GGML_BACKEND_TARGET} hip::host) endif() add_executable(test_qwen35_tensor_parallel test/test_qwen35_tensor_parallel.cpp) target_include_directories(test_qwen35_tensor_parallel PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu index 851de693c..10c8434ca 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu @@ -1319,6 +1319,161 @@ static void topk_block_radix_cuda( CUDA_CHECK(cudaGetLastError()); } +// Long-context DS4 indexer shape: select 512 compressed KV rows from as many +// as 32K scores. A full segmented sort materializes and orders every score even +// though only 512 survive. Instead, select 512 candidates independently from +// each 4K tile, then select the final 512 from the at-most-4K candidates. Any +// global top-512 element must be in its tile's local top 512, so this is exact. +// +// Both stages retain the original global indices as values. TOP_K only requires +// the selected set; ordering among equal scores is intentionally unspecified. +// Keeping tiles and candidates in increasing source order nevertheless gives +// BlockRadixSort the same natural tie order as the flat input. +template +static __global__ void k_topk_block_radix_tiles_f32_i32( + const float * x, + int * candidates, + int ncols, + int ntiles, + int k) { + constexpr int BLOCK_THREADS = 256; + constexpr int TILE_COLS = BLOCK_THREADS * ITEMS_PER_THREAD; + using block_sort = hipcub::BlockRadixSort< + uint32_t, BLOCK_THREADS, ITEMS_PER_THREAD, int>; + __shared__ typename block_sort::TempStorage storage; + + const int tile_block = (int) blockIdx.x; + const int row = tile_block / ntiles; + const int tile = tile_block - row * ntiles; + const int tile_begin = tile * TILE_COLS; + const int first = tile_begin + (int) threadIdx.x * ITEMS_PER_THREAD; + const float * x_row = x + (size_t) row * ncols; + uint32_t keys[ITEMS_PER_THREAD]; + int indices[ITEMS_PER_THREAD]; +#pragma unroll + for (int item = 0; item < ITEMS_PER_THREAD; ++item) { + const int col = first + item; + if (col < ncols) { + const uint32_t bits = (uint32_t) __float_as_int(x_row[col]); + const uint32_t ordered = (bits & 0x80000000u) + ? ~bits : (bits ^ 0x80000000u); + keys[item] = ordered == 0 ? 1 : ordered; + } else { + keys[item] = 0; + } + indices[item] = col; + } + + block_sort(storage).SortDescending(keys, indices); + +#pragma unroll + for (int item = 0; item < ITEMS_PER_THREAD; ++item) { + const int rank = (int) threadIdx.x * ITEMS_PER_THREAD + item; + if (rank < k) { + candidates[((size_t) row * ntiles + tile) * k + rank] = indices[item]; + } + } +} + +template +static __global__ void k_topk_block_radix_merge_f32_i32( + const float * x, + const int * candidates, + int * dst, + int ncols, + int ncandidates, + int k) { + constexpr int BLOCK_THREADS = 256; + using block_sort = hipcub::BlockRadixSort< + uint32_t, BLOCK_THREADS, ITEMS_PER_THREAD, int>; + __shared__ typename block_sort::TempStorage storage; + + const int row = (int) blockIdx.x; + const int first = (int) threadIdx.x * ITEMS_PER_THREAD; + const float * x_row = x + (size_t) row * ncols; + const int * candidate_row = candidates + (size_t) row * ncandidates; + uint32_t keys[ITEMS_PER_THREAD]; + int indices[ITEMS_PER_THREAD]; +#pragma unroll + for (int item = 0; item < ITEMS_PER_THREAD; ++item) { + const int rank = first + item; + if (rank < ncandidates) { + const int col = candidate_row[rank]; + if (col < ncols) { + const uint32_t bits = (uint32_t) __float_as_int(x_row[col]); + const uint32_t ordered = (bits & 0x80000000u) + ? ~bits : (bits ^ 0x80000000u); + keys[item] = ordered == 0 ? 1 : ordered; + indices[item] = col; + } else { + // The final 4K tile may contain fewer than k real columns. + // Its local selection then contains padded indices; keep them + // below every real score and never dereference them. + keys[item] = 0; + indices[item] = ncols; + } + } else { + keys[item] = 0; + indices[item] = ncols; + } + } + + block_sort(storage).SortDescending(keys, indices); + +#pragma unroll + for (int item = 0; item < ITEMS_PER_THREAD; ++item) { + const int rank = first + item; + if (rank < k) { + dst[(size_t) row * k + rank] = indices[item]; + } + } +} + +static void topk_hierarchical_block_radix_cuda( + ggml_cuda_pool & pool, + const float * x, + int * dst, + int ncols, + int nrows, + int k, + cudaStream_t stream) { + constexpr int TILE_COLS = 4096; + constexpr int TILE_ITEMS_PER_THREAD = 16; + const int ntiles = (ncols + TILE_COLS - 1) / TILE_COLS; + const int ncandidates = ntiles * k; + + GGML_ASSERT(k == 512); + GGML_ASSERT(ntiles >= 2 && ntiles <= 8); + ggml_cuda_pool_alloc candidates_alloc( + pool, (size_t) nrows * ncandidates); + int * candidates = candidates_alloc.get(); + + const dim3 tile_blocks((unsigned) (nrows * ntiles), 1, 1); + constexpr int threads = 256; + k_topk_block_radix_tiles_f32_i32 + <<>>( + x, candidates, ncols, ntiles, k); + + const dim3 merge_blocks((unsigned) nrows, 1, 1); + if (ntiles <= 2) { + k_topk_block_radix_merge_f32_i32<4><<>>( + x, candidates, dst, ncols, ncandidates, k); + } else if (ntiles <= 3) { + k_topk_block_radix_merge_f32_i32<6><<>>( + x, candidates, dst, ncols, ncandidates, k); + } else if (ntiles <= 4) { + k_topk_block_radix_merge_f32_i32<8><<>>( + x, candidates, dst, ncols, ncandidates, k); + } else if (ntiles <= 6) { + k_topk_block_radix_merge_f32_i32<12><<>>( + x, candidates, dst, ncols, ncandidates, k); + } else { + k_topk_block_radix_merge_f32_i32<16><<>>( + x, candidates, dst, ncols, ncandidates, k); + } + CUDA_CHECK(cudaGetLastError()); +} + #endif // GGML_CUDA_USE_HIPCUB void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { @@ -1346,9 +1501,14 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { #elif defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) // CUB_TOP_K_AVAILABLE #ifdef GGML_CUDA_USE_HIPCUB if (ds4_env_flag_enabled("GGML_DS4_TOPK_BLOCK_RADIX") && - k == 512 && ncols > 1024 && ncols <= 5120) { - topk_block_radix_cuda( - src0_d, dst_d, (int) ncols, (int) nrows, (int) k, stream); + k == 512 && ncols > 1024 && ncols <= 32768) { + if (ncols <= 5120) { + topk_block_radix_cuda( + src0_d, dst_d, (int) ncols, (int) nrows, (int) k, stream); + } else { + topk_hierarchical_block_radix_cuda( + pool, src0_d, dst_d, (int) ncols, (int) nrows, (int) k, stream); + } return; } #endif diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index e732fa6bc..7d73d8842 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -1223,14 +1223,20 @@ static int ds4_try_fused_verify_step( ds4_fv_set(fg->i32_bundle, i32v.data(), sizeof(int32_t) * i32v.size()); ds4_fv_set(fg->i64_bundle, i64v.data(), sizeof(int64_t) * i64v.size()); - // causal mask values + // Causal mask values. At long context this tensor is multi-megabyte, + // while only the ring slots overwritten by this q-wide batch and the + // padded compressed tail change between verifier steps. Build those + // negative spans directly. Large masks are zeroed on the GPU, then only + // those negative spans cross the host/backend boundary. { + using MaskRange = std::pair; // [begin, end) const size_t mask_count = (size_t) ggml_nelements(fg->mask_bundle); std::vector & maskv = ex->mask_values; - if (maskv.size() != mask_count) { - maskv.resize(mask_count); - } - std::fill(maskv.begin(), maskv.end(), 0.0f); + std::vector next_negative; + next_negative.reserve((size_t) w.n_layer * (size_t) q * 4u); + const auto add_range = [&](size_t begin, size_t end) { + if (begin < end) next_negative.emplace_back(begin, end); + }; size_t off = 0; for (int il = 0; il < w.n_layer; ++il) { const int ratio = (int) w.compress_ratios[il]; @@ -1249,34 +1255,97 @@ static int ds4_try_fused_verify_step( const size_t n_attn = (size_t) w.n_swa + padded + (lane_q > 1 ? (size_t) lane_q : 0u); for (int i = 0; i < lane_q; ++i) { - float * col = maskv.data() + off + (size_t) i * n_attn; + const size_t col = off + (size_t) i * n_attn; const int pos_i = lane_kv_start + i; - for (int r = 0; r < w.n_swa; ++r) { - // position held by slot r AFTER this batch's ring writes - const int p_r = (e <= w.n_swa) ? (r < e ? r : -1) - : (e - 1) - ((e - 1 - r) % w.n_swa); - if (p_r < 0 || p_r > pos_i) col[r] = -1.0e30f; + + // Raw ring after this batch's writes. Before the ring fills, + // every slot after pos_i is invalid or future. Afterwards, + // only the q-1-i future writes are hidden; their ring slots + // form one range, or two when they wrap. + if (e <= w.n_swa) { + add_range(col + (size_t) (pos_i + 1), + col + (size_t) w.n_swa); + } else { + const int future = e - 1 - pos_i; + if (future > 0) { + const int begin = (pos_i + 1) % w.n_swa; + const int first = std::min(future, w.n_swa - begin); + add_range(col + (size_t) begin, + col + (size_t) (begin + first)); + add_range(col, + col + (size_t) (future - first)); + } } + const int vis = (ratio > 0 && lc.comp_kv) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, pos_i) : 0; - for (int c = 0; c < padded; ++c) { - if (c >= vis) col[(size_t) w.n_swa + c] = -1.0e30f; - } + add_range(col + (size_t) w.n_swa + (size_t) vis, + col + (size_t) w.n_swa + (size_t) padded); + if (lane_q > 1) { - for (int j = 0; j < lane_q; ++j) { - const bool visible = - (j > i) && - (lane_kv_start + j >= w.n_swa); - if (!visible) { - col[(size_t) w.n_swa + padded + j] = -1.0e30f; - } - } + const int first_visible = std::max( + i + 1, w.n_swa - lane_kv_start); + const int masked = std::clamp(first_visible, 0, lane_q); + add_range(col + (size_t) w.n_swa + (size_t) padded, + col + (size_t) w.n_swa + (size_t) padded + + (size_t) masked); } } off += n_attn * lane_q; } - GGML_ASSERT(off == maskv.size()); - ds4_fv_set(fg->mask_bundle, maskv.data(), sizeof(float) * maskv.size()); + GGML_ASSERT(off == mask_count); + + const auto normalize_ranges = [](std::vector & ranges) { + std::sort(ranges.begin(), ranges.end()); + size_t out = 0; + for (const MaskRange & range : ranges) { + if (out > 0 && range.first <= ranges[out - 1].second) { + ranges[out - 1].second = std::max( + ranges[out - 1].second, range.second); + } else { + ranges[out++] = range; + } + } + ranges.resize(out); + }; + normalize_ranges(next_negative); + + static const bool sparse_mask_update = [] { + const char * value = std::getenv( + "DFLASH_DS4_INCREMENTAL_VERIFY_MASK"); + return !value || std::atoi(value) != 0; + }(); + static const size_t sparse_mask_min_bytes = [] { + constexpr size_t default_bytes = 4u * 1024u * 1024u; + const char * value = std::getenv( + "DFLASH_DS4_INCREMENTAL_VERIFY_MASK_MIN_BYTES"); + if (!value || !*value) return default_bytes; + const unsigned long long parsed = std::strtoull(value, nullptr, 10); + return parsed > 0 ? (size_t) parsed : default_bytes; + }(); + const bool use_sparse_mask_update = sparse_mask_update && + mask_count * sizeof(float) >= sparse_mask_min_bytes; + if (!use_sparse_mask_update) { + maskv.assign(mask_count, 0.0f); + for (const MaskRange & range : next_negative) { + std::fill(maskv.begin() + range.first, + maskv.begin() + range.second, -1.0e30f); + } + ds4_fv_set(fg->mask_bundle, maskv.data(), + sizeof(float) * maskv.size()); + } else { + if (maskv.size() != mask_count) maskv.resize(mask_count); + ggml_backend_tensor_memset( + fg->mask_bundle, 0, 0, mask_count * sizeof(float)); + for (const MaskRange & range : next_negative) { + std::fill(maskv.begin() + range.first, + maskv.begin() + range.second, -1.0e30f); + ggml_backend_tensor_set( + fg->mask_bundle, maskv.data() + range.first, + range.first * sizeof(float), + (range.second - range.first) * sizeof(float)); + } + } } if (token_ids) { diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 87f7f0d45..cec5b9b6c 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -39,6 +39,7 @@ #include #include #include +#include #include #if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) @@ -4834,7 +4835,8 @@ struct Ds4FusedVerifyCache { ggml_tensor * capture = nullptr; // f32 [n_embd*ncap,q], token-major ggml_tensor * argmax = nullptr; // i32 [q], optional greedy output // Reused host staging for the context-sized additive attention mask. - // Keeping it per slot removes one allocation from every verify step. + // Keeping it per slot removes allocation churn in both full and + // sparse-range mask update modes. std::vector mask_values; int q = 0; diff --git a/server/test/bench_ds4_topk.cpp b/server/test/bench_ds4_topk.cpp new file mode 100644 index 000000000..0e317f717 --- /dev/null +++ b/server/test/bench_ds4_topk.cpp @@ -0,0 +1,184 @@ +#include "ggml-backend.h" +#include "ggml-cuda.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include + +namespace { + +struct RunResult { + std::vector indices; + double milliseconds = 0.0; +}; + +std::vector make_scores(int ncols, int nrows) { + std::vector scores((size_t) ncols * nrows); + for (int row = 0; row < nrows; ++row) { + for (int col = 0; col < ncols; ++col) { + // An odd multiplier is a permutation modulo 2^24, so every score + // in a row is distinct and exactly representable as float. + const uint32_t value = + ((uint32_t) col * 2654435761u + (uint32_t) row * 2246822519u) & + 0x00ffffffu; + scores[(size_t) row * ncols + col] = + (float) ((int32_t) value - 0x00800000); + } + } + return scores; +} + +bool compute_topk( + ggml_backend_t backend, + int ncols, + int nrows, + int k, + bool block_radix, + int warmup, + int iterations, + RunResult & result) { + ggml_init_params params{}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + std::fprintf(stderr, "ggml_init failed\n"); + return false; + } + + ggml_tensor * scores = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ncols, nrows); + ggml_tensor * selected = ggml_top_k(ctx, scores, k); + ggml_set_input(scores); + ggml_set_output(selected); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, selected); + + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buffer) { + std::fprintf(stderr, "backend tensor allocation failed\n"); + ggml_free(ctx); + return false; + } + + const std::vector host_scores = make_scores(ncols, nrows); + ggml_backend_tensor_set( + scores, host_scores.data(), 0, host_scores.size() * sizeof(float)); + setenv("GGML_DS4_TOPK_BLOCK_RADIX", block_radix ? "1" : "0", 1); + + bool ok = true; + for (int i = 0; i < warmup; ++i) { + ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + if (!ok) { + break; + } + } + ggml_backend_synchronize(backend); + + const auto begin = std::chrono::steady_clock::now(); + for (int i = 0; ok && i < iterations; ++i) { + ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + } + ggml_backend_synchronize(backend); + const auto end = std::chrono::steady_clock::now(); + result.milliseconds = + std::chrono::duration(end - begin).count() / + iterations; + + if (ok) { + result.indices.resize((size_t) nrows * k); + ggml_backend_tensor_get( + selected, + result.indices.data(), + 0, + result.indices.size() * sizeof(int32_t)); + } + + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + return ok; +} + +bool same_selected_set( + const std::vector & reference, + const std::vector & candidate, + int ncols, + int nrows, + int k) { + if (reference.size() != candidate.size()) { + return false; + } + for (int row = 0; row < nrows; ++row) { + const auto begin = (size_t) row * k; + std::vector expected( + reference.begin() + begin, reference.begin() + begin + k); + std::vector actual( + candidate.begin() + begin, candidate.begin() + begin + k); + if (std::any_of(actual.begin(), actual.end(), [ncols](int32_t index) { + return index < 0 || index >= ncols; + })) { + return false; + } + std::sort(expected.begin(), expected.end()); + std::sort(actual.begin(), actual.end()); + if (expected != actual) { + return false; + } + } + return true; +} + +} // namespace + +int main() { + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, "ggml_backend_cuda_init failed\n"); + return 1; + } + + const bool previous_graphs_disabled = + ggml_backend_cuda_set_graphs_disabled_override(true); + constexpr int k = 512; + constexpr int nrows = 4; + const int shapes[] = { + 5121, 8192, 8193, 12288, 12289, 16384, 28673, 30720, 32768 + }; + bool all_ok = true; + + for (const int ncols : shapes) { + RunResult full_sort; + RunResult hierarchical; + const int iterations = ncols == 30720 ? 100 : 5; + const bool ran = + compute_topk( + backend, ncols, nrows, k, false, 2, iterations, full_sort) && + compute_topk( + backend, ncols, nrows, k, true, 2, iterations, hierarchical); + const bool exact = ran && same_selected_set( + full_sort.indices, hierarchical.indices, ncols, nrows, k); + const double speedup = hierarchical.milliseconds > 0.0 + ? full_sort.milliseconds / hierarchical.milliseconds + : 0.0; + std::printf( + "ncols=%d rows=%d k=%d full_sort=%.3f ms hierarchical=%.3f ms " + "speedup=%.2fx exact_set=%s\n", + ncols, + nrows, + k, + full_sort.milliseconds, + hierarchical.milliseconds, + speedup, + exact ? "yes" : "NO"); + all_ok = all_ok && exact; + } + + unsetenv("GGML_DS4_TOPK_BLOCK_RADIX"); + ggml_backend_cuda_set_graphs_disabled_override(previous_graphs_disabled); + ggml_backend_free(backend); + return all_ok ? 0 : 1; +} From 5eacd1f3962950615668ef54c9acd268b99309f3 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 27 Aug 2026 12:16:22 +0200 Subject: [PATCH 3/7] feat(spec): adapt verify width from acceptance feedback --- server/CMakeLists.txt | 1 + server/docs/ENVIRONMENT.md | 2 + server/src/common/adaptive_spec_width.h | 100 ++++++++++++++++++ server/src/common/dflash_spec_decode.cpp | 18 +++- .../src/deepseek4/deepseek4_dspark_spec.cpp | 48 +++++---- server/src/gemma4/gemma4_backend.cpp | 18 +++- server/src/laguna/laguna_backend.cpp | 32 +++--- server/src/laguna/laguna_backend.h | 4 - server/src/qwen35/qwen35_backend.cpp | 29 +++-- server/src/qwen35moe/qwen35moe_backend.cpp | 15 ++- server/test/test_adaptive_spec_width.cpp | 73 +++++++++++++ 11 files changed, 280 insertions(+), 60 deletions(-) create mode 100644 server/src/common/adaptive_spec_width.h create mode 100644 server/test/test_adaptive_spec_width.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 1fc77c42c..2da748b69 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1583,6 +1583,7 @@ if(DFLASH27B_TESTS) test/test_server_unit.cpp test/test_anchor_params.cpp test/test_derived_scalars.cpp + test/test_adaptive_spec_width.cpp test/test_adaptive_keep_ratio.cpp test/test_skip_park_guard.cpp test/test_bandit_integration.cpp diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index d83371e66..7431460d6 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -16,6 +16,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. | Variable | Default | Purpose | |---|---|---| +| `DFLASH_ADAPTIVE_SPEC_WIDTH` | 1 | BURN-IN KILL SWITCH: =0 disables the shared acceptance-feedback verify-width controller. Fixed per-backend width overrides still take precedence. | | `DFLASH_DRAFT_KV` | 1 | KILL SWITCH (remove after burn-in): =0 restores the legacy per-step drafter window recompute instead of the ring cache. | | `DFLASH_LAGUNA_SWA_RING` | 1 | KILL SWITCH (remove after burn-in): =0 keeps SWA layers on pool-sized caches under KVFlash. | | `DFLASH_PROF` | unset | DEBUG: comma list of profilers (step,verify,prefill). Replaces DFLASH_LAGUNA_{STEP,VERIFY,PREFILL}_PROF. | @@ -72,6 +73,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH27B_PREFILL_UBATCH` - layer_split_daemon.cpp, qwen35_backend.cpp, qwen35_layer_split_adapter.cpp - `DFLASH_ADAPTIVE_K_DENSE` - mmid_adaptive_k.h - `DFLASH_ADAPTIVE_K_TAU` - mmid_adaptive_k.h +- `DFLASH_ADAPTIVE_SPEC_WIDTH` - adaptive_spec_width.h - `DFLASH_ADAPTIVE_WIDTH_MIN` - adaptive_verify_width.h - `DFLASH_ADAPTIVE_WIDTH_THETA` - adaptive_verify_width.h - `DFLASH_COLD_THREADS` - moe_expert_compute_cpu.cpp diff --git a/server/src/common/adaptive_spec_width.h b/server/src/common/adaptive_spec_width.h new file mode 100644 index 000000000..8b700242e --- /dev/null +++ b/server/src/common/adaptive_spec_width.h @@ -0,0 +1,100 @@ +#pragma once + +// Shared feedback controller for chain speculative decoding. +// +// A rejection reveals the exact useful draft depth, so back off with an EMA. +// An all-accepted draft is censored: it only proves that the useful depth was +// at least the offered width. Averaging that lower bound would strand the +// controller at a narrow width, so clean drafts probe upward additively. +// +// Widths are seed-inclusive throughout this API. For example, width 4 means +// one always-committed seed plus three speculative candidates. +// +// Enabled by default. Set DFLASH_ADAPTIVE_SPEC_WIDTH=0 to retain each +// backend's fixed-width behavior. Backend-specific fixed-width overrides +// still take precedence. + +#include +#include +#include +#include + +namespace dflash::common { + +inline bool adaptive_spec_width_globally_enabled() { + static const bool enabled = [] { + const char * value = std::getenv("DFLASH_ADAPTIVE_SPEC_WIDTH"); + return value == nullptr || value[0] == '\0' || std::strcmp(value, "0") != 0; + }(); + return enabled; +} + +class AdaptiveSpecWidth { +public: + AdaptiveSpecWidth(int max_width, int min_width = 2, bool enabled = true, + float initial_accepted_candidates = 2.0f, + float backoff_alpha = 0.25f, + float full_accept_probe = 1.0f) + : max_width_(std::max(1, max_width)), + min_width_(std::clamp(min_width, 1, std::max(1, max_width))), + enabled_(enabled), + initial_accepted_candidates_(std::clamp( + initial_accepted_candidates, 0.0f, + static_cast(std::max(0, max_width_ - 1)))), + backoff_alpha_(std::clamp(backoff_alpha, 0.0f, 1.0f)), + full_accept_probe_(std::max(0.0f, full_accept_probe)), + accepted_candidates_ema_(initial_accepted_candidates_) {} + + // Apply both this feedback cap and an optional model-specific cap. + int next_width(int proposed_width) const { + const int proposed = std::clamp(proposed_width, 1, max_width_); + if (!enabled_ || max_width_ <= 1) return proposed; + + const int feedback_width = std::clamp( + static_cast(std::lround(accepted_candidates_ema_)) + 1, + min_width_, max_width_); + return std::min(proposed, feedback_width); + } + + int next_width() const { return next_width(max_width_); } + + // accepted_width and offered_width both include the seed row. + void observe(int accepted_width, int offered_width) { + if (!enabled_ || offered_width <= 1 || max_width_ <= 1) return; + + const int offered = std::clamp(offered_width, 1, max_width_); + const int accepted = std::clamp(accepted_width, 1, offered); + const int offered_candidates = offered - 1; + const int accepted_candidates = accepted - 1; + + if (accepted_candidates >= offered_candidates) { + // Censored lower bound: do not average it downward; probe wider. + accepted_candidates_ema_ = std::min( + static_cast(max_width_ - 1), + accepted_candidates_ema_ + full_accept_probe_); + } else { + // The first rejection exposes the exact accepted prefix length. + accepted_candidates_ema_ = + (1.0f - backoff_alpha_) * accepted_candidates_ema_ + + backoff_alpha_ * static_cast(accepted_candidates); + } + } + + void reset() { accepted_candidates_ema_ = initial_accepted_candidates_; } + + bool enabled() const { return enabled_; } + float accepted_candidates_ema() const { return accepted_candidates_ema_; } + int min_width() const { return min_width_; } + int max_width() const { return max_width_; } + +private: + int max_width_; + int min_width_; + bool enabled_; + float initial_accepted_candidates_; + float backoff_alpha_; + float full_accept_probe_; + float accepted_candidates_ema_; +}; + +} // namespace dflash::common diff --git a/server/src/common/dflash_spec_decode.cpp b/server/src/common/dflash_spec_decode.cpp index 9bf903359..66faf63cc 100644 --- a/server/src/common/dflash_spec_decode.cpp +++ b/server/src/common/dflash_spec_decode.cpp @@ -6,6 +6,7 @@ #include "io_utils.h" #include "dflash_draft_graph.h" // build_draft_step #include "step_graph.h" +#include "adaptive_spec_width.h" #include #include @@ -66,6 +67,8 @@ bool run_dflash_spec_decode( const int hidden = draft_weights.n_embd; const int q_len = draft_weights.block_size; + AdaptiveSpecWidth width_controller( + q_len, 2, adaptive_spec_width_globally_enabled()); StepGraph draft_sg; StepGraphGuard draft_sg_guard{draft_sg}; @@ -84,6 +87,7 @@ bool run_dflash_spec_decode( int n_generated = 0; int n_draft_steps = 0; int n_accept_sum = 0; + int n_offered_sum = 0; int n_hint_proposed = 0; int n_hint_accepted = 0; const ChainRollbackPolicy rollback_policy = @@ -162,6 +166,10 @@ bool run_dflash_spec_decode( return false; } draft_tok[0] = last_tok; + const int verify_width = width_controller.next_width(q_len); + if ((int)draft_tok.size() > verify_width) { + draft_tok.resize((size_t)verify_width); + } // ── Tool call hint injection ────────────────────────────────────── // Override draft tokens with pre-known hint tokens for near-100% @@ -169,7 +177,7 @@ bool run_dflash_spec_decode( int hint_filled = 0; if (hint_tokens && n_generated < (int)hint_tokens->size()) { const int hint_avail = (int)hint_tokens->size() - n_generated; - hint_filled = std::min(hint_avail, q_len - 1); + hint_filled = std::min(hint_avail, verify_width - 1); for (int i = 0; i < hint_filled; i++) { draft_tok[1 + i] = (*hint_tokens)[n_generated + i]; } @@ -200,16 +208,18 @@ bool run_dflash_spec_decode( // Acceptance: longest matching prefix between draft and target argmax. int accept_n = 1; - for (int i = 0; i < q_len - 1; i++) { + for (int i = 0; i < verify_width - 1; i++) { if (draft_tok[i + 1] == target_tok[i]) accept_n++; else break; } + width_controller.observe(accept_n, verify_width); + n_offered_sum += verify_width; // Track hint acceptance telemetry. if (hint_filled > 0) { n_hint_proposed += hint_filled; n_hint_accepted += std::min(hint_filled, accept_n - 1); } - int bonus_tok = (accept_n < q_len) ? target_tok[accept_n - 1] : -1; + int bonus_tok = (accept_n < verify_width) ? target_tok[accept_n - 1] : -1; int commit_n = accept_n + (bonus_tok >= 0 ? 1 : 0); if (commit_n > need_commit_budget) { commit_n = need_commit_budget; @@ -302,7 +312,7 @@ bool run_dflash_spec_decode( if (!use_remote_draft && draft_backend) ggml_backend_synchronize(draft_backend); auto t_dec1 = std::chrono::steady_clock::now(); const double decode_s = std::chrono::duration(t_dec1 - t_dec0).count(); - const int total_draft_pos = std::max(1, n_draft_steps * q_len); + const int total_draft_pos = std::max(1, n_offered_sum); const double accept_pct = 100.0 * (double)n_accept_sum / (double)total_draft_pos; if (accept_rate_out) { *accept_rate_out = total_draft_pos > 0 diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index 1d687b586..492727e1b 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -25,6 +25,7 @@ #include "deepseek4_internal.h" #include "deepseek4_roctx.h" #include "internal.h" +#include "common/adaptive_spec_width.h" #include "common/dspark_head.h" #include "ggml.h" @@ -704,27 +705,21 @@ bool run_deepseek4_dspark_spec_decode( "[ds4-spec] draft overlap probe requested without an independent " "in-process backend; probe disabled\n"); } - // Laguna-style adaptive verify width: EWMA of accepted candidates, width = - // ewma + 2 (avg_commit << block means the wide tail is usually wasted). - // Keep this process-local: a shared filesystem control would let an - // unrelated user or benchmark change live inference behavior. + // Shared adaptive verify width. Keep the existing DS4 switch as a + // backend-specific override; the common controller also has one global + // opt-out for matched fixed-width A/B runs. bool adaptive_width = true; if (const char * raw = std::getenv("DFLASH_DS4_ADAPTIVE_WIDTH")) { adaptive_width = raw[0] && std::strcmp(raw, "0") != 0; } - // The adaptive policy was calibrated only through q=4. Q5 is an explicit - // fixed-width mode and must not be silently narrowed. - if (q5_verify) { - adaptive_width = false; - } + adaptive_width = adaptive_width && adaptive_spec_width_globally_enabled(); + // The confidence artifact was calibrated through q=4. Q5 still adapts, + // but uses target acceptance feedback rather than extrapolating that head. const bool use_confidence_width = adaptive_width && !seq_verify_mode && + !q5_verify && drafter.confidence_w != nullptr && drafter.confidence_b != nullptr && (drafter.confidence_dim == n_embd || drafter.confidence_dim == n_embd + drafter.markov_rank); - if (timing && use_confidence_width) { - std::fprintf(stderr, "[ds4-spec] adaptive width policy=confidence\n"); - } - double ewma_accept = 1.5; // The conservative fast path remains capped at the compression ratio. // The explicit wide path handles a second ratio-4 boundary in-graph and @@ -750,6 +745,14 @@ bool run_deepseek4_dspark_spec_decode( DS4_CONSERVATIVE_VERIFY_MAX_TOKENS); q_cap = DS4_CONSERVATIVE_VERIFY_MAX_TOKENS; } + AdaptiveSpecWidth width_controller( + q_cap, 2, adaptive_width && !seq_verify_mode); + if (timing && width_controller.enabled()) { + std::fprintf(stderr, "[ds4-spec] adaptive width policy=%s\n", + use_confidence_width + ? "confidence (acceptance fallback)" + : "acceptance"); + } // Snapshot backend for the legacy full-snapshot rollback path. ggml_backend_t snap_backend = ggml_backend_cpu_init(); @@ -883,9 +886,12 @@ bool run_deepseek4_dspark_spec_decode( : std::min( q_cap, DS4_CONSERVATIVE_VERIFY_MAX_TOKENS - (pos & 3)); - if (adaptive_width && !use_confidence_width && !seq_verify_mode) { - const int w_cap = (int) ewma_accept + 2; - if (w_cap < q_step_cap) q_step_cap = w_cap; + // A calibrated confidence head already predicts this individual + // step. Do not stack the slower acceptance-regime cap on top of it; + // acceptance feedback remains the fallback and drives q5/artifacts + // without confidence metadata. + if (!use_confidence_width) { + q_step_cap = width_controller.next_width(q_step_cap); } if (q_step_cap >= 2) { std::memcpy(padded_hidden.data() + n_embd, local_hidden.data(), @@ -938,10 +944,10 @@ bool run_deepseek4_dspark_spec_decode( } if ((int) draft_tok.size() > selected_q) draft_tok.resize((size_t) selected_q); } else if (use_confidence_width && !seq_verify_mode) { - // The fused head should always return confidence for a compatible - // artifact. Preserve the old policy if a backend cannot do so. - const int selected_q = (int) ewma_accept + 2; - if ((int) draft_tok.size() > selected_q) draft_tok.resize((size_t) selected_q); + const int selected_q = width_controller.next_width((int)draft_tok.size()); + if ((int)draft_tok.size() > selected_q) { + draft_tok.resize((size_t)selected_q); + } } if ((int) draft_tok.size() > q_step_cap) draft_tok.resize(q_step_cap); const int q = (int) draft_tok.size(); // seed + candidates @@ -1129,7 +1135,7 @@ bool run_deepseek4_dspark_spec_decode( lt = bonus; // deferred bonus becomes next seed accept_sum += matched; offered_sum += q - 1; - ewma_accept = 0.7 * ewma_accept + 0.3 * (double) matched; + width_controller.observe(accept, q); steps++; if (timing && (steps <= 4 || (steps & 31) == 0)) { std::fprintf(stderr, diff --git a/server/src/gemma4/gemma4_backend.cpp b/server/src/gemma4/gemma4_backend.cpp index d9cdd9ad6..a45fb0e6a 100644 --- a/server/src/gemma4/gemma4_backend.cpp +++ b/server/src/gemma4/gemma4_backend.cpp @@ -11,6 +11,7 @@ #include "common/io_utils.h" #include "common/dflash_feature_ring.h" #include "common/dflash_draft_graph.h" +#include "common/adaptive_spec_width.h" #include "common/step_graph.h" #include "ggml-cuda.h" @@ -485,6 +486,8 @@ bool Gemma4Backend::do_spec_decode(int committed, int n_gen, DFlashTarget * target = dflash_target_; const int q_len = dw_.block_size; + AdaptiveSpecWidth width_controller( + q_len, 2, adaptive_spec_width_globally_enabled()); StepGraph draft_sg; @@ -499,6 +502,7 @@ bool Gemma4Backend::do_spec_decode(int committed, int n_gen, int n_generated = 0; int n_draft_steps = 0; int n_accept_sum = 0; + int n_offered_sum = 0; auto t_dec0 = std::chrono::steady_clock::now(); @@ -545,7 +549,7 @@ bool Gemma4Backend::do_spec_decode(int committed, int n_gen, tail_hook, forced_close_out); auto t_dec1 = std::chrono::steady_clock::now(); const double decode_s = std::chrono::duration(t_dec1 - t_dec0).count(); - const int total_draft_pos = std::max(1, n_draft_steps * q_len); + const int total_draft_pos = std::max(1, n_offered_sum); const double accept_pct = 100.0 * (double)n_accept_sum / (double)total_draft_pos; std::fprintf(stderr, "[gemma4-spec] tail-off-stats tokens=%d time=%.3f s " @@ -623,6 +627,10 @@ bool Gemma4Backend::do_spec_decode(int committed, int n_gen, return false; } draft_tok[0] = last_tok; + const int verify_width = width_controller.next_width(q_len); + if ((int)draft_tok.size() > verify_width) { + draft_tok.resize((size_t)verify_width); + } // 4. Verify: run target forward over all draft tokens. // Gemma4 is a pure transformer — after verify, KV entries at accepted @@ -638,11 +646,13 @@ bool Gemma4Backend::do_spec_decode(int committed, int n_gen, // 5. Acceptance: longest matching prefix int accept_n = 1; - for (int i = 0; i < q_len - 1; i++) { + for (int i = 0; i < verify_width - 1; i++) { if (draft_tok[i + 1] == target_tok[i]) accept_n++; else break; } - int bonus_tok = (accept_n < q_len) ? target_tok[accept_n - 1] : -1; + width_controller.observe(accept_n, verify_width); + n_offered_sum += verify_width; + int bonus_tok = (accept_n < verify_width) ? target_tok[accept_n - 1] : -1; int commit_n = accept_n + (bonus_tok >= 0 ? 1 : 0); if (commit_n > need_commit_budget) { commit_n = need_commit_budget; @@ -700,7 +710,7 @@ bool Gemma4Backend::do_spec_decode(int committed, int n_gen, auto t_dec1 = std::chrono::steady_clock::now(); const double decode_s = std::chrono::duration(t_dec1 - t_dec0).count(); - const int total_draft_pos = std::max(1, n_draft_steps * q_len); + const int total_draft_pos = std::max(1, n_offered_sum); const double accept_pct = 100.0 * (double)n_accept_sum / (double)total_draft_pos; std::fprintf(stderr, "[gemma4-spec] tokens=%d time=%.3f s speed=%.2f tok/s " "steps=%d accepted=%d/%d (%.1f%%) avg_commit=%.2f\n", diff --git a/server/src/laguna/laguna_backend.cpp b/server/src/laguna/laguna_backend.cpp index cf42b013c..b78817564 100644 --- a/server/src/laguna/laguna_backend.cpp +++ b/server/src/laguna/laguna_backend.cpp @@ -32,6 +32,7 @@ #include "common/step_graph.h" #include "ggml-cuda.h" +#include "../common/adaptive_spec_width.h" #include "../common/adaptive_verify_width.h" #include "ggml-alloc.h" #include "common/snapshot_backend.h" @@ -512,9 +513,9 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, // proposes block_size tokens but avg_commit is ~2.9, so verifying the whole // block is wasteful. Measured (laguna-xs2 Q4_K_M, RTX 3090): width 8 -> 110 // tok/s, 4 -> 138, 3 -> 150 (== AR). We verify only the first q_len of the - // drafted block; the accept rule is unchanged, so this stays lossless. AUTO - // tracks an EWMA of the accepted length (held constant per request so the - // verify graph stays CUDA-graph-stable); --verify-width forces a fixed width. + // drafted block; the accept rule is unchanged, so this stays lossless. + // --verify-width forces a fixed width; AUTO uses the shared per-request + // censored-feedback controller below. const bool sampled_verify = laguna_sampled_verify_enabled(sampler_, true); int verify_width = args_.verify_width; if (const char * e = std::getenv("DFLASH_LAGUNA_VERIFY_WIDTH")) { @@ -535,17 +536,21 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, // [TAG_ADAPTIVE_WIDTH] default width policy: with the per-step // drafter-confidence trim active (on by default for greedy chains), run // from a base of 8 rows and let the trim shrink each step. The legacy - // accept-EWMA AUTO remains the fallback when the trim is off (theta 0) - // and for sampled verify, which has no candidate probabilities. + // acceptance-feedback AUTO remains the fallback when the trim is off + // (theta 0) and for sampled verify, which has no candidate probabilities. const bool width_trim = adaptive_verify_width_theta() > 0.0f && !sampled_verify && !args_.ddtree_mode; int chain_w = adaptive_width - ? (width_trim ? 8 : std::min((int)(spec_ewma_accept_ + 0.5) + 1, auto_w_max)) + ? (width_trim ? 8 : auto_w_max) : verify_width; if (chain_w < 2) chain_w = 2; if (chain_w > std::min(block_size, 8)) chain_w = std::min(block_size, 8); // DDTree sizes its batch via its budget; chain uses the width chosen above. const int base_q_len = args_.ddtree_mode ? block_size : chain_w; + AdaptiveSpecWidth width_controller( + base_q_len, 2, + adaptive_width && !args_.ddtree_mode && + adaptive_spec_width_globally_enabled()); const bool ignore_eos = (std::getenv("DFLASH_IGNORE_EOS") != nullptr); if (dflash_target_) { @@ -1093,6 +1098,12 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, target_tok.resize((size_t)q_len); } } + const int feedback_width = width_controller.next_width(q_len); + if (feedback_width < q_len) { + q_len = feedback_width; + draft_tok.resize((size_t)q_len); + target_tok.resize((size_t)q_len); + } int verify_last_tok = -1; if (step_prof) prof_lap(); @@ -1134,6 +1145,7 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, } bonus_tok = (accept_n < q_len) ? target_tok[(size_t)accept_n - 1] : -1; } + width_controller.observe(accept_n, q_len); int commit_n = accept_n + (bonus_tok >= 0 ? 1 : 0); if (commit_n > need_commit_budget) { commit_n = need_commit_budget; @@ -1262,14 +1274,6 @@ bool LagunaBackend::do_spec_decode(int committed, int n_gen, *accept_rate_out = (float)(n_accept_sum / (double)total_draft_pos); } - // [TAG_LAGUNA_VERIFY_WIDTH] Update the persisted accepted-length EWMA so the - // AUTO width converges to the throughput optimum for the active draft (chain - // only; DDTree sizes via its budget and accounts accepts differently). - if (adaptive_width && !args_.ddtree_mode && n_draft_steps > 0) { - const double mean_accept = (double)n_accept_sum / (double)n_draft_steps; - spec_ewma_accept_ = 0.7 * spec_ewma_accept_ + 0.3 * mean_accept; - } - if (dflash_target_) { dflash_target_->set_keep_verify_logits(false); } diff --git a/server/src/laguna/laguna_backend.h b/server/src/laguna/laguna_backend.h index e006c6446..90334678f 100644 --- a/server/src/laguna/laguna_backend.h +++ b/server/src/laguna/laguna_backend.h @@ -116,10 +116,6 @@ class LagunaBackend : public ModelBackend { DraftKvState draft_kv_{}; LagunaDFlashTarget * dflash_target_ = nullptr; bool draft_parked_ = false; - // [TAG_LAGUNA_VERIFY_WIDTH] EWMA of the accepted block length, persisted - // across requests. Drives the AUTO chain verify width (seeded for width 3). - double spec_ewma_accept_ = 1.5; - // PFlash drafter (lazy-loaded on first compress command). DrafterContext drafter_ctx_{}; bool drafter_loaded_ = false; diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 55b4fa4e1..7a97be3b6 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -1,6 +1,7 @@ #include "qwen35_backend.h" #include "concurrency/qwen35_seq_engine.h" #include "common/chain_rollback_policy.h" +#include "common/adaptive_spec_width.h" #include "common/draft_block_size.h" #include "placement/skip_park_guard.h" #include "qwen35_dflash_target.h" @@ -2703,12 +2704,14 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // visible to the rows we keep, and would break the fixed block_size // contract the IPC drafter validates against. // Clamped to q_len: a checkpoint whose published block is below the - // narrowing floor never widens past its own block, and the accept-rate - // denominator (n_spec_steps * verify_cap) stays equal to the positions - // actually drafted. + // narrowing floor never widens past its own block. Accept-rate accounting + // below sums the widths actually verified. const int verify_cap = committed >= kLongCtxNarrowTokens ? std::min(q_len, std::max(kLongCtxMinVerify, q_len / 2)) : q_len; + AdaptiveSpecWidth width_controller( + q_len, 2, + !cfg_.ddtree_mode && adaptive_spec_width_globally_enabled()); if (verify_cap != q_len) { static std::atomic s_narrowed_logged{false}; if (!s_narrowed_logged.exchange(true)) { @@ -2748,6 +2751,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, int n_generated = 0; int n_draft_steps = 0; int n_accept_sum = 0; + int n_spec_offered_sum = 0; int n_hint_proposed = 0; int n_hint_accepted = 0; int target_forwards = 0; @@ -2808,7 +2812,6 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, float accepted_ema = 2.0f * adaptive.accept_threshold(); int ar_burst_left = 0; int n_ar_burst_steps = 0; - int n_spec_steps = 0; // steps that actually proposed q_len drafts bool probe_step = false; // first spec step after a burst // Live step-time EMAs (seconds) for the break-even ratio; 0 = not yet measured. double t_spec_step_ema = 0.0; @@ -3580,6 +3583,13 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, draft_tok.resize((size_t)verify_cap); } } + if (!ar_step) { + const int feedback_width = width_controller.next_width(v_len); + if (feedback_width < v_len) { + v_len = feedback_width; + draft_tok.resize((size_t)v_len); + } + } // 3b. Tool call hint injection: override draft tokens with pre-known // structural tokens for near-100% acceptance. @@ -3698,6 +3708,10 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, } bonus_tok = (accept_n < v_len) ? target_tok[accept_n - 1] : -1; } + if (!ar_step) { + width_controller.observe(accept_n, v_len); + n_spec_offered_sum += v_len; + } // Track hint acceptance telemetry. if (hint_fill > 0) { n_hint_proposed += hint_fill; @@ -3945,7 +3959,6 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, // that steers the PFlash residency bandit. if (!ar_step) { n_accept_sum += std::min(accept_n, emitted); - n_spec_steps++; } n_draft_steps++; @@ -3983,7 +3996,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, if (floor_to_ar) { step_graph_destroy(draft_sg); cache_.last_tok = out_tokens.empty() ? last_tok : out_tokens.back(); - const int total_draft_pos = std::max(1, n_spec_steps * verify_cap); + const int total_draft_pos = std::max(1, n_spec_offered_sum); out_accept_rate = (float)((double)n_accept_sum / (double)total_draft_pos); const int ar_n_gen = n_gen - n_generated; @@ -4028,7 +4041,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, cache_.cur_pos = committed; step_graph_destroy(draft_sg); cache_.last_tok = out_tokens.empty() ? last_tok : out_tokens.back(); - const int total_draft_pos = std::max(1, n_spec_steps * verify_cap); + const int total_draft_pos = std::max(1, n_spec_offered_sum); out_accept_rate = (float)((double)n_accept_sum / (double)total_draft_pos); const int ar_n_gen = n_gen - n_generated; @@ -4059,7 +4072,7 @@ bool Qwen35Backend::do_spec_decode(int committed, int n_gen, auto t_dec1 = std::chrono::steady_clock::now(); const double decode_s = std::chrono::duration(t_dec1 - t_dec0).count(); - const int total_draft_pos = std::max(1, n_spec_steps * verify_cap); + const int total_draft_pos = std::max(1, n_spec_offered_sum); const double accept_pct = 100.0 * (double)n_accept_sum / (double)total_draft_pos; out_accept_rate = (float)((double)n_accept_sum / (double)total_draft_pos); std::fprintf(stderr, "[spec-decode] tokens=%d time=%.3f s speed=%.2f tok/s " diff --git a/server/src/qwen35moe/qwen35moe_backend.cpp b/server/src/qwen35moe/qwen35moe_backend.cpp index 1f6ea7074..c6fe48384 100644 --- a/server/src/qwen35moe/qwen35moe_backend.cpp +++ b/server/src/qwen35moe/qwen35moe_backend.cpp @@ -7,6 +7,7 @@ #include "../common/kvflash_placement.h" #include "common/ggml_graph_precision.h" #include "common/sampler.h" +#include "common/adaptive_spec_width.h" #include "common/dflash_spec_decode.h" #include "dflash_draft_graph.h" #include "dflash_feature_ring.h" @@ -1995,7 +1996,9 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, const char * e = std::getenv("DFLASH_VERIFY_WIDTH"); return e ? std::max(1, std::min(q_len, std::atoi(e))) : 0; }(); - int observed_max_accept = 1; + AdaptiveSpecWidth width_controller( + q_len, std::min(6, q_len), + forced_verify_width == 0 && adaptive_spec_width_globally_enabled()); int32_t last_tok = target_cache().last_tok; std::vector act_cur((size_t)hidden); @@ -2012,6 +2015,7 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, int n_generated = 0; int n_draft_steps = 0; int n_accept_sum = 0; + int n_offered_sum = 0; // Allocate DeltaNet rollback snapshot tensors (no-op if already present). // Without these, snapshot_ssm_state/restore_ssm_state silently do nothing @@ -2028,7 +2032,7 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, const int need_commit_budget = n_gen - n_generated; const int verify_width = forced_verify_width > 0 ? forced_verify_width - : std::min(q_len, std::max(6, observed_max_accept + 2)); + : width_controller.next_width(q_len); // 1. Build noise input for draft noise_ids[0] = last_tok; @@ -2137,8 +2141,9 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, if (draft_tok[i + 1] == target_tok[i]) accept_n++; else break; } + width_controller.observe(accept_n, verify_width); + n_offered_sum += verify_width; int bonus_tok = (accept_n < verify_width) ? target_tok[accept_n - 1] : -1; - observed_max_accept = std::max(observed_max_accept, accept_n); int commit_n = accept_n + (bonus_tok >= 0 ? 1 : 0); if (commit_n > need_commit_budget) { commit_n = need_commit_budget; @@ -2199,7 +2204,7 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, const int fallback_steps = hybrid_spec_min_steps_before_ar(); if (!io.is_cancelled() && !hit_eos && fallback_steps > 0 && n_draft_steps >= fallback_steps && n_generated < n_gen) { - const int total_draft_pos_so_far = std::max(1, n_draft_steps * q_len); + const int total_draft_pos_so_far = std::max(1, n_offered_sum); const float accept_rate_value = (float)((double)n_accept_sum / (double)total_draft_pos_so_far); const float min_accept = hybrid_spec_min_accept_rate(); @@ -2225,7 +2230,7 @@ bool Qwen35MoeBackend::do_hybrid_spec_decode(int committed, int n_gen, auto t_dec1 = std::chrono::steady_clock::now(); const double decode_s = std::chrono::duration(t_dec1 - t_dec0).count(); - const int total_draft_pos = std::max(1, n_draft_steps * q_len); + const int total_draft_pos = std::max(1, n_offered_sum); const double accept_pct = 100.0 * (double)n_accept_sum / (double)total_draft_pos; if (accept_rate_out) { *accept_rate_out = total_draft_pos > 0 diff --git a/server/test/test_adaptive_spec_width.cpp b/server/test/test_adaptive_spec_width.cpp new file mode 100644 index 000000000..e13332d49 --- /dev/null +++ b/server/test/test_adaptive_spec_width.cpp @@ -0,0 +1,73 @@ +#include "CppUnitTestFramework.hpp" +#include "common/adaptive_spec_width.h" + +#include + +using namespace dflash::common; + +namespace { +struct AdaptiveSpecWidthFixture {}; + +bool near(float lhs, float rhs, float tolerance = 1e-6f) { + return std::fabs(lhs - rhs) <= tolerance; +} +} // namespace + +TEST_CASE(AdaptiveSpecWidthFixture, starts_with_two_candidates) { + AdaptiveSpecWidth width(16); + CHECK(width.next_width() == 3); +} + +TEST_CASE(AdaptiveSpecWidthFixture, clean_drafts_probe_upward) { + AdaptiveSpecWidth width(8); + CHECK(width.next_width() == 3); + width.observe(3, 3); + CHECK(width.next_width() == 4); + width.observe(4, 4); + CHECK(width.next_width() == 5); +} + +TEST_CASE(AdaptiveSpecWidthFixture, censored_sample_is_not_averaged_down) { + AdaptiveSpecWidth width(8, 2, true, 5.0f); + width.observe(3, 3); + CHECK(near(width.accepted_candidates_ema(), 6.0f)); + CHECK(width.next_width() == 7); +} + +TEST_CASE(AdaptiveSpecWidthFixture, rejection_backs_off_gently) { + AdaptiveSpecWidth width(8, 2, true, 4.0f); + width.observe(2, 6); // one accepted candidate, then a rejection + CHECK(near(width.accepted_candidates_ema(), 3.25f)); + CHECK(width.next_width() == 4); +} + +TEST_CASE(AdaptiveSpecWidthFixture, respects_model_and_feedback_caps) { + AdaptiveSpecWidth width(8, 2, true, 5.0f); + CHECK(width.next_width(4) == 4); + CHECK(width.next_width(99) == 6); + CHECK(width.next_width(1) == 1); +} + +TEST_CASE(AdaptiveSpecWidthFixture, respects_minimum_width) { + AdaptiveSpecWidth width(16, 6, true, 0.0f); + CHECK(width.next_width() == 6); + width.observe(1, 6); + CHECK(width.next_width() == 6); +} + +TEST_CASE(AdaptiveSpecWidthFixture, disabled_controller_preserves_proposal) { + AdaptiveSpecWidth width(16, 2, false); + CHECK(width.next_width() == 16); + CHECK(width.next_width(7) == 7); + width.observe(1, 16); + CHECK(width.next_width() == 16); +} + +TEST_CASE(AdaptiveSpecWidthFixture, reset_restores_initial_estimate) { + AdaptiveSpecWidth width(8); + width.observe(1, 3); + CHECK(!near(width.accepted_candidates_ema(), 2.0f)); + width.reset(); + CHECK(near(width.accepted_candidates_ema(), 2.0f)); + CHECK(width.next_width() == 3); +} From ecae4618198a7b0ac42bda8a99b2cbe4dd803b85 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 27 Aug 2026 12:53:17 +0200 Subject: [PATCH 4/7] perf(moe): add model-neutral fused expert combine --- server/CMakeLists.txt | 10 ++ server/deps/llama.cpp/ggml/include/ggml.h | 11 ++ .../llama.cpp/ggml/src/ggml-cuda/moe-fused.cu | 8 +- server/deps/llama.cpp/ggml/src/ggml.c | 9 +- server/docs/ENVIRONMENT.md | 1 + server/src/common/moe_hybrid_ffn_eval.cpp | 4 +- server/src/deepseek4/deepseek4_graph.cpp | 69 ++++++---- server/src/laguna/laguna_target_graph.cpp | 2 +- server/test/test_moe_combine.cpp | 125 ++++++++++++++++++ 9 files changed, 207 insertions(+), 32 deletions(-) create mode 100644 server/test/test_moe_combine.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 2da748b69..c439bacb0 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -1967,6 +1967,16 @@ if(DFLASH27B_TESTS) add_dependencies(check test_concat_transpose) endif() endif() + if((DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR + DFLASH27B_GPU_BACKEND STREQUAL "hip") + AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_moe_combine.cpp") + dflash_add_ggml_gpu_executable( + test_moe_combine test/test_moe_combine.cpp) + add_test(NAME moe_combine COMMAND test_moe_combine) + if(TARGET check) + add_dependencies(check test_moe_combine) + endif() + endif() if((DFLASH27B_GPU_BACKEND STREQUAL "cuda" OR DFLASH27B_GPU_BACKEND STREQUAL "hip") AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/bench_paged_attention.cpp") diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index 0e1bcebf4..5018af331 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2560,6 +2560,17 @@ extern "C" { int64_t ff_dim, int64_t n_expert_used); + // Apply one routing weight to every expert output column and reduce the + // expert axis. `experts` is [n_embd, n_used, n_tokens] and + // `expert_weights` is [n_used, n_tokens]. Backends may fuse the weighting + // and reduction into one kernel. + GGML_API struct ggml_tensor * ggml_moe_combine( + struct ggml_context * ctx, + struct ggml_tensor * experts, + struct ggml_tensor * expert_weights); + + // Compatibility alias for out-of-tree Laguna callers. New model code + // should use the model-neutral ggml_moe_combine API. GGML_API struct ggml_tensor * ggml_laguna_moe_combine( struct ggml_context * ctx, struct ggml_tensor * experts, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu index 0541c58f2..16b08d5f5 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu @@ -263,7 +263,7 @@ static __global__ void moe_fused_kernel( } } -static __global__ void laguna_moe_combine_kernel( +static __global__ void moe_combine_kernel( const char * __restrict__ experts, const char * __restrict__ weights, char * __restrict__ output, @@ -564,7 +564,7 @@ static void ggml_cuda_op_ds4_moe_owner( const int total = n_embd * n_tokens; const int block = 256; const int grid = (total + block - 1) / block; - laguna_moe_combine_kernel<<>>( + moe_combine_kernel<<>>( (const char *) experts.data, (const char *) weights->data, (char *) dst->data, @@ -637,7 +637,7 @@ static void ggml_cuda_op_ds4_moe_owner_split( const int total = n_embd * n_tokens; const int block = 256; const int grid = (total + block - 1) / block; - laguna_moe_combine_kernel<<>>( + moe_combine_kernel<<>>( (const char *) experts.data, (const char *) weights->data, (char *) dst->data, @@ -764,7 +764,7 @@ void ggml_cuda_op_moe_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst) const int block = 256; const int grid = (total + block - 1) / block; - laguna_moe_combine_kernel<<>>( + moe_combine_kernel<<>>( (const char *) experts->data, (const char *) weights->data, (char *) dst->data, diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index ef6f51532..d4907d65c 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -8605,7 +8605,7 @@ struct ggml_tensor * ggml_moe_fused( return result; } -struct ggml_tensor * ggml_laguna_moe_combine( +struct ggml_tensor * ggml_moe_combine( struct ggml_context * ctx, struct ggml_tensor * experts, struct ggml_tensor * expert_weights) { @@ -8629,6 +8629,13 @@ struct ggml_tensor * ggml_laguna_moe_combine( return result; } +struct ggml_tensor * ggml_laguna_moe_combine( + struct ggml_context * ctx, + struct ggml_tensor * experts, + struct ggml_tensor * expert_weights) { + return ggml_moe_combine(ctx, experts, expert_weights); +} + struct ggml_tensor * ggml_ds4_moe_owner( struct ggml_context * ctx, struct ggml_tensor * input, diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index 7431460d6..666f7427b 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -27,6 +27,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_MMID_GROUPED` | unset | Grouped MUL_MAT_ID kernel for small verify batches; candidate for CLI promotion. | | `DFLASH_MMID_GROUPED_TYPES` | 7 | Grouped-kernel type mask; bit 3 (`8`) opts ROCmFP2/ROCmFP3 into the path. | | `DFLASH_MMID_GROUPED_DEVICE` | -1 | Optional zero-based device restriction; unset/-1 applies to every eligible device. | +| `DFLASH_MOE_FUSED_COMBINE` | unset | DEBUG/A-B: use the model-neutral fused routed-expert weighting and summation kernel. Available to DeepSeek4, Laguna, and common hybrid-MoE graphs; changes floating-point association, so keep disabled until model-level qualification. | | `DFLASH_DS4_MOE_TP` / `DFLASH_DS4_MOE_TP_INPROC` | unset | BURN-IN: enable DeepSeek4 route-owner expert parallelism in one process. | | `DFLASH_DS4_MOE_TP_BACKEND` / `DFLASH_MOE_TP_BACKEND` | peer runtime in a mixed build; compiled runtime otherwise | Select the in-process cold expert owner backend. | | `DFLASH_DS4_MOE_TP_GPU` | peer backend device 0 in a mixed build; other local device otherwise | Device index within the cold DeepSeek4 expert backend. | diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index d5e80cec0..d4ed148dd 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -872,7 +872,7 @@ static bool build_batched_routed_graph( // Weight and sum over experts: [n_embd, n_used, n_tokens] * [1, n_used, n_tokens] if (!defer_route_reduction && allow_fused_combine && (force_fused_combine || moe_hybrid_graph_policy().fused_combine)) { - *out_routed = track(ggml_laguna_moe_combine(ctx, experts, wts)); + *out_routed = track(ggml_moe_combine(ctx, experts, wts)); return *out_routed != nullptr; } @@ -3155,7 +3155,7 @@ static bool eval_moe_owner_expert_major_batched( ggml_tensor * route_major = ggml_get_rows(ctx, packed_table, inverse_routes); ggml_tensor * routed_out = - ggml_laguna_moe_combine(ctx, route_major, route_weights); + ggml_moe_combine(ctx, route_major, route_weights); combined_out = combined_out ? ggml_add(ctx, routed_out, combined_out) : routed_out; diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index cec5b9b6c..a82bd9f79 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -478,12 +478,19 @@ static bool build_cached_decode_ffn_graph( weights = ggml_scale(out.sg.ctx, weights, w.expert_weight_scale); } - ggml_tensor * weights_3d = ggml_reshape_3d(out.sg.ctx, weights, 1, n_used, n_tokens); - ggml_tensor * routed_out = ggml_mul(out.sg.ctx, down_e, weights_3d); - routed_out = ggml_cont( - out.sg.ctx, ggml_permute(out.sg.ctx, routed_out, 1, 0, 2, 3)); - routed_out = ggml_sum_rows(out.sg.ctx, routed_out); - routed_out = ggml_reshape_2d(out.sg.ctx, routed_out, w.n_embd, n_tokens); + ggml_tensor * routed_out = nullptr; + if (moe_hybrid_graph_policy().fused_combine) { + routed_out = ggml_moe_combine(out.sg.ctx, down_e, weights); + } else { + ggml_tensor * weights_3d = ggml_reshape_3d( + out.sg.ctx, weights, 1, n_used, n_tokens); + routed_out = ggml_mul(out.sg.ctx, down_e, weights_3d); + routed_out = ggml_cont( + out.sg.ctx, ggml_permute(out.sg.ctx, routed_out, 1, 0, 2, 3)); + routed_out = ggml_sum_rows(out.sg.ctx, routed_out); + routed_out = ggml_reshape_2d( + out.sg.ctx, routed_out, w.n_embd, n_tokens); + } ffn_out = ggml_add(out.sg.ctx, shared_out, routed_out); } else { @@ -3337,11 +3344,19 @@ static ggml_tensor * build_moe_ffn( ggml_tensor * down_e = ggml_mul_mat_id(ctx, L.ffn_down_exps, mid_e, routing.selected); down_e = ggml_reshape_3d(ctx, down_e, n_embd, n_used, n_tokens); - ggml_tensor * weights_3d = ggml_reshape_3d(ctx, routing.weights, 1, n_used, n_tokens); - routed_out = ggml_mul(ctx, down_e, weights_3d); - routed_out = ggml_cont(ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); - routed_out = ggml_sum_rows(ctx, routed_out); - routed_out = ggml_reshape_2d(ctx, routed_out, n_embd, n_tokens); + if (moe_hybrid_graph_policy().fused_combine) { + routed_out = ggml_moe_combine( + ctx, down_e, routing.weights); + } else { + ggml_tensor * weights_3d = ggml_reshape_3d( + ctx, routing.weights, 1, n_used, n_tokens); + routed_out = ggml_mul(ctx, down_e, weights_3d); + routed_out = ggml_cont( + ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); + routed_out = ggml_sum_rows(ctx, routed_out); + routed_out = ggml_reshape_2d( + ctx, routed_out, n_embd, n_tokens); + } } return ggml_add(ctx, shared_out, routed_out); @@ -5067,20 +5082,26 @@ static ggml_tensor * ds4_build_hash_routed_ffn( weights = ggml_scale(ctx, weights, w.expert_weight_scale); } - ggml_tensor * weights_3d = ggml_reshape_3d( - ctx, weights, 1, n_used, n_tokens); - ggml_tensor * routed_out = ggml_mul(ctx, down_e, weights_3d); - if (n_tokens == 1) { - // Preserve the established q=1 graph and reduction order. - routed_out = ggml_cont(ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); - routed_out = ggml_sum_rows(ctx, routed_out); - routed_out = ggml_reshape_2d(ctx, routed_out, w.n_embd, 1); + ggml_tensor * routed_out = nullptr; + if (moe_hybrid_graph_policy().fused_combine) { + routed_out = ggml_moe_combine(ctx, down_e, weights); } else { - ggml_tensor * sum_shape = ggml_new_tensor_3d( - ctx, GGML_TYPE_F32, w.n_embd, 1, n_tokens); - routed_out = ggml_repeat_back(ctx, routed_out, sum_shape); - routed_out = ggml_reshape_2d( - ctx, routed_out, w.n_embd, n_tokens); + ggml_tensor * weights_3d = ggml_reshape_3d( + ctx, weights, 1, n_used, n_tokens); + routed_out = ggml_mul(ctx, down_e, weights_3d); + if (n_tokens == 1) { + // Preserve the established q=1 graph and reduction order. + routed_out = ggml_cont( + ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); + routed_out = ggml_sum_rows(ctx, routed_out); + routed_out = ggml_reshape_2d(ctx, routed_out, w.n_embd, 1); + } else { + ggml_tensor * sum_shape = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, w.n_embd, 1, n_tokens); + routed_out = ggml_repeat_back(ctx, routed_out, sum_shape); + routed_out = ggml_reshape_2d( + ctx, routed_out, w.n_embd, n_tokens); + } } return ggml_add(ctx, shared_out, routed_out); } diff --git a/server/src/laguna/laguna_target_graph.cpp b/server/src/laguna/laguna_target_graph.cpp index b52f0e461..21310a9dc 100644 --- a/server/src/laguna/laguna_target_graph.cpp +++ b/server/src/laguna/laguna_target_graph.cpp @@ -490,7 +490,7 @@ static ggml_tensor * build_laguna_moe_block_full(ggml_context * ctx, ggml_cgraph ggml_tensor * routed = nullptr; if (fused_combine) { - routed = ggml_laguna_moe_combine(ctx, experts, weights_2d); + routed = ggml_moe_combine(ctx, experts, weights_2d); } else { experts = ggml_mul(ctx, experts, weights_3d); diff --git a/server/test/test_moe_combine.cpp b/server/test/test_moe_combine.cpp new file mode 100644 index 000000000..353fa625f --- /dev/null +++ b/server/test/test_moe_combine.cpp @@ -0,0 +1,125 @@ +// Model-neutral correctness coverage for the fused MoE route-weight and +// expert-reduction primitive used by DS4, Laguna, and common hybrid graphs. +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cuda.h" +#include "ggml.h" + +#include +#include +#include +#include +#include + +int main() { + constexpr int n_embd = 257; + constexpr int n_used = 6; + constexpr int n_tokens = 5; + + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, "GPU backend unavailable\n"); + return 1; + } + + ggml_init_params params{}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + ggml_backend_free(backend); + return 1; + } + + ggml_tensor * experts = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, n_embd, n_used, n_tokens); + ggml_tensor * weights = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_used, n_tokens); + ggml_set_input(experts); + ggml_set_input(weights); + ggml_tensor * output = ggml_moe_combine(ctx, experts, weights); + ggml_set_output(output); + + bool ok = output->type == GGML_TYPE_F32 && + output->ne[0] == n_embd && output->ne[1] == n_tokens && + output->ne[2] == 1 && output->ne[3] == 1; + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 16, false); + ggml_build_forward_expand(graph, output); + ggml_gallocr_t allocator = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + ok = ok && allocator && ggml_gallocr_alloc_graph(allocator, graph); + + std::vector expert_data( + (size_t) n_embd * n_used * n_tokens); + std::vector weight_data((size_t) n_used * n_tokens); + std::vector expected((size_t) n_embd * n_tokens); + std::vector actual(expected.size()); + + for (int t = 0; t < n_tokens; ++t) { + for (int e = 0; e < n_used; ++e) { + const size_t route = (size_t) t * n_used + e; + weight_data[route] = ((e + 2 * t) % 5 == 0) + ? 0.0f : 0.05f * (float) (e + 1); + for (int h = 0; h < n_embd; ++h) { + const size_t index = route * n_embd + h; + expert_data[index] = weight_data[route] == 0.0f + ? std::numeric_limits::quiet_NaN() + : 0.001f * (float) ((h % 29) - 14) + + 0.01f * (float) (e - t); + } + } + } + + for (int t = 0; t < n_tokens; ++t) { + for (int h = 0; h < n_embd; ++h) { + float sum = 0.0f; + for (int e = 0; e < n_used; ++e) { + const size_t route = (size_t) t * n_used + e; + const float weight = weight_data[route]; + if (weight == 0.0f) { + if (e == 0) sum = 0.0f; + continue; + } + const float product = expert_data[route * n_embd + h] * weight; + sum = e == 0 ? product : sum + product; + } + expected[(size_t) t * n_embd + h] = sum; + } + } + + if (ok) { + ggml_backend_tensor_set( + experts, expert_data.data(), 0, expert_data.size() * sizeof(float)); + ggml_backend_tensor_set( + weights, weight_data.data(), 0, weight_data.size() * sizeof(float)); + ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + } + if (ok) { + ggml_backend_tensor_get( + output, actual.data(), 0, actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + const float tolerance = 1.0e-6f * + std::max(1.0f, std::abs(expected[i])); + if (!std::isfinite(actual[i]) || + std::abs(actual[i] - expected[i]) > tolerance) { + std::fprintf( + stderr, "mismatch at %zu actual=%g expected=%g\n", + i, actual[i], expected[i]); + ok = false; + break; + } + } + } + + if (ok) { + std::printf( + "moe combine embd=%d experts=%d tokens=%d PASS\n", + n_embd, n_used, n_tokens); + } + + if (allocator) ggml_gallocr_free(allocator); + ggml_free(ctx); + ggml_backend_free(backend); + return ok ? 0 : 1; +} From 2f136181b0b387592e413a81dedc0cd79991e128 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Thu, 27 Aug 2026 16:39:18 +0200 Subject: [PATCH 5/7] perf(ds4): skip redundant direct-topk visibility work --- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 166 ++++++++++++------ 1 file changed, 111 insertions(+), 55 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index 78215b4d3..6b9fed8c3 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -1090,7 +1090,7 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_kernel( // every visible row keeps its original owner thread, dot-product order, // reduction tree, softmax order, and value-accumulation position. template + int VALUES_PER_THREAD, bool ALL_SELECTED_ROWS_VISIBLE> __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( float * dst, const float * q, @@ -1119,6 +1119,7 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( constexpr int D = 512; constexpr int N_THREADS = 256; static_assert(VALUES_PER_THREAD == 2 || VALUES_PER_THREAD == 4); + static_assert(!ALL_SELECTED_ROWS_VISIBLE || INDEXED_MASK); const int t = (int) blockIdx.x; const int h_begin = (int) blockIdx.y * HEADS_PER_BLOCK; const int tid = (int) threadIdx.x; @@ -1132,11 +1133,17 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( float * q_rope_tail = reinterpret_cast( value_bounds + (size_t) HEADS_PER_BLOCK * 4); - const int * token_visibility = visibility_bounds + (size_t) t * 4; - const int mask_raw_first = token_visibility[0]; - const int mask_raw_last = token_visibility[1]; - const int mask_comp_first = token_visibility[2]; - const int mask_comp_last = token_visibility[3]; + const int * token_visibility = ALL_SELECTED_ROWS_VISIBLE + ? nullptr + : visibility_bounds + (size_t) t * 4; + const int mask_raw_first = ALL_SELECTED_ROWS_VISIBLE + ? 0 : token_visibility[0]; + const int mask_raw_last = ALL_SELECTED_ROWS_VISIBLE + ? raw_rows - 1 : token_visibility[1]; + const int mask_comp_first = ALL_SELECTED_ROWS_VISIBLE + ? raw_rows : token_visibility[2]; + const int mask_comp_last = ALL_SELECTED_ROWS_VISIBLE + ? n_kv - 1 : token_visibility[3]; const int * token_indexed_rows = nullptr; const int * token_owner_offsets = nullptr; const int * token_owner_ranks = nullptr; @@ -1184,20 +1191,23 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( local_max[j] = h < n_heads && sinks ? sinks[h] : -3.402823466e38f; } - // Build the full kernel's exact per-head non-zero envelope while the - // softmax weights are emitted. This avoids rescanning every context row - // after softmax without changing the subsequent V accumulation interval. - if (tid < HEADS_PER_BLOCK * 4) { - const int slot = tid & 3; - value_bounds[tid] = slot == 0 - ? raw_rows - : slot == 1 - ? -1 - : slot == 2 - ? (INDEXED_MASK ? indexed_count : n_kv) - : -1; + // Direct top-k decode has already filtered and compacted every selected + // compressed row, and a full raw ring has no masked rows. In that exact + // shape the bounds are known and mask loads plus four shared atomics per + // surviving score are redundant. Other shapes retain the original path. + if constexpr (!ALL_SELECTED_ROWS_VISIBLE) { + if (tid < HEADS_PER_BLOCK * 4) { + const int slot = tid & 3; + value_bounds[tid] = slot == 0 + ? raw_rows + : slot == 1 + ? -1 + : slot == 2 + ? (INDEXED_MASK ? indexed_count : n_kv) + : -1; + } + __syncthreads(); } - __syncthreads(); // Preserve each thread's original r = tid + 256*k order. In indexed mode, // owner_ranks is the exact selected subsequence of that traversal, so the @@ -1244,9 +1254,13 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( } } - const float mask_v = ds4_fa_load( - mask + (size_t) t * n_kv + r); - const bool visible = mask_v > -1.0e20f; + float mask_v = 0.0f; + bool visible = true; + if constexpr (!ALL_SELECTED_ROWS_VISIBLE) { + mask_v = ds4_fa_load( + mask + (size_t) t * n_kv + r); + visible = mask_v > -1.0e20f; + } float dot[HEADS_PER_BLOCK] = {}; if (visible) { const KV * kr = k + (size_t) r * D; @@ -1328,10 +1342,12 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( const float weight = expf(*score - max_score[j]); *score = weight; local_sum[j] += weight; - if (weight != 0.0f) { - int * bounds = value_bounds + 4 * j + (raw_value ? 0 : 2); - atomicMin(bounds + 0, bound_value); - atomicMax(bounds + 1, bound_value); + if constexpr (!ALL_SELECTED_ROWS_VISIBLE) { + if (weight != 0.0f) { + int * bounds = value_bounds + 4 * j + (raw_value ? 0 : 2); + atomicMin(bounds + 0, bound_value); + atomicMax(bounds + 1, bound_value); + } } } } @@ -1364,9 +1380,6 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( inv_denom[j] = 1.0f / reduction[(size_t) j * N_THREADS]; } - // Finish the shared envelope before the value phase reads it. - __syncthreads(); - int raw_first = raw_rows; int raw_last = -1; int comp_first = INDEXED_MASK ? indexed_count : n_kv; @@ -1375,17 +1388,33 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( int head_raw_last[HEADS_PER_BLOCK]; int head_comp_first[HEADS_PER_BLOCK]; int head_comp_last[HEADS_PER_BLOCK]; + if constexpr (ALL_SELECTED_ROWS_VISIBLE) { #pragma unroll - for (int j = 0; j < HEADS_PER_BLOCK; ++j) { - const int * bounds = value_bounds + 4 * j; - head_raw_first[j] = bounds[0]; - head_raw_last[j] = bounds[1]; - head_comp_first[j] = bounds[2]; - head_comp_last[j] = bounds[3]; - raw_first = min(raw_first, head_raw_first[j]); - raw_last = max(raw_last, head_raw_last[j]); - comp_first = min(comp_first, head_comp_first[j]); - comp_last = max(comp_last, head_comp_last[j]); + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + head_raw_first[j] = 0; + head_raw_last[j] = raw_rows - 1; + head_comp_first[j] = 0; + head_comp_last[j] = indexed_count - 1; + } + raw_first = 0; + raw_last = raw_rows - 1; + comp_first = 0; + comp_last = indexed_count - 1; + } else { + // Finish the shared envelope before the value phase reads it. + __syncthreads(); +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int * bounds = value_bounds + 4 * j; + head_raw_first[j] = bounds[0]; + head_raw_last[j] = bounds[1]; + head_comp_first[j] = bounds[2]; + head_comp_last[j] = bounds[3]; + raw_first = min(raw_first, head_raw_first[j]); + raw_last = max(raw_last, head_raw_last[j]); + comp_first = min(comp_first, head_comp_first[j]); + comp_last = max(comp_last, head_comp_last[j]); + } } // One active thread owns adjacent value dimensions. This retains each @@ -1619,7 +1648,8 @@ static bool ds4_launch_flash_attn_d512_grouped( return true; } -template +template static bool ds4_launch_flash_attn_d512_grouped_compact( ggml_tensor * dst, const ggml_tensor * Q, @@ -1649,7 +1679,8 @@ static bool ds4_launch_flash_attn_d512_grouped_compact( const float * forward_rope_coefficients, size_t shmem, cudaStream_t stream) { - GGML_ASSERT(mask && visibility_bounds); + GGML_ASSERT(mask && + (visibility_bounds || ALL_SELECTED_ROWS_VISIBLE)); if constexpr (INDEXED_MASK) { GGML_ASSERT(indexed_rows && indexed_counts && indexed_owner_offsets && indexed_owner_ranks); @@ -1659,7 +1690,8 @@ static bool ds4_launch_flash_attn_d512_grouped_compact( (unsigned) (n_heads / HEADS_PER_BLOCK), 1); if (kv_f16 && mask->type == GGML_TYPE_F16) { ds4_flash_attn_d512_shared_kv_grouped_compact_kernel< - half, half, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD> + half, half, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD, + ALL_SELECTED_ROWS_VISIBLE> <<>>( (float *) dst->data, (const float *) Q->data, q_stride_token, q_stride_head, @@ -1674,7 +1706,8 @@ static bool ds4_launch_flash_attn_d512_grouped_compact( forward_rope_coefficients); } else if (kv_f32 && mask->type == GGML_TYPE_F32) { ds4_flash_attn_d512_shared_kv_grouped_compact_kernel< - float, float, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD> + float, float, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD, + ALL_SELECTED_ROWS_VISIBLE> <<>>( (float *) dst->data, (const float *) Q->data, q_stride_token, q_stride_head, @@ -1689,7 +1722,8 @@ static bool ds4_launch_flash_attn_d512_grouped_compact( forward_rope_coefficients); } else if (kv_f32 && mask->type == GGML_TYPE_F16) { ds4_flash_attn_d512_shared_kv_grouped_compact_kernel< - float, half, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD> + float, half, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD, + ALL_SELECTED_ROWS_VISIBLE> <<>>( (float *) dst->data, (const float *) Q->data, q_stride_token, q_stride_head, @@ -1971,9 +2005,14 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( (indexed_mask || group4_shmem > 24 * 1024) && compact_group4_shmem <= 24 * 1024; if (compact_group4) { + const bool direct_topk_decode = indexed_mask && indexer_topk && + n_tokens == 1 && raw_rows == raw_window && + inverse_rope.kv_start >= raw_rows && + getenv("GGML_DS4_FA_DIRECT_TOPK_VISIBLE") != nullptr; ggml_cuda_pool_alloc visibility_bounds_alloc(ctx.pool()); - int * visibility_bounds = visibility_bounds_alloc.alloc( - (size_t) n_tokens * 4); + int * visibility_bounds = direct_topk_decode + ? nullptr + : visibility_bounds_alloc.alloc((size_t) n_tokens * 4); ggml_cuda_pool_alloc indexed_rows_alloc(ctx.pool()); ggml_cuda_pool_alloc indexed_counts_alloc(ctx.pool()); ggml_cuda_pool_alloc indexed_owner_offsets_alloc(ctx.pool()); @@ -2037,17 +2076,34 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( } CUDA_CHECK(cudaGetLastError()); } - if (mask->type == GGML_TYPE_F16) { - ds4_fa_visibility_bounds_kernel<<>>( - (const half *) mask->data, visibility_bounds, - n_tokens, n_kv, raw_rows); - } else { - ds4_fa_visibility_bounds_kernel<<>>( - (const float *) mask->data, visibility_bounds, - n_tokens, n_kv, raw_rows); + if (!direct_topk_decode) { + if (mask->type == GGML_TYPE_F16) { + ds4_fa_visibility_bounds_kernel<<>>( + (const half *) mask->data, visibility_bounds, + n_tokens, n_kv, raw_rows); + } else { + ds4_fa_visibility_bounds_kernel<<>>( + (const float *) mask->data, visibility_bounds, + n_tokens, n_kv, raw_rows); + } + CUDA_CHECK(cudaGetLastError()); } - CUDA_CHECK(cudaGetLastError()); if (indexed_mask) { + if (direct_topk_decode) { + return ds4_launch_flash_attn_d512_grouped_compact< + group4, true, 4, true>( + dst, Q, K, V, mask, sinks, kv_f16, kv_f32, + n_tokens, n_heads, n_kv, scale, raw_rows, + raw_window, compact_score_stride, nullptr, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, + indexed_capacity, + q_stride_token, q_stride_head, + inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients, + compact_group4_shmem, stream); + } return ds4_launch_flash_attn_d512_grouped_compact< group4, true, 4>( dst, Q, K, V, mask, sinks, kv_f16, kv_f32, From 0f1264cf284677bef21fb0fecf4c106a51760cfe Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 28 Aug 2026 17:25:00 +0200 Subject: [PATCH 6/7] hip: split indexed MLA decode across KV blocks --- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 396 ++++++++++++++++++ 1 file changed, 396 insertions(+) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index 6b9fed8c3..7d33d5f1a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -1584,6 +1584,343 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( } } +// Split-KV decode for indexed MLA. A single grouped block leaves most of a +// wide RDNA GPU idle when q is small: DS4 has 64 heads, so the four-head +// kernel exposes only 16 blocks per layer. Split the bounded raw+top-k row +// set across four blocks, retain online-softmax state per split, then combine +// the states in a second kernel. The implementation is expressed in terms of +// the D512 latent-attention contract rather than model weights, so future MLA +// models using the same layout can reuse it. +template +__global__ static void ds4_flash_attn_d512_indexed_split_stage1_kernel( + float * partial, + const float * q, + size_t q_stride_token, + size_t q_stride_head, + const KV * k, + const KV * v, + const Mask * mask, + const float * sinks, + int n_tokens, + int n_heads, + int n_kv, + float scale, + int raw_rows, + int split_stride, + const int * visibility_bounds, + const int * indexed_rows, + const int * indexed_counts, + int indexed_capacity, + ds4_inverse_rope_params inverse_rope, + const float * forward_rope_coefficients) { + constexpr int D = 512; + constexpr int N_THREADS = 256; + constexpr int VALUES_PER_THREAD = 4; + + const int token = (int) blockIdx.x; + const int head_begin = (int) blockIdx.y * HEADS_PER_BLOCK; + const int split = (int) blockIdx.z; + const int tid = (int) threadIdx.x; + if (token >= n_tokens || head_begin >= n_heads) return; + + extern __shared__ float scratch[]; + float * scores = scratch; + float * reduction = scores + (size_t) HEADS_PER_BLOCK * split_stride; + float * q_rope_tail = reduction + (size_t) HEADS_PER_BLOCK * N_THREADS; + + const int raw_first = visibility_bounds + ? visibility_bounds[(size_t) token * 4 + 0] : 0; + const int raw_last = visibility_bounds + ? visibility_bounds[(size_t) token * 4 + 1] : raw_rows - 1; + const int raw_count = raw_last >= raw_first + ? raw_last - raw_first + 1 : 0; + const int indexed_count = indexed_counts[token]; + const int total_rows = raw_count + indexed_count; + const int rows_per_split = (total_rows + N_SPLITS - 1) / N_SPLITS; + const int split_begin = split * rows_per_split; + const int split_end = min(total_rows, split_begin + rows_per_split); + const int * token_indexed_rows = indexed_rows + + (size_t) token * indexed_capacity; + + const float * qh[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + qh[j] = q + (size_t) token * q_stride_token + + (size_t) (head_begin + j) * q_stride_head; + } + + if (inverse_rope.forward_q_enabled) { + for (int index = tid; index < HEADS_PER_BLOCK * 64; + index += N_THREADS) { + const int j = index / 64; + const int tail_d = index % 64; + const int pair = tail_d >> 1; + const float x0 = qh[j][D - 64 + 2 * pair + 0]; + const float x1 = qh[j][D - 64 + 2 * pair + 1]; + const size_t coefficient = + ((size_t) token * 32 + (size_t) pair) * 2; + const float cos_theta = forward_rope_coefficients[coefficient + 0]; + const float sin_theta = forward_rope_coefficients[coefficient + 1]; + q_rope_tail[(size_t) j * 64 + tail_d] = (tail_d & 1) + ? x0 * sin_theta + x1 * cos_theta + : x0 * cos_theta - x1 * sin_theta; + } + __syncthreads(); + } + + float local_max[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + local_max[j] = split == 0 && sinks + ? sinks[head_begin + j] : -3.402823466e38f; + } + + for (int slot = split_begin + tid; slot < split_end; + slot += N_THREADS) { + const int row = slot < raw_count + ? raw_first + slot + : token_indexed_rows[slot - raw_count]; + const float mask_v = ds4_fa_load( + mask + (size_t) token * n_kv + row); + const KV * kr = k + (size_t) row * D; + float dot[HEADS_PER_BLOCK] = {}; +#pragma unroll + for (int d = 0; d < D; ++d) { + const float kv = ds4_fa_load(kr + d); +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const float qv = + inverse_rope.forward_q_enabled && d >= D - 64 + ? q_rope_tail[(size_t) j * 64 + d - (D - 64)] + : qh[j][d]; + dot[j] += qv * kv; + } + } + const int score_index = slot - split_begin; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const float score = dot[j] * scale + mask_v; + scores[(size_t) j * split_stride + score_index] = score; + local_max[j] = fmaxf(local_max[j], score); + } + } + +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + reduction[(size_t) j * N_THREADS + tid] = local_max[j]; + } + __syncthreads(); + for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { + if (tid < stride) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * row = reduction + (size_t) j * N_THREADS; + row[tid] = fmaxf(row[tid], row[tid + stride]); + } + } + __syncthreads(); + } + + float split_max[HEADS_PER_BLOCK]; + float local_sum[HEADS_PER_BLOCK] = {}; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + split_max[j] = reduction[(size_t) j * N_THREADS]; + } + for (int slot = split_begin + tid; slot < split_end; + slot += N_THREADS) { + const int score_index = slot - split_begin; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * score = scores + (size_t) j * split_stride + score_index; + const float weight = expf(*score - split_max[j]); + *score = weight; + local_sum[j] += weight; + } + } + if (tid == 0 && split == 0 && sinks) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + local_sum[j] += expf(sinks[head_begin + j] - split_max[j]); + } + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + reduction[(size_t) j * N_THREADS + tid] = local_sum[j]; + } + __syncthreads(); + for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { + if (tid < stride) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * row = reduction + (size_t) j * N_THREADS; + row[tid] += row[tid + stride]; + } + } + __syncthreads(); + } + + const int d0 = VALUES_PER_THREAD * tid; + if (d0 < D) { + float acc0[HEADS_PER_BLOCK] = {}; + float acc1[HEADS_PER_BLOCK] = {}; + float acc2[HEADS_PER_BLOCK] = {}; + float acc3[HEADS_PER_BLOCK] = {}; + for (int slot = split_begin; slot < split_end; ++slot) { + const int row = slot < raw_count + ? raw_first + slot + : token_indexed_rows[slot - raw_count]; + float vv0, vv1, vv2, vv3; + ds4_fa_load_quad(v + (size_t) row * D + d0, + vv0, vv1, vv2, vv3); + const int score_index = slot - split_begin; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const float weight = + scores[(size_t) j * split_stride + score_index]; + acc0[j] += weight * vv0; + acc1[j] += weight * vv1; + acc2[j] += weight * vv2; + acc3[j] += weight * vv3; + } + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * out = partial + + (((size_t) token * n_heads + head_begin + j) * N_SPLITS + + split) * (D + 2); + out[d0 + 0] = acc0[j]; + out[d0 + 1] = acc1[j]; + out[d0 + 2] = acc2[j]; + out[d0 + 3] = acc3[j]; + } + } + if (tid < HEADS_PER_BLOCK) { + float * out = partial + + (((size_t) token * n_heads + head_begin + tid) * N_SPLITS + + split) * (D + 2); + out[D + 0] = split_max[tid]; + out[D + 1] = reduction[(size_t) tid * N_THREADS]; + } +} + +template +__global__ static void ds4_flash_attn_d512_indexed_split_stage2_kernel( + float * dst, + const float * partial, + int n_tokens, + int n_heads, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients) { + constexpr int D = 512; + const int token = (int) blockIdx.x; + const int head = (int) blockIdx.y; + const int tid = (int) threadIdx.x; + if (token >= n_tokens || head >= n_heads) return; + + __shared__ float global_max; + __shared__ float inverse_denom; + const float * head_partial = partial + + ((size_t) token * n_heads + head) * N_SPLITS * (D + 2); + if (tid == 0) { + float max_value = -3.402823466e38f; +#pragma unroll + for (int split = 0; split < N_SPLITS; ++split) { + max_value = fmaxf(max_value, + head_partial[(size_t) split * (D + 2) + D]); + } + float denom = 0.0f; +#pragma unroll + for (int split = 0; split < N_SPLITS; ++split) { + const float * state = head_partial + (size_t) split * (D + 2); + denom += state[D + 1] * expf(state[D] - max_value); + } + global_max = max_value; + inverse_denom = 1.0f / denom; + } + __syncthreads(); + + const int d0 = 2 * tid; + float x0 = 0.0f; + float x1 = 0.0f; +#pragma unroll + for (int split = 0; split < N_SPLITS; ++split) { + const float * state = head_partial + (size_t) split * (D + 2); + const float rescale = expf(state[D] - global_max); + x0 += state[d0 + 0] * rescale; + x1 += state[d0 + 1] * rescale; + } + x0 *= inverse_denom; + x1 *= inverse_denom; + + float * out = dst + ((size_t) token * n_heads + head) * D + d0; + if (inverse_rope.enabled && d0 >= D - 64) { + const int pair = (d0 - (D - 64)) / 2; + const size_t coefficient = + ((size_t) token * 32 + (size_t) pair) * 2; + const float cos_theta = inverse_rope_coefficients[coefficient + 0]; + const float sin_theta = inverse_rope_coefficients[coefficient + 1]; + float y0; + float y1; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, y0, y1); + out[0] = y0; + out[1] = y1; + } else { + out[0] = x0; + out[1] = x1; + } +} + +template +static void ds4_launch_flash_attn_d512_indexed_split( + float * dst, + float * partial, + const float * q, + size_t q_stride_token, + size_t q_stride_head, + const KV * k, + const KV * v, + const Mask * mask, + const float * sinks, + int n_tokens, + int n_heads, + int n_kv, + float scale, + int raw_rows, + int split_stride, + const int * visibility_bounds, + const int * indexed_rows, + const int * indexed_counts, + int indexed_capacity, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients, + const float * forward_rope_coefficients, + cudaStream_t stream) { + constexpr int HEADS_PER_BLOCK = 4; + const size_t shmem = + ((size_t) HEADS_PER_BLOCK * split_stride + + (size_t) HEADS_PER_BLOCK * 256 + + (inverse_rope.forward_q_enabled + ? (size_t) HEADS_PER_BLOCK * 64 : 0)) * sizeof(float); + const dim3 stage1_grid( + (unsigned) n_tokens, + (unsigned) (n_heads / HEADS_PER_BLOCK), + (unsigned) N_SPLITS); + ds4_flash_attn_d512_indexed_split_stage1_kernel< + KV, Mask, HEADS_PER_BLOCK, N_SPLITS> + <<>>( + partial, q, q_stride_token, q_stride_head, k, v, mask, sinks, + n_tokens, n_heads, n_kv, scale, raw_rows, split_stride, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, forward_rope_coefficients); + const dim3 stage2_grid((unsigned) n_tokens, (unsigned) n_heads, 1); + ds4_flash_attn_d512_indexed_split_stage2_kernel + <<>>( + dst, partial, n_tokens, n_heads, inverse_rope, + inverse_rope_coefficients); +} + template static bool ds4_launch_flash_attn_d512_grouped( ggml_tensor * dst, @@ -2088,6 +2425,65 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( } CUDA_CHECK(cudaGetLastError()); } + // Experimental AITER-style split-KV schedule, implemented directly in + // the native HIP backend. Keep it opt-in until matched output and + // throughput checks show that the extra reduction pays for itself. + constexpr int split_kv_max_decode_tokens = 8; + if (indexed_mask && n_tokens <= split_kv_max_decode_tokens && + (getenv("GGML_CUDA_MLA_SPLIT_KV") != nullptr || + getenv("GGML_DS4_FA_SPLIT_KV") != nullptr)) { + constexpr int split_count = 2; + const int split_stride = + (raw_window + indexed_capacity + split_count - 1) / + split_count; + ggml_cuda_pool_alloc partial_alloc(ctx.pool()); + float * partial = partial_alloc.alloc( + (size_t) n_tokens * n_heads * split_count * (512 + 2)); + if (kv_f16 && mask->type == GGML_TYPE_F16) { + ds4_launch_flash_attn_d512_indexed_split< + half, half, split_count>( + (float *) dst->data, partial, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, (const half *) V->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, split_stride, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, forward_rope_coefficients, + stream); + } else if (kv_f32 && mask->type == GGML_TYPE_F32) { + ds4_launch_flash_attn_d512_indexed_split< + float, float, split_count>( + (float *) dst->data, partial, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + (const float *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, split_stride, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, forward_rope_coefficients, + stream); + } else if (kv_f32 && mask->type == GGML_TYPE_F16) { + ds4_launch_flash_attn_d512_indexed_split< + float, half, split_count>( + (float *) dst->data, partial, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, split_stride, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, forward_rope_coefficients, + stream); + } else { + return false; + } + CUDA_CHECK(cudaGetLastError()); + return true; + } if (indexed_mask) { if (direct_topk_decode) { return ds4_launch_flash_attn_d512_grouped_compact< From fb2f462cb75b76fb542e2cb4c9344c723f04cab9 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sat, 29 Aug 2026 00:32:08 +0200 Subject: [PATCH 7/7] perf(rocmfpx): reuse activations across paired rows --- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 19 ++- .../ggml/src/ggml-cuda/rocmfp2_mix.cu | 102 +++++++++++-- .../ggml/src/ggml-cuda/rocmfp3_mix.cu | 134 +++++++++++++----- server/test/bench_rocmfp_mix_gateup_glu.cpp | 55 +++++-- 4 files changed, 249 insertions(+), 61 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index 7d33d5f1a..6d6c991f8 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -2425,13 +2425,22 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( } CUDA_CHECK(cudaGetLastError()); } - // Experimental AITER-style split-KV schedule, implemented directly in - // the native HIP backend. Keep it opt-in until matched output and - // throughput checks show that the extra reduction pays for itself. + // AITER-style split-KV schedule, implemented directly in the native HIP + // backend. Matched Strix Halo profiling showed a bit-identical output, + // about -58% attention time and +2-3% decode throughput, so make it the + // gfx1151 default. Other devices remain opt-in until measured. constexpr int split_kv_max_decode_tokens = 8; + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const bool split_kv_default = + cc == GGML_CUDA_CC_OFFSET_AMD + 0x1151; + const bool split_kv_forced = + getenv("GGML_CUDA_MLA_SPLIT_KV") != nullptr || + getenv("GGML_DS4_FA_SPLIT_KV") != nullptr; + const bool split_kv_disabled = + getenv("GGML_CUDA_MLA_NO_SPLIT_KV") != nullptr || + getenv("GGML_DS4_FA_NO_SPLIT_KV") != nullptr; if (indexed_mask && n_tokens <= split_kv_max_decode_tokens && - (getenv("GGML_CUDA_MLA_SPLIT_KV") != nullptr || - getenv("GGML_DS4_FA_SPLIT_KV") != nullptr)) { + !split_kv_disabled && (split_kv_forced || split_kv_default)) { constexpr int split_count = 2; const int split_stride = (raw_window + indexed_capacity + split_count - 1) / diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu index 637f6d27f..741f08425 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu @@ -535,6 +535,78 @@ __device__ __forceinline__ void mix_block_accum( } } +// Accumulate the same activation block into two output rows. Keeping the two +// row folds in one j-loop makes activation reuse explicit: xc[col0 + j] has one +// live range and feeds both independent accumulators. Each accumulator still +// sees exactly the same ascending-j sequence and expression as +// mix_block_accum(), so this does not reassociate either dot product. +__device__ __forceinline__ void mix_block_accum2( + const uint8_t * __restrict__ b0, const uint8_t * __restrict__ b1, + const float * __restrict__ xc, int col0, + int mode, const float * __restrict__ lut, float & acc0, float & acc1) { +#if defined(__HIP_PLATFORM_AMD__) && defined(__gfx1151__) + uint64_t codes0, codes1; + uint16_t meta0, meta1; + MIX_MEMCPY(&codes0, b0, sizeof(codes0)); + MIX_MEMCPY(&meta0, b0 + MIX_QS, sizeof(meta0)); + MIX_MEMCPY(&codes1, b1, sizeof(codes1)); + MIX_MEMCPY(&meta1, b1 + MIX_QS, sizeof(meta1)); + const uint8_t m00 = (uint8_t) meta0; + const uint8_t m01 = (uint8_t) (meta0 >> 8); + const uint8_t m10 = (uint8_t) meta1; + const uint8_t m11 = (uint8_t) (meta1 >> 8); +#else + const uintptr_t addr0 = (uintptr_t) b0; + const uintptr_t addr1 = (uintptr_t) b1; + const uint8_t * base80 = (const uint8_t *) MIX_ASSUME_ALIGNED( + (const void *) (addr0 & ~(uintptr_t) 7), 8); + const uint8_t * base81 = (const uint8_t *) MIX_ASSUME_ALIGNED( + (const void *) (addr1 & ~(uintptr_t) 7), 8); + const int sh0 = (int) (addr0 & 7) * 8; + const int sh1 = (int) (addr1 & 7) * 8; + uint64_t lo0, hi0, lo1, hi1; + MIX_MEMCPY(&lo0, base80, 8); + MIX_MEMCPY(&hi0, base80 + 8, 8); + MIX_MEMCPY(&lo1, base81, 8); + MIX_MEMCPY(&hi1, base81 + 8, 8); + const uint64_t codes0 = (sh0 == 0) ? lo0 : ((lo0 >> sh0) | (hi0 << (64 - sh0))); + const uint64_t codes1 = (sh1 == 0) ? lo1 : ((lo1 >> sh1) | (hi1 << (64 - sh1))); + const uint8_t m00 = (uint8_t) (hi0 >> sh0); + const uint8_t m01 = (uint8_t) (hi0 >> (sh0 + 8)); + const uint8_t m10 = (uint8_t) (hi1 >> sh1); + const uint8_t m11 = (uint8_t) (hi1 >> (sh1 + 8)); +#endif + if (mode == 0) { + const float s00 = mix_ue4m3(m00), s01 = mix_ue4m3(m01); + const float s10 = mix_ue4m3(m10), s11 = mix_ue4m3(m11); + #pragma unroll + for (int j = 0; j < MIX_QK; ++j) { + const float x = xc[col0 + j]; + const float rs0 = (j < MIX_QK/2) ? s00 : s01; + const float rs1 = (j < MIX_QK/2) ? s10 : s11; + acc0 += rs0 * mix_fp2_fixed(mix_fp2_code_u64(codes0, j)) * x; + acc1 += rs1 * mix_fp2_fixed(mix_fp2_code_u64(codes1, j)) * x; + } + } else { + const float s00 = mix_ue4m3(m00 & 0x7F), s01 = mix_ue4m3(m01 & 0x7F); + const float s10 = mix_ue4m3(m10 & 0x7F), s11 = mix_ue4m3(m11 & 0x7F); + const float * bk00 = lut + (m00 >> 7) * MIX_K; + const float * bk01 = lut + (m01 >> 7) * MIX_K; + const float * bk10 = lut + (m10 >> 7) * MIX_K; + const float * bk11 = lut + (m11 >> 7) * MIX_K; + #pragma unroll + for (int j = 0; j < MIX_QK; ++j) { + const float x = xc[col0 + j]; + const float rs0 = (j < MIX_QK/2) ? s00 : s01; + const float rs1 = (j < MIX_QK/2) ? s10 : s11; + const float * rbk0 = (j < MIX_QK/2) ? bk00 : bk01; + const float * rbk1 = (j < MIX_QK/2) ? bk10 : bk11; + acc0 += rs0 * rbk0[mix_fp2_code_u64(codes0, j)] * x; + acc1 += rs1 * rbk1[mix_fp2_code_u64(codes1, j)] * x; + } + } +} + // The lane's block loop is unrolled by MIX_UNROLL into a SINGLE accumulator kept // in the exact original block order (acc += dot(blk), stride MIX_WARP), so the // f32 output is bit-for-bit identical to the un-unrolled path — required because @@ -688,13 +760,15 @@ __global__ void mix_matvec_rocmfp2_slice_kernel( #pragma unroll for (int u = 0; u < MIX_UNROLL; ++u) { const int b = blk + u * MIX_WARP; - mix_block_accum(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, s_lut, acc1); + mix_block_accum2(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, + rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, + xcol, b * MIX_QK, mode, s_lut, acc0, acc1); } } for (; blk < nb; blk += MIX_WARP) { - mix_block_accum(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc1); + mix_block_accum2(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, + rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, + xcol, blk * MIX_QK, mode, s_lut, acc0, acc1); } #pragma unroll for (int off = MIX_WARP/2; off > 0; off >>= 1) { @@ -828,20 +902,24 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( #pragma unroll for (int u = 0; u < MIX_UNROLL; ++u) { const int b = blk + u * MIX_WARP; - mix_block_accum(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, s_lut, acc1); + mix_block_accum2(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, + rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, + xcol, b * MIX_QK, mode, s_lut, acc0, acc1); if (DUAL_GLU) { - mix_block_accum(growbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); + mix_block_accum2(growbase0 + (int64_t) b * MIX_BLOCK_BYTES, + growbase1 + (int64_t) b * MIX_BLOCK_BYTES, + xcol, b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0, gacc1); } } } for (; blk < nb; blk += MIX_WARP) { - mix_block_accum(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc1); + mix_block_accum2(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, + rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, + xcol, blk * MIX_QK, mode, s_lut, acc0, acc1); if (DUAL_GLU) { - mix_block_accum(growbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); + mix_block_accum2(growbase0 + (int64_t) blk * MIX_BLOCK_BYTES, + growbase1 + (int64_t) blk * MIX_BLOCK_BYTES, + xcol, blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0, gacc1); } } #pragma unroll diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu index d6905334f..bb0a38b8e 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu @@ -481,6 +481,64 @@ __device__ __forceinline__ void mix_block_accum( } } +// Two-row variant of mix_block_accum. The activation value for each j is +// loaded once and feeds two independent row accumulators. Each accumulator +// retains the original ascending-j expression and rounding order. +__device__ __forceinline__ void mix_block_accum2( + const uint8_t * __restrict__ b0, const uint8_t * __restrict__ b1, + const float * __restrict__ xc, int col0, + int mode, const float * __restrict__ lut, float & acc0, float & acc1) { +#if defined(__HIP_PLATFORM_AMD__) && defined(__gfx1151__) + MixFp3Words qs0, qs1; + uint16_t meta0, meta1; + MIX_MEMCPY(&qs0.lo, b0, sizeof(qs0.lo)); + MIX_MEMCPY(&qs0.hi, b0 + sizeof(qs0.lo), sizeof(qs0.hi)); + MIX_MEMCPY(&meta0, b0 + MIX_QS, sizeof(meta0)); + MIX_MEMCPY(&qs1.lo, b1, sizeof(qs1.lo)); + MIX_MEMCPY(&qs1.hi, b1 + sizeof(qs1.lo), sizeof(qs1.hi)); + MIX_MEMCPY(&meta1, b1 + MIX_QS, sizeof(meta1)); + const uint8_t m00 = (uint8_t) meta0, m01 = (uint8_t) (meta0 >> 8); + const uint8_t m10 = (uint8_t) meta1, m11 = (uint8_t) (meta1 >> 8); +#else + const uint8_t * ba0 = (const uint8_t *) MIX_ASSUME_ALIGNED(b0, 2); + const uint8_t * ba1 = (const uint8_t *) MIX_ASSUME_ALIGNED(b1, 2); + uint8_t qs0[MIX_BLOCK_BYTES], qs1[MIX_BLOCK_BYTES]; + MIX_MEMCPY(qs0, ba0, MIX_BLOCK_BYTES); + MIX_MEMCPY(qs1, ba1, MIX_BLOCK_BYTES); + const uint8_t m00 = qs0[MIX_QS + 0], m01 = qs0[MIX_QS + 1]; + const uint8_t m10 = qs1[MIX_QS + 0], m11 = qs1[MIX_QS + 1]; +#endif + if (mode == 0) { + const float s00 = mix_ue4m3(m00), s01 = mix_ue4m3(m01); + const float s10 = mix_ue4m3(m10), s11 = mix_ue4m3(m11); + #pragma unroll + for (int j = 0; j < MIX_QK; ++j) { + const float x = xc[col0 + j]; + const float rs0 = (j < MIX_QK/2) ? s00 : s01; + const float rs1 = (j < MIX_QK/2) ? s10 : s11; + acc0 = fmaf(rs0 * mix_fp3_fixed(mix_fp3_code(qs0, j)), x, acc0); + acc1 = fmaf(rs1 * mix_fp3_fixed(mix_fp3_code(qs1, j)), x, acc1); + } + } else { + const float s00 = mix_ue4m3(m00 & 0x7F), s01 = mix_ue4m3(m01 & 0x7F); + const float s10 = mix_ue4m3(m10 & 0x7F), s11 = mix_ue4m3(m11 & 0x7F); + const float * bk00 = lut + (m00 >> 7) * MIX_K; + const float * bk01 = lut + (m01 >> 7) * MIX_K; + const float * bk10 = lut + (m10 >> 7) * MIX_K; + const float * bk11 = lut + (m11 >> 7) * MIX_K; + #pragma unroll + for (int j = 0; j < MIX_QK; ++j) { + const float x = xc[col0 + j]; + const float rs0 = (j < MIX_QK/2) ? s00 : s01; + const float rs1 = (j < MIX_QK/2) ? s10 : s11; + const float * rbk0 = (j < MIX_QK/2) ? bk00 : bk01; + const float * rbk1 = (j < MIX_QK/2) ? bk10 : bk11; + acc0 = fmaf(rs0 * rbk0[mix_fp3_code(qs0, j)], x, acc0); + acc1 = fmaf(rs1 * rbk1[mix_fp3_code(qs1, j)], x, acc1); + } + } +} + // The lane's block loop is unrolled by MIX_UNROLL into a SINGLE accumulator kept // in the exact original block order (acc += dot(blk), stride MIX_WARP), so the // f32 output is bit-for-bit identical to the un-unrolled path — required because @@ -610,30 +668,31 @@ __global__ void mix_matvec_rocmfp3_slice_kernel( // src1 is [in, ne11, ntok]; the get_rows-equivalent row for (slot, token) // is token*ne11 + slot%ne11 — i.e. token column + the slot%ne11 broadcast. const float * xcol = src1 + (int64_t) token * src1_s2 + (int64_t) (slot % ne11) * src1_s1; - // Two output rows in one warp. Each row is folded by the SAME mix_block_accum - // that the single-row path uses (byte-identical inlined body, same fixed j - // order, same acc-add chain) so acc0/acc1 are bit-for-bit identical to the - // single-row kernel's output for those rows. The two calls per block share the - // same __restrict__ xcol + col0, so the compiler CSEs the strided activation - // loads to one issue per element — halving activation LSU issue on this partly - // load-instruction-bound matvec — WITHOUT reordering either row's summation. + // Two output rows in one warp. mix_block_accum2 explicitly loads each + // activation once and feeds both independent row accumulators. Each row keeps + // the same fixed j order and acc-add chain as the single-row path. float acc0 = 0.0f, acc1 = 0.0f; int blk = lane; for (; blk + 3 * MIX_WARP < nb; blk += MIX_UNROLL * MIX_WARP) { const int b0 = blk, b1 = blk + MIX_WARP; const int b2 = blk + 2 * MIX_WARP, b3 = blk + 3 * MIX_WARP; - mix_block_accum(rowbase0 + (int64_t) b0 * MIX_BLOCK_BYTES, xcol, b0 * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b0 * MIX_BLOCK_BYTES, xcol, b0 * MIX_QK, mode, s_lut, acc1); - mix_block_accum(rowbase0 + (int64_t) b1 * MIX_BLOCK_BYTES, xcol, b1 * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b1 * MIX_BLOCK_BYTES, xcol, b1 * MIX_QK, mode, s_lut, acc1); - mix_block_accum(rowbase0 + (int64_t) b2 * MIX_BLOCK_BYTES, xcol, b2 * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b2 * MIX_BLOCK_BYTES, xcol, b2 * MIX_QK, mode, s_lut, acc1); - mix_block_accum(rowbase0 + (int64_t) b3 * MIX_BLOCK_BYTES, xcol, b3 * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b3 * MIX_BLOCK_BYTES, xcol, b3 * MIX_QK, mode, s_lut, acc1); + mix_block_accum2(rowbase0 + (int64_t) b0 * MIX_BLOCK_BYTES, + rowbase1 + (int64_t) b0 * MIX_BLOCK_BYTES, + xcol, b0 * MIX_QK, mode, s_lut, acc0, acc1); + mix_block_accum2(rowbase0 + (int64_t) b1 * MIX_BLOCK_BYTES, + rowbase1 + (int64_t) b1 * MIX_BLOCK_BYTES, + xcol, b1 * MIX_QK, mode, s_lut, acc0, acc1); + mix_block_accum2(rowbase0 + (int64_t) b2 * MIX_BLOCK_BYTES, + rowbase1 + (int64_t) b2 * MIX_BLOCK_BYTES, + xcol, b2 * MIX_QK, mode, s_lut, acc0, acc1); + mix_block_accum2(rowbase0 + (int64_t) b3 * MIX_BLOCK_BYTES, + rowbase1 + (int64_t) b3 * MIX_BLOCK_BYTES, + xcol, b3 * MIX_QK, mode, s_lut, acc0, acc1); } for (; blk < nb; blk += MIX_WARP) { - mix_block_accum(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc1); + mix_block_accum2(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, + rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, + xcol, blk * MIX_QK, mode, s_lut, acc0, acc1); } #pragma unroll for (int off = MIX_WARP/2; off > 0; off >>= 1) { @@ -720,13 +779,9 @@ __global__ void mix_matvec_rocmfp3_moe_kernel( // src1 is [in, ne11, ntok]; the get_rows-equivalent row for (slot, token) // is token*ne11 + slot%ne11 — i.e. token column + the slot%ne11 broadcast. const float * xcol = src1 + (int64_t) token * src1_s2 + (int64_t) (slot % ne11) * src1_s1; - // Two output rows in one warp. Each row is folded by the SAME mix_block_accum - // that the single-row path uses (byte-identical inlined body, same fixed j - // order, same acc-add chain) so acc0/acc1 are bit-for-bit identical to the - // single-row kernel's output for those rows. The two calls per block share the - // same __restrict__ xcol + col0, so the compiler CSEs the strided activation - // loads to one issue per element — halving activation LSU issue on this partly - // load-instruction-bound matvec — WITHOUT reordering either row's summation. + // Explicitly share each activation load across the two independent output-row + // folds in the ordinary projection. The fused-GLU specialization below keeps + // its lower-register path because four paired accumulators reduce occupancy. float acc0 = 0.0f, acc1 = 0.0f; float gacc0 = 0.0f, gacc1 = 0.0f; // same block order as acc*, so bit-identical per row int blk = lane; @@ -735,24 +790,39 @@ __global__ void mix_matvec_rocmfp3_moe_kernel( #pragma unroll for (int u = 0; u < MIX_MOE_UNROLL; ++u) { const int b = blk + u * MIX_WARP; - mix_block_accum(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, - b * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, - b * MIX_QK, mode, s_lut, acc1); if (FUSE_GLU) { + // Four independent row folds have too much register pressure for + // the paired helper on gfx1151. Keep the exact lower-pressure + // path for the fused gate/up specialization. + mix_block_accum(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, + b * MIX_QK, mode, s_lut, acc0); + mix_block_accum(rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, + b * MIX_QK, mode, s_lut, acc1); mix_block_accum(growbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); mix_block_accum(growbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); + } else { + mix_block_accum2(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, + rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, + xcol, b * MIX_QK, mode, s_lut, acc0, acc1); } } } for (; blk < nb; blk += MIX_WARP) { - mix_block_accum(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc1); if (FUSE_GLU) { - mix_block_accum(growbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); + mix_block_accum(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, + blk * MIX_QK, mode, s_lut, acc0); + mix_block_accum(rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, + blk * MIX_QK, mode, s_lut, acc1); + mix_block_accum(growbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, + blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); + mix_block_accum(growbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, + blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); + } else { + mix_block_accum2(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, + rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, + xcol, blk * MIX_QK, mode, s_lut, acc0, acc1); } } #pragma unroll diff --git a/server/test/bench_rocmfp_mix_gateup_glu.cpp b/server/test/bench_rocmfp_mix_gateup_glu.cpp index 692070f2e..28782cbda 100644 --- a/server/test/bench_rocmfp_mix_gateup_glu.cpp +++ b/server/test/bench_rocmfp_mix_gateup_glu.cpp @@ -72,6 +72,15 @@ static double median(std::vector v) { return v[v.size() / 2]; } +static uint64_t fnv1a64(const std::vector & values) { + uint64_t hash = 1469598103934665603ull; + const uint8_t * bytes = reinterpret_cast(values.data()); + for (size_t i = 0; i < sizeof(float) * values.size(); ++i) { + hash = (hash ^ bytes[i]) * 1099511628211ull; + } + return hash; +} + int main(int argc, char ** argv) { int ndev = 0; if (hipGetDeviceCount(&ndev) != hipSuccess || ndev == 0) { @@ -85,8 +94,11 @@ int main(int argc, char ** argv) { const int in = 4096, out = 2048, n_experts = 8, n_used = 6; const int ntok = argc > 1 ? std::atoi(argv[1]) : 4; const bool fp3 = argc > 2 && std::strcmp(argv[2], "q3") == 0; - if (ntok <= 0 || ntok > 16 || (argc > 2 && !fp3 && std::strcmp(argv[2], "q2") != 0)) { - std::fprintf(stderr, "usage: %s [tokens:1..16] [q2|q3]\n", argv[0]); + const bool fixed_levels = argc > 3 && std::strcmp(argv[3], "fixed") == 0; + if (ntok <= 0 || ntok > 16 || + (argc > 2 && !fp3 && std::strcmp(argv[2], "q2") != 0) || + (argc > 3 && !fixed_levels && std::strcmp(argv[3], "learned") != 0)) { + std::fprintf(stderr, "usage: %s [tokens:1..16] [q2|q3] [learned|fixed]\n", argv[0]); return 2; } const int block_bytes = fp3 ? 14 : 10; @@ -102,7 +114,7 @@ int main(int argc, char ** argv) { } std::vector books((size_t) n_experts * 2 * levels); for (size_t i = 0; i < books.size(); ++i) books[i] = f32_to_bf16(-0.5f + 0.2f * (float) (i % 5)); - std::vector modes(n_experts, 1); + std::vector modes(n_experts, fixed_levels ? 0 : 1); void * d_up = nullptr, * d_gate = nullptr; float * d_x = nullptr, * d_a = nullptr, * d_b = nullptr; @@ -202,9 +214,10 @@ int main(int argc, char ** argv) { for (int r = 0; r < REPS; ++r) { u.push_back(time_unfused()); f.push_back(time_fused()); } const double mu = median(u), mf = median(f); - std::fprintf(stderr, "geometry: qtype=%s in=%d out=%d top_k=%d ntok=%d " + std::fprintf(stderr, "geometry: qtype=%s levels=%s in=%d out=%d top_k=%d ntok=%d " "(%d iters x %d interleaved reps)\n", - fp3 ? "q3-mix" : "q2-mix", in, out, n_used, ntok, ITERS, REPS); + fp3 ? "q3-mix" : "q2-mix", fixed_levels ? "fixed" : "learned", + in, out, n_used, ntok, ITERS, REPS); std::fprintf(stderr, " unfused (2 launches, swiglu NOT counted): %8.4f ms/step [", mu); for (double v : u) std::fprintf(stderr, " %.4f", v); @@ -217,14 +230,32 @@ int main(int argc, char ** argv) { std::fprintf(stderr, " per-layer saving %.4f ms -> over 43 layers %.3f ms/step\n", mu - mf, (mu - mf) * 43.0); - std::vector result(yn); - HIP_OK(hipMemcpy(result.data(), d_a, sizeof(float) * yn, hipMemcpyDeviceToHost)); - uint64_t result_hash = 1469598103934665603ull; - const uint8_t * result_bytes = reinterpret_cast(result.data()); - for (size_t i = 0; i < sizeof(float) * yn; ++i) { - result_hash = (result_hash ^ result_bytes[i]) * 1099511628211ull; + // Hash both raw projections and the fused output after timing. This catches + // mode-specific numerical drift even when the generated-token hash happens + // to remain unchanged. + if (!mul_mat_id(d_up, d_x, d_ids, d_a, in, out, n_used, ntok, 1, + ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, stream) || + !mul_mat_id(d_gate, d_x, d_ids, d_b, in, out, n_used, ntok, 1, + ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, stream)) { + std::fprintf(stderr, "FAIL: raw projection dispatch rejected\n"); + return 1; + } + HIP_OK(hipStreamSynchronize(stream)); + std::vector up_result(yn), gate_result(yn), fused_result(yn); + HIP_OK(hipMemcpy(up_result.data(), d_a, sizeof(float) * yn, hipMemcpyDeviceToHost)); + HIP_OK(hipMemcpy(gate_result.data(), d_b, sizeof(float) * yn, hipMemcpyDeviceToHost)); + if (!mul_mat_id_glu(d_up, d_gate, d_x, d_ids, d_a, in, out, n_used, + ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, 7.0f, stream)) { + std::fprintf(stderr, "FAIL: fused projection dispatch rejected\n"); + return 1; } - std::fprintf(stderr, " result fnv1a64: %016llx\n", (unsigned long long) result_hash); + HIP_OK(hipStreamSynchronize(stream)); + HIP_OK(hipMemcpy(fused_result.data(), d_a, sizeof(float) * yn, hipMemcpyDeviceToHost)); + std::fprintf(stderr, " raw up/gate fnv1a64: %016llx %016llx\n", + (unsigned long long) fnv1a64(up_result), + (unsigned long long) fnv1a64(gate_result)); + std::fprintf(stderr, " result fnv1a64: %016llx\n", + (unsigned long long) fnv1a64(fused_result)); unregister_mix(d_gate); unregister_mix(d_up);