From 248e01070178b06a8c94af3bb775c5d56b27fa9e Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 10 Jul 2026 18:59:55 +0200 Subject: [PATCH 01/15] feat(ipc): add DeepSeek4 DSpark shared worker --- server/CMakeLists.txt | 1 + server/docs/DS4.md | 32 ++ server/src/common/backend_ipc.cpp | 5 + server/src/common/backend_ipc.h | 1 + server/src/common/dflash_draft_ipc.cpp | 142 +++++-- server/src/common/dflash_draft_ipc.h | 19 +- server/src/deepseek4/deepseek4_backend.cpp | 41 +- server/src/deepseek4/deepseek4_backend.h | 3 + server/src/deepseek4/deepseek4_dspark.h | 13 +- .../deepseek4_dspark_draft_ipc_daemon.cpp | 354 ++++++++++++++++++ .../src/deepseek4/deepseek4_dspark_spec.cpp | 44 ++- server/src/ipc/backend_ipc_main.cpp | 9 + server/test/test_server_unit.cpp | 5 + 13 files changed, 636 insertions(+), 33 deletions(-) create mode 100644 server/src/deepseek4/deepseek4_dspark_draft_ipc_daemon.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index b2d5f833f..581fea034 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -260,6 +260,7 @@ add_library(dflash_common STATIC src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp src/deepseek4/deepseek4_dspark.cpp src/deepseek4/deepseek4_dspark_spec.cpp + src/deepseek4/deepseek4_dspark_draft_ipc_daemon.cpp src/flashprefill_q8.cpp src/kv_cache.cpp src/kv_quant.cpp diff --git a/server/docs/DS4.md b/server/docs/DS4.md index aae2b62b1..25c6769ee 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -3,6 +3,8 @@ This document describes the current DeepSeek V4 Flash implementation in DFlash. DeepSeek4 supports a monolithic HIP backend for single-device Strix Halo systems and a layer-split backend for local or mixed-device deployments. +The opt-in DSpark speculative path can execute proposal blocks in a separate +HIP process while the target remains in the primary process. ## Model Architecture @@ -30,6 +32,7 @@ DeepSeek V4 Flash is a 43-layer MoE model with: | Model weights and metadata | `src/deepseek4/deepseek4_internal.h`, `src/deepseek4/deepseek4_loader.cpp` | | HC pre/post CUDA kernel | `src/deepseek4/deepseek4_hc_cuda.cu`, `.h` | | Remote target-shard daemon | `src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp` | +| DSpark runtime and remote draft daemon | `src/deepseek4/deepseek4_dspark*.{h,cpp}` | | Shared target-shard IPC infrastructure | `src/common/target_shard_ipc.*`, `src/placement/remote_target_shard_config.h` | | Backend IPC CLI entry | `src/ipc/backend_ipc_main.cpp` | @@ -103,6 +106,28 @@ For heterogeneous setups, the CUDA-built server can keep the prefix layers on th This path uses `TargetShardIpcSession`, `deepseek4_target_shard_ipc_daemon.cpp`, and `BackendIpcMode::DeepSeek4TargetShard` rather than the old expert-worker protocol. +### Local target + remote HIP DSpark draft + +The speculative IPC mode keeps the complete DeepSeek4 target, KV cache, tied +embedding, DSpark head, and sampler in the parent. A separate +`BackendIpcMode::DeepSeek4DSparkDraft` process executes the DSpark proposal +graph on the selected HIP GPU. The parent currently retains its locally loaded +draft weights for metadata and fallback; removing that duplicate residency is +a separate optimization. + +Each speculative step transfers: + +- target feature captures in token-major + `[n_tokens, n_target_layers, n_embd]` layout; +- the embedded seed/MASK proposal block; +- the proposal hidden states returned by the HIP worker. + +All target-layer captures are uploaded as one feature block and one +acknowledgement. On POSIX hosts this mode defaults to the existing memfd-backed +shared payload transport; proposal input and output reuse the same mapping. +`stream` remains available as an A/B and compatibility fallback. This is shared +host IPC, not direct `hipIpcMemHandle` peer memory. + ## Shard Boundary State The shard boundary transfers the **full HC state tensor**, not a separate expert-routing payload: @@ -130,6 +155,13 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. |----------|---------| | `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_SPEC` | Enable the DeepSeek4 DSpark speculative runtime. | +| `DFLASH_DS4_DRAFT` | DSpark draft GGUF loaded locally for metadata/head use and by the remote worker. | +| `DFLASH_DS4_DRAFT_IPC_BIN` | Backend IPC daemon executable used for the remote DSpark worker. | +| `DFLASH_DS4_DRAFT_IPC_GPU` | HIP device index for the remote DSpark proposal blocks. | +| `DFLASH_DS4_DRAFT_IPC_WORK_DIR` | Scratch directory for the child process. | +| `DFLASH_DS4_DRAFT_IPC_REQUIRED` | Fail initialization instead of falling back to a local draft when remote startup fails. | +| `DFLASH_DRAFT_IPC_TRANSPORT` | Payload transport: `auto`, `shared`, or `stream`. DeepSeek4 DSpark defaults to `auto`; other draft modes retain their existing default. | `DFLASH_DS4_TIMING` enables the existing timing banners: diff --git a/server/src/common/backend_ipc.cpp b/server/src/common/backend_ipc.cpp index 4c99daae5..c52dd591d 100644 --- a/server/src/common/backend_ipc.cpp +++ b/server/src/common/backend_ipc.cpp @@ -34,6 +34,7 @@ const char * backend_ipc_mode_name(BackendIpcMode mode) { case BackendIpcMode::LagunaTargetShard: return "laguna-target-shard"; case BackendIpcMode::MoeExpertCompute: return "moe-expert-compute"; case BackendIpcMode::DeepSeek4TargetShard: return "deepseek4-target-shard"; + case BackendIpcMode::DeepSeek4DSparkDraft: return "deepseek4-dspark-draft"; } return "unknown"; } @@ -67,6 +68,10 @@ bool parse_backend_ipc_mode(const std::string & value, BackendIpcMode & out) { out = BackendIpcMode::DeepSeek4TargetShard; return true; } + if (value == "deepseek4-dspark-draft") { + out = BackendIpcMode::DeepSeek4DSparkDraft; + return true; + } return false; } diff --git a/server/src/common/backend_ipc.h b/server/src/common/backend_ipc.h index e9ce21b3c..47a8a5c08 100644 --- a/server/src/common/backend_ipc.h +++ b/server/src/common/backend_ipc.h @@ -29,6 +29,7 @@ enum class BackendIpcMode { LagunaTargetShard, MoeExpertCompute, DeepSeek4TargetShard, + DeepSeek4DSparkDraft, }; const char * backend_ipc_mode_name(BackendIpcMode mode); diff --git a/server/src/common/dflash_draft_ipc.cpp b/server/src/common/dflash_draft_ipc.cpp index 387be58b1..ff6cf4f66 100644 --- a/server/src/common/dflash_draft_ipc.cpp +++ b/server/src/common/dflash_draft_ipc.cpp @@ -25,14 +25,15 @@ namespace dflash::common { namespace { -BackendIpcPayloadTransport draft_ipc_transport_from_env() { +BackendIpcPayloadTransport draft_ipc_transport_from_env( + BackendIpcPayloadTransport fallback) { const char * raw = std::getenv("DFLASH_DRAFT_IPC_TRANSPORT"); if (!raw || !*raw) { - return BackendIpcPayloadTransport::Stream; + return fallback; } - BackendIpcPayloadTransport transport = BackendIpcPayloadTransport::Stream; + BackendIpcPayloadTransport transport = fallback; if (!parse_backend_ipc_payload_transport(raw, transport)) { - return BackendIpcPayloadTransport::Stream; + return fallback; } return transport; } @@ -47,16 +48,23 @@ bool checked_mul_size(size_t a, size_t b, size_t & out) { size_t dflash_draft_ipc_required_shared_bytes(int hidden_size, int block_size, - int ring_cap) { - if (hidden_size <= 0 || block_size <= 0 || ring_cap <= 0) { + int ring_cap, + int n_target_layers) { + if (hidden_size <= 0 || block_size <= 0 || ring_cap <= 0 || + n_target_layers <= 0) { return 0; } - const size_t max_tokens = - (size_t)std::max(block_size, ring_cap); - size_t elements = 0; + size_t noise_elements = 0; + size_t feature_elements = 0; size_t bytes = 0; - if (!checked_mul_size(max_tokens, (size_t)hidden_size, elements) || - !checked_mul_size(elements, sizeof(float), bytes)) { + if (!checked_mul_size((size_t)block_size, (size_t)hidden_size, + noise_elements) || + !checked_mul_size((size_t)ring_cap, (size_t)n_target_layers, + feature_elements) || + !checked_mul_size(feature_elements, (size_t)hidden_size, + feature_elements) || + !checked_mul_size(std::max(noise_elements, feature_elements), + sizeof(float), bytes)) { return 0; } return bytes; @@ -100,22 +108,32 @@ bool DFlashDraftIpcClient::start( const std::string & draft_path, int draft_gpu, int ring_cap, - const std::string & work_dir) { + const std::string & work_dir, + BackendIpcMode mode) { #if defined(_WIN32) - (void)bin; (void)draft_path; (void)draft_gpu; (void)ring_cap; (void)work_dir; + (void)bin; (void)draft_path; (void)draft_gpu; (void)ring_cap; (void)work_dir; (void)mode; std::fprintf(stderr, "DFlash draft IPC is only implemented on POSIX hosts\n"); return false; #else close(); - if (bin.empty() || draft_path.empty() || ring_cap <= 0) return false; + if (bin.empty() || draft_path.empty() || ring_cap <= 0 || + (mode != BackendIpcMode::DFlashDraft && + mode != BackendIpcMode::DeepSeek4DSparkDraft)) return false; BackendIpcLaunchConfig launch; launch.bin = bin; - launch.mode = BackendIpcMode::DFlashDraft; + launch.mode = mode; launch.payload_path = draft_path; launch.work_dir = work_dir; - launch.payload_transport = draft_ipc_transport_from_env(); + const BackendIpcPayloadTransport default_transport = + mode == BackendIpcMode::DeepSeek4DSparkDraft + ? BackendIpcPayloadTransport::Auto + : BackendIpcPayloadTransport::Stream; + launch.payload_transport = draft_ipc_transport_from_env(default_transport); + const int shared_feature_layers = + mode == BackendIpcMode::DeepSeek4DSparkDraft ? n_target_layers_ : 1; launch.shared_payload_bytes = draft_ipc_shared_bytes_from_env( - dflash_draft_ipc_required_shared_bytes(hidden_size_, block_size_, ring_cap)); + dflash_draft_ipc_required_shared_bytes( + hidden_size_, block_size_, ring_cap, shared_feature_layers)); launch.args.push_back("--ring-cap=" + std::to_string(ring_cap)); launch.args.push_back("--draft-gpu=" + std::to_string(std::max(0, draft_gpu))); if (!process_.start(launch)) { @@ -123,13 +141,77 @@ bool DFlashDraftIpcClient::start( return false; } ring_cap_ = ring_cap; + mode_ = mode; active_ = true; - std::printf("[draft-ipc] ready bin=%s gpu=%d ring_cap=%d work_dir=%s\n", - bin.c_str(), draft_gpu, ring_cap, process_.work_dir().c_str()); + std::printf("[draft-ipc] ready mode=%s bin=%s gpu=%d ring_cap=%d work_dir=%s\n", + backend_ipc_mode_name(mode), bin.c_str(), draft_gpu, ring_cap, + process_.work_dir().c_str()); return true; #endif } +bool DFlashDraftIpcClient::send_feature_block( + int start_pos, + int n_tokens, + const float * features, + size_t feature_count) { +#if defined(_WIN32) + (void)start_pos; (void)n_tokens; (void)features; (void)feature_count; + return false; +#else + FILE * cmd = process_.command_stream(); + const int stream_fd = process_.stream_fd(); + const int payload_fd = process_.payload_fd(); + if (!active_ || mode_ != BackendIpcMode::DeepSeek4DSparkDraft || + !cmd || stream_fd < 0 || start_pos < 0 || n_tokens <= 0 || + n_tokens > ring_cap_) { + return false; + } + const size_t expected = (size_t)n_tokens * (size_t)n_target_layers_ * + (size_t)hidden_size_; + if (!features || feature_count != expected) return false; + const size_t bytes = feature_count * sizeof(float); + + if (process_.resolved_payload_transport() == BackendIpcPayloadTransport::Shared) { + uint64_t seq = 0; + if (!process_.write_shared_payload(features, bytes, seq)) { + std::fprintf(stderr, + "draft-ipc feature_block shared payload too large bytes=%zu capacity=%zu\n", + bytes, process_.shared_payload_capacity()); + return false; + } + std::fprintf(cmd, "feature_block_shared %d %d %zu %" PRIu64 "\n", + start_pos, n_tokens, bytes, seq); + std::fflush(cmd); + int32_t status = -1; + const bool ok = read_exact_fd(stream_fd, &status, sizeof(status)) && + status == 0; + if (!ok) { + std::fprintf(stderr, + "draft-ipc feature_block_shared failed status=%d\n", + status); + } + return ok; + } + + if (payload_fd < 0) return false; + std::fprintf(cmd, "feature_block_pipe %d %d %zu\n", + start_pos, n_tokens, bytes); + std::fflush(cmd); + if (!write_exact_fd(payload_fd, features, bytes)) { + std::fprintf(stderr, "draft-ipc feature_block payload write failed\n"); + return false; + } + int32_t status = -1; + const bool ok = read_exact_fd(stream_fd, &status, sizeof(status)) && + status == 0; + if (!ok) { + std::fprintf(stderr, "draft-ipc feature_block failed status=%d\n", status); + } + return ok; +#endif +} + bool DFlashDraftIpcClient::send_feature_slice( int capture_idx, int start_pos, @@ -216,11 +298,19 @@ bool DFlashDraftIpcClient::propose( const int payload_fd = process_.payload_fd(); if (!active_ || !cmd || stream_fd < 0 || committed < 0 || ctx_len <= 0 || ctx_len > ring_cap_) { + std::fprintf(stderr, + "draft-ipc propose rejected active=%d cmd=%p stream_fd=%d committed=%d ctx_len=%d ring_cap=%d\n", + (int)active_, (void *)cmd, stream_fd, committed, ctx_len, ring_cap_); return false; } const size_t noise_expected = (size_t)hidden_size_ * block_size_; - if (noise_embed.size() != noise_expected) return false; + if (noise_embed.size() != noise_expected) { + std::fprintf(stderr, + "draft-ipc propose noise size mismatch got=%zu expected=%zu hidden=%d block=%d\n", + noise_embed.size(), noise_expected, hidden_size_, block_size_); + return false; + } const size_t bytes = noise_embed.size() * sizeof(float); if (process_.resolved_payload_transport() == BackendIpcPayloadTransport::Shared) { uint64_t seq = 0; @@ -230,15 +320,20 @@ bool DFlashDraftIpcClient::propose( bytes, process_.shared_payload_capacity()); return false; } - std::fprintf(cmd, "propose_shared %d %d %zu %" PRIu64 "\n", + const bool bidirectional = + mode_ == BackendIpcMode::DeepSeek4DSparkDraft; + std::fprintf(cmd, bidirectional + ? "propose_shared_bidir %d %d %zu %" PRIu64 "\n" + : "propose_shared %d %d %zu %" PRIu64 "\n", committed, ctx_len, bytes, seq); std::fflush(cmd); int32_t status = -1; bool ok = read_exact_fd(stream_fd, &status, sizeof(status)) && status == 0; if (ok) { hidden_out.assign(noise_expected, 0.0f); - ok = read_exact_fd(stream_fd, hidden_out.data(), - hidden_out.size() * sizeof(float)); + ok = bidirectional + ? process_.read_shared_payload(hidden_out.data(), bytes, seq) + : read_exact_fd(stream_fd, hidden_out.data(), bytes); } if (!ok) { std::fprintf(stderr, "draft-ipc propose_shared failed status=%d\n", status); @@ -348,6 +443,7 @@ void DFlashDraftIpcClient::close() { process_.close(); active_ = false; ring_cap_ = 0; + mode_ = BackendIpcMode::Invalid; } // ── Remote draft feature copy helper ──────────────────────────────── diff --git a/server/src/common/dflash_draft_ipc.h b/server/src/common/dflash_draft_ipc.h index 4830732b4..94e215a44 100644 --- a/server/src/common/dflash_draft_ipc.h +++ b/server/src/common/dflash_draft_ipc.h @@ -17,6 +17,7 @@ #include "ggml.h" #include "ggml-backend.h" +#include #include #include #include @@ -41,13 +42,28 @@ class DFlashDraftIpcClient { const std::string & draft_path, int draft_gpu, int ring_cap, - const std::string & work_dir); + const std::string & work_dir, + BackendIpcMode mode = BackendIpcMode::DFlashDraft); bool send_feature_slice(int capture_idx, int start_pos, int n_tokens, const std::vector & slice); + // Token-major [n_tokens, n_target_layers, hidden_size] upload. This is the + // preferred DSpark path because all capture layers cross IPC in one + // command and one acknowledgement. + bool send_feature_block(int start_pos, + int n_tokens, + const float * features, + size_t feature_count); + bool send_feature_block(int start_pos, + int n_tokens, + const std::vector & features) { + return send_feature_block( + start_pos, n_tokens, features.data(), features.size()); + } + bool propose(int committed, int ctx_len, const std::vector & noise_embed, @@ -71,6 +87,7 @@ class DFlashDraftIpcClient { int hidden_size_ = DFLASH27B_TARGET_HIDDEN; int block_size_ = DFLASH27B_DRAFT_BLOCK_SIZE; int n_target_layers_ = DFLASH27B_DRAFT_N_TARGET_LAYERS; + BackendIpcMode mode_ = BackendIpcMode::Invalid; }; // ── Remote draft feature copy helper ──────────────────────────────── diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 33cc81541..000ebc08d 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -437,10 +437,43 @@ bool DeepSeek4Backend::load_spec_drafter() { spec_drafter_parked_ = false; std::fprintf(stderr, "[deepseek4] DSpark spec-decode ENABLED (drafter=%s)\n", spec_draft_path_.c_str()); + return start_spec_remote_drafter(); +} + +bool DeepSeek4Backend::start_spec_remote_drafter() { + spec_remote_drafter_.reset(); + + const char * ipc_bin = std::getenv("DFLASH_DS4_DRAFT_IPC_BIN"); + if (!ipc_bin || !*ipc_bin) return true; + if (!spec_drafter_) return false; + + int ipc_gpu = 0; + if (const char * gpu = std::getenv("DFLASH_DS4_DRAFT_IPC_GPU")) { + ipc_gpu = std::max(0, std::atoi(gpu)); + } + const char * ipc_work = std::getenv("DFLASH_DS4_DRAFT_IPC_WORK_DIR"); + auto remote = std::make_unique( + w_.n_embd, spec_drafter_->block_size, + spec_drafter_->n_target_layers); + if (!remote->start( + ipc_bin, spec_draft_path_, ipc_gpu, w_.n_swa, + ipc_work && *ipc_work ? ipc_work : "", + BackendIpcMode::DeepSeek4DSparkDraft)) { + std::fprintf(stderr, + "[deepseek4] DSpark remote draft IPC failed; using local draft\n"); + return !env_flag_enabled("DFLASH_DS4_DRAFT_IPC_REQUIRED"); + } + + spec_remote_drafter_ = std::move(remote); + std::fprintf(stderr, + "[deepseek4] DSpark remote draft IPC ENABLED gpu=%d\n", + ipc_gpu); return true; } void DeepSeek4Backend::release_spec_drafter(bool mark_parked) { + // Stop the remote process before releasing the local metadata/fallback. + spec_remote_drafter_.reset(); if (spec_drafter_) { free_deepseek4_dspark_drafter(*spec_drafter_); } @@ -498,7 +531,10 @@ bool DeepSeek4Backend::init() { "[deepseek4] DSpark spec-decode requires monolithic model " "placement; disabled for hybrid expert placement\n"); } else { - (void) load_spec_drafter(); + const bool loaded = load_spec_drafter(); + if (!loaded && env_flag_enabled("DFLASH_DS4_DRAFT_IPC_REQUIRED")) { + return false; + } } } else { std::fprintf(stderr, "[deepseek4] DFLASH_DS4_SPEC set but DFLASH_DS4_DRAFT gguf missing\n"); @@ -930,7 +966,8 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, if (out_io.cancelled) return false; out_io.emit(tok); return !out_io.cancelled; - })) { + }, + spec_remote_drafter_.get())) { result.fail(GenerateErrorCode::DecodeFailed, "DSpark speculative decode failed"); return result; diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 5c37c77a1..6195be091 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -14,6 +14,7 @@ #include "../common/moe_hybrid_stream.h" #include "deepseek4_internal.h" #include "deepseek4_dspark.h" +#include "common/dflash_draft_ipc.h" #include "ggml.h" #include "ggml-backend.h" @@ -82,9 +83,11 @@ class DeepSeek4Backend : public ModelBackend { bool spec_drafter_parked_ = false; std::string spec_draft_path_; std::unique_ptr spec_drafter_; + std::unique_ptr spec_remote_drafter_; std::vector spec_feat_window_; bool load_spec_drafter(); + bool start_spec_remote_drafter(); void release_spec_drafter(bool mark_parked); // Prefill prompt tokens in chunks, return absolute committed position. diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index adc08167c..bdefaac64 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -37,6 +37,8 @@ namespace dflash::common { +class DFlashDraftIpcClient; + // The drafter weights. `core` reuses DeepSeek4Weights for the n_layer decoder // blocks + per-layer tensors + metadata + out_norm + output_hc_* tail; its // tok_embd/output stay null (tied to target). The DSpark-specific tensors below @@ -175,6 +177,15 @@ bool run_deepseek4_dspark_spec_decode( int win_len, std::vector & out_tokens, float * accept_rate_out, - const std::function & on_token = {}); + const std::function & on_token = {}, + DFlashDraftIpcClient * remote_draft = nullptr); + +int run_deepseek4_dspark_draft_ipc_daemon(const char * draft_path, + int ring_cap, + int draft_gpu, + int stream_fd, + int payload_fd, + int shared_payload_fd = -1, + size_t shared_payload_bytes = 0); } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_dspark_draft_ipc_daemon.cpp b/server/src/deepseek4/deepseek4_dspark_draft_ipc_daemon.cpp new file mode 100644 index 000000000..24cdcbef3 --- /dev/null +++ b/server/src/deepseek4/deepseek4_dspark_draft_ipc_daemon.cpp @@ -0,0 +1,354 @@ +// Out-of-process DeepSeek4 DSpark proposal worker. + +#include "deepseek4_dspark.h" + +#include "common/dflash_draft_ipc.h" +#include "common/io_utils.h" + +#include "ggml-backend.h" +#include "ggml-cuda.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +# include +# include +#endif + +namespace dflash::common { + +int run_deepseek4_dspark_draft_ipc_daemon( + const char * draft_path, + int ring_cap, + int draft_gpu, + int stream_fd, + int payload_fd, + int shared_payload_fd, + size_t shared_payload_bytes) { +#if defined(_WIN32) + (void) draft_path; + (void) ring_cap; + (void) draft_gpu; + (void) stream_fd; + (void) payload_fd; + (void) shared_payload_fd; + (void) shared_payload_bytes; + return 2; +#else + const bool shared_requested = + shared_payload_fd >= 0 || shared_payload_bytes > 0; + if (!draft_path || !*draft_path || ring_cap <= 0 || + stream_fd < 0 || (payload_fd < 0 && !shared_requested) || + (shared_requested && + (shared_payload_fd < 0 || shared_payload_bytes == 0))) { + std::fprintf(stderr, "[ds4-dspark-ipc] bad daemon configuration\n"); + if (stream_fd >= 0) stream_status(stream_fd, -1); + return 2; + } + + void * shared_payload = nullptr; + void * shared_payload_data = nullptr; + size_t shared_payload_map_bytes = 0; + if (shared_requested) { + if (!backend_ipc_shared_payload_map_bytes( + shared_payload_bytes, shared_payload_map_bytes)) { + std::fprintf(stderr, "[ds4-dspark-ipc] invalid shared payload size\n"); + stream_status(stream_fd, -1); + return 1; + } + shared_payload = ::mmap(nullptr, shared_payload_map_bytes, + PROT_READ | PROT_WRITE, MAP_SHARED, + shared_payload_fd, 0); + if (shared_payload == MAP_FAILED) { + std::fprintf(stderr, "[ds4-dspark-ipc] shared payload mmap failed\n"); + stream_status(stream_fd, -1); + return 1; + } + shared_payload_data = static_cast(shared_payload) + + backend_ipc_shared_payload_header_bytes(); + } + auto unmap_shared = [&]() { + if (shared_payload && shared_payload != MAP_FAILED) { + ::munmap(shared_payload, shared_payload_map_bytes); + shared_payload = nullptr; + shared_payload_data = nullptr; + } + }; + auto shared_request_valid = [&](size_t bytes, uint64_t sequence) { + const auto * header = + static_cast(shared_payload); + return sequence != 0 && shared_payload && shared_payload_data && + backend_ipc_payload_in_bounds( + 0, bytes, shared_payload_bytes) && + header->sequence == sequence && header->bytes == bytes; + }; + + ggml_backend_t backend = ggml_backend_cuda_init(std::max(0, draft_gpu)); + if (!backend) { + std::fprintf(stderr, "[ds4-dspark-ipc] backend init failed gpu=%d\n", + draft_gpu); + stream_status(stream_fd, -1); + unmap_shared(); + return 1; + } + + DSparkDrafter drafter; + if (!load_deepseek4_dspark_drafter(draft_path, backend, drafter)) { + std::fprintf(stderr, "[ds4-dspark-ipc] draft load failed: %s\n", + deepseek4_dspark_last_error()); + stream_status(stream_fd, -1); + ggml_backend_free(backend); + unmap_shared(); + return 1; + } + + const int hidden = drafter.core.n_embd; + const int block = drafter.block_size; + const int n_target_layers = drafter.n_target_layers; + const int feature_row = hidden * n_target_layers; + if (hidden <= 0 || block <= 0 || n_target_layers <= 0 || feature_row <= 0) { + std::fprintf(stderr, "[ds4-dspark-ipc] invalid draft dimensions\n"); + stream_status(stream_fd, -1); + free_deepseek4_dspark_drafter(drafter); + ggml_backend_free(backend); + unmap_shared(); + return 1; + } + + std::vector feature_ring((size_t) ring_cap * feature_row, 0.0f); + std::vector noise_embed((size_t) hidden * block); + std::vector context; + std::vector hidden_out; + + auto ring_slot = [ring_cap](int position) { + int slot = position % ring_cap; + return slot < 0 ? slot + ring_cap : slot; + }; + + auto store_feature_slice = [&](int capture_idx, int start_pos, int n_tokens, + const std::vector & slice) { + if (capture_idx < 0 || capture_idx >= n_target_layers || + start_pos < 0 || n_tokens <= 0 || + slice.size() != (size_t) n_tokens * hidden) { + return false; + } + for (int i = 0; i < n_tokens; ++i) { + float * dst = feature_ring.data() + + (size_t) ring_slot(start_pos + i) * feature_row + + (size_t) capture_idx * hidden; + std::memcpy(dst, slice.data() + (size_t) i * hidden, + sizeof(float) * (size_t) hidden); + } + return true; + }; + + auto store_feature_block = [&](int start_pos, int n_tokens, + const float * features, + size_t feature_count) { + const size_t expected = (size_t) n_tokens * (size_t) feature_row; + if (start_pos < 0 || n_tokens <= 0 || n_tokens > ring_cap || + !features || feature_count != expected) { + return false; + } + for (int i = 0; i < n_tokens; ++i) { + float * dst = feature_ring.data() + + (size_t) ring_slot(start_pos + i) * feature_row; + std::memcpy(dst, features + (size_t) i * feature_row, + sizeof(float) * (size_t) feature_row); + } + return true; + }; + + auto run_proposal = [&](int committed, int ctx_len, + const float * noise) { + if (committed < 0 || ctx_len <= 0 || ctx_len > ring_cap || !noise) { + return false; + } + context.resize((size_t) ctx_len * feature_row); + const int start_pos = committed - ctx_len; + for (int i = 0; i < ctx_len; ++i) { + const float * src = feature_ring.data() + + (size_t) ring_slot(start_pos + i) * feature_row; + std::memcpy(context.data() + (size_t) i * feature_row, src, + sizeof(float) * (size_t) feature_row); + } + return deepseek4_dspark_draft_forward( + backend, drafter, noise, context.data(), ctx_len, + committed, hidden_out); + }; + + std::fprintf(stderr, + "[ds4-dspark-ipc] ready gpu=%d ring_cap=%d hidden=%d block=%d captures=%d transport=%s\n", + draft_gpu, ring_cap, hidden, block, n_target_layers, + shared_payload ? "shared" : "stream"); + if (!stream_status(stream_fd, 0)) { + free_deepseek4_dspark_drafter(drafter); + ggml_backend_free(backend); + unmap_shared(); + return 1; + } + + std::string line; + while (std::getline(std::cin, line)) { + std::istringstream iss(line); + std::string cmd; + iss >> cmd; + if (cmd == "quit" || cmd == "exit") break; + + if (cmd == "feature_slice_pipe") { + int capture_idx = -1; + int start_pos = -1; + int n_tokens = 0; + size_t bytes = 0; + iss >> capture_idx >> start_pos >> n_tokens >> bytes; + const size_t expected = (size_t) std::max(0, n_tokens) * hidden * + sizeof(float); + bool ok = iss && bytes == expected && n_tokens > 0; + std::vector slice(ok ? bytes / sizeof(float) : 0); + if (ok) ok = read_exact_fd(payload_fd, slice.data(), bytes); + if (ok) ok = store_feature_slice(capture_idx, start_pos, n_tokens, slice); + if (!stream_status(stream_fd, ok ? 0 : -1)) break; + continue; + } + + if (cmd == "feature_slice_shared") { + int capture_idx = -1; + int start_pos = -1; + int n_tokens = 0; + size_t bytes = 0; + uint64_t sequence = 0; + iss >> capture_idx >> start_pos >> n_tokens >> bytes >> sequence; + const size_t expected = (size_t) std::max(0, n_tokens) * hidden * + sizeof(float); + bool ok = iss && bytes == expected && n_tokens > 0 && + capture_idx >= 0 && capture_idx < n_target_layers && + shared_request_valid(bytes, sequence); + std::vector slice(ok ? bytes / sizeof(float) : 0); + if (ok) std::memcpy(slice.data(), shared_payload_data, bytes); + if (ok) { + ok = store_feature_slice( + capture_idx, start_pos, n_tokens, slice); + } + if (!stream_status(stream_fd, ok ? 0 : -1)) break; + continue; + } + + if (cmd == "feature_block_pipe") { + int start_pos = -1; + int n_tokens = 0; + size_t bytes = 0; + iss >> start_pos >> n_tokens >> bytes; + const size_t expected = (size_t) std::max(0, n_tokens) * + (size_t) feature_row * sizeof(float); + bool ok = iss && payload_fd >= 0 && bytes == expected && + n_tokens > 0 && n_tokens <= ring_cap; + std::vector features(ok ? bytes / sizeof(float) : 0); + if (ok) ok = read_exact_fd(payload_fd, features.data(), bytes); + if (ok) { + ok = store_feature_block( + start_pos, n_tokens, features.data(), features.size()); + } + if (!stream_status(stream_fd, ok ? 0 : -1)) break; + continue; + } + + if (cmd == "feature_block_shared") { + int start_pos = -1; + int n_tokens = 0; + size_t bytes = 0; + uint64_t sequence = 0; + iss >> start_pos >> n_tokens >> bytes >> sequence; + const size_t expected = (size_t) std::max(0, n_tokens) * + (size_t) feature_row * sizeof(float); + bool ok = iss && bytes == expected && n_tokens > 0 && + n_tokens <= ring_cap && + shared_request_valid(bytes, sequence); + if (ok) { + ok = store_feature_block( + start_pos, n_tokens, + static_cast(shared_payload_data), + bytes / sizeof(float)); + } + if (!stream_status(stream_fd, ok ? 0 : -1)) break; + continue; + } + + if (cmd == "propose_pipe") { + int committed = -1; + int ctx_len = 0; + size_t bytes = 0; + iss >> committed >> ctx_len >> bytes; + const size_t expected = noise_embed.size() * sizeof(float); + bool ok = iss && bytes == expected; + if (!ok) { + std::fprintf(stderr, + "[ds4-dspark-ipc] bad propose header committed=%d ctx=%d bytes=%zu expected=%zu\n", + committed, ctx_len, bytes, expected); + } + if (ok && !read_exact_fd(payload_fd, noise_embed.data(), bytes)) { + std::fprintf(stderr, "[ds4-dspark-ipc] propose payload read failed\n"); + ok = false; + } + if (ok && !run_proposal( + committed, ctx_len, noise_embed.data())) { + std::fprintf(stderr, + "[ds4-dspark-ipc] proposal graph failed committed=%d ctx=%d\n", + committed, ctx_len); + ok = false; + } + if (!stream_status(stream_fd, ok ? 0 : -1)) break; + if (ok && !write_exact_fd(stream_fd, hidden_out.data(), + hidden_out.size() * sizeof(float))) { + break; + } + continue; + } + + if (cmd == "propose_shared_bidir") { + int committed = -1; + int ctx_len = 0; + size_t bytes = 0; + uint64_t sequence = 0; + iss >> committed >> ctx_len >> bytes >> sequence; + const size_t expected = noise_embed.size() * sizeof(float); + bool ok = iss && bytes == expected && + shared_request_valid(bytes, sequence); + if (ok) { + ok = run_proposal( + committed, ctx_len, + static_cast(shared_payload_data)); + } + if (ok && hidden_out.size() * sizeof(float) == bytes) { + std::memcpy(shared_payload_data, hidden_out.data(), bytes); + auto * header = + static_cast(shared_payload); + header->bytes = bytes; + header->sequence = sequence; + } else { + ok = false; + } + if (!stream_status(stream_fd, ok ? 0 : -1)) break; + continue; + } + + std::fprintf(stderr, "[ds4-dspark-ipc] unknown command: %s\n", + line.c_str()); + if (!stream_status(stream_fd, -1)) break; + } + + free_deepseek4_dspark_drafter(drafter); + ggml_backend_free(backend); + unmap_shared(); + return 0; +#endif +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index d22bf0ac9..5a09f5ada 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -25,6 +25,7 @@ #include "deepseek4_internal.h" #include "internal.h" #include "common/dspark_head.h" +#include "common/dflash_draft_ipc.h" #include "ggml.h" #include "ggml-backend.h" @@ -414,7 +415,8 @@ bool run_deepseek4_dspark_spec_decode( int win_len, std::vector & out_tokens, float * accept_rate_out, - const std::function & on_token) { + const std::function & on_token, + DFlashDraftIpcClient * remote_draft) { const int n_embd = target_w.n_embd; const int n_tgt = drafter.n_target_layers; const int block = drafter.block_size; @@ -434,7 +436,10 @@ bool run_deepseek4_dspark_spec_decode( if (std::fscanf(f, "%d", &v) == 1) adaptive_width = (v != 0); std::fclose(f); } - const bool use_confidence_width = adaptive_width && !seq_verify_mode && + // The v1 IPC protocol returns normalized draft states only. Use its EWMA + // width policy until the protocol also carries pre-norm confidence input. + const bool remote_active = remote_draft && remote_draft->active(); + const bool use_confidence_width = !remote_active && adaptive_width && !seq_verify_mode && drafter.confidence_w != nullptr && drafter.confidence_b != nullptr && (drafter.confidence_dim == n_embd || drafter.confidence_dim == n_embd + drafter.markov_rank); @@ -495,6 +500,18 @@ bool run_deepseek4_dspark_spec_decode( } int feat_count = win_have; // number of valid feature columns ending at committed-1 + if (remote_draft && remote_draft->active() && win_have > 0) { + const int start_pos = committed - win_have; + if (!remote_draft->send_feature_block( + start_pos, win_have, feat_win.data(), + (size_t) win_have * feat_row)) { + std::fprintf(stderr, + "[ds4-spec] remote draft prompt feature upload failed\n"); + ggml_backend_free(snap_backend); + return false; + } + } + auto push_feature = [&](const float * col) { // Shift-append one feature column (keep last n_swa). if (feat_count >= n_swa) { @@ -543,10 +560,14 @@ bool run_deepseek4_dspark_spec_decode( // Drafter forward -> normalized states for token logits plus the // pre-output-norm states expected by the confidence head. - if (!deepseek4_dspark_draft_forward(backend, drafter, noise_embed.data(), - ctx_len > 0 ? feat_win.data() : nullptr, - ctx_len, pos, local_hidden, - use_confidence_width ? &confidence_hidden : nullptr)) { + const bool draft_ok = remote_draft && remote_draft->active() + ? remote_draft->propose(pos, ctx_len, noise_embed, local_hidden) + : deepseek4_dspark_draft_forward( + backend, drafter, noise_embed.data(), + ctx_len > 0 ? feat_win.data() : nullptr, + ctx_len, pos, local_hidden, + use_confidence_width ? &confidence_hidden : nullptr); + if (!draft_ok) { std::fprintf(stderr, "[ds4-spec] drafter forward failed\n"); ok = false; break; @@ -735,6 +756,17 @@ bool run_deepseek4_dspark_spec_decode( t0 = SpecClock::now(); const std::vector & feats = target.last_features(); const int fN = full_snap ? target.last_verify_n() : accept; + if (remote_draft && remote_draft->active() && fN > 0) { + const size_t feature_count = (size_t) fN * feat_row; + if (feats.size() < feature_count || + !remote_draft->send_feature_block( + pos, fN, feats.data(), feature_count)) { + std::fprintf(stderr, + "[ds4-spec] remote draft feature update failed\n"); + ggml_backend_free(snap_backend); + return false; + } + } for (int i = 0; i < fN; i++) push_feature(feats.data() + (size_t) i * feat_row); tm_feat += spec_ms_since(t0); diff --git a/server/src/ipc/backend_ipc_main.cpp b/server/src/ipc/backend_ipc_main.cpp index ca5ba1565..e06983540 100644 --- a/server/src/ipc/backend_ipc_main.cpp +++ b/server/src/ipc/backend_ipc_main.cpp @@ -4,6 +4,7 @@ #include "moe_expert_compute.h" #include "dflash_draft_ipc.h" +#include "deepseek4/deepseek4_dspark.h" #include "deepseek4/deepseek4_layer_split_adapter.h" #include "gemma4/gemma4_layer_split_adapter.h" #include "laguna/laguna_layer_split_adapter.h" @@ -130,6 +131,8 @@ int main(int argc, char ** argv) { " or: %s --backend-ipc-mode=deepseek4-target-shard " "--stream-fd=FD --target-gpus=N[,N...] --layer-begins=N[,N...] " "--layer-ends=N[,N...] --max-ctx=N\n" + " or: %s --backend-ipc-mode=deepseek4-dspark-draft " + "--ring-cap=N --stream-fd=FD --payload-fd=FD [--draft-gpu=N]\n" " or: %s --backend-ipc-mode=moe-expert-compute " "--stream-fd=FD --target-gpu=N --placement=PATH\n", argv[0], @@ -137,6 +140,8 @@ int main(int argc, char ** argv) { argv[0], argv[0], argv[0], + argv[0], + argv[0], argv[0]); return 2; } @@ -351,6 +356,10 @@ int main(int argc, char ** argv) { return run_deepseek4_target_shard_ipc_daemon( payload_path, target_gpus, layer_begins, layer_ends, max_ctx, stream_fd, payload_fd); + case BackendIpcMode::DeepSeek4DSparkDraft: + return run_deepseek4_dspark_draft_ipc_daemon( + payload_path, ring_cap, draft_gpu, stream_fd, payload_fd, + shared_payload_fd, shared_payload_bytes); } std::fprintf(stderr, "[backend-ipc-daemon] unsupported mode\n"); return 2; diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index bc9668fac..a498ceb4c 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -3076,6 +3076,11 @@ static void test_backend_ipc_payload_transport_parse() { TEST_ASSERT(mode == BackendIpcMode::PFlashCompress); TEST_ASSERT(parse_backend_ipc_mode("qwen35-target-shard", mode)); TEST_ASSERT(mode == BackendIpcMode::Qwen35TargetShard); + TEST_ASSERT(parse_backend_ipc_mode("deepseek4-dspark-draft", mode)); + TEST_ASSERT(mode == BackendIpcMode::DeepSeek4DSparkDraft); + TEST_ASSERT(std::strcmp( + backend_ipc_mode_name(BackendIpcMode::DeepSeek4DSparkDraft), + "deepseek4-dspark-draft") == 0); TEST_ASSERT(parse_backend_ipc_mode("moe-expert-compute", mode)); TEST_ASSERT(mode == BackendIpcMode::MoeExpertCompute); TEST_ASSERT(!parse_backend_ipc_mode("moe-ffn", mode)); From 321ba38b942fad1cf46dca83a8504565e5be853c Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 19 Jul 2026 22:18:56 -0700 Subject: [PATCH 02/15] feat(ds4): add in-process heterogeneous expert parallel runtime --- .../llama.cpp/ggml/include/ggml-backend.h | 29 + server/deps/llama.cpp/ggml/include/ggml.h | 56 + .../deps/llama.cpp/ggml/src/ggml-backend.cpp | 445 ++++++- server/deps/llama.cpp/ggml/src/ggml.c | 159 +++ server/src/common/backend_ipc.cpp | 10 +- server/src/common/backend_ipc.h | 41 + server/src/common/moe_expert_compute.cpp | 10 +- server/src/common/moe_expert_compute.h | 39 + server/src/common/moe_expert_compute_ipc.cpp | 417 +++++- server/src/common/moe_hybrid_ffn_eval.cpp | 741 +++++++++-- server/src/common/moe_hybrid_ffn_eval.h | 64 + server/src/common/moe_hybrid_storage.cpp | 419 +++++- server/src/common/moe_hybrid_storage.h | 31 +- server/src/deepseek4/deepseek4_backend.cpp | 681 +++++++++- server/src/deepseek4/deepseek4_backend.h | 12 +- .../src/deepseek4/deepseek4_fused_verify.inc | 1118 +++++++++++++++-- server/src/deepseek4/deepseek4_graph.cpp | 1054 +++++++++++++--- server/src/deepseek4/deepseek4_internal.h | 20 +- server/src/deepseek4/deepseek4_loader.cpp | 116 +- 19 files changed, 4974 insertions(+), 488 deletions(-) diff --git a/server/deps/llama.cpp/ggml/include/ggml-backend.h b/server/deps/llama.cpp/ggml/include/ggml-backend.h index 4a8f6d428..debc793a5 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-backend.h +++ b/server/deps/llama.cpp/ggml/include/ggml-backend.h @@ -323,6 +323,9 @@ extern "C" { // Get the number of splits of the last graph GGML_API int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched); + GGML_API int ggml_backend_sched_get_n_splits_for_backend( + ggml_backend_sched_t sched, + ggml_backend_t backend); GGML_API int ggml_backend_sched_get_n_copies(ggml_backend_sched_t sched); GGML_API ggml_backend_buffer_type_t ggml_backend_sched_get_buffer_type(ggml_backend_sched_t sched, ggml_backend_t backend); @@ -340,6 +343,13 @@ extern "C" { GGML_API enum ggml_status ggml_backend_sched_graph_compute_async(ggml_backend_sched_t sched, struct ggml_cgraph * graph); GGML_API void ggml_backend_sched_synchronize(ggml_backend_sched_t sched); + // The caller has synchronized all backends and is recording this compute + // into an outer device graph. Buffer-reuse waits may be omitted, while + // any copy requiring a synchronous fallback becomes an explicit error. + GGML_API void ggml_backend_sched_set_whole_graph_capture( + ggml_backend_sched_t sched, + bool enabled); + // Reset all assignments and allocators - must be called before changing the node backends or allocating a new graph. // This in effect deallocates all tensors that were previously allocated and leaves them with dangling pointers. // The correct way to use this API is to discard the deallocated tensors and create new ones. @@ -348,6 +358,25 @@ extern "C" { // Set a callback to be called for each resulting node during graph compute GGML_API void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data); + // End a backend split before a late node introduces a new cross-backend + // input, so independent earlier nodes can be enqueued before the copy/event + // wait. Must be configured before allocating the graph. + GGML_API void ggml_backend_sched_set_late_cross_input_split(ggml_backend_sched_t sched, bool enabled); + + // Apply the same split rule only at an explicitly marked node. This keeps + // unrelated cross-backend consumers in their established scheduler + // segments. Nodes must be registered before allocating the graph. + GGML_API void ggml_backend_sched_add_late_cross_input_split_node( + ggml_backend_sched_t sched, + const struct ggml_tensor * node); + + // Mark a GGML_OP_MOE_FUSED mode -3 node as a deferred peer copy. The + // scheduler keeps src[0] on its owner backend, records a dedicated event + // after the producing split, and injects the native handle into the op. + GGML_API void ggml_backend_sched_add_deferred_peer_copy_node( + ggml_backend_sched_t sched, + const struct ggml_tensor * node); + // // Meta backend // diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index 1d8f22bb8..a09031452 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2426,6 +2426,51 @@ extern "C" { struct ggml_tensor * experts, struct ggml_tensor * expert_weights); + // Coarse DeepSeek-V4 routed-owner op. gate_up contains concatenated gate + // and up output rows; the backend performs fused gate/up MMVQ + clamped + // SwiGLU, down MMVQ, route weighting, and local reduction. + GGML_API struct ggml_tensor * ggml_ds4_moe_owner( + struct ggml_context * ctx, + struct ggml_tensor * input, + struct ggml_tensor * gate_up_w, + struct ggml_tensor * down_w, + struct ggml_tensor * expert_ids, + struct ggml_tensor * expert_weights, + int64_t ff_dim, + float swiglu_clamp, + float down_scale); + + // Coarse owner variant for checkpoints that store gate and up tensors + // separately. The backend fuses their MMVQ traversal and SwiGLU while + // preserving the checkpoint's three external value scales. + GGML_API struct ggml_tensor * ggml_ds4_moe_owner_split( + struct ggml_context * ctx, + struct ggml_tensor * input, + struct ggml_tensor * gate_w, + struct ggml_tensor * up_w, + struct ggml_tensor * down_w, + struct ggml_tensor * expert_ids, + struct ggml_tensor * expert_weights, + int64_t ff_dim, + float swiglu_clamp, + float gate_scale, + float up_scale, + float down_scale); + + // Reorder a small [route, token] owner-local ID matrix so occurrences of + // the same expert across verification tokens share a kernel block. Each + // encoded ID retains its original route slot for exact output scatter. + GGML_API struct ggml_tensor * ggml_ds4_moe_align_ids( + struct ggml_context * ctx, + struct ggml_tensor * expert_ids); + + // Copy an F32 peer-GPU tensor only after a scheduler-provided device event + // has completed. The scheduler writes the native event handle into the op + // parameters and deliberately leaves src on its owner backend. + GGML_API struct ggml_tensor * ggml_ds4_deferred_peer_copy( + struct ggml_context * ctx, + struct ggml_tensor * src); + // Fused DeepSeek4 hyper-connection helpers (decode, n_tokens == 1). // ggml_ds4_hc_pre: mix[2*n_hc+n_hc^2] + base + hc_state[n_embd*n_hc] -> // dst[n_embd + 2*n_hc + n_hc^2] = { working, split(pre,post,comb) } @@ -2448,6 +2493,17 @@ extern "C" { struct ggml_tensor * split, int n_hc); + // Exact mode-1 variant for heterogeneous MoE verification. It computes + // block_out[d] = peer_block[d] + main_block[d] inside the HC-post kernel, + // eliminating the standalone reduction and tokenwise CONT copies. + GGML_API struct ggml_tensor * ggml_ds4_hc_post_split( + struct ggml_context * ctx, + struct ggml_tensor * residual_hc, + struct ggml_tensor * main_block, + struct ggml_tensor * peer_block, + struct ggml_tensor * split, + int n_hc); + // ggml_ds4_hc_out: output-stage merge of hc streams into one embedding GGML_API struct ggml_tensor * ggml_ds4_hc_out( struct ggml_context * ctx, diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp index 01073f631..47242375b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp @@ -807,6 +807,15 @@ struct ggml_backend_sched_split { struct ggml_cgraph graph; }; +struct ggml_backend_sched_deferred_peer_copy { + struct ggml_tensor * node; + struct ggml_tensor * source; + ggml_backend_event_t event; + int producer_backend_id; + int producer_split_id; + int consumer_split_id; +}; + struct ggml_backend_sched { bool is_reset; // true if the scheduler has been reset since the last graph split bool is_alloc; @@ -853,6 +862,20 @@ struct ggml_backend_sched { size_t context_buffer_size; bool op_offload; + bool whole_graph_capture; + + // When a new cross-backend dependency first appears after executable + // nodes in a split, end the current split before the dependent node. This + // lets independent work already in the split be enqueued before the + // cross-backend copy/event wait is attached to the following split. + bool late_cross_input_split; + const struct ggml_tensor ** late_cross_input_split_nodes; + int n_late_cross_input_split_nodes; + int late_cross_input_split_nodes_capacity; + + struct ggml_backend_sched_deferred_peer_copy * deferred_peer_copies; + int n_deferred_peer_copies; + int deferred_peer_copies_capacity; int debug; @@ -868,6 +891,37 @@ struct ggml_backend_sched { #define tensor_id_copy(id, backend_id, copy_id) sched->hv_tensor_copies[(id) * sched->n_backends * sched->n_copies + (backend_id) * sched->n_copies + (copy_id)] #define tensor_copy(tensor, backend_id, copy_id) tensor_id_copy(hash_id(tensor), backend_id, copy_id) +static bool ggml_backend_sched_is_late_cross_input_split_node( + const ggml_backend_sched_t sched, + const struct ggml_tensor * node) { + for (int i = 0; i < sched->n_late_cross_input_split_nodes; ++i) { + if (sched->late_cross_input_split_nodes[i] == node) { + return true; + } + } + return false; +} + +static struct ggml_backend_sched_deferred_peer_copy * +ggml_backend_sched_find_deferred_peer_copy( + const ggml_backend_sched_t sched, + const struct ggml_tensor * node) { + for (int i = 0; i < sched->n_deferred_peer_copies; ++i) { + if (sched->deferred_peer_copies[i].node == node) { + return &sched->deferred_peer_copies[i]; + } + } + return nullptr; +} + +static bool ggml_backend_sched_deferred_peer_split_enabled() { + static const bool enabled = [] { + const char * raw = getenv("DFLASH_DS4_TP_DEVICE_JOIN_SPLIT"); + return raw && *raw && strcmp(raw, "0") != 0; + }(); + return enabled; +} + // returns the priority of the backend, lower id is higher priority static int ggml_backend_sched_backend_id(ggml_backend_sched_t sched, ggml_backend_t backend) { for (int i = 0; i < sched->n_backends; i++) { @@ -1292,6 +1346,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra split->i_start = 0; split->n_inputs = 0; int cur_backend_id = split->backend_id; + bool split_has_compute = false; for (; i < graph->n_nodes; i++) { struct ggml_tensor * node = graph->nodes[i]; @@ -1305,6 +1360,19 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra // check if we should start a new split based on the sources of the current node bool need_new_split = false; + // HIP graph replay on the R9700 currently gates an entire graph + // when a cross-device event wait is captured inside it, even if + // independent hot/shared nodes precede the wait. The opt-in + // three-stage schedule ends the main-owner prefix here: + // Strix cold -> R9700 hot/shared -> R9700 wait/copy/join. + // The external stream wait is queued only after the useful prefix + // graph, allowing both owner branches to execute concurrently. + if (ggml_backend_sched_deferred_peer_split_enabled() && + node_backend_id == cur_backend_id && + split_has_compute && + ggml_backend_sched_find_deferred_peer_copy(sched, node)) { + need_new_split = true; + } if (node_backend_id == cur_backend_id && split->n_inputs > 0) { for (int j = 0; j < GGML_MAX_SRC; j++) { struct ggml_tensor * src = node->src[j]; @@ -1334,6 +1402,37 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra } } + // The normal scheduler copies every foreign input before it + // launches any node in a split. If a foreign input is first used + // by a late join, attaching it to the existing split needlessly + // makes independent work at the start of that split wait. End the + // split immediately before that join instead. Existing copies do + // not require a new split because their dependency was already + // established by an earlier split. + const bool split_late_cross_input = + sched->late_cross_input_split || + ggml_backend_sched_is_late_cross_input_split_node(sched, node); + if (split_late_cross_input && + node_backend_id == cur_backend_id && + split_has_compute && !need_new_split) { + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (src == NULL) { + continue; + } + const size_t src_id = hash_id(src); + const int src_backend_id = sched->hv_tensor_backend_ids[src_id]; + const bool supported = + ggml_backend_sched_buffer_supported(sched, src, cur_backend_id); + if (src_backend_id != cur_backend_id && + tensor_id_copy(src_id, cur_backend_id, 0) == NULL && + !supported) { + need_new_split = true; + break; + } + } + } + if (node_backend_id != cur_backend_id || need_new_split) { split->i_end = i; i_split++; @@ -1348,6 +1447,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra split->i_start = i; split->n_inputs = 0; cur_backend_id = node_backend_id; + split_has_compute = false; } // find inputs that are not on the same backend @@ -1383,7 +1483,17 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra } } - if (src_backend_id != cur_backend_id && !ggml_backend_sched_buffer_supported(sched, src, cur_backend_id)) { + const bool deferred_peer_src = + j == 0 && + ggml_backend_sched_find_deferred_peer_copy(sched, node) != nullptr; + if (src_backend_id != cur_backend_id && + !ggml_backend_sched_buffer_supported(sched, src, cur_backend_id)) { + if (deferred_peer_src) { + // The destination backend accesses the peer allocation + // directly after the mode -3 op waits on its dedicated + // event. Do not create or substitute a scheduler copy. + continue; + } // create a copy of the input in the split's backend if (tensor_id_copy(src_id, cur_backend_id, 0) == NULL) { ggml_backend_t backend = sched->backends[cur_backend_id]; @@ -1404,11 +1514,80 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra node->src[j] = tensor_id_copy(src_id, cur_backend_id, sched->cur_copy); } } + split_has_compute = true; } split->i_end = graph->n_nodes; sched->n_splits = i_split + 1; } + // Resolve every deferred peer copy after split boundaries are final. A + // unique event per join is required: reusing one backend event across + // dozens of captured waits would let a later record retarget an earlier + // graph replay. + for (int i = 0; i < sched->n_deferred_peer_copies; ++i) { + struct ggml_backend_sched_deferred_peer_copy * deferred = + &sched->deferred_peer_copies[i]; + struct ggml_tensor * node = deferred->node; + GGML_ASSERT(node->op == GGML_OP_MOE_FUSED); + GGML_ASSERT(ggml_get_op_params_i32(node, 0) == -3); + GGML_ASSERT(deferred->source); + + const int producer_backend_id = tensor_backend_id(deferred->source); + const int consumer_backend_id = tensor_backend_id(node); + GGML_ASSERT(producer_backend_id >= 0); + GGML_ASSERT(consumer_backend_id >= 0); + GGML_ASSERT(producer_backend_id != consumer_backend_id); + + if (deferred->event == nullptr) { + deferred->event = ggml_backend_event_new( + sched->backends[producer_backend_id]->device); + GGML_ASSERT(deferred->event); + } + deferred->producer_backend_id = producer_backend_id; + deferred->producer_split_id = -1; + deferred->consumer_split_id = -1; + + int producer_node_id = -1; + int consumer_node_id = -1; + for (int node_id = 0; node_id < graph->n_nodes; ++node_id) { + if (graph->nodes[node_id] == deferred->source) { + producer_node_id = node_id; + } + if (graph->nodes[node_id] == node) { + consumer_node_id = node_id; + } + } + GGML_ASSERT(producer_node_id >= 0); + GGML_ASSERT(consumer_node_id >= 0); + for (int split_id = 0; split_id < sched->n_splits; ++split_id) { + const struct ggml_backend_sched_split * split = + &sched->splits[split_id]; + if (producer_node_id >= split->i_start && + producer_node_id < split->i_end) { + deferred->producer_split_id = split_id; + } + if (consumer_node_id >= split->i_start && + consumer_node_id < split->i_end) { + deferred->consumer_split_id = split_id; + } + } + GGML_ASSERT(deferred->producer_split_id >= 0); + GGML_ASSERT(deferred->consumer_split_id >= 0); + GGML_ASSERT( + sched->splits[deferred->producer_split_id].backend_id == + producer_backend_id); + + void * event_context = deferred->event->context; + static_assert( + 2 * sizeof(int32_t) + sizeof(event_context) <= + GGML_MAX_OP_PARAMS, + "deferred event handle does not fit op_params"); + memcpy(&node->op_params[2], &event_context, sizeof(event_context)); + ggml_set_op_params_i32( + node, 7, + ggml_backend_sched_deferred_peer_split_enabled() ? 1 : 0); + } + if (sched->debug) { ggml_backend_sched_print_assignments(sched, graph); } @@ -1580,6 +1759,43 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s int split_backend_id = split->backend_id; ggml_backend_t split_backend = sched->backends[split_backend_id]; + // Diagnostic control for cross-device event semantics. This preserves + // the exact graph and peer-copy op but resolves its producer events on + // the host before launching the consuming split. + static const bool deferred_host_wait = [] { + const char * raw = + getenv("DFLASH_DS4_TP_DEVICE_JOIN_HOST_WAIT"); + return raw && *raw && strcmp(raw, "0") != 0; + }(); + static const bool deferred_host_copy = [] { + const char * raw = + getenv("DFLASH_DS4_TP_DEVICE_JOIN_HOST_COPY"); + return raw && *raw && strcmp(raw, "0") != 0; + }(); + if (deferred_host_wait || deferred_host_copy) { + for (int i = 0; i < sched->n_deferred_peer_copies; ++i) { + struct ggml_backend_sched_deferred_peer_copy * deferred = + &sched->deferred_peer_copies[i]; + if (deferred->consumer_split_id == split_id) { + ggml_backend_event_synchronize(deferred->event); + if (deferred_host_copy) { + ggml_backend_tensor_copy( + deferred->source, deferred->node); + } + } + } + } + if (ggml_backend_sched_deferred_peer_split_enabled() && + !deferred_host_wait && !deferred_host_copy) { + for (int i = 0; i < sched->n_deferred_peer_copies; ++i) { + struct ggml_backend_sched_deferred_peer_copy * deferred = + &sched->deferred_peer_copies[i]; + if (deferred->consumer_split_id == split_id) { + ggml_backend_event_wait(split_backend, deferred->event); + } + } + } + // copy the input tensors to the split backend for (int input_id = 0; input_id < split->n_inputs; input_id++) { ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]); @@ -1588,18 +1804,31 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s if (input->flags & GGML_TENSOR_FLAG_INPUT) { // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done - if (sched->events[split_backend_id][sched->cur_copy] != NULL) { - ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); + if (sched->whole_graph_capture) { + if (!split_backend->iface.cpy_tensor_async || + !split_backend->iface.cpy_tensor_async( + input_backend, split_backend, input, input_cpy)) { + GGML_LOG_ERROR( + "%s: outer capture requires async graph-input copy\n", + __func__); + return GGML_STATUS_FAILED; + } } else { - ggml_backend_synchronize(split_backend); + if (sched->events[split_backend_id][sched->cur_copy] != NULL) { + ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); + } else { + ggml_backend_synchronize(split_backend); + } + ggml_backend_tensor_copy(input, input_cpy); } - ggml_backend_tensor_copy(input, input_cpy); } else { // wait for the split backend to finish using the input before overwriting it - if (sched->events[split_backend_id][sched->cur_copy] != NULL) { - ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]); - } else { - ggml_backend_synchronize(split_backend); + if (!sched->whole_graph_capture) { + if (sched->events[split_backend_id][sched->cur_copy] != NULL) { + ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]); + } else { + ggml_backend_synchronize(split_backend); + } } // when offloading MoE weights, we can reduce the amount of data copied by copying only the experts that are used @@ -1611,6 +1840,13 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s //|| (node->src[1] == input_cpy && node->op == GGML_OP_ADD_ID) /* GGML_OP_ADD_ID weights are small and not worth splitting */ )) { + if (sched->whole_graph_capture) { + GGML_LOG_ERROR( + "%s: host expert offload is not capturable\n", + __func__); + return GGML_STATUS_FAILED; + } + const int64_t n_expert = node->op == GGML_OP_MUL_MAT_ID ? input->ne[2] : input->ne[1]; const size_t expert_size = node->op == GGML_OP_MUL_MAT_ID ? input->nb[2] : input->nb[1]; @@ -1691,6 +1927,12 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) { + if (sched->whole_graph_capture) { + GGML_LOG_ERROR( + "%s: outer capture requires async split copy\n", + __func__); + return GGML_STATUS_FAILED; + } ggml_backend_synchronize(input_backend); if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); @@ -1742,8 +1984,62 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } - // record the event of this copy - if (split->n_inputs > 0) { + // Publish cold-owner completion before a later main-backend graph + // reaches its in-graph event wait. Recording is asynchronous and does + // not block the host from immediately enqueueing independent work. + for (int i = 0; i < sched->n_deferred_peer_copies; ++i) { + struct ggml_backend_sched_deferred_peer_copy * deferred = + &sched->deferred_peer_copies[i]; + if (deferred->producer_split_id == split_id) { + GGML_ASSERT(deferred->producer_backend_id == split_backend_id); + if (sched->whole_graph_capture) { + const int consumer_backend_id = + sched->splits[deferred->consumer_split_id].backend_id; + ggml_backend_t consumer_backend = + sched->backends[consumer_backend_id]; + // The CUDA/HIP backend recognizes the deferred mode -3 + // destination and emits a producer-side write into + // coherent staging. Publishing the event immediately + // below orders that write before the consumer's already- + // late wait without putting a wait on the main stream + // ahead of hot/shared work. + if (!consumer_backend->iface.cpy_tensor_async || + !consumer_backend->iface.cpy_tensor_async( + split_backend, consumer_backend, + deferred->source, deferred->node)) { + GGML_LOG_ERROR( + "%s: outer capture requires deferred source-write copy\n", + __func__); + return GGML_STATUS_FAILED; + } + // Tell mode -3 to copy from coherent staging into its + // local destination instead of issuing a peer read. + ggml_set_op_params_i32(deferred->node, 6, 3); + } + ggml_backend_event_record(deferred->event, split_backend); + } + } + + // Graph dispatch has already inspected every node in this split. Do + // not leak the capture-only no-op marker into a later eager warmup or + // a freshly recaptured request generation. + if (sched->whole_graph_capture) { + for (int i = 0; i < sched->n_deferred_peer_copies; ++i) { + struct ggml_backend_sched_deferred_peer_copy * deferred = + &sched->deferred_peer_copies[i]; + if (deferred->consumer_split_id == split_id && + ggml_get_op_params_i32(deferred->node, 6) == 3) { + ggml_set_op_params_i32(deferred->node, 6, 0); + } + } + } + + // Record the rotating-buffer reuse event only for ordinary scheduler + // submissions. Whole-step capture deliberately keeps cur_copy fixed + // and suppresses the corresponding reuse waits, so recording this + // same event after every split only creates redundant external graph + // nodes (and triggers a HIP graph-capture failure on gfx1151/gfx1201). + if (!sched->whole_graph_capture && split->n_inputs > 0) { if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_record(sched->events[split_backend_id][sched->cur_copy], split_backend); } @@ -1835,6 +2131,11 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) { ggml_free(sched->ctx); ggml_hash_set_free(&sched->hash_set); free(sched->splits); + free(sched->late_cross_input_split_nodes); + for (int i = 0; i < sched->n_deferred_peer_copies; ++i) { + ggml_backend_event_free(sched->deferred_peer_copies[i].event); + } + free(sched->deferred_peer_copies); free(sched->hv_tensor_backend_ids); free(sched->hv_tensor_copies); free(sched->node_backend_ids); @@ -1904,6 +2205,29 @@ bool ggml_backend_sched_alloc_graph(ggml_backend_sched_t sched, struct ggml_cgra return false; } + // The deferred node intentionally has no foreign src edge by allocation + // time. Inject the stable peer allocation pointer only after gallocr has + // initialized every tensor buffer. + for (int i = 0; i < sched->n_deferred_peer_copies; ++i) { + struct ggml_backend_sched_deferred_peer_copy * deferred = + &sched->deferred_peer_copies[i]; + GGML_ASSERT(deferred->source->data); + void * source_data = deferred->source->data; + static_assert( + 4 * sizeof(int32_t) + sizeof(source_data) <= + GGML_MAX_OP_PARAMS, + "deferred source pointer does not fit op_params"); + memcpy( + &deferred->node->op_params[4], + &source_data, sizeof(source_data)); + + const char * host_copy = + getenv("DFLASH_DS4_TP_DEVICE_JOIN_HOST_COPY"); + ggml_set_op_params_i32( + deferred->node, 6, + host_copy && *host_copy && strcmp(host_copy, "0") != 0 ? 1 : 0); + } + sched->is_alloc = true; return true; @@ -1943,17 +2267,116 @@ void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) { } } +void ggml_backend_sched_set_whole_graph_capture( + ggml_backend_sched_t sched, bool enabled) { + GGML_ASSERT(sched); + sched->whole_graph_capture = enabled; +} + void ggml_backend_sched_set_eval_callback(ggml_backend_sched_t sched, ggml_backend_sched_eval_callback callback, void * user_data) { GGML_ASSERT(sched); sched->callback_eval = callback; sched->callback_eval_user_data = user_data; } +void ggml_backend_sched_set_late_cross_input_split(ggml_backend_sched_t sched, bool enabled) { + GGML_ASSERT(sched); + GGML_ASSERT(!sched->is_alloc); + sched->late_cross_input_split = enabled; +} + +void ggml_backend_sched_add_late_cross_input_split_node( + ggml_backend_sched_t sched, + const struct ggml_tensor * node) { + GGML_ASSERT(sched); + GGML_ASSERT(node); + GGML_ASSERT(!sched->is_alloc); + + if (ggml_backend_sched_is_late_cross_input_split_node(sched, node)) { + return; + } + if (sched->n_late_cross_input_split_nodes == + sched->late_cross_input_split_nodes_capacity) { + const int new_capacity = + sched->late_cross_input_split_nodes_capacity > 0 + ? 2 * sched->late_cross_input_split_nodes_capacity + : 64; + const struct ggml_tensor ** nodes = + (const struct ggml_tensor **) realloc( + sched->late_cross_input_split_nodes, + (size_t) new_capacity * sizeof(*nodes)); + GGML_ASSERT(nodes); + sched->late_cross_input_split_nodes = nodes; + sched->late_cross_input_split_nodes_capacity = new_capacity; + } + sched->late_cross_input_split_nodes[ + sched->n_late_cross_input_split_nodes++] = node; +} + +void ggml_backend_sched_add_deferred_peer_copy_node( + ggml_backend_sched_t sched, + const struct ggml_tensor * node) { + GGML_ASSERT(sched); + GGML_ASSERT(node); + GGML_ASSERT(!sched->is_alloc); + GGML_ASSERT(node->op == GGML_OP_MOE_FUSED); + GGML_ASSERT(ggml_get_op_params_i32(node, 0) == -3); + + if (ggml_backend_sched_find_deferred_peer_copy(sched, node)) { + return; + } + if (sched->n_deferred_peer_copies == + sched->deferred_peer_copies_capacity) { + const int new_capacity = + sched->deferred_peer_copies_capacity > 0 + ? 2 * sched->deferred_peer_copies_capacity + : 64; + struct ggml_backend_sched_deferred_peer_copy * entries = + (struct ggml_backend_sched_deferred_peer_copy *) realloc( + sched->deferred_peer_copies, + (size_t) new_capacity * sizeof(*entries)); + GGML_ASSERT(entries); + sched->deferred_peer_copies = entries; + sched->deferred_peer_copies_capacity = new_capacity; + } + struct ggml_backend_sched_deferred_peer_copy * deferred = + &sched->deferred_peer_copies[sched->n_deferred_peer_copies++]; + *deferred = {}; + deferred->node = const_cast(node); + GGML_ASSERT(deferred->node->src[0]); + deferred->source = deferred->node->src[0]; + // The full graph is already expanded when nodes are registered. Remove + // the foreign edge before backend splitting/allocation; source lifetime is + // explicit via GGML_TENSOR_FLAG_OUTPUT and the scheduler entry above. + deferred->node->src[0] = nullptr; + deferred->producer_backend_id = -1; + deferred->producer_split_id = -1; + deferred->consumer_split_id = -1; +} + int ggml_backend_sched_get_n_splits(ggml_backend_sched_t sched) { GGML_ASSERT(sched); return sched->n_splits; } +int ggml_backend_sched_get_n_splits_for_backend( + ggml_backend_sched_t sched, ggml_backend_t backend) { + GGML_ASSERT(sched); + int backend_id = -1; + for (int i = 0; i < sched->n_backends; ++i) { + if (sched->backends[i] == backend) { + backend_id = i; + break; + } + } + if (backend_id < 0) return 0; + int count = 0; + for (int i = 0; i < sched->n_splits; ++i) { + if (sched->splits[i].backend_id == backend_id) ++count; + } + return count; +} + int ggml_backend_sched_get_n_copies(ggml_backend_sched_t sched) { GGML_ASSERT(sched); return sched->n_copies; diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index db026f34e..c17cf1921 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -2094,6 +2094,15 @@ static struct ggml_tensor * ggml_add_impl( struct ggml_tensor * a, struct ggml_tensor * b, bool inplace) { + if (!ggml_can_repeat(b, a)) { + fprintf(stderr, + "[ggml-add-shape] a='%s' type=%s ne=[%" PRId64 ",%" PRId64 ",%" PRId64 ",%" PRId64 "] " + "b='%s' type=%s ne=[%" PRId64 ",%" PRId64 ",%" PRId64 ",%" PRId64 "] inplace=%d\n", + a->name, ggml_type_name(a->type), a->ne[0], a->ne[1], a->ne[2], a->ne[3], + b->name, ggml_type_name(b->type), b->ne[0], b->ne[1], b->ne[2], b->ne[3], + (int) inplace); + ggml_print_backtrace_symbols(); + } GGML_ASSERT(ggml_can_repeat(b, a)); struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); @@ -8076,6 +8085,124 @@ struct ggml_tensor * ggml_laguna_moe_combine( return result; } +struct ggml_tensor * ggml_ds4_moe_owner( + struct ggml_context * ctx, + struct ggml_tensor * input, + struct ggml_tensor * gate_up_w, + struct ggml_tensor * down_w, + struct ggml_tensor * expert_ids, + struct ggml_tensor * expert_weights, + int64_t ff_dim, + float swiglu_clamp, + float down_scale) { + GGML_ASSERT(input->type == GGML_TYPE_F32); + GGML_ASSERT(expert_ids->type == GGML_TYPE_I32); + GGML_ASSERT(expert_weights->type == GGML_TYPE_F32); + GGML_ASSERT(input->ne[1] == expert_ids->ne[1]); + GGML_ASSERT(input->ne[1] == expert_weights->ne[1]); + GGML_ASSERT(expert_ids->ne[0] == expert_weights->ne[0]); + GGML_ASSERT(gate_up_w->ne[0] == input->ne[0]); + GGML_ASSERT(gate_up_w->ne[1] == 2 * ff_dim); + GGML_ASSERT(down_w->ne[0] == ff_dim); + GGML_ASSERT(down_w->ne[1] == input->ne[0]); + GGML_ASSERT(gate_up_w->ne[2] == down_w->ne[2]); + + const int64_t ne[4] = { input->ne[0], input->ne[1], 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 2, ne); + + result->op = GGML_OP_MOE_FUSED; + result->src[0] = input; + result->src[1] = gate_up_w; + result->src[2] = down_w; + result->src[3] = expert_ids; + result->src[4] = expert_weights; + + ggml_set_op_params_i32(result, 0, -2); + ggml_set_op_params_i32(result, 1, (int32_t) ff_dim); + ggml_set_op_params_f32(result, 2, swiglu_clamp); + ggml_set_op_params_f32(result, 3, down_scale); + + return result; +} + +struct ggml_tensor * ggml_ds4_moe_owner_split( + struct ggml_context * ctx, + struct ggml_tensor * input, + struct ggml_tensor * gate_w, + struct ggml_tensor * up_w, + struct ggml_tensor * down_w, + struct ggml_tensor * expert_ids, + struct ggml_tensor * expert_weights, + int64_t ff_dim, + float swiglu_clamp, + float gate_scale, + float up_scale, + float down_scale) { + GGML_ASSERT(input->type == GGML_TYPE_F32); + GGML_ASSERT(expert_ids->type == GGML_TYPE_I32); + GGML_ASSERT(expert_weights->type == GGML_TYPE_F32); + GGML_ASSERT(input->ne[1] == expert_ids->ne[1]); + GGML_ASSERT(input->ne[1] == expert_weights->ne[1]); + GGML_ASSERT(expert_ids->ne[0] == expert_weights->ne[0]); + GGML_ASSERT(gate_w->ne[0] == input->ne[0]); + GGML_ASSERT(up_w->ne[0] == input->ne[0]); + GGML_ASSERT(gate_w->ne[1] == ff_dim && up_w->ne[1] == ff_dim); + GGML_ASSERT(down_w->ne[0] == ff_dim); + GGML_ASSERT(down_w->ne[1] == input->ne[0]); + GGML_ASSERT(gate_w->ne[2] == up_w->ne[2]); + GGML_ASSERT(gate_w->ne[2] == down_w->ne[2]); + + const int64_t ne[4] = { input->ne[0], input->ne[1], 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 2, ne); + + result->op = GGML_OP_MOE_FUSED; + result->src[0] = input; + result->src[1] = gate_w; + result->src[2] = up_w; + result->src[3] = down_w; + result->src[4] = expert_ids; + result->src[5] = expert_weights; + + ggml_set_op_params_i32(result, 0, -4); + ggml_set_op_params_i32(result, 1, (int32_t) ff_dim); + ggml_set_op_params_f32(result, 2, swiglu_clamp); + ggml_set_op_params_f32(result, 3, gate_scale); + ggml_set_op_params_f32(result, 4, up_scale); + ggml_set_op_params_f32(result, 5, down_scale); + + return result; +} + +struct ggml_tensor * ggml_ds4_moe_align_ids( + struct ggml_context * ctx, + struct ggml_tensor * expert_ids) { + GGML_ASSERT(expert_ids->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_n_dims(expert_ids) == 2); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, expert_ids); + result->op = GGML_OP_MOE_FUSED; + result->src[0] = expert_ids; + ggml_set_op_params_i32(result, 0, -5); + return result; +} + +struct ggml_tensor * ggml_ds4_deferred_peer_copy( + struct ggml_context * ctx, + struct ggml_tensor * src) { + GGML_ASSERT(src->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(src)); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, src); + result->op = GGML_OP_MOE_FUSED; + result->src[0] = src; + + // op_params[2..3] is populated by the backend scheduler with the native + // event handle after backend assignment and before graph allocation. + ggml_set_op_params_i32(result, 0, -3); + + return result; +} + struct ggml_tensor * ggml_ds4_hc_pre( struct ggml_context * ctx, struct ggml_tensor * mix, @@ -8144,6 +8271,38 @@ struct ggml_tensor * ggml_ds4_hc_post( return result; } +struct ggml_tensor * ggml_ds4_hc_post_split( + struct ggml_context * ctx, + struct ggml_tensor * residual_hc, + struct ggml_tensor * main_block, + struct ggml_tensor * peer_block, + struct ggml_tensor * split, + int n_hc) { + GGML_ASSERT(residual_hc->type == GGML_TYPE_F32); + GGML_ASSERT(main_block->type == GGML_TYPE_F32); + GGML_ASSERT(peer_block->type == GGML_TYPE_F32); + GGML_ASSERT(split->type == GGML_TYPE_F32); + GGML_ASSERT(n_hc > 0 && n_hc <= 8); + const int64_t mix_dim = 2*(int64_t)n_hc + (int64_t)n_hc*n_hc; + GGML_ASSERT(ggml_nelements(split) == mix_dim); + GGML_ASSERT(ggml_nelements(residual_hc) % n_hc == 0); + const int64_t n_embd = ggml_nelements(residual_hc) / n_hc; + GGML_ASSERT(ggml_nelements(main_block) == n_embd); + GGML_ASSERT(ggml_nelements(peer_block) == n_embd); + + struct ggml_tensor * result = ggml_new_tensor_1d( + ctx, GGML_TYPE_F32, (int64_t) n_embd * n_hc); + result->op = GGML_OP_DS4_HC; + result->src[0] = residual_hc; + result->src[1] = main_block; + result->src[2] = split; + result->src[3] = peer_block; + ggml_set_op_params_i32(result, 0, 3); + ggml_set_op_params_i32(result, 1, (int32_t) n_embd); + ggml_set_op_params_i32(result, 2, (int32_t) n_hc); + return result; +} + struct ggml_tensor * ggml_ds4_hc_out( struct ggml_context * ctx, struct ggml_tensor * mix, diff --git a/server/src/common/backend_ipc.cpp b/server/src/common/backend_ipc.cpp index c52dd591d..15a4d9a6a 100644 --- a/server/src/common/backend_ipc.cpp +++ b/server/src/common/backend_ipc.cpp @@ -321,8 +321,11 @@ bool BackendIpcProcess::write_shared_payload_segments( } } seq = ++shared_payload_seq_; - header->bytes = (uint64_t)bytes; - header->sequence = seq; + // The command pipe publishes this request to another process. Publish the + // mapped payload header explicitly so the daemon cannot observe the prior + // response header after receiving the new command. + backend_ipc_publish_shared_payload_header( + header, seq, static_cast(bytes)); return true; } @@ -331,7 +334,8 @@ bool BackendIpcProcess::read_shared_payload(void * data, size_t bytes, uint64_t if (bytes > 0 && !data) return false; const auto * header = static_cast(shared_payload_map_); - if (header->sequence != seq || header->bytes != (uint64_t)bytes) { + if (!backend_ipc_shared_payload_header_matches( + header, seq, static_cast(bytes))) { return false; } const void * payload = static_cast( diff --git a/server/src/common/backend_ipc.h b/server/src/common/backend_ipc.h index 47a8a5c08..36c2f47d0 100644 --- a/server/src/common/backend_ipc.h +++ b/server/src/common/backend_ipc.h @@ -7,6 +7,7 @@ #pragma once +#include #include #include #include @@ -65,6 +66,46 @@ struct BackendIpcSharedPayloadHeader { uint64_t bytes = 0; }; +inline void backend_ipc_publish_shared_payload_header( + BackendIpcSharedPayloadHeader * header, + uint64_t sequence, + uint64_t bytes) { +#if defined(__GNUC__) || defined(__clang__) + __atomic_store_n(&header->bytes, bytes, __ATOMIC_RELAXED); + __atomic_store_n(&header->sequence, sequence, __ATOMIC_RELEASE); +#else + header->bytes = bytes; + std::atomic_thread_fence(std::memory_order_release); + header->sequence = sequence; +#endif +} + +inline void backend_ipc_load_shared_payload_header( + const BackendIpcSharedPayloadHeader * header, + uint64_t & sequence, + uint64_t & bytes) { +#if defined(__GNUC__) || defined(__clang__) + sequence = __atomic_load_n(&header->sequence, __ATOMIC_ACQUIRE); + bytes = __atomic_load_n(&header->bytes, __ATOMIC_RELAXED); +#else + sequence = header->sequence; + std::atomic_thread_fence(std::memory_order_acquire); + bytes = header->bytes; +#endif +} + +inline bool backend_ipc_shared_payload_header_matches( + const BackendIpcSharedPayloadHeader * header, + uint64_t sequence, + uint64_t bytes) { + if (!header || sequence == 0) return false; + uint64_t published_sequence = 0; + uint64_t published_bytes = 0; + backend_ipc_load_shared_payload_header( + header, published_sequence, published_bytes); + return published_sequence == sequence && published_bytes == bytes; +} + inline size_t backend_ipc_shared_payload_header_bytes() { return sizeof(BackendIpcSharedPayloadHeader); } diff --git a/server/src/common/moe_expert_compute.cpp b/server/src/common/moe_expert_compute.cpp index 0bcddbf11..fc42fc787 100644 --- a/server/src/common/moe_expert_compute.cpp +++ b/server/src/common/moe_expert_compute.cpp @@ -47,6 +47,7 @@ std::string make_runtime_key(const MoeExpertComputeRuntimeConfig & cfg) { key += "|n_used=" + std::to_string(cfg.n_expert_used); key += "|n_embd=" + std::to_string(cfg.n_embd); key += "|n_ff=" + std::to_string(cfg.n_ff_exp); + key += "|require_remote=" + std::to_string(cfg.require_remote ? 1 : 0); if (const char * ipc_bin = nonempty_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_BIN")) { key += "|ipc_bin="; key += ipc_bin; @@ -202,7 +203,7 @@ bool ensure_moe_expert_compute_runtime( const char * work_dir = nonempty_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_WORK_DIR"); const int remote_gpu = parse_nonnegative_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU", 0); - const bool required = + const bool required = cfg.require_remote || parse_nonnegative_env("DFLASH_MOE_EXPERT_COMPUTE_IPC_REQUIRED", 0) != 0; MoeExpertComputeIpcStartResult remote = make_moe_expert_compute_ipc( ipc_bin, cfg.target_path, remote_gpu, hybrid.placement, @@ -216,6 +217,13 @@ bool ensure_moe_expert_compute_runtime( return false; } runtime.compute = std::move(remote.compute); + } else if (cfg.require_remote) { + if (err) *err = "remote MoE expert compute is required but IPC binary is unset"; + std::fprintf(stderr, "%s %s\n", + cfg.log_prefix ? cfg.log_prefix : "[moe-expert-compute]", + err ? err->c_str() : "remote IPC binary is unset"); + runtime.reset(); + return false; } if (!runtime.compute) { runtime.compute = make_cpu_moe_expert_compute(cfg.n_ff_exp); diff --git a/server/src/common/moe_expert_compute.h b/server/src/common/moe_expert_compute.h index 7cee6a599..ce08ee1a0 100644 --- a/server/src/common/moe_expert_compute.h +++ b/server/src/common/moe_expert_compute.h @@ -95,6 +95,44 @@ struct MoeExpertCompute { } return true; } + + // Batched variant with a per-token selected width. ids/weights use a + // fixed max_selected stride, while selected_counts controls the active + // prefix of each row. IPC implementations use this to keep one request + // per layer without padding expert work or grouping rows on the client. + virtual bool compute_batch_ragged( + const MoeExpertLayer & layer, + const float * input, + const int32_t * ids, + const float * weights, + const int32_t * selected_counts, + int n_tokens, + int max_selected, + int n_embd, + int n_ff, + float * output) { + if (n_tokens < 0 || max_selected <= 0 || n_embd <= 0 || !output || + (n_tokens > 0 && (!input || !ids || !weights || !selected_counts))) { + return false; + } + for (int t = 0; t < n_tokens; ++t) { + const int n_selected = selected_counts[t]; + float * token_output = output + (size_t)t * (size_t)n_embd; + if (n_selected < 0 || n_selected > max_selected) return false; + if (n_selected == 0) { + std::fill(token_output, token_output + n_embd, 0.0f); + continue; + } + if (!compute(layer, + input + (size_t)t * (size_t)n_embd, + ids + (size_t)t * (size_t)max_selected, + weights + (size_t)t * (size_t)max_selected, + n_selected, n_embd, n_ff, token_output)) { + return false; + } + } + return true; + } }; // Create CPU-based fused MoE expert compute. @@ -136,6 +174,7 @@ struct MoeExpertComputeRuntimeConfig { int n_embd = 0; int n_ff_exp = 0; bool enabled = true; + bool require_remote = false; const char * log_prefix = "[moe-expert-compute]"; }; diff --git a/server/src/common/moe_expert_compute_ipc.cpp b/server/src/common/moe_expert_compute_ipc.cpp index a74f47344..b1b9edcac 100644 --- a/server/src/common/moe_expert_compute_ipc.cpp +++ b/server/src/common/moe_expert_compute_ipc.cpp @@ -14,6 +14,7 @@ #include "ggml_graph_precision.h" #include "internal.h" #include "laguna_internal.h" +#include "../deepseek4/deepseek4_internal.h" #include "moe_hybrid_types_impl.h" #include "ggml-backend.h" @@ -31,6 +32,7 @@ #include #include #include +#include #include #include #include @@ -183,6 +185,30 @@ struct RemoteMoeRuntime { MoeHybridStorage hybrid; }; +MoeHybridConfig make_moe_hybrid_config(const DeepSeek4Weights & weights) { + MoeHybridConfig cfg; + cfg.n_embd = weights.n_embd; + cfg.n_expert = weights.n_expert; + cfg.n_expert_used = weights.n_expert_used; + cfg.n_ff_exp = weights.n_ff_exp; + cfg.n_ff_shexp = weights.n_ff_exp; + cfg.n_layer = weights.n_layer; + cfg.first_moe_layer = 0; + cfg.swiglu_clamp = weights.swiglu_clamp_exp; + return cfg; +} + +MoeLayerDesc make_moe_layer_desc(const DeepSeek4Layer & layer) { + MoeLayerDesc desc; + desc.ffn_gate_exps = layer.ffn_gate_exps; + desc.ffn_up_exps = layer.ffn_up_exps; + desc.ffn_down_exps = layer.ffn_down_exps; + desc.ffn_gate_shexp = layer.ffn_gate_shexp; + desc.ffn_up_shexp = layer.ffn_up_shexp; + desc.ffn_down_shexp = layer.ffn_down_shexp; + return desc; +} + bool build_cached_batched_cold_graph( CachedBatchedMoeExpertGraph & out, ggml_backend_t backend, @@ -460,6 +486,7 @@ size_t moe_expert_required_shared_bytes(int n_embd, size_t input_bytes = 0; size_t ids_bytes = 0; size_t weights_bytes = 0; + size_t counts_bytes = 0; size_t total = 0; size_t input_elems = 0; size_t selected_elems = 0; @@ -471,7 +498,10 @@ size_t moe_expert_required_shared_bytes(int n_embd, input_bytes) || !moe_expert_compute_checked_mul_size(selected_elems, sizeof(int32_t), ids_bytes) || !moe_expert_compute_checked_mul_size(selected_elems, sizeof(float), weights_bytes) || - !backend_ipc_checked_add_size(input_bytes, ids_bytes, total) || + !moe_expert_compute_checked_mul_size( + (size_t)batch_limit, sizeof(int32_t), counts_bytes) || + !backend_ipc_checked_add_size(input_bytes, counts_bytes, total) || + !backend_ipc_checked_add_size(total, ids_bytes, total) || !backend_ipc_checked_add_size(total, weights_bytes, total)) { return 0; } @@ -683,6 +713,17 @@ bool load_remote_moe_runtime(const char * target_path, backend, placement, gctx, file_bytes, file_size, data_start, arch, weights, free_laguna_target_weights, out, err); } + if (arch == "deepseek4") { + DeepSeek4Weights weights; + if (!load_deepseek4_gguf_partial(target_path, backend, plan, weights)) { + if (err) *err = "failed to load DeepSeek4 expert metadata"; + free_deepseek4_weights(weights); + return false; + } + return build_remote_moe_runtime_from_weights( + backend, placement, gctx, file_bytes, file_size, + data_start, arch, weights, free_deepseek4_weights, out, err); + } if (err) *err = "unsupported MoE expert compute arch: " + arch; return false; @@ -697,7 +738,8 @@ bool validate_shared_payload_request(const void * shared_payload, static_cast(shared_payload); if (!shared_payload || !shared_payload_data || seq == 0 || !backend_ipc_payload_in_bounds(0, bytes, shared_payload_capacity) || - header->sequence != seq || header->bytes != (uint64_t)bytes) { + !backend_ipc_shared_payload_header_matches( + header, seq, static_cast(bytes))) { return false; } return true; @@ -713,8 +755,8 @@ bool commit_shared_payload_response(void * shared_payload, return false; } auto * header = static_cast(shared_payload); - header->bytes = (uint64_t)bytes; - header->sequence = seq; + backend_ipc_publish_shared_payload_header( + header, seq, static_cast(bytes)); return true; } @@ -806,6 +848,7 @@ class MoeExpertComputeIpc : public MoeExpertCompute { std::fill(output, output + n_embd, 0.0f); return true; } + std::lock_guard request_lock(request_mutex_); if (!active_ || layer.layer_idx < 0 || n_embd != n_embd_) { if (required_) return false; return fallback_->compute(layer, input, ids, weights, n_cold, n_embd, n_ff, output); @@ -848,6 +891,7 @@ class MoeExpertComputeIpc : public MoeExpertCompute { std::fill(output, output + (size_t)n_tokens * (size_t)n_embd, 0.0f); return true; } + std::lock_guard request_lock(request_mutex_); if (!active_ || layer.layer_idx < 0 || n_embd != n_embd_) { if (required_) return false; return fallback_->compute_batch(layer, input, ids, weights, @@ -884,7 +928,75 @@ class MoeExpertComputeIpc : public MoeExpertCompute { return true; } + bool compute_batch_ragged(const MoeExpertLayer & layer, + const float * input, + const int32_t * ids, + const float * weights, + const int32_t * selected_counts, + int n_tokens, + int max_selected, + int n_embd, + int n_ff, + float * output) override { + if (!output || n_tokens < 0 || max_selected <= 0 || + max_selected > n_expert_used_ || n_embd <= 0 || + (n_tokens > 0 && (!input || !ids || !weights || !selected_counts))) { + return false; + } + if (n_tokens == 0) return true; + + bool has_selected = false; + for (int t = 0; t < n_tokens; ++t) { + const int count = selected_counts[t]; + if (count < 0 || count > max_selected) return false; + has_selected = has_selected || count > 0; + } + if (!has_selected) { + std::fill(output, + output + (size_t)n_tokens * (size_t)n_embd, 0.0f); + return true; + } + + std::lock_guard request_lock(request_mutex_); + if (!active_ || layer.layer_idx < 0 || n_embd != n_embd_) { + if (required_) return false; + return fallback_->compute_batch_ragged( + layer, input, ids, weights, selected_counts, + n_tokens, max_selected, n_embd, n_ff, output); + } + for (int t = 0; t < n_tokens; ++t) { + for (int i = 0; i < selected_counts[t]; ++i) { + const size_t idx = + (size_t)t * (size_t)max_selected + (size_t)i; + const int32_t local = ids[idx]; + if (local < 0 || + (size_t)local >= layer.cold_global_by_local.size()) { + if (required_) return false; + return fallback_->compute_batch_ragged( + layer, input, ids, weights, selected_counts, + n_tokens, max_selected, n_embd, n_ff, output); + } + } + } + + if (!compute_remote_batch_ragged( + layer.layer_idx, input, ids, weights, selected_counts, + n_tokens, max_selected, n_embd, output)) { + active_ = false; + if (required_) { + std::fprintf(stderr, + "[moe-expert-compute-ipc] required ragged remote compute failed\n"); + return false; + } + return fallback_->compute_batch_ragged( + layer, input, ids, weights, selected_counts, + n_tokens, max_selected, n_embd, n_ff, output); + } + return true; + } + bool healthy() const override { + std::lock_guard request_lock(request_mutex_); return active_; } @@ -1082,7 +1194,121 @@ class MoeExpertComputeIpc : public MoeExpertCompute { int32_t status = -1; const auto wait_t0 = MoeExpertClock::now(); if (!read_exact_fd(stream_fd, &status, sizeof(status)) || status != 0) { - std::fprintf(stderr, "[moe-expert-compute-ipc] batch compute failed status=%d\n", status); + std::fprintf(stderr, + "[moe-expert-compute-ipc] batch compute failed " + "layer=%d tokens=%d selected=%d status=%d\n", + layer_idx, n_tokens, n_selected, status); + return false; + } + const auto wait_t1 = MoeExpertClock::now(); + const size_t output_bytes = + (size_t)n_tokens * (size_t)n_embd * sizeof(float); + const bool ok = use_shared + ? process_.read_shared_payload(output, output_bytes, seq) + : read_exact_fd(stream_fd, output, output_bytes); + const auto recv_t1 = MoeExpertClock::now(); + if (profile_) { + stats_.calls++; + stats_.payload_bytes += payload_bytes; + stats_.pack_us += moe_expert_compute_elapsed_us(pack_t0, pack_t1); + stats_.send_us += moe_expert_compute_elapsed_us(send_t0, send_t1); + stats_.wait_us += moe_expert_compute_elapsed_us(wait_t0, wait_t1); + stats_.recv_us += moe_expert_compute_elapsed_us(wait_t1, recv_t1); + stats_.maybe_print(false); + } + return ok; +#endif + } + + bool compute_remote_batch_ragged(int layer_idx, + const float * input, + const int32_t * local_ids, + const float * weights, + const int32_t * selected_counts, + int n_tokens, + int max_selected, + int n_embd, + float * output) { +#if defined(_WIN32) + (void)layer_idx; (void)input; (void)local_ids; (void)weights; + (void)selected_counts; (void)n_tokens; (void)max_selected; + (void)n_embd; (void)output; + return false; +#else + FILE * cmd = process_.command_stream(); + const int stream_fd = process_.stream_fd(); + const int payload_fd = process_.payload_fd(); + if (!cmd || stream_fd < 0 || layer_idx < 0 || !input || + !local_ids || !weights || !selected_counts || !output || + n_tokens <= 0 || max_selected <= 0 || + max_selected > n_expert_used_) { + return false; + } + + const ggml_type request_input_type = GGML_TYPE_F32; + const void * ipc_input = nullptr; + size_t input_bytes = 0; + if (!moe_expert_convert_input_to_ipc_type( + request_input_type, input, n_tokens, n_embd, + input_scratch_, ipc_input, input_bytes)) { + return false; + } + const size_t counts_bytes = + (size_t)n_tokens * sizeof(int32_t); + const size_t selected_elems = + (size_t)n_tokens * (size_t)max_selected; + const size_t ids_bytes = selected_elems * sizeof(int32_t); + const size_t weights_bytes = selected_elems * sizeof(float); + size_t payload_bytes = 0; + if (!backend_ipc_checked_add_size(input_bytes, counts_bytes, payload_bytes) || + !backend_ipc_checked_add_size(payload_bytes, ids_bytes, payload_bytes) || + !backend_ipc_checked_add_size(payload_bytes, weights_bytes, payload_bytes)) { + return false; + } + const bool use_shared = + process_.resolved_payload_transport() == BackendIpcPayloadTransport::Shared; + const auto pack_t0 = MoeExpertClock::now(); + uint64_t seq = 0; + const BackendIpcPayloadSegment segments[] = { + {ipc_input, input_bytes}, + {selected_counts, counts_bytes}, + {local_ids, ids_bytes}, + {weights, weights_bytes}, + }; + const auto pack_t1 = MoeExpertClock::now(); + + const auto send_t0 = MoeExpertClock::now(); + if (use_shared) { + if (!process_.write_shared_payload_segments(segments, 4, seq)) { + return false; + } + std::fprintf(cmd, + "compute_batch_ragged_local_shared %d %d %d %zu %" PRIu64 "\n", + layer_idx, n_tokens, max_selected, payload_bytes, seq); + std::fflush(cmd); + } else if (payload_fd >= 0) { + std::fprintf(cmd, + "compute_batch_ragged_local_pipe %d %d %d %zu\n", + layer_idx, n_tokens, max_selected, payload_bytes); + std::fflush(cmd); + if (!write_exact_fd(payload_fd, ipc_input, input_bytes) || + !write_exact_fd(payload_fd, selected_counts, counts_bytes) || + !write_exact_fd(payload_fd, local_ids, ids_bytes) || + !write_exact_fd(payload_fd, weights, weights_bytes)) { + return false; + } + } else { + return false; + } + const auto send_t1 = MoeExpertClock::now(); + + int32_t status = -1; + const auto wait_t0 = MoeExpertClock::now(); + if (!read_exact_fd(stream_fd, &status, sizeof(status)) || status != 0) { + std::fprintf(stderr, + "[moe-expert-compute-ipc] ragged batch compute failed " + "layer=%d tokens=%d max_selected=%d status=%d\n", + layer_idx, n_tokens, max_selected, status); return false; } const auto wait_t1 = MoeExpertClock::now(); @@ -1116,6 +1342,7 @@ class MoeExpertComputeIpc : public MoeExpertCompute { bool active_ = false; bool required_ = false; bool profile_ = false; + mutable std::mutex request_mutex_; }; } // namespace @@ -1274,15 +1501,23 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, std::vector output((size_t)runtime.n_embd); std::vector batch_local_ids; std::vector payload_ids; + std::vector payload_counts; std::vector payload_weights; MoeExpertDaemonStats profile_stats; MoeExpertDaemonStats profile_stats_prefill; MoeExpertDaemonStats profile_stats_decode; const bool profile = moe_expert_compute_profile_enabled(); + const char * token_loop_env = + std::getenv("DFLASH_MOE_EXPERT_COMPUTE_DAEMON_TOKEN_LOOP"); + const bool token_loop_batches = + token_loop_env && *token_loop_env && *token_loop_env != '0'; int warmup_builds = 0; std::fprintf(stderr, - "[moe-expert-compute-daemon] ready gpu=%d remote_hot=%d layers=%d warmup_graphs=%d warmup_ms=0.000\n", - target_gpu, runtime.hybrid.placement.total_hot, runtime.n_layer, warmup_builds); + "[moe-expert-compute-daemon] ready gpu=%d remote_hot=%d " + "layers=%d warmup_graphs=%d warmup_ms=0.000 token_loop=%d\n", + target_gpu, runtime.hybrid.placement.total_hot, + runtime.n_layer, warmup_builds, + token_loop_batches ? 1 : 0); emit_status(stream_fd, 0); std::string line; @@ -1303,6 +1538,7 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, bool pipe_payload = false; bool shared_payload_cmd = false; bool ids_are_local = false; + bool ragged_payload = false; ggml_type input_type = GGML_TYPE_F32; const auto read_t0 = MoeExpertClock::now(); if (cmd == "compute_pipe") { @@ -1357,7 +1593,20 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, shared_payload_cmd = true; ids_are_local = true; input_type = (ggml_type)input_type_i; + } else if (cmd == "compute_batch_ragged_local_pipe") { + iss >> layer_idx >> n_tokens >> n_selected >> bytes; + pipe_payload = true; + ids_are_local = true; + ragged_payload = true; + } else if (cmd == "compute_batch_ragged_local_shared") { + iss >> layer_idx >> n_tokens >> n_selected >> bytes >> shared_seq; + shared_payload_cmd = true; + ids_are_local = true; + ragged_payload = true; } else { + std::fprintf(stderr, + "[moe-expert-compute-daemon] unknown command: %s\n", + line.c_str()); emit_status(stream_fd, -1); continue; } @@ -1367,6 +1616,7 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, size_t input_bytes = 0; size_t ids_bytes = 0; size_t weights_bytes = 0; + size_t counts_bytes = 0; size_t expected = 0; bool header_ok = false; if (n_tokens > 0 && n_selected > 0 && @@ -1389,17 +1639,23 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, ids_bytes) && moe_expert_compute_checked_mul_size(selected_elems, sizeof(float), weights_bytes) && - backend_ipc_checked_add_size(input_bytes, ids_bytes, expected) && + (!ragged_payload || moe_expert_compute_checked_mul_size( + (size_t)n_tokens, sizeof(int32_t), counts_bytes)) && + backend_ipc_checked_add_size(input_bytes, counts_bytes, expected) && + backend_ipc_checked_add_size(expected, ids_bytes, expected) && backend_ipc_checked_add_size(expected, weights_bytes, expected) && bytes == expected; } if (pipe_payload && header_ok) { input.resize(input_bytes); + payload_counts.resize(ragged_payload ? (size_t)n_tokens : 0); payload_ids.resize((size_t)n_tokens * (size_t)n_selected); payload_weights.resize((size_t)n_tokens * (size_t)n_selected); payload_ok = payload_fd >= 0 && read_exact_fd(payload_fd, input.data(), input_bytes) && + (!ragged_payload || read_exact_fd( + payload_fd, payload_counts.data(), counts_bytes)) && read_exact_fd(payload_fd, payload_ids.data(), ids_bytes) && read_exact_fd(payload_fd, payload_weights.data(), weights_bytes); } else if (pipe_payload) { @@ -1414,6 +1670,22 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, const auto read_t1 = MoeExpertClock::now(); if (!header_ok || !payload_ok) { + uint64_t published_seq = 0; + uint64_t published_bytes = 0; + if (shared_payload_cmd && shared_payload) { + backend_ipc_load_shared_payload_header( + static_cast(shared_payload), + published_seq, published_bytes); + } + std::fprintf(stderr, + "[moe-expert-compute-daemon] invalid request " + "cmd=%s layer=%d tokens=%d selected=%d bytes=%zu " + "expected=%zu header=%d payload=%d seq=%" PRIu64 + " published_seq=%" PRIu64 " published_bytes=%" PRIu64 "\n", + cmd.c_str(), layer_idx, n_tokens, n_selected, + bytes, expected, header_ok ? 1 : 0, + payload_ok ? 1 : 0, shared_seq, + published_seq, published_bytes); emit_status(stream_fd, -1); continue; } @@ -1423,10 +1695,12 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, output.resize((size_t)n_tokens * (size_t)runtime.n_embd); } const uint8_t * input_data = nullptr; + const int32_t * selected_count_data = nullptr; const int32_t * payload_id_data = nullptr; const float * in_weights = nullptr; if (pipe_payload) { input_data = input.data(); + selected_count_data = ragged_payload ? payload_counts.data() : nullptr; payload_id_data = payload_ids.data(); in_weights = payload_weights.data(); } else { @@ -1434,6 +1708,11 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, static_cast(shared_payload_data); input_data = payload_data + off; off += input_bytes; + if (ragged_payload) { + selected_count_data = + reinterpret_cast(payload_data + off); + off += counts_bytes; + } payload_id_data = reinterpret_cast(payload_data + off); off += ids_bytes; in_weights = reinterpret_cast(payload_data + off); @@ -1443,13 +1722,33 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, const auto unpack_t1 = MoeExpertClock::now(); const int remote_hot_count = remote_hot_expert_count(storage); + bool counts_ok = !ragged_payload || selected_count_data != nullptr; + for (int t = 0; ragged_payload && counts_ok && t < n_tokens; ++t) { + counts_ok = selected_count_data[t] >= 0 && + selected_count_data[t] <= n_selected; + } + if (!counts_ok) { + std::fprintf(stderr, + "[moe-expert-compute-daemon] invalid ragged counts " + "layer=%d tokens=%d max_selected=%d\n", + layer_idx, n_tokens, n_selected); + emit_status(stream_fd, -1); + continue; + } + auto & layer_graphs = graphs[(size_t)layer_idx]; if ((size_t)n_selected >= layer_graphs.size()) { + std::fprintf(stderr, + "[moe-expert-compute-daemon] selected width out of range " + "layer=%d selected=%d capacity=%zu\n", + layer_idx, n_selected, layer_graphs.size()); emit_status(stream_fd, -1); continue; } const RemoteMoeLayerRuntime & L = runtime.layers[(size_t)layer_idx]; - if (n_tokens > 1 || input_type != GGML_TYPE_F32) { + if (!ragged_payload && + ((n_tokens > 1 && !token_loop_batches) || + input_type != GGML_TYPE_F32)) { auto & layer_batch_graphs = batch_graphs[(size_t)layer_idx]; const uint64_t graph_key = moe_expert_compute_batch_graph_key(n_tokens, n_selected, @@ -1467,6 +1766,11 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, L.down_scale, L.gate_up_scale, runtime.n_embd, runtime.n_ff_exp, n_selected, n_tokens, input_type)) { + std::fprintf(stderr, + "[moe-expert-compute-daemon] batched graph build failed " + "layer=%d tokens=%d selected=%d type=%s\n", + layer_idx, n_tokens, n_selected, + ggml_type_name(input_type)); emit_status(stream_fd, -1); continue; } @@ -1505,6 +1809,10 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, } } if (!ids_ok) { + std::fprintf(stderr, + "[moe-expert-compute-daemon] batched local-id validation failed " + "layer=%d tokens=%d selected=%d remote_hot=%d\n", + layer_idx, n_tokens, n_selected, remote_hot_count); emit_status(stream_fd, -1); continue; } @@ -1517,6 +1825,10 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, auto st = ggml_backend_graph_compute(backend, graph.gf); const auto compute_one_t1 = MoeExpertClock::now(); if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[moe-expert-compute-daemon] batched graph compute failed " + "layer=%d tokens=%d selected=%d status=%d\n", + layer_idx, n_tokens, n_selected, (int)st); emit_status(stream_fd, -1); continue; } @@ -1576,35 +1888,51 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, continue; } - CachedFfnGraph & graph = layer_graphs[(size_t)n_selected]; - if (!graph.valid() || graph.n_hot != n_selected) { - const auto build_t0 = MoeExpertClock::now(); - if (!build_cached_cold_graph(graph, backend, - storage.gate_hot, storage.up_hot, - storage.down_hot, storage.gate_up_hot, - L.gate_scale, L.up_scale, - L.down_scale, L.gate_up_scale, - runtime.n_embd, runtime.n_ff_exp, - n_selected)) { - emit_status(stream_fd, -1); - continue; - } - const auto build_t1 = MoeExpertClock::now(); - if (profile) { - profile_stats_decode.graph_builds++; - profile_stats_decode.graph_build_us += - moe_expert_compute_elapsed_us(build_t0, build_t1); - } - } - bool compute_ok = true; uint64_t set_us_accum = 0; uint64_t compute_us_accum = 0; uint64_t get_us_accum = 0; const size_t one_input_bytes = (size_t)runtime.n_embd * sizeof(float); - const size_t one_ids_bytes = (size_t)n_selected * sizeof(int32_t); - const size_t one_weights_bytes = (size_t)n_selected * sizeof(float); for (int t = 0; t < n_tokens; ++t) { + const int token_selected = ragged_payload + ? selected_count_data[t] : n_selected; + void * response_data = shared_payload_cmd + ? static_cast( + static_cast(shared_payload_data) + + (size_t)t * one_input_bytes) + : static_cast( + output.data() + (size_t)t * (size_t)runtime.n_embd); + if (token_selected == 0) { + std::memset(response_data, 0, one_input_bytes); + continue; + } + + CachedFfnGraph & graph = layer_graphs[(size_t)token_selected]; + if (!graph.valid() || graph.n_hot != token_selected) { + const auto build_t0 = MoeExpertClock::now(); + if (!build_cached_cold_graph( + graph, backend, + storage.gate_hot, storage.up_hot, + storage.down_hot, storage.gate_up_hot, + L.gate_scale, L.up_scale, + L.down_scale, L.gate_up_scale, + runtime.n_embd, runtime.n_ff_exp, + token_selected)) { + std::fprintf(stderr, + "[moe-expert-compute-daemon] single graph build failed " + "layer=%d tokens=%d selected=%d\n", + layer_idx, n_tokens, token_selected); + compute_ok = false; + break; + } + const auto build_t1 = MoeExpertClock::now(); + if (profile) { + profile_stats_decode.graph_builds++; + profile_stats_decode.graph_build_us += + moe_expert_compute_elapsed_us(build_t0, build_t1); + } + } + bool ids_ok = true; const size_t token_off = (size_t)t * (size_t)n_selected; const int32_t * token_ids = payload_id_data + token_off; @@ -1613,7 +1941,7 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, if (!ids_are_local) { graph_ids_data = local_ids.data(); } - for (int i = 0; i < n_selected; ++i) { + for (int i = 0; i < token_selected; ++i) { const int32_t id = token_ids[i]; const int32_t local = ids_are_local ? id @@ -1628,6 +1956,10 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, } } if (!ids_ok) { + std::fprintf(stderr, + "[moe-expert-compute-daemon] single local-id validation failed " + "layer=%d token=%d selected=%d remote_hot=%d\n", + layer_idx, t, token_selected, remote_hot_count); compute_ok = false; break; } @@ -1637,22 +1969,20 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, (size_t)t * one_input_bytes, 0, one_input_bytes); ggml_backend_tensor_set(graph.ids, graph_ids_data, 0, - one_ids_bytes); + (size_t)token_selected * sizeof(int32_t)); ggml_backend_tensor_set(graph.weights, token_weights, 0, - one_weights_bytes); + (size_t)token_selected * sizeof(float)); const auto set_one_t1 = MoeExpertClock::now(); auto st = ggml_backend_graph_compute(backend, graph.gf); const auto compute_one_t1 = MoeExpertClock::now(); if (st != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[moe-expert-compute-daemon] single graph compute failed " + "layer=%d token=%d selected=%d status=%d\n", + layer_idx, t, token_selected, (int)st); compute_ok = false; break; } - void * response_data = shared_payload_cmd - ? static_cast( - static_cast(shared_payload_data) + - (size_t)t * one_input_bytes) - : static_cast( - output.data() + (size_t)t * (size_t)runtime.n_embd); ggml_backend_tensor_get(graph.output, response_data, 0, one_input_bytes); const auto get_one_t1 = MoeExpertClock::now(); @@ -1671,6 +2001,11 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, !commit_shared_payload_response(shared_payload, shared_payload_data, shared_payload_capacity, input_bytes, shared_seq)) { + std::fprintf(stderr, + "[moe-expert-compute-daemon] single shared response commit failed " + "layer=%d tokens=%d selected=%d bytes=%zu seq=%" PRIu64 "\n", + layer_idx, n_tokens, n_selected, + input_bytes, shared_seq); emit_status(stream_fd, -1); continue; } diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index 671010c88..9d7cb69a7 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -39,6 +39,127 @@ static uint64_t elapsed_us(HybridClock::time_point start, HybridClock::time_poin return (uint64_t) std::chrono::duration_cast(end - start).count(); } +static bool compact_materialized_experts_enabled() { + static const bool enabled = [] { + const char * raw = std::getenv("DFLASH_MOE_COMPACT_MATERIALIZED"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +// The legacy ROCmFP2 verifier graph expands a q-token routed FFN into q +// independent single-token subgraphs. That was required before the CUDA/HIP +// backend gained its grouped MUL_MAT_ID MMVQ kernel, but it multiplies graph +// nodes, scheduler copies, and launches by the verify width. Keep the old +// lowering as the default while the grouped path is qualified on each ROCm +// architecture; opt in with DFLASH_DS4_TP_GROUPED_MMVQ=1. +static bool grouped_mmvq_moe_enabled() { + static const bool enabled = [] { + const char * raw = std::getenv("DFLASH_DS4_TP_GROUPED_MMVQ"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +static bool gpu_i32_repeat_enabled() { + static const bool enabled = [] { + const char * raw = std::getenv("LUCE_CUDA_I32_REPEAT"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +// The regular DeepSeek graph already uses ggml_laguna_moe_combine for the +// route-weighted expert reduction. Keep the heterogeneous graph A/B-able +// while replacing its MUL + shape-only REPEAT_BACK sequence with the same +// exact owner-local kernel. +static bool fused_moe_combine_enabled() { + static const bool enabled = [] { + const char * raw = std::getenv("DFLASH_MOE_FUSED_COMBINE"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +// DeepSeek V4 stores routed gate and up projections as one tensor with the +// output rows concatenated. The generic graph computes that full tensor, +// materializes two contiguous views, clamps both, and only then runs SwiGLU. +// Present the two weight halves as views instead: the CUDA/HIP graph optimizer +// can fuse both MUL_MAT_ID operations and DS4 SwiGLU into one MMVQ launch. +// This is exact when the external gate/up scale is one (the ROCmFP checkpoint +// used by the heterogeneous path); other scale values retain the old graph. +static bool fused_gate_up_mmvq_enabled() { + static const bool enabled = [] { + const char * raw = std::getenv("DFLASH_DS4_TP_FUSED_GATE_UP"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +static bool coarse_owner_op_enabled() { + static const bool enabled = [] { + const char * raw = std::getenv("DFLASH_DS4_TP_COARSE_OWNER"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +static bool coarse_owner_split_op_enabled() { + static const bool enabled = [] { + const char * raw = std::getenv("DFLASH_DS4_TP_COARSE_OWNER_SPLIT"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +static bool align_shared_moe_ids_enabled() { + static const bool enabled = [] { + const char * raw = std::getenv("DFLASH_CUDA_MMVQ_MOE_ALIGN_SHARED_IDS"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +static bool device_join_enabled() { + static const bool enabled = [] { + const char * raw = std::getenv("DFLASH_DS4_TP_DEVICE_JOIN"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +static bool route_prefork_enabled() { + static const bool enabled = [] { + const char * raw = std::getenv("DFLASH_DS4_TP_ROUTE_PREFORK"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +static void add_hybrid_telemetry(MoeHybridFfnTelemetry & dst, + const MoeHybridFfnTelemetry & src) { + dst.ffn_wall_us += src.ffn_wall_us; + dst.partition_us += src.partition_us; + dst.hot_us += src.hot_us; + dst.cold_us += src.cold_us; + dst.shared_us += src.shared_us; + dst.combine_us += src.combine_us; + dst.hot_graph_build_us += src.hot_graph_build_us; + dst.hot_input_us += src.hot_input_us; + dst.hot_compute_us += src.hot_compute_us; + dst.hot_read_us += src.hot_read_us; + dst.cold_graph_build_us += src.cold_graph_build_us; + dst.cold_input_us += src.cold_input_us; + dst.cold_compute_us += src.cold_compute_us; + dst.cold_read_us += src.cold_read_us; + dst.hot_graph_builds += src.hot_graph_builds; + dst.hot_graph_hits += src.hot_graph_hits; + dst.cold_graph_builds += src.cold_graph_builds; + dst.cold_graph_hits += src.cold_graph_hits; + dst.hot_selected += src.hot_selected; + dst.cold_selected += src.cold_selected; +} + static int env_int_or_default(const char * name, int fallback) { const char * raw = std::getenv(name); if (!raw || !*raw) return fallback; @@ -458,13 +579,102 @@ static bool build_batched_routed_graph( ggml_tensor * wts, int n_embd, int n_ff_exp, int n_used, int n_tokens, float swiglu_clamp, - ggml_tensor ** out_routed) + ggml_tensor ** out_routed, + bool tokenwise = false, + std::vector * backend_nodes = nullptr) { + const auto track = [&](ggml_tensor * t) -> ggml_tensor * { + if (backend_nodes && t) backend_nodes->push_back(t); + return t; + }; + if (tokenwise && n_tokens > 1) { + ggml_tensor * joined = nullptr; + for (int t = 0; t < n_tokens; ++t) { + ggml_tensor * inp_col = ggml_cont(ctx, ggml_view_2d( + ctx, inp, n_embd, 1, inp->nb[1], (size_t) t * inp->nb[1])); + ggml_tensor * sel_col = ggml_cont(ctx, ggml_view_2d( + ctx, sel, n_used, 1, sel->nb[1], (size_t) t * sel->nb[1])); + ggml_tensor * wts_col = ggml_cont(ctx, ggml_view_2d( + ctx, wts, n_used, 1, wts->nb[1], (size_t) t * wts->nb[1])); + ggml_tensor * routed_col = nullptr; + if (!build_batched_routed_graph( + ctx, gate_tensor, up_tensor, down_tensor, gate_up_tensor, + gate_scale, up_scale, down_scale, gate_up_scale, + inp_col, sel_col, wts_col, + n_embd, n_ff_exp, n_used, 1, swiglu_clamp, + &routed_col, false, backend_nodes)) { + return false; + } + joined = joined ? track(ggml_concat(ctx, joined, routed_col, 1)) + : routed_col; + } + *out_routed = joined; + return joined != nullptr; + } + ggml_tensor * cur_3d = ggml_reshape_3d(ctx, inp, n_embd, 1, n_tokens); ggml_tensor * gu = nullptr; - if (gate_up_tensor) { - ggml_tensor * gate_up_e = apply_scale2(ctx, - ggml_mul_mat_id(ctx, gate_up_tensor, cur_3d, sel), gate_up_scale); + const bool coarse_split_requested = + coarse_owner_op_enabled() && coarse_owner_split_op_enabled(); + const bool coarse_split_eligible = + gate_tensor && up_tensor && + gate_tensor->type == GGML_TYPE_Q2_0_ROCMFP2 && + up_tensor->type == GGML_TYPE_Q2_0_ROCMFP2 && + down_tensor->type == GGML_TYPE_Q3_0_ROCMFPX; + if (coarse_split_requested) { + static bool logged_active = false; + static bool logged_ineligible = false; + bool & logged = coarse_split_eligible ? logged_active : logged_ineligible; + if (!logged) { + std::fprintf(stderr, + "[ds4-tp] split-owner %s gate=%s up=%s down=%s gate_up=%s tokens=%d routes=%d\n", + coarse_split_eligible ? "active" : "ineligible", + gate_tensor ? ggml_type_name(gate_tensor->type) : "none", + up_tensor ? ggml_type_name(up_tensor->type) : "none", + down_tensor ? ggml_type_name(down_tensor->type) : "none", + gate_up_tensor ? ggml_type_name(gate_up_tensor->type) : "none", + n_tokens, n_used); + logged = true; + } + } + if (coarse_split_requested && coarse_split_eligible) { + *out_routed = track(ggml_ds4_moe_owner_split( + ctx, inp, gate_tensor, up_tensor, down_tensor, sel, wts, + n_ff_exp, swiglu_clamp, + gate_scale, up_scale, down_scale)); + return *out_routed != nullptr; + } else if (coarse_owner_op_enabled() && + gate_up_tensor && + gate_up_scale == 1.0f && + gate_up_tensor->type == GGML_TYPE_Q2_0_ROCMFP2 && + down_tensor->type == GGML_TYPE_Q3_0_ROCMFPX) { + *out_routed = track(ggml_ds4_moe_owner( + ctx, inp, gate_up_tensor, down_tensor, sel, wts, + n_ff_exp, swiglu_clamp, down_scale)); + return *out_routed != nullptr; + } else if (gate_up_tensor && + fused_gate_up_mmvq_enabled() && + gate_up_scale == 1.0f) { + GGML_ASSERT(gate_up_tensor->ne[1] == 2 * n_ff_exp); + ggml_tensor * gate_w = ggml_view_3d( + ctx, gate_up_tensor, + gate_up_tensor->ne[0], n_ff_exp, gate_up_tensor->ne[2], + gate_up_tensor->nb[1], gate_up_tensor->nb[2], 0); + ggml_tensor * up_w = ggml_view_3d( + ctx, gate_up_tensor, + gate_up_tensor->ne[0], n_ff_exp, gate_up_tensor->ne[2], + gate_up_tensor->nb[1], gate_up_tensor->nb[2], + (size_t) n_ff_exp * gate_up_tensor->nb[1]); + ggml_tensor * gate_e = track( + ggml_mul_mat_id(ctx, gate_w, cur_3d, sel)); + ggml_tensor * up_e = track( + ggml_mul_mat_id(ctx, up_w, cur_3d, sel)); + gu = track(swiglu_clamp > 1.0e-6f + ? ggml_swiglu_ds4_split(ctx, gate_e, up_e, swiglu_clamp) + : ggml_swiglu_split(ctx, gate_e, up_e)); + } else if (gate_up_tensor) { + ggml_tensor * gate_up_e = track(apply_scale2(ctx, + ggml_mul_mat_id(ctx, gate_up_tensor, cur_3d, sel), gate_up_scale)); ggml_tensor * gate_e = ggml_view_3d(ctx, gate_up_e, n_ff_exp, gate_up_e->ne[1], gate_up_e->ne[2], gate_up_e->nb[1], gate_up_e->nb[2], 0); @@ -472,27 +682,312 @@ static bool build_batched_routed_graph( n_ff_exp, gate_up_e->ne[1], gate_up_e->ne[2], gate_up_e->nb[1], gate_up_e->nb[2], (size_t)n_ff_exp * ggml_element_size(gate_up_e)); - gate_e = ggml_cont(ctx, gate_e); - up_e = ggml_cont(ctx, up_e); - gu = swiglu_maybe_clamped(ctx, gate_e, up_e, swiglu_clamp); + gate_e = track(ggml_cont(ctx, gate_e)); + up_e = track(ggml_cont(ctx, up_e)); + gu = track(swiglu_maybe_clamped(ctx, gate_e, up_e, swiglu_clamp)); } else { - ggml_tensor * gate_e = apply_scale2(ctx, - ggml_mul_mat_id(ctx, gate_tensor, cur_3d, sel), gate_scale); - ggml_tensor * up_e = apply_scale2(ctx, - ggml_mul_mat_id(ctx, up_tensor, cur_3d, sel), up_scale); - gu = swiglu_maybe_clamped(ctx, gate_e, up_e, swiglu_clamp); + ggml_tensor * gate_e = track(apply_scale2(ctx, + ggml_mul_mat_id(ctx, gate_tensor, cur_3d, sel), gate_scale)); + ggml_tensor * up_e = track(apply_scale2(ctx, + ggml_mul_mat_id(ctx, up_tensor, cur_3d, sel), up_scale)); + gu = track(swiglu_maybe_clamped(ctx, gate_e, up_e, swiglu_clamp)); } - ggml_tensor * experts = apply_scale2(ctx, - ggml_mul_mat_id(ctx, down_tensor, gu, sel), down_scale); + ggml_tensor * experts = track(apply_scale2(ctx, + ggml_mul_mat_id(ctx, down_tensor, gu, sel), down_scale)); // Weight and sum over experts: [n_embd, n_used, n_tokens] * [1, n_used, n_tokens] + if (fused_moe_combine_enabled()) { + *out_routed = track(ggml_laguna_moe_combine(ctx, experts, wts)); + return *out_routed != nullptr; + } + ggml_tensor * w_view = ggml_reshape_3d(ctx, wts, 1, n_used, n_tokens); - experts = ggml_mul(ctx, experts, w_view); + experts = track(ggml_mul(ctx, experts, w_view)); + + // repeat_back uses this tensor for shape only, but the scheduler still + // treats it as a leaf. Keep it on the branch backend; otherwise every MoE + // branch acquires a tiny CPU split solely for an uninitialized shape leaf. + ggml_tensor * sum_shape = track( + ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens)); + ggml_tensor * moe_sum = track(ggml_repeat_back(ctx, experts, sum_shape)); + *out_routed = track(ggml_reshape_2d(ctx, moe_sum, n_embd, n_tokens)); + return true; +} - ggml_tensor * sum_shape = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); - ggml_tensor * moe_sum = ggml_repeat_back(ctx, experts, sum_shape); - *out_routed = ggml_reshape_2d(ctx, moe_sum, n_embd, n_tokens); +bool build_moe_hybrid_ffn_graph( + ggml_context * ctx, + ggml_cgraph * schedule_graph, + const MoeHybridConfig & cfg, + const MoeLayerDesc & desc, + MoeHybridLayerStorage & storage, + ggml_tensor * inp, + ggml_tensor * global_ids, + ggml_tensor * router_weights, + int n_tokens, + MoeHybridGraphInputs & out, + bool include_shared) { + + out.output = nullptr; + out.main_output = nullptr; + out.peer_output = nullptr; + if (!ctx || !inp || !global_ids || !router_weights || n_tokens <= 0 || + cfg.n_embd <= 0 || cfg.n_ff_exp <= 0 || cfg.n_expert <= 0 || + cfg.n_expert_used <= 0) { + return false; + } + + const int n_used = cfg.n_expert_used; + // Both owner remaps consume the same normalized top-k route weights. + // Expose the canonical tensor so the heterogeneous scheduler can keep it + // on the main GPU. Otherwise expanding the cold branch first lets backend + // assignment migrate this shared dependency to Strix, forcing the hot + // R9700 branch to wait for a reverse peer copy before it can launch. + out.router_weights = router_weights; + auto build_remap = [&](const std::vector & local_by_global, + ggml_tensor ** local_lut, + ggml_tensor ** valid_lut, + ggml_tensor ** local_ids, + ggml_tensor ** masked_weights, + std::vector * backend_nodes) { + auto track = [backend_nodes](ggml_tensor * tensor) { + if (tensor && backend_nodes) backend_nodes->push_back(tensor); + return tensor; + }; + if (!*local_lut) { + *local_lut = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, 1, cfg.n_expert); + ggml_set_input(*local_lut); + // These inputs are consumed late in a whole-model graph. Preserve + // their allocation from graph start; otherwise gallocr may reuse + // the tiny buffer as activation scratch before its layer executes. + ggml_set_output(*local_lut); + } + if (!*valid_lut) { + *valid_lut = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 1, cfg.n_expert); + ggml_set_input(*valid_lut); + ggml_set_output(*valid_lut); + } + + // q1 can address the 2-D LUT directly. Historically q>1 left its tiny + // I32 repeat unpinned for CPU fallback. With the exact GPU I32 repeat + // enabled, track it with the rest of the owner-local remap nodes. + ggml_tensor * local_lut_batched = *local_lut; + if (n_tokens > 1) { + ggml_tensor * repeated = ggml_repeat_4d( + ctx, *local_lut, 1, cfg.n_expert, n_tokens, 1); + local_lut_batched = + gpu_i32_repeat_enabled() ? track(repeated) : repeated; + } + ggml_tensor * mapped = track(ggml_get_rows( + ctx, local_lut_batched, global_ids)); + mapped = track(ggml_reshape_2d(ctx, mapped, n_used, n_tokens)); + *local_ids = track(ggml_cont(ctx, mapped)); + ggml_tensor * valid_lut_batched = *valid_lut; + if (n_tokens > 1) { + valid_lut_batched = track(ggml_repeat_4d( + ctx, *valid_lut, 1, cfg.n_expert, n_tokens, 1)); + } + ggml_tensor * valid = track(ggml_get_rows( + ctx, valid_lut_batched, global_ids)); + valid = track(ggml_reshape_2d(ctx, valid, n_used, n_tokens)); + *masked_weights = track(ggml_mul(ctx, router_weights, valid)); + return (int)local_by_global.size() == cfg.n_expert; + }; + + ggml_tensor * hot_ids = nullptr; + ggml_tensor * hot_weights = nullptr; + if (!build_remap(storage.hot_local_by_global, + &out.hot_local_lut, &out.hot_valid_lut, + &hot_ids, &hot_weights, &out.hot_remap_nodes)) { + return false; + } + + ggml_tensor * cold_ids = nullptr; + ggml_tensor * cold_weights = nullptr; + if (!build_remap(storage.cold_local_by_global, + &out.cold_local_lut, &out.cold_valid_lut, + &cold_ids, &cold_weights, &out.cold_remap_nodes)) { + return false; + } + + ggml_tensor * shard_ids = nullptr; + ggml_tensor * shard_weights = nullptr; + const bool has_cold_shard = storage.expert_shard_channels > 0 && + storage.gate_shard_hot && storage.up_shard_hot && + storage.down_shard_hot; + if (has_cold_shard && + !build_remap(storage.cold_local_by_global, + &out.shard_local_lut, &out.shard_valid_lut, + &shard_ids, &shard_weights, &out.shard_remap_nodes)) { + return false; + } + + // q-token verification often routes adjacent tokens to the same expert + // at different top-k ranks. Align those owner-local IDs before MMVQ so + // equal weights are consumed by warps in one block. The encoded original + // route slot is decoded by the dedicated MoE kernel, which scatters every + // result back before the unchanged weighted reduction. + if (n_tokens > 1 && align_shared_moe_ids_enabled()) { + hot_ids = ggml_ds4_moe_align_ids(ctx, hot_ids); + cold_ids = ggml_ds4_moe_align_ids(ctx, cold_ids); + out.hot_remap_nodes.push_back(hot_ids); + out.cold_remap_nodes.push_back(cold_ids); + if (has_cold_shard) { + shard_ids = ggml_ds4_moe_align_ids(ctx, shard_ids); + out.shard_remap_nodes.push_back(shard_ids); + } + } + + + ggml_tensor * shard = nullptr; + if (has_cold_shard) { + const bool tokenwise = + storage.gate_shard_hot->type == GGML_TYPE_Q2_0_ROCMFP2 && + !(n_tokens > 1 && grouped_mmvq_moe_enabled()); + if (!build_batched_routed_graph( + ctx, + storage.gate_shard_hot, storage.up_shard_hot, + storage.down_shard_hot, nullptr, + desc.ffn_gate_exps_s, desc.ffn_up_exps_s, + desc.ffn_down_exps_s, 1.0f, + inp, shard_ids, shard_weights, + cfg.n_embd, storage.expert_shard_channels, + n_used, n_tokens, cfg.swiglu_clamp, + &shard, tokenwise, &out.shard_nodes)) { + return false; + } + } + + ggml_tensor * hot = nullptr; + if ((storage.gate_up_hot || (storage.gate_hot && storage.up_hot)) && + storage.down_hot) { + const ggml_tensor * hot_gate = storage.gate_up_hot + ? storage.gate_up_hot : storage.gate_hot; + const bool tokenwise = + hot_gate->type == GGML_TYPE_Q2_0_ROCMFP2 && + !(n_tokens > 1 && grouped_mmvq_moe_enabled()); + if (!build_batched_routed_graph( + ctx, + storage.gate_hot, storage.up_hot, storage.down_hot, + storage.gate_up_hot, + desc.ffn_gate_exps_s, desc.ffn_up_exps_s, + desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, + inp, hot_ids, hot_weights, + cfg.n_embd, cfg.n_ff_exp, n_used, n_tokens, + cfg.swiglu_clamp, &hot, tokenwise, + &out.hot_nodes)) { + return false; + } + } + + ggml_tensor * cold = nullptr; + if ((storage.gate_up_cold || (storage.gate_cold && storage.up_cold)) && + storage.down_cold) { + const ggml_tensor * cold_gate = storage.gate_up_cold + ? storage.gate_up_cold : storage.gate_cold; + const bool tokenwise = + cold_gate->type == GGML_TYPE_Q2_0_ROCMFP2 && + !(n_tokens > 1 && grouped_mmvq_moe_enabled()); + if (!build_batched_routed_graph( + ctx, + storage.gate_cold, storage.up_cold, storage.down_cold, + storage.gate_up_cold, + desc.ffn_gate_exps_s, desc.ffn_up_exps_s, + desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, + inp, cold_ids, cold_weights, + cfg.n_embd, + storage.expert_shard_channels > 0 + ? cfg.n_ff_exp - storage.expert_shard_channels + : cfg.n_ff_exp, + n_used, n_tokens, + cfg.swiglu_clamp, &cold, tokenwise, + &out.cold_nodes)) { + return false; + } + } + + ggml_tensor * main_branch = hot; + if (shard) { + main_branch = main_branch ? ggml_add(ctx, main_branch, shard) : shard; + out.shard_nodes.push_back(main_branch); + } + ggml_tensor * shared = include_shared + ? build_shared_expert_subgraph(ctx, desc, inp, cfg.swiglu_clamp) + : nullptr; + if (shared) main_branch = main_branch ? ggml_add(ctx, main_branch, shared) : shared; + + // The generic scheduler copies every cross-backend input before launching + // any node in a split. If hot/shared and the final add are one contiguous + // main-backend split, the add's cold input makes that whole split wait for + // the peer, serializing the nominally parallel branches. + // + // Expand cold now, then visit hot/shared, then a new peer-owned CONT fence, + // and finally the main-backend add. This yields: + // peer cold compute -> main hot/shared compute -> peer fence -> main join + // The scheduler can enqueue cold first and hot/shared second; the fence + // separates the final join so its event wait is inserted after hot/shared. + ggml_tensor * combined = nullptr; + if (schedule_graph && cold && device_join_enabled()) { + // Materialize both route IDs and normalized route weights before the + // cold split is expanded. Cross-device copies are not guaranteed to + // remain asynchronous on this heterogeneous ROCm pair. If the tiny + // weight copy is discovered after cold expert execution, the fallback + // copy synchronizes the Strix stream and prevents the host from + // enqueueing independent R9700 hot work until cold has completed. + // + // q4 verification calls this builder twice (4 routes + padded 2), so + // retain every derived route tensor rather than only the canonical + // six-wide routing output. + if (route_prefork_enabled()) { + out.route_prefork_nodes.push_back(global_ids); + out.route_prefork_nodes.push_back(router_weights); + ggml_build_forward_expand(schedule_graph, global_ids); + ggml_build_forward_expand(schedule_graph, router_weights); + } + // Enforce fork order without adding another backend graph: + // cold owner -> hot/shared -> in-graph event wait/copy -> add. + // The deferred copy remains in the same main-backend split as the + // hot branch, so its wait is reached only after useful main work. + // Keep the peer result live explicitly. The generic allocator normally + // derives lifetime from children on the same execution backend; this + // custom foreign-buffer edge intentionally bypasses that copy path. + ggml_set_output(cold); + ggml_build_forward_expand(schedule_graph, cold); + if (main_branch) { + ggml_build_forward_expand(schedule_graph, main_branch); + } + ggml_tensor * cold_ready = + ggml_ds4_deferred_peer_copy(ctx, cold); + // The scheduler may prefill this tensor in the host-copy diagnostic + // before its containing main split launches. Reserve a stable buffer + // from graph start so earlier hot-branch scratch cannot alias it. + ggml_set_input(cold_ready); + ggml_set_output(cold_ready); + out.deferred_peer_copy_nodes.push_back(cold_ready); + out.main_output = main_branch; + out.peer_output = cold_ready; + combined = main_branch ? ggml_add(ctx, cold_ready, main_branch) + : cold_ready; + } else if (schedule_graph && cold) { + ggml_build_forward_expand(schedule_graph, cold); + ggml_tensor * cold_fence = ggml_cont(ctx, cold); + out.cold_nodes.push_back(cold_fence); + combined = main_branch ? ggml_add(ctx, main_branch, cold_fence) + : cold_fence; + } else { + // Preserve the established default dependency order exactly. + if (cold && main_branch) { + combined = ggml_add(ctx, cold, main_branch); + out.join_nodes.push_back(combined); + } else { + combined = cold ? cold : main_branch; + } + } + if (!combined) return false; + + out.output = ggml_cont(ctx, combined); return true; } @@ -718,16 +1213,21 @@ bool build_cached_hot_batched_graph( out.inp = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, n_embd, n_tokens); ggml_set_input(out.inp); - out.sel = ggml_new_tensor_2d(out.ctx, GGML_TYPE_I32, n_used, n_tokens); - ggml_set_input(out.sel); - out.wts = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, n_used, n_tokens); - ggml_set_input(out.wts); ggml_tensor * routed = nullptr; - build_batched_routed_graph(out.ctx, - storage.gate_hot, storage.up_hot, storage.down_hot, storage.gate_up_hot, - desc.ffn_gate_exps_s, desc.ffn_up_exps_s, desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, - out.inp, out.sel, out.wts, n_embd, n_ff_exp, n_used, n_tokens, cfg.swiglu_clamp, &routed); + const bool has_hot_stack = + (storage.gate_up_hot || (storage.gate_hot && storage.up_hot)) && + storage.down_hot; + if (has_hot_stack) { + out.sel = ggml_new_tensor_2d(out.ctx, GGML_TYPE_I32, n_used, n_tokens); + ggml_set_input(out.sel); + out.wts = ggml_new_tensor_2d(out.ctx, GGML_TYPE_F32, n_used, n_tokens); + ggml_set_input(out.wts); + build_batched_routed_graph(out.ctx, + storage.gate_hot, storage.up_hot, storage.down_hot, storage.gate_up_hot, + desc.ffn_gate_exps_s, desc.ffn_up_exps_s, desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, + out.inp, out.sel, out.wts, n_embd, n_ff_exp, n_used, n_tokens, cfg.swiglu_clamp, &routed); + } // Shared expert (always on GPU) ggml_tensor * combined = routed; @@ -868,6 +1368,17 @@ bool eval_moe_hybrid_ffn_single( (fixed_slot_mode == 1 && has_cold) ? std::max(n_cold, fixed_slot_limit) : (fixed_slot_mode == 2 && has_cold) ? std::max(n_cold, min_adaptive_slots) : n_cold; + const size_t graph_cache_size = (size_t)cfg.n_expert_used + 1; + if (storage.hot_graph_by_width.size() < graph_cache_size) { + storage.hot_graph_by_width.resize(graph_cache_size); + } + if (storage.cold_graph_by_width.size() < graph_cache_size) { + storage.cold_graph_by_width.resize(graph_cache_size); + } + CachedFfnGraph & hot_graph = + storage.hot_graph_by_width[(size_t)n_hot_graph]; + CachedFfnGraph & cold_graph = + storage.cold_graph_by_width[(size_t)n_cold_graph]; ggml_backend_t cold_backend = storage.cold_backend ? storage.cold_backend : cpu_backend; const bool cold_on_gpu = has_cold && storage.cold_backend_kind == MoeHybridColdBackend::Gpu && @@ -897,10 +1408,10 @@ bool eval_moe_hybrid_ffn_single( hot_weights_data = hot_weights_padded.data(); } // Lazily build cached hot graph on first use - if (!storage.hot_graph.valid() || storage.hot_graph.n_hot != n_hot_graph) { + if (!hot_graph.valid() || hot_graph.n_hot != n_hot_graph) { if (telemetry) telemetry->hot_graph_builds++; const auto graph_build_t0 = HybridClock::now(); - build_cached_hot_graph(storage.hot_graph, gpu_backend, + build_cached_hot_graph(hot_graph, gpu_backend, storage.gate_hot, storage.up_hot, storage.down_hot, storage.gate_up_hot, desc.ffn_gate_exps_s, desc.ffn_up_exps_s, desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, desc, cfg.n_embd, cfg.n_ff_exp, n_hot_graph, @@ -909,19 +1420,19 @@ bool eval_moe_hybrid_ffn_single( } else if (telemetry) { telemetry->hot_graph_hits++; } - if (storage.hot_graph.valid() && storage.hot_graph.n_hot == n_hot_graph) { + if (hot_graph.valid() && hot_graph.n_hot == n_hot_graph) { const auto input_t0 = HybridClock::now(); - ggml_backend_tensor_set(storage.hot_graph.inp, cur_host, 0, sizeof(float) * (size_t)cfg.n_embd); - if (storage.hot_graph.ids && n_hot_graph > 0) { - ggml_backend_tensor_set(storage.hot_graph.ids, hot_ids_data, 0, sizeof(int32_t) * (size_t)n_hot_graph); + ggml_backend_tensor_set(hot_graph.inp, cur_host, 0, sizeof(float) * (size_t)cfg.n_embd); + if (hot_graph.ids && n_hot_graph > 0) { + ggml_backend_tensor_set(hot_graph.ids, hot_ids_data, 0, sizeof(int32_t) * (size_t)n_hot_graph); } - if (storage.hot_graph.weights && n_hot_graph > 0) { - ggml_backend_tensor_set(storage.hot_graph.weights, hot_weights_data, 0, sizeof(float) * (size_t)n_hot_graph); + if (hot_graph.weights && n_hot_graph > 0) { + ggml_backend_tensor_set(hot_graph.weights, hot_weights_data, 0, sizeof(float) * (size_t)n_hot_graph); } if (telemetry) telemetry->hot_input_us += elapsed_us(input_t0, HybridClock::now()); // Launch GPU async — kernel runs while the cold backend runs. const auto compute_t0 = HybridClock::now(); - ggml_backend_graph_compute_async(gpu_backend, storage.hot_graph.gf); + ggml_backend_graph_compute_async(gpu_backend, hot_graph.gf); if (telemetry) telemetry->hot_compute_us += elapsed_us(compute_t0, HybridClock::now()); hot_async_launched = true; if (cold_on_gpu) { @@ -965,10 +1476,10 @@ bool eval_moe_hybrid_ffn_single( cold_ids_data = cold_ids_padded.data(); cold_weights_data = cold_weights_padded.data(); } - if (!storage.cold_graph.valid() || storage.cold_graph.n_hot != n_cold_graph) { + if (!cold_graph.valid() || cold_graph.n_hot != n_cold_graph) { if (telemetry) telemetry->cold_graph_builds++; const auto graph_build_t0 = HybridClock::now(); - build_cached_cold_graph(storage.cold_graph, cold_backend, + build_cached_cold_graph(cold_graph, cold_backend, storage.gate_cold, storage.up_cold, storage.down_cold, storage.gate_up_cold, desc.ffn_gate_exps_s, desc.ffn_up_exps_s, desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, cfg.n_embd, cfg.n_ff_exp, n_cold_graph, cfg.swiglu_clamp); @@ -976,14 +1487,14 @@ bool eval_moe_hybrid_ffn_single( } else if (telemetry) { telemetry->cold_graph_hits++; } - if (storage.cold_graph.valid() && storage.cold_graph.n_hot == n_cold_graph) { + if (cold_graph.valid() && cold_graph.n_hot == n_cold_graph) { const auto input_t0 = HybridClock::now(); - ggml_backend_tensor_set(storage.cold_graph.inp, cur_host, 0, sizeof(float) * (size_t)cfg.n_embd); - ggml_backend_tensor_set(storage.cold_graph.ids, cold_ids_data, 0, sizeof(int32_t) * (size_t)n_cold_graph); - ggml_backend_tensor_set(storage.cold_graph.weights, cold_weights_data, 0, sizeof(float) * (size_t)n_cold_graph); + ggml_backend_tensor_set(cold_graph.inp, cur_host, 0, sizeof(float) * (size_t)cfg.n_embd); + ggml_backend_tensor_set(cold_graph.ids, cold_ids_data, 0, sizeof(int32_t) * (size_t)n_cold_graph); + ggml_backend_tensor_set(cold_graph.weights, cold_weights_data, 0, sizeof(float) * (size_t)n_cold_graph); if (telemetry) telemetry->cold_input_us += elapsed_us(input_t0, HybridClock::now()); const auto compute_t0 = HybridClock::now(); - auto st = ggml_backend_graph_compute(cold_backend, storage.cold_graph.gf); + auto st = ggml_backend_graph_compute(cold_backend, cold_graph.gf); if (telemetry) telemetry->cold_compute_us += elapsed_us(compute_t0, HybridClock::now()); if (st != GGML_STATUS_SUCCESS) { if (hot_async_launched) ggml_backend_synchronize(gpu_backend); @@ -992,7 +1503,7 @@ bool eval_moe_hybrid_ffn_single( } const auto read_t0 = HybridClock::now(); cold.resize((size_t)cfg.n_embd); - ggml_backend_tensor_get(storage.cold_graph.output, cold.data(), 0, sizeof(float) * (size_t)cfg.n_embd); + ggml_backend_tensor_get(cold_graph.output, cold.data(), 0, sizeof(float) * (size_t)cfg.n_embd); if (telemetry) telemetry->cold_read_us += elapsed_us(read_t0, HybridClock::now()); } else { if (!run_routed_subset(cold_backend, @@ -1010,11 +1521,11 @@ bool eval_moe_hybrid_ffn_single( const auto cold_t1 = HybridClock::now(); // ── Sync GPU and read result ── - if ((has_hot || has_shared) && storage.hot_graph.valid() && storage.hot_graph.n_hot == n_hot_graph) { + if ((has_hot || has_shared) && hot_graph.valid() && hot_graph.n_hot == n_hot_graph) { const auto read_t0 = HybridClock::now(); ggml_backend_synchronize(gpu_backend); hot_and_shared.resize((size_t)cfg.n_embd); - ggml_backend_tensor_get(storage.hot_graph.output, hot_and_shared.data(), 0, sizeof(float) * (size_t)cfg.n_embd); + ggml_backend_tensor_get(hot_graph.output, hot_and_shared.data(), 0, sizeof(float) * (size_t)cfg.n_embd); if (telemetry) telemetry->hot_read_us += elapsed_us(read_t0, HybridClock::now()); } const auto hot_t1 = HybridClock::now(); @@ -1307,29 +1818,50 @@ static bool eval_moe_hybrid_ffn_batched_core( CachedHotBatchedGraph & hg = storage.hot_batched_mixed[n_tokens]; const bool hg_ok = (hg.valid() && hg.n_tokens == n_tokens) || build_cached_hot_batched_graph(hg, gpu_backend, storage, desc, cfg, n_tokens); + const bool remote_cold = fp_has_cold && expert_compute && expert_layer; CachedHotBatchedGraph * cg = nullptr; bool cg_ok = true; - if (fp_has_cold) { + if (fp_has_cold && !remote_cold) { cg = &storage.cold_batched_mixed[n_tokens]; + ggml_backend_t cached_cold_backend = + storage.cold_backend ? storage.cold_backend : cpu_backend; cg_ok = (cg->valid() && cg->n_tokens == n_tokens) - || build_cached_cold_batched_graph(*cg, cpu_backend, storage, desc, cfg, n_tokens); + || build_cached_cold_batched_graph( + *cg, cached_cold_backend, storage, desc, cfg, n_tokens); } if (hg_ok && cg_ok) { // Hot (GPU, async): shared expert + routed hot (zero-weight dummy slots // keep an all-cold batch's shared-expert contribution). ggml_backend_tensor_set(hg.inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); - ggml_backend_tensor_set(hg.sel, hot_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); - ggml_backend_tensor_set(hg.wts, hot_wts.data(), 0, sizeof(float) * (size_t)total_slots); + if (hg.sel) { + ggml_backend_tensor_set(hg.sel, hot_sel.data(), 0, + sizeof(int32_t) * (size_t)total_slots); + } + if (hg.wts) { + ggml_backend_tensor_set(hg.wts, hot_wts.data(), 0, + sizeof(float) * (size_t)total_slots); + } ggml_backend_graph_compute_async(gpu_backend, hg.gf); std::vector cold_partial; - if (cg) { + if (remote_cold) { + if (!eval_moe_hybrid_remote_cold_batched( + cfg, storage, cur_host, selected_ids, selected_weights, + n_tokens, cold_partial, err, + expert_compute, expert_layer)) { + ggml_backend_synchronize(gpu_backend); + return false; + } + } else if (cg) { cold_partial.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); ggml_backend_tensor_set(cg->inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); ggml_backend_tensor_set(cg->sel, cold_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); ggml_backend_tensor_set(cg->wts, cold_wts.data(), 0, sizeof(float) * (size_t)total_slots); - if (ggml_backend_graph_compute(cpu_backend, cg->gf) != GGML_STATUS_SUCCESS) { + ggml_backend_t cached_cold_backend = + storage.cold_backend ? storage.cold_backend : cpu_backend; + if (ggml_backend_graph_compute( + cached_cold_backend, cg->gf) != GGML_STATUS_SUCCESS) { ggml_backend_synchronize(gpu_backend); if (err) *err = "batched cold cached compute failed"; return false; @@ -1339,7 +1871,7 @@ static bool eval_moe_hybrid_ffn_batched_core( ggml_backend_synchronize(gpu_backend); ggml_backend_tensor_get(hg.output, out.data(), 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); - if (cg) { + if (remote_cold || cg) { const size_t ntot = (size_t)n_embd * (size_t)n_tokens; for (size_t i = 0; i < ntot; ++i) out[i] += cold_partial[i]; } @@ -1647,82 +2179,12 @@ static bool eval_moe_hybrid_remote_cold_batched( return false; } - const int cold_batch = moe_hybrid_expert_compute_ipc_batch_limit(n_tokens); - std::vector token_group; - std::vector group_input; - std::vector group_ids; - std::vector group_wts; - std::vector group_output; - for (int n_cold = 1; n_cold <= max_cold_selected; ++n_cold) { - token_group.clear(); - for (int t = 0; t < n_tokens; ++t) { - if (cold_counts[(size_t)t] == n_cold) { - token_group.push_back(t); - } - } - for (size_t base = 0; base < token_group.size(); base += (size_t)cold_batch) { - const int tc = (int)std::min((size_t)cold_batch, token_group.size() - base); - bool contiguous_tokens = true; - for (int gi = 1; gi < tc; ++gi) { - if (token_group[base + (size_t)gi] != - token_group[base] + gi) { - contiguous_tokens = false; - break; - } - } - const int first_token = token_group[base]; - const float * compute_input = contiguous_tokens - ? cur_host + (size_t)first_token * (size_t)n_embd - : nullptr; - float * compute_output = contiguous_tokens - ? out.data() + (size_t)first_token * (size_t)n_embd - : nullptr; - if (!contiguous_tokens) { - group_input.resize((size_t)tc * (size_t)n_embd); - compute_input = group_input.data(); - } - group_ids.resize((size_t)tc * (size_t)n_cold); - group_wts.resize((size_t)tc * (size_t)n_cold); - if (!contiguous_tokens) { - group_output.resize((size_t)tc * (size_t)n_embd); - compute_output = group_output.data(); - } - - for (int gi = 0; gi < tc; ++gi) { - const int t = token_group[base + (size_t)gi]; - if (!contiguous_tokens) { - std::memcpy(group_input.data() + (size_t)gi * (size_t)n_embd, - cur_host + (size_t)t * (size_t)n_embd, - sizeof(float) * (size_t)n_embd); - } - for (int i = 0; i < n_cold; ++i) { - const size_t src = (size_t)t * (size_t)n_used + (size_t)i; - const size_t dst = (size_t)gi * (size_t)n_cold + (size_t)i; - group_ids[dst] = cold_sel[src]; - group_wts[dst] = cold_wts[src]; - } - } - - if (!expert_compute->compute_batch(*expert_layer, - compute_input, - group_ids.data(), - group_wts.data(), - tc, n_cold, - n_embd, n_ff_exp, - compute_output)) { - if (err) *err = "hybrid batched remote cold compute failed"; - return false; - } - - if (!contiguous_tokens) { - for (int gi = 0; gi < tc; ++gi) { - const int t = token_group[base + (size_t)gi]; - std::memcpy(out.data() + (size_t)t * (size_t)n_embd, - group_output.data() + (size_t)gi * (size_t)n_embd, - sizeof(float) * (size_t)n_embd); - } - } - } + if (!expert_compute->compute_batch_ragged( + *expert_layer, cur_host, cold_sel.data(), cold_wts.data(), + cold_counts.data(), n_tokens, n_used, n_embd, n_ff_exp, + out.data())) { + if (err) *err = "hybrid ragged remote cold compute failed"; + return false; } return true; } @@ -1931,6 +2393,29 @@ bool eval_moe_hybrid_ffn_batched( const MoeExpertLayer * expert_layer, MoeHybridFfnTelemetry * telemetry) { if (telemetry) *telemetry = {}; + const bool materialized_cold = storage.down_cold || storage.gate_up_cold; + if (compact_materialized_experts_enabled() && materialized_cold && + !expert_compute && n_tokens > 0 && n_tokens <= 4) { + out.assign((size_t)cfg.n_embd * (size_t)n_tokens, 0.0f); + std::vector token_out; + for (int t = 0; t < n_tokens; ++t) { + MoeHybridFfnTelemetry token_telemetry; + if (!eval_moe_hybrid_ffn_single( + gpu_backend, cfg, desc, storage, cpu_backend, + cur_host + (size_t)t * (size_t)cfg.n_embd, + selected_ids + (size_t)t * (size_t)cfg.n_expert_used, + selected_weights + (size_t)t * (size_t)cfg.n_expert_used, + cfg.n_expert_used, token_out, + telemetry ? &token_telemetry : nullptr, err)) { + return false; + } + if (telemetry) add_hybrid_telemetry(*telemetry, token_telemetry); + std::memcpy(out.data() + (size_t)t * (size_t)cfg.n_embd, + token_out.data(), + sizeof(float) * (size_t)cfg.n_embd); + } + return true; + } const int n_hot_stack = storage.gate_up_hot ? (int)storage.gate_up_hot->ne[2] : storage.gate_hot ? (int)storage.gate_hot->ne[2] : 0; diff --git a/server/src/common/moe_hybrid_ffn_eval.h b/server/src/common/moe_hybrid_ffn_eval.h index 239a1529e..08c764990 100644 --- a/server/src/common/moe_hybrid_ffn_eval.h +++ b/server/src/common/moe_hybrid_ffn_eval.h @@ -110,6 +110,70 @@ struct MoeHybridFfnTelemetry { int cold_selected = 0; }; +// Inputs owned by a scheduler-allocated hybrid FFN graph. The lookup tensors +// map global router IDs to each backend's compact expert stack and mask the +// slots owned by the other backend without a host-side routing round trip. +struct MoeHybridGraphInputs { + ggml_tensor * router_weights = nullptr; + std::vector router_nodes; + // q>1 decomposes the six selected routes into a four-wide head and a + // padded two-wide tail. Keep those derived ID/weight tensors on the main + // owner and schedule them before either expert branch. Otherwise the + // scheduler discovers the cold branch first and inserts a second + // main->peer copy in the middle of cold execution, which synchronizes the + // peer stream before the hot branch can be submitted. + std::vector route_prefork_nodes; + ggml_tensor * hot_local_lut = nullptr; + ggml_tensor * hot_valid_lut = nullptr; + ggml_tensor * cold_local_lut = nullptr; + ggml_tensor * cold_valid_lut = nullptr; + ggml_tensor * shard_local_lut = nullptr; + ggml_tensor * shard_valid_lut = nullptr; + ggml_tensor * output = nullptr; + // Exact owner-local partials exposed for consumers that can fold the + // hot+cold reduction into their own kernel. peer_output is the stable + // main-backend activation produced by the existing deferred peer copy; + // neither tensor changes expert placement or routing semantics. + ggml_tensor * main_output = nullptr; + ggml_tensor * peer_output = nullptr; + // Backend-affinity hints consumed after the multi-backend scheduler is + // created. Keeping every intermediate of a routed branch on its weight + // backend avoids gate/up -> activation -> down ping-pong copies. + std::vector hot_remap_nodes; + std::vector cold_remap_nodes; + std::vector shard_remap_nodes; + std::vector hot_nodes; + std::vector shard_nodes; + std::vector cold_nodes; + // Main-backend nodes that first consume a completed cold branch. Hash + // layers can append two joins because six routes are lowered as 4 + 2. + std::vector join_nodes; + // Main-backend peer-copy ops whose src[0] stays on the cold owner. The + // scheduler attaches a dedicated producer event to each node. + std::vector deferred_peer_copy_nodes; +}; + +// Append a device-resident hot+cold+shared MoE FFN to an existing graph. +// `global_ids` and `router_weights` are [n_expert_used, n_tokens]. Weight +// tensors in `storage` determine scheduler placement on the two GPU backends. +// When `schedule_graph` is non-null, the cold branch is expanded immediately +// and a peer-owned fence is inserted before the final main-backend join. This +// forces three scheduler splits (cold, hot/shared, join), preventing the join's +// cold-result copy from blocking hot/shared launch. Consumers may use +// main_output + peer_output to fuse the exact final add into their next op. +bool build_moe_hybrid_ffn_graph( + ggml_context * ctx, + ggml_cgraph * schedule_graph, + const MoeHybridConfig & cfg, + const MoeLayerDesc & desc, + MoeHybridLayerStorage & storage, + ggml_tensor * inp, + ggml_tensor * global_ids, + ggml_tensor * router_weights, + int n_tokens, + MoeHybridGraphInputs & out, + bool include_shared = true); + int moe_hybrid_expert_compute_batch_limit(); int moe_hybrid_expert_compute_ipc_batch_limit(int n_tokens); int moe_hybrid_prefill_hot_sub_batch_limit(); diff --git a/server/src/common/moe_hybrid_storage.cpp b/server/src/common/moe_hybrid_storage.cpp index 37a177a96..1e1d6f858 100644 --- a/server/src/common/moe_hybrid_storage.cpp +++ b/server/src/common/moe_hybrid_storage.cpp @@ -6,6 +6,7 @@ #include "ggml-cuda.h" #include +#include #include #if defined(DFLASH27B_BACKEND_CUDA) @@ -136,12 +137,114 @@ static ggml_tensor * new_like_with_expert_count(ggml_context * ctx, ggml_tensor return ggml_new_tensor(ctx, src->type, 4, ne); } +static int cold_expert_shard_channels() { + const char * value = std::getenv("DFLASH_DS4_TP_COLD_SHARD_CHANNELS"); + return value && *value ? std::max(0, std::atoi(value)) : 0; +} + +static ggml_tensor * new_output_channel_slice( + ggml_context * ctx, ggml_tensor * src, int channels, int experts) { + if (!src || channels <= 0 || experts <= 0) return nullptr; + const int64_t ne[4] = { src->ne[0], channels, experts, 1 }; + return ggml_new_tensor(ctx, src->type, 4, ne); +} + +static ggml_tensor * new_input_channel_slice( + ggml_context * ctx, ggml_tensor * src, int channels, int experts) { + if (!src || channels <= 0 || experts <= 0) return nullptr; + const int64_t ne[4] = { channels, src->ne[1], experts, 1 }; + return ggml_new_tensor(ctx, src->type, 4, ne); +} + +// gate/up are [input, ffn, expert], so a contiguous range of FFN output rows +// can be copied directly for every expert. +static bool read_output_channel_slices_from_mem( + const uint8_t * tensor_data, + size_t tensor_size, + ggml_tensor * src, + const std::vector & expert_ids, + int channel_start, + int channel_count, + std::vector & out, + std::string * err) { + if (!tensor_data || !src || expert_ids.empty() || channel_count <= 0) { + out.clear(); + return true; + } + if (channel_start < 0 || channel_start + channel_count > src->ne[1]) { + if (err) *err = "gate/up channel slice out of bounds"; + return false; + } + const size_t row_bytes = (size_t) src->nb[1]; + const size_t slice_bytes = row_bytes * (size_t) channel_count; + out.resize(slice_bytes * expert_ids.size()); + for (size_t i = 0; i < expert_ids.size(); ++i) { + const size_t src_offset = (size_t) expert_ids[i] * (size_t) src->nb[2] + + (size_t) channel_start * row_bytes; + if (src_offset + slice_bytes > tensor_size) { + if (err) *err = "gate/up channel slice exceeds file tensor"; + return false; + } + std::memcpy(out.data() + i * slice_bytes, + tensor_data + src_offset, slice_bytes); + } + return true; +} + +// down is [ffn, output, expert]. Slice quantization-block-aligned input +// columns independently in every output row. +static bool read_input_channel_slices_from_mem( + const uint8_t * tensor_data, + size_t tensor_size, + ggml_tensor * src, + const std::vector & expert_ids, + int channel_start, + int channel_count, + std::vector & out, + std::string * err) { + if (!tensor_data || !src || expert_ids.empty() || channel_count <= 0) { + out.clear(); + return true; + } + const int64_t block = ggml_blck_size(src->type); + if (channel_start < 0 || channel_count <= 0 || + channel_start + channel_count > src->ne[0] || + channel_start % block != 0 || channel_count % block != 0) { + if (err) *err = "down channel slice is not quantization-block aligned"; + return false; + } + const size_t start_bytes = channel_start == 0 + ? 0 : ggml_row_size(src->type, channel_start); + const size_t slice_row_bytes = ggml_row_size(src->type, channel_count); + const size_t slice_expert_bytes = slice_row_bytes * (size_t) src->ne[1]; + out.resize(slice_expert_bytes * expert_ids.size()); + for (size_t i = 0; i < expert_ids.size(); ++i) { + for (int64_t row = 0; row < src->ne[1]; ++row) { + const size_t src_offset = (size_t) expert_ids[i] * (size_t) src->nb[2] + + (size_t) row * (size_t) src->nb[1] + + start_bytes; + if (src_offset + slice_row_bytes > tensor_size) { + if (err) *err = "down channel slice exceeds file tensor"; + return false; + } + std::memcpy(out.data() + i * slice_expert_bytes + + (size_t) row * slice_row_bytes, + tensor_data + src_offset, slice_row_bytes); + } + } + return true; +} + } // namespace MoeHybridStorage::~MoeHybridStorage() { for (auto & layer : layers) { layer.hot_graph.free(); layer.cold_graph.free(); + for (auto & graph : layer.hot_graph_by_width) graph.free(); + for (auto & graph : layer.cold_graph_by_width) graph.free(); + layer.hot_graph_by_width.clear(); + layer.cold_graph_by_width.clear(); layer.hot_batched_graph.free(); for (auto & g : layer.hot_batched_mixed) g.free(); for (auto & g : layer.cold_batched_mixed) g.free(); @@ -166,6 +269,9 @@ MoeHybridStorage::~MoeHybridStorage() { layer.up_hot = nullptr; layer.down_hot = nullptr; layer.gate_up_hot = nullptr; + layer.gate_shard_hot = nullptr; + layer.up_shard_hot = nullptr; + layer.down_shard_hot = nullptr; layer.gate_cold = nullptr; layer.up_cold = nullptr; layer.down_cold = nullptr; @@ -205,7 +311,8 @@ bool build_moe_hybrid_storage(const MoeHybridConfig & cfg, const MoeHybridPlacement & placement, const std::vector & layer_descs, MoeHybridStorage & out, - std::string * err) { + std::string * err, + ggml_backend_t cold_gpu_backend) { if (!placement.matches(cfg)) { if (err) *err = "placement does not match config"; return false; @@ -226,7 +333,9 @@ bool build_moe_hybrid_storage(const MoeHybridConfig & cfg, out.cold_backend_kind = cfg.cold_expert_backend; out.materialized_hot_experts = cfg.materialize_hot_experts; out.materialized_cold_experts = cfg.materialize_cold_experts; - out.cold_backend = (cfg.cold_expert_backend == MoeHybridColdBackend::Gpu) ? gpu_backend : out.cpu_backend; + out.cold_backend = cfg.cold_expert_backend == MoeHybridColdBackend::Gpu + ? (cold_gpu_backend ? cold_gpu_backend : gpu_backend) + : out.cpu_backend; if (!out.cold_backend) { if (err) *err = "failed to select cold expert backend"; return false; @@ -406,7 +515,8 @@ bool build_moe_hybrid_storage_from_file( MoeHybridStorage & out, std::string * err, int cache_slots, - bool allocate_cold) { + bool allocate_cold, + ggml_backend_t cold_gpu_backend) { if (!placement.matches(cfg)) { if (err) *err = "placement does not match config"; @@ -428,7 +538,9 @@ bool build_moe_hybrid_storage_from_file( out.cold_backend_kind = cfg.cold_expert_backend; out.materialized_hot_experts = cfg.materialize_hot_experts; out.materialized_cold_experts = cfg.materialize_cold_experts; - out.cold_backend = (cfg.cold_expert_backend == MoeHybridColdBackend::Gpu) ? gpu_backend : out.cpu_backend; + out.cold_backend = cfg.cold_expert_backend == MoeHybridColdBackend::Gpu + ? (cold_gpu_backend ? cold_gpu_backend : gpu_backend) + : out.cpu_backend; if (!out.cold_backend) { if (err) *err = "failed to select cold expert backend"; return false; @@ -484,9 +596,34 @@ bool build_moe_hybrid_storage_from_file( return false; } + const int requested_shard = + out.cold_backend_kind == MoeHybridColdBackend::Gpu + ? cold_expert_shard_channels() : 0; + if (requested_shard > 0) { + if (dst.fused_gate_up || !desc.ffn_gate_exps || + !desc.ffn_up_exps || !desc.ffn_down_exps) { + if (err) *err = "cold expert channel sharding requires separate gate/up/down tensors"; + return false; + } + const int64_t down_block = ggml_blck_size(desc.ffn_down_exps->type); + if (requested_shard >= cfg.n_ff_exp || + requested_shard % down_block != 0 || + desc.ffn_gate_exps->ne[1] != cfg.n_ff_exp || + desc.ffn_up_exps->ne[1] != cfg.n_ff_exp || + desc.ffn_down_exps->ne[0] != cfg.n_ff_exp) { + if (err) *err = "invalid or unaligned cold expert channel shard"; + return false; + } + if (!cfg.materialize_hot_experts || !cfg.materialize_cold_experts) { + if (err) *err = "cold expert channel sharding requires both GPU owners materialized"; + return false; + } + dst.expert_shard_channels = requested_shard; + } + const int hot_count = (int)dst.hot_expert_ids.size(); const int cold_count = (int)dst.cold_expert_ids.size(); - const int spare = (cold_count > 0 && cache_slots > 0) + const int spare = (dst.expert_shard_channels == 0 && cold_count > 0 && cache_slots > 0) ? std::min(cache_slots, cold_count) : 0; const int hot_alloc = hot_count + spare; dst.hot_active = hot_count; @@ -495,9 +632,10 @@ bool build_moe_hybrid_storage_from_file( dst.spare_lru.assign((size_t)spare, 0); // Allocate hot expert tensors on GPU - if (hot_count > 0 && cfg.materialize_hot_experts) { + if ((hot_count > 0 || (dst.expert_shard_channels > 0 && cold_count > 0)) && + cfg.materialize_hot_experts) { ggml_init_params ip{}; - ip.mem_size = 16 * ggml_tensor_overhead(); + ip.mem_size = 24 * ggml_tensor_overhead(); ip.mem_buffer = nullptr; ip.no_alloc = true; dst.hot_ctx = ggml_init(ip); @@ -505,14 +643,25 @@ bool build_moe_hybrid_storage_from_file( if (err) *err = "failed to init hot_ctx"; return false; } - if (dst.fused_gate_up) { + if (hot_count > 0 && dst.fused_gate_up) { dst.gate_up_hot = new_like_with_expert_count(dst.hot_ctx, desc.ffn_gate_up_exps, hot_alloc); dst.down_hot = new_like_with_expert_count(dst.hot_ctx, desc.ffn_down_exps, hot_alloc); - } else { + } else if (hot_count > 0) { dst.gate_hot = new_like_with_expert_count(dst.hot_ctx, desc.ffn_gate_exps, hot_alloc); dst.up_hot = new_like_with_expert_count(dst.hot_ctx, desc.ffn_up_exps, hot_alloc); dst.down_hot = new_like_with_expert_count(dst.hot_ctx, desc.ffn_down_exps, hot_alloc); } + if (dst.expert_shard_channels > 0 && cold_count > 0) { + dst.gate_shard_hot = new_output_channel_slice( + dst.hot_ctx, desc.ffn_gate_exps, + dst.expert_shard_channels, cold_count); + dst.up_shard_hot = new_output_channel_slice( + dst.hot_ctx, desc.ffn_up_exps, + dst.expert_shard_channels, cold_count); + dst.down_shard_hot = new_input_channel_slice( + dst.hot_ctx, desc.ffn_down_exps, + dst.expert_shard_channels, cold_count); + } dst.hot_buf = ggml_backend_alloc_ctx_tensors(dst.hot_ctx, gpu_backend); if (!dst.hot_buf) { char msg[128]; @@ -523,7 +672,7 @@ bool build_moe_hybrid_storage_from_file( } std::vector slice_buf; - if (dst.fused_gate_up) { + if (hot_count > 0 && dst.fused_gate_up) { if (!read_expert_slices_from_mem(fd.gate_up_exps.data, fd.gate_up_exps.size, dst.hot_expert_ids, dst.gate_up_expert_bytes, slice_buf, err)) return false; @@ -532,7 +681,7 @@ bool build_moe_hybrid_storage_from_file( dst.hot_expert_ids, dst.down_expert_bytes, slice_buf, err)) return false; ggml_backend_tensor_set(dst.down_hot, slice_buf.data(), 0, slice_buf.size()); - } else { + } else if (hot_count > 0) { if (!read_expert_slices_from_mem(fd.gate_exps.data, fd.gate_exps.size, dst.hot_expert_ids, dst.gate_expert_bytes, slice_buf, err)) return false; @@ -546,6 +695,32 @@ bool build_moe_hybrid_storage_from_file( return false; ggml_backend_tensor_set(dst.down_hot, slice_buf.data(), 0, slice_buf.size()); } + if (dst.expert_shard_channels > 0 && cold_count > 0) { + if (!read_output_channel_slices_from_mem( + fd.gate_exps.data, fd.gate_exps.size, + desc.ffn_gate_exps, dst.cold_expert_ids, + 0, dst.expert_shard_channels, slice_buf, err)) { + return false; + } + ggml_backend_tensor_set(dst.gate_shard_hot, + slice_buf.data(), 0, slice_buf.size()); + if (!read_output_channel_slices_from_mem( + fd.up_exps.data, fd.up_exps.size, + desc.ffn_up_exps, dst.cold_expert_ids, + 0, dst.expert_shard_channels, slice_buf, err)) { + return false; + } + ggml_backend_tensor_set(dst.up_shard_hot, + slice_buf.data(), 0, slice_buf.size()); + if (!read_input_channel_slices_from_mem( + fd.down_exps.data, fd.down_exps.size, + desc.ffn_down_exps, dst.cold_expert_ids, + 0, dst.expert_shard_channels, slice_buf, err)) { + return false; + } + ggml_backend_tensor_set(dst.down_shard_hot, + slice_buf.data(), 0, slice_buf.size()); + } } // Allocate cold expert tensors on the selected cold backend. @@ -562,6 +737,14 @@ bool build_moe_hybrid_storage_from_file( if (dst.fused_gate_up) { dst.gate_up_cold = new_like_with_expert_count(dst.cold_ctx, desc.ffn_gate_up_exps, cold_count); dst.down_cold = new_like_with_expert_count(dst.cold_ctx, desc.ffn_down_exps, cold_count); + } else if (dst.expert_shard_channels > 0) { + const int cold_channels = cfg.n_ff_exp - dst.expert_shard_channels; + dst.gate_cold = new_output_channel_slice( + dst.cold_ctx, desc.ffn_gate_exps, cold_channels, cold_count); + dst.up_cold = new_output_channel_slice( + dst.cold_ctx, desc.ffn_up_exps, cold_channels, cold_count); + dst.down_cold = new_input_channel_slice( + dst.cold_ctx, desc.ffn_down_exps, cold_channels, cold_count); } else { dst.gate_cold = new_like_with_expert_count(dst.cold_ctx, desc.ffn_gate_exps, cold_count); dst.up_cold = new_like_with_expert_count(dst.cold_ctx, desc.ffn_up_exps, cold_count); @@ -587,6 +770,33 @@ bool build_moe_hybrid_storage_from_file( dst.cold_expert_ids, dst.down_expert_bytes, slice_buf, err)) return false; ggml_backend_tensor_set(dst.down_cold, slice_buf.data(), 0, slice_buf.size()); + } else if (dst.expert_shard_channels > 0) { + const int start = dst.expert_shard_channels; + const int count = cfg.n_ff_exp - start; + if (!read_output_channel_slices_from_mem( + fd.gate_exps.data, fd.gate_exps.size, + desc.ffn_gate_exps, dst.cold_expert_ids, + start, count, slice_buf, err)) { + return false; + } + ggml_backend_tensor_set(dst.gate_cold, + slice_buf.data(), 0, slice_buf.size()); + if (!read_output_channel_slices_from_mem( + fd.up_exps.data, fd.up_exps.size, + desc.ffn_up_exps, dst.cold_expert_ids, + start, count, slice_buf, err)) { + return false; + } + ggml_backend_tensor_set(dst.up_cold, + slice_buf.data(), 0, slice_buf.size()); + if (!read_input_channel_slices_from_mem( + fd.down_exps.data, fd.down_exps.size, + desc.ffn_down_exps, dst.cold_expert_ids, + start, count, slice_buf, err)) { + return false; + } + ggml_backend_tensor_set(dst.down_cold, + slice_buf.data(), 0, slice_buf.size()); } else { if (!read_expert_slices_from_mem(fd.gate_exps.data, fd.gate_exps.size, dst.cold_expert_ids, dst.gate_expert_bytes, slice_buf, err)) @@ -664,6 +874,186 @@ int moe_hybrid_cache_swap_in(MoeHybridLayerStorage & st, int global_expert, return hslot; } +bool moe_hybrid_reassign_hot_experts_from_mmap( + MoeHybridStorage & storage, + ggml_backend_t hot_backend, + const std::vector> & hot_ids_by_layer, + std::string * err) { + if (!hot_backend || !storage.has_mmap() || + hot_ids_by_layer.size() != storage.layers.size() || + storage.layer_regions.size() != storage.layers.size()) { + if (err) *err = "dynamic hotset requires matching mmap-backed storage"; + return false; + } + const auto * base = static_cast(storage.mmap_data); + std::vector packed; + auto upload = [&](const ExpertFileRegion & region, + ggml_tensor * dst, + size_t expert_bytes, + const std::vector & ids) -> bool { + if (!dst || expert_bytes == 0 || ids.empty() || region.size == 0) { + return false; + } + packed.resize(expert_bytes * ids.size()); + for (size_t i = 0; i < ids.size(); ++i) { + const int32_t expert = ids[i]; + if (expert < 0) return false; + const size_t source = region.offset + + (size_t) expert * expert_bytes; + if (source + expert_bytes > storage.mmap_size || + (size_t) expert * expert_bytes + expert_bytes > region.size) { + return false; + } + std::memcpy( + packed.data() + i * expert_bytes, + base + source, expert_bytes); + } + ggml_backend_tensor_set(dst, packed.data(), 0, packed.size()); + return true; + }; + auto upload_one = [&](const ExpertFileRegion & region, + ggml_tensor * dst, + size_t expert_bytes, + int32_t expert, + int32_t dst_slot) -> bool { + if (!dst || expert_bytes == 0 || region.size == 0 || + expert < 0 || dst_slot < 0) { + return false; + } + const size_t source = region.offset + + (size_t) expert * expert_bytes; + if (source + expert_bytes > storage.mmap_size || + (size_t) expert * expert_bytes + expert_bytes > region.size || + (size_t) (dst_slot + 1) * expert_bytes > + ggml_nbytes(dst)) { + return false; + } + packed.resize(expert_bytes); + std::memcpy(packed.data(), base + source, expert_bytes); + ggml_backend_tensor_set(dst, packed.data(), + (size_t) dst_slot * expert_bytes, + expert_bytes); + return true; + }; + + for (size_t il = 0; il < storage.layers.size(); ++il) { + MoeHybridLayerStorage & layer = storage.layers[il]; + const LayerExpertRegions & regions = storage.layer_regions[il]; + const std::vector & ids = hot_ids_by_layer[il]; + if (ids.size() != (size_t) layer.hot_active || + layer.cache_slots != 0 || layer.expert_shard_channels != 0) { + if (err) *err = "dynamic hotset size does not match fixed hot stack"; + return false; + } + + // Keep the cold tensor shape at 232 experts. Every expert entering + // the hot set frees exactly one cold slot; overwrite that slot with an + // expert leaving the old hot set. This maintains an exact disjoint + // partition without the 8.57-GiB duplicate store that otherwise puts + // the Strix UMA system under paging pressure. + std::vector desired( + layer.hot_local_by_global.size(), 0); + for (int32_t expert : ids) { + if (expert < 0 || + expert >= (int32_t) desired.size()) { + if (err) *err = "dynamic hot expert id out of range"; + return false; + } + desired[(size_t) expert] = 1; + } + std::vector entering_hot; + std::vector leaving_hot; + for (int32_t expert : ids) { + if (layer.hot_local_by_global[(size_t) expert] < 0) { + entering_hot.push_back(expert); + } + } + for (int32_t expert : layer.hot_expert_ids) { + if (!desired[(size_t) expert]) { + leaving_hot.push_back(expert); + } + } + if (entering_hot.size() != leaving_hot.size()) { + if (err) *err = "dynamic hot/cold swap cardinality mismatch"; + return false; + } + for (size_t si = 0; si < entering_hot.size(); ++si) { + const int32_t incoming = entering_hot[si]; + const int32_t outgoing = leaving_hot[si]; + const int32_t cold_slot = + layer.cold_local_by_global[(size_t) incoming]; + if (cold_slot < 0 || + cold_slot >= (int32_t) layer.cold_expert_ids.size()) { + if (err) *err = "dynamic hot expert has no cold source slot"; + return false; + } + if (layer.fused_gate_up) { + if (!upload_one(regions.gate_up_exps, layer.gate_up_cold, + layer.gate_up_expert_bytes, outgoing, cold_slot) || + !upload_one(regions.down_exps, layer.down_cold, + layer.down_expert_bytes, outgoing, cold_slot)) { + if (err) *err = "failed to swap fused cold expert"; + return false; + } + } else { + if (!upload_one(regions.gate_exps, layer.gate_cold, + layer.gate_expert_bytes, outgoing, cold_slot) || + !upload_one(regions.up_exps, layer.up_cold, + layer.up_expert_bytes, outgoing, cold_slot) || + !upload_one(regions.down_exps, layer.down_cold, + layer.down_expert_bytes, outgoing, cold_slot)) { + if (err) *err = "failed to swap cold expert"; + return false; + } + } + layer.cold_expert_ids[(size_t) cold_slot] = outgoing; + layer.cold_local_by_global[(size_t) incoming] = -1; + layer.cold_local_by_global[(size_t) outgoing] = cold_slot; + } + + if (layer.fused_gate_up) { + if (!upload(regions.gate_up_exps, layer.gate_up_hot, + layer.gate_up_expert_bytes, ids) || + !upload(regions.down_exps, layer.down_hot, + layer.down_expert_bytes, ids)) { + if (err) *err = "failed to upload fused dynamic hot experts"; + return false; + } + } else { + if (!upload(regions.gate_exps, layer.gate_hot, + layer.gate_expert_bytes, ids) || + !upload(regions.up_exps, layer.up_hot, + layer.up_expert_bytes, ids) || + !upload(regions.down_exps, layer.down_hot, + layer.down_expert_bytes, ids)) { + if (err) *err = "failed to upload dynamic hot experts"; + return false; + } + } + layer.hot_expert_ids = ids; + std::fill(layer.hot_local_by_global.begin(), + layer.hot_local_by_global.end(), -1); + std::memset(layer.expert_vram_mask, 0, sizeof(layer.expert_vram_mask)); + for (size_t slot = 0; slot < ids.size(); ++slot) { + const int32_t expert = ids[slot]; + if (expert < 0 || expert >= (int32_t) layer.hot_local_by_global.size()) { + if (err) *err = "dynamic hot expert id out of range"; + return false; + } + layer.hot_local_by_global[(size_t) expert] = (int32_t) slot; + if (expert < 256) { + layer.expert_vram_mask[expert >> 6] |= + 1ULL << (expert & 63); + } + } + } + ggml_backend_synchronize(hot_backend); + if (storage.cold_backend && storage.cold_backend != hot_backend) { + ggml_backend_synchronize(storage.cold_backend); + } + return true; +} + MoeSparkBudget spark_budget_split(uint64_t expert_budget, uint64_t total_expert_bytes, int n_expert, uint64_t core_kv_safety, uint64_t target_bytes) { @@ -698,10 +1088,13 @@ bool build_moe_hybrid_storage_from_file_with_mmap( size_t mmap_total_size, MoeHybridStorage & out, std::string * err, - int cache_slots) { + int cache_slots, + ggml_backend_t cold_gpu_backend) { // First build storage normally (hot GPU + cold CPU buffers). - if (!build_moe_hybrid_storage_from_file(cfg, gpu_backend, placement, layer_descs, file_data, out, err, cache_slots)) { + if (!build_moe_hybrid_storage_from_file( + cfg, gpu_backend, placement, layer_descs, file_data, + out, err, cache_slots, true, cold_gpu_backend)) { return false; } diff --git a/server/src/common/moe_hybrid_storage.h b/server/src/common/moe_hybrid_storage.h index 1237f7515..623677fc4 100644 --- a/server/src/common/moe_hybrid_storage.h +++ b/server/src/common/moe_hybrid_storage.h @@ -14,6 +14,8 @@ namespace dflash::common { +struct MoeHybridRoutingStats; + // File region for one expert tensor (offset into mmap). struct ExpertFileRegion { size_t offset = 0; @@ -79,6 +81,15 @@ struct MoeHybridLayerStorage { ggml_tensor * down_hot = nullptr; ggml_tensor * gate_up_hot = nullptr; + // Optional prompt-independent tensor-parallel shard for every cold + // expert. The main GPU owns the leading `expert_shard_channels` FFN + // channels while the cold GPU owns the exact complementary channels in + // gate_cold/up_cold/down_cold. Full hot experts remain unsharded. + ggml_tensor * gate_shard_hot = nullptr; + ggml_tensor * up_shard_hot = nullptr; + ggml_tensor * down_shard_hot = nullptr; + int expert_shard_channels = 0; + ggml_context * cold_ctx = nullptr; ggml_backend_buffer_t cold_buf = nullptr; ggml_tensor * gate_cold = nullptr; @@ -132,6 +143,8 @@ struct MoeHybridLayerStorage { // Cached FFN graphs for common-case expert counts. CachedFfnGraph hot_graph; CachedFfnGraph cold_graph; + std::vector hot_graph_by_width; + std::vector cold_graph_by_width; // Cached batched hot-only graph for prefill sub-batches (n_tokens=4). CachedHotBatchedGraph hot_batched_graph; @@ -200,13 +213,23 @@ bool build_moe_hybrid_storage(const MoeHybridConfig & cfg, const MoeHybridPlacement & placement, const std::vector & layer_descs, MoeHybridStorage & out, - std::string * err = nullptr); + std::string * err = nullptr, + ggml_backend_t cold_gpu_backend = nullptr); // Swap a cold expert into a spare GPU cache slot (LRU evict). Returns the new // hot-local index, or -1 on failure. No-op (returns existing) if already hot. int moe_hybrid_cache_swap_in(MoeHybridLayerStorage & st, int global_expert, ggml_backend_t gpu_backend); +// Replace the fixed hot stacks in-place from the persistent expert GGUF mmap. +// Tensor shapes and addresses remain unchanged, so cached graphs stay valid; +// only the global->local LUTs need refreshing before their next execution. +bool moe_hybrid_reassign_hot_experts_from_mmap( + MoeHybridStorage & storage, + ggml_backend_t hot_backend, + const std::vector> & hot_ids_by_layer, + std::string * err = nullptr); + // Build hybrid storage by loading expert data directly from file (mmap). bool build_moe_hybrid_storage_from_file( const MoeHybridConfig & cfg, @@ -217,7 +240,8 @@ bool build_moe_hybrid_storage_from_file( MoeHybridStorage & out, std::string * err = nullptr, int cache_slots = 0, - bool allocate_cold = true); + bool allocate_cold = true, + ggml_backend_t cold_gpu_backend = nullptr); // Spark: split a VRAM budget into a pinned-hot tier + an auto-sized expert // cache ring. target_bytes==0 keeps the current budget (use the card); @@ -242,6 +266,7 @@ bool build_moe_hybrid_storage_from_file_with_mmap( size_t mmap_total_size, MoeHybridStorage & out, std::string * err = nullptr, - int cache_slots = 0); + int cache_slots = 0, + ggml_backend_t cold_gpu_backend = nullptr); } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 000ebc08d..1cb21d2a5 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -14,10 +14,12 @@ #include #include +#include #include #include #include #include +#include namespace dflash::common { @@ -60,6 +62,19 @@ static void configure_gfx1151_dspark_mmvq_default(int gpu) { #endif } +static bool ds4_inprocess_moe_tp_enabled() { + return env_flag_enabled("DFLASH_DS4_MOE_TP_INPROC"); +} + +static int ds4_moe_tp_gpu(int local_gpu) { + const char * raw = std::getenv("DFLASH_DS4_MOE_TP_GPU"); + if (!raw || !*raw) { + raw = std::getenv("DFLASH_MOE_EXPERT_COMPUTE_IPC_GPU"); + } + if (raw && *raw) return std::max(0, std::atoi(raw)); + return local_gpu == 0 ? 1 : 0; +} + static double gib(uint64_t bytes) { return (double) bytes / 1024.0 / 1024.0 / 1024.0; } @@ -74,6 +89,11 @@ static void add_step_tel(DeepSeek4StepTelemetry & dst, const DeepSeek4StepTeleme dst.attn_build_us += src.attn_build_us; dst.attn_compute_us += src.attn_compute_us; dst.attn_read_us += src.attn_read_us; + dst.full_graph_build_us += src.full_graph_build_us; + dst.full_graph_alloc_us += src.full_graph_alloc_us; + dst.full_graph_set_us += src.full_graph_set_us; + dst.full_graph_compute_us += src.full_graph_compute_us; + dst.full_graph_read_us += src.full_graph_read_us; dst.hc_post_attn_us += src.hc_post_attn_us; dst.hc_pre_ffn_us += src.hc_pre_ffn_us; dst.ffn_build_us += src.ffn_build_us; @@ -117,6 +137,7 @@ static void log_step_tel(const char * phase, 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_alloc=%.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 " @@ -127,6 +148,8 @@ static void log_step_tel(const char * phase, "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_alloc_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), @@ -268,6 +291,215 @@ static void fill_prefix_hot_placement(const DeepSeek4Weights & w, } } +static bool fill_profiled_hot_placement(const DeepSeek4Weights & w, + int hot_per_layer, + const char * profile_path, + MoeHybridPlacement & out, + std::string * err) { + MoeHybridRoutingStats stats; + if (!MoeHybridRoutingStats::load_csv(profile_path, stats, err)) { + return false; + } + if (stats.n_layer != w.n_layer || stats.n_expert != w.n_expert) { + if (err) { + *err = "routing profile shape does not match DeepSeek V4 target"; + } + return false; + } + + out = {}; + out.n_layer = w.n_layer; + out.n_expert = w.n_expert; + out.n_expert_used = w.n_expert_used; + out.hot_counts.assign((size_t)w.n_layer, hot_per_layer); + out.hot_expert_ids.resize((size_t)w.n_layer); + out.total_hot = hot_per_layer * w.n_layer; + for (int il = 0; il < w.n_layer; ++il) { + std::vector ranked = stats.hot_experts(il, hot_per_layer); + auto & ids = out.hot_expert_ids[(size_t)il]; + ids.assign(ranked.begin(), ranked.end()); + } + return true; +} + +// Assign the same total number of resident experts as the uniform placement, +// but distribute those slots across layers to minimize the predicted owner +// critical path. Uniform expert counts are a poor fit for heterogeneous EP: +// routing skew varies substantially by layer, while every layer joins on the +// slower of its R9700 hot/shared and Strix cold branches. +// +// The cost model intentionally uses measured bandwidth rather than advertised +// peak bandwidth. It is only an allocation objective; actual placement still +// uses authoritative router statistics and evaluates every selected expert. +static bool fill_time_balanced_profiled_hot_placement( + const DeepSeek4Weights & w, + int hot_per_layer, + const char * profile_path, + MoeHybridPlacement & out, + std::string * err) { + MoeHybridRoutingStats stats; + if (!MoeHybridRoutingStats::load_csv(profile_path, stats, err)) { + return false; + } + if (stats.n_layer != w.n_layer || stats.n_expert != w.n_expert) { + if (err) *err = "routing profile shape does not match DeepSeek V4 target"; + return false; + } + + const int total_hot_budget = hot_per_layer * w.n_layer; + if (total_hot_budget <= 0) { + if (err) *err = "time-balanced placement requires a positive hot budget"; + return false; + } + + const auto env_double = [](const char * name, double fallback) { + const char * raw = std::getenv(name); + if (!raw || !*raw) return fallback; + const double value = std::atof(raw); + return value > 0.0 ? value : fallback; + }; + const auto env_int = [](const char * name, int fallback) { + const char * raw = std::getenv(name); + if (!raw || !*raw) return fallback; + return std::atoi(raw); + }; + + // Sustained read probes on lucebox5, in GB/s. Only their ratio affects + // the optimizer, but retaining physical units makes the model auditable. + const double main_bw = env_double( + "DFLASH_DS4_TP_MAIN_OWNER_GBPS", 519.91); + const double cold_bw = env_double( + "DFLASH_DS4_TP_COLD_OWNER_GBPS", 242.36); + const double shared_equiv = env_double( + "DFLASH_DS4_TP_SHARED_EXPERT_EQUIV", 1.0); + const int min_hot = std::clamp(env_int( + "DFLASH_DS4_TP_MIN_HOT_PER_LAYER", 4), 0, w.n_expert); + const int max_hot = std::clamp(env_int( + "DFLASH_DS4_TP_MAX_HOT_PER_LAYER", 64), min_hot, w.n_expert); + + if ((int64_t) min_hot * w.n_layer > total_hot_budget || + (int64_t) max_hot * w.n_layer < total_hot_budget) { + if (err) *err = "time-balanced placement bounds cannot satisfy hot budget"; + return false; + } + + std::vector> ranked((size_t) w.n_layer); + std::vector> coverage( + (size_t) w.n_layer, + std::vector((size_t) max_hot + 1, 0.0)); + for (int il = 0; il < w.n_layer; ++il) { + ranked[(size_t) il] = stats.ranked_experts(il); + const uint64_t total = stats.layer_totals[(size_t) il]; + if (total == 0) { + if (err) *err = "routing profile contains an empty layer"; + return false; + } + uint64_t cumulative = 0; + for (int k = 1; k <= max_hot; ++k) { + cumulative += stats.count(il, ranked[(size_t) il][(size_t) k - 1]); + coverage[(size_t) il][(size_t) k] = + (double) cumulative / (double) total; + } + } + + const auto layer_cost = [&](int il, int hot_count) { + const double hot_fraction = coverage[(size_t) il][(size_t) hot_count]; + const double hot_routes = (double) w.n_expert_used * hot_fraction; + const double cold_routes = + (double) w.n_expert_used * (1.0 - hot_fraction); + const double main_time = (shared_equiv + hot_routes) / main_bw; + const double cold_time = cold_routes / cold_bw; + return std::max(main_time, cold_time); + }; + + const double inf = std::numeric_limits::infinity(); + std::vector dp((size_t) total_hot_budget + 1, inf); + std::vector next((size_t) total_hot_budget + 1, inf); + std::vector> choice( + (size_t) w.n_layer, + std::vector((size_t) total_hot_budget + 1, -1)); + dp[0] = 0.0; + + for (int il = 0; il < w.n_layer; ++il) { + std::fill(next.begin(), next.end(), inf); + const int layers_left = w.n_layer - il - 1; + for (int used = 0; used <= total_hot_budget; ++used) { + if (!std::isfinite(dp[(size_t) used])) continue; + for (int k = min_hot; k <= max_hot; ++k) { + const int new_used = used + k; + if (new_used > total_hot_budget) break; + const int remaining = total_hot_budget - new_used; + if (remaining < layers_left * min_hot || + remaining > layers_left * max_hot) { + continue; + } + const double candidate = dp[(size_t) used] + layer_cost(il, k); + if (candidate < next[(size_t) new_used]) { + next[(size_t) new_used] = candidate; + choice[(size_t) il][(size_t) new_used] = (int16_t) k; + } + } + } + dp.swap(next); + } + + if (!std::isfinite(dp[(size_t) total_hot_budget])) { + if (err) *err = "time-balanced placement optimizer found no feasible assignment"; + return false; + } + + std::vector hot_counts((size_t) w.n_layer, 0); + int used = total_hot_budget; + for (int il = w.n_layer - 1; il >= 0; --il) { + const int k = choice[(size_t) il][(size_t) used]; + if (k < min_hot || k > max_hot) { + if (err) *err = "time-balanced placement backtracking failed"; + return false; + } + hot_counts[(size_t) il] = k; + used -= k; + } + + double uniform_cost = 0.0; + double balanced_cost = 0.0; + for (int il = 0; il < w.n_layer; ++il) { + uniform_cost += layer_cost(il, hot_per_layer); + balanced_cost += layer_cost(il, hot_counts[(size_t) il]); + } + + out = {}; + out.n_layer = w.n_layer; + out.n_expert = w.n_expert; + out.n_expert_used = w.n_expert_used; + out.hot_counts = hot_counts; + out.hot_expert_ids.resize((size_t) w.n_layer); + out.total_hot = total_hot_budget; + for (int il = 0; il < w.n_layer; ++il) { + const int k = hot_counts[(size_t) il]; + auto & ids = out.hot_expert_ids[(size_t) il]; + ids.reserve((size_t) k); + for (int rank = 0; rank < k; ++rank) { + ids.push_back((int32_t) ranked[(size_t) il][(size_t) rank]); + } + } + + const auto [min_it, max_it] = std::minmax_element( + hot_counts.begin(), hot_counts.end()); + std::fprintf(stderr, + "[deepseek4] time-balanced placement: slots=%d layer_range=%d..%d " + "main/cold=%.2f/%.2f GB/s predicted_owner_reduction=%.2f%% counts=", + total_hot_budget, *min_it, *max_it, main_bw, cold_bw, + uniform_cost > 0.0 + ? 100.0 * (1.0 - balanced_cost / uniform_cost) + : 0.0); + for (int il = 0; il < w.n_layer; ++il) { + std::fprintf(stderr, "%s%d", il == 0 ? "[" : ",", + hot_counts[(size_t) il]); + } + std::fprintf(stderr, "]\n"); + return true; +} + static bool compute_ds4_hybrid_budget_info(const DeepSeek4Weights & w, int gpu, int max_ctx, @@ -334,6 +566,17 @@ static MoeHybridConfig make_ds4_parent_cpu_tail_cfg(const DeepSeek4Weights & w) return cfg; } +static MoeLayerDesc make_ds4_expert_layer_desc(const DeepSeek4Layer & layer) { + MoeLayerDesc desc; + desc.ffn_gate_exps = layer.ffn_gate_exps; + desc.ffn_up_exps = layer.ffn_up_exps; + desc.ffn_down_exps = layer.ffn_down_exps; + desc.ffn_gate_shexp = layer.ffn_gate_shexp; + desc.ffn_up_shexp = layer.ffn_up_shexp; + desc.ffn_down_shexp = layer.ffn_down_shexp; + return desc; +} + } // namespace DeepSeek4Backend::DeepSeek4Backend(const DeepSeek4BackendConfig & cfg) @@ -349,22 +592,18 @@ bool DeepSeek4Backend::load_model() { ? compiled_placement_backend() : cfg_.device.backend; - // The fused graph references every expert tensor directly, so it cannot - // run on the hybrid expert representation. Honor an explicit fused-decode - // request by trying the monolithic load first. Large HIP models otherwise - // keep using the hybrid path: a managed allocation can stall or be killed - // on integrated UMA systems before an OOM fallback is possible. - if (target_backend == PlacementBackend::Hip && cfg_.fused_decode) { + // HIP single-device launches should avoid the monolithic full-model load: + // a managed ~80 GiB allocation can stall or be killed on integrated UMA + // systems before we ever reach the existing OOM fallback path. + if (target_backend == PlacementBackend::Hip && + env_flag_enabled("DFLASH_DS4_FORCE_FULL_LOAD")) { std::fprintf(stderr, - "[deepseek4] fused decode requested; loading monolithic HIP model\n"); + "[deepseek4] HIP full-model load explicitly enabled\n"); if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { std::fprintf(stderr, - "[deepseek4] monolithic HIP load failed; trying hybrid mode\n"); - if (!init_hybrid_model()) { - std::fprintf(stderr, "[deepseek4] hybrid mode also failed: %s\n", - cfg_.model_path); - return false; - } + "[deepseek4] explicit HIP full-model load failed: %s\n", + cfg_.model_path); + return false; } } else if (target_backend == PlacementBackend::Hip) { std::fprintf(stderr, @@ -399,17 +638,58 @@ bool DeepSeek4Backend::load_model() { bool DeepSeek4Backend::load_spec_drafter() { if (spec_draft_path_.empty()) return true; - if (parked_ || moe_hybrid_) { + if (parked_) { std::fprintf(stderr, - "[deepseek4] cannot load DSpark drafter without a resident " - "monolithic target\n"); + "[deepseek4] cannot load DSpark drafter while target is parked\n"); return false; } + ggml_backend_t draft_backend = backend_; + int draft_gpu = cfg_.device.gpu; + if (const char * gpu = std::getenv("DFLASH_DS4_DRAFT_GPU")) { + draft_gpu = std::max(0, std::atoi(gpu)); + } + const bool separate_draft_stream = + env_flag_enabled("DFLASH_DS4_DRAFT_SEPARATE_STREAM"); + if (draft_gpu != cfg_.device.gpu || separate_draft_stream) { + spec_backend_ = ggml_backend_cuda_init(draft_gpu); + if (!spec_backend_) { + std::fprintf(stderr, + "[deepseek4] failed to initialize DSpark GPU %d\n", + draft_gpu); + return false; + } + draft_backend = spec_backend_; + const bool low_priority = separate_draft_stream && + env_flag_enabled("DFLASH_DS4_DRAFT_LOW_PRIORITY"); + const bool priority_configured = low_priority && + ggml_backend_cuda_set_low_priority_stream(spec_backend_); + std::fprintf(stderr, + "[deepseek4] DSpark backend gpu=%d target_gpu=%d " + "separate_stream=%d low_priority=%d\n", + draft_gpu, cfg_.device.gpu, + (int) separate_draft_stream, + (int) priority_configured); + } + auto drafter = std::make_unique(); - if (!load_deepseek4_dspark_drafter(spec_draft_path_, backend_, *drafter)) { + if (!load_deepseek4_dspark_drafter( + spec_draft_path_, draft_backend, *drafter)) { std::fprintf(stderr, "[deepseek4] DSpark drafter load FAILED: %s\n", deepseek4_dspark_last_error()); + if (spec_backend_) { + ggml_backend_free(spec_backend_); + spec_backend_ = nullptr; + } + return false; + } + + if (spec_backend_ && !clone_deepseek4_dspark_heads(*drafter, backend_)) { + std::fprintf(stderr, + "[deepseek4] failed to clone DSpark sampling heads to target GPU\n"); + free_deepseek4_dspark_drafter(*drafter); + ggml_backend_free(spec_backend_); + spec_backend_ = nullptr; return false; } @@ -429,6 +709,10 @@ bool DeepSeek4Backend::load_spec_drafter() { w_.n_embd, w_.n_vocab, w_.n_layer, d.core.n_embd, d.vocab_size); free_deepseek4_dspark_drafter(*drafter); + if (spec_backend_) { + ggml_backend_free(spec_backend_); + spec_backend_ = nullptr; + } return false; } @@ -478,6 +762,10 @@ void DeepSeek4Backend::release_spec_drafter(bool mark_parked) { free_deepseek4_dspark_drafter(*spec_drafter_); } spec_drafter_.reset(); + if (spec_backend_) { + ggml_backend_free(spec_backend_); + spec_backend_ = nullptr; + } spec_enabled_ = false; spec_feat_window_.clear(); spec_drafter_parked_ = mark_parked && !spec_draft_path_.empty(); @@ -508,9 +796,31 @@ bool DeepSeek4Backend::init() { return false; } - if (moe_hybrid_) { - // Expert IPC removed — layer split replaces expert split. - // The DeepSeek4Backend single-GPU path now runs all experts locally. + if (env_flag_enabled("DFLASH_DS4_MOE_TP") && !init_moe_tensor_parallel()) { + return false; + } + + if (const char * stats_path = std::getenv("DFLASH_DS4_ROUTING_STATS_OUT")) { + if (*stats_path) { + routing_stats_ = std::make_shared(); + if (!routing_stats_->init(w_.n_layer, w_.n_expert, w_.n_expert_used)) { + std::fprintf(stderr, "[deepseek4] failed to initialize routing stats\n"); + return false; + } + routing_stats_out_path_ = stats_path; + std::fprintf(stderr, "[deepseek4] routing stats enabled output=%s\n", + routing_stats_out_path_.c_str()); + } + } + if (env_flag_enabled("DFLASH_DS4_TP_DYNAMIC_HOTSET") && !routing_stats_) { + routing_stats_ = std::make_shared(); + if (!routing_stats_->init(w_.n_layer, w_.n_expert, w_.n_expert_used)) { + std::fprintf(stderr, + "[deepseek4] failed to initialize dynamic-hotset routing stats\n"); + return false; + } + std::fprintf(stderr, + "[deepseek4-moe-tp] request-adaptive hotset enabled\n"); } const int active_experts = @@ -526,15 +836,9 @@ bool DeepSeek4Backend::init() { const char * dp = std::getenv("DFLASH_DS4_DRAFT"); if (dp && *dp) { spec_draft_path_ = dp; - if (moe_hybrid_) { - std::fprintf(stderr, - "[deepseek4] DSpark spec-decode requires monolithic model " - "placement; disabled for hybrid expert placement\n"); - } else { - const bool loaded = load_spec_drafter(); - if (!loaded && env_flag_enabled("DFLASH_DS4_DRAFT_IPC_REQUIRED")) { - return false; - } + const bool loaded = load_spec_drafter(); + if (!loaded && env_flag_enabled("DFLASH_DS4_DRAFT_IPC_REQUIRED")) { + return false; } } else { std::fprintf(stderr, "[deepseek4] DFLASH_DS4_SPEC set but DFLASH_DS4_DRAFT gguf missing\n"); @@ -543,6 +847,61 @@ bool DeepSeek4Backend::init() { return true; } +bool DeepSeek4Backend::init_moe_tensor_parallel() { + if (!moe_hybrid_) { + std::fprintf(stderr, + "[deepseek4-moe-tp] requires a partial local expert placement\n"); + return false; + } + + if (ds4_inprocess_moe_tp_enabled()) { + if (!expert_backend_ || !moe_hybrid_->materialized_cold_experts || + moe_hybrid_->cold_backend != expert_backend_) { + std::fprintf(stderr, + "[deepseek4-moe-tp] in-process expert backend is not ready\n"); + return false; + } + expert_runtime_.reset(); + std::fprintf(stderr, + "[deepseek4-moe-tp] enabled mode=in-process local_gpu=%d " + "expert_gpu=%d local_experts=%d remote_experts=%d\n", + cfg_.device.gpu, ds4_moe_tp_gpu(cfg_.device.gpu), + moe_placement_.total_hot, + w_.n_layer * w_.n_expert - moe_placement_.total_hot); + return true; + } + + std::vector layer_descs((size_t)w_.n_layer); + for (int il = 0; il < w_.n_layer; ++il) { + layer_descs[(size_t)il] = make_ds4_expert_layer_desc(w_.layers[(size_t)il]); + } + + MoeExpertComputeRuntimeConfig runtime_cfg; + runtime_cfg.target_path = cfg_.model_path; + runtime_cfg.n_layer = w_.n_layer; + runtime_cfg.n_expert = w_.n_expert; + runtime_cfg.n_expert_used = w_.n_expert_used; + runtime_cfg.n_embd = w_.n_embd; + runtime_cfg.n_ff_exp = w_.n_ff_exp; + runtime_cfg.enabled = true; + runtime_cfg.require_remote = true; + runtime_cfg.log_prefix = "[deepseek4-moe-tp]"; + + std::string err; + if (!ensure_moe_expert_compute_runtime(expert_runtime_, runtime_cfg, + *moe_hybrid_, layer_descs, &err)) { + std::fprintf(stderr, "[deepseek4-moe-tp] initialization failed: %s\n", + err.c_str()); + return false; + } + + std::fprintf(stderr, + "[deepseek4-moe-tp] enabled local_experts=%d remote_experts=%d\n", + moe_placement_.total_hot, + w_.n_layer * w_.n_expert - moe_placement_.total_hot); + return true; +} + bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & w, int max_ctx, MoeHybridPlacement & out, @@ -552,8 +911,66 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & return false; } - const int hot_per_layer = budget.max_hot_per_layer; - fill_prefix_hot_placement(w, hot_per_layer, out); + const bool all_cold = env_flag_enabled("DFLASH_DS4_MOE_TP_ALL_COLD"); + int hot_per_layer = all_cold ? 0 : budget.max_hot_per_layer; + int shard_channels = 0; + if (!all_cold && ds4_inprocess_moe_tp_enabled()) { + if (const char * raw = std::getenv("DFLASH_DS4_TP_COLD_SHARD_CHANNELS")) { + shard_channels = std::max(0, std::atoi(raw)); + } + } + if (shard_channels > 0) { + if (shard_channels >= w.n_ff_exp || w.n_ff_exp <= 0) { + if (err) *err = "invalid DFLASH_DS4_TP_COLD_SHARD_CHANNELS"; + return false; + } + // Every non-hot expert keeps a prefix shard on the main GPU. Reserve + // that prompt-independent footprint first, then spend the remainder + // upgrading selected experts from a shard to their full tensors. + const uint64_t shard_base = + budget.mem.total_expert_bytes * (uint64_t) shard_channels / + (uint64_t) w.n_ff_exp; + const uint64_t upgrade_round = + budget.mem.bytes_per_uniform_round * + (uint64_t) (w.n_ff_exp - shard_channels) / + (uint64_t) w.n_ff_exp; + if (budget.expert_budget <= shard_base || upgrade_round == 0) { + if (err) *err = "expert budget cannot fit the universal cold-expert shard"; + return false; + } + hot_per_layer = std::min( + w.n_expert, + (int) ((budget.expert_budget - shard_base) / upgrade_round)); + std::fprintf(stderr, + "[deepseek4-tp] cold channel shard=%d/%d reserve=%.2f GiB full_hot/layer=%d\n", + shard_channels, w.n_ff_exp, gib(shard_base), hot_per_layer); + } + if (all_cold) { + std::fprintf(stderr, + "[deepseek4-moe-tp] all routed experts assigned to the cold backend\n"); + } + if (const char * profile_path = std::getenv("DFLASH_DS4_HOTNESS_CSV")) { + if (*profile_path) { + const bool time_balanced = + env_flag_enabled("DFLASH_DS4_TP_TIME_BALANCED_PLACEMENT"); + const bool placed = time_balanced + ? fill_time_balanced_profiled_hot_placement( + w, hot_per_layer, profile_path, out, err) + : fill_profiled_hot_placement( + w, hot_per_layer, profile_path, out, err); + if (!placed) { + return false; + } + std::fprintf(stderr, + "[deepseek4] hybrid placement profile=%s mode=%s\n", + profile_path, + time_balanced ? "time-balanced" : "uniform"); + } else { + fill_prefix_hot_placement(w, hot_per_layer, out); + } + } else { + fill_prefix_hot_placement(w, hot_per_layer, out); + } Ds4ExpertMemoryInfo placed_mem; if (!compute_ds4_expert_memory_info(w, &out, placed_mem, err)) { @@ -601,10 +1018,34 @@ bool DeepSeek4Backend::init_hybrid_model() { } auto hybrid = std::make_shared(); - const MoeHybridConfig hybrid_cfg = make_ds4_parent_worker_cfg(w_); + MoeHybridConfig hybrid_cfg = make_ds4_parent_worker_cfg(w_); + const bool inprocess_tp = + env_flag_enabled("DFLASH_DS4_MOE_TP") && ds4_inprocess_moe_tp_enabled(); + if (inprocess_tp) { + const int expert_gpu = ds4_moe_tp_gpu(cfg_.device.gpu); + if (expert_gpu == cfg_.device.gpu) { + std::fprintf(stderr, + "[deepseek4-moe-tp] in-process expert GPU must differ from local GPU\n"); + return false; + } + expert_backend_ = ggml_backend_cuda_init(expert_gpu); + if (!expert_backend_) { + std::fprintf(stderr, + "[deepseek4-moe-tp] failed to initialize in-process expert GPU %d\n", + expert_gpu); + return false; + } + hybrid_cfg.materialize_cold_experts = true; + hybrid_cfg.cold_expert_backend = MoeHybridColdBackend::Gpu; + } if (!build_deepseek4_moe_hybrid_storage_from_file_with_mmap( - cfg_.model_path, backend_, w_, moe_placement_, &hybrid_cfg, *hybrid, &err)) { + cfg_.model_path, backend_, w_, moe_placement_, &hybrid_cfg, + *hybrid, &err, expert_backend_)) { std::fprintf(stderr, "[deepseek4] failed to build hybrid expert storage: %s\n", err.c_str()); + if (expert_backend_) { + ggml_backend_free(expert_backend_); + expert_backend_ = nullptr; + } return false; } @@ -664,8 +1105,13 @@ bool DeepSeek4Backend::park(const std::string & what) { } last_logits_.clear(); free_deepseek4_cache(cache_); + expert_runtime_.reset(); stream_engine_.destroy(); moe_hybrid_.reset(); + if (expert_backend_) { + ggml_backend_free(expert_backend_); + expert_backend_ = nullptr; + } moe_placement_ = {}; free_deepseek4_weights(w_); parked_ = true; @@ -689,6 +1135,10 @@ bool DeepSeek4Backend::unpark(const std::string & what) { free_deepseek4_weights(w_); stream_engine_.destroy(); moe_hybrid_.reset(); + if (expert_backend_) { + ggml_backend_free(expert_backend_); + expert_backend_ = nullptr; + } moe_placement_ = {}; return false; } @@ -706,6 +1156,21 @@ bool DeepSeek4Backend::unpark(const std::string & what) { return false; } + if (env_flag_enabled("DFLASH_DS4_MOE_TP") && + !init_moe_tensor_parallel()) { + free_deepseek4_cache(cache_); + free_deepseek4_weights(w_); + expert_runtime_.reset(); + stream_engine_.destroy(); + moe_hybrid_.reset(); + if (expert_backend_) { + ggml_backend_free(expert_backend_); + expert_backend_ = nullptr; + } + moe_placement_ = {}; + return false; + } + parked_ = false; std::printf("[deepseek4] target unparked (VRAM restored)\n"); std::fflush(stdout); @@ -785,6 +1250,8 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, if (timing) step_tel.embed_us = elapsed_us(embed_t0, Clock::now()); std::vector logits; + bool ok = false; + std::vector hc_state; Ds4VerifyHooks spec_hooks; std::vector spec_cap; Ds4VerifyHooks * hp = nullptr; @@ -793,18 +1260,38 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, spec_hooks.capture_out = &spec_cap; hp = &spec_hooks; } - const bool ok = deepseek4_step( - backend_, w_, cache_, embed.data(), n_tok, pos, logits, - moe_hybrid_.get(), tokens.data() + i, - moe_hybrid_ ? &stream_engine_ : nullptr, - timing ? &step_tel : nullptr, - routing_stats_.get(), hp); + if (moe_hybrid_ && (expert_runtime_.compute || expert_backend_)) { + ok = deepseek4_step_layer_range( + backend_, w_, cache_, hc_state, + embed.data(), n_tok, pos, + 0, w_.n_layer, &logits, + tokens.data() + i, + timing ? &step_tel : nullptr, + /*allow_decode_graph_reuse=*/true, hp, + moe_hybrid_.get(), + expert_runtime_.compute ? &expert_runtime_ : nullptr, + routing_stats_.get()); + } else if (moe_hybrid_) { + ok = deepseek4_step(backend_, w_, cache_, embed.data(), n_tok, pos, logits, + moe_hybrid_.get(), tokens.data() + i, + &stream_engine_, + timing ? &step_tel : nullptr, + routing_stats_.get(), + hp, + expert_runtime_.compute ? &expert_runtime_ : nullptr); + } else { + ok = deepseek4_step_layer_range(backend_, w_, cache_, hc_state, + embed.data(), n_tok, pos, + 0, w_.n_layer, &logits, + tokens.data() + i, + timing ? &step_tel : nullptr, + /*allow_decode_graph_reuse=*/true, hp); + } if (ok && hp && !spec_cap.empty()) { const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; const int first_capture = std::max(0, spec_capture_from - i); for (int t = first_capture; t < n_tok; ++t) { - spec_feat_window_.insert( - spec_feat_window_.end(), + spec_feat_window_.insert(spec_feat_window_.end(), spec_cap.begin() + (size_t) t * feat_row, spec_cap.begin() + (size_t) (t + 1) * feat_row); } @@ -866,12 +1353,37 @@ bool DeepSeek4Backend::do_decode(int committed, int n_gen, if (timing) step_tel.embed_us = elapsed_us(embed_t0, Clock::now()); const int pos = std::max(0, committed + generated - 1); - if (!deepseek4_step(backend_, w_, cache_, embed.data(), 1, - pos, logits, - moe_hybrid_.get(), &tok_to_eval, - moe_hybrid_ ? &stream_engine_ : nullptr, - timing ? &step_tel : nullptr, - routing_stats_.get())) { + bool ok = false; + if (moe_hybrid_ && (expert_runtime_.compute || expert_backend_)) { + std::vector hc_state; + ok = deepseek4_step_layer_range( + backend_, w_, cache_, hc_state, + embed.data(), 1, pos, + 0, w_.n_layer, &logits, + &tok_to_eval, + timing ? &step_tel : nullptr, + /*allow_decode_graph_reuse=*/true, nullptr, + moe_hybrid_.get(), + expert_runtime_.compute ? &expert_runtime_ : nullptr, + routing_stats_.get()); + } else if (moe_hybrid_) { + ok = deepseek4_step(backend_, w_, cache_, embed.data(), 1, + pos, logits, + moe_hybrid_.get(), &tok_to_eval, + &stream_engine_, + timing ? &step_tel : nullptr, + routing_stats_.get(), + nullptr, + expert_runtime_.compute ? &expert_runtime_ : nullptr); + } else { + std::vector hc_state; + ok = deepseek4_step_layer_range(backend_, w_, cache_, hc_state, + embed.data(), 1, pos, + 0, w_.n_layer, &logits, + &tok_to_eval, + timing ? &step_tel : nullptr); + } + if (!ok) { std::fprintf(stderr, "[deepseek4] decode step failed\n"); return false; } @@ -915,6 +1427,32 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, DaemonIO out_io = io.with_token_callback(req.on_token); auto t0 = Clock::now(); + const bool dynamic_hotset = + env_flag_enabled("DFLASH_DS4_TP_DYNAMIC_HOTSET"); + if (dynamic_hotset) { + if (!moe_hybrid_ || !expert_backend_ || !routing_stats_) { + result.error = "dynamic hotset requires in-process heterogeneous MoE TP"; + return result; + } + // A previous request may have installed its own placement. Restore + // the calibrated general placement before observing this prompt, then + // collect a request-local route histogram from q1 prefill only. + std::string err; + if (!moe_hybrid_reassign_hot_experts_from_mmap( + *moe_hybrid_, backend_, moe_placement_.hot_expert_ids, &err)) { + std::fprintf(stderr, + "[deepseek4-moe-tp] dynamic-hotset restore failed: %s\n", + err.c_str()); + result.error = "dynamic hotset restore"; + return result; + } + if (!routing_stats_->init(w_.n_layer, w_.n_expert, + w_.n_expert_used)) { + result.error = "dynamic hotset stats reset"; + return result; + } + } + // Prefill int committed = do_prefill(req.prompt, out_io); if (committed < 0) { @@ -929,6 +1467,44 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, return result; } + if (dynamic_hotset) { + const auto adapt_t0 = Clock::now(); + std::vector> request_hot( + (size_t) w_.n_layer); + uint64_t covered = 0; + uint64_t total = 0; + for (int il = 0; il < w_.n_layer; ++il) { + const int slots = moe_hybrid_->layers[(size_t) il].hot_active; + std::vector ranked = routing_stats_->hot_experts(il, slots); + if ((int) ranked.size() != slots) { + result.error = "dynamic hotset empty prefill routes"; + return result; + } + auto & ids = request_hot[(size_t) il]; + ids.assign(ranked.begin(), ranked.end()); + total += routing_stats_->layer_totals[(size_t) il]; + for (int expert : ranked) { + covered += routing_stats_->count(il, expert); + } + } + std::string err; + if (!moe_hybrid_reassign_hot_experts_from_mmap( + *moe_hybrid_, backend_, request_hot, &err)) { + std::fprintf(stderr, + "[deepseek4-moe-tp] dynamic-hotset install failed: %s\n", + err.c_str()); + result.error = "dynamic hotset install"; + return result; + } + const double adapt_s = elapsed_s(adapt_t0); + result.prefill_s += adapt_s; + std::fprintf(stderr, + "[deepseek4-moe-tp] dynamic-hotset prompt_routes=%" PRIu64 + " coverage=%.3f install=%.3fs\n", + total, total ? (double) covered / (double) total : 0.0, + adapt_s); + } + if (req.n_gen <= 0) { result.succeed(); maybe_save_routing_stats(); @@ -967,7 +1543,11 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, out_io.emit(tok); return !out_io.cancelled; }, - spec_remote_drafter_.get())) { + spec_remote_drafter_.get(), + (expert_runtime_.compute || expert_backend_) + ? moe_hybrid_.get() : nullptr, + expert_runtime_.compute ? &expert_runtime_ : nullptr, + routing_stats_.get())) { result.fail(GenerateErrorCode::DecodeFailed, "DSpark speculative decode failed"); return result; @@ -1062,8 +1642,13 @@ void DeepSeek4Backend::shutdown() { free_deepseek4_snapshot(snapshots_[i]); } free_deepseek4_cache(cache_); + expert_runtime_.reset(); stream_engine_.destroy(); moe_hybrid_.reset(); + if (expert_backend_) { + ggml_backend_free(expert_backend_); + expert_backend_ = nullptr; + } routing_stats_.reset(); routing_stats_out_path_.clear(); moe_placement_ = {}; diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 6195be091..809a1b057 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -7,6 +7,7 @@ #pragma once #include "common/model_backend.h" +#include "common/moe_expert_compute.h" #include "common/sampler.h" #include "../common/moe_hybrid_placement.h" #include "../common/moe_hybrid_routing_stats.h" @@ -61,10 +62,15 @@ class DeepSeek4Backend : public ModelBackend { void shutdown() override; + const MoeHybridRoutingStats * get_routing_stats() const override { + return routing_stats_.get(); + } + private: DeepSeek4BackendConfig cfg_; ggml_backend_t backend_ = nullptr; ggml_backend_t snap_backend_ = nullptr; + ggml_backend_t expert_backend_ = nullptr; DeepSeek4Weights w_; DeepSeek4Cache cache_; bool parked_ = false; @@ -82,6 +88,7 @@ class DeepSeek4Backend : public ModelBackend { bool spec_enabled_ = false; bool spec_drafter_parked_ = false; std::string spec_draft_path_; + ggml_backend_t spec_backend_ = nullptr; std::unique_ptr spec_drafter_; std::unique_ptr spec_remote_drafter_; std::vector spec_feat_window_; @@ -103,6 +110,7 @@ class DeepSeek4Backend : public ModelBackend { bool load_model(); bool init_hybrid_model(); + bool init_moe_tensor_parallel(); bool compute_uniform_hybrid_placement(const DeepSeek4Weights & w, int max_ctx, MoeHybridPlacement & out, @@ -112,9 +120,7 @@ class DeepSeek4Backend : public ModelBackend { std::shared_ptr moe_hybrid_; MoeHybridPlacement moe_placement_; MoeHybridStreamEngine stream_engine_; - // Expert IPC removed — layer split replaces expert split. - // Kept for compilation compatibility; init_hybrid_model() is no longer called - // from the layer-split path. + MoeExpertComputeRuntime expert_runtime_; std::shared_ptr routing_stats_; std::string routing_stats_out_path_; }; diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index 83fda9bdf..587e2f91e 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -1,10 +1,12 @@ -// ─── Fused batched VERIFY graph (n_tokens = q in [2,4]) ───────────────── +// ─── Fused whole-model graph (n_tokens = q in [1,4]) ──────────────────── // One whole-model graph per (q, flush, padded-comp) shape: batched attention // (causal mask input) + batched MoE, per-token fused HC chains, in-graph -// drafter-feature capture and q-wide logits. The spec loop guarantees a +// optional drafter-feature capture and q-wide logits. The spec loop guarantees a // ratio-4 boundary can only sit at the batch's LAST token (dynamic q), so the // compressor pools at most once per step, exactly like an AR step at the last -// position. Enable with DFLASH_DS4_FUSED_VERIFY=1. +// position. It also serves opt-in q=1 hybrid decode, because the older q=1 +// fused graph is single-backend. Enable with DFLASH_DS4_FUSED_VERIFY=1 and, +// for q=1 hybrid decode, DFLASH_DS4_FUSED_HYBRID_DECODE=1. // tensor_set that tolerates inputs the graph never consumed (gallocr leaves // them unallocated, e.g. comp-emission bundles in non-flush graphs). @@ -19,9 +21,117 @@ static bool ds4_fused_verify_enabled() { return enabled == 1; } +// Zero-output-change audit for Luce Spark's proposed layer-ahead cache. A +// positive value asks the graph to predict routes that many layers ahead. +// The widest candidate list is retained once; post-compute accounting reports +// recall at the authoritative route width and at wider prefixes. +static int ds4_tp_route_predict_audit_lookahead() { + static const int value = [] { + const char * raw = std::getenv("DFLASH_DS4_TP_PREDICT_AUDIT"); + return raw ? std::max(0, std::min(4, std::atoi(raw))) : 0; + }(); + return value; +} + +static int ds4_tp_route_predict_audit_width() { + static const int value = [] { + const char * raw = std::getenv("DFLASH_DS4_TP_PREDICT_WIDTH"); + return raw ? std::max(1, std::min(32, std::atoi(raw))) : 12; + }(); + return value; +} + +// Opt-in heterogeneous q4 pipeline. The ordinary verifier batches all four +// tokens through attention and then places a hot+cold expert join at every +// transformer layer. On R9700 + Strix this exposes the tail of the slower +// owner once per transformer layer. The wavefront keeps two q2 lanes in +// flight instead: +// lane B's attention/FFN prefix is appended before lane A's deferred join, +// then lane A advances to the next layer before lane B joins. Both lanes +// remain exact causal target work; this is not drafting or prediction. +static bool ds4_tp_q2_wavefront_requested() { + static const bool enabled = + ds4_env_flag("DFLASH_DS4_TP_Q2_WAVEFRONT"); + return enabled; +} + +static bool ds4_tp_q2_wavefront_active( + int q, int, const MoeHybridStorage * hybrid) { + // build_compressor_step already supports one boundary at any index inside + // a causal batch. Exactly one q2 lane owns that boundary; its span-A + // write, pool/rotation, and optional span-B write preserve q4 semantics. + return ds4_tp_q2_wavefront_requested() && hybrid && q == 4; +} + +static bool ds4_fused_verify_trace_node( + ggml_tensor * tensor, bool ask, void *) { + static const bool trace_all = + ds4_env_flag("DFLASH_DS4_TP_SCHED_TRACE_ALL"); + const bool routed = tensor && + (tensor->op == GGML_OP_MUL_MAT_ID || + tensor->op == GGML_OP_GET_ROWS); + if (!tensor || (!trace_all && !routed)) return false; + const ggml_tensor * src0 = tensor->src[0]; + const ggml_tensor * ids = routed + ? (tensor->op == GGML_OP_MUL_MAT_ID ? tensor->src[2] : tensor->src[1]) + : nullptr; + const char * buffer = src0 && src0->buffer + ? ggml_backend_buffer_name(src0->buffer) : "unallocated"; + std::fprintf(stderr, + "[ds4-tp-trace] %s op=%s name='%s' src0=%s " + "ne=[%lld,%lld,%lld] ids=[%lld,%lld] buffer=%s " + "ids_buffer=%s ids_ptr=%p src1_ptr=%p\n", + ask ? "begin" : "done", + ggml_op_desc(tensor), tensor->name, + src0 ? ggml_type_name(src0->type) : "none", + src0 ? (long long) src0->ne[0] : 0LL, + src0 ? (long long) src0->ne[1] : 0LL, + src0 ? (long long) src0->ne[2] : 0LL, + ids ? (long long) ids->ne[0] : 0LL, + ids ? (long long) ids->ne[1] : 0LL, + buffer, + ids && ids->buffer ? ggml_backend_buffer_name(ids->buffer) + : "unallocated", + ids ? ids->data : nullptr, + tensor->src[1] ? tensor->src[1]->data : nullptr); + if (((ask && tensor->op == GGML_OP_MUL_MAT_ID) || + (!ask && tensor->op == GGML_OP_GET_ROWS)) && + ids && ids->buffer && + ids->type == GGML_TYPE_I32 && + ds4_env_flag("DFLASH_DS4_TP_SCHED_TRACE_IDS")) { + const size_t count = std::min( + (size_t) ggml_nelements(ids), 32u); + std::array values{}; + // TRACE_ALL synchronizes the preceding node, so the peer-local remap + // has completed before this diagnostic read. + ggml_backend_tensor_get(ids, values.data(), 0, + count * sizeof(values[0])); + std::fprintf(stderr, "[ds4-tp-trace] ids="); + for (size_t i = 0; i < count; ++i) { + std::fprintf(stderr, "%s%d", i ? "," : "", values[i]); + } + std::fprintf(stderr, "\n"); + if (!ask && tensor->op == GGML_OP_GET_ROWS && + tensor->buffer && tensor->type == GGML_TYPE_I32) { + const size_t out_count = std::min( + (size_t) ggml_nelements(tensor), 32u); + values.fill(0); + ggml_backend_tensor_get(tensor, values.data(), 0, + out_count * sizeof(values[0])); + std::fprintf(stderr, "[ds4-tp-trace] get_rows_out="); + for (size_t i = 0; i < out_count; ++i) { + std::fprintf(stderr, "%s%d", i ? "," : "", values[i]); + } + std::fprintf(stderr, "\n"); + } + } + return true; +} + struct Ds4FusedVerifyCache { const ggml_context * owner_ctx = nullptr; ggml_backend_t backend = nullptr; + ggml_backend_t peer_backend = nullptr; bool disabled = false; uint64_t counter = 0; std::array slots; @@ -35,6 +145,7 @@ struct Ds4FusedVerifyCache { ggml_tensor * st4 = nullptr; // i64 [1,q] ggml_tensor * st128 = nullptr; // i64 [1,q] ggml_tensor * capture = nullptr; // f32 [n_embd*ncap*q], order [ci][t] + ggml_tensor * argmax = nullptr; // i32 [q], optional greedy verify output int q = 0; void reset() { *this = Extra{}; } }; @@ -42,10 +153,250 @@ struct Ds4FusedVerifyCache { void destroy() { for (auto & s : slots) s.destroy(); for (auto & e : extra) e.reset(); + peer_backend = nullptr; } }; static Ds4FusedVerifyCache fused_verify_graph_cache; +// q=1 target-feature capture walks many prompt-position shapes. Keep those +// graphs out of the q>=2 speculative verifier LRU or every fresh prefill +// evicts the expensive warm q=3 verifier working set. +static Ds4FusedVerifyCache fused_capture_graph_cache; + +// Cached verifier graphs retain their input tensors, while a request-adaptive +// placement changes only the global->local expert maps. Refresh those small +// LUT inputs without rebuilding or reallocating the graph. Mask any hot/cold +// duplication defensively so every selected route is evaluated exactly once. +static void ds4_fused_verify_refresh_hybrid_luts( + const DeepSeek4Weights & w, + MoeHybridStorage * hybrid, + DeepSeek4FusedDecodeGraph & fg, + bool q2_wavefront) { + if (!hybrid) return; + const int32_t invalid_route = + ds4_env_flag("DFLASH_DS4_TP_MASKED_ROUTES") ? -1 : 0; + for (size_t hi = 0; hi < fg.hybrid_inputs.size(); ++hi) { + const int il = q2_wavefront ? (int) (hi / 2u) : (int) hi; + const MoeHybridLayerStorage & layer = hybrid->layers[(size_t) il]; + const MoeHybridGraphInputs & inputs = fg.hybrid_inputs[hi]; + std::vector hot_lut((size_t) w.n_expert, invalid_route); + std::vector cold_lut((size_t) w.n_expert, invalid_route); + std::vector hot_valid((size_t) w.n_expert, 0.0f); + std::vector cold_valid((size_t) w.n_expert, 0.0f); + for (int ie = 0; ie < w.n_expert; ++ie) { + const int32_t hot = layer.hot_local_by_global[(size_t) ie]; + const int32_t cold = layer.cold_local_by_global[(size_t) ie]; + if (hot >= 0) { + hot_lut[(size_t) ie] = hot; + hot_valid[(size_t) ie] = 1.0f; + } + if (cold >= 0 && hot < 0) { + cold_lut[(size_t) ie] = cold; + cold_valid[(size_t) ie] = 1.0f; + } + } + ds4_fv_set(inputs.hot_local_lut, hot_lut.data(), + sizeof(int32_t) * hot_lut.size()); + ds4_fv_set(inputs.hot_valid_lut, hot_valid.data(), + sizeof(float) * hot_valid.size()); + ds4_fv_set(inputs.cold_local_lut, cold_lut.data(), + sizeof(int32_t) * cold_lut.size()); + ds4_fv_set(inputs.cold_valid_lut, cold_valid.data(), + sizeof(float) * cold_valid.size()); + ds4_fv_set(inputs.shard_local_lut, cold_lut.data(), + sizeof(int32_t) * cold_lut.size()); + ds4_fv_set(inputs.shard_valid_lut, cold_valid.data(), + sizeof(float) * cold_valid.size()); + } +} + +// Each hybrid cache slot owns a multi-backend scheduler, and each scheduler's +// gallocr retains its per-backend scratch buffers until the slot is destroyed. +// Keeping all 12 shape variants alive can therefore exhaust the 32 GiB target +// GPU after only a few distinct speculative shapes. Two variants fit on the +// R9700 while preserving the common q=3/q=4 reuse; allow experiments to lower +// this to one (or raise it deliberately) without affecting the single-backend +// verifier cache. +static size_t ds4_fused_verify_hybrid_slot_limit() { + static const size_t limit = []() { + const char * raw = std::getenv("DFLASH_DS4_TP_FUSED_CACHE_SLOTS"); + const long requested = raw ? std::strtol(raw, nullptr, 10) : 2; + return (size_t) std::max( + 1, std::min(requested, + (long) fused_verify_graph_cache.slots.size())); + }(); + return limit; +} + +static size_t ds4_fused_capture_hybrid_slot_limit() { + static const size_t limit = []() { + const char * raw = std::getenv("DFLASH_DS4_TP_CAPTURE_CACHE_SLOTS"); + const long requested = raw ? std::strtol(raw, nullptr, 10) : 4; + return (size_t) std::max( + 1, std::min(requested, + (long) fused_capture_graph_cache.slots.size())); + }(); + return limit; +} + +static bool ds4_tp_whole_step_graphs_enabled() { + static const bool enabled = + ds4_env_flag("DFLASH_DS4_TP_WHOLE_STEP_GRAPHS"); + return enabled; +} + +// Replace the host's per-split launch sequence with one executable graph per +// GPU after two ordinary warmups. All kernels, copies, and reduction order +// remain unchanged; only their submission boundary moves to the device. The +// capture is rejected if the scheduler assigned any split to a non-CUDA +// backend because such work could not be replayed by the two device graphs. +static ggml_status ds4_fused_verify_compute_whole_step( + DeepSeek4FusedDecodeGraph & fg, int q) { + if (!fg.sched || q != 4 || !ds4_tp_whole_step_graphs_enabled()) { + return fg.sched + ? ggml_backend_sched_graph_compute(fg.sched, fg.sg.gf) + : GGML_STATUS_FAILED; + } + + std::array gpu_backends{}; + int gpu_count = 0; + int gpu_splits = 0; + const int n_backends = ggml_backend_sched_get_n_backends(fg.sched); + for (int i = 0; i < n_backends; ++i) { + ggml_backend_t candidate = + ggml_backend_sched_get_backend(fg.sched, i); + const int splits = ggml_backend_sched_get_n_splits_for_backend( + fg.sched, candidate); + if (ggml_backend_is_cuda(candidate)) { + if (gpu_count < (int) gpu_backends.size()) { + gpu_backends[(size_t) gpu_count++] = candidate; + } + gpu_splits += splits; + } else if (splits != 0) { + static bool logged_cpu_split = false; + if (!logged_cpu_split) { + std::fprintf(stderr, + "[ds4-whole-graph] disabled: non-GPU backend has %d splits\n", + splits); + logged_cpu_split = true; + } + return ggml_backend_sched_graph_compute(fg.sched, fg.sg.gf); + } + } + if (gpu_count != 2 || + gpu_splits != ggml_backend_sched_get_n_splits(fg.sched)) { + return ggml_backend_sched_graph_compute(fg.sched, fg.sg.gf); + } + + if (fg.whole_step_graph_count == 2) { + // Waits are captured compute kernels on this ROCm pair, not hardware + // semaphore packets. Enqueue the R9700 publisher first so MES never + // starts with an unsignaled Strix spin queue. R9700 dense work leaves + // ample time to enqueue Strix before the first reverse dependency. + for (int i = 0; i < 2; ++i) { + if (!ggml_backend_cuda_whole_graph_launch( + fg.whole_step_backends[(size_t) i], + fg.whole_step_graphs[(size_t) i])) { + return GGML_STATUS_FAILED; + } + } + ggml_backend_sched_synchronize(fg.sched); + return GGML_STATUS_SUCCESS; + } + + if (fg.whole_step_warmups++ < 2) { + return ggml_backend_sched_graph_compute(fg.sched, fg.sg.gf); + } + + // Both streams must be capturing before the scheduler emits the first + // cross-device event. Begin with the cold owner to match replay order. + if (!ggml_backend_cuda_whole_graph_capture_prepare( + gpu_backends[0], gpu_backends[1])) { + std::fprintf(stderr, + "[ds4-whole-graph] doorbell preparation failed; using scheduler\n"); + return ggml_backend_sched_graph_compute(fg.sched, fg.sg.gf); + } + std::fprintf(stderr, "[ds4-whole-graph] capture begin peer\n"); + const bool began_peer = + ggml_backend_cuda_whole_graph_capture_begin(gpu_backends[1]); + std::fprintf(stderr, "[ds4-whole-graph] capture begin main peer=%d\n", + (int) began_peer); + const bool began_main = began_peer && + ggml_backend_cuda_whole_graph_capture_begin(gpu_backends[0]); + if (!began_main) { + if (began_peer) { + auto abandoned = + ggml_backend_cuda_whole_graph_capture_end(gpu_backends[1]); + ggml_backend_cuda_whole_graph_free( + gpu_backends[1], abandoned); + } + return ggml_backend_sched_graph_compute(fg.sched, fg.sg.gf); + } + + ggml_backend_sched_set_whole_graph_capture(fg.sched, true); + std::fprintf(stderr, "[ds4-whole-graph] scheduler capture submit\n"); + const ggml_status captured_status = + ggml_backend_sched_graph_compute_async(fg.sched, fg.sg.gf); + std::fprintf(stderr, "[ds4-whole-graph] scheduler capture status=%d\n", + (int) captured_status); + ggml_backend_sched_set_whole_graph_capture(fg.sched, false); + std::fprintf(stderr, "[ds4-whole-graph] capture end peer\n"); + auto peer_graph = + ggml_backend_cuda_whole_graph_capture_end(gpu_backends[1]); + std::fprintf(stderr, "[ds4-whole-graph] capture end main peer_graph=%p\n", + peer_graph); + auto main_graph = + ggml_backend_cuda_whole_graph_capture_end(gpu_backends[0]); + if (captured_status != GGML_STATUS_SUCCESS || + !main_graph || !peer_graph) { + ggml_backend_cuda_whole_graph_free(gpu_backends[0], main_graph); + ggml_backend_cuda_whole_graph_free(gpu_backends[1], peer_graph); + std::fprintf(stderr, + "[ds4-whole-graph] capture failed status=%d main=%p peer=%p\n", + (int) captured_status, main_graph, peer_graph); + return captured_status == GGML_STATUS_SUCCESS + ? GGML_STATUS_FAILED : captured_status; + } + + fg.whole_step_backends = gpu_backends; + fg.whole_step_graphs = {main_graph, peer_graph}; + fg.whole_step_graph_count = 2; + std::fprintf(stderr, + "[ds4-whole-graph] active q=4 splits=%d\n", + ggml_backend_sched_get_n_splits(fg.sched)); + for (int i = 0; i < 2; ++i) { + if (!ggml_backend_cuda_whole_graph_launch( + fg.whole_step_backends[(size_t) i], + fg.whole_step_graphs[(size_t) i])) { + return GGML_STATUS_FAILED; + } + } + ggml_backend_sched_synchronize(fg.sched); + return GGML_STATUS_SUCCESS; +} + +static ggml_tensor * ds4_fused_weights_for_selected( + ggml_context * ctx, + ggml_tensor * cur, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + ggml_tensor * selected, + int n_tokens) { + ggml_tensor * logits = ggml_mul_mat(ctx, L.ffn_gate_inp, cur); + ggml_tensor * probs = ggml_sqrt(ctx, ggml_softplus(ctx, logits)); + ggml_tensor * probs_3d = ggml_reshape_3d( + ctx, probs, 1, w.n_expert, n_tokens); + ggml_tensor * weights = ggml_get_rows(ctx, probs_3d, selected); + const int n_used = (int) selected->ne[0]; + weights = ggml_reshape_2d(ctx, weights, n_used, n_tokens); + ggml_tensor * sum = ggml_clamp( + ctx, ggml_sum_rows(ctx, weights), 6.103515625e-5f, INFINITY); + weights = ggml_div(ctx, weights, sum); + if (w.expert_weight_scale != 1.0f) { + weights = ggml_scale(ctx, weights, w.expert_weight_scale); + } + return weights; +} static bool ds4_build_fused_verify_graph( DeepSeek4FusedDecodeCache & mc, // fn mirrors (shared with AR fused) @@ -61,16 +412,28 @@ static bool ds4_build_fused_verify_graph( int q, bool have_token_ids, const std::vector & capture_ids, + MoeHybridStorage * hybrid, std::vector && shape_key) { + fg.clear_whole_step_graphs(); + if (fg.sched) { + ggml_backend_sched_free(fg.sched); + fg.sched = nullptr; + } step_graph_free(fg.sg); fg.reset_nodes(); ex.reset(); + const bool q2_wavefront = + ds4_tp_q2_wavefront_active(q, kv_start, hybrid); fg.hash_ids.assign((size_t) w.n_layer, nullptr); + if (hybrid) { + fg.hybrid_inputs.assign( + (size_t) w.n_layer * (q2_wavefront ? 2u : 1u), {}); + } const int n_embd = w.n_embd; const int n_hc = w.n_hc; - const int token_pos = kv_start + q - 1; // last batch position - const size_t arena_size = 256u * 1024 * 1024; + const size_t arena_size = + (q2_wavefront ? 512u : 256u) * 1024 * 1024; if (fg.sg.meta_arena.size() < arena_size) fg.sg.meta_arena.resize(arena_size); ggml_init_params params{}; params.mem_size = fg.sg.meta_arena.size(); @@ -79,7 +442,8 @@ 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; - fg.sg.gf = ggml_new_graph_custom(ctx, 65536, false); + const size_t graph_capacity = q2_wavefront ? 131072u : 65536u; + fg.sg.gf = ggml_new_graph_custom(ctx, graph_capacity, false); ggml_cgraph * gf = fg.sg.gf; fg.inp_embed = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, q); @@ -99,17 +463,32 @@ static bool ds4_build_fused_verify_graph( fg.i64_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1 * (int64_t) w.n_layer); ggml_set_input(fg.i64_bundle); - // mask bundle: per layer [n_swa + padded + q] rows × q tokens + // Mask bundle: q>1 preserves the q ring rows that the batch overwrites, + // so its attention span is [n_swa + padded + q]. q=1 reads the current + // 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, token_pos); - padded = ds4_padded_comp_rows(n_comp, (int) cache.layers[(size_t) il].comp_kv->ne[1]); + const int n_lanes = q2_wavefront ? 2 : 1; + for (int lane = 0; lane < n_lanes; ++lane) { + const int lane_q = q2_wavefront ? 2 : q; + const int lane_start = q2_wavefront ? 2 * lane : 0; + const int lane_token_pos = kv_start + lane_start + lane_q - 1; + 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, + lane_token_pos); + padded = ds4_padded_comp_rows( + n_comp, + (int) cache.layers[(size_t) il].comp_kv->ne[1]); + } + const int preserved_rows = lane_q > 1 ? lane_q : 0; + mask_total += + (int64_t) (w.n_swa + padded + preserved_rows) * lane_q; } - mask_total += (int64_t) (w.n_swa + padded + q) * q; } fg.mask_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, mask_total); ggml_set_input(fg.mask_bundle); @@ -132,11 +511,30 @@ static bool ds4_build_fused_verify_graph( const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; const int ratio = (int) w.compress_ratios[il]; + // In wavefront mode the two lane builds deliberately stay inside the + // same layer loop. The hybrid builder expands each lane's cold and + // hot prefixes but leaves its deferred-copy join unexpanded. Building + // lane B therefore appends useful main-GPU work before lane A's join; + // at the next outer iteration lane A's dependency pulls that join in + // before layer il+1. This produces the diagonal A/B pipeline without + // changing transformer dependencies. + const int n_lanes = q2_wavefront ? 2 : 1; + for (int lane = 0; lane < n_lanes; ++lane) { + const int lane_q = q2_wavefront ? 2 : q; + const int lane_start = q2_wavefront ? 2 * lane : 0; + const int lane_kv_start = kv_start + lane_start; + const int lane_token_pos = lane_kv_start + lane_q - 1; + const size_t hybrid_idx = q2_wavefront + ? (size_t) il * 2 + (size_t) lane + : (size_t) il; + // ── HC pre (attention), per token ── - std::vector hc_flat(q), split_attn(q); + std::vector hc_flat(lane_q), split_attn(lane_q); ggml_tensor * attn_in = nullptr; - for (int t = 0; t < q; ++t) { - hc_flat[(size_t) t] = ggml_reshape_1d(ctx, hc_cur[(size_t) t], (int64_t) n_embd * n_hc); + for (int t = 0; t < lane_q; ++t) { + const size_t token_idx = (size_t) lane_start + (size_t) t; + hc_flat[(size_t) t] = ggml_reshape_1d( + ctx, hc_cur[token_idx], (int64_t) n_embd * n_hc); ggml_tensor * working = ds4_build_fused_hc_pre(ctx, w, hc_flat[(size_t) t], mc.fn_attn_f16[(size_t) il], L.hc_attn_base, hlw.attn, &split_attn[(size_t) t]); @@ -147,40 +545,63 @@ static bool ds4_build_fused_verify_graph( // ── Batched attention ── DeepSeek4AttentionGraphInputs ain{}; - ain.rope_pos = ex.pos_q; - ain.neg_pos = ex.neg_q; - ain.raw_kv_rows = ex.rawrows; + const size_t i32_lane_off = + (size_t) lane_start * sizeof(int32_t); + const size_t i64_lane_off = + (size_t) lane_start * sizeof(int64_t); + ain.rope_pos = q2_wavefront + ? ggml_view_1d(ctx, ex.pos_q, lane_q, i32_lane_off) + : ex.pos_q; + ain.neg_pos = q2_wavefront + ? ggml_view_1d(ctx, ex.neg_q, lane_q, i32_lane_off) + : ex.neg_q; + ain.raw_kv_rows = q2_wavefront + ? ggml_view_1d(ctx, ex.rawrows, lane_q, i64_lane_off) + : ex.rawrows; if (ratio > 0) { - ain.attn_ape_row = (ratio == 4) ? ex.ape4 : ex.ape128; - ain.attn_state_rows = (ratio == 4) ? ex.st4 : ex.st128; + ggml_tensor * ape_all = (ratio == 4) ? ex.ape4 : ex.ape128; + ggml_tensor * state_all = (ratio == 4) ? ex.st4 : ex.st128; + ain.attn_ape_row = q2_wavefront + ? ggml_view_1d(ctx, ape_all, lane_q, i32_lane_off) + : ape_all; + ain.attn_state_rows = q2_wavefront + ? ggml_view_1d(ctx, state_all, lane_q, i64_lane_off) + : state_all; ain.attn_comp_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 2 + 1) * sizeof(int32_t)); ain.attn_comp_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), (size_t) il * sizeof(int64_t)); } if (ratio == 4) { - ain.index_ape_row = ex.ape4; - ain.index_state_rows = ex.st4; + ain.index_ape_row = q2_wavefront + ? ggml_view_1d(ctx, ex.ape4, lane_q, i32_lane_off) + : ex.ape4; + ain.index_state_rows = q2_wavefront + ? ggml_view_1d(ctx, ex.st4, lane_q, i64_lane_off) + : ex.st4; ain.index_comp_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 2 + 1) * sizeof(int32_t)); ain.index_comp_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), (size_t) il * sizeof(int64_t)); } int padded = 0; if (ratio > 0 && lc.comp_kv) { - const int n_comp = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos); + 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]); } - const int64_t n_attn = (int64_t) w.n_swa + padded + q; - ain.attn_row_mask = ggml_view_2d(ctx, fg.mask_bundle, n_attn, q, + const int64_t n_attn = (int64_t) w.n_swa + padded + + (lane_q > 1 ? lane_q : 0); + ain.attn_row_mask = ggml_view_2d(ctx, fg.mask_bundle, n_attn, lane_q, n_attn * sizeof(float), (size_t) mask_off * sizeof(float)); ain.padded_comp = padded; - mask_off += n_attn * q; + mask_off += n_attn * lane_q; std::vector i32b; 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, - kv_start, q, &ain, i32b, i32ab, i64ab); + lane_kv_start, lane_q, &ain, + i32b, i32ab, i64ab); if (!attn_out) return false; if (!i32b.empty() || !i32ab.empty() || !i64ab.empty()) { std::fprintf(stderr, "[ds4-fused-verify] layer %d dynamic bindings; cannot fuse\n", il); @@ -189,15 +610,18 @@ static bool ds4_build_fused_verify_graph( // ── HC post (attention) + HC pre (FFN), per token ── ggml_tensor * ffn_in = nullptr; - std::vector split_ffn(q); - for (int t = 0; t < q; ++t) { + std::vector split_ffn(lane_q); + for (int t = 0; t < lane_q; ++t) { + const size_t token_idx = (size_t) lane_start + (size_t) t; ggml_tensor * ao = ggml_view_2d(ctx, attn_out, n_embd, 1, attn_out->nb[1], (size_t) t * attn_out->nb[1]); ggml_tensor * ao_flat = ggml_reshape_1d(ctx, ggml_cont(ctx, ao), n_embd); ggml_tensor * hc_next = ggml_ds4_hc_post(ctx, hc_flat[(size_t) t], ao_flat, split_attn[(size_t) t], n_hc); - hc_cur[(size_t) t] = ggml_reshape_2d(ctx, hc_next, n_embd, n_hc); - hc_flat[(size_t) t] = ggml_reshape_1d(ctx, hc_cur[(size_t) t], (int64_t) n_embd * n_hc); + hc_cur[token_idx] = ggml_reshape_2d( + ctx, hc_next, n_embd, n_hc); + hc_flat[(size_t) t] = ggml_reshape_1d( + ctx, hc_cur[token_idx], (int64_t) n_embd * n_hc); ggml_tensor * fworking = ds4_build_fused_hc_pre(ctx, w, hc_flat[(size_t) t], mc.fn_ffn_f16[(size_t) il], L.hc_ffn_base, hlw.ffn, &split_ffn[(size_t) t]); @@ -209,41 +633,229 @@ static bool ds4_build_fused_verify_graph( // ── Batched FFN ── ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, L.ffn_norm, w.rms_eps); ggml_tensor * ffn_out = nullptr; + const MoeHybridGraphInputs * fused_hc_join_inputs = nullptr; const bool hash_routed = il < w.n_hash_layer && L.ffn_gate_tid2eid && have_token_ids && hash_tables[(size_t) il].loaded; + ggml_tensor * selected = nullptr; + ggml_tensor * router_weights = nullptr; + std::vector router_nodes; if (hash_routed) { - // ds4_build_hash_routed_ffn is single-token; run per token, concat. - const int n_used = ds4_effective_expert_count(w); - ggml_tensor * hids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_used, q); - ggml_set_input(hids); - fg.hash_ids[(size_t) il] = hids; - for (int t = 0; t < q; ++t) { - ggml_tensor * fcol = ggml_view_2d(ctx, ffn_normed, n_embd, 1, - ffn_normed->nb[1], (size_t) t * ffn_normed->nb[1]); - ggml_tensor * hcol = ggml_view_2d(ctx, hids, n_used, 1, - hids->nb[1], (size_t) t * hids->nb[1]); - ggml_tensor * fo = ds4_build_hash_routed_ffn(ctx, w, L, ggml_cont(ctx, fcol), hcol); - if (!fo) { ffn_out = nullptr; break; } - ffn_out = ffn_out ? ggml_concat(ctx, ffn_out, fo, 1) : fo; + const int route_width = ds4_effective_expert_topk(w); + // Keep one q-wide input for host filling and give each wavefront + // lane a zero-copy two-column view. + ggml_tensor * hids_all = fg.hash_ids[(size_t) il]; + if (!hids_all) { + hids_all = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, route_width, q); + ggml_set_input(hids_all); + fg.hash_ids[(size_t) il] = hids_all; + } + selected = q2_wavefront + ? ggml_view_2d( + ctx, hids_all, route_width, lane_q, + hids_all->nb[1], + (size_t) lane_start * hids_all->nb[1]) + : hids_all; + if (hybrid) { + router_weights = ds4_fused_weights_for_selected( + ctx, ffn_normed, w, L, selected, lane_q); + } else { + for (int t = 0; t < lane_q; ++t) { + ggml_tensor * fcol = ggml_view_2d(ctx, ffn_normed, n_embd, 1, + ffn_normed->nb[1], (size_t) t * ffn_normed->nb[1]); + ggml_tensor * hcol = ggml_view_2d( + ctx, selected, route_width, 1, + selected->nb[1], (size_t) t * selected->nb[1]); + ggml_tensor * fo = ds4_build_hash_routed_ffn(ctx, w, L, ggml_cont(ctx, fcol), hcol); + if (!fo) { ffn_out = nullptr; break; } + ffn_out = ffn_out ? ggml_concat(ctx, ffn_out, fo, 1) : fo; + } } + } else if (hybrid) { + Ds4MoeRouting routing = + build_moe_routing(ctx, ffn_normed, w, L, lane_q); + selected = routing.selected; + router_weights = routing.weights; + router_nodes = routing.nodes; } else { - ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, q); + ffn_out = build_moe_ffn( + ctx, ffn_normed, w, L, il, lane_q); + } + + // Diagnostic branch only: predict future score-routed expert IDs from + // the current normalized FFN input. The authoritative routing above + // remains the sole input to the model FFN, so this cannot affect + // logits. q2 wavefront is excluded to keep the audit's source/target + // correspondence unambiguous. + const int predict_ahead = ds4_tp_route_predict_audit_lookahead(); + if (hybrid && q > 1 && predict_ahead > 0 && !q2_wavefront) { + for (int ahead = 1; ahead <= predict_ahead; ++ahead) { + const int target_layer = il + ahead; + if (target_layer >= w.n_layer || + target_layer < w.n_hash_layer) { + continue; + } + const DeepSeek4Layer & future = + w.layers[(size_t) target_layer]; + if (!future.ffn_gate_inp) continue; + Ds4MoeRoutePrediction prediction = + build_moe_route_prediction( + ctx, ffn_normed, w, future, + ds4_tp_route_predict_audit_width()); + if (!prediction.selected) return false; + ggml_set_output(prediction.selected); + ggml_build_forward_expand(gf, prediction.selected); + DeepSeek4FusedDecodeGraph::RoutePredictionOutput saved; + saved.source_layer = il; + saved.target_layer = target_layer; + saved.lane_start = lane_start; + saved.n_tokens = lane_q; + saved.width = (int) prediction.selected->ne[0]; + saved.selected = prediction.selected; + saved.nodes = std::move(prediction.nodes); + fg.route_predictions.push_back(std::move(saved)); + } + } + if (hybrid) { + MoeHybridConfig hybrid_cfg = make_ds4_moe_hybrid_config(w); + hybrid_cfg.n_expert_used = (int) selected->ne[0]; + MoeLayerDesc desc = make_ds4_moe_layer_desc(L); + MoeHybridGraphInputs & inputs = fg.hybrid_inputs[hybrid_idx]; + inputs.router_nodes = router_nodes; + const MoeHybridLayerStorage & hybrid_layer = + hybrid->layers[(size_t) il]; + const bool has_hot_routed = + (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. + // 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. + const bool native_route_width = + ds4_env_flag("DFLASH_DS4_TP_NATIVE_ROUTE_WIDTH"); + if (lane_q > 1 && hybrid_cfg.n_expert_used > 4 && has_hot_routed && + !native_route_width) { + const int route_width = hybrid_cfg.n_expert_used; + 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; + } + if (!build_moe_hybrid_ffn_graph( + ctx, + (ds4_env_flag("DFLASH_DS4_TP_PEER_FENCE") || + ds4_env_flag("DFLASH_DS4_TP_DEVICE_JOIN")) ? gf : nullptr, + hybrid_cfg, desc, + hybrid->layers[(size_t) il], ffn_normed, + first_ids, first_weights, lane_q, inputs, true)) { + return false; + } + ffn_out = inputs.output; + + if (!build_moe_hybrid_ffn_graph( + ctx, + (ds4_env_flag("DFLASH_DS4_TP_PEER_FENCE") || + ds4_env_flag("DFLASH_DS4_TP_DEVICE_JOIN")) ? gf : nullptr, + hybrid_cfg, desc, + hybrid->layers[(size_t) il], ffn_normed, + padded_ids, padded_weights, lane_q, inputs, false)) { + return false; + } + ffn_out = ggml_add(ctx, ffn_out, inputs.output); + } else { + if (!build_moe_hybrid_ffn_graph( + ctx, + (ds4_env_flag("DFLASH_DS4_TP_PEER_FENCE") || + ds4_env_flag("DFLASH_DS4_TP_DEVICE_JOIN")) ? gf : nullptr, + hybrid_cfg, desc, + hybrid->layers[(size_t) il], + ffn_normed, selected, router_weights, + lane_q, inputs)) { + return false; + } + ffn_out = inputs.output; + if (lane_q > 1 && native_route_width && + ds4_env_flag("DFLASH_DS4_TP_FUSED_HC_JOIN") && + inputs.main_output && inputs.peer_output) { + fused_hc_join_inputs = &inputs; + } + } } if (!ffn_out) return false; // ── HC post (FFN), per token; capture at drafter layers ── - for (int t = 0; t < q; ++t) { - ggml_tensor * fo = ggml_view_2d(ctx, ffn_out, n_embd, 1, - ffn_out->nb[1], (size_t) t * ffn_out->nb[1]); - ggml_tensor * fo_flat = ggml_reshape_1d(ctx, ggml_cont(ctx, fo), n_embd); - ggml_tensor * hc_next = ggml_ds4_hc_post(ctx, hc_flat[(size_t) t], fo_flat, - split_ffn[(size_t) t], n_hc); - hc_cur[(size_t) t] = ggml_reshape_2d(ctx, hc_next, n_embd, n_hc); + for (int t = 0; t < lane_q; ++t) { + const size_t token_idx = (size_t) lane_start + (size_t) t; + ggml_tensor * hc_next = nullptr; + if (fused_hc_join_inputs) { + ggml_tensor * main_part = ggml_view_2d( + ctx, fused_hc_join_inputs->main_output, n_embd, 1, + fused_hc_join_inputs->main_output->nb[1], + (size_t) t * fused_hc_join_inputs->main_output->nb[1]); + ggml_tensor * peer_part = ggml_view_2d( + ctx, fused_hc_join_inputs->peer_output, n_embd, 1, + fused_hc_join_inputs->peer_output->nb[1], + (size_t) t * fused_hc_join_inputs->peer_output->nb[1]); + hc_next = ggml_ds4_hc_post_split( + ctx, hc_flat[(size_t) t], main_part, peer_part, + split_ffn[(size_t) t], n_hc); + } else { + ggml_tensor * fo = ggml_view_2d( + ctx, ffn_out, n_embd, 1, ffn_out->nb[1], + (size_t) t * ffn_out->nb[1]); + ggml_tensor * fo_flat = ggml_reshape_1d( + ctx, ggml_cont(ctx, fo), n_embd); + hc_next = ggml_ds4_hc_post( + ctx, hc_flat[(size_t) t], fo_flat, + split_ffn[(size_t) t], n_hc); + } + hc_cur[token_idx] = ggml_reshape_2d( + ctx, hc_next, n_embd, n_hc); } for (size_t ci = 0; ci < capture_ids.size(); ++ci) { if (capture_ids[ci] != il) continue; - for (int t = 0; t < q; ++t) { - ggml_tensor * hs = hc_cur[(size_t) t]; // [n_embd, n_hc] + for (int t = 0; t < lane_q; ++t) { + const size_t token_idx = + (size_t) lane_start + (size_t) t; + ggml_tensor * hs = hc_cur[token_idx]; // [n_embd, n_hc] ggml_tensor * hsT = ggml_cont(ctx, ggml_transpose(ctx, hs)); // [n_hc, n_embd] ggml_tensor * summed = ggml_sum_rows(ctx, hsT); // [1, n_embd] ggml_tensor * mean = ggml_scale(ctx, summed, 1.0f / (float) n_hc); @@ -251,6 +863,7 @@ static bool ds4_build_fused_verify_graph( ggml_reshape_1d(ctx, ggml_cont(ctx, mean), n_embd); } } + } // lane } // ── Output: per-token HC merge → batched out_norm + lm_head ── @@ -271,6 +884,11 @@ static bool ds4_build_fused_verify_graph( fg.logits = ggml_mul_mat(ctx, w.output, out_normed); // [n_vocab, q] ggml_set_output(fg.logits); ggml_build_forward_expand(gf, fg.logits); + if (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); + } if (!capture_pieces.empty()) { for (ggml_tensor * piece : capture_pieces) { @@ -288,12 +906,200 @@ static bool ds4_build_fused_verify_graph( ggml_build_forward_expand(gf, ex.capture); } - if (!fg.sg.alloc) { - fg.sg.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + // One-shot graph inventory for profiling residual single-token FP4 work. + // This is deliberately build-time only: it does not install a scheduler + // callback or synchronize kernels, so enabling it cannot distort the + // execution timeline being diagnosed. + if (q == 4 && ds4_env_flag("DFLASH_DS4_TP_Q4_INVENTORY")) { + const int n_nodes = ggml_graph_n_nodes(fg.sg.gf); + std::fprintf(stderr, + "[ds4-q4-inventory] graph_nodes=%d q=%d\n", n_nodes, q); + for (int ni = 0; ni < n_nodes; ++ni) { + ggml_tensor * node = ggml_graph_node(fg.sg.gf, ni); + if (!node || + (node->op != GGML_OP_MUL_MAT && + node->op != GGML_OP_MUL_MAT_ID) || + !node->src[0] || !node->src[1] || + node->src[0]->type != GGML_TYPE_Q4_0_ROCMFP4_FAST) { + continue; + } + const ggml_tensor * weight = node->src[0]; + const ggml_tensor * rhs = node->src[1]; + std::fprintf(stderr, + "[ds4-q4-inventory] node=%d op=%s weight='%s' " + "w=[%lld,%lld,%lld] rhs=[%lld,%lld,%lld] out=[%lld,%lld,%lld]\n", + ni, ggml_op_desc(node), weight->name, + (long long) weight->ne[0], (long long) weight->ne[1], + (long long) weight->ne[2], (long long) rhs->ne[0], + (long long) rhs->ne[1], (long long) rhs->ne[2], + (long long) node->ne[0], (long long) node->ne[1], + (long long) node->ne[2]); + } } - if (!fg.sg.alloc || !ggml_gallocr_alloc_graph(fg.sg.alloc, fg.sg.gf)) { - std::fprintf(stderr, "[ds4-fused-verify] graph alloc failed\n"); - return false; + + if (hybrid) { + ggml_backend_t peer = hybrid->cold_backend; + if (!peer || peer == backend) { + std::fprintf(stderr, + "[ds4-fused-verify] hybrid peer backend unavailable\n"); + return false; + } + if (!hybrid->cpu_backend) { + std::fprintf(stderr, + "[ds4-fused-verify] scheduler CPU fallback unavailable\n"); + return false; + } + 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; + } + const bool late_join_split = + ds4_env_flag("DFLASH_DS4_TP_LATE_JOIN_SPLIT"); + const bool targeted_join_split = + ds4_env_flag("DFLASH_DS4_TP_TARGETED_JOIN_SPLIT"); + const bool gpu_i32_repeat = + ds4_env_flag("LUCE_CUDA_I32_REPEAT"); + const bool report_split_count = + ds4_env_flag("DFLASH_DS4_TP_SPLIT_COUNT"); + ggml_backend_sched_set_late_cross_input_split( + fg.sched, late_join_split); + if (ds4_env_flag("DFLASH_DS4_TP_SCHED_TRACE")) { + ggml_backend_sched_set_eval_callback( + fg.sched, ds4_fused_verify_trace_node, nullptr); + } + auto pin_main = [&](ggml_tensor * tensor) { + if (tensor) { + ggml_backend_sched_set_tensor_backend(fg.sched, tensor, backend); + } + }; + auto pin_peer = [&](ggml_tensor * tensor) { + if (tensor) { + ggml_backend_sched_set_tensor_backend(fg.sched, tensor, peer); + } + }; + pin_main(fg.inp_embed); + pin_main(ex.pos_q); + pin_main(ex.neg_q); + pin_main(ex.rawrows); + pin_main(ex.ape4); + pin_main(ex.ape128); + pin_main(ex.st4); + pin_main(ex.st128); + pin_main(fg.i32_bundle); + pin_main(fg.i64_bundle); + pin_main(fg.mask_bundle); + for (ggml_tensor * hids : fg.hash_ids) pin_main(hids); + for (const MoeHybridGraphInputs & inputs : fg.hybrid_inputs) { + if (ds4_env_flag("DFLASH_DS4_TP_MAIN_ROUTE_WEIGHTS")) { + for (ggml_tensor * node : inputs.router_nodes) { + pin_main(node); + } + pin_main(inputs.router_weights); + } + if (ds4_env_flag("DFLASH_DS4_TP_ROUTE_PREFORK")) { + for (ggml_tensor * node : inputs.route_prefork_nodes) { + pin_main(node); + } + } + pin_main(inputs.hot_local_lut); + pin_main(inputs.hot_valid_lut); + pin_main(inputs.shard_local_lut); + pin_main(inputs.shard_valid_lut); + // q1 has no repeat. The opt-in exact I32 GPU repeat lets q>1 keep + // both the LUT and its repeated output on the cold owner as well. + if (q == 1 || gpu_i32_repeat) { + pin_peer(inputs.cold_local_lut); + } else { + pin_main(inputs.cold_local_lut); + } + pin_peer(inputs.cold_valid_lut); + for (ggml_tensor * node : inputs.hot_remap_nodes) { + pin_main(node); + } + for (ggml_tensor * node : inputs.cold_remap_nodes) { + pin_peer(node); + } + for (ggml_tensor * node : inputs.shard_remap_nodes) { + pin_main(node); + } + for (ggml_tensor * node : inputs.hot_nodes) { + pin_main(node); + } + for (ggml_tensor * node : inputs.cold_nodes) { + pin_peer(node); + } + for (ggml_tensor * node : inputs.shard_nodes) { + pin_main(node); + } + for (ggml_tensor * node : inputs.join_nodes) { + pin_main(node); + if (targeted_join_split) { + ggml_backend_sched_add_late_cross_input_split_node( + fg.sched, node); + } + } + for (ggml_tensor * node : inputs.deferred_peer_copy_nodes) { + pin_main(node); + ggml_backend_sched_add_deferred_peer_copy_node( + fg.sched, node); + } + pin_main(inputs.main_output); + pin_main(inputs.peer_output); + pin_main(inputs.output); + } + for (const auto & prediction : fg.route_predictions) { + for (ggml_tensor * node : prediction.nodes) pin_main(node); + pin_main(prediction.selected); + } + if (ds4_env_flag("DFLASH_DS4_TP_ROUTE_STATS") || + ds4_env_flag("DFLASH_DS4_TP_DYNAMIC_HOTSET") || + (q > 1 && ds4_tp_route_predict_audit_lookahead() > 0) || + ds4_env_flag("DFLASH_DS4_TP_CACHE_AUDIT")) { + // Preserve selected-ID buffers through the end of the diagnostic + // graph so the post-compute route-reuse audit can read them. The + // native six-route graph records one {IDs, weights} pair/layer; + // the legacy graph records two pairs, hence the stride of two. + for (const MoeHybridGraphInputs & inputs : fg.hybrid_inputs) { + for (size_t i = 0; i < inputs.route_prefork_nodes.size(); i += 2) { + ggml_set_output(inputs.route_prefork_nodes[i]); + } + } + } + pin_main(fg.logits); + pin_main(ex.argmax); + pin_main(ex.capture); + + if (!ggml_backend_sched_alloc_graph(fg.sched, fg.sg.gf)) { + std::fprintf(stderr, + "[ds4-fused-verify] scheduler graph alloc failed\n"); + return false; + } + if (late_join_split || targeted_join_split || report_split_count) { + const char * split_mode = targeted_join_split + ? "targeted-join" + : (late_join_split ? "late-join" : "diagnostic"); + std::fprintf(stderr, + "[ds4-fused-verify] %s scheduler splits=%d q=%d\n", + split_mode, + ggml_backend_sched_get_n_splits(fg.sched), q); + } + + ds4_fused_verify_refresh_hybrid_luts( + w, hybrid, fg, q2_wavefront); + } else { + if (!fg.sg.alloc) { + fg.sg.alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + } + if (!fg.sg.alloc || !ggml_gallocr_alloc_graph(fg.sg.alloc, fg.sg.gf)) { + std::fprintf(stderr, "[ds4-fused-verify] graph alloc failed\n"); + return false; + } } fg.shape_key = std::move(shape_key); return true; @@ -315,7 +1121,9 @@ static int ds4_try_fused_verify_step( std::vector & out_logits, const int32_t * token_ids, Ds4VerifyHooks * hooks, - DeepSeek4StepTelemetry * telemetry) { + DeepSeek4StepTelemetry * telemetry, + MoeHybridStorage * hybrid, + MoeHybridRoutingStats * routing_stats) { if (vc.disabled || mc.disabled) return 0; if (!hooks || !hooks->capture_layer_ids) return 0; if (!hc_out_weights.loaded || hc_out_weights.scale_data.empty() || @@ -330,10 +1138,13 @@ static int ds4_try_fused_verify_step( return 0; } } - if (vc.owner_ctx != w.ctx || vc.backend != backend) { + ggml_backend_t peer_backend = hybrid ? hybrid->cold_backend : nullptr; + if (vc.owner_ctx != w.ctx || vc.backend != backend || + vc.peer_backend != peer_backend) { vc.destroy(); vc.owner_ctx = w.ctx; vc.backend = backend; + vc.peer_backend = peer_backend; } if (mc.owner_ctx != w.ctx || mc.backend != backend) { mc.destroy(); @@ -347,11 +1158,15 @@ static int ds4_try_fused_verify_step( const int q = n_tokens; const int token_pos = kv_start + q - 1; + const bool q2_wavefront = + ds4_tp_q2_wavefront_active(q, kv_start, hybrid); std::vector key; - key.reserve((size_t) w.n_layer + hooks->capture_layer_ids->size() + 4); + key.reserve((size_t) w.n_layer + hooks->capture_layer_ids->size() + 6); key.push_back(w.n_swa); key.push_back(q); key.push_back(token_ids ? 1 : 0); + key.push_back(hybrid ? 1 : 0); + key.push_back(q2_wavefront ? 1 : 0); key.push_back((int64_t) hooks->capture_layer_ids->size()); for (int layer : *hooks->capture_layer_ids) key.push_back(layer); for (int il = 0; il < w.n_layer; ++il) { @@ -368,31 +1183,70 @@ static int ds4_try_fused_verify_step( if (((kv_start + t + 1) % ratio) == 0) { b_idx = t; break; } } } - key.push_back(((int64_t) padded << 4) | (int64_t) (b_idx + 1)); + if (q2_wavefront) { + int padded_a = 0; + if (ratio > 0 && lc.comp_kv) { + const int n_comp_a = ds4_comp_rows_used( + lc.comp_kv, lc.n_comp, ratio, kv_start + 1); + padded_a = ds4_padded_comp_rows( + n_comp_a, (int) lc.comp_kv->ne[1]); + } + key.push_back( + ((int64_t) padded_a << 32) | + ((int64_t) padded << 8) | + (int64_t) (b_idx + 1)); + } else { + key.push_back( + ((int64_t) padded << 4) | (int64_t) (b_idx + 1)); + } } vc.counter++; + const bool capture_cache = &vc == &fused_capture_graph_cache; + const size_t slot_limit = hybrid + ? (capture_cache ? ds4_fused_capture_hybrid_slot_limit() + : ds4_fused_verify_hybrid_slot_limit()) + : vc.slots.size(); + static bool reported_verify_slot_limit = false; + static bool reported_capture_slot_limit = false; + bool & reported_slot_limit = capture_cache + ? reported_capture_slot_limit : reported_verify_slot_limit; + if (hybrid && !reported_slot_limit) { + std::fprintf(stderr, + "[ds4-fused-verify] hybrid %s cache slots=%zu\n", + capture_cache ? "capture" : "verify", slot_limit); + reported_slot_limit = true; + } DeepSeek4FusedDecodeGraph * fg = nullptr; Ds4FusedVerifyCache::Extra * ex = nullptr; - for (size_t i = 0; i < vc.slots.size(); ++i) { + 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 < vc.slots.size(); ++i) { + 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]; + + // A cache miss may select an occupied LRU slot. Tear the previous + // graph down before rebuilding it: its scheduler owns backend scratch + // buffers and HIP graph replay state still refers to those addresses. + // Reinitializing the ggml context over the same metadata arena while + // that state is live can replay kernels against stale allocations. + fg->destroy(); + ex->reset(); + const auto build_t0 = Ds4TimingClock::now(); 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, std::move(key))) { - step_graph_free(fg->sg); - fg->reset_nodes(); + *hooks->capture_layer_ids, hybrid, + std::move(key))) { + fg->destroy(); ex->reset(); vc.disabled = true; return 0; @@ -401,6 +1255,23 @@ static int ds4_try_fused_verify_step( } fg->last_use = vc.counter; + // A fresh request clears KV/HC/compressor storage in place. Kernel + // addresses remain stable, but a whole-step executable also retains its + // captured cross-device dependency state. Never carry that executable + // across the cache-clear boundary; recapture after the normal two q4 + // warmups and continue replaying for the rest of this request. + if (q == 4 && ds4_tp_whole_step_graphs_enabled() && + fg->whole_step_sequence_id != cache.sequence_id) { + if (fg->whole_step_graph_count != 0) { + std::fprintf(stderr, + "[ds4-whole-graph] request boundary old=%llu new=%llu; recapturing\n", + (unsigned long long) fg->whole_step_sequence_id, + (unsigned long long) cache.sequence_id); + } + fg->clear_whole_step_graphs(); + fg->whole_step_sequence_id = cache.sequence_id; + } + // ── Fill inputs ── const auto set_t0 = Ds4TimingClock::now(); ds4_fv_set(fg->inp_embed, embed, sizeof(float) * (size_t) w.n_embd * q); @@ -442,19 +1313,31 @@ static int ds4_try_fused_verify_step( { std::vector maskv((size_t) ggml_nelements(fg->mask_bundle), 0.0f); size_t off = 0; - const int e = kv_start + q; for (int il = 0; il < w.n_layer; ++il) { const int ratio = (int) w.compress_ratios[il]; DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + const int n_lanes = q2_wavefront ? 2 : 1; + for (int lane = 0; lane < n_lanes; ++lane) { + const int lane_q = q2_wavefront ? 2 : q; + const int lane_start = q2_wavefront ? 2 * lane : 0; + const int lane_kv_start = kv_start + lane_start; + const int lane_token_pos = lane_kv_start + lane_q - 1; + // At lane A's attention only its two writes have happened. At + // lane B, e covers all four writes. This makes the physical-ring + // mask and the preserved overwritten-row suffix equivalent to the + // original q4 causal graph. + const int e = lane_kv_start + lane_q; int padded = 0; if (ratio > 0 && lc.comp_kv) { - const int n_comp = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos); + 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]); } - const size_t n_attn = (size_t) w.n_swa + padded + q; - for (int i = 0; i < q; ++i) { + const size_t n_attn = (size_t) w.n_swa + padded + + (lane_q > 1 ? (size_t) lane_q : 0u); + for (int i = 0; i < lane_q; ++i) { float * col = maskv.data() + off + (size_t) i * n_attn; - const int pos_i = kv_start + i; + const int pos_i = lane_kv_start + i; for (int r = 0; r < w.n_swa; ++r) { // position held by slot r AFTER this batch's ring writes const int p_r = (e <= w.n_swa) ? (r < e ? r : -1) @@ -466,13 +1349,21 @@ static int ds4_try_fused_verify_step( for (int c = 0; c < padded; ++c) { if (c >= vis) col[(size_t) w.n_swa + c] = -1.0e30f; } - for (int j = 0; j < q; ++j) { - const bool visible = (j > i) && (kv_start + j >= w.n_swa); - if (!visible) col[(size_t) w.n_swa + padded + j] = -1.0e30f; + if (lane_q > 1) { + for (int j = 0; j < lane_q; ++j) { + const bool visible = + (j > i) && + (lane_kv_start + j >= w.n_swa); + if (!visible) { + col[(size_t) w.n_swa + padded + j] = -1.0e30f; + } + } } } - off += n_attn * q; + off += n_attn * lane_q; + } // lane } + GGML_ASSERT(off == maskv.size()); ds4_fv_set(fg->mask_bundle, maskv.data(), sizeof(float) * maskv.size()); } @@ -481,7 +1372,7 @@ static int ds4_try_fused_verify_step( ggml_tensor * hids = fg->hash_ids[(size_t) il]; if (!hids) continue; const HashRoutingTableCpu & table = hash_tables[(size_t) il]; - const int n_used = ds4_effective_expert_count(w); + const int n_used = (int) hids->ne[0]; hash_scratch.resize((size_t) n_used * q); for (int i = 0; i < q; ++i) { const int32_t * row = hash_routing_row( @@ -504,9 +1395,56 @@ static int ds4_try_fused_verify_step( telemetry->full_graph_set_us += ds4_elapsed_us(set_t0, Ds4TimingClock::now()); } + // A graph can be reused across requests after the fixed hot tensor stacks + // have been repopulated. The weights keep stable addresses; only these + // tiny map inputs must follow the new request-adaptive placement. + if (hybrid && ds4_env_flag("DFLASH_DS4_TP_DYNAMIC_HOTSET")) { + ds4_fused_verify_refresh_hybrid_luts( + w, hybrid, *fg, ds4_tp_q2_wavefront_active(q, kv_start, hybrid)); + } + // ── Compute + read back ── const auto compute_t0 = Ds4TimingClock::now(); - if (ggml_backend_graph_compute(backend, fg->sg.gf) != GGML_STATUS_SUCCESS) { + + // A q=1 feature-capture graph contains HC width-4 matmuls even though it + // advances only one token. On gfx1151, selecting MMVQ for those width-4 + // nodes is correct in eager mode but its HIP graph replay faults against + // unified-memory expert mappings. Keep the proven width-3 ceiling only + // while q=1 capture graphs are evaluated/captured. q>=2 verification uses + // the process ceiling (LUCE_MMVQ_MAX_NCOLS=4 in the optimized launch). + const bool capture_eager = capture_cache && + ds4_env_flag("DFLASH_DS4_TP_CAPTURE_EAGER"); + if (capture_cache) { + ggml_backend_cuda_set_mmvq_max_ncols_override(3); + } + if (capture_eager) { + ggml_backend_cuda_set_graphs_disabled_override(true); + } + // q>=2 verifier graph slots own fixed tensor storage. Position, mask, + // routing-map, and embedding values are refreshed through stable input + // buffers above; the graph topology and all allocation addresses remain + // immutable until the slot is destroyed. The generic CUDA/HIP property + // scan can nevertheless observe backend-mutated metadata and reset replay + // at request boundaries. Mirror the qualified draft fast path behind an + // independent opt-in while retaining normal warmup/capture semantics. + const bool force_verify_graph_replay = q > 1 && + ds4_env_flag("DFLASH_DS4_VERIFY_FORCE_GRAPH_REPLAY"); + if (force_verify_graph_replay) { + ggml_cuda_set_skip_props_check(true); + } + const enum ggml_status compute_status = fg->sched + ? ds4_fused_verify_compute_whole_step(*fg, q) + : ggml_backend_graph_compute(backend, fg->sg.gf); + if (force_verify_graph_replay) { + ggml_cuda_set_skip_props_check(false); + } + if (capture_eager) { + ggml_backend_cuda_set_graphs_disabled_override(false); + } + if (capture_cache) { + ggml_backend_cuda_set_mmvq_max_ncols_override(0); + } + if (compute_status != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[ds4-fused-verify] compute failed\n"); return -1; } @@ -516,15 +1454,25 @@ static int ds4_try_fused_verify_step( const auto read_t0 = Ds4TimingClock::now(); const int ncap = (int) hooks->capture_layer_ids->size(); - if (hooks->all_logits_out) { + const bool argmax_only = + hooks->prefer_argmax_only && hooks->argmax_out && ex->argmax; + if (argmax_only) { + hooks->argmax_out->resize((size_t) q); + ggml_backend_tensor_get(ex->argmax, hooks->argmax_out->data(), 0, + sizeof(int32_t) * (size_t) q); + } else if (hooks->all_logits_out) { hooks->all_logits_out->resize((size_t) w.n_vocab * q); ggml_backend_tensor_get(fg->logits, hooks->all_logits_out->data(), 0, sizeof(float) * (size_t) w.n_vocab * q); } - out_logits.resize((size_t) w.n_vocab); - ggml_backend_tensor_get(fg->logits, out_logits.data(), - (size_t) (q - 1) * (size_t) w.n_vocab * sizeof(float), - sizeof(float) * (size_t) w.n_vocab); + if (!argmax_only) { + out_logits.resize((size_t) w.n_vocab); + ggml_backend_tensor_get(fg->logits, out_logits.data(), + (size_t) (q - 1) * (size_t) w.n_vocab * sizeof(float), + sizeof(float) * (size_t) w.n_vocab); + } else { + out_logits.clear(); + } if (hooks->capture_out && ex->capture && ncap > 0) { std::vector flat((size_t) w.n_embd * ncap * q); ggml_backend_tensor_get(ex->capture, flat.data(), 0, sizeof(float) * flat.size()); diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 3ef4a25c9..8ca8622bb 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -12,6 +12,7 @@ #include "deepseek4_hc_cuda.h" #include "internal.h" #include "../common/step_graph.h" +#include "../common/moe_expert_compute.h" #include "../common/moe_hybrid_ffn_eval.h" #include "../common/moe_hybrid_routing_stats.h" #include "../common/moe_hybrid_stream.h" @@ -20,6 +21,9 @@ #include "ggml.h" #include "ggml-alloc.h" #include "ggml-backend.h" +#include "ggml-cuda.h" + +extern "C" void ggml_backend_cuda_set_graphs_disabled_override(bool disabled); #include #include @@ -32,6 +36,7 @@ #include #include #include +#include #include #if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) @@ -54,11 +59,22 @@ static bool ds4_env_flag(const char * name) { } static int ds4_effective_expert_count(const DeepSeek4Weights & w) { - const int requested = w.routed_expert_top_k; - if (requested > 0 && requested < w.n_expert_used) { - return requested; - } - return w.n_expert_used; + int requested = w.routed_expert_top_k; + if (const char * value = std::getenv("DFLASH_DS4_TOPK")) { + const int env_requested = std::atoi(value); + if (env_requested > 0) requested = env_requested; + } + const int effective = requested > 0 && requested < w.n_expert_used + ? requested + : w.n_expert_used; + static std::atomic logged{false}; + if (effective != w.n_expert_used && !logged.exchange(true)) { + std::fprintf(stderr, + "[deepseek4] effective routed experts=%d " + "(checkpoint=%d; hash and learned layers)\n", + effective, w.n_expert_used); + } + return effective; } static size_t ds4_attn_step_meta_size(int n_tokens) { @@ -139,6 +155,20 @@ static void add_ffn_telemetry(DeepSeek4StepTelemetry * dst, dst->cold_selected += src.cold_selected; } +static void observe_active_routing(MoeHybridRoutingStats * stats, + int layer, + const int32_t * ids, + const float * weights, + int n_ids) { + if (!stats || !ids || !weights || n_ids <= 0) return; + int32_t active[16]; + int n_active = 0; + for (int i = 0; i < n_ids && n_active < 16; ++i) { + if (weights[i] != 0.0f) active[n_active++] = ids[i]; + } + if (n_active > 0) stats->observe(layer, active, n_active); +} + } // namespace int deepseek4_previous_raw_ring_spans( @@ -2030,13 +2060,49 @@ static bool build_cached_decode_hc_post_graph( struct Ds4MoeRouting { ggml_tensor * selected = nullptr; ggml_tensor * weights = nullptr; + std::vector nodes; +}; + +// Diagnostic-only route look-ahead. The prediction is deliberately kept +// separate from Ds4MoeRouting: it never supplies weights to the authoritative +// FFN and therefore cannot change model output. It applies a future layer's +// real gate to an earlier layer's normalized FFN input, matching the +// cross-layer-gate prefetch construction used by Fate. +struct Ds4MoeRoutePrediction { + ggml_tensor * selected = nullptr; + std::vector nodes; }; +static Ds4MoeRoutePrediction build_moe_route_prediction( + ggml_context * ctx, + ggml_tensor * earlier_cur, + const DeepSeek4Weights & w, + const DeepSeek4Layer & future_layer, + int prediction_width) { + Ds4MoeRoutePrediction out; + auto track = [&](ggml_tensor * tensor) { + if (tensor) out.nodes.push_back(tensor); + return tensor; + }; + ggml_tensor * logits = track( + ggml_mul_mat(ctx, future_layer.ffn_gate_inp, earlier_cur)); + ggml_tensor * selection = track( + ggml_sqrt(ctx, track(ggml_softplus(ctx, logits)))); + if (future_layer.ffn_exp_probs_b) { + selection = track( + ggml_add(ctx, selection, future_layer.ffn_exp_probs_b)); + } + const int width = std::max( + 1, std::min(prediction_width, w.n_expert)); + out.selected = track(ggml_top_k(ctx, selection, width)); + return out; +} + static MoeHybridConfig make_ds4_moe_hybrid_config(const DeepSeek4Weights & w) { MoeHybridConfig cfg; cfg.n_embd = w.n_embd; cfg.n_expert = w.n_expert; - cfg.n_expert_used = w.n_expert_used; + cfg.n_expert_used = ds4_effective_expert_count(w); cfg.n_ff_exp = w.n_ff_exp; cfg.n_ff_shexp = w.n_ff_exp; cfg.n_layer = w.n_layer; @@ -2087,9 +2153,12 @@ static bool eval_ds4_hybrid( std::vector & ffn_out_host, ggml_gallocr_t * hot_alloc, ggml_gallocr_t * cold_alloc, + MoeExpertCompute * expert_compute, + const MoeExpertLayer * expert_layer, DeepSeek4StepTelemetry * step_tel) { const auto ffn_t0 = Ds4TimingClock::now(); - if (!storage.down_cold && !storage.gate_up_cold) { + if (!storage.down_cold && !storage.gate_up_cold && + !(expert_compute && expert_layer)) { if (!hybrid_owner || !stream_engine || !stream_engine->is_ready() || !hybrid_owner->has_mmap() || layer < 0 || layer >= (int) hybrid_owner->layer_regions.size()) { @@ -2175,7 +2244,7 @@ static bool eval_ds4_hybrid( backend, cpu_backend, hybrid_cfg, desc, storage, ffn_normed_host, selected_host, weights_host, n_tokens, ffn_out_host, nullptr, hot_alloc, cold_alloc, - nullptr, nullptr, + expert_compute, expert_layer, step_tel ? &ffn_tel : nullptr); if (ffn_ok) { if (step_tel) { @@ -2185,6 +2254,13 @@ static bool eval_ds4_hybrid( return true; } + if (expert_compute && expert_layer) { + std::fprintf(stderr, + "[deepseek4-moe-tp] remote expert evaluation failed at layer %d\n", + layer); + return false; + } + ffn_out_host.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); std::vector single_out; for (int ti = 0; ti < n_tokens; ++ti) { @@ -2213,28 +2289,33 @@ static Ds4MoeRouting build_moe_routing( const DeepSeek4Layer & L, int n_tokens) { Ds4MoeRouting out; - ggml_tensor * logits = ggml_mul_mat(ctx, L.ffn_gate_inp, cur); + auto track = [&](ggml_tensor * tensor) { + if (tensor) out.nodes.push_back(tensor); + return tensor; + }; + ggml_tensor * logits = track(ggml_mul_mat(ctx, L.ffn_gate_inp, cur)); // DS4 routes with sqrt(softplus(logit)). Optional bias affects only the // top-k expert selection, while expert weights come from the unbiased // router probabilities and are normalized after selection. - ggml_tensor * probs = ggml_sqrt(ctx, ggml_softplus(ctx, logits)); + ggml_tensor * softplus = track(ggml_softplus(ctx, logits)); + ggml_tensor * probs = track(ggml_sqrt(ctx, softplus)); ggml_tensor * selection = probs; if (L.ffn_exp_probs_b) { - selection = ggml_add(ctx, selection, L.ffn_exp_probs_b); + selection = track(ggml_add(ctx, selection, L.ffn_exp_probs_b)); } const int k_used = ds4_effective_expert_count(w); - out.selected = ggml_top_k(ctx, selection, k_used); + out.selected = track(ggml_top_k(ctx, selection, k_used)); ggml_tensor * probs_3d = ggml_reshape_3d(ctx, probs, 1, w.n_expert, n_tokens); - out.weights = ggml_get_rows(ctx, probs_3d, out.selected); + out.weights = track(ggml_get_rows(ctx, probs_3d, out.selected)); out.weights = ggml_reshape_2d(ctx, out.weights, k_used, n_tokens); - ggml_tensor * w_sum = ggml_sum_rows(ctx, out.weights); - w_sum = ggml_clamp(ctx, w_sum, 6.103515625e-5f, INFINITY); - out.weights = ggml_div(ctx, out.weights, w_sum); + ggml_tensor * w_sum = track(ggml_sum_rows(ctx, out.weights)); + w_sum = track(ggml_clamp(ctx, w_sum, 6.103515625e-5f, INFINITY)); + out.weights = track(ggml_div(ctx, out.weights, w_sum)); if (w.expert_weight_scale != 1.0f) { - out.weights = ggml_scale(ctx, out.weights, w.expert_weight_scale); + out.weights = track(ggml_scale(ctx, out.weights, w.expert_weight_scale)); } return out; } @@ -3073,7 +3154,8 @@ static bool deepseek4_step_hybrid( const int32_t * token_ids, MoeHybridStreamEngine * stream_engine, DeepSeek4StepTelemetry * telemetry, - MoeHybridRoutingStats * routing_stats) { + MoeHybridRoutingStats * routing_stats, + MoeExpertComputeRuntime * expert_runtime) { const auto step_t0 = Ds4TimingClock::now(); const int n_embd = w.n_embd; const int n_hc = w.n_hc; @@ -3282,9 +3364,9 @@ static bool deepseek4_step_hybrid( return false; } - const int n_used = ds4_effective_expert_count(w); - std::vector selected_host((size_t)n_used * (size_t)n_tokens); - std::vector weights_host((size_t)n_used * (size_t)n_tokens); + const int route_width = ds4_effective_expert_count(w); + std::vector selected_host((size_t)route_width * (size_t)n_tokens); + std::vector weights_host((size_t)route_width * (size_t)n_tokens); const auto route_select_t0 = Ds4TimingClock::now(); for (int ti = 0; ti < n_tokens; ++ti) { const int32_t tok = token_ids[ti]; @@ -3297,39 +3379,45 @@ static bool deepseek4_step_hybrid( return false; } float sum = 0.0f; - for (int ei = 0; ei < n_used; ++ei) { + for (int ei = 0; ei < route_width; ++ei) { const int32_t expert = row[ei]; - selected_host[(size_t)ti * (size_t)n_used + (size_t)ei] = expert; + selected_host[(size_t)ti * (size_t)route_width + (size_t)ei] = expert; float prob = 0.0f; if (expert >= 0 && expert < w.n_expert) { prob = probs_host[(size_t)ti * (size_t)w.n_expert + (size_t)expert]; } - weights_host[(size_t)ti * (size_t)n_used + (size_t)ei] = prob; + weights_host[(size_t)ti * (size_t)route_width + (size_t)ei] = prob; sum += prob; } sum = std::max(sum, 6.103515625e-5f); - for (int ei = 0; ei < n_used; ++ei) { - float & weight = weights_host[(size_t)ti * (size_t)n_used + (size_t)ei]; + for (int ei = 0; ei < route_width; ++ei) { + float & weight = weights_host[(size_t)ti * (size_t)route_width + (size_t)ei]; weight = weight / sum * w.expert_weight_scale; } } if (telemetry) telemetry->route_select_us += ds4_elapsed_us(route_select_t0, Ds4TimingClock::now()); if (routing_stats) { for (int ti = 0; ti < n_tokens; ++ti) { - routing_stats->observe(il, - selected_host.data() + (size_t)ti * (size_t)n_used, - n_used); + observe_active_routing(routing_stats, il, + selected_host.data() + (size_t)ti * (size_t)route_width, + weights_host.data() + (size_t)ti * (size_t)route_width, + route_width); } } MoeHybridConfig hybrid_cfg = make_ds4_moe_hybrid_config(w); MoeLayerDesc desc = make_ds4_moe_layer_desc(L); auto & storage = moe_hybrid.layers[(size_t) il]; + MoeExpertCompute * expert_compute = + expert_runtime ? expert_runtime->compute_ptr() : nullptr; + const MoeExpertLayer * expert_layer = + expert_runtime ? expert_runtime->layer_ptr((size_t)il) : nullptr; if (!eval_ds4_hybrid( backend, cpu_backend, hybrid_cfg, desc, &moe_hybrid, storage, stream_engine, - il, n_embd, n_used, + il, n_embd, route_width, ffn_normed_host.data(), selected_host.data(), weights_host.data(), - n_tokens, ffn_out_host, &hot_alloc, &cold_alloc, telemetry)) { + n_tokens, ffn_out_host, &hot_alloc, &cold_alloc, + expert_compute, expert_layer, telemetry)) { if (hot_alloc) ggml_gallocr_free(hot_alloc); if (cold_alloc) ggml_gallocr_free(cold_alloc); return false; @@ -3376,9 +3464,9 @@ static bool deepseek4_step_hybrid( std::vector ffn_normed_host((size_t)n_embd * (size_t)n_tokens); std::vector probs_host((size_t)w.n_expert * (size_t)n_tokens); - const int n_used = ds4_effective_expert_count(w); - std::vector selected_host((size_t)n_used * (size_t)n_tokens); - std::vector weights_host((size_t)n_used * (size_t)n_tokens); + const int route_width = ds4_effective_expert_count(w); + std::vector selected_host((size_t)route_width * (size_t)n_tokens); + std::vector weights_host((size_t)route_width * (size_t)n_tokens); const auto route_read_t0 = Ds4TimingClock::now(); ggml_backend_tensor_get(ffn_normed, ffn_normed_host.data(), 0, sizeof(float) * ffn_normed_host.size()); ggml_backend_tensor_get(router_probs, probs_host.data(), 0, sizeof(float) * probs_host.size()); @@ -3395,18 +3483,18 @@ static bool deepseek4_step_hybrid( } for (int ti = 0; ti < n_tokens; ++ti) { const float * probs = probs_host.data() + (size_t)ti * (size_t)w.n_expert; - std::vector top((size_t)n_used, -1); + std::vector top((size_t)route_width, -1); for (int expert = 0; expert < w.n_expert; ++expert) { const float score = probs[expert] + (!bias_host.empty() ? bias_host[(size_t)expert] : 0.0f); - for (int slot = 0; slot < n_used; ++slot) { + for (int slot = 0; slot < route_width; ++slot) { const int32_t cur_expert = top[(size_t)slot]; const float cur_score = cur_expert >= 0 ? probs[cur_expert] + (!bias_host.empty() ? bias_host[(size_t)cur_expert] : 0.0f) : -INFINITY; if (cur_expert < 0 || score > cur_score) { - for (int m = n_used - 1; m > slot; --m) { + for (int m = route_width - 1; m > slot; --m) { top[(size_t)m] = top[(size_t)m - 1]; } top[(size_t)slot] = expert; @@ -3415,36 +3503,43 @@ static bool deepseek4_step_hybrid( } } float sum = 0.0f; - for (int slot = 0; slot < n_used; ++slot) { + for (int slot = 0; slot < route_width; ++slot) { const int32_t expert = top[(size_t)slot]; - selected_host[(size_t)ti * (size_t)n_used + (size_t)slot] = expert; + selected_host[(size_t)ti * (size_t)route_width + (size_t)slot] = expert; const float weight = expert >= 0 ? probs[expert] : 0.0f; - weights_host[(size_t)ti * (size_t)n_used + (size_t)slot] = weight; + weights_host[(size_t)ti * (size_t)route_width + (size_t)slot] = weight; sum += weight; } sum = std::max(sum, 6.103515625e-5f); - for (int slot = 0; slot < n_used; ++slot) { - float & weight = weights_host[(size_t)ti * (size_t)n_used + (size_t)slot]; + for (int slot = 0; slot < route_width; ++slot) { + float & weight = weights_host[(size_t)ti * (size_t)route_width + (size_t)slot]; weight = weight / sum * w.expert_weight_scale; } } if (telemetry) telemetry->route_select_us += ds4_elapsed_us(route_select_t0, Ds4TimingClock::now()); if (routing_stats) { for (int ti = 0; ti < n_tokens; ++ti) { - routing_stats->observe(il, - selected_host.data() + (size_t)ti * (size_t)n_used, - n_used); + observe_active_routing(routing_stats, il, + selected_host.data() + (size_t)ti * (size_t)route_width, + weights_host.data() + (size_t)ti * (size_t)route_width, + route_width); } } MoeHybridConfig hybrid_cfg = make_ds4_moe_hybrid_config(w); + hybrid_cfg.n_expert_used = route_width; MoeLayerDesc desc = make_ds4_moe_layer_desc(L); auto & storage = moe_hybrid.layers[(size_t) il]; + MoeExpertCompute * expert_compute = + expert_runtime ? expert_runtime->compute_ptr() : nullptr; + const MoeExpertLayer * expert_layer = + expert_runtime ? expert_runtime->layer_ptr((size_t)il) : nullptr; if (!eval_ds4_hybrid( backend, cpu_backend, hybrid_cfg, desc, &moe_hybrid, storage, stream_engine, - il, n_embd, n_used, + il, n_embd, route_width, ffn_normed_host.data(), selected_host.data(), weights_host.data(), - n_tokens, ffn_out_host, &hot_alloc, &cold_alloc, telemetry)) { + n_tokens, ffn_out_host, &hot_alloc, &cold_alloc, + expert_compute, expert_layer, telemetry)) { if (hot_alloc) ggml_gallocr_free(hot_alloc); if (cold_alloc) ggml_gallocr_free(cold_alloc); return false; @@ -3548,18 +3643,21 @@ bool deepseek4_step( MoeHybridStreamEngine * stream_engine, DeepSeek4StepTelemetry * telemetry, MoeHybridRoutingStats * routing_stats, - Ds4VerifyHooks * verify_hooks) { + Ds4VerifyHooks * verify_hooks, + MoeExpertComputeRuntime * expert_runtime) { if (w.moe_hybrid && moe_hybrid != nullptr) { return deepseek4_step_hybrid(backend, w, cache, *moe_hybrid, embed, n_tokens, kv_start, out_logits, - token_ids, stream_engine, telemetry, routing_stats); + token_ids, stream_engine, telemetry, routing_stats, + expert_runtime); } std::vector hc_state; return deepseek4_step_layer_range( backend, w, cache, hc_state, embed, n_tokens, kv_start, 0, w.n_layer, &out_logits, token_ids, telemetry, - /*allow_decode_graph_reuse=*/verify_hooks == nullptr, verify_hooks); + /*allow_decode_graph_reuse=*/verify_hooks == nullptr, verify_hooks, + /*moe_hybrid=*/nullptr, expert_runtime, routing_stats); } // ─── Fused single-graph decode (n_tokens == 1) ────────────────────────── @@ -3577,6 +3675,15 @@ bool deepseek4_step( // on, enabling graph replay for the bulk of decode steps. struct DeepSeek4FusedDecodeGraph { + struct RoutePredictionOutput { + int source_layer = -1; + int target_layer = -1; + int lane_start = 0; + int n_tokens = 0; + int width = 0; + ggml_tensor * selected = nullptr; + std::vector nodes; + }; std::vector shape_key; uint64_t last_use = 0; StepGraph sg; @@ -3585,7 +3692,27 @@ struct DeepSeek4FusedDecodeGraph { ggml_tensor * i64_bundle = nullptr; ggml_tensor * mask_bundle = nullptr; // additive score mask (0 / -1e30), may be null std::vector hash_ids; + std::vector hybrid_inputs; + std::vector route_predictions; ggml_tensor * logits = nullptr; + ggml_backend_sched_t sched = nullptr; + uint32_t whole_step_warmups = 0; + uint64_t whole_step_sequence_id = 0; + int whole_step_graph_count = 0; + std::array whole_step_backends{}; + std::array whole_step_graphs{}; + + void clear_whole_step_graphs() { + for (int i = 0; i < whole_step_graph_count; ++i) { + ggml_backend_cuda_whole_graph_free( + whole_step_backends[(size_t) i], + whole_step_graphs[(size_t) i]); + } + whole_step_graphs = {}; + whole_step_backends = {}; + whole_step_graph_count = 0; + whole_step_warmups = 0; + } void reset_nodes() { inp_embed = nullptr; @@ -3594,6 +3721,9 @@ struct DeepSeek4FusedDecodeGraph { mask_bundle = nullptr; logits = nullptr; hash_ids.clear(); + hybrid_inputs.clear(); + route_predictions.clear(); + whole_step_sequence_id = 0; shape_key.clear(); last_use = 0; } @@ -3603,6 +3733,11 @@ struct DeepSeek4FusedDecodeGraph { } void destroy() { + clear_whole_step_graphs(); + if (sched) { + ggml_backend_sched_free(sched); + sched = nullptr; + } step_graph_destroy(sg); reset_nodes(); } @@ -3679,6 +3814,7 @@ static ggml_tensor * ds4_build_hash_routed_ffn( ggml_tensor * probs = ggml_sqrt(ctx, ggml_softplus(ctx, logits)); const int n_used = (int) hash_ids->ne[0]; + GGML_ASSERT(n_used > 0 && n_used <= w.n_expert_used); const int n_ff_exp = w.n_ff_exp; ggml_tensor * cur_3d = ggml_reshape_3d(ctx, ffn_normed, w.n_embd, 1, 1); ggml_tensor * gate_e = ggml_mul_mat_id(ctx, L.ffn_gate_exps, cur_3d, hash_ids); @@ -4150,6 +4286,170 @@ static int ds4_try_fused_decode_step( return 1; } +static bool eval_ds4_layer_range_hybrid_ffn( + ggml_backend_t backend, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + int layer, + int n_tokens, + const float * ffn_working_host, + const ggml_tensor * ffn_in_backend, + const int32_t * token_ids, + const HashRoutingTableCpu & hash_table, + MoeHybridStorage & hybrid, + MoeExpertComputeRuntime * expert_runtime, + MoeHybridRoutingStats * routing_stats, + std::vector & out, + DeepSeek4StepTelemetry * telemetry) { + const int n_embd = w.n_embd; + const bool hash_routed = + layer < w.n_hash_layer && L.ffn_gate_tid2eid && + token_ids && hash_table.loaded; + const int route_width = ds4_effective_expert_count(w); + + const auto route_build_t0 = Ds4TimingClock::now(); + ggml_init_params params{}; + params.mem_size = 16 * 1024 * 1024; + params.mem_buffer = nullptr; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) return false; + ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); + ggml_set_input(inp); + ggml_tensor * normed = build_rms_norm(ctx, inp, L.ffn_norm, w.rms_eps); + ggml_tensor * logits = ggml_mul_mat(ctx, L.ffn_gate_inp, normed); + ggml_tensor * probs = ggml_sqrt(ctx, ggml_softplus(ctx, logits)); + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, normed); + ggml_build_forward_expand(gf, probs); + ggml_gallocr_t alloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!alloc || !ggml_gallocr_alloc_graph(alloc, gf)) { + if (alloc) ggml_gallocr_free(alloc); + ggml_free(ctx); + return false; + } + if (telemetry) { + telemetry->route_build_us += + ds4_elapsed_us(route_build_t0, Ds4TimingClock::now()); + } + + if (ffn_in_backend) { + ggml_backend_tensor_copy( + const_cast(ffn_in_backend), inp); + } else { + ggml_backend_tensor_set(inp, ffn_working_host, 0, + sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + } + const auto route_compute_t0 = Ds4TimingClock::now(); + const bool route_ok = + ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + if (telemetry) { + telemetry->route_compute_us += + ds4_elapsed_us(route_compute_t0, Ds4TimingClock::now()); + } + + std::vector normed_host((size_t)n_embd * (size_t)n_tokens); + std::vector probs_host((size_t)w.n_expert * (size_t)n_tokens); + if (route_ok) { + const auto route_read_t0 = Ds4TimingClock::now(); + ggml_backend_tensor_get(normed, normed_host.data(), 0, + sizeof(float) * normed_host.size()); + ggml_backend_tensor_get(probs, probs_host.data(), 0, + sizeof(float) * probs_host.size()); + if (telemetry) { + telemetry->route_read_us += + ds4_elapsed_us(route_read_t0, Ds4TimingClock::now()); + } + } + ggml_gallocr_free(alloc); + ggml_free(ctx); + if (!route_ok) return false; + + std::vector selected((size_t)route_width * (size_t)n_tokens); + std::vector weights((size_t)route_width * (size_t)n_tokens); + std::vector bias; + if (!hash_routed && L.ffn_exp_probs_b) { + bias.resize((size_t)w.n_expert); + ggml_backend_tensor_get(L.ffn_exp_probs_b, bias.data(), 0, + sizeof(float) * bias.size()); + } + + const auto route_select_t0 = Ds4TimingClock::now(); + for (int t = 0; t < n_tokens; ++t) { + const float * token_probs = + probs_host.data() + (size_t)t * (size_t)w.n_expert; + int32_t * token_ids_out = + selected.data() + (size_t)t * (size_t)route_width; + float * token_weights = + weights.data() + (size_t)t * (size_t)route_width; + + if (hash_routed) { + const int32_t tok = token_ids[t]; + if (tok < 0 || tok >= w.n_vocab) return false; + const int32_t * row = + hash_table.ids.data() + + (size_t)tok * (size_t)w.n_expert_used; + std::memcpy(token_ids_out, row, + sizeof(int32_t) * (size_t)route_width); + } else { + std::fill(token_ids_out, token_ids_out + route_width, -1); + for (int expert = 0; expert < w.n_expert; ++expert) { + const float score = token_probs[expert] + + (!bias.empty() ? bias[(size_t)expert] : 0.0f); + for (int slot = 0; slot < route_width; ++slot) { + const int32_t current = token_ids_out[slot]; + const float current_score = current >= 0 + ? token_probs[current] + + (!bias.empty() ? bias[(size_t)current] : 0.0f) + : -INFINITY; + if (current < 0 || score > current_score) { + for (int move = route_width - 1; move > slot; --move) { + token_ids_out[move] = token_ids_out[move - 1]; + } + token_ids_out[slot] = expert; + break; + } + } + } + } + + float sum = 0.0f; + for (int slot = 0; slot < route_width; ++slot) { + const int32_t expert = token_ids_out[slot]; + token_weights[slot] = + expert >= 0 && expert < w.n_expert ? token_probs[expert] : 0.0f; + sum += token_weights[slot]; + } + sum = std::max(sum, 6.103515625e-5f); + for (int slot = 0; slot < route_width; ++slot) { + token_weights[slot] = + token_weights[slot] / sum * w.expert_weight_scale; + } + observe_active_routing(routing_stats, layer, + token_ids_out, token_weights, route_width); + } + if (telemetry) { + telemetry->route_select_us += + ds4_elapsed_us(route_select_t0, Ds4TimingClock::now()); + } + + MoeHybridConfig cfg = make_ds4_moe_hybrid_config(w); + cfg.n_expert_used = route_width; + MoeLayerDesc desc = make_ds4_moe_layer_desc(L); + MoeExpertCompute * expert_compute = + expert_runtime ? expert_runtime->compute_ptr() : nullptr; + const MoeExpertLayer * expert_layer = + expert_runtime ? expert_runtime->layer_ptr((size_t)layer) : nullptr; + return eval_ds4_hybrid( + backend, hybrid.cpu_backend, cfg, desc, &hybrid, + hybrid.layers[(size_t)layer], nullptr, + layer, n_embd, route_width, + normed_host.data(), selected.data(), weights.data(), + n_tokens, out, nullptr, nullptr, + expert_compute, expert_layer, telemetry); +} + bool deepseek4_step_layer_range( ggml_backend_t backend, @@ -4165,7 +4465,10 @@ bool deepseek4_step_layer_range( const int32_t * token_ids, DeepSeek4StepTelemetry * telemetry, bool allow_decode_graph_reuse, - Ds4VerifyHooks * verify_hooks) { + Ds4VerifyHooks * verify_hooks, + MoeHybridStorage * moe_hybrid, + MoeExpertComputeRuntime * expert_runtime, + MoeHybridRoutingStats * routing_stats) { const auto step_t0 = Ds4TimingClock::now(); // ── Partial layer-range forward with HC ───────────────────────────── @@ -4221,7 +4524,8 @@ bool deepseek4_step_layer_range( chunk, kv_start + off, layer_begin, layer_end, out_logits ? &chunk_out : nullptr, token_ids ? token_ids + off : nullptr, - telemetry, allow_decode_graph_reuse, chunk_hooks_ptr)) { + telemetry, allow_decode_graph_reuse, chunk_hooks_ptr, + moe_hybrid, expert_runtime, routing_stats)) { return false; } hc_all.insert(hc_all.end(), chunk_hc.begin(), chunk_hc.end()); @@ -4338,14 +4642,45 @@ bool deepseek4_step_layer_range( const int n_expert_used = ds4_effective_expert_count(w); scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, n_expert_used); - if (n_tokens >= 2 && n_tokens <= 4 && verify_hooks && layer_begin == 0 && is_last_shard && - out_logits && ds4_backend_is_gpu(backend) && w.fused_decode && - ds4_fused_verify_enabled()) { + const bool fused_hybrid_ready = + moe_hybrid && !expert_runtime && + moe_hybrid->materialized_cold_experts && + moe_hybrid->cold_backend_kind == MoeHybridColdBackend::Gpu && + moe_hybrid->cold_backend && moe_hybrid->cold_backend != backend; + // The batched verifier graph is also the only whole-model graph that can + // currently own tensors on both GPU backends. Reuse it for q=1 hybrid + // decode so native AR does not fall back to 43 host-synchronized FFN + // calls. DSpark prefill also needs this q=1 path: feature-capture hooks + // must not silently force it back to the numerically different per-layer + // implementation, otherwise the first speculative seed can diverge from + // native AR before the drafter has run at all. + const bool fused_hybrid_decode = + fused_hybrid_ready && n_tokens == 1 && allow_decode_graph_reuse && + ds4_env_flag("DFLASH_DS4_FUSED_HYBRID_DECODE"); + std::vector fused_hybrid_decode_capture_ids; + Ds4VerifyHooks fused_hybrid_decode_hooks; + if (fused_hybrid_decode && !verify_hooks) { + fused_hybrid_decode_hooks.capture_layer_ids = + &fused_hybrid_decode_capture_ids; + } + Ds4VerifyHooks * fused_graph_hooks = + (fused_hybrid_decode && !verify_hooks) + ? &fused_hybrid_decode_hooks : verify_hooks; + if ((!moe_hybrid || fused_hybrid_ready) && + ((n_tokens >= 2 && n_tokens <= 4 && verify_hooks) || + fused_hybrid_decode) && + layer_begin == 0 && is_last_shard && + out_logits && ds4_backend_is_gpu(backend) && ds4_fused_verify_enabled()) { + const bool q1_feature_capture = + n_tokens == 1 && verify_hooks && verify_hooks->capture_out; + Ds4FusedVerifyCache & graph_cache = q1_feature_capture + ? fused_capture_graph_cache : fused_verify_graph_cache; const int vrc = ds4_try_fused_verify_step( - fused_verify_graph_cache, fused_decode_graph_cache, backend, w, cache, + graph_cache, fused_decode_graph_cache, backend, w, cache, hc_layer_weights_range, hc_output_weights_range, hash_routing_tables_range, scratch.hash_expert_ids, embed, n_tokens, kv_start, *out_logits, token_ids, - verify_hooks, telemetry); + fused_graph_hooks, telemetry, fused_hybrid_ready ? moe_hybrid : nullptr, + routing_stats); if (vrc < 0) return false; if (vrc > 0) { const int np = kv_start + n_tokens; @@ -4361,8 +4696,10 @@ bool deepseek4_step_layer_range( } } std::vector fused_debug_logits; - if (n_tokens == 1 && allow_decode_graph_reuse && layer_begin == 0 && is_last_shard && - out_logits && ds4_backend_is_gpu(backend) && w.fused_decode) { + if (!moe_hybrid && n_tokens == 1 && allow_decode_graph_reuse && layer_begin == 0 && is_last_shard && + !(verify_hooks && verify_hooks->capture_layer_ids && + verify_hooks->capture_out) && + out_logits && ds4_backend_is_gpu(backend) && ds4_fused_decode_enabled()) { const int rc = ds4_try_fused_decode_step( fused_decode_graph_cache, backend, w, cache, hc_layer_weights_range, hc_output_weights_range, hash_routing_tables_range, scratch.hash_expert_ids, @@ -4788,78 +5125,108 @@ bool deepseek4_step_layer_range( hash_routed = il < w.n_hash_layer && L.ffn_gate_tid2eid && token_ids && hash_routing_tables_range[(size_t)il].loaded; - if (hash_routed) { - const int n_used = n_expert_used; - hash_expert_ids_host.resize((size_t)n_used * (size_t)n_tokens); - for (int ti = 0; ti < n_tokens; ti++) { - const int32_t tok = token_ids[ti]; - const int32_t * row = hash_routing_row( - hash_routing_tables_range[(size_t)il], tok, w.n_expert_used); - if (!row) { - std::fprintf(stderr, - "[deepseek4] token id %d outside hash table for layer %d\n", - tok, il); - return false; + ggml_tensor * ffn_out = nullptr; + if (moe_hybrid) { + if (!eval_ds4_layer_range_hybrid_ffn( + backend, w, L, il, n_tokens, + ffn_working.data(), ffn_in_backend, + token_ids, hash_routing_tables_range[(size_t)il], + *moe_hybrid, expert_runtime, routing_stats, + ffn_out_host, telemetry)) { + std::fprintf(stderr, + "[deepseek4-moe-tp] layer-range FFN failed layer %d\n", + il); + return false; + } + } else { + if (hash_routed) { + const int n_used = n_expert_used; + hash_expert_ids_host.resize((size_t)n_used * (size_t)n_tokens); + for (int ti = 0; ti < n_tokens; ti++) { + const int32_t tok = token_ids[ti]; + const int32_t * row = hash_routing_row( + hash_routing_tables_range[(size_t)il], tok, + w.n_expert_used); + if (!row) { + std::fprintf(stderr, + "[deepseek4] token id %d outside hash table for layer %d\n", + tok, il); + return false; + } + memcpy(hash_expert_ids_host.data() + (size_t)ti * n_used, + row, (size_t)n_used * sizeof(int32_t)); } - memcpy(hash_expert_ids_host.data() + (size_t)ti * n_used, - row, (size_t)n_used * sizeof(int32_t)); } - } - auto & cached = cached_decode_ffn_graphs[(size_t)il]; - if (!cached.valid() || - cached.owner_ctx != w.ctx || - cached.backend != backend || - cached.layer_idx != il || - cached.n_tokens != n_tokens || - cached.n_expert_used != n_expert_used || - cached.hash_routed != hash_routed) { - const auto ffn_build_t0 = Ds4TimingClock::now(); - if (!build_cached_decode_ffn_graph(cached, backend, w, L, il, n_tokens, hash_routed)) { - std::fprintf(stderr, "[deepseek4] cached ffn graph alloc failed layer %d\n", il); - return false; + auto & cached = cached_decode_ffn_graphs[(size_t)il]; + if (!cached.valid() || + cached.owner_ctx != w.ctx || + cached.backend != backend || + cached.layer_idx != il || + cached.n_tokens != n_tokens || + cached.n_expert_used != n_expert_used || + cached.hash_routed != hash_routed) { + const auto ffn_build_t0 = Ds4TimingClock::now(); + if (!build_cached_decode_ffn_graph(cached, backend, w, L, il, n_tokens, hash_routed)) { + std::fprintf(stderr, "[deepseek4] cached ffn graph alloc failed layer %d\n", il); + return false; + } + if (telemetry) telemetry->ffn_build_us += ds4_elapsed_us(ffn_build_t0, Ds4TimingClock::now()); } - if (telemetry) telemetry->ffn_build_us += ds4_elapsed_us(ffn_build_t0, Ds4TimingClock::now()); - } - ggml_tensor * ffn_out = cached.sg.hidden_states; - if (ffn_in_backend) { - ggml_backend_tensor_copy(ffn_in_backend, cached.sg.inp_embed); - } else { - ggml_backend_tensor_set(cached.sg.inp_embed, ffn_working.data(), 0, - sizeof(float) * ffn_working.size()); - } - if (cached.hash_ids) { - ggml_backend_tensor_set(cached.hash_ids, hash_expert_ids_host.data(), 0, - sizeof(int32_t) * hash_expert_ids_host.size()); - } + ffn_out = cached.sg.hidden_states; + if (ffn_in_backend) { + ggml_backend_tensor_copy(ffn_in_backend, cached.sg.inp_embed); + } else { + ggml_backend_tensor_set(cached.sg.inp_embed, ffn_working.data(), 0, + sizeof(float) * ffn_working.size()); + } + if (cached.hash_ids) { + ggml_backend_tensor_set(cached.hash_ids, hash_expert_ids_host.data(), 0, + sizeof(int32_t) * hash_expert_ids_host.size()); + } - const auto ffn_compute_t0 = Ds4TimingClock::now(); - auto status = ggml_backend_graph_compute(backend, cached.sg.gf); - if (telemetry) telemetry->ffn_compute_us += ds4_elapsed_us(ffn_compute_t0, Ds4TimingClock::now()); - if (status != GGML_STATUS_SUCCESS) { - std::fprintf(stderr, "[deepseek4] cached ffn compute failed layer %d\n", il); - return false; + const auto ffn_compute_t0 = Ds4TimingClock::now(); + auto status = ggml_backend_graph_compute(backend, cached.sg.gf); + if (telemetry) telemetry->ffn_compute_us += ds4_elapsed_us(ffn_compute_t0, Ds4TimingClock::now()); + if (status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "[deepseek4] cached ffn compute failed layer %d\n", il); + return false; + } + if (!(use_backend_decode_hc_graph || use_backend_decode_hc_direct)) { + const auto ffn_read_t0 = Ds4TimingClock::now(); + ggml_backend_tensor_get(ffn_out, ffn_out_host.data(), 0, + sizeof(float) * ffn_out_host.size()); + if (telemetry) telemetry->ffn_read_us += + ds4_elapsed_us(ffn_read_t0, Ds4TimingClock::now()); + } } if (use_backend_decode_hc_graph || use_backend_decode_hc_direct) { if (hc_state_backend != cached_decode_hc_post_graph.residual_hc) { - ggml_backend_tensor_copy(hc_state_backend, cached_decode_hc_post_graph.residual_hc); + ggml_backend_tensor_copy(hc_state_backend, + cached_decode_hc_post_graph.residual_hc); + } + if (moe_hybrid) { + ggml_backend_tensor_set(cached_decode_hc_post_graph.block_out, + ffn_out_host.data(), 0, + sizeof(float) * ffn_out_host.size()); + } else { + ggml_backend_tensor_copy(ffn_out, + cached_decode_hc_post_graph.block_out); } - ggml_backend_tensor_copy(ffn_out, cached_decode_hc_post_graph.block_out); - ggml_backend_tensor_copy(ffn_post_backend, cached_decode_hc_post_graph.post); - ggml_backend_tensor_copy(ffn_comb_backend, cached_decode_hc_post_graph.comb); + ggml_backend_tensor_copy(ffn_post_backend, + cached_decode_hc_post_graph.post); + ggml_backend_tensor_copy(ffn_comb_backend, + cached_decode_hc_post_graph.comb); const auto hc_post_ffn_t0 = Ds4TimingClock::now(); if (ggml_backend_graph_compute(backend, cached_decode_hc_post_graph.sg.gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[deepseek4] cached hc-post compute failed layer %d ffn\n", il); return false; } hc_state_backend = cached_decode_hc_post_graph.sg.hidden_states; - if (telemetry) telemetry->hc_post_ffn_us += ds4_elapsed_us(hc_post_ffn_t0, Ds4TimingClock::now()); - } else { - const auto ffn_read_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_get(ffn_out, ffn_out_host.data(), 0, sizeof(float) * ffn_out_host.size()); - if (telemetry) telemetry->ffn_read_us += ds4_elapsed_us(ffn_read_t0, Ds4TimingClock::now()); + if (telemetry) telemetry->hc_post_ffn_us += + ds4_elapsed_us(hc_post_ffn_t0, Ds4TimingClock::now()); } // ── HC post (FFN) ─────────────────────────────────────── @@ -5334,12 +5701,14 @@ static ggml_tensor * build_dspark_attention( ggml_context * ctx, ggml_tensor * cur, // [n_embd, block] (post attn_norm) ggml_tensor * main_x, // [n_embd, ctx_len] (post main_norm, shared) + ggml_tensor * cached_ctx_kv, // optional [head_dim, ctx_len], already normed+RoPE const DeepSeek4Weights & w, const DeepSeek4Layer & L, int ctx_len, ggml_tensor * pos_block, // I32[block] absolute positions committed..committed+block-1 ggml_tensor * neg_block, // I32[block] -(block positions) - ggml_tensor * pos_ctx) { // I32[ctx_len] absolute positions committed-ctx_len..committed-1 + ggml_tensor * pos_ctx, // I32[ctx_len] absolute positions committed-ctx_len..committed-1 + ggml_tensor * attn_mask) { // F32[ctx_len+block], 0 or -inf const int n_embd = w.n_embd; const int head_dim = w.head_dim; const int n_head = w.n_head; @@ -5372,10 +5741,13 @@ static ggml_tensor * build_dspark_attention( ggml_tensor * kv_attn = kv_b; int n_attn = block; if (ctx_len > 0) { - ggml_tensor * kv_c = build_rms_norm(ctx, ggml_mul_mat(ctx, L.attn_kv, main_x), L.attn_kv_a_norm, eps); - kv_c = build_tail_rope_2d(ctx, kv_c, pos_ctx, n_rot, head_dim, ctx_len, - rope_freq, rope_scale, rope_ext, rope_attn, - w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_orig); + ggml_tensor * kv_c = cached_ctx_kv; + if (!kv_c) { + kv_c = build_rms_norm(ctx, ggml_mul_mat(ctx, L.attn_kv, main_x), L.attn_kv_a_norm, eps); + kv_c = build_tail_rope_2d(ctx, kv_c, pos_ctx, n_rot, head_dim, ctx_len, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_orig); + } kv_attn = ggml_concat(ctx, kv_c, kv_b, 1); // [head_dim, ctx_len+block] n_attn = ctx_len + block; } @@ -5384,6 +5756,12 @@ static ggml_tensor * build_dspark_attention( ggml_tensor * q_flat = ggml_reshape_2d(ctx, q, head_dim, n_head * block); ggml_tensor * scores = ggml_mul_mat(ctx, kv_attn, q_flat); // [n_attn, n_head*block] scores = ggml_scale(ctx, scores, 1.0f / sqrtf((float) head_dim)); + if (attn_mask) { + // Runtime mask keeps a fixed n_swa-wide graph mathematically + // equivalent to the shorter valid prompt window. Padded context + // columns receive -inf; block columns remain visible. + scores = ggml_add(ctx, scores, ggml_repeat(ctx, attn_mask, scores)); + } ggml_tensor * probs = nullptr; if (L.attn_sinks) { ggml_tensor * sink = ggml_reshape_2d(ctx, L.attn_sinks, 1, n_head); @@ -5431,22 +5809,209 @@ static void ds4_read_f32(ggml_tensor * t, float * dst, int k) { // drafter instance) and re-set only the inputs per call. namespace { +// Cache DSpark context KV projections across speculative steps. Only newly +// committed target-feature columns are projected on the draft GPU; the final +// normed+RoPE KV window is small enough (<1 MiB at n_swa=128) to retain in a +// host-side ring and upload as one compact draft-graph input. +struct DsparkContextKvProjector { + int n_cols = -1; + const void * drafter = nullptr; + std::vector arena; + ggml_context * ctx = nullptr; + ggml_gallocr_t alloc = nullptr; + ggml_cgraph * gf = nullptr; + ggml_tensor * inp_features = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * out = nullptr; + + int end_pos = -1; + int valid = 0; + std::vector host_kv; + std::vector projected; +}; + +thread_local DsparkContextKvProjector g_dspark_ctx_kv; + +static bool dspark_project_context_columns( + ggml_backend_t backend, + const DSparkDrafter & d, + const float * features, + int n_cols, + int first_pos, + std::vector & out) { + if (!backend || !features || n_cols <= 0) return false; + const DeepSeek4Weights & w = d.core; + const int n_embd = w.n_embd; + const int fc_in = d.n_target_layers * n_embd; + const int head_dim = w.head_dim; + DsparkContextKvProjector & P = g_dspark_ctx_kv; + + if (!P.ctx || P.n_cols != n_cols || P.drafter != (const void *) &d) { + if (P.ctx) { + ggml_free(P.ctx); + P.ctx = nullptr; + } + P.gf = nullptr; + P.inp_features = nullptr; + P.positions = nullptr; + P.out = nullptr; + if (P.arena.empty()) P.arena.resize(32u * 1024 * 1024); + ggml_init_params ip{}; + ip.mem_size = P.arena.size(); + ip.mem_buffer = P.arena.data(); + ip.no_alloc = true; + P.ctx = ggml_init(ip); + if (!P.ctx) return false; + P.gf = ggml_new_graph_custom(P.ctx, 4096, false); + P.inp_features = ggml_new_tensor_2d( + P.ctx, GGML_TYPE_F32, fc_in, n_cols); + ggml_set_input(P.inp_features); + P.positions = ggml_new_tensor_1d( + P.ctx, GGML_TYPE_I32, n_cols); + ggml_set_input(P.positions); + + ggml_tensor * feature_norm = + ggml_rms_norm(P.ctx, P.inp_features, w.rms_eps); + ggml_tensor * main_x = build_rms_norm( + P.ctx, ggml_mul_mat(P.ctx, d.main_proj, feature_norm), + d.main_norm, w.rms_eps); + ggml_tensor * stacked = nullptr; + for (int il = 0; il < w.n_layer; ++il) { + const DeepSeek4Layer & L = w.layers[(size_t) il]; + ggml_tensor * kv = build_rms_norm( + P.ctx, ggml_mul_mat(P.ctx, L.attn_kv, main_x), + L.attn_kv_a_norm, w.rms_eps); + kv = build_tail_rope_2d( + P.ctx, kv, P.positions, w.n_rot, head_dim, n_cols, + w.rope_freq_base, 1.0f, 0.0f, 1.0f, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, + (int) w.rope_orig_ctx); + kv = ggml_reshape_3d(P.ctx, kv, head_dim, n_cols, 1); + stacked = stacked ? ggml_concat(P.ctx, stacked, kv, 2) : kv; + } + P.out = ggml_cont(P.ctx, stacked); + ggml_set_output(P.out); + ggml_build_forward_expand(P.gf, P.out); + if (!P.alloc) { + P.alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + } + if (!P.alloc || !ggml_gallocr_alloc_graph(P.alloc, P.gf)) { + ggml_free(P.ctx); + P.ctx = nullptr; + P.gf = nullptr; + return false; + } + P.n_cols = n_cols; + P.drafter = (const void *) &d; + } + + ggml_backend_tensor_set( + P.inp_features, features, 0, + sizeof(float) * (size_t) fc_in * n_cols); + std::vector pos((size_t) n_cols); + for (int i = 0; i < n_cols; ++i) pos[(size_t) i] = first_pos + i; + ggml_backend_tensor_set( + P.positions, pos.data(), 0, sizeof(int32_t) * pos.size()); + if (ggml_backend_graph_compute(backend, P.gf) != GGML_STATUS_SUCCESS) { + return false; + } + out.resize((size_t) head_dim * n_cols * w.n_layer); + ggml_backend_tensor_get( + P.out, out.data(), 0, sizeof(float) * out.size()); + return true; +} + +static bool dspark_update_context_kv_cache( + ggml_backend_t backend, + const DSparkDrafter & d, + const float * ctx_features, + int ctx_len, + int committed, + const std::vector ** out_host_kv) { + const DeepSeek4Weights & w = d.core; + const int n_swa = w.n_swa; + const int head_dim = w.head_dim; + const int fc_in = d.n_target_layers * w.n_embd; + DsparkContextKvProjector & P = g_dspark_ctx_kv; + if (ctx_len <= 0) { + P.end_pos = committed; + P.valid = 0; + P.host_kv.assign((size_t) w.n_layer * n_swa * head_dim, 0.0f); + *out_host_kv = &P.host_kv; + return true; + } + if (!ctx_features || ctx_len > n_swa) return false; + + int n_new = committed - P.end_pos; + const bool rebuild = + P.drafter != (const void *) &d || P.end_pos < 0 || + n_new <= 0 || n_new > ctx_len || + std::min(n_swa, P.valid + n_new) != ctx_len || + P.host_kv.size() != (size_t) w.n_layer * n_swa * head_dim; + if (rebuild) { + P.host_kv.assign((size_t) w.n_layer * n_swa * head_dim, 0.0f); + P.valid = 0; + n_new = ctx_len; + } + + const int first_new_pos = committed - n_new; + const float * new_features = + ctx_features + (size_t) (ctx_len - n_new) * fc_in; + if (!dspark_project_context_columns( + backend, d, new_features, n_new, first_new_pos, P.projected)) { + return false; + } + + const int keep = std::max(0, std::min(P.valid, n_swa - n_new)); + const int drop = P.valid - keep; + for (int il = 0; il < w.n_layer; ++il) { + float * layer_cache = P.host_kv.data() + + (size_t) il * n_swa * head_dim; + if (keep > 0 && drop > 0) { + std::memmove( + layer_cache, + layer_cache + (size_t) drop * head_dim, + sizeof(float) * (size_t) keep * head_dim); + } + const float * layer_new = P.projected.data() + + (size_t) il * n_new * head_dim; + std::memcpy( + layer_cache + (size_t) keep * head_dim, + layer_new, + sizeof(float) * (size_t) n_new * head_dim); + } + P.valid = keep + n_new; + P.end_pos = committed; + P.drafter = (const void *) &d; + if (P.valid != ctx_len) return false; + *out_host_kv = &P.host_kv; + return true; +} + struct DsparkDraftCache { - int ctx_len = -1; + int ctx_len = -1; // allocated graph context width + int valid_ctx_len = -1; // valid columns in the last uploaded context int block = -1; const void * drafter = nullptr; ggml_backend_t backend = nullptr; + bool fixed_context = false; + bool context_kv_cache = false; std::vector arena; ggml_context * ctx = nullptr; ggml_gallocr_t alloc = nullptr; ggml_cgraph * gf = nullptr; ggml_tensor * inp_noise = nullptr; ggml_tensor * inp_ctx = nullptr; + ggml_tensor * inp_ctx_kv = nullptr; ggml_tensor * pos_block = nullptr; ggml_tensor * neg_block = nullptr; ggml_tensor * pos_ctx = nullptr; + ggml_tensor * attn_mask = nullptr; ggml_tensor * out = nullptr; ggml_tensor * confidence_out = nullptr; + std::vector padded_ctx; + std::vector host_attn_mask; std::vector> dbg_taps; // HC scales are immutable weights: read from the backend once. std::vector> s_attn, s_ffn; @@ -5468,16 +6033,29 @@ void reset_deepseek4_dspark_runtime_cache() { cache.ctx = nullptr; } cache = DsparkDraftCache{}; + + DsparkContextKvProjector & projector = g_dspark_ctx_kv; + if (projector.alloc) { + ggml_gallocr_free(projector.alloc); + projector.alloc = nullptr; + } + if (projector.ctx) { + ggml_free(projector.ctx); + projector.ctx = nullptr; + } + projector = DsparkContextKvProjector{}; } -bool deepseek4_dspark_draft_forward(ggml_backend_t backend, +static bool deepseek4_dspark_draft_forward_impl( + ggml_backend_t backend, const DSparkDrafter & d, const float * noise_embed, const float * ctx_features, int ctx_len, int committed, - std::vector & out_hidden, - std::vector * confidence_hidden) { + std::vector * out_hidden, + std::vector * confidence_hidden, + bool upload_context) { const DeepSeek4Weights & w = d.core; const int n_embd = w.n_embd; const int n_hc = w.n_hc; @@ -5486,10 +6064,38 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, const int mix_dim = 2 * n_hc + n_hc * n_hc; const float hc_eps = w.hc_eps; if (ctx_len < 0) ctx_len = 0; + const int valid_ctx_len = ctx_len; + const bool context_kv_cache = + ds4_env_flag("DFLASH_DS4_DRAFT_CONTEXT_KV_CACHE"); + const bool fixed_context = context_kv_cache || + ds4_env_flag("DFLASH_DS4_DRAFT_FIXED_CONTEXT"); + const int graph_ctx_len = fixed_context + ? std::max(valid_ctx_len, w.n_swa) + : valid_ctx_len; + + const std::vector * cached_host_kv = nullptr; + if (context_kv_cache && upload_context && valid_ctx_len > 0 && + !dspark_update_context_kv_cache( + backend, d, ctx_features, valid_ctx_len, committed, + &cached_host_kv)) { + return false; + } DsparkDraftCache & C = g_dspark_draft_cache; const bool DS4_DBG = std::getenv("DFLASH_DS4_DSPARK_DEBUG") != nullptr; + // Context reuse is deliberately strict: never submit a graph with an + // uninitialized or differently-shaped context tensor. The normal warm + // forward must have populated this exact cache first. + if (!upload_context && + (!C.ctx || C.ctx_len != graph_ctx_len || + C.valid_ctx_len != valid_ctx_len || C.block != block || + C.fixed_context != fixed_context || + C.context_kv_cache != context_kv_cache || + C.drafter != (const void *) &d || C.backend != backend)) { + return false; + } + if (C.drafter != (const void *) &d || C.backend != backend) { // HC scales (host) per layer + output — immutable, read once per drafter. C.s_attn.assign((size_t) w.n_layer, {0.0f, 0.0f, 0.0f}); @@ -5503,7 +6109,9 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, C.s_out = so[0]; } - if (!C.ctx || C.ctx_len != ctx_len || C.block != block || + if (!C.ctx || C.ctx_len != graph_ctx_len || C.block != block || + C.fixed_context != fixed_context || + C.context_kv_cache != context_kv_cache || C.drafter != (const void *) &d || C.backend != backend) { // ── (Re)build the graph ───────────────────────────────────────── if (C.backend != backend && C.alloc) { @@ -5533,18 +6141,35 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, // Inputs. C.inp_noise = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, block); ggml_set_input(C.inp_noise); - C.inp_ctx = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, fc_in, ctx_len > 0 ? ctx_len : 1); - ggml_set_input(C.inp_ctx); + C.inp_ctx = nullptr; + C.inp_ctx_kv = nullptr; + if (context_kv_cache) { + C.inp_ctx_kv = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, w.head_dim, + graph_ctx_len > 0 ? graph_ctx_len : 1, w.n_layer); + ggml_set_input(C.inp_ctx_kv); + } else { + C.inp_ctx = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, fc_in, + graph_ctx_len > 0 ? graph_ctx_len : 1); + ggml_set_input(C.inp_ctx); + } C.pos_block = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, block); ggml_set_input(C.pos_block); C.neg_block = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, block); ggml_set_input(C.neg_block); - C.pos_ctx = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ctx_len > 0 ? ctx_len : 1); - ggml_set_input(C.pos_ctx); + C.pos_ctx = nullptr; + if (!context_kv_cache) { + C.pos_ctx = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, + graph_ctx_len > 0 ? graph_ctx_len : 1); + ggml_set_input(C.pos_ctx); + } + C.attn_mask = ggml_new_tensor_1d( + ctx, GGML_TYPE_F32, graph_ctx_len + block); + ggml_set_input(C.attn_mask); // main_x = main_norm(main_proj(ctx_features)). Shared across layers. ggml_tensor * main_x = nullptr; - if (ctx_len > 0) { + if (graph_ctx_len > 0 && !context_kv_cache) { // Captured target features have large magnitude (rms ~1e3 — HC streams // accumulate over 40+ layers). main_proj is rocmfp4-quantized and its // activation quantization overflows on inputs that big -> NaN. Since @@ -5589,8 +6214,19 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, ggml_tensor * attn_in = work_cols[0]; for (int p = 1; p < block; p++) attn_in = ggml_concat(ctx, attn_in, work_cols[p], 1); ggml_tensor * attn_normed = build_rms_norm(ctx, attn_in, L.attn_norm, w.rms_eps); - ggml_tensor * attn_out = build_dspark_attention(ctx, attn_normed, main_x, w, L, - ctx_len, C.pos_block, C.neg_block, C.pos_ctx); + ggml_tensor * layer_ctx_kv = nullptr; + if (context_kv_cache && graph_ctx_len > 0) { + layer_ctx_kv = ggml_view_2d( + ctx, C.inp_ctx_kv, w.head_dim, graph_ctx_len, + C.inp_ctx_kv->nb[1], + (size_t) il * C.inp_ctx_kv->nb[2]); + } + ggml_tensor * attn_out = build_dspark_attention( + ctx, attn_normed, main_x, + layer_ctx_kv, w, L, + graph_ctx_len, C.pos_block, + C.neg_block, C.pos_ctx, + C.attn_mask); dbg_tap(std::string("attn_L") + std::to_string(il), attn_out); // ── HC post (attention), per block position ───────────────── ggml_tensor * hc_next = nullptr; @@ -5666,33 +6302,116 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, ggml_free(C.ctx); C.ctx = nullptr; C.gf = nullptr; return false; } - C.ctx_len = ctx_len; + if (ds4_env_flag("DFLASH_DS4_DRAFT_GRAPH_STATS")) { + std::fprintf(stderr, + "[ds4-draft-graph] ctx=%d valid=%d nodes=%d scratch=%.2f MiB\n", + graph_ctx_len, valid_ctx_len, ggml_graph_n_nodes(gf), + (double) ggml_gallocr_get_buffer_size(C.alloc, 0) / + (1024.0 * 1024.0)); + } + C.ctx_len = graph_ctx_len; + C.valid_ctx_len = -1; C.block = block; + C.fixed_context = fixed_context; + C.context_kv_cache = context_kv_cache; C.drafter = (const void *) &d; C.backend = backend; } // ── Set inputs + compute (cached graph) ───────────────────────────── ggml_backend_tensor_set(C.inp_noise, noise_embed, 0, sizeof(float) * (size_t) n_embd * block); - if (ctx_len > 0) { - ggml_backend_tensor_set(C.inp_ctx, ctx_features, 0, sizeof(float) * (size_t) fc_in * ctx_len); - std::vector pc(ctx_len); - for (int i = 0; i < ctx_len; i++) pc[i] = committed - ctx_len + i; - ggml_backend_tensor_set(C.pos_ctx, pc.data(), 0, sizeof(int32_t) * ctx_len); + if (upload_context) { + if (graph_ctx_len > 0) { + if (context_kv_cache) { + if (!cached_host_kv) return false; + const int head_dim = w.head_dim; + const int source_stride = w.n_swa * head_dim; + const float * upload = cached_host_kv->data(); + if (valid_ctx_len < graph_ctx_len || graph_ctx_len != w.n_swa) { + C.padded_ctx.assign( + (size_t) w.n_layer * graph_ctx_len * head_dim, 0.0f); + for (int il = 0; il < w.n_layer; ++il) { + std::memcpy( + C.padded_ctx.data() + + ((size_t) il * graph_ctx_len + + (graph_ctx_len - valid_ctx_len)) * head_dim, + cached_host_kv->data() + + (size_t) il * source_stride, + sizeof(float) * (size_t) valid_ctx_len * head_dim); + } + upload = C.padded_ctx.data(); + } + ggml_backend_tensor_set( + C.inp_ctx_kv, upload, 0, + sizeof(float) * (size_t) w.n_layer * + graph_ctx_len * head_dim); + } else { + const float * upload = ctx_features; + if (valid_ctx_len < graph_ctx_len) { + C.padded_ctx.assign((size_t) fc_in * graph_ctx_len, 0.0f); + if (valid_ctx_len > 0 && ctx_features) { + std::memcpy( + C.padded_ctx.data() + + (size_t)(graph_ctx_len - valid_ctx_len) * fc_in, + ctx_features, + sizeof(float) * (size_t) fc_in * valid_ctx_len); + } + upload = C.padded_ctx.data(); + } + if (!upload) return false; + ggml_backend_tensor_set( + C.inp_ctx, upload, 0, + sizeof(float) * (size_t) fc_in * graph_ctx_len); + std::vector pc(graph_ctx_len); + for (int i = 0; i < graph_ctx_len; i++) { + pc[i] = committed - graph_ctx_len + i; + } + ggml_backend_tensor_set( + C.pos_ctx, pc.data(), 0, + sizeof(int32_t) * graph_ctx_len); + } + } + C.host_attn_mask.assign((size_t) graph_ctx_len + block, 0.0f); + const int n_pad = graph_ctx_len - valid_ctx_len; + for (int i = 0; i < n_pad; ++i) { + C.host_attn_mask[(size_t)i] = + -std::numeric_limits::infinity(); + } + ggml_backend_tensor_set( + C.attn_mask, C.host_attn_mask.data(), 0, + sizeof(float) * C.host_attn_mask.size()); + C.valid_ctx_len = valid_ctx_len; } std::vector pb(block), nb(block); for (int i = 0; i < block; i++) { pb[i] = committed + i; nb[i] = -(committed + i); } ggml_backend_tensor_set(C.pos_block, pb.data(), 0, sizeof(int32_t) * block); ggml_backend_tensor_set(C.neg_block, nb.data(), 0, sizeof(int32_t) * block); - const ggml_status st = ggml_backend_graph_compute(backend, C.gf); + // This cache owns a single immutable graph: topology, tensor addresses, + // and shapes remain fixed until the cache is explicitly rebuilt above. + // ggml's conservative property scan sees backend-populated tensor metadata + // change and repeatedly drops otherwise valid HIP-graph replay. Once the + // backend has captured the warm graph, bypass that scan for this call only. + // The backend still performs its normal warmup/capture before the bypass + // can take effect. Keep this opt-in until exact output and replay stats + // have both been qualified on the deployment GPUs. + const bool force_graph_replay = + ds4_env_flag("DFLASH_DS4_DRAFT_FORCE_GRAPH_REPLAY"); + if (force_graph_replay) ggml_cuda_set_skip_props_check(true); + const ggml_status st = out_hidden + ? ggml_backend_graph_compute(backend, C.gf) + : ggml_backend_graph_compute_async(backend, C.gf); + if (force_graph_replay) ggml_cuda_set_skip_props_check(false); if (st != GGML_STATUS_SUCCESS) { // Invalidate: a failed compute leaves no reusable state guarantees. ggml_free(C.ctx); C.ctx = nullptr; C.gf = nullptr; C.ctx_len = -1; return false; } - out_hidden.resize((size_t) n_embd * block); - ggml_backend_tensor_get(C.out, out_hidden.data(), 0, sizeof(float) * out_hidden.size()); + if (!out_hidden) return true; + + out_hidden->resize((size_t) n_embd * block); + ggml_backend_tensor_get( + C.out, out_hidden->data(), 0, sizeof(float) * out_hidden->size()); if (confidence_hidden) { confidence_hidden->resize((size_t) n_embd * block); ggml_backend_tensor_get(C.confidence_out, confidence_hidden->data(), 0, @@ -5717,4 +6436,59 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, return true; } +bool deepseek4_dspark_draft_forward(ggml_backend_t backend, + const DSparkDrafter & d, + const float * noise_embed, + const float * ctx_features, + int ctx_len, + int committed, + std::vector & out_hidden, + std::vector * confidence_hidden) { + return deepseek4_dspark_draft_forward_impl( + backend, d, noise_embed, ctx_features, ctx_len, committed, + &out_hidden, confidence_hidden, true); +} + +bool deepseek4_dspark_draft_forward_async(ggml_backend_t backend, + const DSparkDrafter & d, + const float * noise_embed, + const float * ctx_features, + int ctx_len, + int committed) { + return deepseek4_dspark_draft_forward_impl( + backend, d, noise_embed, ctx_features, ctx_len, committed, + nullptr, nullptr, true); +} + +bool deepseek4_dspark_draft_forward_async_reuse_context( + ggml_backend_t backend, + const DSparkDrafter & d, + const float * noise_embed, + int ctx_len, + int committed) { + return deepseek4_dspark_draft_forward_impl( + backend, d, noise_embed, nullptr, ctx_len, committed, + nullptr, nullptr, false); +} + +bool deepseek4_dspark_draft_read_async_output( + ggml_backend_t backend, + std::vector & out_hidden) { + DsparkDraftCache & C = g_dspark_draft_cache; + if (!backend || !C.ctx || !C.out || C.block <= 0 || !C.drafter) { + return false; + } + const DSparkDrafter * d = + static_cast(C.drafter); + const size_t count = (size_t) d->core.n_embd * (size_t) C.block; + out_hidden.resize(count); + ggml_backend_tensor_get( + C.out, out_hidden.data(), 0, sizeof(float) * count); + return true; +} + +void deepseek4_dspark_draft_wait(ggml_backend_t backend) { + ggml_backend_synchronize(backend); +} + } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index a09cd1f62..275694cee 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -30,6 +30,7 @@ namespace dflash::common { struct MoeHybridPlacement; struct MoeHybridConfig; struct MoeHybridRoutingStats; +struct MoeExpertComputeRuntime; class MoeHybridStreamEngine; struct DeepSeek4StepTelemetry { @@ -149,6 +150,9 @@ struct DeepSeek4Weights { ggml_context * ctx = nullptr; ggml_backend_t backend = nullptr; ggml_backend_buffer_t buf = nullptr; + // Optional row-split buffer for selected dense projections. The buffer + // owns per-device allocations while the tensor metadata stays in ctx. + ggml_backend_buffer_t dense_split_buf = nullptr; // Global tensors ggml_tensor * tok_embd = nullptr; // [n_embd, n_vocab] @@ -260,6 +264,9 @@ struct DeepSeek4LayerCache { struct DeepSeek4Cache { int cur_pos = 0; + // Monotonic request generation. Whole-step HIP graphs may retain kernel + // state safely within one sequence, but must not survive a cache clear. + uint64_t sequence_id = 0; int max_ctx = 0; int n_layer = 0; @@ -347,7 +354,8 @@ bool deepseek4_step( MoeHybridStreamEngine * stream_engine = nullptr, DeepSeek4StepTelemetry * telemetry = nullptr, MoeHybridRoutingStats * routing_stats = nullptr, - Ds4VerifyHooks * verify_hooks = nullptr); + Ds4VerifyHooks * verify_hooks = nullptr, + MoeExpertComputeRuntime * expert_runtime = nullptr); // Optional hooks for the DSpark spec-decode batched verify (deepseek4_dspark). // When set on a multi-token deepseek4_step_layer_range call they add: per-layer @@ -357,6 +365,8 @@ struct Ds4VerifyHooks { const std::vector * capture_layer_ids = nullptr; // e.g. {40,41,42} std::vector * capture_out = nullptr; // [n_cap*n_embd * n_tokens] std::vector * all_logits_out = nullptr; // [n_vocab * n_tokens] + std::vector * argmax_out = nullptr; // [n_tokens], optional GPU result + bool prefer_argmax_only = false; // skip logits D2H when available }; bool deepseek4_step_layer_range( @@ -373,7 +383,10 @@ bool deepseek4_step_layer_range( const int32_t * token_ids = nullptr, DeepSeek4StepTelemetry * telemetry = nullptr, bool allow_decode_graph_reuse = true, - Ds4VerifyHooks * verify_hooks = nullptr); + Ds4VerifyHooks * verify_hooks = nullptr, + MoeHybridStorage * moe_hybrid = nullptr, + MoeExpertComputeRuntime * expert_runtime = nullptr, + MoeHybridRoutingStats * routing_stats = nullptr); bool build_deepseek4_moe_hybrid_storage_from_file( const std::string & path, @@ -399,7 +412,8 @@ bool build_deepseek4_moe_hybrid_storage_from_file_with_mmap( const MoeHybridPlacement & placement, const MoeHybridConfig * cfg_override, MoeHybridStorage & out, - std::string * err = nullptr); + std::string * err = nullptr, + ggml_backend_t cold_gpu_backend = nullptr); // Snapshot struct DeepSeek4Snapshot { diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index ef3aa9fbf..5437985d8 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -198,6 +198,32 @@ static bool should_upload_ds4_tensor(const char * name, return !(plan.skip_expert_tensors && is_expert_tensor(name)); } +static int ds4_dense_tp_mask() { + const char * value = std::getenv("DFLASH_DS4_DENSE_TP_MASK"); + if (!value || !value[0]) return 0; + return std::max(0, std::atoi(value)); +} + +static bool should_split_ds4_dense_tensor(const char * name, int mask) { + if (mask == 0 || !name) return false; + // Each selected tensor is consumed as src0 of a plain 2-D MUL_MAT. Expert + // tensors and the grouped output-A view are deliberately excluded: the + // CUDA/HIP split-buffer implementation cannot split 3-D weight views. + if ((mask & 1) && std::strcmp(name, "output.weight") == 0) return true; + if ((mask & 2) && std::strstr(name, ".attn_q_b.weight")) return true; + if ((mask & 4) && std::strstr(name, ".attn_output_b.weight")) return true; + if ((mask & 8) && (std::strstr(name, ".attn_q_a.weight") || + std::strstr(name, ".attn_kv.weight") || + std::strstr(name, ".indexer.attn_q_b.weight") || + std::strstr(name, ".attn_compressor_kv.weight") || + std::strstr(name, ".attn_compressor_gate.weight") || + std::strstr(name, ".indexer_compressor_kv.weight") || + std::strstr(name, ".indexer_compressor_gate.weight"))) { + return true; + } + return false; +} + struct DS4TensorAlloc { ggml_tensor * tensor = nullptr; size_t tensor_offset = 0; @@ -205,6 +231,7 @@ struct DS4TensorAlloc { size_t file_size = 0; size_t buffer_offset = 0; bool upload_to_backend = true; + bool dense_split = false; }; } // namespace @@ -417,10 +444,30 @@ bool load_deepseek4_gguf_partial(const std::string & path, const size_t data_offset = gguf_get_data_offset(gctx); ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(backend); const size_t alignment = ggml_backend_buft_get_alignment(buft); + const int dense_tp_mask = ds4_dense_tp_mask(); + ggml_backend_buffer_type_t split_buft = nullptr; + size_t split_alignment = 1; + if (dense_tp_mask != 0 && ggml_backend_is_cuda(backend) && + ggml_backend_cuda_get_device_count() >= 2) { + float strix_fraction = 0.28f; + if (const char * value = std::getenv("DFLASH_DS4_DENSE_TP_STRIX_FRACTION")) { + const float parsed = std::strtof(value, nullptr); + if (parsed > 0.0f && parsed < 1.0f) strix_fraction = parsed; + } + float tensor_split[GGML_CUDA_MAX_DEVICES] = {}; + tensor_split[0] = 1.0f - strix_fraction; + tensor_split[1] = strix_fraction; + split_buft = ggml_backend_cuda_split_buffer_type(0, tensor_split); + split_alignment = ggml_backend_buft_get_alignment(split_buft); + std::fprintf(stderr, + "[deepseek4-dense-tp] mask=%d row split R9700=%.3f Strix=%.3f\n", + dense_tp_mask, 1.0f - strix_fraction, strix_fraction); + } std::vector allocs; allocs.reserve(n_tensors); size_t total_buf_size = 0; + size_t split_total_buf_size = 0; size_t tok_embd_alloc_idx = SIZE_MAX; for (int ti = 0; ti < n_tensors; ti++) { @@ -438,10 +485,20 @@ bool load_deepseek4_gguf_partial(const std::string & path, a.tensor_offset = tensor_offset; a.file_size = gguf_get_tensor_size(gctx, ti); a.upload_to_backend = upload_to_backend; + a.dense_split = upload_to_backend && split_buft && + should_split_ds4_dense_tensor(tname, dense_tp_mask); if (upload_to_backend) { - total_buf_size = align_up_size(total_buf_size, alignment); - a.buffer_offset = total_buf_size; - total_buf_size += ggml_backend_buft_get_alloc_size(buft, t); + if (a.dense_split) { + split_total_buf_size = align_up_size( + split_total_buf_size, split_alignment); + a.buffer_offset = split_total_buf_size; + split_total_buf_size += + ggml_backend_buft_get_alloc_size(split_buft, t); + } else { + total_buf_size = align_up_size(total_buf_size, alignment); + a.buffer_offset = total_buf_size; + total_buf_size += ggml_backend_buft_get_alloc_size(buft, t); + } } allocs.push_back(a); if (std::strcmp(tname, "token_embd.weight") == 0) { @@ -451,6 +508,7 @@ bool load_deepseek4_gguf_partial(const std::string & path, // ── Allocate GPU buffer ───────────────────────────────────────────── ggml_backend_buffer_t buf = nullptr; + ggml_backend_buffer_t split_buf = nullptr; if (total_buf_size > 0) { buf = ggml_backend_alloc_buffer(backend, total_buf_size); if (!buf) { @@ -461,14 +519,36 @@ bool load_deepseek4_gguf_partial(const std::string & path, } ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); } + if (split_total_buf_size > 0) { + split_buf = ggml_backend_buft_alloc_buffer( + split_buft, split_total_buf_size); + if (!split_buf) { + set_last_error("failed to allocate dense TP split buffer (" + + std::to_string(split_total_buf_size) + " bytes)"); + if (buf) ggml_backend_buffer_free(buf); + gguf_free(gctx); + if (meta_ctx) ggml_free(meta_ctx); + return false; + } + ggml_backend_buffer_set_usage( + split_buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + } // ── Assign tensors from meta_ctx to the backend buffer ────────────── // Use ggml_backend_tensor_alloc to properly set the buffer association. char * buf_base = buf ? (char *)ggml_backend_buffer_get_base(buf) : nullptr; + char * split_buf_base = split_buf + ? (char *)ggml_backend_buffer_get_base(split_buf) : nullptr; for (auto & a : allocs) { - if (!a.upload_to_backend || !buf) continue; - if (ggml_backend_tensor_alloc(buf, a.tensor, buf_base + a.buffer_offset) != GGML_STATUS_SUCCESS) { + if (!a.upload_to_backend) continue; + ggml_backend_buffer_t tensor_buf = a.dense_split ? split_buf : buf; + char * tensor_base = a.dense_split ? split_buf_base : buf_base; + if (!tensor_buf || + ggml_backend_tensor_alloc(tensor_buf, a.tensor, + tensor_base + a.buffer_offset) != + GGML_STATUS_SUCCESS) { set_last_error("ggml_backend_tensor_alloc failed"); + if (split_buf) ggml_backend_buffer_free(split_buf); if (buf) ggml_backend_buffer_free(buf); gguf_free(gctx); ggml_free(meta_ctx); @@ -481,6 +561,7 @@ bool load_deepseek4_gguf_partial(const std::string & path, std::string mmap_err; if (!mmap.open_ro(path, mmap_err)) { set_last_error("mmap: " + mmap_err); + if (split_buf) ggml_backend_buffer_free(split_buf); if (buf) ggml_backend_buffer_free(buf); gguf_free(gctx); ggml_free(meta_ctx); @@ -496,6 +577,7 @@ bool load_deepseek4_gguf_partial(const std::string & path, a.file_size, mmap.len)); mmap.close_map(); + if (split_buf) ggml_backend_buffer_free(split_buf); if (buf) ggml_backend_buffer_free(buf); gguf_free(gctx); ggml_free(meta_ctx); @@ -518,7 +600,7 @@ bool load_deepseek4_gguf_partial(const std::string & path, size_t i; while ((i = next.fetch_add(1)) < allocs.size()) { auto & a = allocs[i]; - if (!a.upload_to_backend) continue; + if (!a.upload_to_backend || a.dense_split) continue; char * dst = (char *) a.tensor->data; size_t done = 0; while (done < a.file_size) { @@ -534,9 +616,17 @@ bool load_deepseek4_gguf_partial(const std::string & path, for (auto & th : pool) th.join(); posix_fadvise(mmap.fd, 0, (off_t) mmap.len, POSIX_FADV_DONTNEED); ggml_backend_synchronize(backend); // make CPU-written managed pages visible to GPU + // Split tensors have no single CPU-visible base address. Upload them + // through the split buffer, which distributes whole rows to each GPU. + for (auto & a : allocs) { + if (!a.upload_to_backend || !a.dense_split) continue; + const void * src_data = (const char *)mmap.addr + a.file_offset; + ggml_backend_tensor_set(a.tensor, src_data, 0, a.file_size); + } if (!read_ok) { set_last_error("parallel weight read failed"); mmap.close_map(); + if (split_buf) ggml_backend_buffer_free(split_buf); if (buf) ggml_backend_buffer_free(buf); gguf_free(gctx); ggml_free(meta_ctx); @@ -567,6 +657,7 @@ bool load_deepseek4_gguf_partial(const std::string & path, emb_mmap.close_map(); } else { set_last_error("embedder mmap: " + emb_err); + if (split_buf) ggml_backend_buffer_free(split_buf); if (buf) ggml_backend_buffer_free(buf); gguf_free(gctx); ggml_free(meta_ctx); @@ -654,12 +745,14 @@ bool load_deepseek4_gguf_partial(const std::string & path, out.ctx = meta_ctx; out.buf = buf; + out.dense_split_buf = split_buf; gguf_free(gctx); // Note: meta_ctx is now owned by out.ctx — do NOT free it here. - std::fprintf(stderr, "[deepseek4] loaded %zu tensors, %.1f MB GPU buffer%s\n", + std::fprintf(stderr, "[deepseek4] loaded %zu tensors, %.1f MB GPU buffer, %.1f MB dense TP split%s\n", allocs.size(), (double)total_buf_size / (1024.0 * 1024.0), + (double)split_total_buf_size / (1024.0 * 1024.0), plan.expert_metadata_only ? " [expert-metadata-only]" : ""); return true; } @@ -790,7 +883,8 @@ bool build_deepseek4_moe_hybrid_storage_from_file_with_mmap( const MoeHybridPlacement & placement, const MoeHybridConfig * cfg_override, MoeHybridStorage & out, - std::string * err) { + std::string * err, + ggml_backend_t cold_gpu_backend) { ggml_context * expert_meta = nullptr; gguf_init_params gip{}; gip.no_alloc = true; @@ -859,7 +953,7 @@ bool build_deepseek4_moe_hybrid_storage_from_file_with_mmap( const MoeHybridConfig cfg = cfg_override ? *cfg_override : make_ds4_moe_hybrid_config(w); const bool ok = build_moe_hybrid_storage_from_file_with_mmap( cfg, backend, placement, layer_descs, layer_file_data, - mmap.addr, mmap.len, out, err); + mmap.addr, mmap.len, out, err, 0, cold_gpu_backend); if (!ok) { mmap.close_map(); @@ -885,6 +979,10 @@ bool build_deepseek4_moe_hybrid_storage_from_file( void free_deepseek4_weights(DeepSeek4Weights & w) { if (w.ctx) { ggml_free(w.ctx); w.ctx = nullptr; } + if (w.dense_split_buf) { + ggml_backend_buffer_free(w.dense_split_buf); + w.dense_split_buf = nullptr; + } if (w.buf) { ggml_backend_buffer_free(w.buf); w.buf = nullptr; } w.layers.clear(); w.embedder.tok_embd_owned.clear(); From 2377ebd9e89206d18a7f4e3931af8e080f5355e2 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 19 Jul 2026 22:19:11 -0700 Subject: [PATCH 03/15] perf(rocm): specialize Lucebox heterogeneous verify kernels --- .../deps/llama.cpp/ggml/include/ggml-cuda.h | 34 + .../llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu | 41 + .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 1275 ++++++++++++++++- .../deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu | 788 ++++++++-- .../llama.cpp/ggml/src/ggml-cuda/moe-fused.cu | 388 ++++- 5 files changed, 2404 insertions(+), 122 deletions(-) diff --git a/server/deps/llama.cpp/ggml/include/ggml-cuda.h b/server/deps/llama.cpp/ggml/include/ggml-cuda.h index 091e91297..6df853376 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-cuda.h +++ b/server/deps/llama.cpp/ggml/include/ggml-cuda.h @@ -24,6 +24,34 @@ GGML_BACKEND_API ggml_backend_t ggml_backend_cuda_init(int device); GGML_BACKEND_API bool ggml_backend_is_cuda(ggml_backend_t backend); +// Configure streams lazily created by this backend context at the device's +// lowest scheduling priority. Must be called before the backend first submits +// work. Other backend contexts on the same device are unaffected. +GGML_BACKEND_API bool ggml_backend_cuda_set_low_priority_stream( + ggml_backend_t backend); + +// Skip the expensive per-node CUDA/HIP graph property comparison on the +// calling thread once a stable graph has already been captured. Callers must +// bracket only immutable-topology graphs whose tensor addresses and shapes do +// not change; input contents may still be updated in place. +GGML_BACKEND_API void ggml_cuda_set_skip_props_check(bool skip); + +// Capture a scheduler-wide device program around multiple ordinary backend +// graph submissions. The handle is opaque and belongs to `backend`; callers +// must free it before freeing that backend. These hooks are intentionally +// opt-in so normal per-split CUDA/HIP graph behavior is unchanged. +typedef void * ggml_backend_cuda_whole_graph_t; +GGML_BACKEND_API bool ggml_backend_cuda_whole_graph_capture_prepare( + ggml_backend_t backend0, ggml_backend_t backend1); +GGML_BACKEND_API bool ggml_backend_cuda_whole_graph_capture_begin( + ggml_backend_t backend); +GGML_BACKEND_API ggml_backend_cuda_whole_graph_t +ggml_backend_cuda_whole_graph_capture_end(ggml_backend_t backend); +GGML_BACKEND_API bool ggml_backend_cuda_whole_graph_launch( + ggml_backend_t backend, ggml_backend_cuda_whole_graph_t graph); +GGML_BACKEND_API void ggml_backend_cuda_whole_graph_free( + ggml_backend_t backend, ggml_backend_cuda_whole_graph_t graph); + // device buffer GGML_BACKEND_API ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device); @@ -40,6 +68,12 @@ GGML_BACKEND_API int ggml_backend_cuda_get_device_count(void); GGML_BACKEND_API void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size); GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total); +// Override the plain quantized MUL_MAT MMVQ column ceiling on the calling +// thread. Pass zero to restore LUCE_MMVQ_MAX_NCOLS. This is intentionally +// thread-local so one graph builder can select a safe topology without +// changing concurrent requests or other CUDA/HIP backends. +GGML_BACKEND_API void ggml_backend_cuda_set_mmvq_max_ncols_override(int max_ncols); + GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size); GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu index 1c6b86972..cf09c8a88 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu @@ -23,6 +23,10 @@ // src2 = hc_state [n_embd*n_hc] // dst = [n_embd]: weights[h] = sigmoid(mix[h]*s0+base[h]) + 1e-6; // dst[d] = sum_h weights[h]*hc_state[h*n_embd+d] +// +// mode 3 (post split): mode 1 with src1 = main block_out, src3 = peer +// block_out. The kernel evaluates peer[d] + main[d] before +// multiplying by post[h], matching the eliminated GGML add. #define DS4_HC_SINKHORN_EPS 1.0e-6f #define DS4_HC_MAX_HC 8 @@ -288,6 +292,33 @@ static __global__ void ds4_hc_post_kernel( dst[i] = acc; } +static __global__ void ds4_hc_post_split_kernel( + const float * __restrict__ residual, + const float * __restrict__ main_block, + const float * __restrict__ peer_block, + const float * __restrict__ split, + float * __restrict__ dst, + int n_embd, + int n_hc) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int total = n_embd * n_hc; + if (i >= total) { + return; + } + const int h = i / n_embd; + const int d = i - h * n_embd; + const float * post = split + n_hc; + const float * comb = split + 2 * n_hc; + // Keep the established reduction order: the graph being replaced uses + // ggml_add(peer, main), followed by HC post multiplication. + const float block_out = peer_block[d] + main_block[d]; + float acc = block_out * post[h]; + for (int src = 0; src < n_hc; ++src) { + acc += comb[h + src * n_hc] * residual[(size_t) src * n_embd + d]; + } + dst[i] = acc; +} + static __global__ void ds4_hc_out_kernel( const float * __restrict__ mix, const float * __restrict__ base, @@ -361,6 +392,16 @@ void ggml_cuda_op_ds4_hc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { (const float *) src2->data, (float *) dst->data, n_embd, n_hc, pre_scale); } break; + case 3: { + const ggml_tensor * src3 = dst->src[3]; + GGML_ASSERT(src3 && src3->type == GGML_TYPE_F32); + const int total = n_embd * n_hc; + const int blocks = (total + 255) / 256; + ds4_hc_post_split_kernel<<>>( + (const float *) src0->data, (const float *) src1->data, + (const float *) src3->data, (const float *) src2->data, + (float *) dst->data, n_embd, n_hc); + } break; default: GGML_ABORT("ds4_hc: unknown mode"); } 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 6f887eafc..905e70493 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 @@ -85,6 +85,7 @@ #include #include #include +#include #include #include #include @@ -92,6 +93,139 @@ static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size"); +static thread_local int ggml_cuda_mmvq_max_ncols_override = 0; +static thread_local bool ggml_cuda_graphs_disabled_override = false; +// Number of device streams currently participating in one scheduler-wide +// capture on this host thread. While non-zero, per-split graph capture/replay +// must be bypassed and cross-device events must remain external graph nodes. +static thread_local int ggml_cuda_whole_graph_capture_depth = 0; + +// HIP cannot reliably capture cross-device event record/wait nodes for the +// gfx1201 + gfx1151 pair (ROCm 7.0 segfaults inside stream submission). +// hipStreamWait/WriteValue32 also silently disappear during graph capture on +// ROCm 7.0. Use captured kernels and one no-reset sequence counter per +// dependency instead. The producer increments its counter with system-release +// semantics; the consumer waits for its next value with system-acquire +// semantics. This avoids both missing graph nodes and stale-value ABA races. +// q4 hybrid verification captures 129 forward and 43 reverse dependencies. +// Leave headroom for scheduler changes without creating thousands of native +// signal allocations per cached graph slot. +static constexpr uint32_t GGML_CUDA_WHOLE_GRAPH_MAX_DOORBELLS = 256; +static constexpr size_t GGML_CUDA_WHOLE_GRAPH_STAGING_BYTES = + 32u * 1024u * 1024u; +struct ggml_cuda_whole_graph_handle; +struct ggml_cuda_whole_graph_session { + std::array devices{-1, -1}; + std::array, 2> signals{}; + std::array signal_tables{}; + std::array reset_streams{}; + // observed[i] is local to the consumer of signals[i]. + std::array observed{}; + std::array next_signals{}; + void * staging_host = nullptr; + uint8_t * staging = nullptr; + size_t staging_used = 0; + struct doorbell { + uint32_t * signal = nullptr; + uint32_t * observed = nullptr; + }; + std::unordered_map event_signals; + std::vector pending_handles; + bool hardware_sync = true; + bool instantiate_failed = false; + int references = 0; + int graph_count = 0; + int launches_in_round = 0; +}; + +static thread_local ggml_cuda_whole_graph_session * + ggml_cuda_whole_graph_capture_session = nullptr; + +static ggml_cuda_whole_graph_session::doorbell +ggml_cuda_whole_graph_signal_acquire(int producer_device) { +#if defined(GGML_USE_HIP) + auto * session = ggml_cuda_whole_graph_capture_session; + if (!session) return {}; + for (size_t i = 0; i < session->devices.size(); ++i) { + if (session->devices[i] != producer_device) continue; + const uint32_t index = session->next_signals[i]++; + GGML_ASSERT(index < GGML_CUDA_WHOLE_GRAPH_MAX_DOORBELLS); + return {session->signals[i][index], + session->observed[i] + index}; + } +#else + GGML_UNUSED(producer_device); +#endif + return {}; +} + +static uint8_t * ggml_cuda_whole_graph_staging_acquire(size_t bytes) { +#if defined(GGML_USE_HIP) + auto * session = ggml_cuda_whole_graph_capture_session; + if (!session) return nullptr; + const size_t offset = (session->staging_used + 255u) & ~size_t(255u); + GGML_ASSERT(offset <= GGML_CUDA_WHOLE_GRAPH_STAGING_BYTES); + GGML_ASSERT(bytes <= GGML_CUDA_WHOLE_GRAPH_STAGING_BYTES - offset); + session->staging_used = offset + bytes; + return session->staging + offset; +#else + GGML_UNUSED(bytes); + return nullptr; +#endif +} + +static ggml_cuda_whole_graph_session::doorbell +ggml_cuda_whole_graph_event_signal( + void * event_context, int producer_device, bool create) { +#if defined(GGML_USE_HIP) + auto * session = ggml_cuda_whole_graph_capture_session; + if (!session) return {}; + const auto found = session->event_signals.find(event_context); + if (found != session->event_signals.end()) return found->second; + if (!create) return {}; + const auto signal = + ggml_cuda_whole_graph_signal_acquire(producer_device); + if (signal.signal) session->event_signals.emplace(event_context, signal); + return signal; +#else + GGML_UNUSED(event_context); + GGML_UNUSED(producer_device); + GGML_UNUSED(create); + return {}; +#endif +} + +extern "C" void ggml_backend_cuda_set_mmvq_max_ncols_override(int max_ncols) { + ggml_cuda_mmvq_max_ncols_override = std::max(0, max_ncols); +} + +extern "C" void ggml_backend_cuda_set_graphs_disabled_override(bool disabled) { + ggml_cuda_graphs_disabled_override = disabled; +} + +static bool ggml_cuda_graphs_disabled_for_device(int device) { + static const uint64_t disabled_mask = []() { + uint64_t mask = 0; + const char * cursor = std::getenv("GGML_CUDA_DISABLE_GRAPHS_DEVICES"); + while (cursor && *cursor) { + char * end = nullptr; + const long id = std::strtol(cursor, &end, 10); + if (end == cursor) { + ++cursor; + continue; + } + if (id >= 0 && id < 64) { + mask |= uint64_t{1} << id; + } + cursor = end; + } + return mask; + }(); + return device >= 0 && device < 64 && + (disabled_mask & (uint64_t{1} << device)) != 0; +} + [[noreturn]] void ggml_cuda_error(const char * stmt, const char * func, const char * file, int line, const char * msg) { int id = -1; // in case cudaGetDevice fails @@ -2494,11 +2628,14 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor // measured crossover on sm_86 (RTX 3090, Q4_K_M/Q6_K dense GEMVs) — MMVQ // wins at ncols<=3, MMQ wins at 4-8 (laguna w6 chain 199->237 tok/s, // qwen3.6 chain 127->137). Override via env for other hardware. - static const int luce_mmvq_max_ncols = []() { + static const int luce_mmvq_max_ncols_env = []() { const char * e = getenv("LUCE_MMVQ_MAX_NCOLS"); const int v = e ? atoi(e) : 3; return v > 0 ? v : MMVQ_MAX_BATCH_SIZE; }(); + const int luce_mmvq_max_ncols = ggml_cuda_mmvq_max_ncols_override > 0 + ? ggml_cuda_mmvq_max_ncols_override + : luce_mmvq_max_ncols_env; bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !bad_padding_clear && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 && src1->ne[1] <= luce_mmvq_max_ncols; @@ -3146,6 +3283,90 @@ static void ggml_backend_cuda_get_tensor_2d_async(ggml_backend_t backend, const data, stride_data, (const char *) tensor->data + offset, stride_tensor, size, n_copies, cudaMemcpyDeviceToHost, cuda_ctx->stream())); } +__global__ static void ggml_cuda_peer_copy_bytes( + const uint8_t * src, uint8_t * dst, size_t n) { + for (size_t i = (size_t) blockIdx.x * blockDim.x + threadIdx.x; + i < n; i += (size_t) blockDim.x * gridDim.x) { + dst[i] = src[i]; + } +#if defined(GGML_USE_HIP) + // The destination belongs to the peer GPU. Make every thread's stores + // system-visible before the following publish kernel advances the + // dependency sequence. + __threadfence_system(); +#endif +} + +#if defined(GGML_USE_HIP) +__global__ static void ggml_cuda_whole_graph_publish(uint32_t * signal) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + __threadfence_system(); + (void) __hip_atomic_fetch_add(signal, 1u, __ATOMIC_RELEASE, + __HIP_MEMORY_SCOPE_SYSTEM); + } +} + +__global__ static void ggml_cuda_whole_graph_wait( + const uint32_t * signal, uint32_t * observed) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + const uint32_t target = *observed + 1u; + uint32_t current; + do { + current = __hip_atomic_load(signal, __ATOMIC_ACQUIRE, + __HIP_MEMORY_SCOPE_SYSTEM); + } while ((int32_t) (current - target) < 0); + *observed = target; + __threadfence_system(); + } +} + +// Fence-only forms used after graph surgery. The command processor performs +// the actual wait/write; these short kernels retain the system-scope release +// and acquire semantics required by payloads in coherent host memory without +// ever spinning or incrementing the signal themselves. +__global__ static void ggml_cuda_whole_graph_release_fence(uint32_t *) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + __threadfence_system(); + } +} + +__global__ static void ggml_cuda_whole_graph_acquire_fence( + const uint32_t * signal, uint32_t * observed) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + const uint32_t current = __hip_atomic_load( + signal, __ATOMIC_ACQUIRE, __HIP_MEMORY_SCOPE_SYSTEM); + *observed = current; + __threadfence_system(); + } +} + +__global__ static void ggml_cuda_whole_graph_reset_signals( + uint32_t * const * signals, uint32_t count) { + const uint32_t index = blockIdx.x * blockDim.x + threadIdx.x; + if (index < count) { + __hip_atomic_store(signals[index], 0u, __ATOMIC_RELEASE, + __HIP_MEMORY_SCOPE_SYSTEM); + } +} + +static void ggml_cuda_whole_graph_publish_async( + hipStream_t stream, + const ggml_cuda_whole_graph_session::doorbell & doorbell) { + GGML_ASSERT(doorbell.signal && doorbell.observed); + ggml_cuda_whole_graph_publish<<<1, 1, 0, stream>>>(doorbell.signal); + CUDA_CHECK(hipGetLastError()); +} + +static void ggml_cuda_whole_graph_wait_async( + hipStream_t stream, + const ggml_cuda_whole_graph_session::doorbell & doorbell) { + GGML_ASSERT(doorbell.signal && doorbell.observed); + ggml_cuda_whole_graph_wait<<<1, 1, 0, stream>>>( + doorbell.signal, doorbell.observed); + CUDA_CHECK(hipGetLastError()); +} +#endif + static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { ggml_backend_buffer_t buf_src = src->view_src ? src->view_src->buffer : src->buffer; ggml_backend_buffer_t buf_dst = dst->view_src ? dst->view_src->buffer : dst->buffer; @@ -3173,6 +3394,77 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ } if (backend_src != backend_dst) { +#if defined(GGML_USE_HIP) + if (ggml_cuda_whole_graph_capture_depth > 0 && + cuda_ctx_src->device != cuda_ctx_dst->device) { + // Deferred MoE joins are deliberately staged at producer-split + // completion, before their existing event is published. Direct + // Strix writes into ordinary R9700 allocations can fault inside + // the full production graph, while destination-side peer reads + // can return stale cache lines on this gfx1151/gfx1201 pair. + // Therefore the producer writes coherent mapped staging here and + // mode -3 copies staging into the R9700-local destination only at + // its already-late consumer split. Hot/shared R9700 work remains + // enqueued before the join is gated. + if (dst->op == GGML_OP_MOE_FUSED && + ggml_get_op_params_i32(dst, 0) == -3) { + const size_t n = ggml_nbytes(dst); + GGML_ASSERT(ggml_nbytes(src) == n); + uint8_t * staging = + ggml_cuda_whole_graph_staging_acquire(n); + GGML_ASSERT(staging); + ggml_cuda_set_device(cuda_ctx_src->device); + const int block = 256; + const int grid = (int) std::min( + (n + block - 1) / block, 65535); + ggml_cuda_peer_copy_bytes<<stream()>>>( + (const uint8_t *) src->data, + staging, n); + CUDA_CHECK(hipGetLastError()); + static_assert( + 8 * sizeof(int32_t) + sizeof(staging) <= + GGML_MAX_OP_PARAMS, + "deferred staging pointer does not fit op_params"); + memcpy(&dst->op_params[8], &staging, sizeof(staging)); + return true; + } + + // Captured hipMemcpyPeerAsync nodes fail when replayed across this + // discrete + unified-memory pair. Destination-side peer reads can + // consume stale cache lines, while direct R9700 peer writes into + // large Strix unified-memory scheduler buffers fault on replay. + // Stage through coherent mapped memory with captured kernels: + // producer -> staging, publish/wait, staging -> destination-local. + // The host never executes or copies payload data during replay. + const auto doorbell = ggml_cuda_whole_graph_signal_acquire( + cuda_ctx_src->device); + GGML_ASSERT(doorbell.signal); + const size_t n = ggml_nbytes(dst); + uint8_t * staging = + ggml_cuda_whole_graph_staging_acquire(n); + GGML_ASSERT(staging); + ggml_cuda_set_device(cuda_ctx_src->device); + const int block = 256; + const int grid = (int) std::min( + (n + block - 1) / block, 65535); + ggml_cuda_peer_copy_bytes<<stream()>>>( + (const uint8_t *) src->data, + staging, n); + CUDA_CHECK(hipGetLastError()); + ggml_cuda_whole_graph_publish_async( + (hipStream_t) cuda_ctx_src->stream(), doorbell); + ggml_cuda_set_device(cuda_ctx_dst->device); + ggml_cuda_whole_graph_wait_async( + (hipStream_t) cuda_ctx_dst->stream(), doorbell); + ggml_cuda_peer_copy_bytes<<stream()>>>( + staging, (uint8_t *) dst->data, n); + CUDA_CHECK(hipGetLastError()); + return true; + } +#endif // copy on src stream if (cuda_ctx_src->device == cuda_ctx_dst->device) { CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); @@ -3184,16 +3476,30 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ #endif // GGML_CUDA_NO_PEER_COPY } - // record event on src stream after the copy + // During a scheduler-wide HIP capture, publish each peer copy through + // a unique device doorbell. Cross-device external event nodes crash + // the current heterogeneous ROCm graph implementation and reusing an + // event would also make layer waits vulnerable to ABA on replay. + // Ordinary eager/per-split submissions retain the existing event. if (!cuda_ctx_src->copy_event) { ggml_cuda_set_device(cuda_ctx_src->device); CUDA_CHECK(cudaEventCreateWithFlags(&cuda_ctx_src->copy_event, cudaEventDisableTiming)); } - - CUDA_CHECK(cudaEventRecord(cuda_ctx_src->copy_event, cuda_ctx_src->stream())); + { + CUDA_CHECK(cudaEventRecord( + cuda_ctx_src->copy_event, cuda_ctx_src->stream())); + } // wait on dst stream for the copy to complete - CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx_dst->stream(), cuda_ctx_src->copy_event, 0)); + CUDA_CHECK(cudaStreamWaitEvent( + cuda_ctx_dst->stream(), cuda_ctx_src->copy_event, +#if defined(GGML_USE_HIP) + ggml_cuda_whole_graph_capture_depth > 0 + ? hipEventWaitExternal : 0 +#else + 0 +#endif + )); } else { // src and dst are on the same backend CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); @@ -3289,15 +3595,44 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx } } - if (res || memcmp(&graph->node_props[i], &prop, sizeof(prop)) != 0) { - // [TAG_FUSED_LOOP] with GGML_CUDA_GRAPH_STATS, name the first node - // whose properties changed — the reason a graph re-captures. - static const bool dbg = getenv("GGML_CUDA_GRAPH_STATS") != nullptr; - if (dbg && !res) { - GGML_LOG_INFO("[graph-mismatch] key=%p node=%d op=%s name=%s\n", - graph_key, i, ggml_op_name(cgraph->nodes[i]->op), - cgraph->nodes[i]->name); - } + const bool prop_changed = + memcmp(&graph->node_props[i], &prop, sizeof(prop)) != 0; + if (!res && prop_changed && getenv("GGML_CUDA_GRAPH_STATS")) { + GGML_LOG_INFO("[graph-mismatch] key=%p node=%d op=%s name=%s\n", + graph_key, i, ggml_op_name(cgraph->nodes[i]->op), + cgraph->nodes[i]->name); + } + if (!res && prop_changed && getenv("GGML_CUDA_GRAPH_DIFF_STATS")) { + const ggml_cuda_graph::node_properties & old = + graph->node_props[i]; + ggml_tensor old_node = old.node; + ggml_tensor new_node = prop.node; + const bool extra_changed = old_node.extra != new_node.extra; + old_node.extra = nullptr; + new_node.extra = nullptr; + const bool node_without_extra_changed = + memcmp(&old_node, &new_node, sizeof(ggml_tensor)) != 0; + const bool src_data_changed = + memcmp(old.node_src_data_ptrs, prop.node_src_data_ptrs, + sizeof(prop.node_src_data_ptrs)) != 0; + const bool src_ne_changed = + memcmp(old.node_src_ne, prop.node_src_ne, + sizeof(prop.node_src_ne)) != 0; + const bool src_nb_changed = + memcmp(old.node_src_nb, prop.node_src_nb, + sizeof(prop.node_src_nb)) != 0; + GGML_LOG_INFO( + "[cuda-graph-diff] key=%p node=%d/%d name=%s op=%s " + "extra=%d node_no_extra=%d src_data=%d src_ne=%d src_nb=%d " + "old_extra=%p new_extra=%p old_data=%p new_data=%p\n", + graph_key, i, cgraph->n_nodes, prop.node.name, + ggml_op_name(prop.node.op), (int) extra_changed, + (int) node_without_extra_changed, (int) src_data_changed, + (int) src_ne_changed, (int) src_nb_changed, + old.node.extra, prop.node.extra, + old.node.data, prop.node.data); + } + if (res || prop_changed) { graph->node_props[i] = prop; res = true; } @@ -4370,7 +4705,10 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - if (graph->is_enabled()) { + if (!ggml_cuda_graphs_disabled_override && + ggml_cuda_whole_graph_capture_depth == 0 && + !ggml_cuda_graphs_disabled_for_device(cuda_ctx->device) && + graph->is_enabled()) { const bool graph_compatible = ggml_cuda_graph_check_compability(cgraph); if (graph_compatible) { const bool can_skip_props_check = ggml_cuda_skip_props_check @@ -4447,14 +4785,36 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, static void ggml_backend_cuda_event_record(ggml_backend_t backend, ggml_backend_event_t event) { ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; - CUDA_CHECK(cudaEventRecord((cudaEvent_t)event->context, cuda_ctx->stream())); +#if defined(GGML_USE_HIP) + if (ggml_cuda_whole_graph_capture_depth > 0) { + const auto doorbell = ggml_cuda_whole_graph_event_signal( + event->context, cuda_ctx->device, true); + GGML_ASSERT(doorbell.signal); + ggml_cuda_whole_graph_publish_async( + (hipStream_t) cuda_ctx->stream(), doorbell); + return; + } +#endif + CUDA_CHECK(cudaEventRecord( + (cudaEvent_t) event->context, cuda_ctx->stream())); } static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; if (ggml_backend_is_cuda(backend)) { - CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx->stream(), (cudaEvent_t)event->context, 0)); +#if defined(GGML_USE_HIP) + if (ggml_cuda_whole_graph_capture_depth > 0) { + const auto doorbell = ggml_cuda_whole_graph_event_signal( + event->context, cuda_ctx->device, false); + GGML_ASSERT(doorbell.signal); + ggml_cuda_whole_graph_wait_async( + (hipStream_t) cuda_ctx->stream(), doorbell); + return; + } +#endif + CUDA_CHECK(cudaStreamWaitEvent( + cuda_ctx->stream(), (cudaEvent_t) event->context, 0)); } else { #if 0 // untested @@ -4741,6 +5101,871 @@ bool ggml_backend_is_cuda(ggml_backend_t backend) { return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_cuda_guid()); } +struct ggml_cuda_whole_graph_handle { + cudaGraph_t graph = nullptr; + cudaGraphExec_t executable = nullptr; + ggml_cuda_whole_graph_session * session = nullptr; + int device = -1; +#if defined(GGML_USE_HIP) + hipCtx_t context = nullptr; + struct memory_operation { + hipGraphNode_t captured_node = nullptr; + hipGraphNode_t fence_node = nullptr; + hipGraphNode_t node = nullptr; + hipStreamBatchMemOpParams operation{}; + }; + std::vector memory_operations; +#endif +}; + +#if defined(GGML_USE_HIP) +// ROCm accepts hipStreamWaitValue32/hipStreamWriteValue32 during stream +// capture but does not add them to the captured graph. Compute-kernel waits +// are visible to capture, but a resident spin wave has no forward-progress +// guarantee when two full-occupancy device graphs wait on one another. Keep +// the captured kernels only as placement markers, then replace every marker +// with an explicit command-processor batch-memory node before instantiation. +static bool ggml_cuda_whole_graph_replace_sync_kernels( + ggml_cuda_whole_graph_handle * handle) { + ggml_cuda_set_device(handle->device); + hipError_t status = hipCtxGetCurrent(&handle->context); + if (status != hipSuccess || handle->context == nullptr) { + GGML_LOG_WARN( + "%s: device=%d context lookup failed: %s\n", __func__, + handle->device, hipGetErrorString(status)); + return false; + } + + size_t node_count = 0; + status = hipGraphGetNodes(handle->graph, nullptr, &node_count); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d node count failed: %s\n", __func__, + handle->device, hipGetErrorString(status)); + return false; + } + std::vector nodes(node_count); + status = hipGraphGetNodes(handle->graph, nodes.data(), &node_count); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d node enumeration failed: %s\n", __func__, + handle->device, hipGetErrorString(status)); + return false; + } + + size_t waits = 0; + size_t writes = 0; + for (hipGraphNode_t node : nodes) { + hipGraphNodeType type{}; + if (hipGraphNodeGetType(node, &type) != hipSuccess || + type != hipGraphNodeTypeKernel) { + continue; + } + hipKernelNodeParams kernel{}; + if (hipGraphKernelNodeGetParams(node, &kernel) != hipSuccess || + kernel.kernelParams == nullptr || kernel.kernelParams[0] == nullptr) { + continue; + } + const bool is_write = + kernel.func == reinterpret_cast(ggml_cuda_whole_graph_publish); + const bool is_wait = + kernel.func == reinterpret_cast(ggml_cuda_whole_graph_wait); + if (!is_write && !is_wait) continue; + + auto * signal = + *reinterpret_cast(kernel.kernelParams[0]); + if (!signal) { + GGML_LOG_WARN("%s: device=%d marker has null signal\n", __func__, + handle->device); + return false; + } + ggml_cuda_whole_graph_handle::memory_operation replacement{}; + replacement.captured_node = node; + if (is_write) { + replacement.operation.writeValue.operation = + hipStreamMemOpWriteValue32; + replacement.operation.writeValue.address = + reinterpret_cast(signal); + replacement.operation.writeValue.value = 1u; + replacement.operation.writeValue.flags = 0; + ++writes; + } else { + replacement.operation.waitValue.operation = + hipStreamMemOpWaitValue32; + replacement.operation.waitValue.address = + reinterpret_cast(signal); + replacement.operation.waitValue.value = 1u; + replacement.operation.waitValue.flags = hipStreamWaitValueGte; + ++waits; + } + handle->memory_operations.push_back(replacement); + } + if (handle->memory_operations.empty()) { + GGML_LOG_WARN("%s: device=%d found no captured sync markers\n", __func__, + handle->device); + return false; + } + + // Create all replacements first. Rewire in a second pass so adjacent + // marker nodes are handled correctly regardless of graph enumeration order. + std::unordered_map replacement_index; + replacement_index.reserve(handle->memory_operations.size()); + for (size_t i = 0; i < handle->memory_operations.size(); ++i) { + auto & replacement = handle->memory_operations[i]; + hipBatchMemOpNodeParams params{}; + params.ctx = handle->context; + params.count = 1; + params.paramArray = &replacement.operation; + params.flags = 0; + status = hipGraphAddBatchMemOpNode( + &replacement.node, handle->graph, nullptr, 0, ¶ms); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d add node failed: %s\n", __func__, + handle->device, hipGetErrorString(status)); + return false; + } + replacement_index.emplace(replacement.captured_node, i); + } + + std::vector from; + std::vector to; + for (const auto & replacement : handle->memory_operations) { + size_t dependency_count = 0; + status = hipGraphNodeGetDependencies( + replacement.captured_node, nullptr, &dependency_count); + if (status != hipSuccess) return false; + std::vector dependencies(dependency_count); + if (dependency_count > 0) { + status = hipGraphNodeGetDependencies( + replacement.captured_node, dependencies.data(), + &dependency_count); + if (status != hipSuccess) return false; + } + for (hipGraphNode_t dependency : dependencies) { + const auto found = replacement_index.find(dependency); + from.push_back(found == replacement_index.end() + ? dependency + : handle->memory_operations[found->second].node); + to.push_back(replacement.node); + } + + size_t dependent_count = 0; + status = hipGraphNodeGetDependentNodes( + replacement.captured_node, nullptr, &dependent_count); + if (status != hipSuccess) return false; + std::vector dependents(dependent_count); + if (dependent_count > 0) { + status = hipGraphNodeGetDependentNodes( + replacement.captured_node, dependents.data(), &dependent_count); + if (status != hipSuccess) return false; + } + for (hipGraphNode_t dependent : dependents) { + // Marker-to-marker edges are added once through the downstream + // marker's dependency list. + if (replacement_index.find(dependent) != replacement_index.end()) { + continue; + } + from.push_back(replacement.node); + to.push_back(dependent); + } + } + if (!from.empty()) { + status = hipGraphAddDependencies( + handle->graph, from.data(), to.data(), from.size()); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d rewire failed: %s\n", __func__, + handle->device, hipGetErrorString(status)); + return false; + } + } + for (const auto & replacement : handle->memory_operations) { + status = hipGraphDestroyNode(replacement.captured_node); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d marker removal failed: %s\n", __func__, + handle->device, hipGetErrorString(status)); + return false; + } + } + GGML_LOG_INFO( + "[ds4-whole-graph] device=%d hardware sync writes=%zu waits=%zu\n", + handle->device, writes, waits); + return true; +} + +static bool ggml_cuda_whole_graph_replace_sync_kernels_v2( + ggml_cuda_whole_graph_handle * handle) { + ggml_cuda_set_device(handle->device); + hipError_t status = hipCtxGetCurrent(&handle->context); + if (status != hipSuccess || handle->context == nullptr) { + GGML_LOG_WARN( + "%s: device=%d context lookup failed: %s\n", __func__, + handle->device, hipGetErrorString(status)); + return false; + } + + size_t node_count = 0; + status = hipGraphGetNodes(handle->graph, nullptr, &node_count); + if (status != hipSuccess) return false; + std::vector nodes(node_count); + status = hipGraphGetNodes(handle->graph, nodes.data(), &node_count); + if (status != hipSuccess) return false; + + handle->memory_operations.clear(); + std::unordered_map marker_index; + size_t waits = 0; + size_t writes = 0; + for (hipGraphNode_t node : nodes) { + hipGraphNodeType type{}; + if (hipGraphNodeGetType(node, &type) != hipSuccess || + type != hipGraphNodeTypeKernel) { + continue; + } + hipKernelNodeParams kernel{}; + if (hipGraphKernelNodeGetParams(node, &kernel) != hipSuccess || + kernel.kernelParams == nullptr || kernel.kernelParams[0] == nullptr) { + continue; + } + const bool is_write = + kernel.func == reinterpret_cast(ggml_cuda_whole_graph_publish); + const bool is_wait = + kernel.func == reinterpret_cast(ggml_cuda_whole_graph_wait); + if (!is_write && !is_wait) continue; + + auto * signal = + *reinterpret_cast(kernel.kernelParams[0]); + if (!signal) return false; + ggml_cuda_whole_graph_handle::memory_operation operation{}; + operation.captured_node = node; + if (is_write) { + operation.operation.writeValue.operation = + hipStreamMemOpWriteValue32; + operation.operation.writeValue.address = + reinterpret_cast(signal); + operation.operation.writeValue.value = 1u; + operation.operation.writeValue.flags = 0; + kernel.func = reinterpret_cast( + ggml_cuda_whole_graph_release_fence); + ++writes; + } else { + operation.operation.waitValue.operation = + hipStreamMemOpWaitValue32; + operation.operation.waitValue.address = + reinterpret_cast(signal); + operation.operation.waitValue.value = 1u; + operation.operation.waitValue.flags = hipStreamWaitValueGte; + kernel.func = reinterpret_cast( + ggml_cuda_whole_graph_acquire_fence); + ++waits; + } + status = hipGraphKernelNodeSetParams(node, &kernel); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d fence mutation failed: %s\n", + __func__, handle->device, + hipGetErrorString(status)); + return false; + } + const size_t index = handle->memory_operations.size(); + handle->memory_operations.push_back(operation); + marker_index.emplace(node, index); + } + if (handle->memory_operations.empty()) { + GGML_LOG_WARN("%s: device=%d found no sync markers\n", __func__, + handle->device); + return false; + } + + // Snapshot the original edge set before adding command-processor nodes. + size_t edge_count = 0; + status = hipGraphGetEdges(handle->graph, nullptr, nullptr, &edge_count); + if (status != hipSuccess) return false; + std::vector edge_from(edge_count); + std::vector edge_to(edge_count); + status = hipGraphGetEdges( + handle->graph, edge_from.data(), edge_to.data(), &edge_count); + if (status != hipSuccess) return false; + + for (auto & operation : handle->memory_operations) { + hipBatchMemOpNodeParams params{}; + params.ctx = handle->context; + params.count = 1; + params.paramArray = &operation.operation; + params.flags = 0; + status = hipGraphAddBatchMemOpNode( + &operation.node, handle->graph, nullptr, 0, ¶ms); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d add node failed: %s\n", __func__, + handle->device, hipGetErrorString(status)); + return false; + } + } + + const auto entry_node = [&](hipGraphNode_t node) { + const auto found = marker_index.find(node); + if (found == marker_index.end()) return node; + const auto & operation = handle->memory_operations[found->second]; + return operation.operation.operation == hipStreamMemOpWaitValue32 + ? operation.node : node; + }; + const auto exit_node = [&](hipGraphNode_t node) { + const auto found = marker_index.find(node); + if (found == marker_index.end()) return node; + const auto & operation = handle->memory_operations[found->second]; + return operation.operation.operation == hipStreamMemOpWriteValue32 + ? operation.node : node; + }; + + std::vector remove_from; + std::vector remove_to; + std::vector add_from; + std::vector add_to; + for (size_t i = 0; i < edge_count; ++i) { + if (marker_index.find(edge_from[i]) == marker_index.end() && + marker_index.find(edge_to[i]) == marker_index.end()) { + continue; + } + remove_from.push_back(edge_from[i]); + remove_to.push_back(edge_to[i]); + add_from.push_back(exit_node(edge_from[i])); + add_to.push_back(entry_node(edge_to[i])); + } + for (const auto & operation : handle->memory_operations) { + if (operation.operation.operation == hipStreamMemOpWriteValue32) { + add_from.push_back(operation.captured_node); + add_to.push_back(operation.node); + } else { + add_from.push_back(operation.node); + add_to.push_back(operation.captured_node); + } + } + if (!remove_from.empty()) { + status = hipGraphRemoveDependencies( + handle->graph, remove_from.data(), remove_to.data(), + remove_from.size()); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d edge removal failed: %s\n", __func__, + handle->device, hipGetErrorString(status)); + return false; + } + } + status = hipGraphAddDependencies( + handle->graph, add_from.data(), add_to.data(), add_from.size()); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d edge insertion failed: %s\n", __func__, + handle->device, hipGetErrorString(status)); + return false; + } + + GGML_LOG_INFO( + "[ds4-whole-graph] device=%d hardware sync v2 writes=%zu waits=%zu " + "fence_kernels=%zu\n", + handle->device, writes, waits, writes + waits); + return true; +} + +static bool ggml_cuda_whole_graph_replace_sync_kernels_v3( + ggml_cuda_whole_graph_handle * handle) { + ggml_cuda_set_device(handle->device); + hipError_t status = hipCtxGetCurrent(&handle->context); + if (status != hipSuccess || handle->context == nullptr) return false; + + size_t node_count = 0; + status = hipGraphGetNodes(handle->graph, nullptr, &node_count); + if (status != hipSuccess) return false; + std::vector nodes(node_count); + status = hipGraphGetNodes(handle->graph, nodes.data(), &node_count); + if (status != hipSuccess) return false; + + handle->memory_operations.clear(); + std::unordered_map marker_index; + size_t waits = 0; + size_t writes = 0; + for (hipGraphNode_t node : nodes) { + hipGraphNodeType type{}; + if (hipGraphNodeGetType(node, &type) != hipSuccess || + type != hipGraphNodeTypeKernel) { + continue; + } + hipKernelNodeParams kernel{}; + if (hipGraphKernelNodeGetParams(node, &kernel) != hipSuccess || + kernel.kernelParams == nullptr || kernel.kernelParams[0] == nullptr) { + continue; + } + const bool is_write = + kernel.func == reinterpret_cast(ggml_cuda_whole_graph_publish); + const bool is_wait = + kernel.func == reinterpret_cast(ggml_cuda_whole_graph_wait); + if (!is_write && !is_wait) continue; + auto * signal = + *reinterpret_cast(kernel.kernelParams[0]); + if (!signal) return false; + + ggml_cuda_whole_graph_handle::memory_operation operation{}; + operation.captured_node = node; + if (is_write) { + operation.operation.writeValue.operation = + hipStreamMemOpWriteValue32; + operation.operation.writeValue.address = + reinterpret_cast(signal); + operation.operation.writeValue.value = 1u; + operation.operation.writeValue.flags = 0; + ++writes; + } else { + operation.operation.waitValue.operation = + hipStreamMemOpWaitValue32; + operation.operation.waitValue.address = + reinterpret_cast(signal); + operation.operation.waitValue.value = 1u; + operation.operation.waitValue.flags = hipStreamWaitValueGte; + ++waits; + } + const size_t index = handle->memory_operations.size(); + handle->memory_operations.push_back(operation); + marker_index.emplace(node, index); + } + if (handle->memory_operations.empty()) return false; + + size_t edge_count = 0; + status = hipGraphGetEdges(handle->graph, nullptr, nullptr, &edge_count); + if (status != hipSuccess) return false; + std::vector edge_from(edge_count); + std::vector edge_to(edge_count); + status = hipGraphGetEdges( + handle->graph, edge_from.data(), edge_to.data(), &edge_count); + if (status != hipSuccess) return false; + + for (auto & operation : handle->memory_operations) { + hipKernelNodeParams fence{}; + status = hipGraphKernelNodeGetParams(operation.captured_node, &fence); + if (status != hipSuccess) return false; + fence.func = operation.operation.operation == hipStreamMemOpWriteValue32 + ? reinterpret_cast(ggml_cuda_whole_graph_release_fence) + : reinterpret_cast(ggml_cuda_whole_graph_acquire_fence); + status = hipGraphAddKernelNode( + &operation.fence_node, handle->graph, nullptr, 0, &fence); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d add fence failed: %s\n", __func__, + handle->device, hipGetErrorString(status)); + return false; + } + + hipBatchMemOpNodeParams params{}; + params.ctx = handle->context; + params.count = 1; + params.paramArray = &operation.operation; + params.flags = 0; + status = hipGraphAddBatchMemOpNode( + &operation.node, handle->graph, nullptr, 0, ¶ms); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d add memory op failed: %s\n", __func__, + handle->device, hipGetErrorString(status)); + return false; + } + } + + const auto entry_node = [&](hipGraphNode_t node) { + const auto found = marker_index.find(node); + if (found == marker_index.end()) return node; + const auto & operation = handle->memory_operations[found->second]; + return operation.operation.operation == hipStreamMemOpWaitValue32 + ? operation.node : operation.fence_node; + }; + const auto exit_node = [&](hipGraphNode_t node) { + const auto found = marker_index.find(node); + if (found == marker_index.end()) return node; + const auto & operation = handle->memory_operations[found->second]; + return operation.operation.operation == hipStreamMemOpWriteValue32 + ? operation.node : operation.fence_node; + }; + + std::vector remove_from; + std::vector remove_to; + std::vector add_from; + std::vector add_to; + for (size_t i = 0; i < edge_count; ++i) { + if (marker_index.find(edge_from[i]) == marker_index.end() && + marker_index.find(edge_to[i]) == marker_index.end()) { + continue; + } + remove_from.push_back(edge_from[i]); + remove_to.push_back(edge_to[i]); + add_from.push_back(exit_node(edge_from[i])); + add_to.push_back(entry_node(edge_to[i])); + } + for (const auto & operation : handle->memory_operations) { + if (operation.operation.operation == hipStreamMemOpWriteValue32) { + add_from.push_back(operation.fence_node); + add_to.push_back(operation.node); + } else { + add_from.push_back(operation.node); + add_to.push_back(operation.fence_node); + } + } + if (!remove_from.empty()) { + status = hipGraphRemoveDependencies( + handle->graph, remove_from.data(), remove_to.data(), + remove_from.size()); + if (status != hipSuccess) return false; + } + status = hipGraphAddDependencies( + handle->graph, add_from.data(), add_to.data(), add_from.size()); + if (status != hipSuccess) return false; + for (const auto & operation : handle->memory_operations) { + status = hipGraphDestroyNode(operation.captured_node); + if (status != hipSuccess) return false; + } + + GGML_LOG_INFO( + "[ds4-whole-graph] device=%d hardware sync v3 writes=%zu waits=%zu " + "new_fences=%zu\n", + handle->device, writes, waits, writes + waits); + return true; +} + +static bool ggml_cuda_whole_graph_reset_round( + ggml_cuda_whole_graph_session * session) { + if (!session->hardware_sync) return true; + // The two graphs have both been synchronized before the scheduler starts + // another decode step. Reset every one-shot doorbell to zero with one + // local kernel per device, then keep all graph memory-op nodes fixed at + // value 1. ROCm 7.2.4 accepts graph-exec batch-memory parameter updates + // but can fault the second production replay after hundreds of such + // mutations; the API is still explicitly beta. Fixed graph nodes plus a + // bounded pre-launch reset avoid both that runtime path and ABA. + for (size_t i = 0; i < session->devices.size(); ++i) { + const uint32_t count = session->next_signals[i]; + if (count == 0) continue; + ggml_cuda_set_device(session->devices[i]); + const uint32_t block = 256; + const uint32_t grid = (count + block - 1) / block; + ggml_cuda_whole_graph_reset_signals<<reset_streams[i]>>>(session->signal_tables[i], count); + const hipError_t status = hipGetLastError(); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d reset launch failed: %s\n", __func__, + session->devices[i], hipGetErrorString(status)); + return false; + } + } + for (size_t i = 0; i < session->devices.size(); ++i) { + if (session->next_signals[i] == 0) continue; + ggml_cuda_set_device(session->devices[i]); + const hipError_t status = hipStreamSynchronize( + session->reset_streams[i]); + if (status != hipSuccess) { + GGML_LOG_WARN("%s: device=%d reset sync failed: %s\n", __func__, + session->devices[i], hipGetErrorString(status)); + return false; + } + } + return true; +} +#endif + +static void ggml_cuda_whole_graph_session_destroy( + ggml_cuda_whole_graph_session * session) { + if (!session) return; +#if defined(GGML_USE_HIP) + for (size_t i = 0; i < session->signals.size(); ++i) { + ggml_cuda_set_device(session->devices[i]); + if (session->reset_streams[i]) { + (void) hipStreamDestroy(session->reset_streams[i]); + } + if (session->signal_tables[i]) { + (void) hipFree(session->signal_tables[i]); + } + for (uint32_t * signal : session->signals[i]) { + if (signal) (void) hipFree(signal); + } + if (session->observed[i]) { + ggml_cuda_set_device(session->devices[1 - i]); + (void) hipFree(session->observed[i]); + } + } + if (session->staging_host) { + (void) hipHostFree(session->staging_host); + } +#endif + delete session; +} + +bool ggml_backend_cuda_whole_graph_capture_prepare( + ggml_backend_t backend0, ggml_backend_t backend1) { +#if defined(USE_CUDA_GRAPH) && defined(GGML_USE_HIP) + if (!ggml_backend_is_cuda(backend0) || + !ggml_backend_is_cuda(backend1) || backend0 == backend1 || + ggml_cuda_whole_graph_capture_depth != 0 || + ggml_cuda_whole_graph_capture_session != nullptr) { + return false; + } + auto * ctx0 = (ggml_backend_cuda_context *) backend0->context; + auto * ctx1 = (ggml_backend_cuda_context *) backend1->context; + if (ctx0->device == ctx1->device) return false; + + auto * session = new ggml_cuda_whole_graph_session(); + session->devices = {ctx0->device, ctx1->device}; + const char * hardware_sync = + std::getenv("DFLASH_DS4_TP_WHOLE_STEP_HW_SYNC"); + session->hardware_sync = !hardware_sync || !*hardware_sync || + std::strcmp(hardware_sync, "0") != 0; + ggml_cuda_set_device(session->devices[0]); + CUDA_CHECK(hipHostMalloc( + &session->staging_host, GGML_CUDA_WHOLE_GRAPH_STAGING_BYTES, + hipHostMallocMapped | hipHostMallocCoherent)); + CUDA_CHECK(hipHostGetDevicePointer( + (void **) &session->staging, session->staging_host, 0)); + for (size_t i = 0; i < session->devices.size(); ++i) { + ggml_cuda_set_device(session->devices[i]); + CUDA_CHECK(hipStreamCreateWithFlags( + &session->reset_streams[i], hipStreamNonBlocking)); + for (uint32_t *& signal : session->signals[i]) { + CUDA_CHECK(hipExtMallocWithFlags( + (void **) &signal, sizeof(uint64_t), hipMallocSignalMemory)); + CUDA_CHECK(hipMemset(signal, 0, sizeof(uint64_t))); + } + CUDA_CHECK(hipMalloc( + (void **) &session->signal_tables[i], + GGML_CUDA_WHOLE_GRAPH_MAX_DOORBELLS * sizeof(uint32_t *))); + CUDA_CHECK(hipMemcpy( + session->signal_tables[i], session->signals[i].data(), + GGML_CUDA_WHOLE_GRAPH_MAX_DOORBELLS * sizeof(uint32_t *), + hipMemcpyHostToDevice)); + CUDA_CHECK(hipDeviceSynchronize()); + ggml_cuda_set_device(session->devices[1 - i]); + CUDA_CHECK(hipMalloc( + (void **) &session->observed[i], + GGML_CUDA_WHOLE_GRAPH_MAX_DOORBELLS * sizeof(uint32_t))); + CUDA_CHECK(hipMemset( + session->observed[i], 0, + GGML_CUDA_WHOLE_GRAPH_MAX_DOORBELLS * sizeof(uint32_t))); + CUDA_CHECK(hipDeviceSynchronize()); + } + ggml_cuda_whole_graph_capture_session = session; + return true; +#else + GGML_UNUSED(backend0); + GGML_UNUSED(backend1); + return false; +#endif +} + +bool ggml_backend_cuda_whole_graph_capture_begin(ggml_backend_t backend) { +#ifdef USE_CUDA_GRAPH + if (!ggml_backend_is_cuda(backend) || + !ggml_cuda_whole_graph_capture_session) return false; + ggml_backend_cuda_context * ctx = + (ggml_backend_cuda_context *) backend->context; + const auto * session = ggml_cuda_whole_graph_capture_session; + if (ctx->device != session->devices[0] && + ctx->device != session->devices[1]) return false; + ggml_cuda_set_device(ctx->device); + const cudaError_t status = cudaStreamBeginCapture( + ctx->stream(), cudaStreamCaptureModeRelaxed); + if (status != cudaSuccess) { + GGML_LOG_WARN("%s: device=%d begin failed: %s\n", __func__, + ctx->device, cudaGetErrorString(status)); + if (ggml_cuda_whole_graph_capture_depth == 0 && + ggml_cuda_whole_graph_capture_session->references == 0) { + ggml_cuda_whole_graph_session_destroy( + ggml_cuda_whole_graph_capture_session); + ggml_cuda_whole_graph_capture_session = nullptr; + } + return false; + } + ++ggml_cuda_whole_graph_capture_depth; + return true; +#else + GGML_UNUSED(backend); + return false; +#endif +} + +ggml_backend_cuda_whole_graph_t +ggml_backend_cuda_whole_graph_capture_end(ggml_backend_t backend) { +#ifdef USE_CUDA_GRAPH + if (!ggml_backend_is_cuda(backend) || + ggml_cuda_whole_graph_capture_depth <= 0) { + return nullptr; + } + ggml_backend_cuda_context * ctx = + (ggml_backend_cuda_context *) backend->context; + auto * session = ggml_cuda_whole_graph_capture_session; + ggml_cuda_set_device(ctx->device); + auto * result = new ggml_cuda_whole_graph_handle(); + const cudaError_t end_status = + cudaStreamEndCapture(ctx->stream(), &result->graph); + --ggml_cuda_whole_graph_capture_depth; + const bool capture_finished = + ggml_cuda_whole_graph_capture_depth == 0; + if (capture_finished) { + ggml_cuda_whole_graph_capture_session = nullptr; + } + if (end_status != cudaSuccess || result->graph == nullptr) { + GGML_LOG_WARN("%s: device=%d end failed: %s\n", __func__, + ctx->device, cudaGetErrorString(end_status)); + if (result->graph) (void) cudaGraphDestroy(result->graph); + delete result; + if (capture_finished && session && session->references == 0) { + ggml_cuda_whole_graph_session_destroy(session); + } + return nullptr; + } + // Instantiate lazily after both device streams have ended capture. Some + // HIP runtimes cannot instantiate one graph while its peer is capturing. + result->session = session; + result->device = ctx->device; + ++session->references; + ++session->graph_count; + session->pending_handles.push_back(result); + if (capture_finished) { + GGML_LOG_INFO( + "[ds4-whole-graph] captured doorbells dev%d=%u dev%d=%u " + "events=%zu staging=%.2f MiB\n", + session->devices[0], session->next_signals[0], + session->devices[1], session->next_signals[1], + session->event_signals.size(), + session->staging_used / (1024.0 * 1024.0)); + // Both streams are now out of capture. Instantiate both executables + // before either waiter graph can be launched; launching one first can + // surface an asynchronous wait failure while HIP builds its peer. + for (auto * pending : session->pending_handles) { + ggml_cuda_set_device(pending->device); +#if defined(GGML_USE_HIP) + if (session->hardware_sync && + !ggml_cuda_whole_graph_replace_sync_kernels_v3(pending)) { + session->instantiate_failed = true; + break; + } +#endif + const cudaError_t instantiate_status = cudaGraphInstantiate( + &pending->executable, pending->graph, nullptr, nullptr, 0); + if (instantiate_status != cudaSuccess) { + GGML_LOG_WARN( + "%s: device=%d instantiate failed: %s\n", __func__, + pending->device, + cudaGetErrorString(instantiate_status)); + session->instantiate_failed = true; + break; + } + } + if (!session->hardware_sync) { + GGML_LOG_INFO( + "[ds4-whole-graph] native system-atomic sequence sync\n"); + } + } + return result; +#else + GGML_UNUSED(backend); + return nullptr; +#endif +} + +bool ggml_backend_cuda_whole_graph_launch( + ggml_backend_t backend, + ggml_backend_cuda_whole_graph_t opaque_graph) { +#ifdef USE_CUDA_GRAPH + if (!ggml_backend_is_cuda(backend) || !opaque_graph) return false; + ggml_backend_cuda_context * ctx = + (ggml_backend_cuda_context *) backend->context; + auto * graph = (ggml_cuda_whole_graph_handle *) opaque_graph; + auto * session = graph->session; + if (session && session->instantiate_failed) return false; + +#if defined(GGML_USE_HIP) + if (session && session->launches_in_round == 0 && + !ggml_cuda_whole_graph_reset_round(session)) { + return false; + } +#endif + ggml_cuda_set_device(ctx->device); + if (!graph->executable) { + const cudaError_t instantiate_status = cudaGraphInstantiate( + &graph->executable, graph->graph, nullptr, nullptr, 0); + if (instantiate_status != cudaSuccess) { + GGML_LOG_WARN("%s: device=%d instantiate failed: %s\n", __func__, + ctx->device, + cudaGetErrorString(instantiate_status)); + return false; + } + } + const cudaError_t status = + cudaGraphLaunch(graph->executable, ctx->stream()); + if (status != cudaSuccess) { + GGML_LOG_WARN("%s: device=%d launch failed: %s\n", __func__, + ctx->device, cudaGetErrorString(status)); + return false; + } + if (session) { + ++session->launches_in_round; + if (session->launches_in_round == session->graph_count) { + session->launches_in_round = 0; + } + } + return true; +#else + GGML_UNUSED(backend); + GGML_UNUSED(opaque_graph); + return false; +#endif +} + +void ggml_backend_cuda_whole_graph_free( + ggml_backend_t backend, + ggml_backend_cuda_whole_graph_t opaque_graph) { +#ifdef USE_CUDA_GRAPH + if (!opaque_graph) return; + auto * graph = (ggml_cuda_whole_graph_handle *) opaque_graph; + if (ggml_backend_is_cuda(backend)) { + ggml_backend_cuda_context * ctx = + (ggml_backend_cuda_context *) backend->context; + ggml_cuda_set_device(ctx->device); + } + if (graph->executable) (void) cudaGraphExecDestroy(graph->executable); + if (graph->graph) (void) cudaGraphDestroy(graph->graph); + auto * session = graph->session; + delete graph; + if (session && --session->references == 0) { + ggml_cuda_whole_graph_session_destroy(session); + } +#else + GGML_UNUSED(backend); + GGML_UNUSED(opaque_graph); +#endif +} + +bool ggml_backend_cuda_set_low_priority_stream(ggml_backend_t backend) { + if (!ggml_backend_is_cuda(backend)) { + return false; + } + ggml_backend_cuda_context * ctx = + (ggml_backend_cuda_context *) backend->context; + for (int device = 0; device < GGML_CUDA_MAX_DEVICES; ++device) { + for (int stream = 0; stream < GGML_CUDA_MAX_STREAMS; ++stream) { + if (ctx->streams[device][stream] != nullptr) { + GGML_LOG_WARN( + "%s: backend stream already exists; priority unchanged\n", + __func__); + return false; + } + } + } + + ggml_cuda_set_device(ctx->device); + int least_priority = 0; + int greatest_priority = 0; + const cudaError_t status = cudaDeviceGetStreamPriorityRange( + &least_priority, &greatest_priority); + if (status != cudaSuccess) { + GGML_LOG_WARN("%s: priority query failed: %s\n", __func__, + cudaGetErrorString(status)); + return false; + } + ctx->low_priority_streams = true; + ctx->stream_priority = least_priority; + GGML_LOG_INFO( + "%s: device=%d least=%d greatest=%d selected=%d\n", + __func__, ctx->device, least_priority, greatest_priority, + ctx->stream_priority); + return least_priority != greatest_priority; +} + int ggml_backend_cuda_get_device_count() { return ggml_cuda_info().device_count; } @@ -5230,9 +6455,21 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g } break; case GGML_OP_REPEAT: { - // the CUDA REPEAT path only implements F32/F16; other types assert at runtime ggml_type src0_type = op->src[0]->type; - return src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_F16; + if (src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_F16) { + return true; + } + // Exact contiguous I32 repeat is opt-in while it is qualified + // on the heterogeneous DS4 verifier path. + static const bool i32_repeat_enabled = [] { + const char * raw = std::getenv("LUCE_CUDA_I32_REPEAT"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return i32_repeat_enabled && + src0_type == GGML_TYPE_I32 && + op->type == GGML_TYPE_I32 && + ggml_is_contiguous(op->src[0]) && + ggml_is_contiguous(op); } break; case GGML_OP_REPEAT_BACK: return op->type == GGML_TYPE_F32 && (op->src[0]->ne[2]*op->src[0]->ne[3]) <= (1 << 15); 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 448ef93ec..bc2135848 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -4,6 +4,7 @@ #include "vecdotq.cuh" #include +#include #include typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs); @@ -72,6 +73,135 @@ static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) { } } +// MMVQ uses VDR=2 and supplies even iqs values {0, 2, 4, 6}. Each lane +// therefore consumes one disjoint 24-bit quarter of the 96-bit FP3 payload. +// Loading only those three contiguous bytes avoids reconstructing the full +// 12-byte block independently in four lanes. The two DP4A operations and the +// final two-scale expression remain in exactly the same order as the generic +// path in vecdotq.cuh. +static __device__ __forceinline__ float vec_dot_rocmfpx_fp3_q8_1_packed24( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, + const int & kbx, const int & iqs) { + + static_assert(VDR_ROCMFP3_Q8_1_MMVQ == 2, + "packed FP3 MMVQ path requires VDR=2"); + const block_rocmfp3 * bq3 = (const block_rocmfp3 *) vbq + kbx; + const int byte_offset = 3 * (iqs >> 1); + const uint32_t packed24 = + (uint32_t) bq3->qs[byte_offset + 0] | + ((uint32_t) bq3->qs[byte_offset + 1] << 8) | + ((uint32_t) bq3->qs[byte_offset + 2] << 16); + + const int val0 = rocmfpx_pack4_fp3_bits12_vec_cuda(packed24 & 0xFFFu); + const int val1 = rocmfpx_pack4_fp3_bits12_vec_cuda(packed24 >> 12); + const int u0 = get_int_b4(bq8_1->qs, iqs + 0); + const int u1 = get_int_b4(bq8_1->qs, iqs + 1); + + int sumi0 = 0; + int sumi1 = 0; + if (iqs < QK_ROCMFP3/8) { + sumi0 = ggml_cuda_dp4a(val0, u0, sumi0); + sumi0 = ggml_cuda_dp4a(val1, u1, sumi0); + } else { + sumi1 = ggml_cuda_dp4a(val0, u0, sumi1); + sumi1 = ggml_cuda_dp4a(val1, u1, sumi1); + } + + const float db = __low2float(bq8_1->ds); + return db * (rocmfpx_ue4m3_to_fp32_finite(bq3->e[0]) * sumi0 + + rocmfpx_ue4m3_to_fp32_finite(bq3->e[1]) * sumi1); +} + +// FP2 MMVQ assigns one lane to each 16-weight half-block. Load that lane's +// four packed bytes with one unaligned-safe memcpy instead of four independent +// byte loads. Byte extraction, DP4A order, scale conversion, and the final +// floating-point multiplication remain identical to vecdotq.cuh. +static __device__ __forceinline__ float vec_dot_rocmfpx_fp2_q8_1_packed32( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, + const int & kbx, const int & iqs) { + + static_assert(VDR_ROCMFP2_Q8_1_MMVQ == 1, + "packed FP2 MMVQ path requires VDR=1"); + const block_rocmfp2 * bq2 = (const block_rocmfp2 *) vbq + kbx; + uint32_t packed32; + memcpy(&packed32, bq2->qs + 4*iqs, sizeof(packed32)); + + int sumi = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint32_t bits8 = (packed32 >> (8*j)) & 0xFFu; + const int val_packed = rocmfpx_pack4_fp2_bits8_vec_cuda(bits8); + const int u = get_int_b4(bq8_1->qs, 4*iqs + j); + sumi = ggml_cuda_dp4a(val_packed, u, sumi); + } + + const float db = __low2float(bq8_1->ds); + return db * rocmfpx_ue4m3_to_fp32_finite(bq2->e[iqs]) * sumi; +} + +// Dense q4 verification applies one weight row to four token activations. +// Decode the ROCmFP4 payload once and retain one independent integer +// accumulator per token. Each token keeps the original DP4A and floating +// multiplication order, so this removes duplicate weight/codebook work +// without changing its result. +static __device__ __forceinline__ float4 vec_dot_rocmfp4_fast_q8_1_x4( + const void * __restrict__ vbq, + const block_q8_1 * __restrict__ bq8_0, + const block_q8_1 * __restrict__ bq8_1, + const block_q8_1 * __restrict__ bq8_2, + const block_q8_1 * __restrict__ bq8_3, + const int & kbx, const int & iqs) { + const block_rocmfp4_fast * bq4 = + (const block_rocmfp4_fast *) vbq + kbx; + const int * q80 = (const int *) bq8_0->qs + iqs; + const int * q81 = (const int *) bq8_1->qs + iqs; + const int * q82 = (const int *) bq8_2->qs + iqs; + const int * q83 = (const int *) bq8_3->qs + iqs; + + int sumi0 = 0; + int sumi1 = 0; + int sumi2 = 0; + int sumi3 = 0; +#pragma unroll + for (int l = 0; l < VDR_ROCMFP4_FAST_Q8_1_MMVQ; ++l) { + const int aux_q4 = rocmfp4_get_qs_i32(bq4->qs, iqs + l); + const int2 v = + rocmfp4_get_int_from_codebook_16(aux_q4, kvalues_rocmfp4); + sumi0 = ggml_cuda_dp4a(v.x, q80[l + 0], sumi0); + sumi0 = ggml_cuda_dp4a(v.y, q80[l + 4], sumi0); + sumi1 = ggml_cuda_dp4a(v.x, q81[l + 0], sumi1); + sumi1 = ggml_cuda_dp4a(v.y, q81[l + 4], sumi1); + sumi2 = ggml_cuda_dp4a(v.x, q82[l + 0], sumi2); + sumi2 = ggml_cuda_dp4a(v.y, q82[l + 4], sumi2); + sumi3 = ggml_cuda_dp4a(v.x, q83[l + 0], sumi3); + sumi3 = ggml_cuda_dp4a(v.y, q83[l + 4], sumi3); + } + + const float d = rocmfp4_ue4m3_to_fp32_half_finite(bq4->e); + return make_float4( + __low2float(bq8_0->ds) * d * sumi0, + __low2float(bq8_1->ds) * d * sumi1, + __low2float(bq8_2->ds) * d * sumi2, + __low2float(bq8_3->ds) * d * sumi3); +} + +template +static __device__ __forceinline__ float vec_dot_q_mmvq( + const void * __restrict__ vbq, + const block_q8_1 * __restrict__ bq8_1, + const int & kbx, const int & iqs) { + static_assert(!c_fp3_packed24 || type == GGML_TYPE_Q3_0_ROCMFPX, + "packed FP3 dispatch requires the ROCmFP3 type"); + if constexpr (c_fp3_packed24) { + return vec_dot_rocmfpx_fp3_q8_1_packed24( + vbq, bq8_1, kbx, iqs); + } else { + constexpr vec_dot_q_cuda_t vec_dot_q_cuda = + get_vec_dot_q_cuda(type); + return vec_dot_q_cuda(vbq, bq8_1, kbx, iqs); + } +} + enum mmvq_parameter_table_id { MMVQ_PARAMETERS_GENERIC = 0, MMVQ_PARAMETERS_GCN, @@ -501,8 +631,25 @@ static bool is_gfx1151(const int cc) { return cc == GGML_CUDA_CC_OFFSET_AMD + 0x1151; } +static bool rocmfp3_packed24_enabled() { + static const bool enabled = []() { + const char * value = std::getenv("DFLASH_CUDA_MMVQ_FP3_PACKED24"); + return value && value[0] == '1' && value[1] == '\0'; + }(); + return enabled; +} + +static bool rocmfp4_x4_enabled() { + static const bool enabled = []() { + const char * value = std::getenv("DFLASH_CUDA_MMVQ_FP4_X4"); + return value && value[0] == '1' && value[1] == '\0'; + }(); + return enabled; +} + template + int fixed_ncols_x = 0, bool unroll_k_loop_2 = false, + bool c_fp3_packed24 = false, bool c_fp4_x4 = false> __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id())*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q( const void * __restrict__ vx, const void * __restrict__ vy, const int32_t * __restrict__ ids, const ggml_cuda_mm_fusion_args_device fusion, float * __restrict__ dst, @@ -520,8 +667,6 @@ static __global__ void mul_mat_vec_q( constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); constexpr int warp_size = ggml_cuda_get_physical_warp_size(); - constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); - const int tid = warp_size*threadIdx.y + threadIdx.x; const int row0 = rows_per_cuda_block*blockIdx.x; const int blocks_per_row_x = (fixed_ncols_x > 0 ? fixed_ncols_x : ncols_x) / qk; @@ -533,6 +678,11 @@ static __global__ void mul_mat_vec_q( "fixed-K decode specialization is limited to profiled ROCmFP2/FP3 types"); static_assert(!unroll_k_loop_2 || type == GGML_TYPE_Q4_0_ROCMFP4_FAST, "partial K-loop unrolling is limited to the profiled 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 || + (type == GGML_TYPE_Q4_0_ROCMFP4_FAST && ncols_dst == 4), + "FP4 x4 MMVQ specialization requires ROCmFP4-fast q4"); const uint32_t channel_dst = blockIdx.y; @@ -541,13 +691,25 @@ static __global__ void mul_mat_vec_q( const uint32_t sample_dst = blockIdx.z; const uint32_t ids_col = ids_tokenwise_samples ? sample_dst : 0; - const uint32_t channel_x = has_ids ? ids[channel_dst + ids_col*ids_stride] : fastdiv(channel_dst, channel_ratio); + const int32_t channel_x_raw = has_ids ? ids[channel_dst + ids_col*ids_stride] : 0; + const bool channel_valid = !has_ids || channel_x_raw >= 0; + // Negative MUL_MAT_ID entries are masked routes. Clamp the address to a + // valid expert before doing any pointer arithmetic; channel_valid then + // suppresses every expert-weight/bias read and forces a zero result. + const uint32_t channel_x = has_ids ? (channel_valid ? (uint32_t) channel_x_raw : 0u) + : fastdiv(channel_dst, channel_ratio); const uint32_t channel_y = has_ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; uint32_t channel_xs[ncols_dst]; + bool channel_valids[ncols_dst]; #pragma unroll for (int j = 0; j < ncols_dst; ++j) { - channel_xs[j] = ids_multi_col ? ids[channel_dst + j*ids_stride] : channel_x; + const int32_t raw = ids_multi_col ? ids[channel_dst + j*ids_stride] : channel_x_raw; + channel_valids[j] = !has_ids || raw >= 0; + channel_xs[j] = channel_valids[j] ? (uint32_t) raw : 0u; + if (!has_ids) { + channel_xs[j] = channel_x; + } } const uint32_t sample_x = fastdiv(sample_dst, sample_ratio); @@ -562,6 +724,8 @@ static __global__ void mul_mat_vec_q( ggml_glu_op active_glu; float glu_param0 = 0.0f; float glu_param1 = 0.0f; + float gate_value_scale = 1.0f; + float x_value_scale = 1.0f; if constexpr (has_fusion) { use_gate = fusion.gate != nullptr; @@ -573,6 +737,8 @@ static __global__ void mul_mat_vec_q( active_glu = fusion.glu_op; glu_param0 = fusion.glu_param0; glu_param1 = fusion.glu_param1; + gate_value_scale = fusion.gate_value_scale; + x_value_scale = fusion.x_value_scale; } @@ -587,6 +753,9 @@ static __global__ void mul_mat_vec_q( (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) { #pragma unroll for (int j = 0; j < ncols_dst; ++j) { + if (!channel_valids[j]) { + continue; + } const uint32_t channel_bias = has_ids ? channel_xs[j] : channel_dst; const uint32_t col_offset = ids_multi_col ? 0 : j*stride_col_dst; x_biases[j] = x_bias_base[channel_bias*stride_channel_dst + col_offset + threadIdx.x]; @@ -599,6 +768,9 @@ static __global__ void mul_mat_vec_q( (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) { #pragma unroll for (int j = 0; j < ncols_dst; ++j) { + if (!channel_valids[j]) { + continue; + } const uint32_t channel_bias = has_ids ? channel_xs[j] : channel_dst; const uint32_t col_offset = ids_multi_col ? 0 : j*stride_col_dst; gate_biases[j] = gate_bias_base[channel_bias*stride_channel_dst + col_offset + threadIdx.x]; @@ -630,14 +802,17 @@ static __global__ void mul_mat_vec_q( #pragma unroll for (int j = 0; j < ncols_dst; ++j) { + if (!channel_valids[j]) { + continue; + } const int kbx_offset = sample_x*stride_sample_x + channel_xs[j]*stride_channel_x + row0*stride_row_x; #pragma unroll for (int i = 0; i < rows_per_cuda_block; ++i) { - tmp[j][i] += vec_dot_q_cuda( + tmp[j][i] += vec_dot_q_mmvq( vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); if constexpr (has_fusion) { if (use_gate) { - tmp_gate[j][i] += vec_dot_q_cuda( + tmp_gate[j][i] += vec_dot_q_mmvq( vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); } } @@ -652,14 +827,17 @@ static __global__ void mul_mat_vec_q( #pragma unroll for (int j = 0; j < ncols_dst; ++j) { + if (!channel_valids[j]) { + continue; + } const int kbx_offset = sample_x*stride_sample_x + channel_xs[j]*stride_channel_x + row0*stride_row_x; #pragma unroll for (int i = 0; i < rows_per_cuda_block; ++i) { - tmp[j][i] += vec_dot_q_cuda( + tmp[j][i] += vec_dot_q_mmvq( vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); if constexpr (has_fusion) { if (use_gate) { - tmp_gate[j][i] += vec_dot_q_cuda( + tmp_gate[j][i] += vec_dot_q_mmvq( vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); } } @@ -671,16 +849,56 @@ static __global__ void mul_mat_vec_q( const int kby = kbx * (qk/QK8_1); const int kqs = vdr * (tid % (qi/vdr)); + if constexpr (c_fp4_x4) { + const int kbx_offset = sample_x*stride_sample_x + + channel_xs[0]*stride_channel_x + row0*stride_row_x; +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + const float4 dots = vec_dot_rocmfp4_fast_q8_1_x4( + vx, + &y[0*stride_col_y + kby], + &y[1*stride_col_y + kby], + &y[2*stride_col_y + kby], + &y[3*stride_col_y + kby], + kbx_offset + i*stride_row_x + kbx, kqs); + tmp[0][i] += dots.x; + tmp[1][i] += dots.y; + tmp[2][i] += dots.z; + tmp[3][i] += dots.w; + if constexpr (has_fusion) { + if (use_gate) { + const float4 gate_dots = + vec_dot_rocmfp4_fast_q8_1_x4( + vgate, + &y[0*stride_col_y + kby], + &y[1*stride_col_y + kby], + &y[2*stride_col_y + kby], + &y[3*stride_col_y + kby], + kbx_offset + i*stride_row_x + kbx, + kqs); + tmp_gate[0][i] += gate_dots.x; + tmp_gate[1][i] += gate_dots.y; + tmp_gate[2][i] += gate_dots.z; + tmp_gate[3][i] += gate_dots.w; + } + } + } + continue; + } + #pragma unroll for (int j = 0; j < ncols_dst; ++j) { + if (!channel_valids[j]) { + continue; + } const int kbx_offset = sample_x*stride_sample_x + channel_xs[j]*stride_channel_x + row0*stride_row_x; #pragma unroll for (int i = 0; i < rows_per_cuda_block; ++i) { - tmp[j][i] += vec_dot_q_cuda( + tmp[j][i] += vec_dot_q_mmvq( vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); if constexpr (has_fusion) { if (use_gate) { - tmp_gate[j][i] += vec_dot_q_cuda( + tmp_gate[j][i] += vec_dot_q_mmvq( vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); } } @@ -754,34 +972,43 @@ static __global__ void mul_mat_vec_q( } if (threadIdx.x < rows_per_cuda_block && (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) { - float result = tmp[j][threadIdx.x]; - if constexpr (has_fusion) { - if (use_bias) { - result += x_biases[j]; - } - if (use_gate) { - float gate_value = tmp_gate[j][threadIdx.x]; - if (use_gate_bias) { - gate_value += gate_biases[j]; + float result = 0.0f; + if (channel_valids[j]) { + result = tmp[j][threadIdx.x]; + if constexpr (has_fusion) { + if (use_bias) { + result += x_biases[j]; } - switch (active_glu) { - case GGML_GLU_OP_SWIGLU: - result *= ggml_cuda_op_silu_single(gate_value); - break; - case GGML_GLU_OP_GEGLU: - result *= ggml_cuda_op_gelu_single(gate_value); - break; - case GGML_GLU_OP_SWIGLU_OAI: { - result = ggml_cuda_op_swiglu_oai_single(gate_value, result, glu_param0, glu_param1); - break; + if (use_gate) { + float gate_value = tmp_gate[j][threadIdx.x]; + if (use_gate_bias) { + gate_value += gate_biases[j]; } - case GGML_GLU_OP_SWIGLU_DS4: { - result = ggml_cuda_op_swiglu_ds4_single(gate_value, result, glu_param0); - break; + switch (active_glu) { + case GGML_GLU_OP_SWIGLU: + result *= ggml_cuda_op_silu_single(gate_value); + break; + case GGML_GLU_OP_GEGLU: + result *= ggml_cuda_op_gelu_single(gate_value); + break; + case GGML_GLU_OP_SWIGLU_OAI: { + result = ggml_cuda_op_swiglu_oai_single(gate_value, result, glu_param0, glu_param1); + break; + } + case GGML_GLU_OP_SWIGLU_DS4: { + if (gate_value_scale != 1.0f) { + gate_value *= gate_value_scale; + } + if (x_value_scale != 1.0f) { + result *= x_value_scale; + } + result = ggml_cuda_op_swiglu_ds4_single(gate_value, result, glu_param0); + break; + } + default: + result = result * gate_value; + break; } - default: - result = result * gate_value; - break; } } } @@ -798,7 +1025,9 @@ static __global__ void mul_mat_vec_q( // Grid: (ceil(nrows_x / c_rows_per_block), nchannels_dst) // Block: (warp_size, ncols_dst) - each warp handles one token independently. // No shared memory reduction needed since each warp works alone. -template +template __launch_bounds__(get_mmvq_mmid_max_batch_for_device()*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q_moe( const void * __restrict__ vx, const void * __restrict__ vy, const int32_t * __restrict__ ids, @@ -807,7 +1036,9 @@ static __global__ void mul_mat_vec_q_moe( const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst, const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, - const uint32_t ncols_dst, const uint32_t ids_stride) { + const uint32_t ncols_dst, const uint32_t nchannels_dst, const uint32_t ids_stride, + const bool skip_duplicate_ids, const bool compact_masked_ids, + const bool aligned_shared_ids) { constexpr int qk = ggml_cuda_type_traits::qk; constexpr int qi = ggml_cuda_type_traits::qi; @@ -815,19 +1046,102 @@ static __global__ void mul_mat_vec_q_moe( constexpr int warp_size = ggml_cuda_get_physical_warp_size(); constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); - - const uint32_t token_idx = threadIdx.y; - const int row0 = c_rows_per_block*blockIdx.x; + static_assert(c_warp_groups == 1 || c_warp_groups == 2, + "unsupported MoE warp-group count"); + + // Hybrid expert ownership leaves negative IDs in either the hot or cold + // graph. The original (warp_size, n_tokens) block keeps invalid token + // warps resident beside valid ones. The sparse-owner variant gives every + // token/route pair its own one-warp block so invalid routes retire without + // consuming the valid block's occupancy. The per-warp K traversal and + // accumulation order are unchanged. + const uint32_t token_idx = sparse_warp_blocks + ? (uint32_t) blockIdx.y % ncols_dst + : (uint32_t) threadIdx.y % ncols_dst; + const uint32_t row_group = sparse_warp_blocks + ? 0u + : (uint32_t) threadIdx.y / ncols_dst; + const int row0 = c_rows_per_block * + ((int) blockIdx.x * c_warp_groups + (int) row_group); const int blocks_per_row_x = ncols_x / qk; constexpr int blocks_per_iter = vdr * warp_size / qi; - const uint32_t channel_dst = blockIdx.y; + const uint32_t compact_slot = sparse_warp_blocks + ? (uint32_t) blockIdx.y / ncols_dst + : (uint32_t) blockIdx.y; - if (token_idx >= ncols_dst) { + // The final grid block can contain a row group whose first output row is + // already outside nrows_x. This branch is uniform within that warp and + // the kernel has no block-wide synchronization. + if (token_idx >= ncols_dst || (uint32_t) row0 >= nrows_x) { return; } - const uint32_t channel_x = ids[channel_dst + token_idx * ids_stride]; + // Pack owner-valid routes into the low block slots without changing the + // logical output layout. In the heterogeneous graph each owner receives + // the same six route positions with non-owned IDs masked to -1. Leaving + // those holes interleaved makes nearly every (route, q-token) block carry + // at least one valid warp, so invalid warps occupy the Strix wave beside + // useful work. The stable permutation below executes valid routes first + // but writes every result back to its original route position. Down MMVQ + // and route weighting therefore retain their original indexing and + // floating-point reduction order. + uint32_t channel_dst = compact_slot; + int32_t channel_x_raw = -1; + const int32_t encoded_id = ids[compact_slot + token_idx * ids_stride]; + const bool has_aligned_id = aligned_shared_ids && + (((uint32_t) encoded_id & 0x7f000000u) == 0x5a000000u); + if (has_aligned_id) { + channel_dst = ((uint32_t) encoded_id >> 16) & 0xffu; + if (encoded_id >= 0) { + channel_x_raw = encoded_id & 0xffff; + } else { + channel_x_raw = -1; + } + } else if (compact_masked_ids) { + uint32_t n_valid = 0; + for (uint32_t c = 0; c < nchannels_dst; ++c) { + const int32_t id = ids[c + token_idx * ids_stride]; + const bool duplicate = skip_duplicate_ids && c >= 2 && id >= 0 && + id == ids[c - 2 + token_idx * ids_stride]; + n_valid += id >= 0 && !duplicate; + } + + const bool want_valid = compact_slot < n_valid; + uint32_t rank = want_valid ? compact_slot : compact_slot - n_valid; + for (uint32_t c = 0; c < nchannels_dst; ++c) { + const int32_t id = ids[c + token_idx * ids_stride]; + const bool duplicate = skip_duplicate_ids && c >= 2 && id >= 0 && + id == ids[c - 2 + token_idx * ids_stride]; + const bool valid = id >= 0 && !duplicate; + if (valid == want_valid) { + if (rank == 0) { + channel_dst = c; + channel_x_raw = valid ? id : -1; + break; + } + --rank; + } + } + } else { + channel_x_raw = ids[channel_dst + token_idx * ids_stride]; + } + // The legacy six-route q4 graph lowers its two-route tail as + // [id4, id5, id4, id5] with weights [w4, w5, 0, 0]. Top-k IDs are unique, + // so an exact two-slot-back duplicate is padding and its expert matvec is + // dead work. Keep the output zero; the downstream zero weight therefore + // preserves the original result while avoiding one third of route work. + int32_t previous_channel_x = -1; + if (skip_duplicate_ids && !has_aligned_id && channel_dst >= 2) { + const int32_t previous_raw = + ids[channel_dst - 2 + token_idx * ids_stride]; + previous_channel_x = previous_raw; + } + const bool duplicate_padding = skip_duplicate_ids && !has_aligned_id && + channel_dst >= 2 && + channel_x_raw >= 0 && channel_x_raw == previous_channel_x; + const bool channel_valid = channel_x_raw >= 0 && !duplicate_padding; + const uint32_t channel_x = channel_valid ? (uint32_t) channel_x_raw : 0u; const uint32_t channel_y = fastmodulo(channel_dst, nchannels_y); const block_q8_1 * y = ((const block_q8_1 *) vy) + channel_y*stride_channel_y + token_idx*stride_col_y; @@ -842,6 +1156,8 @@ static __global__ void mul_mat_vec_q_moe( ggml_glu_op active_glu; float glu_param0 = 0.0f; float glu_param1 = 0.0f; + float gate_value_scale = 1.0f; + float x_value_scale = 1.0f; if constexpr (has_fusion) { use_gate = fusion.gate != nullptr; @@ -853,22 +1169,54 @@ static __global__ void mul_mat_vec_q_moe( active_glu = fusion.glu_op; glu_param0 = fusion.glu_param0; glu_param1 = fusion.glu_param1; + gate_value_scale = fusion.gate_value_scale; + x_value_scale = fusion.x_value_scale; } // partial sum for each thread float tmp[c_rows_per_block] = {0.0f}; float tmp_gate[c_rows_per_block] = {0.0f}; - for (int kbx = threadIdx.x / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { - const int kby = kbx * (qk/QK8_1); - const int kqs = vdr * (threadIdx.x % (qi/vdr)); + if (channel_valid) { + for (int kbx = threadIdx.x / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (threadIdx.x % (qi/vdr)); #pragma unroll - for (int i = 0; i < c_rows_per_block; ++i) { - tmp[i] += vec_dot_q_cuda(vx, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs); - if constexpr (has_fusion) { - if (use_gate) { - tmp_gate[i] += vec_dot_q_cuda(vgate, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs); + for (int i = 0; i < c_rows_per_block; ++i) { + if constexpr (type == GGML_TYPE_Q2_0_ROCMFP2 && + c_fp2_packed32) { + tmp[i] += vec_dot_rocmfpx_fp2_q8_1_packed32( + vx, &y[kby], kbx_offset + i*stride_row_x + kbx, + kqs); + } else if constexpr (type == GGML_TYPE_Q3_0_ROCMFPX && + c_fp3_packed24) { + tmp[i] += vec_dot_rocmfpx_fp3_q8_1_packed24( + vx, &y[kby], kbx_offset + i*stride_row_x + kbx, + kqs); + } else { + tmp[i] += vec_dot_q_cuda( + vx, &y[kby], kbx_offset + i*stride_row_x + kbx, + kqs); + } + if constexpr (has_fusion) { + if (use_gate) { + if constexpr (type == GGML_TYPE_Q2_0_ROCMFP2 && + c_fp2_packed32) { + tmp_gate[i] += vec_dot_rocmfpx_fp2_q8_1_packed32( + vgate, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } else if constexpr (type == GGML_TYPE_Q3_0_ROCMFPX && + c_fp3_packed24) { + tmp_gate[i] += vec_dot_rocmfpx_fp3_q8_1_packed24( + vgate, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } else { + tmp_gate[i] += vec_dot_q_cuda( + vgate, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } + } } } } @@ -887,32 +1235,41 @@ static __global__ void mul_mat_vec_q_moe( // Write results if (threadIdx.x < c_rows_per_block && (c_rows_per_block == 1 || uint32_t(row0 + threadIdx.x) < nrows_x)) { - float result = tmp[threadIdx.x]; - if constexpr (has_fusion) { - if (use_bias) { - result += x_bias[channel_x*stride_channel_dst + row0 + threadIdx.x]; - } - if (use_gate) { - float gate_value = tmp_gate[threadIdx.x]; - if (use_gate_bias) { - gate_value += gate_bias[channel_x*stride_channel_dst + row0 + threadIdx.x]; + float result = 0.0f; + if (channel_valid) { + result = tmp[threadIdx.x]; + if constexpr (has_fusion) { + if (use_bias) { + result += x_bias[channel_x*stride_channel_dst + row0 + threadIdx.x]; } - switch (active_glu) { - case GGML_GLU_OP_SWIGLU: - result *= ggml_cuda_op_silu_single(gate_value); - break; - case GGML_GLU_OP_GEGLU: - result *= ggml_cuda_op_gelu_single(gate_value); - break; - case GGML_GLU_OP_SWIGLU_OAI: - result = ggml_cuda_op_swiglu_oai_single(gate_value, result, glu_param0, glu_param1); - break; - case GGML_GLU_OP_SWIGLU_DS4: - result = ggml_cuda_op_swiglu_ds4_single(gate_value, result, glu_param0); - break; - default: - result = result * gate_value; - break; + if (use_gate) { + float gate_value = tmp_gate[threadIdx.x]; + if (use_gate_bias) { + gate_value += gate_bias[channel_x*stride_channel_dst + row0 + threadIdx.x]; + } + switch (active_glu) { + case GGML_GLU_OP_SWIGLU: + result *= ggml_cuda_op_silu_single(gate_value); + break; + case GGML_GLU_OP_GEGLU: + result *= ggml_cuda_op_gelu_single(gate_value); + break; + case GGML_GLU_OP_SWIGLU_OAI: + result = ggml_cuda_op_swiglu_oai_single(gate_value, result, glu_param0, glu_param1); + break; + case GGML_GLU_OP_SWIGLU_DS4: + if (gate_value_scale != 1.0f) { + gate_value *= gate_value_scale; + } + if (x_value_scale != 1.0f) { + result *= x_value_scale; + } + result = ggml_cuda_op_swiglu_ds4_single(gate_value, result, glu_param0); + break; + default: + result = result * gate_value; + break; + } } } } @@ -1204,7 +1561,8 @@ static std::pair calc_launch_params( } template + int fixed_ncols_x = 0, bool unroll_k_loop_2 = false, + bool fp3_packed24 = false, bool fp4_x4 = false> static void mul_mat_vec_q_switch_fusion( const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, @@ -1216,14 +1574,18 @@ static void mul_mat_vec_q_switch_fusion( const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; if (has_fusion) { - mul_mat_vec_q<<>> + mul_mat_vec_q<<>> (vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, ids_tokenwise_samples); return; } - mul_mat_vec_q<<>> + mul_mat_vec_q<<>> (vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, ids_tokenwise_samples); @@ -1280,6 +1642,75 @@ static void mul_mat_vec_rocmfp4_unroll2_launch( ids_stride, stream); } +template +static void mul_mat_vec_q_moe_launch_rpb( + const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, + const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst, + const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, + const uint32_t ncols_dst, const uint32_t ids_stride, + const int warp_size, const int nchannels_dst, cudaStream_t stream, + const bool sparse_warp_blocks, const bool skip_duplicate_ids, + const bool compact_masked_ids, const bool aligned_shared_ids) { + + static_assert(rows_per_block == 1 || rows_per_block == 2 || + rows_per_block == 4 || rows_per_block == 8, + "unsupported MoE rows-per-block variant"); + static_assert(warp_groups == 1 || warp_groups == 2, + "unsupported MoE warp-group variant"); + const int active_warp_groups = sparse_warp_blocks ? 1 : warp_groups; + const int64_t rows_per_grid_block = + (int64_t) rows_per_block * active_warp_groups; + const int64_t nblocks_rows = + (nrows_x + rows_per_grid_block - 1) / rows_per_grid_block; + const dim3 block_nums( + nblocks_rows, + sparse_warp_blocks ? nchannels_dst * ncols_dst : nchannels_dst); + const dim3 block_dims( + warp_size, + sparse_warp_blocks ? 1 : ncols_dst * active_warp_groups); + + const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; + if (has_fusion) { + if (sparse_warp_blocks) { + mul_mat_vec_q_moe<<>>( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst, + ncols_dst, nchannels_dst, ids_stride, + skip_duplicate_ids, compact_masked_ids, aligned_shared_ids); + } else { + mul_mat_vec_q_moe<<>>( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst, + ncols_dst, nchannels_dst, ids_stride, + skip_duplicate_ids, compact_masked_ids, aligned_shared_ids); + } + } else { + if (sparse_warp_blocks) { + mul_mat_vec_q_moe<<>>( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst, + ncols_dst, nchannels_dst, ids_stride, + skip_duplicate_ids, compact_masked_ids, aligned_shared_ids); + } else { + mul_mat_vec_q_moe<<>>( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, + stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, stride_channel_dst, + ncols_dst, nchannels_dst, ids_stride, + skip_duplicate_ids, compact_masked_ids, aligned_shared_ids); + } + } +} + template static void mul_mat_vec_q_moe_launch( const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, @@ -1289,25 +1720,133 @@ static void mul_mat_vec_q_moe_launch( const uint32_t ncols_dst, const uint32_t ids_stride, const int warp_size, const int nchannels_dst, cudaStream_t stream) { - constexpr int rows_per_block = 2; // 2 gives best perf based on tuning - const int64_t nblocks_rows = (nrows_x + rows_per_block - 1) / rows_per_block; - const dim3 block_nums(nblocks_rows, nchannels_dst); - const dim3 block_dims(warp_size, ncols_dst); + static const bool sparse_warp_blocks = []() { + const char * e = std::getenv("DFLASH_CUDA_MMVQ_MOE_SPARSE_WARP_BLOCKS"); + return e && e[0] == '1' && e[1] == '\0'; + }(); + static const bool skip_duplicate_ids = []() { + const char * e = std::getenv("DFLASH_CUDA_MMVQ_MOE_SKIP_DUPLICATE_IDS"); + return e && e[0] == '1' && e[1] == '\0'; + }(); + static const bool compact_masked_ids = []() { + const char * e = std::getenv("DFLASH_CUDA_MMVQ_MOE_COMPACT_MASKED_IDS"); + return e && e[0] == '1' && e[1] == '\0'; + }(); + static const bool aligned_shared_ids = []() { + const char * e = std::getenv("DFLASH_CUDA_MMVQ_MOE_ALIGN_SHARED_IDS"); + return e && e[0] == '1' && e[1] == '\0'; + }(); + static const int tuned_rows_per_block = []() { + const char * e = std::getenv("DFLASH_CUDA_MMVQ_MOE_ROWS_PER_BLOCK"); + if (!e || !e[0]) return 2; + const int value = std::atoi(e); + return value == 1 || value == 2 || value == 4 || value == 8 + ? value : 2; + }(); + static const bool q2_warp_groups = []() { + const char * e = + std::getenv("DFLASH_CUDA_MMVQ_MOE_Q2_WARP_GROUPS"); + return e && e[0] == '2' && e[1] == '\0'; + }(); + static const bool q4_warp_groups = []() { + const char * e = + std::getenv("DFLASH_CUDA_MMVQ_MOE_Q4_WARP_GROUPS"); + return e && e[0] == '2' && e[1] == '\0'; + }(); + static const bool fp3_packed24 = []() { + const char * e = + std::getenv("DFLASH_CUDA_MMVQ_MOE_FP3_PACKED24"); + return e && e[0] == '1' && e[1] == '\0'; + }(); + static const bool fp2_packed32 = []() { + const char * e = + std::getenv("DFLASH_CUDA_MMVQ_MOE_FP2_PACKED32"); + return e && e[0] == '1' && e[1] == '\0'; + }(); + +#define GGML_MOE_LAUNCH_RPB(RPB) \ + mul_mat_vec_q_moe_launch_rpb( \ + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, \ + stride_row_x, stride_col_y, stride_col_dst, \ + stride_channel_x, stride_channel_y, stride_channel_dst, \ + ncols_dst, ids_stride, warp_size, nchannels_dst, stream, \ + sparse_warp_blocks, skip_duplicate_ids, compact_masked_ids, \ + aligned_shared_ids) + +#define GGML_MOE_LAUNCH_TWO_WARP_GROUPS() \ + mul_mat_vec_q_moe_launch_rpb( \ + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, \ + stride_row_x, stride_col_y, stride_col_dst, \ + stride_channel_x, stride_channel_y, stride_channel_dst, \ + ncols_dst, ids_stride, warp_size, nchannels_dst, stream, \ + sparse_warp_blocks, skip_duplicate_ids, compact_masked_ids, \ + aligned_shared_ids) + +#define GGML_MOE_LAUNCH_FP3_PACKED24(WARP_GROUPS) \ + mul_mat_vec_q_moe_launch_rpb( \ + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, \ + stride_row_x, stride_col_y, stride_col_dst, \ + stride_channel_x, stride_channel_y, stride_channel_dst, \ + ncols_dst, ids_stride, warp_size, nchannels_dst, stream, \ + sparse_warp_blocks, skip_duplicate_ids, compact_masked_ids, \ + aligned_shared_ids) + +#define GGML_MOE_LAUNCH_FP2_PACKED32(WARP_GROUPS) \ + mul_mat_vec_q_moe_launch_rpb( \ + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, \ + stride_row_x, stride_col_y, stride_col_dst, \ + stride_channel_x, stride_channel_y, stride_channel_dst, \ + ncols_dst, ids_stride, warp_size, nchannels_dst, stream, \ + sparse_warp_blocks, skip_duplicate_ids, compact_masked_ids, \ + aligned_shared_ids) - const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; - if (has_fusion) { - mul_mat_vec_q_moe<<>>( - vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, - stride_row_x, stride_col_y, stride_col_dst, - stride_channel_x, stride_channel_y, stride_channel_dst, - ncols_dst, ids_stride); + if constexpr (type == GGML_TYPE_Q2_0_ROCMFP2) { + if (fp2_packed32 && tuned_rows_per_block == 2 && + !sparse_warp_blocks) { + if ((q2_warp_groups && ncols_dst == 2) || + (q4_warp_groups && ncols_dst == 4)) { + GGML_MOE_LAUNCH_FP2_PACKED32(2); + } else { + GGML_MOE_LAUNCH_FP2_PACKED32(1); + } + return; + } + } + + if constexpr (type == GGML_TYPE_Q3_0_ROCMFPX) { + if (fp3_packed24 && tuned_rows_per_block == 2 && + !sparse_warp_blocks) { + if ((q2_warp_groups && ncols_dst == 2) || + (q4_warp_groups && ncols_dst == 4)) { + GGML_MOE_LAUNCH_FP3_PACKED24(2); + } else { + GGML_MOE_LAUNCH_FP3_PACKED24(1); + } + return; + } + } + + if constexpr (type == GGML_TYPE_Q2_0_ROCMFP2 || + type == GGML_TYPE_Q3_0_ROCMFPX) { + if (((q2_warp_groups && ncols_dst == 2) || + (q4_warp_groups && ncols_dst == 4)) && + tuned_rows_per_block == 2 && !sparse_warp_blocks) { + GGML_MOE_LAUNCH_TWO_WARP_GROUPS(); + } else { + switch (tuned_rows_per_block) { + case 1: GGML_MOE_LAUNCH_RPB(1); break; + case 4: GGML_MOE_LAUNCH_RPB(4); break; + case 8: GGML_MOE_LAUNCH_RPB(8); break; + default: GGML_MOE_LAUNCH_RPB(2); break; + } + } } else { - mul_mat_vec_q_moe<<>>( - vx, vy, ids, fusion, dst, ncols_x, nchannels_y, nrows_x, - stride_row_x, stride_col_y, stride_col_dst, - stride_channel_x, stride_channel_y, stride_channel_dst, - ncols_dst, ids_stride); + GGML_MOE_LAUNCH_RPB(2); } +#undef GGML_MOE_LAUNCH_TWO_WARP_GROUPS +#undef GGML_MOE_LAUNCH_FP2_PACKED32 +#undef GGML_MOE_LAUNCH_FP3_PACKED24 +#undef GGML_MOE_LAUNCH_RPB } template @@ -1450,9 +1989,50 @@ static void mul_mat_vec_q_switch_ncols_dst( return; } + if constexpr (type == GGML_TYPE_Q4_0_ROCMFP4_FAST) { + if (!has_ids && ncols_dst == 4 && rocmfp4_x4_enabled()) { + constexpr int c_ncols_dst = 4; + std::pair dims = calc_launch_params( + c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, + warp_size, table_id); + mul_mat_vec_q_switch_fusion( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, + stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, + stride_channel_dst, sample_ratio_fd, stride_sample_x, + stride_sample_y, stride_sample_dst, dims.first, dims.second, + 0, ids_stride, stream); + return; + } + } + + // The generic q4 verifier kernels previously made every FP3 lane load the + // complete 12-byte payload even though VDR=2 assigns that lane one + // disjoint three-byte quarter. Keep a separately instantiated, opt-in + // ncols=4 path so the exact packed-dot A/B does not perturb q1 or the + // dedicated multi-token expert kernel. + if constexpr (type == GGML_TYPE_Q3_0_ROCMFPX) { + if (!has_ids && ncols_dst == 4 && rocmfp3_packed24_enabled()) { + constexpr int c_ncols_dst = 4; + std::pair dims = calc_launch_params( + c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, + warp_size, table_id); + mul_mat_vec_q_switch_fusion( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, + stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, + stride_channel_dst, sample_ratio_fd, stride_sample_x, + stride_sample_y, stride_sample_dst, dims.first, dims.second, + 0, ids_stride, stream); + return; + } + } + // Decode-only gfx1151 specializations. Each one preserves the original // per-lane K traversal and accumulation order, so exact validated shapes - // can use the faster kernel without a serving-time tuning flag. + // use the faster kernel without a serving-time tuning flag. if constexpr (type == GGML_TYPE_Q4_0_ROCMFP4_FAST) { if (is_gfx1151(cc) && ncols_dst == 1) { mul_mat_vec_rocmfp4_unroll2_launch( @@ -1817,6 +2397,8 @@ void ggml_cuda_mul_mat_vec_q( fusion_local.glu_op = fusion->glu_op; fusion_local.glu_param0 = fusion->glu_param0; fusion_local.glu_param1 = fusion->glu_param1; + fusion_local.gate_value_scale = fusion->gate_value_scale; + fusion_local.x_value_scale = fusion->x_value_scale; } // If src0 is a temporary compute buffer, clear any potential padding. @@ -1838,7 +2420,15 @@ void ggml_cuda_mul_mat_vec_q( const size_t q8_bytes = ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1; ggml_cuda_pool_alloc src1_q8_1(ctx.pool()); char * src1_q8_d = nullptr; - if (luce_q8_memo_on) { + // The src1->q8_1 quantization depends only on src0->type and src1's dims/strides, + // not on ids (ids only affects the matmul kernel's channel/dst strides below), so + // the memo is valid for MUL_MAT_ID too: gate/up and the shared expert re-quantize + // the same ffn_norm activation and can share one q8 buffer within an evaluation. + // Synthetic tensors used inside a coarse fused operator do not own a + // persistent graph allocation. Never retain their stack address as a + // cross-node memo key. + const bool use_q8_memo = luce_q8_memo_on && src1->buffer != nullptr; + if (use_q8_memo) { for (const auto & e : ctx.luce_q8_memo) { if (e.src1_node == (const void *) src1 && e.src1_data == (const void *) src1_d && e.src0_type == (int) src0->type && @@ -1850,7 +2440,7 @@ void ggml_cuda_mul_mat_vec_q( } if (src1_q8_d == nullptr) { char * q8_dst; - if (luce_q8_memo_on) { + if (use_q8_memo) { ggml_backend_cuda_context::luce_q8_memo_entry ent; ent.src1_node = (const void *) src1; ent.src1_data = (const void *) src1_d; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu index d86ab7f20..8505499a7 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu @@ -1,6 +1,9 @@ #include "moe-fused.cuh" #include "ggml-cuda/vecdotq.cuh" #include "ggml-cuda/dequantize.cuh" +#include "ggml-cuda/mmvq.cuh" + +#include static __device__ __forceinline__ float silu_f32(float x) { return x / (1.0f + expf(-x)); @@ -272,7 +275,8 @@ static __global__ void laguna_moe_combine_kernel( const size_t weights_nb0, const size_t weights_nb1, const size_t output_nb0, - const size_t output_nb1) { + const size_t output_nb1, + const float value_scale) { const int idx = blockIdx.x * blockDim.x + threadIdx.x; const int total = n_embd * n_tokens; if (idx >= total) return; @@ -288,7 +292,9 @@ static __global__ void laguna_moe_combine_kernel( const float w = *(const float *)(weights + (size_t)e * weights_nb0 + (size_t)t * weights_nb1); - const float prod = __fmul_rn(v, w); + const float scaled = value_scale == 1.0f + ? v : __fmul_rn(v, value_scale); + const float prod = __fmul_rn(scaled, w); sum = (e == 0) ? prod : __fadd_rn(sum, prod); } *(float *)(output + @@ -296,8 +302,381 @@ static __global__ void laguna_moe_combine_kernel( (size_t)t * output_nb1) = sum; } +static __global__ void ds4_peer_copy_f32_kernel( + const float * __restrict__ src, + float * __restrict__ dst, + const int64_t n) { + const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) { + dst[i] = src[i]; + } +} + +// Align equal owner-local expert IDs across q-token warps. The high 16 bits +// of every valid output encode the original route slot; invalid entries use +// the sign bit plus the original slot. The dedicated MoE MMVQ kernel decodes +// this metadata and scatters its result back to the original route layout. +// A single thread is deliberate: this is at most a 4 x 6 assignment problem +// and runs once per owner/layer, outside the weight-streaming kernels. +static __global__ void ds4_align_moe_ids_kernel( + const int32_t * __restrict__ ids, + int32_t * __restrict__ aligned, + int n_routes, + int n_tokens, + int ids_stride, + int aligned_stride) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + + constexpr int max_routes = 6; + constexpr int max_tokens = 4; + if (n_routes > max_routes || n_tokens > max_tokens) { + for (int t = 0; t < n_tokens; ++t) { + for (int r = 0; r < n_routes; ++r) { + const int32_t id = ids[t * ids_stride + r]; + aligned[t * aligned_stride + r] = id >= 0 + ? (int32_t) (0x5a000000u | ((uint32_t) r << 16) | + ((uint32_t) id & 0xffffu)) + : (int32_t) (0xda00ffffu | ((uint32_t) r << 16)); + } + } + return; + } + + int route_for_slot[max_tokens][max_routes]; + for (int t = 0; t < max_tokens; ++t) { + for (int s = 0; s < max_routes; ++s) { + route_for_slot[t][s] = -1; + } + } + for (int r = 0; r < n_routes; ++r) { + route_for_slot[0][r] = r; + } + + for (int t = 1; t < n_tokens; ++t) { + bool used_slot[max_routes] = {}; + bool assigned_route[max_routes] = {}; + + // Prefer assignments that preserve the most previous occurrences of + // this expert in one slot. This handles conflicts deterministically + // when two earlier tokens placed different experts in the same slot. + for (int pass = 0; pass < n_routes; ++pass) { + int best_route = -1; + int best_slot = -1; + int best_count = 0; + for (int r = 0; r < n_routes; ++r) { + if (assigned_route[r]) continue; + const int32_t id = ids[t * ids_stride + r]; + if (id < 0) continue; + for (int s = 0; s < n_routes; ++s) { + if (used_slot[s]) continue; + int count = 0; + for (int p = 0; p < t; ++p) { + const int prev_route = route_for_slot[p][s]; + if (prev_route >= 0 && + ids[p * ids_stride + prev_route] == id) { + ++count; + } + } + if (count > best_count || + (count == best_count && count > 0 && + (best_route < 0 || r < best_route || + (r == best_route && s < best_slot)))) { + best_route = r; + best_slot = s; + best_count = count; + } + } + } + if (best_count == 0) break; + route_for_slot[t][best_slot] = best_route; + used_slot[best_slot] = true; + assigned_route[best_route] = true; + } + + // Keep unmatched routes in their original slot when possible, then + // fill the remaining permutation slots in ascending order. + for (int r = 0; r < n_routes; ++r) { + if (assigned_route[r]) continue; + int slot = !used_slot[r] ? r : -1; + if (slot < 0) { + for (int s = 0; s < n_routes; ++s) { + if (!used_slot[s]) { + slot = s; + break; + } + } + } + route_for_slot[t][slot] = r; + used_slot[slot] = true; + assigned_route[r] = true; + } + } + + for (int t = 0; t < n_tokens; ++t) { + for (int s = 0; s < n_routes; ++s) { + const int original_route = route_for_slot[t][s]; + const int32_t id = ids[t * ids_stride + original_route]; + aligned[t * aligned_stride + s] = id >= 0 + ? (int32_t) (0x5a000000u | + ((uint32_t) original_route << 16) | + ((uint32_t) id & 0xffffu)) + : (int32_t) (0xda00ffffu | + ((uint32_t) original_route << 16)); + } + } +} + +static ggml_tensor make_contiguous_f32_tensor( + float * data, + int64_t ne0, + int64_t ne1, + int64_t ne2) { + ggml_tensor tensor = {}; + tensor.type = GGML_TYPE_F32; + tensor.data = data; + tensor.ne[0] = ne0; + tensor.ne[1] = ne1; + tensor.ne[2] = ne2; + tensor.ne[3] = 1; + tensor.nb[0] = sizeof(float); + tensor.nb[1] = tensor.nb[0] * ne0; + tensor.nb[2] = tensor.nb[1] * ne1; + tensor.nb[3] = tensor.nb[2] * ne2; + return tensor; +} + +static void ggml_cuda_op_ds4_moe_owner( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst) { + const ggml_tensor * input = dst->src[0]; // [n_embd, n_tokens] + const ggml_tensor * gate_up_w = dst->src[1]; // [n_embd, 2*n_ff, n_expert] + const ggml_tensor * down_w = dst->src[2]; // [n_ff, n_embd, n_expert] + const ggml_tensor * expert_ids = dst->src[3]; // [n_used, n_tokens] + const ggml_tensor * weights = dst->src[4]; // [n_used, n_tokens] + + const int n_embd = (int) input->ne[0]; + const int n_tokens = (int) input->ne[1]; + const int n_used = (int) expert_ids->ne[0]; + const int n_ff = ggml_get_op_params_i32(dst, 1); + const float clamp = ggml_get_op_params_f32(dst, 2); + const float down_scale = ggml_get_op_params_f32(dst, 3); + + GGML_ASSERT(gate_up_w->ne[0] == n_embd); + GGML_ASSERT(gate_up_w->ne[1] == 2 * n_ff); + GGML_ASSERT(down_w->ne[0] == n_ff); + GGML_ASSERT(down_w->ne[1] == n_embd); + GGML_ASSERT(expert_ids->ne[1] == n_tokens); + GGML_ASSERT(weights->ne[0] == n_used && weights->ne[1] == n_tokens); + GGML_ASSERT(dst->ne[0] == n_embd && dst->ne[1] == n_tokens); + + // The checkpoint concatenates gate rows followed by up rows inside each + // expert. Keep the original expert stride while viewing each half. + ggml_tensor gate_w = *gate_up_w; + ggml_tensor up_w = *gate_up_w; + gate_w.ne[1] = n_ff; + up_w.ne[1] = n_ff; + up_w.data = (char *) gate_up_w->data + (size_t) n_ff * gate_up_w->nb[1]; + + // MUL_MAT_ID expects token columns in dimension 2. The input is physically + // identical; only the descriptor changes from [D,T] to [D,1,T]. + ggml_tensor input_3d = *input; + input_3d.ne[1] = 1; + input_3d.ne[2] = n_tokens; + input_3d.ne[3] = 1; + input_3d.nb[1] = input->nb[1]; + input_3d.nb[2] = input->nb[1]; + input_3d.nb[3] = input->nb[1] * n_tokens; + // This descriptor is stack-local. Disable q8 memoization for it. + input_3d.buffer = nullptr; + + ggml_cuda_pool_alloc gu_alloc( + ctx.pool(), (size_t) n_ff * n_used * n_tokens); + ggml_tensor gu = make_contiguous_f32_tensor( + gu_alloc.ptr, n_ff, n_used, n_tokens); + + ggml_cuda_mm_fusion_args_host gate_up_fusion{}; + gate_up_fusion.gate = &gate_w; + gate_up_fusion.glu_op = clamp > 1.0e-6f + ? GGML_GLU_OP_SWIGLU_DS4 : GGML_GLU_OP_SWIGLU; + gate_up_fusion.glu_param0 = clamp; + ggml_cuda_mul_mat_vec_q( + ctx, &up_w, &input_3d, expert_ids, &gu, &gate_up_fusion); + + ggml_cuda_pool_alloc experts_alloc( + ctx.pool(), (size_t) n_embd * n_used * n_tokens); + ggml_tensor experts = make_contiguous_f32_tensor( + experts_alloc.ptr, n_embd, n_used, n_tokens); + ggml_cuda_mul_mat_vec_q( + ctx, down_w, &gu, expert_ids, &experts, nullptr); + + const int total = n_embd * n_tokens; + const int block = 256; + const int grid = (total + block - 1) / block; + laguna_moe_combine_kernel<<>>( + (const char *) experts.data, + (const char *) weights->data, + (char *) dst->data, + n_embd, n_used, n_tokens, + experts.nb[0], experts.nb[1], experts.nb[2], + weights->nb[0], weights->nb[1], + dst->nb[0], dst->nb[1], + down_scale); +} + +static void ggml_cuda_op_ds4_moe_owner_split( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst) { + const ggml_tensor * input = dst->src[0]; // [n_embd, n_tokens] + const ggml_tensor * gate_w = dst->src[1]; // [n_embd, n_ff, n_expert] + const ggml_tensor * up_w = dst->src[2]; // [n_embd, n_ff, n_expert] + const ggml_tensor * down_w = dst->src[3]; // [n_ff, n_embd, n_expert] + const ggml_tensor * expert_ids = dst->src[4]; // [n_used, n_tokens] + const ggml_tensor * weights = dst->src[5]; // [n_used, n_tokens] + + const int n_embd = (int) input->ne[0]; + const int n_tokens = (int) input->ne[1]; + const int n_used = (int) expert_ids->ne[0]; + const int n_ff = ggml_get_op_params_i32(dst, 1); + const float clamp = ggml_get_op_params_f32(dst, 2); + const float gate_scale = ggml_get_op_params_f32(dst, 3); + const float up_scale = ggml_get_op_params_f32(dst, 4); + const float down_scale = ggml_get_op_params_f32(dst, 5); + + GGML_ASSERT(gate_w->ne[0] == n_embd && gate_w->ne[1] == n_ff); + GGML_ASSERT(up_w->ne[0] == n_embd && up_w->ne[1] == n_ff); + GGML_ASSERT(down_w->ne[0] == n_ff && down_w->ne[1] == n_embd); + GGML_ASSERT(expert_ids->ne[1] == n_tokens); + GGML_ASSERT(weights->ne[0] == n_used && weights->ne[1] == n_tokens); + GGML_ASSERT(dst->ne[0] == n_embd && dst->ne[1] == n_tokens); + + // MUL_MAT_ID consumes token columns in dimension 2. This descriptor is + // stack-local, so it deliberately does not participate in q8 memoization. + ggml_tensor input_3d = *input; + input_3d.ne[1] = 1; + input_3d.ne[2] = n_tokens; + input_3d.ne[3] = 1; + input_3d.nb[1] = input->nb[1]; + input_3d.nb[2] = input->nb[1]; + input_3d.nb[3] = input->nb[1] * n_tokens; + input_3d.buffer = nullptr; + + ggml_cuda_pool_alloc gu_alloc( + ctx.pool(), (size_t) n_ff * n_used * n_tokens); + ggml_tensor gu = make_contiguous_f32_tensor( + gu_alloc.ptr, n_ff, n_used, n_tokens); + + ggml_cuda_mm_fusion_args_host gate_up_fusion{}; + gate_up_fusion.gate = gate_w; + gate_up_fusion.glu_op = clamp > 1.0e-6f + ? GGML_GLU_OP_SWIGLU_DS4 : GGML_GLU_OP_SWIGLU; + gate_up_fusion.glu_param0 = clamp; + gate_up_fusion.gate_value_scale = gate_scale; + gate_up_fusion.x_value_scale = up_scale; + ggml_cuda_mul_mat_vec_q( + ctx, up_w, &input_3d, expert_ids, &gu, &gate_up_fusion); + + ggml_cuda_pool_alloc experts_alloc( + ctx.pool(), (size_t) n_embd * n_used * n_tokens); + ggml_tensor experts = make_contiguous_f32_tensor( + experts_alloc.ptr, n_embd, n_used, n_tokens); + ggml_cuda_mul_mat_vec_q( + ctx, down_w, &gu, expert_ids, &experts, nullptr); + + const int total = n_embd * n_tokens; + const int block = 256; + const int grid = (total + block - 1) / block; + laguna_moe_combine_kernel<<>>( + (const char *) experts.data, + (const char *) weights->data, + (char *) dst->data, + n_embd, n_used, n_tokens, + experts.nb[0], experts.nb[1], experts.nb[2], + weights->nb[0], weights->nb[1], + dst->nb[0], dst->nb[1], + down_scale); +} + void ggml_cuda_op_moe_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - if (ggml_get_op_params_i32(dst, 0) == -1) { + const int mode = ggml_get_op_params_i32(dst, 0); + if (mode == -5) { + const ggml_tensor * ids = dst->src[0]; + GGML_ASSERT(ids && ids->type == GGML_TYPE_I32); + GGML_ASSERT(dst->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(ids) && ggml_is_contiguous(dst)); + ds4_align_moe_ids_kernel<<<1, 1, 0, ctx.stream()>>>( + (const int32_t *) ids->data, + (int32_t *) dst->data, + (int) ids->ne[0], + (int) ids->ne[1], + (int) (ids->nb[1] / sizeof(int32_t)), + (int) (dst->nb[1] / sizeof(int32_t))); + return; + } + if (mode == -3) { + GGML_ASSERT(ggml_is_contiguous(dst)); + + void * event_context = nullptr; + memcpy(&event_context, &dst->op_params[2], sizeof(event_context)); + GGML_ASSERT(event_context); + + const int transfer_mode = ggml_get_op_params_i32(dst, 6); + if (transfer_mode == 3) { + // Whole-step heterogeneous graph path: the Strix producer wrote + // this activation into coherent mapped staging before publishing + // the scheduler event. The scheduler's late event wait precedes + // this node, so only a local R9700 destination write is needed. + void * staging_data = nullptr; + memcpy(&staging_data, &dst->op_params[8], sizeof(staging_data)); + GGML_ASSERT(staging_data); + const int64_t n = ggml_nelements(dst); + const int block = 256; + const int grid = (int) ((n + block - 1) / block); + ds4_peer_copy_f32_kernel<<>>( + (const float *) staging_data, (float *) dst->data, n); + return; + } + + if (transfer_mode != 0) { + // Diagnostic path: the scheduler prefilled dst with a synchronous + // peer copy. Leave the graph node as a dependency-only no-op. + return; + } + + void * source_data = nullptr; + memcpy(&source_data, &dst->op_params[4], sizeof(source_data)); + GGML_ASSERT(source_data); + + if (ggml_get_op_params_i32(dst, 7) == 0) { + // Same-split control: capture the wait at its exact position in + // the main GPU graph. On the current ROCm runtime this preserves + // correctness but realizes almost no cross-device overlap. + CUDA_CHECK(cudaStreamWaitEvent( + ctx.stream(), (cudaEvent_t) event_context, 0)); + } + + // Direct peer reads are required here. On this R9700 + gfx1151 pair, + // embedding hipMemcpyPeerAsync in the captured consumer graph returned + // stale/corrupt values even with a host-synchronized producer event. + // The small contiguous activation is only 16-64 KiB, and the peer-read + // kernel is exact with both eager execution and HIP graph replay. + const int64_t n = ggml_nelements(dst); + const int block = 256; + const int grid = (int) ((n + block - 1) / block); + ds4_peer_copy_f32_kernel<<>>( + (const float *) source_data, (float *) dst->data, n); + return; + } + if (mode == -2) { + ggml_cuda_op_ds4_moe_owner(ctx, dst); + return; + } + if (mode == -4) { + ggml_cuda_op_ds4_moe_owner_split(ctx, dst); + return; + } + if (mode == -1) { const ggml_tensor * experts = dst->src[0]; // [n_embd, n_used, n_tokens] const ggml_tensor * weights = dst->src[1]; // [n_used, n_tokens] const int n_embd = (int) experts->ne[0]; @@ -314,7 +693,8 @@ void ggml_cuda_op_moe_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst) n_embd, n_used, n_tokens, experts->nb[0], experts->nb[1], experts->nb[2], weights->nb[0], weights->nb[1], - dst->nb[0], dst->nb[1]); + dst->nb[0], dst->nb[1], + 1.0f); return; } From dc4a9c473f1f8b6bf5fc59620c9f1bc24c47e2d3 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sun, 19 Jul 2026 22:19:23 -0700 Subject: [PATCH 04/15] perf(ds4): optimize speculative verify and draft state --- server/src/deepseek4/deepseek4_dspark.cpp | 66 ++ server/src/deepseek4/deepseek4_dspark.h | 64 +- .../src/deepseek4/deepseek4_dspark_spec.cpp | 586 +++++++++++++++--- 3 files changed, 626 insertions(+), 90 deletions(-) diff --git a/server/src/deepseek4/deepseek4_dspark.cpp b/server/src/deepseek4/deepseek4_dspark.cpp index 81300e48f..18825ec37 100644 --- a/server/src/deepseek4/deepseek4_dspark.cpp +++ b/server/src/deepseek4/deepseek4_dspark.cpp @@ -538,8 +538,74 @@ bool load_deepseek4_dspark_drafter(const std::string & path, return true; } +bool clone_deepseek4_dspark_heads(DSparkDrafter & d, + ggml_backend_t backend) { + if (!backend || !d.markov_w1 || !d.markov_w2) { + set_err("cannot clone missing DSpark heads"); + return false; + } + if (d.head_buf || d.head_ctx) { + set_err("DSpark heads are already cloned"); + return false; + } + + ggml_init_params ip{}; + ip.mem_size = 1024 * 1024; + ip.mem_buffer = nullptr; + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) { + set_err("DSpark head mirror ggml_init failed"); + return false; + } + + auto duplicate = [&](ggml_tensor * src) -> ggml_tensor * { + return src ? ggml_dup_tensor(ctx, src) : nullptr; + }; + ggml_tensor * markov_w1 = duplicate(d.markov_w1); + ggml_tensor * markov_w2 = duplicate(d.markov_w2); + ggml_tensor * confidence_w = duplicate(d.confidence_w); + ggml_tensor * confidence_b = duplicate(d.confidence_b); + if (!markov_w1 || !markov_w2) { + ggml_free(ctx); + set_err("DSpark head mirror metadata allocation failed"); + return false; + } + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (!buf) { + ggml_free(ctx); + set_err("DSpark head mirror backend allocation failed"); + return false; + } + ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + std::vector staging; + auto copy_tensor = [&](ggml_tensor * src, ggml_tensor * dst) { + if (!src || !dst) return; + const size_t bytes = ggml_nbytes(src); + staging.resize(bytes); + ggml_backend_tensor_get(src, staging.data(), 0, bytes); + ggml_backend_tensor_set(dst, staging.data(), 0, bytes); + }; + copy_tensor(d.markov_w1, markov_w1); + copy_tensor(d.markov_w2, markov_w2); + copy_tensor(d.confidence_w, confidence_w); + copy_tensor(d.confidence_b, confidence_b); + + d.head_ctx = ctx; + d.head_buf = buf; + d.markov_w1 = markov_w1; + d.markov_w2 = markov_w2; + d.confidence_w = confidence_w; + d.confidence_b = confidence_b; + return true; +} + void free_deepseek4_dspark_drafter(DSparkDrafter & d) { reset_deepseek4_dspark_runtime_cache(); + if (d.head_buf) { ggml_backend_buffer_free(d.head_buf); d.head_buf = nullptr; } + if (d.head_ctx) { ggml_free(d.head_ctx); d.head_ctx = nullptr; } if (d.core.buf) { ggml_backend_buffer_free(d.core.buf); d.core.buf = nullptr; } if (d.core.ctx) { ggml_free(d.core.ctx); d.core.ctx = nullptr; } d = DSparkDrafter{}; diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index bdefaac64..b086ab9bb 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -46,6 +46,12 @@ class DFlashDraftIpcClient; struct DSparkDrafter { DeepSeek4Weights core; + // Optional target-backend mirrors of the small Markov/confidence heads. + // The decoder blocks may live on a separate draft GPU while these heads + // stay beside the target's tied lm_head. + ggml_context * head_ctx = nullptr; + ggml_backend_buffer_t head_buf = nullptr; + // Captured-feature fusion (stage 0 only in the checkpoint, but stored global). ggml_tensor * main_proj = nullptr; // dflash.fc.weight [n_tgt*n_embd, n_embd] ggml_tensor * main_norm = nullptr; // dflash.hidden_norm.weight [n_embd] @@ -73,6 +79,11 @@ bool load_deepseek4_dspark_drafter(const std::string & path, ggml_backend_t backend, DSparkDrafter & out); +// Copy only the DSpark Markov/confidence heads to `backend` and repoint the +// drafter fields. Decoder-block tensors remain on `core.backend`. +bool clone_deepseek4_dspark_heads(DSparkDrafter & d, + ggml_backend_t backend); + void free_deepseek4_dspark_drafter(DSparkDrafter & d); // Drop thread-local graph/allocation state that retains drafter tensor and @@ -108,6 +119,34 @@ bool deepseek4_dspark_draft_forward(ggml_backend_t backend, std::vector & out_hidden, std::vector * confidence_hidden = nullptr); +// Submit the same cached drafter graph without synchronizing the backend. +// The caller must use a backend/stream that is independent from target +// verification and call deepseek4_dspark_draft_wait() before reusing inputs. +bool deepseek4_dspark_draft_forward_async(ggml_backend_t backend, + const DSparkDrafter & d, + const float * noise_embed, + const float * ctx_features, + int ctx_len, + int committed); +// Submit with the feature/context tensors retained from the preceding forward. +// Only the noise block and block positions are updated. This is valid only +// when the same drafter graph and ctx_len are already warm on this backend. +// It is intended for an exact-verifier draft-ahead branch whose context may be +// one speculative step stale; stale draft inputs can affect acceptance, never +// target correctness. +bool deepseek4_dspark_draft_forward_async_reuse_context( + ggml_backend_t backend, + const DSparkDrafter & d, + const float * noise_embed, + int ctx_len, + int committed); +// Copy the output of a completed async drafter forward. The caller must wait +// for the drafter backend before calling this function. +bool deepseek4_dspark_draft_read_async_output( + ggml_backend_t backend, + std::vector & out_hidden); +void deepseek4_dspark_draft_wait(ggml_backend_t backend); + // Batched target verify forward WITH feature capture (defined in // deepseek4_graph.cpp so it can reuse the target sub-builders). Runs the DS4 // target over `n_tokens` embeddings at absolute position `kv_start` in one @@ -131,7 +170,10 @@ bool deepseek4_dspark_verify_forward(ggml_backend_t backend, std::vector * logits_out, std::vector & capture_out, DeepSeek4StepTelemetry * telemetry = nullptr, - bool allow_graph_reuse = false); + bool allow_graph_reuse = false, + MoeHybridStorage * moe_hybrid = nullptr, + MoeExpertComputeRuntime * expert_runtime = nullptr, + MoeHybridRoutingStats * routing_stats = nullptr); // Minimal speculative-decode rollback state. Rejected positions must restore // the physical SWA rows they overwrote after the ring wraps; otherwise a later @@ -140,13 +182,28 @@ bool deepseek4_dspark_verify_forward(ggml_backend_t backend, struct DeepSeek4SpecRollback { int raw_pos = 0; int raw_count = 0; + struct PinnedSpan { + std::size_t offset = 0; + std::size_t size = 0; + }; struct Layer { std::vector attn_kv, attn_sc, idx_kv, idx_sc; std::size_t raw_row_bytes = 0; std::vector raw_rows; + PinnedSpan pinned_attn_kv, pinned_attn_sc, pinned_idx_kv, pinned_idx_sc; + PinnedSpan pinned_raw_rows; }; std::vector layers; std::vector hc_state; + PinnedSpan pinned_hc; + ggml_backend_buffer_t pinned_buf = nullptr; + uint8_t * pinned_base = nullptr; + ggml_backend_t async_backend = nullptr; + + DeepSeek4SpecRollback() = default; + ~DeepSeek4SpecRollback(); + DeepSeek4SpecRollback(const DeepSeek4SpecRollback &) = delete; + DeepSeek4SpecRollback & operator=(const DeepSeek4SpecRollback &) = delete; }; void deepseek4_spec_rollback_save(const DeepSeek4Cache & cache, @@ -178,7 +235,10 @@ bool run_deepseek4_dspark_spec_decode( std::vector & out_tokens, float * accept_rate_out, const std::function & on_token = {}, - DFlashDraftIpcClient * remote_draft = nullptr); + DFlashDraftIpcClient * remote_draft = nullptr, + MoeHybridStorage * moe_hybrid = nullptr, + MoeExpertComputeRuntime * expert_runtime = nullptr, + MoeHybridRoutingStats * routing_stats = nullptr); int run_deepseek4_dspark_draft_ipc_daemon(const char * draft_path, int ring_cap, diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index 5a09f5ada..28bb9fbd5 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -31,6 +31,7 @@ #include "ggml-backend.h" #include "ggml-cpu.h" +#include #include #include #include @@ -45,9 +46,14 @@ class DeepSeek4DFlashTarget : public DFlashTarget { public: DeepSeek4DFlashTarget(const DeepSeek4Weights & w, DeepSeek4Cache & cache, ggml_backend_t backend, ggml_backend_t snap_backend, - std::vector capture_ids, int mask_tok) + std::vector capture_ids, int mask_tok, + MoeHybridStorage * moe_hybrid, + MoeExpertComputeRuntime * expert_runtime, + MoeHybridRoutingStats * routing_stats) : w_(w), cache_(cache), backend_(backend), snap_backend_(snap_backend), - capture_ids_(std::move(capture_ids)), mask_tok_(mask_tok) {} + capture_ids_(std::move(capture_ids)), mask_tok_(mask_tok), + moe_hybrid_(moe_hybrid), expert_runtime_(expert_runtime), + routing_stats_(routing_stats) {} ~DeepSeek4DFlashTarget() override { clear_snapshot(); } @@ -86,7 +92,9 @@ class DeepSeek4DFlashTarget : public DFlashTarget { tokens.data() + t, 1, base_pos + t, am1, keep_logits_ ? &logits1 : nullptr, feat1, telemetry_, - /*allow_graph_reuse=*/false)) { + /*allow_graph_reuse=*/false, + moe_hybrid_, expert_runtime_, + routing_stats_)) { return false; } if (am1.empty()) return false; @@ -108,7 +116,9 @@ class DeepSeek4DFlashTarget : public DFlashTarget { embed_buf_.data(), tokens.data(), n, base_pos, am, keep_logits_ ? &verify_logits_ : nullptr, verify_features_, telemetry_, - /*allow_graph_reuse=*/n > 1)) { + /*allow_graph_reuse=*/n > 1, + moe_hybrid_, expert_runtime_, + routing_stats_)) { return false; } if (am.empty()) return false; @@ -206,6 +216,9 @@ class DeepSeek4DFlashTarget : public DFlashTarget { std::vector embed_buf_; std::vector verify_logits_; std::vector verify_features_; + MoeHybridStorage * moe_hybrid_ = nullptr; + MoeExpertComputeRuntime * expert_runtime_ = nullptr; + MoeHybridRoutingStats * routing_stats_ = nullptr; }; namespace { @@ -238,52 +251,190 @@ constexpr float kConfidenceQ3Threshold = 0.40f; constexpr float kConfidenceQ4Threshold = 0.30f; // ── Light rollback state ──────────────────────────────────────────────── +// Save the ratio-4 rolling state, HC state, and the raw SWA rows that a q<=4 +// verify may overwrite. Pinned host storage lets the GPU copy this compact +// rollback state on its stream before the verifier without a host fence. // prev-half = first 4 rows of a [comp_width, 8] ratio-4 rolling state. // ratio-128 states ([comp_width, 128]) are pure position rings -> skip. -void save_prev_half(ggml_tensor * t, std::vector & buf) { +size_t prev_half_bytes(const ggml_tensor * t) { + return t && t->ne[1] == 8 ? (size_t) t->nb[1] * 4 : 0; +} + +size_t align_up_rollback(size_t value, size_t alignment) { + return (value + alignment - 1) / alignment * alignment; +} + +void assign_pinned_span(DeepSeek4SpecRollback::PinnedSpan & span, size_t bytes, + size_t & total) { + if (bytes == 0) { + span = {}; + return; + } + total = align_up_rollback(total, 64); + span.offset = total; + span.size = bytes; + total += bytes; +} + +bool init_pinned_rollback(const DeepSeek4Cache & cache, DeepSeek4SpecRollback & rb, + ggml_backend_t backend) { + if (rb.pinned_buf) return true; + + rb.layers.resize(cache.layers.size()); + size_t total = 0; + for (size_t il = 0; il < cache.layers.size(); ++il) { + const DeepSeek4LayerCache & lc = cache.layers[il]; + DeepSeek4SpecRollback::Layer & s = rb.layers[il]; + assign_pinned_span( + s.pinned_attn_kv, prev_half_bytes(lc.attn_compressor.state_kv), total); + assign_pinned_span( + s.pinned_attn_sc, prev_half_bytes(lc.attn_compressor.state_score), total); + assign_pinned_span( + s.pinned_idx_kv, prev_half_bytes(lc.indexer_compressor.state_kv), total); + assign_pinned_span( + s.pinned_idx_sc, prev_half_bytes(lc.indexer_compressor.state_score), total); + s.raw_row_bytes = lc.raw_kv + ? ggml_row_size(lc.raw_kv->type, lc.raw_kv->ne[0]) : 0; + assign_pinned_span(s.pinned_raw_rows, s.raw_row_bytes * 4, total); + } + assign_pinned_span( + rb.pinned_hc, cache.hc_state ? ggml_nbytes(cache.hc_state) : 0, total); + if (total == 0) return false; + + ggml_backend_dev_t device = ggml_backend_get_device(backend); + ggml_backend_buffer_type_t buft = + device ? ggml_backend_dev_host_buffer_type(device) : nullptr; + if (!buft) return false; + + rb.pinned_buf = ggml_backend_buft_alloc_buffer(buft, total); + if (!rb.pinned_buf) return false; + rb.pinned_base = + static_cast(ggml_backend_buffer_get_base(rb.pinned_buf)); + if (!rb.pinned_base) { + ggml_backend_buffer_free(rb.pinned_buf); + rb.pinned_buf = nullptr; + return false; + } + return true; +} + +void save_prev_half(ggml_backend_t backend, ggml_tensor * t, + std::vector & buf, bool async_copy) { if (!t || t->ne[1] != 8) { buf.clear(); return; } const size_t bytes = (size_t) t->nb[1] * 4; if (buf.size() != bytes) buf.resize(bytes); - ggml_backend_tensor_get(t, buf.data(), 0, bytes); + if (async_copy) { + ggml_backend_tensor_get_async(backend, t, buf.data(), 0, bytes); + } else { + ggml_backend_tensor_get(t, buf.data(), 0, bytes); + } } -void restore_prev_half(ggml_tensor * t, const std::vector & buf) { +void restore_prev_half(ggml_backend_t backend, ggml_tensor * t, + const std::vector & buf, bool async_copy) { if (!t || buf.empty()) return; - ggml_backend_tensor_set(t, buf.data(), 0, buf.size()); + if (async_copy) { + ggml_backend_tensor_set_async(backend, t, buf.data(), 0, buf.size()); + } else { + ggml_backend_tensor_set(t, buf.data(), 0, buf.size()); + } +} + +void save_prev_half_pinned(ggml_backend_t backend, ggml_tensor * t, + uint8_t * base, + const DeepSeek4SpecRollback::PinnedSpan & span) { + if (!t || span.size == 0) return; + GGML_ASSERT(span.size == prev_half_bytes(t)); + ggml_backend_tensor_get_async( + backend, t, base + span.offset, 0, span.size); +} + +void restore_prev_half_pinned(ggml_backend_t backend, ggml_tensor * t, + const uint8_t * base, + const DeepSeek4SpecRollback::PinnedSpan & span) { + if (!t || span.size == 0) return; + GGML_ASSERT(span.size == prev_half_bytes(t)); + ggml_backend_tensor_set_async( + backend, t, base + span.offset, 0, span.size); } -void spec_rollback_save_impl(const DeepSeek4Cache & cache, - DeepSeek4SpecRollback & rb, - int raw_pos, int raw_count) { +void spec_rollback_save(const DeepSeek4Cache & cache, DeepSeek4SpecRollback & rb, + ggml_backend_t backend, bool async_copy, + bool pinned_copy, int raw_pos, int raw_count) { rb.raw_pos = raw_pos; - rb.raw_count = std::max(0, raw_count); + rb.raw_count = std::clamp(raw_count, 0, 4); rb.layers.resize(cache.layers.size()); + if (async_copy || pinned_copy) { + rb.async_backend = backend; + } + const bool use_pinned = + pinned_copy && init_pinned_rollback(cache, rb, backend); for (size_t il = 0; il < cache.layers.size(); ++il) { const DeepSeek4LayerCache & lc = cache.layers[il]; DeepSeek4SpecRollback::Layer & s = rb.layers[il]; - save_prev_half(lc.attn_compressor.state_kv, s.attn_kv); - save_prev_half(lc.attn_compressor.state_score, s.attn_sc); - save_prev_half(lc.indexer_compressor.state_kv, s.idx_kv); - save_prev_half(lc.indexer_compressor.state_score, s.idx_sc); - s.raw_row_bytes = 0; - s.raw_rows.clear(); - if (lc.raw_kv && lc.raw_kv->ne[1] > 0 && rb.raw_count > 0) { - s.raw_row_bytes = ggml_row_size(lc.raw_kv->type, lc.raw_kv->ne[0]); + if (use_pinned) { + save_prev_half_pinned( + backend, lc.attn_compressor.state_kv, + rb.pinned_base, s.pinned_attn_kv); + save_prev_half_pinned( + backend, lc.attn_compressor.state_score, + rb.pinned_base, s.pinned_attn_sc); + save_prev_half_pinned( + backend, lc.indexer_compressor.state_kv, + rb.pinned_base, s.pinned_idx_kv); + save_prev_half_pinned( + backend, lc.indexer_compressor.state_score, + rb.pinned_base, s.pinned_idx_sc); + } else { + save_prev_half(backend, lc.attn_compressor.state_kv, s.attn_kv, async_copy); + save_prev_half(backend, lc.attn_compressor.state_score, s.attn_sc, async_copy); + save_prev_half(backend, lc.indexer_compressor.state_kv, s.idx_kv, async_copy); + save_prev_half(backend, lc.indexer_compressor.state_score, s.idx_sc, async_copy); + } + + s.raw_row_bytes = lc.raw_kv + ? ggml_row_size(lc.raw_kv->type, lc.raw_kv->ne[0]) : 0; + if (!lc.raw_kv || lc.raw_kv->ne[1] <= 0 || + s.raw_row_bytes == 0 || rb.raw_count == 0) { + s.raw_rows.clear(); + continue; + } + if (!use_pinned) { s.raw_rows.resize(s.raw_row_bytes * (size_t) rb.raw_count); - for (int t = 0; t < rb.raw_count; ++t) { - int row = (rb.raw_pos + t) % (int) lc.raw_kv->ne[1]; - if (row < 0) row += (int) lc.raw_kv->ne[1]; + } + for (int t = 0; t < rb.raw_count; ++t) { + int row = (rb.raw_pos + t) % (int) lc.raw_kv->ne[1]; + if (row < 0) row += (int) lc.raw_kv->ne[1]; + uint8_t * dst = use_pinned + ? rb.pinned_base + s.pinned_raw_rows.offset + + (size_t) t * s.raw_row_bytes + : s.raw_rows.data() + (size_t) t * s.raw_row_bytes; + if (use_pinned || async_copy) { + ggml_backend_tensor_get_async( + backend, lc.raw_kv, dst, + (size_t) row * lc.raw_kv->nb[1], s.raw_row_bytes); + } else { ggml_backend_tensor_get( - lc.raw_kv, - s.raw_rows.data() + (size_t) t * s.raw_row_bytes, + lc.raw_kv, dst, (size_t) row * lc.raw_kv->nb[1], s.raw_row_bytes); } } } if (cache.hc_state) { const size_t bytes = ggml_nbytes(cache.hc_state); - if (rb.hc_state.size() != bytes) rb.hc_state.resize(bytes); - ggml_backend_tensor_get(cache.hc_state, rb.hc_state.data(), 0, bytes); + if (use_pinned) { + ggml_backend_tensor_get_async( + backend, cache.hc_state, + rb.pinned_base + rb.pinned_hc.offset, 0, bytes); + } else { + if (rb.hc_state.size() != bytes) rb.hc_state.resize(bytes); + if (async_copy) { + ggml_backend_tensor_get_async( + backend, cache.hc_state, rb.hc_state.data(), 0, bytes); + } else { + ggml_backend_tensor_get(cache.hc_state, rb.hc_state.data(), 0, bytes); + } + } } } @@ -292,11 +443,11 @@ void spec_rollback_save_impl(const DeepSeek4Cache & cache, // prev-half rows with a chunk containing rejected tokens, so put the // pre-verify rows back. (A boundary strictly inside the committed range is a // legitimate flush and must be kept.) -void spec_rollback_apply_impl(const DeepSeek4SpecRollback & rb, - const DeepSeek4Weights & w, - DeepSeek4Cache & cache, - int commit_pos, - bool restore_prev) { +void spec_rollback_apply(const DeepSeek4SpecRollback & rb, const DeepSeek4Weights & w, + DeepSeek4Cache & cache, int commit_pos, bool restore_prev, + ggml_backend_t backend, bool async_copy, + bool pinned_copy) { + const bool use_pinned = pinned_copy && rb.pinned_buf && rb.pinned_base; cache.cur_pos = commit_pos; for (size_t il = 0; il < cache.layers.size(); ++il) { DeepSeek4LayerCache & lc = cache.layers[il]; @@ -305,10 +456,25 @@ void spec_rollback_apply_impl(const DeepSeek4SpecRollback & rb, if (ratio == 4) lc.n_index_comp = commit_pos / 4; if (restore_prev && il < rb.layers.size()) { const DeepSeek4SpecRollback::Layer & s = rb.layers[il]; - restore_prev_half(lc.attn_compressor.state_kv, s.attn_kv); - restore_prev_half(lc.attn_compressor.state_score, s.attn_sc); - restore_prev_half(lc.indexer_compressor.state_kv, s.idx_kv); - restore_prev_half(lc.indexer_compressor.state_score, s.idx_sc); + if (use_pinned) { + restore_prev_half_pinned( + backend, lc.attn_compressor.state_kv, + rb.pinned_base, s.pinned_attn_kv); + restore_prev_half_pinned( + backend, lc.attn_compressor.state_score, + rb.pinned_base, s.pinned_attn_sc); + restore_prev_half_pinned( + backend, lc.indexer_compressor.state_kv, + rb.pinned_base, s.pinned_idx_kv); + restore_prev_half_pinned( + backend, lc.indexer_compressor.state_score, + rb.pinned_base, s.pinned_idx_sc); + } else { + restore_prev_half(backend, lc.attn_compressor.state_kv, s.attn_kv, async_copy); + restore_prev_half(backend, lc.attn_compressor.state_score, s.attn_sc, async_copy); + restore_prev_half(backend, lc.indexer_compressor.state_kv, s.idx_kv, async_copy); + restore_prev_half(backend, lc.indexer_compressor.state_score, s.idx_sc, async_copy); + } } if (il < rb.layers.size() && lc.raw_kv && lc.raw_kv->ne[1] > 0) { const DeepSeek4SpecRollback::Layer & s = rb.layers[il]; @@ -316,19 +482,47 @@ void spec_rollback_apply_impl(const DeepSeek4SpecRollback & rb, for (int t = first_rejected; t < rb.raw_count && s.raw_row_bytes > 0 && - s.raw_rows.size() >= (size_t) (t + 1) * s.raw_row_bytes; ++t) { int row = (rb.raw_pos + t) % (int) lc.raw_kv->ne[1]; if (row < 0) row += (int) lc.raw_kv->ne[1]; - ggml_backend_tensor_set( - lc.raw_kv, - s.raw_rows.data() + (size_t) t * s.raw_row_bytes, - (size_t) row * lc.raw_kv->nb[1], s.raw_row_bytes); + const uint8_t * src = nullptr; + if (use_pinned && + s.pinned_raw_rows.size >= + (size_t) (t + 1) * s.raw_row_bytes) { + src = rb.pinned_base + s.pinned_raw_rows.offset + + (size_t) t * s.raw_row_bytes; + } else if (s.raw_rows.size() >= + (size_t) (t + 1) * s.raw_row_bytes) { + src = s.raw_rows.data() + (size_t) t * s.raw_row_bytes; + } + if (!src) continue; + if (use_pinned || async_copy) { + ggml_backend_tensor_set_async( + backend, lc.raw_kv, src, + (size_t) row * lc.raw_kv->nb[1], s.raw_row_bytes); + } else { + ggml_backend_tensor_set( + lc.raw_kv, src, + (size_t) row * lc.raw_kv->nb[1], s.raw_row_bytes); + } } } } - if (restore_prev && cache.hc_state && !rb.hc_state.empty()) { - ggml_backend_tensor_set(cache.hc_state, rb.hc_state.data(), 0, rb.hc_state.size()); + if (restore_prev && cache.hc_state) { + if (use_pinned && rb.pinned_hc.size > 0) { + ggml_backend_tensor_set_async( + backend, cache.hc_state, + rb.pinned_base + rb.pinned_hc.offset, 0, rb.pinned_hc.size); + } else if (!rb.hc_state.empty()) { + if (async_copy) { + ggml_backend_tensor_set_async( + backend, cache.hc_state, + rb.hc_state.data(), 0, rb.hc_state.size()); + } else { + ggml_backend_tensor_set( + cache.hc_state, rb.hc_state.data(), 0, rb.hc_state.size()); + } + } } } @@ -340,11 +534,22 @@ double spec_ms_since(SpecClock::time_point t0) { } // namespace +DeepSeek4SpecRollback::~DeepSeek4SpecRollback() { + if (async_backend) { + ggml_backend_synchronize(async_backend); + } + if (pinned_buf) { + ggml_backend_buffer_free(pinned_buf); + } +} + void deepseek4_spec_rollback_save(const DeepSeek4Cache & cache, DeepSeek4SpecRollback & rollback, int raw_pos, int raw_count) { - spec_rollback_save_impl(cache, rollback, raw_pos, raw_count); + spec_rollback_save(cache, rollback, nullptr, + /*async_copy=*/false, /*pinned_copy=*/false, + raw_pos, raw_count); } void deepseek4_spec_rollback_apply(const DeepSeek4SpecRollback & rollback, @@ -352,7 +557,9 @@ void deepseek4_spec_rollback_apply(const DeepSeek4SpecRollback & rollback, DeepSeek4Cache & cache, int commit_pos, bool restore_prev) { - spec_rollback_apply_impl(rollback, weights, cache, commit_pos, restore_prev); + spec_rollback_apply(rollback, weights, cache, commit_pos, restore_prev, + nullptr, /*async_copy=*/false, + /*pinned_copy=*/false); } // Batched target verify + capture: wraps the existing multi-token @@ -371,22 +578,35 @@ bool deepseek4_dspark_verify_forward(ggml_backend_t backend, std::vector * logits_out, std::vector & capture_out, DeepSeek4StepTelemetry * telemetry, - bool allow_graph_reuse) { + bool allow_graph_reuse, + MoeHybridStorage * moe_hybrid, + MoeExpertComputeRuntime * expert_runtime, + MoeHybridRoutingStats * routing_stats) { std::vector hc_state; std::vector all_logits; std::vector last_logits; + std::vector gpu_argmax; Ds4VerifyHooks hooks; hooks.capture_layer_ids = &capture_layer_ids; hooks.capture_out = &capture_out; hooks.all_logits_out = &all_logits; + hooks.argmax_out = &gpu_argmax; + hooks.prefer_argmax_only = + spec_env_flag("DFLASH_DS4_GPU_ARGMAX_VERIFY") && logits_out == nullptr; if (!deepseek4_step_layer_range(backend, w, cache, hc_state, embed, n_tokens, kv_start, 0, w.n_layer, &last_logits, token_ids, telemetry, allow_graph_reuse, - &hooks)) { + &hooks, moe_hybrid, expert_runtime, + routing_stats)) { std::fprintf(stderr, "[ds4-verify] step_layer_range returned false (n_tokens=%d kv_start=%d)\n", n_tokens, kv_start); return false; } + if (hooks.prefer_argmax_only && + (int) gpu_argmax.size() == n_tokens) { + argmax_out = std::move(gpu_argmax); + return true; + } if ((int) all_logits.size() < w.n_vocab * n_tokens) { std::fprintf(stderr, "[ds4-verify] all_logits too small: got=%zu need=%d (cap=%zu)\n", all_logits.size(), w.n_vocab * n_tokens, capture_out.size()); @@ -416,7 +636,10 @@ bool run_deepseek4_dspark_spec_decode( std::vector & out_tokens, float * accept_rate_out, const std::function & on_token, - DFlashDraftIpcClient * remote_draft) { + DFlashDraftIpcClient * remote_draft, + MoeHybridStorage * moe_hybrid, + MoeExpertComputeRuntime * expert_runtime, + MoeHybridRoutingStats * routing_stats) { const int n_embd = target_w.n_embd; const int n_tgt = drafter.n_target_layers; const int block = drafter.block_size; @@ -427,6 +650,32 @@ bool run_deepseek4_dspark_spec_decode( const bool timing = spec_env_flag("DFLASH_DS4_TIMING"); const bool full_snap = spec_env_flag("DFLASH_DS4_FULL_SNAP"); const bool seq_verify_mode = spec_env_flag("DFLASH_DS4_SEQ_VERIFY"); + const bool async_rollback = spec_env_flag("DFLASH_DS4_ASYNC_ROLLBACK"); + const bool pinned_rollback = spec_env_flag("DFLASH_DS4_PINNED_ROLLBACK"); + const bool draft_overlap_probe = + spec_env_flag("DFLASH_DS4_DRAFT_OVERLAP_PROBE"); + const bool draft_overlap_reuse_context = + spec_env_flag("DFLASH_DS4_DRAFT_OVERLAP_REUSE_CONTEXT"); + const bool draft_ahead = + spec_env_flag("DFLASH_DS4_DRAFT_AHEAD"); + ggml_backend_t drafter_backend = + drafter.core.backend ? drafter.core.backend : backend; + const bool draft_overlap_probe_active = + draft_overlap_probe && !remote_draft && drafter_backend != backend; + bool draft_overlap_probe_enabled = draft_overlap_probe_active; + if (draft_overlap_probe && !draft_overlap_probe_active) { + std::fprintf(stderr, + "[ds4-spec] draft overlap probe requested without an independent " + "in-process backend; probe disabled\n"); + } + const bool draft_ahead_active = + draft_ahead && !remote_draft && drafter_backend != backend; + bool draft_ahead_enabled = draft_ahead_active; + if (draft_ahead && !draft_ahead_active) { + std::fprintf(stderr, + "[ds4-spec] draft-ahead requested without an independent " + "in-process backend; draft-ahead disabled\n"); + } // Laguna-style adaptive verify width: EWMA of accepted candidates, width = // ewma + 2 (avg_commit << block means the wide tail is usually wasted). // /tmp/ds4_awidth: 1 = on, 0 = off (default on). @@ -481,7 +730,8 @@ bool run_deepseek4_dspark_spec_decode( if (!snap_backend) { std::fprintf(stderr, "[ds4-spec] no CPU snapshot backend\n"); return false; } DeepSeek4DFlashTarget target(target_w, target_cache, backend, snap_backend, - drafter.capture_layer_ids, drafter.mask_token_id); + drafter.capture_layer_ids, drafter.mask_token_id, + moe_hybrid, expert_runtime, routing_stats); DraftWeights dw = make_dspark_shim(drafter); DeepSeek4SpecRollback rollback; DeepSeek4StepTelemetry tel{}; @@ -540,9 +790,18 @@ bool run_deepseek4_dspark_spec_decode( std::vector padded_confidence_hidden((size_t) n_embd * (block + 1), 0.0f); std::vector draft_tok, tgt_am; std::vector draft_confidence; + std::vector cached_ahead_hidden; + std::vector ahead_noise_embed((size_t) n_embd * block); + bool cached_ahead_ready = false; + int cached_ahead_pos = -1; + int cached_ahead_seed = -1; + long ahead_launches = 0; + long ahead_seed_hits = 0; + long ahead_reuses = 0; // Cumulative phase timings (ms). double tm_draft = 0, tm_head = 0, tm_save = 0, tm_verify = 0, tm_apply = 0, tm_feat = 0; + double tm_probe_submit = 0, tm_probe_wait = 0; const SpecClock::time_point run_t0 = SpecClock::now(); while (n_generated < n_gen) { @@ -551,26 +810,41 @@ bool run_deepseek4_dspark_spec_decode( // Noise block = [seed] + [MASK]*(block-1). SpecClock::time_point t0 = SpecClock::now(); if (q_cap >= 2) { - noise_ids[0] = lt; - for (int i = 1; i < block; i++) noise_ids[i] = drafter.mask_token_id; - if (!target.embed_tokens(noise_ids.data(), block, noise_embed.data())) { - ok = false; - break; - } + const bool use_ahead = + cached_ahead_ready && cached_ahead_pos == pos && + cached_ahead_seed == lt && + cached_ahead_hidden.size() == (size_t) n_embd * block; + if (use_ahead) { + local_hidden.swap(cached_ahead_hidden); + cached_ahead_ready = false; + ++ahead_reuses; + } else { + cached_ahead_ready = false; + cached_ahead_hidden.clear(); + noise_ids[0] = lt; + for (int i = 1; i < block; i++) { + noise_ids[i] = drafter.mask_token_id; + } + if (!target.embed_tokens( + noise_ids.data(), block, noise_embed.data())) { + break; + } - // Drafter forward -> normalized states for token logits plus the - // pre-output-norm states expected by the confidence head. - const bool draft_ok = remote_draft && remote_draft->active() - ? remote_draft->propose(pos, ctx_len, noise_embed, local_hidden) - : deepseek4_dspark_draft_forward( - backend, drafter, noise_embed.data(), - ctx_len > 0 ? feat_win.data() : nullptr, - ctx_len, pos, local_hidden, - use_confidence_width ? &confidence_hidden : nullptr); - if (!draft_ok) { - std::fprintf(stderr, "[ds4-spec] drafter forward failed\n"); - ok = false; - break; + // Drafter forward -> block normed hidden states. + const bool draft_ok = remote_draft && remote_draft->active() + ? remote_draft->propose( + pos, ctx_len, noise_embed, local_hidden) + : deepseek4_dspark_draft_forward( + drafter_backend, + drafter, noise_embed.data(), + ctx_len > 0 ? feat_win.data() : nullptr, + ctx_len, pos, local_hidden, + use_confidence_width ? &confidence_hidden : nullptr); + if (!draft_ok) { + std::fprintf(stderr, "[ds4-spec] drafter forward failed\n"); + ggml_backend_free(snap_backend); + return false; + } } } tm_draft += spec_ms_since(t0); @@ -593,15 +867,29 @@ bool run_deepseek4_dspark_spec_decode( bool ds_ok = false; // Batched-verify exactness: the batch must not cross a ratio-4 // boundary except at its last token (state rows stay distinct and the - // comp emission matches AR). Boundaries sit at p % 4 == 3. The legacy - // sequential measurement path is independently capped to q=4 above. - int q_step_cap = seq_verify_mode + // comp emission matches AR). Boundaries sit at p % 4 == 3. The + // sequential verify has no such limit. + // The fused verifier handles a ratio-4 compressor boundary at any + // position inside the batch. Keep the legacy boundary clamp only for + // the dynamic batched path; otherwise "fixed q3" degenerates into q1 + // whenever pos % 4 == 3 and pays a full target verify for one token. + static const bool fused_verify_mode = [] { + const char * v = std::getenv("DFLASH_DS4_FUSED_VERIFY"); + return v && *v && *v != '0'; + }(); + int q_step_cap = (seq_verify_mode || fused_verify_mode) ? std::min(q_cap, 4) : std::min(q_cap, 4 - (pos & 3)); if (adaptive_width && !use_confidence_width && !seq_verify_mode) { const int w_cap = (int) ewma_accept + 2; if (w_cap < q_step_cap) q_step_cap = w_cap; } + // Draft-ahead needs one unverified token beyond the verifier width. + // It predicts only the next seed; the target remains authoritative. + const int head_step_cap = draft_ahead_enabled && q_step_cap >= 2 + ? std::min(q_step_cap + 1, block + 1) + : q_step_cap; + int predicted_ahead_seed = -1; if (q_step_cap >= 2) { std::memcpy(padded_hidden.data() + n_embd, local_hidden.data(), sizeof(float) * (size_t) n_embd * block); @@ -612,23 +900,30 @@ bool run_deepseek4_dspark_spec_decode( } ds_ok = dspark_markov_correct_greedy_chain_fused( dw, backend, target.lm_head_tensor(), padded_hidden.data(), - q_step_cap, lt, draft_tok, + head_step_cap, lt, draft_tok, use_confidence_width ? &draft_confidence : nullptr, - use_confidence_width ? padded_confidence_hidden.data() : nullptr); + use_confidence_width + ? padded_confidence_hidden.data() : nullptr); if (!ds_ok) { ds_ok = dspark_markov_correct_greedy_chain(dw, backend, target, - padded_hidden.data(), q_step_cap, lt, 0.0f, draft_tok); + padded_hidden.data(), head_step_cap, lt, 0.0f, draft_tok); } if (!ds_ok || (int) draft_tok.size() < 2) { // Fallback: plain projection of the block hiddens. std::vector pj; - if (!target.project_hidden_to_tokens(local_hidden.data(), q_step_cap - 1, pj)) { - ok = false; + if (!target.project_hidden_to_tokens( + local_hidden.data(), head_step_cap - 1, pj)) { break; } draft_tok.clear(); draft_tok.push_back(lt); - for (int i = 0; i < q_step_cap - 1; i++) draft_tok.push_back(pj[i]); + for (int i = 0; i < head_step_cap - 1; i++) { + draft_tok.push_back(pj[(size_t) i]); + } + } + if ((int) draft_tok.size() > q_step_cap) { + predicted_ahead_seed = draft_tok[(size_t) q_step_cap]; + draft_tok.resize((size_t) q_step_cap); } } else { draft_tok.push_back(lt); // q=1: seed only, no speculation @@ -663,6 +958,61 @@ bool run_deepseek4_dspark_spec_decode( q > 2 ? draft_tok[2] : -1, q > 3 ? draft_tok[3] : -1); } + // Predict the next seed with one extra Markov step, then run its + // drafter forward on the independent low-priority R9700 stream while + // the target verifies this step. The feature window is deliberately + // one verifier step stale. A miss is discarded; every proposal still + // passes through the exact target verifier. + bool ahead_inflight = false; + bool ahead_output_ready = false; + int inflight_ahead_seed = -1; + int inflight_ahead_pos = -1; + std::vector completed_ahead_hidden; + if (draft_ahead_enabled && q == q_step_cap && q >= 2 && + predicted_ahead_seed >= 0) { + const SpecClock::time_point probe_t0 = SpecClock::now(); + noise_ids[0] = predicted_ahead_seed; + for (int i = 1; i < block; ++i) { + noise_ids[i] = drafter.mask_token_id; + } + if (target.embed_tokens( + noise_ids.data(), block, ahead_noise_embed.data())) { + inflight_ahead_seed = predicted_ahead_seed; + inflight_ahead_pos = pos + q; + ahead_inflight = deepseek4_dspark_draft_forward_async( + drafter_backend, drafter, ahead_noise_embed.data(), + ctx_len > 0 ? feat_win.data() : nullptr, + ctx_len, inflight_ahead_pos); + } + tm_probe_submit += spec_ms_since(probe_t0); + if (ahead_inflight) { + ++ahead_launches; + } else { + draft_ahead_enabled = false; + std::fprintf(stderr, + "[ds4-spec] draft-ahead launch failed; disabling\n"); + } + } + + // Feasibility-only control: duplicate the current draft and discard + // it. Never run this at the same time as real draft-ahead. + bool probe_inflight = false; + if (!draft_ahead_enabled && draft_overlap_probe_enabled && q_cap >= 2) { + const SpecClock::time_point probe_t0 = SpecClock::now(); + probe_inflight = draft_overlap_reuse_context + ? deepseek4_dspark_draft_forward_async_reuse_context( + drafter_backend, drafter, noise_embed.data(), ctx_len, pos) + : deepseek4_dspark_draft_forward_async( + drafter_backend, drafter, noise_embed.data(), + ctx_len > 0 ? feat_win.data() : nullptr, ctx_len, pos); + tm_probe_submit += spec_ms_since(probe_t0); + if (!probe_inflight) { + draft_overlap_probe_enabled = false; + std::fprintf(stderr, + "[ds4-spec] draft overlap probe launch failed; disabling\n"); + } + } + // ── Rollback state save (cheap) or legacy full snapshot ── t0 = SpecClock::now(); if (full_snap) { @@ -672,7 +1022,10 @@ bool run_deepseek4_dspark_spec_decode( break; } } else { - deepseek4_spec_rollback_save(target_cache, rollback, pos, q); + spec_rollback_save( + target_cache, rollback, backend, + async_rollback || pinned_rollback, pinned_rollback, + pos, q); } tm_save += spec_ms_since(t0); @@ -683,20 +1036,39 @@ bool run_deepseek4_dspark_spec_decode( // ── ONE batched verify (writes cache + captures features for all q) ── t0 = SpecClock::now(); int verify_last = -1; - if (!target.verify_batch(draft_tok, pos, verify_last, &tgt_am)) { + const bool verify_ok = + target.verify_batch(draft_tok, pos, verify_last, &tgt_am); + tm_verify += spec_ms_since(t0); + if (ahead_inflight || probe_inflight) { + const SpecClock::time_point probe_t0 = SpecClock::now(); + deepseek4_dspark_draft_wait(drafter_backend); + if (ahead_inflight) { + ahead_output_ready = + deepseek4_dspark_draft_read_async_output( + drafter_backend, completed_ahead_hidden); + if (!ahead_output_ready) { + draft_ahead_enabled = false; + std::fprintf(stderr, + "[ds4-spec] draft-ahead output read failed; disabling\n"); + } + } + tm_probe_wait += spec_ms_since(probe_t0); + } + if (!verify_ok) { if (full_snap) { if (!target.restore_kv()) { std::fprintf(stderr, "[ds4-spec] restore after verify failure failed\n"); } } else { - deepseek4_spec_rollback_apply( - rollback, target_w, target_cache, pos, boundary_crossed); + spec_rollback_apply( + rollback, target_w, target_cache, pos, boundary_crossed, + backend, async_rollback || pinned_rollback, + pinned_rollback); } std::fprintf(stderr, "[ds4-spec] verify failed\n"); ok = false; break; } - tm_verify += spec_ms_since(t0); // Accept the longest matching prefix. accept counts the seed (slot 0) // plus each candidate the target agrees with. @@ -709,6 +1081,29 @@ bool run_deepseek4_dspark_spec_decode( const int bonus = tgt_am[accept - 1]; // target's token at the accept point const int commit_pos = pos + accept; // seed + accepted candidates in KV + if (ahead_output_ready) { + const bool seed_hit = + accept == q && inflight_ahead_pos == commit_pos && + inflight_ahead_seed == bonus; + if (seed_hit) { + ++ahead_seed_hits; + cached_ahead_hidden = std::move(completed_ahead_hidden); + cached_ahead_pos = commit_pos; + cached_ahead_seed = bonus; + cached_ahead_ready = true; + } else { + cached_ahead_ready = false; + cached_ahead_hidden.clear(); + } + if (timing && steps < 8) { + std::fprintf(stderr, + "[ds4-draft-ahead] step=%ld predicted_seed=%d bonus=%d " + "full_accept=%d hit=%d\n", + steps, inflight_ahead_seed, bonus, + accept == q ? 1 : 0, seed_hit ? 1 : 0); + } + } + if (timing && steps < 8 && q >= 2) { // Alignment probe: draft candidate i should match tgt_am[i-1]. A // consistent draft[i]==tgt_am[i] pattern instead = off-by-one. @@ -745,8 +1140,10 @@ bool run_deepseek4_dspark_spec_decode( // The prev-half flush is bad only if the boundary sits at-or-past // the commit point (its chunk then contains rejected tokens). const bool restore_prev = boundary_crossed && first_boundary >= commit_pos; - deepseek4_spec_rollback_apply( - rollback, target_w, target_cache, commit_pos, restore_prev); + spec_rollback_apply( + rollback, target_w, target_cache, commit_pos, restore_prev, + backend, async_rollback || pinned_rollback, + pinned_rollback); } // accept == q on the fast path: cur_pos/n_comp already exact, keep. tm_apply += spec_ms_since(t0); @@ -792,10 +1189,12 @@ bool run_deepseek4_dspark_spec_decode( if (timing && (steps <= 4 || (steps & 31) == 0)) { std::fprintf(stderr, "[ds4-spec-t] step=%ld q=%d acc=%d | draft=%.1f head=%.1f save=%.1f " - "verify=%.1f apply=%.1f feat=%.1f ms (cum means)\n", + "verify=%.1f probe(submit/wait)=%.1f/%.1f apply=%.1f feat=%.1f " + "ms (cum means)\n", steps, q, accept, tm_draft / steps, tm_head / steps, tm_save / steps, - tm_verify / steps, tm_apply / steps, tm_feat / steps); + tm_verify / steps, tm_probe_submit / steps, + tm_probe_wait / steps, tm_apply / steps, tm_feat / steps); } if (hit_eos || stop_requested) break; } @@ -813,14 +1212,25 @@ bool run_deepseek4_dspark_spec_decode( steps ? (double) accept_sum / steps : 0.0, steps ? (double) offered_sum / steps : 0.0, q_cap, (int) full_snap); + if (draft_ahead) { + std::fprintf(stderr, + "[ds4-draft-ahead] launches=%ld seed_hits=%ld reused=%ld " + "seed_hit_rate=%.3f reuse_rate=%.3f\n", + ahead_launches, ahead_seed_hits, ahead_reuses, + ahead_launches > 0 + ? (double) ahead_seed_hits / (double) ahead_launches : 0.0, + steps > 0 ? (double) ahead_reuses / (double) steps : 0.0); + } if (steps > 0) { std::fprintf(stderr, "[ds4-spec-t] TOTAL %.1f ms, %ld steps (%.1f ms/step), %d tok (%.1f tok/s) | " - "means: draft=%.1f head=%.1f save=%.1f verify=%.1f apply=%.1f feat=%.1f ms\n", + "means: draft=%.1f head=%.1f save=%.1f verify=%.1f " + "probe_submit=%.1f probe_wait=%.1f apply=%.1f feat=%.1f ms\n", total_ms, steps, total_ms / steps, n_generated, total_ms > 0 ? n_generated * 1000.0 / total_ms : 0.0, tm_draft / steps, tm_head / steps, tm_save / steps, - tm_verify / steps, tm_apply / steps, tm_feat / steps); + tm_verify / steps, tm_probe_submit / steps, + tm_probe_wait / steps, tm_apply / steps, tm_feat / steps); } if (timing && steps > 0) { const double s = 1000.0 * steps; // us -> ms per-step means From a56f45e1747ef2a37eaafd271374100da2f7a919 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 20 Jul 2026 09:52:25 -0700 Subject: [PATCH 05/15] fix(rocm): support DS4 integer routing graphs --- .../llama.cpp/ggml/src/ggml-cuda/binbcast.cu | 55 +++++++++++++++++++ .../llama.cpp/ggml/src/ggml-cuda/common.cuh | 15 ++++- .../llama.cpp/ggml/src/ggml-cuda/concat.cu | 3 + .../ggml/src/ggml-cuda/vendors/hip.h | 2 + 4 files changed, 74 insertions(+), 1 deletion(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/binbcast.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/binbcast.cu index adb4d5f0c..96e000b11 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/binbcast.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/binbcast.cu @@ -23,6 +23,44 @@ static __device__ __forceinline__ float op_div(const float a, const float b) { return a / b; } +// REPEAT normally goes through the floating-point bin-broadcast path. DS4 +// verification also repeats a small I32 expert-routing LUT across the token +// dimension. Preserve those indices exactly instead of converting them +// through float, which is not lossless for arbitrary int32_t values. +static __global__ void k_repeat_i32_contiguous( + const int32_t * src, + int32_t * dst, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + int64_t ne00, + int64_t ne01, + int64_t ne02, + int64_t ne03, + int64_t n) { + const int64_t stride = (int64_t) blockDim.x * gridDim.x; + for (int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + i < n; i += stride) { + int64_t rem = i; + const int64_t i0 = rem % ne0; + rem /= ne0; + const int64_t i1 = rem % ne1; + rem /= ne1; + const int64_t i2 = rem % ne2; + rem /= ne2; + const int64_t i3 = rem % ne3; + + const int64_t s0 = i0 % ne00; + const int64_t s1 = i1 % ne01; + const int64_t s2 = i2 % ne02; + const int64_t s3 = i3 % ne03; + const int64_t src_i = + (((s3 * ne02 + s2) * ne01 + s1) * ne00 + s0); + dst[i] = src[src_i]; + } +} + template src[0]; + if (src->type == GGML_TYPE_I32 && dst->type == GGML_TYPE_I32) { + GGML_ASSERT(ggml_is_contiguous(src)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + constexpr int threads = 256; + const int64_t n = ggml_nelements(dst); + const int64_t blocks_needed = (n + threads - 1) / threads; + const int blocks = (int) (blocks_needed < 65535 ? blocks_needed : 65535); + k_repeat_i32_contiguous<<>>( + (const int32_t *) src->data, + (int32_t *) dst->data, + dst->ne[0], dst->ne[1], dst->ne[2], dst->ne[3], + src->ne[0], src->ne[1], src->ne[2], src->ne[3], + n); + return; + } ggml_cuda_op_bin_bcast>(dst, dst->src[0], dst, nullptr, dst->src[0]->data, dst->data, ctx.stream()); } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh index 59e06c049..b398c5a17 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh @@ -1417,6 +1417,8 @@ struct ggml_backend_cuda_context { cublasHandle_t cublas_handles[GGML_CUDA_MAX_DEVICES] = {nullptr}; int curr_stream_no = 0; + bool low_priority_streams = false; + int stream_priority = 0; #ifdef USE_CUDA_GRAPH // Map from first_node_ptr to cuda_graph - allows multiple graphs per context @@ -1466,7 +1468,14 @@ struct ggml_backend_cuda_context { cudaStream_t stream(int device, int stream) { if (streams[device][stream] == nullptr) { ggml_cuda_set_device(device); - CUDA_CHECK(cudaStreamCreateWithFlags(&streams[device][stream], cudaStreamNonBlocking)); + if (low_priority_streams) { + CUDA_CHECK(cudaStreamCreateWithPriority( + &streams[device][stream], cudaStreamNonBlocking, + stream_priority)); + } else { + CUDA_CHECK(cudaStreamCreateWithFlags( + &streams[device][stream], cudaStreamNonBlocking)); + } } return streams[device][stream]; } @@ -1512,6 +1521,8 @@ struct ggml_cuda_mm_fusion_args_host { ggml_glu_op glu_op; float glu_param0 = 0.0f; float glu_param1 = 0.0f; + float gate_value_scale = 1.0f; + float x_value_scale = 1.0f; }; struct ggml_cuda_mm_fusion_args_device { const void * x_bias = nullptr; @@ -1520,4 +1531,6 @@ struct ggml_cuda_mm_fusion_args_device { ggml_glu_op glu_op; float glu_param0 = 0.0f; float glu_param1 = 0.0f; + float gate_value_scale = 1.0f; + float x_value_scale = 1.0f; }; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/concat.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/concat.cu index 51a3efa53..774b6e099 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/concat.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/concat.cu @@ -236,6 +236,9 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { case GGML_TYPE_I8: concat_cuda_typed(ctx, src0, src1, dst, dim); break; + case GGML_TYPE_I32: + concat_cuda_typed(ctx, src0, src1, dst, dim); + break; default: GGML_ABORT("unsupported concat type %s", ggml_type_name(src0->type)); } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/vendors/hip.h b/server/deps/llama.cpp/ggml/src/ggml-cuda/vendors/hip.h index 935539f88..6b0bf33f7 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/vendors/hip.h +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/vendors/hip.h @@ -116,6 +116,8 @@ #define CUmemAllocationProp hipMemAllocationProp #define cuDeviceGetAttribute hipDeviceGetAttribute #define cudaStreamCreateWithFlags hipStreamCreateWithFlags +#define cudaStreamCreateWithPriority hipStreamCreateWithPriority +#define cudaDeviceGetStreamPriorityRange hipDeviceGetStreamPriorityRange #define cudaStreamDestroy hipStreamDestroy #define cudaStreamFireAndForget hipStreamFireAndForget #define cudaStreamNonBlocking hipStreamNonBlocking From 94f2fa7e14ab2f1052e3598f138b3f64b2767a73 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 20 Jul 2026 09:52:41 -0700 Subject: [PATCH 06/15] fix(ds4): preserve fused q4 verification across boundaries --- server/src/deepseek4/deepseek4_backend.cpp | 15 +++++--- .../src/deepseek4/deepseek4_dspark_spec.cpp | 3 +- .../src/deepseek4/deepseek4_fused_verify.inc | 2 +- server/src/deepseek4/deepseek4_graph.cpp | 35 +++++++++++++++---- server/src/deepseek4/deepseek4_internal.h | 1 + server/tests/test_deepseek4_unit.cpp | 1 + 6 files changed, 43 insertions(+), 14 deletions(-) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 1cb21d2a5..2a8f563ea 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -1431,7 +1431,8 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, env_flag_enabled("DFLASH_DS4_TP_DYNAMIC_HOTSET"); if (dynamic_hotset) { if (!moe_hybrid_ || !expert_backend_ || !routing_stats_) { - result.error = "dynamic hotset requires in-process heterogeneous MoE TP"; + result.fail(GenerateErrorCode::BackendSpecific, + "dynamic hotset requires in-process heterogeneous MoE TP"); return result; } // A previous request may have installed its own placement. Restore @@ -1443,12 +1444,14 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, std::fprintf(stderr, "[deepseek4-moe-tp] dynamic-hotset restore failed: %s\n", err.c_str()); - result.error = "dynamic hotset restore"; + result.fail(GenerateErrorCode::BackendSpecific, + "dynamic hotset restore"); return result; } if (!routing_stats_->init(w_.n_layer, w_.n_expert, w_.n_expert_used)) { - result.error = "dynamic hotset stats reset"; + result.fail(GenerateErrorCode::BackendSpecific, + "dynamic hotset stats reset"); return result; } } @@ -1477,7 +1480,8 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, const int slots = moe_hybrid_->layers[(size_t) il].hot_active; std::vector ranked = routing_stats_->hot_experts(il, slots); if ((int) ranked.size() != slots) { - result.error = "dynamic hotset empty prefill routes"; + result.fail(GenerateErrorCode::BackendSpecific, + "dynamic hotset empty prefill routes"); return result; } auto & ids = request_hot[(size_t) il]; @@ -1493,7 +1497,8 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, std::fprintf(stderr, "[deepseek4-moe-tp] dynamic-hotset install failed: %s\n", err.c_str()); - result.error = "dynamic hotset install"; + result.fail(GenerateErrorCode::BackendSpecific, + "dynamic hotset install"); return result; } const double adapt_s = elapsed_s(adapt_t0); diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index 28bb9fbd5..91a926819 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -480,8 +480,7 @@ void spec_rollback_apply(const DeepSeek4SpecRollback & rb, const DeepSeek4Weight const DeepSeek4SpecRollback::Layer & s = rb.layers[il]; const int first_rejected = std::max(0, commit_pos - rb.raw_pos); for (int t = first_rejected; - t < rb.raw_count && - s.raw_row_bytes > 0 && + t < rb.raw_count && s.raw_row_bytes > 0; ++t) { int row = (rb.raw_pos + t) % (int) lc.raw_kv->ne[1]; if (row < 0) row += (int) lc.raw_kv->ne[1]; diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index 587e2f91e..ef7d6fb47 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -640,7 +640,7 @@ static bool ds4_build_fused_verify_graph( ggml_tensor * router_weights = nullptr; std::vector router_nodes; if (hash_routed) { - const int route_width = ds4_effective_expert_topk(w); + const int route_width = ds4_effective_expert_count(w); // Keep one q-wide input for host filling and give each wavefront // lane a zero-copy two-column view. ggml_tensor * hids_all = fg.hash_ids[(size_t) il]; diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 8ca8622bb..cf50b04d9 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -3674,6 +3674,12 @@ bool deepseek4_step( // while a variant recurs, which is what the ggml-cuda/HIP graph cache keys // on, enabling graph replay for the bulk of decode steps. +static bool ds4_fused_decode_enabled() { + static const bool enabled = + ds4_env_flag("DFLASH_DS4_FUSED_DECODE"); + return enabled; +} + struct DeepSeek4FusedDecodeGraph { struct RoutePredictionOutput { int source_layer = -1; @@ -4477,13 +4483,25 @@ bool deepseek4_step_layer_range( const int hc_dim = n_hc * n_embd; const bool is_last_shard = (layer_end >= w.n_layer); + const bool fused_hybrid_ready = + moe_hybrid && !expert_runtime && + moe_hybrid->materialized_cold_experts && + moe_hybrid->cold_backend_kind == MoeHybridColdBackend::Gpu && + moe_hybrid->cold_backend && moe_hybrid->cold_backend != backend; + const bool fused_verify_candidate = + (!moe_hybrid || fused_hybrid_ready) && + n_tokens >= 2 && n_tokens <= 4 && verify_hooks && + layer_begin == 0 && is_last_shard && out_logits && + ds4_backend_is_gpu(backend) && ds4_fused_verify_enabled(); + // A dynamic batch may be supplied by callers other than the DSpark // verifier. Split it whenever it spans a learned-compressor boundary: // each sub-forward then writes at most one window and, if present, its // boundary is the final token. This preserves the same pool/rotate order // as sequential execution while retaining safe batched prefixes. const int first_chunk = deepseek4_safe_compressor_batch_tokens(w, kv_start, n_tokens); - if (first_chunk > 0 && first_chunk < n_tokens) { + if (first_chunk > 0 && first_chunk < n_tokens && + !fused_verify_candidate) { const int input_width = layer_begin == 0 ? n_embd : hc_dim; std::vector hc_all; std::vector shard_out_all; @@ -4642,11 +4660,6 @@ bool deepseek4_step_layer_range( const int n_expert_used = ds4_effective_expert_count(w); scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, n_expert_used); - const bool fused_hybrid_ready = - moe_hybrid && !expert_runtime && - moe_hybrid->materialized_cold_experts && - moe_hybrid->cold_backend_kind == MoeHybridColdBackend::Gpu && - moe_hybrid->cold_backend && moe_hybrid->cold_backend != backend; // The batched verifier graph is also the only whole-model graph that can // currently own tensors on both GPU backends. Reuse it for q=1 hybrid // decode so native AR does not fall back to 43 host-synchronized FFN @@ -4694,6 +4707,16 @@ bool deepseek4_step_layer_range( if (telemetry) telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); return true; } + // The generic dynamic graph cannot safely span a learned-compressor + // boundary. A fused candidate deliberately bypassed the splitter + // above because the fused graph models that boundary explicitly. If + // graph construction was unavailable, fail closed instead of running + // an unsafe dynamic batch or silently degrading into q=3 + q=1. + if (first_chunk > 0 && first_chunk < n_tokens) { + std::fprintf(stderr, + "[ds4-fused-verify] boundary-spanning graph unavailable\n"); + return false; + } } std::vector fused_debug_logits; if (!moe_hybrid && n_tokens == 1 && allow_decode_graph_reuse && layer_begin == 0 && is_last_shard && diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 275694cee..c9032bba9 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -66,6 +66,7 @@ struct DeepSeek4StepTelemetry { uint64_t sample_us = 0; uint64_t emit_us = 0; uint64_t full_graph_build_us = 0; + uint64_t full_graph_alloc_us = 0; uint64_t full_graph_set_us = 0; uint64_t full_graph_compute_us = 0; uint64_t full_graph_read_us = 0; diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index ec7c33bd5..34031f293 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -14,6 +14,7 @@ #include #include +#include #include #include From 9bb5a99249e032a9ea8c4ceaaaa89fd3879d9b92 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 20 Jul 2026 09:52:57 -0700 Subject: [PATCH 07/15] perf(rocm): keep graph mismatch tracing opt-in --- server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 5 ----- 1 file changed, 5 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 905e70493..8c488ccc9 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 @@ -3597,11 +3597,6 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx const bool prop_changed = memcmp(&graph->node_props[i], &prop, sizeof(prop)) != 0; - if (!res && prop_changed && getenv("GGML_CUDA_GRAPH_STATS")) { - GGML_LOG_INFO("[graph-mismatch] key=%p node=%d op=%s name=%s\n", - graph_key, i, ggml_op_name(cgraph->nodes[i]->op), - cgraph->nodes[i]->name); - } if (!res && prop_changed && getenv("GGML_CUDA_GRAPH_DIFF_STATS")) { const ggml_cuda_graph::node_properties & old = graph->node_props[i]; From 30423af328f54715055914e3fc4cfb6048b20a1e Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 20 Jul 2026 09:53:07 -0700 Subject: [PATCH 08/15] docs(ds4): document Lucebox heterogeneous profile --- server/docs/DS4.md | 74 ++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 62 insertions(+), 12 deletions(-) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index 25c6769ee..f547dabd9 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -1,10 +1,10 @@ # DeepSeek V4 Flash — DFlash Integration This document describes the current DeepSeek V4 Flash implementation in -DFlash. DeepSeek4 supports a monolithic HIP backend for single-device Strix -Halo systems and a layer-split backend for local or mixed-device deployments. -The opt-in DSpark speculative path can execute proposal blocks in a separate -HIP process while the target remains in the primary process. +DFlash. DeepSeek4 supports a monolithic HIP backend, a layer-split backend, +and an in-process heterogeneous expert-parallel backend for Lucebox systems +with an RDNA4 discrete GPU and a Strix Halo GPU. The DSpark speculative path +can run locally on a selected HIP device or in a separate HIP process. ## Model Architecture @@ -33,6 +33,7 @@ DeepSeek V4 Flash is a 43-layer MoE model with: | HC pre/post CUDA kernel | `src/deepseek4/deepseek4_hc_cuda.cu`, `.h` | | Remote target-shard daemon | `src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp` | | DSpark runtime and remote draft daemon | `src/deepseek4/deepseek4_dspark*.{h,cpp}` | +| Heterogeneous expert storage/evaluation | `src/common/moe_hybrid_{storage,ffn_eval}.*`, `src/common/moe_expert_compute.*` | | Shared target-shard IPC infrastructure | `src/common/target_shard_ipc.*`, `src/placement/remote_target_shard_config.h` | | Backend IPC CLI entry | `src/ipc/backend_ipc_main.cpp` | @@ -46,7 +47,9 @@ DeepSeek V4 Flash is a 43-layer MoE model with: 4. **Shard boundary handoff** — Non-final shards return the updated **full HC state tensor** (`n_tokens × n_hc × n_embd`) to the next shard. 5. **Tail shard completion** — The last shard resumes at its `layer_begin`, runs the remaining layers, then performs the final HC merge, RMSNorm, and `lm_head` projection to produce logits. -The production DeepSeek4 path does **not** use the retired per-expert worker split. The MoE computation stays inside the shard that owns each layer. +The layer-split path keeps MoE computation inside the shard that owns each +layer. The heterogeneous expert-parallel path below is different: every layer +can dispatch routed experts concurrently to two local HIP backends. ## Execution Modes @@ -76,6 +79,45 @@ validated Strix Halo profile: --ds4-expert-top-k 4 ``` +### Lucebox heterogeneous expert parallel + +The in-process Lucebox path keeps dense target work, selected hot experts, and +the DSpark drafter on the R9700 (`hip:0`). It materializes the remaining +experts on Strix Halo (`hip:1`). At every MoE layer, routing is partitioned by +expert ownership, both HIP backends are submitted before the join, and the +owner results are reduced back into the target graph. Hot placement is bounded +by the R9700 memory budget and may use a measured routing profile. + +This is route-level expert parallelism, not layer pipelining and not the +retired per-expert IPC worker. It avoids host IPC and preserves a single target +KV/cache owner. The optional DSpark worker remains available as a compatibility +mode, but the fastest two-GPU profile uses the local drafter on `hip:0`. + +Build one HIP binary for both Lucebox architectures: + +```bash +cmake -S server -B server/build-hip-dual \ + -DDFLASH27B_GPU_BACKEND=hip \ + -DDFLASH27B_HIP_ARCHITECTURES='gfx1151;gfx1201' \ + -DGGML_HIP_GRAPHS=ON \ + -DCMAKE_BUILD_TYPE=Release +cmake --build server/build-hip-dual -j +``` + +The promoted q=4 profile additionally enables fused target verification, +device-side owner joining, small-batch MoE MMVQ kernels, pinned rollback, and +DSpark context-KV reuse. `DFLASH_DS4_TOPK=4` is an explicit approximate +inference policy; omit it when model-default top-6 quality is required. + +On 2026-07-19, the clean dual-architecture build produced steady 50.7 and +50.8 tok/s decode runs on the deterministic integer-list workload (128 output +tokens, q=4, acceptance 1.00). Mean draft time was 6.1 ms and mean target +verification was 70.7--70.9 ms. With the same prompt and the same 1.00 +acceptance, incorrectly splitting boundary-spanning verification into q=3 and +q=1 measured 37.5 tok/s and 98.5 ms verification. These figures establish the +execution-path A/B; they are workload-specific and are not a quality or +held-out benchmark claim. + ### 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. @@ -162,13 +204,22 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. | `DFLASH_DS4_DRAFT_IPC_WORK_DIR` | Scratch directory for the child process. | | `DFLASH_DS4_DRAFT_IPC_REQUIRED` | Fail initialization instead of falling back to a local draft when remote startup fails. | | `DFLASH_DRAFT_IPC_TRANSPORT` | Payload transport: `auto`, `shared`, or `stream`. DeepSeek4 DSpark defaults to `auto`; other draft modes retain their existing default. | +| `DFLASH_DS4_MOE_TP` | Enable routed-expert partitioning between local and expert owners. | +| `DFLASH_DS4_MOE_TP_INPROC` | Use the in-process dual-HIP implementation instead of an IPC expert worker. | +| `DFLASH_DS4_MOE_TP_GPU` | HIP index of the expert-owner GPU; on Lucebox this is normally Strix Halo (`1`). | +| `DFLASH_EXPERT_BUDGET_MB` | R9700 memory budget used to select locally resident hot experts. | +| `DFLASH_DS4_HOTNESS_CSV` | Optional per-layer expert routing profile used for hot placement. | +| `DFLASH_DS4_FUSED_VERIFY` | Enable the persistent q-wide target verification graph. | +| `DFLASH_DS4_TP_DEVICE_JOIN` | Join expert-owner results on device instead of through a host reduction. | +| `DFLASH_DS4_DRAFT_CONTEXT_KV_CACHE` | Reuse DSpark context projection/KV state between speculative steps. | `DFLASH_DS4_TIMING` enables the existing timing banners: - parent / local shard: `[deepseek4-split-timing]` - remote Halo shard: `[deepseek4-target-timing]` -DeepSeek4 no longer uses the old expert-split environment variables or expert-worker tuning knobs. Those retired knobs were removed from the codebase rather than left behind as unsupported debug switches. +The old per-expert IPC worker is retired. The `DFLASH_DS4_MOE_TP*` variables +above configure the newer in-process route-owner implementation. ## DSpark Speculative Decode @@ -223,12 +274,11 @@ set `DFLASH_DS4_SEQ_VERIFY=1` for the slower token-at-a-time verification diagnostic. Neither fused verification nor the separate `--ds4-expert-top-k 4` approximation should be presented as byte-identical AR. -DSpark currently requires monolithic target placement. On HIP, -`--ds4-fused-decode` selects that placement; if the target falls back to hybrid -expert placement, the server logs that DSpark is disabled and continues with -the normal autoregressive path. `--ds4-expert-top-k 4` is a separate, -approximate inference policy used by the validated Strix Halo profile; omit it -to retain the model's default six routed experts. +DSpark can verify against the in-process hybrid expert placement. The target +cache and sampler remain in the parent, while routed target experts execute on +their configured owners. `--ds4-expert-top-k 4` (or the equivalent environment +setting) remains a separate approximate policy; omit it to retain the model's +default six routed experts. On HIP `gfx1151`, enabling DSpark defaults `LUCE_MMVQ_MAX_NCOLS` to `4` when the variable is unset. This keeps the four-row verifier on MMVQ. On a 128 GiB From dfa1f24331aad05e11bba6024d6c1dd1b0ad1435 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 20 Jul 2026 12:19:46 -0700 Subject: [PATCH 09/15] fix(ds4): harden heterogeneous runtime paths Close lifecycle, IPC ordering, routing diagnostics, optional kernel, and fallback correctness gaps found during the PR 505 production review. Preserve the qualified top-4 fast path while making disabled experiment paths fail safely and report real measurements. --- .../llama.cpp/ggml/include/ggml-backend.h | 2 +- .../deps/llama.cpp/ggml/include/ggml-cuda.h | 2 +- .../deps/llama.cpp/ggml/src/ggml-backend.cpp | 20 +- .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 5 +- .../deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu | 81 +++---- server/deps/llama.cpp/ggml/src/ggml.c | 14 +- server/src/common/dflash_draft_ipc.cpp | 33 ++- server/src/common/dflash_draft_ipc_daemon.cpp | 6 +- server/src/common/moe_expert_compute_ipc.cpp | 18 +- server/src/common/moe_hybrid_ffn_eval.cpp | 61 +++-- server/src/common/moe_hybrid_ffn_eval.h | 7 +- server/src/common/moe_hybrid_storage.cpp | 134 +++++++++++ server/src/common/target_shard_ipc_daemon.cpp | 3 +- server/src/deepseek4/deepseek4_backend.cpp | 50 ++-- server/src/deepseek4/deepseek4_dspark.h | 3 +- .../deepseek4_dspark_draft_ipc_daemon.cpp | 81 +++++-- .../src/deepseek4/deepseek4_dspark_spec.cpp | 41 +++- .../src/deepseek4/deepseek4_fused_verify.inc | 223 +++++++++++++++++- server/src/deepseek4/deepseek4_graph.cpp | 61 ++++- server/src/deepseek4/deepseek4_internal.h | 5 +- server/src/deepseek4/deepseek4_loader.cpp | 27 ++- .../qwen35/qwen35_target_shard_ipc_daemon.cpp | 3 +- server/test/test_ds4_dspark_load.cpp | 20 ++ server/test/test_server_unit.cpp | 14 +- server/tests/test_deepseek4_unit.cpp | 5 + 25 files changed, 738 insertions(+), 181 deletions(-) diff --git a/server/deps/llama.cpp/ggml/include/ggml-backend.h b/server/deps/llama.cpp/ggml/include/ggml-backend.h index debc793a5..9a3d5489c 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-backend.h +++ b/server/deps/llama.cpp/ggml/include/ggml-backend.h @@ -375,7 +375,7 @@ extern "C" { // after the producing split, and injects the native handle into the op. GGML_API void ggml_backend_sched_add_deferred_peer_copy_node( ggml_backend_sched_t sched, - const struct ggml_tensor * node); + struct ggml_tensor * node); // // Meta backend diff --git a/server/deps/llama.cpp/ggml/include/ggml-cuda.h b/server/deps/llama.cpp/ggml/include/ggml-cuda.h index 6df853376..e11f462f0 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-cuda.h +++ b/server/deps/llama.cpp/ggml/include/ggml-cuda.h @@ -34,7 +34,7 @@ GGML_BACKEND_API bool ggml_backend_cuda_set_low_priority_stream( // calling thread once a stable graph has already been captured. Callers must // bracket only immutable-topology graphs whose tensor addresses and shapes do // not change; input contents may still be updated in place. -GGML_BACKEND_API void ggml_cuda_set_skip_props_check(bool skip); +GGML_BACKEND_API void ggml_backend_cuda_set_skip_props_check(bool skip); // Capture a scheduler-wide device program around multiple ordinary backend // graph submissions. The handle is opaque and belongs to `backend`; callers diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp index 47242375b..2e621ec55 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp @@ -2150,6 +2150,22 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) { void ggml_backend_sched_reset(ggml_backend_sched_t sched) { GGML_ASSERT(sched); + // Per-graph registrations retain tensor pointers and native events. A + // reset explicitly discards that graph, so carrying either list into the + // next allocation would dereference stale metadata (and leak the events). + sched->n_late_cross_input_split_nodes = 0; + if (sched->n_deferred_peer_copies > 0) { + // A caller may reset after an asynchronous submission. Do not destroy + // producer events while either backend can still reference them. + for (int backend_id = 0; backend_id < sched->n_backends; ++backend_id) { + ggml_backend_synchronize(sched->backends[backend_id]); + } + } + for (int i = 0; i < sched->n_deferred_peer_copies; ++i) { + ggml_backend_event_free(sched->deferred_peer_copies[i].event); + sched->deferred_peer_copies[i] = {}; + } + sched->n_deferred_peer_copies = 0; // reset state for the next run if (!sched->is_reset) { ggml_hash_set_reset(&sched->hash_set); @@ -2315,7 +2331,7 @@ void ggml_backend_sched_add_late_cross_input_split_node( void ggml_backend_sched_add_deferred_peer_copy_node( ggml_backend_sched_t sched, - const struct ggml_tensor * node) { + struct ggml_tensor * node) { GGML_ASSERT(sched); GGML_ASSERT(node); GGML_ASSERT(!sched->is_alloc); @@ -2342,7 +2358,7 @@ void ggml_backend_sched_add_deferred_peer_copy_node( struct ggml_backend_sched_deferred_peer_copy * deferred = &sched->deferred_peer_copies[sched->n_deferred_peer_copies++]; *deferred = {}; - deferred->node = const_cast(node); + deferred->node = node; GGML_ASSERT(deferred->node->src[0]); deferred->source = deferred->node->src[0]; // The full graph is already expanded when nodes are registered. Remove 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 8c488ccc9..98158d49b 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 @@ -68,7 +68,6 @@ #include "ggml.h" #include -#include #include #include #include @@ -4674,9 +4673,9 @@ static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, co static thread_local bool ggml_cuda_skip_props_check = false; -extern "C" void ggml_cuda_set_skip_props_check(bool skip); +extern "C" void ggml_backend_cuda_set_skip_props_check(bool skip); -extern "C" void ggml_cuda_set_skip_props_check(bool skip) { +extern "C" void ggml_backend_cuda_set_skip_props_check(bool skip) { ggml_cuda_skip_props_check = skip; } 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 bc2135848..f16d0c490 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -428,6 +428,12 @@ static bool mmid_grouped_env() { return on; } +static bool mmvq_env_flag(const char * name, bool default_value = false) { + const char * value = std::getenv(name); + if (!value || !value[0]) return default_value; + return std::strcmp(value, "0") != 0; +} + static bool mmid_grouped_type_ok(ggml_type type) { // bit0 = Q4_K, bit1 = Q6_K, bit2 = Q4_0/Q8_0/Q5_K. Default: all validated // types (7); DFLASH_MMID_GROUPED_TYPES is a debug override. @@ -462,10 +468,8 @@ int get_mmvq_mmid_max_batch(ggml_type type, int cc) { // Dedicated multi-token MoE kernel: extend the MUL_MAT_ID ceiling to 16 // tokens on NVIDIA Turing+ for types whose base ceiling is already the // maximum. Types with tuned lower ceilings (per PR 20905) keep them. - static const bool moe_kernel_enabled = []() { - const char * e = std::getenv("DFLASH_CUDA_MMVQ_MOE_KERNEL"); - return !(e && e[0] == '0' && e[1] == '\0'); - }(); + static const bool moe_kernel_enabled = + mmvq_env_flag("DFLASH_CUDA_MMVQ_MOE_KERNEL", true); // NVIDIA: Volta, Ada Lovelace, and Blackwell always use MMVQ for MUL_MAT_ID. if (GGML_CUDA_CC_IS_NVIDIA(cc)) { if (cc == GGML_CUDA_CC_VOLTA || cc >= GGML_CUDA_CC_ADA_LOVELACE) { @@ -1037,8 +1041,7 @@ static __global__ void mul_mat_vec_q_moe( const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst, const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint32_t ncols_dst, const uint32_t nchannels_dst, const uint32_t ids_stride, - const bool skip_duplicate_ids, const bool compact_masked_ids, - const bool aligned_shared_ids) { + const bool compact_masked_ids, const bool aligned_shared_ids) { constexpr int qk = ggml_cuda_type_traits::qk; constexpr int qi = ggml_cuda_type_traits::qi; @@ -1102,18 +1105,14 @@ static __global__ void mul_mat_vec_q_moe( uint32_t n_valid = 0; for (uint32_t c = 0; c < nchannels_dst; ++c) { const int32_t id = ids[c + token_idx * ids_stride]; - const bool duplicate = skip_duplicate_ids && c >= 2 && id >= 0 && - id == ids[c - 2 + token_idx * ids_stride]; - n_valid += id >= 0 && !duplicate; + n_valid += id >= 0; } const bool want_valid = compact_slot < n_valid; uint32_t rank = want_valid ? compact_slot : compact_slot - n_valid; for (uint32_t c = 0; c < nchannels_dst; ++c) { const int32_t id = ids[c + token_idx * ids_stride]; - const bool duplicate = skip_duplicate_ids && c >= 2 && id >= 0 && - id == ids[c - 2 + token_idx * ids_stride]; - const bool valid = id >= 0 && !duplicate; + const bool valid = id >= 0; if (valid == want_valid) { if (rank == 0) { channel_dst = c; @@ -1126,21 +1125,9 @@ static __global__ void mul_mat_vec_q_moe( } else { channel_x_raw = ids[channel_dst + token_idx * ids_stride]; } - // The legacy six-route q4 graph lowers its two-route tail as - // [id4, id5, id4, id5] with weights [w4, w5, 0, 0]. Top-k IDs are unique, - // so an exact two-slot-back duplicate is padding and its expert matvec is - // dead work. Keep the output zero; the downstream zero weight therefore - // preserves the original result while avoiding one third of route work. - int32_t previous_channel_x = -1; - if (skip_duplicate_ids && !has_aligned_id && channel_dst >= 2) { - const int32_t previous_raw = - ids[channel_dst - 2 + token_idx * ids_stride]; - previous_channel_x = previous_raw; - } - const bool duplicate_padding = skip_duplicate_ids && !has_aligned_id && - channel_dst >= 2 && - channel_x_raw >= 0 && channel_x_raw == previous_channel_x; - const bool channel_valid = channel_x_raw >= 0 && !duplicate_padding; + // Expert IDs alone cannot distinguish a zero-weight padding duplicate + // from a valid repeated route. Never suppress a route without its weight. + const bool channel_valid = channel_x_raw >= 0; const uint32_t channel_x = channel_valid ? (uint32_t) channel_x_raw : 0u; const uint32_t channel_y = fastmodulo(channel_dst, nchannels_y); @@ -1651,8 +1638,8 @@ static void mul_mat_vec_q_moe_launch_rpb( const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint32_t ncols_dst, const uint32_t ids_stride, const int warp_size, const int nchannels_dst, cudaStream_t stream, - const bool sparse_warp_blocks, const bool skip_duplicate_ids, - const bool compact_masked_ids, const bool aligned_shared_ids) { + const bool sparse_warp_blocks, const bool compact_masked_ids, + const bool aligned_shared_ids) { static_assert(rows_per_block == 1 || rows_per_block == 2 || rows_per_block == 4 || rows_per_block == 8, @@ -1680,7 +1667,7 @@ static void mul_mat_vec_q_moe_launch_rpb( stride_row_x, stride_col_y, stride_col_dst, stride_channel_x, stride_channel_y, stride_channel_dst, ncols_dst, nchannels_dst, ids_stride, - skip_duplicate_ids, compact_masked_ids, aligned_shared_ids); + compact_masked_ids, aligned_shared_ids); } else { mul_mat_vec_q_moe<<>>( @@ -1688,7 +1675,7 @@ static void mul_mat_vec_q_moe_launch_rpb( stride_row_x, stride_col_y, stride_col_dst, stride_channel_x, stride_channel_y, stride_channel_dst, ncols_dst, nchannels_dst, ids_stride, - skip_duplicate_ids, compact_masked_ids, aligned_shared_ids); + compact_masked_ids, aligned_shared_ids); } } else { if (sparse_warp_blocks) { @@ -1698,7 +1685,7 @@ static void mul_mat_vec_q_moe_launch_rpb( stride_row_x, stride_col_y, stride_col_dst, stride_channel_x, stride_channel_y, stride_channel_dst, ncols_dst, nchannels_dst, ids_stride, - skip_duplicate_ids, compact_masked_ids, aligned_shared_ids); + compact_masked_ids, aligned_shared_ids); } else { mul_mat_vec_q_moe<<>>( @@ -1706,7 +1693,7 @@ static void mul_mat_vec_q_moe_launch_rpb( stride_row_x, stride_col_y, stride_col_dst, stride_channel_x, stride_channel_y, stride_channel_dst, ncols_dst, nchannels_dst, ids_stride, - skip_duplicate_ids, compact_masked_ids, aligned_shared_ids); + compact_masked_ids, aligned_shared_ids); } } } @@ -1724,18 +1711,12 @@ static void mul_mat_vec_q_moe_launch( const char * e = std::getenv("DFLASH_CUDA_MMVQ_MOE_SPARSE_WARP_BLOCKS"); return e && e[0] == '1' && e[1] == '\0'; }(); - static const bool skip_duplicate_ids = []() { - const char * e = std::getenv("DFLASH_CUDA_MMVQ_MOE_SKIP_DUPLICATE_IDS"); - return e && e[0] == '1' && e[1] == '\0'; - }(); static const bool compact_masked_ids = []() { const char * e = std::getenv("DFLASH_CUDA_MMVQ_MOE_COMPACT_MASKED_IDS"); return e && e[0] == '1' && e[1] == '\0'; }(); - static const bool aligned_shared_ids = []() { - const char * e = std::getenv("DFLASH_CUDA_MMVQ_MOE_ALIGN_SHARED_IDS"); - return e && e[0] == '1' && e[1] == '\0'; - }(); + static const bool aligned_shared_ids = + mmvq_env_flag("DFLASH_CUDA_MMVQ_MOE_ALIGN_SHARED_IDS"); static const int tuned_rows_per_block = []() { const char * e = std::getenv("DFLASH_CUDA_MMVQ_MOE_ROWS_PER_BLOCK"); if (!e || !e[0]) return 2; @@ -1770,8 +1751,7 @@ static void mul_mat_vec_q_moe_launch( stride_row_x, stride_col_y, stride_col_dst, \ stride_channel_x, stride_channel_y, stride_channel_dst, \ ncols_dst, ids_stride, warp_size, nchannels_dst, stream, \ - sparse_warp_blocks, skip_duplicate_ids, compact_masked_ids, \ - aligned_shared_ids) + sparse_warp_blocks, compact_masked_ids, aligned_shared_ids) #define GGML_MOE_LAUNCH_TWO_WARP_GROUPS() \ mul_mat_vec_q_moe_launch_rpb( \ @@ -1779,8 +1759,7 @@ static void mul_mat_vec_q_moe_launch( stride_row_x, stride_col_y, stride_col_dst, \ stride_channel_x, stride_channel_y, stride_channel_dst, \ ncols_dst, ids_stride, warp_size, nchannels_dst, stream, \ - sparse_warp_blocks, skip_duplicate_ids, compact_masked_ids, \ - aligned_shared_ids) + sparse_warp_blocks, compact_masked_ids, aligned_shared_ids) #define GGML_MOE_LAUNCH_FP3_PACKED24(WARP_GROUPS) \ mul_mat_vec_q_moe_launch_rpb( \ @@ -1788,8 +1767,7 @@ static void mul_mat_vec_q_moe_launch( stride_row_x, stride_col_y, stride_col_dst, \ stride_channel_x, stride_channel_y, stride_channel_dst, \ ncols_dst, ids_stride, warp_size, nchannels_dst, stream, \ - sparse_warp_blocks, skip_duplicate_ids, compact_masked_ids, \ - aligned_shared_ids) + sparse_warp_blocks, compact_masked_ids, aligned_shared_ids) #define GGML_MOE_LAUNCH_FP2_PACKED32(WARP_GROUPS) \ mul_mat_vec_q_moe_launch_rpb( \ @@ -1797,8 +1775,7 @@ static void mul_mat_vec_q_moe_launch( stride_row_x, stride_col_y, stride_col_dst, \ stride_channel_x, stride_channel_y, stride_channel_dst, \ ncols_dst, ids_stride, warp_size, nchannels_dst, stream, \ - sparse_warp_blocks, skip_duplicate_ids, compact_masked_ids, \ - aligned_shared_ids) + sparse_warp_blocks, compact_masked_ids, aligned_shared_ids) if constexpr (type == GGML_TYPE_Q2_0_ROCMFP2) { if (fp2_packed32 && tuned_rows_per_block == 2 && @@ -1974,10 +1951,8 @@ static void mul_mat_vec_q_switch_ncols_dst( return; } - static const bool use_moe_kernel = []() { - const char * e = std::getenv("DFLASH_CUDA_MMVQ_MOE_KERNEL"); - return !(e && e[0] == '0' && e[1] == '\0'); - }(); + static const bool use_moe_kernel = + mmvq_env_flag("DFLASH_CUDA_MMVQ_MOE_KERNEL", true); if (has_ids && ncols_dst > 1 && (use_moe_kernel || ncols_dst > MMVQ_MAX_BATCH_SIZE)) { // Multi-token MUL_MAT_ID path - dedicated MoE kernel diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index c17cf1921..2fb1df855 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -2094,15 +2094,6 @@ static struct ggml_tensor * ggml_add_impl( struct ggml_tensor * a, struct ggml_tensor * b, bool inplace) { - if (!ggml_can_repeat(b, a)) { - fprintf(stderr, - "[ggml-add-shape] a='%s' type=%s ne=[%" PRId64 ",%" PRId64 ",%" PRId64 ",%" PRId64 "] " - "b='%s' type=%s ne=[%" PRId64 ",%" PRId64 ",%" PRId64 ",%" PRId64 "] inplace=%d\n", - a->name, ggml_type_name(a->type), a->ne[0], a->ne[1], a->ne[2], a->ne[3], - b->name, ggml_type_name(b->type), b->ne[0], b->ne[1], b->ne[2], b->ne[3], - (int) inplace); - ggml_print_backtrace_symbols(); - } GGML_ASSERT(ggml_can_repeat(b, a)); struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); @@ -8178,6 +8169,7 @@ struct ggml_tensor * ggml_ds4_moe_align_ids( struct ggml_tensor * expert_ids) { GGML_ASSERT(expert_ids->type == GGML_TYPE_I32); GGML_ASSERT(ggml_n_dims(expert_ids) == 2); + GGML_ASSERT(ggml_is_contiguous(expert_ids)); struct ggml_tensor * result = ggml_dup_tensor(ctx, expert_ids); result->op = GGML_OP_MOE_FUSED; @@ -8282,6 +8274,10 @@ struct ggml_tensor * ggml_ds4_hc_post_split( GGML_ASSERT(main_block->type == GGML_TYPE_F32); GGML_ASSERT(peer_block->type == GGML_TYPE_F32); GGML_ASSERT(split->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(residual_hc)); + GGML_ASSERT(ggml_is_contiguous(main_block)); + GGML_ASSERT(ggml_is_contiguous(peer_block)); + GGML_ASSERT(ggml_is_contiguous(split)); GGML_ASSERT(n_hc > 0 && n_hc <= 8); const int64_t mix_dim = 2*(int64_t)n_hc + (int64_t)n_hc*n_hc; GGML_ASSERT(ggml_nelements(split) == mix_dim); diff --git a/server/src/common/dflash_draft_ipc.cpp b/server/src/common/dflash_draft_ipc.cpp index ff6cf4f66..e2e13ba83 100644 --- a/server/src/common/dflash_draft_ipc.cpp +++ b/server/src/common/dflash_draft_ipc.cpp @@ -117,6 +117,7 @@ bool DFlashDraftIpcClient::start( #else close(); if (bin.empty() || draft_path.empty() || ring_cap <= 0 || + hidden_size_ <= 0 || block_size_ <= 0 || n_target_layers_ <= 0 || (mode != BackendIpcMode::DFlashDraft && mode != BackendIpcMode::DeepSeek4DSparkDraft)) return false; BackendIpcLaunchConfig launch; @@ -167,10 +168,14 @@ bool DFlashDraftIpcClient::send_feature_block( n_tokens > ring_cap_) { return false; } - const size_t expected = (size_t)n_tokens * (size_t)n_target_layers_ * - (size_t)hidden_size_; + size_t expected = 0; + if (!checked_mul_size((size_t)n_tokens, (size_t)n_target_layers_, expected) || + !checked_mul_size(expected, (size_t)hidden_size_, expected)) { + return false; + } if (!features || feature_count != expected) return false; - const size_t bytes = feature_count * sizeof(float); + size_t bytes = 0; + if (!checked_mul_size(feature_count, sizeof(float), bytes)) return false; if (process_.resolved_payload_transport() == BackendIpcPayloadTransport::Shared) { uint64_t seq = 0; @@ -228,9 +233,13 @@ bool DFlashDraftIpcClient::send_feature_slice( capture_idx >= n_target_layers_ || start_pos < 0 || n_tokens <= 0) { return false; } - const size_t expected = (size_t)n_tokens * hidden_size_; + size_t expected = 0; + if (!checked_mul_size((size_t)n_tokens, (size_t)hidden_size_, expected)) { + return false; + } if (slice.size() != expected) return false; - const size_t bytes = slice.size() * sizeof(float); + size_t bytes = 0; + if (!checked_mul_size(slice.size(), sizeof(float), bytes)) return false; if (process_.resolved_payload_transport() == BackendIpcPayloadTransport::Shared) { uint64_t seq = 0; if (!process_.write_shared_payload(slice.data(), bytes, seq)) { @@ -297,21 +306,27 @@ bool DFlashDraftIpcClient::propose( const int stream_fd = process_.stream_fd(); const int payload_fd = process_.payload_fd(); if (!active_ || !cmd || stream_fd < 0 || committed < 0 || - ctx_len <= 0 || ctx_len > ring_cap_) { + ctx_len <= 0 || ctx_len > ring_cap_ || + (mode_ == BackendIpcMode::DeepSeek4DSparkDraft && + ctx_len > committed)) { std::fprintf(stderr, "draft-ipc propose rejected active=%d cmd=%p stream_fd=%d committed=%d ctx_len=%d ring_cap=%d\n", (int)active_, (void *)cmd, stream_fd, committed, ctx_len, ring_cap_); return false; } - const size_t noise_expected = - (size_t)hidden_size_ * block_size_; + size_t noise_expected = 0; + if (!checked_mul_size( + (size_t)hidden_size_, (size_t)block_size_, noise_expected)) { + return false; + } if (noise_embed.size() != noise_expected) { std::fprintf(stderr, "draft-ipc propose noise size mismatch got=%zu expected=%zu hidden=%d block=%d\n", noise_embed.size(), noise_expected, hidden_size_, block_size_); return false; } - const size_t bytes = noise_embed.size() * sizeof(float); + size_t bytes = 0; + if (!checked_mul_size(noise_embed.size(), sizeof(float), bytes)) return false; if (process_.resolved_payload_transport() == BackendIpcPayloadTransport::Shared) { uint64_t seq = 0; if (!process_.write_shared_payload(noise_embed.data(), bytes, seq)) { diff --git a/server/src/common/dflash_draft_ipc_daemon.cpp b/server/src/common/dflash_draft_ipc_daemon.cpp index 8b2f5cfe6..c39ef9e17 100644 --- a/server/src/common/dflash_draft_ipc_daemon.cpp +++ b/server/src/common/dflash_draft_ipc_daemon.cpp @@ -394,7 +394,8 @@ int run_dflash_draft_ipc_daemon(const char * draft_path, capture_idx < 0 || capture_idx >= n_tgt_layers || start_pos < 0 || n_tokens <= 0 || bytes != expected_bytes || !backend_ipc_payload_in_bounds(0, bytes, shared_payload_capacity) || - header->sequence != sequence || header->bytes != (uint64_t)bytes) { + !backend_ipc_shared_payload_header_matches( + header, sequence, static_cast(bytes))) { std::fprintf(stderr, "[draft-ipc-daemon] bad feature_slice_shared: %s\n", line.c_str()); stream_status(stream_fd, -1); @@ -447,7 +448,8 @@ int run_dflash_draft_ipc_daemon(const char * draft_path, committed < 0 || ctx_len <= 0 || ctx_len > feature_ring.cap || bytes != expected_bytes || !backend_ipc_payload_in_bounds(0, bytes, shared_payload_capacity) || - header->sequence != sequence || header->bytes != (uint64_t)bytes) { + !backend_ipc_shared_payload_header_matches( + header, sequence, static_cast(bytes))) { std::fprintf(stderr, "[draft-ipc-daemon] bad propose_shared: %s\n", line.c_str()); stream_status(stream_fd, -1); diff --git a/server/src/common/moe_expert_compute_ipc.cpp b/server/src/common/moe_expert_compute_ipc.cpp index b1b9edcac..57d36c386 100644 --- a/server/src/common/moe_expert_compute_ipc.cpp +++ b/server/src/common/moe_expert_compute_ipc.cpp @@ -181,6 +181,7 @@ struct RemoteMoeRuntime { int n_expert = 0; int n_expert_used = 0; int n_ff_exp = 0; + float swiglu_clamp = 0.0f; std::vector layers; MoeHybridStorage hybrid; }; @@ -224,6 +225,7 @@ bool build_cached_batched_cold_graph( int n_ff_exp, int n_selected, int n_tokens, + float swiglu_clamp, ggml_type input_type = GGML_TYPE_F32) { out.free(); @@ -263,7 +265,10 @@ bool build_cached_batched_cold_graph( (size_t)n_ff_exp * ggml_element_size(gate_up_e)); gate_e = ggml_cont(out.ctx, gate_e); up_e = ggml_cont(out.ctx, up_e); - gu = ggml_swiglu_split(out.ctx, gate_e, up_e); + gu = swiglu_clamp > 1.0e-6f + ? ggml_swiglu_ds4_split( + out.ctx, gate_e, up_e, swiglu_clamp) + : ggml_swiglu_split(out.ctx, gate_e, up_e); } else { ggml_tensor * gate_e = moe_expert_apply_scale2(out.ctx, ggml_mul_mat_id(out.ctx, gate_tensor, cur_3d, out.ids), @@ -271,7 +276,10 @@ bool build_cached_batched_cold_graph( ggml_tensor * up_e = moe_expert_apply_scale2(out.ctx, ggml_mul_mat_id(out.ctx, up_tensor, cur_3d, out.ids), up_scale); - gu = ggml_swiglu_split(out.ctx, gate_e, up_e); + gu = swiglu_clamp > 1.0e-6f + ? ggml_swiglu_ds4_split( + out.ctx, gate_e, up_e, swiglu_clamp) + : ggml_swiglu_split(out.ctx, gate_e, up_e); } ggml_tensor * experts = moe_expert_apply_scale2(out.ctx, @@ -648,6 +656,7 @@ bool build_remote_moe_runtime_from_weights( out.n_ff_exp = weights.n_ff_exp; MoeHybridConfig cfg = make_moe_hybrid_config(weights); + out.swiglu_clamp = cfg.swiglu_clamp; std::vector layer_descs((size_t)weights.n_layer); out.layers.resize((size_t)weights.n_layer); for (int il = 0; il < weights.n_layer; ++il) { @@ -1765,7 +1774,8 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, L.gate_scale, L.up_scale, L.down_scale, L.gate_up_scale, runtime.n_embd, runtime.n_ff_exp, - n_selected, n_tokens, input_type)) { + n_selected, n_tokens, runtime.swiglu_clamp, + input_type)) { std::fprintf(stderr, "[moe-expert-compute-daemon] batched graph build failed " "layer=%d tokens=%d selected=%d type=%s\n", @@ -1917,7 +1927,7 @@ int run_moe_expert_compute_ipc_daemon(const char * target_path, L.gate_scale, L.up_scale, L.down_scale, L.gate_up_scale, runtime.n_embd, runtime.n_ff_exp, - token_selected)) { + token_selected, runtime.swiglu_clamp)) { std::fprintf(stderr, "[moe-expert-compute-daemon] single graph build failed " "layer=%d tokens=%d selected=%d\n", diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index 9d7cb69a7..996df7e6a 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -3,6 +3,7 @@ #include "ggml-alloc.h" #include "ggml-backend.h" +#include "ggml-cuda.h" #include #include @@ -115,7 +116,16 @@ static bool coarse_owner_split_op_enabled() { static bool align_shared_moe_ids_enabled() { static const bool enabled = [] { const char * raw = std::getenv("DFLASH_CUDA_MMVQ_MOE_ALIGN_SHARED_IDS"); - return raw && *raw && std::strcmp(raw, "0") != 0; + const bool requested = raw && *raw && std::strcmp(raw, "0") != 0; + const char * kernel = std::getenv("DFLASH_CUDA_MMVQ_MOE_KERNEL"); + const bool dedicated_kernel = !kernel || !*kernel || + std::strcmp(kernel, "0") != 0; + if (requested && !dedicated_kernel) { + std::fprintf(stderr, + "[ds4-tp] shared-ID alignment disabled because the dedicated " + "MMVQ MoE kernel is disabled\n"); + } + return requested && dedicated_kernel; }(); return enabled; } @@ -581,7 +591,8 @@ static bool build_batched_routed_graph( float swiglu_clamp, ggml_tensor ** out_routed, bool tokenwise = false, - std::vector * backend_nodes = nullptr) + std::vector * backend_nodes = nullptr, + bool allow_fused_combine = false) { const auto track = [&](ggml_tensor * t) -> ggml_tensor * { if (backend_nodes && t) backend_nodes->push_back(t); @@ -602,7 +613,8 @@ static bool build_batched_routed_graph( gate_scale, up_scale, down_scale, gate_up_scale, inp_col, sel_col, wts_col, n_embd, n_ff_exp, n_used, 1, swiglu_clamp, - &routed_col, false, backend_nodes)) { + &routed_col, false, backend_nodes, + allow_fused_combine)) { return false; } joined = joined ? track(ggml_concat(ctx, joined, routed_col, 1)) @@ -697,7 +709,7 @@ static bool build_batched_routed_graph( ggml_mul_mat_id(ctx, down_tensor, gu, sel), down_scale)); // Weight and sum over experts: [n_embd, n_used, n_tokens] * [1, n_used, n_tokens] - if (fused_moe_combine_enabled()) { + if (allow_fused_combine && fused_moe_combine_enabled()) { *out_routed = track(ggml_laguna_moe_combine(ctx, experts, wts)); return *out_routed != nullptr; } @@ -720,13 +732,14 @@ bool build_moe_hybrid_ffn_graph( ggml_cgraph * schedule_graph, const MoeHybridConfig & cfg, const MoeLayerDesc & desc, - MoeHybridLayerStorage & storage, + const MoeHybridLayerStorage & storage, ggml_tensor * inp, ggml_tensor * global_ids, ggml_tensor * router_weights, int n_tokens, MoeHybridGraphInputs & out, - bool include_shared) { + bool include_shared, + bool allow_fused_combine) { out.output = nullptr; out.main_output = nullptr; @@ -855,7 +868,8 @@ bool build_moe_hybrid_ffn_graph( inp, shard_ids, shard_weights, cfg.n_embd, storage.expert_shard_channels, n_used, n_tokens, cfg.swiglu_clamp, - &shard, tokenwise, &out.shard_nodes)) { + &shard, tokenwise, &out.shard_nodes, + allow_fused_combine)) { return false; } } @@ -877,7 +891,7 @@ bool build_moe_hybrid_ffn_graph( inp, hot_ids, hot_weights, cfg.n_embd, cfg.n_ff_exp, n_used, n_tokens, cfg.swiglu_clamp, &hot, tokenwise, - &out.hot_nodes)) { + &out.hot_nodes, allow_fused_combine)) { return false; } } @@ -903,7 +917,7 @@ bool build_moe_hybrid_ffn_graph( : cfg.n_ff_exp, n_used, n_tokens, cfg.swiglu_clamp, &cold, tokenwise, - &out.cold_nodes)) { + &out.cold_nodes, allow_fused_combine)) { return false; } } @@ -1192,7 +1206,7 @@ bool build_cached_cold_graph( bool build_cached_hot_batched_graph( CachedHotBatchedGraph & out, ggml_backend_t gpu_backend, - MoeHybridLayerStorage & storage, + const MoeHybridLayerStorage & storage, const MoeLayerDesc & desc, const MoeHybridConfig & cfg, int n_tokens) { @@ -1226,7 +1240,9 @@ bool build_cached_hot_batched_graph( build_batched_routed_graph(out.ctx, storage.gate_hot, storage.up_hot, storage.down_hot, storage.gate_up_hot, desc.ffn_gate_exps_s, desc.ffn_up_exps_s, desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, - out.inp, out.sel, out.wts, n_embd, n_ff_exp, n_used, n_tokens, cfg.swiglu_clamp, &routed); + out.inp, out.sel, out.wts, n_embd, n_ff_exp, n_used, n_tokens, + cfg.swiglu_clamp, &routed, false, nullptr, + ggml_backend_is_cuda(gpu_backend)); } // Shared expert (always on GPU) @@ -1257,7 +1273,7 @@ bool build_cached_hot_batched_graph( static bool build_cached_cold_batched_graph( CachedHotBatchedGraph & out, ggml_backend_t cpu_backend, - MoeHybridLayerStorage & storage, + const MoeHybridLayerStorage & storage, const MoeLayerDesc & desc, const MoeHybridConfig & cfg, int n_tokens) { @@ -1286,7 +1302,9 @@ static bool build_cached_cold_batched_graph( build_batched_routed_graph(out.ctx, storage.gate_cold, storage.up_cold, storage.down_cold, storage.gate_up_cold, desc.ffn_gate_exps_s, desc.ffn_up_exps_s, desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, - out.inp, out.sel, out.wts, n_embd, n_ff_exp, n_used, n_tokens, cfg.swiglu_clamp, &routed); + out.inp, out.sel, out.wts, n_embd, n_ff_exp, n_used, n_tokens, + cfg.swiglu_clamp, &routed, false, nullptr, + ggml_backend_is_cuda(cpu_backend)); if (!routed) { out.free(); return false; } out.output = routed; @@ -1818,10 +1836,11 @@ static bool eval_moe_hybrid_ffn_batched_core( CachedHotBatchedGraph & hg = storage.hot_batched_mixed[n_tokens]; const bool hg_ok = (hg.valid() && hg.n_tokens == n_tokens) || build_cached_hot_batched_graph(hg, gpu_backend, storage, desc, cfg, n_tokens); - const bool remote_cold = fp_has_cold && expert_compute && expert_layer; + const bool remote_cold = !skip_cold && fp_has_cold && + expert_compute && expert_layer; CachedHotBatchedGraph * cg = nullptr; bool cg_ok = true; - if (fp_has_cold && !remote_cold) { + if (!skip_cold && fp_has_cold && !remote_cold) { cg = &storage.cold_batched_mixed[n_tokens]; ggml_backend_t cached_cold_backend = storage.cold_backend ? storage.cold_backend : cpu_backend; @@ -1961,7 +1980,9 @@ static bool eval_moe_hybrid_ffn_batched_core( build_batched_routed_graph(hot_ctx, storage.gate_hot, storage.up_hot, storage.down_hot, storage.gate_up_hot, desc.ffn_gate_exps_s, desc.ffn_up_exps_s, desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, - inp, sel, wts, n_embd, n_ff_exp, n_used, n_tokens, cfg.swiglu_clamp, &routed); + inp, sel, wts, n_embd, n_ff_exp, n_used, n_tokens, + cfg.swiglu_clamp, &routed, false, nullptr, + ggml_backend_is_cuda(gpu_backend)); } // Shared expert (always on GPU) @@ -2053,7 +2074,9 @@ static bool eval_moe_hybrid_ffn_batched_core( build_batched_routed_graph(cold_ctx, storage.gate_cold, storage.up_cold, storage.down_cold, storage.gate_up_cold, desc.ffn_gate_exps_s, desc.ffn_up_exps_s, desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, - inp, sel, wts, n_embd, n_ff_exp, n_used, n_tokens, cfg.swiglu_clamp, &cold_routed); + inp, sel, wts, n_embd, n_ff_exp, n_used, n_tokens, + cfg.swiglu_clamp, &cold_routed, false, nullptr, + ggml_backend_is_cuda(cold_backend)); ggml_cgraph * cold_gf = ggml_new_graph_custom(cold_ctx, 4096, false); ggml_set_output(cold_routed); @@ -2314,7 +2337,9 @@ bool eval_moe_hot_only_batched( build_batched_routed_graph(ctx, storage.gate_hot, storage.up_hot, storage.down_hot, storage.gate_up_hot, desc.ffn_gate_exps_s, desc.ffn_up_exps_s, desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, - inp, sel, wts, n_embd, n_ff_exp, n_used, n_tokens, cfg.swiglu_clamp, &routed); + inp, sel, wts, n_embd, n_ff_exp, n_used, n_tokens, + cfg.swiglu_clamp, &routed, false, nullptr, + ggml_backend_is_cuda(gpu_backend)); // Shared expert (always on GPU) ggml_tensor * combined = routed; diff --git a/server/src/common/moe_hybrid_ffn_eval.h b/server/src/common/moe_hybrid_ffn_eval.h index 08c764990..50b612c29 100644 --- a/server/src/common/moe_hybrid_ffn_eval.h +++ b/server/src/common/moe_hybrid_ffn_eval.h @@ -166,13 +166,14 @@ bool build_moe_hybrid_ffn_graph( ggml_cgraph * schedule_graph, const MoeHybridConfig & cfg, const MoeLayerDesc & desc, - MoeHybridLayerStorage & storage, + const MoeHybridLayerStorage & storage, ggml_tensor * inp, ggml_tensor * global_ids, ggml_tensor * router_weights, int n_tokens, MoeHybridGraphInputs & out, - bool include_shared = true); + bool include_shared = true, + bool allow_fused_combine = false); int moe_hybrid_expert_compute_batch_limit(); int moe_hybrid_expert_compute_ipc_batch_limit(int n_tokens); @@ -301,7 +302,7 @@ bool build_cached_cold_graph( bool build_cached_hot_batched_graph( CachedHotBatchedGraph & out, ggml_backend_t gpu_backend, - MoeHybridLayerStorage & storage, + const MoeHybridLayerStorage & storage, const MoeLayerDesc & desc, const MoeHybridConfig & cfg, int n_tokens); diff --git a/server/src/common/moe_hybrid_storage.cpp b/server/src/common/moe_hybrid_storage.cpp index 1e1d6f858..bca3c9473 100644 --- a/server/src/common/moe_hybrid_storage.cpp +++ b/server/src/common/moe_hybrid_storage.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #if defined(DFLASH27B_BACKEND_CUDA) #include @@ -886,6 +887,139 @@ bool moe_hybrid_reassign_hot_experts_from_mmap( return false; } const auto * base = static_cast(storage.mmap_data); + auto source_valid = [&](const ExpertFileRegion & region, + size_t expert_bytes, + int32_t expert) { + if (expert < 0 || expert_bytes == 0 || region.size < expert_bytes || + region.offset > storage.mmap_size) { + return false; + } + const size_t ie = static_cast(expert); + if (ie > (region.size - expert_bytes) / expert_bytes || + ie > (std::numeric_limits::max() - region.offset) / + expert_bytes) { + return false; + } + const size_t source = region.offset + ie * expert_bytes; + return source <= storage.mmap_size && + expert_bytes <= storage.mmap_size - source; + }; + auto destination_valid = [](ggml_tensor * dst, + size_t expert_bytes, + size_t slots) { + return dst && expert_bytes > 0 && + slots <= ggml_nbytes(dst) / expert_bytes; + }; + + // Validate the complete request before changing any layer. A malformed + // later layer must not leave earlier GPU stacks and routing maps updated + // to a placement the caller believes was rejected. + for (size_t il = 0; il < storage.layers.size(); ++il) { + const MoeHybridLayerStorage & layer = storage.layers[il]; + const LayerExpertRegions & regions = storage.layer_regions[il]; + const std::vector & ids = hot_ids_by_layer[il]; + const size_t n_expert = layer.hot_local_by_global.size(); + if (ids.size() != static_cast(layer.hot_active) || + layer.cache_slots != 0 || layer.expert_shard_channels != 0 || + layer.cold_local_by_global.size() != n_expert) { + if (err) *err = "dynamic hotset size does not match fixed hot stack"; + return false; + } + + std::vector desired(n_expert, 0); + for (int32_t expert : ids) { + if (expert < 0 || static_cast(expert) >= n_expert) { + if (err) *err = "dynamic hot expert id out of range"; + return false; + } + if (desired[static_cast(expert)] != 0) { + if (err) *err = "dynamic hot expert ids must be unique"; + return false; + } + desired[static_cast(expert)] = 1; + } + + std::vector entering_hot; + std::vector leaving_hot; + for (int32_t expert : ids) { + if (layer.hot_local_by_global[static_cast(expert)] < 0) { + entering_hot.push_back(expert); + } + } + for (int32_t expert : layer.hot_expert_ids) { + if (expert < 0 || static_cast(expert) >= n_expert) { + if (err) *err = "existing dynamic hot expert id out of range"; + return false; + } + if (!desired[static_cast(expert)]) { + leaving_hot.push_back(expert); + } + } + if (entering_hot.size() != leaving_hot.size()) { + if (err) *err = "dynamic hot/cold swap cardinality mismatch"; + return false; + } + + auto stack_valid = [&](const ExpertFileRegion & region, + ggml_tensor * dst, + size_t expert_bytes) { + if (!destination_valid(dst, expert_bytes, ids.size())) return false; + return std::all_of(ids.begin(), ids.end(), [&](int32_t expert) { + return source_valid(region, expert_bytes, expert); + }); + }; + const bool hot_stack_valid = layer.fused_gate_up + ? stack_valid(regions.gate_up_exps, layer.gate_up_hot, + layer.gate_up_expert_bytes) && + stack_valid(regions.down_exps, layer.down_hot, + layer.down_expert_bytes) + : stack_valid(regions.gate_exps, layer.gate_hot, + layer.gate_expert_bytes) && + stack_valid(regions.up_exps, layer.up_hot, + layer.up_expert_bytes) && + stack_valid(regions.down_exps, layer.down_hot, + layer.down_expert_bytes); + if (!hot_stack_valid) { + if (err) *err = "dynamic hot expert source or destination is invalid"; + return false; + } + + for (size_t si = 0; si < entering_hot.size(); ++si) { + const int32_t incoming = entering_hot[si]; + const int32_t outgoing = leaving_hot[si]; + const int32_t cold_slot = + layer.cold_local_by_global[static_cast(incoming)]; + if (cold_slot < 0 || + static_cast(cold_slot) >= layer.cold_expert_ids.size()) { + if (err) *err = "dynamic hot expert has no cold source slot"; + return false; + } + auto swap_valid = [&](const ExpertFileRegion & region, + ggml_tensor * dst, + size_t expert_bytes) { + return source_valid(region, expert_bytes, outgoing) && + destination_valid( + dst, expert_bytes, + static_cast(cold_slot) + 1); + }; + const bool cold_swap_valid = layer.fused_gate_up + ? swap_valid(regions.gate_up_exps, layer.gate_up_cold, + layer.gate_up_expert_bytes) && + swap_valid(regions.down_exps, layer.down_cold, + layer.down_expert_bytes) + : swap_valid(regions.gate_exps, layer.gate_cold, + layer.gate_expert_bytes) && + swap_valid(regions.up_exps, layer.up_cold, + layer.up_expert_bytes) && + swap_valid(regions.down_exps, layer.down_cold, + layer.down_expert_bytes); + if (!cold_swap_valid) { + if (err) *err = "dynamic cold expert source or destination is invalid"; + return false; + } + } + } + std::vector packed; auto upload = [&](const ExpertFileRegion & region, ggml_tensor * dst, diff --git a/server/src/common/target_shard_ipc_daemon.cpp b/server/src/common/target_shard_ipc_daemon.cpp index 1ee6f02c5..20045be49 100644 --- a/server/src/common/target_shard_ipc_daemon.cpp +++ b/server/src/common/target_shard_ipc_daemon.cpp @@ -140,7 +140,8 @@ int run_target_shard_ipc_daemon_loop( if (shared_payload && shared_payload != MAP_FAILED && shared_payload_data && seq != 0 && n_tokens > 0 && bytes == expected_bytes && backend_ipc_payload_in_bounds(0, bytes, shared_payload_capacity) && - header->sequence == seq && header->bytes == (uint64_t)bytes) { + backend_ipc_shared_payload_header_matches( + header, seq, static_cast(bytes))) { host_act.assign(bytes / sizeof(float), 0.0f); std::memcpy(host_act.data(), shared_payload_data, bytes); payload_ok = true; diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 2a8f563ea..66877ca76 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -90,7 +90,6 @@ static void add_step_tel(DeepSeek4StepTelemetry & dst, const DeepSeek4StepTeleme dst.attn_compute_us += src.attn_compute_us; dst.attn_read_us += src.attn_read_us; dst.full_graph_build_us += src.full_graph_build_us; - dst.full_graph_alloc_us += src.full_graph_alloc_us; dst.full_graph_set_us += src.full_graph_set_us; dst.full_graph_compute_us += src.full_graph_compute_us; dst.full_graph_read_us += src.full_graph_read_us; @@ -116,10 +115,6 @@ static void add_step_tel(DeepSeek4StepTelemetry & dst, const DeepSeek4StepTeleme dst.output_us += src.output_us; dst.sample_us += src.sample_us; dst.emit_us += src.emit_us; - dst.full_graph_build_us += src.full_graph_build_us; - dst.full_graph_set_us += src.full_graph_set_us; - dst.full_graph_compute_us += src.full_graph_compute_us; - dst.full_graph_read_us += src.full_graph_read_us; dst.hot_selected += src.hot_selected; dst.cold_selected += src.cold_selected; } @@ -137,18 +132,17 @@ static void log_step_tel(const char * phase, 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_alloc=%.1fms full_set=%.1fms full_compute=%.1fms full_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 " - "full_build=%.1fms full_set=%.1fms full_compute=%.1fms full_read=%.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_alloc_us), ms(t.full_graph_set_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), @@ -162,8 +156,6 @@ static void log_step_tel(const char * phase, 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), - ms(t.full_graph_build_us), ms(t.full_graph_set_us), ms(t.full_graph_compute_us), - ms(t.full_graph_read_us), t.hot_selected, t.cold_selected); } @@ -333,6 +325,7 @@ static bool fill_profiled_hot_placement(const DeepSeek4Weights & w, // uses authoritative router statistics and evaluates every selected expert. static bool fill_time_balanced_profiled_hot_placement( const DeepSeek4Weights & w, + int active_routes, int hot_per_layer, const char * profile_path, MoeHybridPlacement & out, @@ -347,7 +340,8 @@ static bool fill_time_balanced_profiled_hot_placement( } const int total_hot_budget = hot_per_layer * w.n_layer; - if (total_hot_budget <= 0) { + if (total_hot_budget <= 0 || active_routes <= 0 || + active_routes > w.n_expert_used) { if (err) *err = "time-balanced placement requires a positive hot budget"; return false; } @@ -404,9 +398,9 @@ static bool fill_time_balanced_profiled_hot_placement( const auto layer_cost = [&](int il, int hot_count) { const double hot_fraction = coverage[(size_t) il][(size_t) hot_count]; - const double hot_routes = (double) w.n_expert_used * hot_fraction; + const double hot_routes = (double) active_routes * hot_fraction; const double cold_routes = - (double) w.n_expert_used * (1.0 - hot_fraction); + (double) active_routes * (1.0 - hot_fraction); const double main_time = (shared_equiv + hot_routes) / main_bw; const double cold_time = cold_routes / cold_bw; return std::max(main_time, cold_time); @@ -487,8 +481,9 @@ static bool fill_time_balanced_profiled_hot_placement( hot_counts.begin(), hot_counts.end()); std::fprintf(stderr, "[deepseek4] time-balanced placement: slots=%d layer_range=%d..%d " - "main/cold=%.2f/%.2f GB/s predicted_owner_reduction=%.2f%% counts=", - total_hot_budget, *min_it, *max_it, main_bw, cold_bw, + "routes=%d main/cold=%.2f/%.2f GB/s " + "predicted_owner_reduction=%.2f%% counts=", + total_hot_budget, *min_it, *max_it, active_routes, main_bw, cold_bw, uniform_cost > 0.0 ? 100.0 * (1.0 - balanced_cost / uniform_cost) : 0.0); @@ -812,6 +807,17 @@ bool DeepSeek4Backend::init() { routing_stats_out_path_.c_str()); } } + if (env_flag_enabled("DFLASH_DS4_TP_ROUTE_STATS") && !routing_stats_) { + routing_stats_ = std::make_shared(); + if (!routing_stats_->init(w_.n_layer, w_.n_expert, + w_.n_expert_used)) { + std::fprintf(stderr, + "[deepseek4] failed to initialize TP routing stats\n"); + return false; + } + std::fprintf(stderr, + "[deepseek4-moe-tp] in-memory routing stats enabled\n"); + } if (env_flag_enabled("DFLASH_DS4_TP_DYNAMIC_HOTSET") && !routing_stats_) { routing_stats_ = std::make_shared(); if (!routing_stats_->init(w_.n_layer, w_.n_expert, w_.n_expert_used)) { @@ -912,6 +918,14 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & } const bool all_cold = env_flag_enabled("DFLASH_DS4_MOE_TP_ALL_COLD"); + int active_routes = cfg_.expert_top_k; + if (const char * raw = std::getenv("DFLASH_DS4_TOPK")) { + const int env_routes = std::atoi(raw); + if (env_routes > 0) active_routes = env_routes; + } + if (active_routes <= 0 || active_routes > w.n_expert_used) { + active_routes = w.n_expert_used; + } int hot_per_layer = all_cold ? 0 : budget.max_hot_per_layer; int shard_channels = 0; if (!all_cold && ds4_inprocess_moe_tp_enabled()) { @@ -955,7 +969,7 @@ bool DeepSeek4Backend::compute_uniform_hybrid_placement(const DeepSeek4Weights & env_flag_enabled("DFLASH_DS4_TP_TIME_BALANCED_PLACEMENT"); const bool placed = time_balanced ? fill_time_balanced_profiled_hot_placement( - w, hot_per_layer, profile_path, out, err) + w, active_routes, hot_per_layer, profile_path, out, err) : fill_profiled_hot_placement( w, hot_per_layer, profile_path, out, err); if (!placed) { @@ -1152,6 +1166,10 @@ bool DeepSeek4Backend::unpark(const std::string & what) { free_deepseek4_weights(w_); stream_engine_.destroy(); moe_hybrid_.reset(); + if (expert_backend_) { + ggml_backend_free(expert_backend_); + expert_backend_ = nullptr; + } moe_placement_ = {}; return false; } diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index b086ab9bb..6e9b7da6e 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -144,7 +144,8 @@ bool deepseek4_dspark_draft_forward_async_reuse_context( // for the drafter backend before calling this function. bool deepseek4_dspark_draft_read_async_output( ggml_backend_t backend, - std::vector & out_hidden); + std::vector & out_hidden, + std::vector * confidence_hidden = nullptr); void deepseek4_dspark_draft_wait(ggml_backend_t backend); // Batched target verify forward WITH feature capture (defined in diff --git a/server/src/deepseek4/deepseek4_dspark_draft_ipc_daemon.cpp b/server/src/deepseek4/deepseek4_dspark_draft_ipc_daemon.cpp index 24cdcbef3..df233e160 100644 --- a/server/src/deepseek4/deepseek4_dspark_draft_ipc_daemon.cpp +++ b/server/src/deepseek4/deepseek4_dspark_draft_ipc_daemon.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -25,6 +26,18 @@ namespace dflash::common { +namespace { + +bool checked_mul_size(size_t a, size_t b, size_t & out) { + if (a != 0 && b > std::numeric_limits::max() / a) { + return false; + } + out = a * b; + return true; +} + +} // namespace + int run_deepseek4_dspark_draft_ipc_daemon( const char * draft_path, int ring_cap, @@ -88,7 +101,8 @@ int run_deepseek4_dspark_draft_ipc_daemon( return sequence != 0 && shared_payload && shared_payload_data && backend_ipc_payload_in_bounds( 0, bytes, shared_payload_bytes) && - header->sequence == sequence && header->bytes == bytes; + backend_ipc_shared_payload_header_matches( + header, sequence, static_cast(bytes)); }; ggml_backend_t backend = ggml_backend_cuda_init(std::max(0, draft_gpu)); @@ -113,8 +127,18 @@ int run_deepseek4_dspark_draft_ipc_daemon( const int hidden = drafter.core.n_embd; const int block = drafter.block_size; const int n_target_layers = drafter.n_target_layers; - const int feature_row = hidden * n_target_layers; - if (hidden <= 0 || block <= 0 || n_target_layers <= 0 || feature_row <= 0) { + size_t feature_row_size = 0; + size_t feature_ring_size = 0; + size_t noise_size = 0; + const bool dimensions_ok = hidden > 0 && block > 0 && + n_target_layers > 0 && + checked_mul_size((size_t) hidden, (size_t) n_target_layers, + feature_row_size) && + feature_row_size <= (size_t) std::numeric_limits::max() && + checked_mul_size((size_t) ring_cap, feature_row_size, + feature_ring_size) && + checked_mul_size((size_t) hidden, (size_t) block, noise_size); + if (!dimensions_ok) { std::fprintf(stderr, "[ds4-dspark-ipc] invalid draft dimensions\n"); stream_status(stream_fd, -1); free_deepseek4_dspark_drafter(drafter); @@ -122,9 +146,10 @@ int run_deepseek4_dspark_draft_ipc_daemon( unmap_shared(); return 1; } + const int feature_row = (int) feature_row_size; - std::vector feature_ring((size_t) ring_cap * feature_row, 0.0f); - std::vector noise_embed((size_t) hidden * block); + std::vector feature_ring(feature_ring_size, 0.0f); + std::vector noise_embed(noise_size); std::vector context; std::vector hidden_out; @@ -169,7 +194,8 @@ int run_deepseek4_dspark_draft_ipc_daemon( auto run_proposal = [&](int committed, int ctx_len, const float * noise) { - if (committed < 0 || ctx_len <= 0 || ctx_len > ring_cap || !noise) { + if (committed < 0 || ctx_len <= 0 || ctx_len > ring_cap || + ctx_len > committed || !noise) { return false; } context.resize((size_t) ctx_len * feature_row); @@ -209,9 +235,12 @@ int run_deepseek4_dspark_draft_ipc_daemon( int n_tokens = 0; size_t bytes = 0; iss >> capture_idx >> start_pos >> n_tokens >> bytes; - const size_t expected = (size_t) std::max(0, n_tokens) * hidden * - sizeof(float); - bool ok = iss && bytes == expected && n_tokens > 0; + size_t expected = 0; + bool ok = iss && n_tokens > 0 && + checked_mul_size((size_t) n_tokens, (size_t) hidden, + expected) && + checked_mul_size(expected, sizeof(float), expected) && + bytes == expected; std::vector slice(ok ? bytes / sizeof(float) : 0); if (ok) ok = read_exact_fd(payload_fd, slice.data(), bytes); if (ok) ok = store_feature_slice(capture_idx, start_pos, n_tokens, slice); @@ -226,9 +255,12 @@ int run_deepseek4_dspark_draft_ipc_daemon( size_t bytes = 0; uint64_t sequence = 0; iss >> capture_idx >> start_pos >> n_tokens >> bytes >> sequence; - const size_t expected = (size_t) std::max(0, n_tokens) * hidden * - sizeof(float); - bool ok = iss && bytes == expected && n_tokens > 0 && + size_t expected = 0; + bool ok = iss && n_tokens > 0 && + checked_mul_size((size_t) n_tokens, (size_t) hidden, + expected) && + checked_mul_size(expected, sizeof(float), expected) && + bytes == expected && capture_idx >= 0 && capture_idx < n_target_layers && shared_request_valid(bytes, sequence); std::vector slice(ok ? bytes / sizeof(float) : 0); @@ -246,10 +278,13 @@ int run_deepseek4_dspark_draft_ipc_daemon( int n_tokens = 0; size_t bytes = 0; iss >> start_pos >> n_tokens >> bytes; - const size_t expected = (size_t) std::max(0, n_tokens) * - (size_t) feature_row * sizeof(float); - bool ok = iss && payload_fd >= 0 && bytes == expected && - n_tokens > 0 && n_tokens <= ring_cap; + size_t expected = 0; + bool ok = iss && payload_fd >= 0 && n_tokens > 0 && + n_tokens <= ring_cap && + checked_mul_size((size_t) n_tokens, feature_row_size, + expected) && + checked_mul_size(expected, sizeof(float), expected) && + bytes == expected; std::vector features(ok ? bytes / sizeof(float) : 0); if (ok) ok = read_exact_fd(payload_fd, features.data(), bytes); if (ok) { @@ -266,10 +301,12 @@ int run_deepseek4_dspark_draft_ipc_daemon( size_t bytes = 0; uint64_t sequence = 0; iss >> start_pos >> n_tokens >> bytes >> sequence; - const size_t expected = (size_t) std::max(0, n_tokens) * - (size_t) feature_row * sizeof(float); - bool ok = iss && bytes == expected && n_tokens > 0 && - n_tokens <= ring_cap && + size_t expected = 0; + bool ok = iss && n_tokens > 0 && n_tokens <= ring_cap && + checked_mul_size((size_t) n_tokens, feature_row_size, + expected) && + checked_mul_size(expected, sizeof(float), expected) && + bytes == expected && shared_request_valid(bytes, sequence); if (ok) { ok = store_feature_block( @@ -330,8 +367,8 @@ int run_deepseek4_dspark_draft_ipc_daemon( std::memcpy(shared_payload_data, hidden_out.data(), bytes); auto * header = static_cast(shared_payload); - header->bytes = bytes; - header->sequence = sequence; + backend_ipc_publish_shared_payload_header( + header, sequence, static_cast(bytes)); } else { ok = false; } diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index 91a926819..1a8a4c886 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -677,11 +677,16 @@ bool run_deepseek4_dspark_spec_decode( } // Laguna-style adaptive verify width: EWMA of accepted candidates, width = // ewma + 2 (avg_commit << block means the wide tail is usually wasted). - // /tmp/ds4_awidth: 1 = on, 0 = off (default on). + // Prefer an explicit process-level policy. Retain /tmp/ds4_awidth only as + // a compatibility fallback for older experiment harnesses. bool adaptive_width = true; - if (std::FILE * f = std::fopen("/tmp/ds4_awidth", "r")) { - int v = 1; - if (std::fscanf(f, "%d", &v) == 1) adaptive_width = (v != 0); + if (const char * raw = std::getenv("DFLASH_DS4_ADAPTIVE_WIDTH")) { + adaptive_width = raw[0] && std::strcmp(raw, "0") != 0; + } else if (std::FILE * f = std::fopen("/tmp/ds4_awidth", "r")) { + int value = 1; + if (std::fscanf(f, "%d", &value) == 1) { + adaptive_width = value != 0; + } std::fclose(f); } // The v1 IPC protocol returns normalized draft states only. Use its EWMA @@ -790,6 +795,7 @@ bool run_deepseek4_dspark_spec_decode( std::vector draft_tok, tgt_am; std::vector draft_confidence; std::vector cached_ahead_hidden; + std::vector cached_ahead_confidence_hidden; std::vector ahead_noise_embed((size_t) n_embd * block); bool cached_ahead_ready = false; int cached_ahead_pos = -1; @@ -812,20 +818,30 @@ bool run_deepseek4_dspark_spec_decode( const bool use_ahead = cached_ahead_ready && cached_ahead_pos == pos && cached_ahead_seed == lt && - cached_ahead_hidden.size() == (size_t) n_embd * block; + cached_ahead_hidden.size() == (size_t) n_embd * block && + (!use_confidence_width || + cached_ahead_confidence_hidden.size() == + (size_t) n_embd * block); if (use_ahead) { local_hidden.swap(cached_ahead_hidden); + if (use_confidence_width) { + confidence_hidden.swap(cached_ahead_confidence_hidden); + } cached_ahead_ready = false; ++ahead_reuses; } else { cached_ahead_ready = false; cached_ahead_hidden.clear(); + cached_ahead_confidence_hidden.clear(); noise_ids[0] = lt; for (int i = 1; i < block; i++) { noise_ids[i] = drafter.mask_token_id; } if (!target.embed_tokens( noise_ids.data(), block, noise_embed.data())) { + std::fprintf(stderr, + "[ds4-spec] draft embedding lookup failed\n"); + ok = false; break; } @@ -912,6 +928,9 @@ bool run_deepseek4_dspark_spec_decode( std::vector pj; if (!target.project_hidden_to_tokens( local_hidden.data(), head_step_cap - 1, pj)) { + std::fprintf(stderr, + "[ds4-spec] draft projection fallback failed\n"); + ok = false; break; } draft_tok.clear(); @@ -967,6 +986,7 @@ bool run_deepseek4_dspark_spec_decode( int inflight_ahead_seed = -1; int inflight_ahead_pos = -1; std::vector completed_ahead_hidden; + std::vector completed_ahead_confidence_hidden; if (draft_ahead_enabled && q == q_step_cap && q >= 2 && predicted_ahead_seed >= 0) { const SpecClock::time_point probe_t0 = SpecClock::now(); @@ -1044,7 +1064,9 @@ bool run_deepseek4_dspark_spec_decode( if (ahead_inflight) { ahead_output_ready = deepseek4_dspark_draft_read_async_output( - drafter_backend, completed_ahead_hidden); + drafter_backend, completed_ahead_hidden, + use_confidence_width + ? &completed_ahead_confidence_hidden : nullptr); if (!ahead_output_ready) { draft_ahead_enabled = false; std::fprintf(stderr, @@ -1087,12 +1109,19 @@ bool run_deepseek4_dspark_spec_decode( if (seed_hit) { ++ahead_seed_hits; cached_ahead_hidden = std::move(completed_ahead_hidden); + if (use_confidence_width) { + cached_ahead_confidence_hidden = + std::move(completed_ahead_confidence_hidden); + } else { + cached_ahead_confidence_hidden.clear(); + } cached_ahead_pos = commit_pos; cached_ahead_seed = bonus; cached_ahead_ready = true; } else { cached_ahead_ready = false; cached_ahead_hidden.clear(); + cached_ahead_confidence_hidden.clear(); } if (timing && steps < 8) { std::fprintf(stderr, diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index ef7d6fb47..596656d05 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -153,7 +153,11 @@ struct Ds4FusedVerifyCache { void destroy() { for (auto & s : slots) s.destroy(); for (auto & e : extra) e.reset(); + owner_ctx = nullptr; + backend = nullptr; peer_backend = nullptr; + disabled = false; + counter = 0; } }; @@ -163,6 +167,15 @@ static Ds4FusedVerifyCache fused_verify_graph_cache; // evicts the expensive warm q=3 verifier working set. static Ds4FusedVerifyCache fused_capture_graph_cache; +void reset_deepseek4_graph_runtime_caches() { + fused_verify_graph_cache.destroy(); + fused_capture_graph_cache.destroy(); + if (g_deepseek4_fused_decode_runtime_cache) { + g_deepseek4_fused_decode_runtime_cache->destroy(); + g_deepseek4_fused_decode_runtime_cache = nullptr; + } +} + // Cached verifier graphs retain their input tensors, while a request-adaptive // placement changes only the global->local expert maps. Refresh those small // LUT inputs without rebuilding or reallocating the graph. Mask any hot/cold @@ -210,6 +223,169 @@ static void ds4_fused_verify_refresh_hybrid_luts( } } +struct Ds4AuthoritativeRouteSnapshot { + const DeepSeek4FusedDecodeGraph::AuthoritativeRouteOutput * route = nullptr; + std::vector ids; + std::vector weights; +}; + +template +static bool ds4_fused_read_route_matrix( + ggml_tensor * tensor, + ggml_type expected_type, + int width, + int n_tokens, + std::vector & out) { + if (!tensor || tensor->type != expected_type || width <= 0 || + n_tokens <= 0 || tensor->ne[0] < width || + tensor->ne[1] < n_tokens || tensor->nb[0] != sizeof(T)) { + return false; + } + const size_t row_bytes = sizeof(T) * (size_t) width; + const size_t tensor_bytes = ggml_nbytes(tensor); + out.resize((size_t) width * n_tokens); + for (int token = 0; token < n_tokens; ++token) { + const size_t offset = (size_t) token * tensor->nb[1]; + if (offset > tensor_bytes || row_bytes > tensor_bytes - offset) { + out.clear(); + return false; + } + ggml_backend_tensor_get( + tensor, out.data() + (size_t) token * width, + offset, row_bytes); + } + return true; +} + +static void ds4_fused_consume_route_diagnostics( + DeepSeek4FusedDecodeGraph & fg, + const MoeHybridStorage * hybrid, + MoeHybridRoutingStats * routing_stats) { + const bool predict_audit = + ds4_tp_route_predict_audit_lookahead() > 0; + const bool cache_audit = + ds4_env_flag("DFLASH_DS4_TP_CACHE_AUDIT"); + if (!routing_stats && !predict_audit && !cache_audit) return; + + std::vector snapshots; + snapshots.reserve(fg.authoritative_routes.size()); + uint64_t cache_hot = 0; + uint64_t cache_total = 0; + for (const auto & route : fg.authoritative_routes) { + Ds4AuthoritativeRouteSnapshot snapshot; + snapshot.route = &route; + if (!ds4_fused_read_route_matrix( + route.selected, GGML_TYPE_I32, route.width, + route.n_tokens, snapshot.ids) || + !ds4_fused_read_route_matrix( + route.weights, GGML_TYPE_F32, route.width, + route.n_tokens, snapshot.weights)) { + std::fprintf(stderr, + "[ds4-route-audit] failed to read layer=%d lane=%d\n", + route.layer, route.lane_start); + continue; + } + for (int token = 0; token < route.n_tokens; ++token) { + const int32_t * ids = snapshot.ids.data() + + (size_t) token * route.width; + const float * weights = snapshot.weights.data() + + (size_t) token * route.width; + observe_active_routing( + routing_stats, route.layer, ids, weights, route.width); + if (cache_audit && hybrid && route.layer >= 0 && + route.layer < (int) hybrid->layers.size()) { + const MoeHybridLayerStorage & layer = + hybrid->layers[(size_t) route.layer]; + for (int i = 0; i < route.width; ++i) { + const int32_t id = ids[i]; + if (weights[i] == 0.0f || id < 0 || + id >= (int32_t) layer.hot_local_by_global.size()) { + continue; + } + ++cache_total; + if (layer.hot_local_by_global[(size_t) id] >= 0) { + ++cache_hot; + } + } + } + } + snapshots.push_back(std::move(snapshot)); + } + + if (cache_audit && cache_total > 0) { + std::fprintf(stderr, + "[ds4-route-cache-audit] hot=%llu total=%llu hit=%.3f\n", + (unsigned long long) cache_hot, + (unsigned long long) cache_total, + (double) cache_hot / (double) cache_total); + } + if (!predict_audit) return; + + uint64_t active_total = 0; + uint64_t base_hits = 0; + uint64_t wide_hits = 0; + uint64_t comparisons = 0; + int base_width = 0; + int wide_width = 0; + for (const auto & prediction : fg.route_predictions) { + const Ds4AuthoritativeRouteSnapshot * target = nullptr; + for (const auto & snapshot : snapshots) { + if (snapshot.route->layer == prediction.target_layer && + snapshot.route->lane_start == prediction.lane_start && + snapshot.route->n_tokens == prediction.n_tokens) { + target = &snapshot; + break; + } + } + if (!target) continue; + std::vector predicted; + if (!ds4_fused_read_route_matrix( + prediction.selected, GGML_TYPE_I32, prediction.width, + prediction.n_tokens, predicted)) { + std::fprintf(stderr, + "[ds4-route-predict-audit] failed to read source=%d target=%d\n", + prediction.source_layer, prediction.target_layer); + continue; + } + const int primary = std::min( + target->route->width, prediction.width); + base_width = std::max(base_width, primary); + wide_width = std::max(wide_width, prediction.width); + ++comparisons; + for (int token = 0; token < prediction.n_tokens; ++token) { + const int32_t * actual = target->ids.data() + + (size_t) token * target->route->width; + const float * weights = target->weights.data() + + (size_t) token * target->route->width; + const int32_t * candidate = predicted.data() + + (size_t) token * prediction.width; + for (int i = 0; i < target->route->width; ++i) { + if (weights[i] == 0.0f || actual[i] < 0) continue; + ++active_total; + bool in_base = false; + bool in_wide = false; + for (int j = 0; j < prediction.width; ++j) { + if (candidate[j] != actual[i]) continue; + in_wide = true; + if (j < primary) in_base = true; + break; + } + base_hits += in_base ? 1u : 0u; + wide_hits += in_wide ? 1u : 0u; + } + } + } + if (active_total > 0) { + std::fprintf(stderr, + "[ds4-route-predict-audit] comparisons=%llu active=%llu " + "recall@%d=%.3f recall@%d=%.3f\n", + (unsigned long long) comparisons, + (unsigned long long) active_total, + base_width, (double) base_hits / (double) active_total, + wide_width, (double) wide_hits / (double) active_total); + } +} + // Each hybrid cache slot owns a multi-backend scheduler, and each scheduler's // gallocr retains its per-backend scratch buffers until the slot is destroyed. // Keeping all 12 shape variants alive can therefore exhaust the 32 GiB target @@ -682,6 +858,17 @@ static bool ds4_build_fused_verify_graph( ctx, ffn_normed, w, L, il, lane_q); } + if (hybrid) { + DeepSeek4FusedDecodeGraph::AuthoritativeRouteOutput route; + route.layer = il; + route.lane_start = lane_start; + route.n_tokens = lane_q; + route.width = (int) selected->ne[0]; + route.selected = selected; + route.weights = router_weights; + fg.authoritative_routes.push_back(route); + } + // Diagnostic branch only: predict future score-routed expert IDs from // the current normalized FFN input. The authoritative routing above // remains the sole input to the model FFN, so this cannot affect @@ -717,6 +904,9 @@ static bool ds4_build_fused_verify_graph( } } if (hybrid) { + const bool allow_fused_combine = + ggml_backend_is_cuda(backend) && + ggml_backend_is_cuda(hybrid->cold_backend); MoeHybridConfig hybrid_cfg = make_ds4_moe_hybrid_config(w); hybrid_cfg.n_expert_used = (int) selected->ne[0]; MoeLayerDesc desc = make_ds4_moe_layer_desc(L); @@ -785,7 +975,8 @@ static bool ds4_build_fused_verify_graph( ds4_env_flag("DFLASH_DS4_TP_DEVICE_JOIN")) ? gf : nullptr, hybrid_cfg, desc, hybrid->layers[(size_t) il], ffn_normed, - first_ids, first_weights, lane_q, inputs, true)) { + first_ids, first_weights, lane_q, inputs, true, + allow_fused_combine)) { return false; } ffn_out = inputs.output; @@ -796,7 +987,8 @@ static bool ds4_build_fused_verify_graph( ds4_env_flag("DFLASH_DS4_TP_DEVICE_JOIN")) ? gf : nullptr, hybrid_cfg, desc, hybrid->layers[(size_t) il], ffn_normed, - padded_ids, padded_weights, lane_q, inputs, false)) { + padded_ids, padded_weights, lane_q, inputs, false, + allow_fused_combine)) { return false; } ffn_out = ggml_add(ctx, ffn_out, inputs.output); @@ -808,7 +1000,7 @@ static bool ds4_build_fused_verify_graph( hybrid_cfg, desc, hybrid->layers[(size_t) il], ffn_normed, selected, router_weights, - lane_q, inputs)) { + lane_q, inputs, true, allow_fused_combine)) { return false; } ffn_out = inputs.output; @@ -1056,18 +1248,22 @@ static bool ds4_build_fused_verify_graph( for (ggml_tensor * node : prediction.nodes) pin_main(node); pin_main(prediction.selected); } + for (const auto & route : fg.authoritative_routes) { + pin_main(route.selected); + pin_main(route.weights); + } if (ds4_env_flag("DFLASH_DS4_TP_ROUTE_STATS") || + ds4_env_flag("DFLASH_DS4_ROUTING_STATS_OUT") || ds4_env_flag("DFLASH_DS4_TP_DYNAMIC_HOTSET") || (q > 1 && ds4_tp_route_predict_audit_lookahead() > 0) || ds4_env_flag("DFLASH_DS4_TP_CACHE_AUDIT")) { - // Preserve selected-ID buffers through the end of the diagnostic - // graph so the post-compute route-reuse audit can read them. The - // native six-route graph records one {IDs, weights} pair/layer; - // the legacy graph records two pairs, hence the stride of two. - for (const MoeHybridGraphInputs & inputs : fg.hybrid_inputs) { - for (size_t i = 0; i < inputs.route_prefork_nodes.size(); i += 2) { - ggml_set_output(inputs.route_prefork_nodes[i]); - } + // Preserve the unsplit authoritative route and weight matrices. + // Reading the pre-fork list is insufficient in the legacy 4+2 + // lowering because its second build overwrites the per-layer + // bookkeeping with padded IDs. + for (const auto & route : fg.authoritative_routes) { + ggml_set_output(route.selected); + ggml_set_output(route.weights); } } pin_main(fg.logits); @@ -1430,13 +1626,13 @@ static int ds4_try_fused_verify_step( const bool force_verify_graph_replay = q > 1 && ds4_env_flag("DFLASH_DS4_VERIFY_FORCE_GRAPH_REPLAY"); if (force_verify_graph_replay) { - ggml_cuda_set_skip_props_check(true); + ggml_backend_cuda_set_skip_props_check(true); } const enum ggml_status compute_status = fg->sched ? ds4_fused_verify_compute_whole_step(*fg, q) : ggml_backend_graph_compute(backend, fg->sg.gf); if (force_verify_graph_replay) { - ggml_cuda_set_skip_props_check(false); + ggml_backend_cuda_set_skip_props_check(false); } if (capture_eager) { ggml_backend_cuda_set_graphs_disabled_override(false); @@ -1453,6 +1649,7 @@ static int ds4_try_fused_verify_step( } const auto read_t0 = Ds4TimingClock::now(); + ds4_fused_consume_route_diagnostics(*fg, hybrid, routing_stats); const int ncap = (int) hooks->capture_layer_ids->size(); const bool argmax_only = hooks->prefer_argmax_only && hooks->argmax_out && ex->argmax; diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index cf50b04d9..328c70eea 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -3681,6 +3681,14 @@ static bool ds4_fused_decode_enabled() { } struct DeepSeek4FusedDecodeGraph { + struct AuthoritativeRouteOutput { + int layer = -1; + int lane_start = 0; + int n_tokens = 0; + int width = 0; + ggml_tensor * selected = nullptr; + ggml_tensor * weights = nullptr; + }; struct RoutePredictionOutput { int source_layer = -1; int target_layer = -1; @@ -3699,6 +3707,7 @@ struct DeepSeek4FusedDecodeGraph { ggml_tensor * mask_bundle = nullptr; // additive score mask (0 / -1e30), may be null std::vector hash_ids; std::vector hybrid_inputs; + std::vector authoritative_routes; std::vector route_predictions; ggml_tensor * logits = nullptr; ggml_backend_sched_t sched = nullptr; @@ -3728,6 +3737,7 @@ struct DeepSeek4FusedDecodeGraph { logits = nullptr; hash_ids.clear(); hybrid_inputs.clear(); + authoritative_routes.clear(); route_predictions.clear(); whole_step_sequence_id = 0; shape_key.clear(); @@ -3779,6 +3789,11 @@ struct DeepSeek4FusedDecodeCache { } }; +// Registered once the layer-range runtime is entered. free_deepseek4_cache() +// resets this while both GPU backends are still alive, before park/shutdown +// releases the expert owner used by a cached multi-backend scheduler. +static DeepSeek4FusedDecodeCache * g_deepseek4_fused_decode_runtime_cache = nullptr; + static ggml_tensor * ds4_fused_hc_base_f32(ggml_context * ctx, ggml_tensor * base) { if (!base) return nullptr; ggml_tensor * b = base; @@ -4610,7 +4625,10 @@ bool deepseek4_step_layer_range( static Ds4DecodeSharedInputs decode_shared_inputs; static int hc_loaded_n_layer = 0; static const ggml_context * hc_loaded_ctx = nullptr; - if (hc_loaded_n_layer != w.n_layer || hc_loaded_ctx != w.ctx) { + static uint64_t hc_loaded_generation = 0; + g_deepseek4_fused_decode_runtime_cache = &fused_decode_graph_cache; + if (hc_loaded_n_layer != w.n_layer || hc_loaded_ctx != w.ctx || + hc_loaded_generation != w.runtime_generation) { reset_hc_layer_weights_cpu(hc_layer_weights_range); reset_hc_weights_cpu(hc_output_weights_range); hc_layer_weights_range.resize((size_t)w.n_layer); @@ -4653,6 +4671,7 @@ bool deepseek4_step_layer_range( load_hc_weights_cpu(hc_output_weights_range, w.output_hc_fn, w.output_hc_scale, w.output_hc_base); hc_loaded_n_layer = w.n_layer; hc_loaded_ctx = w.ctx; + hc_loaded_generation = w.runtime_generation; } // Per-layer execution with CPU-side HC @@ -5482,6 +5501,7 @@ bool create_deepseek4_cache(ggml_backend_t backend, } void free_deepseek4_cache(DeepSeek4Cache & c) { + reset_deepseek4_graph_runtime_caches(); if (c.ctx) { ggml_free(c.ctx); c.ctx = nullptr; } if (c.buf) { ggml_backend_buffer_free(c.buf); c.buf = nullptr; } c.layers.clear(); @@ -5489,6 +5509,13 @@ void free_deepseek4_cache(DeepSeek4Cache & c) { } void reset_deepseek4_cache(DeepSeek4Cache & c) { + // Graph executables may retain cross-device dependency state even though + // their tensor addresses are stable. Give every fresh request a distinct + // generation so those executables cannot be replayed across a cache clear. + ++c.sequence_id; + if (c.sequence_id == 0) { + ++c.sequence_id; + } c.cur_pos = 0; for (DeepSeek4LayerCache & lc : c.layers) { lc.n_comp = 0; @@ -5732,7 +5759,6 @@ static ggml_tensor * build_dspark_attention( ggml_tensor * neg_block, // I32[block] -(block positions) ggml_tensor * pos_ctx, // I32[ctx_len] absolute positions committed-ctx_len..committed-1 ggml_tensor * attn_mask) { // F32[ctx_len+block], 0 or -inf - 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; @@ -5839,6 +5865,7 @@ namespace { struct DsparkContextKvProjector { int n_cols = -1; const void * drafter = nullptr; + ggml_backend_t backend = nullptr; std::vector arena; ggml_context * ctx = nullptr; ggml_gallocr_t alloc = nullptr; @@ -5869,7 +5896,12 @@ static bool dspark_project_context_columns( const int head_dim = w.head_dim; DsparkContextKvProjector & P = g_dspark_ctx_kv; - if (!P.ctx || P.n_cols != n_cols || P.drafter != (const void *) &d) { + if (!P.ctx || P.n_cols != n_cols || P.drafter != (const void *) &d || + P.backend != backend) { + if (P.alloc && P.backend != backend) { + ggml_gallocr_free(P.alloc); + P.alloc = nullptr; + } if (P.ctx) { ggml_free(P.ctx); P.ctx = nullptr; @@ -5927,6 +5959,7 @@ static bool dspark_project_context_columns( } P.n_cols = n_cols; P.drafter = (const void *) &d; + P.backend = backend; } ggml_backend_tensor_set( @@ -5968,7 +6001,8 @@ static bool dspark_update_context_kv_cache( int n_new = committed - P.end_pos; const bool rebuild = - P.drafter != (const void *) &d || P.end_pos < 0 || + P.drafter != (const void *) &d || P.backend != backend || + P.end_pos < 0 || n_new <= 0 || n_new > ctx_len || std::min(n_swa, P.valid + n_new) != ctx_len || P.host_kv.size() != (size_t) w.n_layer * n_swa * head_dim; @@ -6097,7 +6131,7 @@ static bool deepseek4_dspark_draft_forward_impl( : valid_ctx_len; const std::vector * cached_host_kv = nullptr; - if (context_kv_cache && upload_context && valid_ctx_len > 0 && + if (context_kv_cache && upload_context && !dspark_update_context_kv_cache( backend, d, ctx_features, valid_ctx_len, committed, &cached_host_kv)) { @@ -6420,11 +6454,11 @@ static bool deepseek4_dspark_draft_forward_impl( // have both been qualified on the deployment GPUs. const bool force_graph_replay = ds4_env_flag("DFLASH_DS4_DRAFT_FORCE_GRAPH_REPLAY"); - if (force_graph_replay) ggml_cuda_set_skip_props_check(true); + if (force_graph_replay) ggml_backend_cuda_set_skip_props_check(true); const ggml_status st = out_hidden ? ggml_backend_graph_compute(backend, C.gf) : ggml_backend_graph_compute_async(backend, C.gf); - if (force_graph_replay) ggml_cuda_set_skip_props_check(false); + if (force_graph_replay) ggml_backend_cuda_set_skip_props_check(false); if (st != GGML_STATUS_SUCCESS) { // Invalidate: a failed compute leaves no reusable state guarantees. ggml_free(C.ctx); C.ctx = nullptr; C.gf = nullptr; C.ctx_len = -1; @@ -6496,9 +6530,12 @@ bool deepseek4_dspark_draft_forward_async_reuse_context( bool deepseek4_dspark_draft_read_async_output( ggml_backend_t backend, - std::vector & out_hidden) { + std::vector & out_hidden, + std::vector * confidence_hidden) { DsparkDraftCache & C = g_dspark_draft_cache; - if (!backend || !C.ctx || !C.out || C.block <= 0 || !C.drafter) { + if (!backend || backend != C.backend || !C.ctx || !C.out || + (confidence_hidden && !C.confidence_out) || + C.block <= 0 || !C.drafter) { return false; } const DSparkDrafter * d = @@ -6507,6 +6544,12 @@ bool deepseek4_dspark_draft_read_async_output( out_hidden.resize(count); ggml_backend_tensor_get( C.out, out_hidden.data(), 0, sizeof(float) * count); + if (confidence_hidden) { + confidence_hidden->resize(count); + ggml_backend_tensor_get( + C.confidence_out, confidence_hidden->data(), 0, + sizeof(float) * count); + } return true; } diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index c9032bba9..92c74c3dc 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -66,7 +66,6 @@ struct DeepSeek4StepTelemetry { uint64_t sample_us = 0; uint64_t emit_us = 0; uint64_t full_graph_build_us = 0; - uint64_t full_graph_alloc_us = 0; uint64_t full_graph_set_us = 0; uint64_t full_graph_compute_us = 0; uint64_t full_graph_read_us = 0; @@ -228,6 +227,7 @@ struct DeepSeek4Weights { // Runtime serving policy. These values are set by the backend after the // GGUF is loaded; they are not model metadata. + uint64_t runtime_generation = 0; // unique successful-load identity int routed_expert_top_k = 0; // 0 = model default (n_expert_used) bool fused_decode = false; }; @@ -319,6 +319,9 @@ bool create_deepseek4_cache(ggml_backend_t backend, void free_deepseek4_cache(DeepSeek4Cache & c); void reset_deepseek4_cache(DeepSeek4Cache & c); +// Release cached fused decode/verify schedulers before their model/backend +// owners are destroyed (park, shutdown, or failed reload). +void reset_deepseek4_graph_runtime_caches(); int deepseek4_previous_raw_ring_spans( int kv_start, int n_swa, diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index 5437985d8..865ee2450 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -210,7 +210,8 @@ static bool should_split_ds4_dense_tensor(const char * name, int mask) { // tensors and the grouped output-A view are deliberately excluded: the // CUDA/HIP split-buffer implementation cannot split 3-D weight views. if ((mask & 1) && std::strcmp(name, "output.weight") == 0) return true; - if ((mask & 2) && std::strstr(name, ".attn_q_b.weight")) return true; + if ((mask & 2) && std::strstr(name, ".attn_q_b.weight") && + !std::strstr(name, ".indexer.attn_q_b.weight")) return true; if ((mask & 4) && std::strstr(name, ".attn_output_b.weight")) return true; if ((mask & 8) && (std::strstr(name, ".attn_q_a.weight") || std::strstr(name, ".attn_kv.weight") || @@ -444,7 +445,18 @@ bool load_deepseek4_gguf_partial(const std::string & path, const size_t data_offset = gguf_get_data_offset(gctx); ggml_backend_buffer_type_t buft = ggml_backend_get_default_buffer_type(backend); const size_t alignment = ggml_backend_buft_get_alignment(buft); - const int dense_tp_mask = ds4_dense_tp_mask(); + int dense_tp_mask = ds4_dense_tp_mask(); + if (dense_tp_mask != 0) { + const char * fused_verify = std::getenv("DFLASH_DS4_FUSED_VERIFY"); + if (fused_verify && fused_verify[0] && + std::strcmp(fused_verify, "0") != 0) { + std::fprintf(stderr, + "[deepseek4-dense-tp] disabling mask=%d: split-buffer dense " + "weights are incompatible with fused verifier graph replay\n", + dense_tp_mask); + dense_tp_mask = 0; + } + } ggml_backend_buffer_type_t split_buft = nullptr; size_t split_alignment = 1; if (dense_tp_mask != 0 && ggml_backend_is_cuda(backend) && @@ -746,6 +758,16 @@ bool load_deepseek4_gguf_partial(const std::string & path, out.ctx = meta_ctx; out.buf = buf; out.dense_split_buf = split_buf; + // Pointer equality is insufficient to identify a reload: allocators may + // reuse the same ggml_context address after park/unpark. Cached graphs use + // this generation to reject tensors from the prior model lifetime. + static std::atomic next_runtime_generation{1}; + out.runtime_generation = + next_runtime_generation.fetch_add(1, std::memory_order_relaxed); + if (out.runtime_generation == 0) { + out.runtime_generation = + next_runtime_generation.fetch_add(1, std::memory_order_relaxed); + } gguf_free(gctx); // Note: meta_ctx is now owned by out.ctx — do NOT free it here. @@ -988,6 +1010,7 @@ void free_deepseek4_weights(DeepSeek4Weights & w) { w.embedder.tok_embd_owned.clear(); w.embedder.tok_embd_bytes = nullptr; w.moe_hybrid = false; + w.runtime_generation = 0; } } // namespace dflash::common diff --git a/server/src/qwen35/qwen35_target_shard_ipc_daemon.cpp b/server/src/qwen35/qwen35_target_shard_ipc_daemon.cpp index 7d6c86f87..7f9162191 100644 --- a/server/src/qwen35/qwen35_target_shard_ipc_daemon.cpp +++ b/server/src/qwen35/qwen35_target_shard_ipc_daemon.cpp @@ -540,7 +540,8 @@ int run_qwen35_target_shard_ipc_daemon(const char * target_path, if (shared_payload && shared_payload != MAP_FAILED && shared_payload_data && seq != 0 && bytes == expected_bytes && backend_ipc_payload_in_bounds(0, bytes, shared_payload_capacity) && - header->sequence == seq && header->bytes == (uint64_t)bytes) { + backend_ipc_shared_payload_header_matches( + header, seq, static_cast(bytes))) { host_act.assign(bytes / sizeof(float), 0.0f); std::memcpy(host_act.data(), shared_payload_data, bytes); payload_ok = true; diff --git a/server/test/test_ds4_dspark_load.cpp b/server/test/test_ds4_dspark_load.cpp index bd1d2b5ec..a0550fb2b 100644 --- a/server/test/test_ds4_dspark_load.cpp +++ b/server/test/test_ds4_dspark_load.cpp @@ -150,6 +150,26 @@ int main(int argc, char ** argv) { need(confidence_finite, "finite pre-norm confidence hidden"); need(confidence_hidden.size() == (size_t) n_embd * block, "pre-norm confidence hidden size"); + + // The speculative draft-ahead path reads both graph outputs after + // an asynchronous launch. Keep its hidden/confidence pairing + // covered by the same end-to-end model smoke as the sync path. + std::vector async_hidden; + std::vector async_confidence_hidden; + const bool async_ok = deepseek4_dspark_draft_forward_async( + backend, d, noise.data(), feats.data(), ctx_len, committed); + if (async_ok) { + deepseek4_dspark_draft_wait(backend); + } + const bool async_read_ok = async_ok && + deepseek4_dspark_draft_read_async_output( + backend, async_hidden, &async_confidence_hidden); + need(async_read_ok, "async draft output read"); + if (async_read_ok) { + need(async_hidden == hidden, "async normalized hidden identity"); + need(async_confidence_hidden == confidence_hidden, + "async confidence hidden identity"); + } } } diff --git a/server/test/test_server_unit.cpp b/server/test/test_server_unit.cpp index a498ceb4c..8247bb31d 100644 --- a/server/test/test_server_unit.cpp +++ b/server/test/test_server_unit.cpp @@ -3117,10 +3117,16 @@ static void test_backend_ipc_shared_payload_map_sizing() { TEST_ASSERT(map_bytes == 1024 + backend_ipc_shared_payload_header_bytes()); BackendIpcSharedPayloadHeader header; - header.sequence = 7; - header.bytes = 1024; - TEST_ASSERT(header.sequence == 7); - TEST_ASSERT(header.bytes == 1024); + backend_ipc_publish_shared_payload_header(&header, 7, 1024); + TEST_ASSERT(backend_ipc_shared_payload_header_matches(&header, 7, 1024)); + TEST_ASSERT(!backend_ipc_shared_payload_header_matches(&header, 0, 1024)); + TEST_ASSERT(!backend_ipc_shared_payload_header_matches(&header, 8, 1024)); + TEST_ASSERT(!backend_ipc_shared_payload_header_matches(&header, 7, 512)); + uint64_t sequence = 0; + uint64_t bytes = 0; + backend_ipc_load_shared_payload_header(&header, sequence, bytes); + TEST_ASSERT(sequence == 7); + TEST_ASSERT(bytes == 1024); TEST_ASSERT(!backend_ipc_shared_payload_map_bytes( std::numeric_limits::max(), map_bytes)); diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 34031f293..1cd261a95 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -1322,6 +1322,7 @@ static void test_reset_deepseek4_cache(ggml_backend_t backend) { DeepSeek4Cache cache; TEST_ASSERT(create_deepseek4_cache(backend, weights, 32, cache)); if (cache.buf) { + TEST_ASSERT(cache.sequence_id == 0); ggml_backend_buffer_clear(cache.buf, 0x7f); cache.cur_pos = 17; cache.layers[0].n_comp = 4; @@ -1331,6 +1332,7 @@ static void test_reset_deepseek4_cache(ggml_backend_t backend) { reset_deepseek4_cache(cache); + TEST_ASSERT(cache.sequence_id == 1); TEST_ASSERT(cache.cur_pos == 0); for (const auto & layer : cache.layers) { TEST_ASSERT(layer.n_comp == 0); @@ -1339,6 +1341,9 @@ static void test_reset_deepseek4_cache(ggml_backend_t backend) { TEST_ASSERT(std::all_of(bytes.begin(), bytes.end(), [](uint8_t value) { return value == 0; })); } + + reset_deepseek4_cache(cache); + TEST_ASSERT(cache.sequence_id == 2); } free_deepseek4_cache(cache); From 3cb30a529c2cdbce047c531f972caf103bccc74c Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 20 Jul 2026 12:19:56 -0700 Subject: [PATCH 10/15] docs(ds4): qualify the promoted top-4 profile --- server/docs/DS4.md | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index f547dabd9..0b7343fb7 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -110,7 +110,8 @@ DSpark context-KV reuse. `DFLASH_DS4_TOPK=4` is an explicit approximate inference policy; omit it when model-default top-6 quality is required. On 2026-07-19, the clean dual-architecture build produced steady 50.7 and -50.8 tok/s decode runs on the deterministic integer-list workload (128 output +50.8 tok/s decode runs on the deterministic integer-list workload (the +53-token prompt includes the explicit default system message; 128 output tokens, q=4, acceptance 1.00). Mean draft time was 6.1 ms and mean target verification was 70.7--70.9 ms. With the same prompt and the same 1.00 acceptance, incorrectly splitting boundary-spanning verification into q=3 and @@ -210,8 +211,11 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. | `DFLASH_EXPERT_BUDGET_MB` | R9700 memory budget used to select locally resident hot experts. | | `DFLASH_DS4_HOTNESS_CSV` | Optional per-layer expert routing profile used for hot placement. | | `DFLASH_DS4_FUSED_VERIFY` | Enable the persistent q-wide target verification graph. | +| `DFLASH_DS4_DENSE_TP_MASK` | Experimental row splitting for dense projections. This is automatically disabled when fused verification is enabled because split-buffer weights are not compatible with verifier graph replay. | | `DFLASH_DS4_TP_DEVICE_JOIN` | Join expert-owner results on device instead of through a host reduction. | | `DFLASH_DS4_DRAFT_CONTEXT_KV_CACHE` | Reuse DSpark context projection/KV state between speculative steps. | +| `DFLASH_DS4_ADAPTIVE_WIDTH` | Set to `1` to select q=2/3/4 from DSpark confidence; `0` keeps the configured fixed verify width. If unset, the legacy `/tmp/ds4_awidth` control remains supported. | +| `DFLASH_DS4_DRAFT_AHEAD` | On an independent in-process draft backend, launch the next proposal concurrently with target verification. The target still validates every candidate. | `DFLASH_DS4_TIMING` enables the existing timing banners: From dd2e6f509fd29236c5160f55eaa214617f0ad72a Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:57:11 -0700 Subject: [PATCH 11/15] perf(ds4): batch peer handoffs and group cold experts --- .../deps/llama.cpp/ggml/src/ggml-backend.cpp | 33 +++- .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 103 ++++++++++ .../deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu | 176 +++++++++++++++--- 3 files changed, 280 insertions(+), 32 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp index 2e621ec55..a21ad3907 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend.cpp @@ -1796,6 +1796,32 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } + // All copied inputs for this split use the same destination backend + // and scheduler copy generation. Waiting for that generation once is + // sufficient before overwriting any of its input buffers. The legacy + // loop waited again before every tensor; without scheduler events each + // wait synchronizes the whole destination stream and serializes a + // multi-input peer handoff. + static const bool batch_peer_copies = [] { + const char * value = + getenv("DFLASH_DS4_TP_BATCH_PEER_COPIES"); + return value && *value && strcmp(value, "0") != 0; + }(); + bool split_copy_generation_ready = false; + auto wait_for_split_copy_generation = [&]() { + if (batch_peer_copies && split_copy_generation_ready) { + return; + } + if (sched->events[split_backend_id][sched->cur_copy] != NULL) { + ggml_backend_event_wait( + split_backend, + sched->events[split_backend_id][sched->cur_copy]); + } else { + ggml_backend_synchronize(split_backend); + } + split_copy_generation_ready = true; + }; + // copy the input tensors to the split backend for (int input_id = 0; input_id < split->n_inputs; input_id++) { ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]); @@ -1819,16 +1845,13 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } else { ggml_backend_synchronize(split_backend); } + split_copy_generation_ready = true; ggml_backend_tensor_copy(input, input_cpy); } } else { // wait for the split backend to finish using the input before overwriting it if (!sched->whole_graph_capture) { - if (sched->events[split_backend_id][sched->cur_copy] != NULL) { - ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]); - } else { - ggml_backend_synchronize(split_backend); - } + wait_for_split_copy_generation(); } // when offloading MoE weights, we can reduce the amount of data copied by copying only the experts that are used 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 98158d49b..78482896b 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 @@ -3366,6 +3366,69 @@ static void ggml_cuda_whole_graph_wait_async( } #endif +#if defined(GGML_USE_HIP) +// A scheduler split can have several inputs produced by the same peer GPU. +// hipMemcpyPeerAsync is already ordered on the producer stream, so publishing +// and waiting after every individual tensor only adds redundant cross-device +// dependencies. Keep the copies ordered on that stream and publish one event +// immediately before the consumer graph is submitted. +// +// This is deliberately opt-in while the heterogeneous path is qualified. It +// changes synchronization granularity only: tensor placement, copy direction, +// and the bytes copied remain identical. +struct ggml_cuda_pending_peer_copy_batch { + ggml_backend_cuda_context * src = nullptr; + ggml_backend_cuda_context * dst = nullptr; + size_t bytes = 0; + int copies = 0; +}; + +static thread_local ggml_cuda_pending_peer_copy_batch + ggml_cuda_pending_peer_copies; + +static bool ggml_cuda_batch_peer_copies_enabled() { + static const bool enabled = [] { + const char * value = getenv("DFLASH_DS4_TP_BATCH_PEER_COPIES"); + return value && *value && strcmp(value, "0") != 0; + }(); + return enabled; +} + +static void ggml_cuda_flush_peer_copy_batch(const char * reason) { + auto & batch = ggml_cuda_pending_peer_copies; + if (batch.copies == 0) { + return; + } + + GGML_ASSERT(batch.src && batch.dst); + ggml_cuda_set_device(batch.src->device); + if (!batch.src->copy_event) { + CUDA_CHECK(cudaEventCreateWithFlags( + &batch.src->copy_event, cudaEventDisableTiming)); + } + CUDA_CHECK(cudaEventRecord( + batch.src->copy_event, batch.src->stream())); + + ggml_cuda_set_device(batch.dst->device); + CUDA_CHECK(cudaStreamWaitEvent( + batch.dst->stream(), batch.src->copy_event, 0)); + + static const bool trace = [] { + const char * value = getenv("DFLASH_DS4_TP_BATCH_PEER_TRACE"); + return value && *value && strcmp(value, "0") != 0; + }(); + static thread_local int reports = 0; + if (trace && reports++ < 256) { + GGML_LOG_INFO( + "[ds4-peer-batch] src=%d dst=%d copies=%d bytes=%zu reason=%s\n", + batch.src->device, batch.dst->device, batch.copies, batch.bytes, + reason); + } + + batch = {}; +} +#endif + static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { ggml_backend_buffer_t buf_src = src->view_src ? src->view_src->buffer : src->buffer; ggml_backend_buffer_t buf_dst = dst->view_src ? dst->view_src->buffer : dst->buffer; @@ -3394,6 +3457,11 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ if (backend_src != backend_dst) { #if defined(GGML_USE_HIP) + // A pending eager batch must be visible before entering a captured + // whole-model transfer path, which supplies its own synchronization. + if (ggml_cuda_whole_graph_capture_depth > 0) { + ggml_cuda_flush_peer_copy_batch("whole-graph-capture"); + } if (ggml_cuda_whole_graph_capture_depth > 0 && cuda_ctx_src->device != cuda_ctx_dst->device) { // Deferred MoE joins are deliberately staged at producer-split @@ -3466,11 +3534,34 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ #endif // copy on src stream if (cuda_ctx_src->device == cuda_ctx_dst->device) { +#if defined(GGML_USE_HIP) + ggml_cuda_flush_peer_copy_batch("same-device-copy"); +#endif CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); } else { #ifdef GGML_CUDA_NO_PEER_COPY return false; #else +#if defined(GGML_USE_HIP) + if (ggml_cuda_batch_peer_copies_enabled() && + ggml_cuda_whole_graph_capture_depth == 0) { + auto & batch = ggml_cuda_pending_peer_copies; + if (batch.copies != 0 && + (batch.src != cuda_ctx_src || batch.dst != cuda_ctx_dst)) { + ggml_cuda_flush_peer_copy_batch("peer-pair-change"); + } + ggml_cuda_set_device(cuda_ctx_src->device); + CUDA_CHECK(cudaMemcpyPeerAsync( + dst->data, cuda_ctx_dst->device, + src->data, cuda_ctx_src->device, + ggml_nbytes(dst), cuda_ctx_src->stream())); + batch.src = cuda_ctx_src; + batch.dst = cuda_ctx_dst; + batch.bytes += ggml_nbytes(dst); + batch.copies++; + return true; + } +#endif CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, cuda_ctx_dst->device, src->data, cuda_ctx_src->device, ggml_nbytes(dst), cuda_ctx_src->stream())); #endif // GGML_CUDA_NO_PEER_COPY } @@ -3501,6 +3592,9 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ )); } else { // src and dst are on the same backend +#if defined(GGML_USE_HIP) + ggml_cuda_flush_peer_copy_batch("same-backend-copy"); +#endif CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); } return true; @@ -3509,6 +3603,9 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ static void ggml_backend_cuda_synchronize(ggml_backend_t backend) { ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; +#if defined(GGML_USE_HIP) + ggml_cuda_flush_peer_copy_batch("backend-synchronize"); +#endif CUDA_CHECK(cudaStreamSynchronize(cuda_ctx->stream())); GGML_UNUSED(backend); @@ -4682,6 +4779,12 @@ extern "C" void ggml_backend_cuda_set_skip_props_check(bool skip) { static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; +#if defined(GGML_USE_HIP) + // This is the consumer boundary for copies accumulated while the generic + // scheduler prepared this split. Publishing here preserves the existing + // copy-before-compute dependency with a single event for the whole batch. + ggml_cuda_flush_peer_copy_batch("graph-compute"); +#endif ggml_cuda_set_device(cuda_ctx->device); // LIFO-free last evaluation's memoized q8_1 activations. 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 f16d0c490..b67b909ee 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -386,8 +386,12 @@ static constexpr __host__ __device__ int get_mmvq_mmid_max_batch_rdna4(ggml_type // Model-agnostic: applies to any MoE with n_expert_used*n_tokens <= 256. // DFLASH_MMID_GROUPED=1 opt-in // DFLASH_MMID_GROUPED_TYPES bitmask, 1 = Q4_K (default), 2 = Q6_K, -// 4 = Q4_0/Q8_0/Q5_K. Q6_K stays on its tuned -// MMQ route above 5 tokens unless enabled. +// 4 = Q4_0/Q8_0/Q5_K, 8 = ROCmFP2/ROCmFP3. +// DFLASH_MMID_GROUPED_DEVICE optional zero-based device index; unset/-1 +// applies the path to every eligible device. +// Q6_K stays on its tuned MMQ route above 5 +// tokens unless enabled. The ROCmFP formats are +// opt-in until qualified on each AMD target. #define MMID_GROUPED_MAX_PAIRS 256 #define MMID_GROUPED_MAX_TPG 8 #define MMID_META_NG 0 @@ -435,8 +439,9 @@ static bool mmvq_env_flag(const char * name, bool default_value = false) { } static bool mmid_grouped_type_ok(ggml_type type) { - // bit0 = Q4_K, bit1 = Q6_K, bit2 = Q4_0/Q8_0/Q5_K. Default: all validated - // types (7); DFLASH_MMID_GROUPED_TYPES is a debug override. + // bit0 = Q4_K, bit1 = Q6_K, bit2 = Q4_0/Q8_0/Q5_K, + // bit3 = Q2_0_ROCMFP2/Q3_0_ROCMFPX. Default: previously validated types + // only (7); DFLASH_MMID_GROUPED_TYPES is an experimental override. static const int mask = []() { const char * e = std::getenv("DFLASH_MMID_GROUPED_TYPES"); if (e == nullptr || e[0] == '\0') { @@ -453,11 +458,22 @@ static bool mmid_grouped_type_ok(ggml_type type) { case GGML_TYPE_Q8_0: case GGML_TYPE_Q5_K: return (mask & 4) != 0; + case GGML_TYPE_Q2_0_ROCMFP2: + case GGML_TYPE_Q3_0_ROCMFPX: + return (mask & 8) != 0; default: return false; } } +static bool mmid_grouped_device_ok() { + static const int selected = []() { + const char * e = std::getenv("DFLASH_MMID_GROUPED_DEVICE"); + return e == nullptr || e[0] == '\0' ? -1 : atoi(e); + }(); + return selected < 0 || selected == ggml_cuda_get_device(); +} + // Host function: returns the max batch size for the current arch+type at runtime. int get_mmvq_mmid_max_batch(ggml_type type, int cc) { // [TAG_MMID_GROUPED] the grouped kernel handles any supported type up to the @@ -1347,7 +1363,8 @@ static __global__ void mmid_group_prep( // their concurrent weight reads are served once from DRAM, then from L1/L2. // Weight traffic approaches (union of routed experts) instead of // (n_expert_used x n_tokens) expert-matrix reads. -template +template __launch_bounds__(MMID_GROUPED_MAX_TPG*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q_moe_grouped( const void * __restrict__ vx, const void * __restrict__ vy, const int32_t * __restrict__ meta, @@ -1397,6 +1414,10 @@ static __global__ void mul_mat_vec_q_moe_grouped( const float * x_bias = nullptr; const float * gate_bias = nullptr; ggml_glu_op active_glu; + float glu_param0 = 0.0f; + float glu_param1 = 0.0f; + float gate_value_scale = 1.0f; + float x_value_scale = 1.0f; if constexpr (has_fusion) { use_gate = fusion.gate != nullptr; @@ -1406,6 +1427,10 @@ static __global__ void mul_mat_vec_q_moe_grouped( x_bias = (const float *) fusion.x_bias; gate_bias = (const float *) fusion.gate_bias; active_glu = fusion.glu_op; + glu_param0 = fusion.glu_param0; + glu_param1 = fusion.glu_param1; + gate_value_scale = fusion.gate_value_scale; + x_value_scale = fusion.x_value_scale; } float tmp[c_rows_per_block] = {0.0f}; @@ -1417,10 +1442,38 @@ static __global__ void mul_mat_vec_q_moe_grouped( #pragma unroll for (int i = 0; i < c_rows_per_block; ++i) { - tmp[i] += vec_dot_q_cuda(vx, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (type == GGML_TYPE_Q2_0_ROCMFP2 && + c_fp2_packed32) { + tmp[i] += vec_dot_rocmfpx_fp2_q8_1_packed32( + vx, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } else if constexpr (type == GGML_TYPE_Q3_0_ROCMFPX && + c_fp3_packed24) { + tmp[i] += vec_dot_rocmfpx_fp3_q8_1_packed24( + vx, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } else { + tmp[i] += vec_dot_q_cuda( + vx, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } if constexpr (has_fusion) { if (use_gate) { - tmp_gate[i] += vec_dot_q_cuda(vgate, &y[kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (type == GGML_TYPE_Q2_0_ROCMFP2 && + c_fp2_packed32) { + tmp_gate[i] += vec_dot_rocmfpx_fp2_q8_1_packed32( + vgate, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } else if constexpr (type == GGML_TYPE_Q3_0_ROCMFPX && + c_fp3_packed24) { + tmp_gate[i] += vec_dot_rocmfpx_fp3_q8_1_packed24( + vgate, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } else { + tmp_gate[i] += vec_dot_q_cuda( + vgate, &y[kby], + kbx_offset + i*stride_row_x + kbx, kqs); + } } } } @@ -1455,7 +1508,18 @@ static __global__ void mul_mat_vec_q_moe_grouped( result *= ggml_cuda_op_gelu_single(gate_value); break; case GGML_GLU_OP_SWIGLU_OAI: - result = ggml_cuda_op_swiglu_oai_single(gate_value, result); + result = ggml_cuda_op_swiglu_oai_single( + gate_value, result, glu_param0, glu_param1); + break; + case GGML_GLU_OP_SWIGLU_DS4: + if (gate_value_scale != 1.0f) { + gate_value *= gate_value_scale; + } + if (x_value_scale != 1.0f) { + result *= x_value_scale; + } + result = ggml_cuda_op_swiglu_ds4_single( + gate_value, result, glu_param0); break; default: result = result * gate_value; @@ -1467,11 +1531,14 @@ static __global__ void mul_mat_vec_q_moe_grouped( } if constexpr (!has_fusion) { - GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, vgate, x_bias, gate_bias, active_glu, tmp_gate); + GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, vgate, x_bias, + gate_bias, active_glu, glu_param0, glu_param1, + gate_value_scale, x_value_scale, tmp_gate); } } -template +template static void mul_mat_vec_q_moe_grouped_launch( const void * vx, const void * vy, const int32_t * meta, const ggml_cuda_mm_fusion_args_device & fusion, float * dst, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, @@ -1486,12 +1553,16 @@ static void mul_mat_vec_q_moe_grouped_launch( const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; if (has_fusion) { - mul_mat_vec_q_moe_grouped<<>>( + mul_mat_vec_q_moe_grouped<<>>( vx, vy, meta, fusion, dst, (uint32_t) np, ncols_x, nchannels_y, nrows_x, stride_row_x, stride_col_y, stride_col_dst, stride_channel_x, stride_channel_y, stride_channel_dst); } else { - mul_mat_vec_q_moe_grouped<<>>( + mul_mat_vec_q_moe_grouped<<>>( vx, vy, meta, fusion, dst, (uint32_t) np, ncols_x, nchannels_y, nrows_x, stride_row_x, stride_col_y, stride_col_dst, stride_channel_x, stride_channel_y, stride_channel_dst); @@ -1508,6 +1579,16 @@ static bool mul_mat_vec_q_grouped_dispatch( const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; const uint3 nchannels_y_fd = init_fastdiv_values((uint32_t) nchannels_y); + static const bool fp3_packed24 = []() { + const char * e = + std::getenv("DFLASH_CUDA_MMVQ_MOE_FP3_PACKED24"); + return e && e[0] == '1' && e[1] == '\0'; + }(); + static const bool fp2_packed32 = []() { + const char * e = + std::getenv("DFLASH_CUDA_MMVQ_MOE_FP2_PACKED32"); + return e && e[0] == '1' && e[1] == '\0'; + }(); switch (type) { case GGML_TYPE_Q4_0: @@ -1530,6 +1611,40 @@ static bool mul_mat_vec_q_grouped_dispatch( mul_mat_vec_q_moe_grouped_launch(vx, vy, meta, fusion, dst, ncols_x, nchannels_y_fd, nrows_x, stride_row_x, stride_col_y, stride_col_dst, stride_channel_x, stride_channel_y, stride_channel_dst, max_groups, warp_size, stream); return true; + case GGML_TYPE_Q2_0_ROCMFP2: + if (fp2_packed32) { + mul_mat_vec_q_moe_grouped_launch< + GGML_TYPE_Q2_0_ROCMFP2, false, true>( + vx, vy, meta, fusion, dst, ncols_x, nchannels_y_fd, + nrows_x, stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, + stride_channel_dst, max_groups, warp_size, stream); + } else { + mul_mat_vec_q_moe_grouped_launch< + GGML_TYPE_Q2_0_ROCMFP2>( + vx, vy, meta, fusion, dst, ncols_x, nchannels_y_fd, + nrows_x, stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, + stride_channel_dst, max_groups, warp_size, stream); + } + return true; + case GGML_TYPE_Q3_0_ROCMFPX: + if (fp3_packed24) { + mul_mat_vec_q_moe_grouped_launch< + GGML_TYPE_Q3_0_ROCMFPX, true, false>( + vx, vy, meta, fusion, dst, ncols_x, nchannels_y_fd, + nrows_x, stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, + stride_channel_dst, max_groups, warp_size, stream); + } else { + mul_mat_vec_q_moe_grouped_launch< + GGML_TYPE_Q3_0_ROCMFPX>( + vx, vy, meta, fusion, dst, ncols_x, nchannels_y_fd, + nrows_x, stride_row_x, stride_col_y, stride_col_dst, + stride_channel_x, stride_channel_y, + stride_channel_dst, max_groups, warp_size, stream); + } + return true; default: return false; } @@ -2460,7 +2575,7 @@ void ggml_cuda_mul_mat_vec_q( // [TAG_MMID_GROUPED] grouped-expert path for small MUL_MAT_ID batches. if (ids && ncols_dst >= 2 && ncols_dst <= MMVQ_MAX_MOE_BATCH_SIZE && (int) (nchannels_dst*ncols_dst) <= MMID_GROUPED_MAX_PAIRS && - mmid_grouped_env() && mmid_grouped_type_ok(src0->type)) { + mmid_grouped_env() && mmid_grouped_type_ok(src0->type) && mmid_grouped_device_ok()) { // Batches above MMID_GROUPED_MAX_PAIRS fall through to the legacy // per-expert kernel instead of aborting the request. const int np = (int) (nchannels_dst*ncols_dst); @@ -2474,20 +2589,27 @@ void ggml_cuda_mul_mat_vec_q( gate_w_stride = (int) (gx->weights->nb[1]/sizeof(float)); gate_tau = gx->tau; } - // Gate on a scratch COPY of ids: consumers of the same ids tensor that - // take the non-grouped path (e.g. a sibling weight whose quant type is - // not grouped-capable) must never see the -1 drop sentinels, which the - // legacy kernel would cast to uint32_t and read out of bounds. The - // in-place gate-weight renorm stays shared: dropped slots get weight - // 0.0f, which keeps any legacy consumer bit-correct (it just computes - // an expert that contributes nothing). - ggml_cuda_pool_alloc ids_gated(ctx.pool(), (size_t) (nchannels_dst*ncols_dst)); - CUDA_CHECK(cudaMemcpy2DAsync(ids_gated.ptr, nchannels_dst*sizeof(int32_t), - ids_d, ids_stride*sizeof(int32_t), - nchannels_dst*sizeof(int32_t), ncols_dst, - cudaMemcpyDeviceToDevice, stream)); - mmid_group_prep<<<1, MMID_GROUPED_MAX_PAIRS, 0, stream>>>( - ids_gated.ptr, mmid_meta.ptr, (int) nchannels_dst, (int) ncols_dst, (int) nchannels_dst, + // Adaptive-k writes -1 drop sentinels, so it must gate a scratch copy: + // sibling weights can still take the legacy kernel, which interprets + // ids as unsigned. Fixed top-k only reads ids and avoids both the + // allocation and the tiny device-to-device copy on every expert op. + ggml_cuda_pool_alloc ids_gated(ctx.pool()); + int32_t * prep_ids = const_cast(ids_d); + int prep_ids_stride = (int) ids_stride; + if (gate_w != nullptr) { + prep_ids = ids_gated.alloc((size_t) np); + prep_ids_stride = (int) nchannels_dst; + CUDA_CHECK(cudaMemcpy2DAsync(prep_ids, nchannels_dst*sizeof(int32_t), + ids_d, ids_stride*sizeof(int32_t), + nchannels_dst*sizeof(int32_t), ncols_dst, + cudaMemcpyDeviceToDevice, stream)); + } + // q4/top4 has just 16 pairs. Launch the smallest whole-warp block + // instead of 256 threads; __syncthreads() still covers every pair. + const int prep_warp = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; + const int prep_threads = ((np + prep_warp - 1)/prep_warp)*prep_warp; + mmid_group_prep<<<1, prep_threads, 0, stream>>>( + prep_ids, mmid_meta.ptr, (int) nchannels_dst, (int) ncols_dst, prep_ids_stride, gate_w, gate_w_stride, gate_tau); CUDA_CHECK(cudaGetLastError()); if (mul_mat_vec_q_grouped_dispatch( From e29854f268b22796c2ab9f69b29b6777b4de93b5 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:58:01 -0700 Subject: [PATCH 12/15] docs(ds4): record the 54.9 tok/s Lucebox profile --- server/docs/DS4.md | 24 ++++++++++++++++++++++++ server/docs/ENVIRONMENT.md | 3 +++ variables.md | 2 ++ 3 files changed, 29 insertions(+) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index 0b7343fb7..644c16a0a 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -119,6 +119,26 @@ q=1 measured 37.5 tok/s and 98.5 ms verification. These figures establish the execution-path A/B; they are workload-specific and are not a quality or held-out benchmark claim. +The 2026-07-20 Lucebox follow-up removes two remaining verification stalls: +it publishes all same-split peer copies with one cross-device dependency, and +it sorts the q4 routed pairs so repeated Strix experts reuse compressed weights +through the GPU caches. Enable the qualified gfx1201 + gfx1151 profile with: + +```bash +export DFLASH_DS4_TP_BATCH_PEER_COPIES=1 +export DFLASH_MMID_GROUPED=1 +export DFLASH_MMID_GROUPED_TYPES=8 +export DFLASH_MMID_GROUPED_DEVICE=1 # hip:1 is Strix Halo in this profile +``` + +The cleaned dual-architecture binary measured 54.7, 54.9, and 54.9 tok/s +after warm-up on the same fixed-q4 workload. All responses matched the +qualified SHA-256 output, acceptance remained 1.00, and mean target +verification fell to 64.9--65.2 ms. This is a target-execution improvement, +not an acceptance-rate gain. Leave `DFLASH_MMID_GROUPED_DEVICE` unset to make +the grouped kernel model-agnostic across eligible devices; device `1` is the +measured Lucebox-specific selector. + ### 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. @@ -213,6 +233,10 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. | `DFLASH_DS4_FUSED_VERIFY` | Enable the persistent q-wide target verification graph. | | `DFLASH_DS4_DENSE_TP_MASK` | Experimental row splitting for dense projections. This is automatically disabled when fused verification is enabled because split-buffer weights are not compatible with verifier graph replay. | | `DFLASH_DS4_TP_DEVICE_JOIN` | Join expert-owner results on device instead of through a host reduction. | +| `DFLASH_DS4_TP_BATCH_PEER_COPIES` | Batch same-split HIP peer copies behind one producer event and consumer wait. | +| `DFLASH_MMID_GROUPED` | Enable the grouped small-batch routed-expert kernel. | +| `DFLASH_MMID_GROUPED_TYPES` | Grouped-kernel type mask; `8` selects ROCmFP2 gate/up and ROCmFP3 down. | +| `DFLASH_MMID_GROUPED_DEVICE` | Optional HIP device selector; the qualified Lucebox profile uses `1` for Strix Halo. | | `DFLASH_DS4_DRAFT_CONTEXT_KV_CACHE` | Reuse DSpark context projection/KV state between speculative steps. | | `DFLASH_DS4_ADAPTIVE_WIDTH` | Set to `1` to select q=2/3/4 from DSpark confidence; `0` keeps the configured fixed verify width. If unset, the legacy `/tmp/ds4_awidth` control remains supported. | | `DFLASH_DS4_DRAFT_AHEAD` | On an independent in-process draft backend, launch the next proposal concurrently with target verification. The target still validates every candidate. | diff --git a/server/docs/ENVIRONMENT.md b/server/docs/ENVIRONMENT.md index f711a920c..479652287 100644 --- a/server/docs/ENVIRONMENT.md +++ b/server/docs/ENVIRONMENT.md @@ -24,6 +24,9 @@ consolidation of this list into CLI flags is tracked as follow-up work. | `DFLASH_ADAPTIVE_K_TAU` | 0 = off | Prefer the CLI: --adaptive-experts [tau]. Cumulative combine-weight threshold for per-token expert gating. | | `DFLASH_ADAPTIVE_K_DENSE` | per-model default | CSV of MoE layers kept dense under adaptive-K (DFlash capture layers). Warned-inert on families that do not thread layer indices yet. | | `DFLASH_MMID_GROUPED` | unset | Grouped MUL_MAT_ID kernel for small verify batches; candidate for CLI promotion. | +| `DFLASH_MMID_GROUPED_TYPES` | 7 | Grouped-kernel type mask; bit 3 (`8`) opts ROCmFP2/ROCmFP3 into the path. | +| `DFLASH_MMID_GROUPED_DEVICE` | -1 | Optional zero-based device restriction; unset/-1 applies to every eligible device. | +| `DFLASH_DS4_TP_BATCH_PEER_COPIES` | unset | Lucebox burn-in switch: publish same-split HIP peer copies with one cross-device dependency. | | `DFLASH_KVFLASH` | unset | Prefer the CLI: `--kvflash` (token count or `auto`). | ## Full inventory (generated) diff --git a/variables.md b/variables.md index b3bb5a232..93446fdcc 100644 --- a/variables.md +++ b/variables.md @@ -243,6 +243,8 @@ Untagged variables are operational tuning knobs. |---|---| | `DFLASH_MMID_GROUPED` | Grouped `MUL_MAT_ID` kernel for small verify batches. | | `DFLASH_MMID_GROUPED_TYPES` | Types eligible for the grouped MMID kernel. | +| `DFLASH_MMID_GROUPED_DEVICE` | Optional device restriction for the grouped MMID kernel. | +| `DFLASH_DS4_TP_BATCH_PEER_COPIES` | Batch a scheduler split's HIP peer copies behind one dependency. | | `DFLASH_MMQ_FULL_BATCH_MIN` / `DFLASH_MMQ_SUB_BATCH` | MMQ batch thresholds. | | `DFLASH_CUDA_MMVQ_TOKENWISE` / `DFLASH_CUDA_MMVQ_MOE_TOKENWISE` / `DFLASH_CUDA_MMVQ_MOE_KERNEL` | MMVQ token-wise / MoE kernel selection. | | `DFLASH_GDN_FORCE_GROUPED_COLS` / `DFLASH_GDN_NO_GROUPED_COLS` | Gated-delta-net grouped-column control. | From 44d3e11f756556eb6410aa0c190f2375a8cf0853 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Sat, 18 Jul 2026 12:13:30 -0700 Subject: [PATCH 13/15] feat(ds4): integrate indexed sparse prefill --- server/deps/llama.cpp/ggml/include/ggml-rpc.h | 4 +- server/deps/llama.cpp/ggml/include/ggml.h | 106 +- .../llama.cpp/ggml/src/ggml-backend-meta.cpp | 17 +- .../llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c | 260 +- .../llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp | 5 + .../llama.cpp/ggml/src/ggml-cuda/argsort.cu | 9 +- .../llama.cpp/ggml/src/ggml-cuda/argsort.cuh | 4 +- .../llama.cpp/ggml/src/ggml-cuda/common.cuh | 7 + .../llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu | 179 +- .../ggml/src/ggml-cuda/ds4-indexer.cu | 409 ++++ .../ggml/src/ggml-cuda/ds4-indexer.cuh | 15 + .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 1865 +++++++++++++++ .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 228 +- .../deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu | 126 +- .../deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh | 281 ++- .../llama.cpp/ggml/src/ggml-cuda/quantize.cu | 65 +- .../llama.cpp/ggml/src/ggml-cuda/quantize.cuh | 6 + .../template-instances/generate_cu_files.py | 1 + .../mmq-instance-q2_0_rocmfp2.cu | 6 + .../mmq-instance-q3_0_rocmfpx.cu | 6 + .../mmq-instance-q4_0_rocmfp4_fast.cu | 6 + .../llama.cpp/ggml/src/ggml-cuda/top-k.cu | 4 +- .../llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh | 14 +- server/deps/llama.cpp/ggml/src/ggml.c | 236 +- server/src/common/backend_factory.cpp | 19 + server/src/common/backend_factory.h | 5 +- server/src/common/prefill_attention_mode.h | 28 + server/src/deepseek4/deepseek4_backend.cpp | 85 +- server/src/deepseek4/deepseek4_backend.h | 2 + .../src/deepseek4/deepseek4_fused_verify.inc | 3 +- server/src/deepseek4/deepseek4_graph.cpp | 2093 +++++++++++++++-- server/src/deepseek4/deepseek4_internal.h | 13 + server/src/deepseek4/deepseek4_loader.cpp | 1 + server/src/server/server_main.cpp | 25 + server/tests/test_deepseek4_unit.cpp | 657 ++++++ 35 files changed, 6366 insertions(+), 424 deletions(-) create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cu create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cuh create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu create mode 100644 server/src/common/prefill_attention_mode.h diff --git a/server/deps/llama.cpp/ggml/include/ggml-rpc.h b/server/deps/llama.cpp/ggml/include/ggml-rpc.h index 3f7ab5859..398f614c6 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-rpc.h +++ b/server/deps/llama.cpp/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 3 #define RPC_PROTO_MINOR_VERSION 6 -#define RPC_PROTO_PATCH_VERSION 2 +#define RPC_PROTO_PATCH_VERSION 4 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index a09031452..706d39ab0 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -603,6 +603,16 @@ extern "C" { GGML_OP_DS4_HC, // Fused DeepSeek4 hyper-connection pre/post/out mixing + GGML_OP_DS4_INDEXER_QAT, // DS4 indexer Hadamard + FP4 activation simulation + + GGML_OP_DS4_INDEXER_SCORE, // Fused DS4 indexer dot/ReLU/head reduction + + GGML_OP_DS4_INDEXER_MASK, // Apply per-token DS4 indexer top-k to an attention mask + + // Keep extension operations appended so established GGML op ordinals + // remain stable within a protocol generation. + GGML_OP_MUL_MAT_GROUPED_SRC, + GGML_OP_COUNT, }; @@ -1446,6 +1456,23 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * b); + // Matrix multiply where b is physically [K/group, N, group] but is + // consumed as logical [K, N]. This is a distinct operation so backends + // that do not implement the grouped layout reject it instead of silently + // treating the physical storage as an ordinary flattened matrix. + // CUDA/HIP MMQ fuses the gather into its Q8 activation quantizer; the CPU + // implementation materializes each logical activation row in scratch. + GGML_API struct ggml_tensor * ggml_mul_mat_grouped_src( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API bool ggml_mul_mat_is_grouped_src( + const struct ggml_tensor * tensor); + + GGML_API int64_t ggml_mul_mat_grouped_src_groups( + const struct ggml_tensor * tensor); + // change the precision of a matrix multiplication // set to GGML_PREC_F32 for higher precision (useful for phi-2) GGML_API void ggml_mul_mat_set_prec( @@ -2377,6 +2404,40 @@ extern "C" { struct ggml_tensor * a, enum ggml_prec prec); + // DS4 layout and block-sparse policy for flash_attn_ext. raw_window is the + // maximum visible span inside the raw-row region. Compressed rows are + // selected in fixed-size blocks, capped to keep_rows. Zero leaves the + // operation dense. A negative value means the mask already contains an + // exact row selection and abs(keep_rows) is its maximum compressed width. + GGML_API void ggml_flash_attn_ext_set_ds4_sparse( + struct ggml_tensor * a, + int raw_rows, + int raw_window, + int keep_rows, + int block_size); + + // Fuse DS4's inverse 64-d tail RoPE into the D=512 flash-attention + // writeback. q_unrotated additionally asks the kernel to apply the forward + // tail RoPE to Q from shared F32. This is exact-only plumbing: both paths + // retain the explicit F32 rounding boundary and YaRN arithmetic used by + // GGML_OP_ROPE. + GGML_API void ggml_flash_attn_ext_set_ds4_inverse_rope( + struct ggml_tensor * a, + int kv_start, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + int n_ctx_orig, + bool q_unrotated); + + // True when flash_attn_ext carries the DS4 sparse-layout or fused-RoPE + // contract. Backends must implement that complete contract or reject it. + GGML_API bool ggml_flash_attn_ext_is_ds4( + const struct ggml_tensor * a); + GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec( const struct ggml_tensor * a); @@ -2471,9 +2532,12 @@ extern "C" { struct ggml_context * ctx, struct ggml_tensor * src); - // Fused DeepSeek4 hyper-connection helpers (decode, n_tokens == 1). - // ggml_ds4_hc_pre: mix[2*n_hc+n_hc^2] + base + hc_state[n_embd*n_hc] -> - // dst[n_embd + 2*n_hc + n_hc^2] = { working, split(pre,post,comb) } + // Fused DeepSeek4 hyper-connection helpers. The first dimension is the + // per-token payload; an optional second dimension batches independent + // tokens without changing their arithmetic. + // ggml_ds4_hc_pre: mix[mix_dim,n_tokens] + base[mix_dim] + + // hc_state[n_embd*n_hc,n_tokens] -> + // dst[n_embd+mix_dim,n_tokens] = { working, split(pre,post,comb) } GGML_API struct ggml_tensor * ggml_ds4_hc_pre( struct ggml_context * ctx, struct ggml_tensor * mix, @@ -2485,7 +2549,8 @@ extern "C" { float post_scale, float comb_scale); - // ggml_ds4_hc_post: residual hc_state + block_out + split -> new hc_state + // ggml_ds4_hc_post: residual hc_state + block_out + split -> new hc_state; + // all non-base tensors may carry the same n_tokens second dimension. GGML_API struct ggml_tensor * ggml_ds4_hc_post( struct ggml_context * ctx, struct ggml_tensor * residual_hc, @@ -2504,7 +2569,8 @@ extern "C" { struct ggml_tensor * split, int n_hc); - // ggml_ds4_hc_out: output-stage merge of hc streams into one embedding + // ggml_ds4_hc_out: output-stage merge of hc streams into one embedding per + // token. mix and hc_state may carry the same n_tokens second dimension. GGML_API struct ggml_tensor * ggml_ds4_hc_out( struct ggml_context * ctx, struct ggml_tensor * mix, @@ -2513,6 +2579,36 @@ extern "C" { int n_hc, float pre_scale); + // Official DS4 ratio-4 indexer transform. Each contiguous 128-wide F32 + // row is Hadamard-rotated and passed through the model's blockwise FP4 + // activation-simulation round trip. The operation is out-of-place so the + // pre-QAT query remains available to the main attention path. + GGML_API struct ggml_tensor * ggml_ds4_indexer_qat( + struct ggml_context * ctx, + struct ggml_tensor * input); + + // Compute the official ratio-4 indexer score matrix without materializing + // [n_comp,n_head,n_tokens] per-head dots. q is F32 + // [128,n_head,n_tokens], head_weights is F32 [n_head,n_tokens], and + // index_comp is F16 [128,n_comp]. head_weights already includes the + // model's 1/sqrt(128*n_head) scale. + GGML_API struct ggml_tensor * ggml_ds4_indexer_score( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * head_weights, + struct ggml_tensor * index_comp, + int kv_start, + int ratio); + + // Preserve the raw rows of base_mask and retain only selected compressed + // rows. selected is I32 [top_k,n_tokens], indexing the compressed span; + // base_mask is F32 [raw_rows+n_comp,n_tokens]. + GGML_API struct ggml_tensor * ggml_ds4_indexer_mask( + struct ggml_context * ctx, + struct ggml_tensor * base_mask, + struct ggml_tensor * selected, + int raw_rows); + // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp index 8e4273212..bf974d722 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp @@ -846,6 +846,12 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co case GGML_OP_MUL_MAT_ID: { split_state = handle_mul_mat(src_ss); } break; + case GGML_OP_MUL_MAT_GROUPED_SRC: { + // The logical activation rows are assembled from a private + // grouped physical layout. Keep the operation local; only the + // CPU and CUDA/HIP backends implement that layout contract. + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; case GGML_OP_OUT_PROD: { split_state = handle_generic(src_ss, /*scalar_only =*/ true); } break; @@ -962,6 +968,16 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co case GGML_OP_GATED_DELTA_NET: { split_state = handle_gated_delta_net(src_ss); } break; + case GGML_OP_DS4_INDEXER_QAT: { + split_state = handle_per_row(src_ss); + } break; + case GGML_OP_DS4_INDEXER_SCORE: + case GGML_OP_DS4_INDEXER_MASK: { + // These fused DS4 indexer operations combine axes from + // multiple inputs. Keep them local unless every source is + // mirrored on each simple backend. + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; case GGML_OP_UNARY: { split_state = handle_generic(src_ss, /*scalar_only =*/ false); } break; @@ -1922,4 +1938,3 @@ ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, siz const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context; return backend_ctx->backend_configs[index].backend; } - 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 df6ee5144..4193c8081 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 @@ -78,6 +78,168 @@ float ggml_table_f32_f16[1 << 16]; // precomputed f32 table for e8m0 half (1 KB) (simd-mappings.h) float ggml_table_f32_e8m0_half[1 << 8]; +static float ggml_ds4_indexer_e2m1_value_cpu(int i) { + static const float values[8] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + }; + return values[i & 7]; +} + +static float ggml_ds4_indexer_e2m1_round_cpu(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float ax = fminf(fabsf(x), 6.0f); + int best = 0; + float best_diff = fabsf(ax - ggml_ds4_indexer_e2m1_value_cpu(0)); + for (int i = 1; i < 8; ++i) { + const float diff = fabsf(ax - ggml_ds4_indexer_e2m1_value_cpu(i)); + if (diff < best_diff || + (diff == best_diff && (i & 1) == 0 && (best & 1) != 0)) { + best = i; + best_diff = diff; + } + } + return sign * ggml_ds4_indexer_e2m1_value_cpu(best); +} + +static void ggml_compute_forward_ds4_indexer_qat( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + const struct ggml_tensor * src = dst->src[0]; + GGML_ASSERT(src && src->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src->ne[0] == 128 && ggml_are_same_shape(src, dst)); + GGML_ASSERT(ggml_is_contiguous(src) && ggml_is_contiguous(dst)); + + const int64_t n_rows = ggml_nrows(src); + for (int64_t row = params->ith; row < n_rows; row += params->nth) { + float values[128]; + const float * src_row = (const float *) ((const char *) src->data + row * src->nb[1]); + float * dst_row = (float *) ((char *) dst->data + row * dst->nb[1]); + memcpy(values, src_row, sizeof(values)); + + for (int stride = 1; stride < 128; stride <<= 1) { + for (int base = 0; base < 128; base += 2 * stride) { + for (int i = 0; i < stride; ++i) { + const float a = values[base + i]; + const float b = values[base + stride + i]; + values[base + i] = a + b; + values[base + stride + i] = a - b; + } + } + } + for (int i = 0; i < 128; ++i) { + values[i] *= 0.08838834764831845f; + } + for (int block = 0; block < 4; ++block) { + float amax = 0.0f; + for (int i = 0; i < 32; ++i) { + amax = fmaxf(amax, fabsf(values[block * 32 + i])); + } + amax = fmaxf(amax, 7.052966104933725e-38f); + const float scale = exp2f(ceilf(log2f(amax / 6.0f))); + for (int i = 0; i < 32; ++i) { + const int index = block * 32 + i; + const float normalized = fminf( + 6.0f, fmaxf(-6.0f, values[index] / scale)); + dst_row[index] = + ggml_ds4_indexer_e2m1_round_cpu(normalized) * scale; + } + } + } +} + +static void ggml_compute_forward_ds4_indexer_score( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + const struct ggml_tensor * q = dst->src[0]; + const struct ggml_tensor * weights = dst->src[1]; + const struct ggml_tensor * comp = dst->src[2]; + GGML_ASSERT(q && weights && comp); + GGML_ASSERT(q->type == GGML_TYPE_F32 && q->ne[0] == 128); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(comp->type == GGML_TYPE_F16 && comp->ne[0] == 128); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(q) && ggml_is_contiguous(weights)); + GGML_ASSERT(ggml_is_contiguous(comp) && ggml_is_contiguous(dst)); + + const int n_head = (int) q->ne[1]; + const int n_tokens = (int) q->ne[2]; + const int n_comp = (int) comp->ne[1]; + const int kv_start = ggml_get_op_params_i32(dst, 0); + const int ratio = ggml_get_op_params_i32(dst, 1); + GGML_ASSERT(weights->ne[0] == n_head && weights->ne[1] == n_tokens); + GGML_ASSERT(dst->ne[0] == n_comp && dst->ne[1] == n_tokens); + GGML_ASSERT(kv_start >= 0 && ratio > 0); + + const float * q_data = (const float *) q->data; + const float * weight_data = (const float *) weights->data; + const ggml_fp16_t * comp_data = (const ggml_fp16_t *) comp->data; + float * dst_data = (float *) dst->data; + for (int token = params->ith; token < n_tokens; token += params->nth) { + const int visible = (kv_start + token + 1) / ratio; + for (int c = 0; c < n_comp; ++c) { + if (c >= visible) { + dst_data[(size_t) token * n_comp + c] = -1.0e30f; + continue; + } + const ggml_fp16_t * k = comp_data + (size_t) c * 128; + float score = 0.0f; + for (int h = 0; h < n_head; ++h) { + const float * qh = q_data + + ((size_t) token * n_head + h) * 128; + float dot = 0.0f; + for (int d = 0; d < 128; ++d) { + dot += qh[d] * GGML_FP16_TO_FP32(k[d]); + } + score += fmaxf(dot, 0.0f) * + weight_data[(size_t) token * n_head + h]; + } + dst_data[(size_t) token * n_comp + c] = score; + } + } +} + +static void ggml_compute_forward_ds4_indexer_mask( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + const struct ggml_tensor * base = dst->src[0]; + const struct ggml_tensor * selected = dst->src[1]; + GGML_ASSERT(base && selected); + GGML_ASSERT(base->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + GGML_ASSERT(selected->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_are_same_shape(base, dst)); + GGML_ASSERT(ggml_is_contiguous(base) && ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_is_contiguous(selected)); + + const int n_attn = (int) base->ne[0]; + const int n_tokens = (int) ggml_nrows(base); + const int top_k = (int) selected->ne[0]; + const int raw_rows = ggml_get_op_params_i32(dst, 0); + const int n_comp = n_attn - raw_rows; + GGML_ASSERT(raw_rows >= 0 && n_comp >= 0); + GGML_ASSERT(ggml_nrows(selected) == n_tokens); + + for (int token = params->ith; token < n_tokens; token += params->nth) { + const float * base_row = + (const float *) ((const char *) base->data + (size_t) token * base->nb[1]); + float * dst_row = + (float *) ((char *) dst->data + (size_t) token * dst->nb[1]); + const int32_t * selected_row = + (const int32_t *) ((const char *) selected->data + + (size_t) token * selected->nb[1]); + memcpy(dst_row, base_row, (size_t) raw_rows * sizeof(float)); + for (int row = raw_rows; row < n_attn; ++row) { + dst_row[row] = -1.0e30f; + } + for (int k = 0; k < top_k; ++k) { + const int comp = selected_row[k]; + if (comp >= 0 && comp < n_comp) { + dst_row[raw_rows + comp] = base_row[raw_rows + comp]; + } + } + } +} + #if defined(__ARM_ARCH) struct ggml_arm_arch_features_type { int sve_cnt; @@ -1164,6 +1326,7 @@ static void ggml_compute_forward_mul_mat_one_chunk( GGML_TENSOR_BINARY_OP_LOCALS const bool src1_cont = ggml_is_contiguous(src1); + const bool grouped_src = ggml_mul_mat_is_grouped_src(dst); ggml_vec_dot_t const vec_dot = type_traits_cpu[type].vec_dot; enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type; @@ -1179,7 +1342,8 @@ static void ggml_compute_forward_mul_mat_one_chunk( return; } - const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; + const void * wdata = (src1->type == vec_dot_type && !grouped_src) + ? src1->data : params->wdata; const size_t row_size = ggml_row_size(vec_dot_type, ne10); assert(ne12 % ne02 == 0); @@ -1253,6 +1417,8 @@ void ggml_compute_forward_mul_mat( enum ggml_type const vec_dot_type = type_traits_cpu[src0->type].vec_dot_type; ggml_from_float_t const from_float = type_traits_cpu[vec_dot_type].from_float; int64_t const vec_dot_num_rows = type_traits_cpu[src0->type].nrows; + const bool grouped_src = ggml_mul_mat_is_grouped_src(dst); + const bool convert_src1 = grouped_src || src1->type != vec_dot_type; GGML_ASSERT(ne0 == ne01); GGML_ASSERT(ne1 == ne11); @@ -1280,7 +1446,7 @@ void ggml_compute_forward_mul_mat( const bool src1_cont = ggml_is_contiguous(src1); - if (src1_cont) { + if (src1_cont && !grouped_src) { for (int64_t i13 = 0; i13 < ne13; i13++) for (int64_t i12 = 0; i12 < ne12; i12++) if (!llamafile_sgemm(params, @@ -1300,41 +1466,59 @@ void ggml_compute_forward_mul_mat( UseGgmlGemm1:; #endif - if (src1->type != vec_dot_type) { + if (convert_src1) { char * wdata = params->wdata; const size_t nbw0 = ggml_type_size(vec_dot_type); const size_t nbw1 = ggml_row_size(vec_dot_type, ne10); const size_t nbw2 = nbw1*ne11; const size_t nbw3 = nbw2*ne12; + const size_t converted_size = ne13*nbw3; - assert(params->wsize >= ne13*nbw3); + assert(params->wsize >= converted_size); GGML_ASSERT(src1->type == GGML_TYPE_F32); - #if 0 - for (int64_t i13 = 0; i13 < ne13; ++i13) { - for (int64_t i12 = 0; i12 < ne12; ++i12) { - for (int64_t i11 = ith; i11 < ne11; i11 += nth) { - from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11), - (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1), - ne10); + if (grouped_src) { + const struct ggml_tensor * physical = src1->view_src; + const int64_t groups = ggml_mul_mat_grouped_src_groups(dst); + GGML_ASSERT(physical && physical->type == GGML_TYPE_F32); + GGML_ASSERT(physical->ne[0] * groups == ne10); + GGML_ASSERT(physical->ne[1] == ne11); + GGML_ASSERT(physical->ne[2] == groups && physical->ne[3] == 1); + GGML_ASSERT(ggml_is_contiguous(physical)); + + const size_t scratch_offset = GGML_PAD(converted_size, CACHE_LINE_SIZE); + const size_t scratch_stride = (size_t) ne10 * sizeof(float); + GGML_ASSERT(params->wsize >= + scratch_offset + (size_t) nth * scratch_stride); + float * scratch = (float *) (wdata + scratch_offset + + (size_t) ith * scratch_stride); + const int64_t group_width = physical->ne[0]; + for (int64_t i11 = ith; i11 < ne11; i11 += nth) { + for (int64_t group = 0; group < groups; ++group) { + const float * src_group = (const float *) + ((const char *) physical->data + + (size_t) i11 * physical->nb[1] + + (size_t) group * physical->nb[2]); + memcpy(scratch + group * group_width, src_group, + (size_t) group_width * sizeof(float)); } + from_float(scratch, wdata + (size_t) i11 * nbw1, ne10); } - } - #else - for (int64_t i13 = 0; i13 < ne13; ++i13) { - for (int64_t i12 = 0; i12 < ne12; ++i12) { - for (int64_t i11 = 0; i11 < ne11; ++i11) { - size_t bs = ggml_blck_size(vec_dot_type); - int64_t ne10_block_start = (ith * ne10/bs) / nth; - int64_t ne10_block_end = ((ith + 1) * ne10/bs) / nth; - from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11 + ne10_block_start*bs*nb10), - (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0), - (ne10_block_end - ne10_block_start) * bs); + } else { + for (int64_t i13 = 0; i13 < ne13; ++i13) { + for (int64_t i12 = 0; i12 < ne12; ++i12) { + for (int64_t i11 = 0; i11 < ne11; ++i11) { + size_t bs = ggml_blck_size(vec_dot_type); + int64_t ne10_block_start = (ith * ne10/bs) / nth; + int64_t ne10_block_end = ((ith + 1) * ne10/bs) / nth; + from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11 + ne10_block_start*bs*nb10), + (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0), + (ne10_block_end - ne10_block_start) * bs); + } } } } - #endif } if (ith == 0) { @@ -1345,8 +1529,8 @@ UseGgmlGemm1:; ggml_barrier(params->threadpool); #if GGML_USE_LLAMAFILE - if (src1->type != vec_dot_type) { - const void* wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; + if (convert_src1) { + const void* wdata = params->wdata; const size_t row_size = ggml_row_size(vec_dot_type, ne10); for (int64_t i13 = 0; i13 < ne13; i13++) @@ -1815,6 +1999,7 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm ggml_compute_forward_l2_norm(params, tensor); } break; case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_GROUPED_SRC: { ggml_compute_forward_mul_mat(params, tensor); } break; @@ -1830,6 +2015,18 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { GGML_ABORT("GGML_OP_DS4_HC is only implemented for CUDA"); } + case GGML_OP_DS4_INDEXER_QAT: + { + ggml_compute_forward_ds4_indexer_qat(params, tensor); + } break; + case GGML_OP_DS4_INDEXER_SCORE: + { + ggml_compute_forward_ds4_indexer_score(params, tensor); + } break; + case GGML_OP_DS4_INDEXER_MASK: + { + ggml_compute_forward_ds4_indexer_mask(params, tensor); + } break; case GGML_OP_OUT_PROD: { ggml_compute_forward_out_prod(params, tensor); @@ -2308,6 +2505,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_GROUP_NORM: case GGML_OP_CONCAT: case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_GROUPED_SRC: case GGML_OP_MUL_MAT_ID: case GGML_OP_OUT_PROD: { @@ -2374,6 +2572,9 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_TIMESTEP_EMBEDDING: case GGML_OP_ARGSORT: case GGML_OP_TOP_K: + case GGML_OP_DS4_INDEXER_QAT: + case GGML_OP_DS4_INDEXER_SCORE: + case GGML_OP_DS4_INDEXER_MASK: case GGML_OP_FLASH_ATTN_EXT: case GGML_OP_FLASH_ATTN_SPARSE: case GGML_OP_FLASH_ATTN_BACK: @@ -2825,11 +3026,18 @@ struct ggml_cplan ggml_graph_plan( cur = ggml_type_size(node->type)*n_tasks; } break; case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_GROUPED_SRC: { const enum ggml_type vec_dot_type = type_traits_cpu[node->src[0]->type].vec_dot_type; + const bool grouped_src = ggml_mul_mat_is_grouped_src(node); - if (node->src[1]->type != vec_dot_type) { + if (node->src[1]->type != vec_dot_type || grouped_src) { cur = ggml_row_size(vec_dot_type, ggml_nelements(node->src[1])); + if (grouped_src) { + cur = GGML_PAD(cur, CACHE_LINE_SIZE); + cur += (size_t) node->src[1]->ne[0] * + sizeof(float) * (size_t) n_tasks; + } } } break; case GGML_OP_MUL_MAT_ID: diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp index 49f840be2..26bab55b9 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -450,6 +450,7 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st op->type != GGML_TYPE_IQ1_S && op->type != GGML_TYPE_IQ1_M; // missing type_traits.from_float case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_GROUPED_SRC: return src1->type == GGML_TYPE_F32 || src1->type == ggml_get_type_traits_cpu(src0->type)->vec_dot_type; case GGML_OP_SOFT_MAX_BACK: { if (op->src[0]->type != GGML_TYPE_F32 || op->src[1]->type != GGML_TYPE_F32) { @@ -465,6 +466,10 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st return src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32; case GGML_OP_GET_ROWS_BACK: return src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16; + case GGML_OP_FLASH_ATTN_EXT: + // DS4 attention carries sparse-layout and fused-RoPE semantics. + // The generic CPU kernel does not implement that contract. + return !ggml_flash_attn_ext_is_ds4(op); case GGML_OP_OUT_PROD: return (src0->type == GGML_TYPE_F32 || (ggml_is_quantized(src0->type) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) && src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu index 0f3f017b5..ac0634a16 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu @@ -6,6 +6,9 @@ # define STRIDED_ITERATOR_AVAILABLE # endif using namespace cub; +#elif defined(GGML_CUDA_USE_HIPCUB) +# include +using namespace hipcub; #endif // GGML_CUDA_USE_CUB static __global__ void init_indices(int * indices, const int ncols, const int nrows) { @@ -26,7 +29,7 @@ static __global__ void init_offsets(int * offsets, const int ncols, const int nr } #endif // STRIDED_ITERATOR_AVAILABLE -#ifdef GGML_CUDA_USE_CUB +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const float * x, int * dst, @@ -138,7 +141,7 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, } } } -#endif // GGML_CUDA_USE_CUB +#endif // GGML_CUDA_USE_CUB || GGML_CUDA_USE_HIPCUB // Bitonic sort implementation template @@ -248,7 +251,7 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { enum ggml_sort_order order = (enum ggml_sort_order) dst->op_params[0]; -#ifdef GGML_CUDA_USE_CUB +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) const int ncols_pad = next_power_of_2(ncols); const size_t shared_mem = ncols_pad * sizeof(int); const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh index 22b7306f2..018dc9fb2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh @@ -2,7 +2,7 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst); -#ifdef GGML_CUDA_USE_CUB +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const float * x, int * dst, @@ -10,7 +10,7 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const int nrows, ggml_sort_order order, cudaStream_t stream); -#endif // GGML_CUDA_USE_CUB +#endif // GGML_CUDA_USE_CUB || GGML_CUDA_USE_HIPCUB void argsort_f32_i32_cuda_bitonic(const float * x, int * dst, const int ncols, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh index b398c5a17..763a087df 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh @@ -109,6 +109,13 @@ # define GGML_CUDA_USE_CUB #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +// rocPRIM ships hipCUB, a CUB-compatible device-sort API. The ordinary +// single-block bitonic argsort cannot represent DS4's long-context indexer once +// the number of compressed rows exceeds 1024, so use hipCUB for that case. +#if defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +# define GGML_CUDA_USE_HIPCUB +#endif // defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + #ifdef __CUDA_ARCH_LIST__ constexpr bool ggml_cuda_has_arch_impl(int) { return false; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu index cf09c8a88..2e8ebd61b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu @@ -2,26 +2,26 @@ // Fused DeepSeek4 hyper-connection ops. // -// mode 0 (pre): src0 = mix [mix_dim] (f32, from fn @ rms_norm(hc_state)) +// mode 0 (pre): src0 = mix [mix_dim,n_tokens] (f32, from fn @ rms_norm(hc_state)) // src1 = base [mix_dim] (f32) -// src2 = hc_state [n_embd*n_hc] (f32, raw residual streams) -// dst = [n_embd + mix_dim]: +// src2 = hc_state [n_embd*n_hc,n_tokens] (f32, raw residual streams) +// dst = [n_embd + mix_dim,n_tokens]: // dst[0..n_embd) = working vector (pre-mixed input) // dst[n_embd..n_embd+mix) = split = {pre[n_hc], post[n_hc], comb[n_hc*n_hc]} // Math matches cpu_hc_sinkhorn + finish_hc_pre_from_mix_into in // deepseek4_graph.cpp (sigmoid gates + Sinkhorn-normalized combine). // -// mode 1 (post): src0 = residual hc_state [n_embd*n_hc] -// src1 = block_out [n_embd] -// src2 = split [mix_dim] (view of a mode-0 dst tail) -// dst = new hc_state [n_embd*n_hc]: +// mode 1 (post): src0 = residual hc_state [n_embd*n_hc,n_tokens] +// src1 = block_out [n_embd,n_tokens] +// src2 = split [mix_dim,n_tokens] (view of a mode-0 dst tail) +// dst = new hc_state [n_embd*n_hc,n_tokens]: // dst[h*n_embd+d] = post[h]*block_out[d] // + sum_src comb[h + src*n_hc] * residual[src*n_embd+d] // -// mode 2 (out): src0 = mix [n_hc] +// mode 2 (out): src0 = mix [n_hc,n_tokens] // src1 = base [n_hc] -// src2 = hc_state [n_embd*n_hc] -// dst = [n_embd]: weights[h] = sigmoid(mix[h]*s0+base[h]) + 1e-6; +// src2 = hc_state [n_embd*n_hc,n_tokens] +// dst = [n_embd,n_tokens]: weights[h] = sigmoid(mix[h]*s0+base[h]) + 1e-6; // dst[d] = sum_h weights[h]*hc_state[h*n_embd+d] // // mode 3 (post split): mode 1 with src1 = main block_out, src3 = peer @@ -187,7 +187,15 @@ static __global__ void ds4_hc_pre_kernel_t( int iters, float pre_scale, float post_scale, - float comb_scale) { + float comb_scale, + size_t mix_stride, + size_t hc_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + mix += (size_t) token * mix_stride; + hc_state += (size_t) token * hc_stride; + dst += (size_t) token * dst_stride; + __shared__ float split[DS4_HC_MAX_MIX]; __shared__ float s_mix[DS4_HC_MAX_MIX]; __shared__ float s_base[DS4_HC_MAX_MIX]; @@ -222,6 +230,78 @@ static __global__ void ds4_hc_pre_kernel_t( } } +// Large prefill batches already expose thousands of token blocks, so they do +// not need every embedding tile to recompute the same serial Sinkhorn. Split +// it into one deterministic solve per token followed by the parallel mixing +// pass. The small-token path keeps the fused kernel above to avoid an extra +// launch during decode and speculative verification. +template +static __global__ void ds4_hc_pre_split_kernel_t( + const float * __restrict__ mix, + const float * __restrict__ base, + float * __restrict__ dst, + int n_embd, + int iters, + float pre_scale, + float post_scale, + float comb_scale, + size_t mix_stride, + size_t dst_stride) { + const int token = (int) blockIdx.x; + mix += (size_t) token * mix_stride; + dst += (size_t) token * dst_stride; + + __shared__ float split[DS4_HC_MAX_MIX]; + __shared__ float s_mix[DS4_HC_MAX_MIX]; + __shared__ float s_base[DS4_HC_MAX_MIX]; + constexpr int mix_dim = 2 * NHC + NHC * NHC; + const int tid = (int) threadIdx.x; + + if (tid < mix_dim) { + s_mix[tid] = mix[tid]; + s_base[tid] = base[tid]; + } + __syncthreads(); + + if (tid == 0) { + ds4_hc_sinkhorn_split_t( + s_mix, s_base, pre_scale, post_scale, comb_scale, iters, split); +#pragma unroll + for (int i = 0; i < mix_dim; ++i) { + dst[n_embd + i] = split[i]; + } + } +} + +template +static __global__ void ds4_hc_pre_mix_kernel_t( + const float * __restrict__ hc_state, + float * __restrict__ dst, + int n_embd, + size_t hc_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + hc_state += (size_t) token * hc_stride; + dst += (size_t) token * dst_stride; + + __shared__ float pre[NHC]; + const int tid = (int) threadIdx.x; + if (tid < NHC) { + pre[tid] = dst[n_embd + tid]; + } + __syncthreads(); + + const int d = (int) blockIdx.x * (int) blockDim.x + tid; + if (d < n_embd) { + float acc = 0.0f; +#pragma unroll + for (int h = 0; h < NHC; ++h) { + acc += pre[h] * hc_state[(size_t) h * n_embd + d]; + } + dst[d] = acc; + } +} + static __global__ void ds4_hc_pre_kernel( const float * __restrict__ mix, const float * __restrict__ base, @@ -232,7 +312,15 @@ static __global__ void ds4_hc_pre_kernel( int iters, float pre_scale, float post_scale, - float comb_scale) { + float comb_scale, + size_t mix_stride, + size_t hc_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + mix += (size_t) token * mix_stride; + hc_state += (size_t) token * hc_stride; + dst += (size_t) token * dst_stride; + __shared__ float split[DS4_HC_MAX_MIX]; __shared__ float s_mix[DS4_HC_MAX_MIX]; __shared__ float s_base[DS4_HC_MAX_MIX]; @@ -275,7 +363,17 @@ static __global__ void ds4_hc_post_kernel( const float * __restrict__ split, float * __restrict__ dst, int n_embd, - int n_hc) { + int n_hc, + size_t residual_stride, + size_t block_out_stride, + size_t split_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + residual += (size_t) token * residual_stride; + block_out += (size_t) token * block_out_stride; + split += (size_t) token * split_stride; + dst += (size_t) token * dst_stride; + const int i = blockIdx.x * blockDim.x + threadIdx.x; const int total = n_embd * n_hc; if (i >= total) { @@ -326,7 +424,15 @@ static __global__ void ds4_hc_out_kernel( float * __restrict__ dst, int n_embd, int n_hc, - float pre_scale) { + float pre_scale, + size_t mix_stride, + size_t hc_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + mix += (size_t) token * mix_stride; + hc_state += (size_t) token * hc_stride; + dst += (size_t) token * dst_stride; + const int d = blockIdx.x * blockDim.x + threadIdx.x; if (d >= n_embd) { return; @@ -352,8 +458,14 @@ void ggml_cuda_op_ds4_hc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int mode = ggml_get_op_params_i32(dst, 0); const int n_embd = ggml_get_op_params_i32(dst, 1); const int n_hc = ggml_get_op_params_i32(dst, 2); + const int n_tokens = (int) dst->ne[1]; GGML_ASSERT(n_hc > 0 && n_hc <= DS4_HC_MAX_HC); + GGML_ASSERT(n_tokens > 0); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(src1->nb[0] == sizeof(float)); + GGML_ASSERT(src2->nb[0] == sizeof(float)); + GGML_ASSERT(dst->nb[0] == sizeof(float)); cudaStream_t stream = ctx.stream(); @@ -364,33 +476,54 @@ void ggml_cuda_op_ds4_hc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const float post_scale = ggml_get_op_params_f32(dst, 5); const float comb_scale = ggml_get_op_params_f32(dst, 6); const int pre_blocks = (n_embd + 255) / 256; - if (n_hc == 4) { - ds4_hc_pre_kernel_t<4><<>>( + const dim3 grid(pre_blocks, n_tokens, 1); + if (n_hc == 4 && n_tokens >= 64) { + ds4_hc_pre_split_kernel_t<4><<>>( + (const float *) src0->data, (const float *) src1->data, + (float *) dst->data, + n_embd, iters, pre_scale, post_scale, comb_scale, + src0->nb[1] / sizeof(float), dst->nb[1] / sizeof(float)); + ds4_hc_pre_mix_kernel_t<4><<>>( + (const float *) src2->data, (float *) dst->data, + n_embd, src2->nb[1] / sizeof(float), + dst->nb[1] / sizeof(float)); + } else if (n_hc == 4) { + ds4_hc_pre_kernel_t<4><<>>( (const float *) src0->data, (const float *) src1->data, (const float *) src2->data, (float *) dst->data, - n_embd, iters, pre_scale, post_scale, comb_scale); + n_embd, iters, pre_scale, post_scale, comb_scale, + src0->nb[1] / sizeof(float), src2->nb[1] / sizeof(float), + dst->nb[1] / sizeof(float)); } else { - ds4_hc_pre_kernel<<>>( + ds4_hc_pre_kernel<<>>( (const float *) src0->data, (const float *) src1->data, (const float *) src2->data, (float *) dst->data, - n_embd, n_hc, iters, pre_scale, post_scale, comb_scale); + n_embd, n_hc, iters, pre_scale, post_scale, comb_scale, + src0->nb[1] / sizeof(float), src2->nb[1] / sizeof(float), + dst->nb[1] / sizeof(float)); } } break; case 1: { const int total = n_embd * n_hc; const int blocks = (total + 255) / 256; - ds4_hc_post_kernel<<>>( + const dim3 grid(blocks, n_tokens, 1); + ds4_hc_post_kernel<<>>( (const float *) src0->data, (const float *) src1->data, (const float *) src2->data, (float *) dst->data, - n_embd, n_hc); + n_embd, n_hc, + src0->nb[1] / sizeof(float), src1->nb[1] / sizeof(float), + src2->nb[1] / sizeof(float), dst->nb[1] / sizeof(float)); } break; case 2: { const float pre_scale = ggml_get_op_params_f32(dst, 4); const int blocks = (n_embd + 255) / 256; - ds4_hc_out_kernel<<>>( + const dim3 grid(blocks, n_tokens, 1); + ds4_hc_out_kernel<<>>( (const float *) src0->data, (const float *) src1->data, (const float *) src2->data, (float *) dst->data, - n_embd, n_hc, pre_scale); + n_embd, n_hc, pre_scale, + src0->nb[1] / sizeof(float), src2->nb[1] / sizeof(float), + dst->nb[1] / sizeof(float)); } break; case 3: { const ggml_tensor * src3 = dst->src[3]; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cu new file mode 100644 index 000000000..f4358a775 --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cu @@ -0,0 +1,409 @@ +#include "ds4-indexer.cuh" + +#if defined(GGML_USE_HIP) +# include +namespace ds4_wmma = rocwmma; +# define DS4_INDEXER_WMMA_AVAILABLE 1 +#elif !defined(GGML_USE_MUSA) +# include +namespace ds4_wmma = nvcuda::wmma; +# define DS4_INDEXER_WMMA_AVAILABLE 1 +#else +# define DS4_INDEXER_WMMA_AVAILABLE 0 +#endif + +#if DS4_INDEXER_WMMA_AVAILABLE +# if defined(GGML_USE_HIP) && HIP_VERSION >= 60500000 +using ds4_indexer_wmma_half = _Float16; +# else +using ds4_indexer_wmma_half = half; +# endif +#endif + +// Keep this operation bit-for-bit aligned with the official DeepSeek V4 +// graph and antirez/ds4's dsv4_indexer_qat implementation. Indexer query and +// compressed-key rows both pass through an orthonormal Hadamard-128 followed +// by one UE4M3-scaled E2M1 activation-simulation block per 32 values. + +static __device__ __forceinline__ float ds4_indexer_e2m1_value(int i) { + switch (i & 7) { + case 0: return 0.0f; + case 1: return 0.5f; + case 2: return 1.0f; + case 3: return 1.5f; + case 4: return 2.0f; + case 5: return 3.0f; + case 6: return 4.0f; + default: return 6.0f; + } +} + +static __device__ __forceinline__ float ds4_indexer_e2m1_round(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float ax = fminf(fabsf(x), 6.0f); + int best = 0; + float best_diff = fabsf(ax - ds4_indexer_e2m1_value(0)); +#pragma unroll + for (int i = 1; i < 8; ++i) { + const float diff = fabsf(ax - ds4_indexer_e2m1_value(i)); + // Round ties to the even E2M1 code, matching the converter/reference. + if (diff < best_diff || + (diff == best_diff && (i & 1) == 0 && (best & 1) != 0)) { + best = i; + best_diff = diff; + } + } + return sign * ds4_indexer_e2m1_value(best); +} + +static __global__ void ds4_indexer_qat_kernel( + float * dst, + const float * src, + int64_t n_rows, + int64_t src_row_stride, + int64_t dst_row_stride) { + constexpr int WIDTH = 128; + constexpr float HADAMARD_SCALE = 0.08838834764831845f; + const int64_t row = (int64_t) blockIdx.x; + const int tid = (int) threadIdx.x; + if (row >= n_rows || tid >= WIDTH) return; + + __shared__ float values[WIDTH]; + __shared__ float abs_values[WIDTH]; + const float * src_row = src + row * src_row_stride; + float * dst_row = dst + row * dst_row_stride; + values[tid] = src_row[tid]; + __syncthreads(); + + for (int stride = 1; stride < WIDTH; stride <<= 1) { + if ((tid & stride) == 0) { + const int base = + (tid & ~(2 * stride - 1)) + (tid & (stride - 1)); + const float a = values[base]; + const float b = values[base + stride]; + values[base] = a + b; + values[base + stride] = a - b; + } + __syncthreads(); + } + + const float value = values[tid] * HADAMARD_SCALE; + const int block = tid >> 5; + const int lane = tid & 31; + const int block_base = block * 32; + abs_values[tid] = fabsf(value); + __syncthreads(); + + for (int stride = 16; stride > 0; stride >>= 1) { + if (lane < stride) { + abs_values[block_base + lane] = fmaxf( + abs_values[block_base + lane], + abs_values[block_base + lane + stride]); + } + __syncthreads(); + } + + const float amax = fmaxf( + abs_values[block_base], 7.052966104933725e-38f); + const float scale = exp2f(ceilf(log2f(amax / 6.0f))); + const float normalized = fminf(6.0f, fmaxf(-6.0f, value / scale)); + dst_row[tid] = ds4_indexer_e2m1_round(normalized) * scale; +} + +void ggml_cuda_op_ds4_indexer_qat( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst) { + const ggml_tensor * src = dst->src[0]; + GGML_ASSERT(src && src->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src->ne[0] == 128 && dst->ne[0] == 128); + GGML_ASSERT(ggml_are_same_shape(src, dst)); + GGML_ASSERT(ggml_is_contiguous(src)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int64_t n_rows = ggml_nrows(src); + const int64_t src_row_stride = src->nb[1] / sizeof(float); + const int64_t dst_row_stride = dst->nb[1] / sizeof(float); + GGML_ASSERT(src_row_stride >= 128 && dst_row_stride >= 128); + + cudaStream_t stream = ctx.stream(); + ds4_indexer_qat_kernel<<<(unsigned) n_rows, 128, 0, stream>>>( + static_cast(dst->data), + static_cast(src->data), + n_rows, src_row_stride, dst_row_stride); + CUDA_CHECK(cudaGetLastError()); +} + +// Compute 16 query tokens against 128 compressed rows per block. QAT values +// are powers-of-two-scaled E2M1 and therefore exactly representable as F16 in +// the model's operating range; the compressed cache is already F16. WMMA +// removes the otherwise enormous [n_comp,64,n_tokens] intermediate while the +// ReLU, head weighting and reduction remain F32. +#if DS4_INDEXER_WMMA_AVAILABLE +static __global__ void ds4_indexer_score_wmma_kernel( + float * scores, + const float * q, + const float * weights, + const half * index_comp, + int n_comp, + int n_tokens, + int kv_start, + int n_head, + int ratio) { + const int tile_c = (int) blockIdx.x * 128; + const int tile_t = (int) blockIdx.y * 16; + const int tid = (int) threadIdx.x; + const int warp = tid >> 5; + + __shared__ half a_sh[16 * 128]; + __shared__ half b_sh[128 * 128]; + __shared__ float c_sh[8 * 16 * 16]; + + float acc[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) acc[i] = 0.0f; + + for (int i = tid; i < 128 * 128; i += 256) { + const int c = i >> 7; + const int d = i & 127; + const int comp = tile_c + c; + b_sh[d + c * 128] = comp < n_comp + ? index_comp[(size_t) comp * 128 + d] + : __float2half(0.0f); + } + __syncthreads(); + + for (int h = 0; h < n_head; ++h) { + for (int pair = tid; pair < 16 * 64; pair += 256) { + const int row = pair >> 6; + const int d = (pair & 63) * 2; + const int token = tile_t + row; + half2 value = __float2half2_rn(0.0f); + if (token < n_tokens) { + const float2 q_value = *reinterpret_cast( + q + ((size_t) token * n_head + h) * 128 + d); + value = __floats2half2_rn(q_value.x, q_value.y); + } + *reinterpret_cast(a_sh + row * 128 + d) = value; + } + __syncthreads(); + + ds4_wmma::fragment a_frag; + ds4_wmma::fragment b_frag; + ds4_wmma::fragment c_frag; + ds4_wmma::fill_fragment(c_frag, 0.0f); + const int col0 = warp * 16; + for (int k0 = 0; k0 < 128; k0 += 16) { + const ds4_indexer_wmma_half * a_wmma = + reinterpret_cast(a_sh); + const ds4_indexer_wmma_half * b_wmma = + reinterpret_cast(b_sh); + ds4_wmma::load_matrix_sync(a_frag, a_wmma + k0, 128); + ds4_wmma::load_matrix_sync( + b_frag, b_wmma + col0 * 128 + k0, 128); + ds4_wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); + } + ds4_wmma::store_matrix_sync( + c_sh + warp * 16 * 16, c_frag, 16, + ds4_wmma::mem_row_major); + __syncthreads(); + + const int token_for_lane = tile_t + (tid >> 4); + const float head_weight = token_for_lane < n_tokens + ? weights[(size_t) token_for_lane * n_head + h] + : 0.0f; + int slot = 0; + for (int i = tid; i < 8 * 16 * 16; i += 256, ++slot) { + acc[slot] += fmaxf(c_sh[i], 0.0f) * head_weight; + } + __syncthreads(); + } + + int slot = 0; + for (int i = tid; i < 8 * 16 * 16; i += 256, ++slot) { + const int wtile = i >> 8; + const int local = i & 255; + const int row = local >> 4; + const int col = local & 15; + const int token = tile_t + row; + const int comp = tile_c + wtile * 16 + col; + if (token < n_tokens && comp < n_comp) { + const int visible = (kv_start + token + 1) / ratio; + scores[(size_t) token * n_comp + comp] = + comp < visible ? acc[slot] : -1.0e30f; + } + } +} +#endif + +static __global__ void ds4_indexer_score_scalar_kernel( + float * scores, + const float * q, + const float * weights, + const half * index_comp, + int n_comp, + int n_tokens, + int kv_start, + int n_head, + int ratio) { + const int comp = (int) blockIdx.x; + const int token = (int) blockIdx.y; + const int tid = (int) threadIdx.x; + if (comp >= n_comp || token >= n_tokens) return; + const int visible = (kv_start + token + 1) / ratio; + if (comp >= visible) { + if (tid == 0) { + scores[(size_t) token * n_comp + comp] = -1.0e30f; + } + return; + } + + __shared__ float partial[256]; + float total = 0.0f; + const half * k = index_comp + (size_t) comp * 128; + for (int h = 0; h < n_head; ++h) { + const float * qh = q + ((size_t) token * n_head + h) * 128; + float dot = 0.0f; + for (int d = tid; d < 128; d += 256) { + dot += qh[d] * __half2float(k[d]); + } + partial[tid] = dot; + __syncthreads(); + for (int stride = 128; stride > 0; stride >>= 1) { + if (tid < stride) partial[tid] += partial[tid + stride]; + __syncthreads(); + } + if (tid == 0) { + total += fmaxf(partial[0], 0.0f) * + weights[(size_t) token * n_head + h]; + } + __syncthreads(); + } + if (tid == 0) scores[(size_t) token * n_comp + comp] = total; +} + +void ggml_cuda_op_ds4_indexer_score( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * weights = dst->src[1]; + const ggml_tensor * comp = dst->src[2]; + GGML_ASSERT(q && weights && comp); + GGML_ASSERT(q->type == GGML_TYPE_F32 && q->ne[0] == 128); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(comp->type == GGML_TYPE_F16 && comp->ne[0] == 128); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(q)); + GGML_ASSERT(ggml_is_contiguous(weights)); + GGML_ASSERT(ggml_is_contiguous(comp)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int n_head = (int) q->ne[1]; + const int n_tokens = (int) q->ne[2]; + const int n_comp = (int) comp->ne[1]; + const int kv_start = ggml_get_op_params_i32(dst, 0); + const int ratio = ggml_get_op_params_i32(dst, 1); + GGML_ASSERT(weights->ne[0] == n_head && weights->ne[1] == n_tokens); + GGML_ASSERT(dst->ne[0] == n_comp && dst->ne[1] == n_tokens); + + cudaStream_t stream = ctx.stream(); + const int warp_size = + ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; +#if DS4_INDEXER_WMMA_AVAILABLE + if (warp_size == 32) { + const dim3 grid((unsigned) ((n_comp + 127) / 128), + (unsigned) ((n_tokens + 15) / 16), 1); + ds4_indexer_score_wmma_kernel<<>>( + static_cast(dst->data), + static_cast(q->data), + static_cast(weights->data), + static_cast(comp->data), + n_comp, n_tokens, kv_start, n_head, ratio); + } else +#endif + { + const dim3 grid((unsigned) n_comp, (unsigned) n_tokens, 1); + ds4_indexer_score_scalar_kernel<<>>( + static_cast(dst->data), + static_cast(q->data), + static_cast(weights->data), + static_cast(comp->data), + n_comp, n_tokens, kv_start, n_head, ratio); + } + CUDA_CHECK(cudaGetLastError()); +} + +static __global__ void ds4_indexer_mask_init_kernel( + float * dst, + const float * base, + int64_t n_elements, + int n_attn, + int raw_rows) { + const int64_t index = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (index >= n_elements) return; + const int row = (int) (index % n_attn); + dst[index] = row < raw_rows ? base[index] : -1.0e30f; +} + +static __global__ void ds4_indexer_mask_scatter_kernel( + float * dst, + const float * base, + const int32_t * selected, + int64_t n_selected, + int n_attn, + int raw_rows, + int n_comp, + int top_k) { + const int64_t index = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (index >= n_selected) return; + const int token = (int) (index / top_k); + const int comp = selected[index]; + if (comp < 0 || comp >= n_comp) return; + const int64_t output_index = + (int64_t) token * n_attn + raw_rows + comp; + // Preserve the base causal mask: when a row is not yet visible, top-k may + // still return it only as an -inf filler for tokens with < k live rows. + dst[output_index] = base[output_index]; +} + +void ggml_cuda_op_ds4_indexer_mask( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst) { + const ggml_tensor * base = dst->src[0]; + const ggml_tensor * selected = dst->src[1]; + GGML_ASSERT(base && selected); + GGML_ASSERT(base->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + GGML_ASSERT(selected->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_are_same_shape(base, dst)); + GGML_ASSERT(ggml_is_contiguous(base) && ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_is_contiguous(selected)); + + const int n_attn = (int) base->ne[0]; + const int top_k = (int) selected->ne[0]; + const int n_tokens = (int) ggml_nrows(base); + const int raw_rows = ggml_get_op_params_i32(dst, 0); + const int n_comp = n_attn - raw_rows; + GGML_ASSERT(raw_rows >= 0 && n_comp >= 0); + GGML_ASSERT(ggml_nrows(selected) == n_tokens); + + const int64_t n_elements = ggml_nelements(base); + const int64_t n_selected = ggml_nelements(selected); + cudaStream_t stream = ctx.stream(); + ds4_indexer_mask_init_kernel<<< + (unsigned) ((n_elements + 255) / 256), 256, 0, stream>>>( + static_cast(dst->data), + static_cast(base->data), + n_elements, n_attn, raw_rows); + ds4_indexer_mask_scatter_kernel<<< + (unsigned) ((n_selected + 255) / 256), 256, 0, stream>>>( + static_cast(dst->data), + static_cast(base->data), + static_cast(selected->data), + n_selected, n_attn, raw_rows, n_comp, top_k); + CUDA_CHECK(cudaGetLastError()); +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cuh new file mode 100644 index 000000000..e4f29ea05 --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cuh @@ -0,0 +1,15 @@ +#pragma once + +#include "common.cuh" + +void ggml_cuda_op_ds4_indexer_qat( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst); + +void ggml_cuda_op_ds4_indexer_score( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst); + +void ggml_cuda_op_ds4_indexer_mask( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index 0eff4f1b7..153739b90 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -7,6 +7,1854 @@ #include "fattn-chunked.cuh" #include "fattn.cuh" +#if defined(GGML_USE_HIP) + +__device__ static float ds4_fa_block_sum(float v) { + __shared__ float smem[256]; + const int tid = threadIdx.x; + smem[tid] = v; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) smem[tid] += smem[tid + stride]; + __syncthreads(); + } + return smem[0]; +} + +__device__ static float ds4_fa_block_max(float v) { + __shared__ float smem[256]; + const int tid = threadIdx.x; + smem[tid] = v; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) smem[tid] = fmaxf(smem[tid], smem[tid + stride]); + __syncthreads(); + } + return smem[0]; +} + +template +__device__ static __forceinline__ float ds4_fa_load(const KV * ptr) { + return (float) *ptr; +} + +template <> +__device__ __forceinline__ float ds4_fa_load(const half * ptr) { + return __half2float(*ptr); +} + +template +__device__ static __forceinline__ void ds4_fa_load_pair( + const KV * ptr, float & v0, float & v1) { + v0 = (float) ptr[0]; + v1 = (float) ptr[1]; +} + +template <> +__device__ __forceinline__ void ds4_fa_load_pair( + const float * ptr, float & v0, float & v1) { + const float2 pair = *reinterpret_cast(ptr); + v0 = pair.x; + v1 = pair.y; +} + +template <> +__device__ __forceinline__ void ds4_fa_load_pair( + const half * ptr, float & v0, float & v1) { + const half2 pair = *reinterpret_cast(ptr); + const float2 unpacked = __half22float2(pair); + v0 = unpacked.x; + v1 = unpacked.y; +} + +template +__device__ static __forceinline__ void ds4_fa_load_quad( + const KV * ptr, float & v0, float & v1, float & v2, float & v3) { + v0 = (float) ptr[0]; + v1 = (float) ptr[1]; + v2 = (float) ptr[2]; + v3 = (float) ptr[3]; +} + +template <> +__device__ __forceinline__ void ds4_fa_load_quad( + const float * ptr, float & v0, float & v1, float & v2, float & v3) { + const float4 values = *reinterpret_cast(ptr); + v0 = values.x; + v1 = values.y; + v2 = values.z; + v3 = values.w; +} + +template <> +__device__ __forceinline__ void ds4_fa_load_quad( + const half * ptr, float & v0, float & v1, float & v2, float & v3) { + const float2 lo = __half22float2( + *reinterpret_cast(ptr + 0)); + const float2 hi = __half22float2( + *reinterpret_cast(ptr + 2)); + v0 = lo.x; + v1 = lo.y; + v2 = hi.x; + v3 = hi.y; +} + +struct ds4_inverse_rope_params { + int enabled; + int forward_q_enabled; + int kv_start; + float freq_scale; + float ext_factor; + float attn_factor; + float corr_low; + float corr_high; + float theta_scale; +}; + +// Keep these expressions aligned with rope.cu. The attention result is first +// stored in shared F32, matching the standalone attention-output store/load +// boundary, before the pair is rotated. +__device__ static __forceinline__ float ds4_rope_theta_fp64( + int32_t p, float theta_scale, int exp_int) { + const double tau = 6.2831853071795864769; + double angle = exp_int == 0 + ? (double) p + : (double) p * pow((double) theta_scale, (double) exp_int); + angle -= tau * floor(angle * (1.0 / tau)); + return (float) angle; +} + +__device__ static __forceinline__ void ds4_inverse_rope_coefficients( + int pair, int token, + const ds4_inverse_rope_params & p, + float & cos_theta, float & sin_theta) { + const int i0 = 2 * pair; + const float theta_extrap = ds4_rope_theta_fp64( + -(p.kv_start + token), p.theta_scale, pair); + const float theta_interp = p.freq_scale * theta_extrap; + float theta = theta_interp; + float mscale = p.attn_factor; + if (p.ext_factor != 0.0f) { + const float ramp_y = (i0 / 2 - p.corr_low) / + max(0.001f, p.corr_high - p.corr_low); + const float ramp_mix = + (1.0f - min(1.0f, max(0.0f, ramp_y))) * p.ext_factor; + theta = theta_interp * (1.0f - ramp_mix) + + theta_extrap * ramp_mix; + mscale *= 1.0f + 0.1f * logf(1.0f / p.freq_scale); + } + cos_theta = cosf(theta) * mscale; + sin_theta = sinf(theta) * mscale; +} + +// Forward counterpart of ds4_inverse_rope_coefficients. Keep the expressions +// aligned with rope_norm in rope.cu; unlike the inverse path, position is +// positive. Compressed-layer YaRN interpolation means the inverse coefficients +// cannot safely be recovered by merely negating sin(theta). +__device__ static __forceinline__ void ds4_forward_rope_coefficients( + int pair, int token, + const ds4_inverse_rope_params & p, + float & cos_theta, float & sin_theta) { + const int i0 = 2 * pair; + const float theta_extrap = ds4_rope_theta_fp64( + p.kv_start + token, p.theta_scale, pair); + const float theta_interp = p.freq_scale * theta_extrap; + float theta = theta_interp; + float mscale = p.attn_factor; + if (p.ext_factor != 0.0f) { + const float ramp_y = (i0 / 2 - p.corr_low) / + max(0.001f, p.corr_high - p.corr_low); + const float ramp_mix = + (1.0f - min(1.0f, max(0.0f, ramp_y))) * p.ext_factor; + theta = theta_interp * (1.0f - ramp_mix) + + theta_extrap * ramp_mix; + mscale *= 1.0f + 0.1f * logf(1.0f / p.freq_scale); + } + cos_theta = cosf(theta) * mscale; + sin_theta = sinf(theta) * mscale; +} + +__device__ static __forceinline__ void ds4_apply_inverse_rope_pair( + float x0, float x1, float cos_theta, float sin_theta, + float & y0, float & y1) { + y0 = x0 * cos_theta - x1 * sin_theta; + y1 = x0 * sin_theta + x1 * cos_theta; +} + +// RoPE coefficients depend on token position and pair, not on the query head. +// Materialize them once per attention call instead of recomputing FP64 pow, +// floor, cos and sin in every head block. F32 storage preserves the same +// explicit coefficient rounding used by the original in-kernel calculation. +__global__ static void ds4_inverse_rope_coefficients_kernel( + float * coefficients, + int n_tokens, + ds4_inverse_rope_params inverse_rope) { + const int index = (int) blockIdx.x * (int) blockDim.x + + (int) threadIdx.x; + const int count = n_tokens * 32; + if (index >= count) return; + const int token = index / 32; + const int pair = index % 32; + float cos_theta; + float sin_theta; + ds4_inverse_rope_coefficients( + pair, token, inverse_rope, cos_theta, sin_theta); + coefficients[2 * index + 0] = cos_theta; + coefficients[2 * index + 1] = sin_theta; +} + +__global__ static void ds4_forward_rope_coefficients_kernel( + float * coefficients, + int n_tokens, + ds4_inverse_rope_params rope) { + const int index = (int) blockIdx.x * (int) blockDim.x + + (int) threadIdx.x; + const int count = n_tokens * 32; + if (index >= count) return; + const int token = index / 32; + const int pair = index % 32; + float cos_theta; + float sin_theta; + ds4_forward_rope_coefficients( + pair, token, rope, cos_theta, sin_theta); + coefficients[2 * index + 0] = cos_theta; + coefficients[2 * index + 1] = sin_theta; +} + +// One mean latent-key vector per compressed-cache block. Raw SWA/current rows +// deliberately stay outside this summary and are always evaluated exactly. +template +__global__ static void ds4_fa_mean_comp_blocks_kernel( + const KV * k, + float * mean_k, + int n_kv, + int raw_rows, + int block_size, + int n_blocks) { + constexpr int D = 512; + const int b = (int) blockIdx.x; + if (b >= n_blocks) return; + const int begin = raw_rows + b * block_size; + const int end = min(n_kv, begin + block_size); + const float inv = 1.0f / (float) max(1, end - begin); + for (int d = (int) threadIdx.x; d < D; d += (int) blockDim.x) { + float sum = 0.0f; + for (int r = begin; r < end; ++r) { + sum += ds4_fa_load(k + (size_t) r * D + d); + } + mean_k[(size_t) b * D + d] = sum * inv; + } +} + +// Find the visible envelope in the raw and non-raw regions once per query +// token. Attention blocks for all query-head groups reuse these four bounds. +// Internal masked rows remain inside the envelope and are still evaluated as +// masked, so this changes storage only, not attention semantics. +template +__global__ static void ds4_fa_visibility_bounds_kernel( + const Mask * mask, + int * bounds, + int n_tokens, + int n_kv, + int raw_rows) { + const int t = (int) blockIdx.x; + const int lane = (int) threadIdx.x; + if (t >= n_tokens || lane >= warpSize) return; + + int raw_first = raw_rows; + int raw_last = -1; + int comp_first = n_kv; + int comp_last = -1; + const Mask * token_mask = mask + (size_t) t * n_kv; + + for (int base = 0; base < raw_rows; base += warpSize) { + const int r = base + lane; + const unsigned long long active = __ballot( + r < raw_rows && + ds4_fa_load(token_mask + r) > -1.0e20f); + if (lane == 0 && active != 0) { + if (raw_first == raw_rows) { + raw_first = base + __ffsll(active) - 1; + } + raw_last = base + 63 - __clzll(active); + } + } + for (int base = raw_rows; base < n_kv; base += warpSize) { + const int r = base + lane; + const unsigned long long active = __ballot( + r < n_kv && + ds4_fa_load(token_mask + r) > -1.0e20f); + if (lane == 0 && active != 0) { + if (comp_first == n_kv) { + comp_first = base + __ffsll(active) - 1; + } + comp_last = base + 63 - __clzll(active); + } + } + if (lane == 0) { + int * token_bounds = bounds + (size_t) t * 4; + token_bounds[0] = raw_first; + token_bounds[1] = raw_last; + token_bounds[2] = comp_first; + token_bounds[3] = comp_last; + } +} + +// Convert an externally selected compressed-row mask into exact lookup tables. +// selected_rows preserves ascending physical-row order for the value pass. +// owner_offsets/owner_ranks group those ascending ranks by the thread that +// owned the physical row in the original r = tid + 256*k traversal. The hot +// score and softmax passes can therefore visit only selected rows while +// retaining every thread's original accumulation order and reduction leaf. +template +__global__ static void ds4_fa_indexed_rows_kernel( + const Mask * mask, + int * selected_rows, + int * selected_counts, + int * owner_offsets, + int * owner_ranks, + int n_tokens, + int n_kv, + int raw_rows, + int capacity) { + const int t = (int) blockIdx.x; + const int tid = (int) threadIdx.x; + if (t >= n_tokens) return; + + constexpr int N_THREADS = 256; + __shared__ int owner_write[N_THREADS]; + const int n_comp_rows = n_kv - raw_rows; + int * token_rows = selected_rows + (size_t) t * capacity; + int * token_owner_offsets = owner_offsets + (size_t) t * (N_THREADS + 1); + int * token_owner_ranks = owner_ranks + (size_t) t * capacity; + token_owner_offsets[tid] = 0; + if (tid == 0) token_owner_offsets[N_THREADS] = 0; + __syncthreads(); + + if (tid == 0) { + const Mask * token_mask = mask + (size_t) t * n_kv; + int count = 0; + for (int c = 0; c < n_comp_rows; ++c) { + if (ds4_fa_load(token_mask + raw_rows + c) <= -1.0e20f) { + continue; + } + if (count < capacity) { + const int r = raw_rows + c; + token_rows[count] = r; + ++token_owner_offsets[(r & (N_THREADS - 1)) + 1]; + } + ++count; + } + count = min(count, capacity); + selected_counts[t] = count; + + int prefix = 0; + for (int owner = 0; owner < N_THREADS; ++owner) { + const int owner_count = token_owner_offsets[owner + 1]; + token_owner_offsets[owner] = prefix; + prefix += owner_count; + } + token_owner_offsets[N_THREADS] = prefix; + } + __syncthreads(); + + owner_write[tid] = token_owner_offsets[tid]; + __syncthreads(); + + if (tid == 0) { + const int count = selected_counts[t]; + for (int rank = 0; rank < count; ++rank) { + const int owner = token_rows[rank] & (N_THREADS - 1); + token_owner_ranks[owner_write[owner]++] = rank; + } + } +} + +template +__global__ static void ds4_flash_attn_d512_shared_kv_kernel( + float * dst, + const float * q, + size_t q_stride_token, + size_t q_stride_head, + const KV * k, + const KV * v, + const Mask * mask, + const float * sinks, + const float * mean_k, + int n_tokens, + int n_heads, + int n_kv, + float scale, + int raw_rows, + int sparse_keep_rows, + int sparse_block_size, + int n_comp_blocks, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients, + const float * forward_rope_coefficients) { + constexpr int D = 512; + const int t = (int) blockIdx.x; + const int h = (int) blockIdx.y; + const int tid = (int) threadIdx.x; + if (t >= n_tokens || h >= n_heads) return; + + extern __shared__ float scratch[]; + float * scores = scratch; + float * block_scores = scores + n_kv; + float * block_keep = block_scores + n_comp_blocks; + float * q_rope_tail = block_keep + n_comp_blocks; + const float * qh = q + (size_t) t * q_stride_token + + (size_t) h * q_stride_head; + + if (inverse_rope.forward_q_enabled) { + for (int pair = tid; pair < 32; pair += (int) blockDim.x) { + const float x0 = qh[D - 64 + 2 * pair + 0]; + const float x1 = qh[D - 64 + 2 * pair + 1]; + const size_t coefficient_index = + ((size_t) t * 32 + (size_t) pair) * 2; + const float cos_theta = forward_rope_coefficients[ + coefficient_index + 0]; + const float sin_theta = forward_rope_coefficients[ + coefficient_index + 1]; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, + q_rope_tail[2 * pair + 0], q_rope_tail[2 * pair + 1]); + } + __syncthreads(); + } + + const int n_comp_rows = n_kv - raw_rows; + const bool sparse = mean_k && n_comp_blocks > 0 && + sparse_keep_rows > 0 && + sparse_keep_rows < n_comp_rows; + if (sparse) { + for (int b = tid; b < n_comp_blocks; b += (int) blockDim.x) { + const int first_row = raw_rows + b * sparse_block_size; + const float mask_v = mask + ? ds4_fa_load( + mask + (size_t) t * n_kv + first_row) + : 0.0f; + const float * kb = mean_k + (size_t) b * D; + float dot = -3.402823466e38f; + if (mask_v > -1.0e20f) { + dot = 0.0f; +#pragma unroll + for (int d = 0; d < D; ++d) { + const float qv = inverse_rope.forward_q_enabled && d >= D - 64 + ? q_rope_tail[d - (D - 64)] : qh[d]; + dot += qv * kb[d]; + } + dot *= scale; + } + block_scores[b] = dot; + } + __syncthreads(); + + const int keep_blocks = min(n_comp_blocks, + (sparse_keep_rows + sparse_block_size - 1) / sparse_block_size); + for (int b = tid; b < n_comp_blocks; b += (int) blockDim.x) { + const float score = block_scores[b]; + int rank = 0; + for (int j = 0; j < n_comp_blocks; ++j) { + const float other = block_scores[j]; + rank += (other > score || (other == score && j < b)) ? 1 : 0; + } + block_keep[b] = rank < keep_blocks ? 1.0f : 0.0f; + } + __syncthreads(); + } + + // Keep the original row-to-thread mapping and reduction order. The mask + // remains the sole visibility authority for dense attention; duplicating + // its DS4 policy here risks changing semantics. Masked rows already avoid + // the expensive D=512 dot product. Sparse mode additionally drops selected + // compressed blocks by design. + float local_max = sinks ? sinks[h] : -3.402823466e38f; + for (int r = tid; r < n_kv; r += blockDim.x) { + bool keep = true; + if (sparse && r >= raw_rows) { + const int b = (r - raw_rows) / sparse_block_size; + keep = b < n_comp_blocks && block_keep[b] != 0.0f; + } + const float mask_v = mask + ? ds4_fa_load(mask + (size_t) t * n_kv + r) + : 0.0f; + float s = -3.402823466e38f; + if (keep && mask_v > -1.0e20f) { + const KV * kr = k + (size_t) r * D; + float dot = 0.0f; +#pragma unroll + for (int d = 0; d < D; ++d) { + const float qv = inverse_rope.forward_q_enabled && d >= D - 64 + ? q_rope_tail[d - (D - 64)] : qh[d]; + dot += qv * ds4_fa_load(kr + d); + } + s = dot * scale + mask_v; + } + scores[r] = s; + local_max = fmaxf(local_max, s); + } + const float max_score = ds4_fa_block_max(local_max); + + float local_sum = 0.0f; + for (int r = tid; r < n_kv; r += blockDim.x) { + const float w = expf(scores[r] - max_score); + scores[r] = w; + local_sum += w; + } + if (tid == 0 && sinks) { + local_sum += expf(sinks[h] - max_score); + } + const float denom = ds4_fa_block_sum(local_sum); + const float inv_denom = 1.0f / denom; + + // One wave locates the non-zero envelope independently in the raw and + // compressed spans. Ballots inspect a whole wave of scores at once and + // require only one block barrier, unlike a four-reduction implementation. + // Values inside each envelope retain their original order, including any + // internal zero, so floating-point accumulation is unchanged. + __shared__ int value_bounds[4]; + if (tid < warpSize) { + const int lane = tid; + if (lane == 0) { + value_bounds[0] = raw_rows; + value_bounds[1] = -1; + value_bounds[2] = n_kv; + value_bounds[3] = -1; + } + for (int base = 0; base < raw_rows; base += warpSize) { + const int r = base + lane; + const unsigned long long active = __ballot( + r < raw_rows && scores[r] != 0.0f); + if (lane == 0 && active != 0) { + if (value_bounds[0] == raw_rows) { + value_bounds[0] = base + __ffsll(active) - 1; + } + value_bounds[1] = base + 63 - __clzll(active); + } + } + for (int base = raw_rows; base < n_kv; base += warpSize) { + const int r = base + lane; + const unsigned long long active = __ballot( + r < n_kv && scores[r] != 0.0f); + if (lane == 0 && active != 0) { + if (value_bounds[2] == n_kv) { + value_bounds[2] = base + __ffsll(active) - 1; + } + value_bounds[3] = base + 63 - __clzll(active); + } + } + } + __syncthreads(); + const int raw_first = value_bounds[0]; + const int raw_last = value_bounds[1]; + const int comp_first = value_bounds[2]; + const int comp_last = value_bounds[3]; + + // Scoring is complete here, so the forward-RoPE tail scratch can be + // reused for inverse-RoPE output. Do not alias scores: different waves + // accumulate value dimensions concurrently, and tail writers would race + // readers of scores[0..63]. + float * rope_tail = q_rope_tail; + for (int d = tid; d < D; d += blockDim.x) { + float acc = 0.0f; + for (int r = raw_first; r <= raw_last; ++r) { + acc += scores[r] * ds4_fa_load( + v + (size_t) r * D + d); + } + for (int r = comp_first; r <= comp_last; ++r) { + acc += scores[r] * ds4_fa_load( + v + (size_t) r * D + d); + } + const float value = acc * inv_denom; + if (inverse_rope.enabled && d >= D - 64) { + rope_tail[d - (D - 64)] = value; + } else { + dst[((size_t) t * (size_t) n_heads + (size_t) h) * D + d] = value; + } + } + if (inverse_rope.enabled) { + __syncthreads(); + if (tid < 32) { + const float x0 = rope_tail[2 * tid + 0]; + const float x1 = rope_tail[2 * tid + 1]; + const size_t coefficient_index = + ((size_t) t * 32 + (size_t) tid) * 2; + const float cos_theta = inverse_rope_coefficients[ + coefficient_index + 0]; + const float sin_theta = inverse_rope_coefficients[ + coefficient_index + 1]; + float y0; + float y1; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, y0, y1); + float * out = dst + + ((size_t) t * (size_t) n_heads + (size_t) h) * D + D - 64; + out[2 * tid + 0] = y0; + out[2 * tid + 1] = y1; + } + } +} + +// DS4 MLA uses one latent K/V head for every query head. Grouping query heads +// in one block lets them share each K/V load while retaining the reference +// kernel's row-to-thread mapping, per-head reduction tree, and accumulation +// order. The grouped path is dense-only; experimental sparse selection keeps +// using the single-head kernel above. +template +__global__ static void ds4_flash_attn_d512_shared_kv_grouped_kernel( + float * dst, + const float * q, + size_t q_stride_token, + size_t q_stride_head, + const KV * k, + const KV * v, + const Mask * mask, + const float * sinks, + int n_tokens, + int n_heads, + int n_kv, + float scale, + int raw_rows, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients, + const float * forward_rope_coefficients) { + constexpr int D = 512; + constexpr int N_THREADS = 256; + const int t = (int) blockIdx.x; + const int h_begin = (int) blockIdx.y * HEADS_PER_BLOCK; + const int tid = (int) threadIdx.x; + if (t >= n_tokens || h_begin >= n_heads) return; + + extern __shared__ float scratch[]; + float * scores = scratch; + float * reduction = scores + (size_t) HEADS_PER_BLOCK * n_kv; + int * value_bounds = reinterpret_cast( + reduction + (size_t) HEADS_PER_BLOCK * N_THREADS); + float * q_rope_tail = reinterpret_cast( + value_bounds + (size_t) HEADS_PER_BLOCK * 4); + + const float * qh[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + qh[j] = q + (size_t) t * q_stride_token + + (size_t) h * q_stride_head; + } + + if (inverse_rope.forward_q_enabled) { + for (int index = tid; index < HEADS_PER_BLOCK * 32; + index += (int) blockDim.x) { + const int j = index / 32; + const int pair = index % 32; + const float x0 = qh[j][D - 64 + 2 * pair + 0]; + const float x1 = qh[j][D - 64 + 2 * pair + 1]; + const size_t coefficient_index = + ((size_t) t * 32 + (size_t) pair) * 2; + const float cos_theta = forward_rope_coefficients[ + coefficient_index + 0]; + const float sin_theta = forward_rope_coefficients[ + coefficient_index + 1]; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, + q_rope_tail[(size_t) j * 64 + 2 * pair + 0], + q_rope_tail[(size_t) j * 64 + 2 * pair + 1]); + } + __syncthreads(); + } + + float local_max[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + local_max[j] = h < n_heads && sinks + ? sinks[h] : -3.402823466e38f; + } + // A thread owns exactly the same rows as in the single-head kernel. Four + // independent dot-product chains consume one shared K value per feature. + for (int r = tid; r < n_kv; r += blockDim.x) { + const float mask_v = mask + ? ds4_fa_load(mask + (size_t) t * n_kv + r) + : 0.0f; + const bool visible = mask_v > -1.0e20f; + float dot[HEADS_PER_BLOCK] = {}; + if (visible) { + const KV * kr = k + (size_t) r * D; +#pragma unroll + for (int d = 0; d < D; ++d) { + const float kv = ds4_fa_load(kr + d); +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const float qv = + inverse_rope.forward_q_enabled && d >= D - 64 + ? q_rope_tail[(size_t) j * 64 + d - (D - 64)] + : qh[j][d]; + dot[j] += qv * kv; + } + } + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + const float s = h < n_heads && visible + ? dot[j] * scale + mask_v : -3.402823466e38f; + scores[(size_t) j * n_kv + r] = s; + local_max[j] = fmaxf(local_max[j], s); + } + } + + // Match ds4_fa_block_max independently for every grouped head. +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + reduction[(size_t) j * N_THREADS + tid] = local_max[j]; + } + __syncthreads(); + for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { + if (tid < stride) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * row = reduction + (size_t) j * N_THREADS; + row[tid] = fmaxf(row[tid], row[tid + stride]); + } + } + __syncthreads(); + } + + float max_score[HEADS_PER_BLOCK]; + float local_sum[HEADS_PER_BLOCK] = {}; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + max_score[j] = reduction[(size_t) j * N_THREADS]; + } + for (int r = tid; r < n_kv; r += blockDim.x) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * score = scores + (size_t) j * n_kv + r; + const float weight = expf(*score - max_score[j]); + *score = weight; + local_sum[j] += weight; + } + } + if (tid == 0 && sinks) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + local_sum[j] += expf(sinks[h_begin + j] - max_score[j]); + } + } + + // Match ds4_fa_block_sum independently for every grouped head. +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + reduction[(size_t) j * N_THREADS + tid] = local_sum[j]; + } + __syncthreads(); + for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { + if (tid < stride) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * row = reduction + (size_t) j * N_THREADS; + row[tid] += row[tid + stride]; + } + } + __syncthreads(); + } + + float inv_denom[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + inv_denom[j] = 1.0f / reduction[(size_t) j * N_THREADS]; + } + + // Underflow can make the non-zero envelope differ by head, so retain one + // pair of raw/compressed bounds per head. Ballot order does not affect any + // arithmetic result. + if (tid < warpSize) { + const int lane = tid; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + int * bounds = value_bounds + 4 * j; + if (lane == 0) { + bounds[0] = raw_rows; + bounds[1] = -1; + bounds[2] = n_kv; + bounds[3] = -1; + } + for (int base = 0; base < raw_rows; base += warpSize) { + const int r = base + lane; + const unsigned long long active = __ballot( + r < raw_rows && + scores[(size_t) j * n_kv + r] != 0.0f); + if (lane == 0 && active != 0) { + if (bounds[0] == raw_rows) { + bounds[0] = base + __ffsll(active) - 1; + } + bounds[1] = base + 63 - __clzll(active); + } + } + for (int base = raw_rows; base < n_kv; base += warpSize) { + const int r = base + lane; + const unsigned long long active = __ballot( + r < n_kv && scores[(size_t) j * n_kv + r] != 0.0f); + if (lane == 0 && active != 0) { + if (bounds[2] == n_kv) { + bounds[2] = base + __ffsll(active) - 1; + } + bounds[3] = base + 63 - __clzll(active); + } + } + } + } + __syncthreads(); + + int raw_first = raw_rows; + int raw_last = -1; + int comp_first = n_kv; + int comp_last = -1; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int * bounds = value_bounds + 4 * j; + raw_first = min(raw_first, bounds[0]); + raw_last = max(raw_last, bounds[1]); + comp_first = min(comp_first, bounds[2]); + comp_last = max(comp_last, bounds[3]); + } + + float * rope_tail = reduction; + for (int d = tid; d < D; d += blockDim.x) { + float acc[HEADS_PER_BLOCK] = {}; + for (int r = raw_first; r <= raw_last; ++r) { + const float vv = ds4_fa_load(v + (size_t) r * D + d); +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int * bounds = value_bounds + 4 * j; + if (r >= bounds[0] && r <= bounds[1]) { + acc[j] += scores[(size_t) j * n_kv + r] * vv; + } + } + } + for (int r = comp_first; r <= comp_last; ++r) { + const float vv = ds4_fa_load(v + (size_t) r * D + d); +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int * bounds = value_bounds + 4 * j; + if (r >= bounds[2] && r <= bounds[3]) { + acc[j] += scores[(size_t) j * n_kv + r] * vv; + } + } + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + if (h < n_heads) { + const float value = acc[j] * inv_denom[j]; + if (inverse_rope.enabled && d >= D - 64) { + rope_tail[(size_t) j * 64 + d - (D - 64)] = value; + } else { + dst[((size_t) t * (size_t) n_heads + (size_t) h) * D + d] = + value; + } + } + } + } + if (inverse_rope.enabled) { + __syncthreads(); + if (tid < HEADS_PER_BLOCK * 32) { + const int j = tid / 32; + const int pair = tid % 32; + const float x0 = rope_tail[(size_t) j * 64 + 2 * pair + 0]; + const float x1 = rope_tail[(size_t) j * 64 + 2 * pair + 1]; + const size_t coefficient_index = + ((size_t) t * 32 + (size_t) pair) * 2; + const float cos_theta = inverse_rope_coefficients[ + coefficient_index + 0]; + const float sin_theta = inverse_rope_coefficients[ + coefficient_index + 1]; + float y0; + float y1; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, y0, y1); + const int h = h_begin + j; + float * out = dst + + ((size_t) t * (size_t) n_heads + (size_t) h) * D + D - 64; + out[2 * pair + 0] = y0; + out[2 * pair + 1] = y1; + } + } +} + +// Dense-prefill variant of the grouped kernel with compact score storage. +// The mask-derived envelopes only change the address used to retain a score; +// every visible row keeps its original owner thread, dot-product order, +// reduction tree, softmax order, and value-accumulation position. +template +__global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( + float * dst, + const float * q, + size_t q_stride_token, + size_t q_stride_head, + const KV * k, + const KV * v, + const Mask * mask, + const float * sinks, + int n_tokens, + int n_heads, + int n_kv, + float scale, + int raw_rows, + int raw_score_capacity, + int score_stride, + const int * visibility_bounds, + const int * indexed_rows, + const int * indexed_counts, + const int * indexed_owner_offsets, + const int * indexed_owner_ranks, + int indexed_capacity, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients, + const float * forward_rope_coefficients) { + constexpr int D = 512; + constexpr int N_THREADS = 256; + static_assert(VALUES_PER_THREAD == 2 || VALUES_PER_THREAD == 4); + const int t = (int) blockIdx.x; + const int h_begin = (int) blockIdx.y * HEADS_PER_BLOCK; + const int tid = (int) threadIdx.x; + if (t >= n_tokens || h_begin >= n_heads) return; + + extern __shared__ float scratch[]; + float * scores = scratch; + float * reduction = scores + (size_t) HEADS_PER_BLOCK * score_stride; + int * value_bounds = reinterpret_cast( + reduction + (size_t) HEADS_PER_BLOCK * N_THREADS); + float * q_rope_tail = reinterpret_cast( + value_bounds + (size_t) HEADS_PER_BLOCK * 4); + + const int * token_visibility = visibility_bounds + (size_t) t * 4; + const int mask_raw_first = token_visibility[0]; + const int mask_raw_last = token_visibility[1]; + const int mask_comp_first = token_visibility[2]; + const int mask_comp_last = token_visibility[3]; + const int * token_indexed_rows = nullptr; + const int * token_owner_offsets = nullptr; + const int * token_owner_ranks = nullptr; + int indexed_count = 0; + if constexpr (INDEXED_MASK) { + token_indexed_rows = indexed_rows + (size_t) t * indexed_capacity; + token_owner_offsets = indexed_owner_offsets + (size_t) t * (N_THREADS + 1); + token_owner_ranks = indexed_owner_ranks + (size_t) t * indexed_capacity; + indexed_count = indexed_counts[t]; + } + + const float * qh[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + qh[j] = q + (size_t) t * q_stride_token + + (size_t) h * q_stride_head; + } + + if (inverse_rope.forward_q_enabled) { + for (int index = tid; index < HEADS_PER_BLOCK * 32; + index += (int) blockDim.x) { + const int j = index / 32; + const int pair = index % 32; + const float x0 = qh[j][D - 64 + 2 * pair + 0]; + const float x1 = qh[j][D - 64 + 2 * pair + 1]; + const size_t coefficient_index = + ((size_t) t * 32 + (size_t) pair) * 2; + const float cos_theta = forward_rope_coefficients[ + coefficient_index + 0]; + const float sin_theta = forward_rope_coefficients[ + coefficient_index + 1]; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, + q_rope_tail[(size_t) j * 64 + 2 * pair + 0], + q_rope_tail[(size_t) j * 64 + 2 * pair + 1]); + } + __syncthreads(); + } + + float local_max[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + local_max[j] = h < n_heads && sinks + ? sinks[h] : -3.402823466e38f; + } + // Build the full kernel's exact per-head non-zero envelope while the + // softmax weights are emitted. This avoids rescanning every context row + // after softmax without changing the subsequent V accumulation interval. + if (tid < HEADS_PER_BLOCK * 4) { + const int slot = tid & 3; + value_bounds[tid] = slot == 0 + ? raw_rows + : slot == 1 + ? -1 + : slot == 2 + ? (INDEXED_MASK ? indexed_count : n_kv) + : -1; + } + + // Preserve each thread's original r = tid + 256*k order. In indexed mode, + // owner_ranks is the exact selected subsequence of that traversal, so the + // hot passes no longer scan every unselected compressed row. + const int raw_owner_first = mask_raw_first + + ((tid - (mask_raw_first & (N_THREADS - 1)) + N_THREADS) & + (N_THREADS - 1)); + const int raw_iteration_count = raw_owner_first <= mask_raw_last + ? 1 + (mask_raw_last - raw_owner_first) / N_THREADS : 0; + const int comp_owner_first = mask_comp_first + + ((tid - (mask_comp_first & (N_THREADS - 1)) + N_THREADS) & + (N_THREADS - 1)); + const int comp_iteration_count = comp_owner_first <= mask_comp_last + ? 1 + (mask_comp_last - comp_owner_first) / N_THREADS : 0; + int owner_begin = 0; + int owner_count = 0; + int score_iteration_count = raw_iteration_count + comp_iteration_count; + if constexpr (INDEXED_MASK) { + owner_begin = token_owner_offsets[tid]; + owner_count = token_owner_offsets[tid + 1] - owner_begin; + score_iteration_count = raw_iteration_count + owner_count; + } + for (int iteration = 0; iteration < score_iteration_count; ++iteration) { + int r; + int score_index; + if constexpr (INDEXED_MASK) { + if (iteration < raw_iteration_count) { + r = raw_owner_first + iteration * N_THREADS; + score_index = r - mask_raw_first; + } else { + const int rank = token_owner_ranks[ + owner_begin + iteration - raw_iteration_count]; + r = token_indexed_rows[rank]; + score_index = raw_score_capacity + rank; + } + } else { + if (iteration < raw_iteration_count) { + r = raw_owner_first + iteration * N_THREADS; + score_index = r - mask_raw_first; + } else { + r = comp_owner_first + + (iteration - raw_iteration_count) * N_THREADS; + score_index = raw_score_capacity + r - mask_comp_first; + } + } + + const float mask_v = ds4_fa_load( + mask + (size_t) t * n_kv + r); + const bool visible = mask_v > -1.0e20f; + float dot[HEADS_PER_BLOCK] = {}; + if (visible) { + const KV * kr = k + (size_t) r * D; +#pragma unroll + for (int d = 0; d < D; ++d) { + const float kv = ds4_fa_load(kr + d); +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const float qv = + inverse_rope.forward_q_enabled && d >= D - 64 + ? q_rope_tail[(size_t) j * 64 + d - (D - 64)] + : qh[j][d]; + dot[j] += qv * kv; + } + } + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + const float s = h < n_heads && visible + ? dot[j] * scale + mask_v : -3.402823466e38f; + scores[(size_t) j * score_stride + score_index] = s; + local_max[j] = fmaxf(local_max[j], s); + } + } + +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + reduction[(size_t) j * N_THREADS + tid] = local_max[j]; + } + __syncthreads(); + for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { + if (tid < stride) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * row = reduction + (size_t) j * N_THREADS; + row[tid] = fmaxf(row[tid], row[tid + stride]); + } + } + __syncthreads(); + } + + float max_score[HEADS_PER_BLOCK]; + float local_sum[HEADS_PER_BLOCK] = {}; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + max_score[j] = reduction[(size_t) j * N_THREADS]; + } + for (int iteration = 0; iteration < score_iteration_count; ++iteration) { + int score_index; + int bound_value; + const bool raw_value = iteration < raw_iteration_count; + if constexpr (INDEXED_MASK) { + if (raw_value) { + const int r = raw_owner_first + iteration * N_THREADS; + score_index = r - mask_raw_first; + bound_value = r; + } else { + const int rank = token_owner_ranks[ + owner_begin + iteration - raw_iteration_count]; + score_index = raw_score_capacity + rank; + bound_value = rank; + } + } else { + if (raw_value) { + const int r = raw_owner_first + iteration * N_THREADS; + score_index = r - mask_raw_first; + bound_value = r; + } else { + const int r = comp_owner_first + + (iteration - raw_iteration_count) * N_THREADS; + score_index = raw_score_capacity + r - mask_comp_first; + bound_value = r; + } + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * score = scores + (size_t) j * score_stride + score_index; + const float weight = expf(*score - max_score[j]); + *score = weight; + local_sum[j] += weight; + if (weight != 0.0f) { + int * bounds = value_bounds + 4 * j + (raw_value ? 0 : 2); + atomicMin(bounds + 0, bound_value); + atomicMax(bounds + 1, bound_value); + } + } + } + if (tid == 0 && sinks) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + local_sum[j] += expf(sinks[h_begin + j] - max_score[j]); + } + } + +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + reduction[(size_t) j * N_THREADS + tid] = local_sum[j]; + } + __syncthreads(); + for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { + if (tid < stride) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * row = reduction + (size_t) j * N_THREADS; + row[tid] += row[tid + stride]; + } + } + __syncthreads(); + } + + float inv_denom[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + inv_denom[j] = 1.0f / reduction[(size_t) j * N_THREADS]; + } + + // Finish the shared envelope before the value phase reads it. + __syncthreads(); + + int raw_first = raw_rows; + int raw_last = -1; + int comp_first = INDEXED_MASK ? indexed_count : n_kv; + int comp_last = -1; + int head_raw_first[HEADS_PER_BLOCK]; + int head_raw_last[HEADS_PER_BLOCK]; + int head_comp_first[HEADS_PER_BLOCK]; + int head_comp_last[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int * bounds = value_bounds + 4 * j; + head_raw_first[j] = bounds[0]; + head_raw_last[j] = bounds[1]; + head_comp_first[j] = bounds[2]; + head_comp_last[j] = bounds[3]; + raw_first = min(raw_first, head_raw_first[j]); + raw_last = max(raw_last, head_raw_last[j]); + comp_first = min(comp_first, head_comp_first[j]); + comp_last = max(comp_last, head_comp_last[j]); + } + + // One active thread owns adjacent value dimensions. This retains each + // dimension's ascending row accumulation order while sharing score loads, + // row-loop control, and a naturally aligned vector V load across the group. + float * rope_tail = reduction; + const int d0 = VALUES_PER_THREAD * tid; + const int d1 = d0 + 1; + const int d2 = d0 + 2; + const int d3 = d0 + 3; + float acc0[HEADS_PER_BLOCK] = {}; + float acc1[HEADS_PER_BLOCK] = {}; + float acc2[HEADS_PER_BLOCK] = {}; + float acc3[HEADS_PER_BLOCK] = {}; + if (d0 < D) { + for (int r = raw_first; r <= raw_last; ++r) { + const int score_index = r - mask_raw_first; + float vv0; + float vv1; + float vv2 = 0.0f; + float vv3 = 0.0f; + if constexpr (VALUES_PER_THREAD == 2) { + ds4_fa_load_pair( + v + (size_t) r * D + d0, vv0, vv1); + } else { + ds4_fa_load_quad( + v + (size_t) r * D + d0, vv0, vv1, vv2, vv3); + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + if (r >= head_raw_first[j] && r <= head_raw_last[j]) { + const float weight = + scores[(size_t) j * score_stride + score_index]; + acc0[j] += weight * vv0; + acc1[j] += weight * vv1; + if constexpr (VALUES_PER_THREAD == 4) { + acc2[j] += weight * vv2; + acc3[j] += weight * vv3; + } + } + } + } + if constexpr (INDEXED_MASK) { + for (int rank = comp_first; rank <= comp_last; ++rank) { + const int r = token_indexed_rows[rank]; + const int score_index = raw_score_capacity + rank; + bool any_nonzero = false; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + any_nonzero = any_nonzero || + scores[(size_t) j * score_stride + score_index] != 0.0f; + } + if (!any_nonzero) continue; + float vv0; + float vv1; + float vv2 = 0.0f; + float vv3 = 0.0f; + if constexpr (VALUES_PER_THREAD == 2) { + ds4_fa_load_pair( + v + (size_t) r * D + d0, vv0, vv1); + } else { + ds4_fa_load_quad( + v + (size_t) r * D + d0, vv0, vv1, vv2, vv3); + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + if (rank >= head_comp_first[j] && rank <= head_comp_last[j]) { + const float weight = + scores[(size_t) j * score_stride + score_index]; + acc0[j] += weight * vv0; + acc1[j] += weight * vv1; + if constexpr (VALUES_PER_THREAD == 4) { + acc2[j] += weight * vv2; + acc3[j] += weight * vv3; + } + } + } + } + } else { + for (int r = comp_first; r <= comp_last; ++r) { + const int score_index = + raw_score_capacity + r - mask_comp_first; + bool any_nonzero = false; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + any_nonzero = any_nonzero || + scores[(size_t) j * score_stride + score_index] != 0.0f; + } + if (!any_nonzero) continue; + float vv0; + float vv1; + float vv2 = 0.0f; + float vv3 = 0.0f; + if constexpr (VALUES_PER_THREAD == 2) { + ds4_fa_load_pair( + v + (size_t) r * D + d0, vv0, vv1); + } else { + ds4_fa_load_quad( + v + (size_t) r * D + d0, vv0, vv1, vv2, vv3); + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + if (r >= head_comp_first[j] && r <= head_comp_last[j]) { + const float weight = + scores[(size_t) j * score_stride + score_index]; + acc0[j] += weight * vv0; + acc1[j] += weight * vv1; + if constexpr (VALUES_PER_THREAD == 4) { + acc2[j] += weight * vv2; + acc3[j] += weight * vv3; + } + } + } + } + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + if (h < n_heads) { + const float value0 = acc0[j] * inv_denom[j]; + const float value1 = acc1[j] * inv_denom[j]; + if (inverse_rope.enabled && d0 >= D - 64) { + rope_tail[(size_t) j * 64 + d0 - (D - 64)] = value0; + rope_tail[(size_t) j * 64 + d1 - (D - 64)] = value1; + if constexpr (VALUES_PER_THREAD == 4) { + const float value2 = acc2[j] * inv_denom[j]; + const float value3 = acc3[j] * inv_denom[j]; + rope_tail[(size_t) j * 64 + d2 - (D - 64)] = value2; + rope_tail[(size_t) j * 64 + d3 - (D - 64)] = value3; + } + } else { + float * out = dst + + ((size_t) t * (size_t) n_heads + (size_t) h) * D + d0; + out[0] = value0; + out[1] = value1; + if constexpr (VALUES_PER_THREAD == 4) { + out[2] = acc2[j] * inv_denom[j]; + out[3] = acc3[j] * inv_denom[j]; + } + } + } + } + } + if (inverse_rope.enabled) { + __syncthreads(); + if (tid < HEADS_PER_BLOCK * 32) { + const int j = tid / 32; + const int pair = tid % 32; + const float x0 = rope_tail[(size_t) j * 64 + 2 * pair + 0]; + const float x1 = rope_tail[(size_t) j * 64 + 2 * pair + 1]; + const size_t coefficient_index = + ((size_t) t * 32 + (size_t) pair) * 2; + const float cos_theta = inverse_rope_coefficients[ + coefficient_index + 0]; + const float sin_theta = inverse_rope_coefficients[ + coefficient_index + 1]; + float y0; + float y1; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, y0, y1); + const int h = h_begin + j; + float * out = dst + + ((size_t) t * (size_t) n_heads + (size_t) h) * D + D - 64; + out[2 * pair + 0] = y0; + out[2 * pair + 1] = y1; + } + } +} + +template +static bool ds4_launch_flash_attn_d512_grouped( + ggml_tensor * dst, + const ggml_tensor * Q, + const ggml_tensor * K, + const ggml_tensor * V, + const ggml_tensor * mask, + const ggml_tensor * sinks, + bool kv_f16, + bool kv_f32, + int n_tokens, + int n_heads, + int n_kv, + float scale, + int raw_rows, + size_t q_stride_token, + size_t q_stride_head, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients, + const float * forward_rope_coefficients, + size_t shmem, + cudaStream_t stream) { + dim3 grid( + (unsigned) n_tokens, + (unsigned) (n_heads / HEADS_PER_BLOCK), 1); + if (kv_f16 && (!mask || mask->type == GGML_TYPE_F16)) { + ds4_flash_attn_d512_shared_kv_grouped_kernel< + half, half, HEADS_PER_BLOCK> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, (const half *) V->data, + mask ? (const half *) mask->data : nullptr, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, inverse_rope, + inverse_rope_coefficients, forward_rope_coefficients); + } else if (kv_f32 && (!mask || mask->type == GGML_TYPE_F32)) { + ds4_flash_attn_d512_shared_kv_grouped_kernel< + float, float, HEADS_PER_BLOCK> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + mask ? (const float *) mask->data : nullptr, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, inverse_rope, + inverse_rope_coefficients, forward_rope_coefficients); + } else if (kv_f32 && mask && mask->type == GGML_TYPE_F16) { + ds4_flash_attn_d512_shared_kv_grouped_kernel< + float, half, HEADS_PER_BLOCK> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, inverse_rope, + inverse_rope_coefficients, forward_rope_coefficients); + } else { + return false; + } + return true; +} + +template +static bool ds4_launch_flash_attn_d512_grouped_compact( + ggml_tensor * dst, + const ggml_tensor * Q, + const ggml_tensor * K, + const ggml_tensor * V, + const ggml_tensor * mask, + const ggml_tensor * sinks, + bool kv_f16, + bool kv_f32, + int n_tokens, + int n_heads, + int n_kv, + float scale, + int raw_rows, + int raw_score_capacity, + int score_stride, + const int * visibility_bounds, + const int * indexed_rows, + const int * indexed_counts, + const int * indexed_owner_offsets, + const int * indexed_owner_ranks, + int indexed_capacity, + size_t q_stride_token, + size_t q_stride_head, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients, + const float * forward_rope_coefficients, + size_t shmem, + cudaStream_t stream) { + GGML_ASSERT(mask && visibility_bounds); + if constexpr (INDEXED_MASK) { + GGML_ASSERT(indexed_rows && indexed_counts && + indexed_owner_offsets && indexed_owner_ranks); + } + dim3 grid( + (unsigned) n_tokens, + (unsigned) (n_heads / HEADS_PER_BLOCK), 1); + if (kv_f16 && mask->type == GGML_TYPE_F16) { + ds4_flash_attn_d512_shared_kv_grouped_compact_kernel< + half, half, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, (const half *) V->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, + raw_score_capacity, score_stride, visibility_bounds, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, indexed_capacity, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients); + } else if (kv_f32 && mask->type == GGML_TYPE_F32) { + ds4_flash_attn_d512_shared_kv_grouped_compact_kernel< + float, float, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + (const float *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, + raw_score_capacity, score_stride, visibility_bounds, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, indexed_capacity, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients); + } else if (kv_f32 && mask->type == GGML_TYPE_F16) { + ds4_flash_attn_d512_shared_kv_grouped_compact_kernel< + float, half, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, + raw_score_capacity, score_stride, visibility_bounds, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, indexed_capacity, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients); + } else { + return false; + } + return true; +} + +static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(const ggml_tensor * dst) { + if (!ggml_flash_attn_ext_is_ds4(dst)) { + return false; + } + + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const ggml_tensor * sinks = dst->src[4]; + const bool kv_f32 = K && V && K->type == GGML_TYPE_F32 && + V->type == GGML_TYPE_F32; + const bool kv_f16 = K && V && K->type == GGML_TYPE_F16 && + V->type == GGML_TYPE_F16; + const bool mask_ok = !mask || mask->type == GGML_TYPE_F16 || + (kv_f32 && mask->type == GGML_TYPE_F32); + float max_bias = 0.0f; + float logit_softcap = 0.0f; + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + if (!Q || !K || !V || + Q->type != GGML_TYPE_F32 || (!kv_f32 && !kv_f16) || !mask_ok || + dst->type != GGML_TYPE_F32 || + Q->ne[0] != 512 || K->ne[0] != 512 || V->ne[0] != 512 || + K->ne[1] != V->ne[1] || + K->ne[2] != 1 || V->ne[2] != 1 || + Q->ne[3] != 1 || K->ne[3] != 1 || V->ne[3] != 1 || + dst->ne[0] != 512 || dst->ne[1] != Q->ne[2] || + dst->ne[2] != Q->ne[1] || dst->ne[3] != 1 || + Q->nb[0] != (int64_t) sizeof(float) || + !ggml_is_contiguous(dst) || + max_bias != 0.0f || logit_softcap != 0.0f) { + return false; + } + const size_t kv_esz = kv_f16 ? sizeof(half) : sizeof(float); + if (K->nb[0] != kv_esz || V->nb[0] != kv_esz || + K->nb[1] != (size_t) K->ne[0] * kv_esz || + V->nb[1] != (size_t) V->ne[0] * kv_esz || + Q->nb[1] % sizeof(float) != 0 || + Q->nb[2] % sizeof(float) != 0 || + (mask && (mask->ne[0] != K->ne[1] || + mask->ne[1] != Q->ne[1] || + mask->ne[2] != 1 || mask->ne[3] != 1 || + mask->nb[0] != ggml_type_size(mask->type) || + !ggml_is_contiguous(mask)))) { + return false; + } + if (sinks && (sinks->type != GGML_TYPE_F32 || + sinks->ne[0] != Q->ne[2] || sinks->ne[1] != 1 || + sinks->ne[2] != 1 || sinks->ne[3] != 1 || + !ggml_is_contiguous(sinks))) { + return false; + } + + const int n_tokens = (int) Q->ne[1]; + const int n_heads = (int) Q->ne[2]; + const int n_kv = (int) K->ne[1]; + if (n_tokens <= 0 || n_heads <= 0 || n_kv <= 0) { + return false; + } + + const int raw_rows = ggml_get_op_params_i32(dst, 4); + const int sparse_keep_rows = ggml_get_op_params_i32(dst, 5); + const unsigned int ds4_layout = + (unsigned int) ggml_get_op_params_i32(dst, 6); + const int raw_window = (int) (ds4_layout >> 16); + const int sparse_block_size = (int) (ds4_layout & 0xffffu); + const int rope_flags = ggml_get_op_params_i32(dst, 7); + if (raw_rows < 0 || raw_rows > n_kv || + (ds4_layout != 0 && (raw_window <= 0 || sparse_block_size <= 0)) || + sparse_keep_rows == INT_MIN || + (sparse_keep_rows != 0 && !mask) || + (rope_flags & ~3) != 0 || + ((rope_flags & 2) != 0 && (rope_flags & 1) == 0)) { + return false; + } + + return true; +} + +static bool ggml_cuda_ds4_flash_attn_d512_f32( + ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + if (!ggml_cuda_ds4_flash_attn_d512_f32_supported(dst)) { + return false; + } + + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const ggml_tensor * sinks = dst->src[4]; + const bool kv_f32 = K->type == GGML_TYPE_F32; + const bool kv_f16 = K->type == GGML_TYPE_F16; + const int n_tokens = (int) Q->ne[1]; + const int n_heads = (int) Q->ne[2]; + const int n_kv = (int) K->ne[1]; + const size_t q_stride_token = Q->nb[1] / sizeof(float); + const size_t q_stride_head = Q->nb[2] / sizeof(float); + + int raw_rows = ggml_get_op_params_i32(dst, 4); + int sparse_keep_rows = ggml_get_op_params_i32(dst, 5); + const unsigned int ds4_layout = + (unsigned int) ggml_get_op_params_i32(dst, 6); + int raw_window = (int) (ds4_layout >> 16); + int sparse_block_size = (int) (ds4_layout & 0xffffu); + raw_rows = max(0, min(raw_rows, n_kv)); + if (raw_window <= 0) raw_window = raw_rows; + raw_window = max(1, min(raw_window, raw_rows)); + if (sparse_block_size <= 0) sparse_block_size = 32; + const int n_comp_rows = n_kv - raw_rows; + const int n_comp_blocks = (n_comp_rows + sparse_block_size - 1) / + sparse_block_size; + const bool sparse = sparse_keep_rows > 0 && + sparse_keep_rows < n_comp_rows && + n_comp_blocks > 0; + const bool indexed_mask = sparse_keep_rows < 0 && n_comp_rows > 0; + const int indexed_capacity = indexed_mask + ? min(-sparse_keep_rows, n_comp_rows) : 0; + + ds4_inverse_rope_params inverse_rope{}; + const int rope_flags = ggml_get_op_params_i32(dst, 7); + inverse_rope.enabled = rope_flags & 1; + inverse_rope.forward_q_enabled = (rope_flags & 2) != 0; + if (rope_flags != 0) { + inverse_rope.kv_start = ggml_get_op_params_i32(dst, 8); + const float freq_base = ggml_get_op_params_f32(dst, 9); + inverse_rope.freq_scale = ggml_get_op_params_f32(dst, 10); + inverse_rope.ext_factor = ggml_get_op_params_f32(dst, 11); + inverse_rope.attn_factor = ggml_get_op_params_f32(dst, 12); + const float beta_fast = ggml_get_op_params_f32(dst, 13); + const float beta_slow = ggml_get_op_params_f32(dst, 14); + const int n_ctx_orig = ggml_get_op_params_i32(dst, 15); + float corr_dims[2]; + ggml_rope_yarn_corr_dims( + 64, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); + inverse_rope.corr_low = corr_dims[0]; + inverse_rope.corr_high = corr_dims[1]; + inverse_rope.theta_scale = powf(freq_base, -2.0f / 64.0f); + } + + cudaStream_t stream = ctx.stream(); + ggml_cuda_pool_alloc inverse_rope_coefficients_alloc(ctx.pool()); + float * inverse_rope_coefficients = nullptr; + if (inverse_rope.enabled) { + inverse_rope_coefficients = inverse_rope_coefficients_alloc.alloc( + (size_t) n_tokens * 32 * 2); + const int coefficient_count = n_tokens * 32; + ds4_inverse_rope_coefficients_kernel<<< + (coefficient_count + 255) / 256, 256, 0, stream>>>( + inverse_rope_coefficients, n_tokens, inverse_rope); + } + ggml_cuda_pool_alloc forward_rope_coefficients_alloc(ctx.pool()); + float * forward_rope_coefficients = nullptr; + if (inverse_rope.forward_q_enabled) { + forward_rope_coefficients = forward_rope_coefficients_alloc.alloc( + (size_t) n_tokens * 32 * 2); + const int coefficient_count = n_tokens * 32; + ds4_forward_rope_coefficients_kernel<<< + (coefficient_count + 255) / 256, 256, 0, stream>>>( + forward_rope_coefficients, n_tokens, inverse_rope); + } + ggml_cuda_pool_alloc mean_k_alloc(ctx.pool()); + float * mean_k = nullptr; + if (sparse) { + mean_k = mean_k_alloc.alloc((size_t) n_comp_blocks * 512); + if (kv_f16) { + ds4_fa_mean_comp_blocks_kernel + <<>>( + (const half *) K->data, mean_k, n_kv, raw_rows, + sparse_block_size, n_comp_blocks); + } else { + ds4_fa_mean_comp_blocks_kernel + <<>>( + (const float *) K->data, mean_k, n_kv, raw_rows, + sparse_block_size, n_comp_blocks); + } + } + dim3 grid((unsigned) n_tokens, (unsigned) n_heads, 1); + const bool needs_rope_tail = inverse_rope.enabled || + inverse_rope.forward_q_enabled; + const size_t shmem = + (size_t) (n_kv + 2 * n_comp_blocks + + (needs_rope_tail ? 64 : 0)) * sizeof(float); + float params[3] = {}; + memcpy(params, dst->op_params, sizeof(params)); + const float scale = params[0]; + + constexpr int group4 = 4; + constexpr int group2 = 2; + const size_t group4_shmem = + ((size_t) group4 * n_kv + (size_t) group4 * 256) * sizeof(float) + + (size_t) group4 * 4 * sizeof(int) + + (inverse_rope.forward_q_enabled ? (size_t) group4 * 64 * sizeof(float) : 0); + const size_t group2_shmem = + ((size_t) group2 * n_kv + (size_t) group2 * 256) * sizeof(float) + + (size_t) group2 * 4 * sizeof(int) + + (inverse_rope.forward_q_enabled ? (size_t) group2 * 64 * sizeof(float) : 0); + const int compact_score_stride = raw_window + + (indexed_mask ? indexed_capacity : n_comp_rows); + const size_t compact_group4_shmem = + ((size_t) group4 * compact_score_stride + (size_t) group4 * 256) * sizeof(float) + + (size_t) group4 * 4 * sizeof(int) + + (inverse_rope.forward_q_enabled ? (size_t) group4 * 64 * sizeof(float) : 0); + // Four heads win while two blocks can remain resident in 48 KiB of LDS. + // Beyond that point, two-head grouping trades some K/V reuse for higher + // occupancy; larger working sets fall back to the single-head kernel. + if (!sparse && n_heads % group4 == 0 && group4_shmem <= 24 * 1024) { + return ds4_launch_flash_attn_d512_grouped( + dst, Q, K, V, mask, sinks, kv_f16, kv_f32, + n_tokens, n_heads, n_kv, scale, raw_rows, + q_stride_token, q_stride_head, + inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients, + group4_shmem, stream); + } + // Long causal-prefill chunks can have thousands of physical raw rows but + // at most raw_window visible rows for any one token. Compacting only the + // score storage lets the same exact four-head kernel remain at two-block + // occupancy. Keep the ordinary path for shapes that already fit, avoiding + // a bounds-scan launch where it cannot improve grouping. + const bool compact_group4 = + !sparse && mask && n_heads % group4 == 0 && + raw_rows > raw_window && group4_shmem > 24 * 1024 && + compact_group4_shmem <= 24 * 1024; + if (compact_group4) { + ggml_cuda_pool_alloc visibility_bounds_alloc(ctx.pool()); + int * visibility_bounds = visibility_bounds_alloc.alloc( + (size_t) n_tokens * 4); + ggml_cuda_pool_alloc indexed_rows_alloc(ctx.pool()); + ggml_cuda_pool_alloc indexed_counts_alloc(ctx.pool()); + ggml_cuda_pool_alloc indexed_owner_offsets_alloc(ctx.pool()); + ggml_cuda_pool_alloc indexed_owner_ranks_alloc(ctx.pool()); + int * indexed_rows = nullptr; + int * indexed_counts = nullptr; + int * indexed_owner_offsets = nullptr; + int * indexed_owner_ranks = nullptr; + if (indexed_mask) { + indexed_rows = indexed_rows_alloc.alloc( + (size_t) n_tokens * indexed_capacity); + indexed_counts = indexed_counts_alloc.alloc((size_t) n_tokens); + indexed_owner_offsets = indexed_owner_offsets_alloc.alloc( + (size_t) n_tokens * 257); + indexed_owner_ranks = indexed_owner_ranks_alloc.alloc( + (size_t) n_tokens * indexed_capacity); + if (mask->type == GGML_TYPE_F16) { + ds4_fa_indexed_rows_kernel<<>>( + (const half *) mask->data, indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, + n_tokens, n_kv, raw_rows, + indexed_capacity); + } else { + ds4_fa_indexed_rows_kernel<<>>( + (const float *) mask->data, indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, + n_tokens, n_kv, raw_rows, + indexed_capacity); + } + CUDA_CHECK(cudaGetLastError()); + } + if (mask->type == GGML_TYPE_F16) { + ds4_fa_visibility_bounds_kernel<<>>( + (const half *) mask->data, visibility_bounds, + n_tokens, n_kv, raw_rows); + } else { + ds4_fa_visibility_bounds_kernel<<>>( + (const float *) mask->data, visibility_bounds, + n_tokens, n_kv, raw_rows); + } + CUDA_CHECK(cudaGetLastError()); + if (indexed_mask) { + return ds4_launch_flash_attn_d512_grouped_compact< + group4, true, 4>( + dst, Q, K, V, mask, sinks, kv_f16, kv_f32, + n_tokens, n_heads, n_kv, scale, raw_rows, + raw_window, compact_score_stride, visibility_bounds, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, indexed_capacity, + q_stride_token, q_stride_head, + inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients, + compact_group4_shmem, stream); + } + return ds4_launch_flash_attn_d512_grouped_compact< + group4, false, 4>( + dst, Q, K, V, mask, sinks, kv_f16, kv_f32, + n_tokens, n_heads, n_kv, scale, raw_rows, + raw_window, compact_score_stride, visibility_bounds, + nullptr, nullptr, nullptr, nullptr, 0, + q_stride_token, q_stride_head, + inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients, + compact_group4_shmem, stream); + } + if (!sparse && n_heads % group2 == 0 && group2_shmem <= 48 * 1024) { + return ds4_launch_flash_attn_d512_grouped( + dst, Q, K, V, mask, sinks, kv_f16, kv_f32, + n_tokens, n_heads, n_kv, scale, raw_rows, + q_stride_token, q_stride_head, + inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients, + group2_shmem, stream); + } + + if (kv_f16 && (!mask || mask->type == GGML_TYPE_F16)) { + ds4_flash_attn_d512_shared_kv_kernel + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, (const half *) V->data, + mask ? (const half *) mask->data : nullptr, + sinks ? (const float *) sinks->data : nullptr, + mean_k, n_tokens, n_heads, n_kv, scale, raw_rows, + sparse_keep_rows, sparse_block_size, n_comp_blocks, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients); + } else if (kv_f32 && (!mask || mask->type == GGML_TYPE_F32)) { + ds4_flash_attn_d512_shared_kv_kernel + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + mask ? (const float *) mask->data : nullptr, + sinks ? (const float *) sinks->data : nullptr, + mean_k, n_tokens, n_heads, n_kv, scale, raw_rows, + sparse_keep_rows, sparse_block_size, n_comp_blocks, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients); + } else if (kv_f32 && mask && mask->type == GGML_TYPE_F16) { + ds4_flash_attn_d512_shared_kv_kernel + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + mean_k, n_tokens, n_heads, n_kv, scale, raw_rows, + sparse_keep_rows, sparse_block_size, n_comp_blocks, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients); + } else { + return false; + } + return true; +} + +#endif // defined(GGML_USE_HIP) + template static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; @@ -570,6 +2418,16 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_set_device(ctx.device); + if (ggml_flash_attn_ext_is_ds4(dst)) { +#if defined(GGML_USE_HIP) + if (!ggml_cuda_ds4_flash_attn_d512_f32(ctx, dst)) { + GGML_ABORT("unsupported DeepSeek4 D=512 flash-attention contract"); + } + return; +#else + GGML_ABORT("DeepSeek4 D=512 flash attention is only available on HIP"); +#endif // defined(GGML_USE_HIP) + } switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error"); @@ -592,5 +2450,12 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst } bool ggml_cuda_flash_attn_ext_supported(int device, const ggml_tensor * dst) { + if (ggml_flash_attn_ext_is_ds4(dst)) { +#if defined(GGML_USE_HIP) + return ggml_cuda_ds4_flash_attn_d512_f32_supported(dst); +#else + return false; +#endif // defined(GGML_USE_HIP) + } return ggml_cuda_get_best_fattn_kernel(device, dst) != BEST_FATTN_KERNEL_NONE; } 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 78482896b..4ca96d5a0 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 @@ -65,6 +65,7 @@ #include "ggml-cuda/fill.cuh" #include "ggml-cuda/moe-fused.cuh" #include "ggml-cuda/ds4-hc.cuh" +#include "ggml-cuda/ds4-indexer.cuh" #include "ggml.h" #include @@ -2442,9 +2443,17 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, const ggml_tensor * ffn_gate, const ggml_tensor * glu, const ggml_tensor * ffn_up_bias = nullptr, - const ggml_tensor * ffn_gate_bias = nullptr) { + const ggml_tensor * ffn_gate_bias = nullptr, + const ggml_tensor * ffn_up_glu = nullptr, + const ggml_tensor * ffn_gate_glu = nullptr) { const bool has_bias = ffn_up_bias != nullptr || ffn_gate_bias != nullptr; + // Shape-only views are common between routed expert projections and GLU. + // Keep validation anchored on the actual matmuls while allowing the GLU + // to consume their equivalent reshape nodes. + ffn_up_glu = ffn_up_glu ? ffn_up_glu : ffn_up; + ffn_gate_glu = ffn_gate_glu ? ffn_gate_glu : ffn_gate; + if (has_bias && (!ffn_up_bias || !ffn_gate_bias)) { return false; } @@ -2484,7 +2493,7 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, } } } else { - if (glu->src[0] != ffn_gate && glu->src[1] != ffn_up) { + if (glu->src[0] != ffn_gate_glu || glu->src[1] != ffn_up_glu) { return false; } } @@ -2609,8 +2618,68 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { return use_mul_mat_vec_q; } +// Execute an unbiased gate/up GLU while retaining the caller's existing +// output tensors. Vector kernels write the GLU directly. Large routed-expert +// MMQ keeps both ordinary matmul writes (and therefore numerical behavior), +// but shares the identical ids sort and F32->Q8 activation quantization. +static bool ggml_cuda_try_fuse_mul_mat_glu( + ggml_backend_cuda_context & ctx, + ggml_tensor * gate, + ggml_tensor * up, + ggml_tensor * glu) { + const ggml_tensor * src0 = up->src[0]; + const ggml_tensor * src1 = up->src[1]; + const ggml_tensor * ids = up->src[2]; + + if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate->src[0]; + ggml_cuda_set_fusion_glu_params(fusion_data, glu); + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, ids, glu, &fusion_data); + return true; + } + + if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate->src[0]; + ggml_cuda_set_fusion_glu_params(fusion_data, glu); + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, glu, &fusion_data); + return true; + } + + if (ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU_DS4) { + if (!ids && (ggml_mul_mat_is_grouped_src(up) || + ggml_mul_mat_is_grouped_src(gate))) { + return false; + } + + const ggml_tensor * weights[] = { up->src[0], gate->src[0] }; + for (const ggml_tensor * weight : weights) { + const bool bad_padding_clear = + ggml_backend_buffer_get_usage(weight->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && + ggml_nbytes(weight) != ggml_backend_buffer_get_alloc_size(weight->buffer, weight) && + weight->view_src; + if (bad_padding_clear) { + return false; + } + } + + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const int64_t ncols = ids ? src1->ne[2] : src1->ne[1]; + if (ggml_cuda_should_use_mmq(src0->type, cc, ncols, src0->ne[2])) { + ggml_cuda_mul_mat_q_pair( + ctx, up->src[0], gate->src[0], src1, ids, up, gate); + ggml_cuda_op_swiglu_ds4(ctx, glu); + return true; + } + } + + return false; +} + static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft); + const bool grouped_src = ggml_mul_mat_is_grouped_src(dst); // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. // But if src0 is also a view of another tensor then this cannot be done safely because it may overwrite valid tensor data. @@ -2682,7 +2751,13 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor bool use_batched_cublas_bf16 = src0->type == GGML_TYPE_BF16 && bf16_mma_hardware_available(cc); bool use_batched_cublas_f32 = src0->type == GGML_TYPE_F32; - if (!split && use_mul_mat_vec_f) { + if (grouped_src) { + // Only MMQ's grouped activation quantizer understands the physical + // [K/group,N,group] source layout. + GGML_ASSERT(!split); + GGML_ASSERT(use_mul_mat_q); + ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + } else if (!split && use_mul_mat_vec_f) { // the custom F16 vector kernel can be used over batched cuBLAS GEMM // but this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention) ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); @@ -3030,6 +3105,15 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_DS4_HC: ggml_cuda_op_ds4_hc(ctx, dst); break; + case GGML_OP_DS4_INDEXER_QAT: + ggml_cuda_op_ds4_indexer_qat(ctx, dst); + break; + case GGML_OP_DS4_INDEXER_SCORE: + ggml_cuda_op_ds4_indexer_score(ctx, dst); + break; + case GGML_OP_DS4_INDEXER_MASK: + ggml_cuda_op_ds4_indexer_mask(ctx, dst); + break; case GGML_OP_GROUP_NORM: ggml_cuda_op_group_norm(ctx, dst); break; @@ -3067,6 +3151,7 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg ggml_cuda_op_rms_norm_back(ctx, dst); break; case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_GROUPED_SRC: ggml_cuda_mul_mat(ctx, dst->src[0], dst->src[1], dst); break; case GGML_OP_MUL_MAT_ID: @@ -4025,6 +4110,12 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, std::initializer_list mul_mat_id_glu_ops = { GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU }; std::initializer_list mul_mat_glu_ops = { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU }; + std::initializer_list mul_mat_id_reshape_glu_ops = { + GGML_OP_MUL_MAT_ID, GGML_OP_RESHAPE, + GGML_OP_MUL_MAT_ID, GGML_OP_RESHAPE, GGML_OP_GLU }; + std::initializer_list mul_mat_reshape_glu_ops = { + GGML_OP_MUL_MAT, GGML_OP_RESHAPE, + GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_GLU }; if ((is_equal(mul_mat_bias_glu_ops, ops) || is_equal(mul_mat_id_bias_glu_ops, ops)) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 4 })) { @@ -4052,6 +4143,25 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, } } + if ((is_equal(mul_mat_id_reshape_glu_ops, ops) || is_equal(mul_mat_reshape_glu_ops, ops)) && + ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 4 })) { + const ggml_tensor * ffn_gate = cgraph->nodes[node_idx]; + const ggml_tensor * ffn_gate_glu = cgraph->nodes[node_idx + 1]; + const ggml_tensor * ffn_up = cgraph->nodes[node_idx + 2]; + const ggml_tensor * ffn_up_glu = cgraph->nodes[node_idx + 3]; + const ggml_tensor * glu = cgraph->nodes[node_idx + 4]; + + if (ffn_gate_glu->src[0] == ffn_gate && + ffn_up_glu->src[0] == ffn_up && + ggml_cuda_should_fuse_mul_mat( + ffn_up, ffn_gate, glu, nullptr, nullptr, + ffn_up_glu, ffn_gate_glu)) { + int out_nodes[] = { node_idx + 4 }; + return ggml_cuda_check_fusion_memory_ranges( + cgraph, node_idx, (int) ops.size(), out_nodes, 1); + } + } + std::initializer_list rope_set_rows_ops = { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }; if (is_equal(rope_set_rows_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) { @@ -4550,6 +4660,32 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud fused_node_count = 5; break; } + } else if (ggml_cuda_can_fuse( + cgraph, i, + { op, GGML_OP_RESHAPE, op, GGML_OP_RESHAPE, GGML_OP_GLU }, {})) { + ggml_tensor * glu = cgraph->nodes[i + 4]; + ggml_tensor * gate_view = glu->src[0]; + ggml_tensor * up_view = glu->src[1]; + ggml_tensor * gate = nullptr; + ggml_tensor * up = nullptr; + + if (gate_view == cgraph->nodes[i + 1] && + up_view == cgraph->nodes[i + 3]) { + gate = cgraph->nodes[i]; + up = cgraph->nodes[i + 2]; + } else if (gate_view == cgraph->nodes[i + 3] && + up_view == cgraph->nodes[i + 1]) { + gate = cgraph->nodes[i + 2]; + up = cgraph->nodes[i]; + } else { + continue; + } + + if (ggml_cuda_try_fuse_mul_mat_glu(*cuda_ctx, gate, up, glu)) { + fused_mul_mat_vec = true; + fused_node_count = 5; + break; + } } else if (ggml_cuda_can_fuse(cgraph, i, { op, op, GGML_OP_GLU }, {})) { ggml_tensor * glu = cgraph->nodes[i + 2]; ggml_tensor * gate = glu->src[0]; @@ -4560,27 +4696,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud if (!ok) continue; - const ggml_tensor * src0 = up->src[0]; - const ggml_tensor * src1 = up->src[1]; - const ggml_tensor * ids = up->src[2]; - - if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate->src[0]; - ggml_cuda_set_fusion_glu_params(fusion_data, glu); - - ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 3; - break; - } - - if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate->src[0]; - ggml_cuda_set_fusion_glu_params(fusion_data, glu); - - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + if (ggml_cuda_try_fuse_mul_mat_glu(*cuda_ctx, gate, up, glu)) { fused_mul_mat_vec = true; fused_node_count = 3; break; @@ -6363,11 +6479,65 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return true; case GGML_OP_DS4_HC: return true; + case GGML_OP_DS4_INDEXER_QAT: + return op->src[0]->type == GGML_TYPE_F32 && + op->src[0]->ne[0] == 128 && + ggml_is_contiguous(op->src[0]); + case GGML_OP_DS4_INDEXER_SCORE: + return op->src[0]->type == GGML_TYPE_F32 && + op->src[0]->ne[0] == 128 && + op->src[0]->ne[3] == 1 && + op->src[1]->type == GGML_TYPE_F32 && + op->src[1]->ne[2] == 1 && + op->src[1]->ne[3] == 1 && + op->src[2]->type == GGML_TYPE_F16 && + op->src[2]->ne[0] == 128 && + op->src[2]->ne[2] == 1 && + op->src[2]->ne[3] == 1 && + ggml_is_contiguous(op->src[0]) && + ggml_is_contiguous(op->src[1]) && + ggml_is_contiguous(op->src[2]); + case GGML_OP_DS4_INDEXER_MASK: + return op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_I32 && + ggml_is_contiguous(op->src[0]) && + ggml_is_contiguous(op->src[1]); case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_GROUPED_SRC: case GGML_OP_MUL_MAT_ID: { struct ggml_tensor * a = op->src[0]; struct ggml_tensor * b = op->src[1]; + if (ggml_mul_mat_is_grouped_src(op)) { + const ggml_tensor * physical = b->view_src; + const int cc = ggml_cuda_info().devices[dev_ctx->device].cc; + const bool bad_padding_clear = + a->buffer && + ggml_backend_buffer_get_usage(a->buffer) == + GGML_BACKEND_BUFFER_USAGE_COMPUTE && + ggml_nbytes(a) != + ggml_backend_buffer_get_alloc_size(a->buffer, a) && + a->view_src; + return op->op == GGML_OP_MUL_MAT_GROUPED_SRC && + a->buffer && + !ggml_backend_buft_is_cuda_split(a->buffer->buft) && + ggml_is_quantized(a->type) && + !bad_padding_clear && + b->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32 && + physical && + physical->type == GGML_TYPE_F32 && + physical->ne[0] % 4 == 0 && + physical->ne[0] * physical->ne[2] == b->ne[0] && + physical->ne[1] == b->ne[1] && + physical->ne[2] == + ggml_mul_mat_grouped_src_groups(op) && + physical->ne[3] == 1 && + b->ne[0] % (4 * QK8_1) == 0 && + ggml_is_contiguous(physical) && + ggml_cuda_should_use_mmq( + a->type, cc, b->ne[1], /*n_experts=*/0); + } if (a->buffer && ggml_backend_buft_is_cuda_split(a->buffer->buft)) { if (a->ne[2] > 1 || a->ne[3] > 1) { return false; @@ -6576,7 +6746,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_F16 || src0_type == GGML_TYPE_BF16 || - src0_type == GGML_TYPE_I8; + src0_type == GGML_TYPE_I8 || + src0_type == GGML_TYPE_I32; } break; case GGML_OP_CONV_TRANSPOSE_1D: { @@ -6666,10 +6837,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return ggml_is_contiguous_rows(op->src[0]); case GGML_OP_TOP_K: case GGML_OP_ARGSORT: -#ifndef GGML_CUDA_USE_CUB - return op->src[0]->ne[0] <= 1024; -#else +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) return true; +#else + return op->src[0]->ne[0] <= 1024; #endif case GGML_OP_SUM_ROWS: case GGML_OP_MEAN: @@ -6724,6 +6895,7 @@ static int64_t get_op_batch_size(const ggml_tensor * op) { case GGML_OP_GET_ROWS: return 0; case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_GROUPED_SRC: return op->ne[1]; case GGML_OP_MUL_MAT_ID: case GGML_OP_ROPE: diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu index 06dfae700..c4a405f28 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu @@ -20,6 +20,15 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con case GGML_TYPE_Q8_0: mul_mat_q_case(ctx, args, stream); break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q2_0_ROCMFP2: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q3_0_ROCMFPX: + mul_mat_q_case(ctx, args, stream); + break; case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: #ifndef GGML_CUDA_BLACKWELL_CONSUMER @@ -77,11 +86,26 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con } } -void ggml_cuda_mul_mat_q( - ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) { +static void ggml_cuda_mul_mat_q_impl( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, + const ggml_tensor * src0_pair, + const ggml_tensor * src1, + const ggml_tensor * ids, + ggml_tensor * dst, + ggml_tensor * dst_pair) { GGML_ASSERT( src1->type == GGML_TYPE_F32); GGML_ASSERT( dst->type == GGML_TYPE_F32); GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. + GGML_ASSERT((src0_pair == nullptr) == (dst_pair == nullptr)); + if (src0_pair) { + GGML_ASSERT(src0_pair->type == src0->type); + GGML_ASSERT(dst_pair->type == dst->type); + GGML_ASSERT(ggml_are_same_shape(src0_pair, src0)); + GGML_ASSERT(ggml_are_same_stride(src0_pair, src0)); + GGML_ASSERT(ggml_are_same_shape(dst_pair, dst)); + GGML_ASSERT(ggml_are_same_stride(dst_pair, dst)); + } GGML_TENSOR_BINARY_OP_LOCALS; @@ -101,14 +125,23 @@ void ggml_cuda_mul_mat_q( const float * src1_d = (const float *) src1->data; float * dst_d = (float *) dst->data; - // If src0 is a temporary compute buffer, clear any potential padding. - if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { - const size_t size_data = ggml_nbytes(src0); - const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0); + // Temporary weight tensors may have an allocation tail read by tiled MMQ. + // Clear it once for each projection before launching either multiply. + const ggml_tensor * weights[] = {src0, src0_pair}; + for (const ggml_tensor * weight : weights) { + if (!weight || + ggml_backend_buffer_get_usage(weight->buffer) != + GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + continue; + } + const size_t size_data = ggml_nbytes(weight); + const size_t size_alloc = + ggml_backend_buffer_get_alloc_size(weight->buffer, weight); if (size_alloc > size_data) { - GGML_ASSERT(ggml_is_contiguously_allocated(src0)); - GGML_ASSERT(!src0->view_src); - CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream)); + GGML_ASSERT(ggml_is_contiguously_allocated(weight)); + GGML_ASSERT(!weight->view_src); + CUDA_CHECK(cudaMemsetAsync((char *) weight->data + size_data, 0, + size_alloc - size_data, stream)); } } @@ -121,14 +154,16 @@ void ggml_cuda_mul_mat_q( const int64_t s03 = src0->nb[3] / ts_src0; const int64_t s3 = dst->nb[3] / ts_dst; - bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) - || GGML_CUDA_CC_IS_CDNA(cc); - // LUCE_MMQ_DP_MAX_NE1: for ne1 at or below this, skip stream-k and use - // data-parallel tiles - the stream-k fixup pass costs ~1ms/step at - // spec-decode verify widths (measured sm_86, w6 chain). 0 = always stream-k. + bool use_stream_k = + (GGML_CUDA_CC_IS_NVIDIA(cc) && + ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) || + GGML_CUDA_CC_IS_CDNA(cc); + // Keep the established small-batch tuning hook from #503. The default + // remains stream-k; positive values opt small verify widths into the + // lower-overhead data-parallel path. static const int luce_mmq_dp_max_ne1 = []() { - const char * e = getenv("LUCE_MMQ_DP_MAX_NE1"); - return e ? atoi(e) : 0; + const char * value = getenv("LUCE_MMQ_DP_MAX_NE1"); + return value ? atoi(value) : 0; }(); if (use_stream_k && ne11 <= luce_mmq_dp_max_ne1) { use_stream_k = false; @@ -136,6 +171,17 @@ void ggml_cuda_mul_mat_q( // TODO: tighter pool buffer size vs q8 path const bool use_native_mxfp4 = blackwell_mma_available(cc) && src0->type == GGML_TYPE_MXFP4; + const bool grouped_src = !ids && ggml_mul_mat_is_grouped_src(dst); + const ggml_tensor * grouped_physical = grouped_src ? src1->view_src : nullptr; + if (grouped_src) { + GGML_ASSERT(grouped_physical && grouped_physical->type == GGML_TYPE_F32); + GGML_ASSERT(grouped_physical->ne[2] == + ggml_mul_mat_grouped_src_groups(dst)); + GGML_ASSERT(grouped_physical->ne[0] * grouped_physical->ne[2] == ne10); + GGML_ASSERT(grouped_physical->ne[1] == ne11); + GGML_ASSERT(grouped_physical->ne[3] == 1); + GGML_ASSERT(!use_native_mxfp4); + } if (!ids) { const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1 + @@ -146,7 +192,14 @@ void ggml_cuda_mul_mat_q( const int64_t s11 = src1->nb[1] / ts_src1; const int64_t s12 = src1->nb[2] / ts_src1; const int64_t s13 = src1->nb[3] / ts_src1; - if (use_native_mxfp4) { + if (grouped_src) { + quantize_mmq_q8_1_grouped_cuda( + src1_d, src1_q8_1.get(), src0->type, + ne10, grouped_physical->ne[0], + grouped_physical->nb[1] / ts_src1, + grouped_physical->nb[2] / ts_src1, + ne10_padded, ne11, stream); + } else if (use_native_mxfp4) { static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1)); quantize_mmq_mxfp4_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); @@ -173,6 +226,12 @@ void ggml_cuda_mul_mat_q( ne03, ne13, s03, s13, s3, use_stream_k, ne1}; ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); + if (src0_pair) { + mmq_args pair_args = args; + pair_args.x = (const char *) src0_pair->data; + pair_args.dst = (float *) dst_pair->data; + ggml_cuda_mul_mat_q_switch_type(ctx, pair_args, stream); + } return; } @@ -238,6 +297,33 @@ void ggml_cuda_mul_mat_q( use_stream_k, ne12}; ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); + if (src0_pair) { + mmq_args pair_args = args; + pair_args.x = (const char *) src0_pair->data; + pair_args.dst = (float *) dst_pair->data; + ggml_cuda_mul_mat_q_switch_type(ctx, pair_args, stream); + } +} + +void ggml_cuda_mul_mat_q( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, + const ggml_tensor * src1, + const ggml_tensor * ids, + ggml_tensor * dst) { + ggml_cuda_mul_mat_q_impl(ctx, src0, nullptr, src1, ids, dst, nullptr); +} + +void ggml_cuda_mul_mat_q_pair( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0_a, + const ggml_tensor * src0_b, + const ggml_tensor * src1, + const ggml_tensor * ids, + ggml_tensor * dst_a, + ggml_tensor * dst_b) { + ggml_cuda_mul_mat_q_impl( + ctx, src0_a, src0_b, src1, ids, dst_a, dst_b); } void ggml_cuda_op_mul_mat_q( @@ -316,6 +402,12 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t case GGML_TYPE_IQ4_NL: mmq_supported = true; break; + case GGML_TYPE_Q2_0_ROCMFP2: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + // ROCmFPX MMQ variants are implemented for gfx1151 only. + mmq_supported = GGML_CUDA_CC_IS_RDNA3_5(cc); + break; default: mmq_supported = false; break; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index d8df400fb..ad26af456 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh @@ -66,6 +66,10 @@ static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) { return MMQ_Q8_1_DS_LAYOUT_DS4; case GGML_TYPE_Q8_0: return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q2_0_ROCMFP2: + case GGML_TYPE_Q3_0_ROCMFPX: + return MMQ_Q8_1_DS_LAYOUT_D4; case GGML_TYPE_MXFP4: return MMQ_Q8_1_DS_LAYOUT_D4; case GGML_TYPE_NVFP4: @@ -101,24 +105,9 @@ struct tile_x_sizes { int sc; }; -// Lucebox: tuned MMQ tile for RDNA3/RDNA4. -// DFlash issues many small mul_mat_q calls (ne[1] ~= ddtree_budget+1, ~23 at -// the default budget=22) where the stock 128x128 / 8-warp tile under-occupies. -// Current settings (mmq_y=128, nwarps=8, mmq_x_max=128): -// - mmq_y=128/nwarps=8: stock value for non-RDNA1 AMD — occupancy for decode -// batches (ne[1]~23) measured at +8.3% vs stock on gfx1201; also the right -// size for prefill (ne[1]=512) where halving block count improves L2 reuse -// (+6.9% prefill speedup on Q4_K m=4096,n=512,k=14336 on gfx1201). -// - mmq_x_max=128: prefill (ne[1]=512) fits in 4 x-tiles instead of 11. -// For decode (ne[1]~23), mmq_x selects 24 regardless of the cap — no impact. -// Original decode measurements (Qwen3.6-27B Q4_K_M, --ddtree-budget=22): -// gfx1100 (RX 7900 XTX): 56.78 -> 60.18 tok/s (+6.0%) -// gfx1201 (R9700): 54.65 -> 59.20 tok/s (+8.3%) -// gfx1151 (Strix Halo): 11.53 -> 12.00 tok/s (+4.1%, 256-token smoke) -// Constraint: nwarps * granularity(mmq_x) / ntx(mmq_x) == mmq_y. -// mmq_x< 128: granularity=16, ntx=1 → 8*16/1=128 ✓ -// mmq_x>=128: granularity=32, ntx=2 → 8*32/2=128 ✓ -// Define LUCEBOX_RDNA_MMQ_TILE_OVERRIDE=0 to disable. +// RDNA uses 128x128, eight-warp MMQ tiles by default. ROCmFPX template +// instances use 64x64, four-warp tiles: their unpacking pressure makes the +// smaller tile faster on gfx1151 without changing other quant formats. #ifndef LUCEBOX_RDNA_MMQ_TILE_OVERRIDE #define LUCEBOX_RDNA_MMQ_TILE_OVERRIDE 1 #endif @@ -130,7 +119,13 @@ struct tile_x_sizes { #endif static int get_mmq_x_max_host(const int cc) { - if (LUCEBOX_RDNA_TILE_HOST(cc)) return 128; + if (LUCEBOX_RDNA_TILE_HOST(cc)) { +#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) + return 64; +#else + return 128; +#endif + } return (amd_mfma_available(cc) || turing_mma_available(cc) || amd_wmma_available(cc)) ? 128 : GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA ? #ifdef GGML_CUDA_FORCE_MMQ @@ -142,8 +137,12 @@ static int get_mmq_x_max_host(const int cc) { static constexpr __device__ int get_mmq_x_max_device() { #if LUCEBOX_RDNA_TILE_DEVICE +#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) + return 64; +#else return 128; #endif +#endif #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) return 128; #else // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) @@ -167,7 +166,13 @@ static constexpr __device__ int get_mmq_x_max_device() { } static int get_mmq_y_host(const int cc) { - if (LUCEBOX_RDNA_TILE_HOST(cc)) return 128; + if (LUCEBOX_RDNA_TILE_HOST(cc)) { +#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) + return 64; +#else + return 128; +#endif + } return GGML_CUDA_CC_IS_AMD(cc) ? (GGML_CUDA_CC_IS_RDNA1(cc) ? 64 : 128) : ((GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64); } @@ -182,8 +187,12 @@ static constexpr __device__ int get_iter_k([[maybe_unused]] const ggml_type type static constexpr __device__ int get_mmq_y_device() { #if LUCEBOX_RDNA_TILE_DEVICE +#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) + return 64; +#else return 128; #endif +#endif #if defined(GGML_USE_HIP) #if defined(RDNA1) return 64; @@ -226,6 +235,9 @@ static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml case GGML_TYPE_Q5_0: return MMQ_DP4A_TXS_Q8_0; case GGML_TYPE_Q5_1: return MMQ_DP4A_TXS_Q8_1; case GGML_TYPE_Q8_0: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_Q2_0_ROCMFP2: return MMQ_DP4A_TXS_Q8_0_16; + case GGML_TYPE_Q3_0_ROCMFPX: return MMQ_DP4A_TXS_Q8_0_16; case GGML_TYPE_MXFP4: return MMQ_DP4A_TXS_Q8_1; case GGML_TYPE_NVFP4: return MMQ_DP4A_TXS_Q8_0_16; case GGML_TYPE_Q2_K: return MMQ_DP4A_TXS_Q2_K; @@ -270,6 +282,9 @@ static constexpr __host__ __device__ int mmq_get_mma_tile_x_k(ggml_type type) { case GGML_TYPE_Q5_0: return MMQ_MMA_TILE_X_K_Q8_0; case GGML_TYPE_Q5_1: return MMQ_MMA_TILE_X_K_Q8_1; case GGML_TYPE_Q8_0: return MMQ_MMA_TILE_X_K_Q8_0; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: return MMQ_MMA_TILE_X_K_Q8_0; + case GGML_TYPE_Q2_0_ROCMFP2: return MMQ_MMA_TILE_X_K_Q3_K; + case GGML_TYPE_Q3_0_ROCMFPX: return MMQ_MMA_TILE_X_K_Q3_K; // tile sizes are the same for Q8_1 and FP4 for blackwell case GGML_TYPE_MXFP4: return MMQ_MMA_TILE_X_K_Q8_1; case GGML_TYPE_NVFP4: return MMQ_MMA_TILE_X_K_NVFP4; @@ -320,7 +335,13 @@ static constexpr __device__ int mmq_get_granularity_device(const int /*mmq_x*/) #if defined(GGML_USE_HIP) static int mmq_get_nwarps_host(const int cc, const int warp_size) { - if (LUCEBOX_RDNA_TILE_HOST(cc)) return 8; + if (LUCEBOX_RDNA_TILE_HOST(cc)) { +#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) + return 4; +#else + return 8; +#endif + } return amd_mfma_available(cc) ? 8 : 256/warp_size; } #else @@ -331,8 +352,12 @@ static int mmq_get_nwarps_host(const int /*cc*/, const int warp_size) { static constexpr __device__ int mmq_get_nwarps_device() { #if LUCEBOX_RDNA_TILE_DEVICE +#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) + return 4; +#else return 8; #endif +#endif #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) return 8; #else @@ -845,6 +870,178 @@ template static __device__ __forceinline__ void loa } } +// Expand packed ROCmFP4 Codebook10 weights into the signed int8 layout used by +// the existing Q8_0 x Q8_1 MMQ kernels. Q4_0_ROCMFP4_FAST has one UE4M3 +// scale per 32-weight block, so after expansion its shared-memory shape is +// identical to Q8_0: 256 int8 weights and eight float scales per row/tile. +template static __device__ __forceinline__ void load_tiles_rocmfp4_fast( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0_ROCMFP4_FAST, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR_ROCMFP4); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI_ROCMFP4; + const int kqsx = txi % QI_ROCMFP4; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_rocmfp4_fast * bxi = (const block_rocmfp4_fast *) x + kbx0 + i*stride + kbx; + const int aux_q4 = rocmfp4_get_qs_i32(bxi->qs, kqsx); + const int2 v = rocmfp4_get_int_from_codebook_16(aux_q4, kvalues_rocmfp4); + const int k0 = kbx * (2 * QI_ROCMFP4) + kqsx; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + k0 + 0] = v.x; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + k0 + QI_ROCMFP4] = v.y; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI_ROCMFP4] = v.y; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = 2*MMQ_TILE_NE_K / QI8_0; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) { + int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_rocmfp4_fast * bxi = (const block_rocmfp4_fast *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kbxd] = rocmfp4_ue4m3_to_fp32_half_finite(bxi->e); +#else + x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + kbxd] = rocmfp4_ue4m3_to_fp32_half_finite(bxi->e); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template +struct rocmfpx_dual_mmq_traits; + +template <> +struct rocmfpx_dual_mmq_traits { + using block_t = block_rocmfp2; + + static __device__ __forceinline__ int pack4(const block_t * block, const int base) { + return rocmfpx_pack4_fp2_vec_cuda(block->qs, base); + } + + static __device__ __forceinline__ float scale(const block_t * block, const int half) { + return rocmfpx_ue4m3_to_fp32_finite(block->e[half]); + } +}; + +template <> +struct rocmfpx_dual_mmq_traits { + using block_t = block_rocmfp3; + + static __device__ __forceinline__ int pack4(const block_t * block, const int base) { + return rocmfpx_pack4_fp3_vec_cuda(block->qs, base); + } + + static __device__ __forceinline__ float scale(const block_t * block, const int half) { + return rocmfpx_ue4m3_to_fp32_finite(block->e[half]); + } +}; + +// Q2/FP3 store one UE4M3 scale per 16 weights. Expand their packed values +// into int8 and reuse the existing scale-per-16 Q8_0 x Q8_1 MMQ kernels. One +// wave loads a full 256-weight row: each lane expands two groups of four. +template +static __device__ __forceinline__ void load_tiles_rocmfpx_dual( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + using traits = rocmfpx_dual_mmq_traits; + using block_t = typename traits::block_t; + + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int groups_per_block = QK_ROCMFPX / 4; + constexpr int blocks_per_tile = MMQ_ITER_K / QK_ROCMFPX; + constexpr int threads_per_row = blocks_per_tile * groups_per_block / 2; + static_assert(threads_per_row == 32, "ROCmFPX MMQ loader expects 32 lanes per row"); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(type, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / (groups_per_block / 2); + const int group = txi % (groups_per_block / 2); + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + if (need_check) { + i = min(i, i_max); + } + + const block_t * block = (const block_t *) x + kbx0 + i*stride + kbx; + const int k0 = kbx*groups_per_block + group; + const int q0 = traits::pack4(block, 4*group); + const int q1 = traits::pack4( + block, 4*(group + groups_per_block/2)); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + k0] = q0; + x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + k0 + groups_per_block/2] = q1; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0] = q0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + groups_per_block/2] = q1; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + } + + constexpr int scales_per_tile = 2*blocks_per_tile; + constexpr int scale_rows_per_warp = warp_size / scales_per_tile; + const int kscale = threadIdx.x % scales_per_tile; + const int scale_block = kscale / 2; + const int scale_half = kscale % 2; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*scale_rows_per_warp) { + int i = i0 + threadIdx.y*scale_rows_per_warp + threadIdx.x/scales_per_tile; + if (need_check) { + i = min(i, i_max); + } + + const block_t * block = (const block_t *) x + kbx0 + i*stride + scale_block; +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q3_K + kscale] = traits::scale(block, scale_half); +#else + x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + kscale] = traits::scale(block, scale_half); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + template static __device__ __forceinline__ void load_tiles_mxfp4_fp4(const char * __restrict__ x, int * __restrict__ x_tile, @@ -3392,6 +3589,30 @@ struct mmq_type_traits { static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; }; +template +struct mmq_type_traits { + static constexpr int vdr = VDR_ROCMFP4_FAST_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_rocmfp4_fast; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_ROCMFP2_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_rocmfpx_dual; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_16_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_16_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_ROCMFP3_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_rocmfpx_dual; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_16_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_16_q8_1_dp4a; +}; + template struct mmq_type_traits { static constexpr int vdr = VDR_MXFP4_Q8_1_MMQ; @@ -3623,7 +3844,7 @@ template #if defined(GGML_USE_HIP) // RDNA4 is compute-bound on MMQ (WMMA path); allow compiler to use more VGPRs // (minBlocks=1 matches NVIDIA Volta+ behavior and reduces register spilling). -#if defined(RDNA4) +#if defined(RDNA4) && !defined(GGML_CUDA_ROCMFPX_MMQ_TILE) __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 1) #elif defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN) __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 2) @@ -4242,6 +4463,9 @@ extern DECL_MMQ_CASE(GGML_TYPE_Q4_1); extern DECL_MMQ_CASE(GGML_TYPE_Q5_0); extern DECL_MMQ_CASE(GGML_TYPE_Q5_1); extern DECL_MMQ_CASE(GGML_TYPE_Q8_0); +extern DECL_MMQ_CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST); +extern DECL_MMQ_CASE(GGML_TYPE_Q2_0_ROCMFP2); +extern DECL_MMQ_CASE(GGML_TYPE_Q3_0_ROCMFPX); extern DECL_MMQ_CASE(GGML_TYPE_MXFP4); extern DECL_MMQ_CASE(GGML_TYPE_NVFP4); extern DECL_MMQ_CASE(GGML_TYPE_Q2_K); @@ -4263,6 +4487,19 @@ extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); void ggml_cuda_mul_mat_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst); +// Execute two same-shape MUL_MAT operations with one activation quantization. +// When ids is non-null, also share the MUL_MAT_ID sort. The two MMQ launches +// and their output order remain unchanged, so callers can apply the ordinary +// GLU kernel byte-identically. +void ggml_cuda_mul_mat_q_pair( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0_a, + const ggml_tensor * src0_b, + const ggml_tensor * src1, + const ggml_tensor * ids, + ggml_tensor * dst_a, + ggml_tensor * dst_b); + void ggml_cuda_op_mul_mat_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cu index 4300ffc14..992e12d85 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cu @@ -176,7 +176,8 @@ template static __global__ void quantize_mmq_q8_1( const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, - const int64_t ne0, const int ne1, const int ne2) { + const int64_t ne0, const int ne1, const int ne2, + const int64_t group_width, const int64_t group_stride) { constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32; constexpr int vals_per_sum = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32; @@ -204,8 +205,21 @@ static __global__ void quantize_mmq_q8_1( const int64_t ib = ib0 + (i0 / (4*QK8_1))*ne1 + blockIdx.x; // block index in channel const int64_t iqs = i0 % (4*QK8_1); // quant index in block - // Load 4 floats per thread and calculate max. abs. value between them: - const float4 xi = i0 < ne00 ? x4[(i03*s03 + i02*s02 + i01*s01 + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f); + // Load 4 floats per thread and calculate max. abs. value between them. + // A grouped source is physically [group_width, token, group] while this + // kernel emits the same flattened-K Q8 layout as a materialized permute. + int64_t src_i00 = i00; + int64_t src_group_offset = 0; + if (group_width > 0) { + const int64_t group = i00 / group_width; + src_i00 = i00 - group * group_width; + src_group_offset = group * group_stride; + } + const int64_t src_index = + i03*s03 + i02*s02 + i01*s01 + src_group_offset + src_i00; + const float4 xi = i0 < ne00 + ? x4[src_index/4] + : make_float4(0.0f, 0.0f, 0.0f, 0.0f); float amax = fabsf(xi.x); amax = fmaxf(amax, fabsf(xi.y)); amax = fmaxf(amax, fabsf(xi.z)); @@ -300,15 +314,15 @@ void quantize_mmq_q8_1_cuda( switch (mmq_get_q8_1_ds_layout(type_src0)) { case MMQ_Q8_1_DS_LAYOUT_D4: quantize_mmq_q8_1 - <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2); + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, 0); break; case MMQ_Q8_1_DS_LAYOUT_DS4: quantize_mmq_q8_1 - <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2); + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, 0); break; case MMQ_Q8_1_DS_LAYOUT_D2S6: quantize_mmq_q8_1 - <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2); + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, 0); break; default: GGML_ABORT("fatal error"); @@ -316,6 +330,45 @@ void quantize_mmq_q8_1_cuda( } } +void quantize_mmq_q8_1_grouped_cuda( + const float * x, void * vy, const ggml_type type_src0, + const int64_t ne00, const int64_t group_width, + const int64_t token_stride, const int64_t group_stride, + const int64_t ne0, const int64_t ne1, cudaStream_t stream) { + GGML_ASSERT(ne00 % 4 == 0); + GGML_ASSERT(group_width > 0 && group_width % 4 == 0); + GGML_ASSERT(ne00 % group_width == 0); + GGML_ASSERT(ne0 % (4*QK8_1) == 0); + + const int64_t block_num_y = + (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / + (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const dim3 num_blocks(ne1, block_num_y, 1); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + switch (mmq_get_q8_1_ds_layout(type_src0)) { + case MMQ_Q8_1_DS_LAYOUT_D4: + quantize_mmq_q8_1 + <<>>( + x, nullptr, vy, ne00, token_stride, 0, 0, + ne0, ne1, 1, group_width, group_stride); + break; + case MMQ_Q8_1_DS_LAYOUT_DS4: + quantize_mmq_q8_1 + <<>>( + x, nullptr, vy, ne00, token_stride, 0, 0, + ne0, ne1, 1, group_width, group_stride); + break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: + quantize_mmq_q8_1 + <<>>( + x, nullptr, vy, ne00, token_stride, 0, 0, + ne0, ne1, 1, group_width, group_stride); + break; + default: + GGML_ABORT("fatal error"); + } +} + void quantize_mmq_mxfp4_cuda(const float * x, const int32_t * ids, void * vy, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cuh index 6a91df635..4768791c4 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cuh @@ -26,6 +26,12 @@ void quantize_mmq_q8_1_cuda( ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); +void quantize_mmq_q8_1_grouped_cuda( + const float * x, void * vy, ggml_type type_src0, + int64_t ne00, int64_t group_width, + int64_t token_stride, int64_t group_stride, + int64_t ne0, int64_t ne1, cudaStream_t stream); + void quantize_mmq_mxfp4_cuda(const float * x, const int32_t * ids, void * vy, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py index 17cee853e..22a7d57f4 100755 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py @@ -33,6 +33,7 @@ TYPES_MMQ = [ "GGML_TYPE_Q4_0", "GGML_TYPE_Q4_1", "GGML_TYPE_Q5_0", "GGML_TYPE_Q5_1", "GGML_TYPE_Q8_0", + "GGML_TYPE_Q4_0_ROCMFP4_FAST", "GGML_TYPE_Q2_0_ROCMFP2", "GGML_TYPE_Q3_0_ROCMFPX", "GGML_TYPE_Q2_K", "GGML_TYPE_Q3_K", "GGML_TYPE_Q4_K", "GGML_TYPE_Q5_K", "GGML_TYPE_Q6_K", "GGML_TYPE_IQ2_XXS", "GGML_TYPE_IQ2_XS", "GGML_TYPE_IQ2_S", "GGML_TYPE_IQ3_XXS", "GGML_TYPE_IQ3_S", "GGML_TYPE_IQ1_S", "GGML_TYPE_IQ4_NL", "GGML_TYPE_IQ4_XS", "GGML_TYPE_MXFP4", "GGML_TYPE_NVFP4" diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu new file mode 100644 index 000000000..8221e1d1e --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q2_0_ROCMFP2); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu new file mode 100644 index 000000000..2380af75c --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q3_0_ROCMFPX); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu new file mode 100644 index 000000000..94a2bb0f5 --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu index 59ce36fb1..adfa25312 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu @@ -35,7 +35,7 @@ static void top_k_cub(ggml_cuda_pool & pool, ncols, k, env)); } -#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE +#elif defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) // CUB_TOP_K_AVAILABLE static int next_power_of_2(int x) { int n = 1; @@ -69,7 +69,7 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { for (int i = 0; i < nrows; i++) { top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream); } -#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE +#elif defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) // CUB_TOP_K_AVAILABLE // Fall back to argsort + copy const int ncols_pad = next_power_of_2(ncols); const size_t shared_mem = ncols_pad * sizeof(int); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh index 07a28864f..6aa453c3b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh @@ -425,12 +425,14 @@ static __device__ __forceinline__ int rocmfpx_pack4_fp3_bits12_vec_cuda(const ui } static __device__ __forceinline__ int rocmfpx_pack4_fp3_vec_cuda(const uint8_t * qs, const int base) { - const char4 v = make_char4( - (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 0)*3, 3)), - (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 1)*3, 3)), - (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 2)*3, 3)), - (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 3)*3, 3))); - return *((const int *) &v); + // MMQ requests four-value groups, so the 12 packed bits span at most two + // adjacent bytes. Decode them with the same lookup used by MMVQ. + const int bit_pos = 3 * base; + const int byte_pos = bit_pos >> 3; + const uint32_t packed = (uint32_t) qs[byte_pos] | + ((uint32_t) qs[byte_pos + 1] << 8); + return rocmfpx_pack4_fp3_bits12_vec_cuda( + (packed >> (bit_pos & 7)) & 0x0fffu); } static __device__ __forceinline__ int rocmfpx_pack4_fp2_bits8_vec_cuda(const uint32_t bits8) { diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 2fb1df855..ce54b6978 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -1128,9 +1128,17 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "MOE_FUSED", "DS4_HC", + + "DS4_INDEXER_QAT", + + "DS4_INDEXER_SCORE", + + "DS4_INDEXER_MASK", + + "MUL_MAT_GROUPED_SRC", }; -static_assert(GGML_OP_COUNT == 100, "GGML_OP_COUNT != 100"); +static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1245,9 +1253,17 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "moe_fused(x)", "ds4_hc(x)", + + "ds4_indexer_qat(x)", + + "ds4_indexer_score(q,w,k)", + + "ds4_indexer_mask(x,topk)", + + "X*grouped(Y)", }; -static_assert(GGML_OP_COUNT == 100, "GGML_OP_COUNT != 100"); +static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -3337,10 +3353,41 @@ struct ggml_tensor * ggml_mul_mat( return result; } +struct ggml_tensor * ggml_mul_mat_grouped_src( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(b->type == GGML_TYPE_F32); + GGML_ASSERT(b->ne[2] > 1 && b->ne[3] == 1); + GGML_ASSERT(ggml_is_contiguous(b)); + GGML_ASSERT(a->ne[0] == b->ne[0] * b->ne[2]); + GGML_ASSERT(a->ne[2] == 1 && a->ne[3] == 1); + + const int64_t flat_k = b->ne[0] * b->ne[2]; + struct ggml_tensor * logical = ggml_view_2d( + ctx, b, flat_k, b->ne[1], + (size_t) flat_k * ggml_type_size(b->type), 0); + struct ggml_tensor * result = ggml_mul_mat(ctx, a, logical); + + result->op = GGML_OP_MUL_MAT_GROUPED_SRC; + ggml_set_op_params_i32(result, 14, (int32_t) b->ne[2]); + return result; +} + +bool ggml_mul_mat_is_grouped_src(const struct ggml_tensor * tensor) { + return tensor && tensor->op == GGML_OP_MUL_MAT_GROUPED_SRC; +} + +int64_t ggml_mul_mat_grouped_src_groups(const struct ggml_tensor * tensor) { + GGML_ASSERT(ggml_mul_mat_is_grouped_src(tensor)); + return ggml_get_op_params_i32(tensor, 14); +} + void ggml_mul_mat_set_prec( struct ggml_tensor * a, enum ggml_prec prec) { - GGML_ASSERT(a->op == GGML_OP_MUL_MAT); + GGML_ASSERT(a->op == GGML_OP_MUL_MAT || + a->op == GGML_OP_MUL_MAT_GROUPED_SRC); const int32_t prec_i32 = (int32_t) prec; @@ -5453,6 +5500,58 @@ void ggml_flash_attn_ext_set_prec( ggml_set_op_params_i32(a, 3, prec_i32); // scale is on first pos, max_bias on second } +void ggml_flash_attn_ext_set_ds4_sparse( + struct ggml_tensor * a, + int raw_rows, + int raw_window, + int keep_rows, + int block_size) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(raw_rows >= 0); + GGML_ASSERT(raw_window > 0 && raw_window <= 0xffff); + GGML_ASSERT(block_size > 0 && block_size <= 0xffff); + ggml_set_op_params_i32(a, 4, raw_rows); + ggml_set_op_params_i32(a, 5, keep_rows); + // Both values are DS4-local and bounded well below 16 bits. Packing them + // keeps the remaining op-parameter slots available for the exact RoPE + // metadata without growing ggml_tensor. + const uint32_t packed_layout = + ((uint32_t) raw_window << 16) | (uint32_t) block_size; + ggml_set_op_params_i32(a, 6, (int32_t) packed_layout); +} + +void ggml_flash_attn_ext_set_ds4_inverse_rope( + struct ggml_tensor * a, + int kv_start, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + int n_ctx_orig, + bool q_unrotated) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(kv_start >= 0); + // Bit 0: inverse RoPE on attention output. Bit 1: Q still needs forward + // RoPE. Both directions share the position and YaRN parameters below. + ggml_set_op_params_i32(a, 7, 1 | (q_unrotated ? 2 : 0)); + ggml_set_op_params_i32(a, 8, kv_start); + ggml_set_op_params_f32(a, 9, freq_base); + ggml_set_op_params_f32(a, 10, freq_scale); + ggml_set_op_params_f32(a, 11, ext_factor); + ggml_set_op_params_f32(a, 12, attn_factor); + ggml_set_op_params_f32(a, 13, beta_fast); + ggml_set_op_params_f32(a, 14, beta_slow); + ggml_set_op_params_i32(a, 15, n_ctx_orig); +} + +bool ggml_flash_attn_ext_is_ds4(const struct ggml_tensor * a) { + return a && a->op == GGML_OP_FLASH_ATTN_EXT && + (ggml_get_op_params_i32(a, 6) != 0 || + ggml_get_op_params_i32(a, 7) != 0); +} + enum ggml_prec ggml_flash_attn_ext_get_prec( const struct ggml_tensor * a) { GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); @@ -8212,13 +8311,21 @@ struct ggml_tensor * ggml_ds4_hc_pre( GGML_ASSERT(ggml_is_contiguous(base)); GGML_ASSERT(ggml_is_contiguous(hc_state)); GGML_ASSERT(n_hc > 0 && n_hc <= 8); + GGML_ASSERT(mix->ne[2] == 1 && mix->ne[3] == 1); + GGML_ASSERT(base->ne[1] == 1 && base->ne[2] == 1 && base->ne[3] == 1); + GGML_ASSERT(hc_state->ne[2] == 1 && hc_state->ne[3] == 1); const int64_t mix_dim = 2*(int64_t)n_hc + (int64_t)n_hc*n_hc; - GGML_ASSERT(ggml_nelements(mix) == mix_dim); - GGML_ASSERT(ggml_nelements(base) >= mix_dim); - GGML_ASSERT(ggml_nelements(hc_state) % n_hc == 0); - const int64_t n_embd = ggml_nelements(hc_state) / n_hc; - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd + mix_dim); + const int64_t n_tokens = mix->ne[1]; + GGML_ASSERT(mix->ne[0] == mix_dim); + GGML_ASSERT(n_tokens > 0); + GGML_ASSERT(base->ne[0] >= mix_dim); + GGML_ASSERT(hc_state->ne[1] == n_tokens); + GGML_ASSERT(hc_state->ne[0] % n_hc == 0); + const int64_t n_embd = hc_state->ne[0] / n_hc; + + struct ggml_tensor * result = n_tokens == 1 + ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd + mix_dim) + : ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd + mix_dim, n_tokens); result->op = GGML_OP_DS4_HC; result->src[0] = mix; result->src[1] = base; @@ -8244,15 +8351,27 @@ struct ggml_tensor * ggml_ds4_hc_post( GGML_ASSERT(split->type == GGML_TYPE_F32); GGML_ASSERT(ggml_is_contiguous(residual_hc)); GGML_ASSERT(ggml_is_contiguous(block_out)); - GGML_ASSERT(ggml_is_contiguous(split)); + // HC-pre stores the working vector and split parameters in one tensor. + // Its split view is dense within each token but inherits the parent row + // stride. The device kernel consumes nb[1] explicitly, so require the + // supported row-major layout rather than an unnecessarily packed tensor. + GGML_ASSERT(split->nb[0] == sizeof(float)); + GGML_ASSERT(split->nb[1] >= split->ne[0] * split->nb[0]); + GGML_ASSERT(split->ne[2] == 1 && split->ne[3] == 1); GGML_ASSERT(n_hc > 0 && n_hc <= 8); + GGML_ASSERT(residual_hc->ne[2] == 1 && residual_hc->ne[3] == 1); + GGML_ASSERT(block_out->ne[2] == 1 && block_out->ne[3] == 1); const int64_t mix_dim = 2*(int64_t)n_hc + (int64_t)n_hc*n_hc; - GGML_ASSERT(ggml_nelements(split) == mix_dim); - GGML_ASSERT(ggml_nelements(residual_hc) % n_hc == 0); - const int64_t n_embd = ggml_nelements(residual_hc) / n_hc; - GGML_ASSERT(ggml_nelements(block_out) == n_embd); - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t) n_embd * n_hc); + const int64_t n_tokens = residual_hc->ne[1]; + GGML_ASSERT(n_tokens > 0); + GGML_ASSERT(split->ne[0] == mix_dim && split->ne[1] == n_tokens); + GGML_ASSERT(residual_hc->ne[0] % n_hc == 0); + const int64_t n_embd = residual_hc->ne[0] / n_hc; + GGML_ASSERT(block_out->ne[0] == n_embd && block_out->ne[1] == n_tokens); + + struct ggml_tensor * result = n_tokens == 1 + ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t) n_embd * n_hc) + : ggml_new_tensor_2d(ctx, GGML_TYPE_F32, (int64_t) n_embd * n_hc, n_tokens); result->op = GGML_OP_DS4_HC; result->src[0] = residual_hc; result->src[1] = block_out; @@ -8313,12 +8432,19 @@ struct ggml_tensor * ggml_ds4_hc_out( GGML_ASSERT(ggml_is_contiguous(base)); GGML_ASSERT(ggml_is_contiguous(hc_state)); GGML_ASSERT(n_hc > 0 && n_hc <= 8); - GGML_ASSERT(ggml_nelements(mix) >= n_hc); - GGML_ASSERT(ggml_nelements(base) >= n_hc); - GGML_ASSERT(ggml_nelements(hc_state) % n_hc == 0); - const int64_t n_embd = ggml_nelements(hc_state) / n_hc; - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); + GGML_ASSERT(mix->ne[2] == 1 && mix->ne[3] == 1); + GGML_ASSERT(base->ne[1] == 1 && base->ne[2] == 1 && base->ne[3] == 1); + GGML_ASSERT(hc_state->ne[2] == 1 && hc_state->ne[3] == 1); + const int64_t n_tokens = hc_state->ne[1]; + GGML_ASSERT(n_tokens > 0); + GGML_ASSERT(mix->ne[0] >= n_hc && mix->ne[1] == n_tokens); + GGML_ASSERT(base->ne[0] >= n_hc); + GGML_ASSERT(hc_state->ne[0] % n_hc == 0); + const int64_t n_embd = hc_state->ne[0] / n_hc; + + struct ggml_tensor * result = n_tokens == 1 + ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd) + : ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); result->op = GGML_OP_DS4_HC; result->src[0] = mix; result->src[1] = base; @@ -8329,3 +8455,69 @@ struct ggml_tensor * ggml_ds4_hc_out( ggml_set_op_params_f32(result, 4, pre_scale); return result; } + +struct ggml_tensor * ggml_ds4_indexer_qat( + struct ggml_context * ctx, + struct ggml_tensor * input) { + GGML_ASSERT(input->type == GGML_TYPE_F32); + GGML_ASSERT(input->ne[0] == 128); + GGML_ASSERT(ggml_is_contiguous(input)); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, input); + result->op = GGML_OP_DS4_INDEXER_QAT; + result->src[0] = input; + return result; +} + +struct ggml_tensor * ggml_ds4_indexer_score( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * head_weights, + struct ggml_tensor * index_comp, + int kv_start, + int ratio) { + GGML_ASSERT(q->type == GGML_TYPE_F32 && q->ne[0] == 128); + GGML_ASSERT(head_weights->type == GGML_TYPE_F32); + GGML_ASSERT(index_comp->type == GGML_TYPE_F16 && index_comp->ne[0] == 128); + GGML_ASSERT(ggml_is_contiguous(q)); + GGML_ASSERT(ggml_is_contiguous(head_weights)); + GGML_ASSERT(ggml_is_contiguous(index_comp)); + GGML_ASSERT(head_weights->ne[0] == q->ne[1]); + GGML_ASSERT(head_weights->ne[1] == q->ne[2]); + GGML_ASSERT(q->ne[3] == 1); + GGML_ASSERT(head_weights->ne[2] == 1 && head_weights->ne[3] == 1); + GGML_ASSERT(index_comp->ne[2] == 1 && index_comp->ne[3] == 1); + GGML_ASSERT(kv_start >= 0 && ratio > 0); + + struct ggml_tensor * result = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, index_comp->ne[1], q->ne[2]); + result->op = GGML_OP_DS4_INDEXER_SCORE; + result->src[0] = q; + result->src[1] = head_weights; + result->src[2] = index_comp; + ggml_set_op_params_i32(result, 0, kv_start); + ggml_set_op_params_i32(result, 1, ratio); + return result; +} + +struct ggml_tensor * ggml_ds4_indexer_mask( + struct ggml_context * ctx, + struct ggml_tensor * base_mask, + struct ggml_tensor * selected, + int raw_rows) { + GGML_ASSERT(base_mask->type == GGML_TYPE_F32); + GGML_ASSERT(selected->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(base_mask)); + GGML_ASSERT(ggml_is_contiguous(selected)); + GGML_ASSERT(raw_rows >= 0 && raw_rows <= base_mask->ne[0]); + GGML_ASSERT(selected->ne[1] == base_mask->ne[1]); + GGML_ASSERT(selected->ne[2] == base_mask->ne[2]); + GGML_ASSERT(selected->ne[3] == base_mask->ne[3]); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, base_mask); + result->op = GGML_OP_DS4_INDEXER_MASK; + result->src[0] = base_mask; + result->src[1] = selected; + ggml_set_op_params_i32(result, 0, raw_rows); + return result; +} diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index df96a1410..e9ee327af 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -48,6 +48,13 @@ std::unique_ptr create_backend(const BackendArgs & args) { std::fprintf(stderr, "[backend_factory] detected arch=%s\n", arch.c_str()); + if (args.ds4_prefill_mode_set && arch != "deepseek4") { + std::fprintf(stderr, + "[backend_factory] --ds4-prefill is only valid for deepseek4 models " + "(detected %s)\n", arch.c_str()); + return nullptr; + } + if (arch == "qwen35") { if (args.device.is_layer_split()) { Qwen35LayerSplitAdapterConfig cfg; @@ -219,6 +226,17 @@ std::unique_ptr create_backend(const BackendArgs & args) { args.device.backend == PlacementBackend::Auto ? compiled_placement_backend() : args.device.backend; + if (prefill_attention_mode_is_approximate(args.ds4_prefill_mode) && + (target_backend != PlacementBackend::Hip || + args.device.is_layer_split() || + args.remote_target_shard.enabled())) { + std::fprintf(stderr, + "[backend_factory] DS4 %s prefill requires a single local HIP " + "target; use --ds4-prefill exact for split, remote, or CUDA " + "placement\n", + prefill_attention_mode_name(args.ds4_prefill_mode)); + return nullptr; + } // HIP single-device launches cannot rely on the CUDA/Halo auto-split // path; use the single-backend loader, which can fall back to hybrid @@ -234,6 +252,7 @@ std::unique_ptr create_backend(const BackendArgs & args) { cfg.chunk = args.chunk; cfg.expert_top_k = args.ds4_expert_top_k; cfg.fused_decode = args.ds4_fused_decode; + cfg.prefill_mode = args.ds4_prefill_mode; auto backend = std::make_unique(cfg); if (!backend->init()) { diff --git a/server/src/common/backend_factory.h b/server/src/common/backend_factory.h index 75c2053f7..7b468000c 100644 --- a/server/src/common/backend_factory.h +++ b/server/src/common/backend_factory.h @@ -15,6 +15,7 @@ #include "placement/placement_config.h" #include "placement/remote_draft_config.h" #include "placement/remote_target_shard_config.h" +#include "prefill_attention_mode.h" #include #include @@ -43,7 +44,9 @@ struct BackendArgs { int stream_fd = -1; // Chunked prefill - int chunk = 512; + int chunk = 512; + PrefillAttentionMode ds4_prefill_mode = PrefillAttentionMode::Exact; + bool ds4_prefill_mode_set = false; // deepseek4-specific decode options int ds4_expert_top_k = 0; // 0 = model default diff --git a/server/src/common/prefill_attention_mode.h b/server/src/common/prefill_attention_mode.h new file mode 100644 index 000000000..6172f3a5d --- /dev/null +++ b/server/src/common/prefill_attention_mode.h @@ -0,0 +1,28 @@ +#pragma once + +namespace dflash::common { + +// Numerics/performance policy for model-specific prefill attention. Exact is +// the tokenwise reference. Dense changes the execution/reduction topology and +// Sparse additionally prunes model-specific cache entries; neither optimized +// mode promises byte-identical logits or generated tokens. +enum class PrefillAttentionMode { + Exact, + Dense, + Sparse, +}; + +inline const char * prefill_attention_mode_name(PrefillAttentionMode mode) { + switch (mode) { + case PrefillAttentionMode::Exact: return "exact"; + case PrefillAttentionMode::Dense: return "dense"; + case PrefillAttentionMode::Sparse: return "sparse"; + } + return "unknown"; +} + +inline bool prefill_attention_mode_is_approximate(PrefillAttentionMode mode) { + return mode != PrefillAttentionMode::Exact; +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 66877ca76..2b0382e38 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -90,6 +90,7 @@ static void add_step_tel(DeepSeek4StepTelemetry & dst, const DeepSeek4StepTeleme dst.attn_compute_us += src.attn_compute_us; dst.attn_read_us += src.attn_read_us; dst.full_graph_build_us += src.full_graph_build_us; + dst.full_graph_alloc_us += src.full_graph_alloc_us; dst.full_graph_set_us += src.full_graph_set_us; dst.full_graph_compute_us += src.full_graph_compute_us; dst.full_graph_read_us += src.full_graph_read_us; @@ -132,7 +133,7 @@ static void log_step_tel(const char * phase, 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 " + "full_build=%.1fms full_alloc=%.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 " @@ -142,7 +143,7 @@ static void log_step_tel(const char * phase, "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_build_us), ms(t.full_graph_alloc_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), @@ -581,6 +582,38 @@ DeepSeek4Backend::~DeepSeek4Backend() { shutdown(); } +bool DeepSeek4Backend::requires_monolithic_model() const { + // Heterogeneous execution deliberately keeps routed experts split across + // the R9700 and Strix. Only an explicit diagnostic override may request + // the ~80-GiB monolithic HIP allocation. + return env_flag_enabled("DFLASH_DS4_FORCE_FULL_LOAD"); +} + +bool DeepSeek4Backend::validate_prefill_mode() const { + if (cfg_.prefill_mode == PrefillAttentionMode::Exact) { + return true; + } + const PlacementBackend target_backend = + cfg_.device.backend == PlacementBackend::Auto + ? compiled_placement_backend() + : cfg_.device.backend; + if (target_backend != PlacementBackend::Hip || + cfg_.device.is_layer_split()) { + std::fprintf(stderr, + "[deepseek4] %s prefill requires a single HIP target\n", + prefill_attention_mode_name(cfg_.prefill_mode)); + return false; + } + if (w_.moe_hybrid || moe_hybrid_) { + std::fprintf(stderr, + "[deepseek4] %s prefill requires every expert to be resident; " + "the selected placement has cold experts\n", + prefill_attention_mode_name(cfg_.prefill_mode)); + return false; + } + return true; +} + bool DeepSeek4Backend::load_model() { const PlacementBackend target_backend = cfg_.device.backend == PlacementBackend::Auto @@ -591,7 +624,7 @@ bool DeepSeek4Backend::load_model() { // a managed ~80 GiB allocation can stall or be killed on integrated UMA // systems before we ever reach the existing OOM fallback path. if (target_backend == PlacementBackend::Hip && - env_flag_enabled("DFLASH_DS4_FORCE_FULL_LOAD")) { + requires_monolithic_model()) { std::fprintf(stderr, "[deepseek4] HIP full-model load explicitly enabled\n"); if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { @@ -784,12 +817,22 @@ bool DeepSeek4Backend::init() { if (!load_model()) { return false; } + if (!validate_prefill_mode()) { + return false; + } + if (prefill_attention_mode_is_approximate(cfg_.prefill_mode)) { + std::fprintf(stderr, + "[deepseek4] warning: %s prefill is approximate and may change " + "generated tokens; use --ds4-prefill exact for reference output\n", + prefill_attention_mode_name(cfg_.prefill_mode)); + } 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; } + cache_.prefill_mode = cfg_.prefill_mode; if (env_flag_enabled("DFLASH_DS4_MOE_TP") && !init_moe_tensor_parallel()) { return false; @@ -833,9 +876,10 @@ bool DeepSeek4Backend::init() { w_.routed_expert_top_k > 0 ? w_.routed_expert_top_k : w_.n_expert_used; std::fprintf(stderr, "[deepseek4] initialized: %d layers, ctx=%d, %d experts " - "(%d/%d routed), fused_decode=%s%s\n", + "(%d/%d routed), fused_decode=%s, prefill=%s%s\n", w_.n_layer, max_ctx, w_.n_expert, active_experts, w_.n_expert_used, w_.fused_decode ? "on" : "off", + prefill_attention_mode_name(cfg_.prefill_mode), moe_hybrid_ ? " [hybrid]" : ""); if (env_flag_enabled("DFLASH_DS4_SPEC")) { @@ -1193,6 +1237,13 @@ bool DeepSeek4Backend::unpark(const std::string & what) { std::printf("[deepseek4] target unparked (VRAM restored)\n"); std::fflush(stdout); } + if (!validate_prefill_mode()) { + free_deepseek4_weights(w_); + stream_engine_.destroy(); + moe_hybrid_.reset(); + moe_placement_ = {}; + return false; + } if (want_draft && spec_drafter_parked_) { if (parked_) { @@ -1205,17 +1256,31 @@ bool DeepSeek4Backend::unpark(const std::string & what) { return false; } } + cache_.prefill_mode = cfg_.prefill_mode; return true; } int DeepSeek4Backend::do_prefill(const std::vector & tokens, const DaemonIO & io, int kv_offset) { - // Keep server prefill token-at-a-time for established numerics and bounded - // dynamic-graph workspace. The lower-level layer-range API independently - // splits arbitrary multi-token callers at every compressor boundary. - constexpr int chunk = 1; + // The all-hot layer-range path supports causal chunked prefill. The + // optimized graph snapshots the previous raw SWA window, attends over + // that snapshot plus the current ubatch, and commits only the final SWA + // tail. Learned compressor boundaries are emitted inside the same graph. + // + // Mixed hot/cold hybrid execution still has single-token HC semantics, so + // retain the reference path there. --chunk 1 is the explicit fallback. + const int requested_chunk = cfg_.chunk > 0 ? cfg_.chunk : w_.n_swa; const int n_total = (int)tokens.size(); + // Bound the layer-major graph to the topology validated by the prefill + // kernels. Smaller tail chunks use the same scheduler or its reference + // fallback. + const int layer_major_cap = DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS; + const int chunk = (moe_hybrid_ || + !prefill_attention_mode_is_approximate(cfg_.prefill_mode)) + ? 1 + : std::max(1, std::min(requested_chunk, + layer_major_cap)); int pos = kv_offset; // New sequence: clear the cache buffer so compressor state double-buffers // and compressed-KV rows start from zeros, exactly like a fresh server. @@ -1285,7 +1350,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, 0, w_.n_layer, &logits, tokens.data() + i, timing ? &step_tel : nullptr, - /*allow_decode_graph_reuse=*/true, hp, + cfg_.prefill_mode != PrefillAttentionMode::Sparse, hp, moe_hybrid_.get(), expert_runtime_.compute ? &expert_runtime_ : nullptr, routing_stats_.get()); @@ -1303,7 +1368,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, 0, w_.n_layer, &logits, tokens.data() + i, timing ? &step_tel : nullptr, - /*allow_decode_graph_reuse=*/true, hp); + cfg_.prefill_mode != PrefillAttentionMode::Sparse, hp); } if (ok && hp && !spec_cap.empty()) { const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 809a1b057..09c9faec1 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -111,6 +111,8 @@ class DeepSeek4Backend : public ModelBackend { bool load_model(); bool init_hybrid_model(); bool init_moe_tensor_parallel(); + bool requires_monolithic_model() const; + bool validate_prefill_mode() const; bool compute_uniform_hybrid_placement(const DeepSeek4Weights & w, int max_ctx, MoeHybridPlacement & out, diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc index 596656d05..ab221605c 100644 --- a/server/src/deepseek4/deepseek4_fused_verify.inc +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -842,7 +842,8 @@ static bool ds4_build_fused_verify_graph( ggml_tensor * hcol = ggml_view_2d( ctx, selected, route_width, 1, selected->nb[1], (size_t) t * selected->nb[1]); - ggml_tensor * fo = ds4_build_hash_routed_ffn(ctx, w, L, ggml_cont(ctx, fcol), hcol); + ggml_tensor * fo = ds4_build_hash_routed_ffn( + ctx, w, L, ggml_cont(ctx, fcol), hcol, 1); if (!fo) { ffn_out = nullptr; break; } ffn_out = ffn_out ? ggml_concat(ctx, ffn_out, fo, 1) : fo; } diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 328c70eea..6d309df47 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -26,6 +26,7 @@ extern "C" void ggml_backend_cuda_set_graphs_disabled_override(bool disabled); #include +#include #include #include #include @@ -215,6 +216,13 @@ struct DeepSeek4F32ArrayBinding { std::vector values; }; +// Attention implementation selected by the DS4 prefill scheduler. Decode +// retains the established explicit reduction path. +enum class DeepSeek4AttentionImpl { + Explicit, + DenseFlash, + SparseFlash, +}; static ggml_tensor * build_rms_norm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * weight, float eps); static ggml_tensor * build_clamped_swiglu(ggml_context * ctx, @@ -380,6 +388,10 @@ struct DeepSeek4LayerRangeScratch { std::vector final_embd; std::vector hash_expert_ids; + void clear() { + *this = {}; + } + void ensure(const ggml_context * ctx, int tokens, int embd, @@ -636,6 +648,320 @@ int deepseek4_safe_compressor_batch_tokens(const DeepSeek4Weights & w, return std::max(1, safe); } +// Build an exact multi-token compressor update for prefill. Complete windows +// are pooled as one batched tensor, so a 2K ubatch does not create hundreds of +// serial softmax subgraphs. The state is assembled functionally from an +// initial snapshot and written back once, avoiding persistent-buffer races. +static bool build_compressor_prefill( + ggml_context * ctx, + ggml_cgraph * gf, + ggml_tensor * cur_all, + ggml_tensor * ape, + ggml_tensor * kv_proj, + ggml_tensor * gate_proj, + ggml_tensor * norm_weight, + DeepSeek4CompressorState & state, + ggml_tensor * comp_cache, + int ratio, + int head_dim, + int kv_start, + int n_tokens, + int n_rot, + float rms_eps, + float compress_rope_freq_base, + float rope_scale_factor, + float rope_yarn_beta_fast, + float rope_yarn_beta_slow, + int rope_orig_ctx, + std::vector & i64_array_inputs, + std::vector & i32_array_inputs, + ggml_tensor ** comp_cache_source_out, + bool indexer_qat) { + if (!cur_all || n_tokens <= 1 || + n_tokens > DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS || + (ratio != 4 && ratio != 128)) { + return false; + } + + const int coff = ratio == 4 ? 2 : 1; + const int comp_width = coff * head_dim; + const int n_state_rows = ratio == 4 ? 2 * ratio : ratio; + + struct Pair { + ggml_tensor * kv = nullptr; + ggml_tensor * score = nullptr; + }; + + auto view_cols = [&](ggml_tensor * src, int width, int first, int count) { + GGML_ASSERT(src && count > 0); + return ggml_cont(ctx, ggml_view_2d(ctx, src, width, count, src->nb[1], + (size_t) first * src->nb[1])); + }; + auto view_pair_cols = [&](const Pair & src, int first, int count) { + return Pair { view_cols(src.kv, comp_width, first, count), + view_cols(src.score, comp_width, first, count) }; + }; + auto concat_tensors = [&](const std::vector & parts) { + GGML_ASSERT(!parts.empty()); + ggml_tensor * out = parts[0]; + for (size_t i = 1; i < parts.size(); ++i) { + out = ggml_concat(ctx, out, parts[i], 1); + } + return ggml_cont(ctx, out); + }; + auto concat_pairs = [&](const std::vector & parts) { + std::vector kv_parts; + std::vector score_parts; + kv_parts.reserve(parts.size()); + score_parts.reserve(parts.size()); + for (const Pair & part : parts) { + kv_parts.push_back(part.kv); + score_parts.push_back(part.score); + } + return Pair { concat_tensors(kv_parts), concat_tensors(score_parts) }; + }; + auto replace_span = [&](const Pair & base, + int first, + const Pair & replacement, + int count, + int width) { + std::vector parts; + if (first > 0) parts.push_back(view_pair_cols(base, 0, first)); + parts.push_back(replacement); + if (first + count < width) { + parts.push_back(view_pair_cols(base, first + count, width - first - count)); + } + return concat_pairs(parts); + }; + + // Project the entire chunk once and add the position-addressed APE score. + Pair projected; + projected.kv = ggml_mul_mat(ctx, kv_proj, cur_all); + projected.score = ggml_mul_mat(ctx, gate_proj, cur_all); + ggml_tensor * ape_rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); + ggml_set_input(ape_rows); + std::vector ape_values((size_t) n_tokens); + for (int i = 0; i < n_tokens; ++i) { + ape_values[(size_t) i] = (kv_start + i) % ratio; + } + i32_array_inputs.push_back({ape_rows, std::move(ape_values)}); + ggml_tensor * ape_cols = ggml_get_rows(ctx, ape, ape_rows); + projected.score = ggml_add(ctx, projected.score, + ds4_cast_if_needed(ctx, ape_cols, GGML_TYPE_F32)); + + // Snapshot before the single writeback. Both compressor output and final + // state depend on these copies, forcing reads to complete before mutation. + Pair initial { ggml_cont(ctx, state.state_kv), + ggml_cont(ctx, state.state_score) }; + ggml_build_forward_expand(gf, initial.kv); + ggml_build_forward_expand(gf, initial.score); + + ggml_tensor * pooled_batch = nullptr; + std::vector comp_rows; + std::vector comp_positions; + + auto pool_groups = [&](ggml_tensor * values_kv, + ggml_tensor * values_score, + int rows, + int groups) { + GGML_ASSERT(values_kv && values_score && rows > 0 && groups > 0); + ggml_tensor * kv3 = ggml_reshape_3d(ctx, values_kv, + head_dim, rows, groups); + ggml_tensor * score3 = ggml_reshape_3d(ctx, values_score, + head_dim, rows, groups); + ggml_tensor * score_t = ggml_cont( + ctx, ggml_permute(ctx, score3, 1, 0, 2, 3)); + ggml_tensor * kv_t = ggml_cont( + ctx, ggml_permute(ctx, kv3, 1, 0, 2, 3)); + ggml_tensor * probs_t = ggml_soft_max(ctx, score_t); + ggml_tensor * weighted_t = ggml_mul(ctx, probs_t, kv_t); + ggml_tensor * pooled_sum = ggml_sum_rows(ctx, weighted_t); + ggml_tensor * pooled = ggml_reshape_2d( + ctx, ggml_cont(ctx, pooled_sum), head_dim, groups); + pooled = ggml_cont(ctx, pooled); + pooled = build_rms_norm(ctx, pooled, norm_weight, rms_eps); + return ggml_reshape_2d(ctx, pooled, head_dim, groups); + }; + + Pair final_state; + if (ratio == 4) { + const int pos_mod = kv_start % ratio; + const int first_count = std::min(ratio - pos_mod, n_tokens); + int consumed = 0; + std::vector complete_parts; + + if (n_tokens >= ratio - pos_mod) { + Pair current = view_pair_cols(initial, ratio, ratio); + Pair first_span = view_pair_cols(projected, 0, first_count); + complete_parts.push_back(replace_span( + current, pos_mod, first_span, first_count, ratio)); + consumed = first_count; + + const int complete_tail = ((n_tokens - consumed) / ratio) * ratio; + if (complete_tail > 0) { + complete_parts.push_back(view_pair_cols( + projected, consumed, complete_tail)); + consumed += complete_tail; + } + } + + if (complete_parts.empty()) { + Pair prev = view_pair_cols(initial, 0, ratio); + Pair current = view_pair_cols(initial, ratio, ratio); + current = replace_span(current, pos_mod, projected, + n_tokens, ratio); + final_state = concat_pairs({prev, current}); + } else { + Pair complete = concat_pairs(complete_parts); + const int groups = (int) complete.kv->ne[1] / ratio; + GGML_ASSERT(groups > 0); + + Pair previous = view_pair_cols(initial, 0, ratio); + if (groups > 1) { + previous = concat_pairs({ + previous, + view_pair_cols(complete, 0, (groups - 1) * ratio), + }); + } + + auto select_half = [&](ggml_tensor * src, int half) { + ggml_tensor * src3 = ggml_reshape_3d( + ctx, src, comp_width, ratio, groups); + return ggml_cont(ctx, ggml_view_3d( + ctx, src3, head_dim, ratio, groups, + src3->nb[1], src3->nb[2], + (size_t) half * head_dim * src3->nb[0])); + }; + ggml_tensor * selected_kv = ggml_concat( + ctx, select_half(previous.kv, 0), + select_half(complete.kv, 1), 1); + ggml_tensor * selected_score = ggml_concat( + ctx, select_half(previous.score, 0), + select_half(complete.score, 1), 1); + pooled_batch = pool_groups( + ggml_cont(ctx, selected_kv), + ggml_cont(ctx, selected_score), 2 * ratio, groups); + + const int first_boundary = kv_start + first_count - 1; + for (int g = 0; g < groups; ++g) { + const int boundary = first_boundary + g * ratio; + const int64_t comp_row = boundary / ratio; + GGML_ASSERT(comp_row >= 0 && comp_row < comp_cache->ne[1]); + comp_rows.push_back(comp_row); + comp_positions.push_back(boundary + 1 - ratio); + } + + Pair last_complete = view_pair_cols( + complete, (groups - 1) * ratio, ratio); + Pair current = last_complete; + const int tail = n_tokens - consumed; + if (tail > 0) { + current = replace_span( + current, 0, view_pair_cols(projected, consumed, tail), + tail, ratio); + } + final_state = concat_pairs({last_complete, current}); + } + } else { + const int pos_mod = kv_start % ratio; + const int to_boundary = ratio - pos_mod; + if (n_tokens < to_boundary) { + final_state = replace_span(initial, pos_mod, projected, n_tokens, ratio); + } else { + std::vector first_parts; + if (pos_mod > 0) { + first_parts.push_back(view_pair_cols(initial, 0, pos_mod)); + } + first_parts.push_back(view_pair_cols(projected, 0, to_boundary)); + Pair first_complete = concat_pairs(first_parts); + + int consumed = to_boundary; + const int complete_tail = ((n_tokens - consumed) / ratio) * ratio; + Pair complete = first_complete; + if (complete_tail > 0) { + complete = concat_pairs({ + first_complete, + view_pair_cols(projected, consumed, complete_tail), + }); + consumed += complete_tail; + } + const int groups = (int) complete.kv->ne[1] / ratio; + pooled_batch = pool_groups(complete.kv, complete.score, + ratio, groups); + const int first_boundary = kv_start + to_boundary - 1; + for (int g = 0; g < groups; ++g) { + const int boundary = first_boundary + g * ratio; + const int64_t comp_row = boundary / ratio; + GGML_ASSERT(comp_row >= 0 && comp_row < comp_cache->ne[1]); + comp_rows.push_back(comp_row); + comp_positions.push_back(boundary + 1 - ratio); + } + + Pair last_complete = view_pair_cols( + complete, (groups - 1) * ratio, ratio); + const int tail = n_tokens - consumed; + if (tail > 0) { + final_state = replace_span( + last_complete, 0, + view_pair_cols(projected, consumed, tail), tail, ratio); + } else { + final_state = last_complete; + } + } + } + + // Persist the exact sequential state using unique row indices. + ggml_tensor * state_rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_state_rows); + ggml_set_input(state_rows); + std::vector state_row_values((size_t) n_state_rows); + for (int i = 0; i < n_state_rows; ++i) state_row_values[(size_t) i] = i; + i64_array_inputs.push_back({state_rows, std::move(state_row_values)}); + final_state.kv = ggml_cont(ctx, final_state.kv); + final_state.score = ggml_cont(ctx, final_state.score); + ggml_tensor * state_kv_source = ggml_set_rows(ctx, state.state_kv, + final_state.kv, state_rows); + ggml_tensor * state_score_source = ggml_set_rows(ctx, state.state_score, + final_state.score, state_rows); + ggml_build_forward_expand(gf, state_kv_source); + ggml_build_forward_expand(gf, state_score_source); + + ggml_tensor * comp_cache_source = comp_cache; + if (pooled_batch) { + ggml_tensor * pooled = pooled_batch; + const int n_pooled = (int) comp_positions.size(); + ggml_tensor * comp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, + n_pooled); + ggml_set_input(comp_pos); + i32_array_inputs.push_back({comp_pos, std::move(comp_positions)}); + + const float rope_scale = rope_scale_factor > 0.0f + ? (1.0f / rope_scale_factor) : 1.0f; + float rope_attn = 1.0f; + if (rope_scale > 0.0f) { + rope_attn /= (1.0f + 0.1f * logf(1.0f / rope_scale)); + } + pooled = build_tail_rope_2d(ctx, pooled, comp_pos, n_rot, head_dim, + n_pooled, + compress_rope_freq_base, rope_scale, + 1.0f, rope_attn, + rope_yarn_beta_fast, rope_yarn_beta_slow, + rope_orig_ctx); + pooled = ggml_cont(ctx, pooled); + if (indexer_qat) { + pooled = ggml_ds4_indexer_qat(ctx, pooled); + } + + ggml_tensor * comp_row_tensor = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, (int64_t) comp_rows.size()); + ggml_set_input(comp_row_tensor); + i64_array_inputs.push_back({comp_row_tensor, std::move(comp_rows)}); + comp_cache_source = ggml_set_rows(ctx, comp_cache, pooled, comp_row_tensor); + ggml_build_forward_expand(gf, comp_cache_source); + } + if (comp_cache_source_out) *comp_cache_source_out = comp_cache_source; + return true; +} + static void build_compressor_step( ggml_context * ctx, ggml_cgraph * gf, @@ -666,12 +992,27 @@ static void build_compressor_step( ggml_tensor * flush_rows_inp = nullptr, ggml_tensor * cur_all = nullptr, int n_tokens_all = 1, - int kv_start_all = -1) { + int kv_start_all = -1, + bool indexer_qat = false) { if (!gf || !cur_last || !ape || !kv_proj || !gate_proj || !norm_weight || !state.state_kv || !state.state_score || !comp_cache || ratio <= 0) { return; } + // Multi-token speculative verification uses the boundary-split path below. + // The layer-major prefill scheduler only enters here for wider batches. + if (cur_all && n_tokens_all > 4 && !state_rows_inp && kv_start_all >= 0 && + build_compressor_prefill(ctx, gf, cur_all, ape, kv_proj, gate_proj, + norm_weight, state, comp_cache, ratio, head_dim, + kv_start_all, n_tokens_all, n_rot, rms_eps, + compress_rope_freq_base, rope_scale_factor, + rope_yarn_beta_fast, rope_yarn_beta_slow, + rope_orig_ctx, i64_array_inputs, + i32_array_inputs, comp_cache_source_out, + indexer_qat)) { + return; + } + // DS4 compression: internal width = coff * head_dim (2x for ratio-4, 1x for ratio-128) const int coff = (ratio == 4) ? 2 : 1; const int comp_width = coff * head_dim; @@ -785,13 +1126,11 @@ static void build_compressor_step( } if (batched_rows && batched_b < 0) { - // no boundary inside the batch: state rows written, nothing to emit + // State rows were written, but this batch did not complete a window. return; } - if (!batched_rows && !flush_rows_inp && ((token_pos + 1) % ratio) != 0) { - // Legacy per-layer graphs only pool at flush boundaries. The fused - // stable-topology graph (flush_rows_inp set) pools every step; the - // partial result lands on the masked running comp row. + if (!batched_rows && ((token_pos + 1) % ratio) != 0) { + // Per-layer graphs only pool at flush boundaries. return; } @@ -859,6 +1198,9 @@ static void build_compressor_step( pooled = build_tail_rope_2d(ctx, pooled, comp_pos, n_rot, head_dim, 1, compress_rope_freq_base, rope_scale, 1.0f, rope_attn, rope_yarn_beta_fast, rope_yarn_beta_slow, rope_orig_ctx); + if (indexer_qat) { + pooled = ggml_ds4_indexer_qat(ctx, ggml_cont(ctx, pooled)); + } ggml_tensor * pooled_f16 = ggml_cast(ctx, pooled, GGML_TYPE_F16); const int comp_row = token_pos / ratio; @@ -882,34 +1224,46 @@ static void build_compressor_step( if (batched_rows) { if (ratio == 4) { - // completed window: rotate current half -> prev half, reading - // through the span-A writes so ordering is explicit. + // Rotate the completed current window into the previous half. + // Reading through the first span makes the dependency explicit. for (int r = 0; r < ratio; ++r) { - ggml_tensor * src_kv = ggml_view_2d(ctx, state_kv_source, comp_width, 1, - state_kv_source->nb[1], - (size_t)(ratio + r) * state_kv_source->nb[1]); - ggml_tensor * dst_kv = ggml_view_2d(ctx, state.state_kv, comp_width, 1, - state.state_kv->nb[1], - (size_t) r * state.state_kv->nb[1]); + ggml_tensor * src_kv = ggml_view_2d( + ctx, state_kv_source, comp_width, 1, + state_kv_source->nb[1], + (size_t) (ratio + r) * state_kv_source->nb[1]); + ggml_tensor * dst_kv = ggml_view_2d( + ctx, state.state_kv, comp_width, 1, + state.state_kv->nb[1], + (size_t) r * state.state_kv->nb[1]); ggml_build_forward_expand(gf, ggml_cpy(ctx, src_kv, dst_kv)); - ggml_tensor * src_sc = ggml_view_2d(ctx, state_score_source, comp_width, 1, - state_score_source->nb[1], - (size_t)(ratio + r) * state_score_source->nb[1]); - ggml_tensor * dst_sc = ggml_view_2d(ctx, state.state_score, comp_width, 1, - state.state_score->nb[1], - (size_t) r * state.state_score->nb[1]); + + ggml_tensor * src_sc = ggml_view_2d( + ctx, state_score_source, comp_width, 1, + state_score_source->nb[1], + (size_t) (ratio + r) * state_score_source->nb[1]); + ggml_tensor * dst_sc = ggml_view_2d( + ctx, state.state_score, comp_width, 1, + state.state_score->nb[1], + (size_t) r * state.state_score->nb[1]); ggml_build_forward_expand(gf, ggml_cpy(ctx, src_sc, dst_sc)); } } if (batched_nB > 0) { - ggml_tensor * kv_v = ggml_cont(ctx, ggml_view_2d(ctx, batched_kv_all, comp_width, batched_nB, - batched_kv_all->nb[1], (size_t) batched_span_off * batched_kv_all->nb[1])); - ggml_tensor * sc_v = ggml_cont(ctx, ggml_view_2d(ctx, batched_sc_all, comp_width, batched_nB, - batched_sc_all->nb[1], (size_t) batched_span_off * batched_sc_all->nb[1])); - ggml_tensor * rows_v = ggml_view_1d(ctx, state_rows_inp, batched_nB, - (size_t) batched_span_off * state_rows_inp->nb[0]); - ggml_build_forward_expand(gf, ggml_set_rows(ctx, state.state_kv, kv_v, rows_v)); - ggml_build_forward_expand(gf, ggml_set_rows(ctx, state.state_score, sc_v, rows_v)); + ggml_tensor * kv_v = ggml_cont(ctx, ggml_view_2d( + ctx, batched_kv_all, comp_width, batched_nB, + batched_kv_all->nb[1], + (size_t) batched_span_off * batched_kv_all->nb[1])); + ggml_tensor * sc_v = ggml_cont(ctx, ggml_view_2d( + ctx, batched_sc_all, comp_width, batched_nB, + batched_sc_all->nb[1], + (size_t) batched_span_off * batched_sc_all->nb[1])); + ggml_tensor * rows_v = ggml_view_1d( + ctx, state_rows_inp, batched_nB, + (size_t) batched_span_off * state_rows_inp->nb[0]); + ggml_build_forward_expand( + gf, ggml_set_rows(ctx, state.state_kv, kv_v, rows_v)); + ggml_build_forward_expand( + gf, ggml_set_rows(ctx, state.state_score, sc_v, rows_v)); } return; } @@ -968,10 +1322,12 @@ static void build_indexer_compressor_step( ggml_tensor * comp_pos_inp, std::vector & i64_array_inputs, std::vector & i32_array_inputs, + ggml_tensor ** index_comp_cache_source_out = nullptr, ggml_tensor * flush_rows_inp = nullptr, ggml_tensor * cur_all = nullptr, int n_tokens_all = 1, - int kv_start_all = -1) { + int kv_start_all = -1, + bool indexer_qat = false) { build_compressor_step(ctx, gf, cur_last, L.indexer_compressor_ape, L.indexer_compressor_kv, @@ -995,19 +1351,24 @@ static void build_indexer_compressor_step( comp_pos_inp, i64_array_inputs, i32_array_inputs, - nullptr, + index_comp_cache_source_out, flush_rows_inp, cur_all, n_tokens_all, - kv_start_all); + kv_start_all, + indexer_qat); } static int ds4_comp_rows_used(const ggml_tensor * comp_cache, int n_cached, int ratio, int token_pos) { if (!comp_cache || ratio <= 0) { return 0; } - const int grew_this_step = ((token_pos + 1) % ratio) == 0 ? 1 : 0; - return std::min(n_cached + grew_this_step, (int) comp_cache->ne[1]); + // n_cached is the committed count before this graph. A multi-token + // prefill graph may cross several compressor boundaries, so derive the + // live count from the query position rather than adding at most one row. + const int through_position = (token_pos + 1) / ratio; + return std::min(std::max(n_cached, through_position), + (int) comp_cache->ne[1]); } // Round the live compressed-row count up to a fixed stride so the fused decode @@ -1022,88 +1383,101 @@ static int ds4_padded_comp_rows(int n_comp, int cap) { return padded < cap ? padded : cap; } -static ggml_tensor * build_indexer_score( +static ggml_tensor * build_indexer_topk( ggml_context * ctx, - ggml_tensor * qr_norm_last, // [n_lora_q, 1] - ggml_tensor * cur_last, // [n_embd, 1] + ggml_tensor * qr_norm, // [n_lora_q, n_tokens] + ggml_tensor * cur, // [n_embd, n_tokens] const DeepSeek4Weights & w, const DeepSeek4Layer & L, - const DeepSeek4LayerCache & lc, - int token_pos, - std::vector & i32_inputs) { - const int n_comp = ds4_comp_rows_used(lc.index_comp_kv, lc.n_index_comp, 4, token_pos); - if (!qr_norm_last || !cur_last || !L.indexer_attn_q_b || !L.indexer_proj || - !lc.index_comp_kv || n_comp <= 0) { + ggml_tensor * index_comp_source, + int n_comp, + int kv_start, + int n_tokens, + ggml_tensor * rope_pos, + std::vector & i32_array_inputs) { + if (!qr_norm || !cur || !L.indexer_attn_q_b || !L.indexer_proj || + !index_comp_source || !rope_pos || n_tokens <= 0 || + n_comp <= w.n_indexer_top_k) { return nullptr; } const int n_indexer_head = w.n_indexer_head; const int head_dim = w.n_indexer_head_dim; + const int top_k = std::min(n_comp, w.n_indexer_top_k); + // A token with <= top_k visible compressed rows needs no ranking: selecting + // [0,top_k) and retaining the ordinary causal mask is exactly equivalent. + // Score only the suffix beginning with the first token that can see row + // top_k. For a zero-prefix ratio-4 2K request this shrinks 2164 score rows + // to just 113. + const int first_scored = std::max( + 0, std::min(n_tokens, 4 * (top_k + 1) - 1 - kv_start)); + const int n_scored = n_tokens - first_scored; + if (n_scored <= 0) return nullptr; + + auto token_slice = [&](ggml_tensor * input, int width) { + if (first_scored == 0) return input; + return ggml_view_2d( + ctx, input, width, n_scored, input->nb[1], + (size_t) first_scored * input->nb[1]); + }; + qr_norm = token_slice(qr_norm, (int) qr_norm->ne[0]); + cur = token_slice(cur, (int) cur->ne[0]); + if (first_scored > 0) { + rope_pos = ggml_view_1d( + ctx, rope_pos, n_scored, + (size_t) first_scored * rope_pos->nb[0]); + } + + // Official ratio-4 indexer graph: q_a-normalized query projection, tail + // RoPE, Hadamard+FP4 QAT, per-head scalar projection, ReLU dot products, + // weighted head reduction and top-512 selection for every query token. + ggml_tensor * index_q = ggml_mul_mat(ctx, L.indexer_attn_q_b, qr_norm); + index_q = ggml_reshape_3d( + ctx, index_q, head_dim, n_indexer_head, n_scored); + + const float rope_scale = w.rope_scale_factor > 0.0f + ? (1.0f / w.rope_scale_factor) : 1.0f; + float rope_attn = 1.0f; + if (rope_scale > 0.0f) { + rope_attn /= 1.0f + 0.1f * logf(1.0f / rope_scale); + } + index_q = build_tail_rope_3d( + ctx, index_q, rope_pos, w.n_rot, head_dim, n_indexer_head, + n_scored, w.compress_rope_freq_base, rope_scale, 1.0f, + rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, + (int) w.rope_orig_ctx); + index_q = ggml_ds4_indexer_qat(ctx, ggml_cont(ctx, index_q)); - // DS4 indexer decode scoring mirrors ds4.c::indexer_allowed_decode_one(): - // 1. Build an indexer query from qr_norm (after q_a + RMSNorm, before q_b). - // 2. Apply full-dim RoPE in indexer head space. - // 3. Project per-head scalar weights from the current hidden state. - // 4. Score every compressed row with ReLU(dot(key_h, query_h)) * weight_h. - // 5. Return the top-k compressed-row indices. - ggml_tensor * index_q = ggml_mul_mat(ctx, L.indexer_attn_q_b, qr_norm_last); - index_q = ggml_reshape_3d(ctx, index_q, head_dim, n_indexer_head, 1); - - // TODO: RoPE on indexer query (same gallocr issue as compressor RoPE) - // Skipping for now — correctness deferred. - index_q = ggml_reshape_2d(ctx, index_q, head_dim, n_indexer_head); - - ggml_tensor * head_weights = ggml_mul_mat(ctx, L.indexer_proj, cur_last); + ggml_tensor * head_weights = ggml_mul_mat(ctx, L.indexer_proj, cur); head_weights = ggml_scale(ctx, head_weights, 1.0f / std::sqrt((float) head_dim * (float) n_indexer_head)); - // index_comp_kv: [n_indexer_head_dim, comp_cap] — each row is 128-dim - // Score each compressed row against all query heads via broadcast - ggml_tensor * comp_view = ggml_view_2d(ctx, lc.index_comp_kv, - head_dim, n_comp, - lc.index_comp_kv->nb[1], 0); - comp_view = ggml_cast(ctx, comp_view, GGML_TYPE_F32); - // comp_view: [head_dim, n_comp] → [head_dim, 1, n_comp] for broadcast - comp_view = ggml_reshape_3d(ctx, comp_view, head_dim, 1, n_comp); - - // index_q: [head_dim, n_indexer_head, 1] → repeat to [head_dim, n_indexer_head, n_comp] - // But ggml_mul needs same shapes, so use matmul approach: - // Reshape q: [head_dim, n_indexer_head] → used directly as A in matmul - // comp: [head_dim, n_comp] - // matmul: A^T @ B = [n_indexer_head, n_comp] dot scores - ggml_tensor * comp_2d = ggml_reshape_2d(ctx, comp_view, head_dim, n_comp); - // mul_mat(index_q, comp_2d): A=[head_dim, n_indexer_head], B=[head_dim, n_comp] - // → result=[n_indexer_head, n_comp] - ggml_tensor * dots = ggml_mul_mat(ctx, index_q, comp_2d); - dots = ggml_relu(ctx, dots); - - // Weight each head's contribution: dots[n_indexer_head, n_comp] * weights[n_indexer_head, 1] - ggml_tensor * weight_rep = ggml_repeat(ctx, head_weights, dots); - ggml_tensor * weighted = ggml_mul(ctx, dots, weight_rep); - // Sum across heads (ne[0]) → [1, n_comp] - ggml_tensor * scores = ggml_sum_rows(ctx, weighted); - scores = ggml_cont(ctx, scores); - scores = ggml_reshape_2d(ctx, scores, n_comp, 1); - - return ggml_top_k(ctx, scores, std::min(n_comp, w.n_indexer_top_k)); -} - -static ggml_tensor * build_selected_comp_context( - ggml_context * ctx, - ggml_tensor * selected_rows, // [head_dim, n_selected] - ggml_tensor * query_seed, // [head_dim, 1] - ggml_tensor * q_template, // [head_dim, n_head, n_tokens] - int head_dim) { - if (!selected_rows || !query_seed || !q_template || selected_rows->ne[1] <= 0) { - return nullptr; - } - - ggml_tensor * score = ggml_mul_mat(ctx, selected_rows, query_seed); - ggml_tensor * probs = ggml_soft_max(ctx, score); - ggml_tensor * rows_t = ggml_cont(ctx, ggml_transpose(ctx, selected_rows)); - ggml_tensor * context = ggml_mul_mat(ctx, rows_t, probs); - context = ggml_reshape_3d(ctx, context, head_dim, 1, 1); - return ggml_repeat(ctx, context, q_template); + ggml_tensor * comp = ggml_view_2d( + ctx, index_comp_source, head_dim, n_comp, + index_comp_source->nb[1], 0); + GGML_ASSERT(comp->type == GGML_TYPE_F16); + GGML_ASSERT(ggml_is_contiguous(comp)); + + // Fuse comp^T@q, ReLU, per-head weighting and head reduction. The generic + // graph would retain [n_comp,64,n_tokens] dots (about 2 GiB at 8K) before + // reducing them; this operation stores only the final score matrix. + ggml_tensor * scores = ggml_ds4_indexer_score( + ctx, index_q, head_weights, comp, kv_start + first_scored, 4); + ggml_tensor * selected = ggml_top_k( + ctx, ggml_cont(ctx, scores), top_k); + if (first_scored == 0) return selected; + + ggml_tensor * identity = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, top_k, first_scored); + ggml_set_input(identity); + std::vector identity_values((size_t) top_k * first_scored); + for (int t = 0; t < first_scored; ++t) { + for (int k = 0; k < top_k; ++k) { + identity_values[(size_t) t * top_k + k] = k; + } + } + i32_array_inputs.push_back({identity, std::move(identity_values)}); + return ggml_concat(ctx, identity, selected, 1); } // ─── MLA Attention Block ──────────────────────────────────────────────── @@ -1122,12 +1496,12 @@ static ggml_tensor * build_mla_attention( std::vector & i32_inputs, std::vector & i32_array_inputs, std::vector & i64_array_inputs, - std::vector * f32_array_inputs = nullptr) { + std::vector * f32_array_inputs = nullptr, + DeepSeek4AttentionImpl attention_impl = DeepSeek4AttentionImpl::Explicit) { const int n_embd = w.n_embd; const int head_dim = w.head_dim; const int n_head = w.n_head; - const int n_lora_q = w.n_lora_q; 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; @@ -1175,9 +1549,16 @@ static ggml_tensor * build_mla_attention( // n_ctx_orig is critical for YaRN correction on compressed layers const int rope_n_ctx_orig = (int)w.rope_orig_ctx; // 65536 - 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); + // 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); + } 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); @@ -1191,8 +1572,12 @@ static ggml_tensor * build_mla_attention( // (bidirectional) behavior for A/B comparison. const bool causal_batch = (n_tokens > 1) && !cached_inputs && f32_array_inputs && !ds4_env_flag("DFLASH_DS4_NO_CAUSAL_VERIFY"); + const bool layer_major_batch = + causal_batch && attention_impl != DeepSeek4AttentionImpl::Explicit; ggml_tensor * old_rows_scratch = nullptr; int n_old_rows = 0; + ggml_tensor * prior_rows_scratch = nullptr; + int n_prior_rows = 0; const bool fused_causal = cached_inputs && cached_inputs->attn_row_mask && n_tokens > 1; if (fused_causal) { // Fused verify: ALWAYS q preserved rows so the topology is stable; @@ -1208,7 +1593,7 @@ static ggml_tensor * build_mla_attention( n_old_rows++; } old_rows_scratch = ds4_cast_if_needed(ctx, old_rows_scratch, GGML_TYPE_F32); - } else if (causal_batch) { + } else if (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++) { @@ -1225,17 +1610,49 @@ 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) { + // 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. + n_prior_rows = std::min(kv_start, w.n_swa); + if (n_prior_rows > 0) { + const int first = kv_start < w.n_swa ? 0 : (kv_start % w.n_swa); + 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]); + return ggml_cont(ctx, span); + }; + prior_rows_scratch = snapshot_span(first, tail); + if (tail < n_prior_rows) { + prior_rows_scratch = ggml_concat( + ctx, prior_rows_scratch, + snapshot_span(0, n_prior_rows - tail), 1); + prior_rows_scratch = ggml_cont(ctx, prior_rows_scratch); + } + ggml_build_forward_expand(gf, prior_rows_scratch); + prior_rows_scratch = ds4_cast_if_needed( + ctx, prior_rows_scratch, GGML_TYPE_F32); + } } // ── 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; - if (cached_inputs && cached_inputs->raw_kv_rows) { + ggml_tensor * raw_kv_rows = cached_inputs + ? cached_inputs->raw_kv_rows + : nullptr; + 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, cached_inputs->raw_kv_rows); + raw_kv_source = ggml_set_rows(ctx, lc.raw_kv, kv_f32, raw_kv_rows); ggml_build_forward_expand(gf, raw_kv_source); } else { - for (int ti = 0; ti < n_tokens; ti++) { + // The attention graph consumes the whole current ubatch directly. + // Persist only its final SWA tail so every physical ring row is written + // once, even when the ubatch is much larger than the 128-row ring. + const int first_write = std::max(0, n_tokens - w.n_swa); + for (int ti = first_write; ti < n_tokens; ti++) { const int pos_ti = kv_start + ti; ggml_tensor * kv_row = ggml_view_2d( ctx, kv, head_dim, 1, kv->nb[1], (size_t)ti * kv->nb[1]); @@ -1250,8 +1667,6 @@ 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 * qr_last = ggml_view_2d( - ctx, qr, n_lora_q, 1, qr->nb[1], (size_t)(n_tokens - 1) * qr->nb[1]); ggml_tensor * comp_kv_source = lc.comp_kv; if (ratio > 0 && L.attn_compressor_kv) { build_compressor_step(ctx, gf, cur_last, @@ -1284,6 +1699,7 @@ static ggml_tensor * build_mla_attention( 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, cached_inputs ? cached_inputs->index_ape_row : nullptr, @@ -1292,11 +1708,13 @@ static ggml_tensor * build_mla_attention( cached_inputs ? cached_inputs->index_comp_pos : nullptr, i64_array_inputs, i32_array_inputs, + &index_comp_kv_source, cached_inputs ? cached_inputs->flush_rows : nullptr, (causal_batch || fused_causal) ? cur : nullptr, n_tokens, - kv_start); - (void)build_indexer_score(ctx, qr_last, cur_last, w, L, lc, token_pos, i32_inputs); + kv_start, + attention_impl == + DeepSeek4AttentionImpl::SparseFlash); } // ── MLA Dot-Product Attention (SWA + compressed KV) ──────────── @@ -1306,11 +1724,22 @@ static ggml_tensor * build_mla_attention( // 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; + ggml_tensor * indexer_topk = nullptr; + if (attention_impl == DeepSeek4AttentionImpl::SparseFlash && + ratio == 4 && f32_array_inputs) { + const int n_index_comp = ds4_comp_rows_used( + lc.index_comp_kv, lc.n_index_comp, 4, token_pos); + indexer_topk = build_indexer_topk( + ctx, qr, cur, w, L, index_comp_kv_source, + n_index_comp, kv_start, n_tokens, rope_pos, + 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 : std::min(kv_start + n_tokens, w.n_swa); + const int n_raw = 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; - const int n_valid_raw = std::min(kv_start + n_tokens, w.n_swa); const int n_attn = n_raw + n_comp_attn + n_old_rows; const float kq_scale = 1.0f / sqrtf((float)head_dim); @@ -1334,9 +1763,26 @@ static ggml_tensor * build_mla_attention( ggml_tensor * ring = ggml_view_2d( ctx, raw_kv_source, head_dim, w.n_swa, raw_kv_source->nb[1], 0); kv_attn = ds4_cast_if_needed(ctx, ring, GGML_TYPE_F32); + } else if (layer_major_batch) { + ggml_tensor * current = ds4_cast_if_needed(ctx, kv, GGML_TYPE_F32); + kv_attn = prior_rows_scratch + ? ggml_concat(ctx, prior_rows_scratch, current, 1) + : current; } else if (n_tokens == 1) { ggml_tensor * cur_kv = ds4_cast_if_needed(ctx, kv, GGML_TYPE_F32); - if (n_raw > 1) { + if (n_raw == w.n_swa && raw_kv_rows) { + // Once the ring is full, use a stable physical row order. The + // cached q=1 graph is first built at position n_swa-1 and then + // reused across every wrap position, so chronological views baked + // into that first topology become stale. Insert the current F32 + // 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); + 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); + } else if (n_raw > 1) { ggml_tensor * prev = nullptr; DeepSeek4RawRingSpan spans[2]; const int n_spans = @@ -1362,96 +1808,328 @@ static ggml_tensor * build_mla_attention( } // kv_attn: [head_dim, n_attn] - // Flatten q to [head_dim, n_head*n_tokens] for batched matmul - ggml_tensor * q_flat = ggml_reshape_2d(ctx, q, head_dim, n_head * n_tokens); - - // Scores: mul_mat(kv_attn, q_flat) = kv_attn^T[n_attn, head_dim] @ q_flat[head_dim, n_head*n_tokens] - // → [n_attn, n_head*n_tokens] - ggml_tensor * scores = ggml_mul_mat(ctx, kv_attn, q_flat); - scores = ggml_scale(ctx, scores, kq_scale); - if (masked_kv && n_tokens > 1) { - // Per-token causal mask [n_attn, n_tokens] from the host-filled bundle. - ggml_tensor * m3 = ggml_reshape_3d(ctx, cached_inputs->attn_row_mask, n_attn, 1, n_tokens); - ggml_tensor * scores3d = ggml_reshape_3d(ctx, scores, n_attn, n_head, n_tokens); - scores3d = ggml_add(ctx, scores3d, m3); - scores = ggml_reshape_2d(ctx, scores3d, n_attn, n_head * n_tokens); - } else if (masked_kv) { - // Broadcast-add the [n_attn,1] additive mask across all query columns. - scores = ggml_add(ctx, scores, cached_inputs->attn_row_mask); - } else if (causal_batch) { - // Per-token causal mask over [ring rows | comp rows | old rows]. - ggml_tensor * cmask = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_attn, 1, n_tokens); - ggml_set_input(cmask); - std::vector mvals((size_t) n_attn * n_tokens, 0.0f); - const int e = kv_start + n_tokens; // exclusive end position - for (int i = 0; i < n_tokens; i++) { - const int pos_i = kv_start + i; - float * col = mvals.data() + (size_t) i * n_attn; - for (int r = 0; r < n_raw; r++) { - // position held by ring slot r AFTER this batch's writes - const int p_r = (e <= w.n_swa) ? r - : (e - 1) - ((e - 1 - r) % w.n_swa); - if (p_r > pos_i) col[r] = -1e30f; + // Build one additive mask tensor and share it between the explicit and + // flash-attention implementations. ggml flash attention expects + // [n_kv,n_query] F16; the explicit path broadcasts the same values over + // heads in F32. + ggml_tensor * score_mask = nullptr; + const bool exact_two_band = + attention_impl == DeepSeek4AttentionImpl::DenseFlash && + causal_batch && + n_tokens > DS4_NUMERICAL_PREFILL_BAND && + n_tokens <= 2 * DS4_NUMERICAL_PREFILL_BAND; + if (!exact_two_band) { + if (masked_kv && n_tokens > 1) { + score_mask = ggml_reshape_2d(ctx, cached_inputs->attn_row_mask, + n_attn, n_tokens); + } else if (masked_kv) { + score_mask = ggml_reshape_2d(ctx, cached_inputs->attn_row_mask, + n_attn, 1); + } else if (layer_major_batch) { + // Per-token causal mask over [prior rows | current rows | comp rows]. + ggml_tensor * cmask = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_attn, 1, n_tokens); + ggml_set_input(cmask); + std::vector mvals((size_t) n_attn * n_tokens, 0.0f); + for (int i = 0; i < n_tokens; i++) { + const int pos_i = kv_start + i; + float * col = mvals.data() + (size_t) i * n_attn; + const int min_pos = pos_i - w.n_swa + 1; + for (int r = 0; r < n_prior_rows; ++r) { + const int prior_pos = kv_start - n_prior_rows + r; + if (prior_pos < min_pos) col[r] = -1e30f; + } + for (int t = 0; t < n_tokens; ++t) { + const int current_pos = kv_start + t; + if (t > i || current_pos < min_pos) { + col[n_prior_rows + t] = -1e30f; + } + } + if (n_comp_attn > 0) { + const int vis = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, pos_i); + for (int c = vis; c < n_comp_attn; c++) col[n_raw + c] = -1e30f; + } + } + f32_array_inputs->push_back({cmask, std::move(mvals)}); + score_mask = ggml_reshape_2d(ctx, cmask, n_attn, n_tokens); + } else if (causal_batch) { + // Speculative verification keeps the physical ring order and + // appends snapshots of rows overwritten by later batch tokens. + ggml_tensor * cmask = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, n_attn, 1, n_tokens); + ggml_set_input(cmask); + std::vector mvals((size_t) n_attn * n_tokens, 0.0f); + const int end = kv_start + n_tokens; + for (int i = 0; i < n_tokens; ++i) { + const int pos_i = kv_start + i; + float * col = mvals.data() + (size_t) i * n_attn; + for (int r = 0; r < n_raw; ++r) { + const int pos_r = end <= w.n_swa + ? r + : (end - 1) - ((end - 1 - r) % w.n_swa); + 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); + for (int c = visible; c < n_comp_attn; ++c) { + col[n_raw + c] = -1e30f; + } + } + int old_row = 0; + for (int t = 0; t < n_tokens; ++t) { + if (kv_start + t < w.n_swa) continue; + if (t <= i) { + col[n_raw + n_comp_attn + old_row] = -1e30f; + } + ++old_row; + } } - if (n_comp_attn > 0) { - const int vis = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, pos_i); - for (int c = vis; c < n_comp_attn; c++) col[n_raw + c] = -1e30f; + f32_array_inputs->push_back({cmask, std::move(mvals)}); + score_mask = ggml_reshape_2d(ctx, cmask, n_attn, n_tokens); + } + } + if (indexer_topk) { + if (!score_mask) { + score_mask = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_attn, n_tokens); + ggml_set_input(score_mask); + f32_array_inputs->push_back({ + score_mask, + std::vector((size_t) n_attn * n_tokens, 0.0f), + }); + } + score_mask = ggml_ds4_indexer_mask( + ctx, ggml_cont(ctx, score_mask), indexer_topk, n_raw); + } + ggml_tensor * context = nullptr; + bool inverse_rope_fused = false; + const bool use_flash = attention_impl != DeepSeek4AttentionImpl::Explicit && + n_tokens > 1; + if (use_flash) { + if (exact_two_band) { + // A larger scheduling band must retain the numerical topology of + // two 2K requests. Prefix queries use the first band's F32 raw KV; + // suffix queries see its final SWA tail after the same F16 cache + // round-trip. HC, projections and MoE still run once over the full + // token batch, avoiding a second expert-weight sweep. + 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, + kv_start + first_count - 1) + : 0; + const int second_comp = n_comp_live; + const int second_prior_count = std::min(first_count, w.n_swa); + + auto view_kv = [&](int first, int count) { + return ggml_view_2d( + ctx, kv, head_dim, count, kv->nb[1], + (size_t) first * kv->nb[1]); + }; + auto append_comp = [&](ggml_tensor * raw, int count) { + if (count <= 0 || !comp_kv_source) return raw; + ggml_tensor * comp = ggml_view_2d( + ctx, comp_kv_source, head_dim, count, + comp_kv_source->nb[1], 0); + comp = ds4_cast_if_needed(ctx, comp, GGML_TYPE_F32); + return ggml_concat(ctx, raw, comp, 1); + }; + auto make_band_mask = [&](int start, int count, int prior, + int comp_count) { + const int raw_count = prior + count; + const int attn_count = raw_count + comp_count; + ggml_tensor * mask3 = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, attn_count, 1, count); + ggml_set_input(mask3); + std::vector values( + (size_t) attn_count * count, 0.0f); + for (int i = 0; i < count; ++i) { + const int pos_i = start + i; + const int min_pos = pos_i - w.n_swa + 1; + float * col = values.data() + (size_t) i * attn_count; + for (int r = 0; r < prior; ++r) { + const int prior_pos = start - prior + r; + if (prior_pos < min_pos) col[r] = -1e30f; + } + for (int t = 0; t < count; ++t) { + const int current_pos = start + t; + if (t > i || current_pos < min_pos) { + col[prior + t] = -1e30f; + } + } + if (comp_count > 0) { + const int visible = ds4_comp_rows_used( + lc.comp_kv, lc.n_comp, ratio, pos_i); + for (int c = visible; c < comp_count; ++c) { + col[raw_count + c] = -1e30f; + } + } + } + f32_array_inputs->push_back({mask3, std::move(values)}); + return ggml_reshape_2d(ctx, mask3, attn_count, count); + }; + auto make_flash = [&](ggml_tensor * q_band, + ggml_tensor * kv_band, + ggml_tensor * mask_band, + int raw_count, + int start) { + const int attn_count = (int) kv_band->ne[1]; + ggml_tensor * k_band = ggml_reshape_3d( + ctx, kv_band, head_dim, attn_count, 1); + ggml_tensor * mask_fa = ds4_cast_if_needed( + ctx, mask_band, GGML_TYPE_F16); + ggml_tensor * result = ggml_flash_attn_ext( + ctx, q_band, k_band, k_band, mask_fa, + kq_scale, 0.0f, 0.0f); + if (L.attn_sinks) { + ggml_flash_attn_ext_add_sinks(result, L.attn_sinks); + } + ggml_flash_attn_ext_set_prec(result, GGML_PREC_F32); + ggml_flash_attn_ext_set_ds4_sparse( + result, raw_count, w.n_swa, 0, 32); + ggml_flash_attn_ext_set_ds4_inverse_rope( + result, start, rope_freq, rope_scale, rope_ext, + rope_attn, w.rope_yarn_beta_fast, + w.rope_yarn_beta_slow, rope_n_ctx_orig, fuse_q_rope); + return result; + }; + + ggml_tensor * q_fa = ggml_permute(ctx, q, 0, 2, 1, 3); + auto view_q = [&](int first, int count) { + return ggml_view_3d( + ctx, q_fa, head_dim, count, n_head, + q_fa->nb[1], q_fa->nb[2], + (size_t) first * q_fa->nb[1]); + }; + + ggml_tensor * first_raw = ds4_cast_if_needed( + ctx, view_kv(0, first_count), GGML_TYPE_F32); + if (prior_rows_scratch) { + first_raw = ggml_concat( + ctx, prior_rows_scratch, first_raw, 1); + } + ggml_tensor * first_kv = append_comp(first_raw, first_comp); + ggml_tensor * first_mask = make_band_mask( + kv_start, first_count, n_prior_rows, first_comp); + ggml_tensor * first_context = make_flash( + view_q(0, first_count), first_kv, first_mask, + n_prior_rows + first_count, kv_start); + + ggml_tensor * rounded_prior = ggml_cast( + ctx, view_kv(first_count - second_prior_count, + second_prior_count), + GGML_TYPE_F16); + rounded_prior = ggml_cast(ctx, rounded_prior, GGML_TYPE_F32); + ggml_tensor * second_raw = ggml_concat( + ctx, rounded_prior, view_kv(first_count, second_count), 1); + ggml_tensor * second_kv = append_comp(second_raw, second_comp); + ggml_tensor * second_mask = make_band_mask( + kv_start + first_count, second_count, + second_prior_count, second_comp); + ggml_tensor * second_context = make_flash( + view_q(first_count, second_count), second_kv, second_mask, + second_prior_count + second_count, + kv_start + first_count); + + context = ggml_concat(ctx, first_context, second_context, 2); + inverse_rope_fused = true; + } else { + // ggml FA convention: Q[D,T,H], K/V[D,K,Hkv]. DS4 MLA has one shared + // latent KV head and uses the same latent vector as both key and value. + // The DS4 D=512 kernel consumes Q strides directly, avoiding a full + // [D,H,T] -> [D,T,H] materialization for every layer. + ggml_tensor * q_fa = ggml_permute(ctx, q, 0, 2, 1, 3); + ggml_tensor * kv_fa = ds4_cast_if_needed(ctx, kv_attn, GGML_TYPE_F32); + ggml_tensor * k_fa = ggml_reshape_3d(ctx, kv_fa, head_dim, n_attn, 1); + ggml_tensor * v_fa = k_fa; + ggml_tensor * mask_fa = score_mask + ? ds4_cast_if_needed(ctx, score_mask, GGML_TYPE_F16) + : nullptr; + context = ggml_flash_attn_ext(ctx, q_fa, k_fa, v_fa, mask_fa, + kq_scale, 0.0f, 0.0f); + if (L.attn_sinks) { + ggml_flash_attn_ext_add_sinks(context, L.attn_sinks); } - // old contents of slot overwritten by batch token j are visible - // exactly to tokens i < j (still inside their SWA window) - int oi = 0; - for (int tj = 0; tj < n_tokens; tj++) { - if (kv_start + tj < w.n_swa) continue; - if (tj <= i) col[n_raw + n_comp_attn + oi] = -1e30f; - oi++; + ggml_flash_attn_ext_set_prec(context, GGML_PREC_F32); + // Always publish the raw/compressed boundary. A zero keep count leaves + // dense attention unchanged while allowing the D=512 value pass to + // skip the two masked envelopes without guessing DS4 cache layout. + ggml_flash_attn_ext_set_ds4_sparse( + context, n_raw, w.n_swa, + indexer_topk + ? -w.n_indexer_top_k + : attention_impl == DeepSeek4AttentionImpl::SparseFlash + ? w.n_indexer_top_k : 0, + 32); + if (attention_impl != DeepSeek4AttentionImpl::Explicit && + head_dim == 512 && n_rot == 64) { + ggml_flash_attn_ext_set_ds4_inverse_rope( + context, kv_start, rope_freq, rope_scale, rope_ext, + rope_attn, w.rope_yarn_beta_fast, + w.rope_yarn_beta_slow, rope_n_ctx_orig, fuse_q_rope); + inverse_rope_fused = true; } } - f32_array_inputs->push_back({cmask, std::move(mvals)}); - ggml_tensor * scores3d = ggml_reshape_3d(ctx, scores, n_attn, n_head, n_tokens); - scores3d = ggml_add(ctx, scores3d, cmask); - scores = ggml_reshape_2d(ctx, scores3d, n_attn, n_head * n_tokens); - } - (void) n_valid_raw; - - // Sink-aware softmax: 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); - if (n_tokens > 1) { - ggml_tensor * sink_shape = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n_head * n_tokens); - sink_scores = ggml_repeat(ctx, sink_scores, sink_shape); - } - ggml_tensor * scores_with_sink = ggml_concat(ctx, scores, sink_scores, 0); - ggml_tensor * probs_with_sink = ggml_soft_max(ctx, scores_with_sink); - probs = ggml_view_2d(ctx, probs_with_sink, n_attn, n_head * n_tokens, - probs_with_sink->nb[1], 0); } else { - probs = ggml_soft_max(ctx, scores); - } - // probs: [n_attn, n_head*n_tokens] - - // Context: kv_T^T[head_dim, n_attn] @ probs[n_attn, n_head*n_tokens] → [head_dim, n_head*n_tokens] - // i.e. mul_mat(kv_T, probs) where kv_T = cont(transpose(kv_attn)) = [n_raw, head_dim] - ggml_tensor * kv_T = ggml_cont(ctx, ggml_transpose(ctx, kv_attn)); - ggml_tensor * context = ggml_mul_mat(ctx, kv_T, probs); - // context: [head_dim, n_head*n_tokens] + // Flatten q to [head_dim, n_head*n_tokens] for batched matmul. + ggml_tensor * q_flat = ggml_reshape_2d(ctx, q, head_dim, + n_head * n_tokens); + ggml_tensor * scores = ggml_mul_mat(ctx, kv_attn, q_flat); + scores = ggml_scale(ctx, scores, kq_scale); + if (score_mask) { + if (n_tokens > 1) { + ggml_tensor * m3 = ggml_reshape_3d(ctx, score_mask, + n_attn, 1, n_tokens); + ggml_tensor * s3 = ggml_reshape_3d(ctx, scores, + n_attn, n_head, n_tokens); + scores = ggml_reshape_2d(ctx, ggml_add(ctx, s3, m3), + n_attn, n_head * n_tokens); + } else { + scores = ggml_add(ctx, scores, score_mask); + } + } - // Reshape back to [head_dim, n_head, n_tokens] - context = ggml_reshape_3d(ctx, context, head_dim, n_head, n_tokens); + // 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); + if (n_tokens > 1) { + ggml_tensor * sink_shape = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 1, n_head * n_tokens); + sink_scores = ggml_repeat(ctx, sink_scores, sink_shape); + } + ggml_tensor * scores_with_sink = ggml_concat(ctx, scores, + sink_scores, 0); + ggml_tensor * probs_with_sink = ggml_soft_max(ctx, + scores_with_sink); + probs = ggml_view_2d(ctx, probs_with_sink, n_attn, + n_head * n_tokens, probs_with_sink->nb[1], 0); + } 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); + context = ggml_reshape_3d(ctx, context, head_dim, n_head, n_tokens); + } // ── Inverse tail RoPE on attention output ─────────────────────── - 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); - ggml_set_input(neg_pos); - std::vector neg_vals(n_tokens); - for (int i = 0; i < n_tokens; i++) neg_vals[i] = -(kv_start + i); - i32_array_inputs.push_back({neg_pos, std::move(neg_vals)}); - } - context = build_tail_rope_3d(ctx, context, neg_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 (!inverse_rope_fused) { + 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); + ggml_set_input(neg_pos); + std::vector neg_vals(n_tokens); + for (int i = 0; i < n_tokens; i++) neg_vals[i] = -(kv_start + i); + i32_array_inputs.push_back({neg_pos, std::move(neg_vals)}); + } + context = build_tail_rope_3d( + ctx, context, neg_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); + } // 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); @@ -1466,16 +2144,28 @@ static ggml_tensor * build_mla_attention( 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_cont(ctx, ggml_permute(ctx, attn_out, 0, 2, 1, 3)); + 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] - // Permute back to [n_lora_o, n_out_group, n_tokens] then flatten - 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); - ggml_tensor * out = ggml_mul_mat(ctx, L.attn_output_b, attn_low); + ggml_tensor * out = nullptr; + if (n_tokens > 1) { + // 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 its numerical + // behavior. Decode is intentionally outside the prefill fast 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); + } return out; } @@ -3127,6 +3817,23 @@ static void reset_hc_layer_weights_cpu(std::vector & weights) weights.clear(); } +struct DeepSeek4HybridRuntime { + const ggml_context * owner_ctx = nullptr; + std::vector hc_layer_weights; + HcWeightsCpu hc_output_weights; + std::vector hash_routing_tables; + + void destroy() { + reset_hc_layer_weights_cpu(hc_layer_weights); + reset_hc_weights_cpu(hc_output_weights); + hash_routing_tables.clear(); + hash_routing_tables.shrink_to_fit(); + owner_ctx = nullptr; + } +}; + +static thread_local DeepSeek4HybridRuntime ds4_hybrid_runtime; + static const void * hc_fn_device_ptr(const HcWeightsCpu &, ggml_tensor * fn) { if (!fn) return nullptr; if (fn->type == GGML_TYPE_F16) return fn->data; @@ -3174,10 +3881,17 @@ static bool deepseek4_step_hybrid( } } - // Lazy-loaded per-layer HC weights on CPU - static std::vector hc_layer_weights; - static HcWeightsCpu hc_output_weights; - static std::vector hash_routing_tables; + // Cache host mirrors by model ownership so unload/reload cannot reuse + // tensor data from a previous model context. + DeepSeek4HybridRuntime & runtime = ds4_hybrid_runtime; + if (runtime.owner_ctx != w.ctx || + runtime.hc_layer_weights.size() != (size_t) w.n_layer) { + runtime.destroy(); + runtime.owner_ctx = w.ctx; + } + auto & hc_layer_weights = runtime.hc_layer_weights; + auto & hc_output_weights = runtime.hc_output_weights; + auto & hash_routing_tables = runtime.hash_routing_tables; if (hc_layer_weights.empty()) { hc_layer_weights.resize((size_t)w.n_layer); hash_routing_tables.resize((size_t)w.n_layer); @@ -3794,6 +4508,66 @@ struct DeepSeek4FusedDecodeCache { // releases the expert owner used by a cached multi-backend scheduler. static DeepSeek4FusedDecodeCache * g_deepseek4_fused_decode_runtime_cache = nullptr; +struct DeepSeek4LayerRangeRuntime { + const ggml_context * owner_ctx = nullptr; + int loaded_n_layer = 0; + uint64_t loaded_generation = 0; + std::vector hc_layer_weights; + HcWeightsCpu hc_output_weights; + std::vector hash_routing_tables; + std::vector cached_attn_allocs; + std::vector cached_decode_attn_hc_pre_graphs; + std::vector cached_decode_ffn_hc_pre_graphs; + DeepSeek4CachedDecodeHcPostGraph cached_decode_hc_post_graph; + std::vector> cached_decode_attn_graphs; + std::vector cached_decode_ffn_graphs; + DeepSeek4CachedDecodeOutputGraph cached_decode_output_graph; + DeepSeek4CachedLayerAlloc cached_dynamic_output_alloc; + DeepSeek4FusedDecodeCache fused_decode_graph_cache; + Ds4DecodeSharedInputs decode_shared_inputs; + DeepSeek4LayerRangeScratch scratch; + + void destroy() { + for (auto & alloc : cached_attn_allocs) { + alloc.free(); + } + cached_attn_allocs.clear(); + for (auto & graph : cached_decode_attn_hc_pre_graphs) { + graph.free(); + } + cached_decode_attn_hc_pre_graphs.clear(); + for (auto & graph : cached_decode_ffn_hc_pre_graphs) { + graph.free(); + } + cached_decode_ffn_hc_pre_graphs.clear(); + cached_decode_hc_post_graph.free(); + for (auto & per_layer : cached_decode_attn_graphs) { + for (auto & graph : per_layer) { + graph.free(); + } + } + cached_decode_attn_graphs.clear(); + for (auto & graph : cached_decode_ffn_graphs) { + graph.free(); + } + cached_decode_ffn_graphs.clear(); + cached_decode_output_graph.free(); + cached_dynamic_output_alloc.free(); + fused_decode_graph_cache.destroy(); + decode_shared_inputs.free(); + reset_hc_layer_weights_cpu(hc_layer_weights); + reset_hc_weights_cpu(hc_output_weights); + hash_routing_tables.clear(); + hash_routing_tables.shrink_to_fit(); + scratch.clear(); + loaded_n_layer = 0; + loaded_generation = 0; + owner_ctx = nullptr; + } +}; + +static thread_local DeepSeek4LayerRangeRuntime ds4_layer_range_runtime; + static ggml_tensor * ds4_fused_hc_base_f32(ggml_context * ctx, ggml_tensor * base) { if (!base) return nullptr; ggml_tensor * b = base; @@ -3829,7 +4603,8 @@ static ggml_tensor * ds4_build_hash_routed_ffn( const DeepSeek4Weights & w, const DeepSeek4Layer & L, ggml_tensor * ffn_normed, - ggml_tensor * hash_ids) { + ggml_tensor * hash_ids, + int n_tokens) { ggml_tensor * shared_out = build_shared_ffn(ctx, ffn_normed, w, L); ggml_tensor * logits = ggml_mul_mat(ctx, L.ffn_gate_inp, ffn_normed); ggml_tensor * probs = ggml_sqrt(ctx, ggml_softplus(ctx, logits)); @@ -3837,18 +4612,20 @@ static ggml_tensor * ds4_build_hash_routed_ffn( const int n_used = (int) hash_ids->ne[0]; GGML_ASSERT(n_used > 0 && n_used <= w.n_expert_used); const int n_ff_exp = w.n_ff_exp; - ggml_tensor * cur_3d = ggml_reshape_3d(ctx, ffn_normed, w.n_embd, 1, 1); + ggml_tensor * cur_3d = ggml_reshape_3d( + ctx, ffn_normed, w.n_embd, 1, n_tokens); ggml_tensor * gate_e = ggml_mul_mat_id(ctx, L.ffn_gate_exps, cur_3d, hash_ids); ggml_tensor * up_e = ggml_mul_mat_id(ctx, L.ffn_up_exps, cur_3d, hash_ids); - gate_e = ggml_reshape_3d(ctx, gate_e, n_ff_exp, n_used, 1); - up_e = ggml_reshape_3d(ctx, up_e, n_ff_exp, n_used, 1); + gate_e = ggml_reshape_3d(ctx, gate_e, n_ff_exp, n_used, n_tokens); + up_e = ggml_reshape_3d(ctx, up_e, n_ff_exp, n_used, n_tokens); ggml_tensor * mid_e = build_clamped_swiglu(ctx, gate_e, up_e, w.swiglu_clamp_exp); ggml_tensor * down_e = ggml_mul_mat_id(ctx, L.ffn_down_exps, mid_e, hash_ids); - down_e = ggml_reshape_3d(ctx, down_e, w.n_embd, n_used, 1); + down_e = ggml_reshape_3d(ctx, down_e, w.n_embd, n_used, n_tokens); - ggml_tensor * probs_3d = ggml_reshape_3d(ctx, probs, 1, w.n_expert, 1); + ggml_tensor * probs_3d = ggml_reshape_3d( + ctx, probs, 1, w.n_expert, n_tokens); ggml_tensor * weights = ggml_get_rows(ctx, probs_3d, hash_ids); - weights = ggml_reshape_2d(ctx, weights, n_used, 1); + weights = ggml_reshape_2d(ctx, weights, n_used, n_tokens); ggml_tensor * w_sum = ggml_sum_rows(ctx, weights); w_sum = ggml_clamp(ctx, w_sum, 6.103515625e-5f, INFINITY); weights = ggml_div(ctx, weights, w_sum); @@ -3856,11 +4633,21 @@ static ggml_tensor * ds4_build_hash_routed_ffn( weights = ggml_scale(ctx, weights, w.expert_weight_scale); } - ggml_tensor * weights_3d = ggml_reshape_3d(ctx, weights, 1, n_used, 1); + ggml_tensor * weights_3d = ggml_reshape_3d( + ctx, weights, 1, n_used, n_tokens); ggml_tensor * routed_out = ggml_mul(ctx, down_e, weights_3d); - routed_out = ggml_cont(ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); - routed_out = ggml_sum_rows(ctx, routed_out); - routed_out = ggml_reshape_2d(ctx, routed_out, w.n_embd, 1); + if (n_tokens == 1) { + // Preserve the established q=1 graph and reduction order. + routed_out = ggml_cont(ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); + routed_out = ggml_sum_rows(ctx, routed_out); + routed_out = ggml_reshape_2d(ctx, routed_out, w.n_embd, 1); + } else { + ggml_tensor * sum_shape = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, w.n_embd, 1, n_tokens); + routed_out = ggml_repeat_back(ctx, routed_out, sum_shape); + routed_out = ggml_reshape_2d( + ctx, routed_out, w.n_embd, n_tokens); + } return ggml_add(ctx, shared_out, routed_out); } @@ -4073,7 +4860,8 @@ static bool ds4_build_fused_decode_graph( ctx, GGML_TYPE_I32, ds4_effective_expert_count(w), 1); ggml_set_input(hids); fg.hash_ids[(size_t) il] = hids; - ffn_out = ds4_build_hash_routed_ffn(ctx, w, L, ffn_normed, hids); + ffn_out = ds4_build_hash_routed_ffn( + ctx, w, L, ffn_normed, hids, 1); } else { ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, 1); } @@ -4469,6 +5257,762 @@ static bool eval_ds4_layer_range_hybrid_ffn( normed_host.data(), selected.data(), weights.data(), n_tokens, out, nullptr, nullptr, expert_compute, expert_layer, telemetry); + +// Exact-order prefill control: retain the layer-major HC/FFN schedule, but run +// the attention subgraph one token at a time. This preserves the q=1 QKV, +// compressor, causal-attention, and output-projection reduction order while +// still allowing the token-independent FFN to use a multi-row ROCMFP graph. +static bool ds4_run_exact_tokenwise_prefill_attention( + ggml_backend_t backend, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + DeepSeek4LayerCache & lc, + int il, + const float * cur, + int n_tokens, + int kv_start, + DeepSeek4AttentionImpl attention_impl, + std::vector & attn_out_host, + DeepSeek4CachedLayerAlloc & attn_alloc, + DeepSeek4StepTelemetry * telemetry) { + if (!backend || !cur || n_tokens <= 1 || kv_start < 0) return false; + + const int n_embd = w.n_embd; + for (int ti = 0; ti < n_tokens; ++ti) { + const auto build_t0 = Ds4TimingClock::now(); + ggml_init_params params{}; + params.mem_size = ds4_attn_step_meta_size(1); + params.mem_buffer = nullptr; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) return false; + + ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, 1); + ggml_set_input(inp); + std::vector i32_inputs; + std::vector i32_array_inputs; + std::vector i64_array_inputs; + std::vector f32_array_inputs; + ggml_cgraph * gf = ggml_new_graph_custom( + ctx, ds4_attn_step_graph_size(1), false); + ggml_tensor * normed = build_rms_norm(ctx, inp, L.attn_norm, w.rms_eps); + ggml_tensor * attn_out = build_mla_attention( + ctx, gf, normed, w, L, lc, il, kv_start + ti, 1, nullptr, + i32_inputs, i32_array_inputs, i64_array_inputs, &f32_array_inputs, + attention_impl); + ggml_set_output(attn_out); + ggml_build_forward_expand(gf, attn_out); + + if (!attn_alloc.valid() || attn_alloc.owner_ctx != w.ctx || + attn_alloc.backend != backend) { + attn_alloc.free(); + attn_alloc.alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + attn_alloc.owner_ctx = w.ctx; + attn_alloc.backend = backend; + } + if (!attn_alloc.alloc || !ggml_gallocr_alloc_graph(attn_alloc.alloc, gf)) { + std::fprintf(stderr, + "[deepseek4] exact prefill attn alloc failed layer %d token %d\n", + il, ti); + ggml_free(ctx); + return false; + } + if (telemetry) { + telemetry->attn_build_us += ds4_elapsed_us(build_t0, Ds4TimingClock::now()); + } + + ggml_backend_tensor_set(inp, cur + (size_t) ti * n_embd, 0, + sizeof(float) * (size_t) n_embd); + for (const auto & b : i32_inputs) { + ggml_backend_tensor_set(b.tensor, &b.value, 0, sizeof(b.value)); + } + for (const auto & b : i32_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(int32_t) * b.values.size()); + } + for (const auto & b : i64_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(int64_t) * b.values.size()); + } + for (const auto & b : f32_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(float) * b.values.size()); + } + + const auto compute_t0 = Ds4TimingClock::now(); + const ggml_status status = ggml_backend_graph_compute(backend, gf); + if (telemetry) { + telemetry->attn_compute_us += ds4_elapsed_us( + compute_t0, Ds4TimingClock::now()); + } + if (status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[deepseek4] exact prefill attn compute failed layer %d token %d\n", + il, ti); + ggml_free(ctx); + return false; + } + + const auto read_t0 = Ds4TimingClock::now(); + ggml_backend_tensor_get(attn_out, + attn_out_host.data() + (size_t) ti * n_embd, + 0, sizeof(float) * (size_t) n_embd); + if (telemetry) { + telemetry->attn_read_us += ds4_elapsed_us(read_t0, Ds4TimingClock::now()); + } + + // Publish compressor rows immediately. The next token in this layer + // must observe a row flushed by the current token, matching the q=1 + // reference when a prefill band crosses a compressor boundary. + const int ratio = (int) w.compress_ratios[il]; + if (ratio > 0) { + const int next_pos = kv_start + ti + 1; + lc.n_comp = std::max(lc.n_comp, next_pos / ratio); + if (ratio == 4) { + lc.n_index_comp = std::max(lc.n_index_comp, + next_pos / ratio); + } + } + ggml_free(ctx); + } + return true; +} + +// Layer-major DS4 prefill. Each layer is one GPU graph containing batched HC, +// attention, MoE and HC post-processing. The HC state is kept in two external +// device tensors and ping-ponged between layers, eliminating the two host +// readbacks per layer in the reference implementation. Attention reads a +// snapshot of the previous SWA window plus the current ubatch; only the final +// SWA tail is committed to the persistent ring. The compressor publishes every +// ratio-4/ratio-128 boundary crossed by the ubatch. +// +struct Ds4LayerMajorCachedLayer { + void * meta_buffer = nullptr; + size_t meta_size = 0; + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + std::vector i32_inputs; + std::vector i32_array_inputs; + std::vector i64_array_inputs; + std::vector f32_array_inputs; + std::vector allocated_tensors; + ggml_tensor * hash_ids = nullptr; + ggml_tensor * logits = nullptr; + + void destroy() { + if (ctx) { + ggml_free(ctx); + ctx = nullptr; + } + if (meta_buffer) { + std::free(meta_buffer); + meta_buffer = nullptr; + } + meta_size = 0; + gf = nullptr; + hash_ids = nullptr; + logits = nullptr; + i32_inputs.clear(); + i32_array_inputs.clear(); + i64_array_inputs.clear(); + f32_array_inputs.clear(); + allocated_tensors.clear(); + } +}; + +struct Ds4LayerMajorGraphCache { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + PrefillAttentionMode mode = PrefillAttentionMode::Exact; + int n_tokens = 0; + int kv_start = -1; + bool ready = false; + ggml_context * state_ctx = nullptr; + ggml_backend_buffer_t state_buf = nullptr; + ggml_tensor * state_a = nullptr; + ggml_tensor * state_b = nullptr; + std::vector layers; + + bool matches(const DeepSeek4Weights & w, ggml_backend_t b, + PrefillAttentionMode m, int tokens, int start) const { + return ready && owner_ctx == w.ctx && backend == b && mode == m && + n_tokens == tokens && kv_start == start && + layers.size() == (size_t) w.n_layer; + } + + void destroy() { + for (auto & layer : layers) layer.destroy(); + layers.clear(); + if (state_buf) { + ggml_backend_buffer_free(state_buf); + state_buf = nullptr; + } + if (state_ctx) { + ggml_free(state_ctx); + state_ctx = nullptr; + } + state_a = nullptr; + state_b = nullptr; + owner_ctx = nullptr; + backend = nullptr; + mode = PrefillAttentionMode::Exact; + n_tokens = 0; + kv_start = -1; + ready = false; + } +}; + +static thread_local std::array + ds4_layer_major_graph_caches; +// A separate gallocr scratch arena per shape consumes several GiB and forces +// the 97-GiB model into managed-memory paging. Every layer executes serially, +// so cached and uncached shapes safely rebind their transient tensors to one +// arena before execution. Retain only the largest/full-chunk topology; tail +// metadata is rebuilt rather than keeping a second long-context graph resident. +static thread_local ggml_gallocr_t ds4_layer_major_shared_alloc = nullptr; +static thread_local const ggml_context * ds4_layer_major_shared_owner = nullptr; +static thread_local ggml_backend_t ds4_layer_major_shared_backend = nullptr; +static thread_local std::vector ds4_layer_major_meta_arena; +static thread_local const ggml_context * ds4_layer_major_meta_owner = nullptr; + +static ggml_gallocr_t ds4_layer_major_get_shared_alloc( + const DeepSeek4Weights & w, + ggml_backend_t backend) { + if (ds4_layer_major_shared_alloc && + (ds4_layer_major_shared_owner != w.ctx || + ds4_layer_major_shared_backend != backend)) { + ggml_gallocr_free(ds4_layer_major_shared_alloc); + ds4_layer_major_shared_alloc = nullptr; + } + if (!ds4_layer_major_shared_alloc) { + ds4_layer_major_shared_alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + ds4_layer_major_shared_owner = w.ctx; + ds4_layer_major_shared_backend = backend; + } + return ds4_layer_major_shared_alloc; +} + +void deepseek4_release_runtime_graphs(const DeepSeek4Weights & w) { + const ggml_context * owner = w.ctx; + if (!owner) { + return; + } + + for (auto & cache : ds4_layer_major_graph_caches) { + if (cache.owner_ctx == owner) { + cache.destroy(); + } + } + + if (ds4_layer_major_shared_owner == owner) { + if (ds4_layer_major_shared_alloc) { + ggml_gallocr_free(ds4_layer_major_shared_alloc); + ds4_layer_major_shared_alloc = nullptr; + } + ds4_layer_major_shared_owner = nullptr; + ds4_layer_major_shared_backend = nullptr; + } + if (ds4_layer_major_meta_owner == owner) { + ds4_layer_major_meta_arena.clear(); + ds4_layer_major_meta_arena.shrink_to_fit(); + ds4_layer_major_meta_owner = nullptr; + } + if (ds4_layer_range_runtime.owner_ctx == owner) { + ds4_layer_range_runtime.destroy(); + } + if (ds4_hybrid_runtime.owner_ctx == owner) { + ds4_hybrid_runtime.destroy(); + } +} + +// Returns 1 on success, 0 when the optimized path is not applicable, and -1 +// after a hard failure. +static int ds4_try_layer_major_prefill( + DeepSeek4FusedDecodeCache & fc, + ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const std::vector & hc_weights, + const HcWeightsCpu & hc_out_weights, + const std::vector & hash_tables, + std::vector & hash_scratch, + const float * embed, + int n_tokens, + int kv_start, + std::vector & out_logits, + const int32_t * token_ids, + DeepSeek4StepTelemetry * telemetry) { + if (!backend || !embed || n_tokens <= 4 || + n_tokens > DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS || + kv_start < 0 || w.moe_hybrid) { + return 0; + } + if (cache.prefill_mode == PrefillAttentionMode::Exact) return 0; + if (!ds4_backend_is_gpu(backend) || !hc_out_weights.loaded || + hc_out_weights.scale_data.empty() || !w.output_hc_fn || + !w.output_hc_base) { + return 0; + } + for (int il = 0; il < w.n_layer; ++il) { + const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; + const DeepSeek4Layer & L = w.layers[(size_t) il]; + if (!hlw.attn.loaded || hlw.attn.scale_data.size() < 3 || + !hlw.ffn.loaded || hlw.ffn.scale_data.size() < 3 || + !L.hc_attn_base || !L.hc_ffn_base) { + return 0; + } + } + + if (fc.owner_ctx != w.ctx || fc.backend != backend) { + fc.destroy(); + fc.owner_ctx = w.ctx; + fc.backend = backend; + } + if (!ds4_fused_ensure_fn_mirrors(fc, backend, w, hc_weights, + hc_out_weights)) { + std::fprintf(stderr, + "[deepseek4-prefill] failed to create HC weight mirrors\n"); + return -1; + } + + const int n_embd = w.n_embd; + const int n_hc = w.n_hc; + const int64_t hc_dim = (int64_t) n_embd * n_hc; + const int64_t mix_dim = 2 * (int64_t) n_hc + (int64_t) n_hc * n_hc; + const int next_pos = kv_start + n_tokens; + + Ds4LayerMajorGraphCache * graph_cache = nullptr; + bool cache_hit = false; + bool cache_build = false; + if (token_ids) { + for (auto & candidate : ds4_layer_major_graph_caches) { + if (candidate.matches(w, backend, cache.prefill_mode, + n_tokens, kv_start)) { + graph_cache = &candidate; + cache_hit = true; + break; + } + } + if (!graph_cache) { + auto & candidate = ds4_layer_major_graph_caches.front(); + // Do not evict a full/larger chunk for an equal-size graph at a + // later position or for a short tail. Both execute with the shared + // scratch arena below, but only the dominant topology stays cached. + const bool same_owner = candidate.owner_ctx == w.ctx && + candidate.backend == backend && + candidate.mode == cache.prefill_mode; + if (!candidate.ready || !same_owner || + n_tokens > candidate.n_tokens) { + graph_cache = &candidate; + graph_cache->destroy(); + graph_cache->owner_ctx = w.ctx; + graph_cache->backend = backend; + graph_cache->mode = cache.prefill_mode; + graph_cache->n_tokens = n_tokens; + graph_cache->kv_start = kv_start; + graph_cache->layers.resize((size_t) w.n_layer); + cache_build = true; + } + } + } + + // Persistent ping-pong state lives outside the per-layer gallocr arena. + ggml_context * state_ctx = cache_hit ? graph_cache->state_ctx : nullptr; + ggml_tensor * state_a = cache_hit ? graph_cache->state_a : nullptr; + ggml_tensor * state_b = cache_hit ? graph_cache->state_b : nullptr; + ggml_backend_buffer_t state_buf = cache_hit ? graph_cache->state_buf : nullptr; + if (!cache_hit) { + ggml_init_params state_params{}; + state_params.mem_size = 4 * ggml_tensor_overhead() + 4096; + state_params.no_alloc = true; + state_ctx = ggml_init(state_params); + if (!state_ctx) { + if (cache_build) graph_cache->destroy(); + return -1; + } + state_a = ggml_new_tensor_2d(state_ctx, GGML_TYPE_F32, + hc_dim, n_tokens); + state_b = ggml_new_tensor_2d(state_ctx, GGML_TYPE_F32, + hc_dim, n_tokens); + state_buf = ggml_backend_alloc_ctx_tensors(state_ctx, backend); + if (!state_buf) { + ggml_free(state_ctx); + if (cache_build) graph_cache->destroy(); + return -1; + } + if (cache_build) { + graph_cache->state_ctx = state_ctx; + graph_cache->state_a = state_a; + graph_cache->state_b = state_b; + graph_cache->state_buf = state_buf; + } + } + + std::vector initial((size_t) hc_dim * n_tokens); + for (int t = 0; t < n_tokens; ++t) { + for (int h = 0; h < n_hc; ++h) { + std::memcpy(initial.data() + (size_t) t * hc_dim + + (size_t) h * n_embd, + embed + (size_t) t * n_embd, + sizeof(float) * (size_t) n_embd); + } + } + ggml_backend_tensor_set(state_a, initial.data(), 0, + sizeof(float) * initial.size()); + initial.clear(); + initial.shrink_to_fit(); + + ggml_gallocr_t alloc = ds4_layer_major_get_shared_alloc(w, backend); + if (!alloc) { + if (cache_build) { + graph_cache->destroy(); + } else { + ggml_backend_buffer_free(state_buf); + ggml_free(state_ctx); + } + return -1; + } + const size_t meta_bytes = 160u * 1024 * 1024; + if (ds4_layer_major_meta_owner != w.ctx) { + ds4_layer_major_meta_arena.clear(); + ds4_layer_major_meta_arena.shrink_to_fit(); + ds4_layer_major_meta_owner = w.ctx; + } + if (!graph_cache && ds4_layer_major_meta_arena.size() < meta_bytes) { + ds4_layer_major_meta_arena.resize(meta_bytes); + } + + auto fail = [&](const char * what, int il) { + std::fprintf(stderr, "[deepseek4-prefill] %s at layer %d\n", what, il); + if (graph_cache) { + graph_cache->destroy(); + } else { + ggml_backend_buffer_free(state_buf); + ggml_free(state_ctx); + } + return -1; + }; + + if (cache_hit) { + for (int il = 0; il < w.n_layer; ++il) { + Ds4LayerMajorCachedLayer & layer = + graph_cache->layers[(size_t) il]; + for (ggml_tensor * tensor : layer.allocated_tensors) { + tensor->data = nullptr; + tensor->buffer = nullptr; + } + const auto alloc_t0 = Ds4TimingClock::now(); + if (!layer.ctx || !layer.gf || + !ggml_gallocr_alloc_graph(alloc, layer.gf)) { + return fail("cached scratch allocation failed", il); + } + if (telemetry) { + telemetry->full_graph_alloc_us += ds4_elapsed_us( + alloc_t0, Ds4TimingClock::now()); + } + for (const auto & b : layer.i32_inputs) { + ggml_backend_tensor_set(b.tensor, &b.value, 0, + sizeof(b.value)); + } + for (const auto & b : layer.i32_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(int32_t) * b.values.size()); + } + for (const auto & b : layer.i64_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(int64_t) * b.values.size()); + } + for (const auto & b : layer.f32_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(float) * b.values.size()); + } + if (layer.hash_ids) { + const int n_used = w.n_expert_used; + hash_scratch.resize((size_t) n_used * n_tokens); + const auto & table = hash_tables[(size_t) il].ids; + for (int t = 0; t < n_tokens; ++t) { + std::memcpy( + hash_scratch.data() + (size_t) t * n_used, + table.data() + (size_t) token_ids[t] * n_used, + sizeof(int32_t) * (size_t) n_used); + } + ggml_backend_tensor_set( + layer.hash_ids, hash_scratch.data(), 0, + sizeof(int32_t) * hash_scratch.size()); + } + + const auto compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, layer.gf) != + GGML_STATUS_SUCCESS) { + return fail("cached compute failed", il); + } + if (telemetry) { + telemetry->full_graph_compute_us += ds4_elapsed_us( + compute_t0, Ds4TimingClock::now()); + } + if (layer.logits) { + out_logits.resize((size_t) w.n_vocab); + ggml_backend_tensor_get( + layer.logits, out_logits.data(), 0, + sizeof(float) * (size_t) w.n_vocab); + } + + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + const int ratio = (int) w.compress_ratios[(size_t) il]; + if (ratio > 0) { + lc.n_comp = std::max(lc.n_comp, next_pos / ratio); + if (ratio == 4) { + lc.n_index_comp = std::max( + lc.n_index_comp, next_pos / ratio); + } + } + } + cache.cur_pos = next_pos; + return out_logits.empty() ? -1 : 1; + } + + ggml_tensor * state_in = state_a; + ggml_tensor * state_out = state_b; + for (int il = 0; il < w.n_layer; ++il) { + const auto build_t0 = Ds4TimingClock::now(); + Ds4LayerMajorCachedLayer * cached_layer = cache_build + ? &graph_cache->layers[(size_t) il] : nullptr; + ggml_init_params params{}; + if (cached_layer) { + cached_layer->meta_size = meta_bytes; + cached_layer->meta_buffer = std::malloc(meta_bytes); + if (!cached_layer->meta_buffer) { + return fail("cached metadata allocation failed", il); + } + params.mem_size = cached_layer->meta_size; + params.mem_buffer = cached_layer->meta_buffer; + } else { + params.mem_size = ds4_layer_major_meta_arena.size(); + params.mem_buffer = ds4_layer_major_meta_arena.data(); + } + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) return fail("metadata allocation failed", il); + if (cached_layer) cached_layer->ctx = ctx; + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 65536, false); + if (!gf) { + if (!cached_layer) ggml_free(ctx); + return fail("graph allocation failed", il); + } + if (cached_layer) cached_layer->gf = gf; + + const DeepSeek4Layer & L = w.layers[(size_t) il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; + + // HC pre -> batched attention. + ggml_tensor * norm_hc = ggml_rms_norm(ctx, state_in, w.hc_eps); + ggml_tensor * mix_attn = ggml_mul_mat(ctx, + fc.fn_attn_f16[(size_t) il], norm_hc); + mix_attn = ggml_reshape_2d(ctx, mix_attn, mix_dim, n_tokens); + ggml_tensor * attn_base = ds4_fused_hc_base_f32(ctx, + L.hc_attn_base); + ggml_tensor * pre_attn = ggml_ds4_hc_pre( + ctx, mix_attn, attn_base, state_in, n_hc, + w.n_hc_sinkhorn_iter, hlw.attn.scale_data[0], + hlw.attn.scale_data[1], hlw.attn.scale_data[2]); + ggml_tensor * attn_in = ggml_view_2d(ctx, pre_attn, n_embd, + n_tokens, pre_attn->nb[1], 0); + ggml_tensor * split_attn = ggml_view_2d( + ctx, pre_attn, mix_dim, n_tokens, pre_attn->nb[1], + (size_t) n_embd * sizeof(float)); + + std::vector i32_inputs; + std::vector i32_array_inputs; + std::vector i64_array_inputs; + std::vector f32_array_inputs; + ggml_tensor * attn_normed = build_rms_norm(ctx, attn_in, + L.attn_norm, w.rms_eps); + const DeepSeek4AttentionImpl attention_impl = + cache.prefill_mode == PrefillAttentionMode::Sparse + ? DeepSeek4AttentionImpl::SparseFlash + : DeepSeek4AttentionImpl::DenseFlash; + ggml_tensor * attn_out = build_mla_attention( + ctx, gf, attn_normed, w, L, lc, il, kv_start, n_tokens, + nullptr, i32_inputs, i32_array_inputs, i64_array_inputs, + &f32_array_inputs, attention_impl); + if (!attn_out) { + if (!cached_layer) ggml_free(ctx); + return fail("attention graph build failed", il); + } + ggml_tensor * hc_after_attn = ggml_ds4_hc_post( + ctx, state_in, attn_out, split_attn, n_hc); + + // HC pre -> batched MoE. + norm_hc = ggml_rms_norm(ctx, hc_after_attn, w.hc_eps); + ggml_tensor * mix_ffn = ggml_mul_mat(ctx, + fc.fn_ffn_f16[(size_t) il], norm_hc); + mix_ffn = ggml_reshape_2d(ctx, mix_ffn, mix_dim, n_tokens); + ggml_tensor * ffn_base = ds4_fused_hc_base_f32(ctx, L.hc_ffn_base); + ggml_tensor * pre_ffn = ggml_ds4_hc_pre( + ctx, mix_ffn, ffn_base, hc_after_attn, n_hc, + w.n_hc_sinkhorn_iter, hlw.ffn.scale_data[0], + hlw.ffn.scale_data[1], hlw.ffn.scale_data[2]); + ggml_tensor * ffn_in = ggml_view_2d(ctx, pre_ffn, n_embd, + n_tokens, pre_ffn->nb[1], 0); + ggml_tensor * split_ffn = ggml_view_2d( + ctx, pre_ffn, mix_dim, n_tokens, pre_ffn->nb[1], + (size_t) n_embd * sizeof(float)); + ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, + L.ffn_norm, w.rms_eps); + ggml_tensor * hash_ids = nullptr; + ggml_tensor * ffn_out = nullptr; + const bool hash_routed = il < w.n_hash_layer && L.ffn_gate_tid2eid && + token_ids && hash_tables[(size_t) il].loaded; + if (hash_routed) { + hash_ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, + w.n_expert_used, n_tokens); + ggml_set_input(hash_ids); + ffn_out = ds4_build_hash_routed_ffn( + ctx, w, L, ffn_normed, hash_ids, n_tokens); + } else { + ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, n_tokens); + } + if (!ffn_out) { + if (!cached_layer) ggml_free(ctx); + return fail("FFN graph build failed", il); + } + ggml_tensor * hc_next = ggml_ds4_hc_post( + ctx, hc_after_attn, ffn_out, split_ffn, n_hc); + + // Persist HC state for the next layer before this layer's gallocr + // scratch buffer is reused. + ggml_tensor * state_copy = ggml_cpy(ctx, hc_next, state_out); + ggml_set_output(state_copy); + ggml_build_forward_expand(gf, state_copy); + + ggml_tensor * logits = nullptr; + if (il + 1 == w.n_layer) { + ggml_tensor * last_hc = ggml_view_2d( + ctx, hc_next, hc_dim, 1, hc_next->nb[1], + (size_t) (n_tokens - 1) * hc_next->nb[1]); + last_hc = ggml_reshape_1d(ctx, last_hc, hc_dim); + ggml_tensor * out_hc_norm = ggml_rms_norm(ctx, last_hc, w.hc_eps); + ggml_tensor * out_mix = ggml_mul_mat(ctx, fc.fn_out_f16, + out_hc_norm); + out_mix = ggml_reshape_1d(ctx, out_mix, n_hc); + ggml_tensor * out_base = ds4_fused_hc_base_f32(ctx, + w.output_hc_base); + ggml_tensor * final_embd = ggml_ds4_hc_out( + ctx, out_mix, out_base, last_hc, n_hc, + hc_out_weights.scale_data[0]); + ggml_tensor * final_2d = ggml_reshape_2d(ctx, final_embd, + n_embd, 1); + ggml_tensor * out_normed = build_rms_norm(ctx, final_2d, + w.out_norm, w.rms_eps); + logits = ggml_mul_mat(ctx, w.output, out_normed); + ggml_set_output(logits); + ggml_build_forward_expand(gf, logits); + } + + if (cached_layer) { + auto remember_unallocated = [&](ggml_tensor * tensor) { + if (tensor && tensor->data == nullptr && + tensor->buffer == nullptr) { + cached_layer->allocated_tensors.push_back(tensor); + } + }; + const int n_graph_nodes = ggml_graph_n_nodes(gf); + for (int i = 0; i < n_graph_nodes; ++i) { + ggml_tensor * node = ggml_graph_node(gf, i); + remember_unallocated(node); + for (int j = 0; j < GGML_MAX_SRC; ++j) { + remember_unallocated(node->src[j]); + } + } + auto & tensors = cached_layer->allocated_tensors; + std::sort(tensors.begin(), tensors.end()); + tensors.erase(std::unique(tensors.begin(), tensors.end()), + tensors.end()); + } + if (!ggml_gallocr_alloc_graph(alloc, gf)) { + if (!cached_layer) ggml_free(ctx); + return fail("scratch allocation failed", il); + } + for (const auto & b : i32_inputs) { + ggml_backend_tensor_set(b.tensor, &b.value, 0, sizeof(b.value)); + } + for (const auto & b : i32_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(int32_t) * b.values.size()); + } + for (const auto & b : i64_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(int64_t) * b.values.size()); + } + for (const auto & b : f32_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(float) * b.values.size()); + } + if (hash_ids) { + const int n_used = w.n_expert_used; + hash_scratch.resize((size_t) n_used * n_tokens); + const auto & table = hash_tables[(size_t) il].ids; + for (int t = 0; t < n_tokens; ++t) { + std::memcpy(hash_scratch.data() + (size_t) t * n_used, + table.data() + (size_t) token_ids[t] * n_used, + sizeof(int32_t) * (size_t) n_used); + } + ggml_backend_tensor_set(hash_ids, hash_scratch.data(), 0, + sizeof(int32_t) * hash_scratch.size()); + } + if (telemetry) { + telemetry->full_graph_build_us += ds4_elapsed_us( + build_t0, Ds4TimingClock::now()); + } + + const auto compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { + if (!cached_layer) ggml_free(ctx); + return fail("compute failed", il); + } + if (telemetry) { + telemetry->full_graph_compute_us += ds4_elapsed_us( + compute_t0, Ds4TimingClock::now()); + } + + if (logits) { + out_logits.resize((size_t) w.n_vocab); + ggml_backend_tensor_get(logits, out_logits.data(), 0, + sizeof(float) * (size_t) w.n_vocab); + } + + const int ratio = (int) w.compress_ratios[(size_t) il]; + if (ratio > 0) { + lc.n_comp = std::max(lc.n_comp, next_pos / ratio); + if (ratio == 4) { + lc.n_index_comp = std::max(lc.n_index_comp, + next_pos / ratio); + } + } + if (cached_layer) { + cached_layer->i32_inputs = std::move(i32_inputs); + cached_layer->i32_array_inputs = std::move(i32_array_inputs); + cached_layer->i64_array_inputs = std::move(i64_array_inputs); + cached_layer->f32_array_inputs = std::move(f32_array_inputs); + cached_layer->hash_ids = hash_ids; + cached_layer->logits = logits; + } else { + ggml_free(ctx); + } + std::swap(state_in, state_out); + } + + if (cache_build) { + graph_cache->ready = true; + } else { + ggml_backend_buffer_free(state_buf); + ggml_free(state_ctx); + } + cache.cur_pos = next_pos; + return out_logits.empty() ? -1 : 1; } @@ -4609,23 +6153,27 @@ bool deepseek4_step_layer_range( memcpy(hc_state.data(), embed, sizeof(float) * (size_t)hc_dim * (size_t)n_tokens); } - // Lazy-load per-layer HC weights on CPU (static to avoid reloading) - static std::vector hc_layer_weights_range; - static HcWeightsCpu hc_output_weights_range; - static std::vector hash_routing_tables_range; - static std::vector cached_attn_allocs; - static std::vector cached_decode_attn_hc_pre_graphs; - static std::vector cached_decode_ffn_hc_pre_graphs; - static DeepSeek4CachedDecodeHcPostGraph cached_decode_hc_post_graph; - static std::vector> cached_decode_attn_graphs; - static std::vector cached_decode_ffn_graphs; - static DeepSeek4CachedDecodeOutputGraph cached_decode_output_graph; - static DeepSeek4CachedLayerAlloc cached_dynamic_output_alloc; - static DeepSeek4FusedDecodeCache fused_decode_graph_cache; - static Ds4DecodeSharedInputs decode_shared_inputs; - static int hc_loaded_n_layer = 0; - static const ggml_context * hc_loaded_ctx = nullptr; - static uint64_t hc_loaded_generation = 0; + // Keep all reusable layer-range resources under one model owner so park, + // unload, and reload have a deterministic teardown boundary. + DeepSeek4LayerRangeRuntime & runtime = ds4_layer_range_runtime; + auto & hc_layer_weights_range = runtime.hc_layer_weights; + auto & hc_output_weights_range = runtime.hc_output_weights; + auto & hash_routing_tables_range = runtime.hash_routing_tables; + auto & cached_attn_allocs = runtime.cached_attn_allocs; + auto & cached_decode_attn_hc_pre_graphs = + runtime.cached_decode_attn_hc_pre_graphs; + auto & cached_decode_ffn_hc_pre_graphs = + runtime.cached_decode_ffn_hc_pre_graphs; + auto & cached_decode_hc_post_graph = runtime.cached_decode_hc_post_graph; + auto & cached_decode_attn_graphs = runtime.cached_decode_attn_graphs; + auto & cached_decode_ffn_graphs = runtime.cached_decode_ffn_graphs; + auto & cached_decode_output_graph = runtime.cached_decode_output_graph; + auto & cached_dynamic_output_alloc = runtime.cached_dynamic_output_alloc; + auto & fused_decode_graph_cache = runtime.fused_decode_graph_cache; + auto & decode_shared_inputs = runtime.decode_shared_inputs; + int & hc_loaded_n_layer = runtime.loaded_n_layer; + uint64_t & hc_loaded_generation = runtime.loaded_generation; + const ggml_context * & hc_loaded_ctx = runtime.owner_ctx; g_deepseek4_fused_decode_runtime_cache = &fused_decode_graph_cache; if (hc_loaded_n_layer != w.n_layer || hc_loaded_ctx != w.ctx || hc_loaded_generation != w.runtime_generation) { @@ -4675,10 +6223,31 @@ bool deepseek4_step_layer_range( } // Per-layer execution with CPU-side HC - static thread_local DeepSeek4LayerRangeScratch scratch; + DeepSeek4LayerRangeScratch & scratch = runtime.scratch; const int n_expert_used = ds4_effective_expert_count(w); scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, n_expert_used); + // Large monolithic batches use PR520's device-resident layer-major + // pipeline. Heterogeneous batches are enabled separately once its expert + // stage has been bound to both owner backends. + if (!moe_hybrid && n_tokens > 4 && + n_tokens <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS && layer_begin == 0 && + is_last_shard && out_logits && ds4_backend_is_gpu(backend)) { + const int prc = ds4_try_layer_major_prefill( + fused_decode_graph_cache, backend, w, cache, + hc_layer_weights_range, hc_output_weights_range, + hash_routing_tables_range, scratch.hash_expert_ids, embed, + n_tokens, kv_start, *out_logits, token_ids, telemetry); + if (prc < 0) return false; + if (prc > 0) { + if (telemetry) { + telemetry->total_us += ds4_elapsed_us( + step_t0, Ds4TimingClock::now()); + } + return true; + } + } + // The batched verifier graph is also the only whole-model graph that can // currently own tensors on both GPU backends. Reuse it for q=1 hybrid // decode so native AR does not fall back to 43 host-synchronized FFN @@ -4905,7 +6474,20 @@ bool deepseek4_step_layer_range( ggml_context * ctx = nullptr; DeepSeek4CachedDecodeAttnGraph * cached_attn = nullptr; - if (reuse_decode_attn) { + const bool exact_tokenwise_prefill = + !reuse_decode_attn && n_tokens > 1; + if (exact_tokenwise_prefill) { + const DeepSeek4AttentionImpl attention_impl = + cache.prefill_mode == PrefillAttentionMode::Sparse + ? DeepSeek4AttentionImpl::SparseFlash + : DeepSeek4AttentionImpl::Explicit; + if (!ds4_run_exact_tokenwise_prefill_attention( + backend, w, L, lc, il, cur.data(), n_tokens, kv_start, + attention_impl, attn_out_host, + cached_attn_allocs[(size_t) il], telemetry)) { + return false; + } + } else if (reuse_decode_attn) { const int n_raw = std::min(kv_start + 1, w.n_swa); const int n_comp_attn = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; const int n_index_comp = (ratio == 4) ? ds4_comp_rows_used(lc.index_comp_kv, lc.n_index_comp, 4, token_pos) : 0; @@ -5006,10 +6588,16 @@ bool deepseek4_step_layer_range( gf = ggml_new_graph_custom(ctx, graph_size, false); ggml_tensor * normed = build_rms_norm(ctx, inp, L.attn_norm, w.rms_eps); + const DeepSeek4AttentionImpl attention_impl = + cache.prefill_mode == PrefillAttentionMode::Sparse + ? DeepSeek4AttentionImpl::SparseFlash + : DeepSeek4AttentionImpl::Explicit; attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, il, kv_start, n_tokens, nullptr, i32_inputs, i32_array_inputs, - i64_array_inputs, &f32_array_inputs); + i64_array_inputs, + &f32_array_inputs, + attention_impl); ggml_set_output(attn_out); ggml_build_forward_expand(gf, attn_out); @@ -5041,6 +6629,7 @@ bool deepseek4_step_layer_range( ggml_backend_tensor_set(b.tensor, b.values.data(), 0, sizeof(float) * b.values.size()); } + if (!exact_tokenwise_prefill) { const auto attn_compute_t0 = Ds4TimingClock::now(); if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[deepseek4] attn compute failed layer %d\n", il); @@ -5069,6 +6658,7 @@ bool deepseek4_step_layer_range( if (telemetry) telemetry->attn_read_us += ds4_elapsed_us(attn_read_t0, Ds4TimingClock::now()); } if (ctx) ggml_free(ctx); + } // ── HC post (attention) ───────────────────────────────── if (!(use_backend_decode_hc_graph || use_backend_decode_hc_direct)) { @@ -5350,7 +6940,10 @@ bool deepseek4_step_layer_range( ggml_context * ctx = ggml_init(params); if (!ctx) return false; - ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); + const bool last_only = n_tokens > 1; + const int output_tokens = last_only ? 1 : n_tokens; + ggml_tensor * inp = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_embd, output_tokens); ggml_set_input(inp); ggml_tensor * normed = build_rms_norm(ctx, inp, w.out_norm, w.rms_eps); ggml_tensor * logits = ggml_mul_mat(ctx, w.output, normed); @@ -5371,14 +6964,19 @@ bool deepseek4_step_layer_range( ggml_free(ctx); return false; } - ggml_backend_tensor_set(inp, final_embd.data(), 0, sizeof(float) * final_embd.size()); + const float * output_input = last_only + ? final_embd.data() + (size_t)(n_tokens - 1) * n_embd + : final_embd.data(); + ggml_backend_tensor_set(inp, output_input, 0, + sizeof(float) * (size_t)n_embd * output_tokens); if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { ggml_free(ctx); return false; } out_logits->resize((size_t)w.n_vocab); - const size_t logits_offset = (size_t)(n_tokens - 1) * (size_t)w.n_vocab * sizeof(float); + const size_t logits_offset = last_only ? 0 : + (size_t)(n_tokens - 1) * (size_t)w.n_vocab * sizeof(float); ggml_backend_tensor_get(logits, out_logits->data(), logits_offset, sizeof(float) * (size_t)w.n_vocab); if (verify_hooks && verify_hooks->all_logits_out) { @@ -5395,11 +6993,12 @@ bool deepseek4_step_layer_range( memcpy(out_logits->data(), hc_state.data(), sizeof(float) * hc_dim * n_tokens); } - // Update compressor state + // Update compressor state. Multi-token prefill may cross one or more + // boundaries even when the chunk itself does not end on a boundary. const int next_pos = kv_start + n_tokens; for (int il = layer_begin; il < layer_end; ++il) { const uint32_t ratio = w.compress_ratios[il]; - if (ratio <= 0 || (next_pos % (int)ratio) != 0) continue; + if (ratio <= 0) continue; cache.layers[il].n_comp = std::max(cache.layers[il].n_comp, next_pos / (int)ratio); if (ratio == 4) { cache.layers[il].n_index_comp = std::max(cache.layers[il].n_index_comp, diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 92c74c3dc..ea7db46ac 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -24,9 +24,15 @@ #include "internal.h" #include "common/layer_split_utils.h" +#include "common/prefill_attention_mode.h" namespace dflash::common { +// Layer-major prefill may schedule two 2K numerical bands while preserving +// the raw-cache rounding boundary between them. +inline constexpr int DS4_NUMERICAL_PREFILL_BAND = 2048; +inline constexpr int DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS = 4096; + struct MoeHybridPlacement; struct MoeHybridConfig; struct MoeHybridRoutingStats; @@ -66,6 +72,7 @@ struct DeepSeek4StepTelemetry { uint64_t sample_us = 0; uint64_t emit_us = 0; uint64_t full_graph_build_us = 0; + uint64_t full_graph_alloc_us = 0; uint64_t full_graph_set_us = 0; uint64_t full_graph_compute_us = 0; uint64_t full_graph_read_us = 0; @@ -272,6 +279,7 @@ struct DeepSeek4Cache { int n_layer = 0; std::vector layers; + PrefillAttentionMode prefill_mode = PrefillAttentionMode::Exact; // HC residual streams: [n_hc * n_embd] persistent state ggml_tensor * hc_state = nullptr; // [n_hc * n_embd] @@ -294,6 +302,7 @@ struct DeepSeek4BackendConfig { DevicePlacement device; int stream_fd = -1; int chunk = 512; // prefill chunk size + PrefillAttentionMode prefill_mode = PrefillAttentionMode::Exact; int max_ctx = 0; // 0 = auto from SWA + compression capacity int expert_top_k = 0; // 0 = use all model-routed experts bool fused_decode = false; // single-graph GPU decode @@ -312,6 +321,10 @@ bool load_deepseek4_gguf_partial(const std::string & path, void free_deepseek4_weights(DeepSeek4Weights & w); +// Release graph allocators and host mirrors that retain model tensor pointers. +// This must run before the owning ggml context is destroyed. +void deepseek4_release_runtime_graphs(const DeepSeek4Weights & w); + bool create_deepseek4_cache(ggml_backend_t backend, const DeepSeek4Weights & w, int max_ctx, diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index 865ee2450..663db3488 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -1000,6 +1000,7 @@ bool build_deepseek4_moe_hybrid_storage_from_file( } void free_deepseek4_weights(DeepSeek4Weights & w) { + deepseek4_release_runtime_graphs(w); if (w.ctx) { ggml_free(w.ctx); w.ctx = nullptr; } if (w.dense_split_buf) { ggml_backend_buffer_free(w.dense_split_buf); diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 9c9c30f76..c1f634ae8 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -221,6 +221,9 @@ static void print_usage(const char * prog) { " --ds4-expert-top-k \n" " Keep and renormalize the highest-ranked N routed experts\n" " (0=model default; single-device DeepSeek4 only)\n" + " --ds4-prefill DeepSeek4 prefill: exact, dense, or sparse\n" + " (default: exact; dense/sparse are experimental\n" + " and may change generated tokens)\n" " --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" @@ -435,6 +438,20 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "[server] --ds4-expert-top-k must be non-negative\n"); return 2; } + } else if (std::strcmp(argv[i], "--ds4-prefill") == 0 && i + 1 < argc) { + const char * mode = argv[++i]; + bargs.ds4_prefill_mode_set = true; + if (std::strcmp(mode, "exact") == 0) { + bargs.ds4_prefill_mode = PrefillAttentionMode::Exact; + } else if (std::strcmp(mode, "dense") == 0) { + bargs.ds4_prefill_mode = PrefillAttentionMode::Dense; + } else if (std::strcmp(mode, "sparse") == 0) { + bargs.ds4_prefill_mode = PrefillAttentionMode::Sparse; + } else { + std::fprintf(stderr, + "[server] --ds4-prefill expects exact, dense, or sparse\n"); + return 2; + } } else if (std::strcmp(argv[i], "--fa-window") == 0 && i + 1 < argc) { bargs.fa_window = std::atoi(argv[++i]); } else if (std::strcmp(argv[i], "--model-name") == 0 && i + 1 < argc) { @@ -798,6 +815,12 @@ int main(int argc, char ** argv) { g_peer_access_opt_in = bargs.device.peer_access; std::fprintf(stderr, "[server] creating backend...\n"); const std::string arch = detect_arch(bargs.model_path); + if (bargs.ds4_prefill_mode_set && arch != "deepseek4") { + std::fprintf(stderr, + "[server] --ds4-prefill is only valid for deepseek4 models " + "(detected '%s')\n", arch.c_str()); + return 2; + } const PlacementBackend target_backend = bargs.device.backend == PlacementBackend::Auto ? compiled_placement_backend() @@ -1109,6 +1132,8 @@ int main(int argc, char ** argv) { } else { std::fprintf(stderr, "[server] │ ds4_expert_topk= model default\n"); } + std::fprintf(stderr, "[server] │ ds4_prefill = %s\n", + prefill_attention_mode_name(bargs.ds4_prefill_mode)); } std::fprintf(stderr, "[server] │ fa_window = %d\n", bargs.fa_window); if (bargs.fa_window > 0) { diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 1cd261a95..0aa7a73ab 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -723,6 +723,357 @@ static void test_grouped_output_projection_shape() { std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); } +static void test_grouped_output_projection_cpu(ggml_backend_t backend) { + std::fprintf(stderr, " test_grouped_output_projection_cpu ..."); + + constexpr int group_width = 4; + constexpr int n_groups = 3; + constexpr int n_tokens = 5; + constexpr int n_outputs = 4; + constexpr int flat_width = group_width * n_groups; + + std::vector weights((size_t) flat_width * n_outputs); + std::vector grouped((size_t) group_width * n_tokens * n_groups); + std::vector expected((size_t) n_outputs * n_tokens, 0.0f); + for (int output = 0; output < n_outputs; ++output) { + for (int k = 0; k < flat_width; ++k) { + weights[(size_t) output * flat_width + k] = + 0.03125f * (float) ((output + 1) * (k - 5)); + } + } + for (int group = 0; group < n_groups; ++group) { + for (int token = 0; token < n_tokens; ++token) { + for (int k = 0; k < group_width; ++k) { + grouped[(size_t) group * n_tokens * group_width + + (size_t) token * group_width + k] = + 10.0f * (float) group + 0.5f * (float) token + + 0.125f * (float) k; + } + } + } + for (int token = 0; token < n_tokens; ++token) { + for (int output = 0; output < n_outputs; ++output) { + float sum = 0.0f; + for (int group = 0; group < n_groups; ++group) { + for (int k = 0; k < group_width; ++k) { + const float value = + grouped[(size_t) group * n_tokens * group_width + + (size_t) token * group_width + k]; + sum += weights[(size_t) output * flat_width + + group * group_width + k] * value; + } + } + expected[(size_t) token * n_outputs + output] = sum; + } + } + + ggml_context * ctx = make_test_context(); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) return; + ggml_tensor * weights_t = + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, flat_width, n_outputs); + ggml_tensor * grouped_t = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, group_width, n_tokens, n_groups); + ggml_set_input(weights_t); + ggml_set_input(grouped_t); + ggml_tensor * output_t = + ggml_mul_mat_grouped_src(ctx, weights_t, grouped_t); + TEST_ASSERT_MSG(output_t->op == GGML_OP_MUL_MAT_GROUPED_SRC, + "grouped projection must use a distinct backend contract"); + ggml_set_output(output_t); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, output_t); + ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_cpu_buffer_type()); + TEST_ASSERT_MSG(ggml_gallocr_alloc_graph(alloc, graph), + "grouped projection graph allocation failed"); + ggml_backend_tensor_set(weights_t, weights.data(), 0, + weights.size() * sizeof(float)); + ggml_backend_tensor_set(grouped_t, grouped.data(), 0, + grouped.size() * sizeof(float)); + TEST_ASSERT_MSG(ggml_backend_graph_compute(backend, graph) == + GGML_STATUS_SUCCESS, + "grouped projection graph compute failed"); + std::vector actual(expected.size()); + ggml_backend_tensor_get(output_t, actual.data(), 0, + actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + TEST_ASSERT_MSG(nearly_equal(actual[i], expected[i]), + "grouped projection output mismatch"); + } + ggml_gallocr_free(alloc); + ggml_free(ctx); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + +static void test_ds4_flash_attention_cpu_rejected(ggml_backend_t backend) { + std::fprintf(stderr, " test_ds4_flash_attention_cpu_rejected ..."); + ggml_context * ctx = make_test_context(); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) return; + ggml_tensor * q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 512, 2, 4); + ggml_tensor * k = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 512, 8, 1); + ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 8, 2); + ggml_tensor * op = ggml_flash_attn_ext( + ctx, q, k, k, mask, 1.0f / std::sqrt(512.0f), 0.0f, 0.0f); + ggml_flash_attn_ext_set_ds4_sparse(op, 4, 4, 0, 4); + TEST_ASSERT(ggml_flash_attn_ext_is_ds4(op)); + TEST_ASSERT_MSG(!ggml_backend_supports_op(backend, op), + "CPU accepted the DS4 flash-attention contract"); + ggml_free(ctx); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + +static float reference_e2m1_round(float value) { + static constexpr float levels[] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + }; + const float sign = value < 0.0f ? -1.0f : 1.0f; + const float magnitude = std::min(std::fabs(value), 6.0f); + int best = 0; + float best_diff = std::fabs(magnitude - levels[0]); + for (int i = 1; i < 8; ++i) { + const float diff = std::fabs(magnitude - levels[i]); + if (diff < best_diff || + (diff == best_diff && (i & 1) == 0 && (best & 1) != 0)) { + best = i; + best_diff = diff; + } + } + return sign * levels[best]; +} + +static void reference_ds4_indexer_qat(float * row) { + for (int stride = 1; stride < 128; stride <<= 1) { + for (int base = 0; base < 128; base += 2 * stride) { + for (int i = 0; i < stride; ++i) { + const float a = row[base + i]; + const float b = row[base + stride + i]; + row[base + i] = a + b; + row[base + stride + i] = a - b; + } + } + } + constexpr float inv_sqrt_128 = 0.08838834764831845f; + for (int i = 0; i < 128; ++i) row[i] *= inv_sqrt_128; + for (int block = 0; block < 4; ++block) { + float amax = 0.0f; + for (int i = 0; i < 32; ++i) { + amax = std::max(amax, std::fabs(row[block * 32 + i])); + } + amax = std::max(amax, 7.052966104933725e-38f); + const float scale = std::exp2(std::ceil(std::log2(amax / 6.0f))); + for (int i = 0; i < 32; ++i) { + const int index = block * 32 + i; + const float normalized = + std::clamp(row[index] / scale, -6.0f, 6.0f); + row[index] = reference_e2m1_round(normalized) * scale; + } + } +} + +static void test_indexer_qat_cpu(ggml_backend_t backend) { + std::fprintf(stderr, " test_indexer_qat_cpu ..."); + constexpr int n_head = 3; + constexpr int n_tokens = 2; + std::vector input((size_t) 128 * n_head * n_tokens); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = 0.35f * std::sin(0.17f * (float) i) + + 0.08f * (float) ((int) (i % 11) - 5); + } + std::vector expected = input; + for (int row = 0; row < n_head * n_tokens; ++row) { + reference_ds4_indexer_qat(expected.data() + (size_t) row * 128); + } + ggml_context * ctx = make_test_context(); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) return; + ggml_tensor * input_t = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, 128, n_head, n_tokens); + ggml_set_input(input_t); + ggml_tensor * output_t = ggml_ds4_indexer_qat(ctx, input_t); + ggml_set_output(output_t); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, output_t); + ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_cpu_buffer_type()); + TEST_ASSERT_MSG(ggml_gallocr_alloc_graph(alloc, graph), + "indexer QAT graph allocation failed"); + ggml_backend_tensor_set(input_t, input.data(), 0, + input.size() * sizeof(float)); + TEST_ASSERT_MSG(ggml_backend_graph_compute(backend, graph) == + GGML_STATUS_SUCCESS, + "indexer QAT graph compute failed"); + std::vector actual(expected.size()); + ggml_backend_tensor_get(output_t, actual.data(), 0, + actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + TEST_ASSERT_MSG(nearly_equal(actual[i], expected[i], 1e-6f, 1e-6f), + "indexer QAT output mismatch"); + } + ggml_gallocr_free(alloc); + ggml_free(ctx); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + +static void test_indexer_score_cpu(ggml_backend_t backend) { + std::fprintf(stderr, " test_indexer_score_cpu ..."); + constexpr int n_head = 3; + constexpr int n_tokens = 6; + constexpr int n_comp = 4; + constexpr int kv_start = 2; + constexpr int ratio = 4; + std::vector q((size_t) 128 * n_head * n_tokens); + std::vector weights((size_t) n_head * n_tokens); + std::vector comp_f32((size_t) 128 * n_comp); + std::vector comp_f16(comp_f32.size()); + for (size_t i = 0; i < q.size(); ++i) { + q[i] = 0.025f * (float) ((int) (i % 29) - 14); + } + for (int token = 0; token < n_tokens; ++token) { + for (int head = 0; head < n_head; ++head) { + weights[(size_t) token * n_head + head] = + 0.25f + 0.1f * (float) head + 0.03f * (float) token; + } + } + for (int comp = 0; comp < n_comp; ++comp) { + for (int d = 0; d < 128; ++d) { + comp_f32[(size_t) comp * 128 + d] = + 0.04f * (float) (((comp + 2) * (d + 3)) % 23 - 11); + } + } + ggml_fp32_to_fp16_row(comp_f32.data(), comp_f16.data(), + (int64_t) comp_f32.size()); + std::vector expected((size_t) n_comp * n_tokens); + for (int token = 0; token < n_tokens; ++token) { + const int visible = (kv_start + token + 1) / ratio; + for (int comp = 0; comp < n_comp; ++comp) { + if (comp >= visible) { + expected[(size_t) token * n_comp + comp] = -1.0e30f; + continue; + } + float score = 0.0f; + for (int head = 0; head < n_head; ++head) { + const float * q_row = q.data() + + ((size_t) token * n_head + head) * 128; + const ggml_fp16_t * k_row = + comp_f16.data() + (size_t) comp * 128; + float dot = 0.0f; + for (int d = 0; d < 128; ++d) { + dot += q_row[d] * ggml_fp16_to_fp32(k_row[d]); + } + score += std::max(dot, 0.0f) * + weights[(size_t) token * n_head + head]; + } + expected[(size_t) token * n_comp + comp] = score; + } + } + + ggml_context * ctx = make_test_context(); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) return; + ggml_tensor * q_t = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, 128, n_head, n_tokens); + ggml_tensor * weights_t = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_head, n_tokens); + ggml_tensor * comp_t = ggml_new_tensor_2d( + ctx, GGML_TYPE_F16, 128, n_comp); + ggml_set_input(q_t); + ggml_set_input(weights_t); + ggml_set_input(comp_t); + ggml_tensor * scores_t = ggml_ds4_indexer_score( + ctx, q_t, weights_t, comp_t, kv_start, ratio); + ggml_set_output(scores_t); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, scores_t); + ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_cpu_buffer_type()); + TEST_ASSERT_MSG(ggml_gallocr_alloc_graph(alloc, graph), + "indexer score graph allocation failed"); + ggml_backend_tensor_set(q_t, q.data(), 0, q.size() * sizeof(float)); + ggml_backend_tensor_set(weights_t, weights.data(), 0, + weights.size() * sizeof(float)); + ggml_backend_tensor_set(comp_t, comp_f16.data(), 0, + comp_f16.size() * sizeof(ggml_fp16_t)); + TEST_ASSERT_MSG(ggml_backend_graph_compute(backend, graph) == + GGML_STATUS_SUCCESS, + "indexer score graph compute failed"); + std::vector actual(expected.size()); + ggml_backend_tensor_get(scores_t, actual.data(), 0, + actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + TEST_ASSERT_MSG(nearly_equal(actual[i], expected[i], 2e-5f, 2e-5f), + "indexer score output mismatch"); + } + ggml_gallocr_free(alloc); + ggml_free(ctx); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + +static void test_indexer_mask_cpu(ggml_backend_t backend) { + std::fprintf(stderr, " test_indexer_mask_cpu ..."); + constexpr int raw_rows = 3; + constexpr int n_comp = 5; + constexpr int n_tokens = 3; + constexpr int top_k = 3; + constexpr int n_attn = raw_rows + n_comp; + std::vector base((size_t) n_attn * n_tokens); + for (int token = 0; token < n_tokens; ++token) { + for (int row = 0; row < n_attn; ++row) { + base[(size_t) token * n_attn + row] = + 100.0f * (float) token + (float) row + 0.25f; + } + } + const std::vector selected = { + 4, 1, 1, + 0, -1, 3, + 2, 5, 4, + }; + std::vector expected((size_t) n_attn * n_tokens, -1.0e30f); + for (int token = 0; token < n_tokens; ++token) { + std::copy_n(base.data() + (size_t) token * n_attn, raw_rows, + expected.data() + (size_t) token * n_attn); + for (int k = 0; k < top_k; ++k) { + const int comp = selected[(size_t) token * top_k + k]; + if (comp >= 0 && comp < n_comp) { + expected[(size_t) token * n_attn + raw_rows + comp] = + base[(size_t) token * n_attn + raw_rows + comp]; + } + } + } + ggml_context * ctx = make_test_context(); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) return; + ggml_tensor * base_t = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_attn, n_tokens); + ggml_tensor * selected_t = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, top_k, n_tokens); + ggml_set_input(base_t); + ggml_set_input(selected_t); + ggml_tensor * mask_t = ggml_ds4_indexer_mask( + ctx, base_t, selected_t, raw_rows); + ggml_set_output(mask_t); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, mask_t); + ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_cpu_buffer_type()); + TEST_ASSERT_MSG(ggml_gallocr_alloc_graph(alloc, graph), + "indexer mask graph allocation failed"); + ggml_backend_tensor_set(base_t, base.data(), 0, + base.size() * sizeof(float)); + ggml_backend_tensor_set(selected_t, selected.data(), 0, + selected.size() * sizeof(int32_t)); + TEST_ASSERT_MSG(ggml_backend_graph_compute(backend, graph) == + GGML_STATUS_SUCCESS, + "indexer mask graph compute failed"); + std::vector actual(expected.size()); + ggml_backend_tensor_get(mask_t, actual.data(), 0, + actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + TEST_ASSERT_MSG(actual[i] == expected[i], + "indexer mask output mismatch"); + } + ggml_gallocr_free(alloc); + ggml_free(ctx); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_hash_routing_lookup() { std::fprintf(stderr, " test_hash_routing_lookup ..."); @@ -1701,6 +2052,304 @@ static void test_output_graph_reuse_microbench(ggml_backend_t backend) { } #if defined(GGML_USE_CUDA) || defined(GGML_USE_HIP) +static void test_ds4_flash_attention_keep_cap_gpu() { + std::fprintf(stderr, " test_ds4_flash_attention_keep_cap_gpu ..."); +#if !defined(GGML_USE_HIP) + std::fprintf(stderr, " skipped (HIP-only contract)\n"); + return; +#endif + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, " skipped (no GPU backend)\n"); + return; + } + + constexpr int head_dim = 512; + constexpr int n_heads = 4; + constexpr int n_tokens = 1; + constexpr int raw_rows = 128; + constexpr int n_comp_rows = 8; + constexpr int n_kv = raw_rows + n_comp_rows; + constexpr int configured_keep_rows = 512; + + ggml_context * ctx = make_test_context(2u << 20); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) { + ggml_backend_free(backend); + std::fprintf(stderr, " FAIL\n"); + return; + } + + ggml_tensor * q = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, head_dim, n_tokens, n_heads); + ggml_tensor * kv = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, head_dim, n_kv, 1); + ggml_tensor * mask = ggml_new_tensor_2d( + ctx, GGML_TYPE_F16, n_kv, n_tokens); + ggml_tensor * output = ggml_flash_attn_ext( + ctx, q, kv, kv, mask, 1.0f / std::sqrt((float) head_dim), + 0.0f, 0.0f); + ggml_flash_attn_ext_set_ds4_sparse( + output, raw_rows, raw_rows, configured_keep_rows, 32); + ggml_set_output(output); + TEST_ASSERT_MSG(ggml_backend_supports_op(backend, output), + "GPU rejected a DS4 keep cap larger than live history"); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, output); + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + const bool allocated = ggml_gallocr_alloc_graph(alloc, graph); + TEST_ASSERT_MSG(allocated, "DS4 flash-attention graph allocation failed"); + if (allocated) { + std::vector q_data((size_t) head_dim * n_tokens * n_heads); + std::vector kv_data((size_t) head_dim * n_kv); + std::vector mask_data((size_t) n_kv * n_tokens, + ggml_fp32_to_fp16(0.0f)); + for (size_t i = 0; i < q_data.size(); ++i) { + q_data[i] = ((int) (i % 19) - 9) * 0.002f; + } + for (size_t i = 0; i < kv_data.size(); ++i) { + kv_data[i] = ((int) (i % 23) - 11) * 0.002f; + } + ggml_backend_tensor_set(q, q_data.data(), 0, + q_data.size() * sizeof(float)); + ggml_backend_tensor_set(kv, kv_data.data(), 0, + kv_data.size() * sizeof(float)); + ggml_backend_tensor_set(mask, mask_data.data(), 0, + mask_data.size() * sizeof(ggml_fp16_t)); + const bool computed = + ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + TEST_ASSERT_MSG(computed, "DS4 flash-attention graph compute failed"); + if (computed) { + std::vector output_data((size_t) ggml_nelements(output)); + ggml_backend_tensor_get(output, output_data.data(), 0, + output_data.size() * sizeof(float)); + for (float value : output_data) { + TEST_ASSERT_MSG(std::isfinite(value), + "DS4 flash-attention output must be finite"); + } + } + } + + ggml_gallocr_free(alloc); + ggml_free(ctx); + ggml_backend_free(backend); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + +static void test_ds4_flash_attention_inverse_rope_fallback_gpu() { + std::fprintf(stderr, + " test_ds4_flash_attention_inverse_rope_fallback_gpu ..."); +#if !defined(GGML_USE_HIP) + std::fprintf(stderr, " skipped (HIP-only contract)\n"); + return; +#endif + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, " skipped (no GPU backend)\n"); + return; + } + + constexpr int head_dim = 512; + constexpr int raw_rows = 128; + constexpr int n_comp_rows = 8; + constexpr int n_kv = raw_rows + n_comp_rows; + + ggml_context * ctx = make_test_context(3u << 20); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) { + ggml_backend_free(backend); + std::fprintf(stderr, " FAIL\n"); + return; + } + + ggml_tensor * q = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, head_dim, 1, 1); + ggml_tensor * k = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, head_dim, n_kv, 1); + ggml_tensor * v = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, head_dim, n_kv, 1); + ggml_tensor * mask = ggml_new_tensor_2d( + ctx, GGML_TYPE_F16, n_kv, 1); + ggml_tensor * output = ggml_flash_attn_ext( + ctx, q, k, v, mask, 1.0f / std::sqrt((float) head_dim), + 0.0f, 0.0f); + // One retained compressed block selects the single-head fallback without + // pruning any live row. Position zero makes inverse RoPE the identity, so + // a row-constant V has an exact, simple reference result. + ggml_flash_attn_ext_set_ds4_sparse( + output, raw_rows, raw_rows, 4, n_comp_rows); + ggml_flash_attn_ext_set_ds4_inverse_rope( + output, 0, 10000.0f, 1.0f, 0.0f, 1.0f, + 32.0f, 1.0f, 8192, false); + ggml_set_output(output); + TEST_ASSERT_MSG(ggml_backend_supports_op(backend, output), + "GPU rejected DS4 inverse-RoPE fallback attention"); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, output); + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + const bool allocated = ggml_gallocr_alloc_graph(alloc, graph); + TEST_ASSERT_MSG(allocated, + "DS4 inverse-RoPE fallback graph allocation failed"); + if (allocated) { + std::vector q_data(head_dim); + std::vector k_data((size_t) head_dim * n_kv); + std::vector v_data((size_t) head_dim * n_kv); + std::vector expected(head_dim); + std::vector mask_data( + n_kv, ggml_fp32_to_fp16(0.0f)); + for (int d = 0; d < head_dim; ++d) { + q_data[(size_t) d] = ((d % 19) - 9) * 0.002f; + expected[(size_t) d] = ((d % 17) - 8) * 0.003f; + } + for (int row = 0; row < n_kv; ++row) { + for (int d = 0; d < head_dim; ++d) { + k_data[(size_t) row * head_dim + d] = + (((row + d) % 23) - 11) * 0.002f; + v_data[(size_t) row * head_dim + d] = expected[(size_t) d]; + } + } + ggml_backend_tensor_set(q, q_data.data(), 0, + q_data.size() * sizeof(float)); + ggml_backend_tensor_set(k, k_data.data(), 0, + k_data.size() * sizeof(float)); + ggml_backend_tensor_set(v, v_data.data(), 0, + v_data.size() * sizeof(float)); + ggml_backend_tensor_set(mask, mask_data.data(), 0, + mask_data.size() * sizeof(ggml_fp16_t)); + const bool computed = + ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + TEST_ASSERT_MSG(computed, + "DS4 inverse-RoPE fallback graph compute failed"); + if (computed) { + std::vector actual(head_dim); + ggml_backend_tensor_get(output, actual.data(), 0, + actual.size() * sizeof(float)); + for (int d = 0; d < head_dim; ++d) { + TEST_ASSERT_MSG( + nearly_equal(actual[(size_t) d], expected[(size_t) d], + 2.0e-5f, 2.0e-5f), + "inverse-RoPE fallback output mismatch"); + } + } + } + + ggml_gallocr_free(alloc); + ggml_free(ctx); + ggml_backend_free(backend); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + +static void test_hc_post_strided_split_gpu() { + std::fprintf(stderr, " test_hc_post_strided_split_gpu ..."); + ggml_backend_t backend = ggml_backend_cuda_init(0); + if (!backend) { + std::fprintf(stderr, " skipped (no GPU backend)\n"); + return; + } + + constexpr int n_embd = 16; + constexpr int n_hc = 4; + constexpr int n_tokens = 3; + constexpr int hc_dim = n_embd * n_hc; + constexpr int mix_dim = 2 * n_hc + n_hc * n_hc; + constexpr int split_stride = mix_dim + 7; + + std::vector residual((size_t) hc_dim * n_tokens); + std::vector block_out((size_t) n_embd * n_tokens); + std::vector split_storage((size_t) split_stride * n_tokens, -99.0f); + for (size_t i = 0; i < residual.size(); ++i) { + residual[i] = ((int) (i % 17) - 8) * 0.03125f; + } + for (size_t i = 0; i < block_out.size(); ++i) { + block_out[i] = ((int) (i % 11) - 5) * 0.0625f; + } + for (int token = 0; token < n_tokens; ++token) { + float * split = split_storage.data() + (size_t) token * split_stride; + for (int i = 0; i < mix_dim; ++i) { + split[i] = ((i + 3 * token) % 13 - 6) * 0.05f; + } + } + + std::vector expected((size_t) hc_dim * n_tokens); + for (int token = 0; token < n_tokens; ++token) { + const float * residual_row = residual.data() + (size_t) token * hc_dim; + const float * block_row = block_out.data() + (size_t) token * n_embd; + const float * split = split_storage.data() + (size_t) token * split_stride; + const float * post = split + n_hc; + const float * comb = split + 2 * n_hc; + float * output = expected.data() + (size_t) token * hc_dim; + for (int h = 0; h < n_hc; ++h) { + for (int d = 0; d < n_embd; ++d) { + float value = block_row[d] * post[h]; + for (int src = 0; src < n_hc; ++src) { + value += comb[h + src * n_hc] * + residual_row[src * n_embd + d]; + } + output[h * n_embd + d] = value; + } + } + } + + ggml_context * ctx = make_test_context(1u << 20); + TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); + if (!ctx) { + ggml_backend_free(backend); + std::fprintf(stderr, " FAIL\n"); + return; + } + + ggml_tensor * residual_t = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, hc_dim, n_tokens); + ggml_tensor * block_t = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_embd, n_tokens); + ggml_tensor * split_storage_t = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, split_stride, n_tokens); + ggml_tensor * split_t = ggml_view_2d( + ctx, split_storage_t, mix_dim, n_tokens, + split_storage_t->nb[1], 0); + TEST_ASSERT(!ggml_is_contiguous(split_t)); + ggml_tensor * output_t = ggml_ds4_hc_post( + ctx, residual_t, block_t, split_t, n_hc); + ggml_set_output(output_t); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, output_t); + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + const bool allocated = ggml_gallocr_alloc_graph(alloc, graph); + TEST_ASSERT_MSG(allocated, "strided HC-post graph allocation failed"); + if (allocated) { + ggml_backend_tensor_set(residual_t, residual.data(), 0, + residual.size() * sizeof(float)); + ggml_backend_tensor_set(block_t, block_out.data(), 0, + block_out.size() * sizeof(float)); + ggml_backend_tensor_set(split_storage_t, split_storage.data(), 0, + split_storage.size() * sizeof(float)); + const bool computed = + ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS; + TEST_ASSERT_MSG(computed, "strided HC-post graph compute failed"); + if (computed) { + std::vector actual(expected.size()); + ggml_backend_tensor_get(output_t, actual.data(), 0, + actual.size() * sizeof(float)); + for (size_t i = 0; i < actual.size(); ++i) { + TEST_ASSERT_MSG(nearly_equal(actual[i], expected[i], + 1.0e-6f, 1.0e-6f), + "strided HC-post output mismatch"); + } + } + } + + ggml_gallocr_free(alloc); + ggml_free(ctx); + ggml_backend_free(backend); + std::fprintf(stderr, g_failures ? " done\n" : " ok\n"); +} + static void test_cpu_hc_sinkhorn_ref(float * out, const float * mix, const float * scale, const float * base, int n_hc, int iters, float eps) { const float pre_scale = scale[0]; @@ -2050,6 +2699,11 @@ int main() { test_moe_routing_correctness(backend); test_rmsnorm_correctness(backend); test_grouped_output_projection_shape(); + test_grouped_output_projection_cpu(backend); + test_ds4_flash_attention_cpu_rejected(backend); + test_indexer_qat_cpu(backend); + test_indexer_score_cpu(backend); + test_indexer_mask_cpu(backend); test_hash_routing_lookup(); test_raw_ring_spans_after_wrap(); test_auto_split_computation(); @@ -2074,6 +2728,9 @@ int main() { test_ffn_graph_reuse_microbench(backend); test_output_graph_reuse_microbench(backend); #if defined(GGML_USE_CUDA) || defined(GGML_USE_HIP) + test_ds4_flash_attention_keep_cap_gpu(); + test_ds4_flash_attention_inverse_rope_fallback_gpu(); + test_hc_post_strided_split_gpu(); test_hc_pre_kernel_gpu(); #endif From 79f075438162842aba3a0401a45d649966ef731e Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 21 Jul 2026 10:35:44 -0700 Subject: [PATCH 14/15] perf(ds4): add device-resident heterogeneous prefill --- server/CMakeLists.txt | 14 + server/deps/llama.cpp/ggml/include/ggml.h | 9 + .../deps/llama.cpp/ggml/src/ggml-cuda/mmid.cu | 6 +- .../deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu | 56 +- .../llama.cpp/ggml/src/ggml-cuda/moe-fused.cu | 95 +- server/deps/llama.cpp/ggml/src/ggml.c | 30 + server/src/common/moe_hybrid_ffn_eval.cpp | 985 +++++++++++++++++- server/src/common/moe_hybrid_ffn_eval.h | 23 +- server/src/common/moe_hybrid_storage.cpp | 28 +- server/src/deepseek4/deepseek4_backend.cpp | 41 +- server/src/deepseek4/deepseek4_graph.cpp | 599 ++++++++++- server/test/test_ggml_rmsnorm_batch.cpp | 93 ++ 12 files changed, 1909 insertions(+), 70 deletions(-) create mode 100644 server/test/test_ggml_rmsnorm_batch.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 581fea034..9ffbf408f 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -895,6 +895,20 @@ if(DFLASH27B_TESTS) endif() add_test(NAME deepseek4_unit COMMAND test_deepseek4_unit) endif() + if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_ggml_rmsnorm_batch.cpp") + add_executable(test_ggml_rmsnorm_batch test/test_ggml_rmsnorm_batch.cpp) + target_include_directories(test_ggml_rmsnorm_batch PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) + target_link_libraries(test_ggml_rmsnorm_batch PRIVATE ggml ${DFLASH27B_GGML_BACKEND_TARGET} ggml-base) + if(DFLASH27B_GPU_BACKEND STREQUAL "hip") + target_compile_definitions(test_ggml_rmsnorm_batch PRIVATE GGML_USE_HIP) + target_link_libraries(test_ggml_rmsnorm_batch PRIVATE hip::host) + else() + target_compile_definitions(test_ggml_rmsnorm_batch PRIVATE GGML_USE_CUDA) + find_package(CUDAToolkit REQUIRED) + target_link_libraries(test_ggml_rmsnorm_batch PRIVATE CUDA::cudart) + endif() + add_test(NAME ggml_rmsnorm_batch COMMAND test_ggml_rmsnorm_batch) + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_ds4_dspark_load.cpp") add_executable(test_ds4_dspark_load test/test_ds4_dspark_load.cpp) target_include_directories(test_ds4_dspark_load PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS} ${CMAKE_CURRENT_SOURCE_DIR}/include) diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index 706d39ab0..a1f6864a5 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -2487,6 +2487,15 @@ extern "C" { struct ggml_tensor * experts, struct ggml_tensor * expert_weights); + // Reduce packed owner-local expert rows directly back to token order. + // inverse_routes maps [route, token] to a row in packed_experts; routes + // with zero weight are ignored before the row is read. + GGML_API struct ggml_tensor * ggml_laguna_moe_packed_combine( + struct ggml_context * ctx, + struct ggml_tensor * packed_experts, + struct ggml_tensor * inverse_routes, + struct ggml_tensor * expert_weights); + // Coarse DeepSeek-V4 routed-owner op. gate_up contains concatenated gate // and up output rows; the backend performs fused gate/up MMVQ + clamped // SwiGLU, down MMVQ, route weighting, and local reduction. diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmid.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmid.cu index 3c61e4595..f3e7f10f4 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmid.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmid.cu @@ -44,7 +44,9 @@ static __global__ void mm_ids_helper( int iex_used = -1; // The index at which the expert is used, if any. for (int iex = threadIdx.x; iex < n_expert_used; iex += warp_size) { const int expert_used = ids[it*si1 + iex]; - nex_prev += expert_used < expert; + // Negative IDs are masked owner routes. They belong to no + // expert and must not reserve holes in the compact arrays. + nex_prev += expert_used >= 0 && expert_used < expert; if (expert_used == expert) { iex_used = iex; } @@ -69,7 +71,7 @@ static __global__ void mm_ids_helper( const int expert_used = (neu_padded == n_expert_used || iex < n_expert_used) && it < n_tokens ? ids[it*si1 + iex] : INT_MAX; const int iex_used = expert_used == expert ? iex : -1; - nex_prev += expert_used < expert; + nex_prev += expert_used >= 0 && expert_used < expert; // Whether the threads at this token position have used the expert: const int it_compact_add_self = warp_reduce_any(iex_used != -1); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu index c4a405f28..aaf70e038 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu @@ -3,6 +3,16 @@ #include "quantize.cuh" #include "mmid.cuh" +// Private contract with the heterogeneous owner builder. The hint is attached +// only to synchronous, large-batch prefill ID tensors and bounds the widest +// expert bucket after negative owner routes have been compacted away. +struct mmid_owner_grid_hint { + uint32_t magic; + int32_t max_expert_rows; + int32_t live_routes; +}; +static constexpr uint32_t MMID_OWNER_GRID_HINT_MAGIC = 0x4D4F4752u; // "MOGR" + static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { switch (args.type_x) { case GGML_TYPE_Q4_0: @@ -252,6 +262,14 @@ static void ggml_cuda_mul_mat_q_impl( ggml_cuda_pool_alloc expert_bounds(ctx.pool(), ne02 + 1); { + // The sorter compacts negative (masked) owner routes out of the expert + // ranges. Quantization still visits these fixed-capacity arrays, so + // make the unused tail point at token zero instead of stale pool data. + CUDA_CHECK(cudaMemsetAsync(ids_src1.get(), 0, + ne_get_rows * sizeof(int32_t), stream)); + CUDA_CHECK(cudaMemsetAsync(ids_dst.get(), 0, + (ne_get_rows + mmq_x_pad) * sizeof(int32_t), + stream)); GGML_ASSERT(ids->nb[0] == ggml_element_size(ids)); const int si1 = ids->nb[1] / ggml_element_size(ids); const int sis1 = nb12 / nb11; @@ -261,11 +279,27 @@ static void ggml_cuda_mul_mat_q_impl( CUDA_CHECK(cudaGetLastError()); } - const size_t nbytes_src1_q8_1 = ne12*n_expert_used*ne10_padded * sizeof(block_q8_1)/QK8_1 + + const mmid_owner_grid_hint * grid_hint = + (const mmid_owner_grid_hint *) ids->extra; + const bool valid_grid_hint = + grid_hint && grid_hint->magic == MMID_OWNER_GRID_HINT_MAGIC && + grid_hint->max_expert_rows > 0 && + grid_hint->max_expert_rows <= ne12 && + grid_hint->live_routes > 0 && + grid_hint->live_routes <= ne_get_rows; + // The ID helper compacts negative owner routes to the tail of ids_src1. + // Quantizing that zero-filled tail repeats the same token activation for + // work that no expert will consume. Preserve the expert-sorted prefix + // expected by MMQ, but size and launch activation quantization for only + // the live cold routes. Unlike a separate owner-row gather this adds no + // full-precision activation traffic. + const int64_t n_routes_quantized = + valid_grid_hint ? grid_hint->live_routes : ne_get_rows; + const size_t nbytes_src1_q8_1 = n_routes_quantized*ne10_padded * sizeof(block_q8_1)/QK8_1 + get_mmq_x_max_host(cc)*sizeof(block_q8_1_mmq); ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); - const int64_t ne11_flat = ne12*n_expert_used; + const int64_t ne11_flat = n_routes_quantized; const int64_t ne12_flat = 1; const int64_t ne13_flat = 1; @@ -288,13 +322,18 @@ static void ggml_cuda_mul_mat_q_impl( ne11 * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int)); const int64_t s13 = ne12*s12; + int64_t ncols_max = ne12; + if (valid_grid_hint) { + ncols_max = grid_hint->max_expert_rows; + } + // Note that ne02 is used instead of ne12 because the number of y channels determines the z dimension of the CUDA grid. const mmq_args args = { src0_d, src0->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), dst_d, - ne00, ne01, ne_get_rows, s01, ne_get_rows, s1, + ne00, ne01, ne_get_rows, s01, n_routes_quantized, s1, ne02, ne02, s02, s12, s2, ne03, ne13, s03, s13, s3, - use_stream_k, ne12}; + use_stream_k, ncols_max}; ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); if (src0_pair) { @@ -405,8 +444,13 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t case GGML_TYPE_Q2_0_ROCMFP2: case GGML_TYPE_Q3_0_ROCMFPX: case GGML_TYPE_Q4_0_ROCMFP4_FAST: - // ROCmFPX MMQ variants are implemented for gfx1151 only. - mmq_supported = GGML_CUDA_CC_IS_RDNA3_5(cc); + // ROCmFPX MMQ uses wave32 WMMA plus portable packed-int dot + // products. Both gfx1151 (RDNA 3.5) and gfx12xx (RDNA 4) provide + // that contract; the latter is required by the grouped DS4 + // prefill projection because no BLAS path understands its + // strided [K/group, N, group] activation layout. + mmq_supported = GGML_CUDA_CC_IS_RDNA3_5(cc) || + GGML_CUDA_CC_IS_RDNA4(cc); break; default: mmq_supported = false; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu index 8505499a7..d10d7b4bd 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/moe-fused.cu @@ -285,13 +285,20 @@ static __global__ void laguna_moe_combine_kernel( const int t = idx / n_embd; float sum = 0.0f; for (int e = 0; e < n_used; ++e) { + const float w = *(const float *)(weights + + (size_t)e * weights_nb0 + + (size_t)t * weights_nb1); + // Masked owner IDs deliberately leave this expert slot unwritten. + // Reading it before multiplying by zero can turn stale NaNs into a + // visible result, because IEEE NaN * 0 remains NaN. + if (w == 0.0f) { + if (e == 0) sum = 0.0f; + continue; + } const float v = *(const float *)(experts + (size_t)h * experts_nb0 + (size_t)e * experts_nb1 + (size_t)t * experts_nb2); - const float w = *(const float *)(weights + - (size_t)e * weights_nb0 + - (size_t)t * weights_nb1); const float scaled = value_scale == 1.0f ? v : __fmul_rn(v, value_scale); const float prod = __fmul_rn(scaled, w); @@ -302,6 +309,56 @@ static __global__ void laguna_moe_combine_kernel( (size_t)t * output_nb1) = sum; } +static __global__ void laguna_moe_packed_combine_kernel( + const char * __restrict__ packed, + const char * __restrict__ inverse, + const char * __restrict__ weights, + char * __restrict__ output, + const int n_embd, + const int n_pairs, + const int n_used, + const int n_tokens, + const size_t packed_nb0, + const size_t packed_nb1, + const size_t inverse_nb0, + const size_t inverse_nb1, + const size_t weights_nb0, + const size_t weights_nb1, + const size_t output_nb0, + const size_t output_nb1) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int total = n_embd * n_tokens; + if (idx >= total) return; + + const int h = idx % n_embd; + const int t = idx / n_embd; + float sum = 0.0f; + for (int e = 0; e < n_used; ++e) { + const float w = *(const float *)(weights + + (size_t)e * weights_nb0 + + (size_t)t * weights_nb1); + if (w == 0.0f) { + if (e == 0) sum = 0.0f; + continue; + } + const int32_t row = *(const int32_t *)(inverse + + (size_t)e * inverse_nb0 + + (size_t)t * inverse_nb1); + if (row < 0 || row >= n_pairs) { + if (e == 0) sum = 0.0f; + continue; + } + const float v = *(const float *)(packed + + (size_t)h * packed_nb0 + + (size_t)row * packed_nb1); + const float prod = __fmul_rn(v, w); + sum = (e == 0) ? prod : __fadd_rn(sum, prod); + } + *(float *)(output + + (size_t)h * output_nb0 + + (size_t)t * output_nb1) = sum; +} + static __global__ void ds4_peer_copy_f32_kernel( const float * __restrict__ src, float * __restrict__ dst, @@ -600,6 +657,38 @@ static void ggml_cuda_op_ds4_moe_owner_split( void ggml_cuda_op_moe_fused(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int mode = ggml_get_op_params_i32(dst, 0); + if (mode == -6) { + const ggml_tensor * packed = dst->src[0]; + const ggml_tensor * inverse = dst->src[1]; + const ggml_tensor * weights = dst->src[2]; + GGML_ASSERT(packed && packed->type == GGML_TYPE_F32); + GGML_ASSERT(inverse && inverse->type == GGML_TYPE_I32); + GGML_ASSERT(weights && weights->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + const int n_embd = (int) packed->ne[0]; + const int n_pairs = (int) packed->ne[1]; + const int n_used = (int) inverse->ne[0]; + const int n_tokens = (int) inverse->ne[1]; + GGML_ASSERT(weights->ne[0] == n_used && + weights->ne[1] == n_tokens); + GGML_ASSERT(dst->ne[0] == n_embd && dst->ne[1] == n_tokens); + + const int total = n_embd * n_tokens; + const int block = 256; + const int grid = (total + block - 1) / block; + laguna_moe_packed_combine_kernel<<>>( + (const char *) packed->data, + (const char *) inverse->data, + (const char *) weights->data, + (char *) dst->data, + n_embd, n_pairs, n_used, n_tokens, + packed->nb[0], packed->nb[1], + inverse->nb[0], inverse->nb[1], + weights->nb[0], weights->nb[1], + dst->nb[0], dst->nb[1]); + return; + } if (mode == -5) { const ggml_tensor * ids = dst->src[0]; GGML_ASSERT(ids && ids->type == GGML_TYPE_I32); diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index ce54b6978..b62ce8f52 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -8175,6 +8175,36 @@ struct ggml_tensor * ggml_laguna_moe_combine( return result; } +struct ggml_tensor * ggml_laguna_moe_packed_combine( + struct ggml_context * ctx, + struct ggml_tensor * packed_experts, + struct ggml_tensor * inverse_routes, + struct ggml_tensor * expert_weights) { + GGML_ASSERT(packed_experts->type == GGML_TYPE_F32); + GGML_ASSERT(inverse_routes->type == GGML_TYPE_I32); + GGML_ASSERT(expert_weights->type == GGML_TYPE_F32); + GGML_ASSERT(inverse_routes->ne[0] == expert_weights->ne[0]); + GGML_ASSERT(inverse_routes->ne[1] == expert_weights->ne[1]); + + const int64_t ne[4] = { + packed_experts->ne[0], inverse_routes->ne[1], 1, 1 }; + struct ggml_tensor * result = + ggml_new_tensor(ctx, GGML_TYPE_F32, 2, ne); + + result->op = GGML_OP_MOE_FUSED; + result->src[0] = packed_experts; + result->src[1] = inverse_routes; + result->src[2] = expert_weights; + + ggml_set_op_params_i32(result, 0, -6); + ggml_set_op_params_i32(result, 1, (int32_t) packed_experts->ne[0]); + ggml_set_op_params_i32(result, 2, (int32_t) packed_experts->ne[1]); + ggml_set_op_params_i32(result, 3, (int32_t) inverse_routes->ne[0]); + ggml_set_op_params_i32(result, 4, (int32_t) inverse_routes->ne[1]); + + return result; +} + struct ggml_tensor * ggml_ds4_moe_owner( struct ggml_context * ctx, struct ggml_tensor * input, diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index 996df7e6a..cd0346459 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -13,8 +13,20 @@ #include #include +extern "C" void ggml_backend_cuda_set_graphs_disabled_override(bool disabled); + namespace dflash::common { +// Private contract with ggml-cuda/mmq.cu. Large owner-local MUL_MAT_ID graphs +// know the exact maximum number of live rows assigned to any expert. Passing +// that bound avoids launching the full n_tokens-wide grid for every expert. +struct MmidOwnerGridHint { + uint32_t magic; + int32_t max_expert_rows; + int32_t live_routes; +}; +static constexpr uint32_t MMID_OWNER_GRID_HINT_MAGIC = 0x4D4F4752u; // "MOGR" + // NVFP4 scale2: if weight has a per-tensor scale, multiply the matmul result // by that scale. No-op when scale==1.0f (non-NVFP4 models). inline ggml_tensor * apply_scale2(ggml_context * ctx, ggml_tensor * mm_result, float scale) { @@ -36,6 +48,43 @@ inline ggml_tensor * swiglu_maybe_clamped(ggml_context * ctx, using HybridClock = std::chrono::steady_clock; +class HybridCudaGraphDisableScope { +public: + explicit HybridCudaGraphDisableScope(bool active) : active_(active) { + if (active_) ggml_backend_cuda_set_graphs_disabled_override(true); + } + ~HybridCudaGraphDisableScope() { + if (active_) ggml_backend_cuda_set_graphs_disabled_override(false); + } + + HybridCudaGraphDisableScope(const HybridCudaGraphDisableScope &) = delete; + HybridCudaGraphDisableScope & operator=(const HybridCudaGraphDisableScope &) = delete; + +private: + bool active_; +}; + +static bool heterogeneous_prefill_eager_enabled() { + const char * raw = std::getenv("DFLASH_DS4_HYBRID_PREFILL_EAGER"); + return raw && *raw && std::strcmp(raw, "0") != 0; +} + +static bool prefill_masked_cold_routes_enabled() { + static const bool enabled = []() { + const char * raw = std::getenv("DFLASH_MOE_PREFILL_MASKED_COLD"); + return !raw || !*raw || std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +static bool prefill_compact_cold_grid_enabled() { + static const bool enabled = []() { + const char * raw = std::getenv("DFLASH_MOE_PREFILL_COMPACT_COLD_GRID"); + return !raw || !*raw || std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + static uint64_t elapsed_us(HybridClock::time_point start, HybridClock::time_point end) { return (uint64_t) std::chrono::duration_cast(end - start).count(); } @@ -592,7 +641,8 @@ static bool build_batched_routed_graph( ggml_tensor ** out_routed, bool tokenwise = false, std::vector * backend_nodes = nullptr, - bool allow_fused_combine = false) + bool allow_fused_combine = false, + bool force_fused_combine = false) { const auto track = [&](ggml_tensor * t) -> ggml_tensor * { if (backend_nodes && t) backend_nodes->push_back(t); @@ -614,7 +664,7 @@ static bool build_batched_routed_graph( inp_col, sel_col, wts_col, n_embd, n_ff_exp, n_used, 1, swiglu_clamp, &routed_col, false, backend_nodes, - allow_fused_combine)) { + allow_fused_combine, force_fused_combine)) { return false; } joined = joined ? track(ggml_concat(ctx, joined, routed_col, 1)) @@ -709,7 +759,8 @@ static bool build_batched_routed_graph( ggml_mul_mat_id(ctx, down_tensor, gu, sel), down_scale)); // Weight and sum over experts: [n_embd, n_used, n_tokens] * [1, n_used, n_tokens] - if (allow_fused_combine && fused_moe_combine_enabled()) { + if (allow_fused_combine && + (force_fused_combine || fused_moe_combine_enabled())) { *out_routed = track(ggml_laguna_moe_combine(ctx, experts, wts)); return *out_routed != nullptr; } @@ -1793,7 +1844,11 @@ static bool eval_moe_hybrid_ffn_batched_core( MoeExpertCompute * expert_compute, const MoeExpertLayer * expert_layer, MoeHybridFfnTelemetry * telemetry, - bool skip_cold = false) { + bool skip_cold = false, + bool skip_hot = false, + ggml_tensor * cur_backend = nullptr, + ggml_tensor * device_output = nullptr, + ggml_backend_t device_output_owner = nullptr) { const auto ffn_wall_t0 = HybridClock::now(); const auto partition_t0 = HybridClock::now(); @@ -1802,6 +1857,33 @@ static bool eval_moe_hybrid_ffn_batched_core( const int n_ff_exp = cfg.n_ff_exp; out.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); if (n_tokens <= 0) return true; + if (!cur_host && !cur_backend) { + if (err) *err = "hybrid batched activation is unavailable"; + return false; + } + if (device_output && (!skip_hot || !device_output_owner || + device_output->type != GGML_TYPE_F32 || + device_output->ne[0] != n_embd || + device_output->ne[1] != n_tokens)) { + if (err) *err = "invalid cold-owner device output"; + return false; + } + auto set_cur_input = [&](ggml_tensor * dst, + ggml_backend_t dst_backend) { + if (cur_backend) { + // Use the backend-aware copy path. On heterogeneous HIP this + // submits on the producer stream and publishes an event to the + // consumer stream; the generic buffer copy uses a per-thread + // stream with no cross-device dependency and can leave the Strix + // destination containing zeros. + ggml_backend_tensor_copy_async( + gpu_backend, dst_backend, cur_backend, dst); + } else { + ggml_backend_tensor_set( + dst, cur_host, 0, + sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + } + }; // ── Fast path: cached hot+cold batched graphs (spec-decode verify/replay) ── // Mixed layers used to rebuild+free their hot and cold ggml graphs on every @@ -1852,7 +1934,7 @@ static bool eval_moe_hybrid_ffn_batched_core( if (hg_ok && cg_ok) { // Hot (GPU, async): shared expert + routed hot (zero-weight dummy slots // keep an all-cold batch's shared-expert contribution). - ggml_backend_tensor_set(hg.inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + set_cur_input(hg.inp, gpu_backend); if (hg.sel) { ggml_backend_tensor_set(hg.sel, hot_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); @@ -1865,6 +1947,11 @@ static bool eval_moe_hybrid_ffn_batched_core( std::vector cold_partial; if (remote_cold) { + if (!cur_host) { + ggml_backend_synchronize(gpu_backend); + if (err) *err = "remote cold prefill requires a host activation"; + return false; + } if (!eval_moe_hybrid_remote_cold_batched( cfg, storage, cur_host, selected_ids, selected_weights, n_tokens, cold_partial, err, @@ -1874,11 +1961,11 @@ static bool eval_moe_hybrid_ffn_batched_core( } } else if (cg) { cold_partial.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); - ggml_backend_tensor_set(cg->inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); - ggml_backend_tensor_set(cg->sel, cold_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); - ggml_backend_tensor_set(cg->wts, cold_wts.data(), 0, sizeof(float) * (size_t)total_slots); ggml_backend_t cached_cold_backend = storage.cold_backend ? storage.cold_backend : cpu_backend; + set_cur_input(cg->inp, cached_cold_backend); + ggml_backend_tensor_set(cg->sel, cold_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); + ggml_backend_tensor_set(cg->wts, cold_wts.data(), 0, sizeof(float) * (size_t)total_slots); if (ggml_backend_graph_compute( cached_cold_backend, cg->gf) != GGML_STATUS_SUCCESS) { ggml_backend_synchronize(gpu_backend); @@ -1908,12 +1995,26 @@ static bool eval_moe_hybrid_ffn_batched_core( const int n_hot_init = std::max(1, storage.hot_active); for (int i = 0; i < total_slots; ++i) hot_sel[i] = i % n_hot_init; std::vector hot_wts(total_slots, 0.0f); + const int n_cold_stack = + std::max(1, (int)(storage.down_cold ? storage.down_cold->ne[2] : 1)); + // Owner-only prefill must not turn routes assigned to the other GPU into + // valid zero-weight expert IDs. MUL_MAT_ID executes gate/up/down before + // route weighting, so those dummies duplicate all nominally offloaded work. + // Negative IDs are compacted by the ROCm MMQ sorter; fused combine avoids + // reading their intentionally unwritten outputs. + const bool mask_skipped_cold = + skip_hot && storage.cold_backend_kind == MoeHybridColdBackend::Gpu && + prefill_masked_cold_routes_enabled(); std::vector cold_sel(total_slots); - for (int i = 0; i < total_slots; ++i) cold_sel[i] = i % std::max(1, (int)(storage.down_cold ? storage.down_cold->ne[2] : 1)); + for (int i = 0; i < total_slots; ++i) { + cold_sel[i] = mask_skipped_cold ? -1 : i % n_cold_stack; + } std::vector cold_wts(total_slots, 0.0f); bool has_hot = false, has_cold = false; int hot_selected = 0; int cold_selected = 0; + std::vector cold_rows_per_expert( + mask_skipped_cold ? (size_t)n_cold_stack : 0, 0); for (int i = 0; i < total_slots; ++i) { const int32_t gid = selected_ids[i]; @@ -1931,15 +2032,35 @@ static bool eval_moe_hybrid_ffn_batched_core( cold_wts[i] = selected_weights[i]; has_cold = true; cold_selected++; + if (!cold_rows_per_expert.empty()) { + cold_rows_per_expert[(size_t)cold_lid]++; + } } } } + int32_t max_cold_expert_rows = 0; + for (int32_t count : cold_rows_per_expert) { + max_cold_expert_rows = std::max(max_cold_expert_rows, count); + } + MmidOwnerGridHint cold_grid_hint{ + MMID_OWNER_GRID_HINT_MAGIC, max_cold_expert_rows, cold_selected}; const auto partition_t1 = HybridClock::now(); if (telemetry) { telemetry->partition_us += elapsed_us(partition_t0, partition_t1); telemetry->hot_selected += hot_selected; telemetry->cold_selected += cold_selected; } + if (mask_skipped_cold) { + static bool logged = false; + if (!logged) { + std::fprintf(stderr, + "[hybrid-ffn] masked cold prefill routes active; " + "Strix skips R9700-owned expert GEMMs " + "max_expert_rows=%d\n", + (int)max_cold_expert_rows); + logged = true; + } + } ggml_backend_t cold_backend = storage.cold_backend ? storage.cold_backend : cpu_backend; const bool cold_on_gpu = has_cold && @@ -1957,7 +2078,7 @@ static bool eval_moe_hybrid_ffn_batched_core( ggml_tensor * hot_output = nullptr; const bool has_shared = (desc.ffn_up_shexp && desc.ffn_gate_shexp && desc.ffn_down_shexp); - if (has_hot || has_shared) { + if (!skip_hot && (has_hot || has_shared)) { ggml_init_params ip{}; ip.mem_size = 128 * 1024 * 1024; ip.mem_buffer = nullptr; @@ -2010,7 +2131,7 @@ static bool eval_moe_hybrid_ffn_batched_core( return false; } - ggml_backend_tensor_set(inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + set_cur_input(inp, gpu_backend); if (has_hot) { ggml_backend_tensor_set(sel, hot_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); ggml_backend_tensor_set(wts, hot_wts.data(), 0, sizeof(float) * (size_t)total_slots); @@ -2033,6 +2154,13 @@ static bool eval_moe_hybrid_ffn_batched_core( // stability, but remote cold expert compute can still run as one // larger batched IPC request for the full prefill chunk. } else if (has_cold && expert_compute && expert_layer) { + if (!cur_host) { + if (hot_async_launched) ggml_backend_synchronize(gpu_backend); + if (!p_hot_alloc && hot_alloc) ggml_gallocr_free(hot_alloc); + if (hot_ctx) ggml_free(hot_ctx); + if (err) *err = "remote cold prefill requires a host activation"; + return false; + } if (!eval_moe_hybrid_remote_cold_batched( cfg, storage, cur_host, selected_ids, selected_weights, n_tokens, cold_partial, err, expert_compute, expert_layer)) { @@ -2063,10 +2191,15 @@ static bool eval_moe_hybrid_ffn_batched_core( return false; } - ggml_tensor * inp = ggml_new_tensor_2d(cold_ctx, GGML_TYPE_F32, n_embd, n_tokens); + ggml_tensor * inp = ggml_new_tensor_2d( + cold_ctx, GGML_TYPE_F32, n_embd, n_tokens); ggml_set_input(inp); ggml_tensor * sel = ggml_new_tensor_2d(cold_ctx, GGML_TYPE_I32, n_used, n_tokens); ggml_set_input(sel); + if (mask_skipped_cold && max_cold_expert_rows > 0 && + prefill_compact_cold_grid_enabled()) { + sel->extra = &cold_grid_hint; + } ggml_tensor * wts = ggml_new_tensor_2d(cold_ctx, GGML_TYPE_F32, n_used, n_tokens); ggml_set_input(wts); @@ -2076,7 +2209,8 @@ static bool eval_moe_hybrid_ffn_batched_core( desc.ffn_gate_exps_s, desc.ffn_up_exps_s, desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, inp, sel, wts, n_embd, n_ff_exp, n_used, n_tokens, cfg.swiglu_clamp, &cold_routed, false, nullptr, - ggml_backend_is_cuda(cold_backend)); + ggml_backend_is_cuda(cold_backend), + /*force_fused_combine=*/mask_skipped_cold); ggml_cgraph * cold_gf = ggml_new_graph_custom(cold_ctx, 4096, false); ggml_set_output(cold_routed); @@ -2099,7 +2233,7 @@ static bool eval_moe_hybrid_ffn_batched_core( return false; } - ggml_backend_tensor_set(inp, cur_host, 0, sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + set_cur_input(inp, cold_backend); ggml_backend_tensor_set(sel, cold_sel.data(), 0, sizeof(int32_t) * (size_t)total_slots); ggml_backend_tensor_set(wts, cold_wts.data(), 0, sizeof(float) * (size_t)total_slots); @@ -2114,8 +2248,17 @@ static bool eval_moe_hybrid_ffn_batched_core( return false; } - ggml_backend_tensor_get(cold_routed, cold_partial.data(), 0, - sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + if (device_output) { + ggml_backend_tensor_copy_async( + cold_backend, device_output_owner, + cold_routed, device_output); + // The cold graph uses a reusable gallocr arena. Keep its output + // storage live until the peer transfer has completed. + ggml_backend_synchronize(device_output_owner); + } else { + ggml_backend_tensor_get(cold_routed, cold_partial.data(), 0, + sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + } if (telemetry) telemetry->cold_us += elapsed_us(cold_t0, HybridClock::now()); if (!p_cold_alloc) ggml_gallocr_free(cold_alloc); ggml_free(cold_ctx); @@ -2131,6 +2274,14 @@ static bool eval_moe_hybrid_ffn_batched_core( if (!p_hot_alloc && hot_alloc) ggml_gallocr_free(hot_alloc); if (hot_ctx) ggml_free(hot_ctx); + if (device_output) { + if (telemetry) { + telemetry->ffn_wall_us += + elapsed_us(ffn_wall_t0, HybridClock::now()); + } + return true; + } + // ── Step 5: Merge hot + cold ── const auto combine_t0 = HybridClock::now(); const size_t total_floats = (size_t)n_embd * (size_t)n_tokens; @@ -2212,6 +2363,588 @@ static bool eval_moe_hybrid_remote_cold_batched( return true; } +// Long heterogeneous prefill has enough rows per expert to use ordinary +// dense matmuls efficiently. Packing routes by owner-local expert avoids the +// reduced-stack MUL_MAT_ID stream-k path (which is both slow for a 24-expert +// stack and unstable for very large batches) and lets each expert's weights be +// reused across all of its prompt rows. +static bool expert_major_prefill_enabled(int n_tokens) { + static const bool enabled = []() { + const char * raw = std::getenv("DFLASH_MOE_EXPERT_MAJOR_PREFILL"); + return !raw || !*raw || std::strcmp(raw, "0") != 0; + }(); + static const int min_tokens = + env_int_or_default("DFLASH_MOE_EXPERT_MAJOR_MIN_TOKENS", 512); + return enabled && n_tokens >= min_tokens; +} + +// Expert-major prefill groups prompt rows by expert so the quantized GEMMs can +// reuse each expert's weights. Keep the inverse permutation and route weights +// on the owner GPU as well: gathering the packed rows back to route-major order +// and reducing there avoids reading every per-route hidden vector to the host, +// doing an O(n_tokens * n_used * n_embd) CPU scatter, and uploading the reduced +// result again. The old host reduction remains as an emergency A/B fallback. +static bool expert_major_gpu_reduce_enabled() { + static const bool enabled = []() { + const char * raw = std::getenv("DFLASH_MOE_EXPERT_MAJOR_GPU_REDUCE"); + return !raw || !*raw || std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +static bool full_cold_expert_major_enabled() { + static const bool enabled = []() { + const char * raw = + std::getenv("DFLASH_MOE_FULL_COLD_EXPERT_MAJOR"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +static bool full_cold_packed_mmid_enabled() { + static const bool enabled = []() { + const char * raw = + std::getenv("DFLASH_MOE_FULL_COLD_PACKED_MMID"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + +static bool eval_moe_owner_expert_major_batched( + ggml_backend_t backend, + const MoeHybridConfig & cfg, + const MoeLayerDesc & desc, + ggml_tensor * gate_tensor, + ggml_tensor * up_tensor, + ggml_tensor * down_tensor, + ggml_tensor * gate_up_tensor, + const std::vector & local_by_global, + const float * cur_host, + const int32_t * selected_ids, + const float * selected_weights, + int n_tokens, + bool include_shared, + std::vector & out, + std::string * err, + ggml_tensor * cur_backend = nullptr, + ggml_backend_t cur_backend_owner = nullptr, + ggml_tensor * device_output = nullptr, + ggml_backend_t device_output_owner = nullptr) { + const int n_embd = cfg.n_embd; + const int n_used = cfg.n_expert_used; + const int n_ff = cfg.n_ff_exp; + out.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); + if (!backend || (!cur_host && !cur_backend) || + !selected_ids || !selected_weights || + n_tokens <= 0 || n_embd <= 0 || n_used <= 0 || n_ff <= 0 || + !down_tensor || (!gate_up_tensor && (!gate_tensor || !up_tensor))) { + if (err) *err = "invalid expert-major owner inputs"; + return false; + } + + ggml_tensor * stack_ref = gate_up_tensor ? gate_up_tensor : gate_tensor; + const int n_stack = (int)stack_ref->ne[2]; + if (n_stack <= 0 || down_tensor->ne[2] != n_stack) { + if (err) *err = "invalid expert-major owner stack"; + return false; + } + + std::vector counts((size_t)n_stack, 0); + size_t n_pairs = 0; + for (int t = 0; t < n_tokens; ++t) { + for (int slot = 0; slot < n_used; ++slot) { + const size_t route = (size_t)t * (size_t)n_used + (size_t)slot; + const int32_t global = selected_ids[route]; + if (global < 0 || (size_t)global >= local_by_global.size() || + selected_weights[route] == 0.0f) { + continue; + } + const int32_t local = local_by_global[(size_t)global]; + if (local >= 0 && local < n_stack) { + counts[(size_t)local]++; + n_pairs++; + } + } + } + + const bool has_shared = include_shared && desc.ffn_gate_shexp && + desc.ffn_up_shexp && desc.ffn_down_shexp; + if (n_pairs == 0 && !has_shared) return true; + + std::vector offsets((size_t)n_stack + 1, 0); + for (int e = 0; e < n_stack; ++e) { + offsets[(size_t)e + 1] = offsets[(size_t)e] + + (size_t)counts[(size_t)e]; + } + std::vector cursor(offsets.begin(), offsets.end() - 1); + // Device input keeps the normalized activation on the R9700 and gathers + // the owner rows on-device. The host fallback remains for decode and A/B. + std::vector packed_input; + if (!cur_backend) { + packed_input.resize(n_pairs * (size_t)n_embd); + } + std::vector packed_tokens(n_pairs); + std::vector packed_weights(n_pairs); + std::vector packed_routes(n_pairs); + std::vector route_to_packed( + (size_t)n_tokens * (size_t)n_used, 0); + std::vector owner_weights( + (size_t)n_tokens * (size_t)n_used, 0.0f); + for (int t = 0; t < n_tokens; ++t) { + for (int slot = 0; slot < n_used; ++slot) { + const size_t route = (size_t)t * (size_t)n_used + (size_t)slot; + const int32_t global = selected_ids[route]; + if (global < 0 || (size_t)global >= local_by_global.size() || + selected_weights[route] == 0.0f) { + continue; + } + const int32_t local = local_by_global[(size_t)global]; + if (local < 0 || local >= n_stack) continue; + const size_t packed = cursor[(size_t)local]++; + if (!cur_backend) { + std::memcpy(packed_input.data() + packed * (size_t)n_embd, + cur_host + (size_t)t * (size_t)n_embd, + sizeof(float) * (size_t)n_embd); + } + packed_tokens[packed] = t; + packed_weights[packed] = selected_weights[route]; + packed_routes[packed] = (int32_t)route; + route_to_packed[route] = (int32_t)packed; + owner_weights[route] = selected_weights[route]; + } + } + + const bool gpu_reduce = expert_major_gpu_reduce_enabled() && n_pairs > 0; + + ggml_init_params ip{}; + ip.mem_size = 128 * 1024 * 1024; + ip.mem_buffer = nullptr; + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) { + if (err) *err = "expert-major ggml_init failed"; + return false; + } + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false); + if (!gf) { + ggml_free(ctx); + if (err) *err = "expert-major graph allocation failed"; + return false; + } + + ggml_tensor * packed_in = nullptr; + ggml_tensor * packed_out = nullptr; + ggml_tensor * owner_input = nullptr; + ggml_tensor * packed_token_ids = nullptr; + ggml_tensor * inverse_routes = nullptr; + ggml_tensor * route_weights = nullptr; + if (cur_backend && (n_pairs > 0 || has_shared)) { + owner_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, + n_embd, n_tokens); + ggml_set_input(owner_input); + } + if (n_pairs > 0) { + if (owner_input) { + packed_token_ids = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, (int64_t)n_pairs); + ggml_set_input(packed_token_ids); + packed_in = ggml_get_rows(ctx, owner_input, packed_token_ids); + } else { + packed_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, + n_embd, (int64_t)n_pairs); + ggml_set_input(packed_in); + } + packed_out = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, + n_embd, (int64_t)n_pairs); + ggml_set_output(packed_out); + if (gpu_reduce) { + inverse_routes = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, n_used, n_tokens); + route_weights = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_used, n_tokens); + ggml_set_input(inverse_routes); + ggml_set_input(route_weights); + } + } + + auto expert_view = [&](ggml_tensor * stack, int local) { + return ggml_view_2d(ctx, stack, stack->ne[0], stack->ne[1], + stack->nb[1], (size_t)local * stack->nb[2]); + }; + + for (int local = 0; local < n_stack; ++local) { + const int count = counts[(size_t)local]; + if (count == 0) continue; + const size_t begin = offsets[(size_t)local]; + ggml_tensor * expert_in = ggml_view_2d( + ctx, packed_in, n_embd, count, packed_in->nb[1], + begin * packed_in->nb[1]); + ggml_tensor * mid = nullptr; + if (gate_up_tensor) { + ggml_tensor * gate_up = apply_scale2( + ctx, ggml_mul_mat(ctx, expert_view(gate_up_tensor, local), + expert_in), + desc.ffn_gate_up_exps_s); + ggml_tensor * gate = ggml_view_2d( + ctx, gate_up, n_ff, count, gate_up->nb[1], 0); + ggml_tensor * up = ggml_view_2d( + ctx, gate_up, n_ff, count, gate_up->nb[1], + (size_t)n_ff * ggml_element_size(gate_up)); + gate = ggml_cont(ctx, gate); + up = ggml_cont(ctx, up); + mid = swiglu_maybe_clamped(ctx, gate, up, cfg.swiglu_clamp); + } else { + ggml_tensor * gate = apply_scale2( + ctx, ggml_mul_mat(ctx, expert_view(gate_tensor, local), + expert_in), + desc.ffn_gate_exps_s); + ggml_tensor * up = apply_scale2( + ctx, ggml_mul_mat(ctx, expert_view(up_tensor, local), + expert_in), + desc.ffn_up_exps_s); + mid = swiglu_maybe_clamped(ctx, gate, up, cfg.swiglu_clamp); + } + ggml_tensor * expert_out = apply_scale2( + ctx, ggml_mul_mat(ctx, expert_view(down_tensor, local), mid), + desc.ffn_down_exps_s); + ggml_tensor * dst = ggml_view_2d( + ctx, packed_out, n_embd, count, packed_out->nb[1], + begin * packed_out->nb[1]); + ggml_tensor * copy = ggml_cpy(ctx, expert_out, dst); + ggml_set_output(copy); + ggml_build_forward_expand(gf, copy); + } + + ggml_tensor * shared_in = nullptr; + ggml_tensor * shared_out = nullptr; + if (has_shared) { + if (owner_input) { + shared_in = owner_input; + } else { + shared_in = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, + n_embd, n_tokens); + ggml_set_input(shared_in); + } + shared_out = build_shared_expert_subgraph( + ctx, desc, shared_in, cfg.swiglu_clamp); + ggml_set_output(shared_out); + ggml_build_forward_expand(gf, shared_out); + } + + + ggml_tensor * combined_out = shared_out; + if (gpu_reduce) { + // packed_out is [n_embd, n_pairs]. GET_ROWS applies the inverse + // expert-major permutation and produces [n_embd, n_used, n_tokens]. + // Invalid/non-owner routes gather row zero but have an exact zero + // weight, so the standard fused MoE reduction masks them out. + // GET_ROWS interprets ids dimension 1 as a batched lookup and requires + // the source to expose the same batch dimension. Every token indexes + // the same packed table, so use a zero-stride view instead of physically + // repeating an O(n_pairs * n_tokens) tensor. + ggml_tensor * packed_table = ggml_view_tensor(ctx, packed_out); + packed_table->ne[2] = n_tokens; + packed_table->nb[2] = 0; + packed_table->ne[3] = 1; + packed_table->nb[3] = 0; + ggml_tensor * route_major = + ggml_get_rows(ctx, packed_table, inverse_routes); + ggml_tensor * routed_out = + ggml_laguna_moe_combine(ctx, route_major, route_weights); + combined_out = combined_out + ? ggml_add(ctx, routed_out, combined_out) + : routed_out; + ggml_set_output(combined_out); + ggml_build_forward_expand(gf, combined_out); + } + + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + if (!alloc || !ggml_gallocr_alloc_graph(alloc, gf)) { + if (alloc) ggml_gallocr_free(alloc); + ggml_free(ctx); + if (err) *err = "expert-major scratch allocation failed"; + return false; + } + if (owner_input) { + if (cur_backend_owner) { + ggml_backend_tensor_copy_async( + cur_backend_owner, backend, cur_backend, owner_input); + } else { + ggml_backend_tensor_copy(cur_backend, owner_input); + } + } + if (packed_token_ids) { + ggml_backend_tensor_set(packed_token_ids, packed_tokens.data(), 0, + sizeof(int32_t) * packed_tokens.size()); + } else if (packed_in) { + ggml_backend_tensor_set(packed_in, packed_input.data(), 0, + sizeof(float) * packed_input.size()); + } + if (shared_in && !owner_input) { + ggml_backend_tensor_set(shared_in, cur_host, 0, + sizeof(float) * (size_t)n_embd * + (size_t)n_tokens); + } + if (gpu_reduce) { + ggml_backend_tensor_set(inverse_routes, route_to_packed.data(), 0, + sizeof(int32_t) * route_to_packed.size()); + ggml_backend_tensor_set(route_weights, owner_weights.data(), 0, + sizeof(float) * owner_weights.size()); + } + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { + ggml_gallocr_free(alloc); + ggml_free(ctx); + if (err) *err = "expert-major graph compute failed"; + return false; + } + + if (device_output) { + if (!device_output_owner || !gpu_reduce || !combined_out || + device_output->type != GGML_TYPE_F32 || + device_output->ne[0] != n_embd || + device_output->ne[1] != n_tokens) { + ggml_gallocr_free(alloc); + ggml_free(ctx); + if (err) *err = "invalid expert-major device output"; + return false; + } + // Publish the owner result directly to the target GPU. Synchronizing + // the destination before releasing this graph keeps the source arena + // alive until the peer/local copy has retired. + ggml_backend_tensor_copy_async( + backend, device_output_owner, combined_out, device_output); + ggml_backend_synchronize(device_output_owner); + } else if (gpu_reduce) { + ggml_backend_tensor_get(combined_out, out.data(), 0, + sizeof(float) * out.size()); + } else { + std::vector packed_result(n_pairs * (size_t)n_embd); + if (packed_out) { + ggml_backend_tensor_get(packed_out, packed_result.data(), 0, + sizeof(float) * packed_result.size()); + } + if (shared_out) { + ggml_backend_tensor_get(shared_out, out.data(), 0, + sizeof(float) * out.size()); + } + for (size_t packed = 0; packed < n_pairs; ++packed) { + const size_t route = (size_t)packed_routes[packed]; + const size_t token = route / (size_t)n_used; + float * dst = out.data() + token * (size_t)n_embd; + const float * src = + packed_result.data() + packed * (size_t)n_embd; + const float weight = packed_weights[packed]; + for (int i = 0; i < n_embd; ++i) { + dst[i] += src[i] * weight; + } + } + } + + ggml_gallocr_free(alloc); + ggml_free(ctx); + return true; +} + +// Pack only this owner's live token/expert pairs before MUL_MAT_ID. The +// ordinary masked graph retains n_tokens*n_used rows and merely marks the +// other owner's rows negative; its activation quantizer still visits that +// fixed-capacity tensor. This path shrinks the tensor to the actual owner +// routes, executes one full-stack MMID graph, and performs the inverse +// permutation plus route reduction on the owner GPU. +static bool eval_moe_owner_packed_mmid_batched( + ggml_backend_t backend, + const MoeHybridConfig & cfg, + const MoeLayerDesc & desc, + ggml_tensor * gate_tensor, + ggml_tensor * up_tensor, + ggml_tensor * down_tensor, + ggml_tensor * gate_up_tensor, + const std::vector & local_by_global, + const float * cur_host, + const int32_t * selected_ids, + const float * selected_weights, + int n_tokens, + std::vector & out, + std::string * err, + ggml_gallocr_t * p_alloc, + ggml_tensor * cur_backend = nullptr, + ggml_backend_t cur_backend_owner = nullptr) { + const int n_embd = cfg.n_embd; + const int n_used = cfg.n_expert_used; + const int n_ff = cfg.n_ff_exp; + const size_t n_routes = (size_t)n_tokens * (size_t)n_used; + out.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); + if (!backend || (!cur_host && !cur_backend) || !selected_ids || + !selected_weights || n_tokens <= 0 || n_embd <= 0 || n_used <= 0 || + n_ff <= 0 || !down_tensor || + (!gate_up_tensor && (!gate_tensor || !up_tensor))) { + if (err) *err = "invalid packed-MMID owner inputs"; + return false; + } + + ggml_tensor * stack_ref = gate_up_tensor ? gate_up_tensor : gate_tensor; + const int n_stack = (int)stack_ref->ne[2]; + if (n_stack <= 0 || down_tensor->ne[2] != n_stack) { + if (err) *err = "invalid packed-MMID owner stack"; + return false; + } + + std::vector packed_tokens; + std::vector packed_ids; + std::vector route_to_packed(n_routes, 0); + std::vector route_weights(n_routes, 0.0f); + std::vector rows_per_expert((size_t)n_stack, 0); + packed_tokens.reserve(n_routes); + packed_ids.reserve(n_routes); + for (int t = 0; t < n_tokens; ++t) { + for (int slot = 0; slot < n_used; ++slot) { + const size_t route = + (size_t)t * (size_t)n_used + (size_t)slot; + const int32_t global = selected_ids[route]; + if (global < 0 || (size_t)global >= local_by_global.size() || + selected_weights[route] == 0.0f) { + continue; + } + const int32_t local = local_by_global[(size_t)global]; + if (local < 0 || local >= n_stack) continue; + const int32_t packed = (int32_t)packed_ids.size(); + packed_tokens.push_back(t); + packed_ids.push_back(local); + route_to_packed[route] = packed; + route_weights[route] = selected_weights[route]; + rows_per_expert[(size_t)local]++; + } + } + if (packed_ids.empty()) return true; + + int32_t max_expert_rows = 0; + for (int32_t count : rows_per_expert) { + max_expert_rows = std::max(max_expert_rows, count); + } + MmidOwnerGridHint grid_hint{ + MMID_OWNER_GRID_HINT_MAGIC, max_expert_rows, + (int32_t)packed_ids.size()}; + const int64_t n_pairs = (int64_t)packed_ids.size(); + std::vector unit_weights((size_t)n_pairs, 1.0f); + + ggml_init_params ip{}; + ip.mem_size = 64 * 1024 * 1024; + ip.mem_buffer = nullptr; + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) { + if (err) *err = "packed-MMID ggml_init failed"; + return false; + } + + ggml_tensor * owner_input = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_embd, n_tokens); + ggml_set_input(owner_input); + ggml_tensor * packed_token_ids = ggml_new_tensor_1d( + ctx, GGML_TYPE_I32, n_pairs); + ggml_set_input(packed_token_ids); + ggml_tensor * packed_input = + ggml_get_rows(ctx, owner_input, packed_token_ids); + ggml_tensor * ids = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, 1, n_pairs); + ggml_set_input(ids); + ids->extra = &grid_hint; + ggml_tensor * weights = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 1, n_pairs); + ggml_set_input(weights); + + ggml_tensor * packed_output = nullptr; + if (!build_batched_routed_graph( + ctx, gate_tensor, up_tensor, down_tensor, gate_up_tensor, + desc.ffn_gate_exps_s, desc.ffn_up_exps_s, + desc.ffn_down_exps_s, desc.ffn_gate_up_exps_s, + packed_input, ids, weights, n_embd, n_ff, 1, (int)n_pairs, + cfg.swiglu_clamp, &packed_output, false, nullptr, + /*allow_fused_combine=*/true)) { + ggml_free(ctx); + if (err) *err = "packed-MMID graph build failed"; + return false; + } + + ggml_tensor * inverse_routes = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, n_used, n_tokens); + ggml_set_input(inverse_routes); + ggml_tensor * final_weights = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_used, n_tokens); + ggml_set_input(final_weights); + + ggml_tensor * combined = + ggml_laguna_moe_packed_combine( + ctx, packed_output, inverse_routes, final_weights); + ggml_set_output(combined); + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 4096, false); + // Register the external leaves explicitly. The dynamically packed graph + // is assembled around a caller-owned device tensor and its shape changes + // per layer, so it must not inherit a fixed graph's leaf set. + ggml_build_forward_expand(gf, owner_input); + ggml_build_forward_expand(gf, packed_token_ids); + ggml_build_forward_expand(gf, ids); + ggml_build_forward_expand(gf, weights); + ggml_build_forward_expand(gf, inverse_routes); + ggml_build_forward_expand(gf, final_weights); + ggml_build_forward_expand(gf, combined); + + // n_pairs changes from layer to layer. The reusable fixed-shape owner + // allocator can leave new dynamic input leaves without a buffer, so keep + // this graph's arena local until a shape-keyed cache is introduced. + (void)p_alloc; + ggml_gallocr_t alloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!alloc || !ggml_gallocr_alloc_graph(alloc, gf)) { + if (alloc) ggml_gallocr_free(alloc); + ggml_free(ctx); + if (err) *err = "packed-MMID scratch allocation failed"; + return false; + } + if (!owner_input->buffer) { + ggml_gallocr_free(alloc); + ggml_free(ctx); + if (err) *err = "packed-MMID owner input was not allocated"; + return false; + } + + if (cur_backend) { + if (cur_backend_owner) { + ggml_backend_tensor_copy_async( + cur_backend_owner, backend, cur_backend, owner_input); + } else { + ggml_backend_tensor_copy(cur_backend, owner_input); + } + } else { + ggml_backend_tensor_set(owner_input, cur_host, 0, + sizeof(float) * (size_t)n_embd * (size_t)n_tokens); + } + ggml_backend_tensor_set(packed_token_ids, packed_tokens.data(), 0, + sizeof(int32_t) * packed_tokens.size()); + ggml_backend_tensor_set(ids, packed_ids.data(), 0, + sizeof(int32_t) * packed_ids.size()); + ggml_backend_tensor_set(weights, unit_weights.data(), 0, + sizeof(float) * unit_weights.size()); + ggml_backend_tensor_set(inverse_routes, route_to_packed.data(), 0, + sizeof(int32_t) * route_to_packed.size()); + ggml_backend_tensor_set(final_weights, route_weights.data(), 0, + sizeof(float) * route_weights.size()); + + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { + ggml_gallocr_free(alloc); + ggml_free(ctx); + if (err) *err = "packed-MMID graph compute failed"; + return false; + } + ggml_backend_tensor_get(combined, out.data(), 0, + sizeof(float) * out.size()); + + ggml_gallocr_free(alloc); + ggml_free(ctx); + return true; +} + // ── Hot-Only Batched Prefill ── // When all selected experts are in VRAM, skip cold entirely: no CPU graph, // no partition into hot/cold, no merge loop. Pure GPU. @@ -2416,7 +3149,9 @@ bool eval_moe_hybrid_ffn_batched( ggml_gallocr_t * p_cold_alloc, MoeExpertCompute * expert_compute, const MoeExpertLayer * expert_layer, - MoeHybridFfnTelemetry * telemetry) { + MoeHybridFfnTelemetry * telemetry, + ggml_tensor * cur_backend, + const MoeHybridDeviceOutputs * device_outputs) { if (telemetry) *telemetry = {}; const bool materialized_cold = storage.down_cold || storage.gate_up_cold; if (compact_materialized_experts_enabled() && materialized_cold && @@ -2448,6 +3183,190 @@ bool eval_moe_hybrid_ffn_batched( : storage.gate_cold ? (int)storage.gate_cold->ne[2] : 0; const bool cold_on_gpu = storage.cold_backend_kind == MoeHybridColdBackend::Gpu; + const bool inprocess_expert_major = + !expert_compute && expert_major_prefill_enabled(n_tokens) && + cold_on_gpu && storage.cold_backend && + storage.cold_backend != gpu_backend && + n_hot_stack > 0 && n_cold_stack > 0 && + n_cold_stack < cfg.n_expert; + if (inprocess_expert_major) { + static bool logged = false; + if (!logged) { + std::fprintf(stderr, + "[hybrid-ffn] heterogeneous expert-major prefill " + "active tokens=%d hot_stack=%d cold_stack=%d\n", + n_tokens, n_hot_stack, n_cold_stack); + logged = true; + } + const auto wall_t0 = HybridClock::now(); + std::vector hot_partial; + std::vector cold_partial; + std::string hot_err; + std::string cold_err; + const auto cold_t0 = HybridClock::now(); + auto cold_future = std::async(std::launch::async, [&]() { + HybridCudaGraphDisableScope graph_scope( + heterogeneous_prefill_eager_enabled()); + return eval_moe_owner_expert_major_batched( + storage.cold_backend, cfg, desc, + storage.gate_cold, storage.up_cold, storage.down_cold, + storage.gate_up_cold, storage.cold_local_by_global, + cur_host, selected_ids, selected_weights, n_tokens, + /*include_shared=*/false, cold_partial, &cold_err, + cur_backend, gpu_backend); + }); + const auto hot_t0 = HybridClock::now(); + const bool hot_ok = eval_moe_owner_expert_major_batched( + gpu_backend, cfg, desc, + storage.gate_hot, storage.up_hot, storage.down_hot, + storage.gate_up_hot, storage.hot_local_by_global, + cur_host, selected_ids, selected_weights, n_tokens, + /*include_shared=*/true, hot_partial, &hot_err, + cur_backend, gpu_backend); + const auto hot_t1 = HybridClock::now(); + const bool cold_ok = cold_future.get(); + const auto done = HybridClock::now(); + if (!hot_ok || !cold_ok) { + if (err) { + *err = !hot_ok ? hot_err : cold_err; + } + return false; + } + const size_t total = (size_t)cfg.n_embd * (size_t)n_tokens; + out.resize(total); + for (size_t i = 0; i < total; ++i) { + out[i] = hot_partial[i] + cold_partial[i]; + } + if (telemetry) { + telemetry->hot_us += elapsed_us(hot_t0, hot_t1); + telemetry->cold_us += elapsed_us(cold_t0, done); + telemetry->ffn_wall_us += elapsed_us(wall_t0, done); + } + return true; + } + + // A prefill-only full cold stack keeps Strix on its proven full-stack + // MUL_MAT_ID path. Do not pair it with hundreds of q<=4 hot sub-batches: + // pack the R9700 routes by expert and launch one owner graph instead. The + // hot-local map takes priority in eval_moe_hybrid_ffn_batched_core(), so + // skip_hot makes the duplicated Strix stack evaluate cold routes only. + const bool inprocess_full_cold_hot_expert_major = + !expert_compute && expert_major_prefill_enabled(n_tokens) && + cold_on_gpu && storage.cold_backend && + storage.cold_backend != gpu_backend && + n_hot_stack > 0 && n_cold_stack == cfg.n_expert; + if (inprocess_full_cold_hot_expert_major) { + const bool cold_expert_major = + full_cold_expert_major_enabled(); + const bool cold_packed_mmid = + full_cold_packed_mmid_enabled(); + const bool device_join = + device_outputs && device_outputs->valid() && cur_backend; + if (device_join && (cold_expert_major || cold_packed_mmid)) { + if (err) { + *err = "device-resident join requires the full-cold MMID path"; + } + return false; + } + static bool logged = false; + if (!logged) { + const char * cold_mode = cold_packed_mmid + ? "packed-MMID" + : (cold_expert_major ? "expert-major" : "MMID batch"); + std::fprintf(stderr, + "[hybrid-ffn] full-cold Strix %s + expert-major " + "R9700 prefill active tokens=%d hot_stack=%d\n", + cold_mode, + n_tokens, n_hot_stack); + logged = true; + } + + const auto wall_t0 = HybridClock::now(); + std::vector hot_partial; + std::vector cold_partial; + std::string hot_err; + std::string cold_err; + std::vector cold_only_local; + if (cold_expert_major || cold_packed_mmid) { + cold_only_local = storage.cold_local_by_global; + const size_t n = std::min(cold_only_local.size(), + storage.hot_local_by_global.size()); + for (size_t global = 0; global < n; ++global) { + if (storage.hot_local_by_global[global] >= 0) { + cold_only_local[global] = -1; + } + } + } + + const auto cold_t0 = HybridClock::now(); + auto cold_future = std::async(std::launch::async, [&]() { + HybridCudaGraphDisableScope graph_scope( + heterogeneous_prefill_eager_enabled()); + if (cold_packed_mmid) { + return eval_moe_owner_packed_mmid_batched( + storage.cold_backend, cfg, desc, + storage.gate_cold, storage.up_cold, + storage.down_cold, storage.gate_up_cold, + cold_only_local, + cur_host, selected_ids, selected_weights, n_tokens, + cold_partial, &cold_err, p_cold_alloc, + cur_backend, gpu_backend); + } + if (cold_expert_major) { + return eval_moe_owner_expert_major_batched( + storage.cold_backend, cfg, desc, + storage.gate_cold, storage.up_cold, + storage.down_cold, storage.gate_up_cold, + cold_only_local, + cur_host, selected_ids, selected_weights, n_tokens, + /*include_shared=*/false, cold_partial, &cold_err, + cur_backend, gpu_backend); + } + return eval_moe_hybrid_ffn_batched_core( + gpu_backend, cpu_backend, cfg, desc, storage, + cur_host, selected_ids, selected_weights, + n_tokens, cold_partial, &cold_err, + nullptr, p_cold_alloc, nullptr, nullptr, nullptr, + /*skip_cold=*/false, /*skip_hot=*/true, cur_backend, + device_join ? device_outputs->cold : nullptr, + device_join ? device_outputs->backend : nullptr); + }); + + const auto hot_t0 = HybridClock::now(); + const bool hot_ok = eval_moe_owner_expert_major_batched( + gpu_backend, cfg, desc, + storage.gate_hot, storage.up_hot, storage.down_hot, + storage.gate_up_hot, storage.hot_local_by_global, + cur_host, selected_ids, selected_weights, n_tokens, + /*include_shared=*/true, hot_partial, &hot_err, + cur_backend, gpu_backend, + device_join ? device_outputs->hot : nullptr, + device_join ? device_outputs->backend : nullptr); + const auto hot_t1 = HybridClock::now(); + const bool cold_ok = cold_future.get(); + const auto done = HybridClock::now(); + if (!hot_ok || !cold_ok) { + if (err) *err = !hot_ok ? hot_err : cold_err; + return false; + } + + if (!device_join) { + const size_t total = (size_t)cfg.n_embd * (size_t)n_tokens; + out.resize(total); + for (size_t i = 0; i < total; ++i) { + out[i] = hot_partial[i] + cold_partial[i]; + } + } else { + out.clear(); + } + if (telemetry) { + telemetry->hot_us += elapsed_us(hot_t0, hot_t1); + telemetry->cold_us += elapsed_us(cold_t0, done); + telemetry->ffn_wall_us += elapsed_us(wall_t0, done); + } + return true; + } + const int hot_sub_batch = std::min(mmq_safe_sub_batch(), moe_hybrid_prefill_hot_sub_batch_limit()); if (!mmq_full_batch_ok(cfg, n_tokens) @@ -2457,17 +3376,39 @@ bool eval_moe_hybrid_ffn_batched( const int n_embd = cfg.n_embd; const int n_used = cfg.n_expert_used; out.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); - if (expert_compute && expert_layer && - parse_moe_expert_compute_ipc_mode() == MoeExpertComputeIpcMode::Batched) { + const bool remote_cold_full_batch = + expert_compute && expert_layer && + parse_moe_expert_compute_ipc_mode() == MoeExpertComputeIpcMode::Batched; + // The in-process heterogeneous path owns a nearly-full cold stack on + // Strix and a small hot stack on the R9700. A full reduced-stack MMQ + // is pathological for the 24-expert hot side, but is stable once the + // cold stack covers most experts and the prompt supplies enough rows. + // Run that one cold batch on Strix while retaining the proven q<=4 + // MMVQ slices for R9700 hot/shared work. + const bool inprocess_cold_full_batch = + !expert_compute && cold_on_gpu && + storage.cold_backend && storage.cold_backend != gpu_backend && + n_cold_stack * 4 >= cfg.n_expert * 3 && n_tokens >= 64; + if (remote_cold_full_batch || inprocess_cold_full_batch) { std::vector hot_partial((size_t)n_embd * (size_t)n_tokens, 0.0f); std::vector cold_partial; std::string cold_err; const auto cold_t0 = HybridClock::now(); auto cold_future = std::async(std::launch::async, [&]() { - return eval_moe_hybrid_remote_cold_batched( - cfg, storage, cur_host, selected_ids, selected_weights, + HybridCudaGraphDisableScope graph_scope( + heterogeneous_prefill_eager_enabled()); + if (remote_cold_full_batch) { + return eval_moe_hybrid_remote_cold_batched( + cfg, storage, cur_host, selected_ids, selected_weights, + n_tokens, cold_partial, &cold_err, + expert_compute, expert_layer); + } + return eval_moe_hybrid_ffn_batched_core( + gpu_backend, cpu_backend, cfg, desc, storage, + cur_host, selected_ids, selected_weights, n_tokens, cold_partial, &cold_err, - expert_compute, expert_layer); + nullptr, p_cold_alloc, nullptr, nullptr, nullptr, + /*skip_cold=*/false, /*skip_hot=*/true); }); std::vector sub_out; diff --git a/server/src/common/moe_hybrid_ffn_eval.h b/server/src/common/moe_hybrid_ffn_eval.h index 50b612c29..bc063c74c 100644 --- a/server/src/common/moe_hybrid_ffn_eval.h +++ b/server/src/common/moe_hybrid_ffn_eval.h @@ -206,7 +206,19 @@ bool eval_moe_batched_prefill_ffn( std::vector & out, std::string * err = nullptr); -// Batched hybrid prefill FFN: hot on GPU, cold on CPU concurrently. +// Optional device-resident owner destinations for long heterogeneous prefill. +// When present, the hot/shared and cold partials are copied directly into +// these target-backend tensors; the caller can add them in its next graph +// without reading either full hidden-state buffer through the CPU. +struct MoeHybridDeviceOutputs { + ggml_backend_t backend = nullptr; + ggml_tensor * hot = nullptr; + ggml_tensor * cold = nullptr; + + bool valid() const { return backend && hot && cold; } +}; + +// Batched hybrid prefill FFN: hot and cold owners execute concurrently. bool eval_moe_hybrid_ffn_batched( ggml_backend_t gpu_backend, ggml_backend_t cpu_backend, @@ -223,7 +235,14 @@ bool eval_moe_hybrid_ffn_batched( ggml_gallocr_t * p_cold_alloc = nullptr, MoeExpertCompute * expert_compute = nullptr, const MoeExpertLayer * expert_layer = nullptr, - MoeHybridFfnTelemetry * telemetry = nullptr); + MoeHybridFfnTelemetry * telemetry = nullptr, + // Optional device-resident [n_embd, n_tokens] activation. Long + // heterogeneous prefill uses this to avoid a GPU -> host -> GPU bounce + // before both expert owners start. `cur_host` may be null when set. + ggml_tensor * cur_backend = nullptr, + // Optional target-backend partial destinations. Supported by the long + // in-process expert-major path; `out` is not materialized when active. + const MoeHybridDeviceOutputs * device_outputs = nullptr); // Hot-only batched prefill: all selected experts are in VRAM. // Skips cold graph build, CPU compute, and merge — pure GPU path. diff --git a/server/src/common/moe_hybrid_storage.cpp b/server/src/common/moe_hybrid_storage.cpp index bca3c9473..f7b652bf4 100644 --- a/server/src/common/moe_hybrid_storage.cpp +++ b/server/src/common/moe_hybrid_storage.cpp @@ -28,6 +28,14 @@ namespace dflash::common { +static bool duplicate_hot_experts_on_cold_gpu() { + static const bool enabled = []() { + const char * raw = std::getenv("DFLASH_MOE_DUPLICATE_HOT_ON_COLD"); + return raw && *raw && std::strcmp(raw, "0") != 0; + }(); + return enabled; +} + int query_gpu_compute_sm() { #if defined(DFLASH27B_BACKEND_CUDA) int device = -1; @@ -341,6 +349,14 @@ bool build_moe_hybrid_storage(const MoeHybridConfig & cfg, if (err) *err = "failed to select cold expert backend"; return false; } + const bool duplicate_hot_on_cold = + out.cold_backend_kind == MoeHybridColdBackend::Gpu && + duplicate_hot_experts_on_cold_gpu(); + if (duplicate_hot_on_cold) { + std::fprintf(stderr, + "[hybrid-storage] duplicating hot experts on cold GPU " + "for a full expert stack\n"); + } for (int il = 0; il < cfg.n_layer; ++il) { const MoeLayerDesc & desc = layer_descs[(size_t)il]; @@ -368,7 +384,7 @@ bool build_moe_hybrid_storage(const MoeHybridConfig & cfg, is_hot[(size_t)expert] = 1; } for (int expert = 0; expert < cfg.n_expert; ++expert) { - if (!is_hot[(size_t)expert]) { + if (duplicate_hot_on_cold || !is_hot[(size_t)expert]) { dst.cold_local_by_global[(size_t)expert] = (int32_t)dst.cold_expert_ids.size(); dst.cold_expert_ids.push_back((int32_t)expert); } @@ -546,6 +562,14 @@ bool build_moe_hybrid_storage_from_file( if (err) *err = "failed to select cold expert backend"; return false; } + const bool duplicate_hot_on_cold = + out.cold_backend_kind == MoeHybridColdBackend::Gpu && allocate_cold && + duplicate_hot_experts_on_cold_gpu(); + if (duplicate_hot_on_cold) { + std::fprintf(stderr, + "[hybrid-storage] duplicating hot experts on cold GPU " + "for a full expert stack\n"); + } for (int il = 0; il < cfg.n_layer; ++il) { const MoeLayerDesc & desc = layer_descs[(size_t)il]; @@ -575,7 +599,7 @@ bool build_moe_hybrid_storage_from_file( } if (allocate_cold) { for (int expert = 0; expert < cfg.n_expert; ++expert) { - if (!is_hot[(size_t)expert]) { + if (duplicate_hot_on_cold || !is_hot[(size_t)expert]) { dst.cold_local_by_global[(size_t)expert] = (int32_t)dst.cold_expert_ids.size(); dst.cold_expert_ids.push_back((int32_t)expert); } diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 2b0382e38..867c50a8d 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -2,6 +2,7 @@ #include "deepseek4_backend.h" #include "deepseek4_internal.h" +#include "common/peer_access.h" #include "common/sampler.h" #if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) @@ -62,6 +63,32 @@ static void configure_gfx1151_dspark_mmvq_default(int gpu) { #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) { + return; + } + + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, gpu) != cudaSuccess || + std::strncmp(prop.gcnArchName, "gfx1201", 7) != 0) { + return; + } + + // The generic HIP fallback is q=1 because reduced-stack MMQ is unsafe on + // older AMD parts. ROCmFPX MMVQ on gfx1201 is qualified through q=4; + // using that width removes 75% of hot-owner launches while retaining the + // stable vector kernel instead of the pathological full-batch MMQ path. + if (::setenv("DFLASH_MMQ_SUB_BATCH", "4", 0) == 0) { + std::fprintf(stderr, + "[deepseek4] gfx1201 hybrid prefill: defaulting hot " + "expert sub-batch to 4\n"); + } +#else + (void) gpu; +#endif +} + static bool ds4_inprocess_moe_tp_enabled() { return env_flag_enabled("DFLASH_DS4_MOE_TP_INPROC"); } @@ -606,10 +633,8 @@ bool DeepSeek4Backend::validate_prefill_mode() const { } if (w_.moe_hybrid || moe_hybrid_) { std::fprintf(stderr, - "[deepseek4] %s prefill requires every expert to be resident; " - "the selected placement has cold experts\n", + "[deepseek4] %s prefill using heterogeneous layer-major experts\n", prefill_attention_mode_name(cfg_.prefill_mode)); - return false; } return true; } @@ -804,6 +829,7 @@ bool DeepSeek4Backend::init() { // DSpark q=4 is faster through MMVQ. Keep AR and other devices unchanged, // and preserve LUCE_MMVQ_MAX_NCOLS as an explicit override. configure_gfx1151_dspark_mmvq_default(cfg_.device.gpu); + configure_gfx1201_hybrid_sub_batch_default(cfg_.device.gpu); backend_ = ggml_backend_cuda_init(cfg_.device.gpu); if (!backend_) { @@ -1086,6 +1112,12 @@ bool DeepSeek4Backend::init_hybrid_model() { "[deepseek4-moe-tp] in-process expert GPU must differ from local GPU\n"); return false; } + if (g_peer_access_opt_in) { + const bool peer_ok = enable_peer_access_pair(cfg_.device.gpu, expert_gpu); + std::fprintf(stderr, + "[deepseek4-moe-tp] peer access GPU %d <-> GPU %d: %s\n", + cfg_.device.gpu, expert_gpu, peer_ok ? "enabled" : "unavailable"); + } expert_backend_ = ggml_backend_cuda_init(expert_gpu); if (!expert_backend_) { std::fprintf(stderr, @@ -1276,8 +1308,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, // kernels. Smaller tail chunks use the same scheduler or its reference // fallback. const int layer_major_cap = DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS; - const int chunk = (moe_hybrid_ || - !prefill_attention_mode_is_approximate(cfg_.prefill_mode)) + const int chunk = !prefill_attention_mode_is_approximate(cfg_.prefill_mode) ? 1 : std::max(1, std::min(requested_chunk, layer_major_cap)); diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 6d309df47..05b7bb62e 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -59,6 +59,22 @@ static bool ds4_env_flag(const char * name) { return value && value[0] && std::strcmp(value, "0") != 0; } +class Ds4CudaGraphDisableScope { +public: + explicit Ds4CudaGraphDisableScope(bool active) : active_(active) { + if (active_) ggml_backend_cuda_set_graphs_disabled_override(true); + } + ~Ds4CudaGraphDisableScope() { + if (active_) ggml_backend_cuda_set_graphs_disabled_override(false); + } + + Ds4CudaGraphDisableScope(const Ds4CudaGraphDisableScope &) = delete; + Ds4CudaGraphDisableScope & operator=(const Ds4CudaGraphDisableScope &) = delete; + +private: + bool active_; +}; + static int ds4_effective_expert_count(const DeepSeek4Weights & w) { int requested = w.routed_expert_top_k; if (const char * value = std::getenv("DFLASH_DS4_TOPK")) { @@ -2222,6 +2238,78 @@ struct DeepSeek4CachedDecodeHcPostGraph { } }; +// Heterogeneous sparse prefill keeps the HC residual on the R9700. The HC +// pre graph is rebuilt for each layer (the projection weights change) while +// retaining one gallocr arena; the HC post topology is layer-independent and +// remains cached for the current batch width. This avoids retaining 86 large +// per-layer batch graphs and, more importantly, removes two full HC +// device-to-host round trips per layer. +struct DeepSeek4PrefillHcPreGraph { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + int n_tokens = 0; + int layer_idx = -1; + bool ffn = false; + StepGraph sg; + ggml_tensor * split = nullptr; + + bool valid() const { + return owner_ctx && backend && n_tokens > 0 && layer_idx >= 0 && + sg.ctx && sg.gf && sg.alloc && sg.inp_embed && + sg.hidden_states && split; + } + + void reset_graph() { + step_graph_free(sg); + owner_ctx = nullptr; + backend = nullptr; + n_tokens = 0; + layer_idx = -1; + ffn = false; + split = nullptr; + } + + void free() { + step_graph_destroy(sg); + owner_ctx = nullptr; + backend = nullptr; + n_tokens = 0; + layer_idx = -1; + ffn = false; + split = nullptr; + } +}; + +struct DeepSeek4PrefillHcPostGraph { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + int n_tokens = 0; + StepGraph sg; + ggml_tensor * residual_hc = nullptr; + ggml_tensor * block_out = nullptr; + ggml_tensor * block_out_cold = nullptr; + ggml_tensor * split = nullptr; + bool owner_join = false; + + bool valid() const { + return owner_ctx && backend && n_tokens > 0 && sg.ctx && sg.gf && + sg.alloc && sg.hidden_states && residual_hc && block_out && + split && (!owner_join || block_out_cold); + } + + void free() { + step_graph_destroy(sg); + owner_ctx = nullptr; + backend = nullptr; + n_tokens = 0; + residual_hc = nullptr; + block_out = nullptr; + block_out_cold = nullptr; + split = nullptr; + owner_join = false; + } +}; + // Per-step decode scalar inputs shared by all cached per-layer decode graphs. // Values depend only on (kv_start, ratio), so one tensor per slot serves every // layer with that ratio. i32 layout per ratio-slot: {rope_pos, neg_pos, @@ -2845,7 +2933,9 @@ static bool eval_ds4_hybrid( ggml_gallocr_t * cold_alloc, MoeExpertCompute * expert_compute, const MoeExpertLayer * expert_layer, - DeepSeek4StepTelemetry * step_tel) { + DeepSeek4StepTelemetry * step_tel, + ggml_tensor * ffn_normed_backend = nullptr, + const MoeHybridDeviceOutputs * device_outputs = nullptr) { const auto ffn_t0 = Ds4TimingClock::now(); if (!storage.down_cold && !storage.gate_up_cold && !(expert_compute && expert_layer)) { @@ -2930,12 +3020,14 @@ static bool eval_ds4_hybrid( } MoeHybridFfnTelemetry ffn_tel; + std::string batched_err; bool ffn_ok = eval_moe_hybrid_ffn_batched( backend, cpu_backend, hybrid_cfg, desc, storage, ffn_normed_host, selected_host, weights_host, - n_tokens, ffn_out_host, nullptr, hot_alloc, cold_alloc, + n_tokens, ffn_out_host, &batched_err, hot_alloc, cold_alloc, expert_compute, expert_layer, - step_tel ? &ffn_tel : nullptr); + step_tel ? &ffn_tel : nullptr, + ffn_normed_backend, device_outputs); if (ffn_ok) { if (step_tel) { step_tel->ffn_eval_us += ds4_elapsed_us(ffn_t0, Ds4TimingClock::now()); @@ -2950,6 +3042,13 @@ static bool eval_ds4_hybrid( layer); return false; } + if (ffn_normed_backend) { + std::fprintf(stderr, + "[deepseek4] device-resident batched FFN failed at layer %d: " + "%s; refusing an unsafe host fallback\n", + layer, batched_err.empty() ? "unknown error" : batched_err.c_str()); + return false; + } ffn_out_host.assign((size_t)n_embd * (size_t)n_tokens, 0.0f); std::vector single_out; @@ -4516,9 +4615,16 @@ struct DeepSeek4LayerRangeRuntime { HcWeightsCpu hc_output_weights; std::vector hash_routing_tables; std::vector cached_attn_allocs; + // Heterogeneous sparse prefill executes layers serially. Its large + // attention graphs therefore share one scratch arena instead of retaining + // n_layer copies (a 2K-token graph is roughly 550 MiB on the R9700). + DeepSeek4CachedLayerAlloc shared_prefill_attn_alloc; std::vector cached_decode_attn_hc_pre_graphs; std::vector cached_decode_ffn_hc_pre_graphs; DeepSeek4CachedDecodeHcPostGraph cached_decode_hc_post_graph; + DeepSeek4PrefillHcPreGraph prefill_hc_pre_graph; + DeepSeek4PrefillHcPostGraph prefill_hc_post_graph; + DeepSeek4PrefillHcPostGraph prefill_moe_hc_post_graph; std::vector> cached_decode_attn_graphs; std::vector cached_decode_ffn_graphs; DeepSeek4CachedDecodeOutputGraph cached_decode_output_graph; @@ -4532,6 +4638,7 @@ struct DeepSeek4LayerRangeRuntime { alloc.free(); } cached_attn_allocs.clear(); + shared_prefill_attn_alloc.free(); for (auto & graph : cached_decode_attn_hc_pre_graphs) { graph.free(); } @@ -4541,6 +4648,9 @@ struct DeepSeek4LayerRangeRuntime { } cached_decode_ffn_hc_pre_graphs.clear(); cached_decode_hc_post_graph.free(); + prefill_hc_pre_graph.free(); + prefill_hc_post_graph.free(); + prefill_moe_hc_post_graph.free(); for (auto & per_layer : cached_decode_attn_graphs) { for (auto & graph : per_layer) { graph.free(); @@ -4709,6 +4819,141 @@ static bool ds4_fused_ensure_fn_mirrors( return true; } +static bool build_prefill_hc_pre_graph( + DeepSeek4PrefillHcPreGraph & out, + ggml_backend_t backend, + const DeepSeek4Weights & w, + ggml_tensor * fn_f16, + ggml_tensor * base, + const float * scale_data, + int layer_idx, + bool ffn, + int n_tokens) { + if (!backend || !fn_f16 || !base || !scale_data || n_tokens <= 0) { + return false; + } + if ((out.owner_ctx && out.owner_ctx != w.ctx) || + (out.backend && out.backend != backend)) { + out.free(); + } else { + // Keep the largest scratch buffer while replacing layer-specific graph + // metadata and tensor handles. + out.reset_graph(); + } + + ggml_init_params params{}; + params.mem_size = 8 * 1024 * 1024; + params.mem_buffer = nullptr; + params.no_alloc = true; + out.sg.ctx = ggml_init(params); + if (!out.sg.ctx) return false; + + const int64_t hc_dim = (int64_t)w.n_embd * w.n_hc; + const int64_t mix_dim = 2 * (int64_t)w.n_hc + + (int64_t)w.n_hc * w.n_hc; + out.sg.inp_embed = ggml_new_tensor_2d( + out.sg.ctx, GGML_TYPE_F32, hc_dim, n_tokens); + ggml_set_input(out.sg.inp_embed); + + ggml_tensor * normed = ggml_rms_norm( + out.sg.ctx, out.sg.inp_embed, w.hc_eps); + ggml_tensor * mix = ggml_mul_mat(out.sg.ctx, fn_f16, normed); + mix = ggml_reshape_2d(out.sg.ctx, mix, mix_dim, n_tokens); + ggml_tensor * base_f32 = ds4_fused_hc_base_f32(out.sg.ctx, base); + ggml_tensor * pre = ggml_ds4_hc_pre( + out.sg.ctx, mix, base_f32, out.sg.inp_embed, w.n_hc, + w.n_hc_sinkhorn_iter, scale_data[0], scale_data[1], scale_data[2]); + // The HC op packs working+split in one row. Materialize both views on + // device so subsequent graph-to-graph copies have identical contiguous + // layouts; ggml_backend_tensor_copy deliberately rejects strided copies. + out.sg.hidden_states = ggml_cont(out.sg.ctx, ggml_view_2d( + out.sg.ctx, pre, w.n_embd, n_tokens, pre->nb[1], 0)); + out.split = ggml_cont(out.sg.ctx, ggml_view_2d( + out.sg.ctx, pre, mix_dim, n_tokens, pre->nb[1], + (size_t)w.n_embd * sizeof(float))); + + out.sg.gf = ggml_new_graph_custom(out.sg.ctx, 4096, false); + ggml_set_output(out.sg.hidden_states); + ggml_set_output(out.split); + ggml_build_forward_expand(out.sg.gf, out.sg.hidden_states); + ggml_build_forward_expand(out.sg.gf, out.split); + if (!out.sg.alloc) { + out.sg.alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + } + if (!out.sg.alloc || !ggml_gallocr_alloc_graph(out.sg.alloc, out.sg.gf)) { + out.free(); + return false; + } + + out.owner_ctx = w.ctx; + out.backend = backend; + out.n_tokens = n_tokens; + out.layer_idx = layer_idx; + out.ffn = ffn; + return true; +} + +static bool build_prefill_hc_post_graph( + DeepSeek4PrefillHcPostGraph & out, + ggml_backend_t backend, + const DeepSeek4Weights & w, + int n_tokens, + bool owner_join = false) { + if (out.valid() && out.owner_ctx == w.ctx && out.backend == backend && + out.n_tokens == n_tokens && out.owner_join == owner_join) { + return true; + } + out.free(); + if (!backend || n_tokens <= 0) return false; + + ggml_init_params params{}; + params.mem_size = 4 * 1024 * 1024; + params.mem_buffer = nullptr; + params.no_alloc = true; + out.sg.ctx = ggml_init(params); + if (!out.sg.ctx) return false; + + const int64_t hc_dim = (int64_t)w.n_embd * w.n_hc; + const int64_t mix_dim = 2 * (int64_t)w.n_hc + + (int64_t)w.n_hc * w.n_hc; + out.residual_hc = ggml_new_tensor_2d( + out.sg.ctx, GGML_TYPE_F32, hc_dim, n_tokens); + out.block_out = ggml_new_tensor_2d( + out.sg.ctx, GGML_TYPE_F32, w.n_embd, n_tokens); + if (owner_join) { + out.block_out_cold = ggml_new_tensor_2d( + out.sg.ctx, GGML_TYPE_F32, w.n_embd, n_tokens); + } + out.split = ggml_new_tensor_2d( + out.sg.ctx, GGML_TYPE_F32, mix_dim, n_tokens); + ggml_set_input(out.residual_hc); + ggml_set_input(out.block_out); + if (out.block_out_cold) ggml_set_input(out.block_out_cold); + ggml_set_input(out.split); + + ggml_tensor * hc_block_out = out.block_out_cold + ? ggml_add(out.sg.ctx, out.block_out, out.block_out_cold) + : out.block_out; + out.sg.hidden_states = ggml_ds4_hc_post( + out.sg.ctx, out.residual_hc, hc_block_out, out.split, w.n_hc); + out.sg.gf = ggml_new_graph_custom(out.sg.ctx, 1024, false); + ggml_set_output(out.sg.hidden_states); + ggml_build_forward_expand(out.sg.gf, out.sg.hidden_states); + out.sg.alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + if (!out.sg.alloc || !ggml_gallocr_alloc_graph(out.sg.alloc, out.sg.gf)) { + out.free(); + return false; + } + + out.owner_ctx = w.ctx; + out.backend = backend; + out.n_tokens = n_tokens; + out.owner_join = owner_join; + return true; +} + static bool ds4_build_fused_decode_graph( DeepSeek4FusedDecodeCache & fc, DeepSeek4FusedDecodeGraph & fg, @@ -5109,12 +5354,41 @@ static bool eval_ds4_layer_range_hybrid_ffn( MoeExpertComputeRuntime * expert_runtime, MoeHybridRoutingStats * routing_stats, std::vector & out, - DeepSeek4StepTelemetry * telemetry) { + DeepSeek4StepTelemetry * telemetry, + const MoeHybridDeviceOutputs * device_outputs = nullptr) { + const bool trace_prefill = ds4_env_flag("DFLASH_DS4_PREFILL_TRACE"); + if (trace_prefill) { + std::fprintf(stderr, + "[deepseek4-prefill-trace] layer=%d ffn route begin tokens=%d\n", + layer, n_tokens); + } const int n_embd = w.n_embd; const bool hash_routed = layer < w.n_hash_layer && L.ffn_gate_tid2eid && token_ids && hash_table.loaded; const int route_width = ds4_effective_expert_count(w); + MoeExpertCompute * expert_compute = + expert_runtime ? expert_runtime->compute_ptr() : nullptr; + const MoeExpertLayer * expert_layer = + expert_runtime ? expert_runtime->layer_ptr((size_t)layer) : nullptr; + const MoeHybridLayerStorage & layer_storage = + hybrid.layers[(size_t)layer]; + ggml_tensor * hot_stack_ref = layer_storage.gate_up_hot + ? layer_storage.gate_up_hot : layer_storage.gate_hot; + ggml_tensor * cold_stack_ref = layer_storage.gate_up_cold + ? layer_storage.gate_up_cold : layer_storage.gate_cold; + const char * device_input_env = + std::getenv("DFLASH_MOE_PREFILL_DEVICE_INPUT"); + const bool device_input_enabled = + !device_input_env || !*device_input_env || + std::strcmp(device_input_env, "0") != 0; + const bool device_ffn_input = + device_input_enabled && + !expert_compute && n_tokens >= 512 && + layer_storage.cold_backend_kind == MoeHybridColdBackend::Gpu && + layer_storage.cold_backend && layer_storage.cold_backend != backend && + hot_stack_ref && hot_stack_ref->ne[2] > 0 && + cold_stack_ref && cold_stack_ref->ne[2] > 0; const auto route_build_t0 = Ds4TimingClock::now(); ggml_init_params params{}; @@ -5138,6 +5412,21 @@ static bool eval_ds4_layer_range_hybrid_ffn( ggml_free(ctx); return false; } + struct RouteGraphLifetime { + ggml_gallocr_t alloc = nullptr; + ggml_context * ctx = nullptr; + void reset() { + if (alloc) { + ggml_gallocr_free(alloc); + alloc = nullptr; + } + if (ctx) { + ggml_free(ctx); + ctx = nullptr; + } + } + ~RouteGraphLifetime() { reset(); } + } route_graph{alloc, ctx}; if (telemetry) { telemetry->route_build_us += ds4_elapsed_us(route_build_t0, Ds4TimingClock::now()); @@ -5153,17 +5442,27 @@ static bool eval_ds4_layer_range_hybrid_ffn( const auto route_compute_t0 = Ds4TimingClock::now(); const bool route_ok = ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS; + if (trace_prefill) { + std::fprintf(stderr, + "[deepseek4-prefill-trace] layer=%d ffn route compute=%s\n", + layer, route_ok ? "ok" : "failed"); + } if (telemetry) { telemetry->route_compute_us += ds4_elapsed_us(route_compute_t0, Ds4TimingClock::now()); } - std::vector normed_host((size_t)n_embd * (size_t)n_tokens); + std::vector normed_host; + if (!device_ffn_input) { + normed_host.resize((size_t)n_embd * (size_t)n_tokens); + } std::vector probs_host((size_t)w.n_expert * (size_t)n_tokens); if (route_ok) { const auto route_read_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_get(normed, normed_host.data(), 0, - sizeof(float) * normed_host.size()); + if (!device_ffn_input) { + ggml_backend_tensor_get(normed, normed_host.data(), 0, + sizeof(float) * normed_host.size()); + } ggml_backend_tensor_get(probs, probs_host.data(), 0, sizeof(float) * probs_host.size()); if (telemetry) { @@ -5171,8 +5470,17 @@ static bool eval_ds4_layer_range_hybrid_ffn( ds4_elapsed_us(route_read_t0, Ds4TimingClock::now()); } } - ggml_gallocr_free(alloc); - ggml_free(ctx); + if (!device_ffn_input) { + route_graph.reset(); + } else { + static bool logged_device_input = false; + if (!logged_device_input) { + std::fprintf(stderr, + "[deepseek4] device-resident heterogeneous prefill " + "input active; skipping normalized activation readback\n"); + logged_device_input = true; + } + } if (!route_ok) return false; std::vector selected((size_t)route_width * (size_t)n_tokens); @@ -5246,17 +5554,28 @@ static bool eval_ds4_layer_range_hybrid_ffn( MoeHybridConfig cfg = make_ds4_moe_hybrid_config(w); cfg.n_expert_used = route_width; MoeLayerDesc desc = make_ds4_moe_layer_desc(L); - MoeExpertCompute * expert_compute = - expert_runtime ? expert_runtime->compute_ptr() : nullptr; - const MoeExpertLayer * expert_layer = - expert_runtime ? expert_runtime->layer_ptr((size_t)layer) : nullptr; - return eval_ds4_hybrid( + if (trace_prefill) { + std::fprintf(stderr, + "[deepseek4-prefill-trace] layer=%d expert owners begin\n", + layer); + } + const bool ok = eval_ds4_hybrid( backend, hybrid.cpu_backend, cfg, desc, &hybrid, hybrid.layers[(size_t)layer], nullptr, layer, n_embd, route_width, - normed_host.data(), selected.data(), weights.data(), + device_ffn_input ? nullptr : normed_host.data(), + selected.data(), weights.data(), n_tokens, out, nullptr, nullptr, - expert_compute, expert_layer, telemetry); + expert_compute, expert_layer, telemetry, + device_ffn_input ? normed : nullptr, + device_ffn_input ? device_outputs : nullptr); + if (trace_prefill) { + std::fprintf(stderr, + "[deepseek4-prefill-trace] layer=%d expert owners=%s\n", + layer, ok ? "ok" : "failed"); + } + return ok; +} // Exact-order prefill control: retain the layer-major HC/FFN schedule, but run // the attention subgraph one token at a time. This preserves the q=1 QKV, @@ -6052,6 +6371,20 @@ bool deepseek4_step_layer_range( n_tokens >= 2 && n_tokens <= 4 && verify_hooks && layer_begin == 0 && is_last_shard && out_logits && ds4_backend_is_gpu(backend) && ds4_fused_verify_enabled(); + const bool heterogeneous_sparse_prefill = + moe_hybrid && cache.prefill_mode == PrefillAttentionMode::Sparse && + n_tokens > 4 && n_tokens <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS && + layer_begin == 0 && is_last_shard && out_logits && + ds4_backend_is_gpu(backend); + // These graphs are rebuilt around an owner join on every layer, so tensor + // metadata addresses can be recycled for different topologies. Until + // the full heterogeneous layer is captured as one stable scheduler graph, + // eager execution prevents a prefill graph entry from being replayed by + // the following decode/request. The override is thread-local and scoped + // to this forward call; decode graph replay is restored on every return. + Ds4CudaGraphDisableScope heterogeneous_prefill_eager_scope( + heterogeneous_sparse_prefill && + ds4_env_flag("DFLASH_DS4_HYBRID_PREFILL_EAGER")); // A dynamic batch may be supplied by callers other than the DSpark // verifier. Split it whenever it spans a learned-compressor boundary: @@ -6060,7 +6393,7 @@ bool deepseek4_step_layer_range( // as sequential execution while retaining safe batched prefixes. const int first_chunk = deepseek4_safe_compressor_batch_tokens(w, kv_start, n_tokens); if (first_chunk > 0 && first_chunk < n_tokens && - !fused_verify_candidate) { + !fused_verify_candidate && !heterogeneous_sparse_prefill) { const int input_width = layer_begin == 0 ? n_embd : hc_dim; std::vector hc_all; std::vector shard_out_all; @@ -6165,6 +6498,9 @@ bool deepseek4_step_layer_range( auto & cached_decode_ffn_hc_pre_graphs = runtime.cached_decode_ffn_hc_pre_graphs; auto & cached_decode_hc_post_graph = runtime.cached_decode_hc_post_graph; + auto & prefill_hc_pre_graph = runtime.prefill_hc_pre_graph; + auto & prefill_hc_post_graph = runtime.prefill_hc_post_graph; + auto & prefill_moe_hc_post_graph = runtime.prefill_moe_hc_post_graph; auto & cached_decode_attn_graphs = runtime.cached_decode_attn_graphs; auto & cached_decode_ffn_graphs = runtime.cached_decode_ffn_graphs; auto & cached_decode_output_graph = runtime.cached_decode_output_graph; @@ -6185,6 +6521,7 @@ bool deepseek4_step_layer_range( alloc.free(); } cached_attn_allocs.assign((size_t)w.n_layer, {}); + runtime.shared_prefill_attn_alloc.free(); for (auto & g : cached_decode_attn_hc_pre_graphs) { g.free(); } @@ -6194,6 +6531,9 @@ bool deepseek4_step_layer_range( } cached_decode_ffn_hc_pre_graphs.assign((size_t)w.n_layer, {}); cached_decode_hc_post_graph.free(); + runtime.prefill_hc_pre_graph.free(); + runtime.prefill_hc_post_graph.free(); + runtime.prefill_moe_hc_post_graph.free(); for (auto & per_layer : cached_decode_attn_graphs) { for (auto & g : per_layer) { g.free(); @@ -6226,6 +6566,14 @@ bool deepseek4_step_layer_range( DeepSeek4LayerRangeScratch & scratch = runtime.scratch; const int n_expert_used = ds4_effective_expert_count(w); scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, n_expert_used); + const bool trace_prefill = + heterogeneous_sparse_prefill && + ds4_env_flag("DFLASH_DS4_PREFILL_TRACE"); + if (trace_prefill) { + std::fprintf(stderr, + "[deepseek4-prefill-trace] step begin pos=%d tokens=%d\n", + kv_start, n_tokens); + } // Large monolithic batches use PR520's device-resident layer-major // pipeline. Heterogeneous batches are enabled separately once its expert @@ -6365,8 +6713,28 @@ bool deepseek4_step_layer_range( const bool use_backend_decode_hc_direct = use_backend_decode_hc && ds4_backend_is_hip(backend); const bool use_backend_decode_hc_graph = use_backend_decode_hc && !use_backend_decode_hc_direct; + const bool use_backend_prefill_hc = + heterogeneous_sparse_prefill && + ds4_env_flag("DFLASH_DS4_HYBRID_PREFILL_GPU_HC"); ggml_tensor * hc_state_backend = nullptr; - if (use_backend_decode_hc_graph || use_backend_decode_hc_direct) { + if (use_backend_prefill_hc) { + if (!ds4_fused_ensure_fn_mirrors( + fused_decode_graph_cache, backend, w, + hc_layer_weights_range, hc_output_weights_range) || + !build_prefill_hc_post_graph( + prefill_hc_post_graph, backend, w, n_tokens) || + (moe_hybrid && !build_prefill_hc_post_graph( + prefill_moe_hc_post_graph, backend, w, n_tokens, + /*owner_join=*/true))) { + std::fprintf(stderr, + "[deepseek4-prefill] batched GPU HC initialization failed\n"); + return false; + } + ggml_backend_tensor_set(prefill_hc_post_graph.residual_hc, + hc_state.data(), 0, + sizeof(float) * hc_state.size()); + hc_state_backend = prefill_hc_post_graph.residual_hc; + } else if (use_backend_decode_hc_graph || use_backend_decode_hc_direct) { if (!cached_decode_hc_post_graph.valid() || cached_decode_hc_post_graph.owner_ctx != w.ctx || cached_decode_hc_post_graph.backend != backend) { @@ -6390,10 +6758,46 @@ bool deepseek4_step_layer_range( const ggml_tensor * attn_comb_backend = nullptr; const ggml_tensor * ffn_post_backend = nullptr; const ggml_tensor * ffn_comb_backend = nullptr; + const ggml_tensor * attn_split_backend = nullptr; + const ggml_tensor * ffn_split_backend = nullptr; + if (trace_prefill) { + std::fprintf(stderr, + "[deepseek4-prefill-trace] layer=%d attention begin\n", + il); + } // ── HC pre (attention) ────────────────────────────────────── const auto hc_pre_attn_t0 = Ds4TimingClock::now(); - if (use_backend_decode_hc_direct) { + if (use_backend_prefill_hc) { + const auto hc_pre_attn_build_t0 = Ds4TimingClock::now(); + if (!build_prefill_hc_pre_graph( + prefill_hc_pre_graph, backend, w, + fused_decode_graph_cache.fn_attn_f16[(size_t)il], + L.hc_attn_base, hc_lw.attn.scale_data.data(), + il, /*ffn=*/false, n_tokens)) { + std::fprintf(stderr, + "[deepseek4-prefill] batched HC-pre build failed " + "layer %d attn\n", il); + return false; + } + if (telemetry) telemetry->hc_pre_build_us += ds4_elapsed_us( + hc_pre_attn_build_t0, Ds4TimingClock::now()); + ggml_backend_tensor_copy(hc_state_backend, + prefill_hc_pre_graph.sg.inp_embed); + const auto hc_pre_attn_compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute( + backend, prefill_hc_pre_graph.sg.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[deepseek4-prefill] batched HC-pre compute failed " + "layer %d attn\n", il); + return false; + } + if (telemetry) telemetry->hc_pre_compute_us += ds4_elapsed_us( + hc_pre_attn_compute_t0, Ds4TimingClock::now()); + attn_in_backend = prefill_hc_pre_graph.sg.hidden_states; + attn_split_backend = prefill_hc_pre_graph.split; + } else if (use_backend_decode_hc_direct) { auto & cached = cached_decode_attn_hc_pre_graphs[(size_t)il]; if (!cached.valid() || cached.owner_ctx != w.ctx || @@ -6475,7 +6879,8 @@ bool deepseek4_step_layer_range( DeepSeek4CachedDecodeAttnGraph * cached_attn = nullptr; const bool exact_tokenwise_prefill = - !reuse_decode_attn && n_tokens > 1; + !reuse_decode_attn && n_tokens > 1 && + cache.prefill_mode == PrefillAttentionMode::Exact; if (exact_tokenwise_prefill) { const DeepSeek4AttentionImpl attention_impl = cache.prefill_mode == PrefillAttentionMode::Sparse @@ -6601,7 +7006,9 @@ bool deepseek4_step_layer_range( ggml_set_output(attn_out); ggml_build_forward_expand(gf, attn_out); - auto & attn_alloc = cached_attn_allocs[(size_t)il]; + auto & attn_alloc = heterogeneous_sparse_prefill + ? runtime.shared_prefill_attn_alloc + : cached_attn_allocs[(size_t)il]; if (!attn_alloc.valid() || attn_alloc.owner_ctx != w.ctx || attn_alloc.backend != backend) { attn_alloc.free(); attn_alloc.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); @@ -6637,7 +7044,42 @@ bool deepseek4_step_layer_range( return false; } if (telemetry) telemetry->attn_compute_us += ds4_elapsed_us(attn_compute_t0, Ds4TimingClock::now()); - if (use_backend_decode_hc_graph || use_backend_decode_hc_direct) { + if (trace_prefill) { + std::fprintf(stderr, + "[deepseek4-prefill-trace] layer=%d attention compute=ok\n", + il); + } + if (use_backend_prefill_hc) { + if (!attn_split_backend) { + std::fprintf(stderr, + "[deepseek4-prefill] missing HC split layer %d attn\n", + il); + if (ctx) ggml_free(ctx); + return false; + } + if (hc_state_backend != prefill_hc_post_graph.residual_hc) { + ggml_backend_tensor_copy( + hc_state_backend, prefill_hc_post_graph.residual_hc); + } + ggml_backend_tensor_copy( + attn_out, prefill_hc_post_graph.block_out); + ggml_backend_tensor_copy( + const_cast(attn_split_backend), + prefill_hc_post_graph.split); + const auto hc_post_attn_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute( + backend, prefill_hc_post_graph.sg.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[deepseek4-prefill] batched HC-post compute " + "failed layer %d attn\n", il); + if (ctx) ggml_free(ctx); + return false; + } + hc_state_backend = prefill_hc_post_graph.sg.hidden_states; + if (telemetry) telemetry->hc_post_attn_us += ds4_elapsed_us( + hc_post_attn_t0, Ds4TimingClock::now()); + } else if (use_backend_decode_hc_graph || use_backend_decode_hc_direct) { if (hc_state_backend != cached_decode_hc_post_graph.residual_hc) { ggml_backend_tensor_copy(hc_state_backend, cached_decode_hc_post_graph.residual_hc); } @@ -6661,7 +7103,8 @@ bool deepseek4_step_layer_range( } // ── HC post (attention) ───────────────────────────────── - if (!(use_backend_decode_hc_graph || use_backend_decode_hc_direct)) { + if (!use_backend_prefill_hc && + !(use_backend_decode_hc_graph || use_backend_decode_hc_direct)) { const auto hc_post_attn_t0 = Ds4TimingClock::now(); hc_post_batch(next_hc, attn_out_host.data(), @@ -6678,7 +7121,41 @@ bool deepseek4_step_layer_range( // ── HC pre (FFN) ──────────────────────────────────────────── const auto hc_pre_ffn_t0 = Ds4TimingClock::now(); - if (use_backend_decode_hc_direct) { + if (trace_prefill) { + std::fprintf(stderr, + "[deepseek4-prefill-trace] layer=%d ffn begin\n", + il); + } + if (use_backend_prefill_hc) { + const auto hc_pre_ffn_build_t0 = Ds4TimingClock::now(); + if (!build_prefill_hc_pre_graph( + prefill_hc_pre_graph, backend, w, + fused_decode_graph_cache.fn_ffn_f16[(size_t)il], + L.hc_ffn_base, hc_lw.ffn.scale_data.data(), + il, /*ffn=*/true, n_tokens)) { + std::fprintf(stderr, + "[deepseek4-prefill] batched HC-pre build failed " + "layer %d ffn\n", il); + return false; + } + if (telemetry) telemetry->hc_pre_build_us += ds4_elapsed_us( + hc_pre_ffn_build_t0, Ds4TimingClock::now()); + ggml_backend_tensor_copy(hc_state_backend, + prefill_hc_pre_graph.sg.inp_embed); + const auto hc_pre_ffn_compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute( + backend, prefill_hc_pre_graph.sg.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[deepseek4-prefill] batched HC-pre compute failed " + "layer %d ffn\n", il); + return false; + } + if (telemetry) telemetry->hc_pre_compute_us += ds4_elapsed_us( + hc_pre_ffn_compute_t0, Ds4TimingClock::now()); + ffn_in_backend = prefill_hc_pre_graph.sg.hidden_states; + ffn_split_backend = prefill_hc_pre_graph.split; + } else if (use_backend_decode_hc_direct) { auto & cached = cached_decode_ffn_hc_pre_graphs[(size_t)il]; if (!cached.valid() || cached.owner_ctx != w.ctx || @@ -6758,13 +7235,41 @@ bool deepseek4_step_layer_range( il < w.n_hash_layer && L.ffn_gate_tid2eid && token_ids && hash_routing_tables_range[(size_t)il].loaded; ggml_tensor * ffn_out = nullptr; + bool ffn_device_join = false; + MoeHybridDeviceOutputs owner_outputs; if (moe_hybrid) { + const MoeHybridLayerStorage & layer_storage = + moe_hybrid->layers[(size_t)il]; + ggml_tensor * cold_stack = layer_storage.gate_up_cold + ? layer_storage.gate_up_cold + : layer_storage.gate_cold; + const bool local_expert_runtime = + !expert_runtime || !expert_runtime->compute_ptr(); + ffn_device_join = + use_backend_prefill_hc && ffn_in_backend && + local_expert_runtime && + prefill_moe_hc_post_graph.valid() && + cold_stack && cold_stack->ne[2] == w.n_expert; + if (ffn_device_join) { + static bool logged_device_join = false; + if (!logged_device_join) { + std::fprintf(stderr, + "[deepseek4-prefill] device-resident " + "hot+cold owner join active\n"); + logged_device_join = true; + } + owner_outputs.backend = backend; + owner_outputs.hot = prefill_moe_hc_post_graph.block_out; + owner_outputs.cold = + prefill_moe_hc_post_graph.block_out_cold; + } if (!eval_ds4_layer_range_hybrid_ffn( backend, w, L, il, n_tokens, ffn_working.data(), ffn_in_backend, token_ids, hash_routing_tables_range[(size_t)il], *moe_hybrid, expert_runtime, routing_stats, - ffn_out_host, telemetry)) { + ffn_out_host, telemetry, + ffn_device_join ? &owner_outputs : nullptr)) { std::fprintf(stderr, "[deepseek4-moe-tp] layer-range FFN failed layer %d\n", il); @@ -6834,7 +7339,43 @@ bool deepseek4_step_layer_range( } } - if (use_backend_decode_hc_graph || use_backend_decode_hc_direct) { + if (use_backend_prefill_hc) { + if (!ffn_split_backend) { + std::fprintf(stderr, + "[deepseek4-prefill] missing HC split layer %d ffn\n", + il); + return false; + } + DeepSeek4PrefillHcPostGraph & hc_post_graph = + ffn_device_join + ? prefill_moe_hc_post_graph + : prefill_hc_post_graph; + if (hc_state_backend != hc_post_graph.residual_hc) { + ggml_backend_tensor_copy( + hc_state_backend, hc_post_graph.residual_hc); + } + if (!ffn_device_join) { + ggml_backend_tensor_set( + hc_post_graph.block_out, + ffn_out_host.data(), 0, + sizeof(float) * ffn_out_host.size()); + } + ggml_backend_tensor_copy( + const_cast(ffn_split_backend), + hc_post_graph.split); + const auto hc_post_ffn_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute( + backend, hc_post_graph.sg.gf) != + GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[deepseek4-prefill] batched HC-post compute " + "failed layer %d ffn\n", il); + return false; + } + hc_state_backend = hc_post_graph.sg.hidden_states; + if (telemetry) telemetry->hc_post_ffn_us += ds4_elapsed_us( + hc_post_ffn_t0, Ds4TimingClock::now()); + } else if (use_backend_decode_hc_graph || use_backend_decode_hc_direct) { if (hc_state_backend != cached_decode_hc_post_graph.residual_hc) { ggml_backend_tensor_copy(hc_state_backend, cached_decode_hc_post_graph.residual_hc); @@ -6862,7 +7403,8 @@ bool deepseek4_step_layer_range( } // ── HC post (FFN) ─────────────────────────────────────── - if (!(use_backend_decode_hc_graph || use_backend_decode_hc_direct)) { + if (!use_backend_prefill_hc && + !(use_backend_decode_hc_graph || use_backend_decode_hc_direct)) { const auto hc_post_ffn_t0 = Ds4TimingClock::now(); hc_post_batch(next_hc, ffn_out_host.data(), @@ -6897,7 +7439,8 @@ bool deepseek4_step_layer_range( } } - if ((use_backend_decode_hc_graph || use_backend_decode_hc_direct) && hc_state_backend) { + if ((use_backend_prefill_hc || use_backend_decode_hc_graph || + use_backend_decode_hc_direct) && hc_state_backend) { ggml_backend_tensor_get(hc_state_backend, hc_state.data(), 0, sizeof(float) * hc_state.size()); } diff --git a/server/test/test_ggml_rmsnorm_batch.cpp b/server/test/test_ggml_rmsnorm_batch.cpp new file mode 100644 index 000000000..18296d3a7 --- /dev/null +++ b/server/test/test_ggml_rmsnorm_batch.cpp @@ -0,0 +1,93 @@ +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cuda.h" +#include "ggml.h" + +#include +#include +#include +#include +#include + +int main(int argc, char ** argv) { + const int n_tokens = argc > 1 ? std::max(1, std::atoi(argv[1])) : 64; + const bool initialize_peer = argc <= 2 || std::atoi(argv[2]) != 0; + const int device = argc > 3 ? std::max(0, std::atoi(argv[3])) : 0; + constexpr int n_embd = 4096; + constexpr float eps = 1.0e-6f; + + ggml_backend_t backend = ggml_backend_cuda_init(device); + if (!backend) { + std::fprintf(stderr, "failed to initialize GPU %d\n", device); + return 2; + } + + // Leave a second HIP device initialized, matching the heterogeneous + // server's process state. The backend must still execute this graph on 0. + ggml_backend_t peer = nullptr; + if (initialize_peer && ggml_backend_cuda_get_device_count() > 1) { + peer = ggml_backend_cuda_init(device == 0 ? 1 : 0); + } + + ggml_init_params params{}; + params.mem_size = 2 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) return 3; + + ggml_tensor * input = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_embd, n_tokens); + ggml_tensor * weight = ggml_new_tensor_1d( + ctx, GGML_TYPE_F32, n_embd); + ggml_set_input(input); + ggml_set_input(weight); + ggml_tensor * output = ggml_mul( + ctx, ggml_rms_norm(ctx, input, eps), weight); + ggml_set_output(output); + + ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); + ggml_build_forward_expand(graph, output); + ggml_gallocr_t alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + if (!alloc || !ggml_gallocr_alloc_graph(alloc, graph)) { + std::fprintf(stderr, "graph allocation failed\n"); + return 4; + } + + std::vector input_data((size_t)n_embd * n_tokens); + std::vector weight_data(n_embd); + for (size_t i = 0; i < input_data.size(); ++i) { + input_data[i] = ((int)(i % 31) - 15) * 0.01f; + } + for (int i = 0; i < n_embd; ++i) { + weight_data[(size_t)i] = 0.75f + (i % 17) * 0.01f; + } + ggml_backend_tensor_set(input, input_data.data(), 0, + input_data.size() * sizeof(float)); + ggml_backend_tensor_set(weight, weight_data.data(), 0, + weight_data.size() * sizeof(float)); + + const enum ggml_status status = ggml_backend_graph_compute(backend, graph); + if (status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "graph compute failed: %d\n", (int)status); + return 5; + } + + std::vector output_data(input_data.size()); + ggml_backend_tensor_get(output, output_data.data(), 0, + output_data.size() * sizeof(float)); + for (float value : output_data) { + if (!std::isfinite(value)) { + std::fprintf(stderr, "non-finite output\n"); + return 6; + } + } + std::printf("PASS device=%d peer=%s tokens=%d embd=%d\n", + device, peer ? "initialized" : "absent", n_tokens, n_embd); + + ggml_gallocr_free(alloc); + ggml_free(ctx); + if (peer) ggml_backend_free(peer); + ggml_backend_free(backend); + return 0; +} From 4d092520a4aba6112d1679260116db606e3e5154 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:29:04 -0700 Subject: [PATCH 15/15] fix(ds4): harden heterogeneous GPU correctness --- .../llama.cpp/ggml/src/ggml-backend-meta.cpp | 26 ++++++-- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 63 +++++++++---------- .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 20 +++++- server/src/common/moe_hybrid_ffn_eval.cpp | 5 +- server/tests/test_deepseek4_unit.cpp | 51 +++++++++++++-- 5 files changed, 118 insertions(+), 47 deletions(-) diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp index bf974d722..850ba097b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp @@ -496,6 +496,21 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co return ret; }; + // Nonlinear/fused operations cannot consume PARTIAL values locally: doing + // so would apply the operation before the meta backend's all-reduce. A + // computed partial producer is reported as MIRRORED when assume_sync is + // true, which places this operation in the subgraph after that reduction. + // A genuinely partial input is rejected instead of being miscomputed. + auto handle_mirrored = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { + for (size_t i = 0; i < GGML_MAX_SRC; ++i) { + if (tensor->src[i] == nullptr || tensor->src[i] == tensor) { + continue; + } + GGML_ASSERT(src_ss[i].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, 1}; + }; + // Some ops process data on a per-row bases: auto handle_per_row = [&](const std::vector & src_ss) -> ggml_backend_meta_split_state { GGML_ASSERT(src_ss[0].axis != GGML_BACKEND_SPLIT_AXIS_0); @@ -969,14 +984,15 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co split_state = handle_gated_delta_net(src_ss); } break; case GGML_OP_DS4_INDEXER_QAT: { - split_state = handle_per_row(src_ss); + // Hadamard/FP4 quantization is nonlinear. Reduce any partial + // producer before applying it independently on each backend. + split_state = handle_mirrored(src_ss); } break; case GGML_OP_DS4_INDEXER_SCORE: case GGML_OP_DS4_INDEXER_MASK: { - // These fused DS4 indexer operations combine axes from - // multiple inputs. Keep them local unless every source is - // mirrored on each simple backend. - split_state = handle_generic(src_ss, /*scalar_only =*/ true); + // SCORE includes ReLU, and MASK consumes its nonlinear result; + // neither may run on unreduced dot-product shards. + split_state = handle_mirrored(src_ss); } break; case GGML_OP_UNARY: { split_state = handle_generic(src_ss, /*scalar_only =*/ false); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index 153739b90..cec06459b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -114,37 +114,50 @@ struct ds4_inverse_rope_params { // Keep these expressions aligned with rope.cu. The attention result is first // stored in shared F32, matching the standalone attention-output store/load // boundary, before the pair is rotated. -__device__ static __forceinline__ float ds4_rope_theta_fp64( +// Return the unreduced angle. Frequency scaling and YaRN interpolation must +// happen before modulo reduction; reducing here changes the angle whenever +// freq_scale is not an integer (the production DS4 configuration uses YaRN). +__device__ static __forceinline__ double ds4_rope_theta_fp64( int32_t p, float theta_scale, int exp_int) { - const double tau = 6.2831853071795864769; - double angle = exp_int == 0 + return exp_int == 0 ? (double) p : (double) p * pow((double) theta_scale, (double) exp_int); - angle -= tau * floor(angle * (1.0 / tau)); - return (float) angle; } -__device__ static __forceinline__ void ds4_inverse_rope_coefficients( - int pair, int token, +__device__ static __forceinline__ void ds4_rope_coefficients_at_position( + int pair, int32_t position, const ds4_inverse_rope_params & p, float & cos_theta, float & sin_theta) { const int i0 = 2 * pair; - const float theta_extrap = ds4_rope_theta_fp64( - -(p.kv_start + token), p.theta_scale, pair); - const float theta_interp = p.freq_scale * theta_extrap; - float theta = theta_interp; + const double theta_extrap = ds4_rope_theta_fp64( + position, p.theta_scale, pair); + const double theta_interp = (double) p.freq_scale * theta_extrap; + double theta = theta_interp; float mscale = p.attn_factor; if (p.ext_factor != 0.0f) { const float ramp_y = (i0 / 2 - p.corr_low) / max(0.001f, p.corr_high - p.corr_low); const float ramp_mix = (1.0f - min(1.0f, max(0.0f, ramp_y))) * p.ext_factor; - theta = theta_interp * (1.0f - ramp_mix) + - theta_extrap * ramp_mix; + theta = theta_interp * (1.0 - (double) ramp_mix) + + theta_extrap * (double) ramp_mix; mscale *= 1.0f + 0.1f * logf(1.0f / p.freq_scale); } - cos_theta = cosf(theta) * mscale; - sin_theta = sinf(theta) * mscale; + + // Match rope_yarn(): preserve FP64 precision through all scaling and + // blend operations, then reduce only at the trig boundary. + const double tau = 6.2831853071795864769; + theta -= tau * floor(theta * (1.0 / tau)); + cos_theta = cosf((float) theta) * mscale; + sin_theta = sinf((float) theta) * mscale; +} + +__device__ static __forceinline__ void ds4_inverse_rope_coefficients( + int pair, int token, + const ds4_inverse_rope_params & p, + float & cos_theta, float & sin_theta) { + ds4_rope_coefficients_at_position( + pair, -(p.kv_start + token), p, cos_theta, sin_theta); } // Forward counterpart of ds4_inverse_rope_coefficients. Keep the expressions @@ -155,23 +168,8 @@ __device__ static __forceinline__ void ds4_forward_rope_coefficients( int pair, int token, const ds4_inverse_rope_params & p, float & cos_theta, float & sin_theta) { - const int i0 = 2 * pair; - const float theta_extrap = ds4_rope_theta_fp64( - p.kv_start + token, p.theta_scale, pair); - const float theta_interp = p.freq_scale * theta_extrap; - float theta = theta_interp; - float mscale = p.attn_factor; - if (p.ext_factor != 0.0f) { - const float ramp_y = (i0 / 2 - p.corr_low) / - max(0.001f, p.corr_high - p.corr_low); - const float ramp_mix = - (1.0f - min(1.0f, max(0.0f, ramp_y))) * p.ext_factor; - theta = theta_interp * (1.0f - ramp_mix) + - theta_extrap * ramp_mix; - mscale *= 1.0f + 0.1f * logf(1.0f / p.freq_scale); - } - cos_theta = cosf(theta) * mscale; - sin_theta = sinf(theta) * mscale; + ds4_rope_coefficients_at_position( + pair, p.kv_start + token, p, cos_theta, sin_theta); } __device__ static __forceinline__ void ds4_apply_inverse_rope_pair( @@ -994,6 +992,7 @@ __global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( ? (INDEXED_MASK ? indexed_count : n_kv) : -1; } + __syncthreads(); // Preserve each thread's original r = tid + 256*k order. In indexed mode, // owner_ranks is the exact selected subsequence of that traversal, so the diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu index 4ca96d5a0..6628dfe80 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 @@ -2631,7 +2631,18 @@ static bool ggml_cuda_try_fuse_mul_mat_glu( const ggml_tensor * src1 = up->src[1]; const ggml_tensor * ids = up->src[2]; - if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) { + // Vector fusion writes the final GLU tensor directly and indexes routing + // ids from the matmul layout. A reshape that changes token/expert axes is + // therefore not a shape-only detail for this kernel. MMQ below remains + // safe because it materializes both matmul outputs and lets GLU consume + // the original reshape nodes. + const bool direct_vector_layout = + ggml_are_same_shape(gate, glu->src[0]) && + ggml_are_same_stride(gate, glu->src[0]) && + ggml_are_same_shape(up, glu->src[1]) && + ggml_are_same_stride(up, glu->src[1]); + + if (direct_vector_layout && ggml_cuda_should_fuse_mul_mat_vec_f(up)) { ggml_cuda_mm_fusion_args_host fusion_data{}; fusion_data.gate = gate->src[0]; ggml_cuda_set_fusion_glu_params(fusion_data, glu); @@ -2639,7 +2650,7 @@ static bool ggml_cuda_try_fuse_mul_mat_glu( return true; } - if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) { + if (direct_vector_layout && ggml_cuda_should_fuse_mul_mat_vec_q(up)) { ggml_cuda_mm_fusion_args_host fusion_data{}; fusion_data.gate = gate->src[0]; ggml_cuda_set_fusion_glu_params(fusion_data, glu); @@ -3621,6 +3632,9 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ if (cuda_ctx_src->device == cuda_ctx_dst->device) { #if defined(GGML_USE_HIP) ggml_cuda_flush_peer_copy_batch("same-device-copy"); + // Flushing an unrelated A->B batch leaves B selected. Restore the + // source device before submitting work to this backend's stream. + ggml_cuda_set_device(cuda_ctx_src->device); #endif CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); } else { @@ -3679,6 +3693,7 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ // src and dst are on the same backend #if defined(GGML_USE_HIP) ggml_cuda_flush_peer_copy_batch("same-backend-copy"); + ggml_cuda_set_device(cuda_ctx_src->device); #endif CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); } @@ -3690,6 +3705,7 @@ static void ggml_backend_cuda_synchronize(ggml_backend_t backend) { #if defined(GGML_USE_HIP) ggml_cuda_flush_peer_copy_batch("backend-synchronize"); + ggml_cuda_set_device(cuda_ctx->device); #endif CUDA_CHECK(cudaStreamSynchronize(cuda_ctx->stream())); diff --git a/server/src/common/moe_hybrid_ffn_eval.cpp b/server/src/common/moe_hybrid_ffn_eval.cpp index cd0346459..c88dab31d 100644 --- a/server/src/common/moe_hybrid_ffn_eval.cpp +++ b/server/src/common/moe_hybrid_ffn_eval.cpp @@ -2700,7 +2700,8 @@ static bool eval_moe_owner_expert_major_batched( } if (device_output) { - if (!device_output_owner || !gpu_reduce || !combined_out || + if (!device_output_owner || !combined_out || + (n_pairs > 0 && !gpu_reduce) || device_output->type != GGML_TYPE_F32 || device_output->ne[0] != n_embd || device_output->ne[1] != n_tokens) { @@ -3154,7 +3155,7 @@ bool eval_moe_hybrid_ffn_batched( const MoeHybridDeviceOutputs * device_outputs) { if (telemetry) *telemetry = {}; const bool materialized_cold = storage.down_cold || storage.gate_up_cold; - if (compact_materialized_experts_enabled() && materialized_cold && + if (cur_host && compact_materialized_experts_enabled() && materialized_cold && !expert_compute && n_tokens > 0 && n_tokens <= 4) { out.assign((size_t)cfg.n_embd * (size_t)n_tokens, 0.0f); std::vector token_out; diff --git a/server/tests/test_deepseek4_unit.cpp b/server/tests/test_deepseek4_unit.cpp index 0aa7a73ab..3967f3b25 100644 --- a/server/tests/test_deepseek4_unit.cpp +++ b/server/tests/test_deepseek4_unit.cpp @@ -2155,6 +2155,14 @@ static void test_ds4_flash_attention_inverse_rope_fallback_gpu() { constexpr int raw_rows = 128; constexpr int n_comp_rows = 8; constexpr int n_kv = raw_rows + n_comp_rows; + constexpr int kv_start = 131071; + constexpr float freq_base = 10000.0f; + constexpr float freq_scale = 1.0f / 16.0f; + constexpr float ext_factor = 1.0f; + constexpr float attn_factor = 1.0f; + constexpr float beta_fast = 32.0f; + constexpr float beta_slow = 1.0f; + constexpr int n_ctx_orig = 8192; ggml_context * ctx = make_test_context(3u << 20); TEST_ASSERT_MSG(ctx != nullptr, "ggml_init failed"); @@ -2176,13 +2184,13 @@ static void test_ds4_flash_attention_inverse_rope_fallback_gpu() { ctx, q, k, v, mask, 1.0f / std::sqrt((float) head_dim), 0.0f, 0.0f); // One retained compressed block selects the single-head fallback without - // pruning any live row. Position zero makes inverse RoPE the identity, so - // a row-constant V has an exact, simple reference result. + // pruning any live row. A row-constant V makes attention itself an exact + // identity, leaving a nontrivial scaled/YaRN inverse RoPE to validate. ggml_flash_attn_ext_set_ds4_sparse( output, raw_rows, raw_rows, 4, n_comp_rows); ggml_flash_attn_ext_set_ds4_inverse_rope( - output, 0, 10000.0f, 1.0f, 0.0f, 1.0f, - 32.0f, 1.0f, 8192, false); + output, kv_start, freq_base, freq_scale, ext_factor, attn_factor, + beta_fast, beta_slow, n_ctx_orig, false); ggml_set_output(output); TEST_ASSERT_MSG(ggml_backend_supports_op(backend, output), "GPU rejected DS4 inverse-RoPE fallback attention"); @@ -2198,18 +2206,49 @@ static void test_ds4_flash_attention_inverse_rope_fallback_gpu() { std::vector q_data(head_dim); std::vector k_data((size_t) head_dim * n_kv); std::vector v_data((size_t) head_dim * n_kv); + std::vector base_value(head_dim); std::vector expected(head_dim); std::vector mask_data( n_kv, ggml_fp32_to_fp16(0.0f)); for (int d = 0; d < head_dim; ++d) { q_data[(size_t) d] = ((d % 19) - 9) * 0.002f; - expected[(size_t) d] = ((d % 17) - 8) * 0.003f; + base_value[(size_t) d] = ((d % 17) - 8) * 0.003f; + } + expected = base_value; + float corr_dims[2]; + ggml_rope_yarn_corr_dims( + 64, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); + const float theta_scale = std::pow(freq_base, -2.0f / 64.0f); + constexpr double tau = 6.2831853071795864769; + for (int pair = 0; pair < 32; ++pair) { + const int i0 = 2 * pair; + const double theta_extrap = -(double) kv_start * + std::pow((double) theta_scale, (double) pair); + const double theta_interp = + (double) freq_scale * theta_extrap; + const float ramp_y = (pair - corr_dims[0]) / + std::max(0.001f, corr_dims[1] - corr_dims[0]); + const float ramp_mix = + (1.0f - std::min(1.0f, std::max(0.0f, ramp_y))) * + ext_factor; + double theta = theta_interp * (1.0 - (double) ramp_mix) + + theta_extrap * (double) ramp_mix; + theta -= tau * std::floor(theta * (1.0 / tau)); + const float mscale = attn_factor * + (1.0f + 0.1f * std::log(1.0f / freq_scale)); + const float cos_theta = std::cos((float) theta) * mscale; + const float sin_theta = std::sin((float) theta) * mscale; + const size_t d0 = (size_t) head_dim - 64 + (size_t) i0; + const float x0 = base_value[d0 + 0]; + const float x1 = base_value[d0 + 1]; + expected[d0 + 0] = x0 * cos_theta - x1 * sin_theta; + expected[d0 + 1] = x0 * sin_theta + x1 * cos_theta; } for (int row = 0; row < n_kv; ++row) { for (int d = 0; d < head_dim; ++d) { k_data[(size_t) row * head_dim + d] = (((row + d) % 23) - 11) * 0.002f; - v_data[(size_t) row * head_dim + d] = expected[(size_t) d]; + v_data[(size_t) row * head_dim + d] = base_value[(size_t) d]; } } ggml_backend_tensor_set(q, q_data.data(), 0,