diff --git a/README.md b/README.md index 64f999206..5176b6444 100644 --- a/README.md +++ b/README.md @@ -93,6 +93,10 @@ These runs use different prompts, quantizations, and inference policies. They sh See [Recommended server setups](server/docs/RECOMMENDED_SETUPS.md) for the model and hardware matrix, including single-GPU and mixed-GPU profiles. +The DS4 guide also documents the Strix long-context sparse-verifier profile and +Qwen3-0.6B PFlash integration. PFlash is lossy prompt compression; keep it off +for exact-retrieval and matched true-context benchmarks. + ## Client Harnesses [`harness/`](harness/) runs Lucebox through popular coding clients and checks server compatibility. diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index f5fa16a60..f073b07e1 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2471,6 +2471,13 @@ extern "C" { int n_ctx_orig, bool q_unrotated); + // Optional runtime positions for both fused RoPE directions. I32 [n_query] + // replaces kv_start + query_index, allowing a cached graph to advance + // without rebuilding its topology or retaining stale position constants. + GGML_API void ggml_flash_attn_ext_set_ds4_rope_positions( + struct ggml_tensor * a, + struct ggml_tensor * positions); + // True when flash_attn_ext carries the DS4 sparse-layout or fused-RoPE // contract. Backends must implement that complete contract or reject it. GGML_API bool ggml_flash_attn_ext_is_ds4( 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 e7286de55..5f138d6a9 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 @@ -351,12 +351,14 @@ static __global__ void ds4_indexer_score_decode_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. The post-WMMA loop still accumulates heads in -// their original order, preserving the established F32 numerical topology. -static __global__ void ds4_indexer_score_wmma_q4_kernel( +// The speculative verifier scores only a few query tokens. The general WMMA +// kernel places them in a 16-row tile and executes the unused rows for every +// head. Pack consecutive heads into the tile instead: +// row = N_TOKENS*head_in_group + token. The post-WMMA loop still accumulates +// heads in their original order, preserving the established F32 numerical +// topology. This is the HIP equivalent of the Vulkan small-CM dispatch. +template +static __global__ void ds4_indexer_score_wmma_small_kernel( float * scores, const float * q, const float * weights, @@ -375,7 +377,14 @@ static __global__ void ds4_indexer_score_wmma_q4_kernel( __shared__ float c_sh[8 * 16 * 16]; __shared__ float weight_sh[16]; - float acc[2] = {0.0f, 0.0f}; + static_assert(N_TOKENS >= 2 && N_TOKENS <= 5, + "small-CM kernel is specialized for verifier widths 2..5"); + constexpr int HEADS_PER_TILE = 16 / N_TOKENS; + constexpr int USED_ROWS = HEADS_PER_TILE * N_TOKENS; + constexpr int ACC_SLOTS = (N_TOKENS + 1) / 2; + float acc[ACC_SLOTS]; +#pragma unroll + for (int slot = 0; slot < ACC_SLOTS; ++slot) acc[slot] = 0.0f; for (int i = tid; i < 128 * 128; i += 256) { const int c = i >> 7; @@ -387,21 +396,30 @@ static __global__ void ds4_indexer_score_wmma_q4_kernel( } __syncthreads(); - for (int head_base = 0; head_base < n_head; head_base += 4) { + for (int head_base = 0; head_base < n_head; + head_base += HEADS_PER_TILE) { 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); + half2 value = __float2half2_rn(0.0f); + if (row < USED_ROWS) { + const int token = row % N_TOKENS; + const int head = head_base + row / N_TOKENS; + if (head < n_head) { + const float2 q_value = + *reinterpret_cast( + q + ((size_t) token * n_head + head) * 128 + d); + value = __floats2half2_rn(q_value.x, q_value.y); + } + } + *reinterpret_cast(a_sh + row * 128 + d) = value; } 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]; + const int token = tid % N_TOKENS; + const int head = head_base + tid / N_TOKENS; + weight_sh[tid] = tid < USED_ROWS && head < n_head + ? weights[(size_t) token * n_head + head] + : 0.0f; } __syncthreads(); @@ -431,16 +449,17 @@ static __global__ void ds4_indexer_score_wmma_q4_kernel( __syncthreads(); int slot = 0; - for (int output = tid; output < 4 * 128; + for (int output = tid; output < N_TOKENS * 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; + for (int head_in_group = 0; + head_in_group < HEADS_PER_TILE; ++head_in_group) { - const int row = 4 * head_in_group + token; + const int row = N_TOKENS * 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]; @@ -450,7 +469,7 @@ static __global__ void ds4_indexer_score_wmma_q4_kernel( } int slot = 0; - for (int output = tid; output < 4 * 128; + for (int output = tid; output < N_TOKENS * 128; output += 256, ++slot) { const int token = output >> 7; const int comp = tile_c + (output & 127); @@ -553,6 +572,17 @@ void ggml_cuda_op_ds4_indexer_score( warp_size == 32 && (!GGML_CUDA_CC_IS_NVIDIA(device_info.cc) || device_info.cc >= GGML_CUDA_CC_VOLTA); + const char * packed_small_name = "GGML_DS4_INDEXER_PACK_SMALL"; + const char * packed_small_env = std::getenv(packed_small_name); + if (!packed_small_env) { + // Backward-compatible alias for the original q=4-only prototype. + packed_small_name = "GGML_DS4_INDEXER_PACK_Q4"; + packed_small_env = std::getenv(packed_small_name); + } + const bool use_packed_small = packed_small_env + ? ds4_env_flag_enabled(packed_small_name) + : GGML_CUDA_CC_IS_RDNA3_5(device_info.cc) || + GGML_CUDA_CC_IS_RDNA4(device_info.cc); #if DS4_INDEXER_WMMA_AVAILABLE if (wmma_capable && n_tokens == 1) { const dim3 grid((unsigned) ((n_comp + 127) / 128), 1, 1); @@ -564,10 +594,39 @@ void ggml_cuda_op_ds4_indexer_score( visibility_mask ? static_cast(visibility_mask->data) : nullptr, n_comp, kv_start, n_head, ratio); - } else if (wmma_capable && n_tokens == 4 && n_head % 4 == 0 && - ds4_env_flag_enabled("GGML_DS4_INDEXER_PACK_Q4")) { + } else if (wmma_capable && n_tokens == 2 && use_packed_small) { + const dim3 grid((unsigned) ((n_comp + 127) / 128), 1, 1); + ds4_indexer_score_wmma_small_kernel<2><<>>( + static_cast(dst->data), + static_cast(q->data), + static_cast(weights->data), + static_cast(comp->data), + visibility_mask + ? static_cast(visibility_mask->data) : nullptr, + n_comp, kv_start, n_head, ratio); + } else if (wmma_capable && n_tokens == 3 && use_packed_small) { + const dim3 grid((unsigned) ((n_comp + 127) / 128), 1, 1); + ds4_indexer_score_wmma_small_kernel<3><<>>( + static_cast(dst->data), + static_cast(q->data), + static_cast(weights->data), + static_cast(comp->data), + visibility_mask + ? static_cast(visibility_mask->data) : nullptr, + n_comp, kv_start, n_head, ratio); + } else if (wmma_capable && n_tokens == 4 && use_packed_small) { + const dim3 grid((unsigned) ((n_comp + 127) / 128), 1, 1); + ds4_indexer_score_wmma_small_kernel<4><<>>( + static_cast(dst->data), + static_cast(q->data), + static_cast(weights->data), + static_cast(comp->data), + visibility_mask + ? static_cast(visibility_mask->data) : nullptr, + n_comp, kv_start, n_head, ratio); + } else if (wmma_capable && n_tokens == 5 && use_packed_small) { const dim3 grid((unsigned) ((n_comp + 127) / 128), 1, 1); - ds4_indexer_score_wmma_q4_kernel<<>>( + ds4_indexer_score_wmma_small_kernel<5><<>>( static_cast(dst->data), static_cast(q->data), static_cast(weights->data), diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index 78215b4d3..213968207 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -103,6 +103,7 @@ struct ds4_inverse_rope_params { int enabled; int forward_q_enabled; int kv_start; + const int32_t * positions; float freq_scale; float ext_factor; float attn_factor; @@ -157,7 +158,8 @@ __device__ static __forceinline__ void ds4_inverse_rope_coefficients( const ds4_inverse_rope_params & p, float & cos_theta, float & sin_theta) { ds4_rope_coefficients_at_position( - pair, -(p.kv_start + token), p, cos_theta, sin_theta); + pair, -(p.positions ? p.positions[token] : p.kv_start + token), + p, cos_theta, sin_theta); } // Forward counterpart of ds4_inverse_rope_coefficients. Keep the expressions @@ -169,7 +171,8 @@ __device__ static __forceinline__ void ds4_forward_rope_coefficients( const ds4_inverse_rope_params & p, float & cos_theta, float & sin_theta) { ds4_rope_coefficients_at_position( - pair, p.kv_start + token, p, cos_theta, sin_theta); + pair, p.positions ? p.positions[token] : p.kv_start + token, + p, cos_theta, sin_theta); } __device__ static __forceinline__ void ds4_apply_inverse_rope_pair( @@ -478,7 +481,7 @@ __global__ static void ds4_fa_indexed_rows_parallel_kernel( // 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 +template __global__ static void ds4_fa_indexed_rows_topk_kernel( const Mask * mask, const int32_t * topk, @@ -494,7 +497,6 @@ __global__ static void ds4_fa_indexed_rows_topk_kernel( 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]; @@ -573,6 +575,26 @@ __global__ static void ds4_fa_indexed_rows_topk_kernel( } } +template +static void ds4_launch_indexed_rows_topk( + 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, + cudaStream_t stream) { + // The learned top-512 stays on its original launch. A batched verifier + // appends a small saved-raw suffix and needs the next sorting bucket. + if (capacity <= 512) { + ds4_fa_indexed_rows_topk_kernel<<>>( + mask, topk, selected_rows, selected_counts, owner_offsets, owner_ranks, + n_tokens, n_kv, raw_rows, capacity); + } else { + ds4_fa_indexed_rows_topk_kernel<<>>( + mask, topk, selected_rows, selected_counts, owner_offsets, owner_ranks, + n_tokens, n_kv, raw_rows, capacity); + } +} + template __global__ static void ds4_flash_attn_d512_shared_kv_kernel( float * dst, @@ -1777,6 +1799,14 @@ 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); + const ggml_tensor * rope_positions = dst->src[6]; + if (rope_positions && + ((rope_flags & 1) == 0 || rope_positions->type != GGML_TYPE_I32 || + rope_positions->ne[0] != n_tokens || rope_positions->ne[1] != 1 || + rope_positions->ne[2] != 1 || rope_positions->ne[3] != 1 || + !ggml_is_contiguous(rope_positions))) { + return false; + } if (sparse_keep_rows == INT_MIN) { return false; } @@ -1789,7 +1819,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(const ggml_tensor * dst) } const int n_comp_rows = n_kv - raw_rows; if (indexer_topk && - (sparse_keep_rows >= 0 || -sparse_keep_rows > 512 || + (sparse_keep_rows >= 0 || -sparse_keep_rows > 1024 || -sparse_keep_rows > n_comp_rows || indexer_topk->type != GGML_TYPE_I32 || indexer_topk->ne[0] != -sparse_keep_rows || @@ -1868,6 +1898,8 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( inverse_rope.forward_q_enabled = (rope_flags & 2) != 0; if (rope_flags != 0) { inverse_rope.kv_start = ggml_get_op_params_i32(dst, 8); + inverse_rope.positions = dst->src[6] + ? static_cast(dst->src[6]->data) : nullptr; const float freq_base = ggml_get_op_params_f32(dst, 9); inverse_rope.freq_scale = ggml_get_op_params_f32(dst, 10); inverse_rope.ext_factor = ggml_get_op_params_f32(dst, 11); @@ -1994,12 +2026,12 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( getenv("GGML_DS4_FA_SERIAL_INDEX_SCAN") == nullptr; if (mask->type == GGML_TYPE_F16) { if (indexer_topk) { - ds4_fa_indexed_rows_topk_kernel<<>>( + ds4_launch_indexed_rows_topk( (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); + n_tokens, n_kv, raw_rows, indexed_capacity, stream); } else if (parallel_index_scan) { ds4_fa_indexed_rows_parallel_kernel<<>>( (const half *) mask->data, indexed_rows, indexed_counts, @@ -2015,12 +2047,12 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32( } } else { if (indexer_topk) { - ds4_fa_indexed_rows_topk_kernel<<>>( + ds4_launch_indexed_rows_topk( (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); + n_tokens, n_kv, raw_rows, indexed_capacity, stream); } else if (parallel_index_scan) { ds4_fa_indexed_rows_parallel_kernel<<>>( (const float *) mask->data, indexed_rows, indexed_counts, diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index a2a92e40d..0f8cca79e 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -5668,6 +5668,19 @@ void ggml_flash_attn_ext_set_ds4_inverse_rope( ggml_set_op_params_i32(a, 15, n_ctx_orig); } +void ggml_flash_attn_ext_set_ds4_rope_positions( + struct ggml_tensor * a, + struct ggml_tensor * positions) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT((ggml_get_op_params_i32(a, 7) & 1) != 0); + GGML_ASSERT(a->src[6] == NULL); + GGML_ASSERT(positions && positions->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(positions)); + GGML_ASSERT(positions->ne[0] == a->src[0]->ne[1]); + GGML_ASSERT(positions->ne[1] == 1 && positions->ne[2] == 1 && positions->ne[3] == 1); + a->src[6] = positions; +} + bool ggml_flash_attn_ext_is_ds4(const struct ggml_tensor * a) { return a && a->op == GGML_OP_FLASH_ATTN_EXT && (ggml_get_op_params_i32(a, 6) != 0 || diff --git a/server/docs/DS4.md b/server/docs/DS4.md index 9e181243f..49e62bc9c 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -547,6 +547,8 @@ Run the converted drafter against a DeepSeek4 target with: ```bash export DFLASH_DS4_SPEC=1 export DFLASH_DS4_FUSED_VERIFY=1 +# Experimental, single HIP target only; may change generated tokens: +# export DFLASH_DS4_SPARSE_DECODE_FLASH=1 export DFLASH_DS4_DRAFT=/path/to/dflash-draft.gguf export DFLASH_DS4_SPEC_Q=4 @@ -557,12 +559,39 @@ export DFLASH_DS4_SPEC_Q=4 ``` `--ds4-fused-verify-f16-kv` feeds the persistent F16 MLA cache directly to -batched explicit verifier attention instead of converting the full cache to -F32 on every speculative step. Key-side accumulation remains F32 through 512 -attention rows to preserve the short-context quality baseline. The option is -currently qualified only for a single HIP target and remains off by default. -It changes verifier floating-point inputs and can change generated tokens, so -re-run workload quality checks before enabling it for another checkpoint. +batched explicit or sparse verifier attention instead of converting the full +cache to F32 on every speculative step. With +`DFLASH_DS4_SPARSE_DECODE_FLASH=1`, the verifier keeps explicit attention for +short histories. Single-lane or ratio-4 layouts switch to sparse attention +once it removes at least half of the compressed rows. Other batched layouts +retain explicit attention because coarse block selection does not preserve +the overwritten raw-row suffix. The ratio-4 learned indexer appends that +suffix to its selected rows and applies the causal mask to it. Key-side +accumulation remains F32 through 512 attention rows to preserve the +short-context quality baseline. The option is currently qualified only for a single HIP target and +remains off by default. It changes verifier floating-point inputs and can +change generated tokens, so re-run workload quality checks before enabling it +for another checkpoint. + +Cached sparse decode and verification retain standalone forward-Q RoPE. +The fused-Q variant failed the strict five-line retrieval control even when +an explicit-attention control used the same selected rows. Keeping this +rotation separate restored the control's output without disabling sparse +flash attention, native F16 KV, or adaptive verification. Fused inverse RoPE +reads runtime token positions so graph replay cannot retain a previous +step's position. Uncached prefill keeps both fused rotations. This targeted +check does not establish universal output identity or qualify all contexts; +the sparse verifier remains opt-in. + +Cached batched verification also gathers overwritten raw rows using the +current runtime ring indices before updating the cache. Those saved rows +remain in the native cache dtype and in the ratio-4 sparse selection, so ring +wraps do not require falling back to full-history explicit attention. + +On RDNA3.5 and RDNA4, speculative widths 2–5 use the packed small-CM rocWMMA +indexer by default. It is bit-identical to the generic indexer in the GPU unit +test and can be disabled with `GGML_DS4_INDEXER_PACK_SMALL=0` for diagnosis. +The legacy `GGML_DS4_INDEXER_PACK_Q4` variable remains an alias. `DFLASH_DS4_FUSED_VERIFY=1` is the opt-in throughput profile. Its persistent whole-model GPU graph uses stable padded reduction shapes, so near-tied greedy @@ -780,6 +809,39 @@ at 32.12 tok/s. These numbers are workload-specific; the confidence policy is enabled only when DSpark is explicitly enabled and the draft artifact contains a compatible confidence head. +## PFlash prompt compression + +DeepSeek4 supports the shared in-process Qwen3-0.6B PFlash scorer. + +PFlash is supported for monolithic serving only. Layer-split and paged-serving +configurations reject prefill compression at startup; paged serving cannot +park a target while it owns live sequence state. + +```bash +./server/build-hip/dflash_server /path/to/deepseek4-target.gguf \ + --target-device hip:0 \ + --prefill-compression auto \ + --prefill-drafter /path/to/Qwen3-0.6B-BF16.gguf \ + --prefill-skip-park +``` + +The HTTP path converts target tokens to text, scores Qwen tokens, decodes the +kept Qwen spans, and tokenizes that text for DeepSeek4. This cross-tokenizer +round trip is required; drafter token IDs are never passed directly to the +target. Omit `--prefill-skip-park` when the target, DSpark drafter, and PFlash +drafter do not fit together. + +PFlash reduces TTFT and the effective context used during generation, but it +is lossy prompt compression. Disable it for matched true-context throughput +or exact-retrieval comparisons. + +The legacy daemon command `compress +[nopark]` retains its shared wire contract: backend-local GPU 0 and a drafter +kept loaded until `free drafter` or parking. That protocol has no GPU-placement +or request-scoped-residency fields. Use the HTTP path or typed `compress` / +`compress_batch` API when those controls are required. Typed keep ratios must +be finite and in [0,1]; zero requests the scorer's minimum retained chunk. + ## Example: CUDA + Halo Layer Split Automatic split (CUDA prefix chosen from free memory, optional manual override via `DFLASH_DS4_CUDA_LAYERS`): diff --git a/server/src/common/feature_gate.cpp b/server/src/common/feature_gate.cpp index f8cf94318..987b923e9 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -52,6 +52,10 @@ std::string check_feature_compatibility( !features.pflash_drafter_configured) { return "--prefill-compression requires --prefill-drafter"; } + if (features.pflash_enabled && arch == "deepseek4" && + args.device.is_layer_split()) { + return "--prefill-compression is not supported with DeepSeek4 layer splitting"; + } // ── target/draft backend mixing × remote draft IPC if (mixed_draft_placement && !args.remote_draft.enabled()) { diff --git a/server/src/common/model_capabilities.h b/server/src/common/model_capabilities.h index d1d541312..fdb4caa1b 100644 --- a/server/src/common/model_capabilities.h +++ b/server/src/common/model_capabilities.h @@ -77,7 +77,7 @@ inline constexpr ArchCapabilities kArchCapabilities[] = { {"laguna", true, false, false, true, kMono, kMono, kMono, kNever, kNever,kNever, kNever}, {"qwen3", false, false, true, false, kNever, kNever, kNever, kNever, kNever,kNever, kNever}, {"gemma4", true, false, false, false, kMono, kNever, kNever, kNever, kBoth,kNever, kNever}, - {"deepseek4", true, false, false, false, kNever, kNever, kNever, kNever, kNever,kNever, kMono}, + {"deepseek4", true, false, true, false, kNever, kNever, kNever, kNever, kNever,kNever, kMono}, }; inline constexpr std::size_t kArchCount = diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 2fb74a01e..cc549ea67 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -4,9 +4,11 @@ #include "deepseek4_backend.h" #include "deepseek4_budget_hook.h" #include "deepseek4_internal.h" +#include "dflash27b.h" #include "deepseek4_snapshot.h" #include "deepseek4_page_layout.h" #include "common/dynamic_backend.h" +#include "common/io_utils.h" #include "common/peer_access.h" #include "common/platform_env.h" #include "common/sampler.h" @@ -29,6 +31,7 @@ #include #include #include +#include namespace dflash::common { @@ -1851,6 +1854,11 @@ bool DeepSeek4Backend::park(ParkTarget target) { std::printf("[deepseek4] DSpark drafter parked (VRAM released)\n"); std::fflush(stdout); } + if (want_draft && pflash_drafter_loaded_) { + release_pflash_drafter(); + std::printf("[deepseek4] PFlash drafter parked (VRAM released)\n"); + std::fflush(stdout); + } if (!want_target_model || parked_) return true; if (cfg_.paged_attention) { std::fprintf(stderr, @@ -2861,17 +2869,154 @@ GenerateResult DeepSeek4Backend::restore_and_generate_impl( return result; } +ModelBackend::CompressResult DeepSeek4Backend::compress( + const CompressRequest & req) { + const auto results = compress_batch({req}); + return results.empty() ? CompressResult{} : results.front(); +} + +std::vector DeepSeek4Backend::compress_batch( + const std::vector & requests) { + std::vector results(requests.size()); + if (requests.empty()) return results; + + const auto valid_request = [](const CompressRequest & request) { + return !request.input_ids.empty() && !request.drafter_path.empty() && + std::isfinite(request.keep_ratio) && + request.keep_ratio >= 0.0f && request.keep_ratio <= 1.0f; + }; + const CompressRequest * load_request = nullptr; + for (const CompressRequest & request : requests) { + if (!valid_request(request)) continue; + if (load_request == nullptr) { + load_request = &request; + } else if (request.drafter_path != load_request->drafter_path || + request.drafter_gpu != load_request->drafter_gpu || + request.skip_park != load_request->skip_park || + request.residency_action != load_request->residency_action) { + // A residency window can host only one drafter configuration. + // Process heterogeneous batches as independent windows instead of + // calling the virtual base fallback, which would recurse through + // DeepSeek4Backend::compress(). + std::vector independent(requests.size()); + for (size_t index = 0; index < requests.size(); ++index) { + const auto one = compress_batch({requests[index]}); + if (!one.empty()) independent[index] = one.front(); + } + return independent; + } + } + if (load_request == nullptr) return results; + + // Parking releases target/cache buffers, including the expert backend. + // Drain their queued work before releasing any of those dependencies. + if (backend_) ggml_backend_synchronize(backend_); + if (spec_backend_) ggml_backend_synchronize(spec_backend_); + if (expert_backend_) ggml_backend_synchronize(expert_backend_); + const bool was_parked = parked_; + if (!load_request->skip_park && !parked_ && + !park(ParkTarget::TargetModel)) { + return results; + } + if (pflash_drafter_loaded_ && + (pflash_drafter_path_ != load_request->drafter_path || + pflash_drafter_gpu_ != load_request->drafter_gpu)) { + release_pflash_drafter(); + } + if (!pflash_drafter_loaded_) { + if (!load_drafter(load_request->drafter_path, 999, + load_request->drafter_gpu, + pflash_drafter_ctx_)) { + std::fprintf(stderr, "[deepseek4-pflash] load failed: %s\n", + dflash27b_last_error()); + release_pflash_drafter(); + if (!load_request->skip_park && !was_parked) { + unpark(ParkTarget::TargetModel); + } + return results; + } + pflash_drafter_loaded_ = true; + pflash_drafter_path_ = load_request->drafter_path; + pflash_drafter_gpu_ = load_request->drafter_gpu; + } + + for (size_t index = 0; index < requests.size(); ++index) { + const CompressRequest & request = requests[index]; + if (!valid_request(request)) continue; + CompressResult & result = results[index]; + result.compressed_ids = drafter_score_and_compress( + pflash_drafter_ctx_, request.input_ids, request.keep_ratio); + result.ok = !result.compressed_ids.empty(); + } + + if (load_request->residency_action == + DraftResidencyAction::ReleaseAfterUse) { + release_pflash_drafter(); + } + if (!load_request->skip_park && !was_parked && + !unpark(ParkTarget::TargetModel)) { + std::fill(results.begin(), results.end(), CompressResult{}); + } + return results; +} + bool DeepSeek4Backend::handle_compress(const std::string & line, - const DaemonIO & io) { - (void)line; (void)io; - std::fprintf(stderr, "[deepseek4] compress not yet supported\n"); - return false; + const DaemonIO & io) { + // Legacy wire format has no GPU/residency fields: retain backend-local + // GPU 0 and KeepLoaded. HTTP uses the typed API for those controls. + std::istringstream iss(line.size() > 9 ? line.substr(9) : std::string{}); + std::string prompt_path; + std::string drafter_path; + int keep_x1000 = 0; + if (!(iss >> prompt_path >> keep_x1000)) { + std::fprintf(stderr, "[deepseek4-pflash] bad compress arguments\n"); + io.emit(-1); + return false; + } + + std::getline(iss >> std::ws, drafter_path); + bool skip_park = false; + const std::string suffix = " nopark"; + if (drafter_path.size() > suffix.size() && + drafter_path.compare(drafter_path.size() - suffix.size(), + suffix.size(), suffix) == 0) { + skip_park = true; + drafter_path.resize(drafter_path.size() - suffix.size()); + } + + CompressRequest req; + req.input_ids = read_int32_file(prompt_path); + req.keep_ratio = (float) keep_x1000 / 1000.0f; + req.drafter_path = std::move(drafter_path); + req.skip_park = skip_park; + CompressResult result = compress(req); + if (!result.ok) { + std::fprintf(stderr, "[deepseek4-pflash] compression failed\n"); + io.emit(-1); + return false; + } + + std::printf("[deepseek4-pflash] %zu -> %zu tokens\n", + req.input_ids.size(), result.compressed_ids.size()); + std::fflush(stdout); + for (int32_t token : result.compressed_ids) io.emit(token); + io.emit(-1); + return true; +} + +void DeepSeek4Backend::release_pflash_drafter() { + // A failed load can own a backend even before the loaded flag is set. + dflash::common::free_drafter(pflash_drafter_ctx_); + pflash_drafter_loaded_ = false; + pflash_drafter_path_.clear(); + pflash_drafter_gpu_ = -1; } void DeepSeek4Backend::free_drafter() { // Keep the configured path so request-scoped residency and an explicit // later `unpark draft` can restore the DSpark model. release_spec_drafter(/*mark_parked=*/true); + release_pflash_drafter(); } void DeepSeek4Backend::maybe_save_routing_stats() { diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 8e71c479f..b09dda136 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -15,6 +15,7 @@ #include "../common/moe_hybrid_stream.h" #include "deepseek4_internal.h" #include "deepseek4_dspark.h" +#include "qwen3/qwen3_drafter.h" #include "deepseek4_seq_engine.h" #include "ggml.h" @@ -78,6 +79,9 @@ class DeepSeek4Backend : public ModelBackend { const GenerateRequest & req, const DaemonIO & io) override; + CompressResult compress(const CompressRequest & req) override; + std::vector compress_batch( + const std::vector & requests) override; bool handle_compress(const std::string & line, const DaemonIO & io) override; void free_drafter() override; @@ -125,12 +129,17 @@ class DeepSeek4Backend : public ModelBackend { ggml_backend_t spec_backend_ = nullptr; std::unique_ptr spec_drafter_; std::vector spec_feat_window_; + DrafterContext pflash_drafter_ctx_; + bool pflash_drafter_loaded_ = false; + std::string pflash_drafter_path_; + int pflash_drafter_gpu_ = -1; // Once a long prompt selects the fragmentation-safe prefill shape, retain // it for later requests so the HIP arenas never switch back under load. int hybrid_prefill_chunk_cap_ = 0; bool load_spec_drafter(); void release_spec_drafter(bool mark_parked); + void release_pflash_drafter(); void keep_spec_feature_tail(std::vector & features, size_t max_rows) const; // True when a wide prefill path returns per-token DSpark features and the diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index 7d7ec2dae..5c38d5ca1 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -475,6 +475,8 @@ static bool ds4_build_fused_verify_graph( ex.pos_q = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.pos_q); ex.neg_q = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.neg_q); ex.rawrows = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, q); ggml_set_input(ex.rawrows); + ex.saved_rawrows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); + ggml_set_input(ex.saved_rawrows); ex.ape4 = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.ape4); 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); @@ -833,6 +835,7 @@ static bool ds4_build_fused_verify_graph( ain.rope_pos = ex.pos_q; ain.neg_pos = ex.neg_q; ain.raw_kv_rows = ex.rawrows; + ain.preserved_raw_rows = ex.saved_rawrows; if (ratio > 0) { int n_flush = 0; for (int t = 0; t < lane_q; ++t) { @@ -881,13 +884,34 @@ static bool ds4_build_fused_verify_graph( std::vector i32b; std::vector i32ab; std::vector i64ab; + std::vector f32ab; + // Sparse attention pays a fixed indexer/compaction cost. Keep the + // established explicit kernel while history is small, then switch per + // layer once sparse work can remove at least half of the compressed + // rows. This also avoids a transient identity binding at the exact + // top-k boundary, preserving fused-graph replay. + const bool sparse_attention = + ds4_env_flag("DFLASH_DS4_SPARSE_DECODE_FLASH") && ratio > 0 && + // The exact indexer carries saved raw rows for batched lanes. + // Coarse block selection has no saved-row contract; retain the + // explicit verifier there instead of pruning raw history. + (lane_q == 1 || ratio == 4) && + padded > 2 * w.n_indexer_top_k; + const DeepSeek4AttentionImpl attention_impl = sparse_attention + ? DeepSeek4AttentionImpl::SparseFlash + : DeepSeek4AttentionImpl::Explicit; ggml_tensor * normed = build_rms_norm(ctx, attn_in, L.attn_norm, w.rms_eps); attn_out = build_mla_attention(ctx, gf, normed, w, L, *lc, il, lane_kv_start, lane_q, &ain, - i32b, i32ab, i64ab); + i32b, i32ab, i64ab, + &f32ab, attention_impl); if (!attn_out) return false; - if (!i32b.empty() || !i32ab.empty() || !i64ab.empty()) { - std::fprintf(stderr, "[ds4-fused-verify] layer %d dynamic bindings; cannot fuse\n", il); + if (!i32b.empty() || !i32ab.empty() || !i64ab.empty() || + !f32ab.empty()) { + std::fprintf(stderr, + "[ds4-fused-verify] layer %d dynamic bindings " + "(%zu/%zu/%zu/%zu); cannot fuse\n", + il, i32b.size(), i32ab.size(), i64ab.size(), f32ab.size()); return false; } } @@ -1251,6 +1275,7 @@ static bool ds4_build_fused_verify_graph( pin_main(ex.pos_q); pin_main(ex.neg_q); pin_main(ex.rawrows); + pin_main(ex.saved_rawrows); pin_main(ex.ape4); pin_main(ex.ape128); pin_main(ex.st4); @@ -1523,6 +1548,8 @@ static int ds4_try_fused_verify_step( ds4_fv_set(ex->neg_q, iv.data(), sizeof(int32_t) * q); for (int i = 0; i < q; ++i) lv[(size_t) i] = (kv_start + i) % w.n_swa; ds4_fv_set(ex->rawrows, lv.data(), sizeof(int64_t) * q); + for (int i = 0; i < q; ++i) iv[(size_t) i] = (int32_t) lv[(size_t) i]; + ds4_fv_set(ex->saved_rawrows, iv.data(), sizeof(int32_t) * q); for (int i = 0; i < q; ++i) iv[(size_t) i] = (kv_start + i) % 4; ds4_fv_set(ex->ape4, iv.data(), sizeof(int32_t) * q); for (int i = 0; i < q; ++i) lv[(size_t) i] = 4 + (kv_start + i) % 4; diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 92438dbe6..74b7175a0 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -48,6 +48,42 @@ namespace dflash::common { +ggml_tensor * deepseek4_preserve_raw_rows( + ggml_context * ctx, ggml_tensor * raw_kv, ggml_tensor * rows) { + GGML_ASSERT(raw_kv && ggml_is_matrix(raw_kv)); + GGML_ASSERT(raw_kv->type == GGML_TYPE_F16 || raw_kv->type == GGML_TYPE_F32); + GGML_ASSERT(rows && rows->type == GGML_TYPE_I32 && ggml_is_vector(rows)); + GGML_ASSERT(rows->ne[0] > 0 && rows->ne[0] <= raw_kv->ne[1]); + // Cached graphs advance around the ring without changing their topology. + // Read the same runtime indices used by the upcoming set_rows, rather + // than baking the first step's physical offsets into view nodes. + auto * saved = ggml_get_rows(ctx, raw_kv, rows); + // GET_ROWS returns F32. Round-trip only this q-row suffix to retain native + // F16 verification; do not convert the full raw/compressed cache. + return raw_kv->type == GGML_TYPE_F16 + ? ggml_cast(ctx, saved, GGML_TYPE_F16) : saved; +} + +ggml_tensor * deepseek4_indexed_attention_rows( + ggml_context * ctx, ggml_tensor * compressed_topk, + int compressed_rows, int preserved_rows) { + GGML_ASSERT(compressed_topk && compressed_topk->type == GGML_TYPE_I32); + GGML_ASSERT(compressed_rows >= 0 && preserved_rows >= 0); + if (preserved_rows == 0) return compressed_topk; + // ARANGE uses F32, so its integer endpoints must be exactly representable. + GGML_ASSERT((int64_t) compressed_rows + preserved_rows <= (1 << 24)); + auto * saved = ggml_cast(ctx, ggml_arange( + ctx, (float) compressed_rows, (float) (compressed_rows + preserved_rows), + 1.0f), GGML_TYPE_I32); + auto * shape = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, preserved_rows, compressed_topk->ne[1]); + saved = ggml_repeat(ctx, saved, shape); + // The suffix is not part of the learned top-k competition. Append all of + // it to each lane's row list; the existing causal mask hides overwritten, + // not-yet-written and future rows. No host binding is needed on replay. + return ggml_concat(ctx, compressed_topk, saved, 0); +} + namespace { using Ds4TimingClock = std::chrono::steady_clock; @@ -334,6 +370,7 @@ struct DeepSeek4AttentionGraphInputs { ggml_tensor * rope_pos = nullptr; ggml_tensor * neg_pos = nullptr; ggml_tensor * raw_kv_rows = nullptr; + ggml_tensor * preserved_raw_rows = nullptr; // I32 runtime ring read indices ggml_tensor * attn_ape_row = nullptr; ggml_tensor * attn_state_rows = nullptr; ggml_tensor * attn_comp_rows = nullptr; @@ -1589,6 +1626,20 @@ static int ds4_padded_gathered_raw_rows(int n_raw) { return std::min(padded, (int) DS4_PAGE_TOKENS - 1); } +ggml_tensor * deepseek4_indexer_visibility_suffix( + ggml_context * ctx, ggml_tensor * mask, int first_scored, int n_scored) { + if (!mask) return nullptr; + GGML_ASSERT(mask->type == GGML_TYPE_F32 && ggml_is_matrix(mask)); + GGML_ASSERT(first_scored >= 0 && n_scored > 0); + GGML_ASSERT(mask->ne[1] == (int64_t) first_scored + n_scored); + if (first_scored == 0) return mask; + // Query, head weights, positions and per-token visibility must all start + // at the same lane after skipping the identity-selected prefix. + return ggml_cont(ctx, ggml_view_2d( + ctx, mask, mask->ne[0], n_scored, mask->nb[1], + (size_t) first_scored * mask->nb[1])); +} + static ggml_tensor * build_indexer_topk( ggml_context * ctx, ggml_tensor * qr_norm, // [n_lora_q, n_tokens] @@ -1629,6 +1680,8 @@ static ggml_tensor * build_indexer_topk( }; qr_norm = token_slice(qr_norm, (int) qr_norm->ne[0]); cur = token_slice(cur, (int) cur->ne[0]); + visibility_mask = deepseek4_indexer_visibility_suffix( + ctx, visibility_mask, first_scored, n_scored); if (first_scored > 0) { rope_pos = ggml_view_1d( ctx, rope_pos, n_scored, @@ -1989,7 +2042,13 @@ static ggml_tensor * build_mla_attention_lane_core( // D=512 flash prefill can rotate Q's 64-d tail inside the exact attention // kernel. This avoids materializing cont(nope), cont(tail), rope(tail), // and concat(nope, tail) while retaining the same F32 rounding boundary. - const bool fuse_q_rope = attention_impl != DeepSeek4AttentionImpl::Explicit && + // Cached decode/verification uses the standalone Q rotation. Fusing it + // changes the adaptive sparse verifier's output even with identical + // selected rows; keep the validated rounding/materialization boundary. + // This does not disable sparse flash attention, native F16 KV, inverse + // RoPE fusion, or the uncached prefill optimization. + const bool fuse_q_rope = !cached_inputs && + attention_impl != DeepSeek4AttentionImpl::Explicit && n_tokens > 1 && head_dim == 512 && n_rot == 64; if (prepared) { projected = *prepared; @@ -2021,16 +2080,12 @@ static ggml_tensor * build_mla_attention_lane_core( if (!gathered_history && fused_causal) { // Fused verify: ALWAYS q preserved rows so the topology is stable; // unwrapped/garbage rows are masked by the host-filled mask values. - for (int ti = 0; ti < n_tokens; ti++) { - ggml_tensor * slot = ggml_view_2d( - ctx, lane.raw_kv, head_dim, 1, lane.raw_kv->nb[1], - (size_t)((kv_start + ti) % w.n_swa) * lane.raw_kv->nb[1]); - ggml_tensor * saved = ggml_cont(ctx, slot); - ggml_build_forward_expand(gf, saved); - old_rows_scratch = old_rows_scratch - ? ggml_concat(ctx, old_rows_scratch, saved, 1) : saved; - n_old_rows++; - } + GGML_ASSERT(cached_inputs->preserved_raw_rows && + cached_inputs->preserved_raw_rows->ne[0] == n_tokens); + old_rows_scratch = deepseek4_preserve_raw_rows( + ctx, lane.raw_kv, cached_inputs->preserved_raw_rows); + ggml_build_forward_expand(gf, old_rows_scratch); + n_old_rows = n_tokens; old_rows_scratch_f16 = old_rows_scratch; old_rows_scratch = ds4_cast_if_needed(ctx, old_rows_scratch, GGML_TYPE_F32); } else if (!gathered_history && causal_batch && !layer_major_batch) { @@ -2227,11 +2282,13 @@ static ggml_tensor * build_mla_attention_lane_core( ? cached_inputs->padded_comp : n_index_comp_live; if (masked_kv && n_index_comp > 0) { - index_visibility_mask = ggml_view_2d( + // Each verifier lane owns a full causal-mask column. Preserve + // the per-lane compressed visibility when compacting it. + index_visibility_mask = ggml_cont(ctx, ggml_view_2d( ctx, cached_inputs->attn_row_mask, - n_index_comp, 1, - (size_t) n_index_comp * sizeof(float), - (size_t) w.n_swa * sizeof(float)); + n_index_comp, n_tokens, + cached_inputs->attn_row_mask->nb[1], + (size_t) w.n_swa * sizeof(float))); } } indexer_topk = build_indexer_topk( @@ -2318,19 +2375,24 @@ static ggml_tensor * build_mla_attention_lane_core( } else { kv_attn = raw_kv_view(0, n_raw); } - const bool fused_explicit_f16_kv = w.fused_verify_f16_kv && + const bool fused_verify_f16_kv = w.fused_verify_f16_kv && masked_kv && n_tokens > 1 && - attention_impl == DeepSeek4AttentionImpl::Explicit && kv_attn->type == GGML_TYPE_F32 && raw_kv_source->type == GGML_TYPE_F16 && (!comp_history_source || comp_history_source->type == GGML_TYPE_F16) && (!old_rows_scratch_f16 || old_rows_scratch_f16->type == GGML_TYPE_F16); - if (fused_explicit_f16_kv) { + const bool fused_explicit_f16_kv = fused_verify_f16_kv && + attention_impl == DeepSeek4AttentionImpl::Explicit; + const bool fused_sparse_f16_kv = fused_verify_f16_kv && + attention_impl == DeepSeek4AttentionImpl::SparseFlash; + if (fused_explicit_f16_kv || fused_sparse_f16_kv) { // DS4's persistent MLA caches are already F16. Feed those tensors - // directly to the established explicit attention matmuls instead of - // casting the entire long-context cache to F32 on every verifier step. + // directly to the attention implementation instead of casting the + // entire long-context cache to F32 on every verifier step. + // Current writes are consumed through their set_rows results, while + // preserved overwritten rows retain the same cached F16 values. kv_attn = ggml_view_2d( ctx, raw_kv_source, head_dim, n_raw, raw_kv_source->nb[1], 0); if (n_comp_attn > 0 && comp_history_source) { @@ -2344,10 +2406,14 @@ static ggml_tensor * build_mla_attention_lane_core( ctx, kv_attn, old_rows_scratch_f16, 1); } static std::atomic explicit_f16_kv_logged{false}; - if (!explicit_f16_kv_logged.exchange(true)) { + static std::atomic sparse_f16_kv_logged{false}; + std::atomic & logged = fused_sparse_f16_kv + ? sparse_f16_kv_logged : explicit_f16_kv_logged; + if (!logged.exchange(true)) { std::fprintf(stderr, - "[deepseek4] fused explicit F16 K/V active: tokens=%d " + "[deepseek4] fused %s F16 K/V active: tokens=%d " "compressed=%d\n", + fused_sparse_f16_kv ? "sparse" : "explicit", n_tokens, n_comp_attn); } } else { @@ -2447,6 +2513,10 @@ static ggml_tensor * build_mla_attention_lane_core( score_mask = ggml_reshape_2d(ctx, cmask, n_attn, n_tokens); } } + if (indexer_topk) { + indexer_topk = deepseek4_indexed_attention_rows( + ctx, indexer_topk, n_comp_attn, n_old_rows); + } const bool direct_indexer_topk = indexer_topk && ds4_env_flag("DFLASH_DS4_DIRECT_INDEXER_TOPK"); if (indexer_topk) { @@ -2605,7 +2675,12 @@ static ggml_tensor * build_mla_attention_lane_core( // The DS4 D=512 kernel consumes Q strides directly, avoiding a full // [D,H,T] -> [D,T,H] materialization for every layer. ggml_tensor * q_fa = ggml_permute(ctx, q, 0, 2, 1, 3); - ggml_tensor * kv_fa = ds4_cast_if_needed(ctx, kv_attn, GGML_TYPE_F32); + // The DS4 D=512 kernel has native F16 K/V specializations. Keep + // fused verifier caches in their persistent representation and + // avoid a full long-context F16 -> F32 conversion every step. + ggml_tensor * kv_fa = fused_sparse_f16_kv + ? kv_attn + : ds4_cast_if_needed(ctx, kv_attn, GGML_TYPE_F32); ggml_tensor * k_fa = ggml_reshape_3d(ctx, kv_fa, head_dim, n_attn, 1); ggml_tensor * v_fa = k_fa; ggml_tensor * mask_fa = score_mask @@ -2623,7 +2698,7 @@ static ggml_tensor * build_mla_attention_lane_core( ggml_flash_attn_ext_set_ds4_sparse( context, n_raw, w.n_swa, indexer_topk - ? -w.n_indexer_top_k + ? -(int) indexer_topk->ne[0] : attention_impl == DeepSeek4AttentionImpl::SparseFlash ? w.n_indexer_top_k : 0, 32); @@ -2637,6 +2712,13 @@ static ggml_tensor * build_mla_attention_lane_core( context, kv_start, rope_freq, rope_scale, rope_ext, rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig, fuse_q_rope); + // Cached AR/verifier graphs reuse a shape at new positions. + // Bind the already-uploaded position tensor so both fused + // rotations advance with the graph instead of using kv_start + // from the first build. Prefill's fixed-position path is unchanged. + if (cached_inputs) { + ggml_flash_attn_ext_set_ds4_rope_positions(context, rope_pos); + } inverse_rope_fused = true; } } @@ -5246,6 +5328,7 @@ struct Ds4FusedVerifyCache { ggml_tensor * pos_q = nullptr; // i32 [q] ggml_tensor * neg_q = nullptr; // i32 [q] ggml_tensor * rawrows = nullptr; // i64 [1,q] + ggml_tensor * saved_rawrows = nullptr; // i32 [q], gather before ring writes ggml_tensor * ape4 = nullptr; // i32 [q] ggml_tensor * ape128 = nullptr; // i32 [q] ggml_tensor * st4 = nullptr; // i64 [1,q] diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 5e362b04d..c27d2e4dd 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -349,6 +349,22 @@ struct DeepSeek4BackendConfig { // ─── Function declarations ────────────────────────────────────────────── +// Select compressed rows plus the saved raw suffix of a batched verifier. +// Indices are relative to the end of the physical raw ring. Causal visibility +// remains in the attention mask; saving a row does not make it visible to all lanes. +ggml_tensor * deepseek4_indexed_attention_rows( + ggml_context * ctx, ggml_tensor * compressed_topk, + int compressed_rows, int preserved_rows); + +// Snapshot the raw rows a cached verifier is about to overwrite. Expand this +// tensor before the ring writes; row indices are supplied again on each replay. +ggml_tensor * deepseek4_preserve_raw_rows( + ggml_context * ctx, ggml_tensor * raw_kv, ggml_tensor * rows); + +// Keep a per-token indexer visibility mask aligned with the scored suffix. +ggml_tensor * deepseek4_indexer_visibility_suffix( + ggml_context * ctx, ggml_tensor * mask, int first_scored, int n_scored); + bool load_deepseek4_gguf(const std::string & path, ggml_backend_t backend, DeepSeek4Weights & out); diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp index 77740b627..90cd7c43d 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -168,6 +168,27 @@ void test_feature_gate_pflash_requires_drafter_and_supported_arch() { args, "qwen35", PlacementBackend::Cuda, features).empty()); } +void test_feature_gate_ds4_pflash_rejects_layer_split() { + BackendArgs args = gate_args_hip_deepseek4(); + BackendFeatureConfig features; + features.pflash_enabled = true; + features.pflash_drafter_configured = true; + CHECK(gate_result(args, "deepseek4", PlacementBackend::Hip, features).empty()); + args.device.layer_split_gpus = {0, 1}; + args.device.layer_split_backends = {PlacementBackend::Hip, PlacementBackend::Hip}; + CHECK(!gate_result(args, "deepseek4", PlacementBackend::Hip, features).empty()); + features.pflash_enabled = false; + CHECK(gate_result(args, "deepseek4", PlacementBackend::Hip, features).empty()); + + // Current main also supports DS4 paged serving, which cannot park its + // live target state for PFlash. Keep that independent admission guard. + args = gate_args_hip_deepseek4(); + args.paged_attention = true; + CHECK(gate_result(args, "deepseek4", PlacementBackend::Hip, features).empty()); + features.pflash_enabled = true; + CHECK(!gate_result(args, "deepseek4", PlacementBackend::Hip, features).empty()); +} + void test_feature_gate_validates_target_split_topology() { BackendArgs weights; weights.model_path = "/nonexistent/model.gguf"; @@ -714,6 +735,7 @@ void test_model_capability_tables() { CHECK(!arch_has_expert_offload("qwen35")); // deepseek4 is mixture-of-experts but has no hot/cold offload path. CHECK(!arch_has_expert_offload("deepseek4")); + CHECK(arch_supports_pflash_compression("deepseek4")); // Every capability predicate must be false for an architecture the // factory cannot build, so no rule can admit an unbuildable model. @@ -751,6 +773,7 @@ TEST_CASE(FeatureGateFixture, feature_gate_suite) { test_feature_gate_draft_block_size_requires_local_draft(); test_draft_block_size_override_respects_checkpoint_horizon(); test_feature_gate_pflash_requires_drafter_and_supported_arch(); + test_feature_gate_ds4_pflash_rejects_layer_split(); test_feature_gate_validates_target_split_topology(); test_feature_gate_tensor_parallel_requirements(); test_feature_gate_ds4_prefill_requires_deepseek4(); diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index e879a2855..3c38946fe 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -1875,12 +1875,18 @@ static void test_dspark_park_all_releases_drafter() { backend.spec_draft_path_ = "/tmp/ds4-dspark-fixture.gguf"; backend.spec_drafter_ = std::make_unique(); backend.spec_enabled_ = true; + backend.pflash_drafter_loaded_ = true; + backend.pflash_drafter_path_ = "/tmp/ds4-pflash-fixture.gguf"; + backend.pflash_drafter_gpu_ = 1; TEST_ASSERT(backend.park(ParkTarget::All)); TEST_ASSERT(backend.parked_); TEST_ASSERT(backend.spec_drafter_ == nullptr); TEST_ASSERT(!backend.spec_enabled_); TEST_ASSERT(backend.spec_drafter_parked_); + TEST_ASSERT(!backend.pflash_drafter_loaded_); + TEST_ASSERT(backend.pflash_drafter_path_.empty()); + TEST_ASSERT(backend.pflash_drafter_gpu_ == -1); backend.free_drafter(); TEST_ASSERT(backend.spec_drafter_ == nullptr); @@ -1890,6 +1896,162 @@ static void test_dspark_park_all_releases_drafter() { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_pflash_rejects_invalid_requests() { + std::fprintf(stderr, " test_pflash_rejects_invalid_requests ..."); + DeepSeek4BackendConfig cfg; + DeepSeek4Backend backend(cfg); + + ModelBackend::CompressRequest empty; + TEST_ASSERT(!backend.compress(empty).ok); + + ModelBackend::CompressRequest invalid_ratio; + invalid_ratio.input_ids = {1, 2, 3}; + invalid_ratio.keep_ratio = -0.1f; + invalid_ratio.drafter_path = "/nonexistent/drafter.gguf"; + const auto results = backend.compress_batch({empty, invalid_ratio}); + TEST_ASSERT(results.size() == 2); + TEST_ASSERT(!results[0].ok); + TEST_ASSERT(!results[1].ok); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + +static void test_pflash_failed_load_releases_backend() { + std::fprintf(stderr, " test_pflash_failed_load_releases_backend ..."); + DeepSeek4BackendConfig cfg; + DeepSeek4Backend backend(cfg); + // A supplied CPU backend exercises the real load-failure/cleanup path + // without a model or a second GPU. Ownership is identical for HIP. + backend.pflash_drafter_ctx_.backend = ggml_backend_cpu_init(); + backend.pflash_drafter_ctx_.gpu = 0; + TEST_ASSERT(backend.pflash_drafter_ctx_.backend != nullptr); + ModelBackend::CompressRequest request; + request.input_ids = {1, 2, 3}; + request.keep_ratio = 0.5f; + request.drafter_path = "/nonexistent/pr664-pflash.gguf"; + request.skip_park = true; + TEST_ASSERT(!backend.compress(request).ok); + TEST_ASSERT(!backend.pflash_drafter_loaded_); + TEST_ASSERT(backend.pflash_drafter_ctx_.backend == nullptr); + TEST_ASSERT(backend.pflash_drafter_ctx_.gpu == -1); + // Also keep a failing baseline run leak-free. + dflash::common::free_drafter(backend.pflash_drafter_ctx_); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + +static void test_pflash_keep_ratio_contract() { + std::fprintf(stderr, " test_pflash_keep_ratio_contract ..."); + struct ParkProbe : DeepSeek4Backend { + ParkProbe() : DeepSeek4Backend(DeepSeek4BackendConfig{}) {} + int park_calls = 0; + bool park(ParkTarget) override { + ++park_calls; + return false; // Stop before model loading; only validate admission. + } + }; + for (float ratio : {0.0f, 0.5f, 1.0f, -0.1f, 1.1f, + std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity()}) { + ParkProbe backend; + ModelBackend::CompressRequest request; + request.input_ids = {1, 2, 3}; + request.drafter_path = "/nonexistent/pflash-admission.gguf"; + request.keep_ratio = ratio; + const bool valid = std::isfinite(ratio) && ratio >= 0.0f && ratio <= 1.0f; + TEST_ASSERT(!backend.compress(request).ok); + TEST_ASSERT(backend.park_calls == (valid ? 1 : 0)); + } + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + +static void test_indexer_visibility_suffix() { + std::fprintf(stderr, " test_indexer_visibility_suffix ..."); + auto * backend = ggml_backend_cpu_init(); + TEST_ASSERT(backend != nullptr); + if (!backend) return; + constexpr int rows = 528, tokens = 5; + int passed = 0; + for (int first = 0; first < tokens; ++first) { + auto * ctx = make_test_context(4u << 20); + TEST_ASSERT(ctx != nullptr); + if (!ctx) continue; + auto * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, rows, tokens); + ggml_set_input(mask); + auto * suffix = deepseek4_indexer_visibility_suffix(ctx, mask, first, tokens - first); + TEST_ASSERT(deepseek4_indexer_visibility_suffix(ctx, nullptr, first, tokens - first) == nullptr); + ggml_set_output(suffix); + auto * graph = ggml_new_graph_custom(ctx, 32, false); + ggml_build_forward_expand(graph, suffix); + auto * buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + TEST_ASSERT(buffer != nullptr); + if (buffer) { + std::vector values((size_t) rows * tokens); + for (int t = 0; t < tokens; ++t) { + for (int r = 0; r < rows; ++r) values[(size_t) t * rows + r] = (float) (t * 1000 + r); + } + ggml_backend_tensor_set(mask, values.data(), 0, ggml_nbytes(mask)); + bool ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + std::vector actual((size_t) ggml_nelements(suffix)); + ggml_backend_tensor_get(suffix, actual.data(), 0, ggml_nbytes(suffix)); + const std::vector expected(values.begin() + (size_t) first * rows, values.end()); + ok &= suffix->ne[0] == rows && suffix->ne[1] == tokens - first; + ok &= ggml_is_contiguous(suffix) && actual == expected; + passed += ok; + if (!ok) std::fprintf(stderr, " FAIL first=%d;", first); + } + ggml_backend_buffer_free(buffer); + ggml_free(ctx); + } + ggml_backend_free(backend); + std::fprintf(stderr, " %d/%d cases passed\n", passed, tokens); + TEST_ASSERT(passed == tokens); +} + +static void test_pflash_legacy_compress_contract() { + std::fprintf(stderr, " test_pflash_legacy_compress_contract ..."); + struct CompressProbe : DeepSeek4Backend { + CompressProbe() : DeepSeek4Backend(DeepSeek4BackendConfig{}) {} + CompressRequest captured{}; + int calls = 0; + CompressResult compress(const CompressRequest & request) override { + ++calls; + captured = request; + CompressResult result; + result.ok = true; + result.compressed_ids = {request.input_ids.front()}; + return result; + } + }; + char path[] = "/tmp/ds4-compress-contract-XXXXXX"; + const int fd = mkstemp(path); + TEST_ASSERT(fd >= 0); + if (fd < 0) return; + const int32_t ids[] = {11, 22, 33}; + const bool written = write(fd, ids, sizeof(ids)) == (ssize_t) sizeof(ids); + close(fd); + TEST_ASSERT(written); + if (!written) { unlink(path); return; } + for (bool skip_park : {false, true}) { + CompressProbe backend; + std::vector output; + DaemonIO io; + io.on_token = [&](int32_t token) { output.push_back(token); return true; }; + const std::string command = std::string("compress ") + path + + " 0 /unused/drafter with spaces.gguf" + (skip_park ? " nopark" : ""); + TEST_ASSERT(backend.handle_compress(command, io)); + TEST_ASSERT(backend.calls == 1); + TEST_ASSERT(backend.captured.input_ids == std::vector({11, 22, 33})); + TEST_ASSERT(backend.captured.keep_ratio == 0.0f); + TEST_ASSERT(backend.captured.drafter_path == "/unused/drafter with spaces.gguf"); + TEST_ASSERT(backend.captured.skip_park == skip_park); + TEST_ASSERT(backend.captured.drafter_gpu == 0); + TEST_ASSERT(backend.captured.residency_action == DraftResidencyAction::KeepLoaded); + TEST_ASSERT(output == std::vector({11})); + } + unlink(path); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_dspark_raw_ring_rollback_after_wrap(ggml_backend_t backend) { std::fprintf(stderr, " test_dspark_raw_ring_rollback_after_wrap ..."); @@ -3164,9 +3326,9 @@ static void test_ds4_flash_attention_keep_cap_gpu() { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } -static void test_ds4_flash_attention_parallel_index_scan_gpu() { +static void test_ds4_flash_attention_parallel_index_scan_gpu(int selected_rows) { std::fprintf(stderr, - " test_ds4_flash_attention_parallel_index_scan_gpu ..."); + " test_ds4_flash_attention_parallel_index_scan_gpu(%d) ...", selected_rows); #if !defined(GGML_USE_HIP) std::fprintf(stderr, " skipped (HIP-only contract)\n"); return; @@ -3184,9 +3346,8 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { constexpr int raw_window = 128; // Keep the ordinary four-head shared-memory footprint above 24 KiB so // this shape is forced through the compact indexed path under test. - constexpr int n_comp_rows = 1280; - constexpr int n_kv = raw_rows + n_comp_rows; - constexpr int selected_rows = 512; + const int n_comp_rows = 2 * selected_rows + 256; + const int n_kv = raw_rows + n_comp_rows; ggml_context * ctx = make_test_context(4u << 20); TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); @@ -3243,6 +3404,19 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { !ggml_backend_supports_op(backend, oversized_output), "GPU accepted direct top-k wider than the live compressed span"); + ggml_tensor * over_capacity_topk = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, 1025, n_tokens); + ggml_tensor * over_capacity_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( + over_capacity_output, raw_rows, raw_window, -1025, 1); + ggml_flash_attn_ext_set_ds4_indexer_topk( + over_capacity_output, over_capacity_topk); + TEST_ASSERT_MSG( + !ggml_backend_supports_op(backend, over_capacity_output), + "GPU accepted direct top-k wider than the sorting capacity"); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); ggml_build_forward_expand(graph, output); ggml_build_forward_expand(graph, direct_output); @@ -3360,28 +3534,243 @@ static void test_ds4_flash_attention_parallel_index_scan_gpu() { 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 ..."); +static void test_ds4_saved_raw_rows_replay_gpu() { + std::fprintf(stderr, " test_ds4_saved_raw_rows_replay_gpu ..."); #if !defined(GGML_USE_HIP) - std::fprintf(stderr, " skipped (HIP-only candidate)\n"); + std::fprintf(stderr, " skipped (HIP-only replay qualification)\n"); return; #endif - ggml_backend_t backend = ggml_backend_cuda_init(0); + auto * backend = ggml_backend_cuda_init(0); + TEST_ASSERT_MSG(backend != nullptr, "GPU backend unavailable"); + if (!backend) return; + constexpr int dim = 512, ring_rows = 128; + int passed = 0, cases = 0; + for (auto type : {GGML_TYPE_F16, GGML_TYPE_F32}) { + for (int width = 2; width <= 5; ++width) { + auto * ctx = make_test_context(4u << 20); + TEST_ASSERT(ctx != nullptr); + if (!ctx) continue; + auto * ring = ggml_new_tensor_2d(ctx, type, dim, ring_rows); + auto * read_rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, width); + auto * write_rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, width); + auto * replacement = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, width); + for (auto * input : {ring, read_rows, write_rows, replacement}) { + ggml_set_input(input); + } + auto * saved = deepseek4_preserve_raw_rows(ctx, ring, read_rows); + TEST_ASSERT(saved->type == type); + ggml_set_output(saved); + auto * graph = ggml_new_graph_custom(ctx, 128, false); + ggml_build_forward_expand(graph, saved); + auto * updated = ggml_set_rows(ctx, ring, replacement, write_rows); + ggml_set_output(updated); + ggml_build_forward_expand(graph, updated); + auto alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + const bool allocated = ggml_gallocr_alloc_graph(alloc, graph); + TEST_ASSERT(allocated); + if (allocated) { + std::vector original((size_t) dim * ring_rows); + for (int r = 0; r < ring_rows; ++r) { + for (int d = 0; d < dim; ++d) { + original[(size_t) r * dim + d] = (float) r + (d % 4) * 0.25f; + } + } + std::vector original_f16(original.size()); + std::transform(original.begin(), original.end(), original_f16.begin(), ggml_fp32_to_fp16); + std::vector new_values((size_t) dim * width, -32.0f); + std::vector reads(width); + std::vector writes(width); + for (int position : {124, 124, 127, 128, 132, 255, 7680, 131072, 124}) { + for (int t = 0; t < width; ++t) reads[t] = (position + t) % ring_rows; + std::copy(reads.begin(), reads.end(), writes.begin()); + ggml_backend_tensor_set(ring, + type == GGML_TYPE_F16 ? (const void *) original_f16.data() : original.data(), + 0, ggml_nbytes(ring)); + ggml_backend_tensor_set(read_rows, reads.data(), 0, ggml_nbytes(read_rows)); + ggml_backend_tensor_set(write_rows, writes.data(), 0, ggml_nbytes(write_rows)); + ggml_backend_tensor_set(replacement, new_values.data(), 0, ggml_nbytes(replacement)); + ScopedCudaGraphOverrides replay( + /*disable_graphs=*/false, /*mmvq_max_ncols=*/0, + /*skip_property_check=*/true); + bool ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + std::vector actual((size_t) dim * width), actual_ring(original.size()); + if (ok && type == GGML_TYPE_F16) { + std::vector half_saved(actual.size()), half_ring(actual_ring.size()); + ggml_backend_tensor_get(saved, half_saved.data(), 0, ggml_nbytes(saved)); + ggml_backend_tensor_get(updated, half_ring.data(), 0, ggml_nbytes(updated)); + std::transform(half_saved.begin(), half_saved.end(), actual.begin(), ggml_fp16_to_fp32); + std::transform(half_ring.begin(), half_ring.end(), actual_ring.begin(), ggml_fp16_to_fp32); + } else if (ok) { + ggml_backend_tensor_get(saved, actual.data(), 0, ggml_nbytes(saved)); + ggml_backend_tensor_get(updated, actual_ring.data(), 0, ggml_nbytes(updated)); + } + for (int t = 0; t < width && ok; ++t) { + for (int d = 0; d < dim; ++d) { + ok &= actual[(size_t) t * dim + d] == original[(size_t) reads[t] * dim + d]; + } + } + for (int r = 0; r < ring_rows && ok; ++r) { + const bool replaced = std::find(reads.begin(), reads.end(), r) != reads.end(); + for (int d = 0; d < dim; ++d) { + ok &= actual_ring[(size_t) r * dim + d] == + (replaced ? -32.0f : original[(size_t) r * dim + d]); + } + } + ++cases; + passed += ok; + if (!ok) { + std::fprintf(stderr, " FAIL q=%d type=%s pos=%d;", + width, ggml_type_name(type), position); + } + } + } + ggml_backend_cuda_graph_invalidate_range( + backend, ggml_get_mem_buffer(ctx), ggml_get_mem_size(ctx)); + ggml_gallocr_free(alloc); + ggml_free(ctx); + } + } + ggml_backend_free(backend); + std::fprintf(stderr, " %d/%d cases passed\n", passed, cases); + TEST_ASSERT_MSG(cases == 72 && passed == cases, "cached verifier saved rows from a stale ring position"); +} + +static bool run_ds4_preserved_raw_rows_case( + ggml_backend_t backend, int width, int compressed_rows, + bool visible_suffix, ggml_type kv_type, bool direct) { + constexpr int dim = 512, heads = 64, raw_rows = 128, top_k = 512; + const int saved = width > 1 ? width : 0; + const int rows = raw_rows + compressed_rows + saved; + ggml_context * ctx = make_test_context(4u << 20); + if (!ctx) return false; + auto * q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, dim, width, heads); + auto * kv = ggml_new_tensor_3d(ctx, kv_type, dim, rows, 1); + auto * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, rows, width); + auto * topk = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, top_k, width); + for (auto * input : {q, kv, mask, topk}) ggml_set_input(input); + auto * selected = deepseek4_indexed_attention_rows( + ctx, topk, compressed_rows, saved); + auto * indexed_mask = ggml_ds4_indexer_mask(ctx, mask, selected, raw_rows); + ggml_set_output(indexed_mask); + auto * result = ggml_flash_attn_ext(ctx, q, kv, kv, + ggml_cast(ctx, direct ? mask : indexed_mask, GGML_TYPE_F16), + 1.0f / std::sqrt(float(dim)), 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(result, GGML_PREC_F32); + ggml_flash_attn_ext_set_ds4_sparse(result, raw_rows, raw_rows, + -(int) selected->ne[0], 32); + if (direct) ggml_flash_attn_ext_set_ds4_indexer_topk(result, selected); + ggml_set_output(result); + auto * graph = ggml_new_graph_custom(ctx, 128, false); + ggml_build_forward_expand(graph, indexed_mask); + ggml_build_forward_expand(graph, result); + auto alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + bool ok = ggml_backend_supports_op(backend, result) && + ggml_gallocr_alloc_graph(alloc, graph); + if (ok) { + std::vector qdata((size_t) dim * width * heads, 0.0f); + std::vector kvdata((size_t) dim * rows, 0.0f); + for (int r = 0; r < rows; ++r) { + kvdata[(size_t) dim * r] = r < raw_rows + compressed_rows ? 0.25f : 1.25f; + } + std::vector maskdata((size_t) rows * width, -1.0e30f); + std::vector iddata((size_t) top_k * width); + for (int t = 0; t < width; ++t) { + float * col = maskdata.data() + (size_t) t * rows; + std::fill(col, col + raw_rows, 0.0f); + if (visible_suffix) { + // Later lanes overwrite these ring slots; earlier lanes need + // their saved values, while the future writes stay hidden. + for (int r = t + 1; r < width; ++r) col[r] = -1.0e30f; + for (int s = t + 1; s < saved; ++s) col[raw_rows + compressed_rows + s] = 0.0f; + } + for (int c = 0; c < top_k; ++c) { + // Non-monotone, separated indices exercise sorting and both + // mask scanners. Some selected rows remain causally hidden. + const int row = (29 * c + 17 * t) % compressed_rows; + if (c < top_k - t % 3) col[raw_rows + row] = 0.0f; + iddata[(size_t) t * top_k + c] = row; + } + } + ggml_backend_tensor_set(q, qdata.data(), 0, ggml_nbytes(q)); + if (kv_type == GGML_TYPE_F16) { + std::vector half(kvdata.size()); + std::transform(kvdata.begin(), kvdata.end(), half.begin(), ggml_fp32_to_fp16); + ggml_backend_tensor_set(kv, half.data(), 0, ggml_nbytes(kv)); + } else { + ggml_backend_tensor_set(kv, kvdata.data(), 0, ggml_nbytes(kv)); + } + ggml_backend_tensor_set(mask, maskdata.data(), 0, ggml_nbytes(mask)); + ggml_backend_tensor_set(topk, iddata.data(), 0, ggml_nbytes(topk)); + ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + if (ok) { + std::vector actual(ggml_nelements(result)), actual_mask(maskdata.size()); + ggml_backend_tensor_get(result, actual.data(), 0, ggml_nbytes(result)); + ggml_backend_tensor_get(indexed_mask, actual_mask.data(), 0, ggml_nbytes(indexed_mask)); + for (int t = 0; t < width; ++t) { + // Zero queries give equal attention weights. Keeping the saved + // suffix replaces exactly the hidden future ring slots. + const int visible = visible_suffix ? std::max(0, saved - t - 1) : 0; + const float expected = 0.25f + float(visible) / (raw_rows + top_k - t % 3); + for (int h = 0; h < heads; ++h) { + for (int d = 0; d < dim; ++d) { + const float value = actual[((size_t) t * heads + h) * dim + d]; + ok &= std::isfinite(value) && + std::abs(value - (d == 0 ? expected : 0.0f)) < 1.0e-6f; + } + } + for (int r = 0; r < rows; ++r) { + const size_t at = (size_t) t * rows + r; + ok &= (actual_mask[at] > -1.0e20f) == (maskdata[at] > -1.0e20f); + } + } + } + } + if (!ok) std::fprintf(stderr, " saved-row FAIL q=%d comp=%d visible=%d kv=%s direct=%d\n", + width, compressed_rows, visible_suffix, ggml_type_name(kv_type), direct); + ggml_gallocr_free(alloc); + ggml_free(ctx); + return ok; +} + +static void test_ds4_preserved_raw_rows_gpu() { + std::fprintf(stderr, " test_ds4_preserved_raw_rows_gpu ..."); +#if !defined(GGML_USE_HIP) + std::fprintf(stderr, " skipped (HIP-only indexed attention)\n"); + return; +#endif + auto * backend = ggml_backend_cuda_init(0); if (!backend) { std::fprintf(stderr, " skipped (no GPU backend)\n"); return; } + int passed = 0, cases = 0; + for (auto type : {GGML_TYPE_F16, GGML_TYPE_F32}) { + for (int comp : {2048, 30720}) { + for (int width = 1; width <= 5; ++width) { + for (bool visible : {false, true}) { + for (bool direct : {false, true}) { + ++cases; + passed += run_ds4_preserved_raw_rows_case(backend, width, comp, visible, type, direct); + } + } + } + } + } + std::fprintf(stderr, " %d/%d cases passed\n", passed, cases); + TEST_ASSERT_MSG(passed == cases, "indexed attention lost saved raw KV rows or exposed future rows"); + ggml_backend_free(backend); +} +static void run_ds4_indexer_score_packed_small_case( + ggml_backend_t backend, int n_tokens) { 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; } @@ -3396,14 +3785,14 @@ static void test_ds4_indexer_score_packed_q4_gpu() { 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"); + "GPU rejected packed-small 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"); + TEST_ASSERT_MSG(allocated, "packed-small 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); @@ -3425,37 +3814,35 @@ static void test_ds4_indexer_score_packed_q4_gpu() { 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"); + setenv("GGML_DS4_INDEXER_PACK_SMALL", "0", 1); TEST_ASSERT_MSG( ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, - "reference q4 indexer score failed"); + "reference small-CM indexer score failed"); ggml_backend_tensor_get(scores, reference.data(), 0, reference.size() * sizeof(float)); - setenv("GGML_DS4_INDEXER_PACK_Q4", "1", 1); + setenv("GGML_DS4_INDEXER_PACK_SMALL", "1", 1); TEST_ASSERT_MSG( ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS, - "packed q4 indexer score failed"); + "packed small-CM 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"); + "packed small-CM indexer changed score bits"); auto measure_us = [&](bool packed) { if (packed) { - setenv("GGML_DS4_INDEXER_PACK_Q4", "1", 1); + setenv("GGML_DS4_INDEXER_PACK_SMALL", "1", 1); } else { - unsetenv("GGML_DS4_INDEXER_PACK_Q4"); + setenv("GGML_DS4_INDEXER_PACK_SMALL", "0", 1); } constexpr int warmups = 3; constexpr int iterations = 30; @@ -3474,22 +3861,172 @@ static void test_ds4_indexer_score_packed_q4_gpu() { }; const double reference_us = measure_us(false); const double packed_us = measure_us(true); - std::fprintf(stderr, " reference=%.1fus packed=%.1fus", + std::fprintf(stderr, " q%d=%.1f->%.1fus", n_tokens, 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); +} + +static void test_ds4_indexer_score_packed_small_gpu() { + std::fprintf(stderr, " test_ds4_indexer_score_packed_small_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; + } + ScopedEnvVar packed_small_guard("GGML_DS4_INDEXER_PACK_SMALL"); + for (int n_tokens = 2; n_tokens <= 5; ++n_tokens) { + run_ds4_indexer_score_packed_small_case(backend, n_tokens); + } ggml_backend_free(backend); std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_ds4_flash_rope_replay_gpu() { + std::fprintf(stderr, " test_ds4_flash_rope_replay_gpu ..."); +#if !defined(GGML_USE_HIP) + std::fprintf(stderr, " skipped (HIP-only contract)\n"); + return; +#endif + auto backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, " skipped (no GPU backend)\n"); + return; + } + constexpr int dim = 512, heads = 4, raw = 128, comp = 640, keep = 512; + constexpr int starts[] = {7680, 7688, 131072, 7680}; + int passed = 0, cases = 0; + for (auto type : {GGML_TYPE_F32, GGML_TYPE_F16}) { + for (int width = 1; width <= 5; ++width) { + for (bool forward : {false, true}) { + ++cases; + auto ctx = make_test_context(2u << 20); + const int rows = raw + comp; + auto q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, dim, width, heads); + auto kv = ggml_new_tensor_3d(ctx, type, dim, rows, 1); + auto mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, rows, width); + auto topk = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, keep, width); + auto positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, width); + auto negative_positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, width); + for (auto input : {q, kv, mask, topk, positions, negative_positions}) ggml_set_input(input); + auto attention = [&](int start) { + auto out = ggml_flash_attn_ext(ctx, q, kv, kv, mask, + 1.0f / std::sqrt(float(dim)), 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); + ggml_flash_attn_ext_set_ds4_sparse(out, raw, raw, -keep, 32); + ggml_flash_attn_ext_set_ds4_indexer_topk(out, topk); + ggml_flash_attn_ext_set_ds4_inverse_rope(out, start, + 10000.0f, 1.0f / 16.0f, 1.0f, 1.0f, 32.0f, 1.0f, 8192, forward); + ggml_set_output(out); + return out; + }; + auto dynamic = attention(starts[0]); + // Replay positions are runtime inputs, not part of the shape + // key. The baseline ignores this input and uses stale RoPE. + ggml_flash_attn_ext_set_ds4_rope_positions(dynamic, positions); + if (type == GGML_TYPE_F32 && width == 2 && forward) { + auto strided = ggml_view_1d(ctx, positions, width, 0); + strided->nb[0] = 2 * sizeof(int32_t); + for (auto invalid : { + ggml_new_tensor_1d(ctx, GGML_TYPE_F32, width), + ggml_new_tensor_1d(ctx, GGML_TYPE_I32, width + 1), + ggml_new_tensor_2d(ctx, GGML_TYPE_I32, width, 2), + strided}) { + // Bypass the setter to exercise backend validation of + // malformed imported graphs without launching them. + dynamic->src[6] = invalid; + TEST_ASSERT_MSG(!ggml_backend_supports_op(backend, dynamic), + "DS4 flash accepted invalid runtime RoPE positions"); + } + dynamic->src[6] = positions; + } + std::vector reference; + auto graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, dynamic); + for (int start : starts) { + reference.push_back(attention(start)); + ggml_build_forward_expand(graph, reference.back()); + } + // Also compare with the actual standalone GGML RoPE path, + // not just another instance of the fused implementation. + auto rope = [&](ggml_tensor * input, ggml_tensor * pos) { + return ggml_rope_ext(ctx, input, pos, nullptr, 64, + GGML_ROPE_TYPE_NORMAL | GGML_ROPE_TYPE_TAIL, 8192, + 10000.0f, 1.0f / 16.0f, 1.0f, 1.0f, 32.0f, 1.0f); + }; + auto native_q = forward + ? ggml_permute(ctx, rope(ggml_permute(ctx, q, 0, 2, 1, 3), positions), 0, 2, 1, 3) + : q; + auto native_attention = ggml_flash_attn_ext(ctx, native_q, kv, kv, mask, + 1.0f / std::sqrt(float(dim)), 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(native_attention, GGML_PREC_F32); + ggml_flash_attn_ext_set_ds4_sparse(native_attention, raw, raw, -keep, 32); + ggml_flash_attn_ext_set_ds4_indexer_topk(native_attention, topk); + auto native = rope(native_attention, negative_positions); + ggml_set_output(native); + ggml_build_forward_expand(graph, native); + auto alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + bool ok = ggml_backend_supports_op(backend, dynamic) && + ggml_gallocr_alloc_graph(alloc, graph); + if (ok) { + std::vector qdata(ggml_nelements(q)), kvdata(ggml_nelements(kv)); + for (size_t i = 0; i < qdata.size(); ++i) + qdata[i] = float(int((i * 17) % 71) - 35) * 0.0125f; + for (size_t i = 0; i < kvdata.size(); ++i) + kvdata[i] = float(int((i * 29 + i / dim) % 83) - 41) * 0.025f; + std::vector masks(ggml_nelements(mask), ggml_fp32_to_fp16(0)); + std::vector ids(ggml_nelements(topk)); + for (int t = 0; t < width; ++t) { + for (int k = 0; k < keep; ++k) + ids[(size_t)t * keep + k] = (29 * k + 17 * t) % comp; + } + ggml_backend_tensor_set(q, qdata.data(), 0, ggml_nbytes(q)); + if (type == GGML_TYPE_F16) { + std::vector half(kvdata.size()); + std::transform(kvdata.begin(), kvdata.end(), half.begin(), ggml_fp32_to_fp16); + ggml_backend_tensor_set(kv, half.data(), 0, ggml_nbytes(kv)); + } else ggml_backend_tensor_set(kv, kvdata.data(), 0, ggml_nbytes(kv)); + ggml_backend_tensor_set(mask, masks.data(), 0, ggml_nbytes(mask)); + ggml_backend_tensor_set(topk, ids.data(), 0, ggml_nbytes(topk)); + for (size_t run = 0; run < reference.size(); ++run) { + std::vector pos(width); + for (int t = 0; t < width; ++t) pos[t] = starts[run] + t; + ggml_backend_tensor_set(positions, pos.data(), 0, ggml_nbytes(positions)); + for (int & value : pos) value = -value; + ggml_backend_tensor_set(negative_positions, pos.data(), 0, ggml_nbytes(negative_positions)); + if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) { + ok = false; + break; + } + std::vector actual(ggml_nelements(dynamic)), expected(actual.size()); + ggml_backend_tensor_get(dynamic, actual.data(), 0, ggml_nbytes(dynamic)); + ggml_backend_tensor_get(reference[run], expected.data(), 0, ggml_nbytes(dynamic)); + for (size_t i = 0; i < actual.size(); ++i) + ok &= std::isfinite(actual[i]) && actual[i] == expected[i]; + ggml_backend_tensor_get(native, expected.data(), 0, ggml_nbytes(native)); + for (size_t i = 0; i < actual.size(); ++i) + ok &= nearly_equal(actual[i], expected[i], 2.0e-5f, 2.0e-5f); + } + } + if (!ok) std::fprintf(stderr, " FAIL width=%d kv=%s forward=%d", + width, ggml_type_name(type), forward); + passed += ok; + ggml_gallocr_free(alloc); + ggml_free(ctx); + } + } + } + std::fprintf(stderr, " %d/%d cases passed\n", passed, cases); + TEST_ASSERT_MSG(passed == cases, "fused RoPE replay used stale token positions"); + ggml_backend_free(backend); +} + static void test_ds4_topk_block_radix_gpu() { std::fprintf(stderr, " test_ds4_topk_block_radix_gpu ..."); #if !defined(GGML_USE_HIP) @@ -4773,6 +5310,11 @@ int main() { test_safe_compressor_batch_tokens(); test_hybrid_prefill_chunk_tokens(); test_dspark_park_all_releases_drafter(); + test_pflash_rejects_invalid_requests(); + test_pflash_failed_load_releases_backend(); + test_pflash_keep_ratio_contract(); + test_indexer_visibility_suffix(); + test_pflash_legacy_compress_contract(); test_dspark_raw_ring_rollback_after_wrap(backend); test_snapshot_save_restore(); test_monolithic_snapshot_preserves_decode_state(); @@ -4789,8 +5331,12 @@ int main() { test_output_graph_reuse_microbench(backend); #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_flash_attention_parallel_index_scan_gpu(512); + test_ds4_flash_attention_parallel_index_scan_gpu(1024); + test_ds4_saved_raw_rows_replay_gpu(); + test_ds4_preserved_raw_rows_gpu(); + test_ds4_indexer_score_packed_small_gpu(); + test_ds4_flash_rope_replay_gpu(); test_ds4_topk_block_radix_gpu(); test_ds4_flash_attention_inverse_rope_fallback_gpu(); test_hc_post_strided_split_gpu();