From b9c006970bf863d44ff186f0b23c02a828dd3bc8 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 8 Jul 2026 15:41:07 +0200 Subject: [PATCH 01/10] feat(ggml): vendor ROCmFPX quant support for DeepSeek V4 Flash Ports the ROCmFP4/ROCmFPX ggml work (previously lucebox-ggml #36, on the old submodule pointer) into the vendored server/deps/llama.cpp tree. - ROCmFP4/ROCmFPX quant types + CPU reference conversions (ggml/rocmfp4, ggml/rocmfpx) and the ggml type-trait registrations. - CUDA/HIP dequant, copy, getrows, MMVQ vecdot, MMVF, unary and FA paths for the new types. - Fused DS4 hyper-connection op (GGML_OP_DS4_HC) with the register-resident sinkhorn kernel; inert unless emitted by the DS4 fused-decode path. - DS4 SwiGLU split op plumbing for the fused FFN matvec paths. Layered on current main, so the main-side ggml work is preserved (fp64 RoPE reduction, #497 RDNA MMQ tile, LUCE_MMQ_DP_MAX_NE1, MMVQ_MAX_MOE_BATCH_SIZE, fused dual set_rows, raw-span guard). Review fixes on top of #36: - ggml_ftype_to_ggml_type: handle the 11 new ROCmFPX ftypes (dominant-type mapping) so the enum switch is -Wswitch/-Werror clean. - FP6 MMVQ vecdot: pad qs[] to avoid a stack over-read of the last window (bit-identical; the over-read bits were already masked out). --- server/deps/llama.cpp/ggml/include/ggml.h | 60 +- server/deps/llama.cpp/ggml/rocmfp4/README.md | 1020 +++++++++++++++ server/deps/llama.cpp/ggml/rocmfp4/rocmfp4.c | 797 ++++++++++++ server/deps/llama.cpp/ggml/rocmfp4/rocmfp4.h | 58 + .../llama.cpp/ggml/rocmfp4/rocmfp4_hip.cu | 85 ++ .../ggml/rocmfp4/rocmfp4_hip_codebook.cuh | 79 ++ .../ggml/rocmfp4/rocmfp4_hip_scale.cuh | 157 +++ server/deps/llama.cpp/ggml/rocmfpx/README.md | 106 ++ server/deps/llama.cpp/ggml/rocmfpx/rocmfpx.c | 1140 +++++++++++++++++ server/deps/llama.cpp/ggml/rocmfpx/rocmfpx.h | 111 ++ .../llama.cpp/ggml/rocmfpx/test_rocmfpx.c | 164 +++ server/deps/llama.cpp/ggml/src/CMakeLists.txt | 4 + server/deps/llama.cpp/ggml/src/ggml-common.h | 6 + .../llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c | 4 + .../deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp | 78 ++ .../llama.cpp/ggml/src/ggml-cuda/common.cuh | 68 + .../llama.cpp/ggml/src/ggml-cuda/concat.cu | 97 +- .../llama.cpp/ggml/src/ggml-cuda/convert.cu | 110 ++ .../deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu | 131 ++ .../ggml/src/ggml-cuda/dequantize.cuh | 123 ++ .../llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu | 367 ++++++ .../llama.cpp/ggml/src/ggml-cuda/ds4-hc.cuh | 10 + .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 136 +- .../llama.cpp/ggml/src/ggml-cuda/getrows.cu | 24 + .../llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu | 48 +- .../deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu | 14 +- .../deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu | 73 +- .../llama.cpp/ggml/src/ggml-cuda/unary.cu | 57 + .../llama.cpp/ggml/src/ggml-cuda/unary.cuh | 10 + .../llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh | 375 ++++++ server/deps/llama.cpp/ggml/src/ggml.c | 179 ++- 31 files changed, 5636 insertions(+), 55 deletions(-) create mode 100644 server/deps/llama.cpp/ggml/rocmfp4/README.md create mode 100644 server/deps/llama.cpp/ggml/rocmfp4/rocmfp4.c create mode 100644 server/deps/llama.cpp/ggml/rocmfp4/rocmfp4.h create mode 100644 server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip.cu create mode 100644 server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_codebook.cuh create mode 100644 server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_scale.cuh create mode 100644 server/deps/llama.cpp/ggml/rocmfpx/README.md create mode 100644 server/deps/llama.cpp/ggml/rocmfpx/rocmfpx.c create mode 100644 server/deps/llama.cpp/ggml/rocmfpx/rocmfpx.h create mode 100644 server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu create mode 100644 server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cuh diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index d76166f16..1d8f22bb8 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -430,7 +430,13 @@ extern "C" { GGML_TYPE_NVFP4 = 40, // NVFP4 (4 blocks, E4M3 scale) GGML_TYPE_Q1_0 = 41, GGML_TYPE_TQ3_0 = 42, // TurboQuant 3.5 bpv (3-bit Lloyd-Max + FWHT rotation) - GGML_TYPE_COUNT = 43, + GGML_TYPE_Q4_0_ROCMFP4 = 100, + GGML_TYPE_Q4_0_ROCMFP4_FAST = 101, + GGML_TYPE_Q6_0_ROCMFPX = 102, + GGML_TYPE_Q8_0_ROCMFPX = 103, + GGML_TYPE_Q3_0_ROCMFPX = 104, + GGML_TYPE_Q2_0_ROCMFP2 = 107, + GGML_TYPE_COUNT = 108, }; // precision @@ -468,6 +474,18 @@ extern "C" { GGML_FTYPE_MOSTLY_MXFP4 = 25, // except 1d tensors GGML_FTYPE_MOSTLY_NVFP4 = 26, // except 1d tensors GGML_FTYPE_MOSTLY_Q1_0 = 27, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4 = 100, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_LEAN = 101, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_COHERENT = 102, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST = 103, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST_COHERENT = 104, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX = 105, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX_LEAN = 106, + GGML_FTYPE_MOSTLY_Q6_0_ROCMFPX = 110, + GGML_FTYPE_MOSTLY_Q8_0_ROCMFPX = 111, + GGML_FTYPE_MOSTLY_Q3_0_ROCMFPX = 112, + GGML_FTYPE_MOSTLY_Q2_0_ROCMFP2 = 118, + GGML_FTYPE_MOSTLY_Q2_0_ROCMFP2_STRIX = 119, }; // available tensor operations: @@ -583,6 +601,8 @@ extern "C" { GGML_OP_MOE_FUSED, // Fused MoE FFN: gate+up+swiglu+down+weighted_sum+shared_expert + GGML_OP_DS4_HC, // Fused DeepSeek4 hyper-connection pre/post/out mixing + GGML_OP_COUNT, }; @@ -618,6 +638,7 @@ extern "C" { GGML_GLU_OP_GEGLU, GGML_GLU_OP_SWIGLU, GGML_GLU_OP_SWIGLU_OAI, + GGML_GLU_OP_SWIGLU_DS4, GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK, @@ -1339,6 +1360,12 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * b); + GGML_API struct ggml_tensor * ggml_swiglu_ds4_split( + struct ggml_context * ctx, + struct ggml_tensor * gate, + struct ggml_tensor * up, + float clamp); + GGML_API struct ggml_tensor * ggml_geglu_erf_split( struct ggml_context * ctx, struct ggml_tensor * a, @@ -2399,6 +2426,37 @@ extern "C" { struct ggml_tensor * experts, struct ggml_tensor * expert_weights); + // 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) } + GGML_API struct ggml_tensor * ggml_ds4_hc_pre( + struct ggml_context * ctx, + struct ggml_tensor * mix, + struct ggml_tensor * base, + struct ggml_tensor * hc_state, + int n_hc, + int sinkhorn_iters, + float pre_scale, + float post_scale, + float comb_scale); + + // ggml_ds4_hc_post: residual hc_state + block_out + split -> new hc_state + GGML_API struct ggml_tensor * ggml_ds4_hc_post( + struct ggml_context * ctx, + struct ggml_tensor * residual_hc, + struct ggml_tensor * block_out, + 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, + struct ggml_tensor * mix, + struct ggml_tensor * base, + struct ggml_tensor * hc_state, + int n_hc, + float pre_scale); + // 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/rocmfp4/README.md b/server/deps/llama.cpp/ggml/rocmfp4/README.md new file mode 100644 index 000000000..39bb79302 --- /dev/null +++ b/server/deps/llama.cpp/ggml/rocmfp4/README.md @@ -0,0 +1,1020 @@ +# ROCmFP4 + +ROCmFP4 is the experimental `Q4_0_ROCMFP4` GGUF tensor type for Strix Halo. +The current AMD-tuned variant stores 32 weights per block as packed +E2M1-derived 4-bit values plus two unsigned E4M3 scale bytes, one per +16-weight half, for 18 bytes per block. The codebook stores half-scale signed +integer levels up to `10`, representing `5.0` raw-scale units after the scale +factor is applied. This keeps outlier pull lower than the original wider test +range while preserving a fast integer-dot backend shape. + +This directory owns the format-specific implementation. The rest of ggml only +registers and dispatches the type. + +Current status: +- The format runs on CPU, Vulkan, and ROCm/HIP in this custom tree. +- `Q4_0_ROCMFP4` is the pure 4.50 BPW dual-scale path. +- `Q4_0_ROCMFP4_LEAN` keeps ROCmFP4 for dense tensors but protects token + embeddings with `Q5_K`. On the Qwen3-4B Strix test pass this closed most of + the Q4_0 coherence gap while staying around normal Q4 size. +- `Q4_0_ROCMFP4_COHERENT` protects token embeddings with `Q6_K` and is the + quality-first ROCmFP4 preset. +- `Q4_0_ROCMFP4_FAST` is the 4.25 BPW single-scale speed path. It is the + smallest and fastest decode variant, but the pure version gives up too much + PPL to be the default quality target. +- `Q4_0_ROCMFP4_FAST_COHERENT` combines the fast 4.25 BPW transformer layout + with `Q6_K` token embeddings. On the Qwen3-4B Strix test pass it is the + current balanced AMD target: smaller than `Q4_0`, faster than `Q4_0` on + decode on both Vulkan and ROCm, and close to `Q4_0` PPL on the short + WikiText-2 check. +- `Q4_0_ROCMFP4_STRIX` is the current quality-biased Strix Halo preset. It + keeps most transformer tensors on `Q4_0_ROCMFP4_FAST`, protects token + embeddings with `Q6_K`, and uses the dual-scale `Q4_0_ROCMFP4` layout for + attention-K and attention-V tensors. On Qwen3-4B it improved the short + WikiText-2 PPL to `13.8865` at `4.49 BPW` while still beating the + same-flags stock `Q4_0` decode baselines on both Vulkan and ROCm. +- `Q4_0_ROCMFP4_STRIX_LEAN` is the compact Strix Halo preset. It keeps the + STRIX all-layer dual-scale attention-K/V protection, uses the FAST + single-scale transformer layout for the dense tensors, and protects token + embeddings/output with `Q5_K` instead of `Q6_K`. On the Qwen3-4B validation + pass it landed at `4.38 BPW`, improved short WikiText-2 PPL versus + `FAST_COHERENT`, and kept Vulkan decode in the `81 tok/s` band. +- A smaller first/last-layer-only K/V protection recipe was tested but not + promoted. It reached `4.48 BPW`, `80.13` Vulkan decode, and `69.85` ROCm + decode, but PPL regressed to `14.0167`, so the all-layer K/V STRIX preset + remains the quality target. +- The ROCm/HIP MMQ path for `Q4_0_ROCMFP4_FAST` uses one scale per 32-weight + block, matching the actual FAST layout instead of duplicating the scale into + two half-block slots. +- The ROCm/HIP vector-dot and MMQ loaders use a ROCmFP4-owned Codebook10 + expander backed by AMD `amdgcn_perm` constants. This avoids the generic + table-load helper on the hot ROCm path. +- The ROCm/HIP hot paths use a ROCmFP4-owned unaligned 32-bit quant-byte load + for packed nibble bytes. ROCmFP4 blocks are 17 or 18 bytes wide, so the older + byte-safe assembly path was conservative but expensive on Strix Halo HIP. + Direct unaligned dword loads improved ROCm `MUL_MAT` and FlashAttention + microbenchmarks while the Qwen3.6 27B MTP guard held `33.4 tok/s` short and + `27.7 tok/s` sustained. It is enabled by default and can be disabled with + `-DGGML_ROCMFP4_UNALIGNED_QS_DWORD_LOAD=0` for isolation. +- The ROCm/HIP FAST MMVQ/MMQ path now uses the same ROCmFP4-owned unaligned + quant-byte load instead of the generic byte-assembly helper. The full + promoted gate measured FAST ROCm `MUL_MAT` at `45.17`, `58.38`, `90.54`, + and `157.83` us for `n=1/2/4/8`, and the Qwen3.6 27B MTP guard improved to + `33.6 tok/s` short and `28.0 tok/s` sustained. The ROCm runtime guard now + tightens the FAST ceilings to protect this band. +- Qwen3.6 35B A3B MTP ROCmFP4 STRIX_LEAN was checked separately at native + `262144` context on ROCm0. In the reasoning-off discovery sweep, `n-max 1` + was best sustained at `72.2 tok/s`, while `n-max 5` was best burst at + `107.3 tok/s` but slower sustained at `64.9 tok/s`. With reasoning on, a + follow-up sweep found `n-max 2` best sustained at `92.6 tok/s` short and + `80.6 tok/s` sustained; `n-max 3` was close at `104.3` short and `80.1` + sustained, while `n-max 5` stayed burst-only at `98.6` short and `73.3` + sustained. KV-cache isolation then found that q8 main KV with q4 draft KV + improved the reasoning-on `n-max 2` profile to `93.7 tok/s` short and + `85.6 tok/s` sustained; full q8 main/draft KV reached `93.2` / `85.2`, + draft-only q8 stayed at `80.5` sustained, and K-only / V-only q8 main KV + regressed to `71.8` / `73.3` sustained. Rechecking draft depth under q8 main + KV moved the best sustained profile to `n-max 3` at `104.3 tok/s` short and + `89.3 tok/s` sustained, while `n-max 4` became the best burst-only profile + at `111.2 tok/s` short and `78.7 tok/s` sustained. The Pi sustained profile + therefore uses `n-max 3` with q8 main KV and q4 draft KV, while a separate + `n-max 4` burst profile is available for short-response experimentation. + The older `n-max 5` burst alias remains available for comparison. + Batch and CPU thread follow-ups did not beat the promoted runtime shape: + `-b 1024 -ub 512`, `-b 2048 -ub 512`, `-b 512 -ub 256`, + `-t 24 -tb 32`, and `-t 12 -tb 32` all measured in the `88.7`-`89.1 tok/s` + sustained band, so `-b 512 -ub 512 -t 16 -tb 32` stays promoted. + Current KV isolation rechecked the promoted shape at `104.3 tok/s` short and + `90.1 tok/s` sustained with q8 main KV, q4 draft KV, `n-max 3`, and + `p-min 0.25`; full q8 draft KV tied at `104.5` / `90.0`, while draft-only + q8 reached only `82.2` sustained. K-only and V-only q8 main KV regressed to + `70.5` and `74.6` sustained, so both accepted K and V need q8 and the draft + KV can stay q4. The updated default 35B guard then passed at `103.3 tok/s` + short and `90.1 tok/s` sustained. + Sampler-chain trims and backend sampling did not beat the promoted sustained + profile: `top_k;top_p;temperature` reached `104.2` / `87.5`, the milder + `penalties;top_k;top_p;min_p;temperature` chain reached `104.1` / `88.9`, + and `--backend-sampling` reached `104.3` / `89.2`, so the default sampler + path stays promoted. + A single-sequence MTP `process()` fast path for `-np 1` was also prototyped, + built, and rejected because the guard dropped to `104.1 tok/s` short and + `88.5 tok/s` sustained; the change was removed. + MoE `rows_per_block` compile-time variants were checked against the same 35B + guard and also rejected: `rows_per_block=4` measured `103.8` short / + `89.1 tok/s` sustained, `rows_per_block=3` measured `86.8 tok/s` sustained, + and `rows_per_block=1` measured `103.6` / `88.7 tok/s`, none beating the + promoted `104.3` / `89.3` band. + The Pi server profile was also started and stopped successfully with + `n_ctx = 262144`, `draft-mtp` initialized, built-in tools enabled, and + `thinking = 1`; ROCm reported no KFD PIDs after shutdown. The promoted + profile is now covered by the focused + `scripts/check-rocmfp4-qwen35-a3b-mtp-regression.sh` guard, which defaults + to `n-max 3`, q8 main KV, q4 draft KV, and `p-min 0.25` with `100.0` short + and `85.0 tok/s` sustained floors. + Follow-up `p-min 0.25/0.50/0.75`, `p-split 0.05/0.20`, and `n-min 1` + checks on the q4-KV baseline all tied the same `80.5`-`80.8 tok/s` sustained + band, while `p-min 0.25` on the q8-main/q4-draft path tied at `85.4 tok/s` + sustained, so the + conservative `n-min 0`, `p-min 0.0`, `p-split 0.10` defaults remained + promoted until the MTP internal sampler was changed from `top_k=1` to + `top_k=10`, making the `p-min` cutoff operate on a meaningful candidate + distribution while the draft loop still selects the top sorted candidate. + A pre-change p-min sweep on the actual promoted `n-max 3`, q8-main/q4-draft + profile measured `104.2` / `89.1` at `p-min 0.25`, `104.1` / `89.0` at + `0.50`, `104.0` / `89.3` at `0.75`, and `104.2` / `89.1` at `0.90`. + Follow-up p-split checks measured `103.7` / `88.9` at `0.05` and + `103.8` / `89.3` at `0.20`; n-min checks measured `103.8` / `88.9` for + `1` and `104.2` / `89.2` for `2`. None beat the promoted + sustained-plus-short profile. Combined `p-min`/`p-split` follow-ups also + tied rather than beat the default: `p-min 0.25` with `p-split 0.30` + repeated at `104.6` / `89.5`, `p-split 0.40` and `0.50` measured + `104.6` / `89.5` and `104.4` / `89.5`, and `p-split 0.70` / `0.90` + measured `104.8` / `89.5` and `104.7` / `89.5`; the same-session default + tied at `104.7` / `89.5`. With `top_k=10`, the 35B A3B profile now promotes + `--spec-draft-p-min 0.25`: default `p-min 0.0` measured `104.3` / `89.3`, + while `p-min 0.25` measured `103.9` / `90.0` after a first sustained run at + `89.8 tok/s`. The dense 27B profile should stay at `p-min 0.0`; the same + `p-min 0.25` filter regressed it to `24.6 tok/s` sustained. The full + all-regression harness can include the 35B guard with + `INCLUDE_QWEN35_A3B_GUARD=1`. + The MTP draft sampler now uses a probability-only top-10 helper for this + path. The draft loop only consumes the sorted top candidate and its + probability, so skipping the unused final RNG sampler selection preserved the + output path while improving the 35B A3B guard to `104.6 tok/s` short and + `90.2 tok/s` sustained. The dense 27B guard remained stable at `33.9` / + `28.1 tok/s`. + A fixed insertion top-10 replacement for `std::partial_sort` was tested and + rejected after the 35B A3B short guard dropped to `73.0 tok/s`, below the + `100.0 tok/s` floor. A `std::nth_element` plus top-k slice sort variant was + also rejected after the same guard dropped to `71.5 tok/s`. + Narrowing the MTP top-k probability accumulator from `double` to `float` was + also rejected: dense 27B still passed at `34.0` / `28.1 tok/s`, but the 35B + A3B short guard repeated below floor at `84.1` then `93.4 tok/s`; restoring + the `double` accumulator recovered the short guard to `104.4 tok/s`. + A `std::partial_sort_copy` top-k buffer variant built and passed the same + guard at `104.4 tok/s` short and `90.2 tok/s` sustained, but a same-session + promoted-build comparison measured `104.1` / `90.2`, so it was not promoted: + it added sampler buffer/clone complexity without a sustained decode gain. + Jackrong Qwopus3.6 27B v2 MTP BF16 was converted to STRIX_LEAN ROCmFP4 at + `4.34 BPW`. At native `262144` context with reasoning on, q4 main/draft KV, + and `n-max 4`, ROCm0 initially measured `34.9` / `29.6 tok/s`; increasing + only the batch to `-b 1024 -ub 512` moved sustained decode to `29.9 tok/s`. + Follow-up batch shapes `1280/512`, `1536/512`, and `1536/768` measured + `29.8`, `29.9`, and `29.8 tok/s` sustained respectively, so the smaller + `1024/512` profile remains promoted. Light acceptance filters `p-min 0.05` + and `0.10` also tied at `29.9 tok/s` sustained without beating the default. + Lower draft-depth checks with the same promoted batch rejected `n-max 1`, + `2`, and `3`, which measured only `19.9`, `26.6`, and `27.3 tok/s` + sustained. Qwopus therefore stays on `n-max 4`, unlike the 35B A3B profile + where `n-max 3` plus q8 main KV is best. + `--backend-sampling` tied sustained decode at `29.9 tok/s` but lowered prompt + throughput, so it is not promoted. Thread split checks at target/draft + `12/32` and `24/32` also tied `29.9 tok/s` sustained; the simpler default + `16/32` thread shape remains the recommended Qwopus profile. + KV isolation confirmed this is not like the 35B A3B q8-main profile: + draft-only q8 KV measured `35.1` / `29.8 tok/s`, and full q8 main plus q8 + draft KV measured `36.6` / `26.0 tok/s`. The full-q8 path improves burst + only and regresses sustained decode. Split accepted-KV checks also regressed: + q8 K only measured `35.7` / `22.3 tok/s`, and q8 V only measured `34.2` / + `24.7 tok/s`. Qwopus therefore keeps q4 main and q4 draft KV. + Vulkan0 measured `40.0` / `27.7 tok/s`, with Vulkan `n-max 3` and `n-max 5` + slower sustained; Vulkan `-b 1024 -ub 512` repeated the same `27.7 tok/s` + sustained band. q8 main KV regressed this model on both ROCm0 and Vulkan0. + A normal-path shortcut that only normalized `data[0].p` and filled the rest + of the top-k probabilities only for debug logging was rejected after the + 35B A3B short guard fell to `92.7 tok/s`, below the `100.0 tok/s` floor. + Disabling internal sampler timing with `sparams.no_perf = true` was also + tested and rejected after the same short guard dropped to `96.2 tok/s`. + Skipping the per-draft `common_sampler_reset()` in the MTP path was rejected + after the 35B A3B short guard dropped to `68.3 tok/s`; that reset remains + required to preserve the expected sampler/logit state for this helper. + A small MTP host-path cleanup now delays `llama_get_embeddings_pre_norm_ith()` + until the draft loop has confirmed that another draft token will be queued. + This avoids unused embedding-row pointer fetches on p-min rejects and final + `n-max` draft tokens. It passed the dense 27B guard at `34.0` / `28.1 tok/s`, + the 35B A3B guard at `104.4` / `90.1 tok/s`, and the Qwopus best ROCm0 + profile at `35.0` / `29.8 tok/s`. The default all-regression gate also + passed after this cleanup and ended with no KFD PIDs running. The 2026-05-25 + serial pass measured Qwen3.6 27B MTP at `33.9` / `27.9 tok/s`, ROCm runtime + FAST `45.66` / `57.81` / `88.27` / `155.05` us and dual-scale `49.16` / + `51.58` / `83.34` / `151.42` us for `n=1/2/4/8`, ROCm FlashAttention + `70.86` / `66.51` us for 64d dual-scale / FAST and `189.45` / `172.73` us + for Qwen-style 128d dual-scale / FAST, and ROCm CPY source-to-dual + `1106.89` / `1008.56` / `1006.60` us with source-to-FAST `1050.49` / + `958.98` / `950.50` us for F32/F16/BF16. + A second attempt to skip the final `n-max` `common_sampler_accept()` call was + rejected: the 35B A3B short check still reached `104.5 tok/s`, but sustained + decode fell to `81.1 tok/s`, below the `85.0 tok/s` floor. Reverting only + that sampler-accept change recovered the guard to `104.3` / `90.0 tok/s`. + Retesting the internal MTP sampler candidate count on the 35B A3B + reasoning-on profile rejected both directions around the promoted `top_k=10` + setting: `top_k=5` fell to `77.3 tok/s` sustained and `top_k=20` fell to + `69.6 tok/s` sustained. +- Reasoning-off checks on the final 35B q8-main/q4-draft profile are a + separate lower-throughput mode. `n-max 1/2/3/4` measured `77.7` / `73.9`, + `90.3` / `75.5`, `100.3` / `71.9`, and `85.7` / `66.1` short/sustained + tok/s respectively. If reasoning is disabled, `n-max 2` is currently the + best sustained profile in this bracket; the promoted fastest profile remains + reasoning-on `n-max 3`. +- Draft-thread-only checks on the promoted 35B profile also tied below the + promoted band. Keeping target threads at `16/32`, draft `8/16`, `16/16`, + and `24/32` all measured around `104.1`-`104.2 tok/s` short and + `89.1 tok/s` sustained, so the default matching draft thread counts remain. +- ROCm/HIP single-token MMVQ uses a full-block vector-dot ratio for the + dual-scale layout (`VDR_ROCMFP4_Q8_1_MMVQ=4`) while keeping the FAST layout + on the previous half-block ratio (`VDR_ROCMFP4_FAST_Q8_1_MMVQ=2`). This + lets dual-scale ROCmFP4 consume one full 32-value block per vector-dot call + without slowing the FAST dense-tensor path used by STRIX_LEAN. The focused + ROCm0 `MUL_MAT` guard improved dual-scale from the prior `78.81` us/run + serial pass to `54.89` us/run, and the Qwen3.6 27B MTP sustained guard + improved from `26.2` to `27.8 tok/s`. A broader version that also moved + FAST to `vdr=4` was rejected because sustained Qwen MTP dropped to + `24.2 tok/s`. A later FAST-only retest after the packed-byte load improved + the focused FAST ROCm guard to `41.37`, `49.29`, `80.91`, and `139.58` + us for `n=1/2/4/8`, but was still rejected because sustained Qwen MTP + dropped to `24.7 tok/s`. The remaining narrower FAST setting, + `GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR=1`, was also rejected after it failed the + focused FAST ROCm `n=1` guard at `60.18` us/run. The knob now rejects + invalid FAST MMVQ values at compile time; only `1`, `2`, and `4` are valid. +- A direct ROCmFP4 `vec_dot_q_cuda_dispatch` wrapper was tested in the + MMVQ kernels to bypass the generic constexpr function pointer call. It built + and passed the focused ROCm guard, measuring FAST `45.16` / `57.52` / + `89.44` / `156.49` us and dual-scale `50.83` / `51.11` / `84.51` / + `143.27` us for `n=1/2/4/8`, but Qwen3.6 27B MTP only reached `33.7 tok/s` + short and `27.9 tok/s` sustained. Because it did not beat the promoted + sustained band, the code change was removed. +- ROCm/HIP batched MMQ keeps the upstream-style `vdr=8` default for both + ROCmFP4 layouts, with ROCmFP4-owned compile-time test knobs + `GGML_ROCMFP4_Q8_1_MMQ_VDR` and `GGML_ROCMFP4_FAST_Q8_1_MMQ_VDR`. + FAST-only `vdr=4`, FAST-only `vdr=16`, and dual-scale `vdr=16` were tested + against the Qwen3.6 27B STRIX_LEAN ROCmFP4 bench and did not improve decode; + the tested runs all stayed at `13.56 tok/s` generation, while dual-scale + `vdr=16` slightly reduced prompt throughput. The default remains `8`. +- ROCm/HIP RDNA3.5 small-batch MMVQ now uses a ROCmFP4-specific two-warp + launch geometry for the ROCmFP4 layouts through `n=2`. Strix Halo previously + inherited the older RDNA2 one-warp table for this path. The promoted + `GGML_ROCMFP4_RDNA35_NWARPS=2` and + `GGML_ROCMFP4_RDNA35_NWARPS_MAX_NCOLS=2` defaults keep the single-token + microbench win and improve the guarded ROCm0 `n=2` shape: the latest serial + pass measured FAST/dual `66.56` / `58.40` us/run for `n=2`, versus the prior + `68.85` / `60.98` us/run band before extending the two-warp route. The + Qwen3.6 27B MTP guard held `33.5 tok/s` short and `27.7 tok/s` sustained, so + this is promoted as a backend micro-optimization with no sustained decode + regression. `GGML_ROCMFP4_RDNA35_NWARPS=4` was rejected because dual-scale + regressed to `57.87` us/run. The remaining 8-warp candidate was also rejected + because FAST `n=1` regressed to `59.33` us/run and failed the focused ROCm + runtime guard before multi-column checks. Extending the promoted two-warp + launch from `n=1` to `n=1..4` improved some focused multi-column ROCm + microbench rows (`dual n=2` reached `57.14` us/run), but Qwen3.6 27B MTP + sustained decode fell to `23.6 tok/s`, so the promoted upper bound stops at + `n=2`. A midpoint `GGML_ROCMFP4_RDNA35_NWARPS_MAX_NCOLS=3` build was also + checked on the Qwen3.6 35B A3B reasoning-on 262k profile; it reached + `87.5 tok/s` sustained versus `89.6 tok/s` for the same-session promoted + build, so `n=2` remains the default. +- ROCm/HIP RDNA3.5 wide-column rows-per-block now has ROCmFP4-owned compile-time + test knobs: `GGML_ROCMFP4_RDNA35_RPB_WIDE`, + `GGML_ROCMFP4_RDNA35_RPB_WIDE_DUAL`, and + `GGML_ROCMFP4_RDNA35_RPB_WIDE_FAST`. The defaults stay at `1`. A full + `RPB_WIDE=2` test improved FAST `n=8` `MUL_MAT` from `167.90` to + `131.72` us/run, but dual-scale collapsed from `148.11` to `1382.72` + us/run. A FAST-only `RPB_WIDE_FAST=2` build kept dual-scale safe + (`145.75` us/run) and improved FAST to `135.31` us/run, but Qwen3.6 + 27B MTP did not improve (`33.3 tok/s` short and `27.6 tok/s` sustained + at `n-max 4`; `n-max 5` remained burst-only at `45.5` short and `24.9` + sustained). Keep the knob off by default until a real decode guard benefits. +- ROCm/HIP RDNA3.5 `MUL_MAT_ID` routing now has a ROCmFP4-only compile-time + guard knob, `GGML_ROCMFP4_RDNA35_MMID_MAX_BATCH`, for testing whether MTP/MoE + batches should leave MMVQ earlier. The default keeps the accepted generic + RDNA3 behavior (`MMVQ_MAX_BATCH_SIZE`). A Strix test with the knob set to `1` + tied the promoted Qwen3.6 27B MTP path at `33.4 tok/s` short and `27.7 tok/s` + sustained. Follow-up threshold `2`, `3`, and `4` builds also tied at + `33.6` / `33.8` / `33.7 tok/s` short and `27.7` / `27.9` / `27.9 tok/s` + sustained. The same `3` threshold was then checked on the MoE-heavy + Qwen3.6 35B A3B MTP ROCmFP4 path and regressed to `95.8 tok/s` short and + `74.0 tok/s` sustained versus the promoted `104.4` / `89.3` band. Reusing + the older threshold `4` exploratory build on the same 35B guard measured + `104.2` / `89.2`, a near tie but still below the sustained promoted band. + No lower routing threshold is promoted. +- Target/draft "dual-stream" MTP overlap was inspected and is not a promoted + optimization. The current common MTP implementation verifies with the target + context, mirrors pre-norm embeddings into the MTP context, then drafts through + serial `llama_decode(ctx_dft, ...)` calls. ggml exposes async graph execution + and pipeline parallelism, but llama enables pipeline parallelism only for + multi-device layer-split offload cases. On single-device Strix Halo ROCm0, + target/draft overlap would require speculative scheduler changes and new + correctness guards rather than a runtime flag. The long-context optimization + focus remains KV-cache traffic, FlashAttention, memory bandwidth, and MTP + acceptance, while dual-scale ROCmFP4 mainly protects coherence. +- The MTP host loop now avoids two small sources of scheduler overhead. It + reuses a per-sequence drafting-state buffer instead of allocating a + `std::vector` on every draft call, and it tracks verify-batch sequence + bounds with one pass over token seq-ids instead of scanning every sequence for + every token. It also avoids copying all target verify rows when no previous + draft is pending, keeping only the carryover hidden row needed for the next + MTP step. The MTP path also uses a single-sequence batch append helper + instead of constructing a temporary sequence vector for each `common_batch_add` + call. It stores only non-final verify rows for partial-accept rollback, + copying the final row directly to `pending_h` for the next MTP step. It also + skips the draft-model decode call entirely when no sequence is drafting. + Debug candidate token formatting is skipped unless debug logging is enabled, + and the debug verbosity state is hoisted once per draft call instead of + rechecked inside the per-token loop. This is mainly a cleanup for + multi-sequence and long-running MTP sessions; the latest single-sequence + Qwen3.6 27B guard stayed in range at `33.7 tok/s` short and `27.9 tok/s` + sustained. +- The simple draft path also reuses an object-owned `uint8_t` drafting-state + buffer instead of allocating `std::vector` on every draft call. This + keeps both draft implementations on the same host-side allocation pattern; + the full serial gate stayed clean, with Qwen3.6 27B MTP at `33.8 tok/s` + short and `27.9 tok/s` sustained. +- The all-in-one regression harness now accepts a candidate `BUILD_DIR` and + derives separate `TEST_BACKEND_OPS_BIN` and `LLAMA_CLI_BIN` paths, so backend + microbench guards and Qwen CLI guards can run from the same candidate build. + DeepSeek remains opt-in only through `INCLUDE_DEEPSEEK_SMOKE=1`. +- ROCm/HIP vector FlashAttention now uses a ROCmFP4-only RDNA K/Q thread-group + default of `1` instead of the generic quantized-K RDNA default of `2`. This + affects only `Q4_0_ROCMFP4` and `Q4_0_ROCMFP4_FAST` K-cache cases; other + quantized FA paths keep the existing backend default. On the guarded ROCm0 + `FLASH_ATTN_EXT` shape, dual-scale improved from `122.33` to `113.03` + us/run, and FAST improved from `115.41` to `109.33` us/run. The Qwen3.6 + 27B MTP guard remained in the promoted band at `33.5 tok/s` short and + `27.6 tok/s` sustained. A wider `GGML_ROCMFP4_FATTN_KQ_NTHREADS=4` variant + was rejected because it regressed the same FA guard to `136.79` us/run + dual-scale and `124.37` us/run FAST. After the V-side default moved to `2`, + retesting `KQ_NTHREADS=2,V_NTHREADS=2` and `KQ_NTHREADS=4,V_NTHREADS=2` + still regressed the focused FA guard, so K/Q remains `1`. +- ROCm/HIP vector FlashAttention also uses a ROCmFP4-only V thread-group + default of `2`, down from the generic `D/4` path used by other quantized V + types. On the same guarded ROCm0 `FLASH_ATTN_EXT` shape, this moved + dual-scale from `113.03` to `85.63` us/run and FAST from `109.33` to + `80.74` us/run. The Qwen3.6 27B MTP guard stayed in range at `33.4 tok/s` + short and `27.6 tok/s` sustained, so this is promoted as the new default. + `GGML_ROCMFP4_FATTN_V_NTHREADS=4` also passed but was slower on the focused + FA guard; `1` was rejected because gfx1151 HIP compilation exceeded the 64 + KiB local-memory limit for ROCmFP4 FA instances. The FA guard now also + includes a Qwen3.6-style 128-head-dim, 8-KV-head, 12x-GQA ROCmFP4 shape + (`hsk=128,hsv=128,nh=8,nr23=[12,1],kv=7680,nb=1`). The accepted default + measured `246.91` us dual-scale and `219.23` us FAST on that shape. + Retesting `V_NTHREADS=4` on this wider shape produced a tiny dual-scale + improvement (`244.01` us) but regressed FAST (`224.21` us) and the existing + 64d guard, so it remains rejected. Retesting `KQ_NTHREADS=2,V_NTHREADS=2` + regressed the Qwen-style shape to `268.53` us dual-scale and `226.32` us + FAST, so K/Q remains `1`. +- ROCm/HIP vector FlashAttention now uses a ROCmFP4-owned single-half + Codebook10 expander in K/Q and V decode paths. These FA call sites already + know whether they need the low or high nibble stream, so they no longer pay + to expand both streams and discard one. The focused ROCm0 FA guard improved + from `86.13` / `81.23` us to `82.24` / `78.15` us for the 64d dual-scale / + FAST shapes, and from `247.06` / `221.19` us to `237.12` / `206.50` us for + the Qwen-style 128d shapes. The Qwen3.6 27B MTP guard held at `33.5 tok/s` + short and `27.7 tok/s` sustained in the full serial gate, so this is + promoted. +- ROCm/HIP vector FlashAttention also specializes the ROCmFP4 K/Q path when + `GGML_ROCMFP4_FATTN_KQ_NTHREADS=1`. Each thread owns the full head dot in + this promoted setting, so the kernel now expands both low/high packed + Codebook10 streams once per ROCmFP4 block and accumulates both half-blocks + together instead of loading the same packed bytes twice. The full serial gate + measured 64d FlashAttention at `81.62` / `78.13` us for dual-scale / FAST, + and the wider Qwen-style 128d guard improved from `237.12` / `206.50` us to + `228.58` / `199.32` us. Qwen3.6 27B MTP held `33.2 tok/s` short and + `27.7 tok/s` sustained, so this is promoted as a real Qwen-relevant FA + micro-optimization. + A post-specialization retest of `GGML_ROCMFP4_FATTN_V_NTHREADS=4` was + rejected. It passed the tightened FA guard and nudged the Qwen-style + dual-scale microbench to `228.33` us, but regressed 64d dual-scale to + `86.30` us and Qwen-style FAST to `201.81` us. Qwen3.6 27B MTP stayed at + `27.7 tok/s` sustained, so the promoted V-side default remains `2`. +- ROCm/HIP vector FlashAttention now uses a ROCmFP4-only V dequant rows-per- + thread default of `8`. The ROCmFP4 V helper supports 8-value chunks, and the + post-unaligned-load retest produced a real FA win: 64d dual-scale / FAST + `68.82` / `66.32` us and Qwen-style 128d dual-scale / FAST `201.13` / + `172.22` us. Qwen3.6 27B MTP held `33.5 tok/s` short and `27.7 tok/s` + sustained, so this is now promoted. It can be isolated with + `-DGGML_ROCMFP4_FATTN_V_ROWS_PER_THREAD=4` to return to the prior setting. + Retesting `GGML_ROCMFP4_FATTN_V_NTHREADS=4` on top of the 8-row default + improved only the Qwen-style dual-scale FA microbench (`187.72` us) while + regressing 64d dual-scale / FAST to `77.35` / `77.25` us and Qwen-style FAST + to `182.56` us. Qwen3.6 27B MTP tied at `33.5` / `27.7 tok/s`, so the + shipped V thread-group default remains `2`. + Retesting `GGML_ROCMFP4_FATTN_V_NTHREADS=8` on top of the 8-row default + regressed every focused FA row, including Qwen-style 128d dual-scale / FAST + at `208.94` / `200.45` us, so it was rejected without a longer MTP run. + Retesting `GGML_ROCMFP4_FATTN_KQ_NTHREADS=2` on top of the 8-row default + also regressed the focused FA guard: 64d dual-scale / FAST measured + `80.68` / `78.27` us and Qwen-style 128d measured `245.84` / `207.04` us. + Because the Qwen-style dual-scale row failed the guard, it was rejected + without a longer MTP run and K/Q remains `1`. +- ROCm/HIP vector FlashAttention now has a narrow dual-scale 128d V-thread + specialization: `GGML_ROCMFP4_FATTN_V_NTHREADS_D128_DUAL=4`. This keeps the + promoted 64d default and all FAST paths on `V_NTHREADS=2`, but uses the + previously promising 4-thread V grouping only for the Qwen-style dual-scale + 128d shape. The full serial gate measured FA at `70.37` / `66.31` us for + 64d dual-scale / FAST and `188.75` / `171.89` us for Qwen-style 128d + dual-scale / FAST, while Qwen3.6 27B MTP held `33.7 tok/s` short and + `27.9 tok/s` sustained with no KFD PIDs left running. + Retesting `GGML_ROCMFP4_FATTN_V_NTHREADS_D128_DUAL=8` was rejected. It + slightly improved the 64d dual-scale row (`69.40` us versus a same-session + `70.92` us promoted build), but regressed the Qwen-style 128d dual-scale row + from `194.40` to `211.10` us/run. Because long-context Qwen is the relevant + guard, the promoted D128 dual-scale V grouping remains `4`. + `GGML_ROCMFP4_FATTN_V_ROWS_PER_THREAD=16` was rejected during compilation: + the ROCmFP4 V dequantizer intentionally supports only `2`, `4`, and `8` + rows per thread, and the fixed-copy helper does not support the resulting + 32-byte move. The source now rejects unsupported values with a direct + ROCmFP4 compile-time error. + Retesting `GGML_ROCMFP4_FATTN_V_ROWS_PER_THREAD=2` on the promoted + D128-specialized build was rejected by the focused ROCm FlashAttention + guard. It regressed 64d dual-scale / FAST to `94.31` / `88.77` us and + Qwen-style 128d dual-scale / FAST to `270.00` / `223.88` us, so the + promoted rows-per-thread default remains `8`. +- Vulkan ROCmFP4 scale decode now uses a shared UE4M3 lookup table with + ROCmFP4's half-scale semantics. This moved the focused Vulkan dual-scale + `MUL_MAT` guard from `82.86`, `120.77`, and `181.28` us/run to `65.05`, + `83.07`, and `122.70` us/run for `n=1/2/4`. On Qwen3.6 27B MTP at 262k + context, the same change moved Vulkan sustained output from the older + `20.4 tok/s` at `--spec-draft-n-max 4` to `25.0 tok/s`; `n-max 3` now + reaches `25.3 tok/s` sustained. ROCm0 remains the promoted backend because + it still holds `33.5 tok/s` short and `27.7 tok/s` sustained. + A post-LUT Vulkan runtime sweep found a better fallback profile with q8 KV + and q8 draft KV: `--spec-draft-n-max 4` reached `34.8 tok/s` short and + `27.0 tok/s` sustained. q8 with `n-max 5` improved only the short burst + (`47.8 tok/s`) and regressed sustained output to `23.0 tok/s`; f16 KV + regressed sustained output to `22.5 tok/s`. + Follow-up isolation showed the Vulkan sustained gain comes from the main KV + cache rather than the draft KV cache: q4 main KV with q8 draft KV reached + only `34.6 tok/s` short and `25.0 tok/s` sustained, while q8 main KV with + q4 draft KV reached `34.7 tok/s` short and `26.9 tok/s` sustained. For a + Vulkan fallback, q8 main KV plus q4 draft KV is therefore the leaner near-tie + profile. Splitting q8 across only one accepted KV side was also rejected: + q8 main K with q4 main V reached `34.8 tok/s` short and `25.4 tok/s` + sustained, while q4 main K with q8 main V reached `34.6 tok/s` short and + `23.7 tok/s` sustained. The Vulkan fallback needs both accepted K and V at + q8. Adding mild acceptance filtering to the lean fallback + (`--spec-draft-p-min 0.25`) tied full q8/q8 at `27.0 tok/s` sustained while + keeping draft KV at q4; adding `--spec-draft-n-min 1`, stricter `p-min 0.75`, + and `p-split` checks stayed at `26.9 tok/s`. Retrying `n-max 5` with q8 main + KV, q4 draft KV, and + `p-min 0.25` still produced only a burst win (`47.8 tok/s` short) while + sustained output stayed at `23.0 tok/s`, so `n-max 4` remains the Vulkan + fallback setting. ROCm0 q4/q4 remains the overall promoted backend/profile. + Replacing the shared Codebook10 table with inline integer decode was tested + and rejected because it regressed the focused Vulkan FAST `n=1` guard to + `99.43` us/run. +- ROCm/HIP fallback dequant, copy, get-rows, GPU-side quantization scoring, + and standalone dequant helpers use ROCmFP4-owned HIP helpers for finite + scales and Codebook10 nibbles instead of relying on generic FP8 handling. + This keeps non-MMQ conversion paths aligned with the custom AMD format. +- A branchless HIP scalar Codebook10 nibble decoder was tested and rejected. + It passed the focused ROCm CPY, FlashAttention, and Qwen MTP guards, but + did not improve end-to-end decode. Qwen3.6 27B MTP measured `33.4 tok/s` + short and `27.7 tok/s` sustained, and the focused FlashAttention guard + slowed to `86.38` / `81.13` us for dual-scale / FAST, so the original + scalar decoder remains the default. +- CPU-side Codebook10 table decode in the quantizer/dequantizer was tested + and rejected. The full table variant failed the CPU quant guard with + dequantization at `49.04` cycles/32 for dual-scale and `84.23` cycles/32 for + FAST. The MSE-loop-only variant passed but slowed normal quantization to + `4183.33` / `4018.68` cycles/32 versus the restored arithmetic baseline at + `4034.33` / `3738.85` cycles/32, so the arithmetic CPU decoder remains the + default. +- UE4M3 scale decode in the ROCm/HIP software path uses a ROCmFP4-owned finite + scale decoder. It avoids `ldexpf`, builds normal FP32 values directly from + exponent/mantissa bits, and skips the generic FP8 NaN handling because + ROCmFP4 row validation already rejects non-finite scale bytes. +- A constant-memory HIP lookup table for the finite UE4M3 scale values was + tested and rejected. The latest isolated `GGML_ROCMFP4_USE_SCALE_LUT=1` + pass failed the focused ROCm runtime guard because FAST `n=1` regressed to + `69.48` us/run; prior checks also regressed ROCm CPY source-to-quant paths + and FlashAttention. The arithmetic finite-scale decoder remains the default. +- ROCm/HIP dequant conversion kernels use the same ROCmFP4 finite scale + decoder, keeping tensor conversion aligned with the hot MMQ/MMVQ backend + path instead of falling back to the generic FP8 helper. +- ROCm/HIP backend CPY now advertises and executes quantized + `Q4_0_ROCMFP4 -> F32` and `Q4_0_ROCMFP4_FAST -> F32` conversion paths. + This keeps diagnostic graph ops and fallback tensor conversion inside the + custom AMD decoder instead of being rejected by backend capability checks. + Contiguous q-to-f32 copies now use a ROCmFP4-specific packed-byte kernel: + one HIP thread reads one packed byte and writes the matching low/high + half-block output values. This gives coalesced output writes for normal + contiguous graph copies while the existing stride-aware block kernel remains + the fallback for views and non-contiguous tensors. +- ROCm/HIP backend CPY also supports `F16 -> Q4_0_ROCMFP4`, + `F16 -> Q4_0_ROCMFP4_FAST`, `BF16 -> Q4_0_ROCMFP4`, and + `BF16 -> Q4_0_ROCMFP4_FAST`. The kernels convert each 32-value half/bfloat + block to local FP32 and then run the same exhaustive ROCmFP4 scale search, + so runtime graph copies keep the coherence-first quantizer instead of + falling back to unsupported behavior. +- ROCm/HIP ROCmFP4 source-to-quant and quant-to-F32 CPY wrappers now use + normal multi-thread HIP launch geometry instead of launching one active + thread per quant block. This keeps F32/F16/BF16 runtime quantization and + ROCmFP4 dequantization on the same exact conversion math while removing a + launch-shape bottleneck in helper/fallback graph paths. +- ROCm/HIP ROCmFP4 quant-to-F32 CPY uses ROCmFP4-specific block dequant + helpers that decode each block scale once before unpacking all 32 values. + This avoids repeatedly decoding identical scale bytes through the generic + two-value dequant helper while preserving the same output values. For the + guarded contiguous `8192x512x2` ROCm0 CPY shape, the packed-byte path moved + dual-scale `Q4_0_ROCMFP4 -> F32` from the old `740` us/run band to + `184.48` us/run and FAST `Q4_0_ROCMFP4_FAST -> F32` to `169.99` us/run. + A split launch geometry is now used: source-to-ROCmFP4 quantization uses + 128-thread workgroups, while quant-to-F32 keeps the accepted 64-thread + packed-byte launch. A whole-path 128-thread launch was roughly tied overall + but slightly regressed quant-to-F32; isolating 128 threads to source + quantization preserves the dequant copy win while shaving the source paths. + A 256-thread launch regressed F16 source-to-quant paths and remains rejected. + After the FAST direct-value scoring win, FAST-only 256-thread and 64-thread + source launch splits were rechecked and rejected as well. The 256-thread + split regressed FAST F16 source-to-quant to `1047.24` us/run, and the + 64-thread split measured FAST F32/F16/BF16 `1055.32`, `955.35`, and + `954.78` us/run, so the shared 128-thread source launch remains promoted for + both ROCmFP4 layouts. +- ROCm/HIP contiguous quant-to-F32 CPY shared-scale staging was tested and + rejected. Decoding one block scale into shared memory and synchronizing each + 64-thread launch reduced duplicate scale decode work, but the synchronization + overhead regressed the guarded `8192x512x2` shape: dual-scale + `Q4_0_ROCMFP4 -> F32` moved from `181.84` to `188.63` us/run, and FAST + `Q4_0_ROCMFP4_FAST -> F32` moved from `170.65` to `180.03` us/run. The + direct packed-byte kernel remains promoted. +- ROCm/HIP backend CPY supports same-type packed-block copies for + `Q4_0_ROCMFP4 -> Q4_0_ROCMFP4` and + `Q4_0_ROCMFP4_FAST -> Q4_0_ROCMFP4_FAST`, including block-aligned views. + The kernel copies the packed 18-byte or 17-byte ROCmFP4 blocks directly, so + graph/view copies preserve exact bytes and avoid dequantize/requantize + fallback behavior. The launcher uses normal multi-thread HIP workgroups + rather than one-thread launches, so large packed-view copies scale with the + number of ROCmFP4 blocks. +- ROCm/HIP runtime quantization now finds the finite UE4M3 candidate nearest + `max_abs / 10` with a monotonic binary search, matching the CPU/Vulkan + reference tie behavior while avoiding the older 126-step linear nearest-scale + scan in the HIP copy utility. It also uses the same conservative lower-scale + pruning as the CPU quantizer, so runtime F32/F16/BF16-to-ROCmFP4 copies skip + smaller scales once their unavoidable clipped max-value error cannot beat the + current best exact scale. The HIP max scan uses the same plain absolute + compare as the CPU quantizer, preserving NaN/Inf handling while avoiding + `fmaxf` in the runtime quantization hot loop. +- ROCm/HIP runtime quantization specializes the exact scale search for the + only block shapes used by ROCmFP4 CPY: `0..15`, `16..31`, and `0..31`. + This keeps the same candidate order and tie behavior while letting HIP + compile the dual-scale and FAST source-to-quant paths with fixed offsets. +- ROCm/HIP dual-scale source-to-quant scoring uses a direct Codebook10 value + helper for 16-value half blocks. Final packed quantization still uses the + normal code-index helper, but scale scoring no longer has to map the selected + nibble back through the generic decode helper on the dual-scale path. +- ROCm/HIP FAST source-to-quant scoring now uses the same direct Codebook10 + value helper during scale search. This preserves the exact nearest-value + thresholds and final packed output, but avoids index-then-decode work in the + 32-value FAST scoring path. On the guarded `8192x512x2` ROCm0 CPY shape, + FAST source quantization improved from F32/F16/BF16 `1218.69`, `1138.32`, + and `1138.78` us/run to a repeat guard pass of `1047.21`, `950.93`, and + `951.00` us/run. +- ROCm/HIP `F32 -> ROCmFP4` runtime CPY stages each 32-value block into local + FP32 before running the exhaustive scale search, matching the existing + F16/BF16 source paths. This preserves exact F32 input values but avoids + rereading the global source pointer throughout scale scoring. On the guarded + `8192x512x2` ROCm0 CPY shape, dual-scale F32 source quantization improved + from roughly `9916` to `1117` us/run, and FAST improved from roughly + `10671` to `1231` us/run. +- ROCm/HIP contiguous-only source-to-quant CPY kernels for + `F32 -> ROCmFP4` were tested and rejected. They avoided the generic + multidimensional offset math, but on the guarded `8192x512x2` shape they + did not beat the existing multi-thread generic CPY path and slightly + regressed `F32 -> Q4_0_ROCMFP4_FAST`, so source-to-quant ROCmFP4 CPY keeps + the same guarded implementation for view and contiguous tensors. +- ROCm/HIP backend GET_ROWS supports both ROCmFP4 layouts. This gives pure + ROCmFP4 tensors the same direct row-gather coverage as stock small-block + quants on ROCm and keeps embedding-row access on the custom finite-scale + decoder. +- ROCm/HIP `MUL_MAT` support now covers `F16` activation tensors for both + ROCmFP4 layouts. The backend stages half activations to contiguous FP32 on + the GPU, including non-contiguous/views, then feeds the existing Q8 + activation quantizer and ROCmFP4 MMVQ/MMQ kernels. This keeps the forward + path on the AMD backend instead of rejecting the op and falling through to a + slower dequantized matrix path. The generic matmul runtime guard explicitly + allows this ROCmFP4 x F16 forward-inference case, so the support probe and + execution wrapper agree for batched activations. +- The standalone HIP dequant skeleton covers both the dual-scale and FAST + single-scale layouts, so future fused ROCm kernels can target the current + balanced FAST artifact without reintroducing the older scale path. +- The standalone HIP dequant launch now maps one thread to one packed + ROCmFP4 byte across a 256-thread global grid. This avoids the older + one-16-thread-block-per-32-values launch shape and gives future fused or + diagnostic dequant paths normal GPU occupancy scaling. +- Vulkan ROCmFP4 shaders also decode UE4M3 scales directly to the half-scale + value used by the codebook, matching CPU/HIP and avoiding repeated `* 0.5` + fixups at dequant and matmul call sites. +- Vulkan ROCmFP4 shaders keep a shared `kvalues_rocmfp4` Codebook10 table. + Arithmetic/direct Codebook10 decode variants compiled and preserved + coherence, but measured slower on Strix Halo Vulkan, so the table path + remains the active backend implementation. +- Vulkan `dequantize4()` for ROCmFP4 and ROCmFP4_FAST decodes each block's + UE4M3 scale once per 4-value vector instead of calling the 2-value helper + twice. This improves tested dual-scale ROCmFP4 small-batch matvec shapes + while keeping the promoted Vulkan runtime guard clean. +- Vulkan `Q4_0_ROCMFP4_FAST` matvec/MMQ kernels have a single-scale dot + specialization. They combine the two half-block dot sums and apply the one + FAST scale once, instead of taking the dual-scale path used by + `Q4_0_ROCMFP4`. +- Vulkan `Q4_0_ROCMFP4_FAST` MMQ stores its block scale as a scalar in + shared/register cache instead of duplicating it into a `vec2`. The focused + Strix Halo Vulkan `MUL_MAT` microbench for + `m=4096,n=1,k=14336,type_a=q4_0_rocmfp4_fast` improved from `62.18` to + `61.45` us/run. +- A Vulkan packed16 view for ROCmFP4 quant bytes was tested and rejected. It + compiled, but the focused FAST MMQ guard regressed to `83.03` us/run from the + accepted `61`-`62` us/run range, so ROCmFP4 Vulkan keeps the byte-view load + path. +- Vulkan backend CPY/SET_ROWS now has generated `F32 -> Q4_0_ROCMFP4`, + `F32 -> Q4_0_ROCMFP4_FAST`, `Q4_0_ROCMFP4 -> F32`, + `Q4_0_ROCMFP4_FAST -> F32`, and indexed SET_ROWS shaders. The SET_ROWS path + uses the same exhaustive finite UE4M3 scale search as the CPU reference, so + quantized K/V cache writes favor coherence over a cheap max-abs shortcut. +- Vulkan `F32 -> ROCmFP4` runtime quantization now uses the same exact ordered + UE4M3 scale search as CPU/HIP: find the scale nearest `max_abs / 10`, expand + outward, and stop a candidate once its partial error cannot beat the current + best scale. The candidate set is unchanged, so this avoids a slower linear + scan without falling back to a lower-quality shortcut. +- Vulkan source-to-ROCmFP4 runtime quantization now also prunes lower scale + candidates once clipping the block max alone cannot beat the current best + error. This mirrors the CPU/HIP scale search bound and keeps the candidate + set exact. The 2026-05-25 Vulkan CPY guard passed after this shader change: + F32/F16/BF16-to-dual measured `9525.39`, `2350.54`, and `2418.09` us/run, + dual-to-F32 `516.65` us/run, F32/F16/BF16-to-FAST `10111.85`, `2923.67`, + and `2949.42` us/run, and FAST-to-F32 `509.65` us/run. The full promoted + gate passed after rebuild with Qwen3.6 27B MTP at `33.9` / `28.0 tok/s` and + no KFD PIDs left running. +- Vulkan backend CPY also supports `F16 -> Q4_0_ROCMFP4`, + `F16 -> Q4_0_ROCMFP4_FAST`, `BF16 -> Q4_0_ROCMFP4`, and + `BF16 -> Q4_0_ROCMFP4_FAST`. The runtime quantization shader can load + `float`, `float16_t`, and BF16 source bits and then runs the same exact + ordered finite UE4M3 scale search. Backend tests keep a bounded NMSE + tolerance only for the half/bfloat runtime quantization cases because those + paths are inherently lossy around source-precision tie points; the + `F32 -> ROCmFP4` checks remain strict. +- Vulkan same-type CPY supports packed-block copies for + `Q4_0_ROCMFP4 -> Q4_0_ROCMFP4` and + `Q4_0_ROCMFP4_FAST -> Q4_0_ROCMFP4_FAST`, including block-aligned + non-contiguous/permuted/view copies. The non-contiguous path uses a + byte-addressed block shader and preserves exact 18-byte or 17-byte + ROCmFP4 blocks. Contiguous same-type dual-scale and FAST copies now both use + the direct byte-copy path, avoiding the generic halfword copy route for + ROCmFP4's custom 18-byte and 17-byte block layouts. +- Vulkan CPY is now covered by a dedicated regression guard so copy-path + changes cannot silently fall back or regress outside the ROCm-only CPY gate. + The guard runs same-type and source/dequant CPY correctness before measuring + the large copy performance shape, and now streams the performance phase + through `tee` so long Vulkan runs do not look idle to the command runner. + The ceilings were tightened after the lower-scale pruning gain, so the old + slow source-to-quant path no longer passes this guard. + On Strix Halo RADV Vulkan, the large guarded shape currently measures + F32/F16/BF16-to-dual at `9525.39`, `2350.54`, and `2418.09` us/run, + dual-to-F32 at `516.65` us/run, F32/F16/BF16-to-FAST at `10111.85`, + `2923.67`, and `2949.42` us/run, and FAST-to-F32 at `509.65` us/run. +- Vulkan scalar FlashAttention can now decode ROCmFP4 and ROCmFP4_FAST K/V + cache blocks. ROCmFP4 K/V is forced to the scalar FA path because the current + custom decode is not a coopmat/native matrix-core FP4 path. +- Vulkan scalar FlashAttention can use the integer-dot MMQ K path for both + ROCmFP4 K-cache layouts. The FAST layout expands each 4-bit Codebook10 value + into packed signed int8 lanes and uses its single UE4M3 scale as the K block + multiplier. The dual-scale layout also uses packed signed int8 lanes, but + splits the dot accumulation by half-block so each 16-value half uses its own + UE4M3 scale. This keeps the quality-biased STRIX K-cache path fast without + applying one scale to a two-scale block. +- Build and runtime verification generated the Vulkan SPIR-V entries for + ROCmFP4 copy/SET_ROWS shaders, linked `libggml-vulkan.so`, and passed + Vulkan ROCmFP4 CPY plus MUL_MAT smoke tests on Strix Halo. +- Row validation rejects invalid scale bytes outside finite unsigned UE4M3 + (`0x00` through `0x7e`) so corrupted custom GGUF tensors fail early. +- Quantization keeps the exhaustive 126-scale UE4M3 search for both normal and + imatrix paths. Candidate-window scale search was tested and improved GGUF + creation speed, but it regressed the Qwen3-4B short WikiText-2 PPL on the + compact FAST path, so it was rejected for coherence. +- The exhaustive scale search now visits the UE4M3 candidate nearest the + block's `max_abs / 10` first using a monotonic binary search, expands + outward, and exits a candidate scale once its partial error cannot beat the + current best scale. This remains exact because every finite scale is still + evaluated; on the Qwen3-4B FAST_COHERENT artifact it produced a + byte-identical GGUF while cutting FAST quantization cost sharply. +- CPU scale decode now uses a 127-entry finite UE4M3 half-scale table for the + exhaustive scale search and row dequantization. The table preserves the same + exact FP32 values as the former bit-construction decoder, but avoids + rebuilding them for every scale candidate. +- CPU scale search now has separate unweighted and imatrix-weighted MSE + scoring paths. The normal path avoids per-value weight branches, while the + imatrix path precomputes the row-energy calibration weight once per + 16/32-value block instead of recomputing the same `sqrtf` term for every + candidate scale. +- CPU normal quantization scans each block once for finite values and uses a + finite-only nearest-Codebook10 helper in the exhaustive MSE loop when the + block is clean. Non-finite blocks still use the guarded helper. This + promoted conversion-speed path measured focused dual-scale / FAST normal + quantization at `3844.38` / `3582.57` cycles per 32 values in the latest + guard. +- CPU imatrix-weighted scale-MSE scoring now uses the same finite-only + nearest-Codebook10 helper when the source block is finite. This preserves the + guarded path for non-finite input while improving FAST imatrix GGUF creation: + same-session pre-candidate FAST imatrix was `5258.07` cycles / 32 values, and + two guarded candidate passes measured `4448.73` and `4447.32`. Dual imatrix + stayed in the noisy guarded band, so the claim is limited to FAST imatrix. +- The weighted and unweighted exhaustive scale loops are split as well, so + the normal quantizer does not branch on imatrix state for every finite + UE4M3 candidate. This keeps the exact candidate order and tie behavior while + improving CPU GGUF creation speed for both ROCmFP4 layouts. +- The CPU quantizers fill every packed nibble byte in each ROCmFP4 block, so + they no longer clear `qs` before packing. This is an exact store cleanup: + correctness stayed clean, `q4_0_rocmfp4` quantize improved to `5244.46` + cycles / 32 values, `q4_0_rocmfp4_fast` improved to `4668.28` cycles / 32 + values. +- CPU scale selection scans block maxima with a plain absolute-value compare + instead of `fmaxf`. NaNs still do not raise the maximum, infinities are still + rejected before scale search, and finite output is unchanged. This moved + `q4_0_rocmfp4` quantize to `5228.64` cycles / 32 values and FAST to + `4516.45` cycles / 32 values. +- CPU dequantization and quantization-error scoring use the ROCmFP4 arithmetic + Codebook10 decoder, which avoids a table fetch per unpacked nibble. CPU + fallback vec-dot keeps the 16-entry codebook table because measured Strix + Halo fallback dot speed was better with the table there. This hybrid keeps + the dequant win without regressing CPU vec-dot. +- CPU fallback vec-dot now reuses each packed ROCmFP4 byte for both low and + high nibbles before table decode. This preserves exact output, moved focused + dual-scale vec-dot from `31.86` to the `29.77`-`29.82` cycles / 32 values + band, and kept FAST in the `27.04`-`27.06` cycles / 32 values band. +- Two additional CPU decode shortcuts were rechecked on 2026-05-24 and + rejected. Replacing arithmetic Codebook10 decode with table decode in + row dequantization and scale-MSE scoring slowed CPU dequantization from the + guarded `33`-cycle band to `51`-`84` cycles / 32 values. Returning a direct + decoded value during MSE scoring instead of index-then-decode also passed + correctness but regressed normal quantization in the measured guard, so the + original arithmetic decoder remains the CPU path. A narrower full-block + weighted-MSE-only retry was also rejected because it regressed the focused + FAST imatrix timing to `5007.73` cycles / 32 values. +- A direct finite-scale table helper for CPU quantizer scale-search candidates + was also rejected. It helped dual-scale normal quantization, but repeat + guard runs regressed FAST normal quantization to `4043.60` and + `4289.66` cycles / 32 values. The guarded scale helper remains in place + because the compact STRIX_LEAN preset depends heavily on FAST tensors. +- NaN-only and branchless finite-scan variants were also rejected for default. + The NaN-only scan was noisy and the same-session `isfinite()` scan measured + better dual-scale normal quantization (`3735.41` vs `3776.53` cycles / 32 + values), while the branchless boolean-and scan regressed FAST normal + quantization to `3773.44` cycles / 32 values. +- Per-value Codebook10 quantization uses exact nearest-neighbor thresholds + instead of a 16-entry scan. The hot quantizer path uses one reciprocal scale + per candidate/block and multiplies each value by that reciprocal instead of + dividing per value. On the Qwen3-4B FAST_COHERENT check this kept PPL tied + with the accepted artifact while cutting GGUF creation time further. +- CPU exhaustive scale search now prunes lower-scale candidates after the + clipped max-value error alone cannot beat the current best error. The + imatrix path uses the positive calibration weight of the max element for the + same conservative bound. This does not change accepted values because every + skipped smaller scale has a larger unavoidable clip error. The weighted path + also combines imatrix weight preparation, max-absolute scan, and max-value + pruning-weight selection into one pre-scan, while both exhaustive loops stop + once no valid lower or upper scale candidates remain. Latest 40-iteration + quant guard run: `q4_0_rocmfp4` normal `4049.28` cycles / 32 values, FAST + normal `3721.48`, imatrix dual-scale `5251.76`, imatrix FAST `4898.75`. +- Adding `GGML_RESTRICT` to the local CPU scale-scoring helper pointers was + tested and rejected. It compiled and passed the quant guard, but did not + improve normal quantization and made the imatrix guard noisier/slower in + the measured pass, so the existing helper signatures were kept. +- Regression guards: + - `scripts/check-rocmfp4-deepseek-regression.sh` is an optional compatibility + smoke guard for a second ROCmFP4-converted model family. It is not part of + the promoted-gain gate because no reproducible DeepSeek speedup has been + established yet. The default all-regression script skips it unless + `INCLUDE_DEEPSEEK_SMOKE=1` is set. + - `scripts/check-rocmfp4-qwen-mtp-regression.sh` verifies the promoted + Qwen3.6 27B MTP ROCmFP4 262k-context `draft-mtp` path. The stable guard + uses ROCm0 and `--spec-draft-n-max 4`. On ROCm0, `n-max 4` improved the + short guard prompt from `27.7` to `33.6 tok/s` versus `n-max 3`, and the + longer forced-output prompt from `25.6` to `26.3 tok/s`; the dual-scale + MMVQ VDR tune then moved the longer prompt to `27.8 tok/s`. `n-max 5` + and `n-max 6` improved only the bursty short prompt after the MMVQ tune + but regressed sustained output to `24.8` and `24.0 tok/s`, respectively. + Follow-up `n-max 5` checks with `--spec-draft-p-min 0.75`, + `--spec-draft-p-min 0.90`, and `--spec-draft-n-min 1` still held the + short prompt at `45.0`-`45.1 tok/s` but kept sustained output at only + `24.7 tok/s`, so `n-max 4` remains promoted. + The guard now checks both short and sustained prompts so high-acceptance + best-case gains do not hide sustained-output regressions. + Additional ROCm0 checks rejected larger batch sizes, smaller ubatch, + q8/f16 KV cache, `--swa-full`, `--no-host`, `--no-op-offload`, + `--no-repack`, polling changes, lower and higher CPU thread counts, + greedy sampling, backend sampling, p-split changes, sampler chain trimming, + and combined MTP+ngram speculation because none improved the sustained + prompt. After the MMVQ VDR tune, `-b 1024 -ub 512`, `-b 2048 -ub 512`, + and `-b 512 -ub 256` still tied the promoted `27.8 tok/s` sustained + result, while `-b 1024 -ub 1024` and + `--poll-batch 1 --spec-draft-poll-batch 1` stayed at `27.7 tok/s`. + q8 and f16 KV cache dropped sustained decode to `25.4` and `26.0 tok/s`. + The MTP draft loop now honors `--spec-draft-p-min`. With that cutoff active, + `n-max 4, p-min 0.75` tied the promoted path at `33.5 tok/s` short and + `27.6 tok/s` sustained, while `n-max 5` stayed burst-only at `45.0`-`45.1` + tok/s short and `24.6`-`24.7 tok/s` sustained for `p-min 0.75` and `0.90`. + Follow-up ROCm0 KV isolation confirmed this is not a draft-cache issue: + q8 main KV with q4 draft KV measured `33.5 tok/s` short and `25.3 tok/s` + sustained, q8 main K only measured `31.5` and `21.4 tok/s`, and q8 main V + only measured `31.1` and `21.9 tok/s`. The promoted ROCm0 path therefore + keeps both accepted and draft KV at q4. + `--spec-draft-p-split 0.05` and `0.20` both tied the `27.7 tok/s` + sustained guard, `-t 24 -tb 32 --spec-draft-threads 24 + --spec-draft-threads-batch 32` tied, and trimming the sampler chain to + `top_k;top_p;temperature` regressed sustained decode to `26.6 tok/s`. + Draft-only q8 KV (`SPEC_DRAFT_TYPE_K=q8_0`, + `SPEC_DRAFT_TYPE_V=q8_0`) with main KV left at q4 tied the promoted path + at `33.6 tok/s` short and `27.7 tok/s` sustained, so the leaner q4 draft + KV default remains promoted. + The heavier Qwen3.6 27B MTP `STRIX_MTP_Q6` ROCmFP4 profile was also + checked with the same ROCm0 guard. It reached `30.1 tok/s` on the short + prompt but only `21.3 tok/s` sustained, below the guard floor, so the + Q6/quality-biased profile is not promoted for throughput. + The guard script supports `SPEC_DRAFT_N_MAX`, `SPEC_DRAFT_N_MIN`, + `SPEC_DRAFT_P_MIN`, `SPEC_DRAFT_P_SPLIT`, `THREADS`, `THREADS_BATCH`, + `SPEC_DRAFT_THREADS`, `SPEC_DRAFT_THREADS_BATCH`, `SAMPLERS`, + `BATCH_SIZE`, `UBATCH_SIZE`, `CACHE_TYPE_K`, `CACHE_TYPE_V`, + `SPEC_DRAFT_TYPE_K`, `SPEC_DRAFT_TYPE_V`, `EXTRA_ARGS`, and + `SPEC_EXTRA_ARGS` environment overrides for + controlled sweeps while defaulting to the promoted `n-max 4`, `n-min 0`, + `p-min 0.0`, `p-split 0.10`, q4 KV, `-t 16`, `-tb 32`, `-b 512`, and + `-ub 512` settings. + - `scripts/check-rocmfp4-quant-regression.sh` runs quant correctness plus + CPU quantizer, dequantizer, and vec-dot cycle ceilings for both ROCmFP4 + block layouts. It checks normal and imatrix quantization, and now also + protects CPU fallback dequant/vec-dot paths so decode-helper experiments + cannot pass while slowing non-GPU fallbacks. Recent pass after keeping the + `isfinite()` finite-block scoring promotion and adding the weighted/imatrix + finite scorer: dual-scale quant `3844.38`, FAST quant `3582.57`, dual + dequant `33.59`, FAST dequant `33.13`, dual vec-dot `29.96`, FAST vec-dot + `27.03`, dual imatrix `5587.43`, and FAST imatrix `4447.32` cycles / 32 + values. + - `scripts/check-rocmfp4-vulkan-runtime-regression.sh` measures focused + Vulkan `MUL_MAT` runtime for ROCmFP4 FAST and dual-scale layouts, catching + shader regressions that can be hidden by end-to-end decode noise. The + guard covers `n=1`, `n=2`, `n=4`, and `n=8` for both layouts. Recent + tightened serial pass measured FAST `55.82`, `71.59`, `105.10`, and + `163.41` us/run for `n=1/2/4/8`, and dual-scale `64.87`, `83.24`, + `118.56`, and `194.27` us/run for `n=1/2/4/8`. The previous serial pass + after adding the shared ROCmFP4 UE4M3 + scale LUT was FAST `53.66`, `71.67`, `105.14` us/run and dual-scale + `65.05`, `83.07`, `122.70` us/run for the same shapes. + - `scripts/check-rocmfp4-vulkan-cpy-regression.sh` measures Vulkan0 CPY for + `F32/F16/BF16 -> Q4_0_ROCMFP4`, + `F32/F16/BF16 -> Q4_0_ROCMFP4_FAST`, and the matching quant-to-F32 + dequant copy paths on the same large shape used by the ROCm CPY guard. + It is included in the all-regression harness to catch accidental Vulkan + copy fallback or shader-routing regressions. + - `scripts/check-rocmfp4-rocm-runtime-regression.sh` measures the same + focused `MUL_MAT` shapes on ROCm0. The guard covers `n=1`, `n=2`, `n=4`, + and `n=8` for both FAST and dual-scale ROCmFP4 so MTP-style multi-column + regressions are caught instead of only protecting single-token matvec. + Recent tightened serial pass after extending the RDNA3.5 two-warp route + through `n=2` measured FAST `51.66`, `66.56`, `101.27`, and `168.84` + us/run for `n=1/2/4/8`, and dual-scale `53.66`, `58.40`, `87.32`, and + `148.26` us/run for `n=1/2/4/8`. The `n=8` guard was added after a + rejected wide-rows candidate regressed dual-scale `n=8` to `1382.72` + us/run. + - `scripts/check-rocmfp4-rocm-cpy-regression.sh` measures ROCm0 CPY for + `F32/F16/BF16 -> Q4_0_ROCMFP4`, + `F32/F16/BF16 -> Q4_0_ROCMFP4_FAST`, and the matching quant-to-F32 + dequant copy paths on a large `8192x512x2` shape. Recent serial-gate pass + after direct FAST value scoring: F32/F16/BF16-to-dual `1111.73`, + `1008.69`, `1006.93` us/run, dual-to-F32 `182.25` us/run, + F32/F16/BF16-to-FAST `1047.21`, `950.93`, `951.00` us/run, and + FAST-to-F32 `170.36` us/run. + - `scripts/check-rocmfp4-all-regression.sh` runs the promoted-gain gate: + quant, Vulkan runtime, Vulkan CPY, ROCm runtime, ROCm FlashAttention, + ROCm CPY, and Qwen MTP guards serially, then checks ROCm KFD PIDs. Runtime + microbenchmarks should use this serial path rather than being run in + parallel with other GPU/UMA workloads. DeepSeek is not part of this + promoted-gain gate because no reproducible DeepSeek speedup has been + established; run the separate compatibility smoke only when explicitly + needed with `INCLUDE_DEEPSEEK_SMOKE=1`. The K/Q block-pair FlashAttention + serial pass held Qwen3.6 27B MTP at `33.2 tok/s` short and `27.7 tok/s` + sustained, with no KFD PIDs left running. The same pass measured ROCm + runtime `MUL_MAT` at FAST `54.16` / `66.13` / `101.73` / `178.36` us and + dual-scale `51.65` / `57.98` / `88.33` / `148.48` us for `n=1/2/4/8`, + and ROCm FlashAttention at `81.62` / `78.13` us for 64d dual-scale / FAST + and `228.58` / `199.32` us for Qwen-style 128d dual-scale / FAST. + - The MTP `accept()` path now skips re-copying `pending_h` when acceptance + lands on the final verify row already staged by `process()`. The focused + Qwen MTP guard held `33.8 tok/s` short and `27.9 tok/s` sustained after + this host-side cleanup. + - The MTP `process()` path now stores only the non-final verify rows needed + for partial-accept rollback, while copying the final target hidden row + directly to `pending_h` for the next MTP step. The focused Qwen MTP guard + held `33.9 tok/s` short and `27.9 tok/s` sustained, and the full serial + gate held `33.7 tok/s` short and `27.9 tok/s` sustained after this cleanup. + - The simple draft and MTP draft paths now return before + `llama_decode(ctx_dft, batch)` when no sequence is actively drafting. The + focused Qwen MTP guard held `33.8 tok/s` short and `28.0 tok/s` sustained, + and the full serial gate held `33.7 tok/s` short and `27.9 tok/s` + sustained after this idle/no-draft cleanup. + - The simple draft path now uses a reusable `uint8_t` drafting-state buffer + instead of allocating `std::vector` per draft call. The focused Qwen + MTP guard held `33.7 tok/s` short and `28.0 tok/s` sustained, and the full + serial gate held `33.8 tok/s` short and `27.9 tok/s` sustained. + - The MTP draft loop now skips its debug-candidate loop entirely unless + debug logging is enabled, hoists the debug verbosity state once per draft + call, and uses direct vector indexing in the hot per-token path. The + focused Qwen MTP guard held `33.8 tok/s` short and `27.9 tok/s` sustained + after this cleanup, so it is kept as a low-risk host-side simplification + rather than claimed as a decode-speed gain. + - The simple and MTP draft paths now use a shared direct one-sequence batch + append helper instead of `common_batch_add(..., { seq_id }, ...)` in the + hot speculative path. This avoids a temporary sequence-id container without + changing multi-sequence batch behavior elsewhere. The focused Qwen MTP + guard held `33.9 tok/s` short and `27.9 tok/s` sustained, and the full + serial gate held `33.8 tok/s` short and `27.9 tok/s` sustained. + - The MTP verify-row buffer now reserves for the configured draft depth and + only grows when needed instead of shrinking/resizing on every verification + pass. The focused Qwen MTP guard held `33.8 tok/s` short and `27.9 tok/s` + sustained, and the full serial gate held `33.9 tok/s` short and + `27.9 tok/s` sustained with no KFD PIDs left running. + - The MTP `process()` path now copies retained verification hidden rows from + the target embedding buffer in one contiguous `memcpy()` instead of one + `llama_get_embeddings_pre_norm_ith()` plus `memcpy()` per row. Rollback + behavior is unchanged. The focused Qwen3.6 27B MTP guard held + `33.7 tok/s` short and `27.9 tok/s` sustained, and the 35B A3B guard held + `104.1 tok/s` short and `89.2 tok/s` sustained. The same build later + passed the full serial all-regression gate with `INCLUDE_QWEN35_A3B_GUARD=1`, + including Qwen3.6 27B MTP at `33.8` / `27.9 tok/s` and Qwen3.6 35B A3B MTP + at `104.1` / `89.3 tok/s`. + - A matching contiguous-pointer cleanup inside the MTP `draft()` loop was + tested and rejected. It replaced per-row + `llama_get_embeddings_pre_norm_ith(ctx_dft, i_batch)` calls with one + `llama_get_embeddings_pre_norm(ctx_dft)` pointer per draft decode + iteration. The 27B MTP guard held at `33.8` / `27.9 tok/s`, but the 35B A3B + repeat measured `104.3` / `88.7` and `104.3` / `89.2 tok/s`, below the + promoted `89.3 tok/s` sustained band, so the code change was removed. + - A single-sequence MTP `draft()` fast path was tested and rejected on + 2026-05-25. It removed the active-sequence bookkeeping loop for the common + `n_seq == 1` case and passed the dense 27B guard at `33.7` / `28.0 tok/s`, + but the 35B A3B sustained guard collapsed to `25.7 tok/s` despite a + passing `103.1 tok/s` short check. Reverting that path restored the 35B A3B + guard to `104.3` / `90.3 tok/s`, so the shared multi-sequence draft loop + remains the promoted implementation. + - A dual-scale-only finite-pack CPU quantizer shortcut was tested and + rejected on 2026-05-24. It passed correctness, but even after isolating + the shared scale chooser it regressed the protected FAST quant path + (`q4_0_rocmfp4_fast` normal quant rose into the `3882`-`4022` cycles / 32 + values band versus the clean `3623` cycles / 32 values baseline). The + final packing loops therefore remain on the guarded nearest-code helper. +- The Strix ROCmFP4 build script now builds the quant regression test + binaries (`test-quantize-fns`, `test-quantize-perf`, and + `test-backend-ops`) alongside `llama-cli`, `llama-quantize`, and + `llama-bench`, so quant and runtime guards work from a clean build + directory. +- `test-quantize-perf --imatrix` benchmarks the ROCmFP4 quality/coherence + path through `ggml_quantize_chunk(..., imatrix)` with synthetic importance + weights. The quant regression guard includes this mode to reject imatrix-side + CPU changes that would not show up in the normal non-imatrix timing. + +High-upside Strix-specific work: +- The next significant speed step is likely a fused ROCmFP4 decode-and-dot + ROCm matvec path. The current ROCmFP4 wins come from compact blocks and + avoiding generic decode/copy overhead; a fused path would keep Codebook10 + values and UE4M3 scales in registers/shared memory through the dot product + instead of materializing wider intermediates. +- A ROCmFP4-aware long-context attention path is the second major target. The + long-context limit is mostly KV traffic, FlashAttention shape efficiency, and + memory bandwidth. The current dual-scale layout protects coherence, but it is + not itself a large speed lever unless the attention kernels consume it + directly and efficiently. +- MTP target/draft overlap is a possible scheduler project, not a flag. The + current common MTP flow is serial on one Strix Halo ROCm device. Real overlap + would need scheduler changes plus acceptance/correctness guards, and it + should only be promoted if it beats the sustained Qwen 262k guard. +- Coherence should stay protected by tensor-aware ROCmFP4 profiles: dual-scale + for sensitive tensors and FAST only where guarded model tests show no + sustained decode or quality regression. Pure speed-only profiles are not + promoted in this tree without the serial regression gate. + +Hardware note: +- This is a special AMD-targeted ggml/llama.cpp quantization and backend + path. It includes custom Vulkan and ROCm/HIP handling for the new GGUF + types, but it is not yet a native rocWMMA FP4 tensor-core implementation. + Current speed gains come from the compact block layout and backend decode + paths; deeper rocWMMA/cooperative-matrix work is future optimization work. +- NVIDIA CUDA is disabled in the Strix-FP4 build (`-DGGML_CUDA=OFF`). Some + upstream llama.cpp HIP backend sources still live under + `ggml/src/ggml-cuda` and are compiled by HIP for AMD, but the ROCmFP4-owned + helper code and user-facing build/run path are ROCm/HIP/Vulkan targeted. + This tree also accepts `GGML_HIP_ENABLE_UNIFIED_MEMORY=1` as the AMD-named + alias for the upstream unified-memory switch. +- The bundled rocWMMA 7.1.0 headers expose gfx12 WMMA paths for FP8/BF8 and + integer 8-bit inputs, but no native FP4 input type or FP4 WMMA/MFMA builtin + is visible locally. A true matrix-core ROCmFP4 path therefore needs a + measured unpack/convert strategy first, such as ROCmFP4 Codebook10 to int8 + WMMA or FP8 WMMA tiles, before claiming native FP4 tensor-core execution. +- rocWMMA FlashAttention is intentionally opt-in via + `GGML_HIP_ROCWMMA_FATTN=ON scripts/build-strix-rocmfp4-mtp.sh`. + It currently compiles with the local rocWMMA headers in + `/path/to/third_party/rocWMMA`, but the Strix Halo benchmark + regressed sustained Qwen MTP decode (`23.3 tok/s` vs the promoted + `26.2 tok/s` default HIP FlashAttention path), so the default build keeps it + disabled. A follow-up on the 35B A3B ROCmFP4 MTP guard measured only + `99.7` / `76.1 tok/s` versus the promoted `104.4` / `89.3` band, so this + remains true for the MoE-heavy profile as well. +- TurboQuant and TriAttention are not present as runtime flags in this isolated + tree. Integrating them would require a source-level merge into `strix-fp4`; + do not claim support or promote them unless they beat the serial ROCmFP4 + regression gate. diff --git a/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4.c b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4.c new file mode 100644 index 000000000..bccefcec1 --- /dev/null +++ b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4.c @@ -0,0 +1,797 @@ +#define GGML_COMMON_DECL_C +#include "../src/ggml-common.h" + +#include "rocmfp4.h" + +#include +#include +#include + +// ggml-base is compiled architecture-neutral (no -mavx2), so SIMD for the hot +// CPU dot product is enabled per-function via a target attribute plus a runtime +// CPU check. This keeps the AVX2 path in one translation unit without exporting +// internals to ggml-cpu. Non-GNU/non-x86 builds fall back to the scalar loop. +#if defined(__GNUC__) && (defined(__x86_64__) || defined(__i386__)) && !defined(ROCMFP4_NO_AVX2) +#include +#define ROCMFP4_X86_AVX2_DISPATCH 1 +#endif + +// ROCmFP4 stores a signed integer FP4-like codebook at half-scale. It is +// E2M1-derived, but the largest magnitude is retuned from 12 to 10 after +// sampling Qwen3 dense tensors; this reduces outlier pull without changing the +// packed 4-bit layout or integer dot-product path. +static const int8_t rocmfp4_codebook[16] = { + 0, 1, 2, 3, 4, 6, 8, 10, + 0, -1, -2, -3, -4, -6, -8,-10, +}; + +static inline int8_t rocmfp4_decode(uint8_t q) { + q &= 0x0f; + const int mag3 = q & 0x07; + const int mag = mag3 <= 4 ? mag3 : 2*mag3 - 4; + return (q & 0x08) ? -mag : mag; +} + +static inline int8_t rocmfp4_decode_table(uint8_t q) { + return rocmfp4_codebook[q & 0x0f]; +} + +// Finite unsigned E4M3 scale bytes decoded to the half-scale values used by +// ROCmFP4. Keeping this as a table avoids rebuilding identical FP32 values for +// every candidate during exhaustive scale search. +#define ROCMFP4_SCALE_SUB(M) ((M) * 0x1p-10f) +#define ROCMFP4_SCALE_E1(M) ((8 + (M)) * 0x1p-10f) +#define ROCMFP4_SCALE_E2(M) ((8 + (M)) * 0x1p-9f) +#define ROCMFP4_SCALE_E3(M) ((8 + (M)) * 0x1p-8f) +#define ROCMFP4_SCALE_E4(M) ((8 + (M)) * 0x1p-7f) +#define ROCMFP4_SCALE_E5(M) ((8 + (M)) * 0x1p-6f) +#define ROCMFP4_SCALE_E6(M) ((8 + (M)) * 0x1p-5f) +#define ROCMFP4_SCALE_E7(M) ((8 + (M)) * 0x1p-4f) +#define ROCMFP4_SCALE_E8(M) ((8 + (M)) * 0x1p-3f) +#define ROCMFP4_SCALE_E9(M) ((8 + (M)) * 0x1p-2f) +#define ROCMFP4_SCALE_E10(M) ((8 + (M)) * 0x1p-1f) +#define ROCMFP4_SCALE_E11(M) ((8 + (M)) * 0x1p0f) +#define ROCMFP4_SCALE_E12(M) ((8 + (M)) * 0x1p1f) +#define ROCMFP4_SCALE_E13(M) ((8 + (M)) * 0x1p2f) +#define ROCMFP4_SCALE_E14(M) ((8 + (M)) * 0x1p3f) +#define ROCMFP4_SCALE_E15(M) ((8 + (M)) * 0x1p4f) + +static const float rocmfp4_scale_ue4m3_half[127] = { + ROCMFP4_SCALE_SUB(0), ROCMFP4_SCALE_SUB(1), ROCMFP4_SCALE_SUB(2), ROCMFP4_SCALE_SUB(3), + ROCMFP4_SCALE_SUB(4), ROCMFP4_SCALE_SUB(5), ROCMFP4_SCALE_SUB(6), ROCMFP4_SCALE_SUB(7), + ROCMFP4_SCALE_E1(0), ROCMFP4_SCALE_E1(1), ROCMFP4_SCALE_E1(2), ROCMFP4_SCALE_E1(3), + ROCMFP4_SCALE_E1(4), ROCMFP4_SCALE_E1(5), ROCMFP4_SCALE_E1(6), ROCMFP4_SCALE_E1(7), + ROCMFP4_SCALE_E2(0), ROCMFP4_SCALE_E2(1), ROCMFP4_SCALE_E2(2), ROCMFP4_SCALE_E2(3), + ROCMFP4_SCALE_E2(4), ROCMFP4_SCALE_E2(5), ROCMFP4_SCALE_E2(6), ROCMFP4_SCALE_E2(7), + ROCMFP4_SCALE_E3(0), ROCMFP4_SCALE_E3(1), ROCMFP4_SCALE_E3(2), ROCMFP4_SCALE_E3(3), + ROCMFP4_SCALE_E3(4), ROCMFP4_SCALE_E3(5), ROCMFP4_SCALE_E3(6), ROCMFP4_SCALE_E3(7), + ROCMFP4_SCALE_E4(0), ROCMFP4_SCALE_E4(1), ROCMFP4_SCALE_E4(2), ROCMFP4_SCALE_E4(3), + ROCMFP4_SCALE_E4(4), ROCMFP4_SCALE_E4(5), ROCMFP4_SCALE_E4(6), ROCMFP4_SCALE_E4(7), + ROCMFP4_SCALE_E5(0), ROCMFP4_SCALE_E5(1), ROCMFP4_SCALE_E5(2), ROCMFP4_SCALE_E5(3), + ROCMFP4_SCALE_E5(4), ROCMFP4_SCALE_E5(5), ROCMFP4_SCALE_E5(6), ROCMFP4_SCALE_E5(7), + ROCMFP4_SCALE_E6(0), ROCMFP4_SCALE_E6(1), ROCMFP4_SCALE_E6(2), ROCMFP4_SCALE_E6(3), + ROCMFP4_SCALE_E6(4), ROCMFP4_SCALE_E6(5), ROCMFP4_SCALE_E6(6), ROCMFP4_SCALE_E6(7), + ROCMFP4_SCALE_E7(0), ROCMFP4_SCALE_E7(1), ROCMFP4_SCALE_E7(2), ROCMFP4_SCALE_E7(3), + ROCMFP4_SCALE_E7(4), ROCMFP4_SCALE_E7(5), ROCMFP4_SCALE_E7(6), ROCMFP4_SCALE_E7(7), + ROCMFP4_SCALE_E8(0), ROCMFP4_SCALE_E8(1), ROCMFP4_SCALE_E8(2), ROCMFP4_SCALE_E8(3), + ROCMFP4_SCALE_E8(4), ROCMFP4_SCALE_E8(5), ROCMFP4_SCALE_E8(6), ROCMFP4_SCALE_E8(7), + ROCMFP4_SCALE_E9(0), ROCMFP4_SCALE_E9(1), ROCMFP4_SCALE_E9(2), ROCMFP4_SCALE_E9(3), + ROCMFP4_SCALE_E9(4), ROCMFP4_SCALE_E9(5), ROCMFP4_SCALE_E9(6), ROCMFP4_SCALE_E9(7), + ROCMFP4_SCALE_E10(0), ROCMFP4_SCALE_E10(1), ROCMFP4_SCALE_E10(2), ROCMFP4_SCALE_E10(3), + ROCMFP4_SCALE_E10(4), ROCMFP4_SCALE_E10(5), ROCMFP4_SCALE_E10(6), ROCMFP4_SCALE_E10(7), + ROCMFP4_SCALE_E11(0), ROCMFP4_SCALE_E11(1), ROCMFP4_SCALE_E11(2), ROCMFP4_SCALE_E11(3), + ROCMFP4_SCALE_E11(4), ROCMFP4_SCALE_E11(5), ROCMFP4_SCALE_E11(6), ROCMFP4_SCALE_E11(7), + ROCMFP4_SCALE_E12(0), ROCMFP4_SCALE_E12(1), ROCMFP4_SCALE_E12(2), ROCMFP4_SCALE_E12(3), + ROCMFP4_SCALE_E12(4), ROCMFP4_SCALE_E12(5), ROCMFP4_SCALE_E12(6), ROCMFP4_SCALE_E12(7), + ROCMFP4_SCALE_E13(0), ROCMFP4_SCALE_E13(1), ROCMFP4_SCALE_E13(2), ROCMFP4_SCALE_E13(3), + ROCMFP4_SCALE_E13(4), ROCMFP4_SCALE_E13(5), ROCMFP4_SCALE_E13(6), ROCMFP4_SCALE_E13(7), + ROCMFP4_SCALE_E14(0), ROCMFP4_SCALE_E14(1), ROCMFP4_SCALE_E14(2), ROCMFP4_SCALE_E14(3), + ROCMFP4_SCALE_E14(4), ROCMFP4_SCALE_E14(5), ROCMFP4_SCALE_E14(6), ROCMFP4_SCALE_E14(7), + ROCMFP4_SCALE_E15(0), ROCMFP4_SCALE_E15(1), ROCMFP4_SCALE_E15(2), ROCMFP4_SCALE_E15(3), + ROCMFP4_SCALE_E15(4), ROCMFP4_SCALE_E15(5), ROCMFP4_SCALE_E15(6), +}; + +#undef ROCMFP4_SCALE_SUB +#undef ROCMFP4_SCALE_E1 +#undef ROCMFP4_SCALE_E2 +#undef ROCMFP4_SCALE_E3 +#undef ROCMFP4_SCALE_E4 +#undef ROCMFP4_SCALE_E5 +#undef ROCMFP4_SCALE_E6 +#undef ROCMFP4_SCALE_E7 +#undef ROCMFP4_SCALE_E8 +#undef ROCMFP4_SCALE_E9 +#undef ROCMFP4_SCALE_E10 +#undef ROCMFP4_SCALE_E11 +#undef ROCMFP4_SCALE_E12 +#undef ROCMFP4_SCALE_E13 +#undef ROCMFP4_SCALE_E14 +#undef ROCMFP4_SCALE_E15 + +static inline float rocmfp4_ue4m3_to_fp32_half(uint8_t e) { + return e <= 0x7e ? rocmfp4_scale_ue4m3_half[e] : 0.0f; +} + +static inline uint8_t rocmfp4_best_index_scaled_finite(float x, float inv_scale_half) { + // Exact nearest-neighbor thresholds for Codebook10: + // 0, +/-1, +/-2, +/-3, +/-4, +/-6, +/-8, +/-10 + // Ties intentionally choose the lower-magnitude code, matching the former + // linear scan because the positive codes and zero appear first. + const float a = fabsf(x * inv_scale_half); + if (a <= 0.5f) { + return 0; + } + + const bool neg = x < 0.0f; + if (a <= 1.5f) { + return neg ? 9 : 1; + } + if (a <= 2.5f) { + return neg ? 10 : 2; + } + if (a <= 3.5f) { + return neg ? 11 : 3; + } + if (a <= 5.0f) { + return neg ? 12 : 4; + } + if (a <= 7.0f) { + return neg ? 13 : 5; + } + if (a <= 9.0f) { + return neg ? 14 : 6; + } + + return neg ? 15 : 7; +} + +static inline uint8_t rocmfp4_best_index_scaled(float x, float inv_scale_half) { + if (!isfinite(x)) { + return 0; + } + + return rocmfp4_best_index_scaled_finite(x, inv_scale_half); +} + +// Fused best-index + decode used only inside the exhaustive scale search. The +// scale search re-scans every block element for every candidate scale byte, so +// avoiding the code -> decode round-trip on the hottest quantize path matters. +// Returns the same signed Codebook10 magnitude that +// rocmfp4_decode(rocmfp4_best_index_scaled_finite(x, inv_scale_half)) produces, +// so quantized output is bit-identical to the previous path. +static inline float rocmfp4_decoded_mag_scaled_finite(float x, float inv_scale_half) { + const float a = fabsf(x * inv_scale_half); + + float mag; + if (a <= 0.5f) { + mag = 0.0f; + } else if (a <= 1.5f) { + mag = 1.0f; + } else if (a <= 2.5f) { + mag = 2.0f; + } else if (a <= 3.5f) { + mag = 3.0f; + } else if (a <= 5.0f) { + mag = 4.0f; + } else if (a <= 7.0f) { + mag = 6.0f; + } else if (a <= 9.0f) { + mag = 8.0f; + } else { + mag = 10.0f; + } + + return x < 0.0f ? -mag : mag; +} + +static inline float rocmfp4_decoded_mag_scaled(float x, float inv_scale_half) { + if (!isfinite(x)) { + return 0.0f; + } + + return rocmfp4_decoded_mag_scaled_finite(x, inv_scale_half); +} + +static inline bool rocmfp4_scale_is_valid(uint8_t e) { + // ROCmFP4 scale bytes are unsigned finite E4M3 values. 0x7f is NaN in the + // unsigned encoding and values with the sign bit set are not valid scales. + return e <= 0x7e; +} + +static float rocmfp4_block_mse_for_scale_unweighted( + const float * x, int n, int e, float best_err) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e); + const float inv_scale_half = 1.0f / scale_half; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfp4_decoded_mag_scaled(x[i], inv_scale_half) * scale_half; + const float d = x[i] - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfp4_block_mse_for_scale_unweighted_finite( + const float * x, int n, int e, float best_err) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e); + const float inv_scale_half = 1.0f / scale_half; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfp4_decoded_mag_scaled_finite(x[i], inv_scale_half) * scale_half; + const float d = x[i] - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfp4_block_mse_for_scale_weighted( + const float * x, int n, const float * mse_weights, int e, float best_err) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e); + const float inv_scale_half = 1.0f / scale_half; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfp4_decoded_mag_scaled(x[i], inv_scale_half) * scale_half; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfp4_block_mse_for_scale_weighted_finite( + const float * x, int n, const float * mse_weights, int e, float best_err) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e); + const float inv_scale_half = 1.0f / scale_half; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfp4_decoded_mag_scaled_finite(x[i], inv_scale_half) * scale_half; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static void rocmfp4_prepare_mse_weights( + float * dst, const float * x, int n, const float * quant_weights, float sigma2, + float * max_abs, float * max_abs_weight, bool * all_finite) { + *max_abs = 0.0f; + *max_abs_weight = 0.0f; + *all_finite = true; + + for (int i = 0; i < n; ++i) { + const float qw = quant_weights[i]; + const float ax = fabsf(x[i]); + const float weight = isfinite(qw) && qw > 0.0f ? qw * sqrtf(sigma2 + x[i]*x[i]) : 0.0f; + *all_finite = *all_finite && isfinite(x[i]); + + if (ax > *max_abs) { + *max_abs = ax; + *max_abs_weight = weight; + } else if (ax == *max_abs && weight > *max_abs_weight) { + *max_abs_weight = weight; + } + + // Match llama.cpp's imatrix weighting style for Q4_0: calibration + // importance is scaled by row energy so large activations remain protected. + dst[i] = weight; + } +} + +static int rocmfp4_nearest_scale_ue4m3(float target_scale_half) { + if (!(target_scale_half > 0.0f) || !isfinite(target_scale_half)) { + return 1; + } + + int lo = 1; + int hi = 126; + while (lo < hi) { + const int mid = lo + (hi - lo) / 2; + if (rocmfp4_ue4m3_to_fp32_half((uint8_t) mid) < target_scale_half) { + lo = mid + 1; + } else { + hi = mid; + } + } + + if (lo == 1) { + return 1; + } + + const float hi_scale = rocmfp4_ue4m3_to_fp32_half((uint8_t) lo); + const float lo_scale = rocmfp4_ue4m3_to_fp32_half((uint8_t) (lo - 1)); + + // Match the former ascending nearest scan: exact midpoint ties keep the + // lower scale byte. + return (target_scale_half - lo_scale <= hi_scale - target_scale_half) ? lo - 1 : lo; +} + +static uint8_t rocmfp4_choose_scale_ue4m3_exhaustive_unweighted( + const float * x, int n, float max_abs, bool all_finite) { + const int start_e = rocmfp4_nearest_scale_ue4m3(max_abs / 10.0f); + + int best_e = 0; + float best_err = FLT_MAX; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e0); + const float clip_delta = max_abs - 10.0f*scale_half; + if (clip_delta > 0.0f && clip_delta*clip_delta > best_err) { + lower_done = true; + } else { + const float err = all_finite ? + rocmfp4_block_mse_for_scale_unweighted_finite(x, n, e0, best_err) : + rocmfp4_block_mse_for_scale_unweighted(x, n, e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = e0; + } + } + } + + const int e1 = start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = all_finite ? + rocmfp4_block_mse_for_scale_unweighted_finite(x, n, e1, best_err) : + rocmfp4_block_mse_for_scale_unweighted(x, n, e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return (uint8_t) best_e; +} + +static uint8_t rocmfp4_choose_scale_ue4m3_exhaustive_weighted( + const float * x, int n, const float * mse_weights, float max_abs, float max_abs_weight, bool all_finite) { + const int start_e = rocmfp4_nearest_scale_ue4m3(max_abs / 10.0f); + + int best_e = 0; + float best_err = FLT_MAX; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e0); + const float clip_delta = max_abs - 10.0f*scale_half; + if (max_abs_weight > 0.0f && clip_delta > 0.0f && max_abs_weight*clip_delta*clip_delta > best_err) { + lower_done = true; + } else { + const float err = all_finite ? + rocmfp4_block_mse_for_scale_weighted_finite(x, n, mse_weights, e0, best_err) : + rocmfp4_block_mse_for_scale_weighted(x, n, mse_weights, e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = e0; + } + } + } + + const int e1 = start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = all_finite ? + rocmfp4_block_mse_for_scale_weighted_finite(x, n, mse_weights, e1, best_err) : + rocmfp4_block_mse_for_scale_weighted(x, n, mse_weights, e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return (uint8_t) best_e; +} + +static uint8_t rocmfp4_choose_scale_ue4m3(const float * x, int n, const float * quant_weights, float sigma2) { + if (quant_weights) { + assert(n <= QK_ROCMFP4); + float mse_weights_buf[QK_ROCMFP4]; + float weighted_max_abs; + float max_abs_weight; + bool all_finite; + rocmfp4_prepare_mse_weights(mse_weights_buf, x, n, quant_weights, sigma2, &weighted_max_abs, &max_abs_weight, &all_finite); + if (!(weighted_max_abs > 0.0f) || !isfinite(weighted_max_abs)) { + return 0; + } + return rocmfp4_choose_scale_ue4m3_exhaustive_weighted(x, n, mse_weights_buf, weighted_max_abs, max_abs_weight, all_finite); + } + + float max_abs = 0.0f; + bool all_finite = true; + for (int i = 0; i < n; ++i) { + all_finite = all_finite && isfinite(x[i]); + const float ax = fabsf(x[i]); + if (ax > max_abs) { + max_abs = ax; + } + } + + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfp4_choose_scale_ue4m3_exhaustive_unweighted(x, n, max_abs, all_finite); +} + +static void rocmfp4_quantize_row_q4_0_weighted( + const float * GGML_RESTRICT x, block_rocmfp4 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP4 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += x[i]*x[i]; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP4; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP4 : NULL; + const uint8_t e0 = rocmfp4_choose_scale_ue4m3(xb, QK_ROCMFP4/2, qw, sigma2); + const uint8_t e1 = rocmfp4_choose_scale_ue4m3(xb + QK_ROCMFP4/2, QK_ROCMFP4/2, qw ? qw + QK_ROCMFP4/2 : NULL, sigma2); + const float scale_half0 = rocmfp4_ue4m3_to_fp32_half(e0); + const float scale_half1 = rocmfp4_ue4m3_to_fp32_half(e1); + const float inv_scale_half0 = scale_half0 > 0.0f ? 1.0f / scale_half0 : 0.0f; + const float inv_scale_half1 = scale_half1 > 0.0f ? 1.0f / scale_half1 : 0.0f; + + y[ib].e[0] = e0; + y[ib].e[1] = e1; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q0 = rocmfp4_best_index_scaled(xb[j], inv_scale_half0); + const uint8_t q1 = rocmfp4_best_index_scaled(xb[j + QK_ROCMFP4/2], inv_scale_half1); + y[ib].qs[j] = q0 | (q1 << 4); + } + } +} + +static void rocmfp4_quantize_row_q4_0_fast_weighted( + const float * GGML_RESTRICT x, block_rocmfp4_fast * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP4 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += x[i]*x[i]; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP4; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP4 : NULL; + const uint8_t e = rocmfp4_choose_scale_ue4m3(xb, QK_ROCMFP4, qw, sigma2); + const float scale_half = rocmfp4_ue4m3_to_fp32_half(e); + const float inv_scale_half = scale_half > 0.0f ? 1.0f / scale_half : 0.0f; + + y[ib].e = e; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q0 = rocmfp4_best_index_scaled(xb[j], inv_scale_half); + const uint8_t q1 = rocmfp4_best_index_scaled(xb[j + QK_ROCMFP4/2], inv_scale_half); + y[ib].qs[j] = q0 | (q1 << 4); + } + } +} + +void rocmfp4_quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_rocmfp4 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP4 == 0); + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP4; + const uint8_t e0 = rocmfp4_choose_scale_ue4m3(xb, QK_ROCMFP4/2, NULL, 0.0f); + const uint8_t e1 = rocmfp4_choose_scale_ue4m3(xb + QK_ROCMFP4/2, QK_ROCMFP4/2, NULL, 0.0f); + const float scale_half0 = rocmfp4_ue4m3_to_fp32_half(e0); + const float scale_half1 = rocmfp4_ue4m3_to_fp32_half(e1); + const float inv_scale_half0 = scale_half0 > 0.0f ? 1.0f / scale_half0 : 0.0f; + const float inv_scale_half1 = scale_half1 > 0.0f ? 1.0f / scale_half1 : 0.0f; + + y[ib].e[0] = e0; + y[ib].e[1] = e1; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q0 = rocmfp4_best_index_scaled(xb[j], inv_scale_half0); + const uint8_t q1 = rocmfp4_best_index_scaled(xb[j + QK_ROCMFP4/2], inv_scale_half1); + y[ib].qs[j] = q0 | (q1 << 4); + } + } +} + +void rocmfp4_quantize_row_q4_0_fast_ref(const float * GGML_RESTRICT x, block_rocmfp4_fast * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP4 == 0); + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP4; + const uint8_t e = rocmfp4_choose_scale_ue4m3(xb, QK_ROCMFP4, NULL, 0.0f); + const float scale_half = rocmfp4_ue4m3_to_fp32_half(e); + const float inv_scale_half = scale_half > 0.0f ? 1.0f / scale_half : 0.0f; + + y[ib].e = e; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q0 = rocmfp4_best_index_scaled(xb[j], inv_scale_half); + const uint8_t q1 = rocmfp4_best_index_scaled(xb[j + QK_ROCMFP4/2], inv_scale_half); + y[ib].qs[j] = q0 | (q1 << 4); + } + } +} + +void rocmfp4_dequantize_row_q4_0(const block_rocmfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP4 == 0); + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float d0 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[1]); + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + y[ib*QK_ROCMFP4 + j] = (float) rocmfp4_decode(x[ib].qs[j] & 0x0f) * d0; + y[ib*QK_ROCMFP4 + j + QK_ROCMFP4/2] = (float) rocmfp4_decode(x[ib].qs[j] >> 4) * d1; + } + } +} + +void rocmfp4_dequantize_row_q4_0_fast(const block_rocmfp4_fast * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP4 == 0); + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float d = rocmfp4_ue4m3_to_fp32_half(x[ib].e); + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + y[ib*QK_ROCMFP4 + j] = (float) rocmfp4_decode(x[ib].qs[j] & 0x0f) * d; + y[ib*QK_ROCMFP4 + j + QK_ROCMFP4/2] = (float) rocmfp4_decode(x[ib].qs[j] >> 4) * d; + } + } +} + +void rocmfp4_quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfp4_quantize_row_q4_0_ref(x, (block_rocmfp4 *) y, k); +} + +void rocmfp4_quantize_row_q4_0_fast(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfp4_quantize_row_q4_0_fast_ref(x, (block_rocmfp4_fast *) y, k); +} + +size_t rocmfp4_quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = ggml_row_size(GGML_TYPE_Q4_0_ROCMFP4, n_per_row); + + if (!imatrix) { + rocmfp4_quantize_row_q4_0_ref(src, (block_rocmfp4 *) dst, nrows*n_per_row); + return nrows * row_size; + } + + char * qrow = (char *) dst; + for (int64_t row = 0; row < nrows; ++row) { + rocmfp4_quantize_row_q4_0_weighted(src, (block_rocmfp4 *) qrow, n_per_row, imatrix); + src += n_per_row; + qrow += row_size; + } + + return nrows * row_size; +} + +size_t rocmfp4_quantize_q4_0_fast(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = ggml_row_size(GGML_TYPE_Q4_0_ROCMFP4_FAST, n_per_row); + + if (!imatrix) { + rocmfp4_quantize_row_q4_0_fast_ref(src, (block_rocmfp4_fast *) dst, nrows*n_per_row); + return nrows * row_size; + } + + char * qrow = (char *) dst; + for (int64_t row = 0; row < nrows; ++row) { + rocmfp4_quantize_row_q4_0_fast_weighted(src, (block_rocmfp4_fast *) qrow, n_per_row, imatrix); + src += n_per_row; + qrow += row_size; + } + + return nrows * row_size; +} + +bool rocmfp4_validate_row_data(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp4) != 0) { + return false; + } + + const block_rocmfp4 * blocks = (const block_rocmfp4 *) data; + const size_t nblocks = nbytes / sizeof(block_rocmfp4); + for (size_t i = 0; i < nblocks; ++i) { + if (!rocmfp4_scale_is_valid(blocks[i].e[0]) || !rocmfp4_scale_is_valid(blocks[i].e[1])) { + return false; + } + } + + return true; +} + +bool rocmfp4_validate_row_data_fast(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp4_fast) != 0) { + return false; + } + + const block_rocmfp4_fast * blocks = (const block_rocmfp4_fast *) data; + const size_t nblocks = nbytes / sizeof(block_rocmfp4_fast); + for (size_t i = 0; i < nblocks; ++i) { + if (!rocmfp4_scale_is_valid(blocks[i].e)) { + return false; + } + } + + return true; +} + +#ifdef ROCMFP4_X86_AVX2_DISPATCH +__attribute__((target("avx2"))) +static inline int rocmfp4_hsum_i32_8_avx2(__m256i v) { + __m128i s = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256(v, 1)); + s = _mm_add_epi32(s, _mm_shuffle_epi32(s, _MM_SHUFFLE(1, 0, 3, 2))); + s = _mm_add_epi32(s, _mm_shuffle_epi32(s, _MM_SHUFFLE(2, 3, 0, 1))); + return _mm_cvtsi128_si32(s); +} + +// Decode one 32-weight block's low and high nibble streams through the +// Codebook10 table with a single PSHUFB, then integer-dot each against its half +// of the q8_0 block. Integer sums are order-independent, so sumi0/sumi1 match +// the scalar reference exactly and the float result is bit-identical. +__attribute__((target("avx2"))) +static inline void rocmfp4_block_isums_avx2( + const uint8_t * qs, const int8_t * q8, int * sumi0, int * sumi1) { + const __m128i tbl = _mm_loadu_si128((const __m128i *) rocmfp4_codebook); + const __m128i q = _mm_loadu_si128((const __m128i *) qs); + const __m128i lo = _mm_and_si128(q, _mm_set1_epi8(0x0F)); + const __m128i hi = _mm_and_si128(_mm_srli_epi16(q, 4), _mm_set1_epi8(0x0F)); + const __m128i dlo = _mm_shuffle_epi8(tbl, lo); + const __m128i dhi = _mm_shuffle_epi8(tbl, hi); + const __m128i ylo = _mm_loadu_si128((const __m128i *) q8); + const __m128i yhi = _mm_loadu_si128((const __m128i *) (q8 + QK_ROCMFP4/2)); + const __m256i pl = _mm256_madd_epi16(_mm256_cvtepi8_epi16(dlo), _mm256_cvtepi8_epi16(ylo)); + const __m256i ph = _mm256_madd_epi16(_mm256_cvtepi8_epi16(dhi), _mm256_cvtepi8_epi16(yhi)); + *sumi0 = rocmfp4_hsum_i32_8_avx2(pl); + *sumi1 = rocmfp4_hsum_i32_8_avx2(ph); +} + +__attribute__((target("avx2"))) +static void rocmfp4_vec_dot_q4_0_q8_0_avx2( + int nb, float * GGML_RESTRICT s, const block_rocmfp4 * GGML_RESTRICT x, const block_q8_0 * GGML_RESTRICT y) { + float sumf = 0.0f; + for (int ib = 0; ib < nb; ++ib) { + const float d0 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[0]) * ggml_fp16_to_fp32(y[ib].d); + const float d1 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[1]) * ggml_fp16_to_fp32(y[ib].d); + int sumi0, sumi1; + rocmfp4_block_isums_avx2(x[ib].qs, y[ib].qs, &sumi0, &sumi1); + sumf += d0 * (float) sumi0 + d1 * (float) sumi1; + } + *s = sumf; +} + +__attribute__((target("avx2"))) +static void rocmfp4_vec_dot_q4_0_fast_q8_0_avx2( + int nb, float * GGML_RESTRICT s, const block_rocmfp4_fast * GGML_RESTRICT x, const block_q8_0 * GGML_RESTRICT y) { + float sumf = 0.0f; + for (int ib = 0; ib < nb; ++ib) { + const float d = rocmfp4_ue4m3_to_fp32_half(x[ib].e) * ggml_fp16_to_fp32(y[ib].d); + int sumi0, sumi1; + rocmfp4_block_isums_avx2(x[ib].qs, y[ib].qs, &sumi0, &sumi1); + sumf += d * (float) (sumi0 + sumi1); + } + *s = sumf; +} +#endif // ROCMFP4_X86_AVX2_DISPATCH + +void rocmfp4_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + assert(nrc == 1); + GGML_UNUSED(nrc); + assert(n % QK_ROCMFP4 == 0); + assert(QK_ROCMFP4 == QK8_0); + + const block_rocmfp4 * GGML_RESTRICT x = (const block_rocmfp4 *) vx; + const block_q8_0 * GGML_RESTRICT y = (const block_q8_0 *) vy; + + const int nb = n / QK_ROCMFP4; + +#ifdef ROCMFP4_X86_AVX2_DISPATCH + if (__builtin_cpu_supports("avx2")) { + rocmfp4_vec_dot_q4_0_q8_0_avx2(nb, s, x, y); + return; + } +#endif + + float sumf = 0.0f; + + for (int ib = 0; ib < nb; ++ib) { + const float d0 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[0]) * ggml_fp16_to_fp32(y[ib].d); + const float d1 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[1]) * ggml_fp16_to_fp32(y[ib].d); + int sumi0 = 0; + int sumi1 = 0; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q = x[ib].qs[j]; + sumi0 += rocmfp4_decode_table(q) * y[ib].qs[j]; + sumi1 += rocmfp4_decode_table(q >> 4) * y[ib].qs[j + QK_ROCMFP4/2]; + } + + sumf += d0 * (float) sumi0 + d1 * (float) sumi1; + } + + *s = sumf; +} + +void rocmfp4_vec_dot_q4_0_fast_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + assert(nrc == 1); + GGML_UNUSED(nrc); + assert(n % QK_ROCMFP4 == 0); + assert(QK_ROCMFP4 == QK8_0); + + const block_rocmfp4_fast * GGML_RESTRICT x = (const block_rocmfp4_fast *) vx; + const block_q8_0 * GGML_RESTRICT y = (const block_q8_0 *) vy; + + const int nb = n / QK_ROCMFP4; + +#ifdef ROCMFP4_X86_AVX2_DISPATCH + if (__builtin_cpu_supports("avx2")) { + rocmfp4_vec_dot_q4_0_fast_q8_0_avx2(nb, s, x, y); + return; + } +#endif + + float sumf = 0.0f; + + for (int ib = 0; ib < nb; ++ib) { + const float d = rocmfp4_ue4m3_to_fp32_half(x[ib].e) * ggml_fp16_to_fp32(y[ib].d); + int sumi = 0; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q = x[ib].qs[j]; + sumi += rocmfp4_decode_table(q) * y[ib].qs[j]; + sumi += rocmfp4_decode_table(q >> 4) * y[ib].qs[j + QK_ROCMFP4/2]; + } + + sumf += d * (float) sumi; + } + + *s = sumf; +} diff --git a/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4.h b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4.h new file mode 100644 index 000000000..9756f6ad4 --- /dev/null +++ b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4.h @@ -0,0 +1,58 @@ +#pragma once + +#include +#include +#include + +#include "ggml.h" + +#ifdef __cplusplus +extern "C" { +#endif + +#define QK_ROCMFP4 32 +#define QR_ROCMFP4 2 +#define QI_ROCMFP4 (QK_ROCMFP4 / (4 * QR_ROCMFP4)) +#define QS_ROCMFP4 32 + +// AMD-tuned compact layout: 16 bytes of packed E2M1-derived 4-bit codes, then +// one unsigned E4M3 scale byte per 16-weight half block. +typedef struct { + uint8_t qs[QK_ROCMFP4/2]; + uint8_t e[2]; +} block_rocmfp4; + +// Speed-focused layout: same 32 packed ROCmFP4 nibbles, but one UE4M3 scale +// for the whole block. This is a separate GGUF type so fast 4.25 BPW artifacts +// never alias the safer dual-scale format above. +typedef struct { + uint8_t qs[QK_ROCMFP4/2]; + uint8_t e; +} block_rocmfp4_fast; + +#if defined(__cplusplus) +static_assert(sizeof(block_rocmfp4) == QK_ROCMFP4/2 + 2*sizeof(uint8_t), "wrong rocmfp4 block size/padding"); +static_assert(sizeof(block_rocmfp4_fast) == QK_ROCMFP4/2 + sizeof(uint8_t), "wrong rocmfp4 fast block size/padding"); +#else +_Static_assert(sizeof(block_rocmfp4) == QK_ROCMFP4/2 + 2*sizeof(uint8_t), "wrong rocmfp4 block size/padding"); +_Static_assert(sizeof(block_rocmfp4_fast) == QK_ROCMFP4/2 + sizeof(uint8_t), "wrong rocmfp4 fast block size/padding"); +#endif + +GGML_API void rocmfp4_quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_rocmfp4 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfp4_dequantize_row_q4_0(const block_rocmfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfp4_quantize_row_q4_0_fast_ref(const float * GGML_RESTRICT x, block_rocmfp4_fast * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfp4_dequantize_row_q4_0_fast(const block_rocmfp4_fast * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + +GGML_API void rocmfp4_quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfp4_quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API void rocmfp4_quantize_row_q4_0_fast(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfp4_quantize_q4_0_fast(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API bool rocmfp4_validate_row_data(const void * data, size_t nbytes); +GGML_API bool rocmfp4_validate_row_data_fast(const void * data, size_t nbytes); + +GGML_API void rocmfp4_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +GGML_API void rocmfp4_vec_dot_q4_0_fast_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +#ifdef __cplusplus +} +#endif diff --git a/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip.cu b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip.cu new file mode 100644 index 000000000..d2c9048c1 --- /dev/null +++ b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip.cu @@ -0,0 +1,85 @@ +#include "rocmfp4.h" + +#include + +#include "rocmfp4_hip_scale.cuh" + +// Standalone ROCm/HIP dequant kernel for integration tests and future fused +// paths. One lane owns one packed byte and writes the matching low/high +// half-block values, so each byte is read once. +extern "C" __global__ void rocmfp4_dequantize_q4_0_f32_kernel( + const block_rocmfp4 * __restrict__ x, + float * __restrict__ y, + int64_t k) { + const int64_t packed_idx = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; + const int64_t nblocks = (k + QK_ROCMFP4 - 1) / QK_ROCMFP4; + const int64_t packed_count = nblocks * (QK_ROCMFP4/2); + + if (packed_idx >= packed_count) { + return; + } + + const int64_t ib = packed_idx / (QK_ROCMFP4/2); + const int tid = packed_idx - ib*(QK_ROCMFP4/2); + const int64_t base = ib*QK_ROCMFP4; + const uint8_t packed = x[ib].qs[tid]; + const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[1]); + + if (base + tid < k) { + y[base + tid] = (float) rocmfp4_decode_i8(packed & 0x0f) * d0; + } + if (base + tid + QK_ROCMFP4/2 < k) { + y[base + tid + QK_ROCMFP4/2] = (float) rocmfp4_decode_i8(packed >> 4) * d1; + } +} + +extern "C" __global__ void rocmfp4_dequantize_q4_0_fast_f32_kernel( + const block_rocmfp4_fast * __restrict__ x, + float * __restrict__ y, + int64_t k) { + const int64_t packed_idx = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; + const int64_t nblocks = (k + QK_ROCMFP4 - 1) / QK_ROCMFP4; + const int64_t packed_count = nblocks * (QK_ROCMFP4/2); + + if (packed_idx >= packed_count) { + return; + } + + const int64_t ib = packed_idx / (QK_ROCMFP4/2); + const int tid = packed_idx - ib*(QK_ROCMFP4/2); + const int64_t base = ib*QK_ROCMFP4; + const uint8_t packed = x[ib].qs[tid]; + const float d = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e); + + if (base + tid < k) { + y[base + tid] = (float) rocmfp4_decode_i8(packed & 0x0f) * d; + } + if (base + tid + QK_ROCMFP4/2 < k) { + y[base + tid + QK_ROCMFP4/2] = (float) rocmfp4_decode_i8(packed >> 4) * d; + } +} + +extern "C" void rocmfp4_hip_dequantize_q4_0_to_f32( + const void * src, + float * dst, + int64_t k, + hipStream_t stream) { + const int64_t nblocks = (k + QK_ROCMFP4 - 1) / QK_ROCMFP4; + const int64_t packed_count = nblocks * (QK_ROCMFP4/2); + const dim3 block(256); + const dim3 grid((unsigned int) ((packed_count + block.x - 1) / block.x)); + rocmfp4_dequantize_q4_0_f32_kernel<<>>((const block_rocmfp4 *) src, dst, k); +} + +extern "C" void rocmfp4_hip_dequantize_q4_0_fast_to_f32( + const void * src, + float * dst, + int64_t k, + hipStream_t stream) { + const int64_t nblocks = (k + QK_ROCMFP4 - 1) / QK_ROCMFP4; + const int64_t packed_count = nblocks * (QK_ROCMFP4/2); + const dim3 block(256); + const dim3 grid((unsigned int) ((packed_count + block.x - 1) / block.x)); + rocmfp4_dequantize_q4_0_fast_f32_kernel<<>>((const block_rocmfp4_fast *) src, dst, k); +} diff --git a/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_codebook.cuh b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_codebook.cuh new file mode 100644 index 000000000..1d9106112 --- /dev/null +++ b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_codebook.cuh @@ -0,0 +1,79 @@ +#pragma once + +#include "rocmfp4_hip_scale.cuh" + +#include +#include + +#ifndef GGML_ROCMFP4_UNALIGNED_QS_DWORD_LOAD +#define GGML_ROCMFP4_UNALIGNED_QS_DWORD_LOAD 1 +#endif + +static __device__ __forceinline__ int rocmfp4_get_qs_i32(const void * x, const int & i32) { +#if defined(GGML_USE_HIP) && GGML_ROCMFP4_UNALIGNED_QS_DWORD_LOAD + return *((const int *) ((const uint8_t *) x + 4*i32)); +#else + const uint8_t * x8 = (const uint8_t *) x; + + int x32 = x8[4*i32 + 0] << 0; + x32 |= x8[4*i32 + 1] << 8; + x32 |= x8[4*i32 + 2] << 16; + x32 |= x8[4*i32 + 3] << 24; + + return x32; +#endif +} + +// AMD-specific fast path for expanding eight packed ROCmFP4 nibbles into two +// int32 DP4A operands. This encodes the Codebook10 table directly as four +// 32-bit constants: +// [0, 1, 2, 3], [4, 6, 8, 10], [0, -1, -2, -3], [-4, -6, -8, -10] +// Avoiding the table pointer keeps the ROCm/HIP MMVQ/MMQ hot path fully local +// to this format. Non-HIP builds still use llama.cpp's generic table expander. +static __device__ __forceinline__ int2 rocmfp4_get_int_from_codebook_16(const int & q4, const int8_t * fallback_table) { +#if defined(GGML_USE_HIP) + constexpr uint32_t values0 = 0x03020100u; + constexpr uint32_t values1 = 0x0a080604u; + constexpr uint32_t values2 = 0xfdfeff00u; + constexpr uint32_t values3 = 0xf6f8fafcu; + + const uint32_t q_even = q4; + const uint32_t q_odd = q4 >> 4; + + const uint32_t v_even_low = __builtin_amdgcn_perm(values1, values0, q_even & 0x07070707u); + const uint32_t v_odd_low = __builtin_amdgcn_perm(values1, values0, q_odd & 0x07070707u); + const uint32_t v_even_high = __builtin_amdgcn_perm(values3, values2, q_even & 0x07070707u); + const uint32_t v_odd_high = __builtin_amdgcn_perm(values3, values2, q_odd & 0x07070707u); + + const uint32_t mask_even = 0x03020100u | ((q_even & 0x08080808u) >> 1); + const uint32_t mask_odd = 0x03020100u | ((q_odd & 0x08080808u) >> 1); + + return make_int2( + __builtin_amdgcn_perm(v_even_high, v_even_low, mask_even), + __builtin_amdgcn_perm(v_odd_high, v_odd_low, mask_odd)); +#else + return get_int_from_table_16(q4, fallback_table); +#endif +} + +// Variant for call sites that already selected either the low or high nibble +// stream and only need one DP4A operand. This avoids the extra odd/even table +// expansion work in ROCmFP4 FlashAttention K/V decode. +static __device__ __forceinline__ int rocmfp4_get_low_int_from_codebook_16(const int & q4, const int8_t * fallback_table) { +#if defined(GGML_USE_HIP) + constexpr uint32_t values0 = 0x03020100u; + constexpr uint32_t values1 = 0x0a080604u; + constexpr uint32_t values2 = 0xfdfeff00u; + constexpr uint32_t values3 = 0xf6f8fafcu; + + const uint32_t q = q4; + + const uint32_t v_low = __builtin_amdgcn_perm(values1, values0, q & 0x07070707u); + const uint32_t v_high = __builtin_amdgcn_perm(values3, values2, q & 0x07070707u); + const uint32_t mask = 0x03020100u | ((q & 0x08080808u) >> 1); + + return __builtin_amdgcn_perm(v_high, v_low, mask); +#else + return get_int_from_table_16(q4, fallback_table).x; +#endif +} diff --git a/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_scale.cuh b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_scale.cuh new file mode 100644 index 000000000..19525a4fb --- /dev/null +++ b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_scale.cuh @@ -0,0 +1,157 @@ +#pragma once + +#include +#include +#include + +#ifndef GGML_ROCMFP4_USE_SCALE_LUT +#define GGML_ROCMFP4_USE_SCALE_LUT 0 +#endif + +#if defined(GGML_USE_HIP) && GGML_ROCMFP4_USE_SCALE_LUT +#define ROCMFP4_SCALE_SUB(M) ((M) * 0x1p-10f) +#define ROCMFP4_SCALE_E1(M) ((8 + (M)) * 0x1p-10f) +#define ROCMFP4_SCALE_E2(M) ((8 + (M)) * 0x1p-9f) +#define ROCMFP4_SCALE_E3(M) ((8 + (M)) * 0x1p-8f) +#define ROCMFP4_SCALE_E4(M) ((8 + (M)) * 0x1p-7f) +#define ROCMFP4_SCALE_E5(M) ((8 + (M)) * 0x1p-6f) +#define ROCMFP4_SCALE_E6(M) ((8 + (M)) * 0x1p-5f) +#define ROCMFP4_SCALE_E7(M) ((8 + (M)) * 0x1p-4f) +#define ROCMFP4_SCALE_E8(M) ((8 + (M)) * 0x1p-3f) +#define ROCMFP4_SCALE_E9(M) ((8 + (M)) * 0x1p-2f) +#define ROCMFP4_SCALE_E10(M) ((8 + (M)) * 0x1p-1f) +#define ROCMFP4_SCALE_E11(M) ((8 + (M)) * 0x1p0f) +#define ROCMFP4_SCALE_E12(M) ((8 + (M)) * 0x1p1f) +#define ROCMFP4_SCALE_E13(M) ((8 + (M)) * 0x1p2f) +#define ROCMFP4_SCALE_E14(M) ((8 + (M)) * 0x1p3f) +#define ROCMFP4_SCALE_E15(M) ((8 + (M)) * 0x1p4f) + +static __device__ __constant__ const float rocmfp4_scale_ue4m3_half_lut[127] = { + ROCMFP4_SCALE_SUB(0), ROCMFP4_SCALE_SUB(1), ROCMFP4_SCALE_SUB(2), ROCMFP4_SCALE_SUB(3), + ROCMFP4_SCALE_SUB(4), ROCMFP4_SCALE_SUB(5), ROCMFP4_SCALE_SUB(6), ROCMFP4_SCALE_SUB(7), + ROCMFP4_SCALE_E1(0), ROCMFP4_SCALE_E1(1), ROCMFP4_SCALE_E1(2), ROCMFP4_SCALE_E1(3), + ROCMFP4_SCALE_E1(4), ROCMFP4_SCALE_E1(5), ROCMFP4_SCALE_E1(6), ROCMFP4_SCALE_E1(7), + ROCMFP4_SCALE_E2(0), ROCMFP4_SCALE_E2(1), ROCMFP4_SCALE_E2(2), ROCMFP4_SCALE_E2(3), + ROCMFP4_SCALE_E2(4), ROCMFP4_SCALE_E2(5), ROCMFP4_SCALE_E2(6), ROCMFP4_SCALE_E2(7), + ROCMFP4_SCALE_E3(0), ROCMFP4_SCALE_E3(1), ROCMFP4_SCALE_E3(2), ROCMFP4_SCALE_E3(3), + ROCMFP4_SCALE_E3(4), ROCMFP4_SCALE_E3(5), ROCMFP4_SCALE_E3(6), ROCMFP4_SCALE_E3(7), + ROCMFP4_SCALE_E4(0), ROCMFP4_SCALE_E4(1), ROCMFP4_SCALE_E4(2), ROCMFP4_SCALE_E4(3), + ROCMFP4_SCALE_E4(4), ROCMFP4_SCALE_E4(5), ROCMFP4_SCALE_E4(6), ROCMFP4_SCALE_E4(7), + ROCMFP4_SCALE_E5(0), ROCMFP4_SCALE_E5(1), ROCMFP4_SCALE_E5(2), ROCMFP4_SCALE_E5(3), + ROCMFP4_SCALE_E5(4), ROCMFP4_SCALE_E5(5), ROCMFP4_SCALE_E5(6), ROCMFP4_SCALE_E5(7), + ROCMFP4_SCALE_E6(0), ROCMFP4_SCALE_E6(1), ROCMFP4_SCALE_E6(2), ROCMFP4_SCALE_E6(3), + ROCMFP4_SCALE_E6(4), ROCMFP4_SCALE_E6(5), ROCMFP4_SCALE_E6(6), ROCMFP4_SCALE_E6(7), + ROCMFP4_SCALE_E7(0), ROCMFP4_SCALE_E7(1), ROCMFP4_SCALE_E7(2), ROCMFP4_SCALE_E7(3), + ROCMFP4_SCALE_E7(4), ROCMFP4_SCALE_E7(5), ROCMFP4_SCALE_E7(6), ROCMFP4_SCALE_E7(7), + ROCMFP4_SCALE_E8(0), ROCMFP4_SCALE_E8(1), ROCMFP4_SCALE_E8(2), ROCMFP4_SCALE_E8(3), + ROCMFP4_SCALE_E8(4), ROCMFP4_SCALE_E8(5), ROCMFP4_SCALE_E8(6), ROCMFP4_SCALE_E8(7), + ROCMFP4_SCALE_E9(0), ROCMFP4_SCALE_E9(1), ROCMFP4_SCALE_E9(2), ROCMFP4_SCALE_E9(3), + ROCMFP4_SCALE_E9(4), ROCMFP4_SCALE_E9(5), ROCMFP4_SCALE_E9(6), ROCMFP4_SCALE_E9(7), + ROCMFP4_SCALE_E10(0), ROCMFP4_SCALE_E10(1), ROCMFP4_SCALE_E10(2), ROCMFP4_SCALE_E10(3), + ROCMFP4_SCALE_E10(4), ROCMFP4_SCALE_E10(5), ROCMFP4_SCALE_E10(6), ROCMFP4_SCALE_E10(7), + ROCMFP4_SCALE_E11(0), ROCMFP4_SCALE_E11(1), ROCMFP4_SCALE_E11(2), ROCMFP4_SCALE_E11(3), + ROCMFP4_SCALE_E11(4), ROCMFP4_SCALE_E11(5), ROCMFP4_SCALE_E11(6), ROCMFP4_SCALE_E11(7), + ROCMFP4_SCALE_E12(0), ROCMFP4_SCALE_E12(1), ROCMFP4_SCALE_E12(2), ROCMFP4_SCALE_E12(3), + ROCMFP4_SCALE_E12(4), ROCMFP4_SCALE_E12(5), ROCMFP4_SCALE_E12(6), ROCMFP4_SCALE_E12(7), + ROCMFP4_SCALE_E13(0), ROCMFP4_SCALE_E13(1), ROCMFP4_SCALE_E13(2), ROCMFP4_SCALE_E13(3), + ROCMFP4_SCALE_E13(4), ROCMFP4_SCALE_E13(5), ROCMFP4_SCALE_E13(6), ROCMFP4_SCALE_E13(7), + ROCMFP4_SCALE_E14(0), ROCMFP4_SCALE_E14(1), ROCMFP4_SCALE_E14(2), ROCMFP4_SCALE_E14(3), + ROCMFP4_SCALE_E14(4), ROCMFP4_SCALE_E14(5), ROCMFP4_SCALE_E14(6), ROCMFP4_SCALE_E14(7), + ROCMFP4_SCALE_E15(0), ROCMFP4_SCALE_E15(1), ROCMFP4_SCALE_E15(2), ROCMFP4_SCALE_E15(3), + ROCMFP4_SCALE_E15(4), ROCMFP4_SCALE_E15(5), ROCMFP4_SCALE_E15(6), +}; + +#undef ROCMFP4_SCALE_SUB +#undef ROCMFP4_SCALE_E1 +#undef ROCMFP4_SCALE_E2 +#undef ROCMFP4_SCALE_E3 +#undef ROCMFP4_SCALE_E4 +#undef ROCMFP4_SCALE_E5 +#undef ROCMFP4_SCALE_E6 +#undef ROCMFP4_SCALE_E7 +#undef ROCMFP4_SCALE_E8 +#undef ROCMFP4_SCALE_E9 +#undef ROCMFP4_SCALE_E10 +#undef ROCMFP4_SCALE_E11 +#undef ROCMFP4_SCALE_E12 +#undef ROCMFP4_SCALE_E13 +#undef ROCMFP4_SCALE_E14 +#undef ROCMFP4_SCALE_E15 +#endif + +static __device__ __forceinline__ float rocmfp4_u32_as_f32(uint32_t bits) { +#if defined(GGML_USE_HIP) + return __uint_as_float(bits); +#else + float result; + memcpy(&result, &bits, sizeof(float)); + return result; +#endif +} + +// ROCmFP4 validates scale bytes before backend execution, so HIP/ROCm hot +// paths can decode finite unsigned E4M3 half-scales directly without the +// generic FP8 NaN handling used by other formats. +static __device__ __forceinline__ float rocmfp4_ue4m3_to_fp32_half_finite(uint8_t x) { +#if defined(GGML_USE_HIP) && GGML_ROCMFP4_USE_SCALE_LUT + return x <= 0x7e ? rocmfp4_scale_ue4m3_half_lut[x] : 0.0f; +#else + const int exp = (x >> 3) & 0xF; + const int man = x & 0x7; + + if (exp == 0) { + return (float) man * (1.0f / 1024.0f); + } + + const uint32_t bits = ((uint32_t) exp + 119u) << 23 | ((uint32_t) man << 20); + return rocmfp4_u32_as_f32(bits); +#endif +} + +static __device__ __forceinline__ float rocmfpx_ue4m3_to_fp32_finite(uint8_t x) { + if (x > 0x7e) { + return 0.0f; + } + + const int exp = (x >> 3) & 0xF; + const int man = x & 0x7; + + if (exp == 0) { + return (float) man * (1.0f / 1024.0f); + } + + const uint32_t bits = ((uint32_t) exp + 119u) << 23 | ((uint32_t) man << 20); + return rocmfp4_u32_as_f32(bits); +} + +static __device__ __forceinline__ uint8_t rocmfpx_nearest_scale_ue4m3_cuda(float target_scale) { + if (!(target_scale > 0.0f) || !isfinite(target_scale)) { + return 0; + } + + uint8_t lo = 1; + uint8_t hi = 0x7e; + while (lo < hi) { + const uint8_t mid = lo + (hi - lo) / 2; + if (rocmfpx_ue4m3_to_fp32_finite(mid) < target_scale) { + lo = mid + 1; + } else { + hi = mid; + } + } + + if (lo == 1) { + return 1; + } + + const float hi_scale = rocmfpx_ue4m3_to_fp32_finite(lo); + const float lo_scale = rocmfpx_ue4m3_to_fp32_finite((uint8_t) (lo - 1)); + return (target_scale - lo_scale <= hi_scale - target_scale) ? (uint8_t) (lo - 1) : lo; +} + +static __device__ __forceinline__ int8_t rocmfp4_decode_i8(uint8_t q) { + q &= 0x0f; + const int mag3 = q & 0x07; + const int mag = mag3 <= 4 ? mag3 : 2*mag3 - 4; + return (q & 0x08) ? -mag : mag; +} diff --git a/server/deps/llama.cpp/ggml/rocmfpx/README.md b/server/deps/llama.cpp/ggml/rocmfpx/README.md new file mode 100644 index 000000000..a92eb4658 --- /dev/null +++ b/server/deps/llama.cpp/ggml/rocmfpx/README.md @@ -0,0 +1,106 @@ +# ROCmFPx Experimental Formats + +This folder contains the reference layer for the proposed ROCmFP3, ROCmFP6, and +ROCmFP8 quantization family. It is intentionally separate from `ggml/rocmfp4/` +so the promoted ROCmFP4 GGUF formats and kernels are not affected while the new +layouts are evaluated. + +## ROCmFP4 Instructions To Preserve + +ROCmFPX is a sibling model-weight quant family, not a new K/V-only compression +scheme. The promoted ROCmFP4 implementation remains the template for how the +family should behave in llama.cpp: + +- Keep 32-weight blocks so CPU, HIP, and Vulkan kernels can reuse the same + Q4/Q8-style reduction shape and GGUF row-size assumptions. +- Use finite unsigned UE4M3 scale bytes only. `0x7f` and sign-bit scale bytes + are invalid, matching the ROCmFP4 validation rule. +- Prefer reconstruction-MSE scale selection over plain max-abs scaling. ROCmFP4 + searches each 16-weight half-block; ROCmFP3 and ROCmFP6 follow the same + half-block policy, while ROCmFP8 currently uses one full-block scale. +- Preserve the ROCmFP4 kernel contract: CPU reference quant/dequant/dot first, + then HIP/Vulkan `CPY`, `GET_ROWS`, `SET_ROWS`, `MUL_MAT`, and `MUL_MAT_ID` + paths, with backend-op coverage before claiming runtime support. +- Keep dequant math explicit and deterministic: integer code times decoded + UE4M3 scale. ROCmFP4 uses the Codebook10 half-scale table; ROCmFPX formats + use their own integer code ranges but must retain the same finite-scale and + integer-dot discipline. + +The ROCmFP4 Codebook10 levels are not reused by FP3/FP6/FP8 directly: +`ROCmFP3` uses `0, +/-1, +/-2, +/-4`, `ROCmFP6` uses signed-magnitude levels up +to `31`, and `ROCmFP8` uses signed int8 levels clamped to `[-127, 127]`. +What is inherited is the block/scale/kernel/dequant contract. + +Current status (June 16, 2026): +- CPU reference quantize/dequantize exists for all three formats. +- `Q3_0_ROCMFPX`, `Q6_0_ROCMFPX`, and `Q8_0_ROCMFPX` are registered as + experimental GGUF tensor types. +- ROCm/HIP and Vulkan kernels support `CPY`, `GET_ROWS`, `SET_ROWS`, and + `MUL_MAT`/`MUL_MAT_ID` for all three formats. +- Qwen3-0.6B BF16 smoke tests pass on CPU, ROCm0, and Vulkan0. +- Default quant presets now include lean coherency routing: + - `Q3_0_ROCMFPX`: selective `Q5_K` on attention Q/O and early K/V, boosted + FFN-down at `Q5_K`, selective FFN-gate at `Q6_0_ROCMFPX`, bulk FFN-up on + `Q3_0_ROCMFPX`, embeddings/output at `Q4_0_ROCMFP4_FAST`. + - `Q6_0_ROCMFPX`: early attention and boosted FFN-down at `Q8_0_ROCMFPX`, + embeddings/output at `Q6_0_ROCMFPX`, bulk gate/up on `Q6_0_ROCMFPX`. + - `Q8_0_ROCMFPX`: pure FP8-family preset. +- Opt-in `*_AGENT` presets boost attention/FFN routing for tool-call / + Hermes / OpenClaw style workloads: + - `Q3_0_ROCMFPX_AGENT`, `Q6_0_ROCMFPX_AGENT`, `Q8_0_ROCMFPX_AGENT`. + - Routing is layered on top of LEAN; default presets are unchanged. +- FP3 and FP6 quantization use reconstruction-MSE scale selection per + 16-weight half-block. + +## Validation Script Index + +```text +scripts/check-rocmfpx-reference.sh # CPU reference math +scripts/check-rocmfpx-qwen-all.sh # core Qwen gates +scripts/check-rocmfpx-all.sh # qwen-all + optional smokes +scripts/check-rocmfpx-summary.sh # full JSON summary runner +scripts/sweep-rocmfpx-backend-ops.sh # test-backend-ops per backend +scripts/sweep-rocmfpx-agent-size-table.sh # LEAN vs AGENT MiB/BPW +scripts/sweep-rocmfpx-perplexity.sh # calibration PPL sweep +scripts/sweep-rocmfpx-decode-tune.sh # decode-tune matrix +scripts/build-rocmfpx-agent-fixtures.sh # proxy Hermes/OpenClaw AGENT GGUFs +``` + +## Layouts + +All formats use 32-weight blocks. + +| Format | Payload | Scale bytes | Block bytes | BPW | Purpose | +|---|---:|---:|---:|---:|---| +| `Q3_0_ROCMFPX` | 32 packed 3-bit codes | 2, one per 16 weights | 14 | 3.50 | Experimental low-bit candidate | +| `Q6_0_ROCMFPX` | 32 packed 6-bit codes | 2, one per 16 weights | 26 | 6.50 | Experimental quality candidate | +| `Q8_0_ROCMFPX` | 32 signed 8-bit codes | 1, one per 32 weights | 33 | 8.25 | Experimental high-quality reference | + +`ROCmFP3` uses a tiny signed codebook: `0, +/-1, +/-2, +/-4`. +`ROCmFP6` uses signed-magnitude integer levels up to `31`. +`ROCmFP8` uses signed int8 levels clamped to `[-127, 127]`. + +## Validation + +Reference math only: + +```bash +scripts/check-rocmfpx-reference.sh +``` + +Focused backend sweep from the experiment worktree: + +```bash +cmake --build build-strix-rocmfp4 --target test-backend-ops -j 8 +timeout 120 build-strix-rocmfp4/bin/test-backend-ops test -o MUL_MAT,GET_ROWS,CPY,SET_ROWS -b CPU +timeout 180 build-strix-rocmfp4/bin/test-backend-ops test -o MUL_MAT,GET_ROWS,CPY,SET_ROWS -b ROCm0 +timeout 180 build-strix-rocmfp4/bin/test-backend-ops test -o MUL_MAT,GET_ROWS,CPY,SET_ROWS -b Vulkan0 +``` + +Qwen3 BF16 coherency and decode-speed gates: + +```bash +MODEL=/home/caf/strix-fp4/models/rocmfpx-bf16-tests/Qwen3-0.6B-Q3_0_ROCMFPX_COHERENT-LEAN.gguf BACKEND=ROCm0 scripts/check-rocmfpx-qwen-coherency.sh +MODEL=/home/caf/strix-fp4/models/rocmfpx-bf16-tests/Qwen3-0.6B-Q3_0_ROCMFPX_COHERENT-LEAN.gguf BACKEND=ROCm0 scripts/check-rocmfpx-qwen-bench.sh +MODEL=/home/caf/strix-fp4/models/rocmfpx-bf16-tests/Qwen3-0.6B-Q3_0_ROCMFPX_COHERENT-LEAN.gguf BACKEND=ROCm0 scripts/check-rocmfpx-qwen-strict-json.sh +``` diff --git a/server/deps/llama.cpp/ggml/rocmfpx/rocmfpx.c b/server/deps/llama.cpp/ggml/rocmfpx/rocmfpx.c new file mode 100644 index 000000000..65dbca726 --- /dev/null +++ b/server/deps/llama.cpp/ggml/rocmfpx/rocmfpx.c @@ -0,0 +1,1140 @@ +#include "rocmfpx.h" + +#include +#include +#include + +// Finite unsigned E4M3 scale bytes decoded to FP32. Precomputed from the same +// exp/mant formula rocmfpx_ue4m3_to_fp32() used to evaluate with ldexpf(): +// exp == 0 -> mant * 2^-10 ; otherwise (8 + mant) * 2^(exp - 11). +// The scale search re-decodes candidate bytes for every block, and dequant +// decodes a scale for every element, so keeping this as a table (identical to +// the former per-call ldexpf result) removes the transcendental from both hot +// paths without changing any produced value. +#define ROCMFPX_SCALE_SUB(M) ((M) * 0x1p-10f) +#define ROCMFPX_SCALE_E(B, M) ((8 + (M)) * (B)) + +static const float rocmfpx_scale_ue4m3[127] = { + ROCMFPX_SCALE_SUB(0), ROCMFPX_SCALE_SUB(1), ROCMFPX_SCALE_SUB(2), ROCMFPX_SCALE_SUB(3), + ROCMFPX_SCALE_SUB(4), ROCMFPX_SCALE_SUB(5), ROCMFPX_SCALE_SUB(6), ROCMFPX_SCALE_SUB(7), + ROCMFPX_SCALE_E(0x1p-10f,0), ROCMFPX_SCALE_E(0x1p-10f,1), ROCMFPX_SCALE_E(0x1p-10f,2), ROCMFPX_SCALE_E(0x1p-10f,3), + ROCMFPX_SCALE_E(0x1p-10f,4), ROCMFPX_SCALE_E(0x1p-10f,5), ROCMFPX_SCALE_E(0x1p-10f,6), ROCMFPX_SCALE_E(0x1p-10f,7), + ROCMFPX_SCALE_E(0x1p-9f,0), ROCMFPX_SCALE_E(0x1p-9f,1), ROCMFPX_SCALE_E(0x1p-9f,2), ROCMFPX_SCALE_E(0x1p-9f,3), + ROCMFPX_SCALE_E(0x1p-9f,4), ROCMFPX_SCALE_E(0x1p-9f,5), ROCMFPX_SCALE_E(0x1p-9f,6), ROCMFPX_SCALE_E(0x1p-9f,7), + ROCMFPX_SCALE_E(0x1p-8f,0), ROCMFPX_SCALE_E(0x1p-8f,1), ROCMFPX_SCALE_E(0x1p-8f,2), ROCMFPX_SCALE_E(0x1p-8f,3), + ROCMFPX_SCALE_E(0x1p-8f,4), ROCMFPX_SCALE_E(0x1p-8f,5), ROCMFPX_SCALE_E(0x1p-8f,6), ROCMFPX_SCALE_E(0x1p-8f,7), + ROCMFPX_SCALE_E(0x1p-7f,0), ROCMFPX_SCALE_E(0x1p-7f,1), ROCMFPX_SCALE_E(0x1p-7f,2), ROCMFPX_SCALE_E(0x1p-7f,3), + ROCMFPX_SCALE_E(0x1p-7f,4), ROCMFPX_SCALE_E(0x1p-7f,5), ROCMFPX_SCALE_E(0x1p-7f,6), ROCMFPX_SCALE_E(0x1p-7f,7), + ROCMFPX_SCALE_E(0x1p-6f,0), ROCMFPX_SCALE_E(0x1p-6f,1), ROCMFPX_SCALE_E(0x1p-6f,2), ROCMFPX_SCALE_E(0x1p-6f,3), + ROCMFPX_SCALE_E(0x1p-6f,4), ROCMFPX_SCALE_E(0x1p-6f,5), ROCMFPX_SCALE_E(0x1p-6f,6), ROCMFPX_SCALE_E(0x1p-6f,7), + ROCMFPX_SCALE_E(0x1p-5f,0), ROCMFPX_SCALE_E(0x1p-5f,1), ROCMFPX_SCALE_E(0x1p-5f,2), ROCMFPX_SCALE_E(0x1p-5f,3), + ROCMFPX_SCALE_E(0x1p-5f,4), ROCMFPX_SCALE_E(0x1p-5f,5), ROCMFPX_SCALE_E(0x1p-5f,6), ROCMFPX_SCALE_E(0x1p-5f,7), + ROCMFPX_SCALE_E(0x1p-4f,0), ROCMFPX_SCALE_E(0x1p-4f,1), ROCMFPX_SCALE_E(0x1p-4f,2), ROCMFPX_SCALE_E(0x1p-4f,3), + ROCMFPX_SCALE_E(0x1p-4f,4), ROCMFPX_SCALE_E(0x1p-4f,5), ROCMFPX_SCALE_E(0x1p-4f,6), ROCMFPX_SCALE_E(0x1p-4f,7), + ROCMFPX_SCALE_E(0x1p-3f,0), ROCMFPX_SCALE_E(0x1p-3f,1), ROCMFPX_SCALE_E(0x1p-3f,2), ROCMFPX_SCALE_E(0x1p-3f,3), + ROCMFPX_SCALE_E(0x1p-3f,4), ROCMFPX_SCALE_E(0x1p-3f,5), ROCMFPX_SCALE_E(0x1p-3f,6), ROCMFPX_SCALE_E(0x1p-3f,7), + ROCMFPX_SCALE_E(0x1p-2f,0), ROCMFPX_SCALE_E(0x1p-2f,1), ROCMFPX_SCALE_E(0x1p-2f,2), ROCMFPX_SCALE_E(0x1p-2f,3), + ROCMFPX_SCALE_E(0x1p-2f,4), ROCMFPX_SCALE_E(0x1p-2f,5), ROCMFPX_SCALE_E(0x1p-2f,6), ROCMFPX_SCALE_E(0x1p-2f,7), + ROCMFPX_SCALE_E(0x1p-1f,0), ROCMFPX_SCALE_E(0x1p-1f,1), ROCMFPX_SCALE_E(0x1p-1f,2), ROCMFPX_SCALE_E(0x1p-1f,3), + ROCMFPX_SCALE_E(0x1p-1f,4), ROCMFPX_SCALE_E(0x1p-1f,5), ROCMFPX_SCALE_E(0x1p-1f,6), ROCMFPX_SCALE_E(0x1p-1f,7), + ROCMFPX_SCALE_E(0x1p0f,0), ROCMFPX_SCALE_E(0x1p0f,1), ROCMFPX_SCALE_E(0x1p0f,2), ROCMFPX_SCALE_E(0x1p0f,3), + ROCMFPX_SCALE_E(0x1p0f,4), ROCMFPX_SCALE_E(0x1p0f,5), ROCMFPX_SCALE_E(0x1p0f,6), ROCMFPX_SCALE_E(0x1p0f,7), + ROCMFPX_SCALE_E(0x1p1f,0), ROCMFPX_SCALE_E(0x1p1f,1), ROCMFPX_SCALE_E(0x1p1f,2), ROCMFPX_SCALE_E(0x1p1f,3), + ROCMFPX_SCALE_E(0x1p1f,4), ROCMFPX_SCALE_E(0x1p1f,5), ROCMFPX_SCALE_E(0x1p1f,6), ROCMFPX_SCALE_E(0x1p1f,7), + ROCMFPX_SCALE_E(0x1p2f,0), ROCMFPX_SCALE_E(0x1p2f,1), ROCMFPX_SCALE_E(0x1p2f,2), ROCMFPX_SCALE_E(0x1p2f,3), + ROCMFPX_SCALE_E(0x1p2f,4), ROCMFPX_SCALE_E(0x1p2f,5), ROCMFPX_SCALE_E(0x1p2f,6), ROCMFPX_SCALE_E(0x1p2f,7), + ROCMFPX_SCALE_E(0x1p3f,0), ROCMFPX_SCALE_E(0x1p3f,1), ROCMFPX_SCALE_E(0x1p3f,2), ROCMFPX_SCALE_E(0x1p3f,3), + ROCMFPX_SCALE_E(0x1p3f,4), ROCMFPX_SCALE_E(0x1p3f,5), ROCMFPX_SCALE_E(0x1p3f,6), ROCMFPX_SCALE_E(0x1p3f,7), + ROCMFPX_SCALE_E(0x1p4f,0), ROCMFPX_SCALE_E(0x1p4f,1), ROCMFPX_SCALE_E(0x1p4f,2), ROCMFPX_SCALE_E(0x1p4f,3), + ROCMFPX_SCALE_E(0x1p4f,4), ROCMFPX_SCALE_E(0x1p4f,5), ROCMFPX_SCALE_E(0x1p4f,6), +}; + +#undef ROCMFPX_SCALE_SUB +#undef ROCMFPX_SCALE_E + +float rocmfpx_ue4m3_to_fp32(uint8_t e) { + return rocmfpx_scale_is_valid(e) ? rocmfpx_scale_ue4m3[e] : 0.0f; +} + +bool rocmfpx_scale_is_valid(uint8_t e) { + return e <= 0x7e; +} + +size_t rocmfpx_row_size_fp2(int64_t k) { + assert(k % QK_ROCMFP2 == 0); + return (size_t) (k / QK_ROCMFP2) * sizeof(block_rocmfp2); +} + +size_t rocmfpx_row_size_fp3(int64_t k) { + assert(k % QK_ROCMFP3 == 0); + return (size_t) (k / QK_ROCMFP3) * sizeof(block_rocmfp3); +} + +size_t rocmfpx_row_size_fp6(int64_t k) { + assert(k % QK_ROCMFP6 == 0); + return (size_t) (k / QK_ROCMFP6) * sizeof(block_rocmfp6); +} + +size_t rocmfpx_row_size_fp8(int64_t k) { + assert(k % QK_ROCMFP8 == 0); + return (size_t) (k / QK_ROCMFP8) * sizeof(block_rocmfp8); +} + +static uint8_t rocmfpx_nearest_scale_ue4m3(float target) { + if (!(target > 0.0f) || !isfinite(target)) { + return 0; + } + + uint8_t best_e = 1; + float best_err = fabsf(rocmfpx_ue4m3_to_fp32(best_e) - target); + + for (int e = 2; e <= 0x7e; ++e) { + const float err = fabsf(rocmfpx_ue4m3_to_fp32((uint8_t) e) - target); + if (err < best_err) { + best_err = err; + best_e = (uint8_t) e; + } + } + + return best_e; +} + +static float rocmfpx_max_abs(const float * x, int n) { + float max_abs = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float ax = fabsf(x[i]); + if (ax > max_abs) { + max_abs = ax; + } + } + + return max_abs; +} + +static void rocmfpx_prepare_mse_weights( + float * dst, const float * x, int n, const float * quant_weights, float sigma2, + float * max_abs, float * max_abs_weight) { + *max_abs = 0.0f; + *max_abs_weight = 0.0f; + + for (int i = 0; i < n; ++i) { + const float ax = fabsf(x[i]); + const float qw = quant_weights[i]; + const float weight = isfinite(qw) && qw > 0.0f && isfinite(x[i]) ? qw * sqrtf(sigma2 + x[i]*x[i]) : 0.0f; + + if (isfinite(x[i])) { + if (ax > *max_abs) { + *max_abs = ax; + *max_abs_weight = weight; + } else if (ax == *max_abs && weight > *max_abs_weight) { + *max_abs_weight = weight; + } + } + + // Match llama.cpp imatrix weighting style: calibration importance is + // scaled by row energy so large activations stay protected. + dst[i] = weight; + } +} + +static void rocmfpx_set_bits(uint8_t * dst, int bit_pos, int nbits, uint32_t code) { + for (int bit = 0; bit < nbits; ++bit) { + const int absolute_bit = bit_pos + bit; + const int byte_index = absolute_bit >> 3; + const int bit_index = absolute_bit & 7; + + if ((code >> bit) & 1u) { + dst[byte_index] |= (uint8_t) (1u << bit_index); + } + } +} + +static uint32_t rocmfpx_get_bits(const uint8_t * src, int bit_pos, int nbits) { + uint32_t code = 0; + + for (int bit = 0; bit < nbits; ++bit) { + const int absolute_bit = bit_pos + bit; + const int byte_index = absolute_bit >> 3; + const int bit_index = absolute_bit & 7; + + code |= (uint32_t) ((src[byte_index] >> bit_index) & 1u) << bit; + } + + return code; +} + +// Starting 2-bit ROCmFP2 codebook. Keep this single definition easy to tune. +// TODO: affine scale+min would likely improve quality, but it would break the +// ROCmFPx family's unsigned-UE4M3 scale contract; revisit for a v2 layout. +static const float kvalues_rocmfp2[4] = ROCMFP2_KVALUES_INIT; + +static uint8_t rocmfpx_quantize_fp2_code(float x, float inv_scale) { + if (!isfinite(x) || inv_scale <= 0.0f) { + return 1; + } + + const float q = x * inv_scale; + uint8_t best_code = 0; + float best_err = fabsf(q - kvalues_rocmfp2[0]); + + for (uint8_t code = 1; code < 4; ++code) { + const float err = fabsf(q - kvalues_rocmfp2[code]); + if (err < best_err) { + best_err = err; + best_code = code; + } + } + + return best_code; +} + +static inline float rocmfpx_fp2_decoded_mag(float x, float inv_scale) { + return kvalues_rocmfp2[rocmfpx_quantize_fp2_code(x, inv_scale)]; +} + +static float rocmfpx_fp2_block_mse_for_scale(const float * x, int n, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float y = rocmfpx_fp2_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfpx_fp2_block_weighted_mse_for_scale(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float y = rocmfpx_fp2_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static uint8_t rocmfpx_choose_scale_fp2_mse_impl(const float * x, int n, const float * mse_weights, float max_abs, float max_abs_weight) { + const uint8_t start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 2.0f); + uint8_t best_e = start_e; + float best_err = INFINITY; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = (int) start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = rocmfpx_ue4m3_to_fp32((uint8_t) e0); + const float clip_delta = max_abs - 2.0f*scale; + const float clip_err = mse_weights ? max_abs_weight*clip_delta*clip_delta : clip_delta*clip_delta; + if (clip_delta > 0.0f && clip_err > best_err) { + lower_done = true; + } else { + const float err = mse_weights ? + rocmfpx_fp2_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e0, best_err) : + rocmfpx_fp2_block_mse_for_scale(x, n, (uint8_t) e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = (uint8_t) e0; + } + } + } + + const int e1 = (int) start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = mse_weights ? + rocmfpx_fp2_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e1, best_err) : + rocmfpx_fp2_block_mse_for_scale(x, n, (uint8_t) e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = (uint8_t) e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return best_e; +} + +static uint8_t rocmfpx_choose_scale_fp2_mse(const float * x, int n) { + const float max_abs = rocmfpx_max_abs(x, n); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp2_mse_impl(x, n, NULL, max_abs, 0.0f); +} + +static uint8_t rocmfpx_choose_scale_fp2_weighted_mse(const float * x, int n, const float * quant_weights, float sigma2) { + assert(n <= QK_ROCMFP2); + float mse_weights[QK_ROCMFP2]; + float max_abs; + float max_abs_weight; + rocmfpx_prepare_mse_weights(mse_weights, x, n, quant_weights, sigma2, &max_abs, &max_abs_weight); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp2_mse_impl(x, n, mse_weights, max_abs, max_abs_weight); +} + +static int rocmfpx_decode_fp3_code(uint8_t code) { + static const int mag[4] = { 0, 1, 2, 4 }; + const int value = mag[code & 3u]; + return (code & 4u) ? -value : value; +} + +static uint8_t rocmfpx_quantize_fp3_code(float x, float inv_scale) { + if (!isfinite(x) || inv_scale <= 0.0f) { + return 0; + } + + const float ax = fabsf(x * inv_scale); + uint8_t mag; + + if (ax <= 0.5f) { + mag = 0; + } else if (ax <= 1.5f) { + mag = 1; + } else if (ax <= 3.0f) { + mag = 2; + } else { + mag = 3; + } + + return mag == 0 ? 0 : (uint8_t) ((x < 0.0f ? 4u : 0u) | mag); +} + +// Fused threshold + decode used only inside the exhaustive scale search, which +// re-scans every element for every candidate scale byte. Returns the same +// signed decoded magnitude that +// rocmfpx_decode_fp3_code(rocmfpx_quantize_fp3_code(x, inv_scale)) produces +// (fp3 magnitudes {0,1,2,4}), so quantized output stays bit-identical. +static inline float rocmfpx_fp3_decoded_mag(float x, float inv_scale) { + const float a = fabsf(x * inv_scale); + float mag; + if (a <= 0.5f) { + return 0.0f; + } else if (a <= 1.5f) { + mag = 1.0f; + } else if (a <= 3.0f) { + mag = 2.0f; + } else { + mag = 4.0f; + } + return x < 0.0f ? -mag : mag; +} + +static float rocmfpx_fp3_block_mse_for_scale(const float * x, int n, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float y = rocmfpx_fp3_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfpx_fp3_block_weighted_mse_for_scale(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float y = rocmfpx_fp3_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static uint8_t rocmfpx_choose_scale_fp3_mse_impl(const float * x, int n, const float * mse_weights, float max_abs, float max_abs_weight) { + const uint8_t start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 4.0f); + uint8_t best_e = start_e; + float best_err = INFINITY; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = (int) start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = rocmfpx_ue4m3_to_fp32((uint8_t) e0); + const float clip_delta = max_abs - 4.0f*scale; + const float clip_err = mse_weights ? max_abs_weight*clip_delta*clip_delta : clip_delta*clip_delta; + if (clip_delta > 0.0f && clip_err > best_err) { + lower_done = true; + } else { + const float err = mse_weights ? + rocmfpx_fp3_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e0, best_err) : + rocmfpx_fp3_block_mse_for_scale(x, n, (uint8_t) e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = (uint8_t) e0; + } + } + } + + const int e1 = (int) start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = mse_weights ? + rocmfpx_fp3_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e1, best_err) : + rocmfpx_fp3_block_mse_for_scale(x, n, (uint8_t) e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = (uint8_t) e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return best_e; +} + +static uint8_t rocmfpx_choose_scale_fp3_mse(const float * x, int n) { + const float max_abs = rocmfpx_max_abs(x, n); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp3_mse_impl(x, n, NULL, max_abs, 0.0f); +} + +static uint8_t rocmfpx_choose_scale_fp3_weighted_mse(const float * x, int n, const float * quant_weights, float sigma2) { + assert(n <= QK_ROCMFP3); + float mse_weights[QK_ROCMFP3]; + float max_abs; + float max_abs_weight; + rocmfpx_prepare_mse_weights(mse_weights, x, n, quant_weights, sigma2, &max_abs, &max_abs_weight); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp3_mse_impl(x, n, mse_weights, max_abs, max_abs_weight); +} + +static int rocmfpx_decode_fp6_code(uint8_t code) { + const int value = code & 31u; + return (code & 32u) ? -value : value; +} + +static uint8_t rocmfpx_quantize_fp6_code(float x, float inv_scale) { + if (!isfinite(x) || inv_scale <= 0.0f) { + return 0; + } + + int mag = (int) lroundf(fabsf(x * inv_scale)); + if (mag > 31) { + mag = 31; + } + + return mag == 0 ? 0 : (uint8_t) ((x < 0.0f ? 32u : 0u) | (uint8_t) mag); +} + +// Fused round + clamp + decode for the fp6 scale search. Returns the same signed +// decoded magnitude as rocmfpx_decode_fp6_code(rocmfpx_quantize_fp6_code(...)) +// (nearest integer in [0,31], signed), keeping quantized output bit-identical. +static inline float rocmfpx_fp6_decoded_mag(float x, float inv_scale) { + int mag = (int) lroundf(fabsf(x * inv_scale)); + if (mag > 31) { + mag = 31; + } + if (mag == 0) { + return 0.0f; + } + return x < 0.0f ? -(float) mag : (float) mag; +} + +static float rocmfpx_fp6_block_mse_for_scale(const float * x, int n, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + const float y = rocmfpx_fp6_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfpx_fp6_block_weighted_mse_for_scale(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + const float y = rocmfpx_fp6_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static uint8_t rocmfpx_choose_scale_fp6_mse_impl(const float * x, int n, const float * mse_weights, float max_abs, float max_abs_weight) { + const uint8_t start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 31.0f); + uint8_t best_e = start_e; + float best_err = INFINITY; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = (int) start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = rocmfpx_ue4m3_to_fp32((uint8_t) e0); + const float clip_delta = max_abs - 31.0f*scale; + const float clip_err = mse_weights ? max_abs_weight*clip_delta*clip_delta : clip_delta*clip_delta; + if (clip_delta > 0.0f && clip_err > best_err) { + lower_done = true; + } else { + const float err = mse_weights ? + rocmfpx_fp6_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e0, best_err) : + rocmfpx_fp6_block_mse_for_scale(x, n, (uint8_t) e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = (uint8_t) e0; + } + } + } + + const int e1 = (int) start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = mse_weights ? + rocmfpx_fp6_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e1, best_err) : + rocmfpx_fp6_block_mse_for_scale(x, n, (uint8_t) e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = (uint8_t) e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return best_e; +} + +static uint8_t rocmfpx_choose_scale_fp6_mse(const float * x, int n) { + const float max_abs = rocmfpx_max_abs(x, n); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp6_mse_impl(x, n, NULL, max_abs, 0.0f); +} + +static uint8_t rocmfpx_choose_scale_fp6_weighted_mse(const float * x, int n, const float * quant_weights, float sigma2) { + assert(n <= QK_ROCMFP6); + float mse_weights[QK_ROCMFP6]; + float max_abs; + float max_abs_weight; + rocmfpx_prepare_mse_weights(mse_weights, x, n, quant_weights, sigma2, &max_abs, &max_abs_weight); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp6_mse_impl(x, n, mse_weights, max_abs, max_abs_weight); +} + +static int8_t rocmfpx_quantize_fp8_code(float x, float inv_scale) { + if (!isfinite(x) || inv_scale <= 0.0f) { + return 0; + } + + int q = (int) lroundf(x * inv_scale); + if (q > 127) { + q = 127; + } else if (q < -127) { + q = -127; + } + + return (int8_t) q; +} + +static float rocmfpx_fp8_block_weighted_mse_for_scale(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const int8_t code = rocmfpx_quantize_fp8_code(x[i], inv_scale); + const float y = (float) code * scale; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static uint8_t rocmfpx_choose_scale_fp8_weighted_mse(const float * x, int n, const float * quant_weights, float sigma2) { + assert(n <= QK_ROCMFP8); + float mse_weights[QK_ROCMFP8]; + float max_abs; + float max_abs_weight; + rocmfpx_prepare_mse_weights(mse_weights, x, n, quant_weights, sigma2, &max_abs, &max_abs_weight); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + const uint8_t start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 127.0f); + uint8_t best_e = start_e; + float best_err = INFINITY; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = (int) start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = rocmfpx_ue4m3_to_fp32((uint8_t) e0); + const float clip_delta = max_abs - 127.0f*scale; + if (clip_delta > 0.0f && max_abs_weight*clip_delta*clip_delta > best_err) { + lower_done = true; + } else { + const float err = rocmfpx_fp8_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = (uint8_t) e0; + } + } + } + + const int e1 = (int) start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = rocmfpx_fp8_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = (uint8_t) e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return best_e; +} + +void rocmfpx_quantize_row_fp2_ref(const float * GGML_RESTRICT x, block_rocmfp2 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP2 == 0); + + const int64_t nb = k / QK_ROCMFP2; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP2; + block_rocmfp2 * yb = y + ib; + + memset(yb->qs, 0, sizeof(yb->qs)); + + for (int half = 0; half < 2; ++half) { + const float * xh = xb + half*(QK_ROCMFP2/2); + yb->e[half] = rocmfpx_choose_scale_fp2_mse(xh, QK_ROCMFP2/2); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int j = 0; j < QK_ROCMFP2/2; ++j) { + const int i = half*(QK_ROCMFP2/2) + j; + const uint8_t code = rocmfpx_quantize_fp2_code(xb[i], inv_scale); + yb->qs[i >> 2] |= (uint8_t) (code << (2*(i & 3))); + } + } + } +} + +static void rocmfpx_quantize_row_fp2_weighted( + const float * GGML_RESTRICT x, block_rocmfp2 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP2 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += isfinite(x[i]) ? x[i]*x[i] : 0.0f; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP2; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP2; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP2 : NULL; + block_rocmfp2 * yb = y + ib; + + memset(yb->qs, 0, sizeof(yb->qs)); + + for (int half = 0; half < 2; ++half) { + const int half_off = half*(QK_ROCMFP2/2); + const float * xh = xb + half_off; + const float * qh = qw ? qw + half_off : NULL; + yb->e[half] = qh ? + rocmfpx_choose_scale_fp2_weighted_mse(xh, QK_ROCMFP2/2, qh, sigma2) : + rocmfpx_choose_scale_fp2_mse(xh, QK_ROCMFP2/2); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int j = 0; j < QK_ROCMFP2/2; ++j) { + const int i = half_off + j; + const uint8_t code = rocmfpx_quantize_fp2_code(xb[i], inv_scale); + yb->qs[i >> 2] |= (uint8_t) (code << (2*(i & 3))); + } + } + } +} + +void rocmfpx_dequantize_row_fp2(const block_rocmfp2 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP2 == 0); + + const int64_t nb = k / QK_ROCMFP2; + for (int64_t ib = 0; ib < nb; ++ib) { + const block_rocmfp2 * xb = x + ib; + float * yb = y + ib*QK_ROCMFP2; + + for (int i = 0; i < QK_ROCMFP2; ++i) { + const float scale = rocmfpx_ue4m3_to_fp32(xb->e[i >= QK_ROCMFP2/2]); + const uint8_t code = (uint8_t) ((xb->qs[i >> 2] >> (2*(i & 3))) & 3u); + yb[i] = kvalues_rocmfp2[code] * scale; + } + } +} + +void rocmfpx_quantize_row_fp2(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp2_ref(x, (block_rocmfp2 *) y, k); +} + +size_t rocmfpx_quantize_fp2(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = rocmfpx_row_size_fp2(n_per_row); + char * qrow = (char *) dst; + + for (int64_t row = 0; row < nrows; ++row) { + if (imatrix) { + rocmfpx_quantize_row_fp2_weighted(src + row*n_per_row, (block_rocmfp2 *) qrow, n_per_row, imatrix); + } else { + rocmfpx_quantize_row_fp2_ref(src + row*n_per_row, (block_rocmfp2 *) qrow, n_per_row); + } + qrow += row_size; + } + + return (size_t) nrows * row_size; +} + +void rocmfpx_quantize_row_fp3_ref(const float * GGML_RESTRICT x, block_rocmfp3 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP3 == 0); + + const int64_t nb = k / QK_ROCMFP3; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP3; + block_rocmfp3 * yb = y + ib; + + memset(yb->qs, 0, sizeof(yb->qs)); + + for (int half = 0; half < 2; ++half) { + const float * xh = xb + half*(QK_ROCMFP3/2); + yb->e[half] = rocmfpx_choose_scale_fp3_mse(xh, QK_ROCMFP3/2); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int j = 0; j < QK_ROCMFP3/2; ++j) { + const int i = half*(QK_ROCMFP3/2) + j; + const uint8_t code = rocmfpx_quantize_fp3_code(xb[i], inv_scale); + rocmfpx_set_bits(yb->qs, i*3, 3, code); + } + } + } +} + +static void rocmfpx_quantize_row_fp3_weighted( + const float * GGML_RESTRICT x, block_rocmfp3 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP3 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += isfinite(x[i]) ? x[i]*x[i] : 0.0f; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP3; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP3; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP3 : NULL; + block_rocmfp3 * yb = y + ib; + + memset(yb->qs, 0, sizeof(yb->qs)); + + for (int half = 0; half < 2; ++half) { + const int half_off = half*(QK_ROCMFP3/2); + const float * xh = xb + half_off; + const float * qh = qw ? qw + half_off : NULL; + yb->e[half] = qh ? + rocmfpx_choose_scale_fp3_weighted_mse(xh, QK_ROCMFP3/2, qh, sigma2) : + rocmfpx_choose_scale_fp3_mse(xh, QK_ROCMFP3/2); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int j = 0; j < QK_ROCMFP3/2; ++j) { + const int i = half_off + j; + const uint8_t code = rocmfpx_quantize_fp3_code(xb[i], inv_scale); + rocmfpx_set_bits(yb->qs, i*3, 3, code); + } + } + } +} + +void rocmfpx_dequantize_row_fp3(const block_rocmfp3 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP3 == 0); + + const int64_t nb = k / QK_ROCMFP3; + for (int64_t ib = 0; ib < nb; ++ib) { + const block_rocmfp3 * xb = x + ib; + float * yb = y + ib*QK_ROCMFP3; + + for (int i = 0; i < QK_ROCMFP3; ++i) { + const float scale = rocmfpx_ue4m3_to_fp32(xb->e[i >= QK_ROCMFP3/2]); + const uint8_t code = (uint8_t) rocmfpx_get_bits(xb->qs, i*3, 3); + yb[i] = (float) rocmfpx_decode_fp3_code(code) * scale; + } + } +} + +void rocmfpx_quantize_row_fp3(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp3_ref(x, (block_rocmfp3 *) y, k); +} + +size_t rocmfpx_quantize_fp3(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = rocmfpx_row_size_fp3(n_per_row); + char * qrow = (char *) dst; + + for (int64_t row = 0; row < nrows; ++row) { + if (imatrix) { + rocmfpx_quantize_row_fp3_weighted(src + row*n_per_row, (block_rocmfp3 *) qrow, n_per_row, imatrix); + } else { + rocmfpx_quantize_row_fp3_ref(src + row*n_per_row, (block_rocmfp3 *) qrow, n_per_row); + } + qrow += row_size; + } + + return (size_t) nrows * row_size; +} + +static void rocmfpx_quantize_row_fp6_weighted( + const float * GGML_RESTRICT x, block_rocmfp6 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP6 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += isfinite(x[i]) ? x[i]*x[i] : 0.0f; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP6; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP6; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP6 : NULL; + block_rocmfp6 * yb = y + ib; + + memset(yb->qs, 0, sizeof(yb->qs)); + + for (int half = 0; half < 2; ++half) { + const int half_off = half*(QK_ROCMFP6/2); + const float * xh = xb + half_off; + const float * qh = qw ? qw + half_off : NULL; + yb->e[half] = qh ? + rocmfpx_choose_scale_fp6_weighted_mse(xh, QK_ROCMFP6/2, qh, sigma2) : + rocmfpx_choose_scale_fp6_mse(xh, QK_ROCMFP6/2); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int j = 0; j < QK_ROCMFP6/2; ++j) { + const int i = half_off + j; + const uint8_t code = rocmfpx_quantize_fp6_code(xb[i], inv_scale); + rocmfpx_set_bits(yb->qs, i*6, 6, code); + } + } + } +} + +void rocmfpx_quantize_row_fp6_ref(const float * GGML_RESTRICT x, block_rocmfp6 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP6 == 0); + + const int64_t nb = k / QK_ROCMFP6; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP6; + block_rocmfp6 * yb = y + ib; + + memset(yb->qs, 0, sizeof(yb->qs)); + + for (int half = 0; half < 2; ++half) { + const float * xh = xb + half*(QK_ROCMFP6/2); + yb->e[half] = rocmfpx_choose_scale_fp6_mse(xh, QK_ROCMFP6/2); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int j = 0; j < QK_ROCMFP6/2; ++j) { + const int i = half*(QK_ROCMFP6/2) + j; + const uint8_t code = rocmfpx_quantize_fp6_code(xb[i], inv_scale); + rocmfpx_set_bits(yb->qs, i*6, 6, code); + } + } + } +} + +void rocmfpx_dequantize_row_fp6(const block_rocmfp6 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP6 == 0); + + const int64_t nb = k / QK_ROCMFP6; + for (int64_t ib = 0; ib < nb; ++ib) { + const block_rocmfp6 * xb = x + ib; + float * yb = y + ib*QK_ROCMFP6; + + for (int i = 0; i < QK_ROCMFP6; ++i) { + const float scale = rocmfpx_ue4m3_to_fp32(xb->e[i >= QK_ROCMFP6/2]); + const uint8_t code = (uint8_t) rocmfpx_get_bits(xb->qs, i*6, 6); + yb[i] = (float) rocmfpx_decode_fp6_code(code) * scale; + } + } +} + +void rocmfpx_quantize_row_fp6(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp6_ref(x, (block_rocmfp6 *) y, k); +} + +size_t rocmfpx_quantize_fp6(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = rocmfpx_row_size_fp6(n_per_row); + char * qrow = (char *) dst; + + for (int64_t row = 0; row < nrows; ++row) { + if (imatrix) { + rocmfpx_quantize_row_fp6_weighted(src + row*n_per_row, (block_rocmfp6 *) qrow, n_per_row, imatrix); + } else { + rocmfpx_quantize_row_fp6_ref(src + row*n_per_row, (block_rocmfp6 *) qrow, n_per_row); + } + qrow += row_size; + } + + return (size_t) nrows * row_size; +} + +static void rocmfpx_quantize_row_fp8_weighted( + const float * GGML_RESTRICT x, block_rocmfp8 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP8 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += isfinite(x[i]) ? x[i]*x[i] : 0.0f; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP8; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP8; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP8 : NULL; + block_rocmfp8 * yb = y + ib; + + yb->e = qw ? rocmfpx_choose_scale_fp8_weighted_mse(xb, QK_ROCMFP8, qw, sigma2) : + rocmfpx_nearest_scale_ue4m3(rocmfpx_max_abs(xb, QK_ROCMFP8) / 127.0f); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int i = 0; i < QK_ROCMFP8; ++i) { + yb->qs[i] = rocmfpx_quantize_fp8_code(xb[i], inv_scale); + } + } +} + +void rocmfpx_quantize_row_fp8_ref(const float * GGML_RESTRICT x, block_rocmfp8 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP8 == 0); + + const int64_t nb = k / QK_ROCMFP8; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP8; + block_rocmfp8 * yb = y + ib; + + const float max_abs = rocmfpx_max_abs(xb, QK_ROCMFP8); + yb->e = rocmfpx_nearest_scale_ue4m3(max_abs / 127.0f); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int i = 0; i < QK_ROCMFP8; ++i) { + yb->qs[i] = rocmfpx_quantize_fp8_code(xb[i], inv_scale); + } + } +} + +void rocmfpx_dequantize_row_fp8(const block_rocmfp8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP8 == 0); + + const int64_t nb = k / QK_ROCMFP8; + for (int64_t ib = 0; ib < nb; ++ib) { + const block_rocmfp8 * xb = x + ib; + float * yb = y + ib*QK_ROCMFP8; + + const float scale = rocmfpx_ue4m3_to_fp32(xb->e); + for (int i = 0; i < QK_ROCMFP8; ++i) { + yb[i] = (float) xb->qs[i] * scale; + } + } +} + +void rocmfpx_quantize_row_fp8(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp8_ref(x, (block_rocmfp8 *) y, k); +} + +size_t rocmfpx_quantize_fp8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = rocmfpx_row_size_fp8(n_per_row); + char * qrow = (char *) dst; + + for (int64_t row = 0; row < nrows; ++row) { + if (imatrix) { + rocmfpx_quantize_row_fp8_weighted(src + row*n_per_row, (block_rocmfp8 *) qrow, n_per_row, imatrix); + } else { + rocmfpx_quantize_row_fp8_ref(src + row*n_per_row, (block_rocmfp8 *) qrow, n_per_row); + } + qrow += row_size; + } + + return (size_t) nrows * row_size; +} + +bool rocmfpx_validate_row_data_fp2(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp2) != 0) { + return false; + } + + const block_rocmfp2 * blocks = (const block_rocmfp2 *) data; + const size_t nb = nbytes / sizeof(block_rocmfp2); + + for (size_t i = 0; i < nb; ++i) { + if (!rocmfpx_scale_is_valid(blocks[i].e[0]) || !rocmfpx_scale_is_valid(blocks[i].e[1])) { + return false; + } + } + + return true; +} + +bool rocmfpx_validate_row_data_fp3(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp3) != 0) { + return false; + } + + const block_rocmfp3 * blocks = (const block_rocmfp3 *) data; + const size_t nb = nbytes / sizeof(block_rocmfp3); + + for (size_t i = 0; i < nb; ++i) { + if (!rocmfpx_scale_is_valid(blocks[i].e[0]) || !rocmfpx_scale_is_valid(blocks[i].e[1])) { + return false; + } + } + + return true; +} + +bool rocmfpx_validate_row_data_fp6(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp6) != 0) { + return false; + } + + const block_rocmfp6 * blocks = (const block_rocmfp6 *) data; + const size_t nb = nbytes / sizeof(block_rocmfp6); + + for (size_t i = 0; i < nb; ++i) { + if (!rocmfpx_scale_is_valid(blocks[i].e[0]) || !rocmfpx_scale_is_valid(blocks[i].e[1])) { + return false; + } + } + + return true; +} + +bool rocmfpx_validate_row_data_fp8(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp8) != 0) { + return false; + } + + const block_rocmfp8 * blocks = (const block_rocmfp8 *) data; + const size_t nb = nbytes / sizeof(block_rocmfp8); + + for (size_t i = 0; i < nb; ++i) { + if (!rocmfpx_scale_is_valid(blocks[i].e)) { + return false; + } + } + + return true; +} diff --git a/server/deps/llama.cpp/ggml/rocmfpx/rocmfpx.h b/server/deps/llama.cpp/ggml/rocmfpx/rocmfpx.h new file mode 100644 index 000000000..f9d6cb896 --- /dev/null +++ b/server/deps/llama.cpp/ggml/rocmfpx/rocmfpx.h @@ -0,0 +1,111 @@ +#pragma once + +#include +#include +#include + +#include "ggml.h" + +#ifdef __cplusplus +extern "C" { +#endif + +#define QK_ROCMFPX 32 + +#define QK_ROCMFP2 QK_ROCMFPX +#define QK_ROCMFP3 QK_ROCMFPX +#define QK_ROCMFP6 QK_ROCMFPX +#define QK_ROCMFP8 QK_ROCMFPX + +#define QS_ROCMFP2 ((QK_ROCMFP2 * 2) / 8) +#define QS_ROCMFP3 ((QK_ROCMFP3 * 3) / 8) +#define QS_ROCMFP6 ((QK_ROCMFP6 * 6) / 8) +#define QS_ROCMFP8 QK_ROCMFP8 + +#define QR_ROCMFP2 1 +#define QI_ROCMFP2 (QS_ROCMFP2 / 4) // 8B qs = 2 int32; decoupled from QR_ROCMFP2 (QR stays 1 for convert.cu dequant) + +#define ROCMFP2_KVALUE_0_I8 (-1) +#define ROCMFP2_KVALUE_1_I8 0 +#define ROCMFP2_KVALUE_2_I8 1 +#define ROCMFP2_KVALUE_3_I8 2 +#define ROCMFP2_KVALUES_INIT { (float) ROCMFP2_KVALUE_0_I8, (float) ROCMFP2_KVALUE_1_I8, (float) ROCMFP2_KVALUE_2_I8, (float) ROCMFP2_KVALUE_3_I8 } + +#define QR_ROCMFP3 1 +#define QI_ROCMFP3 (QK_ROCMFP3 / (4 * QR_ROCMFP3)) + +#define QR_ROCMFP6 1 +#define QI_ROCMFP6 (QK_ROCMFP6 / (4 * QR_ROCMFP6)) + +#define QR_ROCMFP8 1 +#define QI_ROCMFP8 (QK_ROCMFP8 / (4 * QR_ROCMFP8)) + +// AMD-native experimental family layouts. The GGUF types are registered, but +// the layouts stay isolated from the promoted ROCmFP4 formats while evaluated. +typedef struct { + uint8_t qs[QS_ROCMFP2]; + uint8_t e[2]; +} block_rocmfp2; + +typedef struct { + uint8_t qs[QS_ROCMFP3]; + uint8_t e[2]; +} block_rocmfp3; + +typedef struct { + uint8_t qs[QS_ROCMFP6]; + uint8_t e[2]; +} block_rocmfp6; + +typedef struct { + int8_t qs[QS_ROCMFP8]; + uint8_t e; +} block_rocmfp8; + +#if defined(__cplusplus) +static_assert(sizeof(block_rocmfp2) == QS_ROCMFP2 + 2*sizeof(uint8_t), "wrong rocmfp2 block size/padding"); +static_assert(sizeof(block_rocmfp3) == QS_ROCMFP3 + 2*sizeof(uint8_t), "wrong rocmfp3 block size/padding"); +static_assert(sizeof(block_rocmfp6) == QS_ROCMFP6 + 2*sizeof(uint8_t), "wrong rocmfp6 block size/padding"); +static_assert(sizeof(block_rocmfp8) == QS_ROCMFP8 + sizeof(uint8_t), "wrong rocmfp8 block size/padding"); +#else +_Static_assert(sizeof(block_rocmfp2) == QS_ROCMFP2 + 2*sizeof(uint8_t), "wrong rocmfp2 block size/padding"); +_Static_assert(sizeof(block_rocmfp3) == QS_ROCMFP3 + 2*sizeof(uint8_t), "wrong rocmfp3 block size/padding"); +_Static_assert(sizeof(block_rocmfp6) == QS_ROCMFP6 + 2*sizeof(uint8_t), "wrong rocmfp6 block size/padding"); +_Static_assert(sizeof(block_rocmfp8) == QS_ROCMFP8 + sizeof(uint8_t), "wrong rocmfp8 block size/padding"); +#endif + +GGML_API float rocmfpx_ue4m3_to_fp32(uint8_t e); +GGML_API bool rocmfpx_scale_is_valid(uint8_t e); +GGML_API size_t rocmfpx_row_size_fp2(int64_t k); +GGML_API size_t rocmfpx_row_size_fp3(int64_t k); +GGML_API size_t rocmfpx_row_size_fp6(int64_t k); +GGML_API size_t rocmfpx_row_size_fp8(int64_t k); + +GGML_API void rocmfpx_quantize_row_fp2_ref(const float * GGML_RESTRICT x, block_rocmfp2 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_dequantize_row_fp2(const block_rocmfp2 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_quantize_row_fp2(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfpx_quantize_fp2(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API void rocmfpx_quantize_row_fp3_ref(const float * GGML_RESTRICT x, block_rocmfp3 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_dequantize_row_fp3(const block_rocmfp3 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_quantize_row_fp3(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfpx_quantize_fp3(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API void rocmfpx_quantize_row_fp6_ref(const float * GGML_RESTRICT x, block_rocmfp6 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_dequantize_row_fp6(const block_rocmfp6 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_quantize_row_fp6(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfpx_quantize_fp6(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API void rocmfpx_quantize_row_fp8_ref(const float * GGML_RESTRICT x, block_rocmfp8 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_dequantize_row_fp8(const block_rocmfp8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_quantize_row_fp8(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfpx_quantize_fp8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API bool rocmfpx_validate_row_data_fp2(const void * data, size_t nbytes); +GGML_API bool rocmfpx_validate_row_data_fp3(const void * data, size_t nbytes); +GGML_API bool rocmfpx_validate_row_data_fp6(const void * data, size_t nbytes); +GGML_API bool rocmfpx_validate_row_data_fp8(const void * data, size_t nbytes); + +#ifdef __cplusplus +} +#endif diff --git a/server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c b/server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c new file mode 100644 index 000000000..5bae72bd7 --- /dev/null +++ b/server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c @@ -0,0 +1,164 @@ +#include "rocmfpx.h" + +#include +#include +#include + +static void fill_row(float * x, int n) { + for (int i = 0; i < n; ++i) { + const float wave = 0.75f*sinf((float) i * 0.37f) + 0.25f*cosf((float) i * 0.13f); + const float ramp = ((float) (i % 11) - 5.0f) * 0.035f; + x[i] = wave + ramp; + } + + x[7] = 3.25f; + x[19] = -2.75f; + x[43] = 1.875f; +} + +static float mse(const float * a, const float * b, int n) { + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float d = a[i] - b[i]; + err += d*d; + } + + return err / (float) n; +} + +static float weighted_mse(const float * a, const float * b, const float * w, int n) { + float err = 0.0f; + float sum_w = 0.0f; + + for (int i = 0; i < n; ++i) { + const float d = a[i] - b[i]; + err += w[i]*d*d; + sum_w += w[i]; + } + + return sum_w > 0.0f ? err / sum_w : 0.0f; +} + +static void check_weighted_imatrix_fp3(void) { + enum { N = QK_ROCMFP3 }; + + float src[N]; + float imatrix[N]; + float plain[N]; + float weighted[N]; + block_rocmfp3 q_plain[N / QK_ROCMFP3]; + block_rocmfp3 q_weighted[N / QK_ROCMFP3]; + + for (int i = 0; i < N; ++i) { + src[i] = (i % 2) ? 0.21f : -0.21f; + imatrix[i] = 100.0f; + } + + src[0] = 9.0f; + imatrix[0] = 0.0f; + + rocmfpx_quantize_fp3(src, q_plain, 1, N, NULL); + rocmfpx_quantize_fp3(src, q_weighted, 1, N, imatrix); + rocmfpx_dequantize_row_fp3(q_plain, plain, N); + rocmfpx_dequantize_row_fp3(q_weighted, weighted, N); + + const float plain_err = weighted_mse(src, plain, imatrix, N); + const float weighted_err = weighted_mse(src, weighted, imatrix, N); + + printf("ROCmFP3 imatrix weighted_mse: plain=%g weighted=%g\n", plain_err, weighted_err); + assert(weighted_err < plain_err); +} + +static void check_weighted_imatrix_fp2(void) { + enum { N = QK_ROCMFP2 }; + + float src[N]; + float imatrix[N]; + float plain[N]; + float weighted[N]; + block_rocmfp2 q_plain[N / QK_ROCMFP2]; + block_rocmfp2 q_weighted[N / QK_ROCMFP2]; + + for (int i = 0; i < N; ++i) { + src[i] = (i % 2) ? 0.21f : -0.21f; + imatrix[i] = 100.0f; + } + + src[0] = 9.0f; + imatrix[0] = 0.0f; + + rocmfpx_quantize_fp2(src, q_plain, 1, N, NULL); + rocmfpx_quantize_fp2(src, q_weighted, 1, N, imatrix); + rocmfpx_dequantize_row_fp2(q_plain, plain, N); + rocmfpx_dequantize_row_fp2(q_weighted, weighted, N); + + const float plain_err = weighted_mse(src, plain, imatrix, N); + const float weighted_err = weighted_mse(src, weighted, imatrix, N); + + printf("ROCmFP2 imatrix weighted_mse: plain=%g weighted=%g\n", plain_err, weighted_err); + assert(weighted_err < plain_err); +} + +int main(void) { + enum { N = 64 }; + + float src[N]; + float fp2[N]; + float fp3[N]; + float fp6[N]; + float fp8[N]; + + block_rocmfp2 q2[N / QK_ROCMFP2]; + block_rocmfp3 q3[N / QK_ROCMFP3]; + block_rocmfp6 q6[N / QK_ROCMFP6]; + block_rocmfp8 q8[N / QK_ROCMFP8]; + + fill_row(src, N); + + rocmfpx_quantize_row_fp2_ref(src, q2, N); + rocmfpx_quantize_row_fp3_ref(src, q3, N); + rocmfpx_quantize_row_fp6_ref(src, q6, N); + rocmfpx_quantize_row_fp8_ref(src, q8, N); + + assert(rocmfpx_validate_row_data_fp2(q2, sizeof(q2))); + assert(rocmfpx_validate_row_data_fp3(q3, sizeof(q3))); + assert(rocmfpx_validate_row_data_fp6(q6, sizeof(q6))); + assert(rocmfpx_validate_row_data_fp8(q8, sizeof(q8))); + + rocmfpx_dequantize_row_fp2(q2, fp2, N); + rocmfpx_dequantize_row_fp3(q3, fp3, N); + rocmfpx_dequantize_row_fp6(q6, fp6, N); + rocmfpx_dequantize_row_fp8(q8, fp8, N); + + const float mse2 = mse(src, fp2, N); + const float mse3 = mse(src, fp3, N); + const float mse6 = mse(src, fp6, N); + const float mse8 = mse(src, fp8, N); + + printf("ROCmFP2: block=%zu row=%zu bpw=%.2f mse=%g\n", + sizeof(block_rocmfp2), rocmfpx_row_size_fp2(N), + 8.0f*(float) sizeof(block_rocmfp2)/(float) QK_ROCMFP2, mse2); + printf("ROCmFP3: block=%zu row=%zu bpw=%.2f mse=%g\n", + sizeof(block_rocmfp3), rocmfpx_row_size_fp3(N), + 8.0f*(float) sizeof(block_rocmfp3)/(float) QK_ROCMFP3, mse3); + printf("ROCmFP6: block=%zu row=%zu bpw=%.2f mse=%g\n", + sizeof(block_rocmfp6), rocmfpx_row_size_fp6(N), + 8.0f*(float) sizeof(block_rocmfp6)/(float) QK_ROCMFP6, mse6); + printf("ROCmFP8: block=%zu row=%zu bpw=%.2f mse=%g\n", + sizeof(block_rocmfp8), rocmfpx_row_size_fp8(N), + 8.0f*(float) sizeof(block_rocmfp8)/(float) QK_ROCMFP8, mse8); + + assert(sizeof(block_rocmfp2) == 10); + assert(isfinite(mse2)); + assert(isfinite(mse3)); + assert(isfinite(mse6)); + assert(isfinite(mse8)); + assert(mse8 < mse6); + assert(mse6 < mse3); + + check_weighted_imatrix_fp2(); + check_weighted_imatrix_fp3(); + + return 0; +} diff --git a/server/deps/llama.cpp/ggml/src/CMakeLists.txt b/server/deps/llama.cpp/ggml/src/CMakeLists.txt index 48fbe208d..315eafe2e 100644 --- a/server/deps/llama.cpp/ggml/src/CMakeLists.txt +++ b/server/deps/llama.cpp/ggml/src/CMakeLists.txt @@ -205,6 +205,10 @@ add_library(ggml-base ggml-threading.cpp ggml-threading.h ggml-quants.c + ../rocmfp4/rocmfp4.c + ../rocmfp4/rocmfp4.h + ../rocmfpx/rocmfpx.c + ../rocmfpx/rocmfpx.h ggml-quants.h gguf.cpp) diff --git a/server/deps/llama.cpp/ggml/src/ggml-common.h b/server/deps/llama.cpp/ggml/src/ggml-common.h index d47c97696..e52a2005d 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-common.h +++ b/server/deps/llama.cpp/ggml/src/ggml-common.h @@ -1132,6 +1132,12 @@ GGML_TABLE_BEGIN(int8_t, kvalues_mxfp4, 16) 0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12, GGML_TABLE_END() +// ROCmFP4 uses an E2M1-derived value set with the largest level retuned from +// 12 to 10, plus dual half-block UE4M3 scales. +GGML_TABLE_BEGIN(int8_t, kvalues_rocmfp4, 16) + 0, 1, 2, 3, 4, 6, 8, 10, 0, -1, -2, -3, -4, -6, -8, -10, +GGML_TABLE_END() + #define NGRID_IQ1S 2048 #define IQ1S_DELTA 0.125f #define IQ1M_DELTA 0.125f 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 f0e14a006..170f1e4a9 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 @@ -1826,6 +1826,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { GGML_ABORT("GGML_OP_MOE_FUSED is only implemented for CUDA"); } + case GGML_OP_DS4_HC: + { + GGML_ABORT("GGML_OP_DS4_HC is only implemented for CUDA"); + } case GGML_OP_OUT_PROD: { ggml_compute_forward_out_prod(params, tensor); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp b/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp index c8f811406..6dff19338 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ops.cpp @@ -3271,6 +3271,80 @@ static void ggml_compute_forward_swiglu( } } +// ggml_compute_forward_swiglu_ds4 + +static void ggml_compute_forward_swiglu_ds4_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; + + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); + GGML_ASSERT(src1); + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + + const int ith = params->ith; + const int nth = params->nth; + + const int nc = src0->ne[0]; + const int nr = ggml_nrows(src0); + + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); + + const float clamp = ggml_get_op_params_f32(dst, 2); + + const int dr = (nr + nth - 1)/nth; + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + + for (int i1 = ir0; i1 < ir1; i1++) { + float * gate_p = (float *) (src0_d + i1*src0_o); + float * up_p = (float *) (src1_d + i1*src1_o); + float * dst_p = (float *) ((char *) dst->data + i1*(dst->nb[1])); + + for (int k = 0; k < nc; k++) { + const float gate = std::min(gate_p[k], clamp); + const float up = std::clamp(up_p[k], -clamp, clamp); + dst_p[k] = up * gate / (1.f + expf(-gate)); + } + +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = dst_p[k]; + GGML_UNUSED(x); + assert(!isnan(x)); + assert(!isinf(x)); + } +#endif // NDEBUG + } +} + +static void ggml_compute_forward_swiglu_ds4( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_swiglu_ds4_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + // ggml_compute_forward_swiglu_oai static void ggml_compute_forward_swiglu_oai_f32( @@ -9801,6 +9875,10 @@ void ggml_compute_forward_glu( { ggml_compute_forward_swiglu_oai(params, dst); } break; + case GGML_GLU_OP_SWIGLU_DS4: + { + ggml_compute_forward_swiglu_ds4(params, dst); + } break; case GGML_GLU_OP_GEGLU_ERF: { ggml_compute_forward_geglu_erf(params, dst); 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 9df1461a5..eb7d125a1 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh @@ -19,6 +19,8 @@ #endif #endif #include "ggml-common.h" +#include "../../rocmfp4/rocmfp4.h" +#include "../../rocmfpx/rocmfpx.h" #include #include @@ -29,6 +31,26 @@ #include #include +#ifndef GGML_ROCMFP6_EXPANDED_DEVICE +#define GGML_ROCMFP6_EXPANDED_DEVICE 0 +#endif + +// Optional device-only ROCmFP6 layout. GGUF/CPU storage remains the packed +// block_rocmfp6 layout; experimental ROCm builds may expand qs to signed +// bytes to avoid bit unpacking in hot matmul/FA kernels. +struct block_rocmfp6_expanded { + int8_t qs[QK_ROCMFP6]; + uint8_t e[2]; +}; + +static_assert(sizeof(block_rocmfp6_expanded) == QK_ROCMFP6 + 2*sizeof(uint8_t), "wrong expanded rocmfp6 block size/padding"); + +#if GGML_ROCMFP6_EXPANDED_DEVICE +using block_rocmfp6_device = block_rocmfp6_expanded; +#else +using block_rocmfp6_device = block_rocmfp6; +#endif + #if defined(GGML_USE_HIP) #include "vendors/hip.h" #elif defined(GGML_USE_MUSA) @@ -959,6 +981,48 @@ struct ggml_cuda_type_traits { static constexpr int qi = QI8_0; }; +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP4; + static constexpr int qr = QR_ROCMFP4; + static constexpr int qi = QI_ROCMFP4; +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP4; + static constexpr int qr = QR_ROCMFP4; + static constexpr int qi = QI_ROCMFP4; +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP2; + static constexpr int qr = QR_ROCMFP2; + static constexpr int qi = QI_ROCMFP2; +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP3; + static constexpr int qr = QR_ROCMFP3; + static constexpr int qi = QI_ROCMFP3; +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP6; + static constexpr int qr = QR_ROCMFP6; + static constexpr int qi = QI_ROCMFP6; +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP8; + static constexpr int qr = QR_ROCMFP8; + static constexpr int qi = QI_ROCMFP8; +}; + template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_MXFP4; @@ -1466,10 +1530,14 @@ struct ggml_cuda_mm_fusion_args_host { const ggml_tensor * gate = nullptr; const ggml_tensor * gate_bias = nullptr; ggml_glu_op glu_op; + float glu_param0 = 0.0f; + float glu_param1 = 0.0f; }; struct ggml_cuda_mm_fusion_args_device { const void * x_bias = nullptr; const void * gate = nullptr; const void * gate_bias = nullptr; ggml_glu_op glu_op; + float glu_param0 = 0.0f; + float glu_param1 = 0.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 e9ffd274b..51a3efa53 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/concat.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/concat.cu @@ -1,7 +1,8 @@ #include "concat.cuh" // contiguous kernels -static __global__ void concat_f32_dim0(const float * x, const float * y, float * dst, const int ne0, const int ne00) { +template +static __global__ void concat_dim0(const T * x, const T * y, T * dst, const int ne0, const int ne00) { int nidx = threadIdx.x + blockIdx.x * blockDim.x; if (nidx >= ne0) { return; @@ -27,7 +28,8 @@ static __global__ void concat_f32_dim0(const float * x, const float * y, float * } } -static __global__ void concat_f32_dim1(const float * x, const float * y, float * dst, const int ne0, const int ne01) { +template +static __global__ void concat_dim1(const T * x, const T * y, T * dst, const int ne0, const int ne01) { int nidx = threadIdx.x + blockIdx.x * blockDim.x; if (nidx >= ne0) { return; @@ -53,7 +55,8 @@ static __global__ void concat_f32_dim1(const float * x, const float * y, float * } } -static __global__ void concat_f32_dim2(const float * x, const float * y, float * dst, const int ne0, const int ne02) { +template +static __global__ void concat_dim2(const T * x, const T * y, T * dst, const int ne0, const int ne02) { int nidx = threadIdx.x + blockIdx.x * blockDim.x; if (nidx >= ne0) { return; @@ -79,24 +82,25 @@ static __global__ void concat_f32_dim2(const float * x, const float * y, float * } } -static void concat_f32_cuda(const float * x, const float * y, float * dst, int ne00, int ne01, int ne02, int ne0, int ne1, int ne2, int dim, cudaStream_t stream) { +template +static void concat_cuda(const T * x, const T * y, T * dst, int ne00, int ne01, int ne02, int ne0, int ne1, int ne2, int dim, cudaStream_t stream) { int num_blocks = (ne0 + CUDA_CONCAT_BLOCK_SIZE - 1) / CUDA_CONCAT_BLOCK_SIZE; dim3 gridDim(num_blocks, ne1, ne2); if (dim == 0) { - concat_f32_dim0<<>>(x, y, dst, ne0, ne00); + concat_dim0<<>>(x, y, dst, ne0, ne00); return; } if (dim == 1) { - concat_f32_dim1<<>>(x, y, dst, ne0, ne01); + concat_dim1<<>>(x, y, dst, ne0, ne01); return; } - concat_f32_dim2<<>>(x, y, dst, ne0, ne02); + concat_dim2<<>>(x, y, dst, ne0, ne02); } // non-contiguous kernel (slow) -template +template static __global__ void __launch_bounds__(CUDA_CONCAT_BLOCK_SIZE) - concat_f32_non_cont( + concat_non_cont( const char * src0, const char * src1, char * dst, @@ -130,54 +134,44 @@ static __global__ void __launch_bounds__(CUDA_CONCAT_BLOCK_SIZE) const int64_t i2 = blockIdx.y; const int64_t i1 = blockIdx.x; - const float * x; + const T * x; for (int64_t i0 = threadIdx.x; i0 < ne0; i0 += blockDim.x) { if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { - x = (const float *)(src0 + (i3 )*nb03 + (i2 )*nb02 + (i1 )*nb01 + (i0 )*nb00); + x = (const T *)(src0 + (i3 )*nb03 + (i2 )*nb02 + (i1 )*nb01 + (i0 )*nb00); } else { if constexpr (dim == 0) { - x = (const float *) (src1 + i3 * nb13 + i2 * nb12 + i1 * nb11 + (i0 - ne00) * nb10); + x = (const T *) (src1 + i3 * nb13 + i2 * nb12 + i1 * nb11 + (i0 - ne00) * nb10); } else if constexpr (dim == 1) { - x = (const float *) (src1 + i3 * nb13 + i2 * nb12 + (i1 - ne01) * nb11 + i0 * nb10); + x = (const T *) (src1 + i3 * nb13 + i2 * nb12 + (i1 - ne01) * nb11 + i0 * nb10); } else if constexpr (dim == 2) { - x = (const float *) (src1 + i3 * nb13 + (i2 - ne02) * nb12 + i1 * nb11 + i0 * nb10); + x = (const T *) (src1 + i3 * nb13 + (i2 - ne02) * nb12 + i1 * nb11 + i0 * nb10); } else if constexpr (dim == 3) { - x = (const float *) (src1 + (i3 - ne03) * nb13 + i2 * nb12 + i1 * nb11 + i0 * nb10); + x = (const T *) (src1 + (i3 - ne03) * nb13 + i2 * nb12 + i1 * nb11 + i0 * nb10); } } - float * y = (float *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); + T * y = (T *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); *y = *x; } } - -void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - +template +static void concat_cuda_typed(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, int dim) { cudaStream_t stream = ctx.stream(); - const int32_t dim = ((int32_t *) dst->op_params)[0]; - - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) { - const float * src0_d = (const float *)src0->data; - const float * src1_d = (const float *)src1->data; - - float * dst_d = (float *)dst->data; + const T * src0_d = (const T *)src0->data; + const T * src1_d = (const T *)src1->data; + T * dst_d = (T *)dst->data; if (dim != 3) { for (int i3 = 0; i3 < dst->ne[3]; i3++) { - concat_f32_cuda( - src0_d + i3 * (src0->nb[3] / 4), - src1_d + i3 * (src1->nb[3] / 4), - dst_d + i3 * ( dst->nb[3] / 4), + concat_cuda( + src0_d + i3 * (src0->nb[3] / sizeof(T)), + src1_d + i3 * (src1->nb[3] / sizeof(T)), + dst_d + i3 * ( dst->nb[3] / sizeof(T)), src0->ne[0], src0->ne[1], src0->ne[2], dst->ne[0], dst->ne[1], dst->ne[2], dim, stream); } @@ -185,13 +179,13 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const size_t size0 = ggml_nbytes(src0); const size_t size1 = ggml_nbytes(src1); - CUDA_CHECK(cudaMemcpyAsync(dst_d, src0_d, size0, cudaMemcpyDeviceToDevice, stream)); - CUDA_CHECK(cudaMemcpyAsync(dst_d + size0/4, src1_d, size1, cudaMemcpyDeviceToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync(dst_d, src0_d, size0, cudaMemcpyDeviceToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync((char *)dst_d + size0, src1_d, size1, cudaMemcpyDeviceToDevice, stream)); } } else { dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]); auto launch_kernel = [&](auto dim) { - concat_f32_non_cont<<>>( + concat_non_cont<<>>( (const char *) src0->data, (const char *) src1->data, (char *) dst->data, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3], @@ -219,3 +213,30 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { } } } + +void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + const int32_t dim = ((int32_t *) dst->op_params)[0]; + + GGML_ASSERT(src0->type == src1->type); + GGML_ASSERT(src0->type == dst->type); + + switch (src0->type) { + case GGML_TYPE_F32: + concat_cuda_typed(ctx, src0, src1, dst, dim); + break; + case GGML_TYPE_F16: + concat_cuda_typed(ctx, src0, src1, dst, dim); + break; + case GGML_TYPE_BF16: + concat_cuda_typed(ctx, src0, src1, dst, dim); + break; + case GGML_TYPE_I8: + 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/convert.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/convert.cu index c675034af..28d042aa0 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/convert.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/convert.cu @@ -1,6 +1,7 @@ #include "convert.cuh" #include "dequantize.cuh" #include "tq3-quant.cuh" +#include "../../rocmfp4/rocmfp4_hip_scale.cuh" #include @@ -487,6 +488,43 @@ static __global__ void dequantize_block_mxfp4(const void * __restrict__ vx, dst_ } } +template +static __global__ void dequantize_block_rocmfp4(const void * __restrict__ vx, dst_t * __restrict__ yy) { + + const int64_t i = blockIdx.x; + const block_rocmfp4 * x = (const block_rocmfp4 *) vx + i*(QK_K/QK_ROCMFP4); + + const int64_t tid = threadIdx.x; + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + i*QK_K + 32*ib + 4*il; + const uint8_t * q4 = x[ib].qs + 4*il; + const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[1]); + for (int j = 0; j < 4; ++j) { + y[j+ 0] = d0 * rocmfp4_decode_i8(q4[j]); + y[j+16] = d1 * rocmfp4_decode_i8(q4[j] >> 4); + } +} + +template +static __global__ void dequantize_block_rocmfp4_fast(const void * __restrict__ vx, dst_t * __restrict__ yy) { + + const int64_t i = blockIdx.x; + const block_rocmfp4_fast * x = (const block_rocmfp4_fast *) vx + i*(QK_K/QK_ROCMFP4); + + const int64_t tid = threadIdx.x; + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + i*QK_K + 32*ib + 4*il; + const uint8_t * q4 = x[ib].qs + 4*il; + const float d = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e); + for (int j = 0; j < 4; ++j) { + y[j+ 0] = d * rocmfp4_decode_i8(q4[j]); + y[j+16] = d * rocmfp4_decode_i8(q4[j] >> 4); + } +} + template static void dequantize_block_cuda(const void * vx, dst_t * y, const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03, @@ -645,6 +683,18 @@ static void dequantize_row_mxfp4_cuda(const void * vx, dst_t * y, const int64_t dequantize_block_mxfp4<<>>(vx, y); } +template +static void dequantize_row_rocmfp4_hip(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = (k + QK_K - 1) / QK_K; + dequantize_block_rocmfp4<<>>(vx, y); +} + +template +static void dequantize_row_rocmfp4_fast_hip(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = (k + QK_K - 1) / QK_K; + dequantize_block_rocmfp4_fast<<>>(vx, y); +} + template static __global__ void dequantize_block_nvfp4( const void * __restrict__ vx, @@ -786,6 +836,18 @@ to_fp16_cuda_t ggml_get_to_fp16_cuda(ggml_type type) { return dequantize_row_mxfp4_cuda; case GGML_TYPE_NVFP4: return dequantize_row_nvfp4_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_row_rocmfp4_hip; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_row_rocmfp4_fast_hip; + case GGML_TYPE_Q2_0_ROCMFP2: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q3_0_ROCMFPX: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q6_0_ROCMFPX: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q8_0_ROCMFPX: + return dequantize_block_cont_cuda; case GGML_TYPE_F32: return convert_unary_cont_cuda; case GGML_TYPE_BF16: @@ -841,6 +903,18 @@ to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type) { return dequantize_row_mxfp4_cuda; case GGML_TYPE_NVFP4: return dequantize_row_nvfp4_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_row_rocmfp4_hip; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_row_rocmfp4_fast_hip; + case GGML_TYPE_Q2_0_ROCMFP2: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q3_0_ROCMFPX: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q6_0_ROCMFPX: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q8_0_ROCMFPX: + return dequantize_block_cont_cuda; case GGML_TYPE_F16: return convert_unary_cont_cuda; case GGML_TYPE_BF16: @@ -864,6 +938,18 @@ to_fp16_nc_cuda_t ggml_get_to_fp16_nc_cuda(ggml_type type) { return dequantize_block_cuda; case GGML_TYPE_Q8_0: return dequantize_block_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_block_cuda; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_block_cuda; + case GGML_TYPE_Q2_0_ROCMFP2: + return dequantize_block_cuda; + case GGML_TYPE_Q3_0_ROCMFPX: + return dequantize_block_cuda; + case GGML_TYPE_Q6_0_ROCMFPX: + return dequantize_block_cuda; + case GGML_TYPE_Q8_0_ROCMFPX: + return dequantize_block_cuda; case GGML_TYPE_BF16: return convert_unary_cuda; default: @@ -885,6 +971,18 @@ to_bf16_nc_cuda_t ggml_get_to_bf16_nc_cuda(ggml_type type) { return dequantize_block_cuda; case GGML_TYPE_Q8_0: return dequantize_block_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_block_cuda; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_block_cuda; + case GGML_TYPE_Q2_0_ROCMFP2: + return dequantize_block_cuda; + case GGML_TYPE_Q3_0_ROCMFPX: + return dequantize_block_cuda; + case GGML_TYPE_Q6_0_ROCMFPX: + return dequantize_block_cuda; + case GGML_TYPE_Q8_0_ROCMFPX: + return dequantize_block_cuda; case GGML_TYPE_F16: return convert_unary_cuda; default: @@ -906,6 +1004,18 @@ to_fp32_nc_cuda_t ggml_get_to_fp32_nc_cuda(ggml_type type) { return dequantize_block_cuda; case GGML_TYPE_Q8_0: return dequantize_block_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_block_cuda; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_block_cuda; + case GGML_TYPE_Q2_0_ROCMFP2: + return dequantize_block_cuda; + case GGML_TYPE_Q3_0_ROCMFPX: + return dequantize_block_cuda; + case GGML_TYPE_Q6_0_ROCMFPX: + return dequantize_block_cuda; + case GGML_TYPE_Q8_0_ROCMFPX: + return dequantize_block_cuda; case GGML_TYPE_BF16: return convert_unary_cuda; default: diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu index 372fc98b1..3d7a698b9 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/cpy.cu @@ -1,6 +1,7 @@ #include "cpy.cuh" #include "dequantize.cuh" #include "cpy-utils.cuh" +#include "../../rocmfp4/rocmfp4_hip_scale.cuh" #if defined(GGML_USE_MUSA) && defined(GGML_MUSA_MUDNN_COPY) #include "ggml-musa/mudnn.cuh" #endif // GGML_USE_MUSA && GGML_MUSA_MUDNN_COPY @@ -119,6 +120,72 @@ static __device__ void cpy_blck_q_f32(const char * cxi, char * cdsti) { } } +static __device__ void cpy_blck_rocmfp4_f32(const char * cxi, char * cdsti) { + const block_rocmfp4 * x = (const block_rocmfp4 *) cxi; + float * cdstf = (float *) cdsti; + + const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(x->e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(x->e[1]); + +#pragma unroll + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q = x->qs[j]; + cdstf[j] = d0 * (float) rocmfp4_decode_i8(q); + cdstf[j + QK_ROCMFP4/2] = d1 * (float) rocmfp4_decode_i8(q >> 4); + } +} + +static __device__ void cpy_blck_rocmfp4_fast_f32(const char * cxi, char * cdsti) { + const block_rocmfp4_fast * x = (const block_rocmfp4_fast *) cxi; + float * cdstf = (float *) cdsti; + + const float d = rocmfp4_ue4m3_to_fp32_half_finite(x->e); + +#pragma unroll + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q = x->qs[j]; + cdstf[j] = d * (float) rocmfp4_decode_i8(q); + cdstf[j + QK_ROCMFP4/2] = d * (float) rocmfp4_decode_i8(q >> 4); + } +} + +static __global__ void cpy_rocmfp4_f32_contiguous(const block_rocmfp4 * cx, float * cdst, const int64_t ne) { + const int64_t packed_idx = (int64_t)blockDim.x * blockIdx.x + threadIdx.x; + const int64_t packed_count = (ne / QK_ROCMFP4) * (QK_ROCMFP4/2); + + if (packed_idx >= packed_count) { + return; + } + + const int64_t ib = packed_idx >> 4; + const int j = packed_idx & 0x0f; + const int64_t base = ib * QK_ROCMFP4; + const uint8_t q = cx[ib].qs[j]; + const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(cx[ib].e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(cx[ib].e[1]); + + cdst[base + j] = d0 * (float) rocmfp4_decode_i8(q); + cdst[base + j + QK_ROCMFP4/2] = d1 * (float) rocmfp4_decode_i8(q >> 4); +} + +static __global__ void cpy_rocmfp4_fast_f32_contiguous(const block_rocmfp4_fast * cx, float * cdst, const int64_t ne) { + const int64_t packed_idx = (int64_t)blockDim.x * blockIdx.x + threadIdx.x; + const int64_t packed_count = (ne / QK_ROCMFP4) * (QK_ROCMFP4/2); + + if (packed_idx >= packed_count) { + return; + } + + const int64_t ib = packed_idx >> 4; + const int j = packed_idx & 0x0f; + const int64_t base = ib * QK_ROCMFP4; + const uint8_t q = cx[ib].qs[j]; + const float d = rocmfp4_ue4m3_to_fp32_half_finite(cx[ib].e); + + cdst[base + j] = d * (float) rocmfp4_decode_i8(q); + cdst[base + j + QK_ROCMFP4/2] = d * (float) rocmfp4_decode_i8(q >> 4); +} + template static __global__ void cpy_f32_q(const char * cx, char * cdst, const int64_t ne, const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t nb00, const int64_t nb01, const int64_t nb02, @@ -308,6 +375,56 @@ static void ggml_cpy_q4_1_f32_cuda( ne10, ne11, ne12, nb10, nb11, nb12, nb13); } +static void ggml_cpy_rocmfp4_f32_hip( + const char * cx, char * cdst, const int64_t ne, + const int64_t ne00, const int64_t ne01, const int64_t ne02, + const int64_t nb00, const int64_t nb01, const int64_t nb02, + const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, + const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, + cudaStream_t stream) { + GGML_ASSERT(ne % QK_ROCMFP4 == 0); + const int64_t num_blocks = ne / QK_ROCMFP4; + GGML_ASSERT(num_blocks < UINT_MAX); + cpy_q_f32<<>>( + cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, + ne10, ne11, ne12, nb10, nb11, nb12, nb13); +} + +static void ggml_cpy_rocmfp4_f32_contiguous_hip( + const char * cx, char * cdst, const int64_t ne, cudaStream_t stream) { + GGML_ASSERT(ne % QK_ROCMFP4 == 0); + const int64_t packed_count = (ne / QK_ROCMFP4) * (QK_ROCMFP4/2); + const int64_t num_blocks = (packed_count + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE; + GGML_ASSERT(num_blocks < UINT_MAX); + cpy_rocmfp4_f32_contiguous<<>>( + (const block_rocmfp4 *) cx, (float *) cdst, ne); +} + +static void ggml_cpy_rocmfp4_fast_f32_hip( + const char * cx, char * cdst, const int64_t ne, + const int64_t ne00, const int64_t ne01, const int64_t ne02, + const int64_t nb00, const int64_t nb01, const int64_t nb02, + const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, + const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, + cudaStream_t stream) { + GGML_ASSERT(ne % QK_ROCMFP4 == 0); + const int64_t num_blocks = ne / QK_ROCMFP4; + GGML_ASSERT(num_blocks < UINT_MAX); + cpy_q_f32<<>>( + cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, + ne10, ne11, ne12, nb10, nb11, nb12, nb13); +} + +static void ggml_cpy_rocmfp4_fast_f32_contiguous_hip( + const char * cx, char * cdst, const int64_t ne, cudaStream_t stream) { + GGML_ASSERT(ne % QK_ROCMFP4 == 0); + const int64_t packed_count = (ne / QK_ROCMFP4) * (QK_ROCMFP4/2); + const int64_t num_blocks = (packed_count + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE; + GGML_ASSERT(num_blocks < UINT_MAX); + cpy_rocmfp4_fast_f32_contiguous<<>>( + (const block_rocmfp4_fast *) cx, (float *) cdst, ne); +} + static void ggml_cpy_f32_q5_0_cuda( const char * cx, char * cdst, const int64_t ne, const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t nb00, const int64_t nb01, const int64_t nb02, @@ -556,6 +673,20 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg } else if (src0->type == GGML_TYPE_Q4_0 && src1->type == GGML_TYPE_F32) { ggml_cpy_q4_0_f32_cuda (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream); + } else if (src0->type == GGML_TYPE_Q4_0_ROCMFP4 && src1->type == GGML_TYPE_F32) { + if (contiguous_srcs) { + ggml_cpy_rocmfp4_f32_contiguous_hip(src0_ddc, src1_ddc, ne, main_stream); + } else { + ggml_cpy_rocmfp4_f32_hip + (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream); + } + } else if (src0->type == GGML_TYPE_Q4_0_ROCMFP4_FAST && src1->type == GGML_TYPE_F32) { + if (contiguous_srcs) { + ggml_cpy_rocmfp4_fast_f32_contiguous_hip(src0_ddc, src1_ddc, ne, main_stream); + } else { + ggml_cpy_rocmfp4_fast_f32_hip + (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream); + } } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_Q4_1) { ggml_cpy_f32_q4_1_cuda (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/dequantize.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/dequantize.cuh index 1801bd048..dc660e1aa 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/dequantize.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/dequantize.cuh @@ -1,5 +1,6 @@ #include "common.cuh" #include "tq3-quant.cuh" +#include "../../rocmfp4/rocmfp4_hip_scale.cuh" static __device__ __forceinline__ void dequantize_q4_0(const void * vx, const int64_t ib, const int iqs, float2 & v){ const block_q4_0 * x = (const block_q4_0 *) vx; @@ -76,3 +77,125 @@ static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const in v.x *= d; v.y *= d; } + +static __device__ __forceinline__ void dequantize_rocmfp4(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const block_rocmfp4 * x = (const block_rocmfp4 *) vx; + + const int q = x[ib].qs[iqs]; + const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[1]); + + v.x = d0 * rocmfp4_decode_i8(q); + v.y = d1 * rocmfp4_decode_i8(q >> 4); +} + +static __device__ __forceinline__ void dequantize_rocmfp4_fast(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const block_rocmfp4_fast * x = (const block_rocmfp4_fast *) vx; + + const int q = x[ib].qs[iqs]; + const float d = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e); + + v.x = d * rocmfp4_decode_i8(q); + v.y = d * rocmfp4_decode_i8(q >> 4); +} + +template +static __device__ __forceinline__ uint32_t rocmfpx_load_qs_window_cuda(const uint8_t * src, const int byte_pos) { + uint32_t v = (uint32_t) src[byte_pos + 0]; + + if (byte_pos + 1 < qs) { + v |= (uint32_t) src[byte_pos + 1] << 8; + } + if (byte_pos + 2 < qs) { + v |= (uint32_t) src[byte_pos + 2] << 16; + } + + return v; +} + +static __device__ __forceinline__ uint32_t rocmfpx_get_fp3_code_cuda(const uint8_t * src, const int i) { + const int bit_pos = i * 3; + const int byte_pos = bit_pos >> 3; + const int shift = bit_pos & 7; + return (rocmfpx_load_qs_window_cuda(src, byte_pos) >> shift) & 7u; +} + +static __device__ __forceinline__ uint32_t rocmfpx_get_fp2_code_cuda(const uint8_t * src, const int i) { + return (src[i >> 2] >> (2*(i & 3))) & 3u; +} + +static __device__ __forceinline__ uint32_t rocmfpx_get_fp6_code_cuda(const uint8_t * src, const int i) { + const int bit_pos = i * 6; + const int byte_pos = bit_pos >> 3; + const int shift = bit_pos & 7; + return (rocmfpx_load_qs_window_cuda(src, byte_pos) >> shift) & 63u; +} + +static __device__ __forceinline__ int rocmfpx_decode_fp3_code_cuda(const uint32_t code) { + const uint32_t mag_code = code & 3u; + const int mag = mag_code == 3u ? 4 : (int) mag_code; + return (code & 4u) ? -mag : mag; +} + +static __device__ __forceinline__ float rocmfpx_decode_fp2_code_cuda(const uint32_t code) { + switch (code & 3u) { + case 0: return (float) ROCMFP2_KVALUE_0_I8; + case 1: return (float) ROCMFP2_KVALUE_1_I8; + case 2: return (float) ROCMFP2_KVALUE_2_I8; + default: return (float) ROCMFP2_KVALUE_3_I8; + } +} + +static __device__ __forceinline__ int rocmfpx_decode_fp6_code_cuda(const uint32_t code) { + const int mag = (int) (code & 31u); + return (code & 32u) ? -mag : mag; +} + +static __device__ __forceinline__ void dequantize_rocmfpx_fp3(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const block_rocmfp3 * x = (const block_rocmfp3 *) vx; + + const int i0 = iqs + 0; + const int i1 = iqs + 1; + const float d0 = rocmfpx_ue4m3_to_fp32_finite(x[ib].e[i0 >= QK_ROCMFP3/2]); + const float d1 = rocmfpx_ue4m3_to_fp32_finite(x[ib].e[i1 >= QK_ROCMFP3/2]); + + v.x = d0 * (float) rocmfpx_decode_fp3_code_cuda(rocmfpx_get_fp3_code_cuda(x[ib].qs, i0)); + v.y = d1 * (float) rocmfpx_decode_fp3_code_cuda(rocmfpx_get_fp3_code_cuda(x[ib].qs, i1)); +} + +static __device__ __forceinline__ void dequantize_rocmfpx_fp2(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const block_rocmfp2 * x = (const block_rocmfp2 *) vx; + + const int i0 = iqs + 0; + const int i1 = iqs + 1; + const float d0 = rocmfpx_ue4m3_to_fp32_finite(x[ib].e[i0 >= QK_ROCMFP2/2]); + const float d1 = rocmfpx_ue4m3_to_fp32_finite(x[ib].e[i1 >= QK_ROCMFP2/2]); + + v.x = d0 * rocmfpx_decode_fp2_code_cuda(rocmfpx_get_fp2_code_cuda(x[ib].qs, i0)); + v.y = d1 * rocmfpx_decode_fp2_code_cuda(rocmfpx_get_fp2_code_cuda(x[ib].qs, i1)); +} + +static __device__ __forceinline__ void dequantize_rocmfpx_fp6(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const block_rocmfp6_device * x = (const block_rocmfp6_device *) vx; + + const int i0 = iqs + 0; + const int i1 = iqs + 1; + const float d0 = rocmfpx_ue4m3_to_fp32_finite(x[ib].e[i0 >= QK_ROCMFP6/2]); + const float d1 = rocmfpx_ue4m3_to_fp32_finite(x[ib].e[i1 >= QK_ROCMFP6/2]); + +#if GGML_ROCMFP6_EXPANDED_DEVICE + v.x = d0 * (float) x[ib].qs[i0]; + v.y = d1 * (float) x[ib].qs[i1]; +#else + v.x = d0 * (float) rocmfpx_decode_fp6_code_cuda(rocmfpx_get_fp6_code_cuda(x[ib].qs, i0)); + v.y = d1 * (float) rocmfpx_decode_fp6_code_cuda(rocmfpx_get_fp6_code_cuda(x[ib].qs, i1)); +#endif +} + +static __device__ __forceinline__ void dequantize_rocmfpx_fp8(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const block_rocmfp8 * x = (const block_rocmfp8 *) vx; + + const float d = rocmfpx_ue4m3_to_fp32_finite(x[ib].e); + v.x = d * (float) x[ib].qs[iqs + 0]; + v.y = d * (float) x[ib].qs[iqs + 1]; +} 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 new file mode 100644 index 000000000..1c6b86972 --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu @@ -0,0 +1,367 @@ +#include "ds4-hc.cuh" + +// Fused DeepSeek4 hyper-connection ops. +// +// mode 0 (pre): src0 = mix [mix_dim] (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]: +// 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]: +// 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] +// src1 = base [n_hc] +// 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] + +#define DS4_HC_SINKHORN_EPS 1.0e-6f +#define DS4_HC_MAX_HC 8 +#define DS4_HC_MAX_MIX (2*DS4_HC_MAX_HC + DS4_HC_MAX_HC*DS4_HC_MAX_HC) + +static __device__ __forceinline__ float ds4_hc_sigmoid(float x) { + return 1.0f / (1.0f + expf(-x)); +} + +static __device__ void ds4_hc_sinkhorn_split( + const float * mix, + const float * base, + float pre_scale, + float post_scale, + float comb_scale, + int n_hc, + int iters, + float * split) { + for (int i = 0; i < n_hc; ++i) { + split[i] = ds4_hc_sigmoid(mix[i] * pre_scale + base[i]) + DS4_HC_SINKHORN_EPS; + } + for (int i = 0; i < n_hc; ++i) { + split[n_hc + i] = 2.0f * ds4_hc_sigmoid(mix[n_hc + i] * post_scale + base[n_hc + i]); + } + + float c[DS4_HC_MAX_HC * DS4_HC_MAX_HC]; + for (int dst_i = 0; dst_i < n_hc; ++dst_i) { + float row_max = -1.0e30f; + for (int src_i = 0; src_i < n_hc; ++src_i) { + const int idx = src_i + dst_i * n_hc; + const float v = mix[2 * n_hc + idx] * comb_scale + base[2 * n_hc + idx]; + c[idx] = v; + row_max = v > row_max ? v : row_max; + } + float row_sum = 0.0f; + for (int src_i = 0; src_i < n_hc; ++src_i) { + const int idx = src_i + dst_i * n_hc; + c[idx] = expf(c[idx] - row_max); + row_sum += c[idx]; + } + const float inv = 1.0f / row_sum; + for (int src_i = 0; src_i < n_hc; ++src_i) { + c[src_i + dst_i * n_hc] = c[src_i + dst_i * n_hc] * inv + DS4_HC_SINKHORN_EPS; + } + } + for (int src_i = 0; src_i < n_hc; ++src_i) { + float sum = 0.0f; + for (int dst_i = 0; dst_i < n_hc; ++dst_i) sum += c[src_i + dst_i * n_hc]; + const float inv = 1.0f / (sum + DS4_HC_SINKHORN_EPS); + for (int dst_i = 0; dst_i < n_hc; ++dst_i) c[src_i + dst_i * n_hc] *= inv; + } + for (int iter = 1; iter < iters; ++iter) { + for (int dst_i = 0; dst_i < n_hc; ++dst_i) { + float sum = 0.0f; + for (int src_i = 0; src_i < n_hc; ++src_i) sum += c[src_i + dst_i * n_hc]; + const float inv = 1.0f / (sum + DS4_HC_SINKHORN_EPS); + for (int src_i = 0; src_i < n_hc; ++src_i) c[src_i + dst_i * n_hc] *= inv; + } + for (int src_i = 0; src_i < n_hc; ++src_i) { + float sum = 0.0f; + for (int dst_i = 0; dst_i < n_hc; ++dst_i) sum += c[src_i + dst_i * n_hc]; + const float inv = 1.0f / (sum + DS4_HC_SINKHORN_EPS); + for (int dst_i = 0; dst_i < n_hc; ++dst_i) c[src_i + dst_i * n_hc] *= inv; + } + } + for (int i = 0; i < n_hc * n_hc; ++i) { + split[2 * n_hc + i] = c[i]; + } +} + +// Fully-unrolled variant: with compile-time NHC the c[] matrix lives in +// registers. The generic version's runtime-bound loops force c[] into scratch +// (private, VRAM-backed) memory, making the serial sinkhorn ~20x slower +// (97us vs 5us measured on gfx1151). +template +static __device__ void ds4_hc_sinkhorn_split_t( + const float * mix, + const float * base, + float pre_scale, + float post_scale, + float comb_scale, + int iters, + float * split) { + #pragma unroll + for (int i = 0; i < NHC; ++i) { + split[i] = ds4_hc_sigmoid(mix[i] * pre_scale + base[i]) + DS4_HC_SINKHORN_EPS; + } + #pragma unroll + for (int i = 0; i < NHC; ++i) { + split[NHC + i] = 2.0f * ds4_hc_sigmoid(mix[NHC + i] * post_scale + base[NHC + i]); + } + float c[NHC * NHC]; + #pragma unroll + for (int dst_i = 0; dst_i < NHC; ++dst_i) { + float row_max = -1.0e30f; + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) { + const int idx = src_i + dst_i * NHC; + const float v = mix[2 * NHC + idx] * comb_scale + base[2 * NHC + idx]; + c[idx] = v; + row_max = v > row_max ? v : row_max; + } + float row_sum = 0.0f; + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) { + const int idx = src_i + dst_i * NHC; + c[idx] = expf(c[idx] - row_max); + row_sum += c[idx]; + } + const float inv = 1.0f / row_sum; + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) { + c[src_i + dst_i * NHC] = c[src_i + dst_i * NHC] * inv + DS4_HC_SINKHORN_EPS; + } + } + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) { + float sum = 0.0f; + #pragma unroll + for (int dst_i = 0; dst_i < NHC; ++dst_i) sum += c[src_i + dst_i * NHC]; + const float inv = 1.0f / (sum + DS4_HC_SINKHORN_EPS); + #pragma unroll + for (int dst_i = 0; dst_i < NHC; ++dst_i) c[src_i + dst_i * NHC] *= inv; + } + for (int iter = 1; iter < iters; ++iter) { + #pragma unroll + for (int dst_i = 0; dst_i < NHC; ++dst_i) { + float sum = 0.0f; + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) sum += c[src_i + dst_i * NHC]; + const float inv = 1.0f / (sum + DS4_HC_SINKHORN_EPS); + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) c[src_i + dst_i * NHC] *= inv; + } + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) { + float sum = 0.0f; + #pragma unroll + for (int dst_i = 0; dst_i < NHC; ++dst_i) sum += c[src_i + dst_i * NHC]; + const float inv = 1.0f / (sum + DS4_HC_SINKHORN_EPS); + #pragma unroll + for (int dst_i = 0; dst_i < NHC; ++dst_i) c[src_i + dst_i * NHC] *= inv; + } + } + #pragma unroll + for (int i = 0; i < NHC * NHC; ++i) { + split[2 * NHC + i] = c[i]; + } +} + +template +static __global__ void ds4_hc_pre_kernel_t( + const float * __restrict__ mix, + const float * __restrict__ base, + const float * __restrict__ hc_state, + float * __restrict__ dst, + int n_embd, + int iters, + float pre_scale, + float post_scale, + float comb_scale) { + __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 = 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); + if (blockIdx.x == 0) { + #pragma unroll + for (int i = 0; i < mix_dim; ++i) { + dst[n_embd + i] = split[i]; + } + } + } + __syncthreads(); + + const int d = (int) blockIdx.x * blockDim.x + tid; + if (d < n_embd) { + float acc = 0.0f; + #pragma unroll + for (int h = 0; h < NHC; ++h) { + acc += split[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, + const float * __restrict__ hc_state, + float * __restrict__ dst, + int n_embd, + int n_hc, + int iters, + float pre_scale, + float post_scale, + float comb_scale) { + __shared__ float split[DS4_HC_MAX_MIX]; + __shared__ float s_mix[DS4_HC_MAX_MIX]; + __shared__ float s_base[DS4_HC_MAX_MIX]; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + const int tid = threadIdx.x; + + // Stage mix/base cooperatively: base lives in managed (UMA) memory where + // serial scalar loads cost ~2us each; one parallel coalesced load instead. + if (tid < mix_dim) { + s_mix[tid] = mix[tid]; + s_base[tid] = base[tid]; + } + __syncthreads(); + + // Each block redoes the (tiny) sinkhorn into shared memory so the mix + // loop below can spread across the whole GPU instead of one CU. + if (tid == 0) { + ds4_hc_sinkhorn_split(s_mix, s_base, pre_scale, post_scale, comb_scale, n_hc, iters, split); + if (blockIdx.x == 0) { + for (int i = 0; i < mix_dim; ++i) { + dst[n_embd + i] = split[i]; + } + } + } + __syncthreads(); + + const int d = (int) blockIdx.x * blockDim.x + tid; + if (d < n_embd) { + float acc = 0.0f; + for (int h = 0; h < n_hc; ++h) { + acc += split[h] * hc_state[(size_t) h * n_embd + d]; + } + dst[d] = acc; + } +} + +static __global__ void ds4_hc_post_kernel( + const float * __restrict__ residual, + const float * __restrict__ block_out, + 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; + float acc = block_out[d] * 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, + const float * __restrict__ hc_state, + float * __restrict__ dst, + int n_embd, + int n_hc, + float pre_scale) { + const int d = blockIdx.x * blockDim.x + threadIdx.x; + if (d >= n_embd) { + return; + } + float acc = 0.0f; + for (int h = 0; h < n_hc; ++h) { + const float wgt = ds4_hc_sigmoid(mix[h] * pre_scale + base[h]) + DS4_HC_SINKHORN_EPS; + acc += wgt * hc_state[(size_t) h * n_embd + d]; + } + dst[d] = acc; +} + +void ggml_cuda_op_ds4_hc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + const ggml_tensor * src2 = dst->src[2]; + + GGML_ASSERT(src0 && src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1 && src1->type == GGML_TYPE_F32); + GGML_ASSERT(src2 && src2->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + 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); + + GGML_ASSERT(n_hc > 0 && n_hc <= DS4_HC_MAX_HC); + + cudaStream_t stream = ctx.stream(); + + switch (mode) { + case 0: { + const int iters = ggml_get_op_params_i32(dst, 3); + const float pre_scale = ggml_get_op_params_f32(dst, 4); + 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 float *) src0->data, (const float *) src1->data, + (const float *) src2->data, (float *) dst->data, + n_embd, iters, pre_scale, post_scale, comb_scale); + } else { + 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); + } + } break; + case 1: { + const int total = n_embd * n_hc; + const int blocks = (total + 255) / 256; + ds4_hc_post_kernel<<>>( + (const float *) src0->data, (const float *) src1->data, + (const float *) src2->data, (float *) dst->data, + n_embd, n_hc); + } 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 float *) src0->data, (const float *) src1->data, + (const float *) src2->data, (float *) dst->data, + n_embd, n_hc, pre_scale); + } break; + default: + GGML_ABORT("ds4_hc: unknown mode"); + } +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cuh new file mode 100644 index 000000000..1a2cdf5eb --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cuh @@ -0,0 +1,10 @@ +#pragma once + +#include "ggml-cuda/common.cuh" + +// Fused DeepSeek4 hyper-connection (HC) decode ops. See ds4-hc.cu for the +// per-mode contract (pre / post / out). Used by the opt-in +// DFLASH_DS4_FUSED_DECODE single-graph decode path; output is deterministic +// but not bit-identical to the CPU HC reference (expf ULP differences +// amplified by the sinkhorn iterations). +void ggml_cuda_op_ds4_hc(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..e991bf4a3 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,133 @@ #include "fattn-chunked.cuh" #include "fattn.cuh" +__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]; +} + +__global__ static void ds4_flash_attn_d512_f32_shared_kv_kernel( + float * dst, + const float * q, + const float * k, + const float * v, + const float * sinks, + int n_tokens, + int n_heads, + int n_kv) { + 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 scores[]; + const float * qh = q + ((size_t) h * (size_t) n_tokens + (size_t) t) * D; + const float scale = rsqrtf((float) D); + + float local_max = sinks ? sinks[h] : -3.402823466e38f; + for (int r = tid; r < n_kv; r += blockDim.x) { + const float * kr = k + (size_t) r * D; + float dot = 0.0f; +#pragma unroll + for (int d = 0; d < D; ++d) { + dot += qh[d] * kr[d]; + } + const float s = dot * scale; + 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; + + for (int d = tid; d < D; d += blockDim.x) { + float acc = 0.0f; + for (int r = 0; r < n_kv; ++r) { + acc += scores[r] * v[(size_t) r * D + d]; + } + dst[((size_t) t * (size_t) n_heads + (size_t) h) * D + d] = acc * inv_denom; + } +} + +static bool ggml_cuda_ds4_flash_attn_d512_f32(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const char * e = getenv("DFLASH_DS4_FLASH_ATTN"); + if (!e || !e[0] || e[0] == (char)48) { + 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]; + if (!Q || !K || !V || mask || + Q->type != GGML_TYPE_F32 || K->type != GGML_TYPE_F32 || V->type != GGML_TYPE_F32 || + dst->type != GGML_TYPE_F32 || + Q->ne[0] != 512 || K->ne[0] != 512 || V->ne[0] != 512 || + 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] || + Q->nb[0] != (int64_t) sizeof(float) || + K->nb[0] != (int64_t) sizeof(float) || + V->nb[0] != (int64_t) sizeof(float) || + dst->nb[0] != (int64_t) sizeof(float)) { + return false; + } + if (sinks && (sinks->type != GGML_TYPE_F32 || sinks->ne[0] != Q->ne[2])) { + 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; + } + + cudaStream_t stream = ctx.stream(); + dim3 grid((unsigned) n_tokens, (unsigned) n_heads, 1); + const size_t shmem = (size_t) n_kv * sizeof(float); + ds4_flash_attn_d512_f32_shared_kv_kernel<<>>( + (float *) dst->data, + (const float *) Q->data, + (const float *) K->data, + (const float *) V->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, + n_heads, + n_kv); + return true; +} + 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; @@ -354,6 +481,10 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const } } } + static const bool ds4_allow_d512_nomask = [] { + const char * e = getenv("DFLASH_DS4_FLASH_ATTN"); + return e && e[0] && e[0] != (char)48; + }(); const int cc = ggml_cuda_info().devices[device].cc; @@ -374,7 +505,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const if (V->ne[0] != K->ne[0]) { return BEST_FATTN_KERNEL_NONE; } - if (!gqa_opt_applies) { + if (!gqa_opt_applies && !ds4_allow_d512_nomask) { return BEST_FATTN_KERNEL_NONE; } break; @@ -570,6 +701,9 @@ 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_cuda_ds4_flash_attn_d512_f32(ctx, dst)) { + return; + } switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error"); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/getrows.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/getrows.cu index 1d7c6d17a..73a28e529 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/getrows.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/getrows.cu @@ -199,6 +199,30 @@ static void ggml_cuda_get_rows_switch_src0_type( get_rows_cuda_q(src0_d, src1_d, dst_d, ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); break; + case GGML_TYPE_Q4_0_ROCMFP4: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q2_0_ROCMFP2: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q3_0_ROCMFPX: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q6_0_ROCMFPX: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q8_0_ROCMFPX: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; case GGML_TYPE_TQ3_0: get_rows_cuda_q(src0_d, src1_d, dst_d, ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); 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 feb9104a9..e5463faf9 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 @@ -64,6 +64,7 @@ #include "ggml-cuda/cumsum.cuh" #include "ggml-cuda/fill.cuh" #include "ggml-cuda/moe-fused.cuh" +#include "ggml-cuda/ds4-hc.cuh" #include "ggml.h" #include @@ -2361,7 +2362,12 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, return false; } - static constexpr std::array valid_glu_ops = { GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_SWIGLU_OAI }; + static constexpr std::array valid_glu_ops = { + GGML_GLU_OP_SWIGLU, + GGML_GLU_OP_GEGLU, + GGML_GLU_OP_SWIGLU_OAI, + GGML_GLU_OP_SWIGLU_DS4, + }; if (std::find(valid_glu_ops.begin(), valid_glu_ops.end(), ggml_get_glu_op(glu)) == valid_glu_ops.end()) { return false; @@ -2417,6 +2423,12 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { return use_mul_mat_vec_f; } +static inline void ggml_cuda_set_fusion_glu_params(ggml_cuda_mm_fusion_args_host & fusion_data, const ggml_tensor * glu) { + fusion_data.glu_op = ggml_get_glu_op(glu); + fusion_data.glu_param0 = ggml_get_op_params_f32(glu, 2); + fusion_data.glu_param1 = ggml_get_op_params_f32(glu, 3); +} + static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { ggml_tensor * src0 = tensor->src[0]; ggml_tensor * src1 = tensor->src[1]; @@ -2850,6 +2862,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_GLU_OP_SWIGLU_OAI: ggml_cuda_op_swiglu_oai(ctx, dst); break; + case GGML_GLU_OP_SWIGLU_DS4: + ggml_cuda_op_swiglu_ds4(ctx, dst); + break; case GGML_GLU_OP_GEGLU_ERF: ggml_cuda_op_geglu_erf(ctx, dst); break; @@ -2869,6 +2884,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_MOE_FUSED: ggml_cuda_op_moe_fused(ctx, dst); break; + case GGML_OP_DS4_HC: + ggml_cuda_op_ds4_hc(ctx, dst); + break; case GGML_OP_GROUP_NORM: ggml_cuda_op_group_norm(ctx, dst); break; @@ -4071,7 +4089,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud fusion_data.gate = gate_n->src[0]; fusion_data.x_bias = up_bias_tensor; fusion_data.gate_bias = gate_bias_tensor; - fusion_data.glu_op = ggml_get_glu_op(glu); + 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; @@ -4084,7 +4102,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud fusion_data.gate = gate_n->src[0]; fusion_data.x_bias = up_bias_tensor; fusion_data.gate_bias = gate_bias_tensor; - fusion_data.glu_op = ggml_get_glu_op(glu); + ggml_cuda_set_fusion_glu_params(fusion_data, glu); ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); fused_mul_mat_vec = true; @@ -4108,7 +4126,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) { ggml_cuda_mm_fusion_args_host fusion_data{}; fusion_data.gate = gate->src[0]; - fusion_data.glu_op = ggml_get_glu_op(glu); + 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; @@ -4119,7 +4137,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) { ggml_cuda_mm_fusion_args_host fusion_data{}; fusion_data.gate = gate->src[0]; - fusion_data.glu_op = ggml_get_glu_op(glu); + ggml_cuda_set_fusion_glu_params(fusion_data, glu); ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); fused_mul_mat_vec = true; @@ -4987,6 +5005,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_GLU_OP_GEGLU: case GGML_GLU_OP_SWIGLU: case GGML_GLU_OP_SWIGLU_OAI: + case GGML_GLU_OP_SWIGLU_DS4: case GGML_GLU_OP_GEGLU_ERF: case GGML_GLU_OP_GEGLU_QUICK: return ggml_is_contiguous_1(op->src[0]); @@ -5000,6 +5019,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return op->src[0]->nb[0] == ggml_type_size(op->src[0]->type); case GGML_OP_MOE_FUSED: return true; + case GGML_OP_DS4_HC: + return true; case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_ID: { @@ -5045,6 +5066,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_TYPE_Q8_0: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q2_0_ROCMFP2: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -5080,6 +5107,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_TYPE_Q5_0: case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q2_0_ROCMFP2: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_TQ3_0: return true; default: @@ -5186,7 +5219,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_CONCAT: { ggml_type src0_type = op->src[0]->type; - return src0_type != GGML_TYPE_I32 && src0_type != GGML_TYPE_I16; + return src0_type == GGML_TYPE_F32 || + src0_type == GGML_TYPE_F16 || + src0_type == GGML_TYPE_BF16 || + src0_type == GGML_TYPE_I8; } break; case GGML_OP_CONV_TRANSPOSE_1D: { diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu index d91472024..6e4db8e8b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvf.cu @@ -51,6 +51,8 @@ static __global__ void mul_mat_vec_f( bool use_bias = false; bool use_gate_bias = false; ggml_glu_op glu_op = ggml_glu_op::GGML_GLU_OP_SWIGLU; + float glu_param0 = 0.0f; + float glu_param1 = 0.0f; const T * gate_x = nullptr; const float * x_bias = nullptr; const float * gate_bias = nullptr; @@ -60,6 +62,8 @@ static __global__ void mul_mat_vec_f( use_bias = fusion.x_bias != nullptr; use_gate_bias = fusion.gate_bias != nullptr; glu_op = fusion.glu_op; + glu_param0 = fusion.glu_param0; + glu_param1 = fusion.glu_param1; if (use_gate) { gate_x = static_cast(fusion.gate); @@ -357,7 +361,11 @@ static __global__ void mul_mat_vec_f( value *= ggml_cuda_op_gelu_single(gate_value); break; case GGML_GLU_OP_SWIGLU_OAI: { - value = ggml_cuda_op_swiglu_oai_single(gate_value, value); + value = ggml_cuda_op_swiglu_oai_single(gate_value, value, glu_param0, glu_param1); + break; + } + case GGML_GLU_OP_SWIGLU_DS4: { + value = ggml_cuda_op_swiglu_ds4_single(gate_value, value, glu_param0); break; } default: @@ -369,7 +377,7 @@ static __global__ void mul_mat_vec_f( dst[tid*stride_col_dst + row] = value; if constexpr (!has_fusion) { - GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, glu_op, gate_x, x_bias, gate_bias, sumf_gate); + GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, glu_op, glu_param0, glu_param1, gate_x, x_bias, gate_bias, sumf_gate); } } @@ -668,6 +676,8 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor fusion_local.gate_bias = fusion->gate_bias->data; } fusion_local.glu_op = fusion->glu_op; + fusion_local.glu_param0 = fusion->glu_param0; + fusion_local.glu_param1 = fusion->glu_param1; } const int64_t s01 = src0->nb[1] / ts_src0; 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 78dcf500f..620a1b3fb 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -17,6 +17,12 @@ static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) case GGML_TYPE_Q8_0: return vec_dot_q8_0_q8_1; case GGML_TYPE_MXFP4: return vec_dot_mxfp4_q8_1; case GGML_TYPE_NVFP4: return vec_dot_nvfp4_q8_1; + case GGML_TYPE_Q4_0_ROCMFP4: return vec_dot_rocmfp4_q8_1; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: return vec_dot_rocmfp4_fast_q8_1; + case GGML_TYPE_Q2_0_ROCMFP2: return vec_dot_rocmfpx_fp2_q8_1; + case GGML_TYPE_Q3_0_ROCMFPX: return vec_dot_rocmfpx_fp3_q8_1; + case GGML_TYPE_Q6_0_ROCMFPX: return vec_dot_rocmfpx_fp6_q8_1; + case GGML_TYPE_Q8_0_ROCMFPX: return vec_dot_rocmfpx_fp8_q8_1; case GGML_TYPE_Q2_K: return vec_dot_q2_K_q8_1; case GGML_TYPE_Q3_K: return vec_dot_q3_K_q8_1; case GGML_TYPE_Q4_K: return vec_dot_q4_K_q8_1; @@ -44,6 +50,12 @@ static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) { case GGML_TYPE_Q8_0: return VDR_Q8_0_Q8_1_MMVQ; case GGML_TYPE_MXFP4: return VDR_MXFP4_Q8_1_MMVQ; case GGML_TYPE_NVFP4: return VDR_NVFP4_Q8_1_MMVQ; + case GGML_TYPE_Q4_0_ROCMFP4: return VDR_ROCMFP4_Q8_1_MMVQ; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: return VDR_ROCMFP4_FAST_Q8_1_MMVQ; + case GGML_TYPE_Q2_0_ROCMFP2: return VDR_ROCMFP2_Q8_1_MMVQ; + case GGML_TYPE_Q3_0_ROCMFPX: return VDR_ROCMFP3_Q8_1_MMVQ; + case GGML_TYPE_Q6_0_ROCMFPX: return VDR_ROCMFP6_Q8_1_MMVQ; + case GGML_TYPE_Q8_0_ROCMFPX: return VDR_ROCMFP8_Q8_1_MMVQ; case GGML_TYPE_Q2_K: return VDR_Q2_K_Q8_1_MMVQ; case GGML_TYPE_Q3_K: return VDR_Q3_K_Q8_1_MMVQ; case GGML_TYPE_Q4_K: return VDR_Q4_K_Q8_1_MMVQ; @@ -454,6 +466,8 @@ static __global__ void mul_mat_vec_q( 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; if constexpr (has_fusion) { use_gate = fusion.gate != nullptr; @@ -463,6 +477,8 @@ static __global__ void mul_mat_vec_q( 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; } @@ -596,7 +612,11 @@ static __global__ void mul_mat_vec_q( 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: { + result = ggml_cuda_op_swiglu_ds4_single(gate_value, result, glu_param0); break; } default: @@ -610,7 +630,7 @@ static __global__ void mul_mat_vec_q( } if constexpr (!has_fusion) { - GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, active_glu, gate_bias, x_bias, tmp_gate); + GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, active_glu, glu_param0, glu_param1, gate_bias, x_bias, tmp_gate); } } @@ -660,6 +680,8 @@ static __global__ void mul_mat_vec_q_moe( 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; if constexpr (has_fusion) { use_gate = fusion.gate != nullptr; @@ -669,6 +691,8 @@ static __global__ void mul_mat_vec_q_moe( 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; } // partial sum for each thread @@ -721,7 +745,10 @@ static __global__ void mul_mat_vec_q_moe( 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: + result = ggml_cuda_op_swiglu_ds4_single(gate_value, result, glu_param0); break; default: result = result * gate_value; @@ -733,7 +760,7 @@ static __global__ void mul_mat_vec_q_moe( } 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, tmp_gate); } } @@ -1082,6 +1109,42 @@ static void mul_mat_vec_q_switch_type( nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); break; + case GGML_TYPE_Q4_0_ROCMFP4: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q2_0_ROCMFP2: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q3_0_ROCMFPX: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q6_0_ROCMFPX: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q8_0_ROCMFPX: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; case GGML_TYPE_Q2_K: mul_mat_vec_q_switch_ncols_dst (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, @@ -1219,6 +1282,8 @@ void ggml_cuda_mul_mat_vec_q( fusion_local.gate_bias = fusion->gate_bias->data; } fusion_local.glu_op = fusion->glu_op; + fusion_local.glu_param0 = fusion->glu_param0; + fusion_local.glu_param1 = fusion->glu_param1; } // If src0 is a temporary compute buffer, clear any potential padding. diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cu index 4ad30fa1f..ed7f3d02a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cu @@ -338,6 +338,63 @@ void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary_gated(ctx, dst); } +// swiglu_ds4 + +template +static __global__ void swiglu_ds4_kernel(const T * gate, const T * up, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, float limit) { + const int64_t i = int64_t(blockDim.x)*blockIdx.x + threadIdx.x; + + if (i >= k) { + return; + } + + const int64_t j0 = (i / n) * o0 + (i % n); + const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n); + + const float gate_v = gate[j0]; + const float up_v = up[j1]; + + dst[i] = ggml_cuda_op_swiglu_ds4_single(gate_v, up_v, limit); +} + +template +static void swiglu_ds4_cuda(const T * gate, const T * up, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, const float limit, cudaStream_t stream) { + const int64_t num_blocks = (k + CUDA_GLU_BLOCK_SIZE - 1) / CUDA_GLU_BLOCK_SIZE; + swiglu_ds4_kernel<<>>(gate, up, dst, k, n, o0, o1, limit); +} + +void ggml_cuda_op_swiglu_ds4(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + void * src0_d = src0->data; + void * src1_d = src1 ? src1->data : src0->data; + const int64_t src0_o = src0->nb[1]; + const int64_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; + void * dst_d = dst->data; + const int64_t nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + cudaStream_t stream = ctx.stream(); + + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(src0->nb[0] == ggml_element_size(src0)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == dst->type); + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == ggml_nrows(src0)); + + GGML_ASSERT(src1); + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src1->nb[0] == ggml_element_size(src1)); + GGML_ASSERT(src1->ne[0] == nc); + GGML_ASSERT(src0->type == src1->type); + + const float limit = ggml_get_op_params_f32(dst, 2); + + swiglu_ds4_cuda((float *) src0_d, (float *) src1_d, (float *)dst_d, ggml_nelements(dst), nc, src0_o / sizeof(float), src1_o / sizeof(float), limit, stream); +} + void ggml_cuda_op_geglu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary_gated(ctx, dst); } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cuh index f1dd2183a..1864ce2dc 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/unary.cuh @@ -81,6 +81,8 @@ void ggml_cuda_op_geglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +void ggml_cuda_op_swiglu_ds4(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + void ggml_cuda_op_swiglu_oai(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_geglu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst); @@ -110,3 +112,11 @@ __device__ __forceinline__ float ggml_cuda_op_swiglu_oai_single(float x, float g out_glu = out_glu * (1.0f + g); return out_glu; } + +__device__ __forceinline__ float ggml_cuda_op_swiglu_ds4_single(float gate, float up, float limit) { + gate = fminf(gate, limit); + up = fmaxf(fminf(up, limit), -limit); + + const float silu = gate / (1.0f + expf(-gate)); + return silu * up; +} 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 40b2b41e7..8e8e79eb2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh @@ -1,6 +1,7 @@ #pragma once #include "common.cuh" +#include "../../rocmfp4/rocmfp4_hip_codebook.cuh" #include @@ -322,6 +323,380 @@ static __device__ __forceinline__ float vec_dot_mxfp4_q8_1( return d * sumi; } +// === ROCMFP vec_dot device functions (ported from ROCmFPX) === +#ifndef GGML_ROCMFP4_Q8_1_MMQ_VDR +#define GGML_ROCMFP4_Q8_1_MMQ_VDR 8 +#endif + +#ifndef GGML_ROCMFP4_FAST_Q8_1_MMQ_VDR +#define GGML_ROCMFP4_FAST_Q8_1_MMQ_VDR GGML_ROCMFP4_Q8_1_MMQ_VDR +#endif + +#ifndef GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR +#define GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR 2 +#endif + +#if GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR != 1 && \ + GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR != 2 && \ + GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR != 4 +#error "GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR must be 1, 2, or 4" +#endif + +#define VDR_ROCMFP4_Q8_1_MMVQ 4 +#define VDR_ROCMFP4_Q8_1_MMQ GGML_ROCMFP4_Q8_1_MMQ_VDR +#define VDR_ROCMFP4_FAST_Q8_1_MMVQ GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR +#define VDR_ROCMFP4_FAST_Q8_1_MMQ GGML_ROCMFP4_FAST_Q8_1_MMQ_VDR +#define VDR_ROCMFP2_Q8_1_MMVQ 1 +#define VDR_ROCMFP3_Q8_1_MMVQ 2 +#ifndef VDR_ROCMFP6_Q8_1_MMVQ +#define VDR_ROCMFP6_Q8_1_MMVQ 4 +#endif +#ifndef GGML_ROCMFP6_FAST_SIGNMAG_PACK +#define GGML_ROCMFP6_FAST_SIGNMAG_PACK 0 +#endif +#ifndef GGML_ROCMFP6_MMVQ_HALF_BLOCK_SPLIT +// Enabled by default: hoisting the half-block scale selection out of the FP6 +// MMVQ dot loop measured ~1.8% faster tg on gfx1151 (Qwen3-0.6B Q6_0_ROCMFPX, +// 211.8 -> 215.7 t/s) with bit-identical results. Set to 0 to restore the +// per-element branch. +#define GGML_ROCMFP6_MMVQ_HALF_BLOCK_SPLIT 1 +#endif +#define VDR_ROCMFP8_Q8_1_MMVQ 2 + +#define VDR_ROCMFP2_Q8_1_MMQ 4 +#define VDR_ROCMFP3_Q8_1_MMQ 4 +#ifndef VDR_ROCMFP6_Q8_1_MMQ +#define VDR_ROCMFP6_Q8_1_MMQ 4 +#endif +#define VDR_ROCMFP8_Q8_1_MMQ 8 + +static __device__ __forceinline__ uint32_t rocmfpx_get_bits_vec_cuda(const uint8_t * src, const int bit_pos, const int nbits) { + uint32_t code = 0; + +#pragma unroll + for (int bit = 0; bit < nbits; ++bit) { + const int src_bit = bit_pos + bit; + code |= ((uint32_t) ((src[src_bit >> 3] >> (src_bit & 7)) & 1u)) << bit; + } + + return code; +} + +static __device__ __forceinline__ int rocmfpx_decode_fp3_code_vec_cuda(const uint32_t code) { + const uint32_t mag_code = code & 3u; + const int mag = mag_code == 3u ? 4 : (int) mag_code; + return (code & 4u) ? -mag : mag; +} + +static __device__ __forceinline__ int rocmfpx_decode_fp2_code_vec_cuda(const uint32_t code) { + switch (code & 3u) { + case 0: return ROCMFP2_KVALUE_0_I8; + case 1: return ROCMFP2_KVALUE_1_I8; + case 2: return ROCMFP2_KVALUE_2_I8; + default: return ROCMFP2_KVALUE_3_I8; + } +} + +static __device__ __forceinline__ int rocmfpx_decode_fp6_code_vec_cuda(const uint32_t code) { +#if GGML_ROCMFP6_FAST_SIGNMAG_PACK + const int mag = (int) (code & 31u); + const int sign = -((int) ((code >> 5) & 1u)); + return (mag ^ sign) - sign; +#else + const int mag = (int) (code & 31u); + return (code & 32u) ? -mag : mag; +#endif +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp6_bits24_vec_cuda(const uint32_t bits24) { + const char4 v = make_char4( + (int8_t) rocmfpx_decode_fp6_code_vec_cuda(bits24 & 63u), + (int8_t) rocmfpx_decode_fp6_code_vec_cuda((bits24 >> 6) & 63u), + (int8_t) rocmfpx_decode_fp6_code_vec_cuda((bits24 >> 12) & 63u), + (int8_t) rocmfpx_decode_fp6_code_vec_cuda((bits24 >> 18) & 63u)); + return *((const int *) &v); +} + +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); +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp2_bits8_vec_cuda(const uint32_t bits8) { +#if defined(GGML_USE_HIP) + const uint32_t values = + ((uint32_t) (uint8_t) (int8_t) ROCMFP2_KVALUE_0_I8) | + ((uint32_t) (uint8_t) (int8_t) ROCMFP2_KVALUE_1_I8 << 8) | + ((uint32_t) (uint8_t) (int8_t) ROCMFP2_KVALUE_2_I8 << 16) | + ((uint32_t) (uint8_t) (int8_t) ROCMFP2_KVALUE_3_I8 << 24); + const uint32_t selectors = + ((bits8 >> 0) & 3u) | + (((bits8 >> 2) & 3u) << 8) | + (((bits8 >> 4) & 3u) << 16) | + (((bits8 >> 6) & 3u) << 24); + return (int) __builtin_amdgcn_perm(0, values, selectors); +#else + const char4 v = make_char4( + (int8_t) rocmfpx_decode_fp2_code_vec_cuda((bits8 >> 0) & 3u), + (int8_t) rocmfpx_decode_fp2_code_vec_cuda((bits8 >> 2) & 3u), + (int8_t) rocmfpx_decode_fp2_code_vec_cuda((bits8 >> 4) & 3u), + (int8_t) rocmfpx_decode_fp2_code_vec_cuda((bits8 >> 6) & 3u)); + return *((const int *) &v); +#endif +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp2_vec_cuda(const uint8_t * qs, const int base) { + return rocmfpx_pack4_fp2_bits8_vec_cuda(qs[base >> 2]); +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp6_vec_cuda(const uint8_t * qs, const int base) { +#if GGML_ROCMFP6_FAST_SIGNMAG_PACK + uint32_t qs0, qs1, qs2, qs3, qs4, qs5; + memcpy(&qs0, qs + 0, 4); + memcpy(&qs1, qs + 4, 4); + memcpy(&qs2, qs + 8, 4); + memcpy(&qs3, qs + 12, 4); + memcpy(&qs4, qs + 16, 4); + memcpy(&qs5, qs + 20, 4); + + const uint32_t words[7] = { qs0, qs1, qs2, qs3, qs4, qs5, 0 }; + const int start_bit = 6 * base; + const int reg_idx = start_bit >> 5; + const int reg_shift = start_bit & 31; + const uint32_t val_low = words[reg_idx]; + const uint32_t val_high = words[reg_idx + 1]; + const uint32_t bits24 = (reg_shift == 0) ? (val_low & 0xFFFFFFu) : + (((val_low >> reg_shift) | (val_high << (32 - reg_shift))) & 0xFFFFFFu); + + return rocmfpx_pack4_fp6_bits24_vec_cuda(bits24); +#else + const char4 v = make_char4( + (int8_t) rocmfpx_decode_fp6_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 0)*6, 6)), + (int8_t) rocmfpx_decode_fp6_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 1)*6, 6)), + (int8_t) rocmfpx_decode_fp6_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 2)*6, 6)), + (int8_t) rocmfpx_decode_fp6_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 3)*6, 6))); + return *((const int *) &v); +#endif +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp6_expanded_vec_cuda(const int8_t * qs, const int base) { + const char4 v = make_char4(qs[base + 0], qs[base + 1], qs[base + 2], qs[base + 3]); + return *((const int *) &v); +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp6_device_vec_cuda(const block_rocmfp6_device * bq6, const int base) { +#if GGML_ROCMFP6_EXPANDED_DEVICE + return rocmfpx_pack4_fp6_expanded_vec_cuda(bq6->qs, base); +#else + return rocmfpx_pack4_fp6_vec_cuda(bq6->qs, base); +#endif +} + +static __device__ __forceinline__ float vec_dot_rocmfp4_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_rocmfp4 * bq4 = (const block_rocmfp4 *) vbq + kbx; + + const int * q8 = (const int *) bq8_1->qs + iqs; + + int sumi0 = 0; + int sumi1 = 0; +#pragma unroll + for (int l = 0; l < VDR_ROCMFP4_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, q8[l + 0], sumi0); + sumi1 = ggml_cuda_dp4a(v.y, q8[l + 4], sumi1); + } + + const float db = __low2float(bq8_1->ds); + return db * (rocmfp4_ue4m3_to_fp32_half_finite(bq4->e[0]) * sumi0 + rocmfp4_ue4m3_to_fp32_half_finite(bq4->e[1]) * sumi1); +} + +static __device__ __forceinline__ float vec_dot_rocmfp4_fast_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_rocmfp4_fast * bq4 = (const block_rocmfp4_fast *) vbq + kbx; + + const int * q8 = (const int *) bq8_1->qs + iqs; + + int sumi = 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); + + sumi = ggml_cuda_dp4a(v.x, q8[l + 0], sumi); + sumi = ggml_cuda_dp4a(v.y, q8[l + 4], sumi); + } + + return __low2float(bq8_1->ds) * rocmfp4_ue4m3_to_fp32_half_finite(bq4->e) * sumi; +} + +static __device__ __forceinline__ float vec_dot_rocmfpx_fp2_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + // int32-consistent MMVQ layout: QI_ROCMFP2=2 (8B qs = 2 int32), VDR=1. + // iqs in {0,1} selects one 16-weight half-block: qs bytes [4*iqs .. 4*iqs+3] + scale e[iqs]. + // qs is a 10-byte-strided, byte-aligned array -> read byte-wise (uint8_t index); NEVER cast qs to int*/int2*. + const block_rocmfp2 * bq2 = (const block_rocmfp2 *) vbq + kbx; + + int sumi = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int val_packed = rocmfpx_pack4_fp2_bits8_vec_cuda((uint32_t) bq2->qs[4*iqs + j]); + 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; +} + +static __device__ __forceinline__ float vec_dot_rocmfpx_fp3_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_rocmfp3 * bq3 = (const block_rocmfp3 *) vbq + kbx; + + uint32_t qs0, qs1, qs2; + memcpy(&qs0, bq3->qs + 0, 4); + memcpy(&qs1, bq3->qs + 4, 4); + memcpy(&qs2, bq3->qs + 8, 4); + + const uint32_t qs[4] = { qs0, qs1, qs2, 0 }; + + int sumi0 = 0; + int sumi1 = 0; + + // The two half-block scales (e[0]/e[1]) split at element QK_ROCMFP3/2. base + // < QK_ROCMFP3/2 is equivalent to (iqs+i) < QK_ROCMFP3/8, so for a VDR + // window that lies entirely in one half the accumulator choice is loop + // invariant and can be hoisted out of the unrolled loop. A straddling window + // (only possible for VDRs that cross the midpoint) still uses the exact + // per-element branch, so results are bit-identical either way. + const bool fp3_first_half = iqs + VDR_ROCMFP3_Q8_1_MMVQ <= QK_ROCMFP3/8; + const bool fp3_second_half = iqs >= QK_ROCMFP3/8; + +#pragma unroll + for (int i = 0; i < VDR_ROCMFP3_Q8_1_MMVQ; ++i) { + const int base = 4 * (iqs + i); + const int start_bit = 12 * (iqs + i); + const int reg_idx = start_bit >> 5; + const int reg_shift = start_bit & 31; + const uint32_t val_low = qs[reg_idx]; + const uint32_t val_high = qs[reg_idx + 1]; + const uint32_t bits12 = (reg_shift == 0) ? (val_low & 0xFFFu) : (((val_low >> reg_shift) | (val_high << (32 - reg_shift))) & 0xFFFu); + + const char4 v = make_char4( + (int8_t) rocmfpx_decode_fp3_code_vec_cuda(bits12 & 7u), + (int8_t) rocmfpx_decode_fp3_code_vec_cuda((bits12 >> 3) & 7u), + (int8_t) rocmfpx_decode_fp3_code_vec_cuda((bits12 >> 6) & 7u), + (int8_t) rocmfpx_decode_fp3_code_vec_cuda((bits12 >> 9) & 7u)); + const int val_packed = *((const int *) &v); + + const int u = get_int_b4(bq8_1->qs, iqs + i); + + if (fp3_first_half) { + sumi0 = ggml_cuda_dp4a(val_packed, u, sumi0); + } else if (fp3_second_half) { + sumi1 = ggml_cuda_dp4a(val_packed, u, sumi1); + } else if (base < QK_ROCMFP3/2) { + sumi0 = ggml_cuda_dp4a(val_packed, u, sumi0); + } else { + sumi1 = ggml_cuda_dp4a(val_packed, u, 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); +} + +static __device__ __forceinline__ float vec_dot_rocmfpx_fp6_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_rocmfp6_device * bq6 = (const block_rocmfp6_device *) vbq + kbx; + +#if !GGML_ROCMFP6_EXPANDED_DEVICE + uint32_t qs0, qs1, qs2, qs3, qs4, qs5; + memcpy(&qs0, bq6->qs + 0, 4); + memcpy(&qs1, bq6->qs + 4, 4); + memcpy(&qs2, bq6->qs + 8, 4); + memcpy(&qs3, bq6->qs + 12, 4); + memcpy(&qs4, bq6->qs + 16, 4); + memcpy(&qs5, bq6->qs + 20, 4); + + // Trailing 0 pad: the last FP6 window reads qs[reg_idx + 1] with reg_idx==5; + // its high bits are always masked out, so 0 keeps the result bit-identical + // while avoiding a stack over-read (matches the FP3 sibling below). + const uint32_t qs[7] = { qs0, qs1, qs2, qs3, qs4, qs5, 0 }; +#endif + + int sumi0 = 0; + int sumi1 = 0; + +#if GGML_ROCMFP6_MMVQ_HALF_BLOCK_SPLIT + const bool fp6_first_half = iqs + VDR_ROCMFP6_Q8_1_MMVQ <= QK_ROCMFP6/8; + const bool fp6_second_half = iqs >= QK_ROCMFP6/8; +#endif + +#pragma unroll + for (int i = 0; i < VDR_ROCMFP6_Q8_1_MMVQ; ++i) { + const int base = 4 * (iqs + i); +#if GGML_ROCMFP6_EXPANDED_DEVICE + const int val_packed = rocmfpx_pack4_fp6_device_vec_cuda(bq6, base); +#else + const int start_bit = 6 * base; + const int reg_idx = start_bit >> 5; + const int reg_shift = start_bit & 31; + const uint32_t val_low = qs[reg_idx]; + const uint32_t val_high = qs[reg_idx + 1]; + const uint32_t bits24 = (reg_shift == 0) ? (val_low & 0xFFFFFFu) : + (((val_low >> reg_shift) | (val_high << (32 - reg_shift))) & 0xFFFFFFu); + + const int val_packed = rocmfpx_pack4_fp6_bits24_vec_cuda(bits24); +#endif + const int u = get_int_b4(bq8_1->qs, iqs + i); + +#if GGML_ROCMFP6_MMVQ_HALF_BLOCK_SPLIT + if (fp6_first_half) { + sumi0 = ggml_cuda_dp4a(val_packed, u, sumi0); + } else if (fp6_second_half) { + sumi1 = ggml_cuda_dp4a(val_packed, u, sumi1); + } else +#endif + if (base < QK_ROCMFP6/2) { + sumi0 = ggml_cuda_dp4a(val_packed, u, sumi0); + } else { + sumi1 = ggml_cuda_dp4a(val_packed, u, sumi1); + } + } + + const float db = __low2float(bq8_1->ds); + return db * (rocmfpx_ue4m3_to_fp32_finite(bq6->e[0]) * sumi0 + rocmfpx_ue4m3_to_fp32_finite(bq6->e[1]) * sumi1); +} + +static __device__ __forceinline__ float vec_dot_rocmfpx_fp8_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_rocmfp8 * bq8 = (const block_rocmfp8 *) vbq + kbx; + + int v[VDR_ROCMFP8_Q8_1_MMVQ]; + int u[VDR_ROCMFP8_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_ROCMFP8_Q8_1_MMVQ; ++i) { + v[i] = get_int_b1(bq8->qs, iqs + i); + u[i] = get_int_b4(bq8_1->qs, iqs + i); + } + + return vec_dot_q8_0_q8_1_impl( + v, u, rocmfpx_ue4m3_to_fp32_finite(bq8->e), __low2half(bq8_1->ds)); +} + #define VDR_NVFP4_Q8_1_MMVQ 4 #define VDR_NVFP4_Q8_1_MMQ 8 diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index fda42ef76..591ab9bda 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -9,6 +9,8 @@ // FIXME: required here for quantization functions #include "ggml-quants.h" +#include "../rocmfp4/rocmfp4.h" +#include "../rocmfpx/rocmfpx.h" #ifdef GGML_USE_CPU_HBM #include @@ -719,6 +721,54 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { .to_float = (ggml_to_float_t) dequantize_row_q8_0, .from_float_ref = (ggml_from_float_t) quantize_row_q8_0_ref, }, + [GGML_TYPE_Q4_0_ROCMFP4] = { + .type_name = "q4_0_rocmfp4", + .blck_size = QK_ROCMFP4, + .type_size = sizeof(block_rocmfp4), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfp4_dequantize_row_q4_0, + .from_float_ref = (ggml_from_float_t) rocmfp4_quantize_row_q4_0_ref, + }, + [GGML_TYPE_Q4_0_ROCMFP4_FAST] = { + .type_name = "q4_0_rocmfp4_fast", + .blck_size = QK_ROCMFP4, + .type_size = sizeof(block_rocmfp4_fast), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfp4_dequantize_row_q4_0_fast, + .from_float_ref = (ggml_from_float_t) rocmfp4_quantize_row_q4_0_fast_ref, + }, + [GGML_TYPE_Q3_0_ROCMFPX] = { + .type_name = "q3_0_rocmfpx", + .blck_size = QK_ROCMFP3, + .type_size = sizeof(block_rocmfp3), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp3, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp3_ref, + }, + [GGML_TYPE_Q2_0_ROCMFP2] = { + .type_name = "q2_0_rocmfp2", + .blck_size = QK_ROCMFP2, + .type_size = sizeof(block_rocmfp2), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp2, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp2_ref, + }, + [GGML_TYPE_Q6_0_ROCMFPX] = { + .type_name = "q6_0_rocmfpx", + .blck_size = QK_ROCMFP6, + .type_size = sizeof(block_rocmfp6), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp6, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp6_ref, + }, + [GGML_TYPE_Q8_0_ROCMFPX] = { + .type_name = "q8_0_rocmfpx", + .blck_size = QK_ROCMFP8, + .type_size = sizeof(block_rocmfp8), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp8, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp8_ref, + }, [GGML_TYPE_Q8_1] = { .type_name = "q8_1", .blck_size = QK8_1, @@ -1074,9 +1124,13 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", "TURBO_WHT", + + "MOE_FUSED", + + "DS4_HC", }; -static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); +static_assert(GGML_OP_COUNT == 100, "GGML_OP_COUNT != 100"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1187,9 +1241,13 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", "turbo_wht(a)", + + "moe_fused(x)", + + "ds4_hc(x)", }; -static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); +static_assert(GGML_OP_COUNT == 100, "GGML_OP_COUNT != 100"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -1225,11 +1283,12 @@ static const char * GGML_GLU_OP_NAME[GGML_GLU_OP_COUNT] = { "GEGLU", "SWIGLU", "SWIGLU_OAI", + "SWIGLU_DS4", "GEGLU_ERF", "GEGLU_QUICK", }; -static_assert(GGML_GLU_OP_COUNT == 6, "GGML_GLU_OP_COUNT != 6"); +static_assert(GGML_GLU_OP_COUNT == 7, "GGML_GLU_OP_COUNT != 7"); static_assert(sizeof(struct ggml_object)%GGML_MEM_ALIGN == 0, "ggml_object size must be a multiple of GGML_MEM_ALIGN"); @@ -1428,6 +1487,20 @@ enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype) { case GGML_FTYPE_MOSTLY_IQ2_S: wtype = GGML_TYPE_IQ2_S; break; case GGML_FTYPE_UNKNOWN: wtype = GGML_TYPE_COUNT; break; case GGML_FTYPE_MOSTLY_Q4_1_SOME_F16: wtype = GGML_TYPE_COUNT; break; + // ROCmFPX presets are mixed-quant recipes (see ggml/rocmfp4/README.md); + // report the dominant transformer-tensor type. + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4: + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_LEAN: + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_COHERENT: wtype = GGML_TYPE_Q4_0_ROCMFP4; break; + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST: + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST_COHERENT: + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX: + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX_LEAN: wtype = GGML_TYPE_Q4_0_ROCMFP4_FAST; break; + case GGML_FTYPE_MOSTLY_Q6_0_ROCMFPX: wtype = GGML_TYPE_Q6_0_ROCMFPX; break; + case GGML_FTYPE_MOSTLY_Q8_0_ROCMFPX: wtype = GGML_TYPE_Q8_0_ROCMFPX; break; + case GGML_FTYPE_MOSTLY_Q3_0_ROCMFPX: wtype = GGML_TYPE_Q3_0_ROCMFPX; break; + case GGML_FTYPE_MOSTLY_Q2_0_ROCMFP2: + case GGML_FTYPE_MOSTLY_Q2_0_ROCMFP2_STRIX: wtype = GGML_TYPE_Q2_0_ROCMFP2; break; } GGML_ASSERT(wtype != GGML_TYPE_COUNT); @@ -3025,6 +3098,16 @@ struct ggml_tensor * ggml_swiglu_split( return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_SWIGLU, false); } +struct ggml_tensor * ggml_swiglu_ds4_split( + struct ggml_context * ctx, + struct ggml_tensor * gate, + struct ggml_tensor * up, + float clamp) { + struct ggml_tensor * result = ggml_glu_impl(ctx, gate, up, GGML_GLU_OP_SWIGLU_DS4, false); + ggml_set_op_params_f32(result, 2, clamp); + return result; +} + // ggml_geglu_erf struct ggml_tensor * ggml_geglu_erf( @@ -7991,3 +8074,93 @@ struct ggml_tensor * ggml_laguna_moe_combine( return result; } + +struct ggml_tensor * ggml_ds4_hc_pre( + struct ggml_context * ctx, + struct ggml_tensor * mix, + struct ggml_tensor * base, + struct ggml_tensor * hc_state, + int n_hc, + int sinkhorn_iters, + float pre_scale, + float post_scale, + float comb_scale) { + GGML_ASSERT(mix->type == GGML_TYPE_F32); + GGML_ASSERT(base->type == GGML_TYPE_F32); + GGML_ASSERT(hc_state->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(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); + result->op = GGML_OP_DS4_HC; + result->src[0] = mix; + result->src[1] = base; + result->src[2] = hc_state; + ggml_set_op_params_i32(result, 0, 0); + ggml_set_op_params_i32(result, 1, (int32_t) n_embd); + ggml_set_op_params_i32(result, 2, (int32_t) n_hc); + ggml_set_op_params_i32(result, 3, (int32_t) sinkhorn_iters); + ggml_set_op_params_f32(result, 4, pre_scale); + ggml_set_op_params_f32(result, 5, post_scale); + ggml_set_op_params_f32(result, 6, comb_scale); + return result; +} + +struct ggml_tensor * ggml_ds4_hc_post( + struct ggml_context * ctx, + struct ggml_tensor * residual_hc, + struct ggml_tensor * block_out, + struct ggml_tensor * split, + int n_hc) { + GGML_ASSERT(residual_hc->type == GGML_TYPE_F32); + GGML_ASSERT(block_out->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(block_out) == 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] = block_out; + result->src[2] = split; + ggml_set_op_params_i32(result, 0, 1); + 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, + struct ggml_tensor * base, + struct ggml_tensor * hc_state, + int n_hc, + float pre_scale) { + GGML_ASSERT(mix->type == GGML_TYPE_F32); + GGML_ASSERT(base->type == GGML_TYPE_F32); + GGML_ASSERT(hc_state->type == GGML_TYPE_F32); + 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); + result->op = GGML_OP_DS4_HC; + result->src[0] = mix; + result->src[1] = base; + result->src[2] = hc_state; + ggml_set_op_params_i32(result, 0, 2); + ggml_set_op_params_i32(result, 1, (int32_t) n_embd); + ggml_set_op_params_i32(result, 2, (int32_t) n_hc); + ggml_set_op_params_f32(result, 4, pre_scale); + return result; +} From a5f1146f49d3a1b218db0b6a06a79631a099798b Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 8 Jul 2026 16:50:22 +0200 Subject: [PATCH 02/10] fix(ggml): make ROCmFP4 codebook self-contained on CUDA The non-HIP fallback of rocmfp4_get_int_from_codebook_16 / rocmfp4_get_low_int_from_codebook_16 called get_int_from_table_16, which is defined in vecdotq.cuh. TUs that pull in this header without vecdotq.cuh (fattn-chunked.cu reaches it via the fattn dequant chain) failed to compile under nvcc: rocmfp4_hip_codebook.cuh: error: identifier "get_int_from_table_16" is undefined The HIP path never hit this (it uses __builtin_amdgcn_perm), so the ROCm CI and the Strix build stayed green while the sm_86 CUDA build broke. Fix: inline the generic table expander (the generic branch of get_int_from_table_16, verbatim) as a static helper in this header, so the fallback no longer depends on include order. Bit-identical; the HIP hot path is unchanged. --- .../ggml/rocmfp4/rocmfp4_hip_codebook.cuh | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_codebook.cuh b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_codebook.cuh index 1d9106112..b5dea2e11 100644 --- a/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_codebook.cuh +++ b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_codebook.cuh @@ -9,6 +9,28 @@ #define GGML_ROCMFP4_UNALIGNED_QS_DWORD_LOAD 1 #endif +#if !defined(GGML_USE_HIP) +// Self-contained 16-entry table expander for the non-HIP fallback. This header +// is pulled into translation units that do not include vecdotq.cuh (where the +// generic get_int_from_table_16 lives) - e.g. fattn-chunked.cu - so relying on +// that symbol breaks the CUDA build on include order. This is the generic +// branch of get_int_from_table_16 verbatim, so results are bit-identical; only +// the HIP path (Strix) uses __builtin_amdgcn_perm and never reaches here. +static __device__ __forceinline__ int2 rocmfp4_table16_fallback(const int & q4, const int8_t * table) { + const int q0_32 = (q4 >> 0) & 0x0F0F0F0F; + const int8_t * q0_8 = (const int8_t *) &q0_32; + const char4 val0_8 = make_char4( + table[q0_8[0]], table[q0_8[1]], table[q0_8[2]], table[q0_8[3]]); + + const int q1_32 = (q4 >> 4) & 0x0F0F0F0F; + const int8_t * q1_8 = (const int8_t *) &q1_32; + const char4 val1_8 = make_char4( + table[q1_8[0]], table[q1_8[1]], table[q1_8[2]], table[q1_8[3]]); + + return make_int2(*((const int *) &val0_8), *((const int *) &val1_8)); +} +#endif + static __device__ __forceinline__ int rocmfp4_get_qs_i32(const void * x, const int & i32) { #if defined(GGML_USE_HIP) && GGML_ROCMFP4_UNALIGNED_QS_DWORD_LOAD return *((const int *) ((const uint8_t *) x + 4*i32)); @@ -52,7 +74,7 @@ static __device__ __forceinline__ int2 rocmfp4_get_int_from_codebook_16(const in __builtin_amdgcn_perm(v_even_high, v_even_low, mask_even), __builtin_amdgcn_perm(v_odd_high, v_odd_low, mask_odd)); #else - return get_int_from_table_16(q4, fallback_table); + return rocmfp4_table16_fallback(q4, fallback_table); #endif } @@ -74,6 +96,6 @@ static __device__ __forceinline__ int rocmfp4_get_low_int_from_codebook_16(const return __builtin_amdgcn_perm(v_high, v_low, mask); #else - return get_int_from_table_16(q4, fallback_table).x; + return rocmfp4_table16_fallback(q4, fallback_table).x; #endif } From 42f1e10579e5536ce253fc48e754fe50b3149797 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 8 Jul 2026 15:46:08 +0200 Subject: [PATCH 03/10] fix(deepseek4): ROCmFPX byte-safe decode + fused-decode path Server-side DeepSeek V4 Flash changes from #494, rebased onto the vendored ROCmFPX ggml tree (no submodule pointer). Byte-identical to the validated #494 server files. Correctness (default ON): - Token-by-token prefill by default; chunked prefill only fits inside the raw SWA ring, so long prompts / multi-turn degraded. Chunked stays available via DFLASH_DS4_CHUNKED_PREFILL for short-prompt benchmarking. - Clear the cache buffer at new-sequence prefill (kv_offset==0) so requests 2..N are byte-stable instead of pooling over leftover compressor state. - Route non-hybrid (all-hot) placement through the HC-complete layer-range path; deepseek4_step's non-hybrid branch is HC-less and generates garbage. - Key the cached decode-attn graph on the flush pattern; the old key collided once the compressed-KV ring filled, reusing a stale-topology graph. - Compressor decode graphs read state/comp-cache through the ggml_set_rows result tensors so the current-step writes are explicit graph dependencies. - Default DeepSeek4 chat prefix when the request has no system message, so the ROCmFPX "Src" GGUF stops behaving like a base model under bare prompts (explicit caller system prompts are preserved). Perf, byte-identical default: - Persistent HC matvec pool (row-split preserves per-row accumulation order) and f16c dot kernels with scalar-order adds. - Per-step decode scalar inputs uploaded in 2 tensor_sets instead of ~430. Opt-in, default OFF (documented as not bit-identical): - DFLASH_DS4_FUSED_DECODE single-graph decode with GGML_OP_DS4_HC. - DFLASH_DS4_FFN_RAW_MMID / DFLASH_DS4_FFN_FUSED_COMBINE accumulation reorder. - DFLASH_DS4_ROCMFPX_HC_GPU GPU HC pre-mix. --- server/src/deepseek4/deepseek4_backend.cpp | 67 +- server/src/deepseek4/deepseek4_graph.cpp | 1745 ++++++++++++++++++-- server/src/deepseek4/deepseek4_hc_cuda.cu | 46 +- server/src/deepseek4/deepseek4_hc_cuda.h | 7 + server/src/deepseek4/deepseek4_internal.h | 8 +- server/src/server/chat_template.cpp | 5 + 6 files changed, 1713 insertions(+), 165 deletions(-) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 9b4616da1..5a3edbf9b 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -535,9 +535,26 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, int kv_offset) { // Hybrid currently implements HC for single-token steps only; keep prefill // token-by-token so the first sampled token is seeded from the correct HC state. - const int chunk = moe_hybrid_ ? 1 : (cfg_.chunk > 0 ? cfg_.chunk : 512); + // + // Chunked prefill (n_tokens > 1) is only exact while the whole prompt fits + // inside the raw SWA ring: the chunk graph writes ring slots pos % n_swa + // while attention reads the ring in the same graph, and the learned + // compressor only runs for the chunk's last token. Prompts beyond n_swa + // degrade into incoherence. Token-by-token prefill matches the reference + // semantics at any length; chunked stays available via + // DFLASH_DS4_CHUNKED_PREFILL=1 for short-prompt benchmarking only. + const bool unsafe_chunked = env_flag_enabled("DFLASH_DS4_CHUNKED_PREFILL"); + const int chunk = (moe_hybrid_ || !unsafe_chunked) ? 1 : (cfg_.chunk > 0 ? cfg_.chunk : 512); const int n_total = (int)tokens.size(); 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. + // Without this, the first flush windows of a request pool over the + // previous request's leftover state rows and outputs from the 2nd/3rd + // request on can drift by a token or two. + if (kv_offset == 0 && cache_.buf) { + ggml_backend_buffer_clear(cache_.buf, 0); + } last_logits_.clear(); const bool timing = env_flag_enabled("DFLASH_DS4_TIMING"); const auto phase_t0 = Clock::now(); @@ -556,13 +573,26 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, DeepSeek4StepTelemetry step_tel; if (timing) step_tel.embed_us = elapsed_us(embed_t0, Clock::now()); - // Run forward pass + // Run forward pass. The non-hybrid (all-hot) placement must use the + // HC-complete layer-range path; deepseek4_step's non-hybrid branch is + // an HC-less stub and produces garbage on this box. std::vector logits; - if (!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())) { + bool ok = false; + 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()); + } else { + std::vector hc_state; + 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); + } + if (!ok) { std::fprintf(stderr, "[deepseek4] prefill step failed at pos=%d\n", pos); return -1; } @@ -619,12 +649,23 @@ 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_) { + 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()); + } 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; } diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 06fbe8a98..280a95946 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -29,8 +29,14 @@ #include #include #include +#include +#include #include +#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) +#include +#endif + namespace dflash::common { namespace { @@ -46,6 +52,106 @@ static bool ds4_env_flag(const char * name) { return value && value[0] && std::strcmp(value, "0") != 0; } +// Opt-in: reorders expert-FFN float accumulation, so output is not +// bit-identical to the reference path. Default OFF. +static bool ds4_ffn_raw_mmid_enabled() { + static int enabled = -1; + if (enabled < 0) { + enabled = ds4_env_flag("DFLASH_DS4_FFN_RAW_MMID") ? 1 : 0; + } + return enabled == 1; +} + +// Opt-in: reorders expert-output combination, so output is not +// bit-identical to the reference path. Default OFF. +static bool ds4_ffn_fused_combine_enabled() { + static int enabled = -1; + if (enabled < 0) { + enabled = ds4_env_flag("DFLASH_DS4_FFN_FUSED_COMBINE") ? 1 : 0; + } + return enabled == 1; +} + +static bool ds4_rocmfpx_hc_gpu_enabled() { + static int enabled = -1; + if (enabled < 0) { + enabled = ds4_env_flag("DFLASH_DS4_ROCMFPX_HC_GPU") ? 1 : 0; + } + return enabled == 1; +} + +static size_t ds4_full_step_graph_size(int n_tokens) { + if (n_tokens <= 1) { + return 16384; + } + if (n_tokens <= 128) { + return 65536; + } + if (n_tokens <= 512) { + return 131072; + } + if (n_tokens <= 1024) { + return 262144; + } + if (n_tokens <= 2048) { + return 524288; + } + if (n_tokens <= 4096) { + return 1048576; + } + return 1572864; +} + +static size_t ds4_full_step_meta_size(int n_tokens) { + const size_t graph_size = ds4_full_step_graph_size(n_tokens); + size_t arena_size = ggml_tensor_overhead() * 65536 + + ggml_graph_overhead_custom(graph_size, false) + + 16 * 1024 * 1024 + + 1024 * 1024 + + 64 * 1024; + if (n_tokens >= 512) { + arena_size += (size_t)n_tokens * 32 * 1024; + } + if (n_tokens >= 1024) { + arena_size += 4 * 1024 * 1024; + } + if (n_tokens > 1024) { + // 2K-token full steps need a materially larger meta arena than the + // 1K-token path once the full graph and its late scratch tensors are + // fully materialized, so keep a coarse but stable reserve here. The + // exact scaling can be tightened in the follow-up chunk-sizing pass. + arena_size += 256 * 1024 * 1024; + } + return arena_size; +} + +static size_t ds4_attn_step_meta_size(int n_tokens) { + size_t arena_size = 48 * 1024 * 1024; + if (n_tokens >= 512) { + arena_size += (size_t)n_tokens * 32 * 1024; + } + return arena_size; +} + +static size_t ds4_attn_step_graph_size(int n_tokens) { + if (n_tokens <= 1) { + return 2048; + } + if (n_tokens <= 512) { + return 32768; + } + if (n_tokens <= 1024) { + return 131072; + } + if (n_tokens <= 2048) { + return 262144; + } + if (n_tokens <= 4096) { + return 524288; + } + return 1048576; +} + template static void ds4_parallel_for_tokens(int n_tokens, int min_parallel_tokens, Fn && fn) { if (n_tokens <= min_parallel_tokens) { @@ -176,6 +282,21 @@ struct DeepSeek4CachedDecodeOutputGraph { } }; +struct DeepSeek4LegacyFullStepCache { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + StepGraph sg; + std::vector meta_arena; + + void free() { + step_graph_destroy(sg); + meta_arena.clear(); + meta_arena.shrink_to_fit(); + owner_ctx = nullptr; + backend = nullptr; + } +}; + struct DeepSeek4AttentionGraphInputs { ggml_tensor * rope_pos = nullptr; ggml_tensor * neg_pos = nullptr; @@ -188,6 +309,15 @@ struct DeepSeek4AttentionGraphInputs { ggml_tensor * index_state_rows = nullptr; ggml_tensor * index_comp_rows = nullptr; ggml_tensor * index_comp_pos = nullptr; + // Fused-decode stable-KV path only: additive score mask over + // [n_swa raw rows ++ padded comp rows]; 0 for valid, -1e30 for padding. + ggml_tensor * attn_row_mask = nullptr; + int padded_comp = 0; // padded compressed-row count (>= n_comp) + // Fused-decode stable-topology compressor: i64[4] target rows for the + // ratio-4 state double-buffer flush copy. [0..3] on flush steps (cur -> + // prev), [4..7] otherwise (self-write no-op). Null = legacy build-time + // flush branch. + ggml_tensor * flush_rows = nullptr; }; struct DeepSeek4CachedDecodeAttnGraph { @@ -198,8 +328,11 @@ struct DeepSeek4CachedDecodeAttnGraph { int n_raw = 0; int n_comp_attn = 0; int n_index_comp = 0; + bool attn_flush = false; + bool index_flush = false; bool compressed = false; bool indexed = false; + bool uses_shared_inputs = false; StepGraph sg; DeepSeek4AttentionGraphInputs inputs; @@ -224,8 +357,11 @@ struct DeepSeek4CachedDecodeAttnGraph { n_raw = 0; n_comp_attn = 0; n_index_comp = 0; + attn_flush = false; + index_flush = false; compressed = false; indexed = false; + uses_shared_inputs = false; } }; @@ -324,14 +460,19 @@ static bool build_cached_decode_ffn_graph( const int n_used = w.n_expert_used; const int n_ff_exp = w.n_ff_exp; + const bool raw_mmid = ds4_ffn_raw_mmid_enabled(); ggml_tensor * cur_3d = ggml_reshape_3d(out.sg.ctx, ffn_normed, w.n_embd, 1, n_tokens); ggml_tensor * gate_e = ggml_mul_mat_id(out.sg.ctx, L.ffn_gate_exps, cur_3d, out.hash_ids); ggml_tensor * up_e = ggml_mul_mat_id(out.sg.ctx, L.ffn_up_exps, cur_3d, out.hash_ids); - gate_e = ggml_reshape_3d(out.sg.ctx, gate_e, n_ff_exp, n_used, n_tokens); - up_e = ggml_reshape_3d(out.sg.ctx, up_e, n_ff_exp, n_used, n_tokens); + if (!raw_mmid) { + gate_e = ggml_reshape_3d(out.sg.ctx, gate_e, n_ff_exp, n_used, n_tokens); + up_e = ggml_reshape_3d(out.sg.ctx, up_e, n_ff_exp, n_used, n_tokens); + } ggml_tensor * mid_e = build_clamped_swiglu(out.sg.ctx, gate_e, up_e, w.swiglu_clamp_exp); ggml_tensor * down_e = ggml_mul_mat_id(out.sg.ctx, L.ffn_down_exps, mid_e, out.hash_ids); - down_e = ggml_reshape_3d(out.sg.ctx, down_e, w.n_embd, n_used, n_tokens); + if (!raw_mmid) { + down_e = ggml_reshape_3d(out.sg.ctx, down_e, w.n_embd, n_used, n_tokens); + } ggml_tensor * probs_3d = ggml_reshape_3d(out.sg.ctx, probs, 1, w.n_expert, n_tokens); ggml_tensor * weights = ggml_get_rows(out.sg.ctx, probs_3d, out.hash_ids); @@ -343,11 +484,16 @@ static bool build_cached_decode_ffn_graph( weights = ggml_scale(out.sg.ctx, weights, w.expert_weight_scale); } - ggml_tensor * weights_3d = ggml_reshape_3d(out.sg.ctx, weights, 1, n_used, n_tokens); - ggml_tensor * routed_out = ggml_mul(out.sg.ctx, down_e, weights_3d); - routed_out = ggml_cont(out.sg.ctx, ggml_permute(out.sg.ctx, routed_out, 1, 0, 2, 3)); - routed_out = ggml_sum_rows(out.sg.ctx, routed_out); - routed_out = ggml_reshape_2d(out.sg.ctx, routed_out, w.n_embd, n_tokens); + ggml_tensor * routed_out = nullptr; + if (ds4_ffn_fused_combine_enabled()) { + routed_out = ggml_laguna_moe_combine(out.sg.ctx, down_e, weights); + } else { + ggml_tensor * weights_3d = ggml_reshape_3d(out.sg.ctx, weights, 1, n_used, n_tokens); + routed_out = ggml_mul(out.sg.ctx, down_e, weights_3d); + ggml_tensor * sum_shape = ggml_new_tensor_3d(out.sg.ctx, GGML_TYPE_F32, w.n_embd, 1, n_tokens); + routed_out = ggml_repeat_back(out.sg.ctx, routed_out, sum_shape); + routed_out = ggml_reshape_2d(out.sg.ctx, routed_out, w.n_embd, n_tokens); + } ffn_out = ggml_add(out.sg.ctx, shared_out, routed_out); } else { @@ -428,12 +574,14 @@ static ggml_tensor * build_clamped_swiglu(ggml_context * ctx, ggml_tensor * gate, ggml_tensor * up, float clamp) { - // DS4 clamps only the upper side of gate, but both sides of up. - gate = ggml_clamp(ctx, gate, -INFINITY, clamp); - up = ggml_clamp(ctx, up, -clamp, clamp); - // silu(gate) * up - gate = ggml_silu(ctx, gate); - return ggml_mul(ctx, gate, up); + return ggml_swiglu_ds4_split(ctx, gate, up, clamp); +} + +static ggml_tensor * ds4_cast_if_needed( + ggml_context * ctx, + ggml_tensor * x, + ggml_type type) { + return x->type == type ? x : ggml_cast(ctx, x, type); } // ─── Helper: Partial RoPE (tail rotation) ─────────────────────────────── @@ -526,7 +674,9 @@ static void build_compressor_step( ggml_tensor * comp_rows_inp, ggml_tensor * comp_pos_inp, std::vector & i64_array_inputs, - std::vector & i32_array_inputs) { + std::vector & i32_array_inputs, + ggml_tensor ** comp_cache_source_out = nullptr, + ggml_tensor * flush_rows_inp = nullptr) { if (!gf || !cur_last || !ape || !kv_proj || !gate_proj || !norm_weight || !state.state_kv || !state.state_score || !comp_cache || ratio <= 0) { return; @@ -541,6 +691,9 @@ static void build_compressor_step( ggml_tensor * kv_cur = ggml_mul_mat(ctx, kv_proj, cur_last); ggml_tensor * sc_cur = ggml_mul_mat(ctx, gate_proj, cur_last); + ggml_tensor * state_kv_source = state.state_kv; + ggml_tensor * state_score_source = state.state_score; + ggml_tensor * comp_cache_source = comp_cache; ggml_tensor * ape_col = nullptr; if (ape_row_inp) { @@ -554,8 +707,10 @@ static void build_compressor_step( sc_cur = ggml_add(ctx, sc_cur, ape_col); if (state_rows_inp) { - ggml_build_forward_expand(gf, ggml_set_rows(ctx, state.state_kv, kv_cur, state_rows_inp)); - ggml_build_forward_expand(gf, ggml_set_rows(ctx, state.state_score, sc_cur, state_rows_inp)); + state_kv_source = ggml_set_rows(ctx, state.state_kv, kv_cur, state_rows_inp); + state_score_source = ggml_set_rows(ctx, state.state_score, sc_cur, state_rows_inp); + ggml_build_forward_expand(gf, state_kv_source); + ggml_build_forward_expand(gf, state_score_source); } else { ggml_tensor * kv_slot = ggml_view_2d( ctx, state.state_kv, comp_width, 1, state.state_kv->nb[1], @@ -567,7 +722,10 @@ static void build_compressor_step( ggml_build_forward_expand(gf, ggml_cpy(ctx, sc_cur, sc_slot)); } - if (((token_pos + 1) % ratio) != 0) { + if (!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. return; } @@ -586,26 +744,26 @@ static void build_compressor_step( ggml_tensor * sv_sc = nullptr; int n_state_rows = ratio; if (ratio == 4) { - const size_t hi_off_kv = (size_t)ratio * state.state_kv->nb[1] + - (size_t)head_dim * state.state_kv->nb[0]; - const size_t hi_off_sc = (size_t)ratio * state.state_score->nb[1] + - (size_t)head_dim * state.state_score->nb[0]; - ggml_tensor * prev_kv = ggml_view_2d(ctx, state.state_kv, head_dim, ratio, - state.state_kv->nb[1], 0); - ggml_tensor * cur_kv_hi = ggml_view_2d(ctx, state.state_kv, head_dim, ratio, - state.state_kv->nb[1], hi_off_kv); - ggml_tensor * prev_sc = ggml_view_2d(ctx, state.state_score, head_dim, ratio, - state.state_score->nb[1], 0); - ggml_tensor * cur_sc_hi = ggml_view_2d(ctx, state.state_score, head_dim, ratio, - state.state_score->nb[1], hi_off_sc); + const size_t hi_off_kv = (size_t)ratio * state_kv_source->nb[1] + + (size_t)head_dim * state_kv_source->nb[0]; + const size_t hi_off_sc = (size_t)ratio * state_score_source->nb[1] + + (size_t)head_dim * state_score_source->nb[0]; + ggml_tensor * prev_kv = ggml_view_2d(ctx, state_kv_source, head_dim, ratio, + state_kv_source->nb[1], 0); + ggml_tensor * cur_kv_hi = ggml_view_2d(ctx, state_kv_source, head_dim, ratio, + state_kv_source->nb[1], hi_off_kv); + ggml_tensor * prev_sc = ggml_view_2d(ctx, state_score_source, head_dim, ratio, + state_score_source->nb[1], 0); + ggml_tensor * cur_sc_hi = ggml_view_2d(ctx, state_score_source, head_dim, ratio, + state_score_source->nb[1], hi_off_sc); sv_kv = ggml_concat(ctx, prev_kv, cur_kv_hi, 1); sv_sc = ggml_concat(ctx, prev_sc, cur_sc_hi, 1); n_state_rows = 2 * ratio; } else { - sv_kv = ggml_view_2d(ctx, state.state_kv, comp_width, n_state_rows, - state.state_kv->nb[1], 0); - sv_sc = ggml_view_2d(ctx, state.state_score, comp_width, n_state_rows, - state.state_score->nb[1], 0); + sv_kv = ggml_view_2d(ctx, state_kv_source, comp_width, n_state_rows, + state_kv_source->nb[1], 0); + sv_sc = ggml_view_2d(ctx, state_score_source, comp_width, n_state_rows, + state_score_source->nb[1], 0); } // Transpose to [n_state_rows, comp_width] so softmax operates per-dimension ggml_tensor * sc_T = ggml_cont(ctx, ggml_transpose(ctx, sv_sc)); @@ -643,7 +801,8 @@ static void build_compressor_step( } if (comp_rows_inp) { - ggml_build_forward_expand(gf, ggml_set_rows(ctx, comp_cache, pooled, comp_rows_inp)); + comp_cache_source = ggml_set_rows(ctx, comp_cache, pooled, comp_rows_inp); + ggml_build_forward_expand(gf, comp_cache_source); } else { ggml_tensor * comp_slot = ggml_view_2d( ctx, comp_cache, head_dim, 1, comp_cache->nb[1], @@ -651,7 +810,24 @@ static void build_compressor_step( ggml_build_forward_expand(gf, ggml_cpy(ctx, pooled_f16, comp_slot)); } - if (ratio == 4) { + if (comp_cache_source_out) { + *comp_cache_source_out = comp_cache_source; + } + + if (ratio == 4 && flush_rows_inp) { + // Stable-topology flush: copy the cur half onto rows given by the + // input (prev half [0..3] at flush, cur half itself [4..7] = no-op + // otherwise). Values are read through the set_rows source so this + // step's state write is ordered first. + ggml_tensor * cur_kv_vals = ggml_cont(ctx, ggml_view_2d( + ctx, state_kv_source, comp_width, ratio, state_kv_source->nb[1], + (size_t) ratio * state_kv_source->nb[1])); + ggml_tensor * cur_sc_vals = ggml_cont(ctx, ggml_view_2d( + ctx, state_score_source, comp_width, ratio, state_score_source->nb[1], + (size_t) ratio * state_score_source->nb[1])); + ggml_build_forward_expand(gf, ggml_set_rows(ctx, state.state_kv, cur_kv_vals, flush_rows_inp)); + ggml_build_forward_expand(gf, ggml_set_rows(ctx, state.state_score, cur_sc_vals, flush_rows_inp)); + } else if (ratio == 4) { for (int r = 0; r < ratio; ++r) { ggml_tensor * src_kv = ggml_view_2d(ctx, state.state_kv, comp_width, 1, state.state_kv->nb[1], @@ -692,7 +868,8 @@ static void build_indexer_compressor_step( ggml_tensor * comp_rows_inp, ggml_tensor * comp_pos_inp, std::vector & i64_array_inputs, - std::vector & i32_array_inputs) { + std::vector & i32_array_inputs, + ggml_tensor * flush_rows_inp = nullptr) { build_compressor_step(ctx, gf, cur_last, L.indexer_compressor_ape, L.indexer_compressor_kv, @@ -715,7 +892,9 @@ static void build_indexer_compressor_step( comp_rows_inp, comp_pos_inp, i64_array_inputs, - i32_array_inputs); + i32_array_inputs, + nullptr, + flush_rows_inp); } static int ds4_comp_rows_used(const ggml_tensor * comp_cache, int n_cached, int ratio, int token_pos) { @@ -726,6 +905,18 @@ static int ds4_comp_rows_used(const ggml_tensor * comp_cache, int n_cached, int return std::min(n_cached + grew_this_step, (int) comp_cache->ne[1]); } +// Round the live compressed-row count up to a fixed stride so the fused decode +// graph topology repeats across steps (enabling CUDA/HIP graph replay). The +// rows in [n_comp, padded) are masked to -1e30 in the score matrix, which +// underflows to exactly 0 in softmax, so a padded read is bit-identical to an +// unpadded read of the first n_comp rows. +static constexpr int DS4_COMP_PAD_STRIDE = 16; +static int ds4_padded_comp_rows(int n_comp, int cap) { + if (n_comp <= 0) return 0; + const int padded = ((n_comp + DS4_COMP_PAD_STRIDE - 1) / DS4_COMP_PAD_STRIDE) * DS4_COMP_PAD_STRIDE; + return padded < cap ? padded : cap; +} + static ggml_tensor * build_indexer_score( ggml_context * ctx, ggml_tensor * qr_norm_last, // [n_lora_q, 1] @@ -887,9 +1078,11 @@ static ggml_tensor * build_mla_attention( // ── Store ALL KV rows in the raw SWA ring ───────────────────── // For decode (n_tokens=1): write single row. For prefill: write all rows. + ggml_tensor * raw_kv_source = lc.raw_kv; if (cached_inputs && cached_inputs->raw_kv_rows) { ggml_tensor * kv_f32 = ggml_is_contiguous(kv) ? kv : ggml_cont(ctx, kv); - ggml_build_forward_expand(gf, 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, cached_inputs->raw_kv_rows); + ggml_build_forward_expand(gf, raw_kv_source); } else { for (int ti = 0; ti < n_tokens; ti++) { const int pos_ti = kv_start + ti; @@ -908,6 +1101,7 @@ static ggml_tensor * build_mla_attention( 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, L.attn_compressor_ape, @@ -931,7 +1125,9 @@ static ggml_tensor * build_mla_attention( cached_inputs ? cached_inputs->attn_comp_rows : nullptr, cached_inputs ? cached_inputs->attn_comp_pos : nullptr, i64_array_inputs, - i32_array_inputs); + i32_array_inputs, + &comp_kv_source, + cached_inputs ? cached_inputs->flush_rows : nullptr); } if (ratio == 4 && L.indexer_compressor_kv) { @@ -941,7 +1137,8 @@ static ggml_tensor * build_mla_attention( cached_inputs ? cached_inputs->index_comp_rows : nullptr, cached_inputs ? cached_inputs->index_comp_pos : nullptr, i64_array_inputs, - i32_array_inputs); + i32_array_inputs, + cached_inputs ? cached_inputs->flush_rows : nullptr); (void)build_indexer_score(ctx, qr_last, cur_last, w, L, lc, token_pos, i32_inputs); } @@ -950,45 +1147,80 @@ static ggml_tensor * build_mla_attention( // raw_kv: [head_dim, n_swa] F16 persistent ring buffer (single KV head, shared) // comp_kv: [head_dim, comp_cap] F16 compressed rows. // n_raw = min(kv_start + n_tokens, n_swa) - const int n_raw = std::min(kv_start + n_tokens, 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 bool masked_kv = (n_tokens == 1) && 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; + // 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_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; const float kq_scale = 1.0f / sqrtf((float)head_dim); - // Get valid KV rows. For single-token decode/prefill, include the current - // in-graph KV row directly; otherwise attention can race the side-effecting - // cache write and see the previous contents of the raw KV slot. - ggml_tensor * kv_f32 = nullptr; - if (n_tokens == 1) { - ggml_tensor * cur_kv_f32 = ggml_cast(ctx, kv, GGML_TYPE_F32); + // Get valid KV rows. For single-token decode, include the current in-graph + // KV row directly; otherwise attention can race the side-effecting cache + // write and see the previous contents of the raw KV slot. + auto raw_kv_view = [&](int row, int count) -> ggml_tensor * { + ggml_tensor * view = ggml_view_2d( + ctx, lc.raw_kv, head_dim, count, lc.raw_kv->nb[1], + (size_t)row * lc.raw_kv->nb[1]); + return ds4_cast_if_needed(ctx, view, GGML_TYPE_F32); + }; + + ggml_tensor * kv_attn = nullptr; + if (masked_kv) { + // Fused stable-KV path: read the full physical ring; rows not yet + // written are masked to -1e30 in the score matrix (exact 0 after + // softmax). Only the fused decode graph sets attn_row_mask. Read + // through the set_rows result so the current row's in-graph write is + // ordered before this read. + 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 (n_tokens == 1) { + ggml_tensor * cur_kv = ds4_cast_if_needed(ctx, kv, GGML_TYPE_F32); if (n_raw > 1) { - ggml_tensor * prev_f32 = ggml_cast(ctx, - ggml_view_2d(ctx, lc.raw_kv, head_dim, n_raw - 1, - lc.raw_kv->nb[1], 0), - GGML_TYPE_F32); - kv_f32 = ggml_concat(ctx, prev_f32, cur_kv_f32, 1); + ggml_tensor * prev = nullptr; + if (kv_start < w.n_swa) { + prev = raw_kv_view(0, n_raw - 1); + } else { + const int raw_row = kv_start % w.n_swa; + const int tail_count = w.n_swa - raw_row - 1; + const int head_count = raw_row; + if (tail_count > 0) { + prev = raw_kv_view(raw_row + 1, tail_count); + } + if (head_count > 0) { + ggml_tensor * head = raw_kv_view(0, head_count); + prev = prev ? ggml_concat(ctx, prev, head, 1) : head; + } + } + kv_attn = prev ? ggml_concat(ctx, prev, cur_kv, 1) : cur_kv; } else { - kv_f32 = cur_kv_f32; + kv_attn = cur_kv; } } else { - kv_f32 = ggml_cast(ctx, ggml_view_2d(ctx, lc.raw_kv, head_dim, n_raw, - lc.raw_kv->nb[1], 0), GGML_TYPE_F32); + kv_attn = raw_kv_view(0, n_raw); } - if (n_comp_attn > 0 && lc.comp_kv) { - ggml_tensor * comp_f32 = ggml_cast(ctx, - ggml_view_2d(ctx, lc.comp_kv, head_dim, n_comp_attn, lc.comp_kv->nb[1], 0), - GGML_TYPE_F32); - kv_f32 = ggml_concat(ctx, kv_f32, comp_f32, 1); + if (n_comp_attn > 0 && comp_kv_source) { + ggml_tensor * comp = ggml_view_2d(ctx, comp_kv_source, head_dim, n_comp_attn, comp_kv_source->nb[1], 0); + comp = ds4_cast_if_needed(ctx, comp, GGML_TYPE_F32); + kv_attn = ggml_concat(ctx, kv_attn, comp, 1); } - // kv_f32: [head_dim, n_attn] + // 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_f32, q_flat) = kv_f32^T[n_attn, head_dim] @ q_flat[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_f32, q_flat); + ggml_tensor * scores = ggml_mul_mat(ctx, kv_attn, q_flat); scores = ggml_scale(ctx, scores, kq_scale); + 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); + } + (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. @@ -1009,8 +1241,8 @@ static ggml_tensor * build_mla_attention( // 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_f32)) = [n_raw, head_dim] - ggml_tensor * kv_T = ggml_cont(ctx, ggml_transpose(ctx, kv_f32)); + // 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] @@ -1109,6 +1341,113 @@ struct DeepSeek4CachedDecodeHcPostGraph { } }; +// 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, +// ape_row, comp_pos, index_ape_row, index_comp_pos}; i64 layout: {raw_kv_row, +// state_row, comp_row, index_state_row, index_comp_row}. +struct Ds4DecodeSharedInputs { + static constexpr int MAX_RATIOS = 4; + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + ggml_context * ctx = nullptr; + ggml_backend_buffer_t buf = nullptr; + int ratios[MAX_RATIOS] = {0}; + int n_ratios = 0; + ggml_tensor * i32_bundle = nullptr; // [6 * n_ratios] + ggml_tensor * i64_bundle = nullptr; // [5 * n_ratios] + // Per-slot views handed to the graph builders. + ggml_tensor * v_rope_pos[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_neg_pos[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_ape_row[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_comp_pos[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_index_ape[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_index_cpos[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_raw_row[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_state_row[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_comp_row[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_index_state[MAX_RATIOS] = {nullptr}; + ggml_tensor * v_index_comp[MAX_RATIOS] = {nullptr}; + + void free() { + if (buf) { ggml_backend_buffer_free(buf); buf = nullptr; } + if (ctx) { ggml_free(ctx); ctx = nullptr; } + owner_ctx = nullptr; backend = nullptr; n_ratios = 0; + i32_bundle = nullptr; i64_bundle = nullptr; + } + + int slot(int ratio) const { + for (int i = 0; i < n_ratios; ++i) if (ratios[i] == ratio) return i; + return -1; + } + + bool ensure(const DeepSeek4Weights & w, ggml_backend_t bk) { + if (ctx && owner_ctx == w.ctx && backend == bk) return true; + free(); + n_ratios = 0; + for (int il = 0; il < w.n_layer; ++il) { + const int r = (int) w.compress_ratios[il]; + if (slot(r) >= 0) continue; + if (n_ratios >= MAX_RATIOS) return false; + ratios[n_ratios++] = r; + } + ggml_init_params p{}; + p.mem_size = ggml_tensor_overhead() * (size_t) (2 + 11 * MAX_RATIOS) + 4096; + p.no_alloc = true; + ctx = ggml_init(p); + if (!ctx) return false; + i32_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 6 * (int64_t) n_ratios); + i64_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 5 * (int64_t) n_ratios); + for (int s = 0; s < n_ratios; ++s) { + v_rope_pos[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t) s * 6 + 0) * sizeof(int32_t)); + v_neg_pos[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t) s * 6 + 1) * sizeof(int32_t)); + v_ape_row[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t) s * 6 + 2) * sizeof(int32_t)); + v_comp_pos[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t) s * 6 + 3) * sizeof(int32_t)); + v_index_ape[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t) s * 6 + 4) * sizeof(int32_t)); + v_index_cpos[s] = ggml_view_1d(ctx, i32_bundle, 1, ((size_t) s * 6 + 5) * sizeof(int32_t)); + v_raw_row[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t) s * 5 + 0) * sizeof(int64_t)); + v_state_row[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t) s * 5 + 1) * sizeof(int64_t)); + v_comp_row[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t) s * 5 + 2) * sizeof(int64_t)); + v_index_state[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t) s * 5 + 3) * sizeof(int64_t)); + v_index_comp[s] = ggml_view_2d(ctx, i64_bundle, 1, 1, sizeof(int64_t), ((size_t) s * 5 + 4) * sizeof(int64_t)); + } + buf = ggml_backend_alloc_ctx_tensors(ctx, bk); + if (!buf) { free(); return false; } + owner_ctx = w.ctx; + backend = bk; + return true; + } + + // Upload all per-step values in two writes. + void set_step(const DeepSeek4Weights & w, int kv_start) { + const int token_pos = kv_start; + int32_t i32v[6 * MAX_RATIOS] = {0}; + int64_t i64v[5 * MAX_RATIOS] = {0}; + for (int s = 0; s < n_ratios; ++s) { + const int ratio = ratios[s]; + i32v[s * 6 + 0] = kv_start; + i32v[s * 6 + 1] = -kv_start; + i64v[s * 5 + 0] = kv_start % w.n_swa; + if (ratio > 0) { + const int pos_mod = token_pos % ratio; + i32v[s * 6 + 2] = pos_mod; + i32v[s * 6 + 3] = token_pos + 1 - ratio; + i64v[s * 5 + 1] = (ratio == 4) ? (int64_t) (ratio + pos_mod) : (int64_t) pos_mod; + i64v[s * 5 + 2] = token_pos / ratio; + } + if (ratio == 4) { + const int pos_mod = token_pos % ratio; + i32v[s * 6 + 4] = pos_mod; + i32v[s * 6 + 5] = token_pos + 1 - ratio; + i64v[s * 5 + 3] = ratio + pos_mod; + i64v[s * 5 + 4] = token_pos / ratio; + } + } + ggml_backend_tensor_set(i32_bundle, i32v, 0, sizeof(int32_t) * 6 * (size_t) n_ratios); + ggml_backend_tensor_set(i64_bundle, i64v, 0, sizeof(int64_t) * 5 * (size_t) n_ratios); + } +}; + static bool build_cached_decode_attn_graph( DeepSeek4CachedDecodeAttnGraph & out, ggml_backend_t backend, @@ -1119,7 +1458,8 @@ static bool build_cached_decode_attn_graph( int kv_start, int raw_attn_count, int comp_attn_count, - int index_comp_count) { + int index_comp_count, + const Ds4DecodeSharedInputs * shared = nullptr) { out.free(); const size_t ctx_size = 48 * 1024 * 1024; @@ -1137,6 +1477,8 @@ static bool build_cached_decode_attn_graph( out.n_raw = raw_attn_count; out.n_comp_attn = comp_attn_count; out.n_index_comp = index_comp_count; + out.attn_flush = ratio > 0 && (((kv_start + 1) % ratio) == 0); + out.index_flush = ratio == 4 && (((kv_start + 1) % ratio) == 0); out.compressed = ratio > 0; out.indexed = ratio == 4; @@ -1144,6 +1486,25 @@ static bool build_cached_decode_attn_graph( ggml_set_input(out.sg.inp_embed); out.sg.gf = ggml_new_graph_custom(out.sg.ctx, 2048, false); + const int shared_slot = shared ? shared->slot(ratio) : -1; + if (shared_slot >= 0) { + out.inputs.rope_pos = shared->v_rope_pos[shared_slot]; + out.inputs.neg_pos = shared->v_neg_pos[shared_slot]; + out.inputs.raw_kv_rows = shared->v_raw_row[shared_slot]; + if (ratio > 0) { + out.inputs.attn_ape_row = shared->v_ape_row[shared_slot]; + out.inputs.attn_comp_pos = shared->v_comp_pos[shared_slot]; + out.inputs.attn_state_rows = shared->v_state_row[shared_slot]; + out.inputs.attn_comp_rows = shared->v_comp_row[shared_slot]; + } + if (ratio == 4) { + out.inputs.index_ape_row = shared->v_index_ape[shared_slot]; + out.inputs.index_comp_pos = shared->v_index_cpos[shared_slot]; + out.inputs.index_state_rows = shared->v_index_state[shared_slot]; + out.inputs.index_comp_rows = shared->v_index_comp[shared_slot]; + } + out.uses_shared_inputs = true; + } else { out.inputs.rope_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); out.inputs.neg_pos = ggml_new_tensor_1d(out.sg.ctx, GGML_TYPE_I32, 1); ggml_set_input(out.inputs.rope_pos); @@ -1171,6 +1532,7 @@ static bool build_cached_decode_attn_graph( ggml_set_input(out.inputs.index_state_rows); ggml_set_input(out.inputs.index_comp_rows); } + } std::vector i32_inputs; std::vector i32_array_inputs; @@ -1275,6 +1637,7 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, bool ffn, ggml_tensor * hc_state, ggml_tensor * fn_tensor, + const void * fn_device_override, ggml_tensor * scale_tensor, ggml_tensor * base_tensor, const float * scale_data, @@ -1284,12 +1647,13 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, int sinkhorn_iters, float hc_eps) { #if defined(DFLASH27B_BACKEND_CUDA) || defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) - if (!working || !post || !comb || !hc_state || !fn_tensor || !scale_data || !base_data || - !working->data || !post->data || !comb->data || !hc_state->data || !fn_tensor->data) { + const void * fn_device = fn_device_override ? fn_device_override : (fn_tensor ? fn_tensor->data : nullptr); + if (!working || !post || !comb || !hc_state || !fn_device || !scale_data || !base_data || + !working->data || !post->data || !comb->data || !hc_state->data) { return false; } const bool can_use_device_params = - ds4_backend_is_cuda(backend) && + ds4_backend_is_gpu(backend) && scale_tensor && base_tensor && scale_tensor->data && base_tensor->data && scale_tensor->buffer && base_tensor->buffer && @@ -1297,7 +1661,7 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, !ggml_backend_buffer_is_host(base_tensor->buffer); if (can_use_device_params) { return deepseek4_cuda_hc_pre_device(hc_state->data, - fn_tensor->data, + fn_device, scale_tensor->data, base_tensor->data, n_embd, @@ -1309,7 +1673,7 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, comb->data); } return deepseek4_cuda_hc_pre_device_params(hc_state->data, - fn_tensor->data, + fn_device, scale_data, base_data, n_embd, @@ -1328,6 +1692,7 @@ static bool ds4_try_gpu_hc_pre_device(ggml_tensor * working, (void) ffn; (void) hc_state; (void) fn_tensor; + (void) fn_device_override; (void) scale_tensor; (void) base_tensor; (void) scale_data; @@ -1714,6 +2079,7 @@ static ggml_tensor * build_moe_ffn( const int n_embd = w.n_embd; const int n_used = w.n_expert_used; const int n_ff_exp = w.n_ff_exp; + const bool raw_mmid = ds4_ffn_raw_mmid_enabled(); ggml_tensor * shared_out = build_shared_ffn(ctx, cur, w, L); ggml_tensor * routed_out = nullptr; @@ -1725,20 +2091,26 @@ static ggml_tensor * build_moe_ffn( ggml_tensor * gate_e = ggml_mul_mat_id(ctx, L.ffn_gate_exps, cur_3d, routing.selected); ggml_tensor * up_e = ggml_mul_mat_id(ctx, L.ffn_up_exps, cur_3d, routing.selected); - 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); + if (!raw_mmid) { + 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, routing.selected); - down_e = ggml_reshape_3d(ctx, down_e, n_embd, n_used, n_tokens); + if (!raw_mmid) { + down_e = ggml_reshape_3d(ctx, down_e, n_embd, n_used, n_tokens); + } - ggml_tensor * weights_3d = ggml_reshape_3d(ctx, routing.weights, 1, n_used, n_tokens); - routed_out = ggml_mul(ctx, down_e, weights_3d); - // Sum over dim-1 (n_used experts): permute [n_embd,n_used,n_tokens] -> [n_used,n_embd,n_tokens], - // then sum_rows reduces dim-0, yielding [1,n_embd,n_tokens], reshape to [n_embd,n_tokens]. - routed_out = ggml_cont(ctx, ggml_permute(ctx, routed_out, 1, 0, 2, 3)); - routed_out = ggml_sum_rows(ctx, routed_out); - routed_out = ggml_reshape_2d(ctx, routed_out, n_embd, n_tokens); + if (ds4_ffn_fused_combine_enabled()) { + routed_out = ggml_laguna_moe_combine(ctx, down_e, routing.weights); + } else { + ggml_tensor * weights_3d = ggml_reshape_3d(ctx, routing.weights, 1, n_used, n_tokens); + routed_out = ggml_mul(ctx, down_e, weights_3d); + ggml_tensor * sum_shape = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); + routed_out = ggml_repeat_back(ctx, routed_out, sum_shape); + routed_out = ggml_reshape_2d(ctx, routed_out, n_embd, n_tokens); + } } return ggml_add(ctx, shared_out, routed_out); @@ -1791,6 +2163,8 @@ struct HcWeightsCpu { std::vector fn_data; // [hc_dim * mix_dim] F16 std::vector scale_data; // [3] std::vector base_data; // [2*n_hc + n_hc*n_hc] + void * fn_f16_device = nullptr; // Persistent F16 mirror for quantized HC fn tensors. + size_t fn_f16_device_bytes = 0; bool loaded = false; }; @@ -1811,7 +2185,7 @@ static void cpu_rms_norm(float * out, const float * x, int n, float eps) { for (int i = 0; i < n; i++) out[i] = x[i] * scale; } -static float cpu_dot_f16_row(const uint16_t * row, const float * x, int cols) { +static float cpu_dot_f16_row_scalar(const uint16_t * row, const float * x, int cols) { float acc = 0.0f; for (int c = 0; c < cols; c++) { acc += ggml_fp16_to_fp32(row[c]) * x[c]; @@ -1819,6 +2193,82 @@ static float cpu_dot_f16_row(const uint16_t * row, const float * x, int cols) { return acc; } +#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) +static bool ds4_cpu_has_f16c() { + static int supported = -1; + if (supported < 0) { + __builtin_cpu_init(); + supported = (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("f16c")) ? 1 : 0; + } + return supported == 1; +} + +__attribute__((target("avx2,f16c"))) +static float cpu_dot_f16_row_f16c(const uint16_t * row, const float * x, int cols) { + float acc = 0.0f; + int c = 0; + alignas(32) float prod[8]; + for (; c + 7 < cols; c += 8) { + const __m128i h = _mm_loadu_si128(reinterpret_cast(row + c)); + const __m256 wf = _mm256_cvtph_ps(h); + const __m256 xf = _mm256_loadu_ps(x + c); + _mm256_store_ps(prod, _mm256_mul_ps(wf, xf)); + acc += prod[0]; + acc += prod[1]; + acc += prod[2]; + acc += prod[3]; + acc += prod[4]; + acc += prod[5]; + acc += prod[6]; + acc += prod[7]; + } + for (; c < cols; ++c) { + acc += ggml_fp16_to_fp32(row[c]) * x[c]; + } + return acc; +} +#endif + +#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) +__attribute__((target("avx2,f16c"))) +static void cpu_dot_f16_rows3_f16c(const uint16_t * r0, const uint16_t * r1, const uint16_t * r2, + const float * x, int cols, + float * o0, float * o1, float * o2) { + float a0 = 0.0f, a1 = 0.0f, a2 = 0.0f; + int c = 0; + alignas(32) float p0[8], p1[8], p2[8]; + for (; c + 7 < cols; c += 8) { + const __m256 xf = _mm256_loadu_ps(x + c); + _mm256_store_ps(p0, _mm256_mul_ps(_mm256_cvtph_ps(_mm_loadu_si128(reinterpret_cast(r0 + c))), xf)); + _mm256_store_ps(p1, _mm256_mul_ps(_mm256_cvtph_ps(_mm_loadu_si128(reinterpret_cast(r1 + c))), xf)); + _mm256_store_ps(p2, _mm256_mul_ps(_mm256_cvtph_ps(_mm_loadu_si128(reinterpret_cast(r2 + c))), xf)); + a0 += p0[0]; a1 += p1[0]; a2 += p2[0]; + a0 += p0[1]; a1 += p1[1]; a2 += p2[1]; + a0 += p0[2]; a1 += p1[2]; a2 += p2[2]; + a0 += p0[3]; a1 += p1[3]; a2 += p2[3]; + a0 += p0[4]; a1 += p1[4]; a2 += p2[4]; + a0 += p0[5]; a1 += p1[5]; a2 += p2[5]; + a0 += p0[6]; a1 += p1[6]; a2 += p2[6]; + a0 += p0[7]; a1 += p1[7]; a2 += p2[7]; + } + for (; c < cols; ++c) { + a0 += ggml_fp16_to_fp32(r0[c]) * x[c]; + a1 += ggml_fp16_to_fp32(r1[c]) * x[c]; + a2 += ggml_fp16_to_fp32(r2[c]) * x[c]; + } + *o0 = a0; *o1 = a1; *o2 = a2; +} +#endif + +static float cpu_dot_f16_row(const uint16_t * row, const float * x, int cols) { +#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) + if (ds4_cpu_has_f16c()) { + return cpu_dot_f16_row_f16c(row, x, cols); + } +#endif + return cpu_dot_f16_row_scalar(row, x, cols); +} + static void cpu_matvec_f16_serial(float * out, const uint16_t * mat, const float * x, int rows, int cols) { // mat: [cols, rows] in row-major F16 (ggml layout: ne[0]=cols, ne[1]=rows) // out[r] = dot(mat_row_r, x) for r in [0, rows) @@ -1839,6 +2289,111 @@ static void cpu_matvec_f16(float * out, const uint16_t * mat, const float * x, i }); } +// Persistent worker pool for the decode-path HC fn matvec. Splitting rows +// across threads leaves each row's accumulation order untouched, so results +// are bit-identical to the serial path; only wall time changes. Decode issues +// ~86 of these 24x16384 matvecs per token, so thread spawn (or condvar wake) +// per call would dominate — workers spin briefly then yield between jobs. +struct Ds4HcMatvecPool { + struct Job { const uint16_t * mat; const float * x; float * out; int rows; int cols; }; + std::mutex client_mu; + std::atomic seq{0}; + std::atomic remaining{0}; + Job job{}; + std::vector workers; + std::atomic stop{false}; + int nth = 0; + + static void cpu_relax() { +#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) + __builtin_ia32_pause(); +#endif + } + + Ds4HcMatvecPool() { + unsigned hw = std::thread::hardware_concurrency(); + nth = hw == 0 ? 4 : (int)(hw < 8 ? hw : 8); + for (int i = 0; i < nth; ++i) { + workers.emplace_back([this, i]() { + uint64_t last = 0; + for (;;) { + uint64_t s; + int spins = 0; + while ((s = seq.load(std::memory_order_acquire)) == last) { + if (stop.load(std::memory_order_relaxed)) return; + if (++spins < 4096) { cpu_relax(); } + else { std::this_thread::yield(); spins = 0; } + } + last = s; + const Job j = job; + const int chunk = (j.rows + nth - 1) / nth; + const int r0 = i * chunk; + const int r1 = j.rows < r0 + chunk ? j.rows : r0 + chunk; + if (row_fn) { + for (int r = r0; r < r1; ++r) row_fn(r); + } else { + int r = r0; +#if (defined(__x86_64__) || defined(_M_X64)) && (defined(__GNUC__) || defined(__clang__)) + if (ds4_cpu_has_f16c()) { + for (; r + 2 < r1; r += 3) { + cpu_dot_f16_rows3_f16c( + j.mat + (size_t) (r + 0) * j.cols, + j.mat + (size_t) (r + 1) * j.cols, + j.mat + (size_t) (r + 2) * j.cols, + j.x, j.cols, + &j.out[r + 0], &j.out[r + 1], &j.out[r + 2]); + } + } +#endif + for (; r < r1; ++r) { + j.out[r] = cpu_dot_f16_row(j.mat + (size_t) r * j.cols, j.x, j.cols); + } + } + remaining.fetch_sub(1, std::memory_order_acq_rel); + } + }); + } + } + ~Ds4HcMatvecPool() { + stop.store(true); + for (auto & t : workers) t.join(); + } + void run(const uint16_t * mat, const float * x, float * out, int rows, int cols) { + std::lock_guard lk(client_mu); + row_fn = nullptr; + job = {mat, x, out, rows, cols}; + remaining.store(nth, std::memory_order_release); + seq.fetch_add(1, std::memory_order_release); + int spins = 0; + while (remaining.load(std::memory_order_acquire) != 0) { + if (++spins < 65536) { cpu_relax(); } + else { std::this_thread::yield(); spins = 0; } + } + } + + // Generic variant: invoke fn(row) for each row in [0, rows), rows split + // across workers with the same static chunking as run(). + std::function row_fn; + void run_custom(int rows, std::function fn) { + std::lock_guard lk(client_mu); + row_fn = std::move(fn); + job = {nullptr, nullptr, nullptr, rows, 0}; + remaining.store(nth, std::memory_order_release); + seq.fetch_add(1, std::memory_order_release); + int spins = 0; + while (remaining.load(std::memory_order_acquire) != 0) { + if (++spins < 65536) { cpu_relax(); } + else { std::this_thread::yield(); spins = 0; } + } + row_fn = nullptr; + } +}; + +static void cpu_matvec_f16_pooled(float * out, const uint16_t * mat, const float * x, int rows, int cols) { + static Ds4HcMatvecPool pool; + pool.run(mat, x, out, rows, cols); +} + static void cpu_hc_sinkhorn(float * out, const float * mix, const float * scale, const float * base, int n_hc, int iters, float eps) { const float pre_scale = scale[0]; @@ -1969,7 +2524,7 @@ static void cpu_hc_pre_into(float * working, if (serial_fn) { cpu_matvec_f16_serial(mix, fn_data, flat, mix_dim, hc_dim); } else { - cpu_matvec_f16(mix, fn_data, flat, mix_dim, hc_dim); + cpu_matvec_f16_pooled(mix, fn_data, flat, mix_dim, hc_dim); } finish_hc_pre_from_mix_into(working, post, comb, hc_state, mix, scale_data, base_data, @@ -2034,7 +2589,8 @@ static void hc_pre_auto_into(float * working, int sinkhorn_iters, float hc_eps, float * flat, - float * mix_scratch) { + float * mix_scratch, + bool serial_fn) { #if defined(DFLASH27B_BACKEND_CUDA) if (ds4_hc_cuda_enabled() && fn_tensor && fn_tensor->data) { float mix[24]; @@ -2053,7 +2609,7 @@ static void hc_pre_auto_into(float * working, cpu_hc_pre_into(working, post, comb, hc_state, weights.fn_data.data(), weights.scale_data.data(), weights.base_data.data(), - n_embd, n_hc, sinkhorn_iters, hc_eps, flat, mix_scratch, true); + n_embd, n_hc, sinkhorn_iters, hc_eps, flat, mix_scratch, serial_fn); } static void hc_pre_batch(std::vector & working, @@ -2087,7 +2643,8 @@ static void hc_pre_batch(std::vector & working, sinkhorn_iters, hc_eps, flat.data(), - mix); + mix, + /*serial_fn=*/n_tokens > 1); } }); } @@ -2116,6 +2673,24 @@ static void hc_post_batch(std::vector & out_hc, int n_hc) { const size_t hc_dim = (size_t)n_embd * (size_t)n_hc; out_hc.resize((size_t)n_tokens * hc_dim); + if (n_tokens == 1) { + // Decode: split the n_hc independent destination streams across the + // persistent pool. Per-element accumulation order is unchanged, so + // the result is bit-identical to the serial loop. + static Ds4HcMatvecPool post_pool; + struct Ctx { const float * block; const float * res; const float * post; const float * comb; float * out; int n_embd; int n_hc; }; + Ctx c{block_out, residual_hc, post, comb, out_hc.data(), n_embd, n_hc}; + post_pool.run_custom(n_hc, [&c](int h) { + for (int d = 0; d < c.n_embd; ++d) { + float acc = c.block[d] * c.post[h]; + for (int src = 0; src < c.n_hc; ++src) { + acc += c.comb[h + src * c.n_hc] * c.res[(size_t)src * c.n_embd + d]; + } + c.out[(size_t)h * c.n_embd + d] = acc; + } + }); + return; + } ds4_parallel_for_tokens(n_tokens, 8, [&](int t0, int t1) { for (int t = t0; t < t1; ++t) { cpu_hc_post(out_hc.data() + (size_t)t * hc_dim, @@ -2163,18 +2738,124 @@ static void hc_output_batch(std::vector & final_embd, }); } +static bool load_tensor_to_f32_cpu(std::vector & dst, ggml_tensor * t) { + if (!t) return false; + + const size_t elems = ggml_nelements(t); + dst.resize(elems); + if (elems == 0) return true; + + if (t->type == GGML_TYPE_F32) { + ggml_backend_tensor_get(t, dst.data(), 0, ggml_nbytes(t)); + return true; + } + + const ggml_type_traits * tr = ggml_get_type_traits(t->type); + if (!tr || !tr->to_float || t->ne[0] <= 0) return false; + + std::vector raw(ggml_nbytes(t)); + ggml_backend_tensor_get(t, raw.data(), 0, raw.size()); + + const int64_t cols = t->ne[0]; + const int64_t rows = (int64_t)elems / cols; + const size_t row_bytes = ggml_row_size(t->type, cols); + for (int64_t r = 0; r < rows; ++r) { + tr->to_float(raw.data() + (size_t)r * row_bytes, + dst.data() + (size_t)r * (size_t)cols, + cols); + } + return true; +} + +static bool load_tensor_to_f16_cpu(std::vector & dst, ggml_tensor * t) { + if (!t) return false; + + const size_t elems = ggml_nelements(t); + dst.resize(elems); + if (elems == 0) return true; + + if (t->type == GGML_TYPE_F16) { + ggml_backend_tensor_get(t, dst.data(), 0, ggml_nbytes(t)); + return true; + } + + std::vector f32; + if (!load_tensor_to_f32_cpu(f32, t)) return false; + ggml_fp32_to_fp16_row(f32.data(), reinterpret_cast(dst.data()), (int64_t)elems); + return true; +} + static void load_hc_weights_cpu(HcWeightsCpu & dst, ggml_tensor * fn, ggml_tensor * scale, ggml_tensor * base) { if (!fn || !scale || !base || dst.loaded) return; - dst.fn_data.resize(ggml_nelements(fn)); - dst.scale_data.resize(ggml_nelements(scale)); - dst.base_data.resize(ggml_nelements(base)); - ggml_backend_tensor_get(fn, dst.fn_data.data(), 0, ggml_nbytes(fn)); - ggml_backend_tensor_get(scale, dst.scale_data.data(), 0, ggml_nbytes(scale)); - ggml_backend_tensor_get(base, dst.base_data.data(), 0, ggml_nbytes(base)); + if (!load_tensor_to_f16_cpu(dst.fn_data, fn) || + !load_tensor_to_f32_cpu(dst.scale_data, scale) || + !load_tensor_to_f32_cpu(dst.base_data, base)) { + dst.fn_data.clear(); + dst.scale_data.clear(); + dst.base_data.clear(); + return; + } dst.loaded = true; } +static void release_hc_fn_device(HcWeightsCpu & w) { +#if defined(DFLASH27B_BACKEND_CUDA) || defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + if (w.fn_f16_device) { + deepseek4_cuda_hc_free(w.fn_f16_device); + } +#endif + w.fn_f16_device = nullptr; + w.fn_f16_device_bytes = 0; +} + +static void reset_hc_weights_cpu(HcWeightsCpu & w) { + release_hc_fn_device(w); + w.fn_data.clear(); + w.scale_data.clear(); + w.base_data.clear(); + w.loaded = false; +} + +static void reset_hc_layer_weights_cpu(std::vector & weights) { + for (HcLayerWeightsCpu & layer : weights) { + reset_hc_weights_cpu(layer.attn); + reset_hc_weights_cpu(layer.ffn); + } + weights.clear(); +} + +static bool ensure_hc_fn_device(HcWeightsCpu & w, ggml_tensor * fn) { + if (!fn || !fn->data) return false; + if (fn->type == GGML_TYPE_F16) return true; + if (!ds4_rocmfpx_hc_gpu_enabled()) return false; + if (!w.loaded || w.fn_data.empty()) return false; + +#if defined(DFLASH27B_BACKEND_CUDA) || defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) + const size_t bytes = w.fn_data.size() * sizeof(uint16_t); + if (w.fn_f16_device && w.fn_f16_device_bytes == bytes) { + return true; + } + release_hc_fn_device(w); + if (!deepseek4_cuda_hc_upload_f16(w.fn_data.data(), bytes, &w.fn_f16_device)) { + w.fn_f16_device = nullptr; + w.fn_f16_device_bytes = 0; + return false; + } + w.fn_f16_device_bytes = bytes; + return true; +#else + return false; +#endif +} + +static const void * hc_fn_device_ptr(const HcWeightsCpu & w, ggml_tensor * fn) { + if (!fn) return nullptr; + if (fn->type == GGML_TYPE_F16) return fn->data; + if (!ds4_rocmfpx_hc_gpu_enabled()) return nullptr; + return w.fn_f16_device; +} + static bool load_hash_routing_cpu(HashRoutingTableCpu & dst, ggml_tensor * table) { if (dst.loaded) return true; if (!table) return false; @@ -2669,6 +3350,7 @@ bool deepseek4_step( MoeHybridStreamEngine * stream_engine, DeepSeek4StepTelemetry * telemetry, MoeHybridRoutingStats * routing_stats) { + const auto step_t0 = Ds4TimingClock::now(); if (w.moe_hybrid && moe_hybrid != nullptr) { return deepseek4_step_hybrid(backend, w, cache, *moe_hybrid, @@ -2678,23 +3360,71 @@ bool deepseek4_step( const int n_embd = w.n_embd; const int n_layer = w.n_layer; + const size_t ctx_size = ds4_full_step_meta_size(n_tokens); + const bool reuse_full_step_decode = + n_tokens == 1 && + ds4_backend_is_gpu(backend) && + !ds4_env_flag("DFLASH_DS4_DISABLE_FULL_STEP_DECODE_REUSE"); + static thread_local DeepSeek4LegacyFullStepCache full_step_cache; + StepGraph * cached_sg = nullptr; + if (reuse_full_step_decode) { + if (full_step_cache.owner_ctx != w.ctx || full_step_cache.backend != backend) { + full_step_cache.free(); + full_step_cache.owner_ctx = w.ctx; + full_step_cache.backend = backend; + } else { + step_graph_free(full_step_cache.sg); + } + cached_sg = &full_step_cache.sg; + if (full_step_cache.meta_arena.size() < ctx_size) { + full_step_cache.meta_arena.resize(ctx_size); + } + } - // Create compute graph context — need large budget for MoE layers - const size_t ctx_size = ggml_tensor_overhead() * 65536 + 16 * 1024 * 1024; ggml_init_params params{}; - params.mem_size = ctx_size; - params.mem_buffer = nullptr; + params.mem_size = cached_sg ? full_step_cache.meta_arena.size() : ctx_size; + params.mem_buffer = cached_sg ? full_step_cache.meta_arena.data() : nullptr; params.no_alloc = true; + const auto full_build_t0 = Ds4TimingClock::now(); ggml_context * ctx = ggml_init(params); if (!ctx) return false; + if (cached_sg) { + cached_sg->ctx = ctx; + } + if (telemetry) { + telemetry->full_graph_build_us += ds4_elapsed_us(full_build_t0, Ds4TimingClock::now()); + } + + ggml_gallocr_t alloc = nullptr; + bool owns_alloc = false; + auto release_full_step = [&]() { + if (cached_sg) { + step_graph_free(*cached_sg); + return; + } + if (alloc && owns_alloc) { + ggml_gallocr_free(alloc); + alloc = nullptr; + } + if (ctx) { + ggml_free(ctx); + ctx = nullptr; + } + }; // Input embeddings ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); ggml_set_name(inp, "inp_embed"); ggml_set_input(inp); + if (cached_sg) { + cached_sg->inp_embed = inp; + } ggml_tensor * cur = inp; - ggml_cgraph * gf = ggml_new_graph_custom(ctx, 32768, false); + ggml_cgraph * gf = ggml_new_graph_custom(ctx, ds4_full_step_graph_size(n_tokens), false); + if (cached_sg) { + cached_sg->gf = gf; + } std::vector i32_inputs; std::vector i32_array_inputs; std::vector i64_array_inputs; @@ -2743,20 +3473,36 @@ bool deepseek4_step( ggml_tensor * logits = ggml_mul_mat(ctx, w.output, cur); ggml_set_name(logits, "logits"); ggml_set_output(logits); + if (cached_sg) { + cached_sg->logits = logits; + } // ── Build and run graph ───────────────────────────────────────────── ggml_build_forward_expand(gf, logits); // Allocate - ggml_gallocr_t alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (cached_sg) { + if (!cached_sg->alloc) { + cached_sg->alloc = + ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + } + alloc = cached_sg->alloc; + } else { + alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + owns_alloc = true; + } + const auto full_alloc_t0 = Ds4TimingClock::now(); if (!ggml_gallocr_alloc_graph(alloc, gf)) { std::fprintf(stderr, "[deepseek4] graph allocation failed\n"); - ggml_gallocr_free(alloc); - ggml_free(ctx); + release_full_step(); return false; } + if (telemetry) { + telemetry->full_graph_alloc_us += ds4_elapsed_us(full_alloc_t0, Ds4TimingClock::now()); + } // Set input data + const auto full_set_t0 = Ds4TimingClock::now(); ggml_backend_tensor_set(inp, embed, 0, n_embd * n_tokens * sizeof(float)); for (const DeepSeek4I32InputBinding & binding : i32_inputs) { ggml_backend_tensor_set(binding.tensor, &binding.value, 0, sizeof(binding.value)); @@ -2769,23 +3515,32 @@ bool deepseek4_step( ggml_backend_tensor_set(binding.tensor, binding.values.data(), 0, sizeof(int64_t) * binding.values.size()); } + if (telemetry) { + telemetry->full_graph_set_us += ds4_elapsed_us(full_set_t0, Ds4TimingClock::now()); + } // Compute + const auto full_compute_t0 = Ds4TimingClock::now(); if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[deepseek4] graph compute failed\n"); - ggml_gallocr_free(alloc); - ggml_free(ctx); + release_full_step(); return false; } + if (telemetry) { + telemetry->full_graph_compute_us += ds4_elapsed_us(full_compute_t0, Ds4TimingClock::now()); + } // Read logits (only last token for generation) + const auto full_read_t0 = Ds4TimingClock::now(); out_logits.resize(w.n_vocab); const size_t logits_offset = (size_t)(n_tokens - 1) * w.n_vocab * sizeof(float); ggml_backend_tensor_get(logits, out_logits.data(), logits_offset, w.n_vocab * sizeof(float)); + if (telemetry) { + telemetry->full_graph_read_us += ds4_elapsed_us(full_read_t0, Ds4TimingClock::now()); + } - ggml_gallocr_free(alloc); - ggml_free(ctx); + release_full_step(); const int next_pos = kv_start + n_tokens; for (int il = 0; il < n_layer; ++il) { @@ -2801,9 +3556,633 @@ bool deepseek4_step( } cache.cur_pos = next_pos; + if (telemetry) { + telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + } + return true; +} + +// ─── Fused single-graph decode (n_tokens == 1) ────────────────────────── +// Chains all layers (HC pre → attention → HC post → HC pre → FFN → HC post) +// plus the output HC merge and lm_head into ONE cached ggml graph, so a +// decode step is a single ggml_backend_graph_compute with one logits +// readback instead of ~90 per-layer graph launches with host round-trips. +// HC Sinkhorn mixing runs in the fused GGML_OP_DS4_HC op (one kernel per +// sublayer instead of ~170 tiny ops for 20 Sinkhorn iterations). +// +// Compressed-KV reads are padded to DS4_FUSED_COMP_PAD rows with an additive +// score mask, and each structural variant (flush pattern) lives in its own +// slot with a private metadata arena. Tensor addresses therefore stay stable +// while a variant recurs, which is what the ggml-cuda/HIP graph cache keys +// on, enabling graph replay for the bulk of decode steps. + +// Opt-in refinement of the fused graph: one topology for flush and non-flush +// steps (pooling every step + input-redirected state flush). Needed for +// CUDA/HIP graph replay experiments; costs ~3% decode speed, and HIP graph +// replay is a net loss on gfx1151 anyway, so default OFF. +static bool ds4_fused_stable_graph_enabled() { + static int enabled = -1; + if (enabled < 0) { + enabled = ds4_env_flag("DFLASH_DS4_FUSED_STABLE_GRAPH") ? 1 : 0; + } + return enabled == 1; +} + +// Opt-in: single-graph decode with GPU hyper-connections. Deterministic and +// near-bit-identical, but expf ULP differences in the sinkhorn iterations can +// diverge from the CPU-HC reference after tens of tokens. Default OFF. +static bool ds4_fused_decode_enabled() { + static int enabled = -1; + if (enabled < 0) { + enabled = ds4_env_flag("DFLASH_DS4_FUSED_DECODE") ? 1 : 0; + } + return enabled == 1; +} + +struct DeepSeek4FusedDecodeGraph { + std::vector shape_key; + uint64_t last_use = 0; + StepGraph sg; + ggml_tensor * inp_embed = nullptr; + ggml_tensor * i32_bundle = nullptr; + ggml_tensor * i64_bundle = nullptr; + ggml_tensor * mask_bundle = nullptr; // additive score mask (0 / -1e30), may be null + ggml_tensor * flush_rows = nullptr; // i64[4]: [0..3] on flush steps, [4..7] otherwise + std::vector hash_ids; + ggml_tensor * logits = nullptr; + + void reset_nodes() { + inp_embed = nullptr; + i32_bundle = nullptr; + i64_bundle = nullptr; + mask_bundle = nullptr; + flush_rows = nullptr; + logits = nullptr; + hash_ids.clear(); + shape_key.clear(); + last_use = 0; + } + + bool built() const { + return sg.ctx && sg.gf && logits; + } + + void destroy() { + step_graph_destroy(sg); + reset_nodes(); + } +}; + +struct DeepSeek4FusedDecodeCache { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + bool disabled = false; + uint64_t counter = 0; + std::array slots; + + // Persistent F16 mirrors of the (quantized) HC fn projection weights so + // the fused graph matches the numerics of the reference HC paths, which + // always dequantize fn to F16 before the mix matvec. + ggml_context * fn_ctx = nullptr; + ggml_backend_buffer_t fn_buf = nullptr; + std::vector fn_attn_f16; + std::vector fn_ffn_f16; + ggml_tensor * fn_out_f16 = nullptr; + + void destroy() { + for (auto & s : slots) s.destroy(); + if (fn_buf) { ggml_backend_buffer_free(fn_buf); fn_buf = nullptr; } + if (fn_ctx) { ggml_free(fn_ctx); fn_ctx = nullptr; } + fn_attn_f16.clear(); + fn_ffn_f16.clear(); + fn_out_f16 = nullptr; + owner_ctx = nullptr; + backend = nullptr; + disabled = false; + counter = 0; + } +}; + +static ggml_tensor * ds4_fused_hc_base_f32(ggml_context * ctx, ggml_tensor * base) { + if (!base) return nullptr; + ggml_tensor * b = base; + if (b->type != GGML_TYPE_F32) { + b = ggml_cast(ctx, b, GGML_TYPE_F32); + } + return ggml_reshape_1d(ctx, b, ggml_nelements(b)); +} + +static ggml_tensor * ds4_build_fused_hc_pre( + ggml_context * ctx, + const DeepSeek4Weights & w, + ggml_tensor * hc_flat, // [n_embd*n_hc] contiguous f32 + ggml_tensor * fn, + ggml_tensor * base, + const HcWeightsCpu & cw, + ggml_tensor ** out_split) { + if (!fn || !base || !cw.loaded || cw.scale_data.size() < 3) return nullptr; + const int mix_dim = 2 * w.n_hc + w.n_hc * w.n_hc; + ggml_tensor * normed = ggml_rms_norm(ctx, hc_flat, w.hc_eps); + ggml_tensor * mix = ggml_mul_mat(ctx, fn, normed); + mix = ggml_reshape_1d(ctx, mix, mix_dim); + ggml_tensor * base_f32 = ds4_fused_hc_base_f32(ctx, base); + ggml_tensor * pre = ggml_ds4_hc_pre(ctx, mix, base_f32, hc_flat, + w.n_hc, w.n_hc_sinkhorn_iter, + cw.scale_data[0], cw.scale_data[1], cw.scale_data[2]); + *out_split = ggml_view_1d(ctx, pre, mix_dim, (size_t) w.n_embd * sizeof(float)); + return ggml_view_1d(ctx, pre, w.n_embd, 0); +} + +static ggml_tensor * ds4_build_hash_routed_ffn( + ggml_context * ctx, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + ggml_tensor * ffn_normed, + ggml_tensor * hash_ids) { + 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)); + + const int n_used = w.n_expert_used; + const int n_ff_exp = w.n_ff_exp; + const bool raw_mmid = ds4_ffn_raw_mmid_enabled(); + 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); + ggml_tensor * up_e = ggml_mul_mat_id(ctx, L.ffn_up_exps, cur_3d, hash_ids); + if (!raw_mmid) { + 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); + } + 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); + if (!raw_mmid) { + down_e = ggml_reshape_3d(ctx, down_e, w.n_embd, n_used, 1); + } + + ggml_tensor * probs_3d = ggml_reshape_3d(ctx, probs, 1, w.n_expert, 1); + ggml_tensor * weights = ggml_get_rows(ctx, probs_3d, hash_ids); + weights = ggml_reshape_2d(ctx, weights, n_used, 1); + 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); + if (w.expert_weight_scale != 1.0f) { + weights = ggml_scale(ctx, weights, w.expert_weight_scale); + } + + ggml_tensor * routed_out = nullptr; + if (ds4_ffn_fused_combine_enabled()) { + routed_out = ggml_laguna_moe_combine(ctx, down_e, weights); + } else { + ggml_tensor * weights_3d = ggml_reshape_3d(ctx, weights, 1, n_used, 1); + routed_out = ggml_mul(ctx, down_e, weights_3d); + ggml_tensor * sum_shape = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w.n_embd, 1, 1); + routed_out = ggml_repeat_back(ctx, routed_out, sum_shape); + routed_out = ggml_reshape_2d(ctx, routed_out, w.n_embd, 1); + } + return ggml_add(ctx, shared_out, routed_out); +} + +static bool ds4_fused_ensure_fn_mirrors( + DeepSeek4FusedDecodeCache & fc, + ggml_backend_t backend, + const DeepSeek4Weights & w, + const std::vector & hc_weights, + const HcWeightsCpu & hc_out_weights) { + if (fc.fn_ctx && fc.fn_buf && fc.fn_attn_f16.size() == (size_t) w.n_layer && fc.fn_out_f16) { + return true; + } + if (fc.fn_buf) { ggml_backend_buffer_free(fc.fn_buf); fc.fn_buf = nullptr; } + if (fc.fn_ctx) { ggml_free(fc.fn_ctx); fc.fn_ctx = nullptr; } + 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; + ggml_init_params params{}; + params.mem_size = ggml_tensor_overhead() * (size_t) (2 * w.n_layer + 4) + 4096; + params.mem_buffer = nullptr; + params.no_alloc = true; + fc.fn_ctx = ggml_init(params); + if (!fc.fn_ctx) return false; + fc.fn_attn_f16.assign((size_t) w.n_layer, nullptr); + fc.fn_ffn_f16.assign((size_t) w.n_layer, nullptr); + for (int il = 0; il < w.n_layer; ++il) { + fc.fn_attn_f16[(size_t) il] = ggml_new_tensor_2d(fc.fn_ctx, GGML_TYPE_F16, hc_dim, mix_dim); + fc.fn_ffn_f16[(size_t) il] = ggml_new_tensor_2d(fc.fn_ctx, GGML_TYPE_F16, hc_dim, mix_dim); + } + fc.fn_out_f16 = ggml_new_tensor_2d(fc.fn_ctx, GGML_TYPE_F16, hc_dim, w.n_hc); + fc.fn_buf = ggml_backend_alloc_ctx_tensors(fc.fn_ctx, backend); + if (!fc.fn_buf) { + ggml_free(fc.fn_ctx); + fc.fn_ctx = nullptr; + return false; + } + for (int il = 0; il < w.n_layer; ++il) { + const auto & a = hc_weights[(size_t) il].attn.fn_data; + const auto & f = hc_weights[(size_t) il].ffn.fn_data; + if ((int64_t) a.size() != hc_dim * mix_dim || (int64_t) f.size() != hc_dim * mix_dim) { + return false; + } + ggml_backend_tensor_set(fc.fn_attn_f16[(size_t) il], a.data(), 0, a.size() * sizeof(uint16_t)); + ggml_backend_tensor_set(fc.fn_ffn_f16[(size_t) il], f.data(), 0, f.size() * sizeof(uint16_t)); + } + const auto & o = hc_out_weights.fn_data; + if ((int64_t) o.size() != hc_dim * w.n_hc) { + return false; + } + ggml_backend_tensor_set(fc.fn_out_f16, o.data(), 0, o.size() * sizeof(uint16_t)); + return true; +} + +static bool ds4_build_fused_decode_graph( + DeepSeek4FusedDecodeCache & fc, + DeepSeek4FusedDecodeGraph & fg, + ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const std::vector & hc_weights, + const HcWeightsCpu & hc_out_weights, + const std::vector & hash_tables, + int kv_start, + bool have_token_ids, + std::vector && shape_key) { + step_graph_free(fg.sg); + fg.reset_nodes(); + fg.hash_ids.assign((size_t) w.n_layer, nullptr); + + const int n_embd = w.n_embd; + const int n_hc = w.n_hc; + const int n_raw = std::min(kv_start + 1, w.n_swa); + const int token_pos = kv_start; + const size_t arena_size = 192u * 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(); + params.mem_buffer = fg.sg.meta_arena.data(); + params.no_alloc = true; + 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, 32768, false); + ggml_cgraph * gf = fg.sg.gf; + + fg.inp_embed = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, 1); + ggml_set_input(fg.inp_embed); + fg.i32_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 6 * (int64_t) w.n_layer); + ggml_set_input(fg.i32_bundle); + fg.i64_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 5 * (int64_t) w.n_layer); + ggml_set_input(fg.i64_bundle); + if (ds4_fused_stable_graph_enabled()) { + fg.flush_rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 4); + ggml_set_input(fg.flush_rows); + } + + // One additive score-mask bundle covering EVERY layer: [n_swa raw rows ++ + // padded comp rows]. All layers take the masked full-ring attention branch + // so the graph topology never depends on the live raw-row count. + 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]); + } + mask_total += w.n_swa + padded; + } + fg.mask_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, mask_total); + ggml_set_input(fg.mask_bundle); + int64_t mask_off = 0; + + // hc state starts as the token embedding replicated into every stream + ggml_tensor * hc_cur = ggml_repeat_4d(ctx, fg.inp_embed, n_embd, n_hc, 1, 1); + + for (int il = 0; il < w.n_layer; ++il) { + const DeepSeek4Layer & L = w.layers[(size_t) il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; + const int ratio = (int) w.compress_ratios[il]; + + // ── HC pre (attention) ───────────────────────────────────── + ggml_tensor * hc_flat = ggml_reshape_1d(ctx, hc_cur, (int64_t) n_embd * n_hc); + ggml_tensor * split_attn = nullptr; + ggml_tensor * working = ds4_build_fused_hc_pre(ctx, w, hc_flat, + fc.fn_attn_f16[(size_t) il], L.hc_attn_base, + hlw.attn, &split_attn); + if (!working) return false; + ggml_tensor * attn_in = ggml_reshape_2d(ctx, working, n_embd, 1); + + // ── Attention (inputs are views into the shared bundles) ── + DeepSeek4AttentionGraphInputs ain{}; + ain.flush_rows = ds4_fused_stable_graph_enabled() ? fg.flush_rows : nullptr; + ain.rope_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 6 + 0) * sizeof(int32_t)); + ain.neg_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 6 + 1) * sizeof(int32_t)); + ain.raw_kv_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), + ((size_t) il * 5 + 0) * sizeof(int64_t)); + if (ratio > 0) { + ain.attn_ape_row = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 6 + 2) * sizeof(int32_t)); + ain.attn_comp_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 6 + 3) * sizeof(int32_t)); + ain.attn_state_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), + ((size_t) il * 5 + 1) * sizeof(int64_t)); + ain.attn_comp_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), + ((size_t) il * 5 + 2) * sizeof(int64_t)); + } + if (ratio == 4) { + ain.index_ape_row = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 6 + 4) * sizeof(int32_t)); + ain.index_comp_pos = ggml_view_1d(ctx, fg.i32_bundle, 1, ((size_t) il * 6 + 5) * sizeof(int32_t)); + ain.index_state_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), + ((size_t) il * 5 + 3) * sizeof(int64_t)); + ain.index_comp_rows = ggml_view_2d(ctx, fg.i64_bundle, 1, 1, sizeof(int64_t), + ((size_t) il * 5 + 4) * 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); + padded = ds4_padded_comp_rows(n_comp, (int) lc.comp_kv->ne[1]); + } + const int64_t n_attn = (int64_t) w.n_swa + padded; + ain.attn_row_mask = ggml_view_2d(ctx, fg.mask_bundle, n_attn, 1, + n_attn * sizeof(float), + (size_t) mask_off * sizeof(float)); + ain.padded_comp = padded; + mask_off += n_attn; + } + + 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, 1, &ain, + i32b, i32ab, i64ab); + if (!attn_out) return false; + if (!i32b.empty() || !i32ab.empty() || !i64ab.empty()) { + std::fprintf(stderr, + "[deepseek4] fused decode: layer %d created %zu/%zu/%zu dynamic bindings; cannot fuse\n", + il, i32b.size(), i32ab.size(), i64ab.size()); + return false; + } + + // ── HC post (attention) ──────────────────────────────────── + ggml_tensor * attn_out_flat = ggml_reshape_1d(ctx, attn_out, n_embd); + ggml_tensor * hc_next = ggml_ds4_hc_post(ctx, hc_flat, attn_out_flat, split_attn, n_hc); + hc_cur = ggml_reshape_2d(ctx, hc_next, n_embd, n_hc); + + // ── HC pre (FFN) ─────────────────────────────────────────── + hc_flat = ggml_reshape_1d(ctx, hc_cur, (int64_t) n_embd * n_hc); + ggml_tensor * split_ffn = nullptr; + ggml_tensor * fworking = ds4_build_fused_hc_pre(ctx, w, hc_flat, + fc.fn_ffn_f16[(size_t) il], L.hc_ffn_base, + hlw.ffn, &split_ffn); + if (!fworking) return false; + ggml_tensor * ffn_in = ggml_reshape_2d(ctx, fworking, n_embd, 1); + + // ── FFN ──────────────────────────────────────────────────── + ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, L.ffn_norm, w.rms_eps); + ggml_tensor * ffn_out = nullptr; + const bool hash_routed = il < w.n_hash_layer && L.ffn_gate_tid2eid && + have_token_ids && hash_tables[(size_t) il].loaded; + if (hash_routed) { + ggml_tensor * hids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, w.n_expert_used, 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); + } else { + ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, 1); + } + if (!ffn_out) return false; + + // ── HC post (FFN) ────────────────────────────────────────── + ggml_tensor * ffn_out_flat = ggml_reshape_1d(ctx, ffn_out, n_embd); + hc_next = ggml_ds4_hc_post(ctx, hc_flat, ffn_out_flat, split_ffn, n_hc); + hc_cur = ggml_reshape_2d(ctx, hc_next, n_embd, n_hc); + } + + // ── Output: HC merge → norm → lm_head ────────────────────────── + ggml_tensor * hc_flat = ggml_reshape_1d(ctx, hc_cur, (int64_t) n_embd * n_hc); + ggml_tensor * onorm = ggml_rms_norm(ctx, hc_flat, w.hc_eps); + ggml_tensor * omix = ggml_mul_mat(ctx, fc.fn_out_f16, onorm); + omix = ggml_reshape_1d(ctx, omix, ggml_nelements(omix)); + ggml_tensor * obase = ds4_fused_hc_base_f32(ctx, w.output_hc_base); + if (!obase || hc_out_weights.scale_data.empty()) return false; + ggml_tensor * final_embd = ggml_ds4_hc_out(ctx, omix, obase, hc_flat, 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); + fg.logits = ggml_mul_mat(ctx, w.output, out_normed); + ggml_set_output(fg.logits); + ggml_build_forward_expand(gf, fg.logits); + + 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, "[deepseek4] fused decode graph alloc failed\n"); + return false; + } + + fg.shape_key = std::move(shape_key); return true; } +// Returns 1 on success (out_logits filled), 0 to fall back to the per-layer +// path, -1 on a hard failure after cache state may have been touched. +static int ds4_try_fused_decode_step( + 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 kv_start, + std::vector & out_logits, + const int32_t * token_ids, + DeepSeek4StepTelemetry * telemetry) { + if (fc.disabled) return 0; + if (!hc_out_weights.loaded || hc_out_weights.scale_data.empty() || + !w.output_hc_fn || !w.output_hc_base) { + fc.disabled = true; + 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_fn || !L.hc_ffn_fn || !L.hc_attn_base || !L.hc_ffn_base) { + fc.disabled = true; + 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] fused decode: HC fn mirror upload failed; using per-layer path\n"); + fc.disabled = true; + return 0; + } + + const int token_pos = kv_start; + const int n_raw = std::min(kv_start + 1, w.n_swa); + std::vector key; + key.reserve((size_t) w.n_layer + 2); + key.push_back(w.n_swa); + key.push_back(token_ids ? 1 : 0); + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + 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); + padded = ds4_padded_comp_rows(n_comp, (int) lc.comp_kv->ne[1]); + } + if (ds4_fused_stable_graph_enabled()) { + key.push_back((int64_t) padded); + } else { + const bool flush = ratio > 0 && (((token_pos + 1) % ratio) == 0); + key.push_back(((int64_t) padded << 1) | (flush ? 1 : 0)); + } + } + + // Pick the slot whose shape key matches; otherwise rebuild the LRU slot. + fc.counter++; + DeepSeek4FusedDecodeGraph * fg = nullptr; + for (auto & s : fc.slots) { + if (s.built() && s.shape_key == key) { + fg = &s; + break; + } + } + if (!fg) { + for (auto & s : fc.slots) { + if (!s.built()) { fg = &s; break; } + } + if (!fg) { + fg = &fc.slots[0]; + for (auto & s : fc.slots) { + if (s.last_use < fg->last_use) fg = &s; + } + } + const auto build_t0 = Ds4TimingClock::now(); + if (!ds4_build_fused_decode_graph(fc, *fg, backend, w, cache, + hc_weights, hc_out_weights, hash_tables, + kv_start, token_ids != nullptr, std::move(key))) { + std::fprintf(stderr, + "[deepseek4] fused decode graph build failed; using per-layer path\n"); + step_graph_free(fg->sg); + fg->reset_nodes(); + fc.disabled = true; + return 0; + } + if (telemetry) telemetry->full_graph_build_us += ds4_elapsed_us(build_t0, Ds4TimingClock::now()); + } + fg->last_use = fc.counter; + + // ── Fill inputs ───────────────────────────────────────────────── + const auto set_t0 = Ds4TimingClock::now(); + ggml_backend_tensor_set(fg->inp_embed, embed, 0, sizeof(float) * (size_t) w.n_embd); + + std::vector i32v((size_t) w.n_layer * 6, 0); + std::vector i64v((size_t) w.n_layer * 5, 0); + const int64_t raw_row = kv_start % w.n_swa; + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + i32v[(size_t) il * 6 + 0] = kv_start; + i32v[(size_t) il * 6 + 1] = -kv_start; + i64v[(size_t) il * 5 + 0] = raw_row; + if (ratio > 0) { + const int pos_mod = token_pos % ratio; + i32v[(size_t) il * 6 + 2] = pos_mod; + i32v[(size_t) il * 6 + 3] = token_pos + 1 - ratio; + i64v[(size_t) il * 5 + 1] = (ratio == 4) ? (int64_t) (ratio + pos_mod) : (int64_t) pos_mod; + i64v[(size_t) il * 5 + 2] = token_pos / ratio; + } + if (ratio == 4) { + const int pos_mod = token_pos % ratio; + i32v[(size_t) il * 6 + 4] = pos_mod; + i32v[(size_t) il * 6 + 5] = token_pos + 1 - ratio; + i64v[(size_t) il * 5 + 3] = ratio + pos_mod; + i64v[(size_t) il * 5 + 4] = token_pos / ratio; + } + } + ggml_backend_tensor_set(fg->i32_bundle, i32v.data(), 0, sizeof(int32_t) * i32v.size()); + ggml_backend_tensor_set(fg->i64_bundle, i64v.data(), 0, sizeof(int64_t) * i64v.size()); + if (fg->flush_rows) { + const bool flush4 = ((token_pos + 1) % 4) == 0; + const int64_t fr[4] = {flush4 ? 0 : 4, flush4 ? 1 : 5, flush4 ? 2 : 6, flush4 ? 3 : 7}; + ggml_backend_tensor_set(fg->flush_rows, fr, 0, sizeof(fr)); + } + + if (fg->mask_bundle) { + std::vector maskv((size_t) ggml_nelements(fg->mask_bundle), 0.0f); + size_t off = 0; + const int n_valid_raw = std::min(kv_start + 1, w.n_swa); + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + int n_comp = 0, padded = 0; + if (ratio > 0 && lc.comp_kv) { + n_comp = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos); + padded = ds4_padded_comp_rows(n_comp, (int) lc.comp_kv->ne[1]); + } + for (int j = n_valid_raw; j < w.n_swa; ++j) { + maskv[off + (size_t) j] = -1.0e30f; + } + for (int j = n_comp; j < padded; ++j) { + maskv[off + (size_t) w.n_swa + (size_t) j] = -1.0e30f; + } + off += (size_t) w.n_swa + (size_t) padded; + } + if (off != maskv.size()) { + std::fprintf(stderr, "[deepseek4] fused decode: mask layout mismatch (%zu vs %zu)\n", + off, maskv.size()); + return -1; + } + ggml_backend_tensor_set(fg->mask_bundle, maskv.data(), 0, sizeof(float) * maskv.size()); + } + + if (token_ids) { + for (int il = 0; il < w.n_layer; ++il) { + ggml_tensor * hids = fg->hash_ids[(size_t) il]; + if (!hids) continue; + const auto & ht = hash_tables[(size_t) il].ids; + const int n_used = w.n_expert_used; + hash_scratch.resize((size_t) n_used); + const int32_t tok = token_ids[0]; + std::memcpy(hash_scratch.data(), ht.data() + (size_t) tok * (size_t) n_used, + (size_t) n_used * sizeof(int32_t)); + ggml_backend_tensor_set(hids, hash_scratch.data(), 0, + sizeof(int32_t) * (size_t) n_used); + } + } + if (telemetry) telemetry->full_graph_set_us += ds4_elapsed_us(set_t0, Ds4TimingClock::now()); + + // ── Compute ───────────────────────────────────────────────────── + const auto compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, fg->sg.gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "[deepseek4] fused decode graph compute failed\n"); + return -1; + } + if (telemetry) telemetry->full_graph_compute_us += ds4_elapsed_us(compute_t0, Ds4TimingClock::now()); + + // ── Read logits ───────────────────────────────────────────────── + const auto read_t0 = Ds4TimingClock::now(); + out_logits.resize((size_t) w.n_vocab); + ggml_backend_tensor_get(fg->logits, out_logits.data(), 0, + sizeof(float) * (size_t) w.n_vocab); + if (telemetry) telemetry->full_graph_read_us += ds4_elapsed_us(read_t0, Ds4TimingClock::now()); + return 1; +} + + bool deepseek4_step_layer_range( ggml_backend_t backend, const DeepSeek4Weights & w, @@ -2816,7 +4195,8 @@ bool deepseek4_step_layer_range( int layer_end, std::vector * out_logits, const int32_t * token_ids, - DeepSeek4StepTelemetry * telemetry) { + DeepSeek4StepTelemetry * telemetry, + bool allow_decode_graph_reuse) { const auto step_t0 = Ds4TimingClock::now(); // NOTE: The old deepseek4_step() lacks HC implementation. @@ -2860,10 +4240,15 @@ bool deepseek4_step_layer_range( 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; - if (hc_loaded_n_layer != w.n_layer) { + static const ggml_context * hc_loaded_ctx = nullptr; + if (hc_loaded_n_layer != w.n_layer || hc_loaded_ctx != w.ctx) { + 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); - hash_routing_tables_range.resize((size_t)w.n_layer); + hash_routing_tables_range.assign((size_t)w.n_layer, {}); for (auto & alloc : cached_attn_allocs) { alloc.free(); } @@ -2889,21 +4274,54 @@ bool deepseek4_step_layer_range( cached_decode_ffn_graphs.assign((size_t)w.n_layer, {}); cached_decode_output_graph.free(); cached_dynamic_output_alloc.free(); + fused_decode_graph_cache.destroy(); + decode_shared_inputs.free(); for (int il = 0; il < w.n_layer; il++) { const DeepSeek4Layer & L = w.layers[(size_t)il]; load_hc_weights_cpu(hc_layer_weights_range[il].attn, L.hc_attn_fn, L.hc_attn_scale, L.hc_attn_base); load_hc_weights_cpu(hc_layer_weights_range[il].ffn, L.hc_ffn_fn, L.hc_ffn_scale, L.hc_ffn_base); + if (ds4_backend_is_gpu(backend) && !ds4_env_flag("DFLASH_DS4_HC_CPU")) { + ensure_hc_fn_device(hc_layer_weights_range[il].attn, L.hc_attn_fn); + ensure_hc_fn_device(hc_layer_weights_range[il].ffn, L.hc_ffn_fn); + } if (il < w.n_hash_layer && L.ffn_gate_tid2eid) { load_hash_routing_cpu(hash_routing_tables_range[(size_t)il], L.ffn_gate_tid2eid); } } 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; } // Per-layer execution with CPU-side HC static thread_local DeepSeek4LayerRangeScratch scratch; scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, w.n_expert_used); + + 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) && 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, + embed, kv_start, *out_logits, token_ids, telemetry); + if (rc < 0) return false; + if (rc > 0 && ds4_env_flag("DFLASH_DS4_FUSED_DEBUG")) { + // Debug: keep the fused logits, then fall through and run the + // per-layer reference for the same token; compare at the end. + fused_debug_logits = *out_logits; + } else if (rc > 0) { + const int np = kv_start + 1; + for (int il = layer_begin; il < layer_end; ++il) { + const uint32_t ratio = w.compress_ratios[il]; + if (ratio <= 0 || (np % (int) ratio) != 0) continue; + cache.layers[il].n_comp = std::max(cache.layers[il].n_comp, np / (int) ratio); + if (ratio == 4) cache.layers[il].n_index_comp = std::max(cache.layers[il].n_index_comp, np / (int) ratio); + } + cache.cur_pos = np; + if (telemetry) telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + return true; + } + } std::vector & cur = scratch.cur; std::vector & ffn_working = scratch.ffn_working; std::vector & hc_post = scratch.hc_post; @@ -2912,12 +4330,33 @@ bool deepseek4_step_layer_range( std::vector & attn_out_host = scratch.attn_out_host; std::vector & ffn_out_host = scratch.ffn_out_host; std::vector & hash_expert_ids_host = scratch.hash_expert_ids; - const bool reuse_decode_graphs = (n_tokens == 1); - const bool use_backend_decode_hc = reuse_decode_graphs && ds4_backend_is_gpu(backend); + const bool reuse_decode_graphs = allow_decode_graph_reuse && (n_tokens == 1); + Ds4DecodeSharedInputs * shared_inputs = nullptr; + if (reuse_decode_graphs && ds4_backend_is_gpu(backend) && + decode_shared_inputs.ensure(w, backend)) { + decode_shared_inputs.set_step(w, kv_start); + shared_inputs = &decode_shared_inputs; + } + + bool backend_decode_hc_supported = true; + for (int il = layer_begin; il < layer_end; ++il) { + const DeepSeek4Layer & L = w.layers[(size_t)il]; + const HcLayerWeightsCpu & hc_lw = hc_layer_weights_range[(size_t)il]; + if (!hc_fn_device_ptr(hc_lw.attn, L.hc_attn_fn) || + !hc_fn_device_ptr(hc_lw.ffn, L.hc_ffn_fn)) { + backend_decode_hc_supported = false; + break; + } + } + const bool use_backend_decode_hc = + reuse_decode_graphs && + ds4_backend_is_gpu(backend) && + backend_decode_hc_supported && + !ds4_env_flag("DFLASH_DS4_HC_CPU"); 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 ggml_tensor * hc_state_backend = nullptr; + ggml_tensor * hc_state_backend = nullptr; 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 || @@ -2946,10 +4385,6 @@ bool deepseek4_step_layer_range( // ── HC pre (attention) ────────────────────────────────────── const auto hc_pre_attn_t0 = Ds4TimingClock::now(); if (use_backend_decode_hc_direct) { - const auto hc_pre_attn_input_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_set(cached_decode_hc_post_graph.residual_hc, - hc_state.data(), 0, sizeof(float) * hc_state.size()); - if (telemetry) telemetry->hc_pre_input_us += ds4_elapsed_us(hc_pre_attn_input_t0, Ds4TimingClock::now()); auto & cached = cached_decode_attn_hc_pre_graphs[(size_t)il]; if (!cached.valid() || cached.owner_ctx != w.ctx || @@ -2970,8 +4405,9 @@ bool deepseek4_step_layer_range( backend, il, false, - cached_decode_hc_post_graph.residual_hc, + hc_state_backend, L.hc_attn_fn, + hc_fn_device_ptr(hc_lw.attn, L.hc_attn_fn), L.hc_attn_scale, L.hc_attn_base, hc_lw.attn.scale_data.data(), @@ -3023,7 +4459,7 @@ bool deepseek4_step_layer_range( // ── Build & run attention graph ───────────────────────────── { const int token_pos = kv_start + n_tokens - 1; - const bool reuse_decode_attn = n_tokens == 1; + const bool reuse_decode_attn = reuse_decode_graphs; ggml_tensor * attn_out = nullptr; ggml_cgraph * gf = nullptr; ggml_context * ctx = nullptr; @@ -3033,6 +4469,8 @@ bool deepseek4_step_layer_range( 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; + const bool attn_flush = ratio > 0 && (((token_pos + 1) % ratio) == 0); + const bool index_flush = ratio == 4 && (((token_pos + 1) % ratio) == 0); auto & per_layer = cached_decode_attn_graphs[(size_t)il]; auto it = std::find_if(per_layer.begin(), per_layer.end(), [&](const DeepSeek4CachedDecodeAttnGraph & candidate) { @@ -3042,7 +4480,9 @@ bool deepseek4_step_layer_range( candidate.layer_idx == il && candidate.n_raw == n_raw && candidate.n_comp_attn == n_comp_attn && - candidate.n_index_comp == n_index_comp; + candidate.n_index_comp == n_index_comp && + candidate.attn_flush == attn_flush && + candidate.index_flush == index_flush; }); if (it == per_layer.end()) { if (per_layer.size() >= 20) { @@ -3053,7 +4493,8 @@ bool deepseek4_step_layer_range( auto & candidate = per_layer.back(); const auto attn_build_t0 = Ds4TimingClock::now(); if (!build_cached_decode_attn_graph(candidate, backend, w, L, lc, il, kv_start, - n_raw, n_comp_attn, n_index_comp)) { + n_raw, n_comp_attn, n_index_comp, + shared_inputs)) { std::fprintf(stderr, "[deepseek4] cached attn graph alloc failed layer %d\n", il); return false; } @@ -3072,13 +4513,14 @@ bool deepseek4_step_layer_range( } else { ggml_backend_tensor_set(cached_attn->sg.inp_embed, cur.data(), 0, sizeof(float) * cur.size()); } + if (!cached_attn->uses_shared_inputs) { ggml_backend_tensor_set(cached_attn->inputs.rope_pos, &rope_pos, 0, sizeof(rope_pos)); ggml_backend_tensor_set(cached_attn->inputs.neg_pos, &neg_pos, 0, sizeof(neg_pos)); ggml_backend_tensor_set(cached_attn->inputs.raw_kv_rows, &raw_row, 0, sizeof(raw_row)); if (ratio > 0) { const int pos_mod = token_pos % ratio; const int32_t ape_row = pos_mod; - const int64_t state_row = ratio + pos_mod; + const int64_t state_row = (ratio == 4) ? (ratio + pos_mod) : pos_mod; const int64_t comp_row = token_pos / ratio; const int32_t comp_pos = token_pos + 1 - ratio; const bool flush_boundary = ((token_pos + 1) % ratio) == 0; @@ -3103,9 +4545,10 @@ bool deepseek4_step_layer_range( ggml_backend_tensor_set(cached_attn->inputs.index_comp_pos, &comp_pos, 0, sizeof(comp_pos)); } } + } } else { const auto attn_build_t0 = Ds4TimingClock::now(); - const size_t ctx_size = 48 * 1024 * 1024; + const size_t ctx_size = ds4_attn_step_meta_size(n_tokens); ggml_init_params params{}; params.mem_size = ctx_size; params.mem_buffer = nullptr; @@ -3118,7 +4561,7 @@ bool deepseek4_step_layer_range( std::vector i32_inputs; std::vector i32_array_inputs; std::vector i64_array_inputs; - const size_t graph_size = n_tokens > 1 ? 32768 : 2048; + const size_t graph_size = ds4_attn_step_graph_size(n_tokens); gf = ggml_new_graph_custom(ctx, graph_size, false); ggml_tensor * normed = build_rms_norm(ctx, inp, L.attn_norm, w.rms_eps); @@ -3162,7 +4605,7 @@ 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) { + 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); } @@ -3180,16 +4623,12 @@ bool deepseek4_step_layer_range( } else { const auto attn_read_t0 = Ds4TimingClock::now(); ggml_backend_tensor_get(attn_out, attn_out_host.data(), 0, sizeof(float) * attn_out_host.size()); - if (use_backend_decode_hc_direct) { - ggml_backend_tensor_get(attn_post_backend, hc_post.data(), 0, sizeof(float) * hc_post.size()); - ggml_backend_tensor_get(attn_comb_backend, hc_comb.data(), 0, sizeof(float) * hc_comb.size()); - } 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) { + if (!(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(), @@ -3207,10 +4646,6 @@ bool deepseek4_step_layer_range( // ── HC pre (FFN) ──────────────────────────────────────────── const auto hc_pre_ffn_t0 = Ds4TimingClock::now(); if (use_backend_decode_hc_direct) { - const auto hc_pre_ffn_input_t0 = Ds4TimingClock::now(); - ggml_backend_tensor_set(cached_decode_hc_post_graph.residual_hc, - hc_state.data(), 0, sizeof(float) * hc_state.size()); - if (telemetry) telemetry->hc_pre_input_us += ds4_elapsed_us(hc_pre_ffn_input_t0, Ds4TimingClock::now()); auto & cached = cached_decode_ffn_hc_pre_graphs[(size_t)il]; if (!cached.valid() || cached.owner_ctx != w.ctx || @@ -3231,8 +4666,9 @@ bool deepseek4_step_layer_range( backend, il, true, - cached_decode_hc_post_graph.residual_hc, + hc_state_backend, L.hc_ffn_fn, + hc_fn_device_ptr(hc_lw.ffn, L.hc_ffn_fn), L.hc_ffn_scale, L.hc_ffn_base, hc_lw.ffn.scale_data.data(), @@ -3335,7 +4771,7 @@ bool deepseek4_step_layer_range( return false; } - if (use_backend_decode_hc_graph) { + 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); } @@ -3352,15 +4788,11 @@ bool deepseek4_step_layer_range( } 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 (use_backend_decode_hc_direct) { - ggml_backend_tensor_get(ffn_post_backend, hc_post.data(), 0, sizeof(float) * hc_post.size()); - ggml_backend_tensor_get(ffn_comb_backend, hc_comb.data(), 0, sizeof(float) * hc_comb.size()); - } if (telemetry) telemetry->ffn_read_us += ds4_elapsed_us(ffn_read_t0, Ds4TimingClock::now()); } // ── HC post (FFN) ─────────────────────────────────────── - if (!use_backend_decode_hc_graph) { + if (!(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(), @@ -3376,7 +4808,7 @@ bool deepseek4_step_layer_range( } } - if (use_backend_decode_hc_graph && hc_state_backend) { + if ((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()); } @@ -3472,6 +4904,22 @@ bool deepseek4_step_layer_range( } cache.cur_pos = next_pos; + if (!fused_debug_logits.empty() && out_logits && !out_logits->empty()) { + const std::vector & ref = *out_logits; + size_t n = std::min(ref.size(), fused_debug_logits.size()); + float maxd = 0.0f; size_t maxi = 0; + size_t aref = 0, afus = 0; + for (size_t i = 1; i < n; ++i) { + if (ref[i] > ref[aref]) aref = i; + if (fused_debug_logits[i] > fused_debug_logits[afus]) afus = i; + const float d = std::fabs(ref[i] - fused_debug_logits[i]); + if (d > maxd) { maxd = d; maxi = i; } + } + std::fprintf(stderr, + "[ds4-fused-dbg] pos=%d argmax ref=%zu(%.4f) fused=%zu(%.4f) %s maxdiff=%.6f@%zu\n", + kv_start, aref, ref[aref], afus, fused_debug_logits[afus], + aref == afus ? "SAME" : "DIFF", maxd, maxi); + } if (telemetry) telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); return true; } @@ -3564,6 +5012,13 @@ bool create_deepseek4_cache(ggml_backend_t backend, return true; } +void free_deepseek4_cache(DeepSeek4Cache & c) { + 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(); + c.hc_state = nullptr; +} + namespace { ggml_tensor * clone_snapshot_tensor(ggml_context * ctx, @@ -3744,16 +5199,10 @@ bool deepseek4_snapshot_restore(const DeepSeek4Snapshot & snap, return true; } -void free_deepseek4_cache(DeepSeek4Cache & c) { - if (c.buf) { ggml_backend_buffer_free(c.buf); c.buf = nullptr; } - if (c.ctx) { ggml_free(c.ctx); c.ctx = nullptr; } - c.layers.clear(); - c.hc_state = nullptr; -} void free_deepseek4_snapshot(DeepSeek4Snapshot & s) { - if (s.buf) { ggml_backend_buffer_free(s.buf); s.buf = nullptr; } if (s.ctx) { ggml_free(s.ctx); s.ctx = nullptr; } + if (s.buf) { ggml_backend_buffer_free(s.buf); s.buf = nullptr; } s.layers.clear(); s.cur_pos = 0; s.hc_state_snap = nullptr; diff --git a/server/src/deepseek4/deepseek4_hc_cuda.cu b/server/src/deepseek4/deepseek4_hc_cuda.cu index 8f784d1f2..b5e592f50 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.cu +++ b/server/src/deepseek4/deepseek4_hc_cuda.cu @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -65,6 +66,15 @@ struct HcCudaScratch { std::mutex g_mu; HcCudaScratch g_scratch; +bool hc_direct_no_sync_enabled() { + static int enabled = -1; + if (enabled < 0) { + const char * value = std::getenv("DFLASH_DS4_HC_DIRECT_NO_SYNC"); + enabled = (value && value[0] && value[0] != '0') ? 1 : 0; + } + return enabled == 1; +} + void hc_log_cuda_error(const char * label, cudaError_t err) { if (err != cudaSuccess) { std::fprintf(stderr, "[deepseek4-hc-direct] %s: %s\n", label, cudaGetErrorString(err)); @@ -304,9 +314,11 @@ bool hc_pre_device_locked(const void * hc_state_device, return fail("finish kernel", err); } #if defined(DFLASH27B_BACKEND_HIP) || defined(GGML_USE_HIP) - err = cudaDeviceSynchronize(); - if (err != cudaSuccess) { - return fail("device sync", err); + if (!hc_direct_no_sync_enabled()) { + err = cudaDeviceSynchronize(); + if (err != cudaSuccess) { + return fail("device sync", err); + } } #endif @@ -551,4 +563,32 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, true); } +bool deepseek4_cuda_hc_upload_f16(const void * host_f16, + size_t bytes, + void ** device_out) { + if (!host_f16 || bytes == 0 || !device_out) { + return false; + } + void * ptr = nullptr; + cudaError_t err = cudaMalloc(&ptr, bytes); + if (err != cudaSuccess) { + hc_log_cuda_error("hc f16 mirror malloc", err); + return false; + } + err = cudaMemcpy(ptr, host_f16, bytes, cudaMemcpyHostToDevice); + if (err != cudaSuccess) { + hc_log_cuda_error("hc f16 mirror upload", err); + cudaFree(ptr); + return false; + } + *device_out = ptr; + return true; +} + +void deepseek4_cuda_hc_free(void * device_ptr) { + if (device_ptr) { + cudaFree(device_ptr); + } +} + } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_hc_cuda.h b/server/src/deepseek4/deepseek4_hc_cuda.h index d2268ac12..b4ac596fe 100644 --- a/server/src/deepseek4/deepseek4_hc_cuda.h +++ b/server/src/deepseek4/deepseek4_hc_cuda.h @@ -1,5 +1,6 @@ #pragma once +#include #include namespace dflash::common { @@ -47,4 +48,10 @@ bool deepseek4_cuda_hc_pre_device_params(const void * hc_state_device, void * post_device, void * comb_device); +bool deepseek4_cuda_hc_upload_f16(const void * host_f16, + size_t bytes, + void ** device_out); + +void deepseek4_cuda_hc_free(void * device_ptr); + } // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index 31c086394..a5d28d23c 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -64,6 +64,11 @@ struct DeepSeek4StepTelemetry { uint64_t output_us = 0; 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; int hot_selected = 0; int cold_selected = 0; }; @@ -330,7 +335,8 @@ bool deepseek4_step_layer_range( int layer_end, std::vector * out_logits, const int32_t * token_ids = nullptr, - DeepSeek4StepTelemetry * telemetry = nullptr); + DeepSeek4StepTelemetry * telemetry = nullptr, + bool allow_decode_graph_reuse = true); bool build_deepseek4_moe_hybrid_storage_from_file( const std::string & path, diff --git a/server/src/server/chat_template.cpp b/server/src/server/chat_template.cpp index ec8f30509..aab9ae022 100644 --- a/server/src/server/chat_template.cpp +++ b/server/src/server/chat_template.cpp @@ -327,6 +327,11 @@ std::string render_chat_template( system_content += msg.content; has_system = true; } + if (!has_system) { + system_content = + "You are a helpful assistant. Answer the user's question directly " + "and carefully. Do not change numbers or facts from the prompt."; + } result = "<｜begin▁of▁sentence｜>"; if (has_tools) { From 6c7ee19f1c6c0629c977d6d0a5ae055dbd6140e2 Mon Sep 17 00:00:00 2001 From: mrciffa Date: Thu, 9 Jul 2026 19:52:55 +0000 Subject: [PATCH 04/10] feat(deepseek4): add opt-in expert top-k knobs --- server/src/deepseek4/deepseek4_graph.cpp | 50 ++++++++++++++++++++++-- 1 file changed, 47 insertions(+), 3 deletions(-) diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 280a95946..3f81ba9aa 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -2036,6 +2036,34 @@ static bool eval_ds4_hybrid( return true; } +// Opt-in serving knobs for the routed-expert FFN, both default-off: +// DFLASH_DS4_TOPK= keep only the leading k of the top-k routed +// experts (weights renormalized over the kept +// set). Bandwidth lever for decode; measured +// +9% AR decode at k=4 on Strix (25.1 vs 23.0 +// tok/s), small quality cost — validate per +// deployment. +// DFLASH_DS4_NEURON_MEANMASK= zero routed-expert activations below +// m * mean(|activation|) per expert and token +// (m=1.0 keeps ~45%). Quality probe for +// intra-expert sparsity; adds a few ops and +// saves no bytes until a sparse kernel lands. +static int ds4_topk_override() { + static const int k = [] { + const char * e = std::getenv("DFLASH_DS4_TOPK"); + return e ? std::atoi(e) : 0; + }(); + return k; +} + +static float ds4_neuron_mask_mult() { + static const float m = [] { + const char * e = std::getenv("DFLASH_DS4_NEURON_MEANMASK"); + return e ? (float) std::atof(e) : 0.0f; + }(); + return m; +} + static Ds4MoeRouting build_moe_routing( ggml_context * ctx, ggml_tensor * cur, @@ -2054,10 +2082,15 @@ static Ds4MoeRouting build_moe_routing( selection = ggml_add(ctx, selection, L.ffn_exp_probs_b); } - out.selected = ggml_top_k(ctx, selection, w.n_expert_used); + int k_used = w.n_expert_used; + const int k_env = ds4_topk_override(); + if (k_env > 0 && k_env < k_used) { + k_used = k_env; + } + out.selected = 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 = ggml_reshape_2d(ctx, out.weights, w.n_expert_used, n_tokens); + 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); @@ -2077,7 +2110,7 @@ static ggml_tensor * build_moe_ffn( int n_tokens) { const int n_embd = w.n_embd; - const int n_used = w.n_expert_used; + int n_used = w.n_expert_used; const int n_ff_exp = w.n_ff_exp; const bool raw_mmid = ds4_ffn_raw_mmid_enabled(); ggml_tensor * shared_out = build_shared_ffn(ctx, cur, w, L); @@ -2087,6 +2120,7 @@ static ggml_tensor * build_moe_ffn( routed_out = ggml_scale(ctx, cur, 0.0f); } else { Ds4MoeRouting routing = build_moe_routing(ctx, cur, w, L, n_tokens); + n_used = (int) routing.selected->ne[0]; // reduced when DFLASH_DS4_TOPK is set ggml_tensor * cur_3d = ggml_reshape_3d(ctx, cur, n_embd, 1, n_tokens); ggml_tensor * gate_e = ggml_mul_mat_id(ctx, L.ffn_gate_exps, cur_3d, routing.selected); ggml_tensor * up_e = ggml_mul_mat_id(ctx, L.ffn_up_exps, cur_3d, routing.selected); @@ -2097,6 +2131,16 @@ static ggml_tensor * build_moe_ffn( } ggml_tensor * mid_e = build_clamped_swiglu(ctx, gate_e, up_e, w.swiglu_clamp_exp); + const float nmask_mult = ds4_neuron_mask_mult(); + if (nmask_mult > 0.0f) { + // Zero activations below mult * mean(|mid|) per (expert, token). + ggml_tensor * amid = ggml_abs(ctx, mid_e); + ggml_tensor * thr = ggml_scale(ctx, ggml_sum_rows(ctx, amid), + nmask_mult / (float) n_ff_exp); + ggml_tensor * mask = ggml_step(ctx, ggml_sub(ctx, amid, thr)); + mid_e = ggml_mul(ctx, mid_e, mask); + } + ggml_tensor * down_e = ggml_mul_mat_id(ctx, L.ffn_down_exps, mid_e, routing.selected); if (!raw_mmid) { down_e = ggml_reshape_3d(ctx, down_e, n_embd, n_used, n_tokens); From b6fcc26ecaad4efd898451d46d420c0e216fbbc6 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 7 Jul 2026 13:59:25 +0200 Subject: [PATCH 05/10] feat(dflash): bridge DeepSeek4 DSpark aux heads --- server/docs/DS4.md | 36 +++++ server/scripts/convert_dflash_to_gguf.py | 167 +++++++++++++++-------- 2 files changed, 143 insertions(+), 60 deletions(-) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index c3e32012f..5a3890934 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -110,6 +110,42 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. 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. +## DSpark Aux Draft Heads + +The current DSpark runtime lives in the shared DFlash speculative stack and is +used by the Laguna backend when the draft GGUF carries `dflash.dspark.*` +tensors. DeepSeek4 DSpark work is still a draft bridge: the DeepSeek4/MTP +artifact stores compatible aux heads under the `mtp.2.*` namespace, and the +converter can now map those names into the existing GGUF contract. + +Supported DeepSeek4/MTP input tensors: + +| DeepSeek4/MTP tensor | GGUF tensor | +|----------------------|-------------| +| `mtp.2.markov_head.markov_w1.weight` | `dflash.dspark.markov.w1` | +| `mtp.2.markov_head.markov_w2.weight` | `dflash.dspark.markov.w2` | +| `mtp.2.confidence_head.proj.weight` | `dflash.dspark.confidence.weight` | +| `mtp.2.confidence_head.proj.bias` | `dflash.dspark.confidence.bias` | + +If the MTP confidence projection is bias-less, the converter writes a zero +bias so the GGUF loader still sees the pair it expects. The Markov head alone +is enough for DSpark greedy-chain correction; confidence gating remains +optional. + +Example conversion with the DS4 MTP shard that contains the DSpark heads: + +```bash +python server/scripts/convert_dflash_to_gguf.py \ + /path/to/dflash-draft/model.safetensors \ + /path/to/dflash-draft.gguf \ + --aux-heads /path/to/hf-ds4-flash-dspark/model-00048-of-00048.safetensors +``` + +This does not by itself make the production DeepSeek4 layer-split backend use +DSpark. It makes the DeepSeek4 DSpark aux artifact consumable by the existing +`dflash.dspark.*` GGUF loader/runtime so the follow-up runtime PR can wire it +into the DS4 decode path with a clean tensor contract. + ## Example: CUDA + Halo Layer Split Automatic split (CUDA prefix chosen from free memory, optional manual override via `DFLASH_DS4_CUDA_LAYERS`): diff --git a/server/scripts/convert_dflash_to_gguf.py b/server/scripts/convert_dflash_to_gguf.py index 9c31f7531..6c1c46291 100644 --- a/server/scripts/convert_dflash_to_gguf.py +++ b/server/scripts/convert_dflash_to_gguf.py @@ -249,36 +249,83 @@ def bytes_to_np(raw: bytes, dtype: str, shape: list[int]) -> np.ndarray: DSPARK_TENSOR_MAP = { - "dspark_markov_head.markov_w1.weight": ("dflash.dspark.markov.w1", gguf.GGMLQuantizationType.F16), - "dspark_markov_head.markov_w2.weight": ("dflash.dspark.markov.w2", gguf.GGMLQuantizationType.F16), + ("dspark_markov_head.markov_w1.weight", + "mtp.2.markov_head.markov_w1.weight"): ("dflash.dspark.markov.w1", gguf.GGMLQuantizationType.F16), + ("dspark_markov_head.markov_w2.weight", + "mtp.2.markov_head.markov_w2.weight"): ("dflash.dspark.markov.w2", gguf.GGMLQuantizationType.F16), } DSPARK_CONFIDENCE_TENSOR_MAP = { - "dspark_confidence_head.weight": ("dflash.dspark.confidence.weight", gguf.GGMLQuantizationType.F16), - "dspark_confidence_head.bias": ("dflash.dspark.confidence.bias", gguf.GGMLQuantizationType.F32), + ("dspark_confidence_head.weight", + "mtp.2.confidence_head.proj.weight"): ("dflash.dspark.confidence.weight", gguf.GGMLQuantizationType.F16), + ("dspark_confidence_head.bias", + "mtp.2.confidence_head.proj.bias"): ("dflash.dspark.confidence.bias", gguf.GGMLQuantizationType.F32), } -def add_domino_aux_heads(writer, arch: str, aux_path: Path | None): - if aux_path is None: - return - if not aux_path.exists(): - print(f"[warn] Domino aux heads not found: {aux_path}") - return +def load_aux_state(aux_path: Path, label: str, wanted_names: set[str] | None = None): + if aux_path.suffix == ".safetensors": + header_size, header = load_safetensors_header(aux_path) + state = {} + for name, info in header.items(): + if name == "__metadata__": + continue + if wanted_names is not None and name not in wanted_names: + continue + raw = read_tensor_bytes(aux_path, header_size, info) + state[name] = bytes_to_np(raw, info["dtype"], info["shape"]) + return state try: import torch except ImportError as exc: - print(f"[error] --aux-heads requires torch: {exc}", file=sys.stderr) + print(f"[error] --aux-heads requires torch for {label} .pt files: {exc}", file=sys.stderr) sys.exit(1) - print(f"[info] reading Domino aux heads from {aux_path}") state = torch.load(aux_path, map_location="cpu") if isinstance(state, dict) and "state_dict" in state and isinstance(state["state_dict"], dict): state = state["state_dict"] if not isinstance(state, dict): - print(f"[error] Domino aux heads file is not a tensor dict: {aux_path}", file=sys.stderr) + print(f"[error] {label} aux heads file is not a tensor dict: {aux_path}", file=sys.stderr) sys.exit(1) + return state + + +def tensor_to_np(t, raw_dtype): + if hasattr(t, "detach"): + arr = t.detach().cpu().float().numpy() + else: + arr = np.asarray(t) + if arr.dtype.kind not in ("f", "i", "u"): + arr = arr.astype(np.float32) + if raw_dtype == gguf.GGMLQuantizationType.F16: + return arr.astype(" set[str]: + out = set() + for names in alias_map: + out.update(names) + return out + + +def add_domino_aux_heads(writer, arch: str, aux_path: Path | None): + if aux_path is None: + return + if not aux_path.exists(): + print(f"[warn] Domino aux heads not found: {aux_path}") + return + + print(f"[info] reading Domino aux heads from {aux_path}") + state = load_aux_state(aux_path, "Domino", set(DOMINO_TENSOR_MAP)) if not any(k in state for k in DOMINO_TENSOR_MAP): return @@ -300,14 +347,7 @@ def add_domino_aux_heads(writer, arch: str, aux_path: Path | None): writer.add_uint32(f"{arch}.dflash.domino.vocab_size", vocab) for st_name, (gguf_name, raw_dtype) in DOMINO_TENSOR_MAP.items(): - t = state[st_name] - if hasattr(t, "detach"): - t = t.detach().cpu() - arr = t.float().numpy() - if raw_dtype == gguf.GGMLQuantizationType.F16: - arr = arr.astype("{raw_dtype.name:4s} {tuple(arr.shape)}") @@ -318,26 +358,22 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): if not aux_path.exists(): return - try: - import torch - except ImportError as exc: - print(f"[error] --aux-heads requires torch: {exc}", file=sys.stderr) - sys.exit(1) - - state = torch.load(aux_path, map_location="cpu") - if isinstance(state, dict) and "state_dict" in state and isinstance(state["state_dict"], dict): - state = state["state_dict"] - if not isinstance(state, dict): - print(f"[error] DSpark aux heads file is not a tensor dict: {aux_path}", file=sys.stderr) - sys.exit(1) - - missing = [k for k in DSPARK_TENSOR_MAP if k not in state] + wanted_names = flat_alias_names(DSPARK_TENSOR_MAP) | flat_alias_names(DSPARK_CONFIDENCE_TENSOR_MAP) + state = load_aux_state(aux_path, "DSpark", wanted_names) + resolved = {} + missing = [] + for names, spec in DSPARK_TENSOR_MAP.items(): + found_name, tensor = find_state_tensor(state, names) + if tensor is None: + missing.append("/".join(names)) + continue + resolved[names] = (found_name, tensor, spec) if missing: return print(f"[info] reading DSpark aux heads from {aux_path}") - w1 = state["dspark_markov_head.markov_w1.weight"] - w2 = state["dspark_markov_head.markov_w2.weight"] + w1 = resolved[("dspark_markov_head.markov_w1.weight", "mtp.2.markov_head.markov_w1.weight")][1] + w2 = resolved[("dspark_markov_head.markov_w2.weight", "mtp.2.markov_head.markov_w2.weight")][1] vocab = int(w1.shape[0]) rank = int(w1.shape[1]) if tuple(w2.shape) != (vocab, rank): @@ -348,21 +384,39 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): writer.add_uint32(f"{arch}.dflash.dspark.markov_rank", rank) writer.add_uint32(f"{arch}.dflash.dspark.vocab_size", vocab) - for st_name, (gguf_name, raw_dtype) in DSPARK_TENSOR_MAP.items(): - t = state[st_name] - if hasattr(t, "detach"): - t = t.detach().cpu() - arr = t.float().numpy().astype("{raw_dtype.name:4s} {tuple(arr.shape)}") + print(f"[tensor] {gguf_name:50s} aux:{st_name} ->{raw_dtype.name:4s} {tuple(arr.shape)}") + + conf_resolved = {} + conf_missing = [] + for names, spec in DSPARK_CONFIDENCE_TENSOR_MAP.items(): + found_name, tensor = find_state_tensor(state, names) + if tensor is None: + conf_missing.append(names) + continue + conf_resolved[names] = (found_name, tensor, spec) + weight_names = ("dspark_confidence_head.weight", "mtp.2.confidence_head.proj.weight") + bias_names = ("dspark_confidence_head.bias", "mtp.2.confidence_head.proj.bias") + if weight_names not in conf_resolved: + if conf_missing: + print("[warn] incomplete DSpark confidence head; Markov head will still load") + return + if bias_names not in conf_resolved: + conf_resolved[bias_names] = ( + "", + np.zeros((1,), dtype=np.float32), + DSPARK_CONFIDENCE_TENSOR_MAP[bias_names], + ) + print("[warn] DSpark confidence head has no bias tensor; writing a zero bias") - conf_missing = [k for k in DSPARK_CONFIDENCE_TENSOR_MAP if k not in state] - if conf_missing: - print(f"[warn] incomplete DSpark confidence head; missing {conf_missing}; Markov head will still load") + if conf_missing and bias_names in conf_missing and len(conf_missing) > 1: + print("[warn] incomplete DSpark confidence head; Markov head will still load") return - conf_w = state["dspark_confidence_head.weight"] - conf_b = state["dspark_confidence_head.bias"] + conf_w = conf_resolved[weight_names][1] + conf_b = conf_resolved[bias_names][1] confidence_dim = int(conf_w.shape[1]) if int(conf_w.shape[0]) != 1 or tuple(conf_b.shape) != (1,): print( @@ -372,17 +426,10 @@ def add_dspark_aux_heads(writer, arch: str, aux_path: Path | None): sys.exit(1) writer.add_uint32(f"{arch}.dflash.dspark.confidence_dim", confidence_dim) writer.add_uint32(f"{arch}.dflash.dspark.confidence.enabled", 1) - for st_name, (gguf_name, raw_dtype) in DSPARK_CONFIDENCE_TENSOR_MAP.items(): - t = state[st_name] - if hasattr(t, "detach"): - t = t.detach().cpu() - arr = t.float().numpy() - if raw_dtype == gguf.GGMLQuantizationType.F16: - arr = arr.astype("{raw_dtype.name:4s} {tuple(arr.shape)}") + print(f"[tensor] {gguf_name:50s} aux:{st_name} ->{raw_dtype.name:4s} {tuple(arr.shape)}") # ────────────────────────────────────────────────────────────────────── @@ -394,9 +441,9 @@ def main(): ap.add_argument("safetensors", type=Path) ap.add_argument("out_gguf", type=Path) ap.add_argument("--aux-heads", type=Path, default=None, - help="optional Domino aux-head .pt file; defaults to dflash_aux_heads.pt next to the safetensors") + help="optional Domino/DSpark aux-head .pt or DS4 MTP .safetensors file; defaults to dflash_aux_heads.pt next to the safetensors") ap.add_argument("--no-aux-heads", action="store_true", - help="do not auto-embed Domino aux-head tensors") + help="do not auto-embed Domino/DSpark aux-head tensors") args = ap.parse_args() if not args.safetensors.exists(): From b989d3f49c9b6b5dae22aa8af7bcd0cbec076978 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 10 Jul 2026 18:36:45 +0200 Subject: [PATCH 06/10] feat(deepseek4): add DSpark speculative runtime --- server/CMakeLists.txt | 18 + server/src/deepseek4/deepseek4_backend.cpp | 69 +- server/src/deepseek4/deepseek4_backend.h | 6 + server/src/deepseek4/deepseek4_dspark.cpp | 300 ++++++++ server/src/deepseek4/deepseek4_dspark.h | 141 ++++ .../src/deepseek4/deepseek4_dspark_spec.cpp | 686 +++++++++++++++++ .../src/deepseek4/deepseek4_fused_verify.inc | 517 +++++++++++++ server/src/deepseek4/deepseek4_graph.cpp | 709 +++++++++++++++++- server/src/deepseek4/deepseek4_internal.h | 13 +- server/test/test_ds4_dspark_load.cpp | 153 ++++ 10 files changed, 2584 insertions(+), 28 deletions(-) create mode 100644 server/src/deepseek4/deepseek4_dspark.cpp create mode 100644 server/src/deepseek4/deepseek4_dspark.h create mode 100644 server/src/deepseek4/deepseek4_dspark_spec.cpp create mode 100644 server/src/deepseek4/deepseek4_fused_verify.inc create mode 100644 server/test/test_ds4_dspark_load.cpp diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 805d91c09..687938351 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -258,6 +258,8 @@ add_library(dflash_common STATIC src/deepseek4/deepseek4_daemon.cpp src/deepseek4/deepseek4_layer_split_adapter.cpp src/deepseek4/deepseek4_target_shard_ipc_daemon.cpp + src/deepseek4/deepseek4_dspark.cpp + src/deepseek4/deepseek4_dspark_spec.cpp src/flashprefill_q8.cpp src/kv_cache.cpp src/kv_quant.cpp @@ -801,6 +803,22 @@ if(DFLASH27B_TESTS) endif() add_test(NAME deepseek4_unit COMMAND test_deepseek4_unit) 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) + if(DFLASH27B_GPU_BACKEND STREQUAL "hip") + target_compile_definitions(test_ds4_dspark_load PRIVATE DFLASH27B_BACKEND_HIP=1 GGML_USE_HIP) + else() + target_compile_definitions(test_ds4_dspark_load PRIVATE DFLASH27B_BACKEND_CUDA=1 GGML_USE_CUDA) + endif() + target_link_libraries(test_ds4_dspark_load PRIVATE dflash_common ggml ggml-cpu ${DFLASH27B_GGML_BACKEND_TARGET} ggml-base) + if(DFLASH27B_GPU_BACKEND STREQUAL "cuda") + find_package(CUDAToolkit REQUIRED) + target_link_libraries(test_ds4_dspark_load PRIVATE CUDA::cudart) + else() + target_link_libraries(test_ds4_dspark_load PRIVATE hip::host) + endif() + endif() if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/smoke_load_draft.cpp") add_executable(smoke_load_draft test/smoke_load_draft.cpp) target_include_directories(smoke_load_draft PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS}) diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 5a3edbf9b..4afe7731c 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -356,6 +356,23 @@ bool DeepSeek4Backend::init() { std::fprintf(stderr, "[deepseek4] initialized: %d layers, ctx=%d, %d experts (%d used)%s\n", w_.n_layer, max_ctx, w_.n_expert, w_.n_expert_used, moe_hybrid_ ? " [hybrid]" : ""); + + if (env_flag_enabled("DFLASH_DS4_SPEC")) { + const char * dp = std::getenv("DFLASH_DS4_DRAFT"); + if (dp && *dp) { + spec_drafter_ = std::make_unique(); + if (load_deepseek4_dspark_drafter(dp, backend_, *spec_drafter_)) { + spec_enabled_ = true; + std::fprintf(stderr, "[deepseek4] DSpark spec-decode ENABLED (drafter=%s)\n", dp); + } else { + std::fprintf(stderr, "[deepseek4] DSpark drafter load FAILED: %s\n", + deepseek4_dspark_last_error()); + spec_drafter_.reset(); + } + } else { + std::fprintf(stderr, "[deepseek4] DFLASH_DS4_SPEC set but DFLASH_DS4_DRAFT gguf missing\n"); + } + } return true; } @@ -556,6 +573,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, ggml_backend_buffer_clear(cache_.buf, 0); } last_logits_.clear(); + if (spec_enabled_ && kv_offset == 0) spec_feat_window_.clear(); const bool timing = env_flag_enabled("DFLASH_DS4_TIMING"); const auto phase_t0 = Clock::now(); DeepSeek4StepTelemetry tel_acc; @@ -586,11 +604,27 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, routing_stats_.get()); } else { std::vector hc_state; + Ds4VerifyHooks spec_hooks; + std::vector spec_cap; + Ds4VerifyHooks * hp = nullptr; + if (spec_enabled_ && spec_drafter_) { + spec_hooks.capture_layer_ids = &spec_drafter_->capture_layer_ids; + spec_hooks.capture_out = &spec_cap; + hp = &spec_hooks; + } 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); + timing ? &step_tel : nullptr, + /*allow_decode_graph_reuse=*/true, hp); + if (hp && !spec_cap.empty()) { + const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; + for (int t = 0; t < n_tok; t++) + 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); + } } if (!ok) { std::fprintf(stderr, "[deepseek4] prefill step failed at pos=%d\n", pos); @@ -724,6 +758,35 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, // Decode auto t1 = Clock::now(); + if (spec_enabled_ && spec_drafter_ && req.n_gen > 0) { + if (last_logits_.empty()) { result.error = "spec: no prefill logits"; return result; } + int seed = 0; + { float mv = last_logits_[0]; + for (int i = 1; i < w_.n_vocab; i++) if (last_logits_[i] > mv) { mv = last_logits_[i]; seed = i; } } + std::vector gen; + gen.push_back(seed); + io.emit(seed); + float accept_rate = 0.0f; + if (!deepseek4_is_eos_tok(seed, w_) && req.n_gen > 1) { + const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; + const int win_len = feat_row > 0 ? (int) (spec_feat_window_.size() / feat_row) : 0; + std::vector spec_toks; + run_deepseek4_dspark_spec_decode(backend_, w_, cache_, *spec_drafter_, + committed, seed, req.n_gen - 1, + win_len > 0 ? spec_feat_window_.data() : nullptr, win_len, + spec_toks, &accept_rate); + for (int t : spec_toks) io.emit(t); + gen.insert(gen.end(), spec_toks.begin(), spec_toks.end()); + } + result.ok = true; + result.tokens = std::move(gen); + result.decode_s = elapsed_s(t1); + std::fprintf(stderr, "[deepseek4] DSpark decode: %zu tok in %.3fs (%.1f tok/s) accept_rate=%.2f\n", + result.tokens.size(), result.decode_s, + result.decode_s > 0 ? result.tokens.size() / result.decode_s : 0.0, accept_rate); + maybe_save_routing_stats(); + return result; + } std::vector gen_tokens; gen_tokens.reserve(req.n_gen); @@ -780,7 +843,9 @@ bool DeepSeek4Backend::handle_compress(const std::string & line, } void DeepSeek4Backend::free_drafter() { - // No drafter in AR-only mode + spec_drafter_.reset(); + spec_enabled_ = false; + spec_feat_window_.clear(); } void DeepSeek4Backend::maybe_save_routing_stats() { diff --git a/server/src/deepseek4/deepseek4_backend.h b/server/src/deepseek4/deepseek4_backend.h index 508637867..0641d9231 100644 --- a/server/src/deepseek4/deepseek4_backend.h +++ b/server/src/deepseek4/deepseek4_backend.h @@ -13,6 +13,7 @@ #include "../common/moe_hybrid_storage.h" #include "../common/moe_hybrid_stream.h" #include "deepseek4_internal.h" +#include "deepseek4_dspark.h" #include "ggml.h" #include "ggml-backend.h" @@ -76,6 +77,11 @@ class DeepSeek4Backend : public ModelBackend { DeepSeek4Snapshot snapshots_[PREFIX_SLOTS]; std::vector last_logits_; + // DSpark speculative decode (opt-in: DFLASH_DS4_SPEC=1 + DFLASH_DS4_DRAFT=). + bool spec_enabled_ = false; + std::unique_ptr spec_drafter_; + std::vector spec_feat_window_; + // Prefill prompt tokens in chunks, return absolute committed position. int do_prefill(const std::vector & tokens, const DaemonIO & io, int kv_offset = 0); diff --git a/server/src/deepseek4/deepseek4_dspark.cpp b/server/src/deepseek4/deepseek4_dspark.cpp new file mode 100644 index 000000000..8975244d1 --- /dev/null +++ b/server/src/deepseek4/deepseek4_dspark.cpp @@ -0,0 +1,300 @@ +// DeepSeek-V4-Flash "DSpark" drafter loader. See deepseek4_dspark.h. +// +// Self-contained (shares nothing with the target loader) so it cannot regress +// the target path. Loads a "deepseek4-dflash-draft" GGUF into a DSparkDrafter: +// the n_layer decoder blocks reuse DeepSeek4Weights leaf-name bindings, and the +// DSpark-specific tensors (dflash.fc / hidden_norm / dspark.*) bind separately. + +#include "deepseek4_dspark.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "gguf.h" + +#include +#include +#include +#include +#include + +#include +#include + +namespace dflash::common { + +namespace { + +std::string g_dspark_err; +void set_err(const std::string & m) { g_dspark_err = m; std::fprintf(stderr, "[ds4-dspark] %s\n", m.c_str()); } + +const char * ARCH = "deepseek4-dflash-draft"; + +uint32_t kv_u32(gguf_context * g, const std::string & key, uint32_t def) { + const int64_t id = gguf_find_key(g, key.c_str()); + if (id < 0) return def; + if (gguf_get_kv_type(g, id) != GGUF_TYPE_UINT32) return def; + return gguf_get_val_u32(g, id); +} +float kv_f32(gguf_context * g, const std::string & key, float def) { + const int64_t id = gguf_find_key(g, key.c_str()); + if (id < 0) return def; + if (gguf_get_kv_type(g, id) != GGUF_TYPE_FLOAT32) return def; + return gguf_get_val_f32(g, id); +} + +// Read an INT32/UINT32 array KV into ints. Returns false if the key is absent. +bool kv_i32_array(gguf_context * g, const std::string & key, std::vector & out) { + const int64_t id = gguf_find_key(g, key.c_str()); + if (id < 0) return false; + if (gguf_get_kv_type(g, id) != GGUF_TYPE_ARRAY) return false; + const enum gguf_type at = gguf_get_arr_type(g, id); + const size_t n = gguf_get_arr_n(g, id); + out.resize(n); + const void * raw = gguf_get_arr_data(g, id); + if (at == GGUF_TYPE_INT32) { + const int32_t * v = static_cast(raw); + for (size_t i = 0; i < n; i++) out[i] = (int)v[i]; + } else if (at == GGUF_TYPE_UINT32) { + const uint32_t * v = static_cast(raw); + for (size_t i = 0; i < n; i++) out[i] = (int)v[i]; + } else { + out.clear(); + return false; + } + return true; +} + +// Suffix after "blk.." — returns "" if `name` isn't a block tensor. +bool block_suffix(const char * name, int & il, std::string & suffix) { + if (std::strncmp(name, "blk.", 4) != 0) return false; + const char * p = name + 4; + if (*p < '0' || *p > '9') return false; + char * end = nullptr; + long v = std::strtol(p, &end, 10); + if (!end || *end != '.' || v < 0) return false; + il = (int)v; + suffix = std::string(end + 1); + return true; +} + +} // namespace + +const char * deepseek4_dspark_last_error() { return g_dspark_err.c_str(); } + +bool load_deepseek4_dspark_drafter(const std::string & path, + ggml_backend_t backend, + DSparkDrafter & out) { + ggml_context * meta = nullptr; + gguf_init_params gip{}; + gip.no_alloc = true; + gip.ctx = &meta; + gguf_context * g = gguf_init_from_file(path.c_str(), gip); + if (!g) { set_err("gguf_init failed: " + path); return false; } + + // ── Arch check ────────────────────────────────────────────────────── + { + const int64_t aid = gguf_find_key(g, "general.architecture"); + if (aid < 0) { set_err("missing general.architecture"); gguf_free(g); if (meta) ggml_free(meta); return false; } + const char * arch = gguf_get_val_str(g, aid); + if (std::string(arch) != ARCH) { + set_err(std::string("unexpected arch: ") + arch + " (expected " + ARCH + ")"); + gguf_free(g); if (meta) ggml_free(meta); return false; + } + } + + const std::string P = std::string(ARCH) + "."; + DeepSeek4Weights & w = out.core; + + // ── Core hparams (mirror the target loader's defaults) ────────────── + const uint32_t n_layer = kv_u32(g, P + "block_count", 3); + w.n_layer = (int)n_layer; + w.n_embd = (int)kv_u32(g, P + "embedding_length", 4096); + w.n_vocab = (int)kv_u32(g, P + "vocab_size", 129280); + w.n_head = (int)kv_u32(g, P + "attention.head_count", 64); + w.n_head_kv = (int)kv_u32(g, P + "attention.head_count_kv", 1); + w.head_dim = (int)kv_u32(g, P + "attention.key_length", 512); + w.n_rot = (int)kv_u32(g, P + "rope.dimension_count", 64); + w.n_lora_q = (int)kv_u32(g, P + "attention.q_lora_rank", 1024); + w.n_lora_o = (int)kv_u32(g, P + "attention.output_lora_rank", 1024); + w.n_out_group = (int)kv_u32(g, P + "attention.output_group_count", 8); + w.n_expert = (int)kv_u32(g, P + "expert_count", 256); + w.n_expert_used = (int)kv_u32(g, P + "expert_used_count", 6); + w.n_expert_shared = (int)kv_u32(g, P + "expert_shared_count", 1); + w.n_ff_exp = (int)kv_u32(g, P + "expert_feed_forward_length", 2048); + w.n_hash_layer = (int)kv_u32(g, P + "hash_layer_count", 3); + w.n_swa = (int)kv_u32(g, P + "attention.sliding_window", 128); + w.n_indexer_head = (int)kv_u32(g, P + "attention.indexer.head_count", 64); + w.n_indexer_head_dim = (int)kv_u32(g, P + "attention.indexer.key_length", 128); + w.n_indexer_top_k = (int)kv_u32(g, P + "attention.indexer.top_k", 512); + w.n_hc = (int)kv_u32(g, P + "hyper_connection.count", 4); + w.n_hc_sinkhorn_iter = (int)kv_u32(g, P + "hyper_connection.sinkhorn_iterations", 20); + w.expert_weight_scale = kv_f32(g, P + "expert_weights_scale", 1.5f); + w.rope_freq_base = kv_f32(g, P + "rope.freq_base", 10000.0f); + w.rope_scale_factor = kv_f32(g, P + "rope.scaling.factor", 16.0f); + w.rope_yarn_beta_fast = kv_f32(g, P + "rope.scaling.yarn_beta_fast", 32.0f); + w.rope_yarn_beta_slow = kv_f32(g, P + "rope.scaling.yarn_beta_slow", 1.0f); + w.compress_rope_freq_base = kv_f32(g, P + "attention.compress_rope_freq_base", 160000.0f); + w.rms_eps = kv_f32(g, P + "attention.layer_norm_rms_epsilon", 1e-6f); + w.hc_eps = kv_f32(g, P + "hyper_connection.epsilon", 1e-6f); + w.swiglu_clamp_exp = kv_f32(g, P + "swiglu_clamp_exp", 10.0f); + w.eos_id = -1; // drafter carries no tokenizer; EOS comes from the target + w.eos_chat_id = -1; + + // DSpark drafter layers have NO KV compression (DSparkAttention asserts + // compress_ratio==0). Force all-zero so no compressor/indexer tensors are + // expected — do NOT run compute_compress_ratios (would give [0,0,4,...]). + w.compress_ratios.assign(n_layer, 0u); + + // The DSpark blocks are NOT hash-routing layers: in the checkpoint their + // real layer ids are n_layers+stage = 43/44/45, all >= num_hash_layers (3), + // and the drafter GGUF carries no ffn_gate_tid2eid. Force 0 so build_moe_ffn + // always takes the score-routed (sqrt-softplus + top-k) path. + w.n_hash_layer = 0; + + // ── DFlash / DSpark metadata ──────────────────────────────────────── + out.n_target_layers = (int)kv_u32(g, P + "dflash.n_target_layers", 3); + out.block_size = (int)kv_u32(g, P + "dflash.block_size", 5); + out.mask_token_id = (int)kv_u32(g, P + "dflash.mask_token_id", 128799); + out.head_hc_enabled = kv_u32(g, P + "dflash.head_hc_enabled", 0) != 0; + if (!kv_i32_array(g, P + "dflash.capture_layer_ids", out.capture_layer_ids)) { + out.capture_layer_ids.clear(); + } + out.dspark_enabled = kv_u32(g, P + "dflash.dspark.enabled", 0) != 0; + out.markov_rank = (int)kv_u32(g, P + "dflash.dspark.markov_rank", 256); + out.vocab_size = (int)kv_u32(g, P + "dflash.dspark.vocab_size", w.n_vocab); + out.confidence_dim = (int)kv_u32(g, P + "dflash.dspark.confidence_dim", 0); + + w.layers.resize(n_layer); + w.backend = backend; + + // ── Allocate all tensors from the meta ctx into one backend buffer ── + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(meta, backend); + if (!buf) { set_err("ggml_backend_alloc_ctx_tensors failed"); gguf_free(g); ggml_free(meta); return false; } + ggml_backend_buffer_set_usage(buf, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + + // ── Stream tensor bytes from the file (pread per tensor) ──────────── + const int fd = ::open(path.c_str(), O_RDONLY); + if (fd < 0) { set_err("open failed: " + path); ggml_backend_buffer_free(buf); gguf_free(g); ggml_free(meta); return false; } + const size_t data_off = gguf_get_data_offset(g); + const int64_t n_tensors = gguf_get_n_tensors(g); + std::vector staging; + bool ok = true; + for (int64_t ti = 0; ti < n_tensors && ok; ti++) { + const char * tname = gguf_get_tensor_name(g, ti); + ggml_tensor * t = ggml_get_tensor(meta, tname); + if (!t) { set_err(std::string("meta tensor missing: ") + tname); ok = false; break; } + const size_t off = data_off + gguf_get_tensor_offset(g, ti); + const size_t size = gguf_get_tensor_size(g, ti); + staging.resize(size); + size_t done = 0; + while (done < size) { + const ssize_t r = ::pread(fd, staging.data() + done, size - done, (off_t)(off + done)); + if (r <= 0) { set_err(std::string("pread failed for ") + tname); ok = false; break; } + done += (size_t)r; + } + if (!ok) break; + ggml_backend_tensor_set(t, staging.data(), 0, size); + } + ::close(fd); + if (!ok) { ggml_backend_buffer_free(buf); gguf_free(g); ggml_free(meta); return false; } + + // ── Bind pointers by name ─────────────────────────────────────────── + for (int64_t ti = 0; ti < n_tensors; ti++) { + const char * name = gguf_get_tensor_name(g, ti); + ggml_tensor * t = ggml_get_tensor(meta, name); + + if (std::strcmp(name, "output_norm.weight") == 0) { w.out_norm = t; continue; } + if (std::strcmp(name, "output_hc_base.weight") == 0) { w.output_hc_base = t; continue; } + if (std::strcmp(name, "output_hc_fn.weight") == 0) { w.output_hc_fn = t; continue; } + if (std::strcmp(name, "output_hc_scale.weight") == 0) { w.output_hc_scale = t; continue; } + if (std::strcmp(name, "dflash.fc.weight") == 0) { out.main_proj = t; continue; } + if (std::strcmp(name, "dflash.hidden_norm.weight") == 0) { out.main_norm = t; continue; } + if (std::strcmp(name, "dflash.dspark.markov.w1") == 0) { out.markov_w1 = t; continue; } + if (std::strcmp(name, "dflash.dspark.markov.w2") == 0) { out.markov_w2 = t; continue; } + if (std::strcmp(name, "dflash.dspark.confidence.weight") == 0) { out.confidence_w = t; continue; } + if (std::strcmp(name, "dflash.dspark.confidence.bias") == 0) { out.confidence_b = t; continue; } + + int il = -1; std::string suffix; + if (!block_suffix(name, il, suffix) || il < 0 || il >= (int)n_layer) continue; + DeepSeek4Layer & L = w.layers[il]; + + if (suffix == "attn_norm.weight") L.attn_norm = t; + else if (suffix == "attn_q_a.weight") L.attn_q_a = t; + else if (suffix == "attn_q_a_norm.weight") L.attn_q_a_norm = t; + else if (suffix == "attn_q_b.weight") L.attn_q_b = t; + else if (suffix == "attn_kv.weight") L.attn_kv = t; + else if (suffix == "attn_kv_a_norm.weight") L.attn_kv_a_norm = t; + else if (suffix == "attn_sinks.weight") L.attn_sinks = t; + else if (suffix == "attn_output_a.weight") L.attn_output_a = t; + else if (suffix == "attn_output_b.weight") L.attn_output_b = t; + else if (suffix == "hc_attn_fn.weight") L.hc_attn_fn = t; + else if (suffix == "hc_attn_scale.weight") L.hc_attn_scale = t; + else if (suffix == "hc_attn_base.weight") L.hc_attn_base = t; + else if (suffix == "ffn_norm.weight") L.ffn_norm = t; + else if (suffix == "ffn_gate_inp.weight") L.ffn_gate_inp = t; + else if (suffix == "exp_probs_b.bias") L.ffn_exp_probs_b = t; + else if (suffix == "ffn_gate_tid2eid.weight") L.ffn_gate_tid2eid = t; + else if (suffix == "ffn_gate_exps.weight") L.ffn_gate_exps = t; + else if (suffix == "ffn_up_exps.weight") L.ffn_up_exps = t; + else if (suffix == "ffn_down_exps.weight") L.ffn_down_exps = t; + else if (suffix == "ffn_gate_shexp.weight") L.ffn_gate_shexp = t; + else if (suffix == "ffn_up_shexp.weight") L.ffn_up_shexp = t; + else if (suffix == "ffn_down_shexp.weight") L.ffn_down_shexp = t; + else if (suffix == "hc_ffn_fn.weight") L.hc_ffn_fn = t; + else if (suffix == "hc_ffn_scale.weight") L.hc_ffn_scale = t; + else if (suffix == "hc_ffn_base.weight") L.hc_ffn_base = t; + } + + w.ctx = meta; + w.buf = buf; + gguf_free(g); // meta_ctx now owned by w.ctx; do not free here + + std::fprintf(stderr, + "[ds4-dspark] loaded %s: n_layer=%d n_embd=%d vocab=%d block_size=%d " + "n_target_layers=%d markov_rank=%d confidence_dim=%d mask_tok=%d dspark=%d head_hc=%d\n", + path.c_str(), w.n_layer, w.n_embd, w.n_vocab, out.block_size, + out.n_target_layers, out.markov_rank, out.confidence_dim, out.mask_token_id, + (int)out.dspark_enabled, (int)out.head_hc_enabled); + return true; +} + +void free_deepseek4_dspark_drafter(DSparkDrafter & d) { + 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{}; +} + +// Validation dump used by the load smoke test. +void deepseek4_dspark_dump(const DSparkDrafter & d) { + const DeepSeek4Weights & w = d.core; + auto shp = [](ggml_tensor * t) -> std::string { + if (!t) return "NULL"; + char b[128]; + std::snprintf(b, sizeof(b), "[%lld,%lld,%lld,%lld] %s", + (long long)t->ne[0], (long long)t->ne[1], (long long)t->ne[2], (long long)t->ne[3], + ggml_type_name(t->type)); + return b; + }; + std::fprintf(stderr, "── DSpark drafter dump ──\n"); + std::fprintf(stderr, " main_proj %s\n", shp(d.main_proj).c_str()); + std::fprintf(stderr, " main_norm %s\n", shp(d.main_norm).c_str()); + std::fprintf(stderr, " out_norm %s\n", shp(w.out_norm).c_str()); + std::fprintf(stderr, " output_hc_fn %s\n", shp(w.output_hc_fn).c_str()); + std::fprintf(stderr, " markov_w1 %s\n", shp(d.markov_w1).c_str()); + std::fprintf(stderr, " markov_w2 %s\n", shp(d.markov_w2).c_str()); + std::fprintf(stderr, " conf_w %s\n", shp(d.confidence_w).c_str()); + std::fprintf(stderr, " conf_b %s\n", shp(d.confidence_b).c_str()); + for (int il = 0; il < w.n_layer; il++) { + const DeepSeek4Layer & L = w.layers[il]; + std::fprintf(stderr, " blk.%d: attn_norm=%s q_a=%s q_b=%s kv=%s o_a=%s o_b=%s sinks=%s\n", + il, shp(L.attn_norm).c_str(), shp(L.attn_q_a).c_str(), shp(L.attn_q_b).c_str(), + shp(L.attn_kv).c_str(), shp(L.attn_output_a).c_str(), shp(L.attn_output_b).c_str(), + shp(L.attn_sinks).c_str()); + std::fprintf(stderr, " ffn_gate_exps=%s up=%s down=%s shexp(g=%s) gate_inp=%s probs_b=%s hc_attn_fn=%s hc_ffn_fn=%s\n", + shp(L.ffn_gate_exps).c_str(), shp(L.ffn_up_exps).c_str(), shp(L.ffn_down_exps).c_str(), + shp(L.ffn_gate_shexp).c_str(), shp(L.ffn_gate_inp).c_str(), shp(L.ffn_exp_probs_b).c_str(), + shp(L.hc_attn_fn).c_str(), shp(L.hc_ffn_fn).c_str()); + } +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h new file mode 100644 index 000000000..9fe20da75 --- /dev/null +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -0,0 +1,141 @@ +// DeepSeek-V4-Flash "DSpark" speculative-decode drafter. +// +// The DSpark drafter is a small (n_layer≈3) DeepSeek-V4 block stack stored under +// the checkpoint's mtp.* namespace and converted to a GGUF with arch +// "deepseek4-dflash-draft" (see scripts/convert_ds4_dspark_draft_to_gguf.py). +// +// Reference forward: deepseek-ai/DeepSeek-V4-Flash-DSpark inference/model.py +// (DSparkBlock / DSparkAttention / DSparkMarkovHead / DSparkConfidenceHead, +// Transformer.forward_spec). Key facts encoded here: +// - The target captures h.mean(dim=2) (mean over the hc_mult HC copies) after +// each of dspark_target_layer_ids (=[40,41,42]) -> main_hidden [.., 3*n_embd]. +// - forward_embed (stage 0): main_x = main_norm(main_proj(main_hidden)); the +// noise block = embed([seed]+[MASK]*(block_size-1)), HC-expanded. +// - DSparkAttention: compress_ratio==0 (no compressor/indexer). Each layer +// projects the shared main_x -> main_kv via its own wkv, and the block's +// block_size query positions attend BIDIRECTIONALLY over +// [sliding-window main-context KV] ++ [block KV]. +// - forward_head (last stage): hc_head collapse -> out_norm -> tied target +// lm_head -> per-position Markov correction + confidence gate. +// +// The drafter's token embedding and lm_head are TIED to the target, so the +// drafter GGUF carries neither token_embd nor output; embedding + projection +// go through the DeepSeek4DFlashTarget adapter. + +#pragma once + +#include "deepseek4_internal.h" + +#include "ggml.h" +#include "ggml-backend.h" + +#include +#include + +namespace dflash::common { + +// 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 +// live in the same ggml context / backend buffer as `core`. +struct DSparkDrafter { + DeepSeek4Weights core; + + // 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] + + // DSpark heads (last stage). + ggml_tensor * markov_w1 = nullptr; // dflash.dspark.markov.w1 [markov_rank, vocab] + ggml_tensor * markov_w2 = nullptr; // dflash.dspark.markov.w2 [markov_rank, vocab] + ggml_tensor * confidence_w = nullptr; // dflash.dspark.confidence.weight [confidence_dim, 1] + ggml_tensor * confidence_b = nullptr; // dflash.dspark.confidence.bias [1] + + int block_size = 5; + int n_target_layers = 3; + int markov_rank = 256; + int vocab_size = 129280; + int confidence_dim = 0; + int mask_token_id = 128799; + bool dspark_enabled = false; + bool head_hc_enabled = false; + std::vector capture_layer_ids; // [40,41,42] +}; + +// Load a "deepseek4-dflash-draft" GGUF into `out`. Returns false on error; +// deepseek4_dspark_last_error() has the message. +bool load_deepseek4_dspark_drafter(const std::string & path, + ggml_backend_t backend, + DSparkDrafter & out); + +void free_deepseek4_dspark_drafter(DSparkDrafter & d); + +const char * deepseek4_dspark_last_error(); + +// One drafter forward. Produces block_size normed hidden states (the input to +// the tied lm_head + Markov head), conditioned on a window of captured target +// features. All host-side f32 for a simple v1 (GPU feature-ring plumbing can +// come later). +// +// seed_tok : the committed anchor token (block position 0's embedding) +// noise_embed : [n_embd * block_size] embeds of [seed]+[MASK]*(block_size-1) +// ctx_features : [n_target_layers*n_embd * ctx_len] captured features, +// ordered oldest..newest, absolute positions +// [committed-ctx_len .. committed-1] +// ctx_len : number of context feature columns (<= n_swa) +// committed : absolute position of the seed (block position 0) +// out_hidden : filled with [n_embd * block_size] = out_norm(hc_head(block)) +// +// Defined in deepseek4_graph.cpp (needs the static DS4 sub-builders). +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); + +// 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 +// forward and returns: +// argmax_out : per-position argmax token id (size n_tokens) +// logits_out : if non-null, full [n_tokens * n_vocab] f32 logits +// capture_out : [n_target_layers*n_embd * n_tokens] f32, per-position mean +// over the hc_mult HC copies after each capture layer +// (concatenated in capture_layer_ids order) — the drafter's +// main_hidden feed. +// Advances/updates the target cache exactly like a decode of these tokens. +bool deepseek4_dspark_verify_forward(ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const std::vector & capture_layer_ids, + const float * embed, + const int32_t * token_ids, + int n_tokens, + int kv_start, + std::vector & argmax_out, + std::vector * logits_out, + std::vector & capture_out, + DeepSeek4StepTelemetry * telemetry = nullptr, + bool allow_graph_reuse = true); + +// Run DSpark speculative decode: draft block_size candidates with `drafter`, +// verify against the DS4 target in one batched forward, accept the matching +// prefix, and loop. Returns generated tokens via `io.emit`. Mirrors the laguna +// DSpark loop. accept_rate_out (optional) gets mean accepted / block. +struct GenerateRequest; // fwd (from common/…); the loop only needs n_gen + committed +bool run_deepseek4_dspark_spec_decode( + ggml_backend_t backend, + const DeepSeek4Weights & target_w, + DeepSeek4Cache & target_cache, + const DSparkDrafter & drafter, + int committed, + int last_tok, + int n_gen, + const float * prompt_feature_window, // [n_target_layers*n_embd * win_len] captured during prefill + int win_len, + std::vector & out_tokens, + float * accept_rate_out); + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp new file mode 100644 index 000000000..f3cec9b3d --- /dev/null +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -0,0 +1,686 @@ +// DeepSeek-V4-Flash DSpark speculative decode: DFlashTarget adapter + spec loop. +// +// The drafter (DSparkDrafter, deepseek4_dspark.{h,cpp}) proposes block_size +// candidates conditioned on captured target features; the DS4 target verifies +// them in one batched forward. Feature capture + verify live in +// deepseek4_dspark_verify_forward (deepseek4_graph.cpp). The DSpark Markov head +// (common/dspark_head.cpp) is target-agnostic and reused verbatim. +// +// Fast spec loop (default): ONE batched verify per step, verify width capped +// at DS4_SPEC_Q=4 tokens (seed + 3 candidates). With q <= ratio(4) the verify +// crosses at most one compression boundary and never aliases rolling-state +// rows, so rejection rollback needs no KV snapshot at all: +// - raw ring rows are position-addressed (pos % n_swa) -> idempotent, +// - comp rows are index-addressed (pos / ratio) -> idempotent, +// - n_comp / n_index_comp are pure functions of commit position, +// - the ONLY non-idempotent state is the ratio-4 compressor prev-half +// (4 rows/state, flushed cur->prev at chunk boundaries), a few KB/layer, +// saved host-side before the verify and restored only when the flush +// happened at-or-past the rollback point. +// The legacy full-snapshot + double-verify path is kept behind +// DFLASH_DS4_FULL_SNAP=1 for A/B validation. + +#include "deepseek4_dspark.h" +#include "deepseek4_internal.h" +#include "internal.h" +#include "common/dspark_head.h" + +#include "ggml.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +// ── DFlashTarget adapter over the DS4 target ──────────────────────────── +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) + : w_(w), cache_(cache), backend_(backend), snap_backend_(snap_backend), + capture_ids_(std::move(capture_ids)), mask_tok_(mask_tok) {} + + bool verify_batch(const std::vector & tokens, int base_pos, int & last_tok, + std::vector * all_argmax = nullptr, + bool capture_ssm_intermediates = false) override { + (void) capture_ssm_intermediates; + const int n = (int) tokens.size(); + embed_buf_.resize((size_t) n * w_.n_embd); + if (!w_.embedder.embed(tokens.data(), n, embed_buf_.data())) { + std::fprintf(stderr, "[ds4-verify] embed FAILED n=%d tok0=%d tok1=%d vocab=%d\n", + n, n > 0 ? tokens[0] : -1, n > 1 ? tokens[1] : -1, w_.n_vocab); + return false; + } + // Sequential exact verify (measurement mode): q single-token forwards + // through the legacy AR decode path. Causal by construction, compressor + // fed every token. Slow; used to measure the drafter's TRUE accept + // rate and produce exact greedy output. Enable: DFLASH_DS4_SEQ_VERIFY=1 + // (pair with DFLASH_DS4_FULL_SNAP=1 so rollback/replay stay exact). + static const bool seq_verify = [] { + const char * v = std::getenv("DFLASH_DS4_SEQ_VERIFY"); + return v && *v && *v != '0'; + }(); + if (seq_verify) { + std::vector am_all; + std::vector feat_all; + std::vector logits_all; + am_all.reserve(n); + for (int t = 0; t < n; t++) { + std::vector am1; + std::vector feat1; + std::vector logits1; + if (!deepseek4_dspark_verify_forward(backend_, w_, cache_, capture_ids_, + embed_buf_.data() + (size_t) t * w_.n_embd, + tokens.data() + t, 1, base_pos + t, am1, + keep_logits_ ? &logits1 : nullptr, + feat1, telemetry_, + /*allow_graph_reuse=*/false)) { + return false; + } + if (am1.empty()) return false; + am_all.push_back(am1[0]); + feat_all.insert(feat_all.end(), feat1.begin(), feat1.end()); + if (keep_logits_) logits_all.insert(logits_all.end(), logits1.begin(), logits1.end()); + } + verify_features_ = std::move(feat_all); + if (keep_logits_) verify_logits_ = std::move(logits_all); + last_tok = am_all.back(); + verify_n_ = n; + if (all_argmax) *all_argmax = std::move(am_all); + return true; + } + std::vector am; + // n==1 must take the dynamic (non-reuse) path: the reused decode graph + // skips the capture/all-logits hooks (backend HC), which this needs. + if (!deepseek4_dspark_verify_forward(backend_, w_, cache_, capture_ids_, + embed_buf_.data(), tokens.data(), n, base_pos, am, + keep_logits_ ? &verify_logits_ : nullptr, + verify_features_, telemetry_, + /*allow_graph_reuse=*/n > 1)) { + return false; + } + if (am.empty()) return false; + last_tok = am.back(); + verify_n_ = n; + if (all_argmax) *all_argmax = std::move(am); + return true; + } + + bool read_verify_logits(int n_tokens, std::vector & out) override { + if (!keep_logits_ || verify_logits_.empty()) return false; + const size_t need = (size_t) n_tokens * w_.n_vocab; + if (verify_logits_.size() < need) return false; + out.assign(verify_logits_.begin(), verify_logits_.begin() + need); + return true; + } + + bool snapshot_kv() override { return deepseek4_snapshot_save(cache_, snap_backend_, snap_); } + bool restore_kv() override { return deepseek4_snapshot_restore(snap_, cache_); } + + bool is_eos(int token) const override { return deepseek4_is_eos_tok(token, w_); } + + bool embed_tokens(const int32_t * tokens, int n, float * out) const override { + return w_.embedder.embed(tokens, n, out); + } + + bool project_hidden_to_tokens(const float * hidden, int n_tokens, + std::vector & tokens_out) override { + std::vector logits; + if (!project_hidden_to_logits(hidden, n_tokens, logits)) return false; + tokens_out.resize(n_tokens); + for (int t = 0; t < n_tokens; t++) { + const float * row = logits.data() + (size_t) t * w_.n_vocab; + int best = 0; float bv = row[0]; + for (int i = 1; i < w_.n_vocab; i++) if (row[i] > bv) { bv = row[i]; best = i; } + tokens_out[t] = best; + } + return true; + } + + // The drafter hidden is already out_norm'd (drafter tail); project with the + // tied target lm_head only (mul_mat, no norm), matching the reference head. + bool project_hidden_to_logits(const float * hidden, int n_tokens, + std::vector & logits_out) override { + if (n_tokens <= 0) return false; + ggml_init_params ip{}; + ip.mem_size = 32u * 1024 * 1024; + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) return false; + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_tensor * h = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, w_.n_embd, n_tokens); + ggml_set_input(h); + ggml_tensor * logits = ggml_mul_mat(ctx, w_.output, h); // [n_vocab, n_tokens] + ggml_set_output(logits); + ggml_build_forward_expand(gf, logits); + 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; + } + ggml_backend_tensor_set(h, hidden, 0, sizeof(float) * (size_t) w_.n_embd * n_tokens); + const ggml_status st = ggml_backend_graph_compute(backend_, gf); + if (st != GGML_STATUS_SUCCESS) { ggml_gallocr_free(alloc); ggml_free(ctx); return false; } + logits_out.resize((size_t) n_tokens * w_.n_vocab); + ggml_backend_tensor_get(logits, logits_out.data(), 0, sizeof(float) * logits_out.size()); + ggml_gallocr_free(alloc); + ggml_free(ctx); + return true; + } + + ggml_tensor * lm_head_tensor() override { return w_.output; } + int hidden_size() const override { return w_.n_embd; } + int mask_token_id() const override { return mask_tok_; } + const std::vector & capture_layer_ids() const override { return capture_ids_; } + + void set_keep_logits(bool b) { keep_logits_ = b; } + void set_telemetry(DeepSeek4StepTelemetry * t) { telemetry_ = t; } + const std::vector & last_features() const { return verify_features_; } + int last_verify_n() const { return verify_n_; } + +private: + const DeepSeek4Weights & w_; + DeepSeek4Cache & cache_; + ggml_backend_t backend_; + ggml_backend_t snap_backend_; + std::vector capture_ids_; + int mask_tok_; + DeepSeek4Snapshot snap_{}; + DeepSeek4StepTelemetry * telemetry_ = nullptr; + bool keep_logits_ = false; + int verify_n_ = 0; + std::vector embed_buf_; + std::vector verify_logits_; + std::vector verify_features_; +}; + +namespace { + +// Build the DraftWeights shim the target-agnostic dspark head expects (only the +// DSpark head fields + n_embd are read). +DraftWeights make_dspark_shim(const DSparkDrafter & d) { + DraftWeights dw{}; + dw.n_embd = d.core.n_embd; + dw.dspark.enabled = d.dspark_enabled; + dw.dspark.markov_rank = d.markov_rank; + dw.dspark.vocab_size = d.vocab_size; + dw.dspark.confidence_dim = d.confidence_dim; + dw.dspark.markov_w1 = d.markov_w1; + dw.dspark.markov_w2 = d.markov_w2; + dw.dspark.confidence_w = d.confidence_w; + dw.dspark.confidence_b = d.confidence_b; + return dw; +} + +bool spec_env_flag(const char * name) { + const char * v = std::getenv(name); + return v && *v && *v != '0'; +} + +// ── Light rollback state ──────────────────────────────────────────────── +// Only the non-position-idempotent cache pieces: the prev-half rows of every +// ratio-4 rolling compressor state (attn + indexer) plus hc_state. Everything +// else (raw ring, comp rows, counters) is position/index-addressed and either +// overwritten idempotently or recomputable from the commit position. +struct SpecRollback { + int cur_pos = 0; + struct Layer { + std::vector attn_kv, attn_sc, idx_kv, idx_sc; + }; + std::vector layers; + std::vector hc_state; +}; + +// 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) { + 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); +} + +void restore_prev_half(ggml_tensor * t, const std::vector & buf) { + if (!t || buf.empty()) return; + ggml_backend_tensor_set(t, buf.data(), 0, buf.size()); +} + +void spec_rollback_save(const DeepSeek4Cache & cache, SpecRollback & rb) { + rb.cur_pos = cache.cur_pos; + rb.layers.resize(cache.layers.size()); + for (size_t il = 0; il < cache.layers.size(); ++il) { + const DeepSeek4LayerCache & lc = cache.layers[il]; + SpecRollback::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); + } + 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); + } +} + +// Truncate the cache to commit_pos. restore_prev is set when the verify +// crossed a ratio-4 boundary at-or-past commit_pos: that flush filled the +// 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(const SpecRollback & rb, const DeepSeek4Weights & w, + DeepSeek4Cache & cache, int commit_pos, bool restore_prev) { + cache.cur_pos = commit_pos; + for (size_t il = 0; il < cache.layers.size(); ++il) { + DeepSeek4LayerCache & lc = cache.layers[il]; + const uint32_t ratio = il < w.compress_ratios.size() ? w.compress_ratios[il] : 0; + if (ratio > 0) lc.n_comp = commit_pos / (int) ratio; + if (ratio == 4) lc.n_index_comp = commit_pos / 4; + if (restore_prev && il < rb.layers.size()) { + const SpecRollback::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 (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()); + } +} + +using SpecClock = std::chrono::steady_clock; + +double spec_ms_since(SpecClock::time_point t0) { + return std::chrono::duration_cast(SpecClock::now() - t0).count() / 1000.0; +} + +} // namespace + +// Batched target verify + capture: wraps the existing multi-token +// deepseek4_step_layer_range (dynamic attention + batched HC), which never +// touches the fused single-token 23 tok/s path, with the Ds4VerifyHooks that +// add per-layer mean-over-HC capture and full per-position logits. +bool deepseek4_dspark_verify_forward(ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const std::vector & capture_layer_ids, + const float * embed, + const int32_t * token_ids, + int n_tokens, + int kv_start, + std::vector & argmax_out, + std::vector * logits_out, + std::vector & capture_out, + DeepSeek4StepTelemetry * telemetry, + bool allow_graph_reuse) { + std::vector hc_state; + std::vector all_logits; + std::vector last_logits; + Ds4VerifyHooks hooks; + hooks.capture_layer_ids = &capture_layer_ids; + hooks.capture_out = &capture_out; + hooks.all_logits_out = &all_logits; + 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)) { + std::fprintf(stderr, "[ds4-verify] step_layer_range returned false (n_tokens=%d kv_start=%d)\n", + n_tokens, kv_start); + return false; + } + 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()); + return false; + } + argmax_out.resize(n_tokens); + for (int t = 0; t < n_tokens; t++) { + const float * row = all_logits.data() + (size_t) t * w.n_vocab; + int best = 0; float bv = row[0]; + for (int i = 1; i < w.n_vocab; i++) if (row[i] > bv) { bv = row[i]; best = i; } + argmax_out[t] = best; + } + if (logits_out) *logits_out = std::move(all_logits); + return true; +} + +bool run_deepseek4_dspark_spec_decode( + ggml_backend_t backend, + const DeepSeek4Weights & target_w, + DeepSeek4Cache & target_cache, + const DSparkDrafter & drafter, + int committed, + int last_tok, + int n_gen, + const float * prompt_feature_window, + int win_len, + std::vector & out_tokens, + float * accept_rate_out) { + const int n_embd = target_w.n_embd; + const int n_tgt = drafter.n_target_layers; + const int block = drafter.block_size; + const int n_swa = target_w.n_swa; + const int feat_row = n_tgt * n_embd; + + const bool debug = spec_env_flag("DFLASH_DS4_DSPARK_DEBUG"); + 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"); + // 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). + 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); + std::fclose(f); + } + double ewma_accept = 1.5; + + // Fast path caps the verify at the compression ratio (4): one boundary max, + // no rolling-state row aliasing -> snapshot-free rollback stays exact. + // The full-snapshot A/B path may use the whole block. + int q_cap = full_snap ? block + 1 : 4; + if (const char * qs = std::getenv("DFLASH_DS4_SPEC_Q")) { + const int v = std::atoi(qs); + if (v >= 2 && v <= block + 1) q_cap = full_snap ? v : std::min(v, 4); + } + if (std::FILE * qf = std::fopen("/tmp/ds4_spec_q", "r")) { + // Per-request override for perf experiments (no server restart needed). + // q=1 disables drafting: pure AR pushed through the batched-verify path + // (diagnoses batched-vs-sequential target divergence). + int v = 0; + if (std::fscanf(qf, "%d", &v) == 1 && v >= 1 && v <= block + 1) { + q_cap = full_snap ? v : std::min(v, 4); + } + std::fclose(qf); + } + + // Snapshot backend for the legacy full-snapshot rollback path. + ggml_backend_t snap_backend = ggml_backend_cpu_init(); + 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); + DraftWeights dw = make_dspark_shim(drafter); + SpecRollback rollback; + DeepSeek4StepTelemetry tel{}; + if (timing) target.set_telemetry(&tel); + + // Host feature window ring [feat_row, n_swa] of absolute positions + // [committed-N .. committed-1]. Seed from the prefill window. + std::vector feat_win((size_t) feat_row * n_swa, 0.0f); + int win_have = win_len > n_swa ? n_swa : win_len; + if (prompt_feature_window && win_have > 0) { + // copy the last win_have columns of the prefill window + const int src_off = (win_len - win_have); + std::memcpy(feat_win.data(), + prompt_feature_window + (size_t) src_off * feat_row, + sizeof(float) * (size_t) feat_row * win_have); + } + int feat_count = win_have; // number of valid feature columns ending at committed-1 + + auto push_feature = [&](const float * col) { + // Shift-append one feature column (keep last n_swa). + if (feat_count >= n_swa) { + std::memmove(feat_win.data(), feat_win.data() + feat_row, + sizeof(float) * (size_t) feat_row * (n_swa - 1)); + std::memcpy(feat_win.data() + (size_t) feat_row * (n_swa - 1), col, + sizeof(float) * feat_row); + } else { + std::memcpy(feat_win.data() + (size_t) feat_row * feat_count, col, + sizeof(float) * feat_row); + feat_count++; + } + }; + + int lt = last_tok; + int pos = committed; // absolute position of the seed (block slot 0) + int n_generated = 0; + long accept_sum = 0, steps = 0; + + std::vector noise_embed((size_t) n_embd * block); + std::vector noise_ids(block); + std::vector local_hidden, padded_hidden((size_t) n_embd * (block + 1), 0.0f); + std::vector draft_tok, tgt_am; + + // Cumulative phase timings (ms). + double tm_draft = 0, tm_head = 0, tm_save = 0, tm_verify = 0, tm_apply = 0, tm_feat = 0; + const SpecClock::time_point run_t0 = SpecClock::now(); + + while (n_generated < n_gen) { + const int ctx_len = feat_count < n_swa ? feat_count : n_swa; + + // 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())) break; + + // Drafter forward -> block normed hidden states. + if (!deepseek4_dspark_draft_forward(backend, drafter, noise_embed.data(), + ctx_len > 0 ? feat_win.data() : nullptr, + ctx_len, pos, local_hidden)) { + std::fprintf(stderr, "[ds4-spec] drafter forward failed\n"); + break; + } + } + tm_draft += spec_ms_since(t0); + + if (debug) { + size_t lh_nan = 0; double lh_ss = 0; + for (float v : local_hidden) { if (!std::isfinite(v)) lh_nan++; else lh_ss += (double) v * v; } + std::fprintf(stderr, "[ds4-spec] hidden nnan=%zu/%zu rms=%.4f ctx_len=%d\n", + lh_nan, local_hidden.size(), + lh_nan < local_hidden.size() ? std::sqrt(lh_ss / (double) local_hidden.size()) : 0.0, + ctx_len); + } + + // DSpark Markov chain over the first q_cap-1 candidates. Reference + // predicts token i+1 from block slot i, so prepend a dummy row 0 and + // let the (row-0-skipping) chain use slots 1..q-1. + t0 = SpecClock::now(); + draft_tok.clear(); + 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 + // sequential verify has no such limit. + 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 && !seq_verify_mode) { + const int w_cap = (int) ewma_accept + 2; + if (w_cap < q_step_cap) q_step_cap = w_cap; + } + if (q_step_cap >= 2) { + std::memcpy(padded_hidden.data() + n_embd, local_hidden.data(), + sizeof(float) * (size_t) n_embd * block); + ds_ok = dspark_markov_correct_greedy_chain_fused(dw, backend, target.lm_head_tensor(), + padded_hidden.data(), q_step_cap, lt, draft_tok); + if (!ds_ok) { + ds_ok = dspark_markov_correct_greedy_chain(dw, backend, target, + padded_hidden.data(), q_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)) 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]); + } + } else { + draft_tok.push_back(lt); // q=1: seed only, no speculation + } + if ((int) draft_tok.size() > q_step_cap) draft_tok.resize(q_step_cap); + const int q = (int) draft_tok.size(); // seed + candidates + tm_head += spec_ms_since(t0); + + if (debug) { + std::fprintf(stderr, "[ds4-spec] dbg ds_ok=%d q=%d lt=%d draft=[%d %d %d %d]\n", + (int) ds_ok, q, lt, + q > 0 ? draft_tok[0] : -1, q > 1 ? draft_tok[1] : -1, + q > 2 ? draft_tok[2] : -1, q > 3 ? draft_tok[3] : -1); + } + + // ── Rollback state save (cheap) or legacy full snapshot ── + t0 = SpecClock::now(); + if (full_snap) { + if (!target.snapshot_kv()) { std::fprintf(stderr, "[ds4-spec] snapshot failed\n"); break; } + } else { + spec_rollback_save(target_cache, rollback); + } + tm_save += spec_ms_since(t0); + + // First ratio-4 boundary position touched by this verify (p % 4 == 3). + const int first_boundary = pos + (3 - (pos & 3)); + const bool boundary_crossed = first_boundary <= pos + q - 1; + + // ── 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)) { + if (full_snap) { + target.restore_kv(); + } else { + spec_rollback_apply(rollback, target_w, target_cache, pos, boundary_crossed); + } + std::fprintf(stderr, "[ds4-spec] verify failed\n"); + 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. + int accept = 1; + for (int i = 0; i < q - 1; i++) { + if (draft_tok[i + 1] == tgt_am[i]) accept++; + else break; + } + const int matched = accept - 1; // accepted candidates + 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 (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. + std::fprintf(stderr, "[ds4-spec-cmp] step=%ld pos=%d draft=[%d %d %d] tgt=[%d %d %d %d] acc=%d\n", + steps, pos, + q > 1 ? draft_tok[1] : -1, q > 2 ? draft_tok[2] : -1, q > 3 ? draft_tok[3] : -1, + tgt_am[0], q > 1 ? tgt_am[1] : -1, q > 2 ? tgt_am[2] : -1, q > 3 ? tgt_am[3] : -1, + accept); + } + + // ── Rollback: truncate to the committed prefix ── + // The bonus token is DEFERRED: it becomes the next step's seed, whose + // KV is written then. + t0 = SpecClock::now(); + if (full_snap) { + // Legacy: full restore + replay the committed tokens through the + // target so ring/compressor/n_comp advance exactly. + std::vector kv_toks; + kv_toks.push_back(lt); + for (int i = 1; i < accept; i++) kv_toks.push_back(draft_tok[i]); + target.restore_kv(); + int replay_last = -1; + std::vector replay_am; + if (!target.verify_batch(kv_toks, pos, replay_last, &replay_am)) { + std::fprintf(stderr, "[ds4-spec] replay verify failed\n"); + break; + } + } else if (accept < q) { + // The prev-half flush is bad only if the boundary sits at-or-past + // the commit point (its chunk then contains rejected tokens). + const bool restore_prev = boundary_crossed && first_boundary >= commit_pos; + spec_rollback_apply(rollback, target_w, target_cache, commit_pos, restore_prev); + } + // accept == q on the fast path: cur_pos/n_comp already exact, keep. + tm_apply += spec_ms_since(t0); + + // Push the committed positions' features (slots 0..accept-1 = positions + // pos..pos+accept-1) into the drafter's context window. + t0 = SpecClock::now(); + const std::vector & feats = target.last_features(); + const int fN = full_snap ? target.last_verify_n() : accept; + for (int i = 0; i < fN; i++) push_feature(feats.data() + (size_t) i * feat_row); + tm_feat += spec_ms_since(t0); + + // Output tokens this step = accepted candidates + bonus. + bool hit_eos = false; + for (int i = 1; i <= accept; i++) { + const int t = (i < accept) ? draft_tok[i] : bonus; + out_tokens.push_back(t); + n_generated++; + if (target.is_eos(t)) { hit_eos = true; break; } + if (n_generated >= n_gen) break; + } + pos = commit_pos; // seed + accepted candidates now in KV + lt = bonus; // deferred bonus becomes next seed + accept_sum += matched; + ewma_accept = 0.7 * ewma_accept + 0.3 * (double) matched; + steps++; + 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", + steps, q, accept, + tm_draft / steps, tm_head / steps, tm_save / steps, + tm_verify / steps, tm_apply / steps, tm_feat / steps); + } + if (hit_eos) break; + } + + const double total_ms = spec_ms_since(run_t0); + if (accept_rate_out && steps > 0) { + const int denom = q_cap > 1 ? q_cap - 1 : 1; + *accept_rate_out = (float) accept_sum / (float) (steps * denom); + } + std::fprintf(stderr, "[ds4-spec] gen=%d steps=%ld mean_accept=%.2f/%d q_cap=%d full_snap=%d\n", + n_generated, steps, steps ? (double) accept_sum / steps : 0.0, q_cap - 1, q_cap, + (int) full_snap); + 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", + 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); + } + if (timing && steps > 0) { + const double s = 1000.0 * steps; // us -> ms per-step means + std::fprintf(stderr, + "[ds4-spec-t] verify tel/step: hc_pre_a=%.1f attn_b=%.1f attn_c=%.1f attn_r=%.1f " + "hc_post_a=%.1f hc_pre_f=%.1f route(b/c/r/s)=%.1f/%.1f/%.1f/%.1f " + "ffn(b/c/r)=%.1f/%.1f/%.1f eval=%.1f hot=%.1f cold=%.1f comb=%.1f part=%.1f " + "ghits=%llu gbuilds=%llu ms\n", + tel.hc_pre_attn_us / s, tel.attn_build_us / s, tel.attn_compute_us / s, + tel.attn_read_us / s, tel.hc_post_attn_us / s, tel.hc_pre_ffn_us / s, + tel.route_build_us / s, tel.route_compute_us / s, tel.route_read_us / s, + tel.route_select_us / s, + tel.ffn_build_us / s, tel.ffn_compute_us / s, tel.ffn_read_us / s, + tel.ffn_eval_us / s, tel.ffn_hot_us / s, tel.ffn_cold_us / s, + tel.ffn_combine_us / s, tel.ffn_partition_us / s, + (unsigned long long) tel.ffn_hot_graph_hits, + (unsigned long long) tel.ffn_hot_graph_builds); + } + ggml_backend_free(snap_backend); + return true; +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_fused_verify.inc b/server/src/deepseek4/deepseek4_fused_verify.inc new file mode 100644 index 000000000..824d609e0 --- /dev/null +++ b/server/src/deepseek4/deepseek4_fused_verify.inc @@ -0,0 +1,517 @@ +// ─── Fused batched VERIFY graph (n_tokens = q in [2,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 +// 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. + +// tensor_set that tolerates inputs the graph never consumed (gallocr leaves +// them unallocated, e.g. comp-emission bundles in non-flush graphs). +static void ds4_fv_set(ggml_tensor * t, const void * data, size_t nbytes) { + if (t && t->buffer) ggml_backend_tensor_set(t, data, 0, nbytes); +} +static bool ds4_fused_verify_enabled() { + static int enabled = -1; + if (enabled < 0) { + enabled = ds4_env_flag("DFLASH_DS4_FUSED_VERIFY") ? 1 : 0; + } + return enabled == 1; +} + +struct Ds4FusedVerifyCache { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + bool disabled = false; + uint64_t counter = 0; + std::array slots; + // verify-specific inputs/outputs per slot (parallel arrays, same index) + struct Extra { + ggml_tensor * pos_q = nullptr; // i32 [q] + ggml_tensor * neg_q = nullptr; // i32 [q] + ggml_tensor * rawrows = nullptr; // i64 [1,q] + ggml_tensor * ape4 = nullptr; // i32 [q] + ggml_tensor * ape128 = nullptr; // i32 [q] + 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] + int q = 0; + void reset() { *this = Extra{}; } + }; + std::array extra; + void destroy() { + for (auto & s : slots) s.destroy(); + for (auto & e : extra) e.reset(); + } +}; + +static Ds4FusedVerifyCache fused_verify_graph_cache; + +static bool ds4_build_fused_verify_graph( + DeepSeek4FusedDecodeCache & mc, // fn mirrors (shared with AR fused) + DeepSeek4FusedDecodeGraph & fg, + Ds4FusedVerifyCache::Extra & ex, + ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const std::vector & hc_weights, + const HcWeightsCpu & hc_out_weights, + const std::vector & hash_tables, + int kv_start, + int q, + bool have_token_ids, + const std::vector & capture_ids, + std::vector && shape_key) { + step_graph_free(fg.sg); + fg.reset_nodes(); + ex.reset(); + fg.hash_ids.assign((size_t) w.n_layer, nullptr); + 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; + 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(); + params.mem_buffer = fg.sg.meta_arena.data(); + params.no_alloc = true; + 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); + ggml_cgraph * gf = fg.sg.gf; + + fg.inp_embed = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, q); + ggml_set_input(fg.inp_embed); + ex.q = q; + ex.pos_q = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.pos_q); + ex.neg_q = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.neg_q); + ex.rawrows = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, q); ggml_set_input(ex.rawrows); + ex.ape4 = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.ape4); + ex.ape128 = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, q); ggml_set_input(ex.ape128); + ex.st4 = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, q); ggml_set_input(ex.st4); + ex.st128 = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, q); ggml_set_input(ex.st128); + // comp emission scalars per layer: i32 {ape_row(unused batched), comp_pos}, + // i64 {comp_row}; reuse the decode-style bundles (2 i32 + 1 i64 per layer). + fg.i32_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 2 * (int64_t) w.n_layer); + ggml_set_input(fg.i32_bundle); + fg.i64_bundle = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1 * (int64_t) w.n_layer); + ggml_set_input(fg.i64_bundle); + + // mask bundle: per layer [n_swa + padded + q] rows × q tokens + 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]); + } + 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); + int64_t mask_off = 0; + + // per-token HC streams + std::vector hc_cur(q); + for (int t = 0; t < q; ++t) { + ggml_tensor * col = ggml_view_2d(ctx, fg.inp_embed, n_embd, 1, + fg.inp_embed->nb[1], (size_t) t * fg.inp_embed->nb[1]); + hc_cur[(size_t) t] = ggml_repeat_4d(ctx, col, n_embd, n_hc, 1, 1); + } + + std::vector capture_pieces; // order [ci][t] + capture_pieces.reserve(capture_ids.size() * (size_t) q); + + for (int il = 0; il < w.n_layer; ++il) { + const DeepSeek4Layer & L = w.layers[(size_t) il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; + const int ratio = (int) w.compress_ratios[il]; + + // ── HC pre (attention), per token ── + std::vector hc_flat(q), split_attn(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); + 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]); + if (!working) return false; + ggml_tensor * w2 = ggml_reshape_2d(ctx, working, n_embd, 1); + attn_in = attn_in ? ggml_concat(ctx, attn_in, w2, 1) : w2; + } + + // ── Batched attention ── + DeepSeek4AttentionGraphInputs ain{}; + ain.rope_pos = ex.pos_q; + ain.neg_pos = ex.neg_q; + ain.raw_kv_rows = 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; + 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_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); + 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, + n_attn * sizeof(float), (size_t) mask_off * sizeof(float)); + ain.padded_comp = padded; + mask_off += n_attn * 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); + 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); + return false; + } + + // ── 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) { + 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); + 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]); + if (!fworking) return false; + ggml_tensor * f2 = ggml_reshape_2d(ctx, fworking, n_embd, 1); + ffn_in = ffn_in ? ggml_concat(ctx, ffn_in, f2, 1) : f2; + } + + // ── Batched FFN ── + ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, L.ffn_norm, w.rms_eps); + ggml_tensor * ffn_out = nullptr; + const bool hash_routed = il < w.n_hash_layer && L.ffn_gate_tid2eid && + have_token_ids && hash_tables[(size_t) il].loaded; + if (hash_routed) { + // ds4_build_hash_routed_ffn is single-token; run per token, concat. + ggml_tensor * hids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, w.n_expert_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, w.n_expert_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; + } + } else { + ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, q); + } + 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 (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] + 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); + capture_pieces.push_back(ggml_reshape_1d(ctx, ggml_cont(ctx, mean), n_embd)); + } + } + } + + // ── Output: per-token HC merge → batched out_norm + lm_head ── + ggml_tensor * final_all = nullptr; + for (int t = 0; t < q; ++t) { + ggml_tensor * hc_flat = ggml_reshape_1d(ctx, hc_cur[(size_t) t], (int64_t) n_embd * n_hc); + ggml_tensor * onorm = ggml_rms_norm(ctx, hc_flat, w.hc_eps); + ggml_tensor * omix = ggml_mul_mat(ctx, mc.fn_out_f16, onorm); + omix = ggml_reshape_1d(ctx, omix, ggml_nelements(omix)); + ggml_tensor * obase = ds4_fused_hc_base_f32(ctx, w.output_hc_base); + if (!obase || hc_out_weights.scale_data.empty()) return false; + ggml_tensor * fe = ggml_ds4_hc_out(ctx, omix, obase, hc_flat, n_hc, + hc_out_weights.scale_data[0]); + ggml_tensor * fe2 = ggml_reshape_2d(ctx, fe, n_embd, 1); + final_all = final_all ? ggml_concat(ctx, final_all, fe2, 1) : fe2; + } + ggml_tensor * out_normed = build_rms_norm(ctx, final_all, w.out_norm, w.rms_eps); + fg.logits = ggml_mul_mat(ctx, w.output, out_normed); // [n_vocab, q] + ggml_set_output(fg.logits); + ggml_build_forward_expand(gf, fg.logits); + + if (!capture_pieces.empty()) { + ggml_tensor * cap = capture_pieces[0]; + for (size_t i = 1; i < capture_pieces.size(); ++i) { + cap = ggml_concat(ctx, cap, capture_pieces[i], 0); + } + ex.capture = cap; + ggml_set_output(ex.capture); + ggml_build_forward_expand(gf, ex.capture); + } + + 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; +} + +static int ds4_try_fused_verify_step( + Ds4FusedVerifyCache & vc, + DeepSeek4FusedDecodeCache & mc, + 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, + Ds4VerifyHooks * hooks, + DeepSeek4StepTelemetry * telemetry) { + 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() || + !w.output_hc_fn || !w.output_hc_base) { vc.disabled = true; 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_fn || !L.hc_ffn_fn || !L.hc_attn_base || !L.hc_ffn_base) { + vc.disabled = true; + return 0; + } + } + if (vc.owner_ctx != w.ctx || vc.backend != backend) { + vc.destroy(); + vc.owner_ctx = w.ctx; + vc.backend = backend; + } + if (mc.owner_ctx != w.ctx || mc.backend != backend) { + mc.destroy(); + mc.owner_ctx = w.ctx; + mc.backend = backend; + } + if (!ds4_fused_ensure_fn_mirrors(mc, backend, w, hc_weights, hc_out_weights)) { + vc.disabled = true; + return 0; + } + + const int q = n_tokens; + const int token_pos = kv_start + q - 1; + std::vector key; + key.reserve((size_t) w.n_layer + 3); + key.push_back(w.n_swa); + key.push_back(q); + key.push_back(token_ids ? 1 : 0); + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + 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); + padded = ds4_padded_comp_rows(n_comp, (int) lc.comp_kv->ne[1]); + } + int b_idx = -1; + if (ratio > 0) { + for (int t = 0; t < q; ++t) { + if (((kv_start + t + 1) % ratio) == 0) { b_idx = t; break; } + } + } + key.push_back(((int64_t) padded << 4) | (int64_t) (b_idx + 1)); + } + + vc.counter++; + DeepSeek4FusedDecodeGraph * fg = nullptr; + Ds4FusedVerifyCache::Extra * ex = nullptr; + for (size_t i = 0; i < vc.slots.size(); ++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) { + 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]; + 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(); + ex->reset(); + vc.disabled = true; + return 0; + } + if (telemetry) telemetry->full_graph_build_us += ds4_elapsed_us(build_t0, Ds4TimingClock::now()); + } + fg->last_use = vc.counter; + + // ── Fill inputs ── + ds4_fv_set(fg->inp_embed, embed, sizeof(float) * (size_t) w.n_embd * q); + std::vector iv(q); + std::vector lv(q); + for (int i = 0; i < q; ++i) iv[(size_t) i] = kv_start + i; + ds4_fv_set(ex->pos_q, iv.data(), sizeof(int32_t) * q); + for (int i = 0; i < q; ++i) iv[(size_t) i] = -(kv_start + i); + ds4_fv_set(ex->neg_q, iv.data(), sizeof(int32_t) * q); + for (int i = 0; i < q; ++i) lv[(size_t) i] = (kv_start + i) % w.n_swa; + ds4_fv_set(ex->rawrows, lv.data(), sizeof(int64_t) * q); + for (int i = 0; i < q; ++i) iv[(size_t) i] = (kv_start + i) % 4; + ds4_fv_set(ex->ape4, iv.data(), sizeof(int32_t) * q); + for (int i = 0; i < q; ++i) lv[(size_t) i] = 4 + (kv_start + i) % 4; + ds4_fv_set(ex->st4, lv.data(), sizeof(int64_t) * q); + for (int i = 0; i < q; ++i) iv[(size_t) i] = (kv_start + i) % 128; + ds4_fv_set(ex->ape128, iv.data(), sizeof(int32_t) * q); + for (int i = 0; i < q; ++i) lv[(size_t) i] = (kv_start + i) % 128; + ds4_fv_set(ex->st128, lv.data(), sizeof(int64_t) * q); + + std::vector i32v((size_t) w.n_layer * 2, 0); + std::vector i64v((size_t) w.n_layer * 1, 0); + for (int il = 0; il < w.n_layer; ++il) { + const int ratio = (int) w.compress_ratios[il]; + if (ratio > 0) { + int pos_b = -1; // boundary position inside the batch (if any) + for (int t = 0; t < q; ++t) { + if (((kv_start + t + 1) % ratio) == 0) { pos_b = kv_start + t; break; } + } + i32v[(size_t) il * 2 + 0] = token_pos % ratio; + i32v[(size_t) il * 2 + 1] = (pos_b >= 0 ? pos_b : token_pos) + 1 - ratio; + i64v[(size_t) il] = (pos_b >= 0 ? pos_b : token_pos) / ratio; + } + } + ds4_fv_set(fg->i32_bundle, i32v.data(), sizeof(int32_t) * i32v.size()); + ds4_fv_set(fg->i64_bundle, i64v.data(), sizeof(int64_t) * i64v.size()); + + // causal mask values + { + 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]; + 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); + 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) { + float * col = maskv.data() + off + (size_t) i * n_attn; + const int pos_i = 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) + : (e - 1) - ((e - 1 - r) % w.n_swa); + if (p_r < 0 || p_r > pos_i) col[r] = -1.0e30f; + } + const int vis = (ratio > 0 && lc.comp_kv) + ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, pos_i) : 0; + 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; + } + } + off += n_attn * q; + } + ds4_fv_set(fg->mask_bundle, maskv.data(), sizeof(float) * maskv.size()); + } + + if (token_ids) { + for (int il = 0; il < w.n_layer; ++il) { + ggml_tensor * hids = fg->hash_ids[(size_t) il]; + if (!hids) continue; + const auto & ht = hash_tables[(size_t) il].ids; + const int n_used = w.n_expert_used; + hash_scratch.resize((size_t) n_used * q); + for (int i = 0; i < q; ++i) { + std::memcpy(hash_scratch.data() + (size_t) i * n_used, + ht.data() + (size_t) token_ids[i] * (size_t) n_used, + (size_t) n_used * sizeof(int32_t)); + } + ggml_backend_tensor_set(hids, hash_scratch.data(), 0, + sizeof(int32_t) * (size_t) n_used * q); + } + } + + // ── Compute + read back ── + const auto compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, fg->sg.gf) != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, "[ds4-fused-verify] compute failed\n"); + return -1; + } + if (telemetry) telemetry->attn_compute_us += ds4_elapsed_us(compute_t0, Ds4TimingClock::now()); + + const int ncap = (int) hooks->capture_layer_ids->size(); + 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 (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()); + hooks->capture_out->assign((size_t) ncap * w.n_embd * q, 0.0f); + for (int ci = 0; ci < ncap; ++ci) { + for (int t = 0; t < q; ++t) { + const float * src = flat.data() + ((size_t) ci * q + t) * w.n_embd; + float * dst = hooks->capture_out->data() + + (size_t) t * ncap * w.n_embd + (size_t) ci * w.n_embd; + std::memcpy(dst, src, sizeof(float) * (size_t) w.n_embd); + } + } + } + return 1; +} diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 3f81ba9aa..0eeea9a0d 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -220,6 +220,11 @@ struct DeepSeek4I64ArrayBinding { std::vector values; }; +struct DeepSeek4F32ArrayBinding { + ggml_tensor * tensor = nullptr; + std::vector values; +}; + 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, @@ -676,7 +681,10 @@ static void build_compressor_step( std::vector & i64_array_inputs, std::vector & i32_array_inputs, ggml_tensor ** comp_cache_source_out = nullptr, - ggml_tensor * flush_rows_inp = nullptr) { + ggml_tensor * flush_rows_inp = nullptr, + ggml_tensor * cur_all = nullptr, + int n_tokens_all = 1, + int kv_start_all = -1) { if (!gf || !cur_last || !ape || !kv_proj || !gate_proj || !norm_weight || !state.state_kv || !state.state_score || !comp_cache || ratio <= 0) { return; @@ -695,18 +703,89 @@ static void build_compressor_step( ggml_tensor * state_score_source = state.state_score; ggml_tensor * comp_cache_source = comp_cache; - ggml_tensor * ape_col = nullptr; - if (ape_row_inp) { - ape_col = ggml_get_rows(ctx, ape, ape_row_inp); - ape_col = ggml_reshape_2d(ctx, ape_col, comp_width, 1); - } else { - ape_col = ggml_view_2d( - ctx, ape, comp_width, 1, ape->nb[1], (size_t)pos_mod * ape->nb[1]); - ape_col = ggml_cast(ctx, ape_col, GGML_TYPE_F32); + // Causal-batch verify: every token's contribution lands in its + // position-addressed state row (the rows are distinct because the batch + // never crosses a ratio window; the boundary may only be the last token). + const bool batched_state = (cur_all != nullptr && n_tokens_all > 1 && + !state_rows_inp && kv_start_all >= 0); + if (batched_state) { + ggml_tensor * kv_all = ggml_mul_mat(ctx, kv_proj, cur_all); + ggml_tensor * sc_all = ggml_mul_mat(ctx, gate_proj, cur_all); + for (int ti = 0; ti < n_tokens_all; ti++) { + const int pm_ti = (kv_start_all + ti) % ratio; + const int row_ti = (ratio == 4) ? (ratio + pm_ti) : pm_ti; + ggml_tensor * kv_ti = ggml_view_2d(ctx, kv_all, comp_width, 1, kv_all->nb[1], + (size_t) ti * kv_all->nb[1]); + ggml_tensor * sc_ti = ggml_view_2d(ctx, sc_all, comp_width, 1, sc_all->nb[1], + (size_t) ti * sc_all->nb[1]); + ggml_tensor * ape_ti = ggml_view_2d(ctx, ape, comp_width, 1, ape->nb[1], + (size_t) pm_ti * ape->nb[1]); + sc_ti = ggml_add(ctx, sc_ti, ggml_cast(ctx, ape_ti, GGML_TYPE_F32)); + ggml_tensor * kv_slot_ti = ggml_view_2d(ctx, state.state_kv, comp_width, 1, + state.state_kv->nb[1], + (size_t) row_ti * state.state_kv->nb[1]); + ggml_tensor * sc_slot_ti = ggml_view_2d(ctx, state.state_score, comp_width, 1, + state.state_score->nb[1], + (size_t) row_ti * state.state_score->nb[1]); + ggml_build_forward_expand(gf, ggml_cpy(ctx, ggml_cast(ctx, kv_ti, state.state_kv->type), kv_slot_ti)); + ggml_build_forward_expand(gf, ggml_cpy(ctx, sc_ti, sc_slot_ti)); + } } - sc_cur = ggml_add(ctx, sc_cur, ape_col); - if (state_rows_inp) { + const bool batched_rows = (state_rows_inp && cur_all != nullptr && n_tokens_all > 1); + int batched_b = -1; // boundary index within the batch (batched_rows) + int batched_nB = 0; // tokens after the boundary + int batched_span_off = 0; + ggml_tensor * batched_kv_all = nullptr; + ggml_tensor * batched_sc_all = nullptr; + ggml_tensor * ape_col = nullptr; + if (!batched_rows) { + if (ape_row_inp) { + ape_col = ggml_get_rows(ctx, ape, ape_row_inp); + ape_col = ggml_reshape_2d(ctx, ape_col, comp_width, 1); + } else { + ape_col = ggml_view_2d( + ctx, ape, comp_width, 1, ape->nb[1], (size_t)pos_mod * ape->nb[1]); + ape_col = ggml_cast(ctx, ape_col, GGML_TYPE_F32); + } + sc_cur = ggml_add(ctx, sc_cur, ape_col); + } + + if (batched_state) { + // state rows already written above (batched) + } else if (batched_rows) { + // Fused verify: batched state writes with ONE boundary allowed at ANY + // batch index b (q <= ratio keeps every pos_mod distinct). Graph order: + // writes[0..b] -> pool(boundary, reads through span A) -> rotate + // cur->prev (ratio-4) -> writes[b+1..]. The pooling and rotation code + // below read state_*_source, which span A set. + ggml_tensor * kv_all = ggml_mul_mat(ctx, kv_proj, cur_all); + ggml_tensor * sc_all = ggml_mul_mat(ctx, gate_proj, cur_all); + ggml_tensor * ape_cols = ggml_get_rows(ctx, ape, ape_row_inp); // [comp_width, q] + sc_all = ggml_add(ctx, sc_all, ape_cols); + for (int ti = 0; ti < n_tokens_all; ++ti) { + if (((kv_start_all + ti + 1) % ratio) == 0) { batched_b = ti; break; } + } + const int nA = (batched_b >= 0) ? (batched_b + 1) : n_tokens_all; + batched_nB = n_tokens_all - nA; + auto write_span = [&](int off, int count, ggml_tensor ** kv_src, ggml_tensor ** sc_src) { + if (count <= 0) return; + ggml_tensor * kv_v = ggml_cont(ctx, ggml_view_2d(ctx, kv_all, comp_width, count, + kv_all->nb[1], (size_t) off * kv_all->nb[1])); + ggml_tensor * sc_v = ggml_cont(ctx, ggml_view_2d(ctx, sc_all, comp_width, count, + sc_all->nb[1], (size_t) off * sc_all->nb[1])); + ggml_tensor * rows_v = ggml_view_1d(ctx, state_rows_inp, count, + (size_t) off * state_rows_inp->nb[0]); + *kv_src = ggml_set_rows(ctx, state.state_kv, kv_v, rows_v); + *sc_src = ggml_set_rows(ctx, state.state_score, sc_v, rows_v); + ggml_build_forward_expand(gf, *kv_src); + ggml_build_forward_expand(gf, *sc_src); + }; + write_span(0, nA, &state_kv_source, &state_score_source); + batched_kv_all = kv_all; + batched_sc_all = sc_all; + batched_span_off = nA; + } else if (state_rows_inp) { state_kv_source = ggml_set_rows(ctx, state.state_kv, kv_cur, state_rows_inp); state_score_source = ggml_set_rows(ctx, state.state_score, sc_cur, state_rows_inp); ggml_build_forward_expand(gf, state_kv_source); @@ -722,7 +801,11 @@ static void build_compressor_step( ggml_build_forward_expand(gf, ggml_cpy(ctx, sc_cur, sc_slot)); } - if (!flush_rows_inp && ((token_pos + 1) % ratio) != 0) { + if (batched_rows && batched_b < 0) { + // no boundary inside the batch: state rows written, nothing to emit + 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. @@ -814,6 +897,39 @@ static void build_compressor_step( *comp_cache_source_out = comp_cache_source; } + if (batched_rows) { + if (ratio == 4) { + // completed window: rotate current half -> prev half, reading + // through the span-A writes so ordering is 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_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_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)); + } + return; + } if (ratio == 4 && flush_rows_inp) { // Stable-topology flush: copy the cur half onto rows given by the // input (prev half [0..3] at flush, cur half itself [4..7] = no-op @@ -869,7 +985,10 @@ static void build_indexer_compressor_step( ggml_tensor * comp_pos_inp, std::vector & i64_array_inputs, std::vector & i32_array_inputs, - ggml_tensor * flush_rows_inp = nullptr) { + ggml_tensor * flush_rows_inp = nullptr, + ggml_tensor * cur_all = nullptr, + int n_tokens_all = 1, + int kv_start_all = -1) { build_compressor_step(ctx, gf, cur_last, L.indexer_compressor_ape, L.indexer_compressor_kv, @@ -894,7 +1013,10 @@ static void build_indexer_compressor_step( i64_array_inputs, i32_array_inputs, nullptr, - flush_rows_inp); + flush_rows_inp, + cur_all, + n_tokens_all, + kv_start_all); } static int ds4_comp_rows_used(const ggml_tensor * comp_cache, int n_cached, int ratio, int token_pos) { @@ -1016,7 +1138,8 @@ static ggml_tensor * build_mla_attention( const DeepSeek4AttentionGraphInputs * cached_inputs, std::vector & i32_inputs, std::vector & i32_array_inputs, - std::vector & i64_array_inputs) { + std::vector & i64_array_inputs, + std::vector * f32_array_inputs = nullptr) { const int n_embd = w.n_embd; const int head_dim = w.head_dim; @@ -1076,6 +1199,51 @@ static ggml_tensor * build_mla_attention( rope_freq, rope_scale, rope_ext, rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + // ── Causal batched step (exact multi-token target semantics) ─── + // The target model is causal: token i must not attend to batch tokens + // j > i, must see the compressed-row count as of its own position, and — + // once the ring has wrapped — must still see the OLD contents of ring + // slots that later batch tokens overwrite. Default ON for multi-token + // steps on this path; DFLASH_DS4_NO_CAUSAL_VERIFY=1 restores the legacy + // (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"); + ggml_tensor * old_rows_scratch = nullptr; + int n_old_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; + // unwrapped/garbage rows are masked by the host-filled mask values. + for (int ti = 0; ti < n_tokens; ti++) { + ggml_tensor * slot = ggml_view_2d( + ctx, lc.raw_kv, head_dim, 1, lc.raw_kv->nb[1], + (size_t)((kv_start + ti) % w.n_swa) * lc.raw_kv->nb[1]); + ggml_tensor * saved = ggml_cont(ctx, slot); + ggml_build_forward_expand(gf, saved); + old_rows_scratch = old_rows_scratch + ? ggml_concat(ctx, old_rows_scratch, saved, 1) : saved; + n_old_rows++; + } + old_rows_scratch = ds4_cast_if_needed(ctx, old_rows_scratch, GGML_TYPE_F32); + } else if (causal_batch) { + // Copy the to-be-overwritten rows FIRST; same-stream build order runs + // these before the ring writes below. + for (int ti = 0; ti < n_tokens; ti++) { + if (kv_start + ti < w.n_swa) continue; // slot never held an older pos + ggml_tensor * slot = ggml_view_2d( + ctx, lc.raw_kv, head_dim, 1, lc.raw_kv->nb[1], + (size_t)((kv_start + ti) % w.n_swa) * lc.raw_kv->nb[1]); + ggml_tensor * saved = ggml_cont(ctx, slot); + ggml_build_forward_expand(gf, saved); + old_rows_scratch = old_rows_scratch + ? ggml_concat(ctx, old_rows_scratch, saved, 1) : saved; + n_old_rows++; + } + if (old_rows_scratch) { + old_rows_scratch = ds4_cast_if_needed(ctx, old_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; @@ -1127,7 +1295,10 @@ static ggml_tensor * build_mla_attention( i64_array_inputs, i32_array_inputs, &comp_kv_source, - cached_inputs ? cached_inputs->flush_rows : nullptr); + cached_inputs ? cached_inputs->flush_rows : nullptr, + (causal_batch || fused_causal) ? cur : nullptr, + n_tokens, + kv_start); } if (ratio == 4 && L.indexer_compressor_kv) { @@ -1138,7 +1309,10 @@ static ggml_tensor * build_mla_attention( cached_inputs ? cached_inputs->index_comp_pos : nullptr, i64_array_inputs, i32_array_inputs, - cached_inputs ? cached_inputs->flush_rows : nullptr); + 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); } @@ -1147,14 +1321,14 @@ static ggml_tensor * build_mla_attention( // raw_kv: [head_dim, n_swa] F16 persistent ring buffer (single KV head, shared) // comp_kv: [head_dim, comp_cap] F16 compressed rows. // n_raw = min(kv_start + n_tokens, n_swa) - const bool masked_kv = (n_tokens == 1) && cached_inputs && cached_inputs->attn_row_mask; + 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; // 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_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; + const int n_attn = n_raw + n_comp_attn + n_old_rows; const float kq_scale = 1.0f / sqrtf((float)head_dim); // Get valid KV rows. For single-token decode, include the current in-graph @@ -1207,6 +1381,9 @@ static ggml_tensor * build_mla_attention( comp = ds4_cast_if_needed(ctx, comp, GGML_TYPE_F32); kv_attn = ggml_concat(ctx, kv_attn, comp, 1); } + if (old_rows_scratch) { + kv_attn = ggml_concat(ctx, kv_attn, old_rows_scratch, 1); + } // kv_attn: [head_dim, n_attn] // Flatten q to [head_dim, n_head*n_tokens] for batched matmul @@ -1216,9 +1393,47 @@ static ggml_tensor * build_mla_attention( // → [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) { + 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; + } + 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; + } + // 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++; + } + } + 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; @@ -2438,6 +2653,17 @@ static void cpu_matvec_f16_pooled(float * out, const uint16_t * mat, const float pool.run(mat, x, out, rows, cols); } +// Token-level persistent-pool parallel-for: same splitting semantics as +// ds4_parallel_for_tokens but without per-call thread spawns (a multi-token +// step issues ~86 batched-HC calls, so spawn cost dominates at small n). +// Inner work must stay serial (serial_fn=true paths) - the pool is not +// reentrant. +static void ds4_pool_for_tokens(int n_tokens, const std::function & fn) { + if (n_tokens <= 1) { fn(0, n_tokens); return; } + static Ds4HcMatvecPool token_pool; + token_pool.run_custom(n_tokens, [&fn](int t) { fn(t, t + 1); }); +} + static void cpu_hc_sinkhorn(float * out, const float * mix, const float * scale, const float * base, int n_hc, int iters, float eps) { const float pre_scale = scale[0]; @@ -2672,7 +2898,7 @@ static void hc_pre_batch(std::vector & working, post.resize((size_t)n_tokens * (size_t)n_hc); comb.resize((size_t)n_tokens * (size_t)n_hc * (size_t)n_hc); - ds4_parallel_for_tokens(n_tokens, 8, [&](int t0, int t1) { + ds4_pool_for_tokens(n_tokens, [&](int t0, int t1) { std::vector flat(hc_dim); float mix[24]; for (int t = t0; t < t1; ++t) { @@ -2735,7 +2961,7 @@ static void hc_post_batch(std::vector & out_hc, }); return; } - ds4_parallel_for_tokens(n_tokens, 8, [&](int t0, int t1) { + ds4_pool_for_tokens(n_tokens, [&](int t0, int t1) { for (int t = t0; t < t1; ++t) { cpu_hc_post(out_hc.data() + (size_t)t * hc_dim, block_out + (size_t)t * n_embd, @@ -2757,7 +2983,7 @@ static void hc_output_batch(std::vector & final_embd, float hc_eps) { const size_t hc_dim = (size_t)n_embd * (size_t)n_hc; final_embd.resize((size_t)n_tokens * (size_t)n_embd); - ds4_parallel_for_tokens(n_tokens, 8, [&](int t0, int t1) { + ds4_pool_for_tokens(n_tokens, [&](int t0, int t1) { std::vector flat(hc_dim); std::vector pre((size_t)n_hc); std::vector hc_weights((size_t)n_hc); @@ -4030,6 +4256,8 @@ static bool ds4_build_fused_decode_graph( return true; } +#include "deepseek4_fused_verify.inc" + // Returns 1 on success (out_logits filled), 0 to fall back to the per-layer // path, -1 on a hard failure after cache state may have been touched. static int ds4_try_fused_decode_step( @@ -4240,7 +4468,8 @@ bool deepseek4_step_layer_range( std::vector * out_logits, const int32_t * token_ids, DeepSeek4StepTelemetry * telemetry, - bool allow_decode_graph_reuse) { + bool allow_decode_graph_reuse, + Ds4VerifyHooks * verify_hooks) { const auto step_t0 = Ds4TimingClock::now(); // NOTE: The old deepseek4_step() lacks HC implementation. @@ -4341,6 +4570,27 @@ bool deepseek4_step_layer_range( static thread_local DeepSeek4LayerRangeScratch scratch; scratch.ensure(w.ctx, n_tokens, n_embd, n_hc, w.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) && ds4_fused_verify_enabled()) { + const int vrc = ds4_try_fused_verify_step( + fused_verify_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); + if (vrc < 0) return false; + if (vrc > 0) { + const int np = kv_start + n_tokens; + for (int il = layer_begin; il < layer_end; ++il) { + const uint32_t vratio = w.compress_ratios[il]; + if (vratio <= 0) continue; + cache.layers[il].n_comp = std::max(cache.layers[il].n_comp, np / (int) vratio); + if (vratio == 4) cache.layers[il].n_index_comp = std::max(cache.layers[il].n_index_comp, np / (int) vratio); + } + cache.cur_pos = np; + if (telemetry) telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + return true; + } + } 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) && ds4_fused_decode_enabled()) { @@ -4605,6 +4855,7 @@ bool deepseek4_step_layer_range( std::vector i32_inputs; std::vector i32_array_inputs; std::vector i64_array_inputs; + std::vector f32_array_inputs; const size_t graph_size = ds4_attn_step_graph_size(n_tokens); gf = ggml_new_graph_custom(ctx, graph_size, false); @@ -4612,7 +4863,7 @@ bool deepseek4_step_layer_range( 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); + i64_array_inputs, &f32_array_inputs); ggml_set_output(attn_out); ggml_build_forward_expand(gf, attn_out); @@ -4640,6 +4891,8 @@ bool deepseek4_step_layer_range( 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 attn_compute_t0 = Ds4TimingClock::now(); @@ -4848,6 +5101,25 @@ bool deepseek4_step_layer_range( n_hc); std::memcpy(hc_state.data(), next_hc.data(), next_hc.size() * sizeof(float)); if (telemetry) telemetry->hc_post_ffn_us += ds4_elapsed_us(hc_post_ffn_t0, Ds4TimingClock::now()); + if (verify_hooks && verify_hooks->capture_layer_ids && verify_hooks->capture_out) { + const std::vector & _ids = *verify_hooks->capture_layer_ids; + for (size_t _ci = 0; _ci < _ids.size(); ++_ci) { + if (_ids[_ci] != il) continue; + const int _ncap = (int) _ids.size(); + std::vector & _cap = *verify_hooks->capture_out; + if ((int) _cap.size() != _ncap * n_embd * n_tokens) + _cap.assign((size_t) _ncap * n_embd * n_tokens, 0.0f); + for (int _t = 0; _t < n_tokens; ++_t) { + float * _dst = _cap.data() + (size_t) _t * _ncap * n_embd + (size_t) _ci * n_embd; + const float * _hs = hc_state.data() + (size_t) _t * hc_dim; + for (int _d = 0; _d < n_embd; ++_d) { + float _acc = 0.0f; + for (int _h = 0; _h < n_hc; ++_h) _acc += _hs[(size_t) _h * n_embd + _d]; + _dst[_d] = _acc / (float) n_hc; + } + } + } + } } } } @@ -4926,6 +5198,11 @@ bool deepseek4_step_layer_range( const size_t logits_offset = (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) { + verify_hooks->all_logits_out->resize((size_t) w.n_vocab * n_tokens); + ggml_backend_tensor_get(logits, verify_hooks->all_logits_out->data(), 0, + sizeof(float) * (size_t) w.n_vocab * n_tokens); + } ggml_free(ctx); } if (telemetry) telemetry->output_us += ds4_elapsed_us(output_t0, Ds4TimingClock::now()); @@ -5253,3 +5530,385 @@ void free_deepseek4_snapshot(DeepSeek4Snapshot & s) { } } // namespace dflash::common + +// ══════════════════════════════════════════════════════════════════════ +// DSpark drafter forward graph (appended to deepseek4_graph.cpp so it can +// reuse the file-static DS4 sub-builders: build_rms_norm, build_tail_rope_*, +// build_moe_ffn, build_shared_ffn). See deepseek4_dspark.h for the contract. +// +// Mirrors deepseek-ai/DeepSeek-V4-Flash-DSpark inference/model.py: +// forward_embed -> main_x = main_norm(main_proj(cat[h40,h41,h42])) +// per layer (DSparkBlock): HC-pre (per block position) -> attn_norm -> +// DSparkAttention (bidirectional over [ctx main-KV ++ block-KV]) -> +// HC-post ; HC-pre -> ffn_norm -> MoE -> HC-post +// tail: hc_head collapse -> out_norm (input to the tied lm_head + Markov) +// +// The ggml_ds4_hc_* ops are single-token, so HC-pre/HC-post run per block +// position; attention batches all block positions together (bidirectional). +// ══════════════════════════════════════════════════════════════════════ + +#include "deepseek4_dspark.h" + +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +// Fresh MLA attention for the drafter: no KV cache, no compression. The 5 +// block queries attend over an explicit [ctx main-context KV ++ block KV] +// tensor with full (bidirectional) visibility, plus the learned per-head sink. +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) + 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 + const int n_embd = w.n_embd; + const int head_dim = w.head_dim; + const int n_head = w.n_head; + const int n_rot = w.n_rot; + const int n_lora_o = w.n_lora_o; + const int n_out_group = w.n_out_group; + const int block = (int) cur->ne[1]; + const float eps = w.rms_eps; + // DSparkAttention has compress_ratio==0 -> base RoPE, YaRN disabled. + const float rope_freq = w.rope_freq_base; + const float rope_scale = 1.0f, rope_ext = 0.0f, rope_attn = 1.0f; + const int rope_orig = (int) w.rope_orig_ctx; + + // ── Q path (block queries) ────────────────────────────────────── + ggml_tensor * qr = build_rms_norm(ctx, ggml_mul_mat(ctx, L.attn_q_a, cur), L.attn_q_a_norm, eps); + ggml_tensor * q = ggml_mul_mat(ctx, L.attn_q_b, qr); // [n_head*head_dim, block] + q = ggml_reshape_3d(ctx, q, head_dim, n_head, block); + q = ggml_rms_norm(ctx, q, eps); // per-head unweighted + q = build_tail_rope_3d(ctx, q, pos_block, n_rot, head_dim, n_head, block, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_orig); + + // ── KV: block positions ───────────────────────────────────────── + ggml_tensor * kv_b = build_rms_norm(ctx, ggml_mul_mat(ctx, L.attn_kv, cur), L.attn_kv_a_norm, eps); + kv_b = build_tail_rope_2d(ctx, kv_b, pos_block, n_rot, head_dim, block, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_orig); + + // ── KV: context positions (from shared main_x) ────────────────── + 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); + kv_attn = ggml_concat(ctx, kv_c, kv_b, 1); // [head_dim, ctx_len+block] + n_attn = ctx_len + block; + } + + // ── Scores + sink softmax (full visibility, no causal mask) ───── + 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)); + ggml_tensor * probs = nullptr; + if (L.attn_sinks) { + ggml_tensor * sink = ggml_reshape_2d(ctx, L.attn_sinks, 1, n_head); + ggml_tensor * sink_shape = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n_head * block); + sink = ggml_repeat(ctx, sink, sink_shape); + ggml_tensor * sws = ggml_concat(ctx, scores, sink, 0); // [n_attn+1, n_head*block] + ggml_tensor * pws = ggml_soft_max(ctx, sws); + probs = ggml_view_2d(ctx, pws, n_attn, n_head * block, pws->nb[1], 0); + } else { + probs = ggml_soft_max(ctx, scores); + } + + // ── Context, inverse RoPE, grouped low-rank output ────────────── + ggml_tensor * kv_T = ggml_cont(ctx, ggml_transpose(ctx, kv_attn)); // [n_attn, head_dim] + ggml_tensor * context = ggml_mul_mat(ctx, kv_T, probs); // [head_dim, n_head*block] + context = ggml_reshape_3d(ctx, context, head_dim, n_head, block); + context = build_tail_rope_3d(ctx, context, neg_block, n_rot, head_dim, n_head, block, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_orig); + ggml_tensor * attn_out = ggml_reshape_2d(ctx, context, head_dim * n_head, block); + const int group_dim = head_dim * (n_head / n_out_group); + attn_out = ggml_reshape_3d(ctx, attn_out, group_dim, n_out_group, block); + attn_out = ggml_cont(ctx, ggml_permute(ctx, attn_out, 0, 2, 1, 3)); // [group_dim, block, n_out_group] + ggml_tensor * out_a_3d = ggml_reshape_3d(ctx, L.attn_output_a, group_dim, n_lora_o, n_out_group); + ggml_tensor * attn_low = ggml_mul_mat(ctx, out_a_3d, attn_out); // [n_lora_o, block, n_out_group] + attn_low = ggml_cont(ctx, ggml_permute(ctx, attn_low, 0, 2, 1, 3)); // [n_lora_o, n_out_group, block] + attn_low = ggml_reshape_2d(ctx, attn_low, n_lora_o * n_out_group, block); + return ggml_mul_mat(ctx, L.attn_output_b, attn_low); // [n_embd, block] +} + +// Read a small F32 GPU tensor (HC scale, [k]) into host floats. +static void ds4_read_f32(ggml_tensor * t, float * dst, int k) { + if (t) ggml_backend_tensor_get(t, dst, 0, sizeof(float) * (size_t) k); + else for (int i = 0; i < k; i++) dst[i] = 0.0f; +} + +} // namespace + +// ── Cached drafter graph ──────────────────────────────────────────────── +// The drafter forward runs every spec step with identical topology (ctx_len +// is constant once the feature window fills at n_swa). Rebuilding the +// multi-thousand-node graph, zero-initializing a fresh 256 MB arena and +// re-planning gallocr each call used to cost more than the 3-layer compute +// itself (~63 ms/step). Cache the built graph keyed by (ctx_len, block, +// drafter instance) and re-set only the inputs per call. +namespace { + +struct DsparkDraftCache { + int ctx_len = -1; + int block = -1; + const void * drafter = nullptr; + 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 * pos_block = nullptr; + ggml_tensor * neg_block = nullptr; + ggml_tensor * pos_ctx = nullptr; + ggml_tensor * out = nullptr; + std::vector> dbg_taps; + // HC scales are immutable weights: read from the backend once. + std::vector> s_attn, s_ffn; + float s_out = 0.0f; +}; + +thread_local DsparkDraftCache g_dspark_draft_cache; + +} // namespace + +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) { + const DeepSeek4Weights & w = d.core; + const int n_embd = w.n_embd; + const int n_hc = w.n_hc; + const int block = d.block_size; + const int fc_in = d.n_target_layers * n_embd; + 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; + + DsparkDraftCache & C = g_dspark_draft_cache; + const bool DS4_DBG = std::getenv("DFLASH_DS4_DSPARK_DEBUG") != nullptr; + + if (C.drafter != (const void *) &d) { + // 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}); + C.s_ffn.assign((size_t) w.n_layer, {0.0f, 0.0f, 0.0f}); + for (int il = 0; il < w.n_layer; il++) { + ds4_read_f32(w.layers[il].hc_attn_scale, C.s_attn[il].data(), 3); + ds4_read_f32(w.layers[il].hc_ffn_scale, C.s_ffn[il].data(), 3); + } + float so[1] = {0.0f}; + ds4_read_f32(w.output_hc_scale, so, 1); + C.s_out = so[0]; + } + + if (!C.ctx || C.ctx_len != ctx_len || C.block != block || C.drafter != (const void *) &d) { + // ── (Re)build the graph ───────────────────────────────────────── + if (C.ctx) { ggml_free(C.ctx); C.ctx = nullptr; } + C.gf = nullptr; + C.dbg_taps.clear(); + if (C.arena.empty()) C.arena.resize(256u * 1024 * 1024); + ggml_init_params ip{}; + ip.mem_size = C.arena.size(); + ip.mem_buffer = C.arena.data(); + ip.no_alloc = true; + ggml_context * ctx = ggml_init(ip); + if (!ctx) return false; + C.ctx = ctx; + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 32768, false); + C.gf = gf; + auto dbg_tap = [&](const std::string & nm, ggml_tensor * t) { + if (!DS4_DBG || !t) return; + ggml_set_output(t); + ggml_build_forward_expand(gf, t); + C.dbg_taps.push_back({nm, t}); + }; + + // 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.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); + + // main_x = main_norm(main_proj(ctx_features)). Shared across layers. + ggml_tensor * main_x = nullptr; + if (ctx_len > 0) { + // 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 + // main_norm (RMSNorm) normalizes main_proj's output and RMSNorm is + // scale-invariant, pre-normalizing the features to unit RMS gives a + // mathematically identical main_x while keeping the rocmfp4 activation + // in a safe range: main_norm(main_proj(f/rms(f))) == main_norm(main_proj(f)). + ggml_tensor * fc_in_normed = ggml_rms_norm(ctx, C.inp_ctx, w.rms_eps); + ggml_tensor * fc_out = ggml_mul_mat(ctx, d.main_proj, fc_in_normed); // [n_embd, ctx_len] + dbg_tap("fc_out", fc_out); + main_x = build_rms_norm(ctx, fc_out, d.main_norm, w.rms_eps); + dbg_tap("main_x", main_x); + } + + // HC state: [n_embd, n_hc, block], init = block embeds replicated over streams. + ggml_tensor * noise3 = ggml_reshape_3d(ctx, C.inp_noise, n_embd, 1, block); + ggml_tensor * hc_cur = ggml_repeat_4d(ctx, noise3, n_embd, n_hc, block, 1); + + auto hc_col = [&](ggml_tensor * hc, int p) -> ggml_tensor * { + // Contiguous [n_embd*n_hc] slab for block position p. + return ggml_view_1d(ctx, hc, (int64_t) n_embd * n_hc, + (size_t) p * hc->nb[2]); + }; + + for (int il = 0; il < w.n_layer; il++) { + const DeepSeek4Layer & L = w.layers[il]; + + // ── HC pre (attention), per block position ────────────────── + std::vector split_attn(block), work_cols(block); + for (int p = 0; p < block; p++) { + ggml_tensor * hcf = hc_col(hc_cur, p); + ggml_tensor * normed = ggml_rms_norm(ctx, hcf, hc_eps); + ggml_tensor * mix = ggml_mul_mat(ctx, L.hc_attn_fn, normed); + mix = ggml_reshape_1d(ctx, mix, mix_dim); + ggml_tensor * base = ggml_reshape_1d(ctx, L.hc_attn_base, mix_dim); + ggml_tensor * pre = ggml_ds4_hc_pre(ctx, mix, base, hcf, n_hc, + w.n_hc_sinkhorn_iter, + C.s_attn[il][0], C.s_attn[il][1], C.s_attn[il][2]); + work_cols[p] = ggml_reshape_2d(ctx, ggml_view_1d(ctx, pre, n_embd, 0), n_embd, 1); + split_attn[p] = ggml_view_1d(ctx, pre, mix_dim, (size_t) n_embd * sizeof(float)); + } + 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); + dbg_tap(std::string("attn_L") + std::to_string(il), attn_out); + // ── HC post (attention), per block position ───────────────── + ggml_tensor * hc_next = nullptr; + for (int p = 0; p < block; p++) { + ggml_tensor * bo = ggml_view_1d(ctx, attn_out, n_embd, (size_t) p * attn_out->nb[1]); + ggml_tensor * hp = ggml_ds4_hc_post(ctx, hc_col(hc_cur, p), bo, split_attn[p], n_hc); + hp = ggml_reshape_3d(ctx, hp, n_embd, n_hc, 1); + hc_next = hc_next ? ggml_concat(ctx, hc_next, hp, 2) : hp; + } + hc_cur = ggml_cont(ctx, hc_next); + + // ── HC pre (FFN), per block position ──────────────────────── + std::vector split_ffn(block), fwork(block); + for (int p = 0; p < block; p++) { + ggml_tensor * hcf = hc_col(hc_cur, p); + ggml_tensor * normed = ggml_rms_norm(ctx, hcf, hc_eps); + ggml_tensor * mix = ggml_mul_mat(ctx, L.hc_ffn_fn, normed); + mix = ggml_reshape_1d(ctx, mix, mix_dim); + ggml_tensor * base = ggml_reshape_1d(ctx, L.hc_ffn_base, mix_dim); + ggml_tensor * pre = ggml_ds4_hc_pre(ctx, mix, base, hcf, n_hc, + w.n_hc_sinkhorn_iter, + C.s_ffn[il][0], C.s_ffn[il][1], C.s_ffn[il][2]); + fwork[p] = ggml_reshape_2d(ctx, ggml_view_1d(ctx, pre, n_embd, 0), n_embd, 1); + split_ffn[p] = ggml_view_1d(ctx, pre, mix_dim, (size_t) n_embd * sizeof(float)); + } + ggml_tensor * ffn_in = fwork[0]; + for (int p = 1; p < block; p++) ffn_in = ggml_concat(ctx, ffn_in, fwork[p], 1); + ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, L.ffn_norm, w.rms_eps); + ggml_tensor * ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, block); + if (!ffn_out) { ggml_free(C.ctx); C.ctx = nullptr; C.gf = nullptr; return false; } + dbg_tap(std::string("ffn_L") + std::to_string(il), ffn_out); + // ── HC post (FFN) ─────────────────────────────────────────── + hc_next = nullptr; + for (int p = 0; p < block; p++) { + ggml_tensor * bo = ggml_view_1d(ctx, ffn_out, n_embd, (size_t) p * ffn_out->nb[1]); + ggml_tensor * hp = ggml_ds4_hc_post(ctx, hc_col(hc_cur, p), bo, split_ffn[p], n_hc); + hp = ggml_reshape_3d(ctx, hp, n_embd, n_hc, 1); + hc_next = hc_next ? ggml_concat(ctx, hc_next, hp, 2) : hp; + } + hc_cur = ggml_cont(ctx, hc_next); + dbg_tap(std::string("hcL") + std::to_string(il), hc_cur); + } + + // ── Tail: hc_head collapse -> out_norm, per block position ────── + ggml_tensor * out = nullptr; + for (int p = 0; p < block; p++) { + ggml_tensor * hcf = hc_col(hc_cur, p); + ggml_tensor * onorm = ggml_rms_norm(ctx, hcf, hc_eps); + ggml_tensor * omix = ggml_mul_mat(ctx, w.output_hc_fn, onorm); + omix = ggml_reshape_1d(ctx, omix, n_hc); + ggml_tensor * obase = ggml_reshape_1d(ctx, w.output_hc_base, n_hc); + ggml_tensor * final_embd = ggml_ds4_hc_out(ctx, omix, obase, hcf, n_hc, C.s_out); + ggml_tensor * final_2d = ggml_reshape_2d(ctx, final_embd, n_embd, 1); + ggml_tensor * hidden_p = build_rms_norm(ctx, final_2d, w.out_norm, w.rms_eps); + out = out ? ggml_concat(ctx, out, hidden_p, 1) : hidden_p; + } + ggml_set_output(out); + ggml_build_forward_expand(gf, out); + C.out = out; + + if (!C.alloc) C.alloc = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (!C.alloc || !ggml_gallocr_alloc_graph(C.alloc, gf)) { + ggml_free(C.ctx); C.ctx = nullptr; C.gf = nullptr; + return false; + } + C.ctx_len = ctx_len; + C.block = block; + C.drafter = (const void *) &d; + } + + // ── 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); + } + 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); + 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 (DS4_DBG) { + for (auto & tp : C.dbg_taps) { + const size_t ne = ggml_nelements(tp.second); + std::vector buf(ne); + ggml_backend_tensor_get(tp.second, buf.data(), 0, sizeof(float) * ne); + double ss = 0.0; size_t nnan = 0; float mn = 1e30f, mx = -1e30f; + for (float v : buf) { + if (!std::isfinite(v)) { nnan++; } + else { ss += (double) v * v; if (v < mn) mn = v; if (v > mx) mx = v; } + } + std::fprintf(stderr, "[ds4-dspark-dbg] %-10s ne=%zu nnan=%zu rms=%.4f min=%.3f max=%.3f\n", + tp.first.c_str(), ne, nnan, std::sqrt(ss / (double) ne), mn, mx); + } + } + + return true; +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index a5d28d23c..66d71c471 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -323,6 +323,16 @@ bool deepseek4_step( DeepSeek4StepTelemetry * telemetry = nullptr, MoeHybridRoutingStats * routing_stats = 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 +// mean-over-HC feature capture and full per-position logits. Null on the normal +// (23 tok/s) decode path so it is completely unaffected. +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] +}; + bool deepseek4_step_layer_range( ggml_backend_t backend, const DeepSeek4Weights & w, @@ -336,7 +346,8 @@ bool deepseek4_step_layer_range( std::vector * out_logits, const int32_t * token_ids = nullptr, DeepSeek4StepTelemetry * telemetry = nullptr, - bool allow_decode_graph_reuse = true); + bool allow_decode_graph_reuse = true, + Ds4VerifyHooks * verify_hooks = nullptr); bool build_deepseek4_moe_hybrid_storage_from_file( const std::string & path, diff --git a/server/test/test_ds4_dspark_load.cpp b/server/test/test_ds4_dspark_load.cpp new file mode 100644 index 000000000..5340e4bb4 --- /dev/null +++ b/server/test/test_ds4_dspark_load.cpp @@ -0,0 +1,153 @@ +// Smoke test: load the DeepSeek-V4-Flash DSpark drafter GGUF and dump bindings. +// +// test_ds4_dspark_load +// +// Verifies the "deepseek4-dflash-draft" GGUF <-> DSparkDrafter loader contract: +// all block tensors + DSpark heads bind, dspark_enabled, capture ids present. + +#include "deepseek4/deepseek4_dspark.h" + +#include "ggml.h" +#include "ggml-backend.h" +#if defined(GGML_USE_CUDA) || defined(GGML_USE_HIP) +#include "ggml-cuda.h" +#define DS4_HAVE_GPU 1 +#endif + +#include +#include +#include +#include +#include + +namespace dflash::common { void deepseek4_dspark_dump(const DSparkDrafter &); } + +int main(int argc, char ** argv) { + using namespace dflash::common; + const char * path = argc > 1 ? argv[1] + : "/home/lucebox2/models/DeepSeek-V4-Flash-DSpark-draft-Q4RMFP4.gguf"; + + ggml_backend_t backend = nullptr; +#ifdef DS4_HAVE_GPU + backend = ggml_backend_cuda_init(0); +#endif + if (!backend) { std::fprintf(stderr, "FAIL: no GPU backend\n"); return 2; } + + DSparkDrafter d; + if (!load_deepseek4_dspark_drafter(path, backend, d)) { + std::fprintf(stderr, "FAIL: load: %s\n", deepseek4_dspark_last_error()); + return 1; + } + deepseek4_dspark_dump(d); + + // Contract checks. + int rc = 0; + auto need = [&](bool cond, const char * what) { + if (!cond) { std::fprintf(stderr, "FAIL: missing %s\n", what); rc = 1; } + }; + need(d.core.n_layer == 3, "n_layer==3"); + need(d.main_proj && d.main_norm, "main_proj/main_norm (dflash.fc/hidden_norm)"); + need(d.markov_w1 && d.markov_w2, "markov heads"); + need(d.dspark_enabled, "dspark_enabled"); + need((int)d.capture_layer_ids.size() == d.n_target_layers, "capture_layer_ids count"); + need(d.core.output == nullptr && d.core.tok_embd == nullptr, "tied embed/lm_head (no output tensors)"); + for (int il = 0; il < d.core.n_layer; il++) { + const auto & L = d.core.layers[il]; + need(L.attn_q_a && L.attn_q_b && L.attn_kv && L.attn_output_a && L.attn_output_b, + "MLA weights"); + need(L.ffn_gate_exps && L.ffn_up_exps && L.ffn_down_exps, "routed experts"); + need(L.ffn_gate_shexp && L.ffn_up_shexp && L.ffn_down_shexp, "shared expert"); + need(L.hc_attn_fn && L.hc_ffn_fn, "HC weights"); + need(d.core.compress_ratios[il] == 0, "compress_ratio==0"); + need(L.attn_compressor_kv == nullptr, "no compressor tensors"); + } + // ── Weight sanity: dequantize main_proj directly (no matmul) ──── + if (rc == 0 && d.main_proj) { + ggml_init_params ip{}; ip.mem_size = 64u*1024*1024; ip.no_alloc = true; + ggml_context * c = ggml_init(ip); + ggml_cgraph * g = ggml_new_graph(c); + ggml_tensor * f32 = ggml_cast(c, d.main_proj, GGML_TYPE_F32); + ggml_set_output(f32); ggml_build_forward_expand(g, f32); + ggml_gallocr_t al = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (al && ggml_gallocr_alloc_graph(al, g) && ggml_backend_graph_compute(backend, g) == GGML_STATUS_SUCCESS) { + size_t ne = ggml_nelements(f32); + std::vector buf(ne); + ggml_backend_tensor_get(f32, buf.data(), 0, sizeof(float)*ne); + size_t nnan=0; double ss=0; float mn=1e30f, mx=-1e30f; + for (float v : buf) { if(!std::isfinite(v)) nnan++; else { ss+=(double)v*v; if(vmx)mx=v; } } + std::fprintf(stderr, "[weight-check] main_proj dequant: ne=%zu nnan=%zu rms=%.5f min=%.4f max=%.4f\n", + ne, nnan, std::sqrt(ss/(double)ne), mn, mx); + } else { std::fprintf(stderr, "[weight-check] main_proj cast failed (type %s not castable?)\n", ggml_type_name(d.main_proj->type)); } + if (al) ggml_gallocr_free(al); ggml_free(c); + } + + // ── Isolate: mul_mat(weight, ones) for main_proj and attn_q_a ─── + if (rc == 0) { + auto mm_check = [&](const char * nm, ggml_tensor * W) { + if (!W) return; + ggml_init_params ip{}; ip.mem_size = 64u*1024*1024; ip.no_alloc = true; + ggml_context * c = ggml_init(ip); + ggml_cgraph * g = ggml_new_graph(c); + ggml_tensor * x = ggml_new_tensor_2d(c, GGML_TYPE_F32, W->ne[0], 1); + ggml_set_input(x); + ggml_tensor * y = ggml_mul_mat(c, W, x); + ggml_set_output(y); ggml_build_forward_expand(g, y); + ggml_gallocr_t al = ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend)); + if (al && ggml_gallocr_alloc_graph(al, g)) { + std::vector ones((size_t)W->ne[0], 1.0f); + ggml_backend_tensor_set(x, ones.data(), 0, sizeof(float)*ones.size()); + if (ggml_backend_graph_compute(backend, g) == GGML_STATUS_SUCCESS) { + size_t ne = ggml_nelements(y); std::vector buf(ne); + ggml_backend_tensor_get(y, buf.data(), 0, sizeof(float)*ne); + size_t nn=0; for(float v:buf) if(!std::isfinite(v)) nn++; + std::fprintf(stderr, "[mm-check] %-12s mul_mat(W[%lld,%lld], ones): nnan=%zu/%zu first=%.3f\n", + nm,(long long)W->ne[0],(long long)W->ne[1],nn,ne,buf.empty()?0:buf[0]); + } + } + if (al) ggml_gallocr_free(al); ggml_free(c); + }; + mm_check("main_proj", d.main_proj); + mm_check("blk0.attn_q_a", d.core.layers[0].attn_q_a); + mm_check("blk0.ffn_gate_shexp", d.core.layers[0].ffn_gate_shexp); + } + + // ── Exercise the drafter forward with dummy inputs ────────────── + // Validates the whole graph runs on the GPU (HC ops, MoE, rocmfp + // matmuls, bidirectional attention, tail) and returns finite output. + if (rc == 0) { + const int n_embd = d.core.n_embd; + const int block = d.block_size; + const int fc_in = d.n_target_layers * n_embd; + const char * cle = std::getenv("DS4_CTX_LEN"); + const int ctx_len = cle ? atoi(cle) : 8; + std::vector noise((size_t) n_embd * block); + std::vector feats((size_t) fc_in * ctx_len); + for (size_t i = 0; i < noise.size(); i++) noise[i] = 0.06f * std::sin(0.31f * (float) i + 1.3f); + for (size_t i = 0; i < feats.size(); i++) feats[i] = 0.05f * std::cos(0.17f * (float) i + 0.4f); + std::vector hidden; + const int committed = 64; // arbitrary >= ctx_len + std::fprintf(stderr, "\n── drafter forward (ctx_len=%d block=%d) ──\n", ctx_len, block); + if (!deepseek4_dspark_draft_forward(backend, d, noise.data(), feats.data(), + ctx_len, committed, hidden)) { + std::fprintf(stderr, "FAIL: drafter forward returned false\n"); + rc = 1; + } else { + bool finite = true; double sumsq = 0.0; + for (float v : hidden) { if (!std::isfinite(v)) finite = false; sumsq += (double) v * v; } + std::fprintf(stderr, "forward out: %zu floats, finite=%d, rms=%.4f, first=[%.4f %.4f %.4f]\n", + hidden.size(), (int) finite, + std::sqrt(sumsq / (double) hidden.size()), + hidden.size() > 0 ? hidden[0] : 0.0f, + hidden.size() > 1 ? hidden[1] : 0.0f, + hidden.size() > 2 ? hidden[2] : 0.0f); + need(finite, "finite forward output"); + need(hidden.size() == (size_t) n_embd * block, "forward output size"); + } + } + + std::fprintf(stderr, rc == 0 ? "\nPASS: DSpark drafter load + forward OK\n" : "\nFAIL\n"); + + free_deepseek4_dspark_drafter(d); + ggml_backend_free(backend); + return rc; +} From 867c9b8a9054fe47322fc2f278d8060d74f94c46 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Fri, 10 Jul 2026 17:45:44 +0200 Subject: [PATCH 07/10] perf(deepseek4): use DSpark confidence for adaptive verify width Return confidence scores from the fused Markov graph in the existing token-id synchronization, then select q=2/q=3/q=4 from cumulative prefix confidence. Compatible DSpark artifacts enable the policy automatically; artifacts without a confidence head retain the existing EWMA controller. No new deployment configuration is introduced.\n\nValidated on gfx1151: GSM+Math 10/10 at 29.25 tok/s weighted, within 0.8% of fixed q=4, with adaptive behavior retained on low-acceptance prompts. --- server/docs/DS4.md | 40 ++++++++++++---- server/src/common/dspark_head.cpp | 46 +++++++++++++++++-- server/src/common/dspark_head.h | 8 ++-- .../src/deepseek4/deepseek4_dspark_spec.cpp | 42 +++++++++++++++-- 4 files changed, 117 insertions(+), 19 deletions(-) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index 5a3890934..b4c259d8e 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -110,13 +110,13 @@ The runtime logs the chosen split with a `[deepseek4-split] auto-split:` banner. 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. -## DSpark Aux Draft Heads +## DSpark Speculative Decode -The current DSpark runtime lives in the shared DFlash speculative stack and is -used by the Laguna backend when the draft GGUF carries `dflash.dspark.*` -tensors. DeepSeek4 DSpark work is still a draft bridge: the DeepSeek4/MTP -artifact stores compatible aux heads under the `mtp.2.*` namespace, and the -converter can now map those names into the existing GGUF contract. +DeepSeek4 uses the shared DFlash DSpark head implementation together with a +DeepSeek4-specific three-layer drafter and fused target verification. The draft +GGUF carries its auxiliary projections under the existing `dflash.dspark.*` +tensor contract. DeepSeek4/MTP checkpoints store compatible heads under the +`mtp.2.*` namespace, which the converter maps as follows. Supported DeepSeek4/MTP input tensors: @@ -141,10 +141,30 @@ python server/scripts/convert_dflash_to_gguf.py \ --aux-heads /path/to/hf-ds4-flash-dspark/model-00048-of-00048.safetensors ``` -This does not by itself make the production DeepSeek4 layer-split backend use -DSpark. It makes the DeepSeek4 DSpark aux artifact consumable by the existing -`dflash.dspark.*` GGUF loader/runtime so the follow-up runtime PR can wire it -into the DS4 decode path with a clean tensor contract. +Run the converted drafter against a DeepSeek4 target with: + +```bash +export DFLASH_DS4_SPEC=1 +export DFLASH_DS4_FUSED_VERIFY=1 +export DFLASH_DS4_DRAFT=/path/to/dflash-draft.gguf +export DFLASH_DS4_SPEC_Q=4 + +./server/build-hip/dflash_server /path/to/deepseek4-target.gguf +``` + +Adaptive width is automatic. When the draft artifact has a compatible +confidence projection, the runtime selects q=2, q=3, or q=4 from the cumulative +confidence of the proposed prefix. It adds the projection to the same fused +Markov graph and reads its scores in the existing token-id synchronization; no +additional host round trip is introduced. Artifacts without a compatible +confidence head transparently retain the existing acceptance-EWMA policy. + +On the gfx1151 validation host, confidence-adaptive width retained 10/10 +GSM+Math accuracy and measured 29.25 tok/s weighted, within 0.8% of fixed q=4 +at 29.49 tok/s. On the low-acceptance stress prompt it measured 21.9/21.8 +tok/s warm, effectively tied with EWMA while avoiding fixed q=4's wasted wide +verification. These numbers are workload-specific; the confidence policy is +opt-in. ## Example: CUDA + Halo Layer Split diff --git a/server/src/common/dspark_head.cpp b/server/src/common/dspark_head.cpp index 4aea171b7..2e1cca95d 100644 --- a/server/src/common/dspark_head.cpp +++ b/server/src/common/dspark_head.cpp @@ -182,6 +182,7 @@ struct MarkovChainGraph { ggml_tensor * base = nullptr; // [vocab, n_positions] std::vector toks; // corrected argmax per depth std::vector corrected; // corrected logits per depth + std::vector confidence; // optional sigmoid score per depth }; // Guards shared by the fused Markov paths: head present, usable inputs, and @@ -218,6 +219,7 @@ bool build_markov_chain_graph(const DraftWeights & dw, ggml_tensor * lm_head, int n_positions, int first_corrected, bool corrected_are_outputs, + bool confidence_are_outputs, std::vector & arena, MarkovChainGraph & out) { const int hdim = dw.n_embd; @@ -252,6 +254,12 @@ bool build_markov_chain_graph(const DraftWeights & dw, ggml_tensor * prev_ids = out.inp_seed; out.toks.assign((size_t)n_corr, nullptr); out.corrected.assign((size_t)n_corr, nullptr); + out.confidence.assign((size_t)n_corr, nullptr); + const bool have_confidence = confidence_are_outputs && + dw.dspark.confidence_w != nullptr && + dw.dspark.confidence_b != nullptr && + (dw.dspark.confidence_dim == hdim || + dw.dspark.confidence_dim == hdim + dw.dspark.markov_rank); for (int i = 0; i < n_corr; ++i) { const int row = first_corrected + i; ggml_tensor * prev_emb = ggml_get_rows(out.ctx, dw.dspark.markov_w1, prev_ids); @@ -268,6 +276,23 @@ bool build_markov_chain_graph(const DraftWeights & dw, ggml_build_forward_expand(out.gf, tok); out.corrected[(size_t)i] = corrected; out.toks[(size_t)i] = tok; + if (have_confidence) { + ggml_tensor * hidden_i = ggml_view_2d( + out.ctx, out.inp_hidden, hdim, 1, out.inp_hidden->nb[1], + (size_t)row * out.inp_hidden->nb[1]); + ggml_tensor * conf_in = hidden_i; + if (dw.dspark.confidence_dim == hdim + dw.dspark.markov_rank) { + conf_in = ggml_concat(out.ctx, hidden_i, prev_emb, 0); + } + ggml_tensor * conf = ggml_mul_mat(out.ctx, dw.dspark.confidence_w, conf_in); + conf = ggml_add( + out.ctx, conf, + ggml_reshape_2d(out.ctx, dw.dspark.confidence_b, 1, 1)); + conf = ggml_sigmoid(out.ctx, conf); + ggml_set_output(conf); + ggml_build_forward_expand(out.gf, conf); + out.confidence[(size_t)i] = conf; + } prev_ids = tok; } return true; @@ -281,7 +306,8 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, const float * local_hidden, int q_len, int32_t last_tok, - std::vector & draft_tok) { + std::vector & draft_tok, + std::vector * confidence_out) { if (q_len <= 1) return false; if (!dspark_fused_usable(dw, backend, lm_head, local_hidden, "dspark_fused")) return false; const int hdim = dw.n_embd; @@ -289,8 +315,12 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, static thread_local std::vector g_arena_chain; MarkovChainGraph g; + const bool want_confidence = confidence_out != nullptr; + if (confidence_out) confidence_out->clear(); if (!build_markov_chain_graph(dw, lm_head, n_cand, /*first_corrected=*/0, - /*corrected_are_outputs=*/false, g_arena_chain, g)) { + /*corrected_are_outputs=*/false, + /*confidence_are_outputs=*/want_confidence, + g_arena_chain, g)) { return false; } @@ -319,14 +349,22 @@ bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, draft_tok[0] = last_tok; // One synchronize instead of n_cand blocking readbacks. int32_t t_out[16]; + float c_out[16] = {}; const int n_get = n_cand < 16 ? n_cand : 16; for (int i = 0; i < n_get; ++i) { ggml_backend_tensor_get_async(backend, g.toks[(size_t)i], &t_out[i], 0, sizeof(int32_t)); + if (want_confidence && g.confidence[(size_t)i]) { + ggml_backend_tensor_get_async( + backend, g.confidence[(size_t)i], &c_out[i], 0, sizeof(float)); + } } ggml_backend_synchronize(backend); for (int i = 0; i < n_get; ++i) { draft_tok[(size_t)i + 1] = t_out[i]; } + if (want_confidence && !g.confidence.empty() && g.confidence[0]) { + confidence_out->assign(c_out, c_out + n_get); + } ggml_free(g.ctx); return true; } @@ -347,7 +385,9 @@ bool dspark_markov_project_topk(const DraftWeights & dw, static thread_local std::vector g_arena_topk; MarkovChainGraph g; if (!build_markov_chain_graph(dw, lm_head, n_tokens, /*first_corrected=*/1, - /*corrected_are_outputs=*/true, g_arena_topk, g)) { + /*corrected_are_outputs=*/true, + /*confidence_are_outputs=*/false, + g_arena_topk, g)) { return false; } diff --git a/server/src/common/dspark_head.h b/server/src/common/dspark_head.h index b76815ce5..ad778e897 100644 --- a/server/src/common/dspark_head.h +++ b/server/src/common/dspark_head.h @@ -20,15 +20,17 @@ bool dspark_markov_correct_greedy_chain(const DraftWeights & dw, // Fused variant: base logits (one lm_head matmul over all candidates) + // unrolled Markov correction chain + in-graph argmax feeding the next // step's get_rows, all in ONE graph on the draft backend. No host logits -// round-trip. Does not implement the confidence gate; callers wanting -// confidence-prefix truncation must use the unfused path. +// round-trip. When confidence_out is non-null and the checkpoint has a +// compatible confidence head, returns one score per candidate from the same +// graph and host synchronization as the token ids. bool dspark_markov_correct_greedy_chain_fused(const DraftWeights & dw, ggml_backend_t backend, ggml_tensor * lm_head, const float * local_hidden, int q_len, int32_t last_tok, - std::vector & draft_tok); + std::vector & draft_tok, + std::vector * confidence_out = nullptr); // DDTree candidate generation with the Markov correction: base logits for // all n_tokens positions in ONE lm_head matmul; rows 1..n-1 get the low-rank diff --git a/server/src/deepseek4/deepseek4_dspark_spec.cpp b/server/src/deepseek4/deepseek4_dspark_spec.cpp index f3cec9b3d..48527e391 100644 --- a/server/src/deepseek4/deepseek4_dspark_spec.cpp +++ b/server/src/deepseek4/deepseek4_dspark_spec.cpp @@ -225,6 +225,12 @@ bool spec_env_flag(const char * name) { return v && *v && *v != '0'; } +// Calibrated cumulative-confidence thresholds for widening the DS4 verify. +// They are part of the policy, not deployment knobs: artifacts without a +// compatible confidence head transparently retain the existing EWMA policy. +constexpr float kConfidenceQ3Threshold = 0.30f; +constexpr float kConfidenceQ4Threshold = 0.25f; + // ── Light rollback state ──────────────────────────────────────────────── // Only the non-position-idempotent cache pieces: the prev-half rows of every // ratio-4 rolling compressor state (attn + indexer) plus hc_state. Everything @@ -384,6 +390,13 @@ 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 && + drafter.confidence_w != nullptr && drafter.confidence_b != nullptr && + (drafter.confidence_dim == n_embd || + drafter.confidence_dim == n_embd + drafter.markov_rank); + if (timing && use_confidence_width) { + std::fprintf(stderr, "[ds4-spec] adaptive width policy=confidence\n"); + } double ewma_accept = 1.5; // Fast path caps the verify at the compression ratio (4): one boundary max, @@ -452,6 +465,7 @@ bool run_deepseek4_dspark_spec_decode( std::vector noise_ids(block); std::vector local_hidden, padded_hidden((size_t) n_embd * (block + 1), 0.0f); std::vector draft_tok, tgt_am; + std::vector draft_confidence; // Cumulative phase timings (ms). double tm_draft = 0, tm_head = 0, tm_save = 0, tm_verify = 0, tm_apply = 0, tm_feat = 0; @@ -491,6 +505,7 @@ bool run_deepseek4_dspark_spec_decode( // let the (row-0-skipping) chain use slots 1..q-1. t0 = SpecClock::now(); draft_tok.clear(); + draft_confidence.clear(); 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 @@ -503,15 +518,17 @@ bool run_deepseek4_dspark_spec_decode( 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 && !seq_verify_mode) { + 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; } if (q_step_cap >= 2) { std::memcpy(padded_hidden.data() + n_embd, local_hidden.data(), sizeof(float) * (size_t) n_embd * block); - ds_ok = dspark_markov_correct_greedy_chain_fused(dw, backend, target.lm_head_tensor(), - padded_hidden.data(), q_step_cap, lt, draft_tok); + ds_ok = dspark_markov_correct_greedy_chain_fused( + dw, backend, target.lm_head_tensor(), padded_hidden.data(), + q_step_cap, lt, draft_tok, + use_confidence_width ? &draft_confidence : 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); @@ -527,6 +544,25 @@ bool run_deepseek4_dspark_spec_decode( } else { draft_tok.push_back(lt); // q=1: seed only, no speculation } + // Confidence estimates are per candidate. Their cumulative product is + // the estimated probability that the target accepts the whole prefix + // unlocked by a wider verify. The defaults were calibrated on q=4 + // traces and keep q=4 for high-confidence prefixes while avoiding its + // extra verify cost on low-acceptance prompts. + if (use_confidence_width && draft_confidence.size() >= 2 && draft_tok.size() >= 3) { + const float p2 = draft_confidence[0] * draft_confidence[1]; + int selected_q = p2 >= kConfidenceQ3Threshold ? 3 : 2; + if (selected_q == 3 && draft_confidence.size() >= 3 && draft_tok.size() >= 4) { + const float p3 = p2 * draft_confidence[2]; + if (p3 >= kConfidenceQ4Threshold) selected_q = 4; + } + if ((int) draft_tok.size() > selected_q) draft_tok.resize((size_t) selected_q); + } else if (use_confidence_width && !seq_verify_mode) { + // The fused head should always return confidence for a compatible + // artifact. Preserve the old policy if a backend cannot do so. + const int selected_q = (int) ewma_accept + 2; + if ((int) draft_tok.size() > selected_q) draft_tok.resize((size_t) selected_q); + } if ((int) draft_tok.size() > q_step_cap) draft_tok.resize(q_step_cap); const int q = (int) draft_tok.size(); // seed + candidates tm_head += spec_ms_since(t0); From 95c9812fd01cdcedafd7f2fc36109243d087bffb Mon Sep 17 00:00:00 2001 From: User Date: Sun, 12 Jul 2026 02:18:04 +0800 Subject: [PATCH 08/10] perf(deepseek4): enable causal Flash prefill --- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 22 +-- server/src/deepseek4/deepseek4_backend.cpp | 50 ++++--- server/src/deepseek4/deepseek4_graph.cpp | 133 +++++++++++------- 3 files changed, 121 insertions(+), 84 deletions(-) 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 e991bf4a3..e5fe90944 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -36,6 +36,7 @@ __global__ static void ds4_flash_attn_d512_f32_shared_kv_kernel( const float * q, const float * k, const float * v, + const __half * mask, const float * sinks, int n_tokens, int n_heads, @@ -58,7 +59,7 @@ __global__ static void ds4_flash_attn_d512_f32_shared_kv_kernel( for (int d = 0; d < D; ++d) { dot += qh[d] * kr[d]; } - const float s = dot * scale; + const float s = dot * scale + __half2float(mask[(size_t) t * n_kv + r]); scores[r] = s; local_max = fmaxf(local_max, s); } @@ -86,17 +87,14 @@ __global__ static void ds4_flash_attn_d512_f32_shared_kv_kernel( } static bool ggml_cuda_ds4_flash_attn_d512_f32(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - const char * e = getenv("DFLASH_DS4_FLASH_ATTN"); - if (!e || !e[0] || e[0] == (char)48) { - 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]; - if (!Q || !K || !V || mask || + if (!Q || !K || !V || !mask || Q->type != GGML_TYPE_F32 || K->type != GGML_TYPE_F32 || V->type != GGML_TYPE_F32 || + mask->type != GGML_TYPE_F16 || !ggml_is_contiguous(mask) || dst->type != GGML_TYPE_F32 || Q->ne[0] != 512 || K->ne[0] != 512 || V->ne[0] != 512 || K->ne[2] != 1 || V->ne[2] != 1 || @@ -105,7 +103,8 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32(ggml_backend_cuda_context & ctx, g Q->nb[0] != (int64_t) sizeof(float) || K->nb[0] != (int64_t) sizeof(float) || V->nb[0] != (int64_t) sizeof(float) || - dst->nb[0] != (int64_t) sizeof(float)) { + dst->nb[0] != (int64_t) sizeof(float) || + mask->ne[0] != K->ne[1] || mask->ne[1] != Q->ne[1]) { return false; } if (sinks && (sinks->type != GGML_TYPE_F32 || sinks->ne[0] != Q->ne[2])) { @@ -127,6 +126,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32(ggml_backend_cuda_context & ctx, g (const float *) Q->data, (const float *) K->data, (const float *) V->data, + (const __half *) mask->data, sinks ? (const float *) sinks->data : nullptr, n_tokens, n_heads, @@ -481,11 +481,6 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const } } } - static const bool ds4_allow_d512_nomask = [] { - const char * e = getenv("DFLASH_DS4_FLASH_ATTN"); - return e && e[0] && e[0] != (char)48; - }(); - const int cc = ggml_cuda_info().devices[device].cc; switch (K->ne[0]) { @@ -505,9 +500,6 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const if (V->ne[0] != K->ne[0]) { return BEST_FATTN_KERNEL_NONE; } - if (!gqa_opt_applies && !ds4_allow_d512_nomask) { - return BEST_FATTN_KERNEL_NONE; - } break; case 576: if (V->ne[0] != 512) { diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 4afe7731c..8ec497310 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -33,6 +33,13 @@ static bool env_flag_enabled(const char * name) { return value && value[0] && std::strcmp(value, "0") != 0; } +static int env_positive_int(const char * name) { + const char * value = std::getenv(name); + if (!value || !value[0]) return 0; + const long parsed = std::strtol(value, nullptr, 10); + return parsed > 0 && parsed <= INT32_MAX ? (int) parsed : 0; +} + static double gib(uint64_t bytes) { return (double) bytes / 1024.0 / 1024.0 / 1024.0; } @@ -319,25 +326,22 @@ bool DeepSeek4Backend::init() { snap_backend_ = ggml_backend_init_by_name("cpu", nullptr); - const PlacementBackend target_backend = - cfg_.device.backend == PlacementBackend::Auto - ? compiled_placement_backend() - : cfg_.device.backend; + const bool force_hybrid = env_flag_enabled("DFLASH_DEEPSEEK4_FORCE_HYBRID"); + if (force_hybrid) { + std::fprintf(stderr, "[deepseek4] force hybrid mode requested; skipping full model load\n"); + } - // 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) { - std::fprintf(stderr, - "[deepseek4] HIP target detected; using hybrid expert load path\n"); + // On UMA, a monolithic managed allocation keeps all routed experts local. + // Preserve the hybrid loader as the OOM fallback and as an explicit override. + if (!force_hybrid && !load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { + std::fprintf(stderr, "[deepseek4] full model load failed, trying hybrid mode...\n"); if (!init_hybrid_model()) { - std::fprintf(stderr, "[deepseek4] hybrid mode failed: %s\n", cfg_.model_path); + std::fprintf(stderr, "[deepseek4] hybrid mode also failed: %s\n", cfg_.model_path); return false; } - } else if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { - std::fprintf(stderr, "[deepseek4] full model load failed, trying hybrid mode...\n"); + } else if (force_hybrid) { if (!init_hybrid_model()) { - std::fprintf(stderr, "[deepseek4] hybrid mode also failed: %s\n", cfg_.model_path); + std::fprintf(stderr, "[deepseek4] forced hybrid mode failed: %s\n", cfg_.model_path); return false; } } @@ -561,6 +565,11 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, // semantics at any length; chunked stays available via // DFLASH_DS4_CHUNKED_PREFILL=1 for short-prompt benchmarking only. const bool unsafe_chunked = env_flag_enabled("DFLASH_DS4_CHUNKED_PREFILL"); + const bool causal_flash_prefill = !moe_hybrid_ && !unsafe_chunked; + const int requested_flash_chunk = env_positive_int("DFLASH_DS4_FLASH_PREFILL_CHUNK"); + const int flash_prefill_chunk = causal_flash_prefill + ? std::min(requested_flash_chunk > 0 ? requested_flash_chunk : 128, w_.n_swa) + : 0; const int chunk = (moe_hybrid_ || !unsafe_chunked) ? 1 : (cfg_.chunk > 0 ? cfg_.chunk : 512); const int n_total = (int)tokens.size(); int pos = kv_offset; @@ -571,6 +580,11 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, // request on can drift by a token or two. if (kv_offset == 0 && cache_.buf) { ggml_backend_buffer_clear(cache_.buf, 0); + cache_.cur_pos = 0; + for (auto & layer : cache_.layers) { + layer.n_comp = 0; + layer.n_index_comp = 0; + } } last_logits_.clear(); if (spec_enabled_ && kv_offset == 0) spec_feat_window_.clear(); @@ -579,10 +593,13 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, DeepSeek4StepTelemetry tel_acc; int steps = 0; - for (int i = 0; i < n_total; i += chunk) { + for (int i = 0; i < n_total;) { if (io.cancelled) return pos; - const int n_tok = std::min(chunk, n_total - i); + int n_tok = std::min(chunk, n_total - i); + if (flash_prefill_chunk > 1) { + n_tok = std::min(flash_prefill_chunk, n_total - i); + } // Embed tokens std::vector embed(w_.n_embd * n_tok); @@ -636,6 +653,7 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, } last_logits_ = std::move(logits); pos += n_tok; + i += n_tok; } if (timing) { log_step_tel("prefill", n_total, steps, elapsed_s(phase_t0), tel_acc); diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 0eeea9a0d..9a95e654b 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -225,6 +225,11 @@ struct DeepSeek4F32ArrayBinding { std::vector values; }; +struct DeepSeek4F16ArrayBinding { + ggml_tensor * tensor = nullptr; + std::vector values; +}; + 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, @@ -1139,7 +1144,8 @@ 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, + std::vector * f16_array_inputs = nullptr) { const int n_embd = w.n_embd; const int head_dim = w.head_dim; @@ -1208,6 +1214,8 @@ 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 use_flash_prefill = causal_batch && f16_array_inputs && + !ds4_env_flag("DFLASH_DS4_CHUNKED_PREFILL"); ggml_tensor * old_rows_scratch = nullptr; int n_old_rows = 0; const bool fused_causal = cached_inputs && cached_inputs->attn_row_mask && n_tokens > 1; @@ -1386,33 +1394,16 @@ 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); + std::vector causal_mask_values; + if (causal_batch) { + // Per-token causal visibility over [ring rows | comp rows | old rows]. + causal_mask_values.assign((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; + float * col = causal_mask_values.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 + // 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; @@ -1421,8 +1412,8 @@ static ggml_tensor * build_mla_attention( 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; } - // old contents of slot overwritten by batch token j are visible - // exactly to tokens i < j (still inside their SWA window) + // Old contents of slot overwritten by batch token j are visible + // exactly to tokens i < j while they remain in the SWA window. int oi = 0; for (int tj = 0; tj < n_tokens; tj++) { if (kv_start + tj < w.n_swa) continue; @@ -1430,39 +1421,71 @@ static ggml_tensor * build_mla_attention( oi++; } } - 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 * context = nullptr; + if (use_flash_prefill) { + ggml_tensor * flash_mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, n_attn, n_tokens); + ggml_set_input(flash_mask); + constexpr uint16_t kF16Zero = 0x0000; + constexpr uint16_t kF16NegInf = 0xfc00; + std::vector flash_mask_values(causal_mask_values.size(), kF16Zero); + for (size_t i = 0; i < causal_mask_values.size(); ++i) { + if (causal_mask_values[i] < 0.0f) flash_mask_values[i] = kF16NegInf; + } + f16_array_inputs->push_back({flash_mask, std::move(flash_mask_values)}); + + ggml_tensor * q_flash = ggml_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); + ggml_tensor * kv_flash = ggml_cont(ctx, kv_attn); + kv_flash = ggml_reshape_3d(ctx, kv_flash, head_dim, n_attn, 1); + context = ggml_flash_attn_ext(ctx, q_flash, kv_flash, kv_flash, flash_mask, + kq_scale, 0.0f, 0.0f); + if (L.attn_sinks) { + ggml_flash_attn_ext_add_sinks(context, L.attn_sinks); } - 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] + // 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 (masked_kv && n_tokens > 1) { + 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) { + scores = ggml_add(ctx, scores, cached_inputs->attn_row_mask); + } else if (causal_batch) { + ggml_tensor * cmask = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_attn, 1, n_tokens); + ggml_set_input(cmask); + f32_array_inputs->push_back({cmask, std::move(causal_mask_values)}); + 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); + } - // 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] + // 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); + } - // Reshape back to [head_dim, n_head, n_tokens] - context = ggml_reshape_3d(ctx, context, head_dim, n_head, n_tokens); + 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; @@ -4856,6 +4879,7 @@ bool deepseek4_step_layer_range( std::vector i32_array_inputs; std::vector i64_array_inputs; std::vector f32_array_inputs; + std::vector f16_array_inputs; const size_t graph_size = ds4_attn_step_graph_size(n_tokens); gf = ggml_new_graph_custom(ctx, graph_size, false); @@ -4863,7 +4887,8 @@ bool deepseek4_step_layer_range( 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, + &f16_array_inputs); ggml_set_output(attn_out); ggml_build_forward_expand(gf, attn_out); @@ -4893,6 +4918,8 @@ bool deepseek4_step_layer_range( 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()); + for (const auto & b : f16_array_inputs) + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, sizeof(uint16_t) * b.values.size()); } const auto attn_compute_t0 = Ds4TimingClock::now(); From c46e9454d1f58017f6a1f7ea457ee908c5dd5325 Mon Sep 17 00:00:00 2001 From: User Date: Sun, 12 Jul 2026 10:57:49 +0800 Subject: [PATCH 09/10] fix(deepseek4): harden causal Flash prefill --- .../ggml/rocmfp4/rocmfp4_hip_scale.cuh | 3 ++ .../llama.cpp/ggml/rocmfpx/test_rocmfpx.c | 23 +++++++++ .../llama.cpp/ggml/src/ggml-cuda/convert.cu | 28 +++++++---- .../llama.cpp/ggml/src/ggml-cuda/fattn.cu | 47 ++++++++++++++---- server/deps/llama.cpp/ggml/src/ggml.c | 6 +++ .../deps/llama.cpp/ggml/tests/CMakeLists.txt | 11 +++++ server/docs/DS4.md | 22 ++++----- server/src/deepseek4/deepseek4_backend.cpp | 48 +++++++++++-------- server/src/deepseek4/deepseek4_dspark.h | 1 - 9 files changed, 135 insertions(+), 54 deletions(-) create mode 100644 server/deps/llama.cpp/ggml/tests/CMakeLists.txt diff --git a/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_scale.cuh b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_scale.cuh index 19525a4fb..9d48a792c 100644 --- a/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_scale.cuh +++ b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_scale.cuh @@ -96,6 +96,9 @@ static __device__ __forceinline__ float rocmfp4_ue4m3_to_fp32_half_finite(uint8_ #if defined(GGML_USE_HIP) && GGML_ROCMFP4_USE_SCALE_LUT return x <= 0x7e ? rocmfp4_scale_ue4m3_half_lut[x] : 0.0f; #else + if (x > 0x7e) { + return 0.0f; + } const int exp = (x >> 3) & 0xF; const int man = x & 0x7; diff --git a/server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c b/server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c index 5bae72bd7..7d355851a 100644 --- a/server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c +++ b/server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c @@ -1,4 +1,6 @@ #include "rocmfpx.h" +#include "rocmfp4.h" +#include "ggml.h" #include #include @@ -100,6 +102,26 @@ static void check_weighted_imatrix_fp2(void) { assert(weighted_err < plain_err); } +static void check_ggml_quantize_dispatch(const float * src) { + enum { N = 64 }; + const enum ggml_type types[] = { + GGML_TYPE_Q4_0_ROCMFP4, + GGML_TYPE_Q4_0_ROCMFP4_FAST, + GGML_TYPE_Q2_0_ROCMFP2, + GGML_TYPE_Q3_0_ROCMFPX, + GGML_TYPE_Q6_0_ROCMFPX, + GGML_TYPE_Q8_0_ROCMFPX, + }; + + uint8_t dst[256]; + for (size_t i = 0; i < sizeof(types)/sizeof(types[0]); ++i) { + const size_t row_size = ggml_row_size(types[i], N); + assert(row_size <= sizeof(dst)); + const size_t written = ggml_quantize_chunk(types[i], src, dst, 0, 1, N, NULL); + assert(written == row_size); + } +} + int main(void) { enum { N = 64 }; @@ -115,6 +137,7 @@ int main(void) { block_rocmfp8 q8[N / QK_ROCMFP8]; fill_row(src, N); + check_ggml_quantize_dispatch(src); rocmfpx_quantize_row_fp2_ref(src, q2, N); rocmfpx_quantize_row_fp3_ref(src, q3, N); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/convert.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/convert.cu index 28d042aa0..138c82e46 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/convert.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/convert.cu @@ -489,7 +489,7 @@ static __global__ void dequantize_block_mxfp4(const void * __restrict__ vx, dst_ } template -static __global__ void dequantize_block_rocmfp4(const void * __restrict__ vx, dst_t * __restrict__ yy) { +static __global__ void dequantize_block_rocmfp4(const void * __restrict__ vx, dst_t * __restrict__ yy, const int64_t ne) { const int64_t i = blockIdx.x; const block_rocmfp4 * x = (const block_rocmfp4 *) vx + i*(QK_K/QK_ROCMFP4); @@ -497,18 +497,22 @@ static __global__ void dequantize_block_rocmfp4(const void * __restrict__ vx, ds const int64_t tid = threadIdx.x; const int64_t il = tid/8; // 0...3 const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + i*QK_K + 32*ib + 4*il; + const int64_t y_base = i*QK_K + 32*ib + 4*il; const uint8_t * q4 = x[ib].qs + 4*il; const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[0]); const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[1]); for (int j = 0; j < 4; ++j) { - y[j+ 0] = d0 * rocmfp4_decode_i8(q4[j]); - y[j+16] = d1 * rocmfp4_decode_i8(q4[j] >> 4); + if (y_base + j < ne) { + yy[y_base + j] = d0 * rocmfp4_decode_i8(q4[j]); + } + if (y_base + 16 + j < ne) { + yy[y_base + 16 + j] = d1 * rocmfp4_decode_i8(q4[j] >> 4); + } } } template -static __global__ void dequantize_block_rocmfp4_fast(const void * __restrict__ vx, dst_t * __restrict__ yy) { +static __global__ void dequantize_block_rocmfp4_fast(const void * __restrict__ vx, dst_t * __restrict__ yy, const int64_t ne) { const int64_t i = blockIdx.x; const block_rocmfp4_fast * x = (const block_rocmfp4_fast *) vx + i*(QK_K/QK_ROCMFP4); @@ -516,12 +520,16 @@ static __global__ void dequantize_block_rocmfp4_fast(const void * __restrict__ v const int64_t tid = threadIdx.x; const int64_t il = tid/8; // 0...3 const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + i*QK_K + 32*ib + 4*il; + const int64_t y_base = i*QK_K + 32*ib + 4*il; const uint8_t * q4 = x[ib].qs + 4*il; const float d = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e); for (int j = 0; j < 4; ++j) { - y[j+ 0] = d * rocmfp4_decode_i8(q4[j]); - y[j+16] = d * rocmfp4_decode_i8(q4[j] >> 4); + if (y_base + j < ne) { + yy[y_base + j] = d * rocmfp4_decode_i8(q4[j]); + } + if (y_base + 16 + j < ne) { + yy[y_base + 16 + j] = d * rocmfp4_decode_i8(q4[j] >> 4); + } } } @@ -686,13 +694,13 @@ static void dequantize_row_mxfp4_cuda(const void * vx, dst_t * y, const int64_t template static void dequantize_row_rocmfp4_hip(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { const int nb = (k + QK_K - 1) / QK_K; - dequantize_block_rocmfp4<<>>(vx, y); + dequantize_block_rocmfp4<<>>(vx, y, k); } template static void dequantize_row_rocmfp4_fast_hip(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { const int nb = (k + QK_K - 1) / QK_K; - dequantize_block_rocmfp4_fast<<>>(vx, y); + dequantize_block_rocmfp4_fast<<>>(vx, y, k); } template 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 e5fe90944..3c716798a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -86,7 +86,7 @@ __global__ static void ds4_flash_attn_d512_f32_shared_kv_kernel( } } -static bool ggml_cuda_ds4_flash_attn_d512_f32(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { +static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(int device, const ggml_tensor * dst) { const ggml_tensor * Q = dst->src[0]; const ggml_tensor * K = dst->src[1]; const ggml_tensor * V = dst->src[2]; @@ -94,20 +94,18 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32(ggml_backend_cuda_context & ctx, g const ggml_tensor * sinks = dst->src[4]; if (!Q || !K || !V || !mask || Q->type != GGML_TYPE_F32 || K->type != GGML_TYPE_F32 || V->type != GGML_TYPE_F32 || - mask->type != GGML_TYPE_F16 || !ggml_is_contiguous(mask) || + mask->type != GGML_TYPE_F16 || + !ggml_is_contiguous(Q) || !ggml_is_contiguous(K) || !ggml_is_contiguous(V) || + !ggml_is_contiguous(mask) || !ggml_is_contiguous(dst) || dst->type != GGML_TYPE_F32 || Q->ne[0] != 512 || K->ne[0] != 512 || V->ne[0] != 512 || 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] || - Q->nb[0] != (int64_t) sizeof(float) || - K->nb[0] != (int64_t) sizeof(float) || - V->nb[0] != (int64_t) sizeof(float) || - dst->nb[0] != (int64_t) sizeof(float) || - mask->ne[0] != K->ne[1] || mask->ne[1] != Q->ne[1]) { + dst->ne[0] != 512 || dst->ne[1] != Q->ne[2] || dst->ne[2] != Q->ne[1] || dst->ne[3] != 1 || + mask->ne[0] != K->ne[1] || mask->ne[1] != Q->ne[1] || mask->ne[2] != 1 || mask->ne[3] != 1) { return false; } - if (sinks && (sinks->type != GGML_TYPE_F32 || sinks->ne[0] != Q->ne[2])) { + if (sinks && (sinks->type != GGML_TYPE_F32 || !ggml_is_contiguous(sinks) || sinks->ne[0] != Q->ne[2])) { return false; } @@ -118,6 +116,26 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32(ggml_backend_cuda_context & ctx, g return false; } + constexpr size_t k_static_smem = 2 * 256 * sizeof(float); + const size_t max_shared = ggml_cuda_info().devices[device].smpb; + return max_shared > k_static_smem && + (size_t) n_kv <= (max_shared - k_static_smem) / sizeof(float); +} + +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(ctx.device, 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 int n_tokens = (int) Q->ne[1]; + const int n_heads = (int) Q->ne[2]; + const int n_kv = (int) K->ne[1]; + cudaStream_t stream = ctx.stream(); dim3 grid((unsigned) n_tokens, (unsigned) n_heads, 1); const size_t shmem = (size_t) n_kv * sizeof(float); @@ -131,6 +149,11 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32(ggml_backend_cuda_context & ctx, g n_tokens, n_heads, n_kv); + const cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + GGML_LOG_ERROR("%s: launch failed: %s\n", __func__, cudaGetErrorString(err)); + return false; + } return true; } @@ -500,6 +523,9 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const if (V->ne[0] != K->ne[0]) { return BEST_FATTN_KERNEL_NONE; } + if (!gqa_opt_applies && (turing_mma_available(cc) || volta_mma_available(cc))) { + return BEST_FATTN_KERNEL_NONE; + } break; case 576: if (V->ne[0] != 512) { @@ -718,5 +744,6 @@ 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) { - return ggml_cuda_get_best_fattn_kernel(device, dst) != BEST_FATTN_KERNEL_NONE; + return ggml_cuda_ds4_flash_attn_d512_f32_supported(device, dst) || + ggml_cuda_get_best_fattn_kernel(device, dst) != BEST_FATTN_KERNEL_NONE; } diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 591ab9bda..7ee365f76 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -7915,10 +7915,16 @@ size_t ggml_quantize_chunk( switch (type) { case GGML_TYPE_Q1_0: result = quantize_q1_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q4_0: result = quantize_q4_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q4_0_ROCMFP4: result = rocmfp4_quantize_q4_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: result = rocmfp4_quantize_q4_0_fast(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q4_1: result = quantize_q4_1(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q5_0: result = quantize_q5_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q5_1: result = quantize_q5_1(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q8_0: result = quantize_q8_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q2_0_ROCMFP2: result = rocmfpx_quantize_fp2(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q3_0_ROCMFPX: result = rocmfpx_quantize_fp3(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q6_0_ROCMFPX: result = rocmfpx_quantize_fp6(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q8_0_ROCMFPX: result = rocmfpx_quantize_fp8(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_MXFP4: result = quantize_mxfp4(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_NVFP4: result = quantize_nvfp4(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q2_K: result = quantize_q2_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; diff --git a/server/deps/llama.cpp/ggml/tests/CMakeLists.txt b/server/deps/llama.cpp/ggml/tests/CMakeLists.txt new file mode 100644 index 000000000..2176f15d5 --- /dev/null +++ b/server/deps/llama.cpp/ggml/tests/CMakeLists.txt @@ -0,0 +1,11 @@ +add_executable(test-rocmfpx ../rocmfpx/test_rocmfpx.c) +target_include_directories(test-rocmfpx PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/../include + ${CMAKE_CURRENT_SOURCE_DIR}/../rocmfpx) +target_link_libraries(test-rocmfpx PRIVATE ggml-base) + +if (NOT MSVC) + target_link_libraries(test-rocmfpx PRIVATE m) +endif() + +add_test(NAME test-rocmfpx COMMAND test-rocmfpx) diff --git a/server/docs/DS4.md b/server/docs/DS4.md index b4c259d8e..91f718d77 100644 --- a/server/docs/DS4.md +++ b/server/docs/DS4.md @@ -116,7 +116,8 @@ DeepSeek4 uses the shared DFlash DSpark head implementation together with a DeepSeek4-specific three-layer drafter and fused target verification. The draft GGUF carries its auxiliary projections under the existing `dflash.dspark.*` tensor contract. DeepSeek4/MTP checkpoints store compatible heads under the -`mtp.2.*` namespace, which the converter maps as follows. +`mtp.2.*` namespace. DeepSeek4 draft artifacts use the following GGUF tensor +names. Supported DeepSeek4/MTP input tensors: @@ -127,21 +128,16 @@ Supported DeepSeek4/MTP input tensors: | `mtp.2.confidence_head.proj.weight` | `dflash.dspark.confidence.weight` | | `mtp.2.confidence_head.proj.bias` | `dflash.dspark.confidence.bias` | -If the MTP confidence projection is bias-less, the converter writes a zero -bias so the GGUF loader still sees the pair it expects. The Markov head alone -is enough for DSpark greedy-chain correction; confidence gating remains +If the MTP confidence projection is bias-less, the artifact must contain a +zero bias so the GGUF loader still sees the pair it expects. The Markov head +alone is enough for DSpark greedy-chain correction; confidence gating remains optional. -Example conversion with the DS4 MTP shard that contains the DSpark heads: +This repository does not yet ship a DeepSeek4 draft converter. Do not use +`convert_dflash_to_gguf.py` for this model: it emits the Qwen draft +architecture and cannot produce a loadable `deepseek4-dflash-draft` artifact. -```bash -python server/scripts/convert_dflash_to_gguf.py \ - /path/to/dflash-draft/model.safetensors \ - /path/to/dflash-draft.gguf \ - --aux-heads /path/to/hf-ds4-flash-dspark/model-00048-of-00048.safetensors -``` - -Run the converted drafter against a DeepSeek4 target with: +Run a compatible drafter against a DeepSeek4 target with: ```bash export DFLASH_DS4_SPEC=1 diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 8ec497310..00a4b7e1d 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -508,24 +508,20 @@ bool DeepSeek4Backend::unpark(const std::string & what) { const bool want_target = (what.empty() || what == "all" || what == "target"); if (!want_target || !parked_) return true; - const PlacementBackend target_backend = - cfg_.device.backend == PlacementBackend::Auto - ? compiled_placement_backend() - : cfg_.device.backend; - - if (target_backend == PlacementBackend::Hip) { + const bool force_hybrid = env_flag_enabled("DFLASH_DEEPSEEK4_FORCE_HYBRID"); + if (!force_hybrid && !load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { + std::fprintf(stderr, "[deepseek4] unpark: full model reload failed, trying hybrid mode...\n"); if (!init_hybrid_model()) { - std::fprintf(stderr, "[deepseek4] unpark: failed to restore hybrid mode\n"); + std::fprintf(stderr, "[deepseek4] unpark: failed to restore target model\n"); free_deepseek4_weights(w_); stream_engine_.destroy(); moe_hybrid_.reset(); moe_placement_ = {}; return false; } - } else if (!load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { - std::fprintf(stderr, "[deepseek4] unpark: full model reload failed, trying hybrid mode...\n"); + } else if (force_hybrid) { if (!init_hybrid_model()) { - std::fprintf(stderr, "[deepseek4] unpark: failed to restore target model\n"); + std::fprintf(stderr, "[deepseek4] unpark: failed to restore forced hybrid mode\n"); free_deepseek4_weights(w_); stream_engine_.destroy(); moe_hybrid_.reset(); @@ -635,8 +631,14 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, tokens.data() + i, timing ? &step_tel : nullptr, /*allow_decode_graph_reuse=*/true, hp); - if (hp && !spec_cap.empty()) { - const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; + if (ok && hp) { + const size_t feat_row = (size_t) spec_drafter_->n_target_layers * (size_t) w_.n_embd; + const size_t expected = (size_t) n_tok * feat_row; + if (feat_row == 0 || spec_cap.size() != expected) { + std::fprintf(stderr, "[deepseek4] speculative capture size mismatch: got=%zu expected=%zu\n", + spec_cap.size(), expected); + return -1; + } for (int t = 0; t < n_tok; t++) spec_feat_window_.insert(spec_feat_window_.end(), spec_cap.begin() + (size_t) t * feat_row, @@ -776,25 +778,31 @@ GenerateResult DeepSeek4Backend::generate_impl(const GenerateRequest & req, // Decode auto t1 = Clock::now(); - if (spec_enabled_ && spec_drafter_ && req.n_gen > 0) { + if (spec_enabled_ && spec_drafter_ && req.n_gen > 0 && req.budget_hook.close_token_ids.empty()) { if (last_logits_.empty()) { result.error = "spec: no prefill logits"; return result; } int seed = 0; { float mv = last_logits_[0]; for (int i = 1; i < w_.n_vocab; i++) if (last_logits_[i] > mv) { mv = last_logits_[i]; seed = i; } } std::vector gen; - gen.push_back(seed); - io.emit(seed); + if (!io.cancelled) { + gen.push_back(seed); + io.emit(seed); + } float accept_rate = 0.0f; - if (!deepseek4_is_eos_tok(seed, w_) && req.n_gen > 1) { - const int feat_row = spec_drafter_->n_target_layers * w_.n_embd; - const int win_len = feat_row > 0 ? (int) (spec_feat_window_.size() / feat_row) : 0; + if (!io.cancelled && !deepseek4_is_eos_tok(seed, w_) && req.n_gen > 1) { + const size_t feat_row = (size_t) spec_drafter_->n_target_layers * (size_t) w_.n_embd; + const int win_len = feat_row > 0 && spec_feat_window_.size() % feat_row == 0 + ? (int) (spec_feat_window_.size() / feat_row) : 0; std::vector spec_toks; run_deepseek4_dspark_spec_decode(backend_, w_, cache_, *spec_drafter_, committed, seed, req.n_gen - 1, win_len > 0 ? spec_feat_window_.data() : nullptr, win_len, spec_toks, &accept_rate); - for (int t : spec_toks) io.emit(t); - gen.insert(gen.end(), spec_toks.begin(), spec_toks.end()); + for (int t : spec_toks) { + if (io.cancelled) break; + io.emit(t); + gen.push_back(t); + } } result.ok = true; result.tokens = std::move(gen); diff --git a/server/src/deepseek4/deepseek4_dspark.h b/server/src/deepseek4/deepseek4_dspark.h index 9fe20da75..77be2a2a9 100644 --- a/server/src/deepseek4/deepseek4_dspark.h +++ b/server/src/deepseek4/deepseek4_dspark.h @@ -77,7 +77,6 @@ const char * deepseek4_dspark_last_error(); // features. All host-side f32 for a simple v1 (GPU feature-ring plumbing can // come later). // -// seed_tok : the committed anchor token (block position 0's embedding) // noise_embed : [n_embd * block_size] embeds of [seed]+[MASK]*(block_size-1) // ctx_features : [n_target_layers*n_embd * ctx_len] captured features, // ordered oldest..newest, absolute positions From 07719343cc8f859211a2ae961b0b4f4edee32177 Mon Sep 17 00:00:00 2001 From: User Date: Sun, 12 Jul 2026 12:13:21 +0800 Subject: [PATCH 10/10] fix(deepseek4): reset weights after failed reload --- server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c | 1 - server/src/deepseek4/deepseek4_backend.cpp | 2 ++ server/src/deepseek4/deepseek4_loader.cpp | 5 +---- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c b/server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c index 7d355851a..9d118022e 100644 --- a/server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c +++ b/server/deps/llama.cpp/ggml/rocmfpx/test_rocmfpx.c @@ -1,5 +1,4 @@ #include "rocmfpx.h" -#include "rocmfp4.h" #include "ggml.h" #include diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 00a4b7e1d..57e53f6a2 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -335,6 +335,7 @@ bool DeepSeek4Backend::init() { // Preserve the hybrid loader as the OOM fallback and as an explicit override. if (!force_hybrid && !load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { std::fprintf(stderr, "[deepseek4] full model load failed, trying hybrid mode...\n"); + free_deepseek4_weights(w_); if (!init_hybrid_model()) { std::fprintf(stderr, "[deepseek4] hybrid mode also failed: %s\n", cfg_.model_path); return false; @@ -511,6 +512,7 @@ bool DeepSeek4Backend::unpark(const std::string & what) { const bool force_hybrid = env_flag_enabled("DFLASH_DEEPSEEK4_FORCE_HYBRID"); if (!force_hybrid && !load_deepseek4_gguf(cfg_.model_path, backend_, w_)) { std::fprintf(stderr, "[deepseek4] unpark: full model reload failed, trying hybrid mode...\n"); + free_deepseek4_weights(w_); if (!init_hybrid_model()) { std::fprintf(stderr, "[deepseek4] unpark: failed to restore target model\n"); free_deepseek4_weights(w_); diff --git a/server/src/deepseek4/deepseek4_loader.cpp b/server/src/deepseek4/deepseek4_loader.cpp index ef3aa9fbf..dc119554a 100644 --- a/server/src/deepseek4/deepseek4_loader.cpp +++ b/server/src/deepseek4/deepseek4_loader.cpp @@ -886,10 +886,7 @@ 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.buf) { ggml_backend_buffer_free(w.buf); w.buf = nullptr; } - w.layers.clear(); - w.embedder.tok_embd_owned.clear(); - w.embedder.tok_embd_bytes = nullptr; - w.moe_hybrid = false; + w = {}; } } // namespace dflash::common