diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8e781239c..06cd1bcd9 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -84,6 +84,7 @@ jobs: test_dflash test_generate test_flash_attn_sparse test_server_unit \ test_deepseek4_unit test_feature_gate test_seq_slot_manager \ test_seq_engine_contract test_seq_batch_plan test_client_send_buffer \ + test_paged_kv_pool test_rope_tail \ test_model_smoke test_batched_gdn test_concat_transpose \ test_kernel_qualification_core test_kernel_qualification_gdn \ test_kernel_qualification_moe_ds4 -j$(nproc) @@ -92,7 +93,7 @@ jobs: run: | cd server/build ctest --output-on-failure \ - -R "server_unit|deepseek4_unit|feature_gate|seq_slot_manager|seq_engine_contract|seq_batch_plan|client_send_buffer|batched_gdn_cpu|test_kernel_qualification_core" \ + -R "server_unit|deepseek4_unit|feature_gate|seq_slot_manager|seq_engine_contract|seq_batch_plan|client_send_buffer|paged_kv_pool|rope_tail|batched_gdn_cpu|test_kernel_qualification_core" \ --no-tests=error - name: Populate venv with cu128 torch + setuptools diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 22552b3ae..a2004279b 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -469,6 +469,8 @@ add_library(dflash_common STATIC src/deepseek4/deepseek4_loader.cpp src/deepseek4/deepseek4_graph.cpp src/deepseek4/deepseek4_roctx.cpp + src/deepseek4/deepseek4_paged_cache.cpp + src/deepseek4/deepseek4_seq_engine.cpp src/deepseek4/deepseek4_backend.cpp src/deepseek4/deepseek4_daemon.cpp src/deepseek4/deepseek4_layer_split_adapter.cpp @@ -505,7 +507,7 @@ add_library(dflash_common STATIC src/common/draft_swa.cpp src/common/dflash_spec_decode.cpp src/common/concurrency/paged_kv_pool.cpp - src/qwen35/concurrency/qwen35_slot_manager.cpp + src/common/concurrency/seq_slot_manager.cpp src/common/layer_split_backend.cpp src/common/layer_split_runtime.cpp src/qwen35/graph_builders.cpp @@ -985,8 +987,23 @@ if(DFLASH27B_TESTS) hip::host) list(APPEND _raw_unit_test_targets test_rocmfpx_mmq) + add_executable(test_gfx1151_width4_dispatch test/test_gfx1151_width4_dispatch.cpp) + set_source_files_properties(test/test_gfx1151_width4_dispatch.cpp PROPERTIES LANGUAGE HIP) + set_target_properties(test_gfx1151_width4_dispatch PROPERTIES HIP_ARCHITECTURES "${_dflash_archs}") + target_include_directories(test_gfx1151_width4_dispatch PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/deps/llama.cpp/ggml/include) + target_link_libraries(test_gfx1151_width4_dispatch PRIVATE + ggml + ggml-base + ${DFLASH27B_GGML_BACKEND_TARGET} + hip::host) + list(APPEND _raw_unit_test_targets test_gfx1151_width4_dispatch) - + add_executable(test_softmax_sink_col test/test_softmax_sink_col.cpp) + set_source_files_properties(test/test_softmax_sink_col.cpp PROPERTIES LANGUAGE HIP) + set_target_properties(test_softmax_sink_col PROPERTIES HIP_ARCHITECTURES "${_dflash_archs}") + target_link_libraries(test_softmax_sink_col PRIVATE ggml ggml-base ggml-cpu ${DFLASH27B_GGML_BACKEND_TARGET} hip::host) + list(APPEND _raw_unit_test_targets test_softmax_sink_col) # Benchmark, not a test: built with the tests but deliberately NOT added to # _raw_unit_test_targets, so ctest stays deterministic. Measures the launch-count @@ -1003,6 +1020,10 @@ if(DFLASH27B_TESTS) target_include_directories(test_qwen35_tensor_parallel PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) target_link_libraries(test_qwen35_tensor_parallel PRIVATE dflash_common) + add_executable(test_rope_tail test/test_rope_tail.cpp) + target_compile_definitions(test_rope_tail PRIVATE DFLASH_ROPE_TEST_GPU=1) + target_link_libraries(test_rope_tail PRIVATE ggml ggml-base ggml-cpu ${DFLASH27B_GGML_BACKEND_TARGET}) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND _dflash_cuda_min_sm GREATER_EQUAL 80 AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_flashprefill_kernels.cpp") add_executable(test_flashprefill_kernels test/test_flashprefill_kernels.cpp) set_target_properties(test_flashprefill_kernels PROPERTIES CUDA_ARCHITECTURES "${_dflash_archs}") @@ -1481,7 +1502,11 @@ if(DFLASH27B_TESTS) # Pure host-side allocator test: no model, ggml, or GPU is required. add_executable(test_paged_kv_pool test/test_paged_kv_pool.cpp + test/test_deepseek4_page_layout.cpp + test/test_deepseek4_paged_cache.cpp + src/deepseek4/deepseek4_paged_cache.cpp src/common/concurrency/paged_kv_pool.cpp) + target_compile_definitions(test_paged_kv_pool PRIVATE DFLASH_DS4_PLAN_ONLY=1) target_include_directories(test_paged_kv_pool PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) list(APPEND _raw_unit_test_targets test_paged_kv_pool) @@ -1490,7 +1515,7 @@ if(DFLASH27B_TESTS) # Host-side slot bookkeeping test (concurrent serving): no GPU. add_executable(test_seq_slot_manager test/test_seq_slot_manager.cpp - src/qwen35/concurrency/qwen35_slot_manager.cpp + src/common/concurrency/seq_slot_manager.cpp src/common/concurrency/paged_kv_pool.cpp) target_include_directories(test_seq_slot_manager PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) @@ -1917,6 +1942,8 @@ if(DFLASH27B_TESTS) test_cuda_comm_api test_mmq_streamk_iq4_xs test_rms_norm_hip + test_rope_tail + test_paged_kv_pool test_rocmfp3_mix_registry test_rocmfp_mix_slice_matvec test_rocmfp_mix_gateup_glu diff --git a/server/README.md b/server/README.md index 8cbb420f3..278b540df 100644 --- a/server/README.md +++ b/server/README.md @@ -299,10 +299,10 @@ See the [current six-expert Strix Halo profile](https://www.lucebox.com/blog/dee | `--prefix-cache-slots ` | `32` | In-memory prefix-cache slots; `0` disables. | | `--agent-turn-cache` | off | Extend prefix caching through generated tool calls. | | `--prefill-cache-slots ` | `0` | Full-prompt cache slots. | -| `--paged-attention` | off | Enable 16-token paged KV blocks for supported Qwen targets. | -| `--max-concurrency ` | `1` | Maximum concurrent decode sequences; values above 1 enable paged attention. | +| `--paged-attention` | off | Enable paged KV for supported Qwen targets (16-token blocks) or DeepSeek4 on Strix Halo and R9700 plus Strix Halo (128-token pages). DeepSeek4 paged serving is AR-only. | +| `--max-concurrency ` | `1` | Maximum concurrent decode sequences. Qwen supports up to 64; DeepSeek4 supports up to 6. Values above 1 enable paged attention. | | `--admission-coalesce-ms ` | `20` | Idle-to-busy batching window from 0 through 1000 ms. | -| `--kv-pool-tokens ` | auto | Shared physical K/V capacity for concurrent serving. | +| `--kv-pool-tokens ` | auto | Shared physical K/V capacity for concurrent serving, rounded to the backend page size. | | `--kv-cache-dir ` | none | Enable persistent disk KV cache in this directory. | | `--kv-cache-budget ` | `4096` | Disk KV-cache size cap. | | `--kv-cache-min-tokens ` | `512` | Minimum prefix length to persist. | diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index 756193d7a..b8ce8e40c 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -252,6 +252,10 @@ #define GGML_ROPE_TYPE_MROPE 8 #define GGML_ROPE_TYPE_VISION 24 #define GGML_ROPE_TYPE_IMROPE 40 // binary: 101000 +// Lucebox: rotate the LAST n_dims of each row (NORMAL pairing) and pass the +// leading ne0 - n_dims through unchanged, so a tail-rotary head (DeepSeek4) +// is one launch instead of split, rotate, and concat. Combine with NORMAL only. +#define GGML_ROPE_TYPE_TAIL 64 #define GGML_MROPE_SECTIONS 4 @@ -1807,6 +1811,16 @@ extern "C" { float scale, float max_bias); + // Lucebox: softmax over [a | sink] per row with the sink as a virtual last + // column, scale folded in (a*scale), no mask. sinks holds one value per row + // of a (ne[0] == ggml_nrows(a)). Bit-identical to scale -> concat(sink) -> + // soft_max -> view, in one launch with a contiguous result of a's shape. + GGML_API struct ggml_tensor * ggml_soft_max_ext_sink_col( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * sinks, + float scale); + GGML_API void ggml_soft_max_add_sinks( struct ggml_tensor * a, struct ggml_tensor * sinks); diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp index 174f6a45d..ae3baf930 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp @@ -2322,7 +2322,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s std::vector ids; std::vector used_ids; size_t batch_staging_cursors[GGML_SCHED_MAX_BACKENDS] = {}; - const bool copy_destinations_ready = sched->backends_synchronized; + const bool staging_arena_reusable = sched->backends_synchronized; + bool copy_destinations_ready[GGML_SCHED_MAX_BACKENDS] = {}; + std::fill_n(copy_destinations_ready, sched->n_backends, + sched->backends_synchronized); // From this point onward an early return must remain conservative: a // backend may have accepted work even if a later split fails. sched->backends_synchronized = false; @@ -2344,7 +2347,8 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s ggml_backend_sched_split_copy_state copy_state{ sched, split_backend_id, split_backend, - copy_destinations_ready, copy_destinations_ready, + staging_arena_reusable, + copy_destinations_ready[split_backend_id], batch_staging_cursors}; // copy the input tensors to the split backend @@ -2537,6 +2541,8 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s j0 = j1; } } + copy_destinations_ready[split_backend_id] = + sched->callback_eval != nullptr; // Publish producer completion before a later consumer-backend graph // reaches its in-graph event wait. Recording is asynchronous and does diff --git a/server/deps/llama.cpp/ggml/src/ggml-cann/aclnn_ops.cpp b/server/deps/llama.cpp/ggml/src/ggml-cann/aclnn_ops.cpp index a950475fc..94e99a440 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cann/aclnn_ops.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cann/aclnn_ops.cpp @@ -2669,6 +2669,7 @@ void ggml_cann_rope(ggml_backend_cann_context & ctx, ggml_tensor * dst) { // const int n_past = ((int32_t *) dst->op_params)[0]; const int n_dims = ((int32_t *) dst->op_params)[1]; const int mode = ((int32_t *) dst->op_params)[2]; + GGML_ASSERT((mode & GGML_ROPE_TYPE_TAIL) == 0); // const int n_ctx = ((int32_t *) dst->op_params)[3]; const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; @@ -3028,6 +3029,7 @@ void ggml_cann_rope_cache_preload(ggml_backend_cann_context & ctx, ggml_tensor * int sections[4]; const int n_dims = ((int32_t *) dst->op_params)[1]; const int mode = ((int32_t *) dst->op_params)[2]; + GGML_ASSERT((mode & GGML_ROPE_TYPE_TAIL) == 0); const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; GGML_TENSOR_UNARY_OP_LOCALS diff --git a/server/deps/llama.cpp/ggml/src/ggml-cann/ggml-cann.cpp b/server/deps/llama.cpp/ggml/src/ggml-cann/ggml-cann.cpp index 5fc484b34..a8bfa82f9 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cann/ggml-cann.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cann/ggml-cann.cpp @@ -2364,6 +2364,11 @@ static enum ggml_status ggml_backend_cann_graph_compute(ggml_backend_t backend, * otherwise false. */ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + if ((op->op == GGML_OP_ROPE || op->op == GGML_OP_ROPE_BACK) && + (op->op_params[2] & GGML_ROPE_TYPE_TAIL)) { + return false; + } + switch (op->op) { case GGML_OP_UNARY: switch (ggml_get_unary_op(op)) { diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c index 6756c8383..f5ae82973 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c @@ -3018,7 +3018,7 @@ struct ggml_cplan ggml_graph_plan( // conversion between F32 and I32 (node->src[0]->type == GGML_TYPE_F32 && node->src[1] && node->src[1]->type == GGML_TYPE_I32) || (node->src[0]->type == GGML_TYPE_I32 && node->src[1] && node->src[1]->type == GGML_TYPE_F32)) { - cur = ggml_type_size(GGML_TYPE_F32) * node->ne[0] * n_tasks; + cur = ggml_type_size(GGML_TYPE_F32) * (node->ne[0] + 1) * n_tasks; // +1: sink-column mode } } break; case GGML_OP_ADD: diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp b/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp index 785fda863..f8d198fd5 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp @@ -5356,6 +5356,12 @@ static void ggml_compute_forward_soft_max_f32( // sinks const float * sk = src2 ? (float *)((char *) src2->data) : nullptr; + // Lucebox sink-as-virtual-column mode (see ggml_soft_max_ext_sink_col) + int32_t sink_col = 0; + memcpy(&sink_col, (int32_t *) dst->op_params + 2, sizeof(int32_t)); + if (sink_col) { + GGML_ASSERT(sk && !src1); + } for (int64_t i03 = 0; i03 < ne03; i03++) { for (int64_t i02 = 0; i02 < ne02; i02++) { @@ -5396,6 +5402,20 @@ static void ggml_compute_forward_soft_max_f32( } #endif // NDEBUG + if (sink_col) { + // Append the row's sink as one more column and run the + // plain softmax over ne00 + 1 values, exactly as the + // concat form would; only the first ne00 outputs are kept. + const int64_t row = i01 + ne01*(i02 + ne02*i03); + wp[ne00] = sk[row]; + float max = -INFINITY; + ggml_vec_max_f32(ne00 + 1, &max, wp); + ggml_float sum = ggml_vec_soft_max_f32(ne00 + 1, wp, wp, max); + assert(sum > 0.0); + sum = 1.0/sum; + ggml_vec_cpy_f32(ne00, dp, wp); + ggml_vec_scale_f32(ne00, dp, sum); + } else { float max = -INFINITY; ggml_vec_max_f32(ne00, &max, wp); @@ -5413,6 +5433,7 @@ static void ggml_compute_forward_soft_max_f32( sum = 1.0/sum; ggml_vec_scale_f32(ne00, dp, sum); + } #ifndef NDEBUG for (int i = 0; i < ne00; ++i) { @@ -5871,9 +5892,13 @@ static void ggml_compute_forward_rope_flt( float corr_dims[2]; ggml_rope_yarn_corr_dims(n_dims, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); - const bool is_imrope = mode == GGML_ROPE_TYPE_IMROPE; // qwen3vl apply interleaved mrope - const bool mrope_used = mode & GGML_ROPE_TYPE_MROPE; // ggml_rope_multi, note: also true for vision (24 & 8 == true) and for imrope - const bool is_vision = mode == GGML_ROPE_TYPE_VISION; + const bool is_tail = mode & GGML_ROPE_TYPE_TAIL; // rotate the last n_dims, pass the head through + const int mode_base = mode & ~GGML_ROPE_TYPE_TAIL; + GGML_ASSERT(!is_tail || mode_base == GGML_ROPE_TYPE_NORMAL); + const int64_t rot_off = is_tail ? ne0 - n_dims : 0; + const bool is_imrope = mode_base == GGML_ROPE_TYPE_IMROPE; // qwen3vl apply interleaved mrope + const bool mrope_used = mode_base & GGML_ROPE_TYPE_MROPE; // ggml_rope_multi, note: also true for vision (24 & 8 == true) and for imrope + const bool is_vision = mode_base == GGML_ROPE_TYPE_VISION; if (mrope_used) { GGML_ASSERT(sections[0] > 0 || sections[1] > 0 || sections[2] > 0); @@ -5909,7 +5934,7 @@ static void ggml_compute_forward_rope_flt( if (last_i2 != i2) { if (!mrope_used) { const int64_t p = pos[i2]; - ggml_rope_cache_init(p, freq_scale, freq_factors, corr_dims, ne0, ext_factor, attn_factor, cache, sin_sign, theta_scale); + ggml_rope_cache_init(p, freq_scale, freq_factors, corr_dims, is_tail ? n_dims : ne0, ext_factor, attn_factor, cache, sin_sign, theta_scale); } else { const int64_t p_t = pos[i2]; @@ -5927,9 +5952,12 @@ static void ggml_compute_forward_rope_flt( T * src = (T *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01); T * dst_data = (T *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1); - switch (mode) { + switch (mode_base) { case GGML_ROPE_TYPE_NORMAL: - rotate_pairs(n_dims, 1, cache, src, dst_data, 1); + // For the tail mode the cache entries [0, n_dims) hold + // the angles of the rotated span, as they would for an + // extracted n_dims-wide tail tensor. + rotate_pairs(n_dims, 1, cache, src + rot_off, dst_data + rot_off, 1); break; case GGML_ROPE_TYPE_NEOX: case GGML_ROPE_TYPE_MROPE: @@ -5945,7 +5973,11 @@ static void ggml_compute_forward_rope_flt( if (!is_vision) { // fill the remain channels with data from src tensor - for (int64_t i0 = n_dims; i0 < ne0; i0 += 2) { + // (the head [0, rot_off) for the tail mode, the tail + // [n_dims, ne0) otherwise) + const int64_t copy_begin = is_tail ? 0 : n_dims; + const int64_t copy_end = is_tail ? rot_off : ne0; + for (int64_t i0 = copy_begin; i0 < copy_end; i0 += 2) { const T * const src = (T *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); T * dst_data = (T *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 865d488e4..b13466c7c 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu @@ -473,19 +473,18 @@ const ggml_cuda_device_info & ggml_cuda_info() { // buffer pool for cuda (legacy) struct ggml_cuda_pool_leg : public ggml_cuda_pool { - // 1024 (upstream 256): LUCE_Q8_MEMO keeps one pooled q8_1 activation - // buffer per quantized matmul alive across a whole graph evaluation - // (~300 on a 64-layer hybrid), and a full pool falls back to freeing - // in-flight buffers with cudaFree. - static const int MAX_BUFFERS = 1024; + // Free buffers keyed by size. alloc() takes the smallest cached buffer + // that fits, the same best-fit choice the original linear scan made, and + // free() returns a buffer to the cache; both are O(log n). A DeepSeek4 + // gathered graph that packs prompt rows memoizes thousands of q8_1 + // activations per evaluation (LUCE_Q8_MEMO keeps one per quantized + // matmul), so the original per-operation scan of the whole pool and its + // cudaFree fallback on a full pool (which serializes the stream) both + // showed up in the round time. 8192 (upstream 256) bounds the cache. + static const int MAX_BUFFERS = 8192; int device; - struct ggml_cuda_buffer { - void * ptr = nullptr; - size_t size = 0; - }; - - ggml_cuda_buffer buffer_pool[MAX_BUFFERS] = {}; + std::multimap free_buffers; size_t pool_size = 0; explicit ggml_cuda_pool_leg(int device) : @@ -494,52 +493,20 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { ~ggml_cuda_pool_leg() { ggml_cuda_set_device(device); - for (int i = 0; i < MAX_BUFFERS; ++i) { - ggml_cuda_buffer & b = buffer_pool[i]; - if (b.ptr != nullptr) { - CUDA_CHECK(cudaFree(b.ptr)); - pool_size -= b.size; - } + for (auto & entry : free_buffers) { + CUDA_CHECK(cudaFree(entry.second)); + pool_size -= entry.first; } + free_buffers.clear(); GGML_ASSERT(pool_size == 0); } void * alloc(size_t size, size_t * actual_size) override { -#ifdef DEBUG_CUDA_MALLOC - int nnz = 0; - size_t max_size = 0; -#endif - size_t best_diff = 1ull << 36; - int ibest = -1; - for (int i = 0; i < MAX_BUFFERS; ++i) { - ggml_cuda_buffer& b = buffer_pool[i]; - if (b.ptr != nullptr) { -#ifdef DEBUG_CUDA_MALLOC - ++nnz; - if (b.size > max_size) max_size = b.size; -#endif - if (b.size >= size) { - size_t diff = b.size - size; - if (diff < best_diff) { - best_diff = diff; - ibest = i; - if (!best_diff) { - void * ptr = b.ptr; - *actual_size = b.size; - b.ptr = nullptr; - b.size = 0; - return ptr; - } - } - } - } - } - if (ibest >= 0) { - ggml_cuda_buffer& b = buffer_pool[ibest]; - void * ptr = b.ptr; - *actual_size = b.size; - b.ptr = nullptr; - b.size = 0; + auto it = free_buffers.lower_bound(size); + if (it != free_buffers.end()) { + void * ptr = it->second; + *actual_size = it->first; + free_buffers.erase(it); return ptr; } void * ptr; @@ -549,21 +516,13 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { CUDA_CHECK(ggml_cuda_device_malloc(&ptr, look_ahead_size, device)); *actual_size = look_ahead_size; pool_size += look_ahead_size; -#ifdef DEBUG_CUDA_MALLOC - GGML_LOG_INFO("%s[%d]: %d buffers, max_size = %u MB, pool_size = %u MB, requested %u MB\n", __func__, device, nnz, - (uint32_t)(max_size / 1024 / 1024), (uint32_t)(pool_size / 1024 / 1024), (uint32_t)(size / 1024 / 1024)); -#endif return ptr; } void free(void * ptr, size_t size) override { - for (int i = 0; i < MAX_BUFFERS; ++i) { - ggml_cuda_buffer& b = buffer_pool[i]; - if (b.ptr == nullptr) { - b.ptr = ptr; - b.size = size; - return; - } + if ((int) free_buffers.size() < MAX_BUFFERS) { + free_buffers.emplace(size, ptr); + return; } GGML_LOG_DEBUG(GGML_CUDA_NAME " buffer pool full, increase MAX_CUDA_BUFFERS\n"); ggml_cuda_set_device(device); @@ -578,17 +537,12 @@ struct ggml_cuda_pool_leg : public ggml_cuda_pool { size_t trim() override { ggml_cuda_set_device(device); size_t freed = 0; - for (int i = 0; i < MAX_BUFFERS; ++i) { - ggml_cuda_buffer & b = buffer_pool[i]; - if (b.ptr == nullptr) { - continue; - } - CUDA_CHECK(cudaFree(b.ptr)); - freed += b.size; - pool_size -= b.size; - b.ptr = nullptr; - b.size = 0; + for (auto & entry : free_buffers) { + CUDA_CHECK(cudaFree(entry.second)); + freed += entry.first; + pool_size -= entry.first; } + free_buffers.clear(); return freed; } }; @@ -2827,10 +2781,10 @@ static bool ggml_cuda_try_fuse_mul_mat_glu( GGML_CUDA_CC_IS_RDNA3_5(cc) && (!mix_mmq || !(mix_mmq[0] == '0' && mix_mmq[1] == '\0')); #endif - if (default_ds4_mix_gate_up_mmq || - ggml_cuda_should_use_mmq( - src0->type, cc, ncols, - ids ? src0->ne[2] : /*n_experts=*/0) && + if ((default_ds4_mix_gate_up_mmq || + ggml_cuda_should_use_mmq( + src0->type, cc, ncols, + ids ? src0->ne[2] : /*n_experts=*/0)) && !(ggml_cuda_mmvq_max_ncols_override > 0 && ncols <= ggml_cuda_mmvq_max_ncols_override)) { ggml_cuda_mul_mat_q_pair( @@ -3988,7 +3942,6 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx if (cgraph->uid != 0 && cgraph->uid == graph->uid) { - GGML_LOG_DEBUG("CUDA Graph id %zu reused\n", cgraph->uid); GGML_ASSERT((int)graph->node_props.size() == cgraph->n_nodes); return false; } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu index 78c79918f..8caac9331 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu @@ -4,6 +4,8 @@ #include "mmvf.cuh" #include "convert.cuh" +#include + template static __global__ void mul_mat_vec_f( const T * __restrict__ x, const float * __restrict__ y, const int32_t * __restrict__ ids, const ggml_cuda_mm_fusion_args_device fusion, float * __restrict__ dst, @@ -786,6 +788,14 @@ void ggml_cuda_op_mul_mat_vec_f( GGML_UNUSED_VARS(ctx, src1, dst, src1_ddq_i, src1_ncols, src1_padded_row_size); } +static bool mmvf_narrow_f16_enabled() { + static const bool enabled = []() { + const char * value = std::getenv("DFLASH_CUDA_MMVF_NARROW_F16"); + return value == nullptr || value[0] == '\0' || std::atoi(value) != 0; + }(); + return enabled; +} + // LUCE_MMVF_MAX_NCOLS_F16: ceiling on ne11 at or below which an F16 mul_mat // prefers this file's mul_mat_vec kernel over rocBLAS on AMD. Unset (or 0) // keeps the per-architecture values chosen below, so this changes nothing on @@ -854,11 +864,25 @@ bool ggml_cuda_should_use_mmvf(enum ggml_type type, int cc, const int64_t * src0 } return ne11 <= 8; } else if (GGML_CUDA_CC_IS_AMD(cc)) { + const int ceiling = luce_mmvf_f16_max_ncols(); + if (ceiling > 0 && fp16_mma_hardware_available(cc)) { + return ne11 <= ceiling; + } + // rocBLAS ships a single 128x128 macro-tile and no split-K in + // its gfx1151 F16 Tensile library, so a narrow F16 weight + // (ne01 well below the tile) with more than three columns runs + // the whole K loop in one workgroup. The DeepSeek4 + // hyper-connection mix projection [16384,24] x 4 measured + // 367 us there against 13 us on MMVF. MMVF instantiates up to + // eight columns and keeps the single-column accumulation + // order per column. DFLASH_CUDA_MMVF_NARROW_F16=0 restores + // the BLAS route. + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && src0_ne[1] <= 32 && + src0_ne[2]*src0_ne[3] == 1 && ne11 <= 8 && + mmvf_narrow_f16_enabled()) { + return true; + } if (fp16_mma_hardware_available(cc)) { - const int ceiling = luce_mmvf_f16_max_ncols(); - if (ceiling > 0) { - return ne11 <= ceiling; - } if (GGML_CUDA_CC_IS_RDNA3(cc)) { return ne11 <= 3; } 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 e6eca3c0b..34d532cb5 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -704,13 +704,18 @@ static bool is_gfx1151(const int cc) { return cc == GGML_CUDA_CC_OFFSET_AMD + 0x1151; } -// Dense four-column projections use the same weight row for every column. -// Decode each packed ROCmFP4 fragment once while preserving the original -// per-column DP4A and floating-point accumulation order. -static __device__ __forceinline__ void vec_dot_rocmfp4_fast_q8_1_4cols( +// Dense projections at two to sixteen columns use the same weight row for +// every column. Decode each packed ROCmFP4 fragment once while preserving the +// original per-column DP4A and floating-point accumulation order, so every +// column equals the single-column kernel bit for bit at any width. +#define ROCMFP4_REUSE_MAX_COLS 16 +template +static __device__ __forceinline__ void vec_dot_rocmfp4_fast_q8_1_ncols( const void * __restrict__ vbq, const block_q8_1 * __restrict__ y, const uint32_t stride_col_y, const int kby, const int kbx, const int iqs, - float (&result)[4]) { + float (&result)[ncols]) { + static_assert(ncols >= 2 && ncols <= ROCMFP4_REUSE_MAX_COLS, + "weight reuse covers two to sixteen dense columns"); const block_rocmfp4_fast * bq4 = (const block_rocmfp4_fast *) vbq + kbx; int2 weights[VDR_ROCMFP4_FAST_Q8_1_MMVQ]; @@ -721,9 +726,9 @@ static __device__ __forceinline__ void vec_dot_rocmfp4_fast_q8_1_4cols( } const float weight_scale = rocmfp4_ue4m3_to_fp32_half_finite(bq4->e); - int sums[4] = {}; + int sums[ncols] = {}; #pragma unroll - for (int j = 0; j < 4; ++j) { + for (int j = 0; j < ncols; ++j) { const block_q8_1 * bq8 = &y[j*stride_col_y + kby]; const int * q8 = (const int *) bq8->qs + iqs; #pragma unroll @@ -793,8 +798,9 @@ static __global__ void mul_mat_vec_q( "partial K-loop unrolling is limited to the profiled FP4FAST path"); static_assert(!reuse_rocmfp4_weights || (type == GGML_TYPE_Q4_0_ROCMFP4_FAST && - ncols_dst == 4 && fixed_ncols_x == 0 && !unroll_k_loop_2), - "weight reuse is limited to the four-column FP4FAST path"); + ncols_dst >= 2 && ncols_dst <= ROCMFP4_REUSE_MAX_COLS && + fixed_ncols_x == 0 && !unroll_k_loop_2), + "weight reuse is limited to the two- to sixteen-column FP4FAST path"); static_assert(!c_fp3_packed24 || type == GGML_TYPE_Q3_0_ROCMFPX, "packed FP3 MMVQ specialization requires ROCmFP3 weights"); static_assert(!c_fp4_x4 || @@ -912,22 +918,22 @@ static __global__ void mul_mat_vec_q( #pragma unroll for (int i = 0; i < rows_per_cuda_block; ++i) { - float dots[4]; - vec_dot_rocmfp4_fast_q8_1_4cols( + float dots[ncols_dst]; + vec_dot_rocmfp4_fast_q8_1_ncols( vx, y, stride_col_y, kby, kbx_offset + i*stride_row_x + kbx, kqs, dots); #pragma unroll - for (int j = 0; j < 4; ++j) { + for (int j = 0; j < ncols_dst; ++j) { tmp[j][i] += dots[j]; } if constexpr (has_fusion) { if (use_gate) { - float gate_dots[4]; - vec_dot_rocmfp4_fast_q8_1_4cols( + float gate_dots[ncols_dst]; + vec_dot_rocmfp4_fast_q8_1_ncols( vgate, y, stride_col_y, kby, kbx_offset + i*stride_row_x + kbx, kqs, gate_dots); #pragma unroll - for (int j = 0; j < 4; ++j) { + for (int j = 0; j < ncols_dst; ++j) { tmp_gate[j][i] += gate_dots[j]; } } @@ -1882,7 +1888,8 @@ static void mul_mat_vec_rocmfp4_unroll2_launch( ids_stride, stream); } -static void mul_mat_vec_rocmfp4_4col_reuse_launch( +template +static void mul_mat_vec_rocmfp4_reuse_launch( const void * vx, const void * vy, const ggml_cuda_mm_fusion_args_device fusion, float * dst, const uint32_t ncols_x, const uint32_t nrows_x, const uint3 nchannels_y, @@ -1895,7 +1902,8 @@ static void mul_mat_vec_rocmfp4_4col_reuse_launch( const uint32_t stride_sample_dst, const uint32_t ids_stride, const int warp_size, const mmvq_parameter_table_id table_id, cudaStream_t stream) { - constexpr int ncols_dst = 4; + static_assert(ncols_dst >= 2 && ncols_dst <= ROCMFP4_REUSE_MAX_COLS, + "weight reuse covers two to sixteen dense columns"); const auto dims = calc_launch_params( ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); mul_mat_vec_q_switch_fusion< @@ -2017,7 +2025,8 @@ static void mul_mat_vec_q_moe_launch( std::getenv("DFLASH_CUDA_MMVQ_MOE_FP3_PACKED24"); return e && e[0] == '1' && e[1] == '\0'; }(); - const bool fp3_packed24 = fp3_packed24_configured && + const bool fp3_packed24 = + fp3_packed24_configured && std::getenv("DFLASH_CUDA_MMVQ_MOE_FP3_PACKED24_RUNTIME_DISABLE") == nullptr; static const bool fp2_packed32 = []() { const char * e = @@ -2117,14 +2126,21 @@ static void mul_mat_vec_q_switch_ncols_dst( const int ids_stride, cudaStream_t stream) { GGML_ASSERT(ncols_x % ggml_blck_size(type) == 0); - GGML_ASSERT(ncols_dst <= (ids ? MMVQ_MAX_MOE_BATCH_SIZE : MMVQ_MAX_BATCH_SIZE)); + + const int device = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[device].cc; + // The gfx1151 ROCmFP4 dense weight-reuse kernel accepts up to sixteen + // columns; every other path keeps the generic MMVQ batch limit. + const bool wide_rocmfp4_reuse = + type == GGML_TYPE_Q4_0_ROCMFP4_FAST && ids == nullptr && is_gfx1151(cc); + GGML_ASSERT(ncols_dst <= (ids ? MMVQ_MAX_MOE_BATCH_SIZE + : wide_rocmfp4_reuse ? ROCMFP4_REUSE_MAX_COLS + : MMVQ_MAX_BATCH_SIZE)); const uint3 nchannels_y_fd = ids ? init_fastdiv_values(nchannels_y) : make_uint3(0, 0, 0); const uint3 channel_ratio_fd = ids ? make_uint3(0, 0, 0) : init_fastdiv_values(nchannels_dst / nchannels_x); const uint3 sample_ratio_fd = init_fastdiv_values(nsamples_dst / nsamples_x); - const int device = ggml_cuda_get_device(); - const int cc = ggml_cuda_info().devices[device].cc; const int warp_size = ggml_cuda_info().devices[device].warp_size; const mmvq_parameter_table_id table_id = get_device_table_id(cc); @@ -2305,14 +2321,40 @@ static void mul_mat_vec_q_switch_ncols_dst( // per-lane K traversal and accumulation order, so exact validated shapes // use the faster kernel without a serving-time tuning flag. if constexpr (type == GGML_TYPE_Q4_0_ROCMFP4_FAST) { - if (is_gfx1151(cc) && !has_ids && ncols_dst == 4) { - mul_mat_vec_rocmfp4_4col_reuse_launch( - vx, vy, fusion, dst, ncols_x, nrows_x, nchannels_y_fd, - nchannels_dst, stride_row_x, stride_col_y, stride_col_dst, - channel_ratio_fd, stride_channel_x, stride_channel_y, - stride_channel_dst, nsamples_dst, sample_ratio_fd, - stride_sample_x, stride_sample_y, stride_sample_dst, - ids_stride, warp_size, table_id, stream); + // Two to sixteen dense columns share one decode of every weight + // fragment. Two and three columns previously took the generic + // per-column kernel; above the generic MMVQ limit a packed prompt + // step previously had to split its projections into four-column + // parts and re-read the weights per part. + if (is_gfx1151(cc) && !has_ids && ncols_dst >= 2 && + ncols_dst <= ROCMFP4_REUSE_MAX_COLS) { +#define GGML_ROCMFP4_REUSE_LAUNCH(NC) \ + case NC: mul_mat_vec_rocmfp4_reuse_launch( \ + vx, vy, fusion, dst, ncols_x, nrows_x, nchannels_y_fd, \ + nchannels_dst, stride_row_x, stride_col_y, stride_col_dst, \ + channel_ratio_fd, stride_channel_x, stride_channel_y, \ + stride_channel_dst, nsamples_dst, sample_ratio_fd, \ + stride_sample_x, stride_sample_y, stride_sample_dst, \ + ids_stride, warp_size, table_id, stream); break + switch (ncols_dst) { + GGML_ROCMFP4_REUSE_LAUNCH(2); + GGML_ROCMFP4_REUSE_LAUNCH(3); + GGML_ROCMFP4_REUSE_LAUNCH(4); + GGML_ROCMFP4_REUSE_LAUNCH(5); + GGML_ROCMFP4_REUSE_LAUNCH(6); + GGML_ROCMFP4_REUSE_LAUNCH(7); + GGML_ROCMFP4_REUSE_LAUNCH(8); + GGML_ROCMFP4_REUSE_LAUNCH(9); + GGML_ROCMFP4_REUSE_LAUNCH(10); + GGML_ROCMFP4_REUSE_LAUNCH(11); + GGML_ROCMFP4_REUSE_LAUNCH(12); + GGML_ROCMFP4_REUSE_LAUNCH(13); + GGML_ROCMFP4_REUSE_LAUNCH(14); + GGML_ROCMFP4_REUSE_LAUNCH(15); + GGML_ROCMFP4_REUSE_LAUNCH(16); + default: GGML_ABORT("unreachable reuse width"); + } +#undef GGML_ROCMFP4_REUSE_LAUNCH return; } if (is_gfx1151(cc) && ncols_dst == 1) { diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu index c61f695bc..637f6d27f 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu @@ -8,6 +8,7 @@ // standalone swiglu_ds4 kernel applies, not a re-derivation of the formula. #include "unary.cuh" #include +#include #include #include #include @@ -56,6 +57,7 @@ struct MixEntry { const nv_bfloat16 * codebooks; // n_experts * 2 * 4 const uint8_t * modes; // n_experts int device; // device the side-data lives on; frees must happen in that context + bool gfx1151; // registration-time dispatch tuning; no hot-path device query }; // Dispatch wrappers keep this lock from lookup through kernel enqueue. It is // recursive because MMQ takes the public dispatch lock and then calls the @@ -83,6 +85,17 @@ static int mix_device_of(const void * p) { return dev; } +static bool mix_device_is_gfx1151(int device) { +#if defined(GGML_USE_HIP) + cudaDeviceProp prop{}; + return cudaGetDeviceProperties(&prop, device) == cudaSuccess && + std::strncmp(prop.gcnArchName, "gfx1151", 7) == 0; +#else + (void) device; + return false; +#endif +} + // RAII device switch: restores the previous device even on the early-return error paths. struct MixDeviceGuard { int prev = -1; @@ -187,7 +200,7 @@ void mix_register_impl(const void * base, size_t nb02, int n_experts, int out, i const size_t expert_bytes = (size_t) out * (size_t) (in / MIX_QK) * MIX_BLOCK_BYTES; MixEntry ne{base, nb02, expert_bytes, n_experts, out, in, codebooks, modes, - device}; + device, mix_device_is_gfx1151(device)}; for (auto & e : g_mix_registry) { if (e.base == base) { // update in place — free the old owned buffers first mix_free_entry_device(e); @@ -464,6 +477,22 @@ __device__ __forceinline__ void mix_block_accum( // block reads up to 6 B past the tensor end. That is now REJECTED AT REGISTRATION // (see the in % 128 guard in ggml_cuda_rocmfp2_mix_register_host) rather than // left to chance, so this loop can stay branch-free. +#if defined(__HIP_PLATFORM_AMD__) && defined(__gfx1151__) + // gfx1151 supports unaligned flat loads. Read the exact 10-byte payload as + // one 8-byte code word plus the 2-byte metadata tail. The old aligned-floor + // path fetched two overlapping 8-byte windows for every 10-byte block, so + // adjacent lanes requested 16 bytes for 10 bytes of useful weights. Keeping + // the packed code word in one register also preserves the existing decode + // and accumulation order exactly. + uint64_t codes; + uint16_t meta; + MIX_MEMCPY(&codes, b, sizeof(codes)); + MIX_MEMCPY(&meta, b + MIX_QS, sizeof(meta)); + const uint8_t m0 = (uint8_t) meta; + const uint8_t m1 = (uint8_t) (meta >> 8); +#else + // Portable fallback: load the two aligned 8-byte windows bracketing this + // 10-byte block, then extract the payload in registers. const uintptr_t addr = (uintptr_t) b; const uint8_t * base8 = (const uint8_t *) MIX_ASSUME_ALIGNED( (const void *) (addr & ~(uintptr_t) 7), 8); @@ -474,6 +503,7 @@ __device__ __forceinline__ void mix_block_accum( const uint64_t codes = (sh == 0) ? lo : ((lo >> sh) | (hi << (64 - sh))); const uint8_t m0 = (uint8_t) (hi >> sh); // block byte MIX_QS+0 const uint8_t m1 = (uint8_t) (hi >> (sh + 8)); // block byte MIX_QS+1 +#endif if (mode == 0) { const float s0 = mix_ue4m3(m0), s1 = mix_ue4m3(m1); #pragma unroll @@ -677,21 +707,27 @@ __global__ void mix_matvec_rocmfp2_slice_kernel( } } -// FUSE_GLU folds the SECOND mul_mat_id of a DeepSeek4 gate/up pair plus the SwiGLU into this -// launch. The unfused shape is two matvec launches writing two [n_ff_exp, n_used, ntok] -// intermediates, then a third kernel reading both back to apply the GLU. qtype 107 never paid -// that: ggml_cuda_try_fuse_mul_mat_glu collapses the trio into one mul_mat_vec_q, which is why -// the profile showed 107 at 15050 launches against 106's 30100 plus a 28 ms swiglu_ds4 pass. +// GLU_MODE folds a DeepSeek4 gate/up pair plus SwiGLU into the MoE matvec path: +// +// 0: ordinary one-tensor matvec; +// 1: one-pass gate+up matvec and SwiGLU (best for q <= 2); +// 2: gate matvec that reads a preceding up result from dst and overwrites it +// with SwiGLU(gate, up) (paired with mode 0 for wider verification). // -// TEMPLATED rather than a second kernel on purpose: both instantiations run the SAME +// The two-pass form keeps each kernel at the lower register footprint of a +// one-tensor matvec. On gfx1151 the one-pass kernel wins through q=2 by +// avoiding a launch, but loses once wider verification exposes its occupancy cost. +// Both forms remove the old third, standalone SwiGLU launch. +// +// TEMPLATED rather than copied kernels on purpose: all instantiations run the SAME // accumulation over the SAME fixed block order, so each dot product is bit-identical to what -// the unfused path computes, and the fused result is bit-identical to +// the unfused path computes, and either fused result is bit-identical to // swiglu_ds4(unfused_gate, unfused_up). A copy-pasted kernel would only *probably* stay that way. // // Naming follows the mmvq fusion convention: the PRIMARY tensor is `up` (src0 of the surviving // mul_mat_id) and `gate` arrives as the extra operand, because // ggml_cuda_op_swiglu_ds4_single(gate, up, limit) is not symmetric -- silu() is applied to gate. -template +template __global__ void mix_matvec_rocmfp2_moe_kernel( const uint8_t * __restrict__ data, size_t nb02, const nv_bfloat16 * __restrict__ codebooks, const uint8_t * __restrict__ modes, @@ -700,12 +736,14 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( int64_t ids_s0, int64_t ids_s1, // element strides (int32) over slot, token int64_t src1_s1, int64_t src1_s2, // element strides (float) over ne11, token int64_t dst_s1, int64_t dst_s2, // element strides (float) over slot, token - // FUSE_GLU only. The gate tensor's own registry entry -- separate codebooks and modes, + // GLU_MODE=1 only. The gate tensor's own registry entry -- separate codebooks and modes, // NOT assumed equal to up's. Producers may emit identical codebooks for the two // halves, but the kernel does not rely on that and staging both costs 32 B of LDS. const uint8_t * __restrict__ gdata, size_t gnb02, const nv_bfloat16 * __restrict__ gcodebooks, const uint8_t * __restrict__ gmodes, float glu_limit) { + constexpr bool DUAL_GLU = GLU_MODE == 1; + constexpr bool FINALIZE_GLU = GLU_MODE == 2; const int warps_per_block = blockDim.x / MIX_WARP; const int warp = blockIdx.x * warps_per_block + (threadIdx.x / MIX_WARP); const int row0 = warp * 2; // two output rows per warp @@ -727,12 +765,12 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( const bool bad_expert = expert < 0 || expert >= n_experts; // Both tables in one array so the staging stays a single guarded write per thread and one // barrier. Gate's table occupies [2*MIX_K, 4*MIX_K). - __shared__ float s_lut[FUSE_GLU ? 4 * MIX_K : 2 * MIX_K]; + __shared__ float s_lut[DUAL_GLU ? 4 * MIX_K : 2 * MIX_K]; if (!bad_expert && (int) threadIdx.x < 2 * MIX_K) { s_lut[threadIdx.x] = __bfloat162float(codebooks[(int64_t) expert * 2 * MIX_K + threadIdx.x]); } - if (FUSE_GLU) { + if (DUAL_GLU) { const int gt = (int) threadIdx.x - 2 * MIX_K; if (!bad_expert && gt >= 0 && gt < 2 * MIX_K) { s_lut[2 * MIX_K + gt] = @@ -762,11 +800,11 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( : rowbase0; // Gate shares the shape, the expert and the row indices -- only the bytes and the table // differ -- so it reuses `nb`, `two`, `row0` and the same activation column below. - const uint8_t * gedata = FUSE_GLU ? gdata + (int64_t) expert * gnb02 : nullptr; - const int gmode = FUSE_GLU ? (int) gmodes[expert] : 0; - const uint8_t * growbase0 = FUSE_GLU ? gedata + (int64_t) row0 * nb * MIX_BLOCK_BYTES + const uint8_t * gedata = DUAL_GLU ? gdata + (int64_t) expert * gnb02 : nullptr; + const int gmode = DUAL_GLU ? (int) gmodes[expert] : 0; + const uint8_t * growbase0 = DUAL_GLU ? gedata + (int64_t) row0 * nb * MIX_BLOCK_BYTES : nullptr; - const uint8_t * growbase1 = FUSE_GLU ? (two ? gedata + (int64_t) (row0 + 1) * nb * MIX_BLOCK_BYTES + const uint8_t * growbase1 = DUAL_GLU ? (two ? gedata + (int64_t) (row0 + 1) * nb * MIX_BLOCK_BYTES : growbase0) : nullptr; // src1 is [in, ne11, ntok]; the get_rows-equivalent row for (slot, token) @@ -792,7 +830,7 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( const int b = blk + u * MIX_WARP; mix_block_accum(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, s_lut, acc0); mix_block_accum(rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { + if (DUAL_GLU) { mix_block_accum(growbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); mix_block_accum(growbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); } @@ -801,7 +839,7 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( for (; blk < nb; blk += MIX_WARP) { mix_block_accum(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc0); mix_block_accum(rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { + if (DUAL_GLU) { mix_block_accum(growbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); mix_block_accum(growbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); } @@ -810,18 +848,24 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( for (int off = MIX_WARP/2; off > 0; off >>= 1) { acc0 += mix_warp_shfl_down(acc0, off); acc1 += mix_warp_shfl_down(acc1, off); - if (FUSE_GLU) { + if (DUAL_GLU) { gacc0 += mix_warp_shfl_down(gacc0, off); gacc1 += mix_warp_shfl_down(gacc1, off); } } if (lane == 0) { const int64_t o = (int64_t) token * dst_s2 + (int64_t) slot * dst_s1 + row0; - if (FUSE_GLU) { + if (DUAL_GLU) { // The SAME function the standalone swiglu_ds4 kernel applies, on inputs bit-identical // to the ones it would have read back from the two intermediates. dst[o] = ggml_cuda_op_swiglu_ds4_single(gacc0, acc0, glu_limit); if (two) dst[o + 1] = ggml_cuda_op_swiglu_ds4_single(gacc1, acc1, glu_limit); + } else if (FINALIZE_GLU) { + // The preceding mode-0 launch left the up projection in dst. Same-stream + // launch ordering makes it visible here; each lane-0 owns distinct rows, + // so reading and replacing those values needs no extra synchronization. + dst[o] = ggml_cuda_op_swiglu_ds4_single(acc0, dst[o], glu_limit); + if (two) dst[o + 1] = ggml_cuda_op_swiglu_ds4_single(acc1, dst[o + 1], glu_limit); } else { dst[o] = acc0; if (two) dst[o + 1] = acc1; @@ -889,13 +933,13 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_id( n_expert_used <= 0 || n_tokens <= 0 || ne11 <= 0) { return false; } - const int warps_per_block = 2; // 64 threads (mirror the mmvq path) + const int warps_per_block = e.gfx1151 ? (n_tokens <= 2 ? 8 : 4) : 2; const int threads = warps_per_block * MIX_WARP; // Two output rows per warp (register-blocked activation reuse), so a workgroup // of `warps_per_block` warps covers 2*warps_per_block rows. const int rows_per_block = 2 * warps_per_block; dim3 grid((out + rows_per_block - 1) / rows_per_block, n_expert_used, n_tokens); - mix_matvec_rocmfp2_moe_kernel<<>>( + mix_matvec_rocmfp2_moe_kernel<0><<>>( (const uint8_t *) e.base, e.nb02, e.codebooks, e.modes, src1, ids, dst, in, out, e.n_experts, ne11, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, @@ -931,15 +975,32 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_id_glu( n_expert_used <= 0 || n_tokens <= 0 || ne11 <= 0) { return false; // not a matched pair; the caller's two-launch path is still correct } - const int warps_per_block = 2; + const bool strix_tuned = eu.gfx1151 && eg.gfx1151; + const int warps_per_block = strix_tuned ? (n_tokens <= 2 ? 8 : 4) : 2; const int threads = warps_per_block * MIX_WARP; const int rows_per_block = 2 * warps_per_block; dim3 grid((out + rows_per_block - 1) / rows_per_block, n_expert_used, n_tokens); - mix_matvec_rocmfp2_moe_kernel<<>>( - (const uint8_t *) eu.base, eu.nb02, eu.codebooks, eu.modes, - src1, ids, dst, in, out, eu.n_experts, ne11, - ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, - (const uint8_t *) eg.base, eg.nb02, eg.codebooks, eg.modes, glu_limit); + if (!strix_tuned || n_tokens <= 2) { + mix_matvec_rocmfp2_moe_kernel<1><<>>( + (const uint8_t *) eu.base, eu.nb02, eu.codebooks, eu.modes, + src1, ids, dst, in, out, eu.n_experts, ne11, + ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, + (const uint8_t *) eg.base, eg.nb02, eg.codebooks, eg.modes, glu_limit); + } else { + // A wide verifier makes the dual-tensor kernel's register pressure more + // expensive than its saved launch. Keep both projections at the ordinary + // kernel's occupancy and fold SwiGLU into the second launch in-place. + mix_matvec_rocmfp2_moe_kernel<0><<>>( + (const uint8_t *) eu.base, eu.nb02, eu.codebooks, eu.modes, + src1, ids, dst, in, out, eu.n_experts, ne11, + ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, + nullptr, 0, nullptr, nullptr, 0.0f); + mix_matvec_rocmfp2_moe_kernel<2><<>>( + (const uint8_t *) eg.base, eg.nb02, eg.codebooks, eg.modes, + src1, ids, dst, in, out, eg.n_experts, ne11, + ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, + nullptr, 0, nullptr, nullptr, glu_limit); + } return true; } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu index df10f4c38..d6905334f 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp3_mix.cu @@ -276,10 +276,18 @@ static bool mix_lookup_expert_base( } __device__ __forceinline__ float mix_ue4m3(uint8_t e) { +#if defined(__HIP_PLATFORM_AMD__) && defined(__gfx1151__) + int exp = e >> 3, mant = e & 7; + const float normal = ldexpf((float) (8 + mant), exp - 11); + const float sub = (float) mant * 0.0009765625f; // 2^-10 + const float value = (exp == 0) ? sub : normal; + return (e > 0x7E) ? 0.0f : value; +#else if (e > 0x7E) return 0.0f; int exp = e >> 3, mant = e & 7; if (exp == 0) return (float) mant * 0.0009765625f; // 2^-10 return ldexpf((float) (8 + mant), exp - 11); +#endif } __device__ __forceinline__ uint32_t mix_fp3_code(const uint8_t * qs, int i) { @@ -290,6 +298,22 @@ __device__ __forceinline__ uint32_t mix_fp3_code(const uint8_t * qs, int i) { return (v >> shift) & 7u; } +struct MixFp3Words { + uint64_t lo; // code bytes 0..7 + uint32_t hi; // code bytes 8..11 +}; + +// Decode the same 96-bit fp3 stream from three registers. The only code that +// crosses the 64-bit boundary is i=21 (bit 63); the unrolled caller makes all +// three cases compile-time constants. This avoids materializing a 14-byte +// private array in the gfx1151 matvec while preserving every packed bit. +__device__ __forceinline__ uint32_t mix_fp3_code(const MixFp3Words & qs, int i) { + const int bit = 3 * i; + if (bit <= 60) return (uint32_t) (qs.lo >> bit) & 7u; + if (bit == 63) return (uint32_t) ((qs.lo >> 63) | ((uint64_t) qs.hi << 1)) & 7u; + return (qs.hi >> (bit - 64)) & 7u; +} + __device__ __forceinline__ float mix_fp3_fixed(uint32_t code) { uint32_t m = code & 3u; int mag = (m == 3u) ? 4 : (int) m; @@ -373,6 +397,14 @@ void dequantize_rocmfp3_mix_to_fp16_cuda(const void * vx, half * y, int64_t k, c // quantized blocks once avoids the ~10x f16 round-trip of the dequant fallback. #define MIX_WARP 32 #define MIX_UNROLL 4 +#if defined(__HIP_PLATFORM_AMD__) && defined(__gfx1151__) +// The two-row MoE kernel already exposes thousands of independent workgroups. +// Keeping only one strided block live at a time cuts the large q3 decode body +// and is faster on gfx1151 than duplicating it four times for local MLP. +#define MIX_MOE_UNROLL 1 +#else +#define MIX_MOE_UNROLL MIX_UNROLL +#endif // Down-shift warp shuffle confined to a 32-lane logical group. width=MIX_WARP // keeps the reduction self-contained on wave64 (GFX8/9, physical wave = 64) and @@ -409,16 +441,25 @@ __device__ __forceinline__ void mix_block_accum( // decode arithmetic and the fixed j accumulation order are untouched, so // acc is bit-for-bit identical to the per-byte path (the correctness gate // hashes the greedy output; any reassociation flips a token). +#if defined(__HIP_PLATFORM_AMD__) && defined(__gfx1151__) + MixFp3Words qs; + uint16_t meta; + MIX_MEMCPY(&qs.lo, b, sizeof(qs.lo)); + MIX_MEMCPY(&qs.hi, b + sizeof(qs.lo), sizeof(qs.hi)); + MIX_MEMCPY(&meta, b + MIX_QS, sizeof(meta)); + const uint8_t m0 = (uint8_t) meta, m1 = (uint8_t) (meta >> 8); +#else const uint8_t * ba = (const uint8_t *) MIX_ASSUME_ALIGNED(b, 2); - uint8_t buf[MIX_BLOCK_BYTES]; - MIX_MEMCPY(buf, ba, MIX_BLOCK_BYTES); - const uint8_t m0 = buf[MIX_QS + 0], m1 = buf[MIX_QS + 1]; + uint8_t qs[MIX_BLOCK_BYTES]; + MIX_MEMCPY(qs, ba, MIX_BLOCK_BYTES); + const uint8_t m0 = qs[MIX_QS + 0], m1 = qs[MIX_QS + 1]; +#endif if (mode == 0) { const float s0 = mix_ue4m3(m0), s1 = mix_ue4m3(m1); #pragma unroll for (int j = 0; j < MIX_QK; ++j) { const float s = (j < MIX_QK/2) ? s0 : s1; - acc += s * mix_fp3_fixed(mix_fp3_code(buf, j)) * xc[col0 + j]; + acc += s * mix_fp3_fixed(mix_fp3_code(qs, j)) * xc[col0 + j]; } } else { const float s0 = mix_ue4m3(m0 & 0x7F), s1 = mix_ue4m3(m1 & 0x7F); @@ -435,7 +476,7 @@ __device__ __forceinline__ void mix_block_accum( for (int j = 0; j < MIX_QK; ++j) { const float s = (j < MIX_QK/2) ? s0 : s1; const float * bk = (j < MIX_QK/2) ? bk0 : bk1; - acc += s * bk[mix_fp3_code(buf, j)] * xc[col0 + j]; + acc += s * bk[mix_fp3_code(qs, j)] * xc[col0 + j]; } } } @@ -689,32 +730,21 @@ __global__ void mix_matvec_rocmfp3_moe_kernel( float acc0 = 0.0f, acc1 = 0.0f; float gacc0 = 0.0f, gacc1 = 0.0f; // same block order as acc*, so bit-identical per row int blk = lane; - for (; blk + 3 * MIX_WARP < nb; blk += MIX_UNROLL * MIX_WARP) { - const int b0 = blk, b1 = blk + MIX_WARP; - const int b2 = blk + 2 * MIX_WARP, b3 = blk + 3 * MIX_WARP; - mix_block_accum(rowbase0 + (int64_t) b0 * MIX_BLOCK_BYTES, xcol, b0 * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b0 * MIX_BLOCK_BYTES, xcol, b0 * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { - mix_block_accum(growbase0 + (int64_t) b0 * MIX_BLOCK_BYTES, xcol, b0 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) b0 * MIX_BLOCK_BYTES, xcol, b0 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); - } - mix_block_accum(rowbase0 + (int64_t) b1 * MIX_BLOCK_BYTES, xcol, b1 * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b1 * MIX_BLOCK_BYTES, xcol, b1 * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { - mix_block_accum(growbase0 + (int64_t) b1 * MIX_BLOCK_BYTES, xcol, b1 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) b1 * MIX_BLOCK_BYTES, xcol, b1 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); - } - mix_block_accum(rowbase0 + (int64_t) b2 * MIX_BLOCK_BYTES, xcol, b2 * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b2 * MIX_BLOCK_BYTES, xcol, b2 * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { - mix_block_accum(growbase0 + (int64_t) b2 * MIX_BLOCK_BYTES, xcol, b2 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) b2 * MIX_BLOCK_BYTES, xcol, b2 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); - } - mix_block_accum(rowbase0 + (int64_t) b3 * MIX_BLOCK_BYTES, xcol, b3 * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b3 * MIX_BLOCK_BYTES, xcol, b3 * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { - mix_block_accum(growbase0 + (int64_t) b3 * MIX_BLOCK_BYTES, xcol, b3 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) b3 * MIX_BLOCK_BYTES, xcol, b3 * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); + for (; blk + (MIX_MOE_UNROLL - 1) * MIX_WARP < nb; + blk += MIX_MOE_UNROLL * MIX_WARP) { + #pragma unroll + for (int u = 0; u < MIX_MOE_UNROLL; ++u) { + const int b = blk + u * MIX_WARP; + mix_block_accum(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, + b * MIX_QK, mode, s_lut, acc0); + mix_block_accum(rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, + b * MIX_QK, mode, s_lut, acc1); + if (FUSE_GLU) { + mix_block_accum(growbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, + b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); + mix_block_accum(growbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, + b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); + } } } for (; blk < nb; blk += MIX_WARP) { @@ -806,7 +836,7 @@ bool ggml_cuda_rocmfp3_mix_mul_mat_id( n_expert_used <= 0 || n_tokens <= 0 || ne11 <= 0) { return false; } - const int warps_per_block = 2; // 64 threads (mirror the mmvq path) + const int warps_per_block = 2; const int threads = warps_per_block * MIX_WARP; // Two output rows per warp (register-blocked activation reuse), so a workgroup // of `warps_per_block` warps covers 2*warps_per_block rows. diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cu index 7d3113a93..fbc7c29eb 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rope.cu @@ -103,7 +103,8 @@ static __global__ void rope_norm(const T * x, const float theta_scale, const float * freq_factors, const int64_t * row_indices, - const int set_rows_stride) { + const int set_rows_stride, + const int rot_offset) { const int i0 = 2*(blockDim.y*blockIdx.y + threadIdx.y); if (i0 >= ne00) { @@ -134,19 +135,24 @@ static __global__ void rope_norm(const T * x, ggml_cuda_memcpy_1<4>(dst + idst, &v); } }; - if (i0 >= n_dims) { + // rot_offset is 0 for the standard head rotation and ne00 - n_dims for + // GGML_ROPE_TYPE_TAIL; j0 is the element's index inside the rotated span, + // so a tail element sees exactly the angle it would see in an extracted + // n_dims-wide tail tensor. + if (i0 < rot_offset || i0 >= rot_offset + n_dims) { store_coaelsced(x[ix + 0], x[ix + 1]); return; } + const int j0 = i0 - rot_offset; - const double theta_base = rope_theta_fp64(pos[i2], theta_scale, i0/2); + const double theta_base = rope_theta_fp64(pos[i2], theta_scale, j0/2); - const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f; + const float freq_factor = has_ff ? freq_factors[j0/2] : 1.0f; float cos_theta; float sin_theta; - rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta); + rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, j0, ext_factor, attn_factor, cos_theta, sin_theta); const float x0 = x[ix + 0]; const float x1 = x[ix + 1]; @@ -392,6 +398,7 @@ static void rope_norm_cuda(const T * x, const float * freq_factors, const int64_t * row_indices, const int set_rows_stride, + const int rot_offset, cudaStream_t stream) { GGML_ASSERT(ne00 % 2 == 0); const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1); @@ -403,11 +410,11 @@ static void rope_norm_cuda(const T * x, if (freq_factors == nullptr) { rope_norm<<>>( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); + attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, rot_offset); } else { rope_norm<<>>( x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); + attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, rot_offset); } } @@ -605,10 +612,14 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, memcpy(&beta_slow, (int32_t *) dst->op_params + 10, sizeof(float)); memcpy(§ions.v, (int32_t *) dst->op_params + 11, sizeof(int)*4); - const bool is_neox = mode & GGML_ROPE_TYPE_NEOX; - const bool is_mrope = mode & GGML_ROPE_TYPE_MROPE; - const bool is_imrope = mode == GGML_ROPE_TYPE_IMROPE; - const bool is_vision = mode == GGML_ROPE_TYPE_VISION; + const bool is_tail = mode & GGML_ROPE_TYPE_TAIL; + const int mode_base = mode & ~GGML_ROPE_TYPE_TAIL; + const bool is_neox = mode_base & GGML_ROPE_TYPE_NEOX; + const bool is_mrope = mode_base & GGML_ROPE_TYPE_MROPE; + const bool is_imrope = mode_base == GGML_ROPE_TYPE_IMROPE; + const bool is_vision = mode_base == GGML_ROPE_TYPE_VISION; + GGML_ASSERT(!is_tail || mode_base == GGML_ROPE_TYPE_NORMAL); + const int rot_offset = is_tail ? (int) ne00 - n_dims : 0; if (is_mrope) { GGML_ASSERT(sections.v[0] > 0 || sections.v[1] > 0 || sections.v[2] > 0); @@ -677,17 +688,17 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, rope_norm_cuda((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, rot_offset, stream); } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) { rope_norm_cuda((const float *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, rot_offset, stream); } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) { rope_norm_cuda((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, rot_offset, stream); } else { GGML_ABORT("fatal error"); } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/softmax.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/softmax.cu index 285c0e954..1ed928d28 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/softmax.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/softmax.cu @@ -43,6 +43,10 @@ struct soft_max_params { float max_bias; float m0; float m1; + // Lucebox sink-as-virtual-column mode: ncols counts one extra virtual + // column per row whose value is sinks[row]; x and dst rows are ncols - 1 + // wide. Reductions then match scale -> concat(sink) -> soft_max exactly. + int sink_col; }; // When ncols_template == 0 the bounds for the loops in this function are not known and can't be unrolled. @@ -69,9 +73,11 @@ static __global__ void soft_max_f32( const int64_t i12 = i02 % p.ne12; const int64_t i13 = i03 % p.ne13; - x += int64_t(rowx)*ncols; + const int ncols_x = p.sink_col ? ncols - 1 : ncols; // width of the x and dst rows + + x += int64_t(rowx)*ncols_x; mask += (i11*p.nb11 + i12*p.nb12 + i13*p.nb13) / sizeof(T) * (mask != nullptr); - dst += int64_t(rowx)*ncols; + dst += int64_t(rowx)*ncols_x; const int block_size = block_size_template == 0 ? blockDim.x : block_size_template; @@ -82,7 +88,7 @@ static __global__ void soft_max_f32( // shared memory buffer to cache values between iterations: float * vals = use_shared ? buf_iw + WARP_SIZE : dst; - float max_val = sinks ? sinks[i02] : -INFINITY; + float max_val = (sinks && !p.sink_col) ? sinks[i02] : -INFINITY; #pragma unroll for (int col0 = 0; col0 < ncols; col0 += block_size) { @@ -92,9 +98,14 @@ static __global__ void soft_max_f32( break; } - const float val = x[col]*p.scale + (mask ? slope*t2f32(mask[col]) : 0.0f); + const float val = (p.sink_col && col == ncols_x) + ? sinks[rowx] + : x[col]*p.scale + (mask ? slope*t2f32(mask[col]) : 0.0f); - vals[col] = val; + // The global-memory cache is dst, which has no slot for the virtual sink. + if (use_shared || col < ncols_x) { + vals[col] = val; + } max_val = max(max_val, val); } @@ -111,15 +122,18 @@ static __global__ void soft_max_f32( break; } - const float val = expf(vals[col] - max_val); + const float cached = (use_shared || col < ncols_x) ? vals[col] : sinks[rowx]; + const float val = expf(cached - max_val); tmp += val; - vals[col] = val; + if (use_shared || col < ncols_x) { + vals[col] = val; + } } // find the sum of exps in the block tmp = block_reduce(tmp, buf_iw); - if (sinks) { + if (sinks && !p.sink_col) { tmp += expf(sinks[i02] - max_val); } @@ -133,6 +147,9 @@ static __global__ void soft_max_f32( return; } + if (p.sink_col && col == ncols_x) { + continue; // the virtual sink column has no output + } dst[col] = vals[col] * inv_sum; } } @@ -435,6 +452,14 @@ void ggml_cuda_op_soft_max(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { params.max_bias = max_bias; params.m0 = m0; params.m1 = m1; + int32_t sink_col = 0; + memcpy(&sink_col, (const int32_t *) dst->op_params + 2, sizeof(int32_t)); + params.sink_col = sink_col; + if (sink_col) { + GGML_ASSERT(src1 == nullptr && src2 != nullptr); + GGML_ASSERT(ggml_nelements(src2) == nrows_x); + params.ncols = ne00 + 1; // one virtual sink column per row + } if (use_f16) { soft_max_f32_cuda(src0_d, (const half *) src1_d, (const float *) src2_d, dst_d, params, stream, ctx); diff --git a/server/deps/llama.cpp/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/server/deps/llama.cpp/ggml/src/ggml-hexagon/ggml-hexagon.cpp index 3d68b8004..b10602691 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -2629,7 +2629,9 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) { case GGML_OP_SQRT: return HTP_OP_SQRT; case GGML_OP_SOFT_MAX: return HTP_OP_SOFTMAX; case GGML_OP_SSM_CONV: return HTP_OP_SSM_CONV; - case GGML_OP_ROPE: return HTP_OP_ROPE; + case GGML_OP_ROPE: + GGML_ASSERT((t->op_params[2] & GGML_ROPE_TYPE_TAIL) == 0); + return HTP_OP_ROPE; case GGML_OP_REPEAT: return HTP_OP_REPEAT; case GGML_OP_CUMSUM: return HTP_OP_CUMSUM; @@ -3030,6 +3032,11 @@ static bool ggml_hexagon_supported_repeat(const struct ggml_hexagon_session * se } static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) { + if ((op->op == GGML_OP_ROPE || op->op == GGML_OP_ROPE_BACK) && + (op->op_params[2] & GGML_ROPE_TYPE_TAIL)) { + return false; + } + auto sess = static_cast(dev->context); // reject ops that match the filter diff --git a/server/deps/llama.cpp/ggml/src/ggml-metal/ggml-metal-device.m b/server/deps/llama.cpp/ggml/src/ggml-metal/ggml-metal-device.m index d148283ca..b4f086c89 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-metal/ggml-metal-device.m +++ b/server/deps/llama.cpp/ggml/src/ggml-metal/ggml-metal-device.m @@ -994,6 +994,11 @@ void ggml_metal_device_get_memory(ggml_metal_device_t dev, size_t * free, size_t } bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_tensor * op) { + if ((op->op == GGML_OP_ROPE || op->op == GGML_OP_ROPE_BACK) && + (op->op_params[2] & GGML_ROPE_TYPE_TAIL)) { + return false; + } + const bool has_simdgroup_mm = dev->props.has_simdgroup_mm; const bool has_simdgroup_reduction = dev->props.has_simdgroup_reduction; const bool has_bfloat = dev->props.has_bfloat; @@ -1094,10 +1099,13 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te case GGML_OP_SUM_ROWS: case GGML_OP_CUMSUM: case GGML_OP_MEAN: - case GGML_OP_SOFT_MAX: case GGML_OP_GROUP_NORM: case GGML_OP_L2_NORM: return has_simdgroup_reduction && ggml_is_contiguous_rows(op->src[0]); + case GGML_OP_SOFT_MAX: + // Sink-column mode is handled by a supporting backend. + return ggml_get_op_params_i32(op, 2) == 0 && + has_simdgroup_reduction && ggml_is_contiguous_rows(op->src[0]); case GGML_OP_COUNT_EQUAL: return has_simdgroup_reduction && op->src[0]->type == GGML_TYPE_I32 && diff --git a/server/deps/llama.cpp/ggml/src/ggml-opencl/ggml-opencl.cpp b/server/deps/llama.cpp/ggml/src/ggml-opencl/ggml-opencl.cpp index a58140230..9896505fd 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -4013,6 +4013,8 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te case GGML_OP_CLAMP: return op->src[0]->type == GGML_TYPE_F32; case GGML_OP_SOFT_MAX: + // Sink-column mode is handled by a supporting backend. + return ggml_get_op_params_i32(op, 2) == 0; case GGML_OP_NORM: return true; case GGML_OP_RMS_NORM: @@ -6186,6 +6188,11 @@ static ggml_backend_buffer_t ggml_backend_opencl_device_buffer_from_ptr(ggml_bac } static bool ggml_backend_opencl_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) { + if ((op->op == GGML_OP_ROPE || op->op == GGML_OP_ROPE_BACK) && + (op->op_params[2] & GGML_ROPE_TYPE_TAIL)) { + return false; + } + return ggml_opencl_supports_op(dev, op); } @@ -12925,6 +12932,7 @@ static void ggml_cl_rope(ggml_backend_t backend, const ggml_tensor * src0, const const int n_past = ((int *) dst->op_params)[0]; const int n_dims = ((int *) dst->op_params)[1]; const int mode = ((int *) dst->op_params)[2]; + GGML_ASSERT((mode & GGML_ROPE_TYPE_TAIL) == 0); const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; float freq_base; diff --git a/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp b/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp index 1e28cf3ea..fd147bd96 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -4665,6 +4665,11 @@ static ggml_backend_buffer_t ggml_backend_sycl_device_buffer_from_host_ptr(ggml_ } static bool ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + if ((op->op == GGML_OP_ROPE || op->op == GGML_OP_ROPE_BACK) && + (op->op_params[2] & GGML_ROPE_TYPE_TAIL)) { + return false; + } + ggml_backend_sycl_device_context *sycl_ctx = (ggml_backend_sycl_device_context *)dev->context; int device = sycl_ctx->device; @@ -4913,7 +4918,8 @@ static bool ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_DIAG_MASK_INF: return true; case GGML_OP_SOFT_MAX: - return true; + // Sink-column mode is handled by a supporting backend. + return ggml_get_op_params_i32(op, 2) == 0; case GGML_OP_SOFT_MAX_BACK: { float max_bias = 0.0f; memcpy(&max_bias, (const float *) op->op_params + 1, sizeof(float)); diff --git a/server/deps/llama.cpp/ggml/src/ggml-sycl/rope.cpp b/server/deps/llama.cpp/ggml/src/ggml-sycl/rope.cpp index 9d83a1e9f..c679aacb0 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-sycl/rope.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-sycl/rope.cpp @@ -496,6 +496,7 @@ void ggml_sycl_op_rope_impl(ggml_backend_sycl_context &ctx, ggml_tensor *dst, const int n_dims = ((int32_t *)dst->op_params)[1]; const int mode = ((int32_t *)dst->op_params)[2]; + GGML_ASSERT((mode & GGML_ROPE_TYPE_TAIL) == 0); const int n_ctx_orig = ((int32_t *)dst->op_params)[4]; mrope_sections sections; diff --git a/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 6bc81a91d..78d5ddce3 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -10999,6 +10999,7 @@ static uint32_t ggml_vk_rms_partials_size(ggml_backend_vk_context * ctx, const g static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor *dst, const ggml_tensor *src0, const bool has_ff, bool backprop, const uint32_t set_rows_stride) { const int n_dims = ((const int32_t *) dst->op_params)[1]; const int mode = ((const int32_t *) dst->op_params)[2]; + GGML_ASSERT((mode & GGML_ROPE_TYPE_TAIL) == 0); // const int n_ctx = ((const int32_t *) dst->op_params)[3]; const int n_ctx_orig = ((const int32_t *) dst->op_params)[4]; const float freq_base = ((const float *) dst->op_params)[5]; @@ -15304,6 +15305,11 @@ static ggml_backend_t ggml_backend_vk_device_init(ggml_backend_dev_t dev, const } static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + if ((op->op == GGML_OP_ROPE || op->op == GGML_OP_ROPE_BACK) && + (op->op_params[2] & GGML_ROPE_TYPE_TAIL)) { + return false; + } + ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context; const vk_device& device = ggml_vk_get_device(ctx->device); @@ -15722,7 +15728,8 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_OP_DIAG_MASK_INF: return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32; case GGML_OP_SOFT_MAX: - return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32 + // Sink-column mode requires per-row sinks and an unscaled sink logit. + return ggml_get_op_params_i32(op, 2) == 0 && ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32 && (!op->src[1] || (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16)); case GGML_OP_SOFT_MAX_BACK: return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32 diff --git a/server/deps/llama.cpp/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/server/deps/llama.cpp/ggml/src/ggml-webgpu/ggml-webgpu.cpp index aa3fe06d5..bf7a66bd2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-webgpu/ggml-webgpu.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-webgpu/ggml-webgpu.cpp @@ -3838,7 +3838,8 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const supports_op = op->type == GGML_TYPE_F32; break; case GGML_OP_SOFT_MAX: - supports_op = op->type == GGML_TYPE_F32; + // Sink-column mode is handled by a supporting backend. + supports_op = ggml_get_op_params_i32(op, 2) == 0 && op->type == GGML_TYPE_F32; break; case GGML_OP_UNARY: { diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 82e2d9a1c..742a5b2fc 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -4266,6 +4266,33 @@ struct ggml_tensor * ggml_soft_max_ext_inplace( return ggml_soft_max_impl(ctx, a, mask, scale, max_bias, true); } +struct ggml_tensor * ggml_soft_max_ext_sink_col( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * sinks, + float scale) { + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(a->type == GGML_TYPE_F32); + GGML_ASSERT(sinks && sinks->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(sinks)); + GGML_ASSERT(ggml_nelements(sinks) == ggml_nrows(a)); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + int32_t params[3] = { 0, 0, /*sink_col*/ 1 }; + const float max_bias = 0.0f; + memcpy(params + 0, &scale, sizeof(float)); + memcpy(params + 1, &max_bias, sizeof(float)); + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_SOFT_MAX; + result->src[0] = a; + result->src[1] = NULL; + result->src[2] = sinks; + + return result; +} + void ggml_soft_max_add_sinks( struct ggml_tensor * a, struct ggml_tensor * sinks) { diff --git a/server/docs/DS4.md b/server/docs/DS4.md index 11be8699d..9e181243f 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -311,6 +311,93 @@ performance profile held 48.1 tok/s median on the deterministic 128-token workload. The all-6-expert reference-exact mode is a correctness profile, not a throughput profile. +### Strix Halo concurrent serving + +DeepSeek4 paged concurrency supports two resident HIP deployments: + +- one monolithic Strix Halo (`gfx1151`) target with every expert on that + device, through 6 lanes; and +- the in-process R9700 (`gfx1201`) target + Strix Halo expert-parallel + deployment, qualified at concurrency 1–4 with the grouped expert setting + below. The target keeps dense work and its selected experts while the + secondary owns the remaining materialized experts. The runtime accepts up + to 6 lanes, but concurrency 5–6 is not qualified for this configuration. + +The heterogeneous mode is route-level expert parallelism. It is not an +explicit `--target-device hip:0,hip:1` layer split and does not use a remote +target shard or host-streamed experts. + +The backend keeps raw MLA rows, compressed rows, and indexer state in a +persistent 128-token paged cache. The shared +HTTP scheduler performs admission, cancellation, slow-client isolation, and +fair continuous batching. DeepSeek4 lowers each scheduler plan into one exact +gathered graph with up to 6 independent lanes. Decode rows share the weight +pass; each selected prompt advances by one exact token because the graph must +not contain two rows from the same sequence. + +```bash +hf download Lucebox/DeepSeek-V4-Flash-0731-ROCmFP3 \ + DeepSeek-V4-Flash-0731-ROCMFPX-MIX-STRIX.gguf \ + --local-dir /path/to/models + +cmake -S server -B server/build-hip \ + -DDFLASH27B_GPU_BACKEND=hip \ + -DDFLASH27B_HIP_ARCHITECTURES=gfx1151 \ + -DDFLASH27B_SERVER=ON \ + -DGGML_HIP_GRAPHS=ON \ + -DCMAKE_BUILD_TYPE=Release +cmake --build server/build-hip -j + +./server/build-hip/dflash_server \ + /path/to/models/DeepSeek-V4-Flash-0731-ROCMFPX-MIX-STRIX.gguf \ + --target-device hip:0 \ + --paged-attention \ + --max-concurrency 6 \ + --kv-pool-tokens 24576 \ + --max-ctx 4096 \ + --ds4-prefill exact \ + --prefix-cache-slots 0 +``` + +For the R9700 + Strix Halo path, build `server/build-hip-dual` for +`gfx1151;gfx1201` as shown in +[In-process heterogeneous expert parallel](#in-process-heterogeneous-expert-parallel). +Then expose the R9700 first and select the static in-process expert split. +The grouped expert setting below is qualified at concurrency 1–4 with +ROCmFP2 gate/up and ROCmFP3 down weights: + +```bash +export HIP_VISIBLE_DEVICES=, +export DFLASH_DS4_MOE_TP=1 +export DFLASH_DS4_MOE_TP_INPROC=1 +export DFLASH_DS4_MOE_TP_GPU=1 +export DFLASH_EXPERT_BUDGET_MB=11700 +export DFLASH_DS4_TP_BATCH_SPLIT_COPIES=1 +export DFLASH_DS4_TP_GROUPED_MMVQ=1 + +./server/build-hip-dual/dflash_server \ + /path/to/models/DeepSeek-V4-Flash.gguf \ + --target-device hip:0 \ + --peer-access \ + --paged-attention \ + --max-concurrency 4 \ + --kv-pool-tokens 24576 \ + --max-ctx 4096 \ + --ds4-prefill exact \ + --prefix-cache-slots 0 +``` + +The heterogeneous paged cache is charged against the R9700 before selecting +resident experts. Increase `--kv-pool-tokens` only when the primary has enough +memory for the larger +shared history pool. + +Paged concurrency fails closed for other primary/secondary architecture pairs, +CUDA or out-of-process expert ownership, explicit layer or remote target +splits, drafts/DSpark, DDTree, PFlash/KVFlash, fused decode, approximate +prefill, windowed attention, mutable expert caching, and prefix-cache parking. +There is no automatic fallback to a slower or asymmetric execution mode. + ### Local single-shard If the adapter decides all 43 layers fit on one CUDA GPU, it loads a single shard locally and no IPC daemon is involved. @@ -367,7 +454,7 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. | Variable | Purpose | |----------|---------| | `DFLASH_DS4_CUDA_LAYERS` | Override the auto-split heuristic and pin the first `N` DeepSeek4 layers to CUDA. The remaining `43 - N` layers run on the Halo shard. | -| `DFLASH_DS4_TIMING` | Enable DS4 timing logs for the layer-split parent and target-shard daemon. Useful for profiling prefill/decode breakdowns; leave unset for normal runs. | +| `DFLASH_DS4_TIMING` | Enable DS4 timing logs for local, paged, and layer-split execution. Paged rounds report full-graph build, input upload, compute, and readback time; leave unset for normal runs. | | `DFLASH_DS4_ROCTX` | HIP-only, default-off semantic ROCTX ranges for an external rocprof trace. The library is loaded dynamically only when set to `1`, `true`, `yes`, or `on`. | | `DFLASH_DS4_SPEC` / `DFLASH_DS4_DRAFT` | Enable DSpark and select its GGUF. | | `DFLASH_DS4_DRAFT_BACKEND` / `DFLASH_DS4_DRAFT_GPU` | Backend and device for the in-process drafter. | @@ -385,12 +472,16 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. | `DFLASH_CUDA_BACKEND_PATH` / `DFLASH_HIP_BACKEND_PATH` | Optional explicit peer backend module path. | | `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_GROUPED_MMVQ` | Opt in to grouped expert MMVQ for `n_tokens > 1`, replacing tokenwise ROCmFP2 gate/up dispatch. The paged R9700 + Strix profile is qualified at concurrency 1–4; the flag itself does not enforce topology or lane limits. `DFLASH_MOE_TP_GROUPED_MMVQ` is the model-neutral name and takes precedence. | +| `DFLASH_DS4_TP_BATCH_SPLIT_COPIES` | Establish destination readiness once per scheduler split without combining backend copy dependencies. The qualified dual-ROCm launcher enables this exact path. | +| `GGML_BATCH_PEER_COPIES` | Additionally combine HIP peer-copy dependency publication. The old `GGML_CUDA_BATCH_PEER_COPIES` spelling remains an alias. Keep these event-batching variables unset for the exact qualified profile. | | `DFLASH_DS4_TP_CRITICAL_PATH_PLACEMENT` | Use the routing profile and measured owner-rate ratio to minimize the predicted two-owner MoE critical path instead of maximizing aggregate hot-hit rate. Requires `DFLASH_DS4_HOTNESS_CSV`. | | `DFLASH_DS4_TP_MAIN_TO_PEER_RATE` | Relative main/peer routed-expert rate used by critical-path placement. It must be finite and greater than zero; the default is `3.4`. | | `DFLASH_DS4_TP_BALANCE_MIN_HOT` | Minimum hot experts retained on every routed layer by critical-path placement. Defaults to `0`. | | `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_CUDA_MMVQ_FP4_X4` | Enable the four-column ROCmFP4 dense x4 kernel and permit the five-column x4+1 kernel. Defaults to `1` for monolithic `gfx1151` paged serving and opt-in HIP q=5 verification; set `0` to restore generic four- and five-column dispatch. | +| `DFLASH_CUDA_MMVQ_FP4_Q5_X4_PLUS1` | Enable five-column ROCmFP4 dense x4+1 dispatch when `DFLASH_CUDA_MMVQ_FP4_X4=1`. Defaults to `1` for monolithic `gfx1151` paged serving and the opt-in q=5 verifier on `gfx1201`; set `0` to restore the generic five-column kernel. | +| `DFLASH_CUDA_MMVQ_MOE_FP3_PACKED24` | Enable packed 24-bit ROCmFP3 expert dispatch. Defaults to `1` for monolithic `gfx1151` paged serving; set `0` to restore generic ROCmFP3 expert dispatch. Other configurations require explicit opt-in. | | `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. | @@ -398,6 +489,7 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. `DFLASH_DS4_TIMING` enables the existing timing banners: +- local and paged serving: `[deepseek4-timing]`. Paged rounds use `paged-prefill`, `paged-decode`, or `paged-mixed` and aggregate all lanes into `full_build`, `full_set`, `full_compute`, and `full_read`. - parent / local shard: `[deepseek4-split-timing]` - remote Halo shard: `[deepseek4-target-timing]` diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index d83371e66..26ad0ba4c 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -38,7 +38,12 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_DS4_DISABLE_GROUPED_OUTPUT_PROJECTION` | unset | DEBUG: restore the materialized output projection when diagnosing grouped-view copies across unlike runtimes. | | `DFLASH_DS4_DRAFT_BACKEND` / `DFLASH_DS4_DRAFT_GPU` | compiled backend / target device | Select the in-process DSpark backend and device. | | `DFLASH_CUDA_BACKEND_PATH` / `DFLASH_HIP_BACKEND_PATH` | auto-discovered beside the executable | Explicit peer module file path for a mixed CUDA+HIP build. | -| `GGML_BATCH_PEER_COPIES` | unset | BURN-IN: batch peer-runtime copies and unlike-runtime host staging with one source wait per split. `GGML_CUDA_BATCH_PEER_COPIES` remains a compatibility alias. | +| `DFLASH_DS4_TP_GROUPED_MMVQ` / `DFLASH_MOE_TP_GROUPED_MMVQ` | unset | OPT-IN: grouped expert MMVQ for `n_tokens > 1` instead of tokenwise ROCmFP2 gate/up dispatch. Qualified for paged R9700 + Strix at concurrency 1–4; the flag itself does not enforce topology or lane limits. The model-neutral name takes precedence. | +| `DFLASH_CUDA_MMVQ_FP4_X4` | 1 for monolithic DS4 `gfx1151` paged serving and opt-in HIP DS4 q=5 verification; unset otherwise | Enable dense ROCmFP4 x4 dispatch. Set `0` to restore generic four- and five-column kernels. | +| `DFLASH_CUDA_MMVQ_FP4_Q5_X4_PLUS1` | 1 for monolithic DS4 `gfx1151` paged serving and opt-in DS4 q=5 verification on `gfx1201`; unset otherwise | Enable dense five-column x4+1 dispatch when `DFLASH_CUDA_MMVQ_FP4_X4=1`. Set `0` to restore the generic five-column kernel. | +| `DFLASH_CUDA_MMVQ_MOE_FP3_PACKED24` | 1 for monolithic DS4 `gfx1151` paged serving; unset otherwise | Enable packed 24-bit ROCmFP3 expert dispatch. Set `0` to restore generic expert dispatch. | +| `DFLASH_DS4_TP_BATCH_SPLIT_COPIES` | unset | OPT-IN: establish destination readiness once per DS4 scheduler split while retaining each backend copy's dependency publication. The qualified dual-ROCm launcher enables it. | +| `GGML_BATCH_PEER_COPIES` | unset | BURN-IN: additionally combine HIP peer-copy dependency publication. `GGML_CUDA_BATCH_PEER_COPIES` remains a compatibility alias. Keep these event-batching variables unset for the exact qualified profile. | | `GGML_SCHED_PROFILE` / `GGML_SCHED_PROFILE_MIN_SPLITS` | unset / 1 | DEBUG: report scheduler splits, copy volume, submission time, and source/destination synchronization time. | | `DFLASH_DS4_TP_FUSED_CACHE_SLOTS` | 2 | BURN-IN: number of heterogeneous verifier schedulers retained; higher values retain substantially more scratch on both GPUs. | | `DFLASH_DS4_VERIFY_FORCE_GRAPH_REPLAY` | unset | OPT-IN: bypass graph property scans only after warmup; scheduler-generation checks remain mandatory. | diff --git a/server/scripts/serve_ds4_dual_rocm_128k.sh b/server/scripts/serve_ds4_dual_rocm_128k.sh index cfd755591..e41aa5cc8 100755 --- a/server/scripts/serve_ds4_dual_rocm_128k.sh +++ b/server/scripts/serve_ds4_dual_rocm_128k.sh @@ -85,7 +85,7 @@ export DFLASH_DS4_TP_FUSED_HC_JOIN=${DFLASH_DS4_TP_FUSED_HC_JOIN:-1} export DFLASH_DS4_TP_MAIN_ROUTE_WEIGHTS=${DFLASH_DS4_TP_MAIN_ROUTE_WEIGHTS:-1} export DFLASH_DS4_TP_COARSE_OWNER=${DFLASH_DS4_TP_COARSE_OWNER:-1} export DFLASH_DS4_TP_COARSE_OWNER_SPLIT=${DFLASH_DS4_TP_COARSE_OWNER_SPLIT:-0} -export GGML_BATCH_PEER_COPIES=${GGML_BATCH_PEER_COPIES:-1} +export DFLASH_DS4_TP_BATCH_SPLIT_COPIES=${DFLASH_DS4_TP_BATCH_SPLIT_COPIES:-1} export DFLASH_CUDA_MMVQ_MOE_ROWS_PER_BLOCK=${DFLASH_CUDA_MMVQ_MOE_ROWS_PER_BLOCK:-2} export DFLASH_DS4_TP_CAPTURE_CACHE_SLOTS=${DFLASH_DS4_TP_CAPTURE_CACHE_SLOTS:-4} export DFLASH_DS4_TP_FUSED_CACHE_SLOTS=${DFLASH_DS4_TP_FUSED_CACHE_SLOTS:-9} diff --git a/server/src/common/backend_args.h b/server/src/common/backend_args.h index 611dd8d77..8bfa0ef57 100644 --- a/server/src/common/backend_args.h +++ b/server/src/common/backend_args.h @@ -68,12 +68,12 @@ struct BackendArgs { // Attention and speculative-decode options. Individual backends consume // only the fields they support. int fa_window = 0; // 0 = full attention. qwen3.6 full-attn layers must see the whole context; a finite window drops the system prompt/tools -> breaks tool calls. - bool paged_attention = false; // 16-token paged K/V blocks for AR decode + bool paged_attention = false; // model-specific paged K/V blocks for AR decode // Concurrent decode slots (--max-concurrency). > 1 requires paged_attention; // the backend serves that many sequences through the seq_* slot API. int max_concurrency = 1; // Total paged K/V pool in tokens shared by all slots (--kv-pool-tokens; - // block-rounded). 0 = derive capacity from available device memory. + // model-page-rounded). 0 = derive capacity from available device memory. long long kv_pool_tokens = 0; int kq_stride_pad = 32; int draft_block_size = 0; // 0 = drafter metadata diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index d462be899..205236f92 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -105,6 +105,9 @@ DFLASH_CHECK_ARCH("deepseek4", DeepSeek4BackendConfig, DeepSeek4LayerSplitAdapte // that shared struct would fail a check that is really about dispatch. DFLASH_CHECK_ARCH_OPTION("qwen35", Qwen35Config, Qwen35LayerSplitAdapterConfig, has_paged_attention, paged_attn); +DFLASH_CHECK_ARCH_OPTION("deepseek4", DeepSeek4BackendConfig, + DeepSeek4LayerSplitAdapterConfig, + has_paged_attention, paged_attn); DFLASH_CHECK_ARCH_OPTION("qwen35", Qwen35Config, Qwen35LayerSplitAdapterConfig, has_draft_block_size, draft_block_size); @@ -445,6 +448,9 @@ std::unique_ptr create_backend( cfg.fused_decode = args.ds4_fused_decode; cfg.fused_verify_f16_kv = args.ds4_fused_verify_f16_kv; cfg.prefill_mode = args.ds4_prefill_mode; + cfg.paged_attention = args.paged_attention; + cfg.max_concurrency = args.max_concurrency; + cfg.kv_pool_tokens = args.kv_pool_tokens; auto backend = std::make_unique(cfg); if (!backend->init()) { diff --git a/server/src/common/concurrency/seq_slot_manager.cpp b/server/src/common/concurrency/seq_slot_manager.cpp new file mode 100644 index 000000000..e43061b5f --- /dev/null +++ b/server/src/common/concurrency/seq_slot_manager.cpp @@ -0,0 +1,339 @@ +#include "common/concurrency/seq_slot_manager.h" + +#include +#include + +namespace dflash::common { + +SeqSlotManager::SeqSlotManager(PagedKvPool & pool, int max_ctx) + : pool_(pool), max_ctx_(max_ctx) { + slots_.assign(pool.max_sequences(), SeqSlot{}); +} + +SeqSlotManager::~SeqSlotManager() { + for (int slot = 0; slot < static_cast(slots_.size()); ++slot) { + retire(slot); + } +} + +int SeqSlotManager::decoding_count() const { + int n = 0; + for (const SeqSlot & s : slots_) { + n += s.decoding() ? 1 : 0; + } + return n; +} + +uint32_t SeqSlotManager::decode_headroom_capacity(int logical_tokens) const { + const uint64_t extended = + static_cast(std::max(0, logical_tokens)) + + pool_.block_size(); + return static_cast(std::min( + static_cast(max_ctx_), extended)); +} + +bool SeqSlotManager::capacity_fits_pool(uint32_t token_capacity) const { + const uint64_t blocks = token_capacity == 0 ? 0 : + 1 + (static_cast(token_capacity) - 1) / + pool_.block_size(); + return blocks <= pool_.physical_block_count(); +} + +PagedKvStatus SeqSlotManager::protect_decode_headroom() { + struct TopUp { + PagedKvSequenceHandle handle; + uint32_t token_capacity = 0; + }; + + std::vector topups; + topups.reserve(slots_.size()); + uint64_t total_additional = 0; + const uint64_t block_size = pool_.block_size(); + for (const SeqSlot & slot : slots_) { + if (!slot.decoding()) continue; + const uint32_t capacity = decode_headroom_capacity(slot.cur_pos); + if (!capacity_fits_pool(capacity)) continue; + + uint32_t owned_blocks = 0; + const PagedKvStatus status = + pool_.owned_block_count(slot.handle, owned_blocks); + if (status != PagedKvStatus::Ok) return status; + const uint64_t target_blocks = capacity == 0 ? 0 : + 1 + (static_cast(capacity) - 1) / block_size; + if (target_blocks <= owned_blocks) continue; + const uint32_t additional = + static_cast(target_blocks - owned_blocks); + total_additional += additional; + topups.push_back({slot.handle, capacity}); + } + + // Preflight the whole cohort before moving a block, so a failed admission + // attempt cannot protect only whichever decoder happened to be visited + // first. + if (total_additional > pool_.free_block_count()) { + return PagedKvStatus::BlocksExhausted; + } + for (const TopUp & topup : topups) { + const PagedKvStatus status = + pool_.reserve_capacity(topup.handle, topup.token_capacity); + if (status != PagedKvStatus::Ok) return status; + } + return PagedKvStatus::Ok; +} + +bool SeqSlotManager::is_active(int slot) const { + return slot >= 0 && slot < (int)slots_.size() && + slots_[(size_t)slot].active(); +} + +bool SeqSlotManager::is_prefilling(int slot) const { + return is_active(slot) && slots_[(size_t)slot].prefilling(); +} + +bool SeqSlotManager::has_prefill_prompt_at_least(int tokens) const { + if (tokens <= 0) return true; + return std::any_of(slots_.begin(), slots_.end(), + [tokens](const SeqSlot & slot) { + return slot.prefilling() && slot.prompt_len >= tokens; + }); +} + +SeqEngine::AdmitResult SeqSlotManager::admit( + uint64_t request_id, const std::vector & prompt, + const SamplerCfg & sampler) { + using AdmitStatus = SeqEngine::AdmitResult::Status; + SeqEngine::AdmitResult r; + if (prompt.empty()) { + r.error = "empty prompt"; + return r; + } + if (prompt.size() > static_cast(max_ctx_)) { + r.error = "prompt exceeds max_ctx"; + return r; + } + const int prompt_len = static_cast(prompt.size()); + + // A prompt larger than the whole pool can NEVER be admitted; waiting + // for other sequences to drain would stall the queue forever and then + // fail anyway. Hard-fail it up front instead of reporting busy. + const uint64_t pool_capacity = + (uint64_t)pool_.physical_block_count() * pool_.block_size(); + if ((uint64_t)prompt_len > pool_capacity) { + r.error = "prompt needs " + std::to_string(prompt_len) + + " KV tokens but the pool holds " + + std::to_string(pool_capacity) + + "; raise --kv-pool-tokens or shorten the prompt"; + return r; + } + + int slot = -1; + for (int i = 0; i < (int)slots_.size(); i++) { + if (!slots_[(size_t)i].active()) { slot = i; break; } + } + if (slot < 0) { + r.status = AdmitStatus::busy; + r.error = "all decode slots are busy"; + return r; + } + + // A newly freed block belongs to any older decoder missing its rolling + // next-page reserve before it can belong to this admission. + const PagedKvStatus headroom_status = protect_decode_headroom(); + if (headroom_status != PagedKvStatus::Ok) { + r.status = headroom_status == PagedKvStatus::BlocksExhausted + ? AdmitStatus::busy : AdmitStatus::failed; + r.error = r.status == AdmitStatus::busy + ? "existing decoders need the available KV headroom" + : paged_kv_status_string(headroom_status); + return r; + } + + PagedKvSequenceHandle handle; + uint32_t reservation_capacity = + decode_headroom_capacity(prompt_len); + if (!capacity_fits_pool(reservation_capacity)) { + // The prompt itself fits, but this physical pool can never hold its + // following page. Preserve useful prompt-only behavior and report + // decode exhaustion later if the sequence reaches that boundary. + reservation_capacity = static_cast(prompt_len); + } + const PagedKvStatus status = pool_.acquire_reserved( + request_id, reservation_capacity, handle); + if (status != PagedKvStatus::Ok) { + r.status = status == PagedKvStatus::SequenceSlotsExhausted || + status == PagedKvStatus::BlocksExhausted + ? AdmitStatus::busy : AdmitStatus::failed; + r.error = status == PagedKvStatus::BlocksExhausted + ? "not enough unreserved KV blocks for the prompt and decode headroom" + : paged_kv_status_string(status); + return r; + } + + SeqSlot & s = slots_[(size_t)slot]; + s.phase = SeqSlotPhase::prefill; + s.handle = handle; + s.cur_pos = 0; + s.prompt_len = prompt_len; + s.sampler = sampler; + s.sample_history = prompt; + // Same predicate the engine uses to pick CPU sampling over GPU argmax: + // a seed only means anything when the sampler actually draws. + if (sampler.needs_logit_processing()) { + s.rng.seed(sampler.seed); + } else { + s.rng.seed(std::random_device{}()); + } + + r.status = AdmitStatus::admitted; + r.slot = slot; + return r; +} + +SeqSlotManager::PrefillChunk SeqSlotManager::append_prefill( + int slot, int n_tokens) { + PrefillChunk out; + if (!is_prefilling(slot) || n_tokens < 1) return out; + + SeqSlot & s = slots_[(size_t)slot]; + if (s.cur_pos > s.prompt_len || + n_tokens > s.prompt_len - s.cur_pos) { + return out; + } + + PagedKvAppendResult app = pool_.append(s.handle, (uint32_t)n_tokens); + if (!app) { + // Admission reserved the whole prompt. Treat exhaustion here as a + // broken invariant, not a retryable condition: retrying a batch of + // all-prefill slots without any decoder able to retire would livelock. + if (app.status == PagedKvStatus::BlocksExhausted) { + std::fprintf(stderr, + "[parallel] reserved prefill capacity missing for slot %d\n", + slot); + } + return out; + } + + out.rows.reserve(app.write_slots.size()); + for (const PagedKvWriteSlot & write : app.write_slots) { + out.rows.push_back((int64_t)write.physical_token_index); + if (write.block_offset == 0) { + if (out.first_new_block < 0) { + out.first_new_block = + (int)(write.logical_position / pool_.block_size()); + } + out.new_blocks.push_back((int32_t)write.physical_block); + } + } + s.cur_pos += n_tokens; + out.ok = true; + return out; +} + +void SeqSlotManager::commit_prefill(int slot) { + if (!is_prefilling(slot)) return; + SeqSlot & s = slots_[(size_t)slot]; + if (s.cur_pos != s.prompt_len) return; + s.phase = SeqSlotPhase::decode; +} + +SeqSlotManager::StepAppend SeqSlotManager::append_tokens( + int slot, const int32_t * fed_tokens, int n_tokens) { + StepAppend out; + if (!is_active(slot) || !slots_[(size_t)slot].decoding() || + !fed_tokens || n_tokens < 1) return out; + SeqSlot & s = slots_[(size_t)slot]; + if (!s.staged_tokens.empty() || s.cur_pos > max_ctx_ || + n_tokens > max_ctx_ - s.cur_pos) { + return out; + } + PagedKvAppendResult app = pool_.append( + s.handle, static_cast(n_tokens)); + if (!app || app.token_count != static_cast(n_tokens)) { + out.busy = app.status == PagedKvStatus::BlocksExhausted; + return out; + } + if (app.write_slots.size() != static_cast(n_tokens) || + app.write_slots.front().logical_position != + static_cast(s.cur_pos) || + app.write_slots.back().logical_position != + static_cast(s.cur_pos + n_tokens - 1)) { + const PagedKvStatus rollback = + pool_.rollback_append(s.handle, static_cast(n_tokens)); + if (rollback != PagedKvStatus::Ok) { + std::fprintf(stderr, + "[parallel] slot %d failed to roll back invalid append: %s\n", + slot, paged_kv_status_string(rollback)); + } + return out; + } + + out.physical_rows.reserve(app.write_slots.size()); + for (const PagedKvWriteSlot & write : app.write_slots) { + out.physical_rows.push_back( + static_cast(write.physical_token_index)); + if (write.block_offset == 0) { + if (out.first_new_block < 0) { + out.first_new_block = static_cast( + write.logical_position / pool_.block_size()); + } + out.new_blocks.push_back( + static_cast(write.physical_block)); + } + } + s.staged_tokens.assign(fed_tokens, fed_tokens + n_tokens); + out.ok = true; + out.count = n_tokens; + out.physical_row = out.physical_rows.front(); + out.position = s.cur_pos; + if (!out.new_blocks.empty()) { + out.new_block = out.new_blocks.front(); + out.new_block_index = out.first_new_block; + } + return out; +} + +SeqSlotManager::StepAppend SeqSlotManager::append_token( + int slot, int32_t fed_token) { + return append_tokens(slot, &fed_token, 1); +} + +void SeqSlotManager::commit_step(int slot) { + if (!is_active(slot)) return; + SeqSlot & s = slots_[(size_t)slot]; + if (s.staged_tokens.empty()) return; + s.sample_history.insert( + s.sample_history.end(), s.staged_tokens.begin(), + s.staged_tokens.end()); + s.cur_pos += static_cast(s.staged_tokens.size()); + s.staged_tokens.clear(); +} + +bool SeqSlotManager::rollback_step(int slot) { + if (!is_active(slot)) return false; + SeqSlot & s = slots_[(size_t)slot]; + if (s.staged_tokens.empty()) return false; + const PagedKvStatus status = pool_.rollback_append( + s.handle, static_cast(s.staged_tokens.size())); + if (status != PagedKvStatus::Ok) { + std::fprintf(stderr, + "[parallel] slot %d staged append rollback failed: %s\n", + slot, paged_kv_status_string(status)); + return false; + } + s.staged_tokens.clear(); + return true; +} + +void SeqSlotManager::retire(int slot) { + if (slot < 0 || slot >= (int)slots_.size()) return; + SeqSlot & s = slots_[(size_t)slot]; + if (!s.active()) return; + const PagedKvStatus status = pool_.release(s.handle); + if (status != PagedKvStatus::Ok && status != PagedKvStatus::StaleHandle) { + std::fprintf(stderr, "[parallel] slot %d release failed: %s\n", + slot, paged_kv_status_string(status)); + } + s = SeqSlot{}; +} + +} // namespace dflash::common diff --git a/server/src/common/concurrency/seq_slot_manager.h b/server/src/common/concurrency/seq_slot_manager.h new file mode 100644 index 000000000..abdf1bef5 --- /dev/null +++ b/server/src/common/concurrency/seq_slot_manager.h @@ -0,0 +1,153 @@ +// SeqSlotManager — complete host-side state for each concurrent serving slot. +// +// Companion of PagedKvPool: the pool hands out sequence handles and physical +// blocks; this class owns everything else a slot needs between admission and +// retirement — the pool-handle lifecycle (including every error path), the +// admission arithmetic (context clamp, prompt reservation, and rolling decode +// headroom), on-demand block allocation, per-slot sampler/RNG/penalty-history +// state, and the position counters. +// +// It deliberately owns NO device state. Prefill/decode allocation returns +// physical rows and block-table deltas as plain vectors. Prompt, KV ownership, +// sampler, and progress live together here; the scheduler keeps +// only its coarse request phase. +// +// Not thread-safe; the single scheduler thread is the only caller. + +#pragma once + +#include "common/concurrency/paged_kv_pool.h" +#include "common/sampler.h" +#include "common/concurrency/seq_engine.h" + +#include +#include +#include +#include + +namespace dflash::common { + +enum class SeqSlotPhase { + free, + prefill, + decode, +}; + +struct SeqSlot { + SeqSlotPhase phase = SeqSlotPhase::free; + PagedKvSequenceHandle handle; + // Prompt tokens are the immutable prefix of sample_history. Decode tokens + // append to the same allocation, avoiding a second full prompt copy. + int prompt_len = 0; + int cur_pos = 0; + SamplerCfg sampler; + std::mt19937_64 rng{0x9E3779B97F4A7C15ull}; + // Penalty history is recorded as fed rather than sampled: the scheduler + // may override a sample before the model consumes it. + std::vector sample_history; + // Decode rows reserved by the current target graph. They become durable + // only after the graph and any speculative promotion succeed. + std::vector staged_tokens; + + int generated_tokens() const { + return sample_history.size() > (size_t)prompt_len + ? (int)(sample_history.size() - (size_t)prompt_len) + : 0; + } + + bool active() const { return phase != SeqSlotPhase::free; } + bool prefilling() const { return phase == SeqSlotPhase::prefill; } + bool decoding() const { return phase == SeqSlotPhase::decode; } +}; + +class SeqSlotManager { +public: + // `max_ctx` is the per-sequence logical bound; slot count comes from the + // pool's max_sequences. The pool must outlive the manager. + SeqSlotManager(PagedKvPool & pool, int max_ctx); + ~SeqSlotManager(); + + SeqSlotManager(const SeqSlotManager &) = delete; + SeqSlotManager & operator=(const SeqSlotManager &) = delete; + + // Claim a free slot and atomically reserve all K/V blocks needed by the + // known prompt plus its next logical decode page when that page can exist + // in both max_ctx and the physical pool. Existing decoders are topped up + // first, so a younger admission cannot steal their next-page headroom. + // Prompts larger than the whole pool hard-fail; temporary capacity pressure + // reports busy. Seeds sampling from sampler.seed, including zero. Greedy + // slots never draw, so their RNG state is intentionally unspecified. + SeqEngine::AdmitResult admit(uint64_t request_id, + const std::vector & prompt, + const SamplerCfg & sampler); + + struct PrefillChunk { + bool ok = false; + std::vector rows; + // Delta to patch into the slot's device block-table column. + std::vector new_blocks; + int first_new_block = -1; + }; + + // Append `n_tokens` more prompt rows for a prefilling slot. Physical block + // ids come from the slot's admission reservation, so any append within the + // admitted prompt is guaranteed not to wait on another sequence. + PrefillChunk append_prefill(int slot, int n_tokens); + + // Record a finished prefill and expose the slot to decode. + void commit_prefill(int slot); + + struct StepAppend { + bool ok = false; + bool busy = false; // no physical block available right now + int64_t physical_row = -1; + std::vector physical_rows; + int count = 0; + int position = -1; // logical position the fed token is written at + int32_t new_block = -1; + int new_block_index = -1; + std::vector new_blocks; + int first_new_block = -1; + }; + + StepAppend append_tokens(int slot, const int32_t * fed_tokens, + int n_tokens); + + // Allocate decode cache rows and stage the fed tokens. Both history and + // cur_pos wait for commit_step(). + StepAppend append_token(int slot, int32_t fed_token); + + // The batched step's compute succeeded: cur_pos++. + void commit_step(int slot); + + bool rollback_step(int slot); + + // Release the slot's blocks and clear its state. Safe on inactive slots + // and after a failed admission/prefill. + void retire(int slot); + + int slot_count() const { return (int)slots_.size(); } + int max_context() const { return max_ctx_; } + int decoding_count() const; + bool is_active(int slot) const; + bool is_prefilling(int slot) const; + bool has_prefill_prompt_at_least(int tokens) const; + SeqSlot & slot(int i) { return slots_[(size_t)i]; } + const SeqSlot & slot(int i) const { return slots_[(size_t)i]; } + +private: + // Logical extent whose block count includes the sequence's current pages + // plus one future page, capped at max_ctx. + uint32_t decode_headroom_capacity(int logical_tokens) const; + bool capacity_fits_pool(uint32_t token_capacity) const; + + // Atomically preflight and top up every decoding slot as one cohort before + // a younger sequence may reserve capacity. + PagedKvStatus protect_decode_headroom(); + + PagedKvPool & pool_; + int max_ctx_ = 0; + std::vector slots_; +}; + +} // namespace dflash::common diff --git a/server/src/common/feature_gate.cpp b/server/src/common/feature_gate.cpp index cae310910..f8cf94318 100644 --- a/server/src/common/feature_gate.cpp +++ b/server/src/common/feature_gate.cpp @@ -193,13 +193,14 @@ std::string check_feature_compatibility( // ── --paged-attention × architecture, placement, and decode features // Paged decode swaps the contiguous K/V cache for a block table owned by - // the monolithic qwen35 backend, so every rule below is about reaching - // that one code path. All are errors rather than warnings: running dense + // a monolithic Qwen or DeepSeek backend. All are errors rather than + // warnings: running dense // instead would hide the memory behavior the flag was chosen for. if (args.paged_attention) { if (!arch_supports_paged_attention(arch, /*is_layer_split=*/false)) { - return "--paged-attention requires a Qwen3.5/Qwen3.6 dense target " - "(architecture '" + arch + "' has no paged decode path)"; + return "--paged-attention requires a dense Qwen3.5/Qwen3.6 or " + "DeepSeek4 target (architecture '" + arch + + "' has no paged decode path)"; } // No rule for "requires a CUDA or HIP build": those are the only two // backends this binary can be configured with, and GGML_OP_PAGED_ATTN @@ -226,6 +227,18 @@ std::string check_feature_compatibility( if (features.kvflash_enabled) { return "--paged-attention cannot be combined with KVFlash"; } + if (arch == "deepseek4") { + if (target_backend != PlacementBackend::Hip) { + return "DeepSeek4 paged attention requires a local HIP target"; + } + if (args.ds4_prefill_mode != PrefillAttentionMode::Exact) { + return "DeepSeek4 paged attention requires --ds4-prefill exact"; + } + if (args.ds4_fused_decode || args.ds4_fused_verify_f16_kv) { + return "DeepSeek4 paged attention requires non-fused " + "autoregressive decode"; + } + } // The pool rounds max_ctx up to a whole number of blocks, so the top // of the range is what can be rounded without overflowing int. if (args.device.max_ctx <= 0 || @@ -236,8 +249,8 @@ std::string check_feature_compatibility( } // ── --max-concurrency × paged attention - // Concurrent decode slots are currently implemented only by the paged - // qwen35 backend. The common scheduler does not require a particular + // Concurrent decode slots are implemented by model-specific paged + // backends. The common scheduler does not require a particular // model-state representation; each backend owns whatever per-slot state // its graph needs alongside one block-table column per sequence. // Everything the paged cluster above rejects is transitively rejected, @@ -249,11 +262,13 @@ std::string check_feature_compatibility( if (!args.paged_attention) { return "--max-concurrency requires --paged-attention"; } - // The paged pool addresses tokens with uint32; 64 slots is far above - // any batch the decode kernel has been sized for and keeps the - // fixed-width decode batch bounded. - if (args.max_concurrency > 64) { - return "--max-concurrency must be at most 64"; + // Qwen's graph is qualified through 64 lanes. DeepSeek's gathered + // whole-model graph has a smaller, separately qualified ceiling. + const int max_slots = arch == "deepseek4" + ? DEEPSEEK4_MAX_PAGED_SEQUENCES : 64; + if (args.max_concurrency > max_slots) { + return "--max-concurrency must be at most " + + std::to_string(max_slots) + " for " + arch; } // Physical capacity is memory-derived and capped independently of the // logical slot count, so max-concurrency no longer multiplies max_ctx diff --git a/server/src/common/model_capabilities.h b/server/src/common/model_capabilities.h index ebc7524c0..d1d541312 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, kNever}, + {"deepseek4", true, false, false, false, kNever, kNever, kNever, kNever, kNever,kNever, kMono}, }; inline constexpr std::size_t kArchCount = diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index d5e80cec0..7ae60e489 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -1306,7 +1306,8 @@ bool build_moe_hybrid_ffn_graph( primary_owner.output, secondary_owner.output, out); if (!combined) return false; - out.output = ggml_cont(ctx, combined); + out.output = ggml_is_contiguous(combined) ? combined + : ggml_cont(ctx, combined); return true; } diff --git a/server/src/common/paged_attention_config.h b/server/src/common/paged_attention_config.h index c7768a5c2..e93cb0874 100644 --- a/server/src/common/paged_attention_config.h +++ b/server/src/common/paged_attention_config.h @@ -14,6 +14,12 @@ namespace dflash::common { constexpr int PAGED_BLOCK_SIZE = 16; +inline constexpr int DEEPSEEK4_MAX_PAGED_SEQUENCES = 6; +// Monolithic exact serving packs chronological prompt rows into the gathered +// graph, sharing this row budget between decode and prompt sequences. +inline constexpr int DEEPSEEK4_MAX_GATHERED_ROWS = 16; +// Routed-expert part width for packed steps: the RDNA3 MUL_MAT_ID MMVQ ceiling. +inline constexpr int DEEPSEEK4_MAX_MOE_ROWS_PER_PART = 8; constexpr int paged_block_count(int max_ctx) { return (max_ctx + PAGED_BLOCK_SIZE - 1) / PAGED_BLOCK_SIZE; diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 7c6f6548c..e31da5950 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -4,6 +4,7 @@ #include "deepseek4_backend.h" #include "deepseek4_budget_hook.h" #include "deepseek4_internal.h" +#include "deepseek4_page_layout.h" #include "common/dynamic_backend.h" #include "common/peer_access.h" #include "common/platform_env.h" @@ -149,6 +150,18 @@ static bool env_int_in_range(const char * name, int fallback, return true; } +static bool is_gfx_device(int gpu, const char * arch) { +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + cudaDeviceProp prop{}; + return cudaGetDeviceProperties(&prop, gpu) == cudaSuccess && + std::strncmp(prop.gcnArchName, arch, std::strlen(arch)) == 0; +#else + (void) gpu; + (void) arch; + return false; +#endif +} + static bool configure_dspark_mmvq_defaults(int gpu) { if (env_flag_enabled("DFLASH_DS4_Q6_VERIFY")) { std::fprintf(stderr, @@ -227,6 +240,35 @@ static bool configure_dspark_mmvq_defaults(int gpu) { return true; } +// Monolithic paged serving packs up to sixteen chronological prompt rows. +// The gfx1151 dense ROCmFP4 weight-reuse kernel preserves single-column +// arithmetic through that width. Keep those projections on MMVQ while +// honoring an explicit LUCE_MMVQ_MAX_NCOLS setting. +static void configure_gfx1151_paged_mmvq_default(int gpu, bool paged_attention) { +#if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + if (!paged_attention || env_flag_enabled("DFLASH_DS4_SPEC") || + !is_gfx_device(gpu, "gfx1151")) { + return; + } + if (std::getenv("LUCE_MMVQ_MAX_NCOLS") == nullptr && + set_environment_variable("LUCE_MMVQ_MAX_NCOLS", "16", false) == 0) { + std::fprintf(stderr, + "[deepseek4] gfx1151 paged serving: defaulting " + "LUCE_MMVQ_MAX_NCOLS=16 (ROCmFP4 weight-reuse MMVQ)\n"); + } + for (const char * name : {"DFLASH_CUDA_MMVQ_FP4_X4", + "DFLASH_CUDA_MMVQ_FP4_Q5_X4_PLUS1", + "DFLASH_CUDA_MMVQ_MOE_FP3_PACKED24"}) { + if (set_environment_variable(name, "1", false) != 0) { + std::fprintf(stderr, "[deepseek4] failed to default %s=1\n", name); + } + } +#else + (void) gpu; + (void) paged_attention; +#endif +} + static void configure_gfx1201_hybrid_sub_batch_default(int gpu) { #if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) if (std::getenv("DFLASH_MMQ_SUB_BATCH") != nullptr) { @@ -368,42 +410,6 @@ static double ms(uint64_t us) { return (double)us / 1000.0; } -static void log_step_tel(const char * phase, - int tokens, - int steps, - double wall_s, - const DeepSeek4StepTelemetry & t) { - const double tok_s = wall_s > 0.0 ? (double)tokens / wall_s : 0.0; - std::fprintf(stderr, - "[deepseek4-timing] %s tokens=%d steps=%d wall=%.3fs %.2f tok/s " - "step=%.1fms embed=%.1fms attn_build=%.1fms attn_compute=%.1fms attn_read=%.1fms " - "full_build=%.1fms full_set=%.1fms full_compute=%.1fms full_read=%.1fms " - "ffn_build=%.1fms ffn_compute=%.1fms ffn_read=%.1fms " - "route_build=%.1fms route_compute=%.1fms route_read=%.1fms route_select=%.1fms " - "ffn=%.1fms hot=%.1fms cold=%.1fms combine=%.1fms partition=%.1fms " - "ffn_hot_graph_build=%llu ffn_hot_graph_hit=%llu ffn_cold_graph_build=%llu ffn_cold_graph_hit=%llu " - "hc_pre=%.1fms hc_pre_build=%.1fms hc_pre_input=%.1fms hc_pre_compute=%.1fms " - "hc_post=%.1fms output=%.1fms sample=%.1fms emit=%.1fms " - "hot_sel=%d cold_sel=%d\n", - phase, tokens, steps, wall_s, tok_s, - ms(t.total_us), ms(t.embed_us), ms(t.attn_build_us), ms(t.attn_compute_us), ms(t.attn_read_us), - ms(t.full_graph_build_us), ms(t.full_graph_set_us), - ms(t.full_graph_compute_us), ms(t.full_graph_read_us), - ms(t.ffn_build_us), ms(t.ffn_compute_us), ms(t.ffn_read_us), - ms(t.route_build_us), ms(t.route_compute_us), ms(t.route_read_us), ms(t.route_select_us), - ms(t.ffn_eval_us), ms(t.ffn_hot_us), ms(t.ffn_cold_us), ms(t.ffn_combine_us), - ms(t.ffn_partition_us), - (unsigned long long)t.ffn_hot_graph_builds, (unsigned long long)t.ffn_hot_graph_hits, - (unsigned long long)t.ffn_cold_graph_builds, (unsigned long long)t.ffn_cold_graph_hits, - ms(t.hc_pre_attn_us + t.hc_pre_ffn_us), - ms(t.hc_pre_build_us), - ms(t.hc_pre_input_us), - ms(t.hc_pre_compute_us), - ms(t.hc_post_attn_us + t.hc_post_ffn_us), - ms(t.output_us), ms(t.sample_us), ms(t.emit_us), - t.hot_selected, t.cold_selected); -} - static uint64_t layer_expert_bytes(const DeepSeek4Layer & layer, int n_expert) { if (n_expert <= 0) return 0; uint64_t bytes = 0; @@ -503,20 +509,16 @@ static uint64_t estimate_ds4_cache_bytes(const DeepSeek4Weights & w, int max_ctx const size_t comp_cap = (size_t) (max_ctx / (int) ratio) + 16; total_bytes += comp_cap * head_dim * sizeof(uint16_t); - const size_t window = (ratio == 4) ? 8 : ratio; - total_bytes += window * head_dim * sizeof(float) * 2; + const size_t state_rows = (ratio == 4) ? 8 : ratio; + const size_t comp_width = head_dim * (ratio == 4 ? 2 : 1); + total_bytes += state_rows * comp_width * sizeof(float) * 2; if (ratio == 4) { - // index_comp_kv is per-head: ne0 = n_indexer_head_dim (see - // deepseek4_graph.cpp). The full 64-head width lives only in the - // fixed-size state_kv/state_score scratch (index_state_rows rows), - // which does not scale with context. The old estimate multiplied - // by n_indexer_head here, overcounting 256K context by ~13x and - // falsely rejecting large contexts in hybrid placement. - total_bytes += comp_cap * (size_t) w.n_indexer_head_dim * sizeof(uint16_t); - total_bytes += window * (size_t) w.n_indexer_head_dim * sizeof(float) * 2; - total_bytes += (size_t) 2 * 2 * ratio * (size_t) w.n_indexer_head * - (size_t) w.n_indexer_head_dim * sizeof(float); + // index_comp_kv is per-head. The full multi-head width lives + // only in fixed-size state scratch and does not scale with context. + const size_t index_dim = (size_t) w.n_indexer_head_dim; + total_bytes += comp_cap * index_dim * sizeof(uint16_t); + total_bytes += state_rows * (2 * index_dim) * sizeof(float) * 2; } } @@ -653,8 +655,9 @@ static bool fill_profiled_hot_placement(const DeepSeek4Weights & w, // uses authoritative router statistics and evaluates every selected expert. static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, ggml_backend_t backend, - int max_ctx, + uint64_t kv_bytes, bool all_cold, + bool paged, Ds4HybridBudgetInfo & out, std::string * err) { out = {}; @@ -675,11 +678,13 @@ static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, out.core_bytes = moe_hybrid_core_bytes_from_memory( "deepseek4", out.gpu_free, out.gpu_total); - out.kv_bytes = estimate_ds4_cache_bytes(w, max_ctx); + out.kv_bytes = kv_bytes; // In all-cold mode the KV cache is owned by the secondary (Strix) - // backend, so it must not consume the primary GPU's expert budget. - const uint64_t main_charge = all_cold ? 0 : out.kv_bytes; + // backend in the legacy contiguous path, so it does not consume the + // primary GPU's expert budget there. Paged serving owns its persistent + // page tensors on the primary target. + const uint64_t main_charge = all_cold && !paged ? 0 : out.kv_bytes; if (out.gpu_total > out.core_bytes + main_charge + out.warm_bytes + out.safety_bytes) { out.expert_budget = out.gpu_total - out.core_bytes - main_charge - out.warm_bytes - out.safety_bytes; } @@ -740,6 +745,43 @@ static MoeLayerDesc make_ds4_expert_layer_desc(const DeepSeek4Layer & layer) { } // namespace +void log_deepseek4_step_telemetry(const char * phase, + int tokens, + int steps, + double wall_s, + const DeepSeek4StepTelemetry & t) { + const double tok_s = wall_s > 0.0 ? (double)tokens / wall_s : 0.0; + std::fprintf(stderr, + "[deepseek4-timing] %s tokens=%d steps=%d wall=%.3fs %.2f tok/s " + "step=%.1fms embed=%.1fms attn_build=%.1fms attn_compute=%.1fms attn_read=%.1fms " + "full_build=%.1fms full_set=%.1fms full_compute=%.1fms full_read=%.1fms " + "ffn_build=%.1fms ffn_compute=%.1fms ffn_read=%.1fms " + "route_build=%.1fms route_compute=%.1fms route_read=%.1fms route_select=%.1fms " + "ffn=%.1fms hot=%.1fms cold=%.1fms combine=%.1fms partition=%.1fms " + "ffn_hot_graph_build=%llu ffn_hot_graph_hit=%llu ffn_cold_graph_build=%llu ffn_cold_graph_hit=%llu " + "hc_pre=%.1fms hc_pre_build=%.1fms hc_pre_input=%.1fms hc_pre_compute=%.1fms " + "hc_post=%.1fms output=%.1fms sample=%.1fms emit=%.1fms " + "hot_sel=%d cold_sel=%d\n", + phase, tokens, steps, wall_s, tok_s, + ms(t.total_us), ms(t.embed_us), ms(t.attn_build_us), ms(t.attn_compute_us), ms(t.attn_read_us), + ms(t.full_graph_build_us), ms(t.full_graph_set_us), + ms(t.full_graph_compute_us), ms(t.full_graph_read_us), + ms(t.ffn_build_us), ms(t.ffn_compute_us), ms(t.ffn_read_us), + ms(t.route_build_us), ms(t.route_compute_us), ms(t.route_read_us), ms(t.route_select_us), + ms(t.ffn_eval_us), ms(t.ffn_hot_us), ms(t.ffn_cold_us), ms(t.ffn_combine_us), + ms(t.ffn_partition_us), + (unsigned long long)t.ffn_hot_graph_builds, (unsigned long long)t.ffn_hot_graph_hits, + (unsigned long long)t.ffn_cold_graph_builds, (unsigned long long)t.ffn_cold_graph_hits, + ms(t.hc_pre_attn_us + t.hc_pre_ffn_us), + ms(t.hc_pre_build_us), + ms(t.hc_pre_input_us), + ms(t.hc_pre_compute_us), + ms(t.hc_post_attn_us + t.hc_post_ffn_us), + ms(t.output_us), ms(t.sample_us), ms(t.emit_us), + t.hot_selected, t.cold_selected); +} + + DeepSeek4Backend::DeepSeek4Backend(const DeepSeek4BackendConfig & cfg) : cfg_(cfg) {} @@ -748,7 +790,8 @@ DeepSeek4Backend::~DeepSeek4Backend() { } bool DeepSeek4Backend::requires_monolithic_model() const { - return cfg_.fused_decode || cfg_.fused_verify_f16_kv || + return cfg_.paged_attention || cfg_.fused_decode || + cfg_.fused_verify_f16_kv || prefill_attention_mode_is_approximate(cfg_.prefill_mode); } @@ -781,10 +824,12 @@ bool DeepSeek4Backend::load_model() { ? compiled_placement_backend() : cfg_.device.backend; - // Fused decode and layer-major prefill normally require monolithic expert - // residency. Heterogeneous TP is the exception: its fused graph owns the - // routed experts across two local GPU backends, so forcing a full load would - // disable the requested split before the TP runtime can initialize. + // Paged concurrency, fused decode, and layer-major prefill normally require + // monolithic expert residency. In-process heterogeneous TP is the explicit + // exception: its fused graph owns a fixed expert split across two local GPU + // backends, so forcing a full load would disable the requested placement + // before the TP runtime can initialize. init() has already rejected paged + // deployments outside the qualified R9700 + Strix Halo topology. const bool force_full = env_flag_enabled("DFLASH_DS4_FORCE_FULL_LOAD"); const bool heterogeneous_tp = env_flag_enabled("DFLASH_DS4_MOE_TP"); const bool need_monolithic = @@ -793,9 +838,10 @@ bool DeepSeek4Backend::load_model() { (force_full || need_monolithic)) { std::fprintf(stderr, "[deepseek4] monolithic execution requested " - "(forced=%s, fused_decode=%s, " + "(forced=%s, paged=%s, fused_decode=%s, " "fused_verify_f16_kv=%s, prefill=%s)\n", force_full ? "yes" : "no", + cfg_.paged_attention ? "on" : "off", cfg_.fused_decode ? "on" : "off", cfg_.fused_verify_f16_kv ? "on" : "off", prefill_attention_mode_name(cfg_.prefill_mode)); @@ -1028,14 +1074,63 @@ bool DeepSeek4Backend::supports_batched_spec_feature_capture( } bool DeepSeek4Backend::init() { + if (cfg_.paged_attention) { + const PlacementBackend target_backend = + cfg_.device.backend == PlacementBackend::Auto + ? compiled_placement_backend() : cfg_.device.backend; + const Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); + if (target_backend != PlacementBackend::Hip) { + std::fprintf(stderr, + "[deepseek4] paged concurrency requires a local HIP target\n"); + return false; + } + if (!tp.requested) { + if (!is_gfx_device(cfg_.device.gpu, "gfx1151")) { + std::fprintf(stderr, + "[deepseek4] monolithic paged concurrency requires one " + "local Strix Halo (gfx1151) target\n"); + return false; + } + } else if (!tp.in_process || !tp.backend_valid || + tp.secondary_backend != PlacementBackend::Hip) { + std::fprintf(stderr, + "[deepseek4] heterogeneous paged concurrency requires " + "in-process HIP expert parallelism\n"); + return false; + } else if (tp.secondary_gpu == cfg_.device.gpu || + !is_gfx_device(cfg_.device.gpu, "gfx1201") || + !is_gfx_device(tp.secondary_gpu, "gfx1151")) { + std::fprintf(stderr, + "[deepseek4] heterogeneous paged concurrency requires an " + "R9700 (gfx1201) target and Strix Halo (gfx1151) secondary\n"); + return false; + } + } + // 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. if (!configure_dspark_mmvq_defaults(cfg_.device.gpu)) { return false; } + configure_gfx1151_paged_mmvq_default(cfg_.device.gpu, cfg_.paged_attention); configure_gfx1201_hybrid_sub_batch_default(cfg_.device.gpu); + if (cfg_.paged_attention && + (cfg_.max_concurrency < 1 || + cfg_.max_concurrency > DEEPSEEK4_MAX_PAGED_SEQUENCES || + cfg_.device.is_layer_split() || + cfg_.prefill_mode != PrefillAttentionMode::Exact || + cfg_.fused_decode || cfg_.fused_verify_f16_kv || + env_flag_enabled("DFLASH_DS4_FUSED_DECODE") || + env_flag_enabled("DFLASH_DS4_SPEC"))) { + std::fprintf(stderr, + "[deepseek4] paged serving requires 1..%d local slots, exact " + "prefill, and autoregressive non-fused decode\n", + DEEPSEEK4_MAX_PAGED_SEQUENCES); + return false; + } + backend_ = ggml_backend_cuda_init(cfg_.device.gpu); if (!backend_) { std::fprintf(stderr, "[deepseek4] failed to create CUDA backend (gpu=%d)\n", @@ -1059,15 +1154,49 @@ bool DeepSeek4Backend::init() { } const int max_ctx = cfg_.max_ctx > 0 ? cfg_.max_ctx : 8192; - if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { - std::fprintf(stderr, "[deepseek4] failed to allocate KV cache (ctx=%d)\n", max_ctx); - return false; + if (cfg_.paged_attention) { + const uint64_t requested = cfg_.kv_pool_tokens > 0 + ? (uint64_t)cfg_.kv_pool_tokens + : (uint64_t)max_ctx * (uint64_t)cfg_.max_concurrency; + uint32_t physical_blocks = 0; + if (!plan_deepseek4_paged_pool_blocks( + (uint32_t)max_ctx, (uint32_t)cfg_.max_concurrency, + cfg_.kv_pool_tokens > 0 ? (uint64_t)cfg_.kv_pool_tokens : 0, + physical_blocks) || + !create_deepseek4_paged_cache( + backend_, w_, (uint32_t)cfg_.max_concurrency, + (uint32_t)max_ctx, physical_blocks, paged_cache_)) { + std::fprintf(stderr, + "[deepseek4] paged cache allocation failed (ctx=%d slots=%d " + "requested_pool_tokens=%llu); reduce --max-ctx/--max-concurrency " + "or set --kv-pool-tokens\n", max_ctx, cfg_.max_concurrency, + (unsigned long long)requested); + return false; + } + } else { + if (!create_deepseek4_cache(backend_, w_, max_ctx, cache_)) { + std::fprintf(stderr, "[deepseek4] failed to allocate KV cache (ctx=%d)\n", max_ctx); + return false; + } + cache_.prefill_mode = cfg_.prefill_mode; } - cache_.prefill_mode = cfg_.prefill_mode; if (env_flag_enabled("DFLASH_DS4_MOE_TP") && !init_moe_tensor_parallel()) { return false; } + if (cfg_.paged_attention && expert_runtime_.compute) { + std::fprintf(stderr, + "[deepseek4] paged serving cannot use the out-of-process expert " + "compute callback; select in-process DFLASH_DS4_MOE_TP or disable paged attention\n"); + return false; + } + if (cfg_.paged_attention && moe_hybrid_ && + !moe_hybrid_->materialized_cold_experts) { + std::fprintf(stderr, + "[deepseek4] paged serving requires statically materialized " + "expert ownership; enable in-process DFLASH_DS4_MOE_TP\n"); + return false; + } if (const char * stats_path = std::getenv("DFLASH_DS4_ROUTING_STATS_OUT")) { if (*stats_path) { @@ -1092,6 +1221,21 @@ bool DeepSeek4Backend::init() { std::fprintf(stderr, "[deepseek4-moe-tp] in-memory routing stats enabled\n"); } + if (cfg_.paged_attention) { + seq_engine_ = std::make_unique( + *this, *paged_cache_.pool, max_ctx, + paged_cache_.plan.max_blocks_per_sequence); + std::fprintf(stderr, + "[deepseek4-parallel] enabled %d slots mode=%s, %u x %d-token physical " + "blocks; prefill is exact reference mode at %s\n", + cfg_.max_concurrency, moe_hybrid_ ? "r9700+strix" : "strix", + paged_cache_.plan.physical_blocks, + DS4_PAGE_TOKENS, + moe_hybrid_ + ? "one prompt token per slot per scheduler iteration" + : "up to sixteen chronological rows for a lone prompt, " + "four per concurrent prompt, and sixteen total per iteration"); + } const int active_experts = w_.routed_expert_top_k > 0 ? w_.routed_expert_top_k : w_.n_expert_used; std::fprintf(stderr, @@ -1102,7 +1246,7 @@ bool DeepSeek4Backend::init() { prefill_attention_mode_name(cfg_.prefill_mode), moe_hybrid_ ? " [hybrid]" : ""); - if (env_flag_enabled("DFLASH_DS4_SPEC")) { + if (!cfg_.paged_attention && env_flag_enabled("DFLASH_DS4_SPEC")) { const char * dp = std::getenv("DFLASH_DS4_DRAFT"); if (dp && *dp) { spec_draft_path_ = dp; @@ -1186,10 +1330,30 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & MoeHybridPlacement * decode_out, std::string * err) const { if (decode_out) *decode_out = {}; + uint64_t kv_bytes = cfg_.paged_attention + ? 0 : estimate_ds4_cache_bytes(w, max_ctx); + if (cfg_.paged_attention) { + uint32_t physical_blocks = 0; + DeepSeek4PagedCachePlan paged_plan; + if (!plan_deepseek4_paged_pool_blocks( + (uint32_t)max_ctx, (uint32_t)cfg_.max_concurrency, + cfg_.kv_pool_tokens > 0 ? (uint64_t)cfg_.kv_pool_tokens : 0, + physical_blocks) || + !plan_deepseek4_paged_cache( + (uint32_t)w.head_dim, (uint32_t)w.n_indexer_head_dim, + (uint32_t)cfg_.max_concurrency, (uint32_t)max_ctx, + physical_blocks, w.compress_ratios, paged_plan) || + kv_bytes > UINT64_MAX - paged_plan.total_persistent_bytes) { + if (err) *err = "failed to plan paged KV memory for hybrid placement"; + return false; + } + kv_bytes += paged_plan.total_persistent_bytes; + } Ds4HybridBudgetInfo budget; const Ds4MoeTpConfig tp = ds4_moe_tp_config(cfg_.device.gpu); - if (!compute_ds4_hybrid_budget_info(w, backend_, max_ctx, - tp.all_on_secondary, budget, err)) { + if (!compute_ds4_hybrid_budget_info( + w, backend_, kv_bytes, tp.all_on_secondary, + cfg_.paged_attention, budget, err)) { return false; } @@ -1670,7 +1834,10 @@ bool DeepSeek4Backend::init_hybrid_model() { void DeepSeek4Backend::print_ready_banner() const { std::printf("[deepseek4-daemon] ready layers=%d ctx=%d experts=%d/%d\n", - w_.n_layer, cache_.max_ctx, w_.n_expert_used, w_.n_expert); + w_.n_layer, + cfg_.paged_attention ? (int)paged_cache_.plan.max_ctx + : cache_.max_ctx, + w_.n_expert_used, w_.n_expert); std::fflush(stdout); } @@ -1684,6 +1851,12 @@ bool DeepSeek4Backend::park(ParkTarget target) { std::fflush(stdout); } if (!want_target_model || parked_) return true; + if (cfg_.paged_attention) { + std::fprintf(stderr, + "[deepseek4] target park is unavailable while paged serving owns " + "live graph and slot state\n"); + return false; + } maybe_save_routing_stats(); for (int i = 0; i < PREFIX_SLOTS; ++i) { @@ -2185,7 +2358,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, finite ? max_value : 0.0f); } if (timing) { - log_step_tel("prefill", n_total, steps, elapsed_s(phase_t0), tel_acc); + log_deepseek4_step_telemetry("prefill", n_total, steps, elapsed_s(phase_t0), tel_acc); } if (spec_enabled_ && spec_drafter_) { deepseek4_release_prefill_scratch(cache_, moe_hybrid_.get()); @@ -2330,7 +2503,7 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, } } if (timing) { - log_step_tel("decode", (int)out_tokens.size(), steps, elapsed_s(phase_t0), tel_acc); + log_deepseek4_step_telemetry("decode", (int)out_tokens.size(), steps, elapsed_s(phase_t0), tel_acc); } return true; } @@ -2628,6 +2801,8 @@ void DeepSeek4Backend::shutdown() { for (int i = 0; i < PREFIX_SLOTS; i++) { snapshot_free(i); } + seq_engine_.reset(); + free_deepseek4_paged_cache(paged_cache_); free_deepseek4_cache(cache_); expert_runtime_.reset(); stream_engine_.destroy(); diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 583132042..ab8e1b64b 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 "deepseek4_seq_engine.h" #include "ggml.h" #include "ggml-backend.h" @@ -75,6 +76,7 @@ class DeepSeek4Backend : public ModelBackend { void free_drafter() override; void shutdown() override; + SeqEngine * seq_engine() override { return seq_engine_.get(); } const MoeHybridRoutingStats * get_routing_stats() const override { return routing_stats_.get(); @@ -87,6 +89,8 @@ class DeepSeek4Backend : public ModelBackend { ggml_backend_t expert_backend_ = nullptr; DeepSeek4Weights w_; DeepSeek4Cache cache_; + DeepSeek4PagedCache paged_cache_; + std::unique_ptr seq_engine_; bool parked_ = false; // Sampler @@ -178,6 +182,7 @@ class DeepSeek4Backend : public ModelBackend { MoeExpertComputeRuntime expert_runtime_; std::shared_ptr routing_stats_; std::string routing_stats_out_path_; + friend class DeepSeek4SeqEngine; }; } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index ee270cc7d..7d7ec2dae 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -48,6 +48,7 @@ static const Ds4MixedMoePolicy & ds4_mixed_moe_policy() { result.deferred_join_split = ds4_env_flag("DFLASH_DS4_TP_DEVICE_JOIN_SPLIT"); result.batch_peer_copies = + ds4_env_flag("DFLASH_DS4_TP_BATCH_SPLIT_COPIES") || ds4_env_flag("GGML_BATCH_PEER_COPIES") || ds4_env_flag("GGML_CUDA_BATCH_PEER_COPIES"); result.report_split_count = @@ -160,8 +161,12 @@ static void ds4_fused_verify_refresh_hybrid_luts( MoeHybridStorage * hybrid, DeepSeek4FusedDecodeGraph & fg) { if (!hybrid) return; + const char * masked_routes = + std::getenv("DFLASH_DS4_TP_MASKED_ROUTES"); const int32_t invalid_route = - ds4_env_flag("DFLASH_DS4_TP_MASKED_ROUTES") ? -1 : 0; + !masked_routes || !*masked_routes || std::strcmp(masked_routes, "0") != 0 + ? -1 + : 0; for (size_t hi = 0; hi < fg.hybrid_inputs.size(); ++hi) { const int il = (int) hi; const MoeHybridLayerStorage & layer = hybrid->layers[(size_t) il]; @@ -237,7 +242,8 @@ static bool ds4_fused_read_route_matrix( static void ds4_fused_consume_route_diagnostics( DeepSeek4FusedDecodeGraph & fg, const MoeHybridStorage * hybrid, - MoeHybridRoutingStats * routing_stats) { + MoeHybridRoutingStats * routing_stats, + const int32_t * active_slots = nullptr) { const bool cache_audit = ds4_env_flag("DFLASH_DS4_TP_CACHE_AUDIT"); if (!routing_stats && !cache_audit) return; @@ -259,6 +265,7 @@ static void ds4_fused_consume_route_diagnostics( continue; } for (int token = 0; token < route.n_tokens; ++token) { + if (active_slots && active_slots[route.lane_start + token] < 0) continue; const int32_t * token_ids = ids.data() + (size_t) token * route.width; const float * token_weights = weights.data() + @@ -364,7 +371,7 @@ static bool ds4_build_fused_verify_graph( Ds4FusedVerifyCache::Extra & ex, ggml_backend_t backend, const DeepSeek4Weights & w, - DeepSeek4Cache & cache, + DeepSeek4Cache * cache, const std::vector & hc_weights, const HcWeightsCpu & hc_out_weights, const std::vector & hash_tables, @@ -373,10 +380,64 @@ static bool ds4_build_fused_verify_graph( bool have_token_ids, const std::vector & capture_ids, MoeHybridStorage * hybrid, - std::vector && shape_key) { - if (fg.sched) { + std::vector && shape_key, + DeepSeek4PagedCache * paged_cache = nullptr, + const std::vector> * paged_rows = nullptr, + bool paged_history_buckets = false) { + const bool paged_mode = paged_cache && paged_rows; + if (!paged_mode && !cache) return false; + if (paged_mode) { + if (q < 1 || q > 16 || + paged_rows->size() != (size_t) w.n_layer) { + return false; + } + for (const auto & layer_rows : *paged_rows) { + if (layer_rows.size() != (size_t) q) return false; + } + } + // A step that packs chronological prompt rows keeps the established + // per-column arithmetic: dense projections run in parts sized by weight + // type (ROCmFP4 sixteen columns on the gfx1151 weight-reuse kernel, + // narrow F16 eight on MMVF, four otherwise), all bit-identical per column + // to the single-lane kernels, and routed experts run in parts of at most + // eight rows on the multi-token expert kernel. A round with one row per + // sequence keeps its unpartitioned graph. + int projection_columns = 0; + int moe_columns = 0; + if (paged_mode && !hybrid) { + std::vector active_slots; + int active_rows = 0; + for (const auto & row : (*paged_rows)[0]) { + if (row.slot < 0) continue; + ++active_rows; + if (std::find(active_slots.begin(), active_slots.end(), row.slot) == active_slots.end()) + active_slots.push_back(row.slot); + } + if (active_rows > (int)active_slots.size()) { + projection_columns = 4; + // The multi-token expert kernel gives every token its own warp + // with a fixed per-row K traversal, so a part's arithmetic does + // not depend on how many tokens share the launch; eight rows is + // the RDNA3 MUL_MAT_ID ceiling for the FP2/FP3 expert types. A + // single live sequence keeps one-row parts: its decode rounds use + // the single-token expert kernel, whose arithmetic differs from + // the multi-token kernel's. + moe_columns = active_slots.size() == 1 + ? 1 : std::min(DEEPSEEK4_MAX_MOE_ROWS_PER_PART, active_rows); + } + } + const size_t graph_capacity = q > 8 ? 131072u : 65536u; + std::array sched_backends{}; + if (hybrid) { + sched_backends = {backend, hybrid->cold_backend, hybrid->cpu_backend}; + } + const bool reuse_sched = + hybrid && fg.sched_reusable(sched_backends, graph_capacity); + if (fg.sched && !reuse_sched) { ggml_backend_sched_free(fg.sched); fg.sched = nullptr; + fg.sched_capacity = 0; + fg.sched_backends = {}; } step_graph_free(fg.sg); fg.reset_nodes(); @@ -401,7 +462,10 @@ static bool ds4_build_fused_verify_graph( fg.sg.ctx = ggml_init(params); if (!fg.sg.ctx) return false; ggml_context * ctx = fg.sg.ctx; - constexpr size_t graph_capacity = 65536u; + // Gathered paged attention builds one q=1 MLA lane per concurrent + // sequence. Above eight lanes the resulting whole-model graph exceeds + // the verifier-era 64K scheduler hash table even though the metadata + // arena still has ample room. fg.sg.gf = ggml_new_graph_custom(ctx, graph_capacity, false); ggml_cgraph * gf = fg.sg.gf; @@ -427,28 +491,71 @@ static bool ds4_build_fused_verify_graph( // row through the in-graph ring update and has no preserved-row suffix, // matching the established single-token fused-decode topology. int64_t mask_total = 0; - for (int il = 0; il < w.n_layer; ++il) { - const int ratio = (int) w.compress_ratios[il]; - int padded = 0; - if (ratio > 0 && cache.layers[(size_t) il].comp_kv) { - const int n_comp = ds4_comp_rows_used( - cache.layers[(size_t) il].comp_kv, - cache.layers[(size_t) il].n_comp, ratio, - kv_start + q - 1); - padded = ds4_padded_comp_rows( - n_comp, (int) cache.layers[(size_t) il].comp_kv->ne[1]); + if (paged_mode && paged_history_buckets) { + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + for (int t = 0; t < q; ++t) { + const auto & rows = (*paged_rows)[(size_t) il][(size_t) t]; + const int emitted = + ratio > 0 && rows.compressed_emitted ? 1 : 0; + mask_total += (int64_t) rows.raw_history.size() + 1 + + rows.compressed_history.size() + emitted; + } + } + } else if (!paged_mode) { + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + int padded = 0; + if (ratio > 0 && cache->layers[(size_t) il].comp_kv) { + const int n_comp = ds4_comp_rows_used( + cache->layers[(size_t) il].comp_kv, + cache->layers[(size_t) il].n_comp, ratio, + kv_start + q - 1); + padded = ds4_padded_comp_rows( + n_comp, (int) cache->layers[(size_t) il].comp_kv->ne[1]); + } + const int preserved_rows = q > 1 ? q : 0; + mask_total += (int64_t) (w.n_swa + padded + preserved_rows) * q; } - const int preserved_rows = q > 1 ? q : 0; - mask_total += (int64_t) (w.n_swa + padded + preserved_rows) * q; } - fg.mask_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, mask_total); + fg.mask_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, + std::max(mask_total, 1)); ggml_set_input(fg.mask_bundle); int64_t mask_off = 0; - // Keep every speculative lane in one HC tensor. The previous graph built - // five independent HC ops and progressively concatenated their outputs at - // every pre/post boundary. HC kernels already support a token dimension, - // so preserve it throughout the verifier and issue one batched op instead. + // Back every gathered lane's scalar and index inputs with shared + // tensors. Their sizes are fixed by the prepared history shape key. + int64_t paged_gather_total = 0; + if (paged_mode) { + for (int il = 0; il < w.n_layer; ++il) { + const int layer_ratio = (int) w.compress_ratios[il]; + for (int t = 0; t < q; ++t) { + const auto & rows = (*paged_rows)[(size_t) il][(size_t) t]; + paged_gather_total += + (int64_t) rows.raw_history.size() + 1; + if (layer_ratio > 0) { + paged_gather_total += (int64_t) std::max( + rows.compressed_history.size(), 1); + } + } + } + ex.paged_i32_n = (int64_t) 5 * w.n_layer * q; + ex.paged_i64_n = (int64_t) 8 * w.n_layer * q; + ex.paged_gather_n = std::max(paged_gather_total, 1); + ex.paged_i32 = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, ex.paged_i32_n); + ggml_set_input(ex.paged_i32); + ex.paged_i64 = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, ex.paged_i64_n); + ggml_set_input(ex.paged_i64); + ex.paged_gather = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, ex.paged_gather_n); + ggml_set_input(ex.paged_gather); + } + int64_t paged_gather_off = 0; + + // Keep every speculative or gathered lane in one HC tensor. HC kernels + // already support a token dimension, so issue one batched op per boundary. ggml_tensor * embed_3d = ggml_reshape_3d( ctx, fg.inp_embed, n_embd, 1, q); ggml_tensor * hc_repeated = ggml_repeat_4d( @@ -457,10 +564,12 @@ static bool ds4_build_fused_verify_graph( ctx, hc_repeated, (int64_t) n_embd * n_hc, q); std::vector capture_layers(capture_ids.size(), nullptr); + std::vector attention_contexts; for (int il = 0; il < w.n_layer; ++il) { const DeepSeek4Layer & L = w.layers[(size_t) il]; - DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + DeepSeek4LayerCache * lc = + paged_mode ? nullptr : &cache->layers[(size_t) il]; const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; const int ratio = (int) w.compress_ratios[il]; @@ -475,13 +584,251 @@ static bool ds4_build_fused_verify_graph( ggml_tensor * split_attn = nullptr; ggml_tensor * attn_working = ds4_build_fused_hc_pre( ctx, w, hc_flat, mc.fn_attn_f16[(size_t) il], L.hc_attn_base, - hlw.attn, &split_attn); + hlw.attn, &split_attn, projection_columns); if (!attn_working) return false; // HC-pre stores working and split values together per token. Materialize // the strided working view once, replacing q kernels and q-1 concats. ggml_tensor * attn_in = ggml_cont(ctx, attn_working); - // ── Batched attention ── + // ── Attention ── + ggml_tensor * attn_out = nullptr; + if (paged_mode) { + // q denotes independent decode lanes here, not a causal sequence. + // Gather each lane's immutable chronological history and run the + // established MLA lane core at q=1; all surrounding HC/MoE/output + // machinery remains q-wide and unchanged. + DeepSeek4PagedLayerCache & plc = paged_cache->layers[(size_t) il]; + ggml_tensor * raw_flat = ggml_reshape_2d( + ctx, plc.raw_kv, w.head_dim, + (int64_t) DS4_PAGE_TOKENS * paged_cache->plan.slots); + ggml_tensor * attn_normed = + build_rms_norm(ctx, attn_in, L.attn_norm, w.rms_eps); + DeepSeek4PreparedProjectedLane batched_proj = + build_mla_qkv_projection(ctx, attn_normed, w, L, q, projection_columns); + // Monolithic paged serving uses gfx1151. Its ROCmFP4 MMVQ + // preserves each column's arithmetic through sixteen rows. + // Only projections are shared: same-slot compressor mutations + // remain chronological inside each lane below. + if (!hybrid && q > 1 && ratio > 0 && L.attn_compressor_kv && + L.attn_compressor_gate && + L.attn_compressor_kv->type == GGML_TYPE_Q4_0_ROCMFP4_FAST && + L.attn_compressor_gate->type == GGML_TYPE_Q4_0_ROCMFP4_FAST) { + batched_proj.compressor_kv = ds4_mul_mat_columns( + ctx, L.attn_compressor_kv, attn_normed, projection_columns); + batched_proj.compressor_score = ds4_mul_mat_columns( + ctx, L.attn_compressor_gate, attn_normed, projection_columns); + } + build_mla_qkv_rope( + ctx, batched_proj, w, ds4_rope_params(w, ratio), q, + ex.pos_q, /*fuse_q_rope=*/false); + ggml_tensor * contexts = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, w.head_dim * w.n_head, q); + if (hybrid) attention_contexts.push_back(contexts); + bool independent = true; + int64_t raw_count = 0, comp_count = 0; + for (int t = 0; t < q; ++t) { + const auto & rows = (*paged_rows)[(size_t) il][(size_t) t]; + for (int prior = 0; prior < t; ++prior) { + if (rows.slot == (*paged_rows)[(size_t) il][(size_t) prior].slot) + independent = false; + } + } + for (int t = 0; t < q; ++t) { + const auto & rows = (*paged_rows)[(size_t) il][(size_t) t]; + raw_count += independent ? rows.raw_history.size() + 1 + : std::max(rows.raw_history.size(), 1); + if (ratio > 0) comp_count += std::max(rows.compressed_history.size(), 1); + } + const int64_t raw_begin = paged_gather_off; + const int64_t comp_begin = raw_begin + raw_count; + int64_t raw_offset = raw_begin, comp_offset = comp_begin; + for (int t = 0; t < q; ++t) { + const auto & rows = (*paged_rows)[(size_t) il][(size_t) t]; + auto & px = ex.paged.emplace_back(); + const int64_t lane_idx = (int64_t) il * q + t; + px.i32_base = lane_idx * 5; + px.i64_base = lane_idx * 7; + auto scalar_i32 = [&](int slot) { + return ggml_view_1d( + ctx, ex.paged_i32, 1, + (size_t) (px.i32_base + slot) * sizeof(int32_t)); + }; + auto scalar_i64 = [&](int slot) { + return ggml_view_1d( + ctx, ex.paged_i64, 1, + (size_t) (px.i64_base + slot) * sizeof(int64_t)); + }; + px.pos = scalar_i32(0); + px.neg_pos = scalar_i32(1); + px.raw_n = independent + ? (int64_t) rows.raw_history.size() + 1 + : (int64_t) std::max(rows.raw_history.size(), 1); + px.raw_off = raw_offset; + px.raw_gather = ggml_view_1d( + ctx, ex.paged_gather, px.raw_n, + (size_t) px.raw_off * sizeof(int32_t)); + raw_offset += px.raw_n; + if (ratio > 0) { + px.comp_n = (int64_t) std::max( + rows.compressed_history.size(), 1); + px.comp_off = comp_offset; + px.comp_gather = ggml_view_1d( + ctx, ex.paged_gather, px.comp_n, + (size_t) px.comp_off * sizeof(int32_t)); + comp_offset += px.comp_n; + if (ratio == 4) { + // The indexer and attention compressors address the + // same chronological compressed rows. Explicit mode + // has no indexer consumer, so share the gather and do + // not emit the otherwise-dead index-history read. + px.index_gather = px.comp_gather; + } + } + px.raw_write = scalar_i64(0); + if (ratio > 0) { + px.comp_write = scalar_i64(1); + px.comp_read = scalar_i32(2); + px.ape = scalar_i32(3); + px.state_row = scalar_i64(2); + px.comp_pos = scalar_i32(4); + } + } + paged_gather_off = comp_offset; + ggml_tensor * raw_stage = nullptr; + ggml_tensor * comp_stage = nullptr; + if (independent) { + ggml_tensor * indices = ggml_view_1d(ctx, ex.paged_gather, raw_count, + (size_t) raw_begin * sizeof(int32_t)); + raw_stage = ggml_get_rows(ctx, raw_flat, indices); + ggml_tensor * current_rows = ggml_view_1d(ctx, ex.paged_i64, q, + (size_t) (7 * w.n_layer * q + il * q) * sizeof(int64_t)); + raw_stage = ggml_set_rows(ctx, raw_stage, ggml_cont(ctx, batched_proj.kv), current_rows); + ggml_build_forward_expand(gf, raw_stage); + if (comp_count > 0) { + indices = ggml_view_1d(ctx, ex.paged_gather, comp_count, + (size_t) comp_begin * sizeof(int32_t)); + comp_stage = ggml_get_rows(ctx, plc.comp_kv, indices); + ggml_build_forward_expand(gf, comp_stage); + } + } + for (int t = 0; t < q; ++t) { + const auto & rows = (*paged_rows)[(size_t) il][(size_t) t]; + auto & px = ex.paged[(size_t) il * q + t]; + ggml_tensor * raw_history = rows.raw_history.empty() ? nullptr + : raw_stage ? ggml_view_2d(ctx, raw_stage, w.head_dim, + rows.raw_history.size(), raw_stage->nb[1], + (size_t) (px.raw_off - raw_begin) * raw_stage->nb[1]) + : ggml_get_rows(ctx, raw_flat, px.raw_gather); + ggml_tensor * comp_history = rows.compressed_history.empty() ? nullptr + : comp_stage ? ggml_view_2d(ctx, comp_stage, w.head_dim, + rows.compressed_history.size(), comp_stage->nb[1], + (size_t) (px.comp_off - comp_begin) * comp_stage->nb[1]) + : ggml_get_rows(ctx, plc.comp_kv, px.comp_gather); + ggml_tensor * index_history = nullptr; + DeepSeek4CompressorState attn_state{}, index_state{}; + if (ratio > 0 && rows.slot >= 0) { + const size_t off = (size_t) rows.slot * plc.attn_compressor.state_kv->nb[2]; + attn_state.state_kv = ggml_view_2d(ctx, plc.attn_compressor.state_kv, + plc.attn_compressor.state_kv->ne[0], plc.attn_compressor.state_kv->ne[1], + plc.attn_compressor.state_kv->nb[1], off); + attn_state.state_score = ggml_view_2d(ctx, plc.attn_compressor.state_score, + plc.attn_compressor.state_score->ne[0], plc.attn_compressor.state_score->ne[1], + plc.attn_compressor.state_score->nb[1], off); + if (ratio == 4) { + const size_t io = (size_t) rows.slot * plc.indexer_compressor.state_kv->nb[2]; + index_state.state_kv = ggml_view_2d(ctx, plc.indexer_compressor.state_kv, + plc.indexer_compressor.state_kv->ne[0], plc.indexer_compressor.state_kv->ne[1], + plc.indexer_compressor.state_kv->nb[1], io); + index_state.state_score = ggml_view_2d(ctx, plc.indexer_compressor.state_score, + plc.indexer_compressor.state_score->ne[0], plc.indexer_compressor.state_score->ne[1], + plc.indexer_compressor.state_score->nb[1], io); + } + } + DeepSeek4MlaLaneBindings lane{}; + lane.history_mode = DeepSeek4MlaLaneBindings::HistoryMode::ChronologicalGathered; + lane.prepared_raw_attention = raw_stage ? ggml_view_2d( + ctx, raw_stage, w.head_dim, rows.raw_history.size() + 1, + raw_stage->nb[1], (size_t) (px.raw_off - raw_begin) * raw_stage->nb[1]) : nullptr; + lane.raw_history = raw_history; lane.n_raw_history = (int) rows.raw_history.size(); + lane.comp_history = comp_history; lane.n_comp_history = (int) rows.compressed_history.size(); + lane.index_comp_history = index_history; + lane.n_index_comp_history = (int) rows.compressed_history.size(); + // Writes use the same flattened physical-row geometry as the + // gather indices (not the slot-local 128-row second axis). + lane.raw_kv = raw_flat; lane.comp_kv = plc.comp_kv; + lane.index_comp_kv = plc.index_comp_kv; + lane.raw_write_rows = px.raw_write; lane.comp_write_rows = px.comp_write; + lane.index_comp_write_rows = px.comp_write; + lane.comp_read_rows = px.comp_read; + lane.index_comp_read_rows = px.comp_read; + lane.write_enabled = rows.slot >= 0; + lane.attn_compressor = &attn_state; lane.indexer_compressor = &index_state; + DeepSeek4AttentionGraphInputs ci{}; + ci.rope_pos = px.pos; ci.neg_pos = px.neg_pos; + if (ratio == 4) { + ci.flush_rows = ggml_view_1d(ctx, ex.paged_i64, 4, + (size_t) (px.i64_base + 3) * sizeof(int64_t)); + } + ci.attn_ape_row = px.ape; ci.attn_state_rows = px.state_row; + ci.attn_comp_pos = px.comp_pos; ci.index_ape_row = px.ape; + ci.index_state_rows = px.state_row; ci.index_comp_pos = px.comp_pos; + if (paged_history_buckets) { + px.mask_off = mask_off; + px.mask_n = (int64_t) rows.raw_history.size() + 1 + + rows.compressed_history.size() + + (ratio > 0 && rows.compressed_emitted ? 1 : 0); + ci.attn_row_mask = ggml_view_2d( + ctx, fg.mask_bundle, px.mask_n, 1, + (size_t) px.mask_n * sizeof(float), + (size_t) px.mask_off * sizeof(float)); + ci.padded_comp = (int) rows.compressed_history.size() + + (ratio > 0 && rows.compressed_emitted ? 1 : 0); + mask_off += px.mask_n; + } + std::vector ib; + std::vector iab; + std::vector lab; + ggml_tensor * col = ggml_view_2d( + ctx, attn_normed, n_embd, 1, attn_normed->nb[1], + (size_t) t * attn_normed->nb[1]); + const DeepSeek4PreparedProjectedLane lane_proj = + ds4_slice_projected_lane(ctx, batched_proj, w, t, px.pos); + ggml_tensor * lane_context = nullptr; + ggml_tensor * one = build_mla_attention_lane_core( + ctx, gf, col, w, L, lane, il, (int) rows.position, 1, + &ci, ib, iab, lab, nullptr, + DeepSeek4AttentionImpl::Explicit, &lane_proj, + &lane_context); + if (!one || !lane_context || !ib.empty() || !iab.empty() || + !lab.empty()) { + std::fprintf(stderr, + "[deepseek4-paged] layer %d lane %d attention build " + "failed (graph=%d i32=%zu arrays=%zu i64=%zu, " + "first_array=%d)\n", + il, t, one != nullptr, ib.size(), iab.size(), lab.size(), + iab.empty() || iab[0].values.empty() + ? INT32_MIN : iab[0].values[0]); + return false; + } + // Complete the chronological read before a later prompt row + // mutates this slot's raw ring or compressor state. + ggml_build_forward_expand(gf, lane_context); + ggml_tensor * destination = ggml_view_2d( + ctx, contexts, w.head_dim * w.n_head, 1, contexts->nb[1], + (size_t) t * contexts->nb[1]); + ggml_build_forward_expand(gf, ggml_cpy(ctx, lane_context, destination)); + } + const auto rope = ds4_rope_params(w, ratio); + attn_out = build_tail_rope_3d( + ctx, ggml_reshape_3d(ctx, contexts, w.head_dim, w.n_head, q), + ex.neg_q, w.n_rot, w.head_dim, w.n_head, q, + rope.freq, rope.scale, rope.ext, rope.attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope.n_ctx_orig); + attn_out = ggml_reshape_2d(ctx, attn_out, w.head_dim * w.n_head, q); + attn_out = build_mla_output_projection( + ctx, attn_out, w, L, q, /*allow_grouped=*/false, projection_columns); + } else { + // ── Batched speculative attention ── DeepSeek4AttentionGraphInputs ain{}; ain.rope_pos = ex.pos_q; ain.neg_pos = ex.neg_q; @@ -519,10 +866,10 @@ static bool ds4_build_fused_verify_graph( (size_t) il * 2 * sizeof(int64_t)); } int padded = 0; - if (ratio > 0 && lc.comp_kv) { + if (ratio > 0 && lc->comp_kv) { const int n_comp = ds4_comp_rows_used( - lc.comp_kv, lc.n_comp, ratio, lane_token_pos); - padded = ds4_padded_comp_rows(n_comp, (int) lc.comp_kv->ne[1]); + lc->comp_kv, lc->n_comp, ratio, lane_token_pos); + padded = ds4_padded_comp_rows(n_comp, (int) lc->comp_kv->ne[1]); } const int64_t n_attn = (int64_t) w.n_swa + padded + (lane_q > 1 ? lane_q : 0); @@ -535,7 +882,7 @@ static bool ds4_build_fused_verify_graph( std::vector i32ab; std::vector i64ab; ggml_tensor * normed = build_rms_norm(ctx, attn_in, L.attn_norm, w.rms_eps); - ggml_tensor * attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, il, + attn_out = build_mla_attention(ctx, gf, normed, w, L, *lc, il, lane_kv_start, lane_q, &ain, i32b, i32ab, i64ab); if (!attn_out) return false; @@ -543,6 +890,7 @@ static bool ds4_build_fused_verify_graph( std::fprintf(stderr, "[ds4-fused-verify] layer %d dynamic bindings; cannot fuse\n", il); return false; } + } // ── Batched HC post (attention) + HC pre (FFN) ── ggml_tensor * attn_batch = ggml_is_contiguous(attn_out) @@ -553,7 +901,7 @@ static bool ds4_build_fused_verify_graph( ggml_tensor * split_ffn = nullptr; ggml_tensor * ffn_working = ds4_build_fused_hc_pre( ctx, w, hc_flat, mc.fn_ffn_f16[(size_t) il], L.hc_ffn_base, - hlw.ffn, &split_ffn); + hlw.ffn, &split_ffn, projection_columns); if (!ffn_working) return false; ggml_tensor * ffn_in = ggml_cont(ctx, ffn_working); @@ -600,8 +948,16 @@ static bool ds4_build_fused_verify_graph( router_weights = routing.weights; router_nodes = routing.nodes; } else { - ffn_out = build_moe_ffn( - ctx, ffn_normed, w, L, il, lane_q); + const int columns = moe_columns > 0 ? moe_columns : lane_q; + for (int first = 0; first < lane_q; first += columns) { + const int count = std::min(columns, lane_q - first); + ggml_tensor * part = moe_columns > 0 + ? ggml_cont(ctx, ggml_view_2d(ctx, ffn_normed, n_embd, count, + ffn_normed->nb[1], (size_t)first * ffn_normed->nb[1])) + : ffn_normed; + ggml_tensor * one = build_moe_ffn(ctx, part, w, L, il, count); + ffn_out = ffn_out ? ggml_concat(ctx, ffn_out, one, 1) : one; + } } if (hybrid) { @@ -655,14 +1011,15 @@ static bool ds4_build_fused_verify_graph( (hybrid_layer.gate_up_hot || (hybrid_layer.gate_hot && hybrid_layer.up_hot)) && hybrid_layer.down_hot; - // Legacy q>1 graphs represented six routes as 4 + (2 padded to - // 4). Native q1 decode already uses width six directly. Keep the - // decomposition as a bisectable fallback while validating the - // dedicated q4 MoE kernel's native route width on both GPUs. + // Legacy q>1 graphs represent six routes as 4 + 2. Native q1 + // decode already uses width six directly. Keep the decomposition + // as a bisectable fallback while validating the dedicated q4 MoE + // kernel's native route width on both GPUs. // The dedicated multi-token MUL_MAT_ID kernel is indexed by token // width, not route width. Once that kernel is enabled, keep all - // six routes in one owner graph: the legacy 4 + padded-2 split - // computes two duplicate expert paths whose weights are zero. + // six routes in one owner graph. Until then, retain the existing + // head-plus-tail reduction order without padding the two-route + // tail with duplicate zero-weight expert paths. const bool native_route_width = !same_gpu_runtime || mixed_policy.native_route_width; @@ -672,59 +1029,32 @@ static bool ds4_build_fused_verify_graph( if (route_width != 6) return false; hybrid_cfg.n_expert_used = 4; - ggml_tensor * first_ids = nullptr; - ggml_tensor * first_weights = nullptr; - ggml_tensor * padded_ids = nullptr; - ggml_tensor * padded_weights = nullptr; - for (int t = 0; t < lane_q; ++t) { - const size_t ids_row = (size_t) t * selected->nb[1]; - const size_t weights_row = (size_t) t * router_weights->nb[1]; - - ggml_tensor * ids_head = ggml_cont(ctx, ggml_view_1d( - ctx, selected, 4, ids_row)); - ggml_tensor * weights_head = ggml_cont(ctx, ggml_view_1d( - ctx, router_weights, 4, weights_row)); - first_ids = first_ids - ? ggml_concat(ctx, first_ids, ids_head, 1) - : ids_head; - first_weights = first_weights - ? ggml_concat(ctx, first_weights, weights_head, 1) - : weights_head; - - ggml_tensor * ids_tail = ggml_cont(ctx, ggml_view_1d( - ctx, selected, 2, ids_row + 4 * selected->nb[0])); - ggml_tensor * weights_tail = ggml_cont(ctx, ggml_view_1d( - ctx, router_weights, 2, - weights_row + 4 * router_weights->nb[0])); - ggml_tensor * ids_tail_padded = ggml_concat( - ctx, ids_tail, ids_tail, 0); - ggml_tensor * weights_tail_padded = ggml_concat( - ctx, weights_tail, ggml_scale(ctx, weights_tail, 0.0f), 0); - padded_ids = padded_ids - ? ggml_concat(ctx, padded_ids, ids_tail_padded, 1) - : ids_tail_padded; - padded_weights = padded_weights - ? ggml_concat(ctx, padded_weights, weights_tail_padded, 1) - : weights_tail_padded; + DeepSeek4Head4Tail2Routes routes; + if (!build_deepseek4_head4_tail2_routes( + ctx, selected, router_weights, lane_q, routes)) { + return false; } if (!build_moe_hybrid_ffn_graph( ctx, hybrid_schedule_graph, hybrid_cfg, desc, hybrid->layers[(size_t) il], ffn_normed, - first_ids, first_weights, lane_q, inputs, true, + routes.head_ids, routes.head_weights, + lane_q, inputs, true, allow_fused_combine, join_mode, MoeHybridRouteBalance::Disabled)) { return false; } ffn_out = inputs.output; + hybrid_cfg.n_expert_used = 2; if (!build_moe_hybrid_ffn_graph( ctx, hybrid_schedule_graph, hybrid_cfg, desc, hybrid->layers[(size_t) il], ffn_normed, - padded_ids, padded_weights, lane_q, inputs, false, + routes.tail_ids, routes.tail_weights, + lane_q, inputs, false, allow_fused_combine, join_mode, MoeHybridRouteBalance::Disabled)) { return false; @@ -749,7 +1079,11 @@ static bool ds4_build_fused_verify_graph( } } } - if (!ffn_out) return false; + if (!ffn_out) { + std::fprintf(stderr, + "[deepseek4-paged] layer %d FFN graph build failed\n", il); + return false; + } // ── Batched HC post (FFN); capture at drafter layers ── if (fused_hc_join_inputs) { @@ -775,20 +1109,21 @@ static bool ds4_build_fused_verify_graph( ctx, mean, n_embd, lane_q); // [n_embd,q] } } + GGML_ASSERT(mask_off == mask_total); // ── Batched output HC merge → out_norm + lm_head ── ggml_tensor * onorm = ggml_rms_norm(ctx, hc_cur, w.hc_eps); - ggml_tensor * omix = ggml_mul_mat(ctx, mc.fn_out_f16, onorm); + ggml_tensor * omix = ds4_mul_mat_columns(ctx, mc.fn_out_f16, onorm, projection_columns); omix = ggml_reshape_2d(ctx, omix, ggml_nelements(omix) / q, q); ggml_tensor * obase = ds4_fused_hc_base_f32(ctx, w.output_hc_base); if (!obase || hc_out_weights.scale_data.empty()) return false; ggml_tensor * final_all = ggml_ds4_hc_out( ctx, omix, obase, hc_cur, n_hc, hc_out_weights.scale_data[0]); ggml_tensor * out_normed = build_rms_norm(ctx, final_all, w.out_norm, w.rms_eps); - fg.logits = ggml_mul_mat(ctx, w.output, out_normed); // [n_vocab, q] + fg.logits = ds4_mul_mat_columns(ctx, w.output, out_normed, projection_columns); // [n_vocab, q] ggml_set_output(fg.logits); ggml_build_forward_expand(gf, fg.logits); - if (ds4_env_flag("DFLASH_DS4_GPU_ARGMAX_VERIFY")) { + if (paged_mode || ds4_env_flag("DFLASH_DS4_GPU_ARGMAX_VERIFY")) { ex.argmax = ggml_argmax(ctx, fg.logits); ggml_set_output(ex.argmax); ggml_build_forward_expand(gf, ex.argmax); @@ -866,14 +1201,22 @@ static bool ds4_build_fused_verify_graph( warned_cross_vendor_join = true; } } - ggml_backend_t backends[3] = { - backend, peer, hybrid->cpu_backend}; - fg.sched = ggml_backend_sched_new( - backends, nullptr, 3, graph_capacity, false, true); - if (!fg.sched) { - std::fprintf(stderr, - "[ds4-fused-verify] scheduler creation failed\n"); - return false; + if (reuse_sched) { + // Drop prior tensor assignments and deferred events while keeping + // the scheduler-owned pinned staging allocations alive. + ggml_backend_sched_reset(fg.sched); + } else { + ggml_backend_t backends[3] = { + backend, peer, hybrid->cpu_backend}; + fg.sched = ggml_backend_sched_new( + backends, nullptr, 3, graph_capacity, false, true); + if (!fg.sched) { + std::fprintf(stderr, + "[ds4-fused-verify] scheduler creation failed\n"); + return false; + } + fg.sched_capacity = graph_capacity; + fg.sched_backends = sched_backends; } const bool late_join_split = mixed_policy.late_join_split; const MoeHybridGraphPolicy & moe_policy = @@ -903,6 +1246,7 @@ static bool ds4_build_fused_verify_graph( ggml_backend_sched_set_tensor_backend(fg.sched, tensor, peer); } }; + for (ggml_tensor * context : attention_contexts) pin_main(context); pin_main(fg.inp_embed); pin_main(ex.pos_q); pin_main(ex.neg_q); @@ -914,6 +1258,15 @@ static bool ds4_build_fused_verify_graph( pin_main(fg.i32_bundle); pin_main(fg.i64_bundle); pin_main(fg.mask_bundle); + pin_main(ex.paged_i32); + pin_main(ex.paged_i64); + pin_main(ex.paged_gather); + for (const auto & px : ex.paged) { + pin_main(px.pos); pin_main(px.neg_pos); pin_main(px.raw_gather); pin_main(px.comp_gather); + pin_main(px.index_gather); + pin_main(px.raw_write); pin_main(px.comp_write); pin_main(px.comp_read); pin_main(px.ape); + pin_main(px.state_row); pin_main(px.comp_pos); + } for (ggml_tensor * hids : fg.hash_ids) pin_main(hids); for (const MoeHybridGraphInputs & inputs : fg.hybrid_inputs) { if (mixed_policy.pin_route_weights) { @@ -1145,7 +1498,7 @@ static int ds4_try_fused_verify_step( ex->reset(); const auto build_t0 = Ds4TimingClock::now(); - if (!ds4_build_fused_verify_graph(mc, *fg, *ex, backend, w, cache, + if (!ds4_build_fused_verify_graph(mc, *fg, *ex, backend, w, &cache, hc_weights, hc_out_weights, hash_tables, kv_start, q, token_ids != nullptr, *hooks->capture_layer_ids, hybrid, diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index ac2fb4ff7..c8b5641e0 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -11,6 +11,7 @@ #include "deepseek4_internal.h" #include "deepseek4_hc_cuda.h" #include "deepseek4_roctx.h" +#include "deepseek4_page_layout.h" #include "internal.h" #include "../common/step_graph.h" #include "../common/cuda_graph_overrides.h" @@ -178,6 +179,42 @@ static void observe_active_routing(MoeHybridRoutingStats * stats, } // namespace +bool build_deepseek4_head4_tail2_routes( + ggml_context * ctx, + ggml_tensor * selected, + ggml_tensor * router_weights, + int n_tokens, + DeepSeek4Head4Tail2Routes & out) { + out = {}; + if (!ctx || !selected || !router_weights || n_tokens <= 0 || + selected->type != GGML_TYPE_I32 || + router_weights->type != GGML_TYPE_F32 || + selected->ne[0] != 6 || router_weights->ne[0] != 6 || + selected->ne[1] != n_tokens || + router_weights->ne[1] != n_tokens || + selected->ne[2] != 1 || selected->ne[3] != 1 || + router_weights->ne[2] != 1 || router_weights->ne[3] != 1) { + return false; + } + + const auto materialize = [ctx, n_tokens]( + ggml_tensor * source, + int first_route, + int width) { + ggml_tensor * view = ggml_view_2d( + ctx, source, width, n_tokens, source->nb[1], + (size_t) first_route * source->nb[0]); + return ggml_cont(ctx, view); + }; + + out.head_ids = materialize(selected, 0, 4); + out.head_weights = materialize(router_weights, 0, 4); + out.tail_ids = materialize(selected, 4, 2); + out.tail_weights = materialize(router_weights, 4, 2); + return out.head_ids && out.head_weights && + out.tail_ids && out.tail_weights; +} + int deepseek4_previous_raw_ring_spans( int kv_start, int n_swa, @@ -576,12 +613,12 @@ static ggml_tensor * ds4_cast_if_needed( // ─── Helper: Partial RoPE (tail rotation) ─────────────────────────────── // DS4 applies RoPE only to the last n_rot dimensions of each head. -// ggml_rope_ext applies to the first n_dims, so we split, rope the tail, concat. -// -// x: [head_dim, n_heads, n_tokens] (3D) — applies tail RoPE to each head. -// pos: [n_tokens] I32 — position for each token. -// Returns: [head_dim, n_heads, n_tokens] with last n_rot dims rotated. - +// DS4 rotates the tail n_rot dims of each head with sequential pairs +// (GGML_ROPE_TYPE_NORMAL). GGML_ROPE_TYPE_TAIL rotates the last n_dims in +// place and passes the head through, so this is one launch instead of two +// contiguity copies, a rotation, and a concat; every rotated element sees the +// same angle it saw in the extracted tail tensor (theta and the YaRN +// correction still derive from n_rot), so the result is bit-identical. static ggml_tensor * build_tail_rope_3d(ggml_context * ctx, ggml_tensor * x, ggml_tensor * pos, @@ -596,23 +633,11 @@ static ggml_tensor * build_tail_rope_3d(ggml_context * ctx, float beta_fast, float beta_slow, int n_ctx_orig) { - const int n_nope = head_dim - n_rot; - // Split: nope [n_nope, n_heads, n_tokens], tail [n_rot, n_heads, n_tokens] - ggml_tensor * nope = ggml_view_3d(ctx, x, n_nope, n_heads, n_tokens, - x->nb[1], x->nb[2], 0); - ggml_tensor * tail = ggml_view_3d(ctx, x, n_rot, n_heads, n_tokens, - x->nb[1], x->nb[2], - (size_t)n_nope * ggml_type_size(x->type)); - // tail is non-contiguous (stride between heads = head_dim, not n_rot) - tail = ggml_cont(ctx, tail); - // Apply rope to the contiguous tail: [n_rot, n_heads, n_tokens] - // DS4 uses standard sequential pairs (i, i+1), which is GGML_ROPE_TYPE_NORMAL - tail = ggml_rope_ext(ctx, tail, pos, nullptr, - n_rot, GGML_ROPE_TYPE_NORMAL, n_ctx_orig, + GGML_ASSERT(x->ne[0] == head_dim && x->ne[1] == n_heads && x->ne[2] == n_tokens); + return ggml_rope_ext(ctx, x, pos, nullptr, + n_rot, GGML_ROPE_TYPE_NORMAL | GGML_ROPE_TYPE_TAIL, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - // Concat nope + tail along dim 0 → [head_dim, n_heads, n_tokens] - return ggml_concat(ctx, ggml_cont(ctx, nope), tail, 0); } // For KV (single head): x is [head_dim, n_tokens] @@ -999,7 +1024,11 @@ static void build_compressor_step( ggml_tensor * cur_all = nullptr, int n_tokens_all = 1, int kv_start_all = -1, - bool indexer_qat = false) { + bool indexer_qat = false, + ggml_tensor ** current_comp_out = nullptr, + bool paged_physical_row = false, + ggml_tensor * prepared_kv = nullptr, + ggml_tensor * prepared_score = nullptr) { if (!gf || !cur_last || !ape || !kv_proj || !gate_proj || !norm_weight || !state.state_kv || !state.state_score || !comp_cache || ratio <= 0) { return; @@ -1026,8 +1055,10 @@ static void build_compressor_step( // For ratio-4: write into second half of state (rows ratio..2*ratio-1) const int row = (ratio == 4) ? (ratio + pos_mod) : pos_mod; - ggml_tensor * kv_cur = ggml_mul_mat(ctx, kv_proj, cur_last); - ggml_tensor * sc_cur = ggml_mul_mat(ctx, gate_proj, cur_last); + // Gathered lanes may share the token-independent projections. State + // writes, pooling and rotation below still execute in lane order. + ggml_tensor * kv_cur = prepared_kv ? prepared_kv : ggml_mul_mat(ctx, kv_proj, cur_last); + ggml_tensor * sc_cur = prepared_score ? prepared_score : ggml_mul_mat(ctx, gate_proj, cur_last); ggml_tensor * state_kv_source = state.state_kv; ggml_tensor * state_score_source = state.state_score; ggml_tensor * comp_cache_source = comp_cache; @@ -1210,10 +1241,14 @@ static void build_compressor_step( if (indexer_qat) { pooled = ggml_ds4_indexer_qat(ctx, ggml_cont(ctx, pooled)); } + if (current_comp_out) { + *current_comp_out = pooled; + } ggml_tensor * pooled_f16 = ggml_cast(ctx, pooled, GGML_TYPE_F16); const int comp_row = token_pos / ratio; - if (comp_row >= (int) comp_cache->ne[1]) { + if ((!comp_rows_inp || !paged_physical_row) && + comp_row >= (int) comp_cache->ne[1]) { return; } @@ -1418,7 +1453,8 @@ static void build_indexer_compressor_step( ggml_tensor * cur_last, const DeepSeek4Weights & w, const DeepSeek4Layer & L, - DeepSeek4LayerCache & lc, + DeepSeek4CompressorState & indexer_compressor, + ggml_tensor * index_comp_kv, int token_pos, ggml_tensor * ape_row_inp, ggml_tensor * state_rows_inp, @@ -1431,14 +1467,16 @@ static void build_indexer_compressor_step( ggml_tensor * cur_all = nullptr, int n_tokens_all = 1, int kv_start_all = -1, - bool indexer_qat = false) { + bool indexer_qat = false, + ggml_tensor ** current_comp_out = nullptr, + bool paged_physical_row = false) { build_compressor_step(ctx, gf, cur_last, L.indexer_compressor_ape, L.indexer_compressor_kv, L.indexer_compressor_gate, L.indexer_compressor_norm, - lc.indexer_compressor, - lc.index_comp_kv, + indexer_compressor, + index_comp_kv, 4, w.n_indexer_head_dim, // indexer head_dim = 128 token_pos, @@ -1460,7 +1498,9 @@ static void build_indexer_compressor_step( cur_all, n_tokens_all, kv_start_all, - indexer_qat); + indexer_qat, + current_comp_out, + paged_physical_row); } static int ds4_comp_rows_used(const ggml_tensor * comp_cache, int n_cached, int ratio, int token_pos) { @@ -1529,6 +1569,13 @@ static int ds4_padded_comp_rows(int n_comp, int cap) { return padded < cap ? padded : cap; } +static int ds4_padded_gathered_raw_rows(int n_raw) { + if (n_raw <= 0) return 0; + constexpr int stride = 16; + const int padded = ((n_raw + stride - 1) / stride) * stride; + return std::min(padded, (int) DS4_PAGE_TOKENS - 1); +} + static ggml_tensor * build_indexer_topk( ggml_context * ctx, ggml_tensor * qr_norm, // [n_lora_q, n_tokens] @@ -1632,13 +1679,254 @@ static ggml_tensor * build_indexer_topk( // ─── MLA Attention Block ──────────────────────────────────────────────── -static ggml_tensor * build_mla_attention( +// All persistent and live-state bindings consumed by one MLA lane. Keeping +// this internal seam tensor-based is intentional: a paged adapter can later +// supply gathered history and slot-specific compressor state without the +// graph builder consulting DeepSeek4LayerCache or host cache counters. +struct DeepSeek4MlaLaneBindings { + enum class HistoryMode { + ContiguousRing, + ChronologicalGathered, + }; + + HistoryMode history_mode = HistoryMode::ContiguousRing; + // In gathered mode these are immutable, chronological attention inputs. + // Counts are explicit so adapters may bind capacity-padded tensors. + ggml_tensor * prepared_raw_attention = nullptr; + ggml_tensor * raw_history = nullptr; + int n_raw_history = 0; + ggml_tensor * comp_history = nullptr; + int n_comp_history = 0; + ggml_tensor * index_comp_history = nullptr; + int n_index_comp_history = 0; + + // Persistent mutation targets are deliberately independent of history. + ggml_tensor * raw_kv = nullptr; + ggml_tensor * comp_kv = nullptr; + ggml_tensor * index_comp_kv = nullptr; + ggml_tensor * raw_write_rows = nullptr; + ggml_tensor * comp_write_rows = nullptr; + ggml_tensor * index_comp_write_rows = nullptr; + ggml_tensor * comp_read_rows = nullptr; // GET_ROWS requires I32 + ggml_tensor * index_comp_read_rows = nullptr; + + // Optional passive outputs let a future adapter scatter current products. + ggml_tensor ** current_raw_out = nullptr; + ggml_tensor ** current_comp_out = nullptr; + ggml_tensor ** current_index_comp_out = nullptr; + // False is the padding/inactive-lane contract: build attention against the + // supplied padded history, but emit no persistent current-row mutations. + bool write_enabled = true; + DeepSeek4CompressorState * attn_compressor = nullptr; + DeepSeek4CompressorState * indexer_compressor = nullptr; + int n_comp_live = 0; + int n_index_comp_live = 0; + int n_comp_committed = 0; +}; + +// Projection/RoPE products handed to the history/update portion of a lane. +// This is deliberately a passive bundle: introducing graph operations in a +// separate builder would risk changing decode graph ordering. +struct DeepSeek4PreparedProjectedLane { + ggml_tensor * normalized_q_lora = nullptr; + ggml_tensor * q = nullptr; + ggml_tensor * kv = nullptr; + ggml_tensor * rope_pos = nullptr; + ggml_tensor * compressor_kv = nullptr; + ggml_tensor * compressor_score = nullptr; +}; + +// Per-layer RoPE parameters. Compressed layers use YaRN scaling, and +// attn_factor cancels the magnitude scaling rope_yarn applies. +struct Ds4RopeParams { + float freq = 0.0f; + float scale = 1.0f; + float ext = 0.0f; + float attn = 1.0f; + int n_ctx_orig = 0; +}; + +static Ds4RopeParams ds4_rope_params(const DeepSeek4Weights & w, int ratio) { + const bool compressed = ratio > 0; + Ds4RopeParams p; + p.freq = compressed ? w.compress_rope_freq_base : w.rope_freq_base; + p.scale = compressed ? (1.0f / w.rope_scale_factor) : 1.0f; + p.ext = compressed ? 1.0f : 0.0f; + if (p.ext != 0.0f && p.scale > 0.0f) { + p.attn /= (1.0f + 0.1f * logf(1.0f / p.scale)); + } + p.n_ctx_orig = (int) w.rope_orig_ctx; + return p; +} + +// Keep a packed prompt step on per-column-exact matmul dispatch. The +// requested width is the generic per-column-exact MMVQ/MMVF width; ROCmFP4 +// dense weights have a gfx1151 weight-reuse kernel that is bit-identical per +// column up to sixteen columns and reads the weights once, and narrow F16 +// weights stay exact on MMVF through eight columns. Higher tensor axes +// (output-projection groups) stay intact. +static int ds4_projection_part_columns(const ggml_tensor * weights, int columns) { + if (columns <= 0) return 0; + if (weights->type == GGML_TYPE_Q4_0_ROCMFP4_FAST) return std::max(columns, 16); + if (weights->type == GGML_TYPE_F16) return std::max(columns, 8); + return columns; +} + +static ggml_tensor * ds4_mul_mat_columns( + ggml_context * ctx, ggml_tensor * weights, ggml_tensor * input, + int columns) { + columns = ds4_projection_part_columns(weights, columns); + if (columns <= 0 || input->ne[1] <= columns) { + return ggml_mul_mat(ctx, weights, input); + } + ggml_tensor * result = nullptr; + for (int64_t first = 0; first < input->ne[1]; first += columns) { + const int64_t count = std::min(columns, input->ne[1] - first); + ggml_tensor * part = ggml_view_4d( + ctx, input, input->ne[0], count, input->ne[2], input->ne[3], + input->nb[1], input->nb[2], input->nb[3], + (size_t)first * input->nb[1]); + ggml_tensor * projected = ggml_mul_mat(ctx, weights, part); + result = result ? ggml_concat(ctx, result, projected, 1) : projected; + } + return result; +} + +// Q/KV projections and their tail RoPE are independent per token. A paged +// caller can evaluate them once at width q and hand each lane a column view, +// avoiding one reread of all three projection weights per active lane. +static DeepSeek4PreparedProjectedLane build_mla_qkv_projection( + ggml_context * ctx, + ggml_tensor * cur, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + int n_tokens, int projection_columns = 0) { + DeepSeek4PreparedProjectedLane out; + ggml_tensor * qr = ds4_mul_mat_columns(ctx, L.attn_q_a, cur, projection_columns); + qr = build_rms_norm(ctx, qr, L.attn_q_a_norm, w.rms_eps); + ggml_tensor * q = ds4_mul_mat_columns(ctx, L.attn_q_b, qr, projection_columns); + q = ggml_reshape_3d(ctx, q, w.head_dim, w.n_head, n_tokens); + q = ggml_rms_norm(ctx, q, w.rms_eps); + + ggml_tensor * kv = ds4_mul_mat_columns(ctx, L.attn_kv, cur, projection_columns); + kv = build_rms_norm(ctx, kv, L.attn_kv_a_norm, w.rms_eps); + + out.normalized_q_lora = qr; + out.q = q; + out.kv = kv; + return out; +} + +static void build_mla_qkv_rope( + ggml_context * ctx, + DeepSeek4PreparedProjectedLane & p, + const DeepSeek4Weights & w, + const Ds4RopeParams & rope, + int n_tokens, + ggml_tensor * rope_pos, + bool fuse_q_rope) { + if (!fuse_q_rope) { + p.q = build_tail_rope_3d(ctx, p.q, rope_pos, w.n_rot, w.head_dim, + w.n_head, n_tokens, rope.freq, rope.scale, + rope.ext, rope.attn, w.rope_yarn_beta_fast, + w.rope_yarn_beta_slow, rope.n_ctx_orig); + } + p.kv = build_tail_rope_2d(ctx, p.kv, rope_pos, w.n_rot, w.head_dim, + n_tokens, rope.freq, rope.scale, rope.ext, + rope.attn, w.rope_yarn_beta_fast, + w.rope_yarn_beta_slow, rope.n_ctx_orig); + p.rope_pos = rope_pos; +} + +// Grouped low-rank output projection. Several independent gathered lanes can +// concatenate their pre-projection contexts and share one q-wide evaluation. +static ggml_tensor * build_mla_output_projection( + ggml_context * ctx, + ggml_tensor * attn_out, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + int n_tokens, + bool allow_grouped, int projection_columns = 0) { + const int group_dim = w.head_dim * (w.n_head / w.n_out_group); + attn_out = ggml_reshape_3d( + ctx, attn_out, group_dim, w.n_out_group, n_tokens); + attn_out = ggml_permute(ctx, attn_out, 0, 2, 1, 3); + if (n_tokens == 1) { + attn_out = ggml_cont(ctx, attn_out); + } + ggml_tensor * out_a_3d = ggml_reshape_3d( + ctx, L.attn_output_a, group_dim, w.n_lora_o, w.n_out_group); + ggml_tensor * attn_low = ds4_mul_mat_columns(ctx, out_a_3d, attn_out, projection_columns); + + const bool grouped_output_projection = + allow_grouped && n_tokens > 1 && + !ds4_env_flag("DFLASH_DS4_DISABLE_GROUPED_OUTPUT_PROJECTION"); + if (grouped_output_projection) { + return ggml_mul_mat_grouped_src(ctx, L.attn_output_b, attn_low); + } + attn_low = ggml_cont(ctx, ggml_permute(ctx, attn_low, 0, 2, 1, 3)); + attn_low = ggml_reshape_2d( + ctx, attn_low, (int64_t) w.n_lora_o * w.n_out_group, n_tokens); + return ds4_mul_mat_columns(ctx, L.attn_output_b, attn_low, projection_columns); +} + +// One lane's column of a batched prologue. These are views only. +static DeepSeek4PreparedProjectedLane ds4_slice_projected_lane( + ggml_context * ctx, + const DeepSeek4PreparedProjectedLane & batched, + const DeepSeek4Weights & w, + int lane, + ggml_tensor * lane_rope_pos) { + DeepSeek4PreparedProjectedLane out; + out.normalized_q_lora = ggml_view_2d( + ctx, batched.normalized_q_lora, batched.normalized_q_lora->ne[0], 1, + batched.normalized_q_lora->nb[1], + (size_t) lane * batched.normalized_q_lora->nb[1]); + out.q = ggml_view_3d( + ctx, batched.q, w.head_dim, w.n_head, 1, + batched.q->nb[1], batched.q->nb[2], + (size_t) lane * batched.q->nb[2]); + out.kv = ggml_view_2d( + ctx, batched.kv, batched.kv->ne[0], 1, batched.kv->nb[1], + (size_t) lane * batched.kv->nb[1]); + out.rope_pos = lane_rope_pos; + if (batched.compressor_kv) { + out.compressor_kv = ggml_view_2d(ctx, batched.compressor_kv, + batched.compressor_kv->ne[0], 1, batched.compressor_kv->nb[1], + (size_t) lane * batched.compressor_kv->nb[1]); + out.compressor_score = ggml_view_2d(ctx, batched.compressor_score, + batched.compressor_score->ne[0], 1, batched.compressor_score->nb[1], + (size_t) lane * batched.compressor_score->nb[1]); + } + return out; +} + +static DeepSeek4MlaLaneBindings deepseek4_contiguous_lane_bindings( + DeepSeek4LayerCache & lc, + int ratio, + int token_pos) { + DeepSeek4MlaLaneBindings lane; + lane.history_mode = DeepSeek4MlaLaneBindings::HistoryMode::ContiguousRing; + lane.raw_kv = lc.raw_kv; + lane.comp_kv = lc.comp_kv; + lane.index_comp_kv = lc.index_comp_kv; + lane.attn_compressor = &lc.attn_compressor; + lane.indexer_compressor = &lc.indexer_compressor; + lane.n_comp_live = ratio > 0 + ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; + lane.n_index_comp_live = ratio == 4 + ? ds4_comp_rows_used(lc.index_comp_kv, lc.n_index_comp, 4, token_pos) : 0; + lane.n_comp_committed = lc.n_comp; + return lane; +} + +static ggml_tensor * build_mla_attention_lane_core( ggml_context * ctx, ggml_cgraph * gf, ggml_tensor * cur, // [n_embd, n_tokens] const DeepSeek4Weights & w, const DeepSeek4Layer & L, - DeepSeek4LayerCache & lc, + const DeepSeek4MlaLaneBindings & lane, int layer_idx, int kv_start, int n_tokens, @@ -1647,47 +1935,36 @@ static ggml_tensor * build_mla_attention( std::vector & i32_array_inputs, std::vector & i64_array_inputs, std::vector * f32_array_inputs = nullptr, - DeepSeek4AttentionImpl attention_impl = DeepSeek4AttentionImpl::Explicit) { + DeepSeek4AttentionImpl attention_impl = DeepSeek4AttentionImpl::Explicit, + const DeepSeek4PreparedProjectedLane * prepared = nullptr, + ggml_tensor ** out_attn_context = nullptr) { const int n_embd = w.n_embd; const int head_dim = w.head_dim; const int n_head = w.n_head; const int n_rot = w.n_rot; - const int n_out_group = w.n_out_group; - const int n_lora_o = w.n_lora_o; const int ratio = w.compress_ratios[layer_idx]; + const bool gathered_history = lane.history_mode == + DeepSeek4MlaLaneBindings::HistoryMode::ChronologicalGathered; - // ── Q path: cur → q_a → norm → q_b → per-head norm ───────────── - // q_a: [n_embd, n_tokens] → [n_lora_q, n_tokens] - ggml_tensor * qr = ggml_mul_mat(ctx, L.attn_q_a, cur); - // qr_norm is reused by the ratio-4 indexer before the main q_b projection. - qr = build_rms_norm(ctx, qr, L.attn_q_a_norm, w.rms_eps); - // q_b: [n_lora_q, n_tokens] → [n_head * head_dim, n_tokens] - ggml_tensor * q = ggml_mul_mat(ctx, L.attn_q_b, qr); - // Reshape to [head_dim, n_head, n_tokens] for per-head ops - q = ggml_reshape_3d(ctx, q, head_dim, n_head, n_tokens); - // Reference DS4 applies unweighted RMSNorm independently to every Q head. - q = ggml_rms_norm(ctx, q, w.rms_eps); - - // ── KV path: cur → kv → norm ─────────────────────────────────── - // kv: [n_embd, n_tokens] → [head_dim, n_tokens] - ggml_tensor * kv = ggml_mul_mat(ctx, L.attn_kv, cur); - kv = build_rms_norm(ctx, kv, L.attn_kv_a_norm, w.rms_eps); + // Existing callers leave prepared null and emit the original prologue in + // place. Only gathered paged concurrency supplies a q-wide projection. + DeepSeek4PreparedProjectedLane projected; + if (!prepared) { + projected = build_mla_qkv_projection(ctx, cur, w, L, n_tokens); + } // ── RoPE on Q and KV (tail rotation on last n_rot dims) ──────── - // DS4 uses per-layer RoPE params: compressed layers get YaRN scaling. - const bool compressed = (ratio > 0); - const float rope_freq = compressed ? w.compress_rope_freq_base : w.rope_freq_base; - const float rope_scale = compressed ? (1.0f / w.rope_scale_factor) : 1.0f; - const float rope_ext = compressed ? 1.0f : 0.0f; - // For YaRN: attn_factor cancels the magnitude scaling in rope_yarn - float rope_attn = 1.0f; - if (rope_ext != 0.0f && rope_scale > 0.0f) { - rope_attn /= (1.0f + 0.1f * logf(1.0f / rope_scale)); - } + const Ds4RopeParams rope = ds4_rope_params(w, ratio); + const float rope_freq = rope.freq; + const float rope_scale = rope.scale; + const float rope_ext = rope.ext; + const float rope_attn = rope.attn; + const int rope_n_ctx_orig = rope.n_ctx_orig; // Position tensor for this token batch - ggml_tensor * rope_pos = cached_inputs ? cached_inputs->rope_pos : nullptr; + ggml_tensor * rope_pos = prepared ? prepared->rope_pos + : (cached_inputs ? cached_inputs->rope_pos : nullptr); if (!rope_pos) { rope_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); ggml_set_input(rope_pos); @@ -1696,22 +1973,20 @@ static ggml_tensor * build_mla_attention( i32_array_inputs.push_back({rope_pos, std::move(pos_vals)}); } - // n_ctx_orig is critical for YaRN correction on compressed layers - const int rope_n_ctx_orig = (int)w.rope_orig_ctx; // 65536 - // 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 && n_tokens > 1 && head_dim == 512 && n_rot == 64; - if (!fuse_q_rope) { - q = build_tail_rope_3d(ctx, q, rope_pos, n_rot, head_dim, n_head, n_tokens, - rope_freq, rope_scale, rope_ext, rope_attn, - w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + if (prepared) { + projected = *prepared; + } else { + build_mla_qkv_rope( + ctx, projected, w, rope, n_tokens, rope_pos, fuse_q_rope); } - kv = build_tail_rope_2d(ctx, kv, rope_pos, n_rot, head_dim, n_tokens, - rope_freq, rope_scale, rope_ext, rope_attn, - w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + ggml_tensor * qr = projected.normalized_q_lora; + ggml_tensor * q = projected.q; + ggml_tensor * kv = projected.kv; // ── Causal batched step (exact multi-token target semantics) ─── // The target model is causal: token i must not attend to batch tokens @@ -1728,15 +2003,15 @@ static ggml_tensor * build_mla_attention( ggml_tensor * old_rows_scratch_f16 = nullptr; int n_old_rows = 0; ggml_tensor * prior_rows_scratch = nullptr; - int n_prior_rows = 0; + int n_prior_rows = gathered_history ? lane.n_raw_history : 0; const bool fused_causal = cached_inputs && cached_inputs->attn_row_mask && n_tokens > 1; - if (fused_causal) { + 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, lc.raw_kv, head_dim, 1, lc.raw_kv->nb[1], - (size_t)((kv_start + ti) % w.n_swa) * lc.raw_kv->nb[1]); + 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 @@ -1745,14 +2020,14 @@ static ggml_tensor * build_mla_attention( } old_rows_scratch_f16 = old_rows_scratch; old_rows_scratch = ds4_cast_if_needed(ctx, old_rows_scratch, GGML_TYPE_F32); - } else if (causal_batch && !layer_major_batch) { + } else if (!gathered_history && causal_batch && !layer_major_batch) { // Copy the to-be-overwritten rows FIRST; same-stream build order runs // these before the ring writes below. for (int ti = 0; ti < n_tokens; ti++) { if (kv_start + ti < w.n_swa) continue; // slot never held an older pos ggml_tensor * slot = ggml_view_2d( - ctx, lc.raw_kv, head_dim, 1, lc.raw_kv->nb[1], - (size_t)((kv_start + ti) % w.n_swa) * lc.raw_kv->nb[1]); + 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 @@ -1762,7 +2037,7 @@ static ggml_tensor * build_mla_attention( if (old_rows_scratch) { old_rows_scratch = ds4_cast_if_needed(ctx, old_rows_scratch, GGML_TYPE_F32); } - } else if (layer_major_batch) { + } else if (!gathered_history && layer_major_batch) { // Snapshot the chronological pre-chunk window before any ring writes. // Attention then consumes [prior F16 rows | current F32 rows], matching // the single-token path and avoiding an F16 round-trip for this chunk. @@ -1772,8 +2047,8 @@ static ggml_tensor * build_mla_attention( const int tail = std::min(n_prior_rows, w.n_swa - first); auto snapshot_span = [&](int row, int count) { ggml_tensor * span = ggml_view_2d( - ctx, lc.raw_kv, head_dim, count, lc.raw_kv->nb[1], - (size_t) row * lc.raw_kv->nb[1]); + ctx, lane.raw_kv, head_dim, count, lane.raw_kv->nb[1], + (size_t) row * lane.raw_kv->nb[1]); return ggml_cont(ctx, span); }; prior_rows_scratch = snapshot_span(first, tail); @@ -1791,13 +2066,18 @@ static ggml_tensor * build_mla_attention( // ── Store ALL KV rows in the raw SWA ring ───────────────────── // For decode (n_tokens=1): write single row. For prefill: write all rows. - ggml_tensor * raw_kv_source = lc.raw_kv; - ggml_tensor * raw_kv_rows = cached_inputs - ? cached_inputs->raw_kv_rows - : nullptr; - if (raw_kv_rows) { + ggml_tensor * raw_kv_source = lane.raw_kv; + ggml_tensor * raw_kv_rows = lane.raw_write_rows + ? lane.raw_write_rows + : (cached_inputs ? cached_inputs->raw_kv_rows : nullptr); + if (lane.current_raw_out) { + *lane.current_raw_out = kv; + } + if (!lane.write_enabled) { + // Inactive/padding lanes intentionally have no cache mutation. + } else if (raw_kv_rows) { ggml_tensor * kv_f32 = ggml_is_contiguous(kv) ? kv : ggml_cont(ctx, kv); - raw_kv_source = ggml_set_rows(ctx, lc.raw_kv, kv_f32, raw_kv_rows); + raw_kv_source = ggml_set_rows(ctx, lane.raw_kv, kv_f32, raw_kv_rows); ggml_build_forward_expand(gf, raw_kv_source); } else { // The attention graph consumes the whole current ubatch directly. @@ -1809,8 +2089,8 @@ static ggml_tensor * build_mla_attention( ggml_tensor * kv_row = ggml_view_2d( ctx, kv, head_dim, 1, kv->nb[1], (size_t)ti * kv->nb[1]); ggml_tensor * kv_slot = ggml_view_2d( - ctx, lc.raw_kv, head_dim, 1, lc.raw_kv->nb[1], - (size_t)(pos_ti % w.n_swa) * lc.raw_kv->nb[1]); + ctx, lane.raw_kv, head_dim, 1, lane.raw_kv->nb[1], + (size_t)(pos_ti % w.n_swa) * lane.raw_kv->nb[1]); ggml_build_forward_expand(gf, ggml_cpy(ctx, ggml_cast(ctx, kv_row, GGML_TYPE_F16), kv_slot)); } } @@ -1819,15 +2099,15 @@ static ggml_tensor * build_mla_attention( // ── Learned compression update ────────────────────────────────── ggml_tensor * cur_last = ggml_view_2d( ctx, cur, n_embd, 1, cur->nb[1], (size_t)(n_tokens - 1) * cur->nb[1]); - ggml_tensor * comp_kv_source = lc.comp_kv; - if (ratio > 0 && L.attn_compressor_kv) { + ggml_tensor * comp_kv_source = lane.comp_kv; + if (lane.write_enabled && ratio > 0 && L.attn_compressor_kv) { build_compressor_step(ctx, gf, cur_last, L.attn_compressor_ape, L.attn_compressor_kv, L.attn_compressor_gate, L.attn_compressor_norm, - lc.attn_compressor, - lc.comp_kv, + *lane.attn_compressor, + lane.comp_kv, ratio, head_dim, token_pos, @@ -1840,7 +2120,8 @@ static ggml_tensor * build_mla_attention( (int)w.rope_orig_ctx, cached_inputs ? cached_inputs->attn_ape_row : nullptr, cached_inputs ? cached_inputs->attn_state_rows : nullptr, - cached_inputs ? cached_inputs->attn_comp_rows : nullptr, + lane.comp_write_rows ? lane.comp_write_rows : + (cached_inputs ? cached_inputs->attn_comp_rows : nullptr), cached_inputs ? cached_inputs->attn_comp_pos : nullptr, i64_array_inputs, i32_array_inputs, @@ -1848,15 +2129,28 @@ static ggml_tensor * build_mla_attention( cached_inputs ? cached_inputs->flush_rows : nullptr, (causal_batch || fused_causal) ? cur : nullptr, n_tokens, - kv_start); - } - - ggml_tensor * index_comp_kv_source = lc.index_comp_kv; - if (ratio == 4 && L.indexer_compressor_kv) { - build_indexer_compressor_step(ctx, gf, cur_last, w, L, lc, token_pos, + kv_start, + false, + lane.current_comp_out, + gathered_history, + prepared ? prepared->compressor_kv : nullptr, + prepared ? prepared->compressor_score : nullptr); + } + + ggml_tensor * index_comp_kv_source = lane.index_comp_kv; + // Gathered paged concurrency always uses Explicit attention, whose + // build_indexer_topk path is disabled. In that mode the indexer compressor + // only writes state that no graph node reads, so omit the dead subgraph. + const bool indexer_compressor_is_dead = + gathered_history && attention_impl != DeepSeek4AttentionImpl::SparseFlash; + if (lane.write_enabled && ratio == 4 && L.indexer_compressor_kv && + !indexer_compressor_is_dead) { + build_indexer_compressor_step(ctx, gf, cur_last, w, L, + *lane.indexer_compressor, lane.index_comp_kv, token_pos, cached_inputs ? cached_inputs->index_ape_row : nullptr, cached_inputs ? cached_inputs->index_state_rows : nullptr, - cached_inputs ? cached_inputs->index_comp_rows : nullptr, + lane.index_comp_write_rows ? lane.index_comp_write_rows : + (cached_inputs ? cached_inputs->index_comp_rows : nullptr), cached_inputs ? cached_inputs->index_comp_pos : nullptr, i64_array_inputs, i32_array_inputs, @@ -1866,7 +2160,9 @@ static ggml_tensor * build_mla_attention( n_tokens, kv_start, attention_impl == - DeepSeek4AttentionImpl::SparseFlash); + DeepSeek4AttentionImpl::SparseFlash, + lane.current_index_comp_out, + gathered_history); } // ── MLA Dot-Product Attention (SWA + compressed KV) ──────────── @@ -1875,43 +2171,66 @@ static ggml_tensor * build_mla_attention( // comp_kv: [head_dim, comp_cap] F16 compressed rows. // n_raw = min(kv_start + n_tokens, n_swa) const bool masked_kv = cached_inputs && cached_inputs->attn_row_mask; - const int n_comp_live = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; + const bool gathered_emits_comp = gathered_history && lane.write_enabled && + ratio > 0 && ((token_pos + 1) % ratio) == 0; + const int n_comp_live = gathered_history + ? lane.n_comp_history + (gathered_emits_comp ? 1 : 0) : lane.n_comp_live; + ggml_tensor * comp_history_source = gathered_history + ? lane.comp_history : comp_kv_source; + ggml_tensor * index_comp_history_source = gathered_history + ? lane.index_comp_history : index_comp_kv_source; + if (gathered_emits_comp) { + // Gather through the post-update source to make the compressor write a + // graph dependency. Reading the F16 cache row preserves ordinary q=1 + // rounding at a boundary instead of feeding the transient F32 pool. + ggml_tensor * emitted = ggml_get_rows( + ctx, comp_kv_source, lane.comp_read_rows); + comp_history_source = lane.comp_history + ? ggml_concat(ctx, lane.comp_history, emitted, 1) : emitted; + if (ratio == 4) { + ggml_tensor * index_emitted = ggml_get_rows( + ctx, index_comp_kv_source, lane.index_comp_read_rows); + index_comp_history_source = lane.index_comp_history + ? ggml_concat(ctx, lane.index_comp_history, index_emitted, 1) + : index_emitted; + } + } ggml_tensor * indexer_topk = nullptr; if (attention_impl == DeepSeek4AttentionImpl::SparseFlash && ratio == 4 && f32_array_inputs) { - const int n_index_comp_live = ds4_comp_rows_used( - lc.index_comp_kv, lc.n_index_comp, 4, token_pos); - // Attention and index compression advance together at ratio 4. Reusing - // the attention mask is safe only while that invariant and the index - // buffer capacity hold; fail at graph construction if state diverges. - GGML_ASSERT(lc.index_comp_kv && index_comp_kv_source); - GGML_ASSERT(n_index_comp_live == n_comp_live); - GGML_ASSERT(!masked_kv || - cached_inputs->padded_comp <= lc.index_comp_kv->ne[1]); - // Use the same padded span as attention in a replayable decode graph. - // The dynamic compressed portion of attn_row_mask is added to the - // indexer scores, so padding stays invisible while live rows can grow - // within the fixed graph shape. - const int n_index_comp = masked_kv - ? cached_inputs->padded_comp - : n_index_comp_live; + int n_index_comp = 0; ggml_tensor * index_visibility_mask = nullptr; - if (masked_kv && n_index_comp > 0) { - index_visibility_mask = 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)); + if (gathered_history) { + n_index_comp = lane.n_index_comp_history + + (gathered_emits_comp ? 1 : 0); + } else { + const int n_index_comp_live = lane.n_index_comp_live; + // Attention and index compression advance together at ratio 4. + GGML_ASSERT(lane.index_comp_kv && index_comp_kv_source); + GGML_ASSERT(n_index_comp_live == n_comp_live); + GGML_ASSERT(!masked_kv || + cached_inputs->padded_comp <= lane.index_comp_kv->ne[1]); + n_index_comp = masked_kv + ? cached_inputs->padded_comp + : n_index_comp_live; + if (masked_kv && n_index_comp > 0) { + index_visibility_mask = 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)); + } } indexer_topk = build_indexer_topk( - ctx, qr, cur, w, L, index_comp_kv_source, + ctx, qr, cur, w, L, index_comp_history_source, n_index_comp, kv_start, n_tokens, rope_pos, index_visibility_mask, i32_array_inputs); } // Stable path reads the full physical ring (masking not-yet-written slots) // and a padded compressed-row span; the plain path reads only valid rows. - const int n_raw = masked_kv ? w.n_swa + const int n_raw = gathered_history ? lane.n_raw_history + n_tokens + : masked_kv ? w.n_swa : layer_major_batch ? n_prior_rows + n_tokens : std::min(kv_start + n_tokens, w.n_swa); const int n_comp_attn = masked_kv ? cached_inputs->padded_comp : n_comp_live; @@ -1923,13 +2242,26 @@ static ggml_tensor * build_mla_attention( // write and see the previous contents of the raw KV slot. auto raw_kv_view = [&](int row, int count) -> ggml_tensor * { ggml_tensor * view = ggml_view_2d( - ctx, lc.raw_kv, head_dim, count, lc.raw_kv->nb[1], - (size_t)row * lc.raw_kv->nb[1]); + ctx, lane.raw_kv, head_dim, count, lane.raw_kv->nb[1], + (size_t)row * lane.raw_kv->nb[1]); return ds4_cast_if_needed(ctx, view, GGML_TYPE_F32); }; ggml_tensor * kv_attn = nullptr; - if (masked_kv) { + if (lane.prepared_raw_attention) { + kv_attn = lane.prepared_raw_attention; + } else if (gathered_history) { + ggml_tensor * current = ds4_cast_if_needed(ctx, kv, GGML_TYPE_F32); + if (lane.n_raw_history > 0 && lane.raw_history) { + ggml_tensor * history = ggml_view_2d( + ctx, lane.raw_history, head_dim, lane.n_raw_history, + lane.raw_history->nb[1], 0); + history = ds4_cast_if_needed(ctx, history, GGML_TYPE_F32); + kv_attn = ggml_concat(ctx, history, current, 1); + } else { + kv_attn = current; + } + } else if (masked_kv) { // Fused stable-KV path: read the full physical ring; rows not yet // written are masked to -1e30 in the score matrix (exact 0 after // softmax). Only the fused decode graph sets attn_row_mask. Read @@ -1953,7 +2285,7 @@ static ggml_tensor * build_mla_attention( // KV at its runtime row in an F32 snapshot instead. The tokenwise // prefill helper takes the same branch and row ordering. ggml_tensor * ring = ggml_view_2d( - ctx, lc.raw_kv, head_dim, w.n_swa, lc.raw_kv->nb[1], 0); + ctx, lane.raw_kv, head_dim, w.n_swa, lane.raw_kv->nb[1], 0); ring = ds4_cast_if_needed(ctx, ring, GGML_TYPE_F32); kv_attn = ggml_set_rows(ctx, ring, cur_kv, raw_kv_rows); ggml_build_forward_expand(gf, kv_attn); @@ -1978,21 +2310,20 @@ static ggml_tensor * build_mla_attention( attention_impl == DeepSeek4AttentionImpl::Explicit && kv_attn->type == GGML_TYPE_F32 && raw_kv_source->type == GGML_TYPE_F16 && - (!comp_kv_source || comp_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) { // 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. - // 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_kv_source) { + if (n_comp_attn > 0 && comp_history_source) { ggml_tensor * comp = ggml_view_2d( - ctx, comp_kv_source, head_dim, n_comp_attn, - comp_kv_source->nb[1], 0); + ctx, comp_history_source, head_dim, n_comp_attn, + comp_history_source->nb[1], 0); kv_attn = ggml_concat(ctx, kv_attn, comp, 1); } if (old_rows_scratch_f16) { @@ -2005,13 +2336,12 @@ static ggml_tensor * build_mla_attention( "[deepseek4] fused explicit F16 K/V active: tokens=%d " "compressed=%d\n", n_tokens, n_comp_attn); - explicit_f16_kv_logged = true; } } else { - if (n_comp_attn > 0 && comp_kv_source) { + if (n_comp_attn > 0 && comp_history_source) { ggml_tensor * comp = ggml_view_2d( - ctx, comp_kv_source, head_dim, n_comp_attn, - comp_kv_source->nb[1], 0); + ctx, comp_history_source, head_dim, n_comp_attn, + comp_history_source->nb[1], 0); comp = ds4_cast_if_needed(ctx, comp, GGML_TYPE_F32); kv_attn = ggml_concat(ctx, kv_attn, comp, 1); } @@ -2058,7 +2388,9 @@ static ggml_tensor * build_mla_attention( } } if (n_comp_attn > 0) { - const int vis = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, pos_i); + const int vis = gathered_history ? n_comp_attn + : ds4_comp_rows_used( + lane.comp_kv, lane.n_comp_committed, ratio, pos_i); for (int c = vis; c < n_comp_attn; c++) col[n_raw + c] = -1e30f; } } @@ -2082,8 +2414,9 @@ static ggml_tensor * build_mla_attention( if (pos_r > pos_i) col[r] = -1e30f; } if (n_comp_attn > 0) { - const int visible = ds4_comp_rows_used( - lc.comp_kv, lc.n_comp, ratio, pos_i); + const int visible = gathered_history ? n_comp_attn + : ds4_comp_rows_used( + lane.comp_kv, lane.n_comp_committed, ratio, pos_i); for (int c = visible; c < n_comp_attn; ++c) { col[n_raw + c] = -1e30f; } @@ -2135,7 +2468,7 @@ static ggml_tensor * build_mla_attention( const int first_count = DS4_NUMERICAL_PREFILL_BAND; const int second_count = n_tokens - first_count; const int first_comp = ratio > 0 - ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, + ? ds4_comp_rows_used(lane.comp_kv, lane.n_comp_committed, ratio, kv_start + first_count - 1) : 0; const int second_comp = n_comp_live; @@ -2179,7 +2512,7 @@ static ggml_tensor * build_mla_attention( } if (comp_count > 0) { const int visible = ds4_comp_rows_used( - lc.comp_kv, lc.n_comp, ratio, pos_i); + lane.comp_kv, lane.n_comp_committed, ratio, pos_i); for (int c = visible; c < comp_count; ++c) { col[raw_count + c] = -1e30f; } @@ -2315,6 +2648,19 @@ static ggml_tensor * build_mla_attention( f32_kv_short_logged = true; } } + // DS4 adds one learned per-head sink logit to the denominator, but the + // sink contributes no value vector. + ggml_tensor * probs = nullptr; + // Keep long histories on the general softmax path; the fused sink + // kernel requires shared memory for every attention column. + if (L.attn_sinks && !score_mask && n_tokens == 1 && n_attn <= 2048) { + // Single-token lanes: scale, sink concat, softmax, and the strided + // view (which forced a 2D staging copy before the PV matmul) fold + // into one launch. The sink is a virtual last column, so every + // reduction matches the concat form bit for bit. + ggml_tensor * sinks = ggml_reshape_1d(ctx, L.attn_sinks, n_head); + probs = ggml_soft_max_ext_sink_col(ctx, scores, sinks, kq_scale); + } else { scores = ggml_scale(ctx, scores, kq_scale); if (score_mask) { if (n_tokens > 1) { @@ -2328,10 +2674,6 @@ static ggml_tensor * build_mla_attention( scores = ggml_add(ctx, scores, score_mask); } } - - // DS4 adds one learned per-head sink logit to the denominator, but the - // sink contributes no value vector. - ggml_tensor * probs = nullptr; if (L.attn_sinks) { ggml_tensor * sink_scores = ggml_reshape_2d(ctx, L.attn_sinks, 1, n_head); @@ -2349,6 +2691,7 @@ static ggml_tensor * build_mla_attention( } else { probs = ggml_soft_max(ctx, scores); } + } ggml_tensor * kv_t = ggml_cont(ctx, ggml_transpose(ctx, kv_attn)); context = ggml_mul_mat(ctx, kv_t, probs); if (explicit_f16_f32_kv_short) { @@ -2361,7 +2704,7 @@ static ggml_tensor * build_mla_attention( } // ── Inverse tail RoPE on attention output ─────────────────────── - if (!inverse_rope_fused) { + if (!inverse_rope_fused && !out_attn_context) { ggml_tensor * neg_pos = cached_inputs ? cached_inputs->neg_pos : nullptr; if (!neg_pos) { neg_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); @@ -2380,46 +2723,41 @@ static ggml_tensor * build_mla_attention( // Flatten to [head_dim*n_head, n_tokens] for output projection ggml_tensor * attn_out = ggml_reshape_2d(ctx, context, head_dim * n_head, n_tokens); - // ── Grouped output projection ────────────────────────────────── - // DS4 output uses grouped low-rank projection: - // attn_out: [head_dim*n_head, n_tokens] → reshape [group_dim, n_tokens, n_groups] - // out_a: [group_dim, n_groups*n_lora_o] → reshape [group_dim, n_lora_o, n_groups] - // batched matmul over n_groups: → [n_lora_o, n_tokens, n_groups] - // → reshape [n_lora_o*n_groups, n_tokens] - // out_b: [n_lora_o*n_groups, n_embd] → final: [n_embd, n_tokens] - const int group_dim = head_dim * (n_head / n_out_group); // 512 * 8 = 4096 - // Reshape attn_out: [32768, n_tokens] → [4096, 8, n_tokens] → permute to [4096, n_tokens, 8] - attn_out = ggml_reshape_3d(ctx, attn_out, group_dim, n_out_group, n_tokens); - attn_out = ggml_permute(ctx, attn_out, 0, 2, 1, 3); - if (n_tokens == 1) { - attn_out = ggml_cont(ctx, attn_out); - } - // attn_out is now [group_dim, n_tokens, n_out_group] - ggml_tensor * out_a_3d = ggml_reshape_3d(ctx, L.attn_output_a, group_dim, n_lora_o, n_out_group); - // out_a_3d: [group_dim, n_lora_o, n_out_group] — ne[2] matches - ggml_tensor * attn_low = ggml_mul_mat(ctx, out_a_3d, attn_out); - // attn_low: [n_lora_o, n_tokens, n_out_group] - ggml_tensor * out = nullptr; - const bool grouped_output_projection = - n_tokens > 1 && - !ds4_env_flag("DFLASH_DS4_DISABLE_GROUPED_OUTPUT_PROJECTION"); - if (grouped_output_projection) { - // Batched ROCmFPX MMQ consumes src1's channel stride directly. This - // avoids materializing both permutations (~256 MiB/layer at 2K). - out = ggml_mul_mat_grouped_src(ctx, L.attn_output_b, attn_low); - } else { - // Preserve the established single-token graph and provide an exact - // fallback for heterogeneous runtimes that cannot retain grouped-view - // metadata across a scheduler copy. At verifier widths (q <= 4), this - // materializes at most 128 KiB per layer rather than the long-prefill - // volume avoided by the grouped path. - attn_low = ggml_cont(ctx, ggml_permute(ctx, attn_low, 0, 2, 1, 3)); - attn_low = ggml_reshape_2d( - ctx, attn_low, n_lora_o * n_out_group, n_tokens); - out = ggml_mul_mat(ctx, L.attn_output_b, attn_low); + if (out_attn_context) { + *out_attn_context = attn_out; + return attn_out; } - return out; + return build_mla_output_projection(ctx, attn_out, w, L, n_tokens, + /*allow_grouped=*/true); +} + +// Legacy contiguous-cache adapter. Both decode and the consecutive q>1 +// verifier/prefill path enter through here, so their graph construction order +// remains exactly the order in build_mla_attention_lane_core. +static ggml_tensor * build_mla_attention( + ggml_context * ctx, + ggml_cgraph * gf, + ggml_tensor * cur, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + DeepSeek4LayerCache & lc, + int layer_idx, + int kv_start, + int n_tokens, + const DeepSeek4AttentionGraphInputs * cached_inputs, + std::vector & i32_inputs, + std::vector & i32_array_inputs, + std::vector & i64_array_inputs, + std::vector * f32_array_inputs = nullptr, + DeepSeek4AttentionImpl attention_impl = DeepSeek4AttentionImpl::Explicit) { + const int ratio = w.compress_ratios[layer_idx]; + DeepSeek4MlaLaneBindings lane = deepseek4_contiguous_lane_bindings( + lc, ratio, kv_start + n_tokens - 1); + return build_mla_attention_lane_core( + ctx, gf, cur, w, L, lane, layer_idx, kv_start, n_tokens, + cached_inputs, i32_inputs, i32_array_inputs, i64_array_inputs, + f32_array_inputs, attention_impl); } struct DeepSeek4CachedDecodeHcPreGraph { @@ -4734,6 +5072,16 @@ struct DeepSeek4FusedDecodeGraph { std::vector authoritative_routes; ggml_tensor * logits = nullptr; ggml_backend_sched_t sched = nullptr; + // The scheduler owns large pinned cross-backend staging buffers. Retain it + // across gathered-paged shape rebuilds when its backend set and capacity + // still match. + size_t sched_capacity = 0; + std::array sched_backends{}; + + bool sched_reusable(const std::array & backends, + size_t capacity) const { + return sched && sched_capacity >= capacity && sched_backends == backends; + } void reset_nodes() { inp_embed = nullptr; @@ -4769,6 +5117,22 @@ struct DeepSeek4FusedDecodeGraph { } } + // Retain shape-independent resources, but first retire native graph + // executables whose keys point into this metadata arena. The allocator and + // scheduler remain alive; the builder resets their per-graph state. + void release_for_rebuild(ggml_backend_t main_backend, + ggml_backend_t peer_backend = nullptr) { + invalidate_native_graphs(main_backend, peer_backend); + // Clear scheduler registrations while their tensor metadata is still + // valid. The builder may reset again after deciding to reuse it; that + // second reset is a no-op but keeps the builder self-contained. + if (sched) { + ggml_backend_sched_reset(sched); + } + step_graph_free(sg); + reset_nodes(); + } + void destroy(ggml_backend_t main_backend, ggml_backend_t peer_backend = nullptr) { // Native graph executables outlive ggml graph metadata in the backend @@ -4779,6 +5143,8 @@ struct DeepSeek4FusedDecodeGraph { ggml_backend_sched_free(sched); sched = nullptr; } + sched_capacity = 0; + sched_backends = {}; step_graph_destroy(sg); reset_nodes(); } @@ -4835,6 +5201,30 @@ struct Ds4FusedVerifyCache { std::array slots; struct Extra { + struct PagedLane { + ggml_tensor * pos = nullptr; + ggml_tensor * neg_pos = nullptr; + ggml_tensor * raw_gather = nullptr; + ggml_tensor * comp_gather = nullptr; + ggml_tensor * index_gather = nullptr; + ggml_tensor * raw_write = nullptr; + ggml_tensor * comp_write = nullptr; + ggml_tensor * comp_read = nullptr; + ggml_tensor * ape = nullptr; + ggml_tensor * state_row = nullptr; + ggml_tensor * comp_pos = nullptr; + // Element offsets into the shared per-dtype upload bundles. + int64_t i32_base = -1; // pos, neg_pos, comp_read, ape, comp_pos + int64_t i64_base = -1; // raw_write, comp_write, state_row + int64_t raw_off = -1; + int64_t raw_n = 0; + int64_t comp_off = -1; + int64_t comp_n = 0; + int64_t index_off = -1; + int64_t index_n = 0; + int64_t mask_off = -1; + int64_t mask_n = 0; + }; ggml_tensor * pos_q = nullptr; // i32 [q] ggml_tensor * neg_q = nullptr; // i32 [q] ggml_tensor * rawrows = nullptr; // i64 [1,q] @@ -4847,6 +5237,16 @@ struct Ds4FusedVerifyCache { // Reused host staging for the context-sized additive attention mask. // Keeping it per slot removes one allocation from every verify step. std::vector mask_values; + std::vector bundle_i32; + std::vector bundle_i64; + std::vector bundle_gather; + std::vector paged; // [layer*q], paged mode only + ggml_tensor * paged_i32 = nullptr; + ggml_tensor * paged_i64 = nullptr; + ggml_tensor * paged_gather = nullptr; + int64_t paged_i32_n = 0; + int64_t paged_i64_n = 0; + int64_t paged_gather_n = 0; int q = 0; void reset() { *this = Extra{}; } @@ -5017,12 +5417,12 @@ static ggml_tensor * ds4_build_fused_hc_pre( ggml_tensor * fn, ggml_tensor * base, const HcWeightsCpu & cw, - ggml_tensor ** out_split) { + ggml_tensor ** out_split, int projection_columns = 0) { if (!fn || !base || !cw.loaded || cw.scale_data.size() < 3) return nullptr; const int mix_dim = 2 * w.n_hc + w.n_hc * w.n_hc; const int64_t n_tokens = hc_flat->ne[1]; ggml_tensor * normed = ggml_rms_norm(ctx, hc_flat, w.hc_eps); - ggml_tensor * mix = ggml_mul_mat(ctx, fn, normed); + ggml_tensor * mix = ds4_mul_mat_columns(ctx, fn, normed, projection_columns); mix = n_tokens == 1 ? ggml_reshape_1d(ctx, mix, mix_dim) : ggml_reshape_2d(ctx, mix, mix_dim, n_tokens); @@ -6879,6 +7279,385 @@ static bool initialize_layer_range_cache( runtime.owns_output = owns_output; return true; } + +struct Ds4PagedGatheredRuntime { + DeepSeek4LayerRangeCache model; + const MoeHybridStorage * hybrid_identity = nullptr; + ggml_backend_t hybrid_cpu_backend = nullptr; + ggml_backend_t hybrid_cold_backend = nullptr; +}; + +void deepseek4_release_paged_gathered_runtime(DeepSeek4PagedCache & cache) { + delete static_cast(cache.gathered_runtime); + cache.gathered_runtime = nullptr; +} + +bool deepseek4_paged_gathered_step( + ggml_backend_t backend, int device, const DeepSeek4Weights & w, + DeepSeek4PagedCache & cache, const float * embeddings, + const int32_t * token_ids, const int64_t * positions, + const int32_t * slots, uint32_t lanes, const int32_t * block_tables, + uint32_t block_table_stride, bool bucket_history, + const uint8_t * logit_lanes, + std::vector & out_logits, std::vector & out_argmax, + MoeHybridStorage * hybrid, + MoeHybridRoutingStats * routing_stats, + DeepSeek4StepTelemetry * telemetry) { + const auto step_t0 = Ds4TimingClock::now(); + if (!backend || !embeddings || !positions || !slots || !block_tables || + !logit_lanes || + lanes < 1 || lanes > (uint32_t) DEEPSEEK4_MAX_GATHERED_ROWS || + cache.layers.size() != (size_t) w.n_layer || + block_table_stride < cache.plan.max_blocks_per_sequence) return false; + for (uint32_t lane = 0; lane < lanes; ++lane) { + if (slots[lane] < 0) continue; + if ((uint32_t) slots[lane] >= cache.plan.slots || positions[lane] < 0 || + (uint64_t) positions[lane] >= cache.plan.max_ctx || + positions[lane] > INT32_MAX) return false; + int64_t previous_position = -1; + for (uint32_t prior = 0; prior < lane; ++prior) { + if (slots[prior] != slots[lane]) continue; + previous_position = positions[prior]; + for (uint32_t logical = 0; logical < block_table_stride; ++logical) { + if (block_tables[(size_t)prior * block_table_stride + logical] != + block_tables[(size_t)lane * block_table_stride + logical]) return false; + } + } + if (previous_position >= 0 && (hybrid || bucket_history || + positions[lane] != previous_position + 1)) return false; + } + // Active logical pages must have valid, exclusive physical ownership. + // Aliasing would make one lane's compressor write mutate another lane's + // chronological history and is therefore malformed addressing. + std::vector physical_owner(cache.plan.physical_blocks, -1); + std::vector physical_logical(cache.plan.physical_blocks, -1); + for (uint32_t lane = 0; lane < lanes; ++lane) { + if (slots[lane] < 0) continue; + const uint64_t last_block = (uint64_t) positions[lane] / DS4_PAGE_TOKENS; + if (last_block >= block_table_stride) return false; + for (uint64_t logical = 0; logical <= last_block; ++logical) { + const int32_t physical = block_tables[(size_t) lane * block_table_stride + logical]; + if (physical < 0 || (uint32_t) physical >= cache.plan.physical_blocks || + (physical_owner[(size_t)physical] >= 0 && + (physical_owner[(size_t)physical] != slots[lane] || + physical_logical[(size_t)physical] != (int64_t)logical))) return false; + physical_owner[(size_t)physical] = slots[lane]; + physical_logical[(size_t)physical] = (int64_t)logical; + } + } + if (hybrid) { + for (size_t il = 0; il < hybrid->layers.size(); ++il) { + if (hybrid->layers[il].cache_slots > 0) { + std::fprintf(stderr, + "[deepseek4-paged] layer %zu uses mutable expert-cache " + "placement, which gathered serving cannot capture\n", il); + return false; + } + } + } + const auto build_t0 = Ds4TimingClock::now(); + const DeepSeek4RoctxRange roctx_range( + "ds4.paged_gathered_step", + {InferencePhase::Batched, static_cast(lanes), 0, w.n_layer, + device}); + auto * rt = static_cast(cache.gathered_runtime); + if (!rt) { + rt = new (std::nothrow) Ds4PagedGatheredRuntime; + if (!rt) return false; + cache.gathered_runtime = rt; + } + if (rt->hybrid_identity != hybrid || + rt->hybrid_cpu_backend != (hybrid ? hybrid->cpu_backend : nullptr) || + rt->hybrid_cold_backend != (hybrid ? hybrid->cold_backend : nullptr)) { + rt->model.fused_verify_graph_cache.destroy(); + rt->hybrid_identity = hybrid; + rt->hybrid_cpu_backend = hybrid ? hybrid->cpu_backend : nullptr; + rt->hybrid_cold_backend = hybrid ? hybrid->cold_backend : nullptr; + } + if (!rt->model.matches(w, backend, device, 0, w.n_layer, true) && + !initialize_layer_range_cache(rt->model, backend, device, w, + 0, w.n_layer, true)) { + std::fprintf(stderr, + "[deepseek4-paged] failed to initialize whole-model graph cache\n"); + return false; + } + + std::vector> prepared((size_t) w.n_layer); + std::vector key = {0x5041474544LL, (int64_t) lanes, + token_ids ? 1 : 0, hybrid ? 1 : 0, + bucket_history ? 1 : 0}; + for (uint32_t lane = 0; lane < lanes; ++lane) key.push_back(slots[lane]); + // The gathered lane rows depend on the layer only through its compress + // ratio (0, 4, or 128), so prepare each ratio once per round and copy; + // the per-layer copies are still padded independently below. + std::vector rows_by_ratio[3]; + bool rows_ready[3] = {false, false, false}; + for (int il = 0; il < w.n_layer; ++il) { + const uint32_t ratio = cache.layers[(size_t) il].ratio; + const int ri = ratio == 0 ? 0 : ratio == 4 ? 1 : 2; + if (!rows_ready[ri]) { + if (!prepare_deepseek4_gathered_lane_rows( + slots, positions, lanes, block_tables, block_table_stride, + cache.plan.physical_blocks, ratio, rows_by_ratio[ri])) return false; + rows_ready[ri] = true; + } + prepared[(size_t) il] = rows_by_ratio[ri]; + for (auto & row : prepared[(size_t) il]) { + if (bucket_history) { + row.raw_history.resize( + (size_t) ds4_padded_gathered_raw_rows( + (int) row.raw_history_valid), + 0); + } + if (bucket_history && ratio > 0) { + row.compressed_history.resize( + (size_t) ds4_padded_comp_rows( + (int) row.compressed_history_valid, + (int) cache.layers[(size_t) il].physical_rows), + 0); + } + key.push_back((int64_t) row.raw_history.size()); + key.push_back((int64_t) row.compressed_history.size()); + // The gathered graph's topology depends on the lane's position only + // through whether this round emits a compressed row: the state + // row, the APE row, the compressed write row, the ring scatter row, + // and the gather lists are all device inputs uploaded every round. + // Keying on the emit flag (as the bucketed path already does) lets + // the exact path reuse its graph on rounds whose history sizes + // repeat, which is every non-emit round once the raw window is + // full, instead of rebuilding every round. + key.push_back(row.slot < 0 ? -1 + : (ratio > 0 && row.compressed_emitted ? 1 : 0)); + } + } + + auto & vc = rt->model.fused_verify_graph_cache; + auto & mc = rt->model.fused_decode_graph_cache; + if (vc.owner_ctx != w.ctx || vc.backend != backend || + vc.peer_backend != (hybrid ? hybrid->cold_backend : nullptr)) { + vc.destroy(); vc.owner_ctx = w.ctx; vc.backend = backend; + vc.peer_backend = hybrid ? hybrid->cold_backend : nullptr; + } + if (mc.owner_ctx != w.ctx || mc.backend != backend) { + mc.destroy(); mc.owner_ctx = w.ctx; mc.backend = backend; + } + if (!ds4_fused_ensure_fn_mirrors(mc, backend, w, + rt->model.hc_layer_weights, rt->model.hc_output_weights)) return false; + vc.counter++; + DeepSeek4FusedDecodeGraph * fg = nullptr; + Ds4FusedVerifyCache::Extra * ex = nullptr; + const size_t slot_limit = hybrid ? ds4_fused_verify_hybrid_slot_limit() + : vc.slots.size(); + for (size_t i = 0; i < slot_limit; ++i) { + if (vc.slots[i].built() && vc.slots[i].shape_key == key) { + fg = &vc.slots[i]; ex = &vc.extra[i]; break; + } + } + if (!fg) { + size_t pick = 0; + for (size_t i = 0; i < slot_limit; ++i) { + if (!vc.slots[i].built()) { pick = i; break; } + if (vc.slots[i].last_use < vc.slots[pick].last_use) pick = i; + } + fg = &vc.slots[pick]; ex = &vc.extra[pick]; + fg->release_for_rebuild(vc.backend, vc.peer_backend); + ex->reset(); + if (!ds4_build_fused_verify_graph( + mc, *fg, *ex, backend, w, nullptr, + rt->model.hc_layer_weights, rt->model.hc_output_weights, + rt->model.hash_routing_tables, 0, (int) lanes, + token_ids != nullptr, {}, hybrid, std::move(key), + &cache, &prepared, bucket_history)) { + std::fprintf(stderr, + "[deepseek4-paged] failed to build gathered graph " + "(lanes=%u)\n", lanes); + fg->destroy(vc.backend, vc.peer_backend); + ex->reset(); + return false; + } + } + if (telemetry) telemetry->full_graph_build_us += ds4_elapsed_us(build_t0, Ds4TimingClock::now()); + const auto set_t0 = Ds4TimingClock::now(); + fg->last_use = vc.counter; + ds4_fv_set(fg->inp_embed, embeddings, + sizeof(float) * (size_t) w.n_embd * lanes); + // The shared Q/KV prologue rotates all gathered lanes at once. Padding + // lanes use position zero, matching their passive prepared row record. + { + std::vector pos_batch(lanes, 0); + std::vector neg_batch(lanes, 0); + for (uint32_t lane = 0; lane < lanes; ++lane) { + if (slots[lane] < 0) continue; + pos_batch[lane] = (int32_t) positions[lane]; + neg_batch[lane] = -(int32_t) positions[lane]; + } + ds4_fv_set(ex->pos_q, pos_batch.data(), sizeof(int32_t) * lanes); + ds4_fv_set(ex->neg_q, neg_batch.data(), sizeof(int32_t) * lanes); + } + + auto & bundle_i32 = ex->bundle_i32; + bundle_i32.resize((size_t) std::max(ex->paged_i32_n, 0), 0); + auto & bundle_i64 = ex->bundle_i64; + bundle_i64.resize((size_t) std::max(ex->paged_i64_n, 0), 0); + auto & bundle_gather = ex->bundle_gather; + bundle_gather.resize((size_t) std::max(ex->paged_gather_n, 0), 0); + std::vector & mask_values = ex->mask_values; + if (bucket_history) { + const size_t mask_count = (size_t) ggml_nelements(fg->mask_bundle); + if (mask_values.size() != mask_count) { + mask_values.resize(mask_count); + } + std::fill(mask_values.begin(), mask_values.end(), 0.0f); + } + size_t pi = 0; + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) cache.layers[(size_t) il].ratio; + for (uint32_t lane = 0; lane < lanes; ++lane, ++pi) { + const auto & row = prepared[(size_t) il][lane]; + const auto & px = ex->paged[pi]; + if (px.i32_base < 0 || px.i64_base < 0) return false; + const int32_t pos = (int32_t) row.position; + bundle_i32[(size_t) px.i32_base + 0] = pos; + bundle_i32[(size_t) px.i32_base + 1] = -pos; + if (px.raw_off < 0 || + px.raw_off + px.raw_n > ex->paged_gather_n || + (int64_t) row.raw_history.size() > px.raw_n) return false; + for (size_t i = 0; i < row.raw_history.size(); ++i) { + bundle_gather[(size_t) px.raw_off + i] = + (int32_t) row.raw_history[i]; + } + for (int i = 0; i < 4; ++i) bundle_i64[(size_t) px.i64_base + 3 + i] = i; + bundle_i64[(size_t) (7 * w.n_layer * lanes + il * lanes + lane)] = + px.raw_off - ex->paged[(size_t) il * lanes].raw_off + row.raw_history.size(); + bundle_i64[(size_t) px.i64_base + 0] = + std::max(row.raw_scatter, 0); + if (ratio > 0 && px.comp_off >= 0) { + if (px.comp_off + px.comp_n > ex->paged_gather_n || + (int64_t) row.compressed_history.size() > px.comp_n) { + return false; + } + for (size_t i = 0; i < row.compressed_history.size(); ++i) { + const int32_t value = + (int32_t) row.compressed_history[i]; + bundle_gather[(size_t) px.comp_off + i] = value; + if (px.index_off >= 0) { + bundle_gather[(size_t) px.index_off + i] = value; + } + } + const int64_t cw = std::max(row.compressed_scatter, 0); + const int32_t ape = pos % ratio; + bundle_i64[(size_t) px.i64_base + 1] = cw; + bundle_i64[(size_t) px.i64_base + 2] = + ratio == 4 ? 4 + ape : ape; + bundle_i32[(size_t) px.i32_base + 2] = (int32_t) cw; + bundle_i32[(size_t) px.i32_base + 3] = ape; + bundle_i32[(size_t) px.i32_base + 4] = pos + 1 - ratio; + } + if (bucket_history) { + if (px.mask_off < 0 || px.mask_n < 1 || + px.mask_off + px.mask_n > (int64_t) mask_values.size()) { + return false; + } + for (size_t i = row.raw_history_valid; + i < row.raw_history.size(); ++i) { + mask_values[(size_t) px.mask_off + i] = -1.0e30f; + } + const size_t comp_base = (size_t) px.mask_off + + row.raw_history.size() + 1; + for (size_t i = row.compressed_history_valid; + i < row.compressed_history.size(); ++i) { + mask_values[comp_base + i] = -1.0e30f; + } + } + } + } + ds4_fv_set(ex->paged_i32, bundle_i32.data(), + bundle_i32.size() * sizeof(int32_t)); + ds4_fv_set(ex->paged_i64, bundle_i64.data(), + bundle_i64.size() * sizeof(int64_t)); + ds4_fv_set(ex->paged_gather, bundle_gather.data(), + bundle_gather.size() * sizeof(int32_t)); + if (bucket_history) { + ds4_fv_set(fg->mask_bundle, mask_values.data(), + mask_values.size() * sizeof(float)); + } + if (token_ids) { + for (int il = 0; il < w.n_layer; ++il) { + ggml_tensor * ids = fg->hash_ids[(size_t) il]; if (!ids) continue; + std::vector values((size_t) ids->ne[0] * lanes); + for (uint32_t lane = 0; lane < lanes; ++lane) { + const int32_t * src = hash_routing_row(rt->model.hash_routing_tables[(size_t) il], + slots[lane] < 0 ? 0 : token_ids[lane], + w.n_expert_used); + if (!src) return false; + std::memcpy(values.data() + lane * ids->ne[0], src, + (size_t) ids->ne[0] * sizeof(int32_t)); + } + ds4_fv_set(ids, values.data(), values.size() * sizeof(int32_t)); + } + } + // Monolithic exact serving reuses a gathered graph for only a few rounds + // (history sizes change at every compressed-row emit), so a HIP graph + // capture of this node count (about 12 ms on gfx1151) never amortizes: + // replay saves only about 0.2 us per launch. Keep the ggml graph reuse + // and execute eagerly; this also skips the per-round node-property scan. + // The heterogeneous path keeps its long-lived bucketed graphs and replay. + if (telemetry) telemetry->full_graph_set_us += ds4_elapsed_us(set_t0, Ds4TimingClock::now()); + const auto compute_t0 = Ds4TimingClock::now(); + ScopedCudaGraphOverrides monolithic_eager_scope(!hybrid); + const enum ggml_status status = fg->sched + ? ggml_backend_sched_graph_compute(fg->sched, fg->sg.gf) + : ggml_backend_graph_compute(backend, fg->sg.gf); + if (status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[deepseek4-paged] gathered graph compute failed: status=%d\n", + (int) status); + return false; + } + if (telemetry) telemetry->full_graph_compute_us += ds4_elapsed_us(compute_t0, Ds4TimingClock::now()); + const auto read_t0 = Ds4TimingClock::now(); + ds4_fused_consume_route_diagnostics(*fg, hybrid, routing_stats, slots); + out_argmax.resize(lanes); + ggml_backend_tensor_get(ex->argmax, out_argmax.data(), 0, + out_argmax.size() * sizeof(int32_t)); + + const uint32_t requested_logits = static_cast(std::count_if( + logit_lanes, logit_lanes + lanes, + [](uint8_t requested) { return requested != 0; })); + if (requested_logits == 0) { + out_logits.clear(); + } else { + out_logits.assign((size_t) w.n_vocab * lanes, 0.0f); + const size_t row_bytes = (size_t) w.n_vocab * sizeof(float); + if (requested_logits == lanes) { + ggml_backend_tensor_get( + fg->logits, out_logits.data(), 0, + out_logits.size() * sizeof(float)); + } else { + for (uint32_t lane = 0; lane < lanes; ++lane) { + if (!logit_lanes[lane]) continue; + ggml_backend_tensor_get( + fg->logits, + out_logits.data() + (size_t) lane * w.n_vocab, + (size_t) lane * fg->logits->nb[1], row_bytes); + } + } + } + for (uint32_t lane = 0; lane < lanes; ++lane) { + if (slots[lane] >= 0) continue; + if (!out_logits.empty()) { + std::fill_n(out_logits.data() + (size_t) lane * w.n_vocab, + w.n_vocab, 0.0f); + } + out_argmax[lane] = -1; + } + if (telemetry) { + telemetry->full_graph_read_us += ds4_elapsed_us(read_t0, Ds4TimingClock::now()); + telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + } + return true; +} + bool deepseek4_step_layer_range( ggml_backend_t backend, int device, diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 66d80e417..79b9105b1 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -15,6 +15,7 @@ #pragma once #include +#include #include #include @@ -24,7 +25,10 @@ #include "internal.h" #include "common/layer_split_utils.h" +#include "common/paged_attention_config.h" #include "common/prefill_attention_mode.h" +#include "common/concurrency/paged_kv_pool.h" +#include "deepseek4_paged_cache.h" namespace dflash::common { @@ -296,6 +300,22 @@ struct DeepSeek4Cache { ggml_backend_buffer_t buf = nullptr; }; +struct DeepSeek4PagedLayerCache : DeepSeek4LayerCache { + uint32_t ratio = 0; + uint64_t physical_rows = 0; +}; + +struct DeepSeek4PagedCache { + std::unique_ptr pool; + DeepSeek4PagedCachePlan plan; + std::vector layers; + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buf = nullptr; + // Dedicated bounded gathered-reference graph cache (opaque here because + // its implementation shares the fused verifier's private machinery). + void * gathered_runtime = nullptr; +}; + struct DeepSeek4Snapshot; struct DeepSeek4RawRingSpan { @@ -303,6 +323,13 @@ struct DeepSeek4RawRingSpan { int count = 0; }; +struct DeepSeek4Head4Tail2Routes { + ggml_tensor * head_ids = nullptr; + ggml_tensor * head_weights = nullptr; + ggml_tensor * tail_ids = nullptr; + ggml_tensor * tail_weights = nullptr; +}; + // ─── Configuration ────────────────────────────────────────────────────── struct DeepSeek4BackendConfig { @@ -315,6 +342,9 @@ struct DeepSeek4BackendConfig { int expert_top_k = 0; // 0 = use all model-routed experts bool fused_decode = false; // single-graph GPU decode bool fused_verify_f16_kv = false; // F16 KV in batched verifier attention + bool paged_attention = false; + int max_concurrency = 1; + long long kv_pool_tokens = 0; }; // ─── Function declarations ────────────────────────────────────────────── @@ -340,6 +370,32 @@ bool create_deepseek4_cache(ggml_backend_t backend, DeepSeek4Cache & out); void free_deepseek4_cache(DeepSeek4Cache & c); +bool create_deepseek4_paged_cache(ggml_backend_t backend, + const DeepSeek4Weights & w, + uint32_t slots, uint32_t max_ctx, + uint32_t physical_blocks, + DeepSeek4PagedCache & out); +void reset_deepseek4_paged_slot(DeepSeek4PagedCache & c, uint32_t slot); +void free_deepseek4_paged_cache(DeepSeek4PagedCache & c); +// Exact gathered-reference decode for up to six independent lanes. Inputs are +// lane-major; negative slots are inactive padding lanes. `logit_lanes` marks +// the lanes that need full host logits. `out_logits` is empty when none do; +// otherwise it is [n_vocab, lanes], with unrequested rows zeroed. +bool deepseek4_paged_gathered_step( + ggml_backend_t backend, int device, const DeepSeek4Weights & w, + DeepSeek4PagedCache & cache, const float * embeddings, + const int32_t * token_ids, const int64_t * positions, + const int32_t * slots, uint32_t lanes, const int32_t * block_tables, + uint32_t block_table_stride, bool bucket_history, + const uint8_t * logit_lanes, + std::vector & out_logits, std::vector & out_argmax, + MoeHybridStorage * moe_hybrid = nullptr, + MoeHybridRoutingStats * routing_stats = nullptr, + DeepSeek4StepTelemetry * telemetry = nullptr); +void log_deepseek4_step_telemetry( + const char * phase, int tokens, int steps, double wall_s, + const DeepSeek4StepTelemetry & telemetry); +void deepseek4_release_paged_gathered_runtime(DeepSeek4PagedCache & cache); void reset_deepseek4_cache(DeepSeek4Cache & c); // Release only reproducible large-batch graph arenas after prefill. KV/model // state and the DSpark feature tail remain live for the following decode. @@ -349,6 +405,12 @@ int deepseek4_previous_raw_ring_spans( int kv_start, int n_swa, DeepSeek4RawRingSpan spans[2]); +bool build_deepseek4_head4_tail2_routes( + ggml_context * ctx, + ggml_tensor * selected, + ggml_tensor * router_weights, + int n_tokens, + DeepSeek4Head4Tail2Routes & out); bool deepseek4_snapshot_save(const DeepSeek4Cache & cache, ggml_backend_t snapshot_backend, DeepSeek4Snapshot & out); diff --git a/server/src/deepseek4/deepseek4_page_layout.h b/server/src/deepseek4/deepseek4_page_layout.h new file mode 100644 index 000000000..b2ee71631 --- /dev/null +++ b/server/src/deepseek4/deepseek4_page_layout.h @@ -0,0 +1,55 @@ +// Host-side address geometry for DeepSeek V4's paged raw and compressed KV. +#pragma once + +#include +#include + +namespace dflash::common { + +inline constexpr uint32_t DS4_PAGE_TOKENS = 128; + +// Raw KV remains slot-indexed: every sequence reuses this 128-row ring. +inline constexpr uint32_t ds4_raw_ring_row(uint64_t logical_token) { + return static_cast(logical_token % DS4_PAGE_TOKENS); +} + +// Number of physically paged compressed rows. Rejects unsupported ratios and +// arithmetic that cannot be represented by the row-index type. +inline bool ds4_compressed_page_capacity(uint64_t physical_blocks, + uint32_t ratio, + uint64_t & rows) { + if (ratio != 4 && ratio != 128) return false; + const uint64_t rows_per_block = DS4_PAGE_TOKENS / ratio; + if (physical_blocks > + std::numeric_limits::max() / rows_per_block) { + return false; + } + rows = physical_blocks * rows_per_block; + return true; +} + +// Computes the destination for a completed compression group. `emitted` is +// false between group boundaries and `row` is left unchanged in that case. +// Physical block IDs need not be contiguous. +inline bool ds4_compressed_page_row(uint64_t logical_token, + uint64_t physical_block, + uint32_t ratio, + uint64_t & row, + bool & emitted) { + if (ratio != 4 && ratio != 128) return false; + emitted = logical_token % ratio == ratio - 1; + if (!emitted) return true; + + const uint64_t rows_per_block = DS4_PAGE_TOKENS / ratio; + if (physical_block > + std::numeric_limits::max() / rows_per_block) { + return false; + } + const uint64_t base = physical_block * rows_per_block; + const uint64_t offset = (logical_token % DS4_PAGE_TOKENS) / ratio; + if (base > std::numeric_limits::max() - offset) return false; + row = base + offset; + return true; +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_paged_cache.cpp b/server/src/deepseek4/deepseek4_paged_cache.cpp new file mode 100644 index 000000000..703402e39 --- /dev/null +++ b/server/src/deepseek4/deepseek4_paged_cache.cpp @@ -0,0 +1,232 @@ +#include "deepseek4_paged_cache.h" + +#include "deepseek4_page_layout.h" + +#ifndef DFLASH_DS4_PLAN_ONLY +#include "deepseek4_internal.h" +#endif + +#include +#include +#include + +namespace dflash::common { +namespace { +bool add_mul(uint64_t & dst, uint64_t a, uint64_t b) { + if (a && b > std::numeric_limits::max() / a) return false; + const uint64_t v = a * b; + if (dst > std::numeric_limits::max() - v) return false; + dst += v; + return true; +} +} + +bool plan_deepseek4_paged_pool_blocks(uint32_t max_ctx, uint32_t slots, + uint64_t requested_tokens, + uint32_t & physical_blocks) { + physical_blocks = 0; + if (!max_ctx || !slots) return false; + const uint64_t blocks = requested_tokens + ? 1 + (requested_tokens - 1) / DS4_PAGE_TOKENS + : (1 + (uint64_t(max_ctx) - 1) / DS4_PAGE_TOKENS) * slots; + if (!blocks) return false; + if (blocks > UINT32_MAX / DS4_PAGE_TOKENS) return false; + physical_blocks = static_cast(blocks); + return true; +} + +bool prepare_deepseek4_gathered_lane_rows( + const int32_t * slots, const int64_t * positions, uint32_t lanes, + const int32_t * block_tables, uint32_t block_table_stride, + uint32_t physical_blocks, uint32_t ratio, + std::vector & out) { + if (!slots || !positions || !block_tables || !block_table_stride || + !physical_blocks || (ratio != 0 && ratio != 4 && ratio != 128)) { + return false; + } + std::vector prepared(lanes); + for (uint32_t lane = 0; lane < lanes; ++lane) { + auto & rows = prepared[lane]; + rows.slot = slots[lane]; + if (rows.slot < 0) continue; // Padding must remain entirely passive. + // Padding lanes do not have a validated position. Keep the default + // zero so inverse RoPE and compressor bookkeeping cannot inherit an + // arbitrary positions[] value once callers use constant-width steps. + rows.position = positions[lane]; + if (rows.position < 0) return false; + const uint64_t pos = static_cast(rows.position); + // The current row is appended in-graph, so retain at most the 127 + // preceding rows that can coexist with it in the 128-row SWA window. + const uint64_t first_raw = pos >= DS4_PAGE_TOKENS + ? pos - DS4_PAGE_TOKENS + 1 : 0; + rows.raw_history.reserve(static_cast(pos - first_raw)); + for (uint64_t p = first_raw; p < pos; ++p) { + rows.raw_history.push_back( + int64_t(rows.slot) * DS4_PAGE_TOKENS + ds4_raw_ring_row(p)); + } + rows.raw_history_valid = + static_cast(rows.raw_history.size()); + rows.raw_scatter = int64_t(rows.slot) * DS4_PAGE_TOKENS + + ds4_raw_ring_row(pos); + // Validate all active pages even when this layer has no compressor. + // Bound the position before reserving compressed history. + const uint64_t current_logical_block = pos / DS4_PAGE_TOKENS; + if (current_logical_block >= block_table_stride) return false; + const int32_t * lane_table = block_tables + size_t(lane) * block_table_stride; + for (uint64_t block = 0; block <= current_logical_block; ++block) { + if (lane_table[block] < 0 || uint32_t(lane_table[block]) >= physical_blocks) { + return false; + } + } + if (!ratio) continue; + const int32_t current_physical = lane_table[current_logical_block]; + + // Every completed group before the current token contributes one + // chronological row. Looking up each logical page (rather than + // assuming contiguous physical pages) is the reference behaviour. + const uint64_t completed = pos / ratio; + rows.compressed_history.reserve(static_cast(completed)); + for (uint64_t group = 0; group < completed; ++group) { + const uint64_t end_token = group * ratio + ratio - 1; + const uint64_t logical_block = end_token / DS4_PAGE_TOKENS; + const int32_t physical = lane_table[logical_block]; + uint64_t row = 0; bool emitted = false; + if (!ds4_compressed_page_row(end_token, uint32_t(physical), ratio, + row, emitted) || !emitted || + row > uint64_t(INT64_MAX)) return false; + rows.compressed_history.push_back(static_cast(row)); + } + rows.compressed_history_valid = + static_cast(rows.compressed_history.size()); + uint64_t scatter = 0; + if (!ds4_compressed_page_row(pos, uint32_t(current_physical), ratio, scatter, + rows.compressed_emitted) || + scatter > uint64_t(INT64_MAX)) return false; + if (rows.compressed_emitted) rows.compressed_scatter = int64_t(scatter); + } + out = std::move(prepared); + return true; +} + +bool plan_deepseek4_paged_cache(uint32_t head_dim, uint32_t indexer_head_dim, + uint32_t slots, uint32_t max_ctx, + uint32_t physical_blocks, + const std::vector & ratios, + DeepSeek4PagedCachePlan & out) { + DeepSeek4PagedCachePlan p; + if (!head_dim || !indexer_head_dim || !slots || !max_ctx || + !physical_blocks || ratios.empty() || + physical_blocks > UINT32_MAX / DS4_PAGE_TOKENS) return false; + p.slots = slots; p.max_ctx = max_ctx; p.physical_blocks = physical_blocks; + p.max_blocks_per_sequence = 1 + (max_ctx - 1) / DS4_PAGE_TOKENS; + p.ratios = ratios; + p.physical_rows.resize(ratios.size()); + for (size_t i = 0; i < ratios.size(); ++i) { + const uint32_t r = ratios[i]; + if (r != 0 && r != 4 && r != 128) return false; + if (!add_mul(p.raw_bytes, uint64_t(head_dim) * DS4_PAGE_TOKENS * 2, slots)) return false; + if (!r) continue; + uint64_t rows = 0; + if (!ds4_compressed_page_capacity(physical_blocks, r, rows)) return false; + p.physical_rows[i] = rows; + if (!add_mul(p.compressed_bytes, uint64_t(head_dim) * 2, rows)) return false; + const uint64_t width = uint64_t(head_dim) * (r == 4 ? 2 : 1); + const uint64_t state_rows = r == 4 ? 8 : 128; + if (!add_mul(p.state_bytes, width * state_rows * 8, slots)) return false; // KV + score F32 + if (r == 4) { + if (!add_mul(p.compressed_bytes, uint64_t(indexer_head_dim) * 2, rows)) return false; + if (!add_mul(p.state_bytes, uint64_t(indexer_head_dim) * 2 * 8 * 8, slots)) return false; + } + } + p.total_persistent_bytes = p.raw_bytes; + if (p.total_persistent_bytes > UINT64_MAX - p.compressed_bytes) return false; + p.total_persistent_bytes += p.compressed_bytes; + if (p.total_persistent_bytes > UINT64_MAX - p.state_bytes) return false; + p.total_persistent_bytes += p.state_bytes; + out = std::move(p); + return true; +} + +#ifndef DFLASH_DS4_PLAN_ONLY +bool create_deepseek4_paged_cache(ggml_backend_t backend, + const DeepSeek4Weights & w, uint32_t slots, + uint32_t max_ctx, uint32_t physical_blocks, + DeepSeek4PagedCache & out) { + free_deepseek4_paged_cache(out); + DeepSeek4PagedCachePlan plan; + if (!backend || w.n_layer <= 0 || w.compress_ratios.size() != size_t(w.n_layer) || + !plan_deepseek4_paged_cache(w.head_dim, w.n_indexer_head_dim, slots, + max_ctx, physical_blocks, w.compress_ratios, plan)) return false; + try { out.pool = std::make_unique(physical_blocks, slots, DS4_PAGE_TOKENS); } + catch (...) { free_deepseek4_paged_cache(out); return false; } + out.plan = plan; + out.layers.resize(w.n_layer); + ggml_init_params ip{ggml_tensor_overhead() * size_t(w.n_layer * 9 + 5) + 4096, nullptr, true}; + out.ctx = ggml_init(ip); + if (!out.ctx) { free_deepseek4_paged_cache(out); return false; } + for (int il = 0; il < w.n_layer; ++il) { + auto & l = out.layers[il]; const uint32_t r = plan.ratios[il]; + l.ratio = r; l.physical_rows = plan.physical_rows[il]; + l.raw_kv = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F16, w.head_dim, DS4_PAGE_TOKENS, slots); + if (!r) continue; + l.comp_kv = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F16, w.head_dim, l.physical_rows); + const int64_t width = int64_t(w.head_dim) * (r == 4 ? 2 : 1), sr = r == 4 ? 8 : 128; + l.attn_compressor.state_kv = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F32, width, sr, slots); + l.attn_compressor.state_score = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F32, width, sr, slots); + if (r == 4) { + l.index_comp_kv = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F16, w.n_indexer_head_dim, l.physical_rows); + const int64_t iw = int64_t(w.n_indexer_head_dim) * 2; + l.indexer_compressor.state_kv = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F32, iw, 8, slots); + l.indexer_compressor.state_score = ggml_new_tensor_3d(out.ctx, GGML_TYPE_F32, iw, 8, slots); + } + } + uint64_t raw_bytes = 0, compressed_bytes = 0, state_bytes = 0; + for (const auto & layer : out.layers) { + raw_bytes += ggml_nbytes(layer.raw_kv); + for (const auto * tensor : {layer.comp_kv, layer.index_comp_kv}) { + if (tensor) compressed_bytes += ggml_nbytes(tensor); + } + for (const auto * tensor : {layer.attn_compressor.state_kv, + layer.attn_compressor.state_score, + layer.indexer_compressor.state_kv, + layer.indexer_compressor.state_score}) { + if (tensor) state_bytes += ggml_nbytes(tensor); + } + } + if (raw_bytes != plan.raw_bytes || compressed_bytes != plan.compressed_bytes || + state_bytes != plan.state_bytes) { + std::fprintf(stderr, "[deepseek4] paged cache tensor sizes disagree with allocation plan\n"); + free_deepseek4_paged_cache(out); + return false; + } + out.buf = ggml_backend_alloc_ctx_tensors(out.ctx, backend); + if (!out.buf) { free_deepseek4_paged_cache(out); return false; } + ggml_backend_buffer_clear(out.buf, 0); + return true; +} + +void reset_deepseek4_paged_slot(DeepSeek4PagedCache & c, uint32_t slot) { + if (!c.buf || slot >= c.plan.slots) return; + auto clear_slot = [slot](ggml_tensor * tensor) { + if (!tensor || tensor->ne[2] <= (int64_t) slot) return; + const size_t bytes = tensor->nb[2]; + std::vector zeros(bytes, 0); + ggml_backend_tensor_set(tensor, zeros.data(), (size_t) slot * bytes, + bytes); + }; + for (DeepSeek4PagedLayerCache & layer : c.layers) { + clear_slot(layer.attn_compressor.state_kv); + clear_slot(layer.attn_compressor.state_score); + clear_slot(layer.indexer_compressor.state_kv); + clear_slot(layer.indexer_compressor.state_score); + } +} + +void free_deepseek4_paged_cache(DeepSeek4PagedCache & c) { + deepseek4_release_paged_gathered_runtime(c); + if (c.buf) { ggml_backend_buffer_free(c.buf); c.buf = nullptr; } + if (c.ctx) { ggml_free(c.ctx); c.ctx = nullptr; } + c.pool.reset(); c.layers.clear(); c.plan = {}; +} +#endif +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_paged_cache.h b/server/src/deepseek4/deepseek4_paged_cache.h new file mode 100644 index 000000000..2d053ecb8 --- /dev/null +++ b/server/src/deepseek4/deepseek4_paged_cache.h @@ -0,0 +1,72 @@ +#pragma once + +#include +#include +#include + +namespace dflash::common { + +// Pure host-side allocation plan. Byte counts describe tensor payloads (ggml +// alignment/padding is deliberately excluded). +struct DeepSeek4PagedCachePlan { + uint32_t slots = 0; + uint32_t max_ctx = 0; + uint32_t physical_blocks = 0; + uint32_t max_blocks_per_sequence = 0; + uint64_t raw_bytes = 0; + uint64_t compressed_bytes = 0; + uint64_t state_bytes = 0; + uint64_t total_persistent_bytes = 0; + std::vector ratios; + std::vector physical_rows; +}; + +// Host metadata for the gathered-reference decode graph. Rows are expressed +// in the flattened persistent tensors: raw rows are [slot, ring-row], while +// compressed rows use the physical page geometry from deepseek4_page_layout.h. +// A negative slot denotes a padding lane and consequently has no scatter rows. +struct DeepSeek4GatheredLaneRows { + int32_t slot = -1; + int64_t position = 0; + std::vector raw_history; + std::vector compressed_history; + uint32_t raw_history_valid = 0; + uint32_t compressed_history_valid = 0; + int64_t raw_scatter = -1; + int64_t compressed_scatter = -1; + bool compressed_emitted = false; +}; + +// Convert the configured pool capacity into 128-token physical pages. A zero +// request selects max_ctx * slots; an explicit request is honored even when it +// is smaller than max_ctx. +bool plan_deepseek4_paged_pool_blocks(uint32_t max_ctx, + uint32_t slots, + uint64_t requested_tokens, + uint32_t & physical_blocks); + +// block_tables is lane-major with block_table_stride entries per lane. +// Physical block IDs may be fragmented and are validated against +// physical_blocks. History excludes the current token; compressed history is +// in chronological group order. Returns false for malformed active lanes. +bool prepare_deepseek4_gathered_lane_rows( + const int32_t * slots, + const int64_t * positions, + uint32_t lanes, + const int32_t * block_tables, + uint32_t block_table_stride, + uint32_t physical_blocks, + uint32_t ratio, + std::vector & out); + +// Ratios must contain only 0, 4, or 128. A ratio-zero layer has a raw ring +// but no compressed storage or compressor state. +bool plan_deepseek4_paged_cache(uint32_t head_dim, + uint32_t indexer_head_dim, + uint32_t slots, + uint32_t max_ctx, + uint32_t physical_blocks, + const std::vector & ratios, + DeepSeek4PagedCachePlan & out); + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_seq_engine.cpp b/server/src/deepseek4/deepseek4_seq_engine.cpp new file mode 100644 index 000000000..abd7cf2cc --- /dev/null +++ b/server/src/deepseek4/deepseek4_seq_engine.cpp @@ -0,0 +1,318 @@ +#include "deepseek4_seq_engine.h" + +#include "deepseek4_backend.h" +#include "common/sampler.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +DeepSeek4SeqEngine::DeepSeek4SeqEngine( + DeepSeek4Backend & backend, PagedKvPool & pool, int max_ctx, + uint32_t table_stride) + : b_(backend), slots_(pool, max_ctx), stride_(table_stride), + host_tables_((size_t)pool.max_sequences() * table_stride, -1) {} + +bool DeepSeek4SeqEngine::token_is_eos(int32_t token) const { + return deepseek4_is_eos_tok(token, b_.w_); +} + +StepPlanLimits DeepSeek4SeqEngine::step_plan_limits( + int decode_rows) const { + // The gathered graph accepts at most six independent sequences. A live + // decoder owns one row. Concurrent prompts advance four rows per turn; + // a lone prompt can use all sixteen rows without delaying another request. + // Each gathered lane reads earlier rows in graph order. + // Heterogeneous execution keeps one prompt token per lane. + const int available = std::max( + 0, DEEPSEEK4_MAX_PAGED_SEQUENCES - decode_rows); + if (b_.moe_hybrid_) return {available, 1, available, 1}; + const int rows = std::max(0, DEEPSEEK4_MAX_GATHERED_ROWS - decode_rows); + const int sequences = std::min(available, rows); + int prompt_rows = 4; + if (decode_rows == 0) { + int prompts = 0; + for (int slot = 0; slot < slots_.slot_count(); ++slot) { + prompts += slots_.is_prefilling(slot); + } + if (prompts == 1) prompt_rows = rows; + } + return {sequences, prompt_rows, rows, /*prefill_allocation_quantum=*/4}; +} + +SeqEngine::AdmitResult DeepSeek4SeqEngine::admit( + uint64_t request_id, const std::vector & prompt, + const SamplerCfg & sampler) { + using AdmitStatus = AdmitResult::Status; + AdmitResult result = slots_.admit( + request_id, prompt, sampler); + if (result.status != AdmitStatus::admitted) return result; + if (result.slot < 0 || result.slot >= slots_.slot_count()) { + result.status = AdmitStatus::failed; + result.error = "invalid DeepSeek4 serving slot"; + return result; + } + std::fill_n(host_tables_.data() + (size_t)result.slot * stride_, + stride_, -1); + reset_deepseek4_paged_slot(b_.paged_cache_, (uint32_t)result.slot); + return result; +} + +bool DeepSeek4SeqEngine::set_block(int slot, int logical, int32_t physical) { + if (slot < 0 || slot >= slots_.slot_count() || logical < 0 || + (uint32_t) logical >= stride_) return false; + host_tables_[(size_t)slot * stride_ + (size_t)logical] = physical; + return true; +} + +void DeepSeek4SeqEngine::fail_prefill( + int slot, std::vector & outputs, + const std::string & error) { + std::fprintf(stderr, "[deepseek4-parallel] prefill slot %d: %s\n", + slot, error.c_str()); + PrefillOutput out; + out.slot = slot; + out.status = PrefillOutput::Status::failed; + out.error = error; + outputs.push_back(std::move(out)); +} + +SeqEngine::StepResult DeepSeek4SeqEngine::step(const StepPlan & plan) { + using Clock = std::chrono::steady_clock; + const auto phase_t0 = Clock::now(); + const char * timing_env = std::getenv("DFLASH_DS4_TIMING"); + const bool timing = timing_env && *timing_env && std::strcmp(timing_env, "0") != 0; + DeepSeek4StepTelemetry telemetry; + StepResult result; + const std::vector & inputs = plan.decode; + const int n_slots = slots_.slot_count(); + + auto fail_step = [&](const std::string & error) { + result.decode.clear(); + result.prefills.clear(); + result.error = error; + return std::move(result); + }; + + if ((int)inputs.size() != slots_.decoding_count()) { + return fail_step("decode plan does not cover every live DeepSeek4 slot"); + } + std::vector decode_seen((size_t)n_slots, 0); + for (const StepInput & input : inputs) { + if (input.slot < 0 || input.slot >= n_slots || input.token < 0 || + decode_seen[(size_t)input.slot] || + !slots_.is_active(input.slot) || + slots_.is_prefilling(input.slot)) { + return fail_step("invalid or duplicate DeepSeek4 decode row"); + } + decode_seen[(size_t)input.slot] = 1; + } + + const StepPlanLimits limits = step_plan_limits((int)inputs.size()); + if ((int)plan.prefills.size() > limits.max_prefill_sequences) { + return fail_step("DeepSeek4 step exceeds the six-lane graph"); + } + std::vector prefill_seen((size_t)n_slots, 0); + int prefill_rows = 0; + for (const PrefillSlice & slice : plan.prefills) { + prefill_rows += slice.max_tokens; + if (slice.slot < 0 || slice.slot >= n_slots || + slice.max_tokens < 1 || + slice.max_tokens > limits.max_prefill_tokens_per_sequence || + prefill_rows > limits.max_prefill_tokens_total || + prefill_seen[(size_t)slice.slot] || + decode_seen[(size_t)slice.slot] || + !slots_.is_prefilling(slice.slot)) { + return fail_step("invalid or duplicate DeepSeek4 prefill row"); + } + prefill_seen[(size_t)slice.slot] = 1; + } + if (inputs.empty() && plan.prefills.empty()) return result; + + std::vector lane_tokens; + std::vector lane_positions; + std::vector lane_slots; + std::vector decode_lanes; + lane_tokens.reserve(inputs.size() + plan.prefills.size()); + lane_positions.reserve(inputs.size() + plan.prefills.size()); + lane_slots.reserve(inputs.size() + plan.prefills.size()); + decode_lanes.reserve(inputs.size()); + result.decode.reserve(inputs.size()); + result.prefills.reserve(plan.prefills.size()); + + for (const StepInput & input : inputs) { + DecodeOutput out; + out.slot = input.slot; + const SeqSlotManager::StepAppend append = + slots_.append_token(input.slot, input.token); + if (!append.ok) { + out.failed = true; + out.error = append.busy + ? "paged KV pool exhausted during DeepSeek4 decode; raise " + "--kv-pool-tokens or lower --max-ctx/--max-concurrency" + : "DeepSeek4 decode K/V append failed"; + decode_lanes.push_back(-1); + result.decode.push_back(std::move(out)); + continue; + } + if (append.new_block >= 0 && + !set_block(input.slot, append.new_block_index, + append.new_block)) { + out.failed = true; + out.error = "DeepSeek4 decode block-table update failed"; + decode_lanes.push_back(-1); + result.decode.push_back(std::move(out)); + continue; + } + decode_lanes.push_back((int)lane_tokens.size()); + lane_tokens.push_back(input.token); + lane_positions.push_back(append.position); + lane_slots.push_back(input.slot); + result.decode.push_back(std::move(out)); + } + + struct PrefillLane { + int slot = -1; + int lane = -1; + bool commit = false; + }; + std::vector prefill_lanes; + prefill_lanes.reserve(plan.prefills.size()); + for (const PrefillSlice & slice : plan.prefills) { + const SeqSlot & before = slots_.slot(slice.slot); + const int remaining = before.prompt_len - before.cur_pos; + const int n_rows = std::min(slice.max_tokens, remaining); + SeqSlotManager::PrefillChunk chunk = + slots_.append_prefill(slice.slot, n_rows); + if (n_rows < 1 || !chunk.ok || chunk.rows.size() != (size_t) n_rows) { + fail_prefill(slice.slot, result.prefills, + "DeepSeek4 prefill K/V append failed"); + continue; + } + bool blocks_ok = true; + for (size_t i = 0; i < chunk.new_blocks.size(); ++i) { + blocks_ok = set_block(slice.slot, chunk.first_new_block + (int) i, + chunk.new_blocks[i]) && blocks_ok; + } + if (!blocks_ok) { + fail_prefill(slice.slot, result.prefills, + "DeepSeek4 prefill block-table update failed"); + continue; + } + const SeqSlot & slot = slots_.slot(slice.slot); + const bool commit = slot.cur_pos == slot.prompt_len; + // One gathered lane per prompt row, in chronological order. Only the + // chunk's last row can need logits. + const int first = slot.cur_pos - n_rows; + for (int pos = first; pos < slot.cur_pos; ++pos) { + lane_tokens.push_back(slot.sample_history[(size_t) pos]); + lane_positions.push_back(pos); + lane_slots.push_back(slice.slot); + } + prefill_lanes.push_back( + {slice.slot, (int) lane_tokens.size() - 1, commit}); + } + + if (lane_tokens.empty()) return result; + if (lane_tokens.size() > (size_t) DEEPSEEK4_MAX_GATHERED_ROWS) { + return fail_step("DeepSeek4 gathered step exceeds sixteen rows"); + } + + const auto embed_t0 = Clock::now(); + std::vector embeddings( + (size_t)b_.w_.n_embd * lane_tokens.size()); + if (!b_.w_.embedder.embed(lane_tokens.data(), (int)lane_tokens.size(), + embeddings.data())) { + return fail_step("DeepSeek4 token embedding failed"); + } + + if (timing) telemetry.embed_us += std::chrono::duration_cast< + std::chrono::microseconds>(Clock::now() - embed_t0).count(); + std::vector compact_tables( + lane_tokens.size() * stride_, -1); + for (size_t lane = 0; lane < lane_slots.size(); ++lane) { + std::copy_n( + host_tables_.data() + (size_t)lane_slots[lane] * stride_, + stride_, compact_tables.data() + lane * stride_); + } + + std::vector logit_lanes(lane_tokens.size(), 0); + for (size_t i = 0; i < inputs.size(); ++i) { + const int lane = decode_lanes[i]; + if (lane >= 0 && + slots_.slot(inputs[i].slot).sampler.needs_logit_processing()) { + logit_lanes[(size_t) lane] = 1; + } + } + for (const PrefillLane & prefill : prefill_lanes) { + if (prefill.commit && + slots_.slot(prefill.slot).sampler.needs_logit_processing()) { + logit_lanes[(size_t) prefill.lane] = 1; + } + } + + std::vector logits; + std::vector argmax; + const bool bucket_history = plan.prefills.empty() && b_.moe_hybrid_; + if (!deepseek4_paged_gathered_step( + b_.backend_, b_.cfg_.device.gpu, b_.w_, b_.paged_cache_, + embeddings.data(), lane_tokens.data(), lane_positions.data(), + lane_slots.data(), (uint32_t)lane_tokens.size(), + compact_tables.data(), stride_, bucket_history, + logit_lanes.data(), logits, argmax, + b_.moe_hybrid_.get(), b_.routing_stats_.get(), + timing ? &telemetry : nullptr)) { + return fail_step("DeepSeek4 gathered paged graph failed"); + } + + const auto sample_t0 = Clock::now(); + auto sample_lane = [&](int slot_id, int lane) { + SeqSlot & slot = slots_.slot(slot_id); + if (!slot.sampler.needs_logit_processing()) { + return argmax[(size_t)lane]; + } + return sample_logits( + logits.data() + (size_t)lane * b_.w_.n_vocab, + b_.w_.n_vocab, slot.sampler, slot.sample_history, slot.rng); + }; + + for (size_t i = 0; i < inputs.size(); ++i) { + const int lane = decode_lanes[i]; + if (lane < 0) continue; + DecodeOutput & out = result.decode[i]; + slots_.commit_step(out.slot); + out.token = sample_lane(out.slot, lane); + } + for (const PrefillLane & prefill : prefill_lanes) { + PrefillOutput out; + out.slot = prefill.slot; + if (prefill.commit) { + out.status = PrefillOutput::Status::completed; + out.token = sample_lane(prefill.slot, prefill.lane); + slots_.commit_prefill(prefill.slot); + } + result.prefills.push_back(std::move(out)); + } + if (timing) { + telemetry.sample_us += std::chrono::duration_cast( + Clock::now() - sample_t0).count(); + const char * phase = plan.prefills.empty() ? "paged-decode" + : inputs.empty() ? "paged-prefill" : "paged-mixed"; + log_deepseek4_step_telemetry(phase, (int)lane_tokens.size(), 1, + std::chrono::duration(Clock::now() - phase_t0).count(), telemetry); + } + return result; +} + +void DeepSeek4SeqEngine::retire(int slot) { + if (!slots_.is_active(slot)) return; + slots_.retire(slot); + reset_deepseek4_paged_slot(b_.paged_cache_, (uint32_t)slot); + std::fill_n(host_tables_.data() + (size_t)slot * stride_, stride_, -1); +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_seq_engine.h b/server/src/deepseek4/deepseek4_seq_engine.h new file mode 100644 index 000000000..9eeadfd81 --- /dev/null +++ b/server/src/deepseek4/deepseek4_seq_engine.h @@ -0,0 +1,41 @@ +#pragma once + +#include "common/concurrency/seq_engine.h" +#include "common/concurrency/seq_slot_manager.h" + +#include +#include + +namespace dflash::common { + +class DeepSeek4Backend; + +// Exact concurrent serving path for DeepSeek4. Model state remains in +// DeepSeek4PagedCache; this class owns only scheduler-facing slot state and +// the host mirror of the model's block table. +class DeepSeek4SeqEngine final : public SeqEngine { +public: + DeepSeek4SeqEngine(DeepSeek4Backend & backend, PagedKvPool & pool, + int max_ctx, uint32_t table_stride); + + int slot_count() const override { return slots_.slot_count(); } + int max_context() const override { return slots_.max_context(); } + AdmitResult admit(uint64_t request_id, const std::vector & prompt, + const SamplerCfg & sampler) override; + StepResult step(const StepPlan & plan) override; + StepPlanLimits step_plan_limits(int decode_rows) const override; + void retire(int slot) override; + bool token_is_eos(int32_t token) const override; + +private: + bool set_block(int slot, int logical, int32_t physical); + void fail_prefill(int slot, std::vector & outputs, + const std::string & error); + + DeepSeek4Backend & b_; + SeqSlotManager slots_; + uint32_t stride_ = 0; + std::vector host_tables_; +}; + +} // namespace dflash::common diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp index ee969b118..02e07f60b 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.cpp +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.cpp @@ -1,333 +1 @@ #include "qwen35_slot_manager.h" - -#include -#include - -namespace dflash::common { - -Qwen35SlotManager::Qwen35SlotManager(PagedKvPool & pool, int max_ctx) - : pool_(pool), max_ctx_(max_ctx) { - slots_.assign(pool.max_sequences(), Qwen35Slot{}); -} - -int Qwen35SlotManager::decoding_count() const { - int n = 0; - for (const Qwen35Slot & s : slots_) { - n += s.decoding() ? 1 : 0; - } - return n; -} - -uint32_t Qwen35SlotManager::decode_headroom_capacity(int logical_tokens) const { - const uint64_t extended = - static_cast(std::max(0, logical_tokens)) + - pool_.block_size(); - return static_cast(std::min( - static_cast(max_ctx_), extended)); -} - -bool Qwen35SlotManager::capacity_fits_pool(uint32_t token_capacity) const { - const uint64_t blocks = token_capacity == 0 ? 0 : - 1 + (static_cast(token_capacity) - 1) / - pool_.block_size(); - return blocks <= pool_.physical_block_count(); -} - -PagedKvStatus Qwen35SlotManager::protect_decode_headroom() { - struct TopUp { - PagedKvSequenceHandle handle; - uint32_t token_capacity = 0; - }; - - std::vector topups; - topups.reserve(slots_.size()); - uint64_t total_additional = 0; - const uint64_t block_size = pool_.block_size(); - for (const Qwen35Slot & slot : slots_) { - if (!slot.decoding()) continue; - const uint32_t capacity = decode_headroom_capacity(slot.cur_pos); - if (!capacity_fits_pool(capacity)) continue; - - uint32_t owned_blocks = 0; - const PagedKvStatus status = - pool_.owned_block_count(slot.handle, owned_blocks); - if (status != PagedKvStatus::Ok) return status; - const uint64_t target_blocks = capacity == 0 ? 0 : - 1 + (static_cast(capacity) - 1) / block_size; - if (target_blocks <= owned_blocks) continue; - const uint32_t additional = - static_cast(target_blocks - owned_blocks); - total_additional += additional; - topups.push_back({slot.handle, capacity}); - } - - // Preflight the whole cohort before moving a block, so a failed admission - // attempt cannot protect only whichever decoder happened to be visited - // first. - if (total_additional > pool_.free_block_count()) { - return PagedKvStatus::BlocksExhausted; - } - for (const TopUp & topup : topups) { - const PagedKvStatus status = - pool_.reserve_capacity(topup.handle, topup.token_capacity); - if (status != PagedKvStatus::Ok) return status; - } - return PagedKvStatus::Ok; -} - -bool Qwen35SlotManager::is_active(int slot) const { - return slot >= 0 && slot < (int)slots_.size() && - slots_[(size_t)slot].active(); -} - -bool Qwen35SlotManager::is_prefilling(int slot) const { - return is_active(slot) && slots_[(size_t)slot].prefilling(); -} - -bool Qwen35SlotManager::has_prefill_prompt_at_least(int tokens) const { - if (tokens <= 0) return true; - return std::any_of(slots_.begin(), slots_.end(), - [tokens](const Qwen35Slot & slot) { - return slot.prefilling() && slot.prompt_len >= tokens; - }); -} - -SeqEngine::AdmitResult Qwen35SlotManager::admit( - uint64_t request_id, const std::vector & prompt, - const SamplerCfg & sampler) { - using AdmitStatus = SeqEngine::AdmitResult::Status; - SeqEngine::AdmitResult r; - if (prompt.empty()) { - r.error = "empty prompt"; - return r; - } - if (prompt.size() > static_cast(max_ctx_)) { - r.error = "prompt exceeds max_ctx"; - return r; - } - const int prompt_len = static_cast(prompt.size()); - - // A prompt larger than the whole pool can NEVER be admitted; waiting - // for other sequences to drain would stall the queue forever and then - // fail anyway. Hard-fail it up front instead of reporting busy. - const uint64_t pool_capacity = - (uint64_t)pool_.physical_block_count() * pool_.block_size(); - if ((uint64_t)prompt_len > pool_capacity) { - r.error = "prompt needs " + std::to_string(prompt_len) + - " KV tokens but the pool holds " + - std::to_string(pool_capacity) + - "; raise --kv-pool-tokens or shorten the prompt"; - return r; - } - - int slot = -1; - for (int i = 0; i < (int)slots_.size(); i++) { - if (!slots_[(size_t)i].active()) { slot = i; break; } - } - if (slot < 0) { - r.status = AdmitStatus::busy; - r.error = "all decode slots are busy"; - return r; - } - - // A newly freed block belongs to any older decoder missing its rolling - // next-page reserve before it can belong to this admission. - const PagedKvStatus headroom_status = protect_decode_headroom(); - if (headroom_status != PagedKvStatus::Ok) { - r.status = headroom_status == PagedKvStatus::BlocksExhausted - ? AdmitStatus::busy : AdmitStatus::failed; - r.error = r.status == AdmitStatus::busy - ? "existing decoders need the available KV headroom" - : paged_kv_status_string(headroom_status); - return r; - } - - PagedKvSequenceHandle handle; - uint32_t reservation_capacity = - decode_headroom_capacity(prompt_len); - if (!capacity_fits_pool(reservation_capacity)) { - // The prompt itself fits, but this physical pool can never hold its - // following page. Preserve useful prompt-only behavior and report - // decode exhaustion later if the sequence reaches that boundary. - reservation_capacity = static_cast(prompt_len); - } - const PagedKvStatus status = pool_.acquire_reserved( - request_id, reservation_capacity, handle); - if (status != PagedKvStatus::Ok) { - r.status = status == PagedKvStatus::SequenceSlotsExhausted || - status == PagedKvStatus::BlocksExhausted - ? AdmitStatus::busy : AdmitStatus::failed; - r.error = status == PagedKvStatus::BlocksExhausted - ? "not enough unreserved KV blocks for the prompt and decode headroom" - : paged_kv_status_string(status); - return r; - } - - Qwen35Slot & s = slots_[(size_t)slot]; - s.phase = Qwen35SlotPhase::prefill; - s.handle = handle; - s.cur_pos = 0; - s.prompt_len = prompt_len; - s.sampler = sampler; - s.sample_history = prompt; - // Same predicate the engine uses to pick CPU sampling over GPU argmax: - // a seed only means anything when the sampler actually draws. - if (sampler.needs_logit_processing() && sampler.seed != 0) { - s.rng.seed(sampler.seed); - } else { - s.rng.seed(std::random_device{}()); - } - - r.status = AdmitStatus::admitted; - r.slot = slot; - return r; -} - -Qwen35SlotManager::PrefillChunk Qwen35SlotManager::append_prefill( - int slot, int n_tokens) { - PrefillChunk out; - if (!is_prefilling(slot) || n_tokens < 1) return out; - - Qwen35Slot & s = slots_[(size_t)slot]; - if (s.cur_pos > s.prompt_len || - n_tokens > s.prompt_len - s.cur_pos) { - return out; - } - - PagedKvAppendResult app = pool_.append(s.handle, (uint32_t)n_tokens); - if (!app) { - // Admission reserved the whole prompt. Treat exhaustion here as a - // broken invariant, not a retryable condition: retrying a batch of - // all-prefill slots without any decoder able to retire would livelock. - if (app.status == PagedKvStatus::BlocksExhausted) { - std::fprintf(stderr, - "[parallel] reserved prefill capacity missing for slot %d\n", - slot); - } - return out; - } - - out.rows.reserve(app.write_slots.size()); - for (const PagedKvWriteSlot & write : app.write_slots) { - out.rows.push_back((int64_t)write.physical_token_index); - if (write.block_offset == 0) { - if (out.first_new_block < 0) { - out.first_new_block = - (int)(write.logical_position / pool_.block_size()); - } - out.new_blocks.push_back((int32_t)write.physical_block); - } - } - s.cur_pos += n_tokens; - out.ok = true; - return out; -} - -void Qwen35SlotManager::commit_prefill(int slot) { - if (!is_prefilling(slot)) return; - Qwen35Slot & s = slots_[(size_t)slot]; - if (s.cur_pos != s.prompt_len) return; - s.phase = Qwen35SlotPhase::decode; -} - -Qwen35SlotManager::StepAppend Qwen35SlotManager::append_tokens( - int slot, const int32_t * fed_tokens, int n_tokens) { - StepAppend out; - if (!is_active(slot) || !slots_[(size_t)slot].decoding() || - !fed_tokens || n_tokens < 1) return out; - Qwen35Slot & s = slots_[(size_t)slot]; - if (!s.staged_tokens.empty() || s.cur_pos > max_ctx_ || - n_tokens > max_ctx_ - s.cur_pos) { - return out; - } - PagedKvAppendResult app = pool_.append( - s.handle, static_cast(n_tokens)); - if (!app || app.token_count != static_cast(n_tokens)) { - out.busy = app.status == PagedKvStatus::BlocksExhausted; - return out; - } - if (app.write_slots.size() != static_cast(n_tokens) || - app.write_slots.front().logical_position != - static_cast(s.cur_pos) || - app.write_slots.back().logical_position != - static_cast(s.cur_pos + n_tokens - 1)) { - const PagedKvStatus rollback = - pool_.rollback_append(s.handle, static_cast(n_tokens)); - if (rollback != PagedKvStatus::Ok) { - std::fprintf(stderr, - "[parallel] slot %d failed to roll back invalid append: %s\n", - slot, paged_kv_status_string(rollback)); - } - return out; - } - - out.physical_rows.reserve(app.write_slots.size()); - for (const PagedKvWriteSlot & write : app.write_slots) { - out.physical_rows.push_back( - static_cast(write.physical_token_index)); - if (write.block_offset == 0) { - if (out.first_new_block < 0) { - out.first_new_block = static_cast( - write.logical_position / pool_.block_size()); - } - out.new_blocks.push_back( - static_cast(write.physical_block)); - } - } - s.staged_tokens.assign(fed_tokens, fed_tokens + n_tokens); - out.ok = true; - out.count = n_tokens; - out.physical_row = out.physical_rows.front(); - out.position = s.cur_pos; - if (!out.new_blocks.empty()) { - out.new_block = out.new_blocks.front(); - out.new_block_index = out.first_new_block; - } - return out; -} - -Qwen35SlotManager::StepAppend Qwen35SlotManager::append_token( - int slot, int32_t fed_token) { - return append_tokens(slot, &fed_token, 1); -} - -void Qwen35SlotManager::commit_step(int slot) { - if (!is_active(slot)) return; - Qwen35Slot & s = slots_[(size_t)slot]; - if (s.staged_tokens.empty()) return; - s.sample_history.insert( - s.sample_history.end(), s.staged_tokens.begin(), - s.staged_tokens.end()); - s.cur_pos += static_cast(s.staged_tokens.size()); - s.staged_tokens.clear(); -} - -bool Qwen35SlotManager::rollback_step(int slot) { - if (!is_active(slot)) return false; - Qwen35Slot & s = slots_[(size_t)slot]; - if (s.staged_tokens.empty()) return false; - const PagedKvStatus status = pool_.rollback_append( - s.handle, static_cast(s.staged_tokens.size())); - if (status != PagedKvStatus::Ok) { - std::fprintf(stderr, - "[parallel] slot %d staged append rollback failed: %s\n", - slot, paged_kv_status_string(status)); - return false; - } - s.staged_tokens.clear(); - return true; -} - -void Qwen35SlotManager::retire(int slot) { - if (slot < 0 || slot >= (int)slots_.size()) return; - Qwen35Slot & s = slots_[(size_t)slot]; - if (!s.active()) return; - const PagedKvStatus status = pool_.release(s.handle); - if (status != PagedKvStatus::Ok && status != PagedKvStatus::StaleHandle) { - std::fprintf(stderr, "[parallel] slot %d release failed: %s\n", - slot, paged_kv_status_string(status)); - } - s = Qwen35Slot{}; -} - -} // namespace dflash::common diff --git a/server/src/qwen35/concurrency/qwen35_slot_manager.h b/server/src/qwen35/concurrency/qwen35_slot_manager.h index 1d7af2db5..65dc0dee1 100644 --- a/server/src/qwen35/concurrency/qwen35_slot_manager.h +++ b/server/src/qwen35/concurrency/qwen35_slot_manager.h @@ -1,149 +1,13 @@ -// Qwen35SlotManager — complete host-side state for each Qwen serving slot. -// -// Companion of PagedKvPool: the pool hands out sequence handles and physical -// blocks; this class owns everything else a slot needs between admission and -// retirement — the pool-handle lifecycle (including every error path), the -// admission arithmetic (context clamp, prompt reservation, and rolling decode -// headroom), on-demand block allocation, per-slot sampler/RNG/penalty-history -// state, and the position counters. -// -// It deliberately owns NO device state. Prefill/decode allocation returns -// physical rows and block-table deltas as plain vectors. Prompt, KV ownership, -// sampler, and progress live together here; the scheduler keeps -// only its coarse request phase. -// -// Not thread-safe; the single scheduler thread is the only caller. - #pragma once -#include "common/concurrency/paged_kv_pool.h" -#include "common/sampler.h" -#include "common/concurrency/seq_engine.h" - -#include -#include -#include -#include +#include "common/concurrency/seq_slot_manager.h" namespace dflash::common { -enum class Qwen35SlotPhase { - free, - prefill, - decode, -}; - -struct Qwen35Slot { - Qwen35SlotPhase phase = Qwen35SlotPhase::free; - PagedKvSequenceHandle handle; - // Prompt tokens are the immutable prefix of sample_history. Decode tokens - // append to the same allocation, avoiding a second full prompt copy. - int prompt_len = 0; - int cur_pos = 0; - SamplerCfg sampler; - std::mt19937_64 rng{0x9E3779B97F4A7C15ull}; - // Penalty history is recorded as fed rather than sampled: the scheduler - // may override a sample before the model consumes it. - std::vector sample_history; - // Decode rows reserved by the current target graph. They become durable - // only after the graph and any speculative promotion succeed. - std::vector staged_tokens; - - int generated_tokens() const { - return sample_history.size() > (size_t)prompt_len - ? (int)(sample_history.size() - (size_t)prompt_len) - : 0; - } - - bool active() const { return phase != Qwen35SlotPhase::free; } - bool prefilling() const { return phase == Qwen35SlotPhase::prefill; } - bool decoding() const { return phase == Qwen35SlotPhase::decode; } -}; - -class Qwen35SlotManager { -public: - // `max_ctx` is the per-sequence logical bound; slot count comes from the - // pool's max_sequences. The pool must outlive the manager. - Qwen35SlotManager(PagedKvPool & pool, int max_ctx); - - // Claim a free slot and atomically reserve all K/V blocks needed by the - // known prompt plus its next logical decode page when that page can exist - // in both max_ctx and the physical pool. Existing decoders are topped up - // first, so a younger admission cannot steal their next-page headroom. - // Prompts larger than the whole pool hard-fail; temporary capacity pressure - // reports busy. Seeds the slot RNG from sampler.seed only when the sampler - // actually draws, else nondeterministically. - SeqEngine::AdmitResult admit(uint64_t request_id, - const std::vector & prompt, - const SamplerCfg & sampler); - - struct PrefillChunk { - bool ok = false; - std::vector rows; - // Delta to patch into the slot's device block-table column. - std::vector new_blocks; - int first_new_block = -1; - }; - - // Append `n_tokens` more prompt rows for a prefilling slot. Physical block - // ids come from the slot's admission reservation, so any append within the - // admitted prompt is guaranteed not to wait on another sequence. - PrefillChunk append_prefill(int slot, int n_tokens); - - // Record a finished prefill and expose the slot to decode. - void commit_prefill(int slot); - - struct StepAppend { - bool ok = false; - bool busy = false; // no physical block available right now - int64_t physical_row = -1; - std::vector physical_rows; - int count = 0; - int position = -1; // logical position the fed token is written at - int32_t new_block = -1; - int new_block_index = -1; - std::vector new_blocks; - int first_new_block = -1; - }; - - StepAppend append_tokens(int slot, const int32_t * fed_tokens, - int n_tokens); - - // Allocate decode cache rows and stage the fed tokens. Both history and - // cur_pos wait for commit_step(). - StepAppend append_token(int slot, int32_t fed_token); - - // The batched step's compute succeeded: cur_pos++. - void commit_step(int slot); - - bool rollback_step(int slot); - - // Release the slot's blocks and clear its state. Safe on inactive slots - // and after a failed admission/prefill. - void retire(int slot); - - int slot_count() const { return (int)slots_.size(); } - int max_context() const { return max_ctx_; } - int decoding_count() const; - bool is_active(int slot) const; - bool is_prefilling(int slot) const; - bool has_prefill_prompt_at_least(int tokens) const; - Qwen35Slot & slot(int i) { return slots_[(size_t)i]; } - const Qwen35Slot & slot(int i) const { return slots_[(size_t)i]; } - -private: - // Logical extent whose block count includes the sequence's current pages - // plus one future page, capped at max_ctx. - uint32_t decode_headroom_capacity(int logical_tokens) const; - bool capacity_fits_pool(uint32_t token_capacity) const; - - // Atomically preflight and top up every decoding slot as one cohort before - // a younger sequence may reserve capacity. - PagedKvStatus protect_decode_headroom(); - - PagedKvPool & pool_; - int max_ctx_ = 0; - std::vector slots_; -}; +// Compatibility aliases keep the Qwen engine source stable while both Qwen +// and DeepSeek share the current model-neutral slot lifecycle implementation. +using Qwen35SlotPhase = SeqSlotPhase; +using Qwen35Slot = SeqSlot; +using Qwen35SlotManager = SeqSlotManager; } // namespace dflash::common diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index e2db605fa..5f6483cf9 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -114,15 +114,17 @@ static void print_usage(const char * prog) { " --fa-window Flash-attention sliding window (default: 0=full).\n" " WARNING: >0 drops system prompt / tool definitions\n" " from attention at long contexts. Use 0 for tools.\n" - " --paged-attention Use 16-token paged KV blocks for Qwen3.6-27B\n" - " autoregressive decode (experimental)\n" + " --paged-attention Use paged autoregressive decode for dense Qwen3.5 and\n" + " Qwen3.6 targets with 16-token blocks, or DeepSeek4\n" + " targets with 128-token blocks. This mode is experimental.\n" " --max-concurrency Maximum concurrent decode sequences\n" " (enables paged attention; default: 1)\n" " --admission-coalesce-ms Idle-to-busy batching window\n" " (default: 20; 0 disables)\n" " --kv-pool-tokens Total paged K/V pool shared by all\n" " --max-concurrency slots, in tokens\n" - " (default: sized from available device memory)\n" + " By default, Qwen sizes the pool from available device\n" + " memory. DeepSeek4 reserves --max-ctx per slot.\n" " --model-name Model name for /v1/models (default: dflash)\n" " --prefix-cache-slots Prefix cache slots (default: 32, 0 disables)\n" " --agent-turn-cache Extend prefix caching through generated tool calls\n" diff --git a/server/test/bench_rocmfp_mix_gateup_glu.cpp b/server/test/bench_rocmfp_mix_gateup_glu.cpp index b7fdf2a5f..4c7423c43 100644 --- a/server/test/bench_rocmfp_mix_gateup_glu.cpp +++ b/server/test/bench_rocmfp_mix_gateup_glu.cpp @@ -1,4 +1,4 @@ -// Microbenchmark: ONE fused gate/up+SwiGLU launch vs TWO unfused matvec launches. +// Microbenchmark: the dynamic fused gate/up+SwiGLU path vs TWO plain matvec launches. // // Tests the claim the whole fusion rests on. Profiling attributed ~102% of the measured 4.6% // adaptive decode penalty to LAUNCH COUNT (30100 vs qtype 107's 15050) rather than to decode @@ -11,9 +11,11 @@ // shape actually costs. The fused arm therefore looks *worse* than reality -- a conservative // comparison, which is the direction to err in. // -// Reports per-iteration wall time over many iterations on one stream, after a warmup, plus the -// relative delta. Single stream and back-to-back launches on purpose: that is the decode -// dependency chain, where launch overhead cannot be hidden. +// The fused wrapper uses one dual-projection launch through q=2 and two +// lower-register launches (with SwiGLU folded into the second) for wider verification. Reports +// per-iteration wall time over many iterations on one stream, after a warmup, +// plus the relative delta. Single stream and back-to-back launches on purpose: +// that is the decode dependency chain, where launch overhead cannot be hidden. #include "ggml-cuda.h" @@ -21,6 +23,7 @@ #include #include +#include #include #include #include @@ -39,9 +42,20 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_id_glu( int64_t ids_s0, int64_t ids_s1, int64_t src1_s1, int64_t src1_s2, int64_t dst_s1, int64_t dst_s2, float glu_limit, hipStream_t stream); +bool ggml_cuda_rocmfp3_mix_mul_mat_id( + const void * vx, const float * src1, const int32_t * ids, float * dst, + int in, int out, int n_expert_used, int n_tokens, int ne11, + int64_t ids_s0, int64_t ids_s1, int64_t src1_s1, int64_t src1_s2, + int64_t dst_s1, int64_t dst_s2, hipStream_t stream); + +bool ggml_cuda_rocmfp3_mix_mul_mat_id_glu( + const void * vx_up, const void * vx_gate, + const float * src1, const int32_t * ids, float * dst, + int in, int out, int n_expert_used, int n_tokens, int ne11, + int64_t ids_s0, int64_t ids_s1, int64_t src1_s1, int64_t src1_s2, + int64_t dst_s1, int64_t dst_s2, float glu_limit, hipStream_t stream); + static constexpr int QK = 32; -static constexpr int BLOCK_BYTES = 10; -static constexpr int K = 4; static uint32_t xs = 0x1234567u; static uint32_t rnd() { xs ^= xs << 13; xs ^= xs >> 17; xs ^= xs << 5; return xs; } @@ -58,26 +72,41 @@ static double median(std::vector v) { return v[v.size() / 2]; } -int main() { +int main(int argc, char ** argv) { int ndev = 0; if (hipGetDeviceCount(&ndev) != hipSuccess || ndev == 0) { std::fprintf(stderr, "SKIP: no HIP device\n"); return 0; } - - // DeepSeek-V4-Flash decode geometry: hidden 7168, n_ff_exp 2048, top-k 4 (--ds4-expert-top-k 4 - // is what the measured runs used), one token per step. - const int in = 7168, out = 2048, n_experts = 8, n_used = 4, ntok = 1; + int device = 0; + hipDeviceProp_t properties{}; + HIP_OK(hipGetDevice(&device)); + HIP_OK(hipGetDeviceProperties(&properties, device)); + const bool strix_tuned = + std::strncmp(properties.gcnArchName, "gfx1151", 7) == 0; + + // DeepSeek-V4-Flash verifier geometry: hidden 4096, n_ff_exp 2048, + // model-default top-k 6. The optional first argument selects the verifier + // width so q=1 and q=4 can be compared without rebuilding. + const int in = 4096, out = 2048, n_experts = 8, n_used = 6; + const int ntok = argc > 1 ? std::atoi(argv[1]) : 4; + const bool fp3 = argc > 2 && std::strcmp(argv[2], "q3") == 0; + if (ntok <= 0 || ntok > 16 || (argc > 2 && !fp3 && std::strcmp(argv[2], "q2") != 0)) { + std::fprintf(stderr, "usage: %s [tokens:1..16] [q2|q3]\n", argv[0]); + return 2; + } + const int block_bytes = fp3 ? 14 : 10; + const int levels = fp3 ? 8 : 4; const int nb = in / QK; - const size_t rows_bytes = (size_t) out * nb * BLOCK_BYTES; + const size_t rows_bytes = (size_t) out * nb * block_bytes; std::vector w(rows_bytes * n_experts); for (auto & b : w) b = (uint8_t) rnd(); - for (size_t blk = 0; blk < w.size() / BLOCK_BYTES; ++blk) { - w[blk * BLOCK_BYTES + 8] = 0x30; - w[blk * BLOCK_BYTES + 9] = 0x30; + for (size_t blk = 0; blk < w.size() / block_bytes; ++blk) { + w[blk * block_bytes + block_bytes - 2] = 0x30; + w[blk * block_bytes + block_bytes - 1] = 0x30; } - std::vector books((size_t) n_experts * 2 * K); + std::vector books((size_t) n_experts * 2 * levels); for (size_t i = 0; i < books.size(); ++i) books[i] = f32_to_bf16(-0.5f + 0.2f * (float) (i % 5)); std::vector modes(n_experts, 1); @@ -100,16 +129,25 @@ int main() { for (int i = 0; i < n_used * ntok; ++i) idsh[i] = i % n_experts; HIP_OK(hipMemcpy(d_ids, idsh.data(), sizeof(int32_t) * idsh.size(), hipMemcpyHostToDevice)); - if (!ggml_cuda_rocmfp2_mix_register_host( + const auto register_mix = fp3 ? ggml_cuda_rocmfp3_mix_register_host + : ggml_cuda_rocmfp2_mix_register_host; + const auto unregister_mix = fp3 ? ggml_cuda_rocmfp3_mix_unregister + : ggml_cuda_rocmfp2_mix_unregister; + const auto mul_mat_id = fp3 ? ggml_cuda_rocmfp3_mix_mul_mat_id + : ggml_cuda_rocmfp2_mix_mul_mat_id; + const auto mul_mat_id_glu = fp3 ? ggml_cuda_rocmfp3_mix_mul_mat_id_glu + : ggml_cuda_rocmfp2_mix_mul_mat_id_glu; + + if (!register_mix( d_up, rows_bytes, n_experts, out, in, books.data(), modes.data())) { std::fprintf(stderr, "FAIL: mixed-tensor registration failed\n"); return 1; } - if (!ggml_cuda_rocmfp2_mix_register_host( + if (!register_mix( d_gate, rows_bytes, n_experts, out, in, books.data(), modes.data())) { - ggml_cuda_rocmfp2_mix_unregister(d_up); + unregister_mix(d_up); std::fprintf(stderr, "FAIL: mixed-tensor registration failed\n"); return 1; } @@ -126,17 +164,17 @@ int main() { auto time_unfused = [&]() { hipEvent_t a, b; hipEventCreate(&a); hipEventCreate(&b); for (int i = 0; i < WARM; ++i) { - ggml_cuda_rocmfp2_mix_mul_mat_id(d_up, d_x, d_ids, d_a, in, out, n_used, ntok, 1, + mul_mat_id(d_up, d_x, d_ids, d_a, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, stream); - ggml_cuda_rocmfp2_mix_mul_mat_id(d_gate, d_x, d_ids, d_b, in, out, n_used, ntok, 1, + mul_mat_id(d_gate, d_x, d_ids, d_b, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, stream); } hipStreamSynchronize(stream); hipEventRecord(a, stream); for (int i = 0; i < ITERS; ++i) { - ggml_cuda_rocmfp2_mix_mul_mat_id(d_up, d_x, d_ids, d_a, in, out, n_used, ntok, 1, + mul_mat_id(d_up, d_x, d_ids, d_a, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, stream); - ggml_cuda_rocmfp2_mix_mul_mat_id(d_gate, d_x, d_ids, d_b, in, out, n_used, ntok, 1, + mul_mat_id(d_gate, d_x, d_ids, d_b, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, stream); } hipEventRecord(b, stream); @@ -148,13 +186,13 @@ int main() { auto time_fused = [&]() { hipEvent_t a, b; hipEventCreate(&a); hipEventCreate(&b); for (int i = 0; i < WARM; ++i) { - ggml_cuda_rocmfp2_mix_mul_mat_id_glu(d_up, d_gate, d_x, d_ids, d_a, in, out, n_used, + mul_mat_id_glu(d_up, d_gate, d_x, d_ids, d_a, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, 7.0f, stream); } hipStreamSynchronize(stream); hipEventRecord(a, stream); for (int i = 0; i < ITERS; ++i) { - ggml_cuda_rocmfp2_mix_mul_mat_id_glu(d_up, d_gate, d_x, d_ids, d_a, in, out, n_used, + mul_mat_id_glu(d_up, d_gate, d_x, d_ids, d_a, in, out, n_used, ntok, 1, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, 7.0f, stream); } hipEventRecord(b, stream); @@ -170,20 +208,32 @@ int main() { for (int r = 0; r < REPS; ++r) { u.push_back(time_unfused()); f.push_back(time_fused()); } const double mu = median(u), mf = median(f); - std::fprintf(stderr, "geometry: in=%d out=%d top_k=%d ntok=%d (%d iters x %d interleaved reps)\n", - in, out, n_used, ntok, ITERS, REPS); + std::fprintf(stderr, "geometry: qtype=%s in=%d out=%d top_k=%d ntok=%d " + "(%d iters x %d interleaved reps)\n", + fp3 ? "q3-mix" : "q2-mix", in, out, n_used, ntok, ITERS, REPS); std::fprintf(stderr, " unfused (2 launches, swiglu NOT counted): %8.4f ms/step [", mu); for (double v : u) std::fprintf(stderr, " %.4f", v); - std::fprintf(stderr, " ]\n fused (1 launch, swiglu included ): %8.4f ms/step [", mf); + const int fused_launches = !fp3 && strix_tuned && ntok > 2 ? 2 : 1; + std::fprintf(stderr, " ]\n fused (%d launch%s, swiglu included): %8.4f ms/step [", + fused_launches, fused_launches == 1 ? "" : "es", mf); for (double v : f) std::fprintf(stderr, " %.4f", v); std::fprintf(stderr, " ]\n"); std::fprintf(stderr, " fused is %+.2f%% vs unfused (negative = faster)\n", 100.0 * (mf / mu - 1.0)); std::fprintf(stderr, " per-layer saving %.4f ms -> over 43 layers %.3f ms/step\n", mu - mf, (mu - mf) * 43.0); - ggml_cuda_rocmfp2_mix_unregister(d_gate); - ggml_cuda_rocmfp2_mix_unregister(d_up); + std::vector result(yn); + HIP_OK(hipMemcpy(result.data(), d_a, sizeof(float) * yn, hipMemcpyDeviceToHost)); + uint64_t result_hash = 1469598103934665603ull; + const uint8_t * result_bytes = reinterpret_cast(result.data()); + for (size_t i = 0; i < sizeof(float) * yn; ++i) { + result_hash = (result_hash ^ result_bytes[i]) * 1099511628211ull; + } + std::fprintf(stderr, " result fnv1a64: %016llx\n", (unsigned long long) result_hash); + + unregister_mix(d_gate); + unregister_mix(d_up); hipStreamDestroy(stream); HIP_OK(hipFree(d_up)); HIP_OK(hipFree(d_gate)); HIP_OK(hipFree(d_x)); HIP_OK(hipFree(d_a)); HIP_OK(hipFree(d_b)); HIP_OK(hipFree(d_ids)); diff --git a/server/test/test_deepseek4_page_layout.cpp b/server/test/test_deepseek4_page_layout.cpp new file mode 100644 index 000000000..93654836e --- /dev/null +++ b/server/test/test_deepseek4_page_layout.cpp @@ -0,0 +1,58 @@ +#include "deepseek4/deepseek4_page_layout.h" +#include "CppUnitTestFramework.hpp" + +#include +#include + +using namespace dflash::common; + +namespace { struct DeepSeek4PageLayoutFixture {}; } + +TEST_CASE(DeepSeek4PageLayoutFixture, raw_ring_wraps) { + CHECK(ds4_raw_ring_row(0) == 0); + CHECK(ds4_raw_ring_row(127) == 127); + CHECK(ds4_raw_ring_row(128) == 0); + CHECK(ds4_raw_ring_row(255) == 127); + +} + +TEST_CASE(DeepSeek4PageLayoutFixture, compressed_page_rows) { + uint64_t row = 999; + bool emitted = true; + CHECK(ds4_compressed_page_row(3, 7, 4, row, emitted)); + CHECK(emitted && row == 7 * 32); + CHECK(ds4_compressed_page_row(4, 7, 4, row, emitted)); + CHECK(!emitted && row == 7 * 32); + CHECK(ds4_compressed_page_row(127, 7, 4, row, emitted)); + CHECK(emitted && row == 7 * 32 + 31); + CHECK(ds4_compressed_page_row(128, 42, 4, row, emitted)); + CHECK(!emitted); + CHECK(ds4_compressed_page_row(131, 42, 4, row, emitted)); + CHECK(emitted && row == 42 * 32); + CHECK(ds4_compressed_page_row(255, 3, 4, row, emitted)); + CHECK(emitted && row == 3 * 32 + 31); + + CHECK(ds4_compressed_page_row(127, 91, 128, row, emitted)); + CHECK(emitted && row == 91); + CHECK(ds4_compressed_page_row(128, 2, 128, row, emitted)); + CHECK(!emitted); + CHECK(ds4_compressed_page_row(255, 2, 128, row, emitted)); + CHECK(emitted && row == 2); + +} + +TEST_CASE(DeepSeek4PageLayoutFixture, invalid_geometry) { + uint64_t row = 999; + bool emitted = true; + uint64_t capacity = 0; + CHECK(ds4_compressed_page_capacity(5, 4, capacity) && capacity == 160); + CHECK(ds4_compressed_page_capacity(5, 128, capacity) && capacity == 5); + CHECK(!ds4_compressed_page_capacity(5, 0, capacity)); + CHECK(!ds4_compressed_page_capacity(5, 16, capacity)); + CHECK(!ds4_compressed_page_capacity( + std::numeric_limits::max(), 4, capacity)); + CHECK(!ds4_compressed_page_row(3, + std::numeric_limits::max(), 4, row, emitted)); + CHECK(!ds4_compressed_page_row(3, 0, 16, row, emitted)); + +} diff --git a/server/test/test_deepseek4_paged_cache.cpp b/server/test/test_deepseek4_paged_cache.cpp new file mode 100644 index 000000000..e7d6b15ad --- /dev/null +++ b/server/test/test_deepseek4_paged_cache.cpp @@ -0,0 +1,156 @@ +#include "deepseek4/deepseek4_paged_cache.h" +#include "deepseek4/deepseek4_page_layout.h" +#include "CppUnitTestFramework.hpp" +#include +using namespace dflash::common; +namespace { struct DeepSeek4PagedCacheFixture {}; } + +TEST_CASE(DeepSeek4PagedCacheFixture, pool_capacity) { + uint32_t blocks = 0; + CHECK(plan_deepseek4_paged_pool_blocks(4096, 6, 0, blocks)); + CHECK(blocks == 192); + CHECK(plan_deepseek4_paged_pool_blocks(1, 6, 0, blocks)); + CHECK(blocks == 6); + CHECK(plan_deepseek4_paged_pool_blocks(129, 6, 0, blocks)); + CHECK(blocks == 12); + CHECK(plan_deepseek4_paged_pool_blocks(4097, 6, 0, blocks)); + CHECK(blocks == 198); + CHECK(plan_deepseek4_paged_pool_blocks(131072, 6, 0, blocks)); + CHECK(blocks == 6144); + CHECK(plan_deepseek4_paged_pool_blocks(4096, 6, 128, blocks)); + CHECK(blocks == 1); + CHECK(plan_deepseek4_paged_pool_blocks(4096, 6, 129, blocks)); + CHECK(blocks == 2); + CHECK(!plan_deepseek4_paged_pool_blocks(0, 6, 128, blocks)); + CHECK(!plan_deepseek4_paged_pool_blocks(4096, 0, 128, blocks)); + const uint64_t overflowing_tokens = + uint64_t(UINT32_MAX / DS4_PAGE_TOKENS) * DS4_PAGE_TOKENS + 1; + CHECK(!plan_deepseek4_paged_pool_blocks( + 4096, 6, overflowing_tokens, blocks)); + +} + +TEST_CASE(DeepSeek4PagedCacheFixture, allocation_plan) { + DeepSeek4PagedCachePlan p, twice; + CHECK(plan_deepseek4_paged_cache(512, 128, 3, 4096, 40, {0, 4, 128}, p)); + CHECK(p.max_blocks_per_sequence == 32 && p.physical_rows[0] == 0); + CHECK(p.physical_rows[1] == 1280 && p.physical_rows[2] == 40); + CHECK(p.raw_bytes == uint64_t(3) * 512 * 128 * 3 * 2); + CHECK(p.total_persistent_bytes == + p.raw_bytes + p.compressed_bytes + p.state_bytes); + CHECK(plan_deepseek4_paged_cache(512, 128, 6, 4096, 40, {0, 4, 128}, twice)); + // Paged rows are shared; only raw rings and compressor state scale by slots. + CHECK(twice.compressed_bytes == p.compressed_bytes); + CHECK(twice.raw_bytes == p.raw_bytes * 2 && twice.state_bytes == p.state_bytes * 2); + CHECK(!plan_deepseek4_paged_cache(512, 128, 1, 4096, 40, {4, 16}, twice)); + CHECK(!plan_deepseek4_paged_cache(512, 128, 1, 4096, + std::numeric_limits::max(), {4}, twice)); + +} + +TEST_CASE(DeepSeek4PagedCacheFixture, gathered_lane_history) { + const int32_t slots[] = {2, 5, -1}; + const int64_t positions[] = {3, 259, 999}; + const int32_t tables[] = {4, 3, 2, 6, 1, 7, 0, 0, 0}; + std::vector rows; + REQUIRE_TRUE(prepare_deepseek4_gathered_lane_rows( + slots, positions, 3, tables, 3, 8, 4, rows)); + CHECK(rows.size() == 3); + CHECK(rows[0].raw_history == std::vector({256, 257, 258})); + CHECK(rows[0].raw_history_valid == 3); + CHECK(rows[0].raw_scatter == 259); + CHECK(rows[0].compressed_emitted && rows[0].compressed_scatter == 4 * 32); + CHECK(rows[0].compressed_history.empty()); + CHECK(rows[0].compressed_history_valid == 0); + CHECK(rows[1].raw_history.size() == 127); + CHECK(rows[1].raw_history.front() == 5 * 128 + 4); + CHECK(rows[1].raw_history.back() == 5 * 128 + 2); + CHECK(rows[1].compressed_history.size() == 64); + CHECK(rows[1].raw_history_valid == 127); + CHECK(rows[1].compressed_history_valid == 64); + CHECK(rows[1].compressed_history.front() == 6 * 32); + CHECK(rows[1].compressed_history[31] == 6 * 32 + 31); + CHECK(rows[1].compressed_history[32] == 1 * 32); + CHECK(rows[1].compressed_history.back() == 1 * 32 + 31); + CHECK(rows[1].compressed_emitted && rows[1].compressed_scatter == 7 * 32); + CHECK(rows[2].raw_history.empty() && rows[2].compressed_history.empty()); + CHECK(rows[2].raw_history_valid == 0 && + rows[2].compressed_history_valid == 0); + CHECK(rows[2].raw_scatter == -1 && rows[2].compressed_scatter == -1); + CHECK(rows[2].position == 0); + +} + +TEST_CASE(DeepSeek4PagedCacheFixture, six_active_lanes) { + std::vector rows; + std::vector six_slots(6); + std::vector six_positions(6, 0); + std::vector six_tables(6); + for (int i = 0; i < 6; ++i) { + six_slots[(size_t) i] = i; + six_tables[(size_t) i] = i; + } + REQUIRE_TRUE(prepare_deepseek4_gathered_lane_rows( + six_slots.data(), six_positions.data(), 6, + six_tables.data(), 1, 6, 4, rows)); + CHECK(rows.size() == 6); + for (int i = 0; i < 6; ++i) { + CHECK(rows[(size_t) i].slot == i); + CHECK(rows[(size_t) i].raw_history.empty()); + CHECK(rows[(size_t) i].raw_scatter == int64_t(i * 128)); + } + +} + +TEST_CASE(DeepSeek4PagedCacheFixture, raw_ring_boundary) { + std::vector rows; + const int32_t boundary_slot[] = {1}; + const int32_t boundary_table[] = {0, 1}; + for (int64_t pos : {127LL, 128LL, 129LL}) { + REQUIRE_TRUE(prepare_deepseek4_gathered_lane_rows( + boundary_slot, &pos, 1, boundary_table, 2, 2, 0, rows)); + CHECK(rows[0].raw_history.size() == 127u); + CHECK(rows[0].raw_history.front() == 128 + (pos == 127 ? 0 : pos - 127)); + CHECK(rows[0].raw_history.back() == 128 + ((pos - 1) % 128)); + } + +} + +TEST_CASE(DeepSeek4PagedCacheFixture, ratio128_history) { + std::vector rows; + const int32_t slots[] = {2}; + const int32_t tables[] = {4, 3, 2}; + const int64_t ratio128_pos[] = {255}; + REQUIRE_TRUE(prepare_deepseek4_gathered_lane_rows( + slots, ratio128_pos, 1, tables, 3, 8, 128, rows)); + CHECK(rows[0].compressed_history.size() == 1); + CHECK(rows[0].compressed_history[0] == 4); + CHECK(rows[0].compressed_emitted && rows[0].compressed_scatter == 3); + +} + +TEST_CASE(DeepSeek4PagedCacheFixture, invalid_active_pages) { + const int32_t slot = 1; + std::vector rows; + for (uint32_t ratio : {0u, 4u, 128u}) { + const int64_t pos = 128; + const int32_t valid[] = {0, 1, -1}; + REQUIRE_TRUE(prepare_deepseek4_gathered_lane_rows( + &slot, &pos, 1, valid, 3, 2, ratio, rows)); + for (const auto & invalid : {std::vector{-1, 1}, {2, 1}, {0, -1}, {0, 2}}) { + REQUIRE_FALSE(prepare_deepseek4_gathered_lane_rows( + &slot, &pos, 1, invalid.data(), 2, 2, ratio, rows)); + CHECK(rows.size() == 1 && rows[0].position == pos); + } + for (int64_t invalid_pos : std::initializer_list{-1, 256, INT64_MAX}) { + REQUIRE_FALSE(prepare_deepseek4_gathered_lane_rows( + &slot, &invalid_pos, 1, valid, 2, 2, ratio, rows)); + } + const int32_t padding = -1; + const int32_t invalid[] = {-1}; + const int64_t invalid_pos = INT64_MAX; + REQUIRE_TRUE(prepare_deepseek4_gathered_lane_rows( + &padding, &invalid_pos, 1, invalid, 1, 2, ratio, rows)); + CHECK(rows[0].raw_scatter == -1 && rows[0].position == 0); + } +} diff --git a/server/test/test_feature_gate.cpp b/server/test/test_feature_gate.cpp index 787e0c6af..77740b627 100644 --- a/server/test/test_feature_gate.cpp +++ b/server/test/test_feature_gate.cpp @@ -377,24 +377,25 @@ void test_feature_gate_layer_split_requires_supported_arch() { CHECK(gate_result(single, "qwen3", PlacementBackend::Cuda).empty()); } -void test_feature_gate_paged_attention_requires_qwen35_monolithic() { +void test_feature_gate_paged_attention_requires_monolithic_backend() { BackendArgs args; args.model_path = "/nonexistent/model.gguf"; args.paged_attention = true; CHECK(gate_result(args, "qwen35", PlacementBackend::Cuda).empty()); CHECK(gate_result(args, "qwen35", PlacementBackend::Hip).empty()); - // Only qwen35 has a paged decode path. qwen35moe shares Qwen35Config, so - // its rejection is this gate's job — the factory's field-presence - // cross-check cannot tell the two apart. - for (const char * arch : {"qwen35moe", "laguna", "qwen3", - "gemma4", "deepseek4"}) { + BackendArgs ds4 = gate_args_hip_deepseek4(); + ds4.paged_attention = true; + CHECK(gate_result(ds4, "deepseek4", PlacementBackend::Hip).empty()); + CHECK(!gate_result(ds4, "deepseek4", PlacementBackend::Cuda).empty()); + + // qwen35moe shares Qwen35Config, so its rejection is this gate's job — + // the factory's field-presence cross-check cannot tell the two apart. + for (const char * arch : {"qwen35moe", "laguna", "qwen3", "gemma4"}) { CHECK(!gate_result(args, arch, PlacementBackend::Cuda).empty()); } - // Only the monolithic qwen35 backend owns a paged K/V pool. Both - // placements are supported qwen35 launches without the flag, so the - // rejection has to come from the paged rule. + // Only monolithic Qwen and DeepSeek backends own paged K/V pools. BackendArgs split = args; CHECK(parse_placement_device_list("cuda:0,cuda:1", split.device)); CHECK(!gate_result(split, "qwen35", PlacementBackend::Cuda).empty()); @@ -506,7 +507,7 @@ void test_feature_gate_parallel_and_kv_pool_rules() { plain.max_concurrency = 1; CHECK(gate_result(plain, "qwen35", PlacementBackend::Cuda).empty()); - // More than one slot exists only in the paged qwen35 backend. + // More than one slot requires a model-specific paged backend. BackendArgs dense; dense.model_path = "/nonexistent/model.gguf"; dense.max_concurrency = 2; @@ -527,6 +528,22 @@ void test_feature_gate_parallel_and_kv_pool_rules() { parallel.max_concurrency = 65; CHECK(!gate_result(parallel, "qwen35", PlacementBackend::Cuda).empty()); + BackendArgs ds4 = gate_args_hip_deepseek4(); + ds4.paged_attention = true; + ds4.max_concurrency = DEEPSEEK4_MAX_PAGED_SEQUENCES; + CHECK(gate_result(ds4, "deepseek4", PlacementBackend::Hip).empty()); + ds4.max_concurrency = DEEPSEEK4_MAX_PAGED_SEQUENCES + 1; + CHECK(!gate_result(ds4, "deepseek4", PlacementBackend::Hip).empty()); + ds4.max_concurrency = 2; + ds4.ds4_prefill_mode = PrefillAttentionMode::Sparse; + CHECK(!gate_result(ds4, "deepseek4", PlacementBackend::Hip).empty()); + ds4.ds4_prefill_mode = PrefillAttentionMode::Exact; + ds4.ds4_fused_decode = true; + CHECK(!gate_result(ds4, "deepseek4", PlacementBackend::Hip).empty()); + ds4.ds4_fused_decode = false; + ds4.ds4_fused_verify_f16_kv = true; + CHECK(!gate_result(ds4, "deepseek4", PlacementBackend::Hip).empty()); + // --kv-pool-tokens sizes the shared pool, so it needs slots to share. BackendArgs pool = paged; pool.kv_pool_tokens = 4096; @@ -711,9 +728,11 @@ void test_model_capability_tables() { CHECK(!arch_supports_draft_swa("qwen36", false)); CHECK(!arch_supports_paged_attention("qwen36", false)); - // Paged decode lives in the monolithic qwen35 backend alone. + // Paged decode lives in the monolithic Qwen and DeepSeek backends. CHECK(arch_supports_paged_attention("qwen35", false)); CHECK(!arch_supports_paged_attention("qwen35", true)); + CHECK(arch_supports_paged_attention("deepseek4", false)); + CHECK(!arch_supports_paged_attention("deepseek4", true)); CHECK(!arch_supports_paged_attention("qwen35moe", false)); CHECK(arch_supports_draft_block_size("qwen35", false)); CHECK(!arch_supports_draft_block_size("qwen35", true)); @@ -739,7 +758,7 @@ TEST_CASE(FeatureGateFixture, feature_gate_suite) { test_feature_gate_ds4_decode_options_require_monolithic_hip(); test_feature_gate_remote_draft_requires_supported_arch(); test_feature_gate_layer_split_requires_supported_arch(); - test_feature_gate_paged_attention_requires_qwen35_monolithic(); + test_feature_gate_paged_attention_requires_monolithic_backend(); test_feature_gate_paged_attention_allows_fixed_local_chains(); test_feature_gate_parallel_and_kv_pool_rules(); test_feature_warnings_silent_when_supported(); diff --git a/server/test/test_gfx1151_width4_dispatch.cpp b/server/test/test_gfx1151_width4_dispatch.cpp new file mode 100644 index 000000000..a9d8d9e77 --- /dev/null +++ b/server/test/test_gfx1151_width4_dispatch.cpp @@ -0,0 +1,187 @@ +// gfx1151 paged dense dispatch regression test. +// +// Monolithic paged DeepSeek4 packs up to sixteen rows through plain mul_mat. +// Two dispatch policies preserve single-column arithmetic: +// +// 1. Narrow F16 weights (the hyper-connection mix projection, ne01 = 24) +// stay on MMVF for ne11 <= 8 instead of falling into the rocBLAS F16 +// Tensile kernel, whose only gfx1151 macro-tile is 128x128 and which +// accumulates in F16. +// 2. ROCmFP4 dense projections through sixteen columns use weight-reuse +// MMVQ instead of the MMQ tile. +// +// For every shape this test compares the N-column result bit for bit with N +// independent single-column computes. Any change that alters the per-column +// accumulation order fails here before it reaches a service measurement. +#include "ggml-backend.h" +#include "ggml-cuda.h" +#include "ggml.h" + +#include + +#include +#include +#include +#include +#include + +namespace { + +uint32_t lcg_state = 0x12345678u; +float lcg_uniform() { + lcg_state = lcg_state * 1664525u + 1013904223u; + return (float) ((lcg_state >> 8) & 0xFFFFFF) / (float) 0x1000000 * 2.0f - 1.0f; +} + +bool compute(ggml_backend_t backend, ggml_type wtype, int64_t k, int64_t m, int64_t n, + const std::vector & weight_bytes, const float * input, + std::vector & out) { + ggml_init_params params{}; + params.mem_size = 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + return false; + } + ggml_tensor * w = ggml_new_tensor_2d(ctx, wtype, k, m); + ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, k, n); + ggml_tensor * y = ggml_mul_mat(ctx, w, x); + ggml_set_input(w); + ggml_set_input(x); + ggml_set_output(y); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, y); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + ggml_backend_tensor_set(w, weight_bytes.data(), 0, weight_bytes.size()); + ggml_backend_tensor_set(x, input, 0, (size_t) k * n * sizeof(float)); + const bool ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + ggml_backend_synchronize(backend); + if (ok) { + out.resize((size_t) m * n); + ggml_backend_tensor_get(y, out.data(), 0, out.size() * sizeof(float)); + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} + +size_t bit_mismatches(const std::vector & a, const std::vector & b) { + if (a.size() != b.size()) { + return a.size() + b.size(); + } + size_t mm = 0; + for (size_t i = 0; i < a.size(); ++i) { + if (std::memcmp(&a[i], &b[i], sizeof(float)) != 0) { + ++mm; + } + } + return mm; +} + +// Compares one N-column compute against N single-column computes. +bool check_columns(ggml_backend_t backend, const char * label, ggml_type wtype, + int64_t k, int64_t m, int n, const std::vector & wbytes, + const std::vector & input) { + std::vector wide; + if (!compute(backend, wtype, k, m, n, wbytes, input.data(), wide)) { + std::printf("FAIL %s N=%d: compute failed\n", label, n); + return false; + } + std::vector joined; + for (int c = 0; c < n; ++c) { + std::vector one; + if (!compute(backend, wtype, k, m, 1, wbytes, input.data() + (size_t) k * c, one)) { + std::printf("FAIL %s N=%d: single-column compute failed\n", label, n); + return false; + } + joined.insert(joined.end(), one.begin(), one.end()); + } + const size_t mm = bit_mismatches(wide, joined); + std::printf("%s %s N=%d: %zu values, %zu bit mismatches vs single-column\n", + mm == 0 ? "PASS" : "FAIL", label, n, wide.size(), mm); + return mm == 0; +} + +bool test_narrow_f16(ggml_backend_t backend) { + const int64_t K = 16384, M = 24; + std::vector wbytes((size_t) K * M * sizeof(ggml_fp16_t)); + ggml_fp16_t * wh = (ggml_fp16_t *) wbytes.data(); + lcg_state = 0x1234567u; + for (size_t i = 0; i < (size_t) K * M; ++i) { + wh[i] = ggml_fp32_to_fp16(lcg_uniform() * 0.05f); + } + std::vector input((size_t) K * 8); + for (float & v : input) { + v = lcg_uniform() * 2.0f; + } + bool ok = true; + for (int n : {2, 3, 4, 5, 6, 8}) { + ok = check_columns(backend, "f16[16384,24]", GGML_TYPE_F16, K, M, n, wbytes, input) && ok; + } + return ok; +} + +bool test_rocmfp4_packed_columns(ggml_backend_t backend) { + const ggml_type T = GGML_TYPE_Q4_0_ROCMFP4_FAST; + const ggml_type_traits * tr = ggml_get_type_traits(T); + struct Shape { int64_t k, m; const char * label; }; + const Shape shapes[] = { + {4096, 1024, "rocmfp4[4096,1024]"}, // attn_q_a, attn_output_a, compressor width + {1024, 4096, "rocmfp4[1024,4096]"}, // attn_output_b + {4096, 512, "rocmfp4[4096,512]"}, // attn_kv + }; + bool ok = true; + for (const Shape & shape : shapes) { + const size_t row_bytes = ggml_row_size(T, shape.k); + std::vector wbytes(row_bytes * shape.m); + std::vector row(shape.k); + lcg_state = 0x9e3779b9u ^ (uint32_t) (shape.k * 31 + shape.m); + for (int64_t r = 0; r < shape.m; ++r) { + for (float & v : row) { + v = lcg_uniform() * 0.02f; + } + tr->from_float_ref(row.data(), wbytes.data() + row_bytes * r, shape.k); + } + std::vector input((size_t) shape.k * 16); + for (float & v : input) { + v = lcg_uniform() * 3.0f; + } + const int previous = ggml_backend_cuda_set_mmvq_max_ncols_override(16); + for (int n = 2; n <= 16; ++n) { + ok = check_columns(backend, shape.label, T, shape.k, shape.m, n, wbytes, input) && ok; + } + ggml_backend_cuda_set_mmvq_max_ncols_override(previous); + } + return ok; +} + +} // namespace + +int main() { + hipDeviceProp_t properties{}; + if (hipGetDeviceProperties(&properties, 0) != hipSuccess) { + std::fprintf(stderr, "failed to query HIP device 0\n"); + return 1; + } + if (std::strncmp(properties.gcnArchName, "gfx1151", 7) != 0) { + std::printf("SKIP: gfx1151 width-4 dispatch test expects gfx1151 (found %s)\n", + properties.gcnArchName); + return 77; + } + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, "failed to initialize HIP backend\n"); + return 1; + } + const bool previous_graphs = ggml_backend_cuda_set_graphs_disabled_override(true); + bool ok = test_narrow_f16(backend); + ok = test_rocmfp4_packed_columns(backend) && ok; + ggml_backend_cuda_set_graphs_disabled_override(previous_graphs); + ggml_backend_free(backend); + std::printf("%s\n", ok ? "PASS gfx1151 width-4 dispatch" : "FAIL gfx1151 width-4 dispatch"); + return ok ? 0 : 1; +} diff --git a/server/test/test_paged_kv_pool.cpp b/server/test/test_paged_kv_pool.cpp index f4a7e2f65..24b261b54 100644 --- a/server/test/test_paged_kv_pool.cpp +++ b/server/test/test_paged_kv_pool.cpp @@ -1,6 +1,5 @@ // Unit tests for PagedKvPool (paged_kv_pool.h). No ggml, no GPU. -#define GENERATE_UNIT_TEST_MAIN #include "CppUnitTestFramework.hpp" #include "common/concurrency/paged_kv_pool.h" #include "../src/common/paged_attention_config.h" diff --git a/server/test/test_rocmfp_mix_gateup_glu.cpp b/server/test/test_rocmfp_mix_gateup_glu.cpp index 8c66aa190..64f684af1 100644 --- a/server/test/test_rocmfp_mix_gateup_glu.cpp +++ b/server/test/test_rocmfp_mix_gateup_glu.cpp @@ -11,11 +11,11 @@ // // WHAT IS PROVEN HERE, and what is not: // -// * The two dot products are bit-identical to the unfused launches. That IS assertable and is -// the property the whole change rests on: FUSE_GLU is a template parameter over ONE -// accumulation body, so the fold order per row cannot drift between instantiations. Checked -// by running the unfused entry point and comparing against the fused result reconstructed -// through the inverse of the clamp-free branch (see `exact_when_unclamped`). +// * The two dot products are bit-identical to the unfused launches. That IS assertable: +// every GLU mode instantiates ONE accumulation body, and the two-pass finalizer reads the +// ordinary up result written by that same body. Checked by running the unfused entry point +// and comparing against the fused result reconstructed through the inverse of the +// clamp-free branch (see `exact_when_unclamped`). // // * The GLU value is compared to a HOST reference at a tight tolerance, not bit-exactly. The // kernel applies ggml_cuda_op_swiglu_ds4_single on device; the device and host expf() are not @@ -122,7 +122,9 @@ TEST_CASE(RocmfpMixGateupGluFixture, fused_gateup_glu) { // in must be a multiple of 128: the wide block load reads 128 weights at a time and would // read past the tensor on the final block (register_host enforces this). - const int in = 256, out = 64, n_experts = 6, n_used = 3, ntok = 2; + // On gfx1151, three tokens exercises the low-register two-pass GLU finalizer. q <= 2 + // uses the one-pass kernel and is checked separately below. + const int in = 256, out = 64, n_experts = 6, n_used = 3, ntok = 3; const int nb = in / QK; const size_t rows_bytes = (size_t) out * nb * BLOCK_BYTES; @@ -224,6 +226,41 @@ TEST_CASE(RocmfpMixGateupGluFixture, fused_gateup_glu) { std::fprintf(stderr, "worst relative deviation from the host reference: %.3e\n", worst); CHECK(worst < 1e-5); + // On gfx1151, q <= 2 selects the one-pass dual-projection kernel. Exercise that + // width separately while the main ntok=3 case above covers the two-pass finalizer. + { + const int one_tok = 1; + const size_t one_yn = (size_t) out * n_used; + CHECK(ggml_cuda_rocmfp2_mix_mul_mat_id(d_up, d_x, d_ids, d_up_out, + in, out, n_used, one_tok, 1, + ids_s0, ids_s1, src1_s1, src1_s2, + dst_s1, dst_s2, nullptr)); + CHECK(ggml_cuda_rocmfp2_mix_mul_mat_id(d_gate, d_x, d_ids, d_gate_out, + in, out, n_used, one_tok, 1, + ids_s0, ids_s1, src1_s1, src1_s2, + dst_s1, dst_s2, nullptr)); + CHECK(ggml_cuda_rocmfp2_mix_mul_mat_id_glu(d_up, d_gate, d_x, d_ids, d_fused, + in, out, n_used, one_tok, 1, + ids_s0, ids_s1, src1_s1, src1_s2, + dst_s1, dst_s2, limit, nullptr)); + HIP_OK(cudaDeviceSynchronize()); + std::vector one_u(one_yn), one_g(one_yn), one_f(one_yn); + HIP_OK(cudaMemcpy(one_u.data(), d_up_out, sizeof(float) * one_yn, + cudaMemcpyDeviceToHost)); + HIP_OK(cudaMemcpy(one_g.data(), d_gate_out, sizeof(float) * one_yn, + cudaMemcpyDeviceToHost)); + HIP_OK(cudaMemcpy(one_f.data(), d_fused, sizeof(float) * one_yn, + cudaMemcpyDeviceToHost)); + double one_worst = 0.0; + for (size_t i = 0; i < one_yn; ++i) { + const float ref = host_swiglu_ds4(one_g[i], one_u[i], limit); + const double denom = std::fmax(1e-6, std::fabs((double) ref)); + one_worst = std::fmax(one_worst, + std::fabs((double) one_f[i] - (double) ref) / denom); + } + CHECK(one_worst < 1e-5); + } + // ---- operand ORDER: swiglu_ds4 applies silu to GATE, so the two are not symmetric ---- float * d_swapped = nullptr; HIP_OK(cudaMalloc(&d_swapped, sizeof(float) * yn)); diff --git a/server/test/test_rocmfpx_mmq.cpp b/server/test/test_rocmfpx_mmq.cpp index fd2e92f8a..d8d5a6971 100644 --- a/server/test/test_rocmfpx_mmq.cpp +++ b/server/test/test_rocmfpx_mmq.cpp @@ -3,6 +3,7 @@ #include "ggml.h" #include "rocmfp4.h" #include "rocmfpx.h" +#include "../src/common/platform_env.h" #include @@ -68,7 +69,8 @@ static bool run_backend( const std::vector & weights_data, const std::vector & input_data, std::vector & output_data, - DispatchPath expected_path) { + DispatchPath expected_path, + int mmvq_ceiling = -1) { ggml_init_params params{}; params.mem_size = 16 * 1024 * 1024; params.no_alloc = true; @@ -113,6 +115,7 @@ static bool run_backend( const size_t mmq_before = ggml_backend_cuda_get_mmq_launch_count(); const int previous_mmvq_max = ggml_backend_cuda_set_mmvq_max_ncols_override( + mmvq_ceiling >= 0 ? mmvq_ceiling : expected_path == DispatchPath::MMVQ ? 8 : 1); const bool previous_graphs_disabled = ggml_backend_cuda_set_graphs_disabled_override(true); @@ -218,7 +221,8 @@ static bool test_case( ggml_backend_t hip_backend, const QuantCase & quant, const Shape & shape, - DispatchPath expected_path) { + DispatchPath expected_path, + int mmvq_ceiling = -1) { const std::vector weights_f32 = make_values((size_t) shape.k * shape.m, 37, 0.015625f); const std::vector input_f32 = @@ -248,7 +252,7 @@ static bool test_case( std::vector actual; if (!run_backend( hip_backend, quant.type, shape, weights_quantized, input_f32, actual, - expected_path)) { + expected_path, mmvq_ceiling)) { std::fprintf( stderr, "%s/%s: HIP %s run failed\n", @@ -261,9 +265,10 @@ static bool test_case( } int main() { - setenv("DFLASH_CUDA_MMVQ_FP2_AFFINE", "1", 1); - setenv("DFLASH_CUDA_MMQ_FP2_AFFINE", "1", 1); - setenv("DFLASH_CUDA_MMQ_FP2_AFFINE_GENERAL", "1", 1); + dflash::common::set_environment_variable("LUCE_MMVQ_MAX_NCOLS", "1", true); + dflash::common::set_environment_variable("DFLASH_CUDA_MMVQ_FP2_AFFINE", "1", true); + dflash::common::set_environment_variable("DFLASH_CUDA_MMQ_FP2_AFFINE", "1", true); + dflash::common::set_environment_variable("DFLASH_CUDA_MMQ_FP2_AFFINE_GENERAL", "1", true); hipDeviceProp_t properties{}; if (hipGetDeviceProperties(&properties, 0) != hipSuccess) { std::fprintf(stderr, "failed to query HIP device 0\n"); @@ -319,6 +324,16 @@ int main() { const DispatchPath expected_path = shape.n <= 8 ? DispatchPath::MMVQ : DispatchPath::MMQ; ok = test_case(hip_backend, quant, shape, expected_path) && ok; + if (quant.type == GGML_TYPE_Q4_0_ROCMFP4_FAST && shape.n == 16 && + std::strncmp(properties.gcnArchName, "gfx1151", 7) == 0) { + // Zero restores the process ceiling; positive overrides win. + for (const int ceiling : {0, 8, 16}) { + const DispatchPath path = ceiling == 16 ? DispatchPath::MMVQ : DispatchPath::MMQ; + std::printf("ROCmFP4 n=16 ceiling=%d (environment=1): expect %s\n", + ceiling, dispatch_path_name(path)); + ok = test_case(hip_backend, quant, shape, path, ceiling) && ok; + } + } } } if (shape_filter && !matched_shape) { diff --git a/server/test/test_rope_tail.cpp b/server/test/test_rope_tail.cpp new file mode 100644 index 000000000..d3f24e968 --- /dev/null +++ b/server/test/test_rope_tail.cpp @@ -0,0 +1,186 @@ +// GGML_ROPE_TYPE_TAIL regression test. +// +// DeepSeek4 rotates the last n_rot dims of each head. The tail mode rotates +// them in place and passes the head through, replacing the previous split, +// rotate, concat idiom (two contiguity copies, one rope, one concat). This +// test builds both forms on the CUDA/HIP backend and on the CPU backend and +// requires bit-identical results, for positive and negative positions and for +// the single-head 2D shape the compressor uses. +#include "ggml-backend.h" +#include "ggml-cpu.h" +#ifdef DFLASH_ROPE_TEST_GPU +#include "ggml-cuda.h" +#endif +#include "CppUnitTestFramework.hpp" +#include "ggml.h" + +#include +#include +#include +#include + +namespace { + +uint32_t lcg_state = 0x12345678u; +float lcg_uniform() { + lcg_state = lcg_state * 1664525u + 1013904223u; + return (float) ((lcg_state >> 8) & 0xFFFFFF) / (float) 0x1000000 * 2.0f - 1.0f; +} + +struct RopeParams { + int n_rot = 64; + int n_ctx_orig = 4096; + float freq_base = 10000.0f; + float freq_scale = 1.0f; + float ext_factor = 0.0f; + float attn_factor = 1.0f; + float beta_fast = 32.0f; + float beta_slow = 1.0f; +}; + +// Returns false on compute failure. Fills out with the rope result of shape +// [head_dim, n_heads, n_tokens], either through the old idiom or the tail mode. +bool run(ggml_backend_t backend, bool tail_mode, int head_dim, int n_heads, int n_tokens, + const std::vector & x, const std::vector & pos, const RopeParams & rp, + std::vector & out, bool use_factors = false) { + ggml_init_params params{}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + return false; + } + ggml_tensor * xt = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_dim, n_heads, n_tokens); + ggml_tensor * pt = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); + ggml_set_input(xt); + ggml_set_input(pt); + ggml_tensor * factors = use_factors ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, rp.n_rot / 2) : nullptr; + if (factors) ggml_set_input(factors); + ggml_tensor * y = nullptr; + if (tail_mode) { + y = ggml_rope_ext(ctx, xt, pt, factors, rp.n_rot, GGML_ROPE_TYPE_NORMAL | GGML_ROPE_TYPE_TAIL, + rp.n_ctx_orig, rp.freq_base, rp.freq_scale, rp.ext_factor, rp.attn_factor, + rp.beta_fast, rp.beta_slow); + } else { + const int n_nope = head_dim - rp.n_rot; + ggml_tensor * nope = ggml_view_3d(ctx, xt, n_nope, n_heads, n_tokens, xt->nb[1], xt->nb[2], 0); + ggml_tensor * tail = ggml_view_3d(ctx, xt, rp.n_rot, n_heads, n_tokens, xt->nb[1], xt->nb[2], + (size_t) n_nope * sizeof(float)); + tail = ggml_cont(ctx, tail); + tail = ggml_rope_ext(ctx, tail, pt, factors, rp.n_rot, GGML_ROPE_TYPE_NORMAL, rp.n_ctx_orig, + rp.freq_base, rp.freq_scale, rp.ext_factor, rp.attn_factor, rp.beta_fast, + rp.beta_slow); + y = ggml_concat(ctx, ggml_cont(ctx, nope), tail, 0); + } + ggml_set_output(y); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, y); + std::vector factor_values(use_factors ? rp.n_rot / 2 : 0); + for (size_t i = 0; i < factor_values.size(); ++i) factor_values[i] = 1.0f + 0.125f * i; + ggml_backend_buffer_t factor_buf = nullptr; + if (factors && ggml_backend_is_cpu(backend)) { + // A separate allocation lets ASan catch reads beyond the n_rot/2 factors. + factor_buf = ggml_backend_cpu_buffer_from_ptr(factor_values.data(), factor_values.size() * sizeof(float)); + ggml_backend_tensor_alloc(factor_buf, factors, factor_values.data()); + } + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_backend_buffer_free(factor_buf); + ggml_free(ctx); + return false; + } + ggml_backend_tensor_set(xt, x.data(), 0, x.size() * sizeof(float)); + ggml_backend_tensor_set(pt, pos.data(), 0, pos.size() * sizeof(int32_t)); + if (factors) { + ggml_backend_tensor_set(factors, factor_values.data(), 0, factor_values.size() * sizeof(float)); + } + const bool ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + ggml_backend_synchronize(backend); + if (ok) { + out.resize((size_t) head_dim * n_heads * n_tokens); + ggml_backend_tensor_get(y, out.data(), 0, out.size() * sizeof(float)); + } + ggml_backend_buffer_free(buf); + ggml_backend_buffer_free(factor_buf); + ggml_free(ctx); + return ok; +} + +size_t bit_mismatches(const std::vector & a, const std::vector & b) { + if (a.size() != b.size()) { + return a.size() + b.size(); + } + size_t mm = 0; + for (size_t i = 0; i < a.size(); ++i) { + if (std::memcmp(&a[i], &b[i], sizeof(float)) != 0) { + ++mm; + } + } + return mm; +} + +bool check(ggml_backend_t backend, const char * label, int head_dim, int n_heads, int n_tokens, + const std::vector & pos, const RopeParams & rp, bool use_factors = false) { + std::vector x((size_t) head_dim * n_heads * n_tokens); + lcg_state = 0x51a7u ^ (uint32_t) (head_dim * 7 + n_heads * 13 + n_tokens); + for (float & v : x) { + v = lcg_uniform() * 3.0f; + } + std::vector old_out, tail_out; + if (!run(backend, false, head_dim, n_heads, n_tokens, x, pos, rp, old_out, use_factors) || + !run(backend, true, head_dim, n_heads, n_tokens, x, pos, rp, tail_out, use_factors)) { + std::printf("FAIL %s: compute failed\n", label); + return false; + } + const size_t mm = bit_mismatches(old_out, tail_out); + // The unrotated head must be a verbatim copy. + size_t head_changed = 0; + const int n_nope = head_dim - rp.n_rot; + for (int t = 0; t < n_tokens; ++t) { + for (int h = 0; h < n_heads; ++h) { + const size_t base = ((size_t) t * n_heads + h) * head_dim; + for (int d = 0; d < head_dim; ++d) { + const bool same = std::memcmp(&tail_out[base + d], &x[base + d], sizeof(float)) == 0; + if (d < n_nope && !same) ++head_changed; + } + } + } + std::printf("%s %s: %zu values, %zu bit mismatches vs split-rotate-concat, head changed %zu\n", + mm == 0 && head_changed == 0 ? "PASS" : "FAIL", label, tail_out.size(), mm, head_changed); + return mm == 0 && head_changed == 0; +} + +} // namespace + +struct RopeTailFixture : CppUnitTestFramework::CommonFixture { + using CommonFixture::CommonFixture; + ggml_backend_t backend = nullptr; + ~RopeTailFixture() { if (backend) ggml_backend_free(backend); } + + void compare() { + REQUIRE_NOT_NULL(backend); + const RopeParams rp; + REQUIRE_TRUE(check(backend, "positive", 512, 64, 4, {5, 100, 1000, 3500}, rp)); + REQUIRE_TRUE(check(backend, "inverse", 512, 64, 4, {-5, -100, -1000, -3500}, rp)); + REQUIRE_TRUE(check(backend, "single head", 512, 1, 1, {127}, rp)); + REQUIRE_TRUE(check(backend, "frequency factors", 512, 8, 4, {5, 100, 1000, 3500}, rp, true)); + RopeParams yarn = rp; + yarn.freq_scale = 0.25f; + yarn.ext_factor = 1.0f; + yarn.attn_factor = 1.1f; + REQUIRE_TRUE(check(backend, "yarn", 512, 8, 4, {5, 100, 1000, 3500}, yarn)); + } +}; + +TEST_CASE(RopeTailFixture, cpu_matches_split_rotate_concat) { + backend = ggml_backend_cpu_init(); + compare(); +} + +#ifdef DFLASH_ROPE_TEST_GPU +TEST_CASE(RopeTailFixture, gpu_matches_split_rotate_concat) { + if (ggml_backend_cuda_get_device_count() == 0) SKIP("no CUDA/HIP device available"); + backend = ggml_backend_cuda_init(0); + compare(); +} +#endif diff --git a/server/test/test_seq_slot_manager.cpp b/server/test/test_seq_slot_manager.cpp index cdc84d2f9..f7bfe6e67 100644 --- a/server/test/test_seq_slot_manager.cpp +++ b/server/test/test_seq_slot_manager.cpp @@ -39,6 +39,21 @@ static bool is_busy(const SeqEngine::AdmitResult & result) { } int main() { + // Destroying the manager retires any requests still active in its pool. + { + PagedKvPool pool(/*physical_block_count=*/4, /*max_sequences=*/2, + /*block_size=*/16); + { + Qwen35SlotManager mgr(pool, /*max_ctx=*/64); + auto active = admit(mgr, 1, prompt_tokens(20), greedy_sampler()); + CHECK(is_admitted(active)); + CHECK(pool.active_sequence_count() == 1); + CHECK(pool.free_block_count() == 1); + } + CHECK(pool.active_sequence_count() == 0); + CHECK(pool.free_block_count() == 4); + } + // 8 blocks x 16 tokens = 128 pool tokens, 2 slots, per-seq max_ctx 64. { PagedKvPool pool(/*physical_block_count=*/8, /*max_sequences=*/2, @@ -483,6 +498,17 @@ int main() { auto c = admit(mgr, 3, prompt_tokens(4), cfg); CHECK(is_admitted(c)); CHECK(mgr.slot(0).rng() != first); + + // Zero is a valid deterministic seed, not a request for entropy. + mgr.retire(0); + cfg.seed = 0; + auto zero_a = admit(mgr, 4, prompt_tokens(4), cfg); + CHECK(is_admitted(zero_a)); + const uint64_t zero_first = mgr.slot(0).rng(); + mgr.retire(0); + auto zero_b = admit(mgr, 5, prompt_tokens(4), cfg); + CHECK(is_admitted(zero_b)); + CHECK(mgr.slot(0).rng() == zero_first); } // A greedy sampler never draws, so its seed is irrelevant and admission diff --git a/server/test/test_softmax_sink_col.cpp b/server/test/test_softmax_sink_col.cpp new file mode 100644 index 000000000..2b3eaa7a1 --- /dev/null +++ b/server/test/test_softmax_sink_col.cpp @@ -0,0 +1,145 @@ +// ggml_soft_max_ext_sink_col regression test. +// +// DeepSeek4 attention adds one learned per-head sink logit to the softmax +// denominator. The fused op treats the sink as a virtual last column with the +// query scale folded in, replacing scale -> concat(sink) -> soft_max -> view. +// This test compares both forms bit for bit on the HIP and CPU backends over +// column counts across block-size steps and the shared-memory fallback boundary. +#include "ggml-backend.h" +#include "ggml-cpu.h" +#include "ggml-cuda.h" +#include "ggml.h" + +#include + +#include +#include +#include +#include + +namespace { + +uint32_t lcg_state = 0x12345678u; +float lcg_uniform() { + lcg_state = lcg_state * 1664525u + 1013904223u; + return (float) ((lcg_state >> 8) & 0xFFFFFF) / (float) 0x1000000 * 2.0f - 1.0f; +} + +bool run(ggml_backend_t backend, bool fused, int ncols, int rows, float scale, + const std::vector & scores, const std::vector & sinks, std::vector & out) { + ggml_init_params params{}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) { + return false; + } + ggml_tensor * st = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ncols, rows); + ggml_tensor * kt = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, rows); + ggml_set_input(st); + ggml_set_input(kt); + ggml_tensor * probs = nullptr; + if (fused) { + probs = ggml_soft_max_ext_sink_col(ctx, st, kt, scale); + } else { + ggml_tensor * scaled = ggml_scale(ctx, st, scale); + ggml_tensor * sink2d = ggml_reshape_2d(ctx, kt, 1, rows); + ggml_tensor * with_sink = ggml_concat(ctx, scaled, sink2d, 0); + ggml_tensor * full = ggml_soft_max(ctx, with_sink); + // The old path views the first ncols columns; materialize the view so + // both forms hand back a contiguous [ncols, rows] tensor. + probs = ggml_cont(ctx, ggml_view_2d(ctx, full, ncols, rows, full->nb[1], 0)); + } + ggml_set_output(probs); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, probs); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + return false; + } + ggml_backend_tensor_set(st, scores.data(), 0, scores.size() * sizeof(float)); + ggml_backend_tensor_set(kt, sinks.data(), 0, sinks.size() * sizeof(float)); + const bool ok = ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + ggml_backend_synchronize(backend); + if (ok) { + out.resize((size_t) ncols * rows); + ggml_backend_tensor_get(probs, out.data(), 0, out.size() * sizeof(float)); + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return ok; +} + +size_t bit_mismatches(const std::vector & a, const std::vector & b) { + if (a.size() != b.size()) { + return a.size() + b.size(); + } + size_t mm = 0; + for (size_t i = 0; i < a.size(); ++i) { + if (std::memcmp(&a[i], &b[i], sizeof(float)) != 0) { + ++mm; + } + } + return mm; +} + +} // namespace + +int main() { + hipDeviceProp_t properties{}; + if (hipGetDeviceProperties(&properties, 0) != hipSuccess) { + std::fprintf(stderr, "failed to query HIP device 0\n"); + return 1; + } + ggml_backend_t hip = ggml_backend_cuda_init(0); + ggml_backend_t cpu = ggml_backend_cpu_init(); + if (!hip || !cpu) { + std::fprintf(stderr, "failed to initialize backends\n"); + return 1; + } + const bool previous_graphs = ggml_backend_cuda_set_graphs_disabled_override(true); + const int rows = 64; // n_head + const float scale = 1.0f / 22.627417f; // 1/sqrt(512) + bool ok = true; + std::vector widths = {1, 20, 31, 32, 33, 63, 64, 65, 100, 127, 128, 129, 200, 255, 256, 257, 511, 512, 513, 700, 1023, 1024, 1025, 1500, 2047, 2048, 3000}; + // Shared scratch includes one warp of reduction storage and the virtual + // sink column. Cross its capacity to exercise output-backed scratch, where + // caching the virtual column would overwrite the next row (or the buffer). + const int shared_cols = (int) (properties.sharedMemPerBlock / sizeof(float)) - properties.warpSize; + widths.insert(widths.end(), {shared_cols - 1, shared_cols, shared_cols + 1, 2 * shared_cols + 1}); + for (int ncols : widths) { + std::vector scores((size_t) ncols * rows), sinks(rows); + lcg_state = 0xabc123u ^ (uint32_t) ncols; + for (float & v : scores) { + v = lcg_uniform() * 40.0f; // pre-scale logits of realistic magnitude + } + for (float & v : sinks) { + v = lcg_uniform() * 3.0f; + } + // Include sink-dominated and negligible-sink rows so large widths still + // expose an omitted sink or an incorrect per-row sink index. + sinks[0] = 80.0f; + sinks[1] = -80.0f; + sinks[2] = 0.0f; + for (auto backend : {hip, cpu}) { + const char * name = backend == hip ? "hip" : "cpu"; + std::vector ref, fused; + if (!run(backend, false, ncols, rows, scale, scores, sinks, ref) || + !run(backend, true, ncols, rows, scale, scores, sinks, fused)) { + std::printf("FAIL %s ncols=%d: compute failed\n", name, ncols); + ok = false; + continue; + } + const size_t mm = bit_mismatches(ref, fused); + std::printf("%s %s ncols=%d: %zu values, %zu bit mismatches vs scale-concat-softmax\n", + mm == 0 ? "PASS" : "FAIL", name, ncols, ref.size(), mm); + ok = ok && mm == 0; + } + } + ggml_backend_cuda_set_graphs_disabled_override(previous_graphs); + ggml_backend_free(hip); + ggml_backend_free(cpu); + std::printf("%s\n", ok ? "PASS softmax sink column" : "FAIL softmax sink column"); + return ok ? 0 : 1; +} diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 0d3d8d361..0933b4e6a 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -688,6 +688,152 @@ static void test_moe_routing_correctness(ggml_backend_t backend) { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_top6_head4_tail2_route_slices(ggml_backend_t backend) { + std::fprintf(stderr, " test_top6_head4_tail2_route_slices ..."); + + constexpr int route_width = 6; + constexpr int n_tokens = 4; + std::vector selected_values(route_width * n_tokens); + std::vector weight_values(route_width * n_tokens); + std::vector expected_head_ids; + std::vector expected_tail_ids; + std::vector expected_head_weights; + std::vector expected_tail_weights; + for (int token = 0; token < n_tokens; ++token) { + for (int route = 0; route < route_width; ++route) { + const size_t index = (size_t) token * route_width + route; + selected_values[index] = token * 10 + route; + weight_values[index] = (float) (token * 10 + route) + 0.25f; + if (route < 4) { + expected_head_ids.push_back(selected_values[index]); + expected_head_weights.push_back(weight_values[index]); + } else { + expected_tail_ids.push_back(selected_values[index]); + expected_tail_weights.push_back(weight_values[index]); + } + } + } + + ggml_context * ctx = make_test_context(); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) { + std::fprintf(stderr, " FAIL\n"); + return; + } + + ggml_tensor * selected = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, route_width, n_tokens); + ggml_tensor * weights = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, route_width, n_tokens); + ggml_set_input(selected); + ggml_set_input(weights); + + DeepSeek4Head4Tail2Routes routes; + TEST_ASSERT(build_deepseek4_head4_tail2_routes( + ctx, selected, weights, n_tokens, routes)); + TEST_ASSERT(routes.head_ids && routes.head_weights && + routes.tail_ids && routes.tail_weights); + if (!routes.head_ids || !routes.head_weights || + !routes.tail_ids || !routes.tail_weights) { + ggml_free(ctx); + std::fprintf(stderr, " FAIL\n"); + return; + } + + TEST_ASSERT(routes.head_ids->ne[0] == 4); + TEST_ASSERT(routes.head_ids->ne[1] == n_tokens); + TEST_ASSERT(routes.head_weights->ne[0] == 4); + TEST_ASSERT(routes.head_weights->ne[1] == n_tokens); + TEST_ASSERT(routes.tail_ids->ne[0] == 2); + TEST_ASSERT(routes.tail_ids->ne[1] == n_tokens); + TEST_ASSERT(routes.tail_weights->ne[0] == 2); + TEST_ASSERT(routes.tail_weights->ne[1] == n_tokens); + + const ggml_tensor * outputs[] = { + routes.head_ids, + routes.head_weights, + routes.tail_ids, + routes.tail_weights, + }; + const ggml_tensor * sources[] = { + selected, + weights, + selected, + weights, + }; + for (size_t i = 0; i < 4; ++i) { + TEST_ASSERT(outputs[i]->op == GGML_OP_CONT); + TEST_ASSERT(outputs[i]->src[0] != nullptr); + if (outputs[i]->src[0]) { + TEST_ASSERT(outputs[i]->src[0]->view_src == sources[i]); + TEST_ASSERT(outputs[i]->src[0]->nb[1] == sources[i]->nb[1]); + } + ggml_set_output(const_cast(outputs[i])); + } + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 32, false); + for (ggml_tensor * output : { + routes.head_ids, + routes.head_weights, + routes.tail_ids, + routes.tail_weights}) { + ggml_build_forward_expand(graph, output); + } + int cont_nodes = 0; + int concat_nodes = 0; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + const ggml_tensor * node = ggml_graph_node(graph, i); + cont_nodes += node->op == GGML_OP_CONT; + concat_nodes += node->op == GGML_OP_CONCAT; + } + TEST_ASSERT(cont_nodes == 4); + TEST_ASSERT(concat_nodes == 0); + + ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_cpu_buffer_type()); + TEST_ASSERT(ggml_gallocr_alloc_graph(alloc, graph)); + ggml_backend_tensor_set(selected, selected_values.data(), 0, + selected_values.size() * sizeof(int32_t)); + ggml_backend_tensor_set(weights, weight_values.data(), 0, + weight_values.size() * sizeof(float)); + TEST_ASSERT(ggml_backend_graph_compute(backend, graph) == + GGML_STATUS_SUCCESS); + + std::vector actual_head_ids(expected_head_ids.size()); + std::vector actual_tail_ids(expected_tail_ids.size()); + std::vector actual_head_weights(expected_head_weights.size()); + std::vector actual_tail_weights(expected_tail_weights.size()); + ggml_backend_tensor_get(routes.head_ids, actual_head_ids.data(), 0, + actual_head_ids.size() * sizeof(int32_t)); + ggml_backend_tensor_get(routes.tail_ids, actual_tail_ids.data(), 0, + actual_tail_ids.size() * sizeof(int32_t)); + ggml_backend_tensor_get(routes.head_weights, + actual_head_weights.data(), 0, + actual_head_weights.size() * sizeof(float)); + ggml_backend_tensor_get(routes.tail_weights, + actual_tail_weights.data(), 0, + actual_tail_weights.size() * sizeof(float)); + TEST_ASSERT(actual_head_ids == expected_head_ids); + TEST_ASSERT(actual_tail_ids == expected_tail_ids); + TEST_ASSERT(actual_head_weights == expected_head_weights); + TEST_ASSERT(actual_tail_weights == expected_tail_weights); + + DeepSeek4Head4Tail2Routes invalid; + ggml_tensor * five_routes = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, 5, n_tokens); + ggml_tensor * integer_weights = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, route_width, n_tokens); + TEST_ASSERT(!build_deepseek4_head4_tail2_routes( + ctx, five_routes, weights, n_tokens, invalid)); + TEST_ASSERT(!build_deepseek4_head4_tail2_routes( + ctx, selected, integer_weights, n_tokens, invalid)); + TEST_ASSERT(!build_deepseek4_head4_tail2_routes( + ctx, selected, weights, 0, invalid)); + + ggml_gallocr_free(alloc); + ggml_free(ctx); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_rmsnorm_correctness(ggml_backend_t backend) { std::fprintf(stderr, " test_rmsnorm_correctness ..."); @@ -4114,6 +4260,94 @@ static void test_hc_scratch_shape_capacity() { #endif #endif +static void test_paged_lone_prompt_limits() { + DeepSeek4Backend backend{DeepSeek4BackendConfig{}}; + PagedKvPool pool(32, 6, 16); + DeepSeek4SeqEngine engine(backend, pool, 256, 16); + TEST_ASSERT(engine.step_plan_limits(0).max_prefill_tokens_per_sequence == 4); + + const auto first = engine.admit(1, std::vector(32, 1), SamplerCfg{}); + TEST_ASSERT(first.status == SeqEngine::AdmitResult::Status::admitted); + auto single = plan_prefill_slices({{first.slot, 0}}, engine.step_plan_limits(0)); + TEST_ASSERT(single.size() == 1 && single[0].max_tokens == 16); + + const auto second = engine.admit(2, std::vector(64, 1), SamplerCfg{}); + TEST_ASSERT(second.status == SeqEngine::AdmitResult::Status::admitted); + auto shared = plan_prefill_slices( + {{first.slot, 0}, {second.slot, 1}}, engine.step_plan_limits(0)); + TEST_ASSERT(shared.size() == 2 && shared[0].max_tokens == 4 && + shared[1].max_tokens == 4); + + TEST_ASSERT(engine.slots_.append_prefill(first.slot, 32).ok); + engine.slots_.commit_prefill(first.slot); + auto mixed = plan_prefill_slices({{second.slot, 1}}, engine.step_plan_limits(1)); + TEST_ASSERT(mixed.size() == 1 && mixed[0].max_tokens == 4); + + engine.retire(first.slot); + single = plan_prefill_slices({{second.slot, 1}}, engine.step_plan_limits(0)); + TEST_ASSERT(single.size() == 1 && single[0].max_tokens == 16); + + backend.moe_hybrid_ = std::make_unique(); + single = plan_prefill_slices({{second.slot, 1}}, engine.step_plan_limits(0)); + TEST_ASSERT(single.size() == 1 && single[0].max_tokens == 1); +} + +static void test_paged_cache_allocation(ggml_backend_t backend) { + std::fprintf(stderr, " paged cache allocation and slot reset..."); + DeepSeek4Weights weights; + weights.n_layer = 3; + weights.head_dim = 16; + weights.n_indexer_head_dim = 8; + weights.compress_ratios = {0, 4, 128}; + DeepSeek4PagedCache cache; + for (uint32_t slots : {1u, 3u}) { + if (!create_deepseek4_paged_cache(backend, weights, slots, 257, 5, cache)) { + TEST_ASSERT_MSG(false, "paged cache creation failed"); + return; + } + uint64_t raw_bytes = 0, compressed_bytes = 0, state_bytes = 0; + for (const auto & layer : cache.layers) { + raw_bytes += ggml_nbytes(layer.raw_kv); + for (const auto * tensor : {layer.comp_kv, layer.index_comp_kv}) { + if (tensor) compressed_bytes += ggml_nbytes(tensor); + } + for (auto * tensor : {layer.attn_compressor.state_kv, + layer.attn_compressor.state_score, + layer.indexer_compressor.state_kv, + layer.indexer_compressor.state_score}) { + if (!tensor) continue; + state_bytes += ggml_nbytes(tensor); + std::vector values(ggml_nelements(tensor), 1.0f); + ggml_backend_tensor_set(tensor, values.data(), 0, ggml_nbytes(tensor)); + } + } + TEST_ASSERT(raw_bytes == cache.plan.raw_bytes); + TEST_ASSERT(compressed_bytes == cache.plan.compressed_bytes); + TEST_ASSERT(state_bytes == cache.plan.state_bytes); + TEST_ASSERT(cache.layers[0].comp_kv == nullptr); + TEST_ASSERT(cache.layers[0].attn_compressor.state_kv == nullptr); + TEST_ASSERT(cache.layers[2].index_comp_kv == nullptr); + reset_deepseek4_paged_slot(cache, slots - 1); + for (const auto & layer : cache.layers) { + for (auto * tensor : {layer.attn_compressor.state_kv, + layer.attn_compressor.state_score, + layer.indexer_compressor.state_kv, + layer.indexer_compressor.state_score}) { + if (!tensor) continue; + std::vector values(ggml_nelements(tensor)); + ggml_backend_tensor_get(tensor, values.data(), 0, ggml_nbytes(tensor)); + const size_t slot_elements = tensor->nb[2] / sizeof(float); + for (size_t i = 0; i < values.size(); ++i) { + TEST_ASSERT(values[i] == (i / slot_elements == slots - 1 ? 0.0f : 1.0f)); + } + } + } + free_deepseek4_paged_cache(cache); + TEST_ASSERT(!cache.ctx && !cache.buf && !cache.pool && cache.layers.empty()); + } + std::fprintf(stderr, " done\n"); +} + int main() { ggml_backend_t backend = ggml_backend_cpu_init(); if (!backend) { @@ -4121,11 +4355,14 @@ int main() { return 1; } + test_paged_lone_prompt_limits(); + test_paged_cache_allocation(backend); test_compressor_pooling_correctness(backend); test_moe_expert_major_default_threshold(); test_chunked_graph_allocator(backend); test_swiglu_ds4_cpu_correctness(backend); test_moe_routing_correctness(backend); + test_top6_head4_tail2_route_slices(backend); test_rmsnorm_correctness(backend); test_grouped_output_projection_shape(); test_grouped_output_projection_cpu(backend);