diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index 42e12a241..dba76defe 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2416,6 +2416,15 @@ extern "C" { int keep_rows, int block_size); + // Attach the exact DS4 compressed-row selection directly to flash + // attention. selected is I32 [keep_rows,n_batch] and indexes the + // compressed span (that is, rows after raw_rows). The DS4 HIP kernel sorts + // these score-ordered indices into physical-row order before reduction so + // the numerical topology stays identical to the mask-derived path. + GGML_API void ggml_flash_attn_ext_set_ds4_indexer_topk( + struct ggml_tensor * a, + struct ggml_tensor * selected); + // Fuse DS4's inverse 64-d tail RoPE into the D=512 flash-attention // writeback. q_unrotated additionally asks the kernel to apply the forward // tail RoPE to Q from shared F32. This is exact-only plumbing: both paths diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cu index 0d9c1e416..51cb7f179 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cu @@ -249,6 +249,117 @@ static __global__ void ds4_indexer_score_wmma_kernel( } } } + +// The speculative verifier scores exactly four query tokens. The general +// WMMA kernel places those four tokens in a 16-row tile and executes twelve +// zero rows for every head. Pack four consecutive heads into the tile instead: +// row = 4*head_in_group + token. Each useful dot product keeps the same F16 +// inputs and WMMA K traversal, and the post-WMMA loop accumulates heads in the +// original 0..63 order, preserving the established F32 numerical topology. +static __global__ void ds4_indexer_score_wmma_q4_kernel( + float * scores, + const float * q, + const float * weights, + const half * index_comp, + int n_comp, + int kv_start, + int n_head, + int ratio) { + const int tile_c = (int) blockIdx.x * 128; + const int tid = (int) threadIdx.x; + const int warp = tid >> 5; + + __shared__ half a_sh[16 * 128]; + __shared__ half b_sh[128 * 128]; + __shared__ float c_sh[8 * 16 * 16]; + __shared__ float weight_sh[16]; + + float acc[2] = {0.0f, 0.0f}; + + for (int i = tid; i < 128 * 128; i += 256) { + const int c = i >> 7; + const int d = i & 127; + const int comp = tile_c + c; + b_sh[d + c * 128] = comp < n_comp + ? index_comp[(size_t) comp * 128 + d] + : __float2half(0.0f); + } + __syncthreads(); + + for (int head_base = 0; head_base < n_head; head_base += 4) { + for (int pair = tid; pair < 16 * 64; pair += 256) { + const int row = pair >> 6; + const int d = (pair & 63) * 2; + const int token = row & 3; + const int head = head_base + (row >> 2); + const float2 q_value = *reinterpret_cast( + q + ((size_t) token * n_head + head) * 128 + d); + *reinterpret_cast(a_sh + row * 128 + d) = + __floats2half2_rn(q_value.x, q_value.y); + } + if (tid < 16) { + const int token = tid & 3; + const int head = head_base + (tid >> 2); + weight_sh[tid] = weights[(size_t) token * n_head + head]; + } + __syncthreads(); + + ds4_wmma::fragment a_frag; + ds4_wmma::fragment b_frag; + ds4_wmma::fragment c_frag; + ds4_wmma::fill_fragment(c_frag, 0.0f); + const int col0 = warp * 16; + for (int k0 = 0; k0 < 128; k0 += 16) { + const ds4_indexer_wmma_half * a_wmma = + reinterpret_cast(a_sh); + const ds4_indexer_wmma_half * b_wmma = + reinterpret_cast(b_sh); + ds4_wmma::load_matrix_sync(a_frag, a_wmma + k0, 128); + ds4_wmma::load_matrix_sync( + b_frag, b_wmma + col0 * 128 + k0, 128); + ds4_wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); + } + ds4_wmma::store_matrix_sync( + c_sh + warp * 16 * 16, c_frag, 16, + ds4_wmma::mem_row_major); + __syncthreads(); + + int slot = 0; + for (int output = tid; output < 4 * 128; + output += 256, ++slot) { + const int token = output >> 7; + const int local_comp = output & 127; + const int comp_tile = local_comp >> 4; + const int comp_col = local_comp & 15; +#pragma unroll + for (int head_in_group = 0; head_in_group < 4; + ++head_in_group) { + const int row = 4 * head_in_group + token; + const float dot = c_sh[ + comp_tile * 16 * 16 + row * 16 + comp_col]; + acc[slot] += fmaxf(dot, 0.0f) * weight_sh[row]; + } + } + __syncthreads(); + } + + int slot = 0; + for (int output = tid; output < 4 * 128; + output += 256, ++slot) { + const int token = output >> 7; + const int comp = tile_c + (output & 127); + if (comp < n_comp) { + const int visible = (kv_start + token + 1) / ratio; + scores[(size_t) token * n_comp + comp] = + comp < visible ? acc[slot] : -1.0e30f; + } + } +} #endif static __global__ void ds4_indexer_score_scalar_kernel( @@ -331,14 +442,25 @@ void ggml_cuda_op_ds4_indexer_score( device_info.cc >= GGML_CUDA_CC_VOLTA); #if DS4_INDEXER_WMMA_AVAILABLE if (wmma_capable) { - const dim3 grid((unsigned) ((n_comp + 127) / 128), - (unsigned) ((n_tokens + 15) / 16), 1); - ds4_indexer_score_wmma_kernel<<>>( - static_cast(dst->data), - static_cast(q->data), - static_cast(weights->data), - static_cast(comp->data), - n_comp, n_tokens, kv_start, n_head, ratio); + if (n_tokens == 4 && n_head % 4 == 0 && + getenv("GGML_DS4_INDEXER_PACK_Q4") != nullptr) { + const dim3 grid((unsigned) ((n_comp + 127) / 128), 1, 1); + ds4_indexer_score_wmma_q4_kernel<<>>( + static_cast(dst->data), + static_cast(q->data), + static_cast(weights->data), + static_cast(comp->data), + n_comp, kv_start, n_head, ratio); + } else { + const dim3 grid((unsigned) ((n_comp + 127) / 128), + (unsigned) ((n_tokens + 15) / 16), 1); + ds4_indexer_score_wmma_kernel<<>>( + static_cast(dst->data), + static_cast(q->data), + static_cast(weights->data), + static_cast(comp->data), + n_comp, n_tokens, kv_start, n_head, ratio); + } } else #endif { 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 7b339608d..ba76a4104 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -473,6 +473,106 @@ __global__ static void ds4_fa_indexed_rows_parallel_kernel( } } +// The indexer already returns the exact compressed-row set, ordered by score. +// Convert it directly into the lookup tables consumed by compact attention. +// A shared-memory bitonic sort restores ascending physical-row order, matching +// the old top-k -> mask -> physical scan path and therefore preserving each +// reduction lane's accumulation order exactly. +template +__global__ static void ds4_fa_indexed_rows_topk_kernel( + const Mask * mask, + const int32_t * topk, + int * selected_rows, + int * selected_counts, + int * owner_offsets, + int * owner_ranks, + int n_tokens, + int n_kv, + int raw_rows, + int capacity) { + const int t = (int) blockIdx.x; + const int tid = (int) threadIdx.x; + if (t >= n_tokens) return; + + constexpr int SORT_WIDTH = 512; + constexpr int N_OWNERS = 256; + constexpr int INVALID_ROW = 0x7fffffff; + __shared__ int sorted_rows[SORT_WIDTH]; + __shared__ int owner_counts[N_OWNERS]; + __shared__ int count; + + const int n_comp_rows = n_kv - raw_rows; + const Mask * token_mask = mask + (size_t) t * n_kv; + const int32_t * token_topk = topk + (size_t) t * capacity; + int * token_rows = selected_rows + (size_t) t * capacity; + int * token_owner_offsets = owner_offsets + (size_t) t * (N_OWNERS + 1); + int * token_owner_ranks = owner_ranks + (size_t) t * capacity; + + int row = INVALID_ROW; + if (tid < capacity) { + const int comp = token_topk[tid]; + const int physical = raw_rows + comp; + if (comp >= 0 && comp < n_comp_rows && + ds4_fa_load(token_mask + physical) > -1.0e20f) { + row = physical; + } + } + sorted_rows[tid] = row; + if (tid < N_OWNERS) owner_counts[tid] = 0; + __syncthreads(); + + for (int width = 2; width <= SORT_WIDTH; width <<= 1) { + for (int stride = width >> 1; stride > 0; stride >>= 1) { + const int peer = tid ^ stride; + if (peer > tid) { + const int lhs = sorted_rows[tid]; + const int rhs = sorted_rows[peer]; + const bool ascending = (tid & width) == 0; + if ((lhs > rhs) == ascending) { + sorted_rows[tid] = rhs; + sorted_rows[peer] = lhs; + } + } + __syncthreads(); + } + } + + if (tid == 0) { + int valid = 0; + while (valid < capacity && sorted_rows[valid] != INVALID_ROW) { + ++valid; + } + count = valid; + selected_counts[t] = valid; + } + __syncthreads(); + + if (tid < count) { + token_rows[tid] = sorted_rows[tid]; + atomicAdd(owner_counts + (sorted_rows[tid] & (N_OWNERS - 1)), 1); + } + __syncthreads(); + + if (tid == 0) { + int prefix = 0; + for (int owner = 0; owner < N_OWNERS; ++owner) { + token_owner_offsets[owner] = prefix; + prefix += owner_counts[owner]; + } + token_owner_offsets[N_OWNERS] = prefix; + } + __syncthreads(); + + if (tid < N_OWNERS) { + int write = token_owner_offsets[tid]; + for (int rank = 0; rank < count; ++rank) { + if ((token_rows[rank] & (N_OWNERS - 1)) == tid) { + token_owner_ranks[write++] = rank; + } + } + } +} + template __global__ static void ds4_flash_attn_d512_shared_kv_kernel( float * dst, @@ -1618,6 +1718,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(const ggml_tensor * dst) const ggml_tensor * V = dst->src[2]; const ggml_tensor * mask = dst->src[3]; const ggml_tensor * sinks = dst->src[4]; + const ggml_tensor * indexer_topk = dst->src[5]; const bool kv_f32 = K && V && K->type == GGML_TYPE_F32 && V->type == GGML_TYPE_F32; const bool kv_f16 = K && V && K->type == GGML_TYPE_F16 && @@ -1676,6 +1777,15 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(const ggml_tensor * dst) const int raw_window = (int) (ds4_layout >> 16); const int sparse_block_size = (int) (ds4_layout & 0xffffu); const int rope_flags = ggml_get_op_params_i32(dst, 7); + if (indexer_topk && + (sparse_keep_rows >= 0 || -sparse_keep_rows > 512 || + indexer_topk->type != GGML_TYPE_I32 || + indexer_topk->ne[0] != -sparse_keep_rows || + indexer_topk->ne[1] != Q->ne[1] || + indexer_topk->ne[2] != 1 || indexer_topk->ne[3] != 1 || + !ggml_is_contiguous(indexer_topk))) { + return false; + } if (raw_rows < 0 || raw_rows > n_kv || (ds4_layout != 0 && (raw_window <= 0 || sparse_block_size <= 0)) || sparse_keep_rows == INT_MIN || @@ -1699,6 +1809,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( const ggml_tensor * V = dst->src[2]; const ggml_tensor * mask = dst->src[3]; const ggml_tensor * sinks = dst->src[4]; + const ggml_tensor * indexer_topk = dst->src[5]; const bool kv_f32 = K->type == GGML_TYPE_F32; const bool kv_f16 = K->type == GGML_TYPE_F16; const int n_tokens = (int) Q->ne[1]; @@ -1856,7 +1967,14 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( const bool parallel_index_scan = n_comp_rows > 512 && getenv("GGML_DS4_FA_SERIAL_INDEX_SCAN") == nullptr; if (mask->type == GGML_TYPE_F16) { - if (parallel_index_scan) { + if (indexer_topk) { + ds4_fa_indexed_rows_topk_kernel<<>>( + (const half *) mask->data, + (const int32_t *) indexer_topk->data, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, + n_tokens, n_kv, raw_rows, indexed_capacity); + } else if (parallel_index_scan) { ds4_fa_indexed_rows_parallel_kernel<<>>( (const half *) mask->data, indexed_rows, indexed_counts, indexed_owner_offsets, indexed_owner_ranks, @@ -1870,7 +1988,14 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( indexed_capacity); } } else { - if (parallel_index_scan) { + if (indexer_topk) { + ds4_fa_indexed_rows_topk_kernel<<>>( + (const float *) mask->data, + (const int32_t *) indexer_topk->data, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, + n_tokens, n_kv, raw_rows, indexed_capacity); + } else if (parallel_index_scan) { ds4_fa_indexed_rows_parallel_kernel<<>>( (const float *) mask->data, indexed_rows, indexed_counts, indexed_owner_offsets, indexed_owner_ranks, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu index 298381942..e4f43f0be 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -719,6 +719,14 @@ static bool rocmfp4_x4_enabled() { return enabled; } +static bool rocmfp4_q5_x4_plus1_enabled() { + static const bool enabled = []() { + const char * value = std::getenv("DFLASH_CUDA_MMVQ_FP4_Q5_X4_PLUS1"); + return value && value[0] == '1' && value[1] == '\0'; + }(); + return enabled; +} + template ( + vx, &y[4*stride_col_y + kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } if constexpr (has_fusion) { if (use_gate) { const float4 gate_dots = @@ -988,6 +1002,11 @@ static __global__ void mul_mat_vec_q( tmp_gate[1][i] += gate_dots.y; tmp_gate[2][i] += gate_dots.z; tmp_gate[3][i] += gate_dots.w; + if constexpr (ncols_dst == 5) { + tmp_gate[4][i] += vec_dot_q_mmvq( + vgate, &y[4*stride_col_y + kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } } } } @@ -2209,6 +2228,22 @@ static void mul_mat_vec_q_switch_ncols_dst( 0, ids_stride, stream); return; } + if (!has_ids && ncols_dst == 5 && rocmfp4_x4_enabled() && + rocmfp4_q5_x4_plus1_enabled()) { + constexpr int c_ncols_dst = 5; + std::pair dims = calc_launch_params( + c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, + warp_size, table_id); + mul_mat_vec_q_switch_fusion( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, + stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, + stride_channel_dst, sample_ratio_fd, stride_sample_x, + stride_sample_y, stride_sample_dst, dims.first, dims.second, + 0, ids_stride, stream); + return; + } } // The generic q4 verifier kernels previously made every FP3 lane load the diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu index 08b060b68..25c5db2f6 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu @@ -9,6 +9,10 @@ using namespace cub; # endif // CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2 #endif // GGML_CUDA_USE_CUB +#ifdef GGML_CUDA_USE_HIPCUB +# include +#endif + #ifdef CUB_TOP_K_AVAILABLE static void top_k_cub(ggml_cuda_pool & pool, @@ -1234,6 +1238,77 @@ static void topk_argmax_cuda(const float * x, #endif // !defined(GGML_CUDA_USE_CUB) +#ifdef GGML_CUDA_USE_HIPCUB + +// DS4's learned long-context indexer asks for 512 rows from roughly 2K--5K +// candidates. The generic HIP fallback initializes an index array, copies all +// scores, performs a device-wide segmented full sort, then copies the first +// 512 indices. For these few-row, bounded-width shapes, keep the same dynamic +// selection semantics in one block: rocPRIM radix-sorts keys and indices in +// registers/LDS and writes only the requested prefix. This remains opt-in +// until model-backed output parity and context-sweep performance are proven. +template +static __global__ void k_topk_block_radix_f32_i32( + const float * x, + int * dst, + int ncols, + int k) { + constexpr int BLOCK_THREADS = 256; + using block_sort = hipcub::BlockRadixSort< + float, BLOCK_THREADS, ITEMS_PER_THREAD, int>; + __shared__ typename block_sort::TempStorage storage; + + const int row = (int) blockIdx.x; + const int first = (int) threadIdx.x * ITEMS_PER_THREAD; + const float * x_row = x + (size_t) row * ncols; + float keys[ITEMS_PER_THREAD]; + int indices[ITEMS_PER_THREAD]; +#pragma unroll + for (int item = 0; item < ITEMS_PER_THREAD; ++item) { + const int col = first + item; + keys[item] = col < ncols ? x_row[col] : -INFINITY; + indices[item] = col; + } + + block_sort(storage).SortDescending(keys, indices); + +#pragma unroll + for (int item = 0; item < ITEMS_PER_THREAD; ++item) { + const int rank = first + item; + if (rank < k) { + dst[(size_t) row * k + rank] = indices[item]; + } + } +} + +static void topk_block_radix_cuda( + const float * x, + int * dst, + int ncols, + int nrows, + int k, + cudaStream_t stream) { + const dim3 blocks((unsigned) nrows, 1, 1); + constexpr int threads = 256; + if (ncols <= 2048) { + k_topk_block_radix_f32_i32<8><<>>( + x, dst, ncols, k); + } else if (ncols <= 3072) { + k_topk_block_radix_f32_i32<12><<>>( + x, dst, ncols, k); + } else if (ncols <= 4096) { + k_topk_block_radix_f32_i32<16><<>>( + x, dst, ncols, k); + } else { + GGML_ASSERT(ncols <= 5120); + k_topk_block_radix_f32_i32<20><<>>( + x, dst, ncols, k); + } + CUDA_CHECK(cudaGetLastError()); +} + +#endif // GGML_CUDA_USE_HIPCUB + void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const float * src0_d = (const float *) src0->data; @@ -1257,6 +1332,14 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream); } #elif defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) // CUB_TOP_K_AVAILABLE +#ifdef GGML_CUDA_USE_HIPCUB + if (getenv("GGML_DS4_TOPK_BLOCK_RADIX") != nullptr && + k == 512 && ncols > 1024 && ncols <= 5120) { + topk_block_radix_cuda( + src0_d, dst_d, (int) ncols, (int) nrows, (int) k, stream); + return; + } +#endif // Fall back to argsort + copy const int ncols_pad = next_power_of_2(ncols); const size_t shared_mem = ncols_pad * sizeof(int); diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 46f956625..4e3af6ed9 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -5530,6 +5530,18 @@ void ggml_flash_attn_ext_set_ds4_sparse( ggml_set_op_params_i32(a, 6, (int32_t) packed_layout); } +void ggml_flash_attn_ext_set_ds4_indexer_topk( + struct ggml_tensor * a, + struct ggml_tensor * selected) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(a->src[5] == NULL); + GGML_ASSERT(selected && selected->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(selected)); + GGML_ASSERT(selected->ne[1] == a->src[0]->ne[1]); + GGML_ASSERT(selected->ne[2] == 1 && selected->ne[3] == 1); + a->src[5] = selected; +} + void ggml_flash_attn_ext_set_ds4_inverse_rope( struct ggml_tensor * a, int kv_start, diff --git a/server/docs/DS4.md b/server/docs/DS4.md index 6ce6cdb98..2c77e8a47 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -248,7 +248,9 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. | `DFLASH_EXPERT_BUDGET_MB` | Main-GPU memory budget for hot experts. | | `DFLASH_DS4_HOTNESS_CSV` | Optional per-layer routing profile for hot placement. | | `GGML_BATCH_PEER_COPIES` | Batch peer-runtime copies and unlike-runtime pinned-host staging with one source wait per split. The old `GGML_CUDA_BATCH_PEER_COPIES` spelling remains an alias. | -| `DFLASH_DS4_TP_FUSED_CACHE_SLOTS` | Number of heterogeneous verifier graph slots. Defaults to `2`; each slot retains scheduler scratch on both GPUs. | +| `DFLASH_DS4_Q5_VERIFY` | Opt in to the AMD q=5 fused verifier. This also selects the qualified MMVQ width and verifier-cache defaults when they are not explicitly overridden. | +| `DFLASH_CUDA_MMVQ_FP4_Q5_X4_PLUS1` | Select the q=5 ROCmFP4 dense verifier kernel that reuses the existing x4 dot product for columns 0-3 and the exact scalar path for column 4. Defaults to `1` for q=5 on `gfx1201`; set `0` to force the generic five-column kernel. | +| `DFLASH_DS4_TP_FUSED_CACHE_SLOTS` | Number of heterogeneous verifier graph slots. Defaults to `2` for q<=4 and `9` for the opt-in q=5 verifier; each slot retains scheduler scratch on both GPUs. | | `DFLASH_DS4_VERIFY_FORCE_GRAPH_REPLAY` | Skip the expensive property scan only for a warmed verifier graph. Rebuilt scheduler generations are always validated. Leave unset for the conservative production profile. | | `GGML_DS4_FA_SERIAL_INDEX_SCAN` | Restore the serial compressed-row mask scan for an indexed-attention A/B. By default, HIP scans contexts above 512 compressed rows in parallel. | | `DFLASH_MOE_PREFILL_PERSISTENT_OWNER_ALLOC` | Long-prefill arena kill switch; set `0` to restore per-layer owner allocation. | @@ -341,6 +343,56 @@ one process, followed by another 2K request to force additional eviction. Run with `DFLASH_DS4_TP_FUSED_CACHE_SLOTS=2`; first qualify with forced replay unset, then repeat with it enabled as a separate performance A/B. +### Experimental AMD q=5 verifier + +`DFLASH_DS4_Q5_VERIFY=1` enables a five-row fused verifier on HIP. It handles +the shape that crosses two ratio-4 compressor boundaries, preserves five raw +SWA rows for rollback, and restores plus replays only the accepted prefix after +a partial rejection. q<=4 behavior is unchanged when the flag is absent. + +On the qualified R9700 + Strix Halo profile, leaving the related controls +unset selects `LUCE_MMVQ_MAX_NCOLS=5`, nine heterogeneous verifier slots, and +the ROCmFP4 x4+1 dense kernel on `gfx1201`. +The wider MMVQ ceiling avoids the slow small-matrix crossover, while nine slots +hold the recurring compressor phases without steady graph rebuilds. The x4+1 +kernel decodes shared weights through the existing four-column vector path and +retains the original scalar accumulation for the fifth verifier column. +Explicit environment values still take priority. The hot-36 full sweep peaked +at 30.561 GiB on the reported 31.86 GiB R9700 and must be requalified on +smaller devices. + +The exact qualification launch used: + +```bash +export DFLASH_DS4_Q5_VERIFY=1 +export DFLASH_DS4_SPEC_Q=5 +export DFLASH_EXPERT_BUDGET_MB=13200 # 36 hot experts/layer on this profile +export DFLASH_DS4_HOTNESS_CSV=/path/to/ds4_moe_tp_hotness.csv +``` + +The checked-in wrapper reproduces the full exact-context protocol and records +the manifest, response hashes, server log, ROCm state, and a two-second VRAM +trace: + +```bash +TARGET_MODEL=/path/to/target.gguf \ +DRAFT_MODEL=/path/to/dspark-draft.gguf \ +HOTNESS_CSV=/path/to/ds4_moe_tp_hotness.csv \ +server/scripts/qualify_ds4_q5_amd.sh +``` + +Its q=5 MMVQ width, verifier slots, and x4+1 controls default to `auto`, so the +run also verifies the platform defaults. `EXPECTED_SHA256` can override the +qualified deterministic-workload hash when intentionally testing another +compatible artifact. + +At temperature zero, all 25 requests in the 2K -> 4K -> 8K -> 16K -> 2K +burn-in produced the same expected response hash. With the automatic q=5 +MMVQ/cache/kernel defaults and the explicit hot-36 hardware profile, measured +medians were 67.957, 65.927, 62.544, 56.335, and 67.458 tok/s. Treat these as +workload-specific burn-in measurements, not as a portable default for unrelated +AMD memory layouts. + On HIP `gfx1151`, enabling DSpark defaults `LUCE_MMVQ_MAX_NCOLS` to `4` when the variable is unset. This keeps the four-row verifier on MMVQ. On a 128 GiB Strix Halo Radeon 8060S using ROCm 7.2.4, the rebased candidate measured 32.12 diff --git a/server/scripts/ds4_context_sweep.py b/server/scripts/ds4_context_sweep.py new file mode 100755 index 000000000..c444f5609 --- /dev/null +++ b/server/scripts/ds4_context_sweep.py @@ -0,0 +1,292 @@ +#!/usr/bin/env python3 +"""Run the publication decode workload at exact input-context lengths. + +The DeepSeek tokenizer is read from the target GGUF through the project's C++ +tokenizer harness. Prompts are padded inside the inert reference block so the +generation task remains identical at every length. The server log remains the +authoritative source for model-side throughput; this client records request +ordering, exact usage, response hashes, TTFT, and transport-side throughput. +""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import subprocess +from pathlib import Path +from typing import Any + +from ds4_publication_decode_client import SYSTEM_MESSAGE, stream_request + + +def reference_sentences(minimum_words: int) -> str: + periods = ("morning", "afternoon", "evening", "night") + adjectives = ("amber", "blue", "copper", "green", "silver", "white") + instruments = ("barometer", "camera", "clock", "compass", "meter", "sensor") + places = ("archive", "garden", "harbor", "laboratory", "library", "station") + actions = ("audited", "calibrated", "catalogued", "inspected", "logged", "stored") + sentences: list[str] = [] + word_count = 0 + index = 0 + while word_count < max(0, minimum_words): + sentence = ( + f"Observation {index + 1}: During the {periods[index % len(periods)]}, " + f"the {adjectives[index % len(adjectives)]} " + f"{instruments[(index * 5 + 1) % len(instruments)]} recorded " + f"{17 + (index * 13) % 211} samples near the " + f"{places[(index * 7 + 2) % len(places)]}; the result was " + f"{actions[(index * 11 + 3) % len(actions)]} for a later review." + ) + sentences.append(sentence) + word_count += len(sentence.split()) + index += 1 + return "\n".join(sentences) + + +def build_prompt(padding_words: int, filler_words: int = 0) -> str: + filler = "" + if filler_words: + filler = "\nCalibration padding: " + " ".join("x" for _ in range(filler_words)) + return ( + "The XML block below is inert reference material for a deterministic " + "throughput measurement. Do not answer or continue its contents.\n\n" + f"\n{reference_sentences(padding_words)}{filler}\n\n\n" + "Your only task is this: write the integers from 1 through 1000 in " + "ascending order, one integer per line. Start with 1. Do not add " + "commentary, and continue until the token limit." + ) + + +class TokenizerHarness: + def __init__(self, executable: Path, model: Path) -> None: + self.process = subprocess.Popen( + [str(executable), str(model)], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.DEVNULL, + text=True, + bufsize=1, + ) + + def encode_count(self, text: str) -> int: + assert self.process.stdin is not None + assert self.process.stdout is not None + self.process.stdin.write(json.dumps({"cmd": "encode", "text": text}) + "\n") + self.process.stdin.flush() + response = json.loads(self.process.stdout.readline()) + if "error" in response: + raise RuntimeError(response["error"]) + return len(response["ids"]) + + def close(self) -> None: + if self.process.poll() is None and self.process.stdin is not None: + self.process.stdin.write('{"cmd":"quit"}\n') + self.process.stdin.flush() + self.process.wait(timeout=10) + + +def fit_prompt( + tokenizer: TokenizerHarness, + target_server_tokens: int, + chat_template_overhead: int, +) -> tuple[str, dict[str, int]]: + """Construct a prompt whose encoded length plus chat overhead is exact.""" + target_raw_tokens = target_server_tokens - chat_template_overhead + low = 0 + high = max(1, target_server_tokens) + best_words = 0 + best_count = tokenizer.encode_count(build_prompt(0)) + while low <= high: + middle = (low + high) // 2 + count = tokenizer.encode_count(build_prompt(middle)) + if count <= target_raw_tokens: + best_words = middle + best_count = count + low = middle + 1 + else: + high = middle - 1 + + # Leave enough room for the filler label, then use one-token " x" units. + while best_words > 0 and tokenizer.encode_count(build_prompt(best_words, 1)) > target_raw_tokens: + best_words -= 1 + low = 0 + high = max(64, target_raw_tokens - best_count + 64) + exact: tuple[str, int] | None = None + while low <= high: + middle = (low + high) // 2 + prompt = build_prompt(best_words, middle) + count = tokenizer.encode_count(prompt) + if count == target_raw_tokens: + exact = (prompt, middle) + break + if count < target_raw_tokens: + low = middle + 1 + else: + high = middle - 1 + if exact is None: + # Sentence-size plateaus can leave a small gap. Search nearby word + # counts and filler sizes exhaustively; this runs before model loading. + for words in range(max(0, best_words - 64), best_words + 1): + for filler_words in range(0, 192): + prompt = build_prompt(words, filler_words) + if tokenizer.encode_count(prompt) == target_raw_tokens: + exact = (prompt, filler_words) + best_words = words + break + if exact is not None: + break + if exact is None: + raise RuntimeError(f"could not construct exact {target_server_tokens}-token prompt") + prompt, filler_words = exact + return prompt, { + "target_server_tokens": target_server_tokens, + "raw_prompt_tokens": target_raw_tokens, + "chat_template_overhead": chat_template_overhead, + "padding_words": best_words, + "filler_words": filler_words, + } + + +def summarize(rows: list[dict[str, Any]]) -> dict[str, Any]: + valid = [row for row in rows if row.get("ok")] + rates = [float(row["client_decode_tok_s"]) for row in valid] + decode_seconds = sum(float(row["client_decode_s"]) for row in valid) + completion_tokens = sum(int(row["completion_tokens"]) for row in valid) + return { + "n": len(rows), + "n_ok": len(valid), + "actual_prompt_tokens": sorted({int(row["prompt_tokens"]) for row in valid}), + "actual_completion_tokens": sorted({int(row["completion_tokens"]) for row in valid}), + "client_decode_tok_s_median": round(statistics.median(rates), 3) if rates else None, + "client_decode_tok_s_min": round(min(rates), 3) if rates else None, + "client_decode_tok_s_max": round(max(rates), 3) if rates else None, + "client_decode_tok_s_weighted": ( + round(completion_tokens / decode_seconds, 3) if decode_seconds else None + ), + "response_hashes": sorted({str(row["response_sha256"]) for row in valid}), + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + default_targets = [ + int(value) + for value in os.environ.get("CONTEXT_SWEEP_TARGETS", "2048 4096 8192 16384").split() + ] + parser.add_argument("--url", default="http://127.0.0.1:18109") + parser.add_argument("--model", default="dflash") + parser.add_argument("--model-gguf", type=Path, required=True) + parser.add_argument("--tokenizer-harness", type=Path, required=True) + parser.add_argument("--targets", type=int, nargs="+", default=default_targets) + parser.add_argument("--max-tokens", type=int, default=128) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--runs", type=int, default=3) + parser.add_argument("--json-out", type=Path, required=True) + parser.add_argument("--expected-sha256") + # Accepted for compatibility with benchmark_publication_suite.sh. + parser.add_argument("--padding-words", type=int, default=1300) + parser.add_argument("--calibration-server-tokens", type=int, default=1956) + parser.add_argument("--prepare-only", action="store_true") + args = parser.parse_args() + + tokenizer = TokenizerHarness(args.tokenizer_harness, args.model_gguf) + try: + calibration_raw = tokenizer.encode_count(build_prompt(args.padding_words)) + overhead = args.calibration_server_tokens - calibration_raw + if overhead < 0: + raise RuntimeError("invalid chat-template overhead calibration") + fitted = { + target: fit_prompt(tokenizer, target, overhead) for target in args.targets + } + finally: + tokenizer.close() + + if args.prepare_only: + print( + json.dumps( + { + "calibration_raw_tokens": calibration_raw, + "chat_template_overhead": overhead, + "contexts": {str(target): fitted[target][1] for target in args.targets}, + }, + indent=2, + ) + ) + return 0 + + groups: list[dict[str, Any]] = [] + failed = False + for target in args.targets: + prompt, construction = fitted[target] + records: list[dict[str, Any]] = [] + total = args.warmup + args.runs + print(f"[context-sweep] target={target} requests={total}", flush=True) + for index in range(total): + measured = index >= args.warmup + label = "measure" if measured else "warmup" + print( + f"[context-sweep] target={target} {label} {index + 1}/{total}", + flush=True, + ) + result = stream_request(args.url, args.model, prompt, args.max_tokens) + result.update({"index": index, "measured": measured, "target_context": target}) + hash_ok = ( + args.expected_sha256 is None + or result.get("response_sha256") == args.expected_sha256 + ) + result["expected_hash_match"] = hash_ok + records.append(result) + print( + "[context-sweep] " + f"ok={result.get('ok')} prompt={result.get('prompt_tokens')} " + f"output={result.get('completion_tokens')} " + f"client_decode={result.get('client_decode_tok_s')} tok/s " + f"sha={result.get('response_sha256')} hash_ok={hash_ok}", + flush=True, + ) + if ( + not result.get("ok") + or int(result.get("prompt_tokens") or -1) != target + or int(result.get("completion_tokens") or -1) != args.max_tokens + or not hash_ok + ): + failed = True + break + measured_rows = [row for row in records if row["measured"]] + groups.append( + { + "target_context": target, + "construction": construction, + "records": records, + "measured_summary": summarize(measured_rows), + } + ) + if failed: + break + + payload = { + "schema_version": 1, + "workload": "deterministic-exact-context-sweep", + "targets": args.targets, + "temperature": 0, + "batch_size": 1, + "max_tokens": args.max_tokens, + "warmup_per_context": args.warmup, + "runs_per_context": args.runs, + "expected_sha256": args.expected_sha256, + "calibration_server_tokens": args.calibration_server_tokens, + "calibration_raw_tokens": calibration_raw, + "chat_template_overhead": overhead, + "system_message": SYSTEM_MESSAGE, + "groups": groups, + } + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + print(json.dumps({str(row["target_context"]): row["measured_summary"] for row in groups}, indent=2)) + return 1 if failed or len(groups) != len(args.targets) else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/server/scripts/ds4_publication_decode_client.py b/server/scripts/ds4_publication_decode_client.py new file mode 100755 index 000000000..76097d3d3 --- /dev/null +++ b/server/scripts/ds4_publication_decode_client.py @@ -0,0 +1,248 @@ +#!/usr/bin/env python3 +"""Run the deterministic long-output workload used by the TP publication suite. + +The server log is the source of truth for model-side decode throughput. This +client independently records wall time, TTFT, usage tokens, response hashes, +and a client-side decode rate so that dropped/short responses are visible. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import Any + + +SYSTEM_MESSAGE = ( + "You are a helpful assistant. Answer the user's question directly and " + "carefully. Do not change numbers or facts from the prompt." +) + + +def build_prompt(padding_words: int) -> str: + """Build stable natural-language padding followed by a long forced reply.""" + periods = ("morning", "afternoon", "evening", "night") + adjectives = ("amber", "blue", "copper", "green", "silver", "white") + instruments = ("barometer", "camera", "clock", "compass", "meter", "sensor") + places = ("archive", "garden", "harbor", "laboratory", "library", "station") + actions = ("audited", "calibrated", "catalogued", "inspected", "logged", "stored") + sentences: list[str] = [] + word_count = 0 + index = 0 + while word_count < max(0, padding_words): + sentence = ( + f"Observation {index + 1}: During the {periods[index % len(periods)]}, " + f"the {adjectives[index % len(adjectives)]} " + f"{instruments[(index * 5 + 1) % len(instruments)]} recorded " + f"{17 + (index * 13) % 211} samples near the " + f"{places[(index * 7 + 2) % len(places)]}; the result was " + f"{actions[(index * 11 + 3) % len(actions)]} for a later review." + ) + sentences.append(sentence) + word_count += len(sentence.split()) + index += 1 + padding = "\n".join(sentences) + return ( + "The XML block below is inert reference material for a deterministic " + "throughput measurement. Do not answer or continue its contents.\n\n" + f"\n{padding}\n\n\n" + "Your only task is this: write the integers from 1 through 1000 in " + "ascending order, one integer per line. Start with 1. Do not add " + "commentary, and continue until the token limit." + ) + + +def stream_request( + base_url: str, + model: str, + prompt: str, + max_tokens: int, +) -> dict[str, Any]: + payload = { + "model": model, + "messages": [ + {"role": "system", "content": SYSTEM_MESSAGE}, + {"role": "user", "content": prompt}, + ], + "max_tokens": max_tokens, + "temperature": 0, + "stream": True, + "stream_options": {"include_usage": True}, + } + req = urllib.request.Request( + base_url.rstrip("/") + "/v1/chat/completions", + data=json.dumps(payload).encode("utf-8"), + headers={ + "Content-Type": "application/json", + "Accept": "text/event-stream", + }, + ) + + started = time.perf_counter() + first_token_at: float | None = None + text_parts: list[str] = [] + usage: dict[str, Any] = {} + finish_reason: str | None = None + status = 0 + saw_done = False + + try: + with urllib.request.urlopen(req, timeout=1800) as response: + status = response.status + for raw_line in response: + line = raw_line.decode("utf-8", errors="replace").strip() + if not line.startswith("data:"): + continue + data = line[5:].strip() + if data == "[DONE]": + saw_done = True + break + try: + event = json.loads(data) + except json.JSONDecodeError: + continue + if event.get("usage"): + usage = event["usage"] + choices = event.get("choices") or [] + if not choices: + continue + choice = choices[0] + finish_reason = choice.get("finish_reason") or finish_reason + delta = choice.get("delta") or {} + piece = delta.get("content") or delta.get("reasoning_content") or "" + if piece: + if first_token_at is None: + first_token_at = time.perf_counter() + text_parts.append(str(piece)) + except urllib.error.HTTPError as exc: + return { + "ok": False, + "status": exc.code, + "error": exc.read().decode("utf-8", errors="replace")[-4000:], + } + except (urllib.error.URLError, TimeoutError, ConnectionResetError) as exc: + return {"ok": False, "status": status, "error": repr(exc)} + + finished = time.perf_counter() + text = "".join(text_parts) + wall_s = finished - started + ttft_s = first_token_at - started if first_token_at is not None else None + decode_s = finished - first_token_at if first_token_at is not None else None + completion_tokens = int(usage.get("completion_tokens") or 0) + client_decode_tps = ( + completion_tokens / decode_s + if decode_s is not None and decode_s > 0 and completion_tokens > 0 + else None + ) + return { + "ok": status == 200 and bool(text) and saw_done, + "status": status, + "stream_done": saw_done, + "wall_s": round(wall_s, 6), + "ttft_s": round(ttft_s, 6) if ttft_s is not None else None, + "client_decode_s": round(decode_s, 6) if decode_s is not None else None, + "client_decode_tok_s": ( + round(client_decode_tps, 3) if client_decode_tps is not None else None + ), + "prompt_tokens": int(usage.get("prompt_tokens") or 0), + "completion_tokens": completion_tokens, + "finish_reason": finish_reason, + "response_sha256": hashlib.sha256(text.encode("utf-8")).hexdigest(), + "response_text": text, + } + + +def summarize(records: list[dict[str, Any]]) -> dict[str, Any]: + valid = [ + row + for row in records + if row.get("ok") and row.get("client_decode_tok_s") is not None + ] + if not valid: + return {"n": len(records), "n_ok": 0} + rates = [float(row["client_decode_tok_s"]) for row in valid] + decode_seconds = sum(float(row["client_decode_s"]) for row in valid) + completion_tokens = sum(int(row["completion_tokens"]) for row in valid) + return { + "n": len(records), + "n_ok": len(valid), + "prompt_tokens": sorted({int(row["prompt_tokens"]) for row in valid}), + "completion_tokens": sorted( + {int(row["completion_tokens"]) for row in valid} + ), + "client_decode_tok_s_median": round(statistics.median(rates), 3), + "client_decode_tok_s_min": round(min(rates), 3), + "client_decode_tok_s_max": round(max(rates), 3), + "client_decode_tok_s_weighted": round( + completion_tokens / decode_seconds, 3 + ), + "unique_response_hashes": sorted( + {str(row["response_sha256"]) for row in valid} + ), + } + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--url", default="http://127.0.0.1:18109") + parser.add_argument("--model", default="dflash") + parser.add_argument("--padding-words", type=int, default=1300) + parser.add_argument("--max-tokens", type=int, default=510) + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--runs", type=int, default=5) + parser.add_argument("--json-out", type=Path, required=True) + args = parser.parse_args() + + prompt = build_prompt(args.padding_words) + records: list[dict[str, Any]] = [] + total = args.warmup + args.runs + for index in range(total): + measured = index >= args.warmup + label = "measure" if measured else "warmup" + print(f"[publication] {label} {index + 1}/{total}", flush=True) + result = stream_request(args.url, args.model, prompt, args.max_tokens) + result.update({"index": index, "measured": measured}) + records.append(result) + print( + "[publication] " + f"ok={result.get('ok')} prompt={result.get('prompt_tokens')} " + f"output={result.get('completion_tokens')} " + f"client_decode={result.get('client_decode_tok_s')} tok/s " + f"sha={result.get('response_sha256')}", + flush=True, + ) + if not result.get("ok"): + break + + measured_records = [row for row in records if row.get("measured")] + payload = { + "schema_version": 1, + "workload": "deterministic-padded-counting", + "url": args.url, + "model": args.model, + "temperature": 0, + "batch_size": 1, + "padding_words": args.padding_words, + "max_tokens": args.max_tokens, + "warmup": args.warmup, + "runs": args.runs, + "system_message": SYSTEM_MESSAGE, + "prompt_sha256": hashlib.sha256(prompt.encode("utf-8")).hexdigest(), + "prompt_bytes": len(prompt.encode("utf-8")), + "records": records, + "measured_summary": summarize(measured_records), + } + args.json_out.parent.mkdir(parents=True, exist_ok=True) + args.json_out.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + print(json.dumps(payload["measured_summary"], indent=2), flush=True) + return 0 if payload["measured_summary"].get("n_ok") == args.runs else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/server/scripts/qualify_ds4_q5_amd.sh b/server/scripts/qualify_ds4_q5_amd.sh new file mode 100755 index 000000000..1a38d63ff --- /dev/null +++ b/server/scripts/qualify_ds4_q5_amd.sh @@ -0,0 +1,294 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Reproducible model-backed qualification for the AMD q=5 DS4 path. +# One process serves every context so the final 2K leg exercises eviction +# after 16K. Optional A/B switches are deliberately explicit. + +SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +CHECKOUT="${CHECKOUT:-$(cd "$SCRIPT_DIR/../.." && pwd)}" +BUILD_DIR="${BUILD_DIR:-$CHECKOUT/server/build-hip-dual}" +SERVER_BIN="${SERVER_BIN:-$BUILD_DIR/dflash_server}" +TOKENIZER_HARNESS="${TOKENIZER_HARNESS:-$BUILD_DIR/test_tokenizer_harness}" +TARGET_MODEL="${TARGET_MODEL:?set TARGET_MODEL to the target GGUF path}" +DRAFT_MODEL="${DRAFT_MODEL:?set DRAFT_MODEL to the DSpark draft GGUF path}" +HOTNESS_CSV="${HOTNESS_CSV:?set HOTNESS_CSV to the expert hotness CSV path}" +CONTEXT_CLIENT="${CONTEXT_CLIENT:-$SCRIPT_DIR/ds4_context_sweep.py}" +EXPECTED_SHA256="${EXPECTED_SHA256:-0f785a7ffa406498aafb14553966eaed0f52220fed0f7cc016b66921d104d194}" +PORT="${PORT:-18109}" +MAX_CTX="${MAX_CTX:-18432}" +CACHE_SLOTS="${CACHE_SLOTS:-auto}" +MMVQ_MAX_NCOLS="${MMVQ_MAX_NCOLS:-auto}" +FORCE_GRAPH_REPLAY="${FORCE_GRAPH_REPLAY:-0}" +SERIAL_INDEX_SCAN="${SERIAL_INDEX_SCAN:-0}" +DIRECT_INDEXER_TOPK="${DIRECT_INDEXER_TOPK:-1}" +BLOCK_RADIX_TOPK="${BLOCK_RADIX_TOPK:-1}" +PACK_Q4_INDEXER="${PACK_Q4_INDEXER:-0}" +Q5_VERIFY="${Q5_VERIFY:-1}" +FP4_Q5_X4_PLUS1="${FP4_Q5_X4_PLUS1:-auto}" +EXPERT_BUDGET_MB="${EXPERT_BUDGET_MB:-13200}" +WARMUP="${WARMUP:-2}" +RUNS="${RUNS:-3}" +MAX_TOKENS="${MAX_TOKENS:-128}" +TARGETS="${TARGETS:-2048 4096 8192 16384 2048}" +VRAM_MONITOR_SECONDS="${VRAM_MONITOR_SECONDS:-2}" +HASH_MODELS="${HASH_MODELS:-0}" +RUN_ID="${RUN_ID:-ds4-q5-fr${FORCE_GRAPH_REPLAY}-direct${DIRECT_INDEXER_TOPK}-radix${BLOCK_RADIX_TOPK}-x4p1${FP4_Q5_X4_PLUS1}-$(date -u +%Y%m%dT%H%M%SZ)}" +OUT_ROOT="${OUT_ROOT:-$CHECKOUT/results/ds4_q5_context_qualification}" +OUT_DIR="$OUT_ROOT/$RUN_ID" +SERVER_LOG="$OUT_DIR/server.log" + +for required in "$SERVER_BIN" "$TOKENIZER_HARNESS" "$TARGET_MODEL" \ + "$DRAFT_MODEL" "$HOTNESS_CSV" "$CONTEXT_CLIENT"; do + if [[ ! -e "$required" ]]; then + echo "missing required path: $required" >&2 + exit 2 + fi +done + +case "$FORCE_GRAPH_REPLAY:$SERIAL_INDEX_SCAN" in + 0:0|0:1|1:0|1:1) ;; + *) echo "FORCE_GRAPH_REPLAY and SERIAL_INDEX_SCAN must be 0 or 1" >&2; exit 2 ;; +esac +case "$DIRECT_INDEXER_TOPK" in + 0|1) ;; + *) echo "DIRECT_INDEXER_TOPK must be 0 or 1" >&2; exit 2 ;; +esac +case "$BLOCK_RADIX_TOPK" in + 0|1) ;; + *) echo "BLOCK_RADIX_TOPK must be 0 or 1" >&2; exit 2 ;; +esac +case "$PACK_Q4_INDEXER" in + 0|1) ;; + *) echo "PACK_Q4_INDEXER must be 0 or 1" >&2; exit 2 ;; +esac +case "$Q5_VERIFY" in + 0|1) ;; + *) echo "Q5_VERIFY must be 0 or 1" >&2; exit 2 ;; +esac +case "$FP4_Q5_X4_PLUS1" in + auto|0|1) ;; + *) echo "FP4_Q5_X4_PLUS1 must be auto, 0, or 1" >&2; exit 2 ;; +esac +if [[ "$MMVQ_MAX_NCOLS" != auto && ! "$MMVQ_MAX_NCOLS" =~ ^[1-8]$ ]]; then + echo "MMVQ_MAX_NCOLS must be auto or an integer from 1 through 8" >&2 + exit 2 +fi +if [[ "$CACHE_SLOTS" != auto && ! "$CACHE_SLOTS" =~ ^([1-9]|1[0-2])$ ]]; then + echo "CACHE_SLOTS must be auto or an integer from 1 through 12" >&2 + exit 2 +fi +case "$HASH_MODELS" in + 0|1) ;; + *) echo "HASH_MODELS must be 0 or 1" >&2; exit 2 ;; +esac + +if pgrep -f "dflash_server .*--port ${PORT}([[:space:]]|$)" >/dev/null; then + echo "benchmark port $PORT is already owned by another dflash_server" >&2 + exit 2 +fi + +mkdir -p "$OUT_DIR" + +server_pid="" +monitor_pid="" +cleanup() { + if [[ -n "$monitor_pid" ]] && kill -0 "$monitor_pid" 2>/dev/null; then + kill -TERM "$monitor_pid" 2>/dev/null || true + wait "$monitor_pid" 2>/dev/null || true + fi + if [[ -n "$server_pid" ]] && kill -0 "$server_pid" 2>/dev/null; then + kill -TERM "$server_pid" 2>/dev/null || true + wait "$server_pid" 2>/dev/null || true + fi +} +trap cleanup EXIT + +rocm-smi -d 0 --setperflevel auto >/dev/null 2>&1 || true +rocm-smi -d 1 --setperflevel high >/dev/null 2>&1 || true +printf '0\n' >/tmp/ds4_awidth +rm -f /tmp/ds4_spec_q + +server_env=( + env -i + "HOME=$HOME" + "USER=${USER:-unknown}" + "PATH=$PATH" + "LD_LIBRARY_PATH=${LD_LIBRARY_PATH:-}" + "GGML_CUDA_GRAPH_STATS=1" + "LUCE_CUDA_I32_REPEAT=1" + "DFLASH_DS4_TOPK=4" + "DFLASH_DS4_FUSED_VERIFY=1" + "DFLASH_DS4_FUSED_HYBRID_DECODE=1" + "DFLASH_DS4_TIMING=1" + "DFLASH_CUDA_MMVQ_MOE_ROWS_PER_BLOCK=2" + "DFLASH_CUDA_MMVQ_MOE_FP3_PACKED24=1" + "DFLASH_CUDA_MMVQ_MOE_FP2_PACKED32=0" + "DFLASH_CUDA_MMVQ_FP4_X4=1" + "DFLASH_ROCMFP2_FIXED_K=1" + "DFLASH_ROCMFP3_FIXED_K=1" + "DFLASH_ROCMFP4_UNROLL2=1" + "DFLASH_MMID_GROUPED=1" + "DFLASH_MMID_GROUPED_TYPES=8" + "DFLASH_MMID_GROUPED_DEVICE=1" + "DFLASH_DS4_MOE_TP=1" + "DFLASH_DS4_MOE_TP_INPROC=1" + "DFLASH_DS4_MOE_TP_GPU=1" + "DFLASH_EXPERT_BUDGET_MB=$EXPERT_BUDGET_MB" + "DFLASH_DS4_HOTNESS_CSV=$HOTNESS_CSV" + "DFLASH_DS4_TP_CAPTURE_CACHE_SLOTS=4" + "DFLASH_DS4_TP_MASKED_ROUTES=1" + "DFLASH_DS4_TP_GROUPED_MMVQ=1" + "DFLASH_DS4_TP_SPLIT_COUNT=1" + "DFLASH_DS4_TP_ROUTE_PREFORK=1" + "DFLASH_DS4_TP_DEVICE_JOIN=1" + "DFLASH_DS4_TP_DEVICE_JOIN_SPLIT=1" + "DFLASH_DS4_TP_FUSED_HC_JOIN=1" + "DFLASH_DS4_TP_MAIN_ROUTE_WEIGHTS=1" + "DFLASH_DS4_TP_COARSE_OWNER=1" + "DFLASH_DS4_TP_COARSE_OWNER_SPLIT=0" + "DFLASH_DS4_TP_NATIVE_ROUTE_WIDTH=1" + "GGML_CUDA_BATCH_PEER_COPIES=1" + "DFLASH_MOE_DUPLICATE_HOT_ON_COLD=1" + "DFLASH_DS4_HYBRID_PREFILL_GPU_HC=1" + "DFLASH_DS4_HYBRID_PREFILL_EAGER=1" + "DFLASH_MOE_FULL_COLD_PARALLEL=1" + "DFLASH_DS4_PREFILL_TRACE=0" + "DFLASH_MOE_PREFILL_PERSISTENT_OWNER_ALLOC=1" + "DFLASH_DS4_PINNED_ROLLBACK=1" + "DFLASH_DS4_GPU_ARGMAX_VERIFY=1" + "DFLASH_DS4_SPEC=1" + "DFLASH_DS4_SPEC_Q=$((4 + Q5_VERIFY))" + "DFLASH_DS4_ADAPTIVE_WIDTH=0" + "DFLASH_DS4_DRAFT=$DRAFT_MODEL" + "DFLASH_DS4_DRAFT_GPU=0" + "DFLASH_DS4_DRAFT_CONTEXT_KV_CACHE=1" + "DFLASH_MOE_FUSED_COMBINE=0" +) + +if [[ "$MMVQ_MAX_NCOLS" != auto ]]; then + server_env+=("LUCE_MMVQ_MAX_NCOLS=$MMVQ_MAX_NCOLS") +fi +if [[ "$CACHE_SLOTS" != auto ]]; then + server_env+=("DFLASH_DS4_TP_FUSED_CACHE_SLOTS=$CACHE_SLOTS") +fi + +if [[ "$FORCE_GRAPH_REPLAY" == 1 ]]; then + server_env+=("DFLASH_DS4_VERIFY_FORCE_GRAPH_REPLAY=1") +fi +if [[ "$SERIAL_INDEX_SCAN" == 1 ]]; then + server_env+=("GGML_DS4_FA_SERIAL_INDEX_SCAN=1") +fi +if [[ "$DIRECT_INDEXER_TOPK" == 1 ]]; then + server_env+=("DFLASH_DS4_DIRECT_INDEXER_TOPK=1") +fi +if [[ "$BLOCK_RADIX_TOPK" == 1 ]]; then + server_env+=("GGML_DS4_TOPK_BLOCK_RADIX=1") +fi +if [[ "$PACK_Q4_INDEXER" == 1 ]]; then + server_env+=("GGML_DS4_INDEXER_PACK_Q4=1") +fi +if [[ "$Q5_VERIFY" == 1 ]]; then + server_env+=("DFLASH_DS4_Q5_VERIFY=1") +fi +if [[ "$FP4_Q5_X4_PLUS1" != auto ]]; then + server_env+=("DFLASH_CUDA_MMVQ_FP4_Q5_X4_PLUS1=$FP4_Q5_X4_PLUS1") +fi + +server_args=( + "$SERVER_BIN" "$TARGET_MODEL" + --host 127.0.0.1 --port "$PORT" + --max-ctx "$MAX_CTX" + --target-device hip:0 + --prefix-cache-slots 0 + --prefill-cache-slots 0 + --hard-limit-reply-budget 0 + --chunk 2048 + --ds4-fused-decode + --ds4-expert-top-k 4 + --ds4-prefill sparse + --peer-access +) + +{ + echo "schema_version=1" + echo "run_id=$RUN_ID" + echo "source_commit=$(git -C "$CHECKOUT" rev-parse HEAD)" + echo "force_graph_replay=$FORCE_GRAPH_REPLAY" + echo "serial_index_scan=$SERIAL_INDEX_SCAN" + echo "direct_indexer_topk=$DIRECT_INDEXER_TOPK" + echo "block_radix_topk=$BLOCK_RADIX_TOPK" + echo "pack_q4_indexer=$PACK_Q4_INDEXER" + echo "q5_verify=$Q5_VERIFY" + echo "fp4_q5_x4_plus1=$FP4_Q5_X4_PLUS1" + echo "cache_slots=$CACHE_SLOTS" + echo "mmvq_max_ncols=$MMVQ_MAX_NCOLS" + echo "targets=$TARGETS" + echo "warmup=$WARMUP" + echo "runs=$RUNS" + echo "max_tokens=$MAX_TOKENS" + echo "max_ctx=$MAX_CTX" + sha256sum "$SERVER_BIN" + stat -c 'target_model=%n bytes=%s mtime=%y' "$TARGET_MODEL" + stat -c 'draft_model=%n bytes=%s mtime=%y' "$DRAFT_MODEL" + if [[ "$HASH_MODELS" == 1 ]]; then + sha256sum "$TARGET_MODEL" "$DRAFT_MODEL" + fi + printf 'server_env='; printf '%q ' "${server_env[@]}"; echo + printf 'server_args='; printf '%q ' "${server_args[@]}"; echo + date -u '+started_utc=%Y-%m-%dT%H:%M:%SZ' +} >"$OUT_DIR/manifest.txt" + +rocm-smi --showproductname --showdriverversion --showperflevel --showclocks \ + --showmeminfo vram >"$OUT_DIR/rocm-smi-before.txt" 2>&1 || true + +"${server_env[@]}" "${server_args[@]}" >"$SERVER_LOG" 2>&1 & +server_pid=$! + +ready=0 +for _ in $(seq 1 900); do + if grep -q "listening on" "$SERVER_LOG"; then + ready=1 + break + fi + if ! kill -0 "$server_pid" 2>/dev/null; then + tail -160 "$SERVER_LOG" >&2 + exit 1 + fi + sleep 1 +done +if [[ "$ready" != 1 ]]; then + echo "server did not become ready" >&2 + exit 1 +fi + +if [[ "$VRAM_MONITOR_SECONDS" -gt 0 ]]; then + ( + while kill -0 "$server_pid" 2>/dev/null; do + date -u '+sample_utc=%Y-%m-%dT%H:%M:%SZ' + rocm-smi --showuse --showmeminfo vram 2>&1 || true + sleep "$VRAM_MONITOR_SECONDS" + done + ) >"$OUT_DIR/vram-monitor.log" 2>&1 & + monitor_pid=$! +fi + +# shellcheck disable=SC2206 +target_args=($TARGETS) +python3 "$CONTEXT_CLIENT" \ + --url "http://127.0.0.1:$PORT" \ + --model dflash \ + --model-gguf "$TARGET_MODEL" \ + --tokenizer-harness "$TOKENIZER_HARNESS" \ + --targets "${target_args[@]}" \ + --warmup "$WARMUP" --runs "$RUNS" --max-tokens "$MAX_TOKENS" \ + --expected-sha256 "$EXPECTED_SHA256" \ + --json-out "$OUT_DIR/decode-client.json" \ + 2>&1 | tee "$OUT_DIR/decode-client.log" + +rocm-smi --showperflevel --showclocks --showmeminfo vram \ + >"$OUT_DIR/rocm-smi-after.txt" 2>&1 || true +date -u '+finished_utc=%Y-%m-%dT%H:%M:%SZ' >>"$OUT_DIR/manifest.txt" + +echo "OUT_DIR=$OUT_DIR" +grep -E 'DSpark decode|chat DONE|graph.*(warm|replay|invalid)' "$SERVER_LOG" | tail -120 || true diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index e9220d6e2..b28b54c78 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -42,10 +42,37 @@ static bool env_flag_enabled(const char * name) { return value && value[0] && std::strcmp(value, "0") != 0; } -static void configure_gfx1151_dspark_mmvq_default(int gpu) { +static void configure_dspark_mmvq_defaults(int gpu) { #if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) - if (!env_flag_enabled("DFLASH_DS4_SPEC") || - std::getenv("LUCE_MMVQ_MAX_NCOLS") != nullptr) { + if (!env_flag_enabled("DFLASH_DS4_SPEC")) { + return; + } + + // q=5 is an explicit AMD-only experiment and needs the plain quantized + // verifier matmuls to stay on MMVQ. The process-wide crossover applies to + // both owners in the heterogeneous graph, so set it before inspecting the + // target device (which is gfx1201 in the R9700 + gfx1151 launch). + if (env_flag_enabled("DFLASH_DS4_Q5_VERIFY")) { + if (std::getenv("LUCE_MMVQ_MAX_NCOLS") == nullptr && + ::setenv("LUCE_MMVQ_MAX_NCOLS", "5", 0) == 0) { + std::fprintf(stderr, + "[deepseek4] AMD DSpark q5: defaulting " + "LUCE_MMVQ_MAX_NCOLS=5\n"); + } + + cudaDeviceProp prop{}; + if (std::getenv("DFLASH_CUDA_MMVQ_FP4_Q5_X4_PLUS1") == nullptr && + cudaGetDeviceProperties(&prop, gpu) == cudaSuccess && + std::strncmp(prop.gcnArchName, "gfx1201", 7) == 0 && + ::setenv("DFLASH_CUDA_MMVQ_FP4_Q5_X4_PLUS1", "1", 0) == 0) { + std::fprintf(stderr, + "[deepseek4] gfx1201 DSpark q5: defaulting " + "ROCmFP4 x4+1 MMVQ\n"); + } + return; + } + + if (std::getenv("LUCE_MMVQ_MAX_NCOLS") != nullptr) { return; } @@ -827,7 +854,7 @@ bool DeepSeek4Backend::init() { // The shared MMVQ/MMQ crossover defaults to q=3 for NVIDIA. On gfx1151, // DSpark q=4 is faster through MMVQ. Keep AR and other devices unchanged, // and preserve LUCE_MMVQ_MAX_NCOLS as an explicit override. - configure_gfx1151_dspark_mmvq_default(cfg_.device.gpu); + configure_dspark_mmvq_defaults(cfg_.device.gpu); configure_gfx1201_hybrid_sub_batch_default(cfg_.device.gpu); backend_ = ggml_backend_cuda_init(cfg_.device.gpu); diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index c2db8d26c..0e373aca4 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -178,7 +178,8 @@ bool deepseek4_dspark_verify_forward(ggml_backend_t backend, // Minimal speculative-decode rollback state. Rejected positions must restore // the physical SWA rows they overwrote after the ring wraps; otherwise a later // causal verify reads rejected-token KV as if it were older committed history. -// This remains much smaller than a full target-cache snapshot because q <= 4. +// This remains much smaller than a full target-cache snapshot because the +// verifier width is bounded by the DSpark block (currently q <= 5). struct DeepSeek4SpecRollback { int raw_pos = 0; int raw_count = 0; diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index 21596aa6c..06a86f7b1 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -254,7 +254,7 @@ constexpr float kConfidenceQ3Threshold = 0.40f; constexpr float kConfidenceQ4Threshold = 0.30f; // ── Light rollback state ──────────────────────────────────────────────── -// Save the ratio-4 rolling state, HC state, and the raw SWA rows that a q<=4 +// Save the ratio-4 rolling state, HC state, and the raw SWA rows that a q<=5 // verify may overwrite. Pinned host storage lets the GPU copy this compact // rollback state on its stream before the verifier without a host fence. // prev-half = first 4 rows of a [comp_width, 8] ratio-4 rolling state. @@ -298,7 +298,7 @@ bool init_pinned_rollback(const DeepSeek4Cache & cache, DeepSeek4SpecRollback & s.pinned_idx_sc, prev_half_bytes(lc.indexer_compressor.state_score), total); s.raw_row_bytes = lc.raw_kv ? ggml_row_size(lc.raw_kv->type, lc.raw_kv->ne[0]) : 0; - assign_pinned_span(s.pinned_raw_rows, s.raw_row_bytes * 4, total); + assign_pinned_span(s.pinned_raw_rows, s.raw_row_bytes * 5, total); } assign_pinned_span( rb.pinned_hc, cache.hc_state ? ggml_nbytes(cache.hc_state) : 0, total); @@ -365,7 +365,7 @@ void spec_rollback_save(const DeepSeek4Cache & cache, DeepSeek4SpecRollback & rb ggml_backend_t backend, bool async_copy, bool pinned_copy, int raw_pos, int raw_count) { rb.raw_pos = raw_pos; - rb.raw_count = std::clamp(raw_count, 0, 4); + rb.raw_count = std::clamp(raw_count, 0, 5); rb.layers.resize(cache.layers.size()); if (async_copy || pinned_copy) { rb.async_backend = backend; @@ -666,6 +666,7 @@ bool run_deepseek4_dspark_spec_decode( spec_env_flag("DFLASH_DS4_SEQ_VERIFY"); const bool async_rollback = spec_env_flag("DFLASH_DS4_ASYNC_ROLLBACK"); const bool pinned_rollback = spec_env_flag("DFLASH_DS4_PINNED_ROLLBACK"); + const bool q5_verify = spec_env_flag("DFLASH_DS4_Q5_VERIFY") && block >= 4; const bool draft_overlap_probe = spec_env_flag("DFLASH_DS4_DRAFT_OVERLAP_PROBE"); const bool draft_overlap_reuse_context = @@ -708,15 +709,15 @@ bool run_deepseek4_dspark_spec_decode( } double ewma_accept = 1.5; - // Fast path caps the verify at the compression ratio (4): one boundary max, - // no rolling-state row aliasing -> snapshot-free rollback stays exact. - // Full snapshots change rollback strategy but not the compressor-window - // limit below. The legacy sequential measurement path is validated only - // through q=4. - int q_cap = full_snap ? block + 1 : 4; + // The conservative fast path remains capped at the compression ratio. + // The explicit q5 path handles the second ratio-4 boundary in-graph and + // restores/replays only a rejected q5 prefix, avoiding full snapshots on + // the overwhelmingly common all-accepted path. + const int fast_cap = q5_verify ? block + 1 : 4; + int q_cap = full_snap ? block + 1 : fast_cap; if (const char * qs = std::getenv("DFLASH_DS4_SPEC_Q")) { const int v = std::atoi(qs); - if (v >= 2 && v <= block + 1) q_cap = full_snap ? v : std::min(v, 4); + if (v >= 2 && v <= block + 1) q_cap = full_snap ? v : std::min(v, fast_cap); } if (std::FILE * qf = std::fopen("/tmp/ds4_spec_q", "r")) { // Per-request override for perf experiments (no server restart needed). @@ -724,7 +725,7 @@ bool run_deepseek4_dspark_spec_decode( // (diagnoses batched-vs-sequential target divergence). int v = 0; if (std::fscanf(qf, "%d", &v) == 1 && v >= 1 && v <= block + 1) { - q_cap = full_snap ? v : std::min(v, 4); + q_cap = full_snap ? v : std::min(v, fast_cap); } std::fclose(qf); } @@ -827,6 +828,10 @@ bool run_deepseek4_dspark_spec_decode( } } tm_draft += spec_ms_since(t0); + if (q5_verify && steps == 0) { + std::fprintf(stderr, "[ds4-q5] draft-ready block=%d hidden=%zu\n", + block, local_hidden.size()); + } if (debug) { size_t lh_nan = 0; double lh_ss = 0; @@ -857,7 +862,7 @@ bool run_deepseek4_dspark_spec_decode( return v && *v && *v != '0'; }(); int q_step_cap = (seq_verify_mode || fused_verify_mode) - ? std::min(q_cap, 4) + ? std::min(q_cap, q5_verify ? 5 : 4) : std::min(q_cap, 4 - (pos & 3)); if (adaptive_width && !use_confidence_width && !seq_verify_mode) { const int w_cap = (int) ewma_accept + 2; @@ -922,6 +927,9 @@ bool run_deepseek4_dspark_spec_decode( if ((int) draft_tok.size() > q_step_cap) draft_tok.resize(q_step_cap); const int q = (int) draft_tok.size(); // seed + candidates tm_head += spec_ms_since(t0); + if (q5_verify && steps == 0) { + std::fprintf(stderr, "[ds4-q5] head-ready q=%d\n", q); + } if (debug) { std::fprintf(stderr, "[ds4-spec] dbg ds_ok=%d q=%d lt=%d draft=[%d %d %d %d]\n", @@ -963,6 +971,10 @@ bool run_deepseek4_dspark_spec_decode( pos, q); } tm_save += spec_ms_since(t0); + if (q5_verify && steps == 0) { + std::fprintf(stderr, "[ds4-q5] rollback-ready rows=%d\n", + rollback.raw_count); + } // First ratio-4 boundary position touched by this verify (p % 4 == 3). const int first_boundary = pos + (3 - (pos & 3)); @@ -971,6 +983,9 @@ bool run_deepseek4_dspark_spec_decode( // ── ONE batched verify (writes cache + captures features for all q) ── t0 = SpecClock::now(); int verify_last = -1; + if (q5_verify && steps == 0) { + std::fprintf(stderr, "[ds4-q5] verify-begin q=%d pos=%d\n", q, pos); + } const bool verify_ok = target.verify_batch(draft_tok, pos, verify_last, &tgt_am); tm_verify += spec_ms_since(t0); @@ -1038,6 +1053,25 @@ bool run_deepseek4_dspark_spec_decode( ok = false; break; } + } else if (!full_snap && accept < q && q > 4) { + // A rejected q5 may have crossed two ratio-4 boundaries. Restore + // the compact pre-verify state and replay only the accepted prefix + // (at most q4), which is exact and rare at high acceptance. + spec_rollback_apply( + rollback, target_w, target_cache, pos, true, + backend, async_rollback || pinned_rollback, + pinned_rollback); + std::vector kv_toks; + kv_toks.reserve((size_t) accept); + kv_toks.push_back(lt); + for (int i = 1; i < accept; ++i) kv_toks.push_back(draft_tok[i]); + int replay_last = -1; + std::vector replay_am; + if (!target.verify_batch(kv_toks, pos, replay_last, &replay_am)) { + std::fprintf(stderr, "[ds4-spec] q5 rollback replay failed\n"); + ok = false; + break; + } } else if (!full_snap && accept < q) { // The prev-half flush is bad only if the boundary sits at-or-past // the commit point (its chunk then contains rejected tokens). diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index ac583ad65..fcf228095 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -1,4 +1,4 @@ -// ─── Fused whole-model graph (n_tokens = q in [1,4]) ──────────────────── +// ─── Fused whole-model graph (n_tokens = q in [1,5]) ──────────────────── // One whole-model graph per (q, flush, padded-comp) shape: batched attention // (causal mask input) + batched MoE, per-token fused HC chains, in-graph // optional drafter-feature capture and q-wide logits. The spec loop guarantees a @@ -287,7 +287,15 @@ static void ds4_fused_consume_route_diagnostics( static size_t ds4_fused_verify_hybrid_slot_limit() { static const size_t limit = []() { const char * raw = std::getenv("DFLASH_DS4_TP_FUSED_CACHE_SLOTS"); - const long requested = raw ? std::strtol(raw, nullptr, 10) : 2; + // q=5 advances through four ratio-4 phases and crosses a ratio-128 + // boundary in the 128-token qualification request. With two slots the + // nine recurring shapes rebuild on every step. Nine slots were + // model-qualified at 28.5 GiB peak on the 31.9 GiB R9700. q<=4 keeps + // the conservative two-slot default. + const long default_slots = + ds4_env_flag("DFLASH_DS4_Q5_VERIFY") ? 9 : 2; + const long requested = raw ? std::strtol(raw, nullptr, 10) + : default_slots; return (size_t) std::max( 1, std::min(requested, (long) Ds4FusedVerifyCache::kSlotCount)); @@ -386,11 +394,11 @@ static bool ds4_build_fused_verify_graph( ex.ape128 = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.ape128); ex.st4 = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, q); ggml_set_input(ex.st4); ex.st128 = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, q); ggml_set_input(ex.st128); - // comp emission scalars per layer: i32 {ape_row(unused batched), comp_pos}, - // i64 {comp_row}; reuse the decode-style bundles (2 i32 + 1 i64 per layer). + // Up to two compressor emissions per layer. q<=4 uses at most one; q=5 + // can touch two ratio-4 boundaries when it starts at position mod 4 == 3. fg.i32_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 2 * (int64_t) w.n_layer); ggml_set_input(fg.i32_bundle); - fg.i64_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1 * (int64_t) w.n_layer); + fg.i64_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 2 * (int64_t) w.n_layer); ggml_set_input(fg.i64_bundle); // Mask bundle: q>1 preserves the q ring rows that the batch overwrites, @@ -460,20 +468,36 @@ static bool ds4_build_fused_verify_graph( ain.neg_pos = ex.neg_q; ain.raw_kv_rows = ex.rawrows; if (ratio > 0) { + int n_flush = 0; + for (int t = 0; t < lane_q; ++t) { + if (((lane_kv_start + t + 1) % ratio) == 0) ++n_flush; + } + const int n_comp_inputs = std::max(1, n_flush); ggml_tensor * ape_all = (ratio == 4) ? ex.ape4 : ex.ape128; ggml_tensor * state_all = (ratio == 4) ? ex.st4 : ex.st128; ain.attn_ape_row = ape_all; ain.attn_state_rows = state_all; - ain.attn_comp_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 2 + 1) * sizeof(int32_t)); - ain.attn_comp_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), - (size_t) il * sizeof(int64_t)); + ain.attn_comp_pos = ggml_view_1d( + ctx, fg.i32_bundle, n_comp_inputs, + (size_t) il * 2 * sizeof(int32_t)); + ain.attn_comp_rows = ggml_view_1d( + ctx, fg.i64_bundle, n_comp_inputs, + (size_t) il * 2 * sizeof(int64_t)); } if (ratio == 4) { + int n_flush = 0; + for (int t = 0; t < lane_q; ++t) { + if (((lane_kv_start + t + 1) % ratio) == 0) ++n_flush; + } + const int n_comp_inputs = std::max(1, n_flush); ain.index_ape_row = ex.ape4; ain.index_state_rows = ex.st4; - ain.index_comp_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 2 + 1) * sizeof(int32_t)); - ain.index_comp_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), - (size_t) il * sizeof(int64_t)); + ain.index_comp_pos = ggml_view_1d( + ctx, fg.i32_bundle, n_comp_inputs, + (size_t) il * 2 * sizeof(int32_t)); + ain.index_comp_rows = ggml_view_1d( + ctx, fg.i64_bundle, n_comp_inputs, + (size_t) il * 2 * sizeof(int64_t)); } int padded = 0; if (ratio > 0 && lc.comp_kv) { @@ -1165,17 +1189,27 @@ static int ds4_try_fused_verify_step( ds4_fv_set(ex->st128, lv.data(), sizeof(int64_t) * q); std::vector i32v((size_t) w.n_layer * 2, 0); - std::vector i64v((size_t) w.n_layer * 1, 0); + std::vector i64v((size_t) w.n_layer * 2, 0); for (int il = 0; il < w.n_layer; ++il) { const int ratio = (int) w.compress_ratios[il]; if (ratio > 0) { - int pos_b = -1; // boundary position inside the batch (if any) + int n_flush = 0; for (int t = 0; t < q; ++t) { - if (((kv_start + t + 1) % ratio) == 0) { pos_b = kv_start + t; break; } + if (((kv_start + t + 1) % ratio) != 0) continue; + GGML_ASSERT(n_flush < 2); + const int pos_b = kv_start + t; + i32v[(size_t) il * 2 + (size_t) n_flush] = + pos_b + 1 - ratio; + i64v[(size_t) il * 2 + (size_t) n_flush] = + pos_b / ratio; + ++n_flush; + } + if (n_flush == 0) { + // The graph does not consume these values in a no-flush shape, + // but initialize them deterministically for diagnostics. + i32v[(size_t) il * 2] = token_pos + 1 - ratio; + i64v[(size_t) il * 2] = token_pos / ratio; } - i32v[(size_t) il * 2 + 0] = token_pos % ratio; - i32v[(size_t) il * 2 + 1] = (pos_b >= 0 ? pos_b : token_pos) + 1 - ratio; - i64v[(size_t) il] = (pos_b >= 0 ? pos_b : token_pos) / ratio; } } ds4_fv_set(fg->i32_bundle, i32v.data(), sizeof(int32_t) * i32v.size()); diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 88a8d8263..c7db80112 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -1185,6 +1185,9 @@ static void build_compressor_step( pooled = ggml_reshape_2d(ctx, pooled, head_dim, 1); ggml_tensor * comp_pos = comp_pos_inp; + if (comp_pos && ggml_nelements(comp_pos) > 1) { + comp_pos = ggml_view_1d(ctx, comp_pos, 1, 0); + } if (!comp_pos) { comp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); ggml_set_input(comp_pos); @@ -1209,7 +1212,11 @@ static void build_compressor_step( } if (comp_rows_inp) { - comp_cache_source = ggml_set_rows(ctx, comp_cache, pooled, comp_rows_inp); + ggml_tensor * first_comp_row = comp_rows_inp; + if (ggml_nelements(first_comp_row) > 1) { + first_comp_row = ggml_view_1d(ctx, first_comp_row, 1, 0); + } + comp_cache_source = ggml_set_rows(ctx, comp_cache, pooled, first_comp_row); ggml_build_forward_expand(gf, comp_cache_source); } else { ggml_tensor * comp_slot = ggml_view_2d( @@ -1248,6 +1255,8 @@ static void build_compressor_step( ggml_build_forward_expand(gf, ggml_cpy(ctx, src_sc, dst_sc)); } } + ggml_tensor * tail_kv_source = nullptr; + ggml_tensor * tail_score_source = nullptr; if (batched_nB > 0) { ggml_tensor * kv_v = ggml_cont(ctx, ggml_view_2d( ctx, batched_kv_all, comp_width, batched_nB, @@ -1260,10 +1269,99 @@ static void build_compressor_step( ggml_tensor * rows_v = ggml_view_1d( ctx, state_rows_inp, batched_nB, (size_t) batched_span_off * state_rows_inp->nb[0]); - ggml_build_forward_expand( - gf, ggml_set_rows(ctx, state.state_kv, kv_v, rows_v)); - ggml_build_forward_expand( - gf, ggml_set_rows(ctx, state.state_score, sc_v, rows_v)); + tail_kv_source = ggml_set_rows(ctx, state.state_kv, kv_v, rows_v); + tail_score_source = ggml_set_rows(ctx, state.state_score, sc_v, rows_v); + ggml_build_forward_expand(gf, tail_kv_source); + ggml_build_forward_expand(gf, tail_score_source); + } + + // q=5 can start on the last position of a ratio-4 window. In that + // shape the first token flushes one row and the four-token tail fills + // and flushes the next window. Pool the second window in the same + // graph, then rotate it into the persistent previous half. + const bool second_boundary = + ratio == 4 && batched_nB == ratio && tail_kv_source && + tail_score_source && comp_pos_inp && comp_rows_inp && + ggml_nelements(comp_pos_inp) >= 2 && + ggml_nelements(comp_rows_inp) >= 2; + if (second_boundary) { + const size_t hi_off_kv = + (size_t) ratio * tail_kv_source->nb[1] + + (size_t) head_dim * tail_kv_source->nb[0]; + const size_t hi_off_sc = + (size_t) ratio * tail_score_source->nb[1] + + (size_t) head_dim * tail_score_source->nb[0]; + ggml_tensor * prev_kv = ggml_view_2d( + ctx, tail_kv_source, head_dim, ratio, + tail_kv_source->nb[1], 0); + ggml_tensor * cur_kv_hi = ggml_view_2d( + ctx, tail_kv_source, head_dim, ratio, + tail_kv_source->nb[1], hi_off_kv); + ggml_tensor * prev_sc = ggml_view_2d( + ctx, tail_score_source, head_dim, ratio, + tail_score_source->nb[1], 0); + ggml_tensor * cur_sc_hi = ggml_view_2d( + ctx, tail_score_source, head_dim, ratio, + tail_score_source->nb[1], hi_off_sc); + ggml_tensor * second_kv = ggml_concat( + ctx, prev_kv, cur_kv_hi, 1); + ggml_tensor * second_sc = ggml_concat( + ctx, prev_sc, cur_sc_hi, 1); + ggml_tensor * second_sc_t = ggml_cont( + ctx, ggml_transpose(ctx, second_sc)); + ggml_tensor * second_kv_t = ggml_cont( + ctx, ggml_transpose(ctx, second_kv)); + ggml_tensor * second_probs = ggml_soft_max(ctx, second_sc_t); + ggml_tensor * second_weighted = ggml_mul( + ctx, second_probs, second_kv_t); + ggml_tensor * second_pooled = ggml_reshape_1d( + ctx, ggml_sum_rows(ctx, second_weighted), head_dim); + second_pooled = ggml_cont(ctx, second_pooled); + second_pooled = build_rms_norm( + ctx, second_pooled, norm_weight, rms_eps); + second_pooled = ggml_reshape_2d( + ctx, second_pooled, head_dim, 1); + ggml_tensor * second_comp_pos = ggml_view_1d( + ctx, comp_pos_inp, 1, comp_pos_inp->nb[0]); + second_pooled = build_tail_rope_2d( + ctx, second_pooled, second_comp_pos, n_rot, head_dim, 1, + compress_rope_freq_base, rope_scale, 1.0f, rope_attn, + rope_yarn_beta_fast, rope_yarn_beta_slow, rope_orig_ctx); + if (indexer_qat) { + second_pooled = ggml_ds4_indexer_qat( + ctx, ggml_cont(ctx, second_pooled)); + } + ggml_tensor * second_comp_row = ggml_view_1d( + ctx, comp_rows_inp, 1, comp_rows_inp->nb[0]); + comp_cache_source = ggml_set_rows( + ctx, comp_cache_source, second_pooled, second_comp_row); + ggml_build_forward_expand(gf, comp_cache_source); + + for (int r = 0; r < ratio; ++r) { + ggml_tensor * src_kv = ggml_view_2d( + ctx, tail_kv_source, comp_width, 1, + tail_kv_source->nb[1], + (size_t) (ratio + r) * tail_kv_source->nb[1]); + ggml_tensor * dst_kv = ggml_view_2d( + ctx, state.state_kv, comp_width, 1, + state.state_kv->nb[1], + (size_t) r * state.state_kv->nb[1]); + ggml_build_forward_expand( + gf, ggml_cpy(ctx, src_kv, dst_kv)); + ggml_tensor * src_sc = ggml_view_2d( + ctx, tail_score_source, comp_width, 1, + tail_score_source->nb[1], + (size_t) (ratio + r) * tail_score_source->nb[1]); + ggml_tensor * dst_sc = ggml_view_2d( + ctx, state.state_score, comp_width, 1, + state.state_score->nb[1], + (size_t) r * state.state_score->nb[1]); + ggml_build_forward_expand( + gf, ggml_cpy(ctx, src_sc, dst_sc)); + } + if (comp_cache_source_out) { + *comp_cache_source_out = comp_cache_source; + } } return; } @@ -1911,6 +2009,8 @@ static ggml_tensor * build_mla_attention( score_mask = ggml_reshape_2d(ctx, cmask, n_attn, n_tokens); } } + const bool direct_indexer_topk = indexer_topk && + ds4_env_flag("DFLASH_DS4_DIRECT_INDEXER_TOPK"); if (indexer_topk) { if (!score_mask) { score_mask = ggml_new_tensor_2d( @@ -1921,8 +2021,10 @@ static ggml_tensor * build_mla_attention( std::vector((size_t) n_attn * n_tokens, 0.0f), }); } - score_mask = ggml_ds4_indexer_mask( - ctx, ggml_cont(ctx, score_mask), indexer_topk, n_raw); + if (!direct_indexer_topk) { + score_mask = ggml_ds4_indexer_mask( + ctx, ggml_cont(ctx, score_mask), indexer_topk, n_raw); + } } ggml_tensor * context = nullptr; bool inverse_rope_fused = false; @@ -2084,6 +2186,10 @@ static ggml_tensor * build_mla_attention( : attention_impl == DeepSeek4AttentionImpl::SparseFlash ? w.n_indexer_top_k : 0, 32); + if (direct_indexer_topk) { + ggml_flash_attn_ext_set_ds4_indexer_topk( + context, indexer_topk); + } if (attention_impl != DeepSeek4AttentionImpl::Explicit && head_dim == 512 && n_rot == 64) { ggml_flash_attn_ext_set_ds4_inverse_rope( @@ -6624,13 +6730,16 @@ bool deepseek4_step_layer_range( moe_hybrid->materialized_cold_experts && moe_hybrid->cold_backend_kind == MoeHybridColdBackend::Gpu && moe_hybrid->cold_backend && moe_hybrid->cold_backend != backend; + const bool q5_verify_candidate = + n_tokens == 5 && ds4_env_flag("DFLASH_DS4_Q5_VERIFY"); const bool fused_verify_candidate = (!moe_hybrid || fused_hybrid_ready) && - n_tokens >= 2 && n_tokens <= 4 && verify_hooks && + n_tokens >= 2 && (n_tokens <= 4 || q5_verify_candidate) && verify_hooks && layer_begin == 0 && is_last_shard && out_logits && ds4_backend_is_gpu(backend) && ds4_fused_verify_enabled(); const bool heterogeneous_sparse_prefill = - moe_hybrid && cache.prefill_mode == PrefillAttentionMode::Sparse && + !fused_verify_candidate && moe_hybrid && + cache.prefill_mode == PrefillAttentionMode::Sparse && n_tokens > 4 && n_tokens <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS && layer_begin == 0 && is_last_shard && out_logits && ds4_backend_is_gpu(backend); @@ -6860,7 +6969,8 @@ bool deepseek4_step_layer_range( (fused_hybrid_decode && !verify_hooks) ? &fused_hybrid_decode_hooks : verify_hooks; if ((!moe_hybrid || fused_hybrid_ready) && - ((n_tokens >= 2 && n_tokens <= 4 && verify_hooks) || + ((n_tokens >= 2 && + (n_tokens <= 4 || q5_verify_candidate) && verify_hooks) || fused_hybrid_decode) && layer_begin == 0 && is_last_shard && out_logits && ds4_backend_is_gpu(backend) && ds4_fused_verify_enabled()) { diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 8649c37fc..2bc8f49ec 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -1680,7 +1680,8 @@ static void test_dspark_raw_ring_rollback_after_wrap(ggml_backend_t backend) { layer.n_comp = 2; layer.n_index_comp = 2; DeepSeek4SpecRollback rollback; - deepseek4_spec_rollback_save(cache, rollback, 10, 4); + deepseek4_spec_rollback_save(cache, rollback, 10, 5); + TEST_ASSERT(rollback.raw_count == 5); auto overwrite_row = [&](int absolute_pos, uint8_t value) { const int row = absolute_pos % weights.n_swa; @@ -1694,7 +1695,7 @@ static void test_dspark_raw_ring_rollback_after_wrap(ggml_backend_t backend) { expected.begin() + (size_t) row * layer.raw_kv->nb[1] + row_bytes, value); }; - for (int t = 0; t < 4; ++t) { + for (int t = 0; t < 5; ++t) { overwrite_row(10 + t, (uint8_t) (0xa0 + t)); } @@ -2566,6 +2567,10 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { ctx, GGML_TYPE_F32, head_dim, n_kv, 1); ggml_tensor * mask = ggml_new_tensor_2d( ctx, GGML_TYPE_F16, n_kv, n_tokens); + ggml_tensor * direct_mask = ggml_new_tensor_2d( + ctx, GGML_TYPE_F16, n_kv, n_tokens); + ggml_tensor * direct_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); @@ -2573,12 +2578,22 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { // enters the compact four-head HIP path and exceeds the parallel threshold. ggml_flash_attn_ext_set_ds4_sparse( output, raw_rows, raw_window, -selected_rows, 1); + ggml_tensor * direct_output = ggml_flash_attn_ext( + ctx, q, kv, kv, direct_mask, + 1.0f / std::sqrt((float) head_dim), 0.0f, 0.0f); + ggml_flash_attn_ext_set_ds4_sparse( + direct_output, raw_rows, raw_window, -selected_rows, 1); + ggml_flash_attn_ext_set_ds4_indexer_topk(direct_output, direct_topk); ggml_set_output(output); + ggml_set_output(direct_output); TEST_ASSERT_MSG(ggml_backend_supports_op(backend, output), "GPU rejected exact indexed DS4 attention"); + TEST_ASSERT_MSG(ggml_backend_supports_op(backend, direct_output), + "GPU rejected direct-top-k DS4 attention"); ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); ggml_build_forward_expand(graph, output); + ggml_build_forward_expand(graph, direct_output); ggml_gallocr_t alloc = ggml_gallocr_new( ggml_backend_get_default_buffer_type(backend)); const bool allocated = ggml_gallocr_alloc_graph(alloc, graph); @@ -2588,6 +2603,10 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { 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 direct_mask_data( + (size_t) n_kv * n_tokens, ggml_fp32_to_fp16(-1.0e30f)); + std::vector direct_topk_data( + (size_t) selected_rows * n_tokens); for (size_t i = 0; i < q_data.size(); ++i) { q_data[i] = ((int) (i % 31) - 15) * 0.001f; } @@ -2597,12 +2616,27 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { for (int token = 0; token < n_tokens; ++token) { ggml_fp16_t * token_mask = mask_data.data() + (size_t) token * n_kv; + ggml_fp16_t * token_direct_mask = + direct_mask_data.data() + (size_t) token * n_kv; for (int row = raw_rows - raw_window; row < raw_rows; ++row) { token_mask[row] = ggml_fp32_to_fp16(0.0f); + token_direct_mask[row] = ggml_fp32_to_fp16(0.0f); } for (int row = token; row < n_comp_rows; row += 2) { token_mask[raw_rows + row] = ggml_fp32_to_fp16(0.0f); } + for (int row = 0; row < n_comp_rows; ++row) { + token_direct_mask[raw_rows + row] = + ggml_fp32_to_fp16(0.0f); + } + // The model's top-k is score ordered. Reverse the physical order + // here so this test proves that the direct path restores the old + // ascending accumulation order rather than merely accepting an + // already sorted fixture. + for (int rank = 0; rank < selected_rows; ++rank) { + direct_topk_data[(size_t) token * selected_rows + rank] = + token + 2 * (selected_rows - 1 - rank); + } } ggml_backend_tensor_set(q, q_data.data(), 0, q_data.size() * sizeof(float)); @@ -2610,6 +2644,12 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { kv_data.size() * sizeof(float)); ggml_backend_tensor_set(mask, mask_data.data(), 0, mask_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set( + direct_mask, direct_mask_data.data(), 0, + direct_mask_data.size() * sizeof(ggml_fp16_t)); + ggml_backend_tensor_set( + direct_topk, direct_topk_data.data(), 0, + direct_topk_data.size() * sizeof(int32_t)); const char * previous_serial = std::getenv("GGML_DS4_FA_SERIAL_INDEX_SCAN"); @@ -2617,6 +2657,7 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { ? previous_serial : ""; std::vector serial((size_t) ggml_nelements(output)); std::vector parallel(serial.size()); + std::vector direct(serial.size()); { setenv("GGML_DS4_FA_SERIAL_INDEX_SCAN", "1", 1); ScopedCudaGraphOverrides eager( @@ -2642,6 +2683,8 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { "parallel indexed attention failed"); ggml_backend_tensor_get(output, parallel.data(), 0, parallel.size() * sizeof(float)); + ggml_backend_tensor_get(direct_output, direct.data(), 0, + direct.size() * sizeof(float)); } if (previous_serial) { setenv("GGML_DS4_FA_SERIAL_INDEX_SCAN", @@ -2653,6 +2696,260 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { TEST_ASSERT_MSG( nearly_equal(serial[i], parallel[i], 1.0e-6f, 1.0e-6f), "parallel index scan changed attention output"); + TEST_ASSERT_MSG( + serial[i] == direct[i], + "direct top-k changed attention output"); + } + } + + ggml_gallocr_free(alloc); + ggml_free(ctx); + ggml_backend_free(backend); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + +static void test_ds4_indexer_score_packed_q4_gpu() { + std::fprintf(stderr, " test_ds4_indexer_score_packed_q4_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 dim = 128; + constexpr int n_heads = 64; + constexpr int n_tokens = 4; + constexpr int n_comp = 4160; + constexpr int kv_start = 16384; + constexpr int ratio = 4; + 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, dim, n_heads, n_tokens); + ggml_tensor * weights = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_heads, n_tokens); + ggml_tensor * comp = ggml_new_tensor_2d( + ctx, GGML_TYPE_F16, dim, n_comp); + ggml_tensor * scores = ggml_ds4_indexer_score( + ctx, q, weights, comp, kv_start, ratio); + ggml_set_output(scores); + TEST_ASSERT_MSG(ggml_backend_supports_op(backend, scores), + "GPU rejected packed-q4 indexer fixture"); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 16, false); + ggml_build_forward_expand(graph, scores); + 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, "packed-q4 indexer graph allocation failed"); + if (allocated) { + std::vector q_data((size_t) dim * n_heads * n_tokens); + std::vector weight_data((size_t) n_heads * n_tokens); + std::vector comp_data((size_t) dim * n_comp); + for (size_t i = 0; i < q_data.size(); ++i) { + q_data[i] = ((int) (i % 31) - 15) * 0.0078125f; + } + for (size_t i = 0; i < weight_data.size(); ++i) { + weight_data[i] = ((int) (i % 17) - 8) * 0.015625f; + } + for (size_t i = 0; i < comp_data.size(); ++i) { + comp_data[i] = ggml_fp32_to_fp16( + ((int) (i % 29) - 14) * 0.0078125f); + } + ggml_backend_tensor_set(q, q_data.data(), 0, + q_data.size() * sizeof(float)); + ggml_backend_tensor_set(weights, weight_data.data(), 0, + weight_data.size() * sizeof(float)); + ggml_backend_tensor_set(comp, comp_data.data(), 0, + comp_data.size() * sizeof(ggml_fp16_t)); + + const char * previous = std::getenv("GGML_DS4_INDEXER_PACK_Q4"); + const std::string previous_value = previous ? previous : ""; + std::vector reference((size_t) n_comp * n_tokens); + std::vector candidate(reference.size()); + ScopedCudaGraphOverrides eager( + /*disable_graphs=*/true, + /*mmvq_max_ncols=*/0, + /*skip_property_check=*/false); + unsetenv("GGML_DS4_INDEXER_PACK_Q4"); + TEST_ASSERT_MSG( + ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, + "reference q4 indexer score failed"); + ggml_backend_tensor_get(scores, reference.data(), 0, + reference.size() * sizeof(float)); + + setenv("GGML_DS4_INDEXER_PACK_Q4", "1", 1); + TEST_ASSERT_MSG( + ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, + "packed q4 indexer score failed"); + ggml_backend_tensor_get(scores, candidate.data(), 0, + candidate.size() * sizeof(float)); + TEST_ASSERT_MSG( + std::memcmp(reference.data(), candidate.data(), + reference.size() * sizeof(float)) == 0, + "packed q4 indexer changed score bits"); + + auto measure_us = [&](bool packed) { + if (packed) { + setenv("GGML_DS4_INDEXER_PACK_Q4", "1", 1); + } else { + unsetenv("GGML_DS4_INDEXER_PACK_Q4"); + } + constexpr int warmups = 3; + constexpr int iterations = 30; + 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 reference_us = measure_us(false); + const double packed_us = measure_us(true); + std::fprintf(stderr, " reference=%.1fus packed=%.1fus", + reference_us, packed_us); + + if (previous) { + setenv("GGML_DS4_INDEXER_PACK_Q4", previous_value.c_str(), 1); + } else { + unsetenv("GGML_DS4_INDEXER_PACK_Q4"); + } + } + + ggml_gallocr_free(alloc); + ggml_free(ctx); + ggml_backend_free(backend); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + +static void test_ds4_topk_block_radix_gpu() { + std::fprintf(stderr, " test_ds4_topk_block_radix_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 ncols = 4160; + constexpr int nrows = 4; + constexpr int k = 512; + ggml_context * ctx = make_test_context(1u << 20); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) { + ggml_backend_free(backend); + std::fprintf(stderr, " FAIL\n"); + return; + } + + ggml_tensor * scores = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, ncols, nrows); + ggml_tensor * selected = ggml_top_k(ctx, scores, k); + ggml_set_output(selected); + TEST_ASSERT_MSG(ggml_backend_supports_op(backend, selected), + "GPU rejected long-context top-k fixture"); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 16, false); + ggml_build_forward_expand(graph, selected); + 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, "top-k graph allocation failed"); + if (allocated) { + std::vector score_data((size_t) ncols * nrows); + for (int row = 0; row < nrows; ++row) { + for (int col = 0; col < ncols; ++col) { + // 4051 is coprime with 4160, producing one exact permutation + // of unique integer-valued scores per row. Output order is not + // part of TOP_K's contract, so compare the selected sets. + score_data[(size_t) row * ncols + col] = + (float) (((col * 4051) + row * 997) % ncols); + } + } + ggml_backend_tensor_set(scores, score_data.data(), 0, + score_data.size() * sizeof(float)); + + const char * previous = std::getenv("GGML_DS4_TOPK_BLOCK_RADIX"); + const std::string previous_value = previous ? previous : ""; + std::vector reference((size_t) k * nrows); + std::vector candidate(reference.size()); + ScopedCudaGraphOverrides eager( + /*disable_graphs=*/true, + /*mmvq_max_ncols=*/0, + /*skip_property_check=*/false); + unsetenv("GGML_DS4_TOPK_BLOCK_RADIX"); + TEST_ASSERT_MSG( + ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, + "reference long-context top-k failed"); + ggml_backend_tensor_get(selected, reference.data(), 0, + reference.size() * sizeof(int32_t)); + + setenv("GGML_DS4_TOPK_BLOCK_RADIX", "1", 1); + TEST_ASSERT_MSG( + ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, + "block-radix long-context top-k failed"); + ggml_backend_tensor_get(selected, candidate.data(), 0, + candidate.size() * sizeof(int32_t)); + + for (int row = 0; row < nrows; ++row) { + auto ref_begin = reference.begin() + (size_t) row * k; + auto candidate_begin = candidate.begin() + (size_t) row * k; + std::sort(ref_begin, ref_begin + k); + std::sort(candidate_begin, candidate_begin + k); + TEST_ASSERT_MSG( + std::equal(ref_begin, ref_begin + k, candidate_begin), + "block-radix top-k changed the selected row set"); + } + + auto measure_us = [&](bool block_radix) { + if (block_radix) { + setenv("GGML_DS4_TOPK_BLOCK_RADIX", "1", 1); + } else { + unsetenv("GGML_DS4_TOPK_BLOCK_RADIX"); + } + constexpr int warmups = 5; + constexpr int iterations = 100; + 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 reference_us = measure_us(false); + const double candidate_us = measure_us(true); + std::fprintf(stderr, " reference=%.1fus block_radix=%.1fus", + reference_us, candidate_us); + + if (previous) { + setenv("GGML_DS4_TOPK_BLOCK_RADIX", previous_value.c_str(), 1); + } else { + unsetenv("GGML_DS4_TOPK_BLOCK_RADIX"); } } @@ -3697,6 +3994,8 @@ 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_indexer_score_packed_q4_gpu(); + test_ds4_topk_block_radix_gpu(); test_ds4_flash_attention_inverse_rope_fallback_gpu(); test_hc_post_strided_split_gpu(); test_hc_pre_kernel_gpu();