From cd8abb806c993dc5a6bbc171953dfabfc675536d Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 30 Aug 2026 13:08:01 +0200 Subject: [PATCH 1/8] perf(ds4): default mixed ROCmFP MMQ prefill on gfx1151 --- .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 6 +-- .../deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh | 5 ++- server/src/deepseek4/deepseek4_backend.cpp | 42 +++++++++++++++++++ server/src/deepseek4/deepseek4_backend.h | 9 ++++ server/tests/test_deepseek4_unit.cpp | 18 ++++++++ 5 files changed, 75 insertions(+), 5 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 82052f611..8a51d33f3 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2827,9 +2827,9 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor : luce_mmvq_max_ncols_env; // The mix qtypes have no generic MMVQ path because their per-expert // codebooks live in an out-of-band registry. Decode uses the dedicated - // fused kernels below. Sparse prefill can opt into their registry-aware - // MMQ loaders; otherwise should_use_mmq rejects them and they retain the - // exact dequantize->cuBLAS fallback. + // fused kernels below. Approximate prefill modes can select their + // registry-aware MMQ loaders; otherwise should_use_mmq rejects them and + // they retain the exact dequantize->cuBLAS fallback. const bool is_rocmfp3_mix = src0->type == GGML_TYPE_Q3_1_ROCMFP3_MIX; const bool is_rocmfp2_mix = src0->type == GGML_TYPE_Q2_1_ROCMFP2_MIX; const bool is_mix_qtype = is_rocmfp3_mix || is_rocmfp2_mix; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index 331b76b10..9d7dd78f4 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh @@ -1101,8 +1101,9 @@ static __device__ __forceinline__ void load_tiles_rocmfpx_dual( // qtype-107, but each expert supplies two learned four-level codebooks. MMQ // needs int8 tiles, so quantize those tiny codebooks once per K tile, then fold // the codebook scale into the block's ordinary UE4M3 scale. The additional -// error is bounded to half an int8 step (measured below 0.4% of codebook range) -// and this path is opt-in because sparse prefill is already approximate. +// error is bounded to half an int8 step (measured below 0.4% of codebook range). +// Backend policy selects this path only for prefill modes that are already +// approximate; exact prefill retains the dequantize-to-F16 fallback. struct rocmfp2_mix_mmq_lut { int packed[2]; float scale[2]; diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index acc9bfb95..e7ebf78a3 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -256,6 +256,35 @@ static void configure_gfx1151_sparse_decode_default(int gpu) { #endif } +static bool configure_gfx1151_mix_mmq_prefill_default( + int gpu, PrefillAttentionMode mode) { +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + // Preserve explicit user policy, including the value 0 kill switch. + if (std::getenv("DFLASH_DS4_MIX_MMQ_PREFILL") != nullptr) { + return true; + } + + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, gpu) != cudaSuccess || + !deepseek4_mix_mmq_prefill_default(mode, prop.gcnArchName)) { + return true; + } + + if (::setenv("DFLASH_DS4_MIX_MMQ_PREFILL", "1", 0) != 0) { + std::fprintf(stderr, + "[deepseek4] failed to enable mixed ROCmFP MMQ prefill\n"); + return false; + } + std::fprintf(stderr, + "[deepseek4] gfx1151 approximate prefill: defaulting mixed " + "ROCmFP MMQ on\n"); +#else + (void) gpu; + (void) mode; +#endif + return true; +} + static void configure_gfx1201_hybrid_sub_batch_default(int gpu) { #if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) if (std::getenv("DFLASH_MMQ_SUB_BATCH") != nullptr) { @@ -769,6 +798,15 @@ static MoeLayerDesc make_ds4_expert_layer_desc(const DeepSeek4Layer & layer) { } // namespace +bool deepseek4_mix_mmq_prefill_default( + PrefillAttentionMode mode, const char * gcn_arch) { + if (!prefill_attention_mode_is_approximate(mode) || gcn_arch == nullptr || + std::strncmp(gcn_arch, "gfx1151", 7) != 0) { + return false; + } + return gcn_arch[7] == '\0' || gcn_arch[7] == ':'; +} + DeepSeek4Backend::DeepSeek4Backend(const DeepSeek4BackendConfig & cfg) : cfg_(cfg) {} @@ -1064,6 +1102,10 @@ bool DeepSeek4Backend::init() { return false; } configure_gfx1151_sparse_decode_default(cfg_.device.gpu); + if (!configure_gfx1151_mix_mmq_prefill_default( + cfg_.device.gpu, cfg_.prefill_mode)) { + return false; + } configure_gfx1201_hybrid_sub_batch_default(cfg_.device.gpu); backend_ = ggml_backend_cuda_init(cfg_.device.gpu); diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 84f0745c6..058f04486 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -42,6 +42,15 @@ int deepseek4_hybrid_prefill_step_tokens( int configured_chunk, int position, int remaining_tokens); + +// Mixed ROCmFP MMQ changes the reduction/quantization topology, so only the +// already-approximate prefill modes may select it automatically. The policy is +// kept separate from the qtype kernels so future model backends can reuse the +// same generic MMQ path after device-level qualification. +bool deepseek4_mix_mmq_prefill_default( + PrefillAttentionMode mode, + const char * gcn_arch); + class DeepSeek4Backend : public ModelBackend { public: explicit DeepSeek4Backend(const DeepSeek4BackendConfig & cfg); diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 94e2a7ac2..1f2a3a701 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -1718,6 +1718,23 @@ static void test_hybrid_prefill_chunk_tokens() { 2048, 32768, 0) == 0); std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } + +static void test_mix_mmq_prefill_default() { + std::fprintf(stderr, " test_mix_mmq_prefill_default ..."); + TEST_ASSERT(!deepseek4_mix_mmq_prefill_default( + PrefillAttentionMode::Exact, "gfx1151")); + TEST_ASSERT(deepseek4_mix_mmq_prefill_default( + PrefillAttentionMode::Dense, "gfx1151")); + TEST_ASSERT(deepseek4_mix_mmq_prefill_default( + PrefillAttentionMode::Sparse, "gfx1151:sramecc+:xnack-")); + TEST_ASSERT(!deepseek4_mix_mmq_prefill_default( + PrefillAttentionMode::Sparse, "gfx1201")); + TEST_ASSERT(!deepseek4_mix_mmq_prefill_default( + PrefillAttentionMode::Sparse, "gfx11510")); + TEST_ASSERT(!deepseek4_mix_mmq_prefill_default( + PrefillAttentionMode::Sparse, nullptr)); + std::fprintf(stderr, " OK\n"); +} static void test_dspark_park_all_releases_drafter() { std::fprintf(stderr, " test_dspark_park_all_releases_drafter ..."); @@ -4174,6 +4191,7 @@ int main() { test_dspark_confidence_uses_separate_hidden(backend); test_safe_compressor_batch_tokens(); test_hybrid_prefill_chunk_tokens(); + test_mix_mmq_prefill_default(); test_dspark_park_all_releases_drafter(); test_pflash_rejects_invalid_requests(); test_dspark_raw_ring_rollback_after_wrap(backend); From 58b10f1e489e14bb10f60db8551ddfc74793394c Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 30 Aug 2026 23:19:41 +0200 Subject: [PATCH 2/8] perf(hip): add opt-in streaming MLA top-k --- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 240 ++++++++++++++++++ server/docs/DS4.md | 9 + server/tests/test_deepseek4_unit.cpp | 152 +++++++++++ 3 files changed, 401 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 3ba63ab4d..49bfbea17 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -1554,6 +1554,202 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( } } +// Streaming indexed MLA for long prefill. The compact grouped kernel above +// stores every score and then reloads V. Once the trained indexer has reduced +// the compressed history to a bounded top-k set, that extra traffic is no +// longer necessary: stage one latent row in LDS, share it across the heads in +// a block, and update online-softmax state while the row is resident. Sixteen +// wave32 heads amortize each LDS load best on gfx1151. The contract is +// backend-generic (D512 MQA with K == V and direct indexed rows); model policy +// remains in the graph/backend layer. +template +__global__ static void ds4_flash_attn_d512_streaming_topk_kernel( + float * dst, + const float * q, + size_t q_stride_token, + size_t q_stride_head, + const KV * kv, + const Mask * mask, + const float * sinks, + int n_tokens, + int n_heads, + int n_kv, + float scale, + 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) { + constexpr int D = 512; + constexpr int WAVE = 32; + constexpr int N_THREADS = HEADS_PER_BLOCK * WAVE; + constexpr int VALUES_PER_LANE = D / WAVE; + static_assert(HEADS_PER_BLOCK == 16); + static_assert(KEYS_PER_STAGE == 16); + static_assert(N_THREADS == 512); + + const int token = (int) blockIdx.x; + const int head_begin = (int) blockIdx.y * HEADS_PER_BLOCK; + const int tid = (int) threadIdx.x; + const int wave = tid / WAVE; + const int lane = tid & (WAVE - 1); + const int head = head_begin + wave; + + // The launch gate makes the grid exact. Keeping every thread live is + // required because each stage has workgroup-wide barriers. + if (token >= n_tokens || head_begin + HEADS_PER_BLOCK > n_heads) return; + + __shared__ KV staged_kv[KEYS_PER_STAGE * D]; + __shared__ int staged_rows[KEYS_PER_STAGE]; + __shared__ float staged_masks[KEYS_PER_STAGE]; + + const int * token_visibility = visibility_bounds + (size_t) token * 4; + const int raw_first = token_visibility[0]; + const int raw_last = token_visibility[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 * token_rows = indexed_rows + + (size_t) token * indexed_capacity; + const Mask * token_mask = mask + (size_t) token * n_kv; + + const float * qh = q + (size_t) token * q_stride_token + + (size_t) head * q_stride_head; + float q_values[VALUES_PER_LANE]; + float accum[VALUES_PER_LANE] = {}; +#pragma unroll + for (int i = 0; i < VALUES_PER_LANE; ++i) { + const int dim = lane + i * WAVE; + float qv = qh[dim]; + if (inverse_rope.forward_q_enabled && dim >= D - 64) { + const int pair = (dim - (D - 64)) / 2; + const float x0 = qh[D - 64 + 2 * pair + 0]; + const float x1 = qh[D - 64 + 2 * pair + 1]; + const size_t coefficient_index = + ((size_t) token * 32 + (size_t) pair) * 2; + float y0; + float y1; + ds4_apply_inverse_rope_pair( + x0, x1, + forward_rope_coefficients[coefficient_index + 0], + forward_rope_coefficients[coefficient_index + 1], + y0, y1); + qv = (dim & 1) == 0 ? y0 : y1; + } + q_values[i] = qv; + } + + float row_max = -3.402823466e38f; + float row_sum = 0.0f; + for (int row_base = 0; row_base < total_rows; + row_base += KEYS_PER_STAGE) { + if (tid < KEYS_PER_STAGE) { + const int selected = row_base + tid; + int row = -1; + if (selected < raw_count) { + row = raw_first + selected; + } else if (selected < total_rows) { + row = token_rows[selected - raw_count]; + } + staged_rows[tid] = row; + staged_masks[tid] = row >= 0 && row < n_kv + ? ds4_fa_load(token_mask + row) + : -3.402823466e38f; + } + __syncthreads(); + + for (int index = tid; index < KEYS_PER_STAGE * D; + index += N_THREADS) { + const int slot = index / D; + const int dim = index - slot * D; + const int row = staged_rows[slot]; + staged_kv[index] = row >= 0 && row < n_kv + ? kv[(size_t) row * D + dim] : KV{}; + } + __syncthreads(); + +#pragma unroll + for (int slot = 0; slot < KEYS_PER_STAGE; ++slot) { + const int selected = row_base + slot; + if (selected >= total_rows) continue; + const float mask_value = staged_masks[slot]; + if (mask_value <= -1.0e20f) continue; + + float partial = 0.0f; +#pragma unroll + for (int i = 0; i < VALUES_PER_LANE; ++i) { + const int dim = lane + i * WAVE; + partial += q_values[i] * + ds4_fa_load(staged_kv + slot * D + dim); + } + partial = warp_reduce_sum(partial); + // XOR reduction can associate operands differently in each lane. + // Broadcast lane zero so every output dimension advances one + // identical softmax state. + partial = __shfl_sync(0xffffffffu, partial, 0, WAVE); + + const float score = partial * scale + mask_value; + const float next_max = fmaxf(row_max, score); + const float old_scale = row_sum == 0.0f + ? 0.0f : expf(row_max - next_max); + const float value_scale = expf(score - next_max); + row_sum = row_sum * old_scale + value_scale; + row_max = next_max; +#pragma unroll + for (int i = 0; i < VALUES_PER_LANE; ++i) { + const int dim = lane + i * WAVE; + const float value = ds4_fa_load( + staged_kv + slot * D + dim); + accum[i] = accum[i] * old_scale + value_scale * value; + } + } + __syncthreads(); + } + + if (sinks) { + const float sink = sinks[head]; + const float next_max = fmaxf(row_max, sink); + const float old_scale = row_sum == 0.0f + ? 0.0f : expf(row_max - next_max); + row_sum = row_sum * old_scale + expf(sink - next_max); +#pragma unroll + for (int i = 0; i < VALUES_PER_LANE; ++i) { + accum[i] *= old_scale; + } + } + + const float inv_sum = row_sum == 0.0f ? 0.0f : 1.0f / row_sum; + float * out = dst + + ((size_t) token * (size_t) n_heads + (size_t) head) * D; +#pragma unroll + for (int i = 0; i < VALUES_PER_LANE; ++i) { + const int dim = lane + i * WAVE; + float value = accum[i] * inv_sum; + if (inverse_rope.enabled && dim >= D - 64) { + const float partner = __shfl_xor_sync( + 0xffffffffu, value, 1, WAVE); + const float x0 = (dim & 1) == 0 ? value : partner; + const float x1 = (dim & 1) == 0 ? partner : value; + const int pair = (dim - (D - 64)) / 2; + const size_t coefficient_index = + ((size_t) token * 32 + (size_t) pair) * 2; + float y0; + float y1; + ds4_apply_inverse_rope_pair( + x0, x1, + inverse_rope_coefficients[coefficient_index + 0], + inverse_rope_coefficients[coefficient_index + 1], + y0, y1); + value = (dim & 1) == 0 ? y0 : y1; + } + out[dim] = value; + } +} + // 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 @@ -2383,6 +2579,50 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( n_tokens, n_kv, raw_rows); } CUDA_CHECK(cudaGetLastError()); + // Vulkan shares each selected latent row across eight subgroup64 + // heads. HIP tuning selected sixteen wave32 heads for the same K == V + // reuse. Keep the native HIP path opt-in until a model-backed A/B + // qualifies its online-softmax association. The shape gate is + // expressed in terms of the D512 indexed-attention + // contract so another model can reuse the kernel without DS4 policy + // leaking into it. + const char * streaming_topk_env = + getenv("GGML_CUDA_MLA_STREAM_TOPK"); + if (!streaming_topk_env) { + streaming_topk_env = getenv("GGML_DS4_FA_STREAM_TOPK"); + } + const bool streaming_topk_enabled = streaming_topk_env && + streaming_topk_env[0] != '\0' && + strcmp(streaming_topk_env, "0") != 0; + constexpr int streaming_min_tokens = 64; + const int active_row_upper_bound = raw_window + indexed_capacity; + const int device_warp_size = + ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; + if (streaming_topk_enabled && indexed_mask && indexer_topk && + kv_f16 && mask->type == GGML_TYPE_F16 && + K->data == V->data && n_heads % 16 == 0 && + device_warp_size == 32 && n_tokens >= streaming_min_tokens && + active_row_upper_bound > 0 && + n_kv >= 3 * active_row_upper_bound) { + constexpr int streaming_heads = 16; + const dim3 streaming_grid( + (unsigned) n_tokens, + (unsigned) (n_heads / streaming_heads), 1); + ds4_flash_attn_d512_streaming_topk_kernel + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients); + CUDA_CHECK(cudaGetLastError()); + return true; + } // 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 diff --git a/server/docs/DS4.md b/server/docs/DS4.md index 2db93acb4..b1d42c2eb 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -487,6 +487,15 @@ On gfx1151, the exact block-radix selector is also the default for the DS4 hipCUB full-sort path. The qualification benchmark checks selected-set parity before reporting selector timing. +For experimental long sparse prefill, set both +`DFLASH_DS4_DIRECT_INDEXER_TOPK=1` and `GGML_CUDA_MLA_STREAM_TOPK=1`. This +enables a reusable D512 K-equals-V streaming attention path that shares each +selected latent row across wave32 heads and avoids materializing scores. It is +currently limited to F16 caches on native wave32 devices and otherwise falls +back to the existing path. Because its online softmax changes floating-point +association, keep it opt-in until the target model passes a matched output and +throughput A/B. `GGML_DS4_FA_STREAM_TOPK` remains a compatibility alias. + Indexed verifier attention with at most eight query rows also uses a two-way split-KV schedule on gfx1151. Set `GGML_CUDA_MLA_NO_SPLIT_KV=1` to restore the single-block schedule. For fused-verifier masks of at least 4 MiB, the runtime diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 1f2a3a701..bdef57893 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -2873,6 +2873,157 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_ds4_flash_attention_streaming_topk_gpu() { + std::fprintf(stderr, + " test_ds4_flash_attention_streaming_topk_gpu ..."); +#if !defined(GGML_USE_HIP) + std::fprintf(stderr, " skipped (HIP-only candidate)\n"); + return; +#endif + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, " skipped (no GPU backend)\n"); + return; + } + + constexpr int head_dim = 512; + constexpr int n_heads = 64; + constexpr int n_tokens = 64; + constexpr int raw_rows = 128; + constexpr int raw_window = 128; + constexpr int selected_rows = 512; + constexpr int n_comp_rows = 1920; + constexpr int n_kv = raw_rows + n_comp_rows; + + ggml_context * ctx = make_test_context(4u << 20); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) { + ggml_backend_free(backend); + std::fprintf(stderr, " FAIL\n"); + return; + } + + ggml_tensor * q = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, head_dim, n_tokens, n_heads); + ggml_tensor * kv = ggml_new_tensor_3d( + ctx, GGML_TYPE_F16, head_dim, n_kv, 1); + ggml_tensor * mask = ggml_new_tensor_2d( + ctx, GGML_TYPE_F16, n_kv, n_tokens); + ggml_tensor * topk = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, selected_rows, n_tokens); + ggml_tensor * output = ggml_flash_attn_ext( + ctx, q, kv, kv, mask, 1.0f / std::sqrt((float) head_dim), + 0.0f, 0.0f); + ggml_flash_attn_ext_set_ds4_sparse( + output, raw_rows, raw_window, -selected_rows, 1); + ggml_flash_attn_ext_set_ds4_indexer_topk(output, topk); + ggml_set_output(output); + TEST_ASSERT_MSG(ggml_backend_supports_op(backend, output), + "GPU rejected streaming top-k attention fixture"); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 32, false); + ggml_build_forward_expand(graph, output); + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + const bool allocated = ggml_gallocr_alloc_graph(alloc, graph); + TEST_ASSERT_MSG(allocated, "streaming top-k graph allocation failed"); + if (allocated) { + std::vector q_data( + (size_t) head_dim * n_tokens * n_heads); + std::vector kv_data((size_t) head_dim * n_kv); + std::vector mask_data( + (size_t) n_kv * n_tokens, ggml_fp32_to_fp16(-1.0e30f)); + std::vector topk_data( + (size_t) selected_rows * n_tokens); + for (size_t i = 0; i < q_data.size(); ++i) { + q_data[i] = ((int) (i % 43) - 21) * 0.001f; + } + for (size_t i = 0; i < kv_data.size(); ++i) { + kv_data[i] = ggml_fp32_to_fp16( + ((int) (i % 47) - 23) * 0.001f); + } + for (int token = 0; token < n_tokens; ++token) { + ggml_fp16_t * token_mask = + mask_data.data() + (size_t) token * n_kv; + for (int row = 0; row < raw_rows; ++row) { + token_mask[row] = ggml_fp32_to_fp16(0.0f); + } + for (int row = 0; row < n_comp_rows; ++row) { + token_mask[raw_rows + row] = ggml_fp32_to_fp16(0.0f); + } + for (int rank = 0; rank < selected_rows; ++rank) { + topk_data[(size_t) token * selected_rows + rank] = + (token * 17 + selected_rows - 1 - rank) % n_comp_rows; + } + } + ggml_backend_tensor_set(q, q_data.data(), 0, + q_data.size() * sizeof(float)); + ggml_backend_tensor_set(kv, kv_data.data(), 0, + kv_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(mask, mask_data.data(), 0, + mask_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set(topk, topk_data.data(), 0, + topk_data.size() * sizeof(int32_t)); + + ScopedEnvVar streaming_guard("GGML_CUDA_MLA_STREAM_TOPK"); + ScopedCudaGraphOverrides eager( + /*disable_graphs=*/true, + /*mmvq_max_ncols=*/0, + /*skip_property_check=*/false); + std::vector reference((size_t) ggml_nelements(output)); + std::vector candidate(reference.size()); + setenv("GGML_CUDA_MLA_STREAM_TOPK", "0", 1); + TEST_ASSERT_MSG( + ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, + "grouped compact attention reference failed"); + ggml_backend_tensor_get(output, reference.data(), 0, + reference.size() * sizeof(float)); + + setenv("GGML_CUDA_MLA_STREAM_TOPK", "1", 1); + TEST_ASSERT_MSG( + ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, + "streaming top-k attention failed"); + ggml_backend_tensor_get(output, candidate.data(), 0, + candidate.size() * sizeof(float)); + + for (size_t i = 0; i < reference.size(); ++i) { + TEST_ASSERT_MSG( + std::isfinite(candidate[i]), + "streaming top-k attention output must be finite"); + TEST_ASSERT_MSG( + nearly_equal(reference[i], candidate[i], 5.0e-4f, 5.0e-4f), + "streaming top-k attention exceeded numeric tolerance"); + } + + auto measure_us = [&](bool streaming) { + setenv("GGML_CUDA_MLA_STREAM_TOPK", streaming ? "1" : "0", 1); + constexpr int warmups = 3; + constexpr int iterations = 20; + for (int i = 0; i < warmups; ++i) { + ggml_backend_graph_compute(backend, graph); + } + ggml_backend_synchronize(backend); + const auto begin = std::chrono::steady_clock::now(); + for (int i = 0; i < iterations; ++i) { + ggml_backend_graph_compute(backend, graph); + } + ggml_backend_synchronize(backend); + const auto end = std::chrono::steady_clock::now(); + return std::chrono::duration(end - begin).count() / + iterations; + }; + const double grouped_us = measure_us(false); + const double streaming_us = measure_us(true); + std::fprintf(stderr, " grouped=%.1fus streaming=%.1fus speedup=%.2fx", + grouped_us, streaming_us, grouped_us / streaming_us); + } + + ggml_gallocr_free(alloc); + ggml_free(ctx); + ggml_backend_free(backend); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void run_ds4_indexer_score_packed_small_case( ggml_backend_t backend, int n_tokens) { constexpr int dim = 128; @@ -4209,6 +4360,7 @@ int main() { #if defined(GGML_USE_CUDA) || defined(GGML_USE_HIP) test_ds4_flash_attention_keep_cap_gpu(); test_ds4_flash_attention_parallel_index_scan_gpu(); + test_ds4_flash_attention_streaming_topk_gpu(); test_ds4_indexer_score_packed_small_gpu(); test_ds4_topk_block_radix_gpu(); test_ds4_flash_attention_inverse_rope_fallback_gpu(); From 8e49d7bb06918b7e536b3ebf1843c89ebd073627 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 31 Aug 2026 01:53:37 +0200 Subject: [PATCH 3/8] perf(hip): tune sparse grouped MoE prefill --- .../deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh | 48 ++- .../llama.cpp/ggml/src/ggml-cuda/moe-fused.cu | 130 +++++++- server/docs/ENVIRONMENT.md | 3 + server/src/common/moe_hybrid_ffn_eval.cpp | 22 +- .../test/test_deepseek4_mmid_grouped_cuda.cpp | 303 +++++++++++++++++- 5 files changed, 470 insertions(+), 36 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index 9d7dd78f4..a354390b7 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh @@ -6,6 +6,7 @@ #include #include +#include using namespace ggml_cuda_mma; @@ -4916,7 +4917,52 @@ void mul_mat_q_case(ggml_backend_cuda_context & ctx, const mmq_args & args, cuda int mmq_x_best = 0; int ntiles_x_best = INT_MAX; - for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) { + static const int forced_mmq_x = []() { + const char * raw = std::getenv("GGML_CUDA_MMQ_X"); + if (!raw || !*raw) return 0; + char * end = nullptr; + const long parsed = std::strtol(raw, &end, 10); + return end && end != raw && *end == '\0' && parsed >= 8 && + parsed <= 128 && parsed % 8 == 0 + ? (int)parsed : 0; + }(); + static const bool adaptive_moe_x_enabled = []() { + const char * raw = std::getenv("GGML_CUDA_MMQ_MOE_ADAPTIVE_X"); + return raw && *raw && !(raw[0] == '0' && raw[1] == '\0'); + }(); + int requested_mmq_x = forced_mmq_x; + if (requested_mmq_x == 0 && adaptive_moe_x_enabled && + cc == GGML_CUDA_CC_OFFSET_AMD + 0x1151 && + args.expert_bounds != nullptr && args.nchannels_x > 0) { + // Grouped MoE routes are sparse across experts. Sizing the X tile from + // the full token width makes almost every workgroup carry padding. Use + // the mean live-route density as a model-neutral trigger, then select + // the measured gfx1151 tile for each unpack format. The grid still + // spans ncols_max, so skewed experts remain fully covered. + const int64_t routes_per_expert = + (args.ncols_y + args.nchannels_x - 1) / args.nchannels_x; + if (routes_per_expert <= 16) { + switch (type) { + case GGML_TYPE_Q2_0_ROCMFP2: requested_mmq_x = 32; break; + case GGML_TYPE_Q3_0_ROCMFPX: requested_mmq_x = 48; break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: requested_mmq_x = 16; break; + default: break; + } + } + } + if (requested_mmq_x > 0 && requested_mmq_x <= mmq_x_max) { + const int granularity = mmq_get_granularity_host(requested_mmq_x, cc); + if (requested_mmq_x % granularity == 0 && + mmq_get_nbytes_shared( + requested_mmq_x, mmq_y, cc, warp_size, nwarps) <= smpbo) { + mmq_x_best = requested_mmq_x; + ntiles_x_best = 1; + } + } + + for (int mmq_x = 8; + mmq_x_best == 0 && mmq_x <= mmq_x_max && ntiles_x_best > 1; + mmq_x += 8) { const int granularity = mmq_get_granularity_host(mmq_x, cc); if (mmq_x % granularity != 0 || mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size, nwarps) > smpbo) { 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..6c1253a58 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 @@ -4,6 +4,8 @@ #include "ggml-cuda/mmvq.cuh" #include +#include +#include #include static __device__ __forceinline__ float silu_f32(float x) { @@ -310,6 +312,115 @@ static __global__ void laguna_moe_combine_kernel( (size_t)t * output_nb1) = sum; } +static __global__ void moe_combine_vec4_kernel( + const char * __restrict__ experts, + const char * __restrict__ weights, + char * __restrict__ output, + const int n_embd_vec4, + const int n_used, + const int n_tokens, + const size_t experts_nb1, + const size_t experts_nb2, + const size_t weights_nb0, + const size_t weights_nb1, + const size_t output_nb1, + const float value_scale) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int total = n_embd_vec4 * n_tokens; + if (idx >= total) return; + + const int h4 = idx % n_embd_vec4; + const int t = idx / n_embd_vec4; + float4 sum = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + for (int e = 0; e < n_used; ++e) { + const float w = *(const float *)(weights + + (size_t)e * weights_nb0 + + (size_t)t * weights_nb1); + if (w == 0.0f) { + if (e == 0) sum = make_float4(0.0f, 0.0f, 0.0f, 0.0f); + continue; + } + const float4 value = *(const float4 *)(experts + + (size_t)h4 * sizeof(float4) + + (size_t)e * experts_nb1 + + (size_t)t * experts_nb2); + const float4 scaled = value_scale == 1.0f + ? value + : make_float4( + __fmul_rn(value.x, value_scale), + __fmul_rn(value.y, value_scale), + __fmul_rn(value.z, value_scale), + __fmul_rn(value.w, value_scale)); + const float4 product = make_float4( + __fmul_rn(scaled.x, w), + __fmul_rn(scaled.y, w), + __fmul_rn(scaled.z, w), + __fmul_rn(scaled.w, w)); + if (e == 0) { + sum = product; + } else { + sum.x = __fadd_rn(sum.x, product.x); + sum.y = __fadd_rn(sum.y, product.y); + sum.z = __fadd_rn(sum.z, product.z); + sum.w = __fadd_rn(sum.w, product.w); + } + } + *(float4 *)(output + + (size_t)h4 * sizeof(float4) + + (size_t)t * output_nb1) = sum; +} + +static void launch_moe_combine( + cudaStream_t stream, + const char * experts, + const char * weights, + char * output, + int n_embd, + int n_used, + int n_tokens, + size_t experts_nb0, + size_t experts_nb1, + size_t experts_nb2, + size_t weights_nb0, + size_t weights_nb1, + size_t output_nb0, + size_t output_nb1, + float value_scale) { + static const bool vec4_requested = []() { + const char * raw = std::getenv("DFLASH_MOE_COMBINE_VEC4"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + static_assert(sizeof(float4) == 4 * sizeof(float)); + constexpr size_t vec4_alignment = sizeof(float4); + const bool aligned = + (reinterpret_cast(experts) % vec4_alignment) == 0 && + (reinterpret_cast(output) % vec4_alignment) == 0 && + experts_nb1 % vec4_alignment == 0 && + experts_nb2 % vec4_alignment == 0 && + output_nb1 % vec4_alignment == 0; + const bool use_vec4 = vec4_requested && n_embd % 4 == 0 && aligned && + experts_nb0 == sizeof(float) && output_nb0 == sizeof(float); + constexpr int block = 256; + if (use_vec4) { + const int n_embd_vec4 = n_embd / 4; + const int total = n_embd_vec4 * n_tokens; + const int grid = (total + block - 1) / block; + moe_combine_vec4_kernel<<>>( + experts, weights, output, n_embd_vec4, n_used, n_tokens, + experts_nb1, experts_nb2, weights_nb0, weights_nb1, + output_nb1, value_scale); + return; + } + + const int total = n_embd * n_tokens; + const int grid = (total + block - 1) / block; + laguna_moe_combine_kernel<<>>( + experts, weights, output, n_embd, n_used, n_tokens, + experts_nb0, experts_nb1, experts_nb2, + weights_nb0, weights_nb1, output_nb0, output_nb1, + value_scale); +} + static __global__ void ds4_peer_copy_f32_kernel( const float * __restrict__ src, float * __restrict__ dst, @@ -561,10 +672,8 @@ static void ggml_cuda_op_ds4_moe_owner( ggml_cuda_mul_mat_vec_q( ctx, down_w, &gu, expert_ids, &experts, nullptr); - const int total = n_embd * n_tokens; - const int block = 256; - const int grid = (total + block - 1) / block; - laguna_moe_combine_kernel<<>>( + launch_moe_combine( + ctx.stream(), (const char *) experts.data, (const char *) weights->data, (char *) dst->data, @@ -634,10 +743,8 @@ static void ggml_cuda_op_ds4_moe_owner_split( ggml_cuda_mul_mat_vec_q( ctx, down_w, &gu, expert_ids, &experts, nullptr); - const int total = n_embd * n_tokens; - const int block = 256; - const int grid = (total + block - 1) / block; - laguna_moe_combine_kernel<<>>( + launch_moe_combine( + ctx.stream(), (const char *) experts.data, (const char *) weights->data, (char *) dst->data, @@ -760,11 +867,8 @@ void ggml_cuda_op_moe_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst) const int n_embd = (int) experts->ne[0]; const int n_used = (int) experts->ne[1]; const int n_tokens = (int) experts->ne[2]; - const int total = n_embd * n_tokens; - - const int block = 256; - const int grid = (total + block - 1) / block; - laguna_moe_combine_kernel<<>>( + launch_moe_combine( + ctx.stream(), (const char *) experts->data, (const char *) weights->data, (char *) dst->data, diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index fba8c7d40..0202b3550 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -49,6 +49,8 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_DS4_ROCTX` | unset | DEBUG: on HIP builds, dynamically load ROCTX and emit semantic DS4 prefill, speculative-decode, and layer-range markers for external rocprof traces. No events, timing, or device synchronization are added. | | `DFLASH_QWEN35_ROCTX` | unset | DEBUG: on HIP builds, dynamically load ROCTX and mark Qwen concurrent steps, graph compute, and argmax readback with live, padded, and packed-prefill shape metadata. | | `DFLASH_GFX1151_HC_MMVF_Q4` | 1 on gfx1151 for the DS4 `[16384,24]` q4 projection | BURN-IN KILL SWITCH: =0 restores the generic hipBLAS dispatch decision. | +| `GGML_CUDA_MMQ_X` | unset | DEBUG: force a supported MMQ output-column tile width (8–128) for architecture tuning; invalid or over-budget values fall back to automatic selection. | +| `GGML_CUDA_MMQ_MOE_ADAPTIVE_X` | unset | BURN-IN: on sparse-route gfx1151 grouped MoE MMQ, choose the measured ROCmFP2/3/4 output tile from routed rows per expert; ordinary matmuls, unmeasured formats, and other devices are unchanged. | | `GGML_CUDA_MLA_SPLIT_KV` / `GGML_DS4_FA_SPLIT_KV` | 1 on gfx1151 indexed decode; unset elsewhere | BURN-IN: force the reusable split-KV MLA schedule. Set `GGML_CUDA_MLA_NO_SPLIT_KV=1` (or legacy `GGML_DS4_FA_NO_SPLIT_KV=1`) to disable it. | | `GGML_DS4_TOPK_BLOCK_RADIX` | 1 on gfx1151 | BURN-IN KILL SWITCH: =0 restores hipCUB full sort for DS4-shaped 512-row top-k selection. | | `GGML_DS4_FA_SERIAL_INDEX_SCAN` | unset | DEBUG/A-B: restore the serial indexed-attention mask scan instead of the long-context HIP parallel scan. | @@ -202,6 +204,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. - `DFLASH_MMQ_SUB_BATCH` - moe_hybrid_ffn_eval.cpp - `DFLASH_MODEL_CARDS_DIR` - model_card.cpp - `DFLASH_MOE_COLD_BACKEND` - deepseek4_loader.cpp +- `DFLASH_MOE_COMBINE_VEC4` - opt in to aligned four-lane GPU route reduction; scalar fallback remains available - `DFLASH_MOE_COMPACT_MATERIALIZED` - moe_hybrid_ffn_eval.cpp - `DFLASH_MOE_DUPLICATE_HOT_ON_COLD` - moe_hybrid_storage.cpp - `DFLASH_MOE_EXPERT_COMPUTE_DAEMON_TOKEN_LOOP` - moe_expert_compute_ipc.cpp diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index d5e80cec0..f0cb2510d 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -2877,11 +2877,23 @@ static bool eval_moe_owner_expert_major_batched( ggml_tensor * down_e = apply_scale2(ctx, ggml_mul_mat_id(ctx, down_tensor, gu, local_ids_tensor), desc.ffn_down_exps_s); - ggml_tensor * weights_3d = ggml_reshape_3d(ctx, owner_weights_tensor, 1, n_used, n_tokens); - ggml_tensor * 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); + ggml_tensor * routed_out = nullptr; + if (moe_hybrid_graph_policy().fused_combine) { + // Keep the grouped-MMID result in route-major form and combine it + // directly. This replaces the materialized weight multiply, + // permutation/copy, and row reduction with one reusable kernel. + routed_out = ggml_laguna_moe_combine( + ctx, down_e, owner_weights_tensor); + } else { + ggml_tensor * weights_3d = ggml_reshape_3d( + ctx, owner_weights_tensor, 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); + } ggml_tensor * combined_out = routed_out; if (has_shared) { diff --git a/server/test/test_deepseek4_mmid_grouped_cuda.cpp b/server/test/test_deepseek4_mmid_grouped_cuda.cpp index 58f062e71..940ddb9da 100644 --- a/server/test/test_deepseek4_mmid_grouped_cuda.cpp +++ b/server/test/test_deepseek4_mmid_grouped_cuda.cpp @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -31,10 +32,25 @@ static bool run_case( bool fused_ds4, bool write_output, std::ofstream & output) { - constexpr int k_dim = 256; - constexpr int n_rows = 128; - constexpr int n_experts = 32; - constexpr int top_k = 8; + int k_dim = 256; + int n_rows = 128; + int n_experts = 32; + int top_k = 8; + const auto env_positive = [](const char * name, int fallback) { + const char * raw = std::getenv(name); + const int parsed = raw ? std::atoi(raw) : 0; + return parsed > 0 ? parsed : fallback; + }; + if (std::getenv("DFLASH_MMID_BENCH_ITERS")) { + k_dim = env_positive("DFLASH_MMID_BENCH_K", k_dim); + n_rows = env_positive("DFLASH_MMID_BENCH_ROWS", n_rows); + n_experts = env_positive("DFLASH_MMID_BENCH_EXPERTS", n_experts); + top_k = env_positive("DFLASH_MMID_BENCH_TOP_K", top_k); + } + if (top_k > n_experts) { + std::fprintf(stderr, "top_k=%d exceeds n_experts=%d\n", top_k, n_experts); + return false; + } ggml_init_params params = {16 * 1024 * 1024, nullptr, true}; ggml_context * ctx = ggml_init(params); @@ -70,9 +86,13 @@ static bool run_case( return false; } + const bool benchmark = std::getenv("DFLASH_MMID_BENCH_ITERS") != nullptr; std::mt19937 rng(20260713u + (unsigned) type * 97u + (unsigned) width); std::uniform_real_distribution dist(-1.0f, 1.0f); - std::vector weights_f((size_t) k_dim * n_rows * n_experts); + std::vector weights_f; + if (!benchmark) { + weights_f.resize((size_t) k_dim * n_rows * n_experts); + } std::vector input_f((size_t) k_dim * width); for (float & value : weights_f) { value = dist(rng); @@ -82,7 +102,11 @@ static bool run_case( } std::vector weights_q(ggml_nbytes(weights)); - const size_t quantized = + // A zero-filled quantized tensor is valid and exercises the identical GPU + // load/dequantize path without constructing multi-gigabyte F32 weights for + // realistic expert-count benchmarks. Correctness runs still use quantized + // randomized weights. + const size_t quantized = benchmark ? weights_q.size() : type == GGML_TYPE_Q2_0_ROCMFP2 ? rocmfpx_quantize_fp2( weights_f.data(), weights_q.data(), n_rows * n_experts, @@ -126,7 +150,27 @@ static bool run_case( ggml_backend_tensor_set(ids, ids_h.data(), 0, ids_h.size() * sizeof(int32_t)); ggml_backend_synchronize(backend); - const ggml_status status = ggml_backend_graph_compute(backend, graph); + ggml_status status = ggml_backend_graph_compute(backend, graph); + int benchmark_iterations = 0; + if (const char * raw = std::getenv("DFLASH_MMID_BENCH_ITERS")) { + benchmark_iterations = std::max(0, std::atoi(raw)); + } + if (status == GGML_STATUS_SUCCESS && benchmark_iterations > 0) { + ggml_backend_synchronize(backend); + const auto start = std::chrono::steady_clock::now(); + for (int i = 0; i < benchmark_iterations && status == GGML_STATUS_SUCCESS; ++i) { + status = ggml_backend_graph_compute(backend, graph); + } + ggml_backend_synchronize(backend); + const auto end = std::chrono::steady_clock::now(); + const double average_us = std::chrono::duration( + end - start).count() / benchmark_iterations; + std::printf( + "[mmid-grouped-test] benchmark type=%s width=%d experts=%d top_k=%d " + "k=%d rows=%d iterations=%d average_us=%.3f\n", + ggml_type_name(type), width, n_experts, top_k, k_dim, n_rows, + benchmark_iterations, average_us); + } std::vector result_h(ggml_nelements(result)); if (status == GGML_STATUS_SUCCESS) { ggml_backend_synchronize(backend); @@ -201,18 +245,28 @@ static int run_child(const char * mode, const char * output_path) { const ggml_type types[] = { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q5_K, GGML_TYPE_Q2_0_ROCMFP2, GGML_TYPE_Q3_0_ROCMFPX, + GGML_TYPE_Q4_0_ROCMFP4_FAST, }; - const int widths[] = {2, 4, 8, 9, 16, 32}; + const int widths[] = {2, 4, 8, 9, 16, 32, 48, 64}; + int width_filter = 0; + if (const char * raw = std::getenv("DFLASH_MMID_TEST_WIDTH")) { + width_filter = std::max(0, std::atoi(raw)); + } bool ok = output.good(); for (ggml_type type : types) { + if (type == GGML_TYPE_Q4_0_ROCMFP4_FAST && width_filter == 0) { + continue; + } for (int width : widths) { - if (width == 32 && + if ((width_filter > 0 && width != width_filter) || + (width >= 32 && type != GGML_TYPE_Q2_0_ROCMFP2 && - type != GGML_TYPE_Q3_0_ROCMFPX) { + type != GGML_TYPE_Q3_0_ROCMFPX && + type != GGML_TYPE_Q4_0_ROCMFP4_FAST)) { continue; } ok = run_case(backend, type, width, false, true, output) && ok; - if (width < 32) { + if (width < 32 && width_filter == 0) { ok = run_case(backend, type, width, true, true, output) && ok; } } @@ -320,6 +374,196 @@ static bool grouped_supported_device() { #endif } +struct CombineRun { + std::vector output; + double average_us = 0.0; + bool ok = false; +}; + +static CombineRun run_combine_graph( + ggml_backend_t backend, + bool fused, + int n_embd, + int n_used, + int n_tokens, + const std::vector & experts_h, + const std::vector & weights_h, + int iterations) { + CombineRun run; + ggml_init_params params = {16 * 1024 * 1024, nullptr, true}; + ggml_context * ctx = ggml_init(params); + if (ctx == nullptr) { + return run; + } + + 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 * result = nullptr; + if (fused) { + result = ggml_laguna_moe_combine(ctx, experts, weights); + } else { + ggml_tensor * weights_3d = ggml_reshape_3d( + ctx, weights, 1, n_used, n_tokens); + result = ggml_mul(ctx, experts, weights_3d); + result = ggml_cont(ctx, ggml_permute(ctx, result, 1, 0, 2, 3)); + result = ggml_sum_rows(ctx, result); + result = ggml_reshape_2d(ctx, result, n_embd, n_tokens); + } + ggml_set_output(result); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, result); + + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + if (!alloc || !ggml_gallocr_alloc_graph(alloc, graph)) { + if (alloc) ggml_gallocr_free(alloc); + ggml_free(ctx); + return run; + } + + ggml_backend_tensor_set( + experts, experts_h.data(), 0, experts_h.size() * sizeof(float)); + ggml_backend_tensor_set( + weights, weights_h.data(), 0, weights_h.size() * sizeof(float)); + ggml_status status = ggml_backend_graph_compute(backend, graph); + ggml_backend_synchronize(backend); + + const auto start = std::chrono::steady_clock::now(); + for (int i = 0; i < iterations && status == GGML_STATUS_SUCCESS; ++i) { + status = ggml_backend_graph_compute(backend, graph); + } + ggml_backend_synchronize(backend); + const auto end = std::chrono::steady_clock::now(); + + run.output.resize((size_t)n_embd * n_tokens); + if (status == GGML_STATUS_SUCCESS) { + ggml_backend_tensor_get( + result, run.output.data(), 0, run.output.size() * sizeof(float)); + run.average_us = std::chrono::duration( + end - start).count() / std::max(iterations, 1); + run.ok = true; + } + + ggml_gallocr_free(alloc); + ggml_free(ctx); + return run; +} + +static bool run_combine_parity_and_benchmark(ggml_backend_t backend) { + const char * bench_raw = std::getenv("DFLASH_MOE_COMBINE_BENCH"); + const bool benchmark = bench_raw && *bench_raw && std::strcmp(bench_raw, "0") != 0; + const int n_embd = benchmark ? 4096 : 260; + const int n_used = 6; + int n_tokens = benchmark ? 3072 : 33; + if (const char * raw = std::getenv("DFLASH_MOE_COMBINE_BENCH_TOKENS")) { + const int requested = std::atoi(raw); + if (requested > 0) n_tokens = requested; + } + const int iterations = benchmark ? 20 : 2; + + std::vector experts((size_t)n_embd * n_used * n_tokens); + std::vector weights((size_t)n_used * n_tokens); + for (size_t i = 0; i < experts.size(); ++i) { + experts[i] = ((int)(i % 251) - 125) * (1.0f / 127.0f); + } + for (int t = 0; t < n_tokens; ++t) { + float total = 0.0f; + for (int e = 0; e < n_used; ++e) { + float value = (float)(e + 1 + t % 7); + if ((t + e) % 11 == 0) value = 0.0f; + weights[(size_t)t * n_used + e] = value; + total += value; + } + for (int e = 0; e < n_used; ++e) { + weights[(size_t)t * n_used + e] /= total; + } + } + + std::vector expected((size_t)n_embd * n_tokens); + 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 float weight = weights[(size_t)t * n_used + e]; + if (weight == 0.0f) { + if (e == 0) sum = 0.0f; + continue; + } + const float product = + experts[((size_t)t * n_used + e) * n_embd + h] * weight; + sum = e == 0 ? product : sum + product; + } + expected[(size_t)t * n_embd + h] = sum; + } + } + + CombineRun legacy; + if (benchmark) { + legacy = run_combine_graph( + backend, false, n_embd, n_used, n_tokens, + experts, weights, iterations); + } + const CombineRun fused = run_combine_graph( + backend, true, n_embd, n_used, n_tokens, + experts, weights, iterations); + if (!fused.ok || fused.output.size() != expected.size() || + (benchmark && (!legacy.ok || legacy.output.size() != expected.size()))) { + return false; + } + + double fused_squared_error = 0.0; + double legacy_squared_error = 0.0; + double reference_power = 0.0; + float fused_max_abs_error = 0.0f; + size_t fused_exact = 0; + for (size_t i = 0; i < expected.size(); ++i) { + if (std::memcmp(&expected[i], &fused.output[i], sizeof(float)) == 0) { + ++fused_exact; + } + const float fused_error = fused.output[i] - expected[i]; + fused_max_abs_error = std::max( + fused_max_abs_error, std::fabs(fused_error)); + fused_squared_error += (double)fused_error * fused_error; + if (benchmark) { + const float legacy_error = legacy.output[i] - expected[i]; + legacy_squared_error += (double)legacy_error * legacy_error; + } + reference_power += (double)expected[i] * expected[i]; + } + const double fused_nmse = + fused_squared_error / std::max(reference_power, 1e-30); + const double legacy_nmse = benchmark + ? legacy_squared_error / std::max(reference_power, 1e-30) : 0.0; + const double speedup = benchmark && fused.average_us > 0.0 + ? legacy.average_us / fused.average_us : 0.0; + if (!std::isfinite(fused_nmse) || fused_nmse > 1e-12) { + for (size_t i = 0; i < std::min(expected.size(), 8); ++i) { + std::fprintf(stderr, + "combine mismatch[%zu] expected=%g fused=%g\n", + i, expected[i], fused.output[i]); + } + } + if (benchmark) { + std::printf( + "[mmid-grouped-test] combine n_embd=%d n_used=%d n_tokens=%d " + "legacy_us=%.3f fused_us=%.3f speedup=%.3fx " + "fused_max_abs=%g fused_nmse=%g legacy_nmse=%g\n", + n_embd, n_used, n_tokens, legacy.average_us, fused.average_us, + speedup, fused_max_abs_error, fused_nmse, legacy_nmse); + } else { + std::printf( + "[mmid-grouped-test] combine-vec4 exact=%zu/%zu " + "max_abs=%g nmse=%g\n", + fused_exact, expected.size(), fused_max_abs_error, fused_nmse); + } + return std::isfinite(fused_nmse) && fused_nmse <= 1e-12; +} + static std::string shell_quote(const std::string & value) { #if defined(_WIN32) std::string quoted = "\""; @@ -370,9 +614,11 @@ int main(int argc, char ** argv) { if (argc == 4 && std::strcmp(argv[1], "--child") == 0) { return run_child(argv[2], argv[3]); } - if (argc != 1) { + const bool combine_only = + argc == 2 && std::strcmp(argv[1], "--combine-only") == 0; + if (argc != 1 && !combine_only) { std::fprintf(stderr, - "usage: %s [--child legacy|grouped|masked-fused OUTPUT]\n", + "usage: %s [--combine-only|--child legacy|grouped|masked-fused OUTPUT]\n", argv[0]); return 2; } @@ -382,6 +628,29 @@ int main(int argc, char ** argv) { return 77; } +#if defined(_WIN32) + if (!std::getenv("DFLASH_MOE_COMBINE_VEC4")) { + _putenv_s("DFLASH_MOE_COMBINE_VEC4", "1"); + } +#else + setenv("DFLASH_MOE_COMBINE_VEC4", "1", 0); +#endif + + ggml_backend_t combine_backend = ggml_backend_cuda_init(0); + if (combine_backend == nullptr) { + std::fprintf(stderr, "GPU backend unavailable for combine parity\n"); + return 1; + } + const bool combine_parity = run_combine_parity_and_benchmark(combine_backend); + ggml_backend_free(combine_backend); + if (!combine_parity) { + std::fprintf(stderr, "grouped MoE fused-combine parity failed\n"); + return 1; + } + if (combine_only) { + return 0; + } + #if defined(_WIN32) const long long pid = (long long) _getpid(); #else @@ -417,7 +686,7 @@ int main(int argc, char ** argv) { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q5_K, GGML_TYPE_Q2_0_ROCMFP2, GGML_TYPE_Q3_0_ROCMFPX, }; - const int widths[] = {2, 4, 8, 9, 16, 32}; + const int widths[] = {2, 4, 8, 9, 16, 32, 48, 64}; size_t offset = 0; size_t compared_bytes = 0; int compared_cases = 0; @@ -427,14 +696,14 @@ int main(int argc, char ** argv) { bool grouped_dispatch = true; for (ggml_type type : types) { for (int width : widths) { - if (width == 32 && + if (width >= 32 && type != GGML_TYPE_Q2_0_ROCMFP2 && type != GGML_TYPE_Q3_0_ROCMFPX) { continue; } const bool legacy_mmvq = has_mmvq_record(legacy_log, type, width); for (bool fused_ds4 : {false, true}) { - if (width == 32 && fused_ds4) { + if (width >= 32 && fused_ds4) { continue; } const size_t case_bytes = (size_t) 128 * 8 * width * sizeof(float); @@ -461,7 +730,7 @@ int main(int argc, char ** argv) { } } } - output_parity = output_parity && offset == legacy.size() && compared_cases == 72; + output_parity = output_parity && offset == legacy.size() && compared_cases == 76; const size_t masked_case_bytes = (size_t) 128 * 8 * 32 * sizeof(float); const bool masked_fused_zero = masked_fused.size() == masked_case_bytes; const bool pass = legacy_status == 0 && grouped_status == 0 && From b0d5b1b816d1a4cf52aca389c4bf1843cd8e54a6 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 31 Aug 2026 02:11:45 +0200 Subject: [PATCH 4/8] perf(hip): vectorize sparse attention staging --- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 92 ++++++++++++++----- server/docs/DS4.md | 6 +- server/docs/ENVIRONMENT.md | 1 + server/tests/test_deepseek4_unit.cpp | 53 ++++++++--- 4 files changed, 117 insertions(+), 35 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 49bfbea17..6911da4c7 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -7,6 +7,8 @@ #include "fattn-chunked.cuh" #include "fattn.cuh" +#include + #if defined(GGML_USE_HIP) __device__ static float ds4_fa_block_sum(float v) { @@ -1563,7 +1565,7 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( // backend-generic (D512 MQA with K == V and direct indexed rows); model policy // remains in the graph/backend layer. template + int KEYS_PER_STAGE = 16, bool STAGE_F32 = false> __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( float * dst, const float * q, @@ -1602,7 +1604,8 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( // required because each stage has workgroup-wide barriers. if (token >= n_tokens || head_begin + HEADS_PER_BLOCK > n_heads) return; - __shared__ KV staged_kv[KEYS_PER_STAGE * D]; + using stage_type = std::conditional_t; + __shared__ __align__(16) stage_type staged_kv[KEYS_PER_STAGE * D]; __shared__ int staged_rows[KEYS_PER_STAGE]; __shared__ float staged_masks[KEYS_PER_STAGE]; @@ -1662,13 +1665,35 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( } __syncthreads(); - for (int index = tid; index < KEYS_PER_STAGE * D; - index += N_THREADS) { - const int slot = index / D; - const int dim = index - slot * D; - const int row = staged_rows[slot]; - staged_kv[index] = row >= 0 && row < n_kv - ? kv[(size_t) row * D + dim] : KV{}; + // Convert aligned half2 pairs once while loading them. Every head in + // the block then consumes the same F32 LDS values without repeating + // half conversion in both the score and value passes. + if constexpr (STAGE_F32 && std::is_same_v) { + constexpr int PAIRS_PER_ROW = D / 2; + for (int pair_index = tid; + pair_index < KEYS_PER_STAGE * PAIRS_PER_ROW; + pair_index += N_THREADS) { + const int slot = pair_index / PAIRS_PER_ROW; + const int pair = pair_index - slot * PAIRS_PER_ROW; + const int row = staged_rows[slot]; + float2 unpacked = make_float2(0.0f, 0.0f); + if (row >= 0 && row < n_kv) { + ds4_fa_load_pair( + kv + (size_t) row * D + 2 * pair, + unpacked.x, unpacked.y); + } + reinterpret_cast(staged_kv)[pair_index] = unpacked; + } + } else { + for (int index = tid; index < KEYS_PER_STAGE * D; + index += N_THREADS) { + const int slot = index / D; + const int dim = index - slot * D; + const int row = staged_rows[slot]; + staged_kv[index] = row >= 0 && row < n_kv + ? static_cast(kv[(size_t) row * D + dim]) + : stage_type{}; + } } __syncthreads(); @@ -1684,7 +1709,8 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( for (int i = 0; i < VALUES_PER_LANE; ++i) { const int dim = lane + i * WAVE; partial += q_values[i] * - ds4_fa_load(staged_kv + slot * D + dim); + ds4_fa_load( + staged_kv + slot * D + dim); } partial = warp_reduce_sum(partial); // XOR reduction can associate operands differently in each lane. @@ -1702,7 +1728,7 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( #pragma unroll for (int i = 0; i < VALUES_PER_LANE; ++i) { const int dim = lane + i * WAVE; - const float value = ds4_fa_load( + const float value = ds4_fa_load( staged_kv + slot * D + dim); accum[i] = accum[i] * old_scale + value_scale * value; } @@ -2604,22 +2630,42 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( device_warp_size == 32 && n_tokens >= streaming_min_tokens && active_row_upper_bound > 0 && n_kv >= 3 * active_row_upper_bound) { + const char * f32_stage_env = + getenv("GGML_CUDA_MLA_STREAM_F32_STAGE"); + const bool f32_stage = f32_stage_env && f32_stage_env[0] != '\0' && + strcmp(f32_stage_env, "0") != 0; constexpr int streaming_heads = 16; const dim3 streaming_grid( (unsigned) n_tokens, (unsigned) (n_heads / streaming_heads), 1); - ds4_flash_attn_d512_streaming_topk_kernel - <<>>( - (float *) dst->data, (const float *) Q->data, - q_stride_token, q_stride_head, - (const half *) K->data, - (const half *) mask->data, - sinks ? (const float *) sinks->data : nullptr, - n_tokens, n_heads, n_kv, scale, - visibility_bounds, indexed_rows, indexed_counts, - indexed_capacity, inverse_rope, - inverse_rope_coefficients, - forward_rope_coefficients); + if (f32_stage) { + ds4_flash_attn_d512_streaming_topk_kernel< + half, half, streaming_heads, 16, true> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients); + } else { + ds4_flash_attn_d512_streaming_topk_kernel + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients); + } CUDA_CHECK(cudaGetLastError()); return true; } diff --git a/server/docs/DS4.md b/server/docs/DS4.md index b1d42c2eb..d62dac430 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -494,7 +494,11 @@ selected latent row across wave32 heads and avoids materializing scores. It is currently limited to F16 caches on native wave32 devices and otherwise falls back to the existing path. Because its online softmax changes floating-point association, keep it opt-in until the target model passes a matched output and -throughput A/B. `GGML_DS4_FA_STREAM_TOPK` remains a compatibility alias. +throughput A/B. `GGML_DS4_FA_STREAM_TOPK` remains a compatibility alias. Add +`GGML_CUDA_MLA_STREAM_F32_STAGE=1` to convert each selected F16 latent once +while staging aligned pairs in shared memory instead of repeating conversion +for every head. The isolated gfx1151 qualification is byte-identical to F16 +staging and reduced alternating-run kernel time by about 9%. Indexed verifier attention with at most eight query rows also uses a two-way split-KV schedule on gfx1151. Set `GGML_CUDA_MLA_NO_SPLIT_KV=1` to restore the diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index 0202b3550..46286e80c 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -51,6 +51,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_GFX1151_HC_MMVF_Q4` | 1 on gfx1151 for the DS4 `[16384,24]` q4 projection | BURN-IN KILL SWITCH: =0 restores the generic hipBLAS dispatch decision. | | `GGML_CUDA_MMQ_X` | unset | DEBUG: force a supported MMQ output-column tile width (8–128) for architecture tuning; invalid or over-budget values fall back to automatic selection. | | `GGML_CUDA_MMQ_MOE_ADAPTIVE_X` | unset | BURN-IN: on sparse-route gfx1151 grouped MoE MMQ, choose the measured ROCmFP2/3/4 output tile from routed rows per expert; ordinary matmuls, unmeasured formats, and other devices are unchanged. | +| `GGML_CUDA_MLA_STREAM_F32_STAGE` | unset | EXPERIMENTAL: with streaming D512 indexed attention, convert aligned F16 pairs once while staging them in shared memory instead of repeating conversion for every head. | | `GGML_CUDA_MLA_SPLIT_KV` / `GGML_DS4_FA_SPLIT_KV` | 1 on gfx1151 indexed decode; unset elsewhere | BURN-IN: force the reusable split-KV MLA schedule. Set `GGML_CUDA_MLA_NO_SPLIT_KV=1` (or legacy `GGML_DS4_FA_NO_SPLIT_KV=1`) to disable it. | | `GGML_DS4_TOPK_BLOCK_RADIX` | 1 on gfx1151 | BURN-IN KILL SWITCH: =0 restores hipCUB full sort for DS4-shaped 512-row top-k selection. | | `GGML_DS4_FA_SERIAL_INDEX_SCAN` | unset | DEBUG/A-B: restore the serial indexed-attention mask scan instead of the long-context HIP parallel scan. | diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index bdef57893..1b3126a58 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -2966,12 +2966,14 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { topk_data.size() * sizeof(int32_t)); ScopedEnvVar streaming_guard("GGML_CUDA_MLA_STREAM_TOPK"); + ScopedEnvVar f32_stage_guard("GGML_CUDA_MLA_STREAM_F32_STAGE"); ScopedCudaGraphOverrides eager( /*disable_graphs=*/true, /*mmvq_max_ncols=*/0, /*skip_property_check=*/false); std::vector reference((size_t) ggml_nelements(output)); - std::vector candidate(reference.size()); + std::vector candidate_f16_stage(reference.size()); + std::vector candidate_f32_stage(reference.size()); setenv("GGML_CUDA_MLA_STREAM_TOPK", "0", 1); TEST_ASSERT_MSG( ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, @@ -2980,23 +2982,36 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { reference.size() * sizeof(float)); setenv("GGML_CUDA_MLA_STREAM_TOPK", "1", 1); + setenv("GGML_CUDA_MLA_STREAM_F32_STAGE", "0", 1); TEST_ASSERT_MSG( ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, - "streaming top-k attention failed"); - ggml_backend_tensor_get(output, candidate.data(), 0, - candidate.size() * sizeof(float)); + "F16-staged streaming top-k attention failed"); + ggml_backend_tensor_get(output, candidate_f16_stage.data(), 0, + candidate_f16_stage.size() * sizeof(float)); + + setenv("GGML_CUDA_MLA_STREAM_F32_STAGE", "1", 1); + TEST_ASSERT_MSG( + ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, + "F32-staged streaming top-k attention failed"); + ggml_backend_tensor_get(output, candidate_f32_stage.data(), 0, + candidate_f32_stage.size() * sizeof(float)); for (size_t i = 0; i < reference.size(); ++i) { TEST_ASSERT_MSG( - std::isfinite(candidate[i]), + std::isfinite(candidate_f32_stage[i]), "streaming top-k attention output must be finite"); TEST_ASSERT_MSG( - nearly_equal(reference[i], candidate[i], 5.0e-4f, 5.0e-4f), + nearly_equal(reference[i], candidate_f32_stage[i], + 5.0e-4f, 5.0e-4f), "streaming top-k attention exceeded numeric tolerance"); + TEST_ASSERT_MSG( + candidate_f16_stage[i] == candidate_f32_stage[i], + "F32 staging changed streaming top-k attention output"); } - auto measure_us = [&](bool streaming) { + auto measure_us = [&](bool streaming, bool f32_stage) { setenv("GGML_CUDA_MLA_STREAM_TOPK", streaming ? "1" : "0", 1); + setenv("GGML_CUDA_MLA_STREAM_F32_STAGE", f32_stage ? "1" : "0", 1); constexpr int warmups = 3; constexpr int iterations = 20; for (int i = 0; i < warmups; ++i) { @@ -3012,10 +3027,26 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { return std::chrono::duration(end - begin).count() / iterations; }; - const double grouped_us = measure_us(false); - const double streaming_us = measure_us(true); - std::fprintf(stderr, " grouped=%.1fus streaming=%.1fus speedup=%.2fx", - grouped_us, streaming_us, grouped_us / streaming_us); + const double grouped_us = measure_us(false, false); + constexpr int timing_rounds = 4; + double streaming_f16_us = 0.0; + double streaming_f32_us = 0.0; + for (int round = 0; round < timing_rounds; ++round) { + if ((round & 1) == 0) { + streaming_f16_us += measure_us(true, false); + streaming_f32_us += measure_us(true, true); + } else { + streaming_f32_us += measure_us(true, true); + streaming_f16_us += measure_us(true, false); + } + } + streaming_f16_us /= timing_rounds; + streaming_f32_us /= timing_rounds; + std::fprintf(stderr, + " grouped=%.1fus streaming_f16=%.1fus" + " streaming_f32=%.1fus speedup=%.2fx", + grouped_us, streaming_f16_us, streaming_f32_us, + grouped_us / streaming_f32_us); } ggml_gallocr_free(alloc); From 1b6cd67618ed8c450d6d2e8af7d05e3c5179029e Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 31 Aug 2026 02:25:01 +0200 Subject: [PATCH 5/8] perf(hip): coalesce ROCmFP3 MMQ loads --- .../deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh | 38 ++++++++++++++++--- 1 file changed, 32 insertions(+), 6 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index a354390b7..45e73686c 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh @@ -1061,17 +1061,43 @@ static __device__ __forceinline__ void load_tiles_rocmfpx_dual( } const block_t * block = (const block_t *) x + kbx0 + i*stride + kbx; - const int k0 = kbx*groups_per_block + group; - const int q0 = traits::pack4(block, 4*group); - const int q1 = traits::pack4( - block, 4*(group + groups_per_block/2)); + int k0; + int q0; + int q1; + int q1_offset; + if constexpr (type == GGML_TYPE_Q3_0_ROCMFPX) { + // Eight FP3 weights occupy exactly three bytes. Assign adjacent + // four-value groups to one lane so the wave reads every packed + // byte once instead of overlapping the two-byte group windows. + const int byte = 3*group; +#if defined(GGML_USE_HIP) + uint32_t packed32; + __builtin_memcpy(&packed32, block->qs + byte, sizeof(packed32)); + const uint32_t bits24 = packed32 & 0x00ffffffu; +#else + const uint32_t bits24 = + (uint32_t) block->qs[byte + 0] | + ((uint32_t) block->qs[byte + 1] << 8) | + ((uint32_t) block->qs[byte + 2] << 16); +#endif + k0 = kbx*groups_per_block + 2*group; + q0 = rocmfpx_pack4_fp3_bits12_vec_cuda(bits24 & 0x0fffu); + q1 = rocmfpx_pack4_fp3_bits12_vec_cuda((bits24 >> 12) & 0x0fffu); + q1_offset = 1; + } else { + k0 = kbx*groups_per_block + group; + q0 = traits::pack4(block, 4*group); + q1 = traits::pack4( + block, 4*(group + groups_per_block/2)); + q1_offset = groups_per_block/2; + } #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + k0] = q0; - x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + k0 + groups_per_block/2] = q1; + x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + k0 + q1_offset] = q1; #else x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0] = q0; - x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + groups_per_block/2] = q1; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + q1_offset] = q1; #endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) } From f140639964f34653393ea2683aa174727a2d50b6 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 31 Aug 2026 02:29:03 +0200 Subject: [PATCH 6/8] perf(hip): pair ROCmFP2 MMQ loads --- .../deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh | 20 ++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index 45e73686c..079d9129f 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh @@ -1085,11 +1085,21 @@ static __device__ __forceinline__ void load_tiles_rocmfpx_dual( q1 = rocmfpx_pack4_fp3_bits12_vec_cuda((bits24 >> 12) & 0x0fffu); q1_offset = 1; } else { - k0 = kbx*groups_per_block + group; - q0 = traits::pack4(block, 4*group); - q1 = traits::pack4( - block, 4*(group + groups_per_block/2)); - q1_offset = groups_per_block/2; + // FP2 stores each four-value group in one byte. Pair adjacent + // groups so HIP can issue one aligned 16-bit load per lane. + const int byte = 2*group; +#if defined(GGML_USE_HIP) + uint16_t bits16; + __builtin_memcpy(&bits16, block->qs + byte, sizeof(bits16)); +#else + const uint16_t bits16 = + (uint16_t) block->qs[byte + 0] | + ((uint16_t) block->qs[byte + 1] << 8); +#endif + k0 = kbx*groups_per_block + 2*group; + q0 = rocmfpx_pack4_fp2_bits8_vec_cuda(bits16 & 0x00ffu); + q1 = rocmfpx_pack4_fp2_bits8_vec_cuda((bits16 >> 8) & 0x00ffu); + q1_offset = 1; } #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) From f1039daac2b8f473f16f3f48b646626169ced205 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 31 Aug 2026 02:49:50 +0200 Subject: [PATCH 7/8] perf(hip): compact sparse MoE prefill work --- .../deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh | 176 ++++++++++++++++++ server/docs/ENVIRONMENT.md | 2 + .../test/test_deepseek4_mmid_grouped_cuda.cpp | 4 +- 3 files changed, 181 insertions(+), 1 deletion(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index 079d9129f..ef05d61d0 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh @@ -4658,6 +4658,112 @@ static __global__ void mul_mat_q( mix_codebooks, mix_modes, fastdiv(zt, channel_ratio)); } +#if defined(GGML_USE_HIP) +template +static __global__ void mul_mat_q_moe_build_tasks( + const int32_t * __restrict__ expert_bounds, + int2 * __restrict__ tasks, + int * __restrict__ task_count, + int n_experts) { + const int expert = blockIdx.x*blockDim.x + threadIdx.x; + if (expert >= n_experts) { + return; + } + const int route_count = + expert_bounds[expert + 1] - expert_bounds[expert]; + const int tile_count = (route_count + mmq_x - 1)/mmq_x; + if (tile_count == 0) { + return; + } + const int task_begin = atomicAdd(task_count, tile_count); + for (int tile = 0; tile < tile_count; ++tile) { + tasks[task_begin + tile] = make_int2(tile, expert); + } +} + +// Sparse grouped MoE has a deliberately wide upper-bound grid: every expert +// receives enough Y tiles for the full token batch even though only top-k +// routes are live. Build a compact device-side task list, then keep a bounded +// set of workgroups resident to consume only non-empty expert tiles. No host +// count readback or synchronization is required. +template +__launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 1) +static __global__ void mul_mat_q_moe_persistent( + const char * __restrict__ x, + const int * __restrict__ y, + const int32_t * __restrict__ ids_dst, + const int32_t * __restrict__ expert_bounds, + const nv_bfloat16 * __restrict__ mix_codebooks, + const uint8_t * __restrict__ mix_modes, + float * __restrict__ dst, + const int2 * __restrict__ tasks, + const int * __restrict__ task_count, + const uint3 blocks_per_ne00, + const int nrows_x, + const int stride_row_x, + const int ncols_y, + const int stride_col_dst, + const uint3 channel_ratio, + const int stride_channel_x) { + if (mmq_x > get_mmq_x_max_device() || + mmq_x % mmq_get_granularity_device(mmq_x) != 0) { + NO_DEVICE_CODE; + return; + } + + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int mmq_y = get_mmq_y_device(); + + extern __shared__ int ids_dst_shared[]; + const int it = blockIdx.x; + const int total_tasks = *task_count; + + for (int task_index = blockIdx.y; task_index < total_tasks; + task_index += gridDim.y) { + const int2 task = tasks[task_index]; + const int jt = task.x; + const int zt = task.y; + + const int col_low = expert_bounds[zt + 0]; + const int col_high = expert_bounds[zt + 1]; + const int col_diff = col_high - col_low; + if (jt*mmq_x >= col_diff) { + continue; + } + + __syncthreads(); +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps*warp_size) { + const int j = j0 + threadIdx.y*warp_size + threadIdx.x; + if (j0 + nwarps*warp_size > mmq_x && j >= mmq_x) { + break; + } + ids_dst_shared[j] = ids_dst[col_low + jt*mmq_x + j]; + } + __syncthreads(); + + const int offset_y = + (col_low + jt*mmq_x)*(sizeof(block_q8_1_mmq)/sizeof(int)); + const int offset_dst = it*mmq_y; + const int offset_x = + fastdiv(zt, channel_ratio)*stride_channel_x + + it*mmq_y*stride_row_x; + const int tile_x_max_i = nrows_x - it*mmq_y - 1; + const int tile_y_max_j = col_diff - jt*mmq_x - 1; + + constexpr bool fixup = false; + mul_mat_q_process_tile( + x, offset_x, y + offset_y, ids_dst_shared, + dst + offset_dst, nullptr, stride_row_x, ncols_y, + stride_col_dst, tile_x_max_i, tile_y_max_j, + 0, blocks_per_ne00.z, mix_codebooks, mix_modes, + fastdiv(zt, channel_ratio)); + __syncthreads(); + } +} +#endif + template __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device()/2, 1) static __global__ void mul_mat_q_stream_k_fixup( @@ -4855,6 +4961,76 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a const uint3 channel_ratio_fd = init_fastdiv_values(channel_ratio); const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio); +#if defined(GGML_USE_HIP) + if constexpr (type == GGML_TYPE_Q2_0_ROCMFP2 || + type == GGML_TYPE_Q3_0_ROCMFPX || + type == GGML_TYPE_Q4_0_ROCMFP4_FAST) { + const char * persistent_env = + std::getenv("GGML_CUDA_MMQ_MOE_PERSISTENT"); + const bool persistent_enabled = persistent_env && *persistent_env && + !(persistent_env[0] == '0' && persistent_env[1] == '\0'); + // The compact queue amortizes its builder and bounded-worker launch at + // prefill widths. Below 256 source columns the ordinary grouped grid is + // already efficient, and is marginally faster for some formats. + if (persistent_enabled && args.ncols_max >= 256 && !args.use_stream_k && + args.ids_dst != nullptr && args.expert_bounds != nullptr && + args.nsamples_y == 1 && + cc == GGML_CUDA_CC_OFFSET_AMD + 0x1151) { + int blocks_per_cu = 32; + if (const char * raw = + std::getenv("GGML_CUDA_MMQ_MOE_PERSISTENT_BLOCKS_PER_CU")) { + const int parsed = std::atoi(raw); + if (parsed >= 1 && parsed <= 32) { + blocks_per_cu = parsed; + } + } + ggml_cuda_pool_alloc tasks(ctx.pool(), args.ncols_y); + ggml_cuda_pool_alloc task_count(ctx.pool(), 1); + CUDA_CHECK(cudaMemsetAsync( + task_count.get(), 0, sizeof(int), stream)); + constexpr int task_builder_threads = 256; + const int task_builder_blocks = + (args.nchannels_y + task_builder_threads - 1)/ + task_builder_threads; + mul_mat_q_moe_build_tasks + <<>>( + args.expert_bounds, tasks.get(), task_count.get(), + args.nchannels_y); + + const int workers = std::max( + 1, std::min((int) args.ncols_y, nsm*blocks_per_cu)); + const dim3 persistent_grid(nty, workers, 1); + CUDA_SET_SHARED_MEMORY_LIMIT( + (mul_mat_q_moe_persistent), + nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT( + (mul_mat_q_moe_persistent), + nbytes_shared); + if (args.nrows_x % mmq_y == 0) { + mul_mat_q_moe_persistent + <<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, + args.mix_codebooks, args.mix_modes, args.dst, + tasks.get(), task_count.get(), + blocks_per_ne00_fd, args.nrows_x, args.stride_row_x, + args.ncols_y, args.nrows_dst, channel_ratio_fd, + args.stride_channel_x); + } else { + mul_mat_q_moe_persistent + <<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, + args.mix_codebooks, args.mix_modes, args.dst, + tasks.get(), task_count.get(), + blocks_per_ne00_fd, args.nrows_x, args.stride_row_x, + args.ncols_y, args.nrows_dst, channel_ratio_fd, + args.stride_channel_x); + } + CUDA_CHECK(cudaGetLastError()); + return; + } + } +#endif + if (!args.use_stream_k) { if (args.nrows_x % mmq_y == 0) { constexpr bool need_check = false; diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index 46286e80c..b304f9532 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -51,6 +51,8 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_GFX1151_HC_MMVF_Q4` | 1 on gfx1151 for the DS4 `[16384,24]` q4 projection | BURN-IN KILL SWITCH: =0 restores the generic hipBLAS dispatch decision. | | `GGML_CUDA_MMQ_X` | unset | DEBUG: force a supported MMQ output-column tile width (8–128) for architecture tuning; invalid or over-budget values fall back to automatic selection. | | `GGML_CUDA_MMQ_MOE_ADAPTIVE_X` | unset | BURN-IN: on sparse-route gfx1151 grouped MoE MMQ, choose the measured ROCmFP2/3/4 output tile from routed rows per expert; ordinary matmuls, unmeasured formats, and other devices are unchanged. | +| `GGML_CUDA_MMQ_MOE_PERSISTENT` | unset | EXPERIMENTAL: on prefill-sized (at least 256-token) sparse grouped ROCmFP2/3/4 MMQ on gfx1151, build a compact device-side expert-tile queue and consume it with bounded persistent workers. Short batches, ordinary matmuls, unmeasured formats, and other devices are unchanged. | +| `GGML_CUDA_MMQ_MOE_PERSISTENT_BLOCKS_PER_CU` | 32 | DEBUG: set the compact grouped-MoE worker budget per gfx1151 CU from 1–32. Invalid values use 32. | | `GGML_CUDA_MLA_STREAM_F32_STAGE` | unset | EXPERIMENTAL: with streaming D512 indexed attention, convert aligned F16 pairs once while staging them in shared memory instead of repeating conversion for every head. | | `GGML_CUDA_MLA_SPLIT_KV` / `GGML_DS4_FA_SPLIT_KV` | 1 on gfx1151 indexed decode; unset elsewhere | BURN-IN: force the reusable split-KV MLA schedule. Set `GGML_CUDA_MLA_NO_SPLIT_KV=1` (or legacy `GGML_DS4_FA_NO_SPLIT_KV=1`) to disable it. | | `GGML_DS4_TOPK_BLOCK_RADIX` | 1 on gfx1151 | BURN-IN KILL SWITCH: =0 restores hipCUB full sort for DS4-shaped 512-row top-k selection. | diff --git a/server/test/test_deepseek4_mmid_grouped_cuda.cpp b/server/test/test_deepseek4_mmid_grouped_cuda.cpp index 940ddb9da..932d38c76 100644 --- a/server/test/test_deepseek4_mmid_grouped_cuda.cpp +++ b/server/test/test_deepseek4_mmid_grouped_cuda.cpp @@ -247,11 +247,13 @@ static int run_child(const char * mode, const char * output_path) { GGML_TYPE_Q5_K, GGML_TYPE_Q2_0_ROCMFP2, GGML_TYPE_Q3_0_ROCMFPX, GGML_TYPE_Q4_0_ROCMFP4_FAST, }; - const int widths[] = {2, 4, 8, 9, 16, 32, 48, 64}; int width_filter = 0; if (const char * raw = std::getenv("DFLASH_MMID_TEST_WIDTH")) { width_filter = std::max(0, std::atoi(raw)); } + const std::vector widths = width_filter > 0 + ? std::vector{width_filter} + : std::vector{2, 4, 8, 9, 16, 32, 48, 64}; bool ok = output.good(); for (ggml_type type : types) { if (type == GGML_TYPE_Q4_0_ROCMFP4_FAST && width_filter == 0) { From 2d09e6500b26c90eec7c1311442497a015db9b73 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 31 Aug 2026 03:06:39 +0200 Subject: [PATCH 8/8] perf(hip): accelerate streaming MLA softmax --- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 40 +++++++++++-- server/docs/DS4.md | 6 +- server/docs/ENVIRONMENT.md | 1 + server/tests/test_deepseek4_unit.cpp | 57 +++++++++++++++---- 4 files changed, 87 insertions(+), 17 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 6911da4c7..57d35d6bf 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -1565,7 +1565,8 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( // backend-generic (D512 MQA with K == V and direct indexed rows); model policy // remains in the graph/backend layer. template + int KEYS_PER_STAGE = 16, bool STAGE_F32 = false, + bool FAST_EXP = false> __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( float * dst, const float * q, @@ -1721,8 +1722,13 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( const float score = partial * scale + mask_value; const float next_max = fmaxf(row_max, score); const float old_scale = row_sum == 0.0f - ? 0.0f : expf(row_max - next_max); - const float value_scale = expf(score - next_max); + ? 0.0f + : (FAST_EXP + ? __expf(row_max - next_max) + : expf(row_max - next_max)); + const float value_scale = FAST_EXP + ? __expf(score - next_max) + : expf(score - next_max); row_sum = row_sum * old_scale + value_scale; row_max = next_max; #pragma unroll @@ -1740,8 +1746,12 @@ __global__ static void ds4_flash_attn_d512_streaming_topk_kernel( const float sink = sinks[head]; const float next_max = fmaxf(row_max, sink); const float old_scale = row_sum == 0.0f - ? 0.0f : expf(row_max - next_max); - row_sum = row_sum * old_scale + expf(sink - next_max); + ? 0.0f + : (FAST_EXP + ? __expf(row_max - next_max) + : expf(row_max - next_max)); + row_sum = row_sum * old_scale + + (FAST_EXP ? __expf(sink - next_max) : expf(sink - next_max)); #pragma unroll for (int i = 0; i < VALUES_PER_LANE; ++i) { accum[i] *= old_scale; @@ -2634,11 +2644,29 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( getenv("GGML_CUDA_MLA_STREAM_F32_STAGE"); const bool f32_stage = f32_stage_env && f32_stage_env[0] != '\0' && strcmp(f32_stage_env, "0") != 0; + const char * fast_exp_env = + getenv("GGML_CUDA_MLA_STREAM_FAST_EXP"); + const bool fast_exp = fast_exp_env && fast_exp_env[0] != '\0' && + strcmp(fast_exp_env, "0") != 0; constexpr int streaming_heads = 16; const dim3 streaming_grid( (unsigned) n_tokens, (unsigned) (n_heads / streaming_heads), 1); - if (f32_stage) { + if (f32_stage && fast_exp) { + ds4_flash_attn_d512_streaming_topk_kernel< + half, half, streaming_heads, 16, true, true> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, + visibility_bounds, indexed_rows, indexed_counts, + indexed_capacity, inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients); + } else if (f32_stage) { ds4_flash_attn_d512_streaming_topk_kernel< half, half, streaming_heads, 16, true> <<>>( diff --git a/server/docs/DS4.md b/server/docs/DS4.md index d62dac430..0b01a3ca0 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -498,7 +498,11 @@ throughput A/B. `GGML_DS4_FA_STREAM_TOPK` remains a compatibility alias. Add `GGML_CUDA_MLA_STREAM_F32_STAGE=1` to convert each selected F16 latent once while staging aligned pairs in shared memory instead of repeating conversion for every head. The isolated gfx1151 qualification is byte-identical to F16 -staging and reduced alternating-run kernel time by about 9%. +staging and reduced alternating-run kernel time by about 9%. Add +`GGML_CUDA_MLA_STREAM_FAST_EXP=1` to use the HIP hardware exponential in that +FP32-staged online softmax. It reduced the remaining kernel time by another +7–9% in isolation; keep it opt-in with the streaming path because it uses an +approximate hardware exponential instead of the default implementation. Indexed verifier attention with at most eight query rows also uses a two-way split-KV schedule on gfx1151. Set `GGML_CUDA_MLA_NO_SPLIT_KV=1` to restore the diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index b304f9532..cd4f3cc73 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -54,6 +54,7 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `GGML_CUDA_MMQ_MOE_PERSISTENT` | unset | EXPERIMENTAL: on prefill-sized (at least 256-token) sparse grouped ROCmFP2/3/4 MMQ on gfx1151, build a compact device-side expert-tile queue and consume it with bounded persistent workers. Short batches, ordinary matmuls, unmeasured formats, and other devices are unchanged. | | `GGML_CUDA_MMQ_MOE_PERSISTENT_BLOCKS_PER_CU` | 32 | DEBUG: set the compact grouped-MoE worker budget per gfx1151 CU from 1–32. Invalid values use 32. | | `GGML_CUDA_MLA_STREAM_F32_STAGE` | unset | EXPERIMENTAL: with streaming D512 indexed attention, convert aligned F16 pairs once while staging them in shared memory instead of repeating conversion for every head. | +| `GGML_CUDA_MLA_STREAM_FAST_EXP` | unset | EXPERIMENTAL: with FP32-staged streaming D512 indexed attention, use the HIP hardware exponential intrinsic for online softmax. Other attention paths are unchanged. | | `GGML_CUDA_MLA_SPLIT_KV` / `GGML_DS4_FA_SPLIT_KV` | 1 on gfx1151 indexed decode; unset elsewhere | BURN-IN: force the reusable split-KV MLA schedule. Set `GGML_CUDA_MLA_NO_SPLIT_KV=1` (or legacy `GGML_DS4_FA_NO_SPLIT_KV=1`) to disable it. | | `GGML_DS4_TOPK_BLOCK_RADIX` | 1 on gfx1151 | BURN-IN KILL SWITCH: =0 restores hipCUB full sort for DS4-shaped 512-row top-k selection. | | `GGML_DS4_FA_SERIAL_INDEX_SCAN` | unset | DEBUG/A-B: restore the serial indexed-attention mask scan instead of the long-context HIP parallel scan. | diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 1b3126a58..1b3676525 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -2936,11 +2936,11 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { std::vector topk_data( (size_t) selected_rows * n_tokens); for (size_t i = 0; i < q_data.size(); ++i) { - q_data[i] = ((int) (i % 43) - 21) * 0.001f; + q_data[i] = ((int) (i % 43) - 21) * 0.05f; } for (size_t i = 0; i < kv_data.size(); ++i) { kv_data[i] = ggml_fp32_to_fp16( - ((int) (i % 47) - 23) * 0.001f); + ((int) (i % 47) - 23) * 0.05f); } for (int token = 0; token < n_tokens; ++token) { ggml_fp16_t * token_mask = @@ -2967,6 +2967,7 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { ScopedEnvVar streaming_guard("GGML_CUDA_MLA_STREAM_TOPK"); ScopedEnvVar f32_stage_guard("GGML_CUDA_MLA_STREAM_F32_STAGE"); + ScopedEnvVar fast_exp_guard("GGML_CUDA_MLA_STREAM_FAST_EXP"); ScopedCudaGraphOverrides eager( /*disable_graphs=*/true, /*mmvq_max_ncols=*/0, @@ -2974,7 +2975,9 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { std::vector reference((size_t) ggml_nelements(output)); std::vector candidate_f16_stage(reference.size()); std::vector candidate_f32_stage(reference.size()); + std::vector candidate_fast_exp(reference.size()); setenv("GGML_CUDA_MLA_STREAM_TOPK", "0", 1); + setenv("GGML_CUDA_MLA_STREAM_FAST_EXP", "0", 1); TEST_ASSERT_MSG( ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, "grouped compact attention reference failed"); @@ -2996,6 +2999,16 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { ggml_backend_tensor_get(output, candidate_f32_stage.data(), 0, candidate_f32_stage.size() * sizeof(float)); + setenv("GGML_CUDA_MLA_STREAM_FAST_EXP", "1", 1); + TEST_ASSERT_MSG( + ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, + "fast-exp streaming top-k attention failed"); + ggml_backend_tensor_get(output, candidate_fast_exp.data(), 0, + candidate_fast_exp.size() * sizeof(float)); + + double fast_exp_squared_error = 0.0; + double f32_stage_power = 0.0; + float fast_exp_max_abs = 0.0f; for (size_t i = 0; i < reference.size(); ++i) { TEST_ASSERT_MSG( std::isfinite(candidate_f32_stage[i]), @@ -3007,11 +3020,29 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { TEST_ASSERT_MSG( candidate_f16_stage[i] == candidate_f32_stage[i], "F32 staging changed streaming top-k attention output"); + TEST_ASSERT_MSG( + std::isfinite(candidate_fast_exp[i]), + "fast-exp streaming output must be finite"); + TEST_ASSERT_MSG( + nearly_equal(reference[i], candidate_fast_exp[i], + 5.0e-4f, 5.0e-4f), + "fast-exp streaming attention exceeded numeric tolerance"); + const double fast_exp_error = + (double) candidate_fast_exp[i] - candidate_f32_stage[i]; + fast_exp_squared_error += fast_exp_error * fast_exp_error; + f32_stage_power += + (double) candidate_f32_stage[i] * candidate_f32_stage[i]; + fast_exp_max_abs = std::max( + fast_exp_max_abs, + std::fabs(candidate_fast_exp[i] - candidate_f32_stage[i])); } + const double fast_exp_nmse = fast_exp_squared_error / + std::max(f32_stage_power, 1.0e-30); - auto measure_us = [&](bool streaming, bool f32_stage) { + auto measure_us = [&](bool streaming, bool f32_stage, bool fast_exp) { setenv("GGML_CUDA_MLA_STREAM_TOPK", streaming ? "1" : "0", 1); setenv("GGML_CUDA_MLA_STREAM_F32_STAGE", f32_stage ? "1" : "0", 1); + setenv("GGML_CUDA_MLA_STREAM_FAST_EXP", fast_exp ? "1" : "0", 1); constexpr int warmups = 3; constexpr int iterations = 20; for (int i = 0; i < warmups; ++i) { @@ -3027,26 +3058,32 @@ static void test_ds4_flash_attention_streaming_topk_gpu() { return std::chrono::duration(end - begin).count() / iterations; }; - const double grouped_us = measure_us(false, false); + const double grouped_us = measure_us(false, false, false); constexpr int timing_rounds = 4; double streaming_f16_us = 0.0; double streaming_f32_us = 0.0; + double fast_exp_us = 0.0; for (int round = 0; round < timing_rounds; ++round) { if ((round & 1) == 0) { - streaming_f16_us += measure_us(true, false); - streaming_f32_us += measure_us(true, true); + streaming_f16_us += measure_us(true, false, false); + streaming_f32_us += measure_us(true, true, false); + fast_exp_us += measure_us(true, true, true); } else { - streaming_f32_us += measure_us(true, true); - streaming_f16_us += measure_us(true, false); + fast_exp_us += measure_us(true, true, true); + streaming_f32_us += measure_us(true, true, false); + streaming_f16_us += measure_us(true, false, false); } } streaming_f16_us /= timing_rounds; streaming_f32_us /= timing_rounds; + fast_exp_us /= timing_rounds; std::fprintf(stderr, " grouped=%.1fus streaming_f16=%.1fus" - " streaming_f32=%.1fus speedup=%.2fx", + " streaming_f32=%.1fus fast_exp=%.1fus speedup=%.2fx" + " fast_nmse=%.3g fast_max_abs=%.3g", grouped_us, streaming_f16_us, streaming_f32_us, - grouped_us / streaming_f32_us); + fast_exp_us, grouped_us / fast_exp_us, + fast_exp_nmse, fast_exp_max_abs); } ggml_gallocr_free(alloc);