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..57d35d6bf 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) { @@ -1554,6 +1556,236 @@ __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; + + 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]; + + 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(); + + // 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(); + +#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 + : (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 + 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 + : (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; + } + } + + 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 +2615,88 @@ 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) { + 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; + 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 && 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> + <<>>( + (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; + } // 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/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..ef05d61d0 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; @@ -1060,17 +1061,53 @@ 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 { + // 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) 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) } @@ -1101,8 +1138,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]; @@ -4620,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( @@ -4817,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; @@ -4915,7 +5129,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/DS4.md b/server/docs/DS4.md index 2db93acb4..0b01a3ca0 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -487,6 +487,23 @@ 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. 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%. 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 single-block schedule. For fused-verifier masks of at least 4 MiB, the runtime diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index fba8c7d40..cd4f3cc73 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -49,6 +49,12 @@ 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_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. | @@ -202,6 +208,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/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/test/test_deepseek4_mmid_grouped_cuda.cpp b/server/test/test_deepseek4_mmid_grouped_cuda.cpp index 58f062e71..932d38c76 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,30 @@ 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}; + 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) { + 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 +376,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 +616,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 +630,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 +688,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 +698,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 +732,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 && diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 94e2a7ac2..1b3676525 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 ..."); @@ -2856,6 +2873,225 @@ 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.05f; + } + for (size_t i = 0; i < kv_data.size(); ++i) { + kv_data[i] = ggml_fp32_to_fp16( + ((int) (i % 47) - 23) * 0.05f); + } + 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"); + 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, + /*skip_property_check=*/false); + 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"); + ggml_backend_tensor_get(output, reference.data(), 0, + 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, + "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)); + + 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]), + "streaming top-k attention output must be finite"); + TEST_ASSERT_MSG( + 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"); + 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, 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) { + 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, 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, false); + streaming_f32_us += measure_us(true, true, false); + fast_exp_us += measure_us(true, true, true); + } else { + 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 fast_exp=%.1fus speedup=%.2fx" + " fast_nmse=%.3g fast_max_abs=%.3g", + grouped_us, streaming_f16_us, streaming_f32_us, + fast_exp_us, grouped_us / fast_exp_us, + fast_exp_nmse, fast_exp_max_abs); + } + + 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; @@ -4174,6 +4410,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); @@ -4191,6 +4428,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();