diff --git a/server/deps/llama.cpp/ggml/include/ggml-rpc.h b/server/deps/llama.cpp/ggml/include/ggml-rpc.h index 3f7ab5859..d6864ed79 100644 --- a/server/deps/llama.cpp/ggml/include/ggml-rpc.h +++ b/server/deps/llama.cpp/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 3 #define RPC_PROTO_MINOR_VERSION 6 -#define RPC_PROTO_PATCH_VERSION 2 +#define RPC_PROTO_PATCH_VERSION 3 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/server/deps/llama.cpp/ggml/include/ggml.h b/server/deps/llama.cpp/ggml/include/ggml.h index d76166f16..e83cfaa93 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,14 @@ 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_DS4_INDEXER_QAT, // DS4 indexer Hadamard + FP4 activation simulation + + GGML_OP_DS4_INDEXER_SCORE, // Fused DS4 indexer dot/ReLU/head reduction + + GGML_OP_DS4_INDEXER_MASK, // Apply per-token DS4 indexer top-k to an attention mask + GGML_OP_COUNT, }; @@ -618,6 +644,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 +1366,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, @@ -1419,6 +1452,20 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * b); + // Matrix multiply where b is physically [K/group, N, group] but is + // consumed as logical [K, N]. CUDA/HIP MMQ fuses the gather into its Q8 + // activation quantizer; this avoids materializing a group/token permute. + GGML_API struct ggml_tensor * ggml_mul_mat_grouped_src( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API bool ggml_mul_mat_is_grouped_src( + const struct ggml_tensor * tensor); + + GGML_API int64_t ggml_mul_mat_grouped_src_groups( + const struct ggml_tensor * tensor); + // change the precision of a matrix multiplication // set to GGML_PREC_F32 for higher precision (useful for phi-2) GGML_API void ggml_mul_mat_set_prec( @@ -2350,6 +2397,35 @@ extern "C" { struct ggml_tensor * a, enum ggml_prec prec); + // DS4 layout and block-sparse policy for flash_attn_ext. raw_window is the + // maximum visible span inside the raw-row region. Compressed rows are + // selected in fixed-size blocks, capped to keep_rows. Zero leaves the + // operation dense. A negative value means the mask already contains an + // exact row selection and abs(keep_rows) is its maximum compressed width. + GGML_API void ggml_flash_attn_ext_set_ds4_sparse( + struct ggml_tensor * a, + int raw_rows, + int raw_window, + int keep_rows, + int block_size); + + // Fuse DS4's inverse 64-d tail RoPE into the D=512 flash-attention + // writeback. q_unrotated additionally asks the kernel to apply the forward + // tail RoPE to Q from shared F32. This is exact-only plumbing: both paths + // retain the explicit F32 rounding boundary and YaRN arithmetic used by + // GGML_OP_ROPE. + GGML_API void ggml_flash_attn_ext_set_ds4_inverse_rope( + struct ggml_tensor * a, + int kv_start, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + int n_ctx_orig, + bool q_unrotated); + GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec( const struct ggml_tensor * a); @@ -2399,6 +2475,72 @@ extern "C" { struct ggml_tensor * experts, struct ggml_tensor * expert_weights); + // Fused DeepSeek4 hyper-connection helpers. The first dimension is the + // per-token payload; an optional second dimension batches independent + // tokens without changing their arithmetic. + // ggml_ds4_hc_pre: mix[mix_dim,n_tokens] + base[mix_dim] + + // hc_state[n_embd*n_hc,n_tokens] -> + // dst[n_embd+mix_dim,n_tokens] = { working, split(pre,post,comb) } + GGML_API struct ggml_tensor * ggml_ds4_hc_pre( + struct ggml_context * ctx, + struct ggml_tensor * mix, + 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; + // all non-base tensors may carry the same n_tokens second dimension. + GGML_API struct ggml_tensor * ggml_ds4_hc_post( + struct ggml_context * ctx, + struct ggml_tensor * residual_hc, + 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 per + // token. mix and hc_state may carry the same n_tokens second dimension. + GGML_API struct ggml_tensor * ggml_ds4_hc_out( + struct ggml_context * ctx, + struct ggml_tensor * mix, + struct ggml_tensor * base, + struct ggml_tensor * hc_state, + int n_hc, + float pre_scale); + + // Official DS4 ratio-4 indexer transform. Each contiguous 128-wide F32 + // row is Hadamard-rotated and passed through the model's blockwise FP4 + // activation-simulation round trip. The operation is out-of-place so the + // pre-QAT query remains available to the main attention path. + GGML_API struct ggml_tensor * ggml_ds4_indexer_qat( + struct ggml_context * ctx, + struct ggml_tensor * input); + + // Compute the official ratio-4 indexer score matrix without materializing + // [n_comp,n_head,n_tokens] per-head dots. q is F32 + // [128,n_head,n_tokens], head_weights is F32 [n_head,n_tokens], and + // index_comp is F16 [128,n_comp]. head_weights already includes the + // model's 1/sqrt(128*n_head) scale. + GGML_API struct ggml_tensor * ggml_ds4_indexer_score( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * head_weights, + struct ggml_tensor * index_comp, + int kv_start, + int ratio); + + // Preserve the raw rows of base_mask and retain only selected compressed + // rows. selected is I32 [top_k,n_tokens], indexing the compressed span; + // base_mask is F32 [raw_rows+n_comp,n_tokens]. + GGML_API struct ggml_tensor * ggml_ds4_indexer_mask( + struct ggml_context * ctx, + struct ggml_tensor * base_mask, + struct ggml_tensor * selected, + int raw_rows); + // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, diff --git a/server/deps/llama.cpp/ggml/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..b5dea2e11 --- /dev/null +++ b/server/deps/llama.cpp/ggml/rocmfp4/rocmfp4_hip_codebook.cuh @@ -0,0 +1,101 @@ +#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 + +#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)); +#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 rocmfp4_table16_fallback(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 rocmfp4_table16_fallback(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-backend-meta.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp index 8e4273212..5c8604854 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp +++ b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp @@ -962,6 +962,16 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(co case GGML_OP_GATED_DELTA_NET: { split_state = handle_gated_delta_net(src_ss); } break; + case GGML_OP_DS4_INDEXER_QAT: { + split_state = handle_per_row(src_ss); + } break; + case GGML_OP_DS4_INDEXER_SCORE: + case GGML_OP_DS4_INDEXER_MASK: { + // These fused DS4 indexer operations combine axes from + // multiple inputs. Keep them local unless every source is + // mirrored on each simple backend. + split_state = handle_generic(src_ss, /*scalar_only =*/ true); + } break; case GGML_OP_UNARY: { split_state = handle_generic(src_ss, /*scalar_only =*/ false); } break; @@ -1922,4 +1932,3 @@ ggml_backend_t ggml_backend_meta_simple_backend(ggml_backend_t meta_backend, siz const ggml_backend_meta_context * backend_ctx = (const ggml_backend_meta_context *) meta_backend->context; return backend_ctx->backend_configs[index].backend; } - diff --git a/server/deps/llama.cpp/ggml/src/ggml-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..ab5eaad7a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c +++ b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c @@ -78,6 +78,168 @@ float ggml_table_f32_f16[1 << 16]; // precomputed f32 table for e8m0 half (1 KB) (simd-mappings.h) float ggml_table_f32_e8m0_half[1 << 8]; +static float ggml_ds4_indexer_e2m1_value_cpu(int i) { + static const float values[8] = { + 0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f, + }; + return values[i & 7]; +} + +static float ggml_ds4_indexer_e2m1_round_cpu(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float ax = fminf(fabsf(x), 6.0f); + int best = 0; + float best_diff = fabsf(ax - ggml_ds4_indexer_e2m1_value_cpu(0)); + for (int i = 1; i < 8; ++i) { + const float diff = fabsf(ax - ggml_ds4_indexer_e2m1_value_cpu(i)); + if (diff < best_diff || + (diff == best_diff && (i & 1) == 0 && (best & 1) != 0)) { + best = i; + best_diff = diff; + } + } + return sign * ggml_ds4_indexer_e2m1_value_cpu(best); +} + +static void ggml_compute_forward_ds4_indexer_qat( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + const struct ggml_tensor * src = dst->src[0]; + GGML_ASSERT(src && src->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src->ne[0] == 128 && ggml_are_same_shape(src, dst)); + GGML_ASSERT(ggml_is_contiguous(src) && ggml_is_contiguous(dst)); + + const int64_t n_rows = ggml_nrows(src); + for (int64_t row = params->ith; row < n_rows; row += params->nth) { + float values[128]; + const float * src_row = (const float *) ((const char *) src->data + row * src->nb[1]); + float * dst_row = (float *) ((char *) dst->data + row * dst->nb[1]); + memcpy(values, src_row, sizeof(values)); + + for (int stride = 1; stride < 128; stride <<= 1) { + for (int base = 0; base < 128; base += 2 * stride) { + for (int i = 0; i < stride; ++i) { + const float a = values[base + i]; + const float b = values[base + stride + i]; + values[base + i] = a + b; + values[base + stride + i] = a - b; + } + } + } + for (int i = 0; i < 128; ++i) { + values[i] *= 0.08838834764831845f; + } + for (int block = 0; block < 4; ++block) { + float amax = 0.0f; + for (int i = 0; i < 32; ++i) { + amax = fmaxf(amax, fabsf(values[block * 32 + i])); + } + amax = fmaxf(amax, 7.052966104933725e-38f); + const float scale = exp2f(ceilf(log2f(amax / 6.0f))); + for (int i = 0; i < 32; ++i) { + const int index = block * 32 + i; + const float normalized = fminf( + 6.0f, fmaxf(-6.0f, values[index] / scale)); + dst_row[index] = + ggml_ds4_indexer_e2m1_round_cpu(normalized) * scale; + } + } + } +} + +static void ggml_compute_forward_ds4_indexer_score( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + const struct ggml_tensor * q = dst->src[0]; + const struct ggml_tensor * weights = dst->src[1]; + const struct ggml_tensor * comp = dst->src[2]; + GGML_ASSERT(q && weights && comp); + GGML_ASSERT(q->type == GGML_TYPE_F32 && q->ne[0] == 128); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(comp->type == GGML_TYPE_F16 && comp->ne[0] == 128); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(q) && ggml_is_contiguous(weights)); + GGML_ASSERT(ggml_is_contiguous(comp) && ggml_is_contiguous(dst)); + + const int n_head = (int) q->ne[1]; + const int n_tokens = (int) q->ne[2]; + const int n_comp = (int) comp->ne[1]; + const int kv_start = ggml_get_op_params_i32(dst, 0); + const int ratio = ggml_get_op_params_i32(dst, 1); + GGML_ASSERT(weights->ne[0] == n_head && weights->ne[1] == n_tokens); + GGML_ASSERT(dst->ne[0] == n_comp && dst->ne[1] == n_tokens); + GGML_ASSERT(kv_start >= 0 && ratio > 0); + + const float * q_data = (const float *) q->data; + const float * weight_data = (const float *) weights->data; + const ggml_fp16_t * comp_data = (const ggml_fp16_t *) comp->data; + float * dst_data = (float *) dst->data; + for (int token = params->ith; token < n_tokens; token += params->nth) { + const int visible = (kv_start + token + 1) / ratio; + for (int c = 0; c < n_comp; ++c) { + if (c >= visible) { + dst_data[(size_t) token * n_comp + c] = -1.0e30f; + continue; + } + const ggml_fp16_t * k = comp_data + (size_t) c * 128; + float score = 0.0f; + for (int h = 0; h < n_head; ++h) { + const float * qh = q_data + + ((size_t) token * n_head + h) * 128; + float dot = 0.0f; + for (int d = 0; d < 128; ++d) { + dot += qh[d] * GGML_FP16_TO_FP32(k[d]); + } + score += fmaxf(dot, 0.0f) * + weight_data[(size_t) token * n_head + h]; + } + dst_data[(size_t) token * n_comp + c] = score; + } + } +} + +static void ggml_compute_forward_ds4_indexer_mask( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + const struct ggml_tensor * base = dst->src[0]; + const struct ggml_tensor * selected = dst->src[1]; + GGML_ASSERT(base && selected); + GGML_ASSERT(base->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + GGML_ASSERT(selected->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_are_same_shape(base, dst)); + GGML_ASSERT(ggml_is_contiguous(base) && ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_is_contiguous(selected)); + + const int n_attn = (int) base->ne[0]; + const int n_tokens = (int) ggml_nrows(base); + const int top_k = (int) selected->ne[0]; + const int raw_rows = ggml_get_op_params_i32(dst, 0); + const int n_comp = n_attn - raw_rows; + GGML_ASSERT(raw_rows >= 0 && n_comp >= 0); + GGML_ASSERT(ggml_nrows(selected) == n_tokens); + + for (int token = params->ith; token < n_tokens; token += params->nth) { + const float * base_row = + (const float *) ((const char *) base->data + (size_t) token * base->nb[1]); + float * dst_row = + (float *) ((char *) dst->data + (size_t) token * dst->nb[1]); + const int32_t * selected_row = + (const int32_t *) ((const char *) selected->data + + (size_t) token * selected->nb[1]); + memcpy(dst_row, base_row, (size_t) raw_rows * sizeof(float)); + for (int row = raw_rows; row < n_attn; ++row) { + dst_row[row] = -1.0e30f; + } + for (int k = 0; k < top_k; ++k) { + const int comp = selected_row[k]; + if (comp >= 0 && comp < n_comp) { + dst_row[raw_rows + comp] = base_row[raw_rows + comp]; + } + } + } +} + #if defined(__ARM_ARCH) struct ggml_arm_arch_features_type { int sve_cnt; @@ -1826,6 +1988,22 @@ 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_DS4_INDEXER_QAT: + { + ggml_compute_forward_ds4_indexer_qat(params, tensor); + } break; + case GGML_OP_DS4_INDEXER_SCORE: + { + ggml_compute_forward_ds4_indexer_score(params, tensor); + } break; + case GGML_OP_DS4_INDEXER_MASK: + { + ggml_compute_forward_ds4_indexer_mask(params, tensor); + } break; case GGML_OP_OUT_PROD: { ggml_compute_forward_out_prod(params, tensor); @@ -2369,6 +2547,9 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_TIMESTEP_EMBEDDING: case GGML_OP_ARGSORT: case GGML_OP_TOP_K: + case GGML_OP_DS4_INDEXER_QAT: + case GGML_OP_DS4_INDEXER_SCORE: + case GGML_OP_DS4_INDEXER_MASK: case GGML_OP_FLASH_ATTN_EXT: case GGML_OP_FLASH_ATTN_SPARSE: case GGML_OP_FLASH_ATTN_BACK: 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/argsort.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu index 0f3f017b5..ac0634a16 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cu @@ -6,6 +6,9 @@ # define STRIDED_ITERATOR_AVAILABLE # endif using namespace cub; +#elif defined(GGML_CUDA_USE_HIPCUB) +# include +using namespace hipcub; #endif // GGML_CUDA_USE_CUB static __global__ void init_indices(int * indices, const int ncols, const int nrows) { @@ -26,7 +29,7 @@ static __global__ void init_offsets(int * offsets, const int ncols, const int nr } #endif // STRIDED_ITERATOR_AVAILABLE -#ifdef GGML_CUDA_USE_CUB +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const float * x, int * dst, @@ -138,7 +141,7 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, } } } -#endif // GGML_CUDA_USE_CUB +#endif // GGML_CUDA_USE_CUB || GGML_CUDA_USE_HIPCUB // Bitonic sort implementation template @@ -248,7 +251,7 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { enum ggml_sort_order order = (enum ggml_sort_order) dst->op_params[0]; -#ifdef GGML_CUDA_USE_CUB +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) const int ncols_pad = next_power_of_2(ncols); const size_t shared_mem = ncols_pad * sizeof(int); const size_t max_shared_mem = ggml_cuda_info().devices[ggml_cuda_get_device()].smpb; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh index 22b7306f2..018dc9fb2 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/argsort.cuh @@ -2,7 +2,7 @@ void ggml_cuda_op_argsort(ggml_backend_cuda_context & ctx, ggml_tensor * dst); -#ifdef GGML_CUDA_USE_CUB +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const float * x, int * dst, @@ -10,7 +10,7 @@ void argsort_f32_i32_cuda_cub(ggml_cuda_pool & pool, const int nrows, ggml_sort_order order, cudaStream_t stream); -#endif // GGML_CUDA_USE_CUB +#endif // GGML_CUDA_USE_CUB || GGML_CUDA_USE_HIPCUB void argsort_f32_i32_cuda_bitonic(const float * x, int * dst, const int ncols, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh index 9df1461a5..ddcb37c4b 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) @@ -107,6 +129,13 @@ # define GGML_CUDA_USE_CUB #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +// rocPRIM ships hipCUB, a CUB-compatible device-sort API. The ordinary +// single-block bitonic argsort cannot represent DS4's long-context indexer once +// the number of compressed rows exceeds 1024, so use hipCUB for that case. +#if defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +# define GGML_CUDA_USE_HIPCUB +#endif // defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + #ifdef __CUDA_ARCH_LIST__ constexpr bool ggml_cuda_has_arch_impl(int) { return false; @@ -959,6 +988,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 +1537,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..774b6e099 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,33 @@ 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; + case GGML_TYPE_I32: + concat_cuda_typed(ctx, src0, src1, dst, dim); + break; + default: + GGML_ABORT("unsupported concat type %s", ggml_type_name(src0->type)); + } +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/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..77ad8005a --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu @@ -0,0 +1,500 @@ +#include "ds4-hc.cuh" + +// Fused DeepSeek4 hyper-connection ops. +// +// mode 0 (pre): src0 = mix [mix_dim,n_tokens] (f32, from fn @ rms_norm(hc_state)) +// src1 = base [mix_dim] (f32) +// src2 = hc_state [n_embd*n_hc,n_tokens] (f32, raw residual streams) +// dst = [n_embd + mix_dim,n_tokens]: +// dst[0..n_embd) = working vector (pre-mixed input) +// dst[n_embd..n_embd+mix) = split = {pre[n_hc], post[n_hc], comb[n_hc*n_hc]} +// Math matches cpu_hc_sinkhorn + finish_hc_pre_from_mix_into in +// deepseek4_graph.cpp (sigmoid gates + Sinkhorn-normalized combine). +// +// mode 1 (post): src0 = residual hc_state [n_embd*n_hc,n_tokens] +// src1 = block_out [n_embd,n_tokens] +// src2 = split [mix_dim,n_tokens] (view of a mode-0 dst tail) +// dst = new hc_state [n_embd*n_hc,n_tokens]: +// dst[h*n_embd+d] = post[h]*block_out[d] +// + sum_src comb[h + src*n_hc] * residual[src*n_embd+d] +// +// mode 2 (out): src0 = mix [n_hc,n_tokens] +// src1 = base [n_hc] +// src2 = hc_state [n_embd*n_hc,n_tokens] +// dst = [n_embd,n_tokens]: weights[h] = sigmoid(mix[h]*s0+base[h]) + 1e-6; +// dst[d] = sum_h weights[h]*hc_state[h*n_embd+d] + +#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, + size_t mix_stride, + size_t hc_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + mix += (size_t) token * mix_stride; + hc_state += (size_t) token * hc_stride; + dst += (size_t) token * dst_stride; + + __shared__ float split[DS4_HC_MAX_MIX]; + __shared__ float s_mix[DS4_HC_MAX_MIX]; + __shared__ float s_base[DS4_HC_MAX_MIX]; + 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; + } +} + +// Large prefill batches already expose thousands of token blocks, so they do +// not need every embedding tile to recompute the same serial Sinkhorn. Split +// it into one deterministic solve per token followed by the parallel mixing +// pass. The small-token path keeps the fused kernel above to avoid an extra +// launch during decode and speculative verification. +template +static __global__ void ds4_hc_pre_split_kernel_t( + const float * __restrict__ mix, + const float * __restrict__ base, + float * __restrict__ dst, + int n_embd, + int iters, + float pre_scale, + float post_scale, + float comb_scale, + size_t mix_stride, + size_t dst_stride) { + const int token = (int) blockIdx.x; + mix += (size_t) token * mix_stride; + dst += (size_t) token * dst_stride; + + __shared__ float split[DS4_HC_MAX_MIX]; + __shared__ float s_mix[DS4_HC_MAX_MIX]; + __shared__ float s_base[DS4_HC_MAX_MIX]; + constexpr int mix_dim = 2 * NHC + NHC * NHC; + const int tid = (int) threadIdx.x; + + if (tid < mix_dim) { + s_mix[tid] = mix[tid]; + s_base[tid] = base[tid]; + } + __syncthreads(); + + if (tid == 0) { + ds4_hc_sinkhorn_split_t( + s_mix, s_base, pre_scale, post_scale, comb_scale, iters, split); +#pragma unroll + for (int i = 0; i < mix_dim; ++i) { + dst[n_embd + i] = split[i]; + } + } +} + +template +static __global__ void ds4_hc_pre_mix_kernel_t( + const float * __restrict__ hc_state, + float * __restrict__ dst, + int n_embd, + size_t hc_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + hc_state += (size_t) token * hc_stride; + dst += (size_t) token * dst_stride; + + __shared__ float pre[NHC]; + const int tid = (int) threadIdx.x; + if (tid < NHC) { + pre[tid] = dst[n_embd + tid]; + } + __syncthreads(); + + const int d = (int) blockIdx.x * (int) blockDim.x + tid; + if (d < n_embd) { + float acc = 0.0f; +#pragma unroll + for (int h = 0; h < NHC; ++h) { + acc += pre[h] * hc_state[(size_t) h * n_embd + d]; + } + dst[d] = acc; + } +} + +static __global__ void ds4_hc_pre_kernel( + const float * __restrict__ mix, + const float * __restrict__ base, + 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, + size_t mix_stride, + size_t hc_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + mix += (size_t) token * mix_stride; + hc_state += (size_t) token * hc_stride; + dst += (size_t) token * dst_stride; + + __shared__ float split[DS4_HC_MAX_MIX]; + __shared__ float s_mix[DS4_HC_MAX_MIX]; + __shared__ float s_base[DS4_HC_MAX_MIX]; + 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, + size_t residual_stride, + size_t block_out_stride, + size_t split_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + residual += (size_t) token * residual_stride; + block_out += (size_t) token * block_out_stride; + split += (size_t) token * split_stride; + dst += (size_t) token * dst_stride; + + const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int total = n_embd * n_hc; + if (i >= total) { + 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, + size_t mix_stride, + size_t hc_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + mix += (size_t) token * mix_stride; + hc_state += (size_t) token * hc_stride; + dst += (size_t) token * dst_stride; + + const int d = blockIdx.x * blockDim.x + threadIdx.x; + if (d >= n_embd) { + return; + } + 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); + const int n_tokens = (int) dst->ne[1]; + + GGML_ASSERT(n_hc > 0 && n_hc <= DS4_HC_MAX_HC); + GGML_ASSERT(n_tokens > 0); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(src1->nb[0] == sizeof(float)); + GGML_ASSERT(src2->nb[0] == sizeof(float)); + GGML_ASSERT(dst->nb[0] == sizeof(float)); + + cudaStream_t stream = ctx.stream(); + + 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; + const dim3 grid(pre_blocks, n_tokens, 1); + if (n_hc == 4 && n_tokens >= 64) { + ds4_hc_pre_split_kernel_t<4><<>>( + (const float *) src0->data, (const float *) src1->data, + (float *) dst->data, + n_embd, iters, pre_scale, post_scale, comb_scale, + src0->nb[1] / sizeof(float), dst->nb[1] / sizeof(float)); + ds4_hc_pre_mix_kernel_t<4><<>>( + (const float *) src2->data, (float *) dst->data, + n_embd, src2->nb[1] / sizeof(float), + dst->nb[1] / sizeof(float)); + } else if (n_hc == 4) { + ds4_hc_pre_kernel_t<4><<>>( + (const float *) src0->data, (const float *) src1->data, + (const float *) src2->data, (float *) dst->data, + n_embd, iters, pre_scale, post_scale, comb_scale, + src0->nb[1] / sizeof(float), src2->nb[1] / sizeof(float), + dst->nb[1] / sizeof(float)); + } 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, + src0->nb[1] / sizeof(float), src2->nb[1] / sizeof(float), + dst->nb[1] / sizeof(float)); + } + } break; + case 1: { + const int total = n_embd * n_hc; + const int blocks = (total + 255) / 256; + const dim3 grid(blocks, n_tokens, 1); + ds4_hc_post_kernel<<>>( + (const float *) src0->data, (const float *) src1->data, + (const float *) src2->data, (float *) dst->data, + n_embd, n_hc, + src0->nb[1] / sizeof(float), src1->nb[1] / sizeof(float), + src2->nb[1] / sizeof(float), dst->nb[1] / sizeof(float)); + } break; + case 2: { + const float pre_scale = ggml_get_op_params_f32(dst, 4); + const int blocks = (n_embd + 255) / 256; + const dim3 grid(blocks, n_tokens, 1); + ds4_hc_out_kernel<<>>( + (const float *) src0->data, (const float *) src1->data, + (const float *) src2->data, (float *) dst->data, + n_embd, n_hc, pre_scale, + src0->nb[1] / sizeof(float), src2->nb[1] / sizeof(float), + dst->nb[1] / sizeof(float)); + } 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/ds4-indexer.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cu new file mode 100644 index 000000000..f4358a775 --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cu @@ -0,0 +1,409 @@ +#include "ds4-indexer.cuh" + +#if defined(GGML_USE_HIP) +# include +namespace ds4_wmma = rocwmma; +# define DS4_INDEXER_WMMA_AVAILABLE 1 +#elif !defined(GGML_USE_MUSA) +# include +namespace ds4_wmma = nvcuda::wmma; +# define DS4_INDEXER_WMMA_AVAILABLE 1 +#else +# define DS4_INDEXER_WMMA_AVAILABLE 0 +#endif + +#if DS4_INDEXER_WMMA_AVAILABLE +# if defined(GGML_USE_HIP) && HIP_VERSION >= 60500000 +using ds4_indexer_wmma_half = _Float16; +# else +using ds4_indexer_wmma_half = half; +# endif +#endif + +// Keep this operation bit-for-bit aligned with the official DeepSeek V4 +// graph and antirez/ds4's dsv4_indexer_qat implementation. Indexer query and +// compressed-key rows both pass through an orthonormal Hadamard-128 followed +// by one UE4M3-scaled E2M1 activation-simulation block per 32 values. + +static __device__ __forceinline__ float ds4_indexer_e2m1_value(int i) { + switch (i & 7) { + case 0: return 0.0f; + case 1: return 0.5f; + case 2: return 1.0f; + case 3: return 1.5f; + case 4: return 2.0f; + case 5: return 3.0f; + case 6: return 4.0f; + default: return 6.0f; + } +} + +static __device__ __forceinline__ float ds4_indexer_e2m1_round(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float ax = fminf(fabsf(x), 6.0f); + int best = 0; + float best_diff = fabsf(ax - ds4_indexer_e2m1_value(0)); +#pragma unroll + for (int i = 1; i < 8; ++i) { + const float diff = fabsf(ax - ds4_indexer_e2m1_value(i)); + // Round ties to the even E2M1 code, matching the converter/reference. + if (diff < best_diff || + (diff == best_diff && (i & 1) == 0 && (best & 1) != 0)) { + best = i; + best_diff = diff; + } + } + return sign * ds4_indexer_e2m1_value(best); +} + +static __global__ void ds4_indexer_qat_kernel( + float * dst, + const float * src, + int64_t n_rows, + int64_t src_row_stride, + int64_t dst_row_stride) { + constexpr int WIDTH = 128; + constexpr float HADAMARD_SCALE = 0.08838834764831845f; + const int64_t row = (int64_t) blockIdx.x; + const int tid = (int) threadIdx.x; + if (row >= n_rows || tid >= WIDTH) return; + + __shared__ float values[WIDTH]; + __shared__ float abs_values[WIDTH]; + const float * src_row = src + row * src_row_stride; + float * dst_row = dst + row * dst_row_stride; + values[tid] = src_row[tid]; + __syncthreads(); + + for (int stride = 1; stride < WIDTH; stride <<= 1) { + if ((tid & stride) == 0) { + const int base = + (tid & ~(2 * stride - 1)) + (tid & (stride - 1)); + const float a = values[base]; + const float b = values[base + stride]; + values[base] = a + b; + values[base + stride] = a - b; + } + __syncthreads(); + } + + const float value = values[tid] * HADAMARD_SCALE; + const int block = tid >> 5; + const int lane = tid & 31; + const int block_base = block * 32; + abs_values[tid] = fabsf(value); + __syncthreads(); + + for (int stride = 16; stride > 0; stride >>= 1) { + if (lane < stride) { + abs_values[block_base + lane] = fmaxf( + abs_values[block_base + lane], + abs_values[block_base + lane + stride]); + } + __syncthreads(); + } + + const float amax = fmaxf( + abs_values[block_base], 7.052966104933725e-38f); + const float scale = exp2f(ceilf(log2f(amax / 6.0f))); + const float normalized = fminf(6.0f, fmaxf(-6.0f, value / scale)); + dst_row[tid] = ds4_indexer_e2m1_round(normalized) * scale; +} + +void ggml_cuda_op_ds4_indexer_qat( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst) { + const ggml_tensor * src = dst->src[0]; + GGML_ASSERT(src && src->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(src->ne[0] == 128 && dst->ne[0] == 128); + GGML_ASSERT(ggml_are_same_shape(src, dst)); + GGML_ASSERT(ggml_is_contiguous(src)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int64_t n_rows = ggml_nrows(src); + const int64_t src_row_stride = src->nb[1] / sizeof(float); + const int64_t dst_row_stride = dst->nb[1] / sizeof(float); + GGML_ASSERT(src_row_stride >= 128 && dst_row_stride >= 128); + + cudaStream_t stream = ctx.stream(); + ds4_indexer_qat_kernel<<<(unsigned) n_rows, 128, 0, stream>>>( + static_cast(dst->data), + static_cast(src->data), + n_rows, src_row_stride, dst_row_stride); + CUDA_CHECK(cudaGetLastError()); +} + +// Compute 16 query tokens against 128 compressed rows per block. QAT values +// are powers-of-two-scaled E2M1 and therefore exactly representable as F16 in +// the model's operating range; the compressed cache is already F16. WMMA +// removes the otherwise enormous [n_comp,64,n_tokens] intermediate while the +// ReLU, head weighting and reduction remain F32. +#if DS4_INDEXER_WMMA_AVAILABLE +static __global__ void ds4_indexer_score_wmma_kernel( + float * scores, + const float * q, + const float * weights, + const half * index_comp, + int n_comp, + int n_tokens, + int kv_start, + int n_head, + int ratio) { + const int tile_c = (int) blockIdx.x * 128; + const int tile_t = (int) blockIdx.y * 16; + const int tid = (int) threadIdx.x; + const int warp = tid >> 5; + + __shared__ half a_sh[16 * 128]; + __shared__ half b_sh[128 * 128]; + __shared__ float c_sh[8 * 16 * 16]; + + float acc[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) acc[i] = 0.0f; + + for (int i = tid; i < 128 * 128; i += 256) { + const int c = i >> 7; + const int d = i & 127; + const int comp = tile_c + c; + b_sh[d + c * 128] = comp < n_comp + ? index_comp[(size_t) comp * 128 + d] + : __float2half(0.0f); + } + __syncthreads(); + + for (int h = 0; h < n_head; ++h) { + for (int pair = tid; pair < 16 * 64; pair += 256) { + const int row = pair >> 6; + const int d = (pair & 63) * 2; + const int token = tile_t + row; + half2 value = __float2half2_rn(0.0f); + if (token < n_tokens) { + const float2 q_value = *reinterpret_cast( + q + ((size_t) token * n_head + h) * 128 + d); + value = __floats2half2_rn(q_value.x, q_value.y); + } + *reinterpret_cast(a_sh + row * 128 + d) = value; + } + __syncthreads(); + + ds4_wmma::fragment a_frag; + ds4_wmma::fragment b_frag; + ds4_wmma::fragment c_frag; + ds4_wmma::fill_fragment(c_frag, 0.0f); + const int col0 = warp * 16; + for (int k0 = 0; k0 < 128; k0 += 16) { + const ds4_indexer_wmma_half * a_wmma = + reinterpret_cast(a_sh); + const ds4_indexer_wmma_half * b_wmma = + reinterpret_cast(b_sh); + ds4_wmma::load_matrix_sync(a_frag, a_wmma + k0, 128); + ds4_wmma::load_matrix_sync( + b_frag, b_wmma + col0 * 128 + k0, 128); + ds4_wmma::mma_sync(c_frag, a_frag, b_frag, c_frag); + } + ds4_wmma::store_matrix_sync( + c_sh + warp * 16 * 16, c_frag, 16, + ds4_wmma::mem_row_major); + __syncthreads(); + + const int token_for_lane = tile_t + (tid >> 4); + const float head_weight = token_for_lane < n_tokens + ? weights[(size_t) token_for_lane * n_head + h] + : 0.0f; + int slot = 0; + for (int i = tid; i < 8 * 16 * 16; i += 256, ++slot) { + acc[slot] += fmaxf(c_sh[i], 0.0f) * head_weight; + } + __syncthreads(); + } + + int slot = 0; + for (int i = tid; i < 8 * 16 * 16; i += 256, ++slot) { + const int wtile = i >> 8; + const int local = i & 255; + const int row = local >> 4; + const int col = local & 15; + const int token = tile_t + row; + const int comp = tile_c + wtile * 16 + col; + if (token < n_tokens && comp < n_comp) { + const int visible = (kv_start + token + 1) / ratio; + scores[(size_t) token * n_comp + comp] = + comp < visible ? acc[slot] : -1.0e30f; + } + } +} +#endif + +static __global__ void ds4_indexer_score_scalar_kernel( + float * scores, + const float * q, + const float * weights, + const half * index_comp, + int n_comp, + int n_tokens, + int kv_start, + int n_head, + int ratio) { + const int comp = (int) blockIdx.x; + const int token = (int) blockIdx.y; + const int tid = (int) threadIdx.x; + if (comp >= n_comp || token >= n_tokens) return; + const int visible = (kv_start + token + 1) / ratio; + if (comp >= visible) { + if (tid == 0) { + scores[(size_t) token * n_comp + comp] = -1.0e30f; + } + return; + } + + __shared__ float partial[256]; + float total = 0.0f; + const half * k = index_comp + (size_t) comp * 128; + for (int h = 0; h < n_head; ++h) { + const float * qh = q + ((size_t) token * n_head + h) * 128; + float dot = 0.0f; + for (int d = tid; d < 128; d += 256) { + dot += qh[d] * __half2float(k[d]); + } + partial[tid] = dot; + __syncthreads(); + for (int stride = 128; stride > 0; stride >>= 1) { + if (tid < stride) partial[tid] += partial[tid + stride]; + __syncthreads(); + } + if (tid == 0) { + total += fmaxf(partial[0], 0.0f) * + weights[(size_t) token * n_head + h]; + } + __syncthreads(); + } + if (tid == 0) scores[(size_t) token * n_comp + comp] = total; +} + +void ggml_cuda_op_ds4_indexer_score( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * weights = dst->src[1]; + const ggml_tensor * comp = dst->src[2]; + GGML_ASSERT(q && weights && comp); + GGML_ASSERT(q->type == GGML_TYPE_F32 && q->ne[0] == 128); + GGML_ASSERT(weights->type == GGML_TYPE_F32); + GGML_ASSERT(comp->type == GGML_TYPE_F16 && comp->ne[0] == 128); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(q)); + GGML_ASSERT(ggml_is_contiguous(weights)); + GGML_ASSERT(ggml_is_contiguous(comp)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int n_head = (int) q->ne[1]; + const int n_tokens = (int) q->ne[2]; + const int n_comp = (int) comp->ne[1]; + const int kv_start = ggml_get_op_params_i32(dst, 0); + const int ratio = ggml_get_op_params_i32(dst, 1); + GGML_ASSERT(weights->ne[0] == n_head && weights->ne[1] == n_tokens); + GGML_ASSERT(dst->ne[0] == n_comp && dst->ne[1] == n_tokens); + + cudaStream_t stream = ctx.stream(); + const int warp_size = + ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; +#if DS4_INDEXER_WMMA_AVAILABLE + if (warp_size == 32) { + const dim3 grid((unsigned) ((n_comp + 127) / 128), + (unsigned) ((n_tokens + 15) / 16), 1); + ds4_indexer_score_wmma_kernel<<>>( + static_cast(dst->data), + static_cast(q->data), + static_cast(weights->data), + static_cast(comp->data), + n_comp, n_tokens, kv_start, n_head, ratio); + } else +#endif + { + const dim3 grid((unsigned) n_comp, (unsigned) n_tokens, 1); + ds4_indexer_score_scalar_kernel<<>>( + static_cast(dst->data), + static_cast(q->data), + static_cast(weights->data), + static_cast(comp->data), + n_comp, n_tokens, kv_start, n_head, ratio); + } + CUDA_CHECK(cudaGetLastError()); +} + +static __global__ void ds4_indexer_mask_init_kernel( + float * dst, + const float * base, + int64_t n_elements, + int n_attn, + int raw_rows) { + const int64_t index = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (index >= n_elements) return; + const int row = (int) (index % n_attn); + dst[index] = row < raw_rows ? base[index] : -1.0e30f; +} + +static __global__ void ds4_indexer_mask_scatter_kernel( + float * dst, + const float * base, + const int32_t * selected, + int64_t n_selected, + int n_attn, + int raw_rows, + int n_comp, + int top_k) { + const int64_t index = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (index >= n_selected) return; + const int token = (int) (index / top_k); + const int comp = selected[index]; + if (comp < 0 || comp >= n_comp) return; + const int64_t output_index = + (int64_t) token * n_attn + raw_rows + comp; + // Preserve the base causal mask: when a row is not yet visible, top-k may + // still return it only as an -inf filler for tokens with < k live rows. + dst[output_index] = base[output_index]; +} + +void ggml_cuda_op_ds4_indexer_mask( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst) { + const ggml_tensor * base = dst->src[0]; + const ggml_tensor * selected = dst->src[1]; + GGML_ASSERT(base && selected); + GGML_ASSERT(base->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + GGML_ASSERT(selected->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_are_same_shape(base, dst)); + GGML_ASSERT(ggml_is_contiguous(base) && ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_is_contiguous(selected)); + + const int n_attn = (int) base->ne[0]; + const int top_k = (int) selected->ne[0]; + const int n_tokens = (int) ggml_nrows(base); + const int raw_rows = ggml_get_op_params_i32(dst, 0); + const int n_comp = n_attn - raw_rows; + GGML_ASSERT(raw_rows >= 0 && n_comp >= 0); + GGML_ASSERT(ggml_nrows(selected) == n_tokens); + + const int64_t n_elements = ggml_nelements(base); + const int64_t n_selected = ggml_nelements(selected); + cudaStream_t stream = ctx.stream(); + ds4_indexer_mask_init_kernel<<< + (unsigned) ((n_elements + 255) / 256), 256, 0, stream>>>( + static_cast(dst->data), + static_cast(base->data), + n_elements, n_attn, raw_rows); + ds4_indexer_mask_scatter_kernel<<< + (unsigned) ((n_selected + 255) / 256), 256, 0, stream>>>( + static_cast(dst->data), + static_cast(base->data), + static_cast(selected->data), + n_selected, n_attn, raw_rows, n_comp, top_k); + CUDA_CHECK(cudaGetLastError()); +} diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cuh new file mode 100644 index 000000000..e4f29ea05 --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cuh @@ -0,0 +1,15 @@ +#pragma once + +#include "common.cuh" + +void ggml_cuda_op_ds4_indexer_qat( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst); + +void ggml_cuda_op_ds4_indexer_score( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst); + +void ggml_cuda_op_ds4_indexer_mask( + ggml_backend_cuda_context & ctx, + ggml_tensor * dst); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu index 0eff4f1b7..5def8828b 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,1793 @@ #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]; +} + +template +__device__ static __forceinline__ float ds4_fa_load(const KV * ptr) { + return (float) *ptr; +} + +template <> +__device__ __forceinline__ float ds4_fa_load(const half * ptr) { + return __half2float(*ptr); +} + +template +__device__ static __forceinline__ void ds4_fa_load_pair( + const KV * ptr, float & v0, float & v1) { + v0 = (float) ptr[0]; + v1 = (float) ptr[1]; +} + +template <> +__device__ __forceinline__ void ds4_fa_load_pair( + const float * ptr, float & v0, float & v1) { + const float2 pair = *reinterpret_cast(ptr); + v0 = pair.x; + v1 = pair.y; +} + +template <> +__device__ __forceinline__ void ds4_fa_load_pair( + const half * ptr, float & v0, float & v1) { + const half2 pair = *reinterpret_cast(ptr); + const float2 unpacked = __half22float2(pair); + v0 = unpacked.x; + v1 = unpacked.y; +} + +template +__device__ static __forceinline__ void ds4_fa_load_quad( + const KV * ptr, float & v0, float & v1, float & v2, float & v3) { + v0 = (float) ptr[0]; + v1 = (float) ptr[1]; + v2 = (float) ptr[2]; + v3 = (float) ptr[3]; +} + +template <> +__device__ __forceinline__ void ds4_fa_load_quad( + const float * ptr, float & v0, float & v1, float & v2, float & v3) { + const float4 values = *reinterpret_cast(ptr); + v0 = values.x; + v1 = values.y; + v2 = values.z; + v3 = values.w; +} + +template <> +__device__ __forceinline__ void ds4_fa_load_quad( + const half * ptr, float & v0, float & v1, float & v2, float & v3) { + const float2 lo = __half22float2( + *reinterpret_cast(ptr + 0)); + const float2 hi = __half22float2( + *reinterpret_cast(ptr + 2)); + v0 = lo.x; + v1 = lo.y; + v2 = hi.x; + v3 = hi.y; +} + +struct ds4_inverse_rope_params { + int enabled; + int forward_q_enabled; + int kv_start; + float freq_scale; + float ext_factor; + float attn_factor; + float corr_low; + float corr_high; + float theta_scale; +}; + +// Keep these expressions aligned with rope.cu. The attention result is first +// stored in shared F32, matching the standalone attention-output store/load +// boundary, before the pair is rotated. +__device__ static __forceinline__ float ds4_rope_theta_fp64( + int32_t p, float theta_scale, int exp_int) { + const double tau = 6.2831853071795864769; + double angle = exp_int == 0 + ? (double) p + : (double) p * pow((double) theta_scale, (double) exp_int); + angle -= tau * floor(angle * (1.0 / tau)); + return (float) angle; +} + +__device__ static __forceinline__ void ds4_inverse_rope_coefficients( + int pair, int token, + const ds4_inverse_rope_params & p, + float & cos_theta, float & sin_theta) { + const int i0 = 2 * pair; + const float theta_extrap = ds4_rope_theta_fp64( + -(p.kv_start + token), p.theta_scale, pair); + const float theta_interp = p.freq_scale * theta_extrap; + float theta = theta_interp; + float mscale = p.attn_factor; + if (p.ext_factor != 0.0f) { + const float ramp_y = (i0 / 2 - p.corr_low) / + max(0.001f, p.corr_high - p.corr_low); + const float ramp_mix = + (1.0f - min(1.0f, max(0.0f, ramp_y))) * p.ext_factor; + theta = theta_interp * (1.0f - ramp_mix) + + theta_extrap * ramp_mix; + mscale *= 1.0f + 0.1f * logf(1.0f / p.freq_scale); + } + cos_theta = cosf(theta) * mscale; + sin_theta = sinf(theta) * mscale; +} + +// Forward counterpart of ds4_inverse_rope_coefficients. Keep the expressions +// aligned with rope_norm in rope.cu; unlike the inverse path, position is +// positive. Compressed-layer YaRN interpolation means the inverse coefficients +// cannot safely be recovered by merely negating sin(theta). +__device__ static __forceinline__ void ds4_forward_rope_coefficients( + int pair, int token, + const ds4_inverse_rope_params & p, + float & cos_theta, float & sin_theta) { + const int i0 = 2 * pair; + const float theta_extrap = ds4_rope_theta_fp64( + p.kv_start + token, p.theta_scale, pair); + const float theta_interp = p.freq_scale * theta_extrap; + float theta = theta_interp; + float mscale = p.attn_factor; + if (p.ext_factor != 0.0f) { + const float ramp_y = (i0 / 2 - p.corr_low) / + max(0.001f, p.corr_high - p.corr_low); + const float ramp_mix = + (1.0f - min(1.0f, max(0.0f, ramp_y))) * p.ext_factor; + theta = theta_interp * (1.0f - ramp_mix) + + theta_extrap * ramp_mix; + mscale *= 1.0f + 0.1f * logf(1.0f / p.freq_scale); + } + cos_theta = cosf(theta) * mscale; + sin_theta = sinf(theta) * mscale; +} + +__device__ static __forceinline__ void ds4_apply_inverse_rope_pair( + float x0, float x1, float cos_theta, float sin_theta, + float & y0, float & y1) { + y0 = x0 * cos_theta - x1 * sin_theta; + y1 = x0 * sin_theta + x1 * cos_theta; +} + +// RoPE coefficients depend on token position and pair, not on the query head. +// Materialize them once per attention call instead of recomputing FP64 pow, +// floor, cos and sin in every head block. F32 storage preserves the same +// explicit coefficient rounding used by the original in-kernel calculation. +__global__ static void ds4_inverse_rope_coefficients_kernel( + float * coefficients, + int n_tokens, + ds4_inverse_rope_params inverse_rope) { + const int index = (int) blockIdx.x * (int) blockDim.x + + (int) threadIdx.x; + const int count = n_tokens * 32; + if (index >= count) return; + const int token = index / 32; + const int pair = index % 32; + float cos_theta; + float sin_theta; + ds4_inverse_rope_coefficients( + pair, token, inverse_rope, cos_theta, sin_theta); + coefficients[2 * index + 0] = cos_theta; + coefficients[2 * index + 1] = sin_theta; +} + +__global__ static void ds4_forward_rope_coefficients_kernel( + float * coefficients, + int n_tokens, + ds4_inverse_rope_params rope) { + const int index = (int) blockIdx.x * (int) blockDim.x + + (int) threadIdx.x; + const int count = n_tokens * 32; + if (index >= count) return; + const int token = index / 32; + const int pair = index % 32; + float cos_theta; + float sin_theta; + ds4_forward_rope_coefficients( + pair, token, rope, cos_theta, sin_theta); + coefficients[2 * index + 0] = cos_theta; + coefficients[2 * index + 1] = sin_theta; +} + +// One mean latent-key vector per compressed-cache block. Raw SWA/current rows +// deliberately stay outside this summary and are always evaluated exactly. +template +__global__ static void ds4_fa_mean_comp_blocks_kernel( + const KV * k, + float * mean_k, + int n_kv, + int raw_rows, + int block_size, + int n_blocks) { + constexpr int D = 512; + const int b = (int) blockIdx.x; + if (b >= n_blocks) return; + const int begin = raw_rows + b * block_size; + const int end = min(n_kv, begin + block_size); + const float inv = 1.0f / (float) max(1, end - begin); + for (int d = (int) threadIdx.x; d < D; d += (int) blockDim.x) { + float sum = 0.0f; + for (int r = begin; r < end; ++r) { + sum += ds4_fa_load(k + (size_t) r * D + d); + } + mean_k[(size_t) b * D + d] = sum * inv; + } +} + +// Find the visible envelope in the raw and non-raw regions once per query +// token. Attention blocks for all query-head groups reuse these four bounds. +// Internal masked rows remain inside the envelope and are still evaluated as +// masked, so this changes storage only, not attention semantics. +template +__global__ static void ds4_fa_visibility_bounds_kernel( + const Mask * mask, + int * bounds, + int n_tokens, + int n_kv, + int raw_rows) { + const int t = (int) blockIdx.x; + const int lane = (int) threadIdx.x; + if (t >= n_tokens || lane >= warpSize) return; + + int raw_first = raw_rows; + int raw_last = -1; + int comp_first = n_kv; + int comp_last = -1; + const Mask * token_mask = mask + (size_t) t * n_kv; + + for (int base = 0; base < raw_rows; base += warpSize) { + const int r = base + lane; + const unsigned long long active = __ballot( + r < raw_rows && + ds4_fa_load(token_mask + r) > -1.0e20f); + if (lane == 0 && active != 0) { + if (raw_first == raw_rows) { + raw_first = base + __ffsll(active) - 1; + } + raw_last = base + 63 - __clzll(active); + } + } + for (int base = raw_rows; base < n_kv; base += warpSize) { + const int r = base + lane; + const unsigned long long active = __ballot( + r < n_kv && + ds4_fa_load(token_mask + r) > -1.0e20f); + if (lane == 0 && active != 0) { + if (comp_first == n_kv) { + comp_first = base + __ffsll(active) - 1; + } + comp_last = base + 63 - __clzll(active); + } + } + if (lane == 0) { + int * token_bounds = bounds + (size_t) t * 4; + token_bounds[0] = raw_first; + token_bounds[1] = raw_last; + token_bounds[2] = comp_first; + token_bounds[3] = comp_last; + } +} + +// Convert an externally selected compressed-row mask into exact lookup tables. +// selected_rows preserves ascending physical-row order for the value pass. +// owner_offsets/owner_ranks group those ascending ranks by the thread that +// owned the physical row in the original r = tid + 256*k traversal. The hot +// score and softmax passes can therefore visit only selected rows while +// retaining every thread's original accumulation order and reduction leaf. +template +__global__ static void ds4_fa_indexed_rows_kernel( + const Mask * mask, + int * selected_rows, + int * selected_counts, + int * owner_offsets, + int * owner_ranks, + int n_tokens, + int n_kv, + int raw_rows, + int capacity) { + const int t = (int) blockIdx.x; + const int tid = (int) threadIdx.x; + if (t >= n_tokens) return; + + constexpr int N_THREADS = 256; + __shared__ int owner_write[N_THREADS]; + const int n_comp_rows = n_kv - raw_rows; + int * token_rows = selected_rows + (size_t) t * capacity; + int * token_owner_offsets = owner_offsets + (size_t) t * (N_THREADS + 1); + int * token_owner_ranks = owner_ranks + (size_t) t * capacity; + token_owner_offsets[tid] = 0; + if (tid == 0) token_owner_offsets[N_THREADS] = 0; + __syncthreads(); + + if (tid == 0) { + const Mask * token_mask = mask + (size_t) t * n_kv; + int count = 0; + for (int c = 0; c < n_comp_rows; ++c) { + if (ds4_fa_load(token_mask + raw_rows + c) <= -1.0e20f) { + continue; + } + if (count < capacity) { + const int r = raw_rows + c; + token_rows[count] = r; + ++token_owner_offsets[(r & (N_THREADS - 1)) + 1]; + } + ++count; + } + count = min(count, capacity); + selected_counts[t] = count; + + int prefix = 0; + for (int owner = 0; owner < N_THREADS; ++owner) { + const int owner_count = token_owner_offsets[owner + 1]; + token_owner_offsets[owner] = prefix; + prefix += owner_count; + } + token_owner_offsets[N_THREADS] = prefix; + } + __syncthreads(); + + owner_write[tid] = token_owner_offsets[tid]; + __syncthreads(); + + if (tid == 0) { + const int count = selected_counts[t]; + for (int rank = 0; rank < count; ++rank) { + const int owner = token_rows[rank] & (N_THREADS - 1); + token_owner_ranks[owner_write[owner]++] = rank; + } + } +} + +template +__global__ static void ds4_flash_attn_d512_shared_kv_kernel( + float * dst, + const float * q, + size_t q_stride_token, + size_t q_stride_head, + const KV * k, + const KV * v, + const Mask * mask, + const float * sinks, + const float * mean_k, + int n_tokens, + int n_heads, + int n_kv, + float scale, + int raw_rows, + int sparse_keep_rows, + int sparse_block_size, + int n_comp_blocks, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients, + const float * forward_rope_coefficients) { + constexpr int D = 512; + const int t = (int) blockIdx.x; + const int h = (int) blockIdx.y; + const int tid = (int) threadIdx.x; + if (t >= n_tokens || h >= n_heads) return; + + extern __shared__ float scratch[]; + float * scores = scratch; + float * block_scores = scores + n_kv; + float * block_keep = block_scores + n_comp_blocks; + float * q_rope_tail = block_keep + n_comp_blocks; + const float * qh = q + (size_t) t * q_stride_token + + (size_t) h * q_stride_head; + + if (inverse_rope.forward_q_enabled) { + for (int pair = tid; pair < 32; pair += (int) blockDim.x) { + const float x0 = qh[D - 64 + 2 * pair + 0]; + const float x1 = qh[D - 64 + 2 * pair + 1]; + const size_t coefficient_index = + ((size_t) t * 32 + (size_t) pair) * 2; + const float cos_theta = forward_rope_coefficients[ + coefficient_index + 0]; + const float sin_theta = forward_rope_coefficients[ + coefficient_index + 1]; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, + q_rope_tail[2 * pair + 0], q_rope_tail[2 * pair + 1]); + } + __syncthreads(); + } + + const int n_comp_rows = n_kv - raw_rows; + const bool sparse = mean_k && n_comp_blocks > 0 && + sparse_keep_rows > 0 && + sparse_keep_rows < n_comp_rows; + if (sparse) { + for (int b = tid; b < n_comp_blocks; b += (int) blockDim.x) { + const int first_row = raw_rows + b * sparse_block_size; + const float mask_v = mask + ? ds4_fa_load( + mask + (size_t) t * n_kv + first_row) + : 0.0f; + const float * kb = mean_k + (size_t) b * D; + float dot = -3.402823466e38f; + if (mask_v > -1.0e20f) { + dot = 0.0f; +#pragma unroll + for (int d = 0; d < D; ++d) { + const float qv = inverse_rope.forward_q_enabled && d >= D - 64 + ? q_rope_tail[d - (D - 64)] : qh[d]; + dot += qv * kb[d]; + } + dot *= scale; + } + block_scores[b] = dot; + } + __syncthreads(); + + const int keep_blocks = min(n_comp_blocks, + (sparse_keep_rows + sparse_block_size - 1) / sparse_block_size); + for (int b = tid; b < n_comp_blocks; b += (int) blockDim.x) { + const float score = block_scores[b]; + int rank = 0; + for (int j = 0; j < n_comp_blocks; ++j) { + const float other = block_scores[j]; + rank += (other > score || (other == score && j < b)) ? 1 : 0; + } + block_keep[b] = rank < keep_blocks ? 1.0f : 0.0f; + } + __syncthreads(); + } + + // Keep the original row-to-thread mapping and reduction order. The mask + // remains the sole visibility authority for dense attention; duplicating + // its DS4 policy here risks changing semantics. Masked rows already avoid + // the expensive D=512 dot product. Sparse mode additionally drops selected + // compressed blocks by design. + float local_max = sinks ? sinks[h] : -3.402823466e38f; + for (int r = tid; r < n_kv; r += blockDim.x) { + bool keep = true; + if (sparse && r >= raw_rows) { + const int b = (r - raw_rows) / sparse_block_size; + keep = b < n_comp_blocks && block_keep[b] != 0.0f; + } + const float mask_v = mask + ? ds4_fa_load(mask + (size_t) t * n_kv + r) + : 0.0f; + float s = -3.402823466e38f; + if (keep && mask_v > -1.0e20f) { + const KV * kr = k + (size_t) r * D; + float dot = 0.0f; +#pragma unroll + for (int d = 0; d < D; ++d) { + const float qv = inverse_rope.forward_q_enabled && d >= D - 64 + ? q_rope_tail[d - (D - 64)] : qh[d]; + dot += qv * ds4_fa_load(kr + d); + } + s = dot * scale + mask_v; + } + scores[r] = s; + local_max = fmaxf(local_max, s); + } + const float max_score = ds4_fa_block_max(local_max); + + float local_sum = 0.0f; + for (int r = tid; r < n_kv; r += blockDim.x) { + const float w = expf(scores[r] - max_score); + scores[r] = w; + local_sum += w; + } + if (tid == 0 && sinks) { + local_sum += expf(sinks[h] - max_score); + } + const float denom = ds4_fa_block_sum(local_sum); + const float inv_denom = 1.0f / denom; + + // One wave locates the non-zero envelope independently in the raw and + // compressed spans. Ballots inspect a whole wave of scores at once and + // require only one block barrier, unlike a four-reduction implementation. + // Values inside each envelope retain their original order, including any + // internal zero, so floating-point accumulation is unchanged. + __shared__ int value_bounds[4]; + if (tid < warpSize) { + const int lane = tid; + if (lane == 0) { + value_bounds[0] = raw_rows; + value_bounds[1] = -1; + value_bounds[2] = n_kv; + value_bounds[3] = -1; + } + for (int base = 0; base < raw_rows; base += warpSize) { + const int r = base + lane; + const unsigned long long active = __ballot( + r < raw_rows && scores[r] != 0.0f); + if (lane == 0 && active != 0) { + if (value_bounds[0] == raw_rows) { + value_bounds[0] = base + __ffsll(active) - 1; + } + value_bounds[1] = base + 63 - __clzll(active); + } + } + for (int base = raw_rows; base < n_kv; base += warpSize) { + const int r = base + lane; + const unsigned long long active = __ballot( + r < n_kv && scores[r] != 0.0f); + if (lane == 0 && active != 0) { + if (value_bounds[2] == n_kv) { + value_bounds[2] = base + __ffsll(active) - 1; + } + value_bounds[3] = base + 63 - __clzll(active); + } + } + } + __syncthreads(); + const int raw_first = value_bounds[0]; + const int raw_last = value_bounds[1]; + const int comp_first = value_bounds[2]; + const int comp_last = value_bounds[3]; + + float * rope_tail = scores; + for (int d = tid; d < D; d += blockDim.x) { + float acc = 0.0f; + for (int r = raw_first; r <= raw_last; ++r) { + acc += scores[r] * ds4_fa_load( + v + (size_t) r * D + d); + } + for (int r = comp_first; r <= comp_last; ++r) { + acc += scores[r] * ds4_fa_load( + v + (size_t) r * D + d); + } + const float value = acc * inv_denom; + if (inverse_rope.enabled && d >= D - 64) { + rope_tail[d - (D - 64)] = value; + } else { + dst[((size_t) t * (size_t) n_heads + (size_t) h) * D + d] = value; + } + } + if (inverse_rope.enabled) { + __syncthreads(); + if (tid < 32) { + const float x0 = rope_tail[2 * tid + 0]; + const float x1 = rope_tail[2 * tid + 1]; + const size_t coefficient_index = + ((size_t) t * 32 + (size_t) tid) * 2; + const float cos_theta = inverse_rope_coefficients[ + coefficient_index + 0]; + const float sin_theta = inverse_rope_coefficients[ + coefficient_index + 1]; + float y0; + float y1; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, y0, y1); + float * out = dst + + ((size_t) t * (size_t) n_heads + (size_t) h) * D + D - 64; + out[2 * tid + 0] = y0; + out[2 * tid + 1] = y1; + } + } +} + +// DS4 MLA uses one latent K/V head for every query head. Grouping query heads +// in one block lets them share each K/V load while retaining the reference +// kernel's row-to-thread mapping, per-head reduction tree, and accumulation +// order. The grouped path is dense-only; experimental sparse selection keeps +// using the single-head kernel above. +template +__global__ static void ds4_flash_attn_d512_shared_kv_grouped_kernel( + float * dst, + const float * q, + size_t q_stride_token, + size_t q_stride_head, + const KV * k, + const KV * v, + const Mask * mask, + const float * sinks, + int n_tokens, + int n_heads, + int n_kv, + float scale, + int raw_rows, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients, + const float * forward_rope_coefficients) { + constexpr int D = 512; + constexpr int N_THREADS = 256; + const int t = (int) blockIdx.x; + const int h_begin = (int) blockIdx.y * HEADS_PER_BLOCK; + const int tid = (int) threadIdx.x; + if (t >= n_tokens || h_begin >= n_heads) return; + + extern __shared__ float scratch[]; + float * scores = scratch; + float * reduction = scores + (size_t) HEADS_PER_BLOCK * n_kv; + int * value_bounds = reinterpret_cast( + reduction + (size_t) HEADS_PER_BLOCK * N_THREADS); + float * q_rope_tail = reinterpret_cast( + value_bounds + (size_t) HEADS_PER_BLOCK * 4); + + const float * qh[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + qh[j] = q + (size_t) t * q_stride_token + + (size_t) h * q_stride_head; + } + + if (inverse_rope.forward_q_enabled) { + for (int index = tid; index < HEADS_PER_BLOCK * 32; + index += (int) blockDim.x) { + const int j = index / 32; + const int pair = index % 32; + const float x0 = qh[j][D - 64 + 2 * pair + 0]; + const float x1 = qh[j][D - 64 + 2 * pair + 1]; + const size_t coefficient_index = + ((size_t) t * 32 + (size_t) pair) * 2; + const float cos_theta = forward_rope_coefficients[ + coefficient_index + 0]; + const float sin_theta = forward_rope_coefficients[ + coefficient_index + 1]; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, + q_rope_tail[(size_t) j * 64 + 2 * pair + 0], + q_rope_tail[(size_t) j * 64 + 2 * pair + 1]); + } + __syncthreads(); + } + + float local_max[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + local_max[j] = h < n_heads && sinks + ? sinks[h] : -3.402823466e38f; + } + // A thread owns exactly the same rows as in the single-head kernel. Four + // independent dot-product chains consume one shared K value per feature. + for (int r = tid; r < n_kv; r += blockDim.x) { + const float mask_v = mask + ? ds4_fa_load(mask + (size_t) t * n_kv + r) + : 0.0f; + const bool visible = mask_v > -1.0e20f; + float dot[HEADS_PER_BLOCK] = {}; + if (visible) { + const KV * kr = k + (size_t) r * D; +#pragma unroll + for (int d = 0; d < D; ++d) { + const float kv = ds4_fa_load(kr + d); +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const float qv = + inverse_rope.forward_q_enabled && d >= D - 64 + ? q_rope_tail[(size_t) j * 64 + d - (D - 64)] + : qh[j][d]; + dot[j] += qv * kv; + } + } + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + const float s = h < n_heads && visible + ? dot[j] * scale + mask_v : -3.402823466e38f; + scores[(size_t) j * n_kv + r] = s; + local_max[j] = fmaxf(local_max[j], s); + } + } + + // Match ds4_fa_block_max independently for every grouped head. +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + reduction[(size_t) j * N_THREADS + tid] = local_max[j]; + } + __syncthreads(); + for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { + if (tid < stride) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * row = reduction + (size_t) j * N_THREADS; + row[tid] = fmaxf(row[tid], row[tid + stride]); + } + } + __syncthreads(); + } + + float max_score[HEADS_PER_BLOCK]; + float local_sum[HEADS_PER_BLOCK] = {}; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + max_score[j] = reduction[(size_t) j * N_THREADS]; + } + for (int r = tid; r < n_kv; r += blockDim.x) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * score = scores + (size_t) j * n_kv + r; + const float weight = expf(*score - max_score[j]); + *score = weight; + local_sum[j] += weight; + } + } + if (tid == 0 && sinks) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + local_sum[j] += expf(sinks[h_begin + j] - max_score[j]); + } + } + + // Match ds4_fa_block_sum independently for every grouped head. +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + reduction[(size_t) j * N_THREADS + tid] = local_sum[j]; + } + __syncthreads(); + for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { + if (tid < stride) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * row = reduction + (size_t) j * N_THREADS; + row[tid] += row[tid + stride]; + } + } + __syncthreads(); + } + + float inv_denom[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + inv_denom[j] = 1.0f / reduction[(size_t) j * N_THREADS]; + } + + // Underflow can make the non-zero envelope differ by head, so retain one + // pair of raw/compressed bounds per head. Ballot order does not affect any + // arithmetic result. + if (tid < warpSize) { + const int lane = tid; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + int * bounds = value_bounds + 4 * j; + if (lane == 0) { + bounds[0] = raw_rows; + bounds[1] = -1; + bounds[2] = n_kv; + bounds[3] = -1; + } + for (int base = 0; base < raw_rows; base += warpSize) { + const int r = base + lane; + const unsigned long long active = __ballot( + r < raw_rows && + scores[(size_t) j * n_kv + r] != 0.0f); + if (lane == 0 && active != 0) { + if (bounds[0] == raw_rows) { + bounds[0] = base + __ffsll(active) - 1; + } + bounds[1] = base + 63 - __clzll(active); + } + } + for (int base = raw_rows; base < n_kv; base += warpSize) { + const int r = base + lane; + const unsigned long long active = __ballot( + r < n_kv && scores[(size_t) j * n_kv + r] != 0.0f); + if (lane == 0 && active != 0) { + if (bounds[2] == n_kv) { + bounds[2] = base + __ffsll(active) - 1; + } + bounds[3] = base + 63 - __clzll(active); + } + } + } + } + __syncthreads(); + + int raw_first = raw_rows; + int raw_last = -1; + int comp_first = n_kv; + int comp_last = -1; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int * bounds = value_bounds + 4 * j; + raw_first = min(raw_first, bounds[0]); + raw_last = max(raw_last, bounds[1]); + comp_first = min(comp_first, bounds[2]); + comp_last = max(comp_last, bounds[3]); + } + + float * rope_tail = reduction; + for (int d = tid; d < D; d += blockDim.x) { + float acc[HEADS_PER_BLOCK] = {}; + for (int r = raw_first; r <= raw_last; ++r) { + const float vv = ds4_fa_load(v + (size_t) r * D + d); +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int * bounds = value_bounds + 4 * j; + if (r >= bounds[0] && r <= bounds[1]) { + acc[j] += scores[(size_t) j * n_kv + r] * vv; + } + } + } + for (int r = comp_first; r <= comp_last; ++r) { + const float vv = ds4_fa_load(v + (size_t) r * D + d); +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int * bounds = value_bounds + 4 * j; + if (r >= bounds[2] && r <= bounds[3]) { + acc[j] += scores[(size_t) j * n_kv + r] * vv; + } + } + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + if (h < n_heads) { + const float value = acc[j] * inv_denom[j]; + if (inverse_rope.enabled && d >= D - 64) { + rope_tail[(size_t) j * 64 + d - (D - 64)] = value; + } else { + dst[((size_t) t * (size_t) n_heads + (size_t) h) * D + d] = + value; + } + } + } + } + if (inverse_rope.enabled) { + __syncthreads(); + if (tid < HEADS_PER_BLOCK * 32) { + const int j = tid / 32; + const int pair = tid % 32; + const float x0 = rope_tail[(size_t) j * 64 + 2 * pair + 0]; + const float x1 = rope_tail[(size_t) j * 64 + 2 * pair + 1]; + const size_t coefficient_index = + ((size_t) t * 32 + (size_t) pair) * 2; + const float cos_theta = inverse_rope_coefficients[ + coefficient_index + 0]; + const float sin_theta = inverse_rope_coefficients[ + coefficient_index + 1]; + float y0; + float y1; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, y0, y1); + const int h = h_begin + j; + float * out = dst + + ((size_t) t * (size_t) n_heads + (size_t) h) * D + D - 64; + out[2 * pair + 0] = y0; + out[2 * pair + 1] = y1; + } + } +} + +// Dense-prefill variant of the grouped kernel with compact score storage. +// The mask-derived envelopes only change the address used to retain a score; +// every visible row keeps its original owner thread, dot-product order, +// reduction tree, softmax order, and value-accumulation position. +template +__global__ static void ds4_flash_attn_d512_shared_kv_grouped_compact_kernel( + float * dst, + const float * q, + size_t q_stride_token, + size_t q_stride_head, + const KV * k, + const KV * v, + const Mask * mask, + const float * sinks, + int n_tokens, + int n_heads, + int n_kv, + float scale, + int raw_rows, + int raw_score_capacity, + int score_stride, + const int * visibility_bounds, + const int * indexed_rows, + const int * indexed_counts, + const int * indexed_owner_offsets, + const int * indexed_owner_ranks, + int indexed_capacity, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients, + const float * forward_rope_coefficients) { + constexpr int D = 512; + constexpr int N_THREADS = 256; + static_assert(VALUES_PER_THREAD == 2 || VALUES_PER_THREAD == 4); + const int t = (int) blockIdx.x; + const int h_begin = (int) blockIdx.y * HEADS_PER_BLOCK; + const int tid = (int) threadIdx.x; + if (t >= n_tokens || h_begin >= n_heads) return; + + extern __shared__ float scratch[]; + float * scores = scratch; + float * reduction = scores + (size_t) HEADS_PER_BLOCK * score_stride; + int * value_bounds = reinterpret_cast( + reduction + (size_t) HEADS_PER_BLOCK * N_THREADS); + float * q_rope_tail = reinterpret_cast( + value_bounds + (size_t) HEADS_PER_BLOCK * 4); + + const int * token_visibility = visibility_bounds + (size_t) t * 4; + const int mask_raw_first = token_visibility[0]; + const int mask_raw_last = token_visibility[1]; + const int mask_comp_first = token_visibility[2]; + const int mask_comp_last = token_visibility[3]; + const int * token_indexed_rows = nullptr; + const int * token_owner_offsets = nullptr; + const int * token_owner_ranks = nullptr; + int indexed_count = 0; + if constexpr (INDEXED_MASK) { + token_indexed_rows = indexed_rows + (size_t) t * indexed_capacity; + token_owner_offsets = indexed_owner_offsets + (size_t) t * (N_THREADS + 1); + token_owner_ranks = indexed_owner_ranks + (size_t) t * indexed_capacity; + indexed_count = indexed_counts[t]; + } + + const float * qh[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + qh[j] = q + (size_t) t * q_stride_token + + (size_t) h * q_stride_head; + } + + if (inverse_rope.forward_q_enabled) { + for (int index = tid; index < HEADS_PER_BLOCK * 32; + index += (int) blockDim.x) { + const int j = index / 32; + const int pair = index % 32; + const float x0 = qh[j][D - 64 + 2 * pair + 0]; + const float x1 = qh[j][D - 64 + 2 * pair + 1]; + const size_t coefficient_index = + ((size_t) t * 32 + (size_t) pair) * 2; + const float cos_theta = forward_rope_coefficients[ + coefficient_index + 0]; + const float sin_theta = forward_rope_coefficients[ + coefficient_index + 1]; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, + q_rope_tail[(size_t) j * 64 + 2 * pair + 0], + q_rope_tail[(size_t) j * 64 + 2 * pair + 1]); + } + __syncthreads(); + } + + float local_max[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + local_max[j] = h < n_heads && sinks + ? sinks[h] : -3.402823466e38f; + } + // Build the full kernel's exact per-head non-zero envelope while the + // softmax weights are emitted. This avoids rescanning every context row + // after softmax without changing the subsequent V accumulation interval. + if (tid < HEADS_PER_BLOCK * 4) { + const int slot = tid & 3; + value_bounds[tid] = slot == 0 + ? raw_rows + : slot == 1 + ? -1 + : slot == 2 + ? (INDEXED_MASK ? indexed_count : n_kv) + : -1; + } + + // Preserve each thread's original r = tid + 256*k order. In indexed mode, + // owner_ranks is the exact selected subsequence of that traversal, so the + // hot passes no longer scan every unselected compressed row. + const int raw_owner_first = mask_raw_first + + ((tid - (mask_raw_first & (N_THREADS - 1)) + N_THREADS) & + (N_THREADS - 1)); + const int raw_iteration_count = raw_owner_first <= mask_raw_last + ? 1 + (mask_raw_last - raw_owner_first) / N_THREADS : 0; + const int comp_owner_first = mask_comp_first + + ((tid - (mask_comp_first & (N_THREADS - 1)) + N_THREADS) & + (N_THREADS - 1)); + const int comp_iteration_count = comp_owner_first <= mask_comp_last + ? 1 + (mask_comp_last - comp_owner_first) / N_THREADS : 0; + int owner_begin = 0; + int owner_count = 0; + int score_iteration_count = raw_iteration_count + comp_iteration_count; + if constexpr (INDEXED_MASK) { + owner_begin = token_owner_offsets[tid]; + owner_count = token_owner_offsets[tid + 1] - owner_begin; + score_iteration_count = raw_iteration_count + owner_count; + } + for (int iteration = 0; iteration < score_iteration_count; ++iteration) { + int r; + int score_index; + if constexpr (INDEXED_MASK) { + if (iteration < raw_iteration_count) { + r = raw_owner_first + iteration * N_THREADS; + score_index = r - mask_raw_first; + } else { + const int rank = token_owner_ranks[ + owner_begin + iteration - raw_iteration_count]; + r = token_indexed_rows[rank]; + score_index = raw_score_capacity + rank; + } + } else { + if (iteration < raw_iteration_count) { + r = raw_owner_first + iteration * N_THREADS; + score_index = r - mask_raw_first; + } else { + r = comp_owner_first + + (iteration - raw_iteration_count) * N_THREADS; + score_index = raw_score_capacity + r - mask_comp_first; + } + } + + const float mask_v = ds4_fa_load( + mask + (size_t) t * n_kv + r); + const bool visible = mask_v > -1.0e20f; + float dot[HEADS_PER_BLOCK] = {}; + if (visible) { + const KV * kr = k + (size_t) r * D; +#pragma unroll + for (int d = 0; d < D; ++d) { + const float kv = ds4_fa_load(kr + d); +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const float qv = + inverse_rope.forward_q_enabled && d >= D - 64 + ? q_rope_tail[(size_t) j * 64 + d - (D - 64)] + : qh[j][d]; + dot[j] += qv * kv; + } + } + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + const float s = h < n_heads && visible + ? dot[j] * scale + mask_v : -3.402823466e38f; + scores[(size_t) j * score_stride + score_index] = s; + local_max[j] = fmaxf(local_max[j], s); + } + } + +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + reduction[(size_t) j * N_THREADS + tid] = local_max[j]; + } + __syncthreads(); + for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { + if (tid < stride) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * row = reduction + (size_t) j * N_THREADS; + row[tid] = fmaxf(row[tid], row[tid + stride]); + } + } + __syncthreads(); + } + + float max_score[HEADS_PER_BLOCK]; + float local_sum[HEADS_PER_BLOCK] = {}; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + max_score[j] = reduction[(size_t) j * N_THREADS]; + } + for (int iteration = 0; iteration < score_iteration_count; ++iteration) { + int score_index; + int bound_value; + const bool raw_value = iteration < raw_iteration_count; + if constexpr (INDEXED_MASK) { + if (raw_value) { + const int r = raw_owner_first + iteration * N_THREADS; + score_index = r - mask_raw_first; + bound_value = r; + } else { + const int rank = token_owner_ranks[ + owner_begin + iteration - raw_iteration_count]; + score_index = raw_score_capacity + rank; + bound_value = rank; + } + } else { + if (raw_value) { + const int r = raw_owner_first + iteration * N_THREADS; + score_index = r - mask_raw_first; + bound_value = r; + } else { + const int r = comp_owner_first + + (iteration - raw_iteration_count) * N_THREADS; + score_index = raw_score_capacity + r - mask_comp_first; + bound_value = r; + } + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * score = scores + (size_t) j * score_stride + score_index; + const float weight = expf(*score - max_score[j]); + *score = weight; + local_sum[j] += weight; + if (weight != 0.0f) { + int * bounds = value_bounds + 4 * j + (raw_value ? 0 : 2); + atomicMin(bounds + 0, bound_value); + atomicMax(bounds + 1, bound_value); + } + } + } + if (tid == 0 && sinks) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + local_sum[j] += expf(sinks[h_begin + j] - max_score[j]); + } + } + +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + reduction[(size_t) j * N_THREADS + tid] = local_sum[j]; + } + __syncthreads(); + for (int stride = N_THREADS / 2; stride > 0; stride >>= 1) { + if (tid < stride) { +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + float * row = reduction + (size_t) j * N_THREADS; + row[tid] += row[tid + stride]; + } + } + __syncthreads(); + } + + float inv_denom[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + inv_denom[j] = 1.0f / reduction[(size_t) j * N_THREADS]; + } + + // Finish the shared envelope before the value phase reads it. + __syncthreads(); + + int raw_first = raw_rows; + int raw_last = -1; + int comp_first = INDEXED_MASK ? indexed_count : n_kv; + int comp_last = -1; + int head_raw_first[HEADS_PER_BLOCK]; + int head_raw_last[HEADS_PER_BLOCK]; + int head_comp_first[HEADS_PER_BLOCK]; + int head_comp_last[HEADS_PER_BLOCK]; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int * bounds = value_bounds + 4 * j; + head_raw_first[j] = bounds[0]; + head_raw_last[j] = bounds[1]; + head_comp_first[j] = bounds[2]; + head_comp_last[j] = bounds[3]; + raw_first = min(raw_first, head_raw_first[j]); + raw_last = max(raw_last, head_raw_last[j]); + comp_first = min(comp_first, head_comp_first[j]); + comp_last = max(comp_last, head_comp_last[j]); + } + + // One active thread owns adjacent value dimensions. This retains each + // dimension's ascending row accumulation order while sharing score loads, + // row-loop control, and a naturally aligned vector V load across the group. + float * rope_tail = reduction; + const int d0 = VALUES_PER_THREAD * tid; + const int d1 = d0 + 1; + const int d2 = d0 + 2; + const int d3 = d0 + 3; + float acc0[HEADS_PER_BLOCK] = {}; + float acc1[HEADS_PER_BLOCK] = {}; + float acc2[HEADS_PER_BLOCK] = {}; + float acc3[HEADS_PER_BLOCK] = {}; + if (d0 < D) { + for (int r = raw_first; r <= raw_last; ++r) { + const int score_index = r - mask_raw_first; + float vv0; + float vv1; + float vv2 = 0.0f; + float vv3 = 0.0f; + if constexpr (VALUES_PER_THREAD == 2) { + ds4_fa_load_pair( + v + (size_t) r * D + d0, vv0, vv1); + } else { + ds4_fa_load_quad( + v + (size_t) r * D + d0, vv0, vv1, vv2, vv3); + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + if (r >= head_raw_first[j] && r <= head_raw_last[j]) { + const float weight = + scores[(size_t) j * score_stride + score_index]; + acc0[j] += weight * vv0; + acc1[j] += weight * vv1; + if constexpr (VALUES_PER_THREAD == 4) { + acc2[j] += weight * vv2; + acc3[j] += weight * vv3; + } + } + } + } + if constexpr (INDEXED_MASK) { + for (int rank = comp_first; rank <= comp_last; ++rank) { + const int r = token_indexed_rows[rank]; + const int score_index = raw_score_capacity + rank; + bool any_nonzero = false; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + any_nonzero = any_nonzero || + scores[(size_t) j * score_stride + score_index] != 0.0f; + } + if (!any_nonzero) continue; + float vv0; + float vv1; + float vv2 = 0.0f; + float vv3 = 0.0f; + if constexpr (VALUES_PER_THREAD == 2) { + ds4_fa_load_pair( + v + (size_t) r * D + d0, vv0, vv1); + } else { + ds4_fa_load_quad( + v + (size_t) r * D + d0, vv0, vv1, vv2, vv3); + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + if (rank >= head_comp_first[j] && rank <= head_comp_last[j]) { + const float weight = + scores[(size_t) j * score_stride + score_index]; + acc0[j] += weight * vv0; + acc1[j] += weight * vv1; + if constexpr (VALUES_PER_THREAD == 4) { + acc2[j] += weight * vv2; + acc3[j] += weight * vv3; + } + } + } + } + } else { + for (int r = comp_first; r <= comp_last; ++r) { + const int score_index = + raw_score_capacity + r - mask_comp_first; + bool any_nonzero = false; +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + any_nonzero = any_nonzero || + scores[(size_t) j * score_stride + score_index] != 0.0f; + } + if (!any_nonzero) continue; + float vv0; + float vv1; + float vv2 = 0.0f; + float vv3 = 0.0f; + if constexpr (VALUES_PER_THREAD == 2) { + ds4_fa_load_pair( + v + (size_t) r * D + d0, vv0, vv1); + } else { + ds4_fa_load_quad( + v + (size_t) r * D + d0, vv0, vv1, vv2, vv3); + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + if (r >= head_comp_first[j] && r <= head_comp_last[j]) { + const float weight = + scores[(size_t) j * score_stride + score_index]; + acc0[j] += weight * vv0; + acc1[j] += weight * vv1; + if constexpr (VALUES_PER_THREAD == 4) { + acc2[j] += weight * vv2; + acc3[j] += weight * vv3; + } + } + } + } + } +#pragma unroll + for (int j = 0; j < HEADS_PER_BLOCK; ++j) { + const int h = h_begin + j; + if (h < n_heads) { + const float value0 = acc0[j] * inv_denom[j]; + const float value1 = acc1[j] * inv_denom[j]; + if (inverse_rope.enabled && d0 >= D - 64) { + rope_tail[(size_t) j * 64 + d0 - (D - 64)] = value0; + rope_tail[(size_t) j * 64 + d1 - (D - 64)] = value1; + if constexpr (VALUES_PER_THREAD == 4) { + const float value2 = acc2[j] * inv_denom[j]; + const float value3 = acc3[j] * inv_denom[j]; + rope_tail[(size_t) j * 64 + d2 - (D - 64)] = value2; + rope_tail[(size_t) j * 64 + d3 - (D - 64)] = value3; + } + } else { + float * out = dst + + ((size_t) t * (size_t) n_heads + (size_t) h) * D + d0; + out[0] = value0; + out[1] = value1; + if constexpr (VALUES_PER_THREAD == 4) { + out[2] = acc2[j] * inv_denom[j]; + out[3] = acc3[j] * inv_denom[j]; + } + } + } + } + } + if (inverse_rope.enabled) { + __syncthreads(); + if (tid < HEADS_PER_BLOCK * 32) { + const int j = tid / 32; + const int pair = tid % 32; + const float x0 = rope_tail[(size_t) j * 64 + 2 * pair + 0]; + const float x1 = rope_tail[(size_t) j * 64 + 2 * pair + 1]; + const size_t coefficient_index = + ((size_t) t * 32 + (size_t) pair) * 2; + const float cos_theta = inverse_rope_coefficients[ + coefficient_index + 0]; + const float sin_theta = inverse_rope_coefficients[ + coefficient_index + 1]; + float y0; + float y1; + ds4_apply_inverse_rope_pair( + x0, x1, cos_theta, sin_theta, y0, y1); + const int h = h_begin + j; + float * out = dst + + ((size_t) t * (size_t) n_heads + (size_t) h) * D + D - 64; + out[2 * pair + 0] = y0; + out[2 * pair + 1] = y1; + } + } +} + +template +static bool ds4_launch_flash_attn_d512_grouped( + ggml_tensor * dst, + const ggml_tensor * Q, + const ggml_tensor * K, + const ggml_tensor * V, + const ggml_tensor * mask, + const ggml_tensor * sinks, + bool kv_f16, + bool kv_f32, + int n_tokens, + int n_heads, + int n_kv, + float scale, + int raw_rows, + size_t q_stride_token, + size_t q_stride_head, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients, + const float * forward_rope_coefficients, + size_t shmem, + cudaStream_t stream) { + dim3 grid( + (unsigned) n_tokens, + (unsigned) (n_heads / HEADS_PER_BLOCK), 1); + if (kv_f16 && (!mask || mask->type == GGML_TYPE_F16)) { + ds4_flash_attn_d512_shared_kv_grouped_kernel< + half, half, HEADS_PER_BLOCK> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, (const half *) V->data, + mask ? (const half *) mask->data : nullptr, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, inverse_rope, + inverse_rope_coefficients, forward_rope_coefficients); + } else if (kv_f32 && (!mask || mask->type == GGML_TYPE_F32)) { + ds4_flash_attn_d512_shared_kv_grouped_kernel< + float, float, HEADS_PER_BLOCK> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + mask ? (const float *) mask->data : nullptr, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, inverse_rope, + inverse_rope_coefficients, forward_rope_coefficients); + } else if (kv_f32 && mask && mask->type == GGML_TYPE_F16) { + ds4_flash_attn_d512_shared_kv_grouped_kernel< + float, half, HEADS_PER_BLOCK> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, inverse_rope, + inverse_rope_coefficients, forward_rope_coefficients); + } else { + return false; + } + return true; +} + +template +static bool ds4_launch_flash_attn_d512_grouped_compact( + ggml_tensor * dst, + const ggml_tensor * Q, + const ggml_tensor * K, + const ggml_tensor * V, + const ggml_tensor * mask, + const ggml_tensor * sinks, + bool kv_f16, + bool kv_f32, + int n_tokens, + int n_heads, + int n_kv, + float scale, + int raw_rows, + int raw_score_capacity, + int score_stride, + const int * visibility_bounds, + const int * indexed_rows, + const int * indexed_counts, + const int * indexed_owner_offsets, + const int * indexed_owner_ranks, + int indexed_capacity, + size_t q_stride_token, + size_t q_stride_head, + ds4_inverse_rope_params inverse_rope, + const float * inverse_rope_coefficients, + const float * forward_rope_coefficients, + size_t shmem, + cudaStream_t stream) { + GGML_ASSERT(mask && visibility_bounds); + if constexpr (INDEXED_MASK) { + GGML_ASSERT(indexed_rows && indexed_counts && + indexed_owner_offsets && indexed_owner_ranks); + } + dim3 grid( + (unsigned) n_tokens, + (unsigned) (n_heads / HEADS_PER_BLOCK), 1); + if (kv_f16 && mask->type == GGML_TYPE_F16) { + ds4_flash_attn_d512_shared_kv_grouped_compact_kernel< + half, half, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, (const half *) V->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, + raw_score_capacity, score_stride, visibility_bounds, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, indexed_capacity, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients); + } else if (kv_f32 && mask->type == GGML_TYPE_F32) { + ds4_flash_attn_d512_shared_kv_grouped_compact_kernel< + float, float, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + (const float *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, + raw_score_capacity, score_stride, visibility_bounds, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, indexed_capacity, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients); + } else if (kv_f32 && mask->type == GGML_TYPE_F16) { + ds4_flash_attn_d512_shared_kv_grouped_compact_kernel< + float, half, HEADS_PER_BLOCK, INDEXED_MASK, VALUES_PER_THREAD> + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, n_heads, n_kv, scale, raw_rows, + raw_score_capacity, score_stride, visibility_bounds, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, indexed_capacity, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients); + } else { + return false; + } + return true; +} + +static bool ggml_cuda_ds4_flash_attn_d512_f32(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const ggml_tensor * sinks = dst->src[4]; + const bool kv_f32 = K && V && K->type == GGML_TYPE_F32 && + V->type == GGML_TYPE_F32; + const bool kv_f16 = K && V && K->type == GGML_TYPE_F16 && + V->type == GGML_TYPE_F16; + const bool mask_ok = !mask || mask->type == GGML_TYPE_F16 || + mask->type == GGML_TYPE_F32; + if (!Q || !K || !V || + Q->type != GGML_TYPE_F32 || (!kv_f32 && !kv_f16) || !mask_ok || + dst->type != GGML_TYPE_F32 || + Q->ne[0] != 512 || K->ne[0] != 512 || V->ne[0] != 512 || + K->ne[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) || + dst->nb[0] != (int64_t) sizeof(float)) { + return false; + } + const size_t kv_esz = kv_f16 ? sizeof(half) : sizeof(float); + if (K->nb[0] != kv_esz || V->nb[0] != kv_esz || + K->nb[1] != (size_t) K->ne[0] * kv_esz || + V->nb[1] != (size_t) V->ne[0] * kv_esz || + Q->nb[1] % sizeof(float) != 0 || + Q->nb[2] % sizeof(float) != 0 || + (mask && (mask->ne[0] != K->ne[1] || + mask->ne[1] != Q->ne[1] || + mask->nb[1] != (size_t) mask->ne[0] * mask->nb[0]))) { + 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]; + const size_t q_stride_token = Q->nb[1] / sizeof(float); + const size_t q_stride_head = Q->nb[2] / sizeof(float); + if (n_tokens <= 0 || n_heads <= 0 || n_kv <= 0) { + return false; + } + + int raw_rows = ggml_get_op_params_i32(dst, 4); + int sparse_keep_rows = ggml_get_op_params_i32(dst, 5); + const unsigned int ds4_layout = + (unsigned int) ggml_get_op_params_i32(dst, 6); + int raw_window = (int) (ds4_layout >> 16); + int sparse_block_size = (int) (ds4_layout & 0xffffu); + raw_rows = max(0, min(raw_rows, n_kv)); + if (raw_window <= 0) raw_window = raw_rows; + raw_window = max(1, min(raw_window, raw_rows)); + if (sparse_block_size <= 0) sparse_block_size = 32; + const int n_comp_rows = n_kv - raw_rows; + const int n_comp_blocks = (n_comp_rows + sparse_block_size - 1) / + sparse_block_size; + const bool sparse = sparse_keep_rows > 0 && + sparse_keep_rows < n_comp_rows && + n_comp_blocks > 0; + const bool indexed_mask = sparse_keep_rows < 0 && n_comp_rows > 0; + const int indexed_capacity = indexed_mask + ? min(-sparse_keep_rows, n_comp_rows) : 0; + + ds4_inverse_rope_params inverse_rope{}; + const int rope_flags = ggml_get_op_params_i32(dst, 7); + inverse_rope.enabled = rope_flags & 1; + inverse_rope.forward_q_enabled = (rope_flags & 2) != 0; + if (rope_flags != 0) { + inverse_rope.kv_start = ggml_get_op_params_i32(dst, 8); + const float freq_base = ggml_get_op_params_f32(dst, 9); + inverse_rope.freq_scale = ggml_get_op_params_f32(dst, 10); + inverse_rope.ext_factor = ggml_get_op_params_f32(dst, 11); + inverse_rope.attn_factor = ggml_get_op_params_f32(dst, 12); + const float beta_fast = ggml_get_op_params_f32(dst, 13); + const float beta_slow = ggml_get_op_params_f32(dst, 14); + const int n_ctx_orig = ggml_get_op_params_i32(dst, 15); + float corr_dims[2]; + ggml_rope_yarn_corr_dims( + 64, n_ctx_orig, freq_base, beta_fast, beta_slow, corr_dims); + inverse_rope.corr_low = corr_dims[0]; + inverse_rope.corr_high = corr_dims[1]; + inverse_rope.theta_scale = powf(freq_base, -2.0f / 64.0f); + } + + cudaStream_t stream = ctx.stream(); + ggml_cuda_pool_alloc inverse_rope_coefficients_alloc(ctx.pool()); + float * inverse_rope_coefficients = nullptr; + if (inverse_rope.enabled) { + inverse_rope_coefficients = inverse_rope_coefficients_alloc.alloc( + (size_t) n_tokens * 32 * 2); + const int coefficient_count = n_tokens * 32; + ds4_inverse_rope_coefficients_kernel<<< + (coefficient_count + 255) / 256, 256, 0, stream>>>( + inverse_rope_coefficients, n_tokens, inverse_rope); + } + ggml_cuda_pool_alloc forward_rope_coefficients_alloc(ctx.pool()); + float * forward_rope_coefficients = nullptr; + if (inverse_rope.forward_q_enabled) { + forward_rope_coefficients = forward_rope_coefficients_alloc.alloc( + (size_t) n_tokens * 32 * 2); + const int coefficient_count = n_tokens * 32; + ds4_forward_rope_coefficients_kernel<<< + (coefficient_count + 255) / 256, 256, 0, stream>>>( + forward_rope_coefficients, n_tokens, inverse_rope); + } + ggml_cuda_pool_alloc mean_k_alloc(ctx.pool()); + float * mean_k = nullptr; + if (sparse) { + mean_k = mean_k_alloc.alloc((size_t) n_comp_blocks * 512); + if (kv_f16) { + ds4_fa_mean_comp_blocks_kernel + <<>>( + (const half *) K->data, mean_k, n_kv, raw_rows, + sparse_block_size, n_comp_blocks); + } else { + ds4_fa_mean_comp_blocks_kernel + <<>>( + (const float *) K->data, mean_k, n_kv, raw_rows, + sparse_block_size, n_comp_blocks); + } + } + dim3 grid((unsigned) n_tokens, (unsigned) n_heads, 1); + const int score_or_rope_rows = inverse_rope.enabled ? max(n_kv, 64) : n_kv; + const size_t shmem = + (size_t) (score_or_rope_rows + 2 * n_comp_blocks + + (inverse_rope.forward_q_enabled ? 64 : 0)) * sizeof(float); + float params[3] = {}; + memcpy(params, dst->op_params, sizeof(params)); + const float scale = params[0]; + + constexpr int group4 = 4; + constexpr int group2 = 2; + const size_t group4_shmem = + ((size_t) group4 * n_kv + (size_t) group4 * 256) * sizeof(float) + + (size_t) group4 * 4 * sizeof(int) + + (inverse_rope.forward_q_enabled ? (size_t) group4 * 64 * sizeof(float) : 0); + const size_t group2_shmem = + ((size_t) group2 * n_kv + (size_t) group2 * 256) * sizeof(float) + + (size_t) group2 * 4 * sizeof(int) + + (inverse_rope.forward_q_enabled ? (size_t) group2 * 64 * sizeof(float) : 0); + const int compact_score_stride = raw_window + + (indexed_mask ? indexed_capacity : n_comp_rows); + const size_t compact_group4_shmem = + ((size_t) group4 * compact_score_stride + (size_t) group4 * 256) * sizeof(float) + + (size_t) group4 * 4 * sizeof(int) + + (inverse_rope.forward_q_enabled ? (size_t) group4 * 64 * sizeof(float) : 0); + // Four heads win while two blocks can remain resident in 48 KiB of LDS. + // Beyond that point, two-head grouping trades some K/V reuse for higher + // occupancy; larger working sets fall back to the single-head kernel. + if (!sparse && n_heads % group4 == 0 && group4_shmem <= 24 * 1024) { + return ds4_launch_flash_attn_d512_grouped( + dst, Q, K, V, mask, sinks, kv_f16, kv_f32, + n_tokens, n_heads, n_kv, scale, raw_rows, + q_stride_token, q_stride_head, + inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients, + group4_shmem, stream); + } + // Long causal-prefill chunks can have thousands of physical raw rows but + // at most raw_window visible rows for any one token. Compacting only the + // score storage lets the same exact four-head kernel remain at two-block + // occupancy. Keep the ordinary path for shapes that already fit, avoiding + // a bounds-scan launch where it cannot improve grouping. + const bool compact_group4 = + !sparse && mask && n_heads % group4 == 0 && + raw_rows > raw_window && group4_shmem > 24 * 1024 && + compact_group4_shmem <= 24 * 1024; + if (compact_group4) { + ggml_cuda_pool_alloc visibility_bounds_alloc(ctx.pool()); + int * visibility_bounds = visibility_bounds_alloc.alloc( + (size_t) n_tokens * 4); + ggml_cuda_pool_alloc indexed_rows_alloc(ctx.pool()); + ggml_cuda_pool_alloc indexed_counts_alloc(ctx.pool()); + ggml_cuda_pool_alloc indexed_owner_offsets_alloc(ctx.pool()); + ggml_cuda_pool_alloc indexed_owner_ranks_alloc(ctx.pool()); + int * indexed_rows = nullptr; + int * indexed_counts = nullptr; + int * indexed_owner_offsets = nullptr; + int * indexed_owner_ranks = nullptr; + if (indexed_mask) { + indexed_rows = indexed_rows_alloc.alloc( + (size_t) n_tokens * indexed_capacity); + indexed_counts = indexed_counts_alloc.alloc((size_t) n_tokens); + indexed_owner_offsets = indexed_owner_offsets_alloc.alloc( + (size_t) n_tokens * 257); + indexed_owner_ranks = indexed_owner_ranks_alloc.alloc( + (size_t) n_tokens * indexed_capacity); + if (mask->type == GGML_TYPE_F16) { + ds4_fa_indexed_rows_kernel<<>>( + (const half *) mask->data, indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, + n_tokens, n_kv, raw_rows, + indexed_capacity); + } else { + ds4_fa_indexed_rows_kernel<<>>( + (const float *) mask->data, indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, + n_tokens, n_kv, raw_rows, + indexed_capacity); + } + CUDA_CHECK(cudaGetLastError()); + } + if (mask->type == GGML_TYPE_F16) { + ds4_fa_visibility_bounds_kernel<<>>( + (const half *) mask->data, visibility_bounds, + n_tokens, n_kv, raw_rows); + } else { + ds4_fa_visibility_bounds_kernel<<>>( + (const float *) mask->data, visibility_bounds, + n_tokens, n_kv, raw_rows); + } + CUDA_CHECK(cudaGetLastError()); + if (indexed_mask) { + return ds4_launch_flash_attn_d512_grouped_compact< + group4, true, 4>( + dst, Q, K, V, mask, sinks, kv_f16, kv_f32, + n_tokens, n_heads, n_kv, scale, raw_rows, + raw_window, compact_score_stride, visibility_bounds, + indexed_rows, indexed_counts, + indexed_owner_offsets, indexed_owner_ranks, indexed_capacity, + q_stride_token, q_stride_head, + inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients, + compact_group4_shmem, stream); + } + return ds4_launch_flash_attn_d512_grouped_compact< + group4, false, 4>( + dst, Q, K, V, mask, sinks, kv_f16, kv_f32, + n_tokens, n_heads, n_kv, scale, raw_rows, + raw_window, compact_score_stride, visibility_bounds, + nullptr, nullptr, nullptr, nullptr, 0, + q_stride_token, q_stride_head, + inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients, + compact_group4_shmem, stream); + } + if (!sparse && n_heads % group2 == 0 && group2_shmem <= 48 * 1024) { + return ds4_launch_flash_attn_d512_grouped( + dst, Q, K, V, mask, sinks, kv_f16, kv_f32, + n_tokens, n_heads, n_kv, scale, raw_rows, + q_stride_token, q_stride_head, + inverse_rope, + inverse_rope_coefficients, + forward_rope_coefficients, + group2_shmem, stream); + } + + if (kv_f16 && (!mask || mask->type == GGML_TYPE_F16)) { + ds4_flash_attn_d512_shared_kv_kernel + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const half *) K->data, (const half *) V->data, + mask ? (const half *) mask->data : nullptr, + sinks ? (const float *) sinks->data : nullptr, + mean_k, n_tokens, n_heads, n_kv, scale, raw_rows, + sparse_keep_rows, sparse_block_size, n_comp_blocks, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients); + } else if (kv_f32 && (!mask || mask->type == GGML_TYPE_F32)) { + ds4_flash_attn_d512_shared_kv_kernel + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + mask ? (const float *) mask->data : nullptr, + sinks ? (const float *) sinks->data : nullptr, + mean_k, n_tokens, n_heads, n_kv, scale, raw_rows, + sparse_keep_rows, sparse_block_size, n_comp_blocks, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients); + } else if (kv_f32 && mask && mask->type == GGML_TYPE_F16) { + ds4_flash_attn_d512_shared_kv_kernel + <<>>( + (float *) dst->data, (const float *) Q->data, + q_stride_token, q_stride_head, + (const float *) K->data, (const float *) V->data, + (const half *) mask->data, + sinks ? (const float *) sinks->data : nullptr, + mean_k, n_tokens, n_heads, n_kv, scale, raw_rows, + sparse_keep_rows, sparse_block_size, n_comp_blocks, + inverse_rope, inverse_rope_coefficients, + forward_rope_coefficients); + } else { + return false; + } + return true; +} + 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 +2141,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 +2165,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 +2361,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..05efb4c58 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,8 @@ #include "ggml-cuda/cumsum.cuh" #include "ggml-cuda/fill.cuh" #include "ggml-cuda/moe-fused.cuh" +#include "ggml-cuda/ds4-hc.cuh" +#include "ggml-cuda/ds4-indexer.cuh" #include "ggml.h" #include @@ -2301,9 +2303,17 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, const ggml_tensor * ffn_gate, const ggml_tensor * glu, const ggml_tensor * ffn_up_bias = nullptr, - const ggml_tensor * ffn_gate_bias = nullptr) { + const ggml_tensor * ffn_gate_bias = nullptr, + const ggml_tensor * ffn_up_glu = nullptr, + const ggml_tensor * ffn_gate_glu = nullptr) { const bool has_bias = ffn_up_bias != nullptr || ffn_gate_bias != nullptr; + // Shape-only views are common between routed expert projections and GLU. + // Keep validation anchored on the actual matmuls while allowing the GLU + // to consume their equivalent reshape nodes. + ffn_up_glu = ffn_up_glu ? ffn_up_glu : ffn_up; + ffn_gate_glu = ffn_gate_glu ? ffn_gate_glu : ffn_gate; + if (has_bias && (!ffn_up_bias || !ffn_gate_bias)) { return false; } @@ -2343,7 +2353,7 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, } } } else { - if (glu->src[0] != ffn_gate && glu->src[1] != ffn_up) { + if (glu->src[0] != ffn_gate_glu || glu->src[1] != ffn_up_glu) { return false; } } @@ -2361,7 +2371,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 +2432,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]; @@ -2457,6 +2478,65 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { return use_mul_mat_vec_q; } +// Execute an unbiased gate/up GLU while retaining the caller's existing +// output tensors. Vector kernels write the GLU directly. Large routed-expert +// MMQ keeps both ordinary matmul writes (and therefore numerical behavior), +// but shares the identical ids sort and F32->Q8 activation quantization. +static bool ggml_cuda_try_fuse_mul_mat_glu( + ggml_backend_cuda_context & ctx, + ggml_tensor * gate, + ggml_tensor * up, + ggml_tensor * glu) { + const ggml_tensor * src0 = up->src[0]; + const ggml_tensor * src1 = up->src[1]; + const ggml_tensor * ids = up->src[2]; + + if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate->src[0]; + ggml_cuda_set_fusion_glu_params(fusion_data, glu); + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, ids, glu, &fusion_data); + return true; + } + + if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate->src[0]; + ggml_cuda_set_fusion_glu_params(fusion_data, glu); + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, glu, &fusion_data); + return true; + } + + if (ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU_DS4) { + if (!ids && (ggml_mul_mat_is_grouped_src(up) || + ggml_mul_mat_is_grouped_src(gate))) { + return false; + } + + const ggml_tensor * weights[] = { up->src[0], gate->src[0] }; + for (const ggml_tensor * weight : weights) { + const bool bad_padding_clear = + ggml_backend_buffer_get_usage(weight->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && + ggml_nbytes(weight) != ggml_backend_buffer_get_alloc_size(weight->buffer, weight) && + weight->view_src; + if (bad_padding_clear) { + return false; + } + } + + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const int64_t ncols = ids ? src1->ne[2] : src1->ne[1]; + if (ggml_cuda_should_use_mmq(src0->type, cc, ncols, src0->ne[2])) { + ggml_cuda_mul_mat_q_pair( + ctx, up->src[0], gate->src[0], src1, ids, up, gate); + ggml_cuda_op_swiglu_ds4(ctx, glu); + return true; + } + } + + return false; +} + static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft); @@ -2850,6 +2930,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 +2952,18 @@ 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_DS4_INDEXER_QAT: + ggml_cuda_op_ds4_indexer_qat(ctx, dst); + break; + case GGML_OP_DS4_INDEXER_SCORE: + ggml_cuda_op_ds4_indexer_score(ctx, dst); + break; + case GGML_OP_DS4_INDEXER_MASK: + ggml_cuda_op_ds4_indexer_mask(ctx, dst); + break; case GGML_OP_GROUP_NORM: ggml_cuda_op_group_norm(ctx, dst); break; @@ -3566,6 +3661,12 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, std::initializer_list mul_mat_id_glu_ops = { GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU }; std::initializer_list mul_mat_glu_ops = { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU }; + std::initializer_list mul_mat_id_reshape_glu_ops = { + GGML_OP_MUL_MAT_ID, GGML_OP_RESHAPE, + GGML_OP_MUL_MAT_ID, GGML_OP_RESHAPE, GGML_OP_GLU }; + std::initializer_list mul_mat_reshape_glu_ops = { + GGML_OP_MUL_MAT, GGML_OP_RESHAPE, + GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_GLU }; if ((is_equal(mul_mat_bias_glu_ops, ops) || is_equal(mul_mat_id_bias_glu_ops, ops)) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 4 })) { @@ -3593,6 +3694,25 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, } } + if ((is_equal(mul_mat_id_reshape_glu_ops, ops) || is_equal(mul_mat_reshape_glu_ops, ops)) && + ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 4 })) { + const ggml_tensor * ffn_gate = cgraph->nodes[node_idx]; + const ggml_tensor * ffn_gate_glu = cgraph->nodes[node_idx + 1]; + const ggml_tensor * ffn_up = cgraph->nodes[node_idx + 2]; + const ggml_tensor * ffn_up_glu = cgraph->nodes[node_idx + 3]; + const ggml_tensor * glu = cgraph->nodes[node_idx + 4]; + + if (ffn_gate_glu->src[0] == ffn_gate && + ffn_up_glu->src[0] == ffn_up && + ggml_cuda_should_fuse_mul_mat( + ffn_up, ffn_gate, glu, nullptr, nullptr, + ffn_up_glu, ffn_gate_glu)) { + int out_nodes[] = { node_idx + 4 }; + return ggml_cuda_check_fusion_memory_ranges( + cgraph, node_idx, (int) ops.size(), out_nodes, 1); + } + } + std::initializer_list rope_set_rows_ops = { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }; if (is_equal(rope_set_rows_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) { @@ -4071,7 +4191,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,13 +4204,39 @@ 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; fused_node_count = 5; break; } + } else if (ggml_cuda_can_fuse( + cgraph, i, + { op, GGML_OP_RESHAPE, op, GGML_OP_RESHAPE, GGML_OP_GLU }, {})) { + ggml_tensor * glu = cgraph->nodes[i + 4]; + ggml_tensor * gate_view = glu->src[0]; + ggml_tensor * up_view = glu->src[1]; + ggml_tensor * gate = nullptr; + ggml_tensor * up = nullptr; + + if (gate_view == cgraph->nodes[i + 1] && + up_view == cgraph->nodes[i + 3]) { + gate = cgraph->nodes[i]; + up = cgraph->nodes[i + 2]; + } else if (gate_view == cgraph->nodes[i + 3] && + up_view == cgraph->nodes[i + 1]) { + gate = cgraph->nodes[i + 2]; + up = cgraph->nodes[i]; + } else { + continue; + } + + if (ggml_cuda_try_fuse_mul_mat_glu(*cuda_ctx, gate, up, glu)) { + fused_mul_mat_vec = true; + fused_node_count = 5; + break; + } } else if (ggml_cuda_can_fuse(cgraph, i, { op, op, GGML_OP_GLU }, {})) { ggml_tensor * glu = cgraph->nodes[i + 2]; ggml_tensor * gate = glu->src[0]; @@ -4101,27 +4247,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud if (!ok) continue; - const ggml_tensor * src0 = up->src[0]; - const ggml_tensor * src1 = up->src[1]; - const ggml_tensor * ids = up->src[2]; - - if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate->src[0]; - fusion_data.glu_op = ggml_get_glu_op(glu); - - ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 3; - break; - } - - if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate->src[0]; - fusion_data.glu_op = ggml_get_glu_op(glu); - - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + if (ggml_cuda_try_fuse_mul_mat_glu(*cuda_ctx, gate, up, glu)) { fused_mul_mat_vec = true; fused_node_count = 3; break; @@ -4987,6 +5113,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 +5127,26 @@ 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_DS4_INDEXER_QAT: + return op->src[0]->type == GGML_TYPE_F32 && + op->src[0]->ne[0] == 128 && + ggml_is_contiguous(op->src[0]); + case GGML_OP_DS4_INDEXER_SCORE: + return op->src[0]->type == GGML_TYPE_F32 && + op->src[0]->ne[0] == 128 && + op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F16 && + op->src[2]->ne[0] == 128 && + ggml_is_contiguous(op->src[0]) && + ggml_is_contiguous(op->src[1]) && + ggml_is_contiguous(op->src[2]); + case GGML_OP_DS4_INDEXER_MASK: + return op->src[0]->type == GGML_TYPE_F32 && + op->src[1]->type == GGML_TYPE_I32 && + ggml_is_contiguous(op->src[0]) && + ggml_is_contiguous(op->src[1]); case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_ID: { @@ -5045,6 +5192,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 +5233,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 +5345,11 @@ 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 || + src0_type == GGML_TYPE_I32; } break; case GGML_OP_CONV_TRANSPOSE_1D: { @@ -5276,10 +5439,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return ggml_is_contiguous_rows(op->src[0]); case GGML_OP_TOP_K: case GGML_OP_ARGSORT: -#ifndef GGML_CUDA_USE_CUB - return op->src[0]->ne[0] <= 1024; -#else +#if defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) return true; +#else + return op->src[0]->ne[0] <= 1024; #endif case GGML_OP_SUM_ROWS: case GGML_OP_MEAN: diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu index 06dfae700..a2c06a20b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cu @@ -20,6 +20,15 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con case GGML_TYPE_Q8_0: mul_mat_q_case(ctx, args, stream); break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q2_0_ROCMFP2: + mul_mat_q_case(ctx, args, stream); + break; + case GGML_TYPE_Q3_0_ROCMFPX: + mul_mat_q_case(ctx, args, stream); + break; case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: #ifndef GGML_CUDA_BLACKWELL_CONSUMER @@ -77,11 +86,26 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con } } -void ggml_cuda_mul_mat_q( - ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) { +static void ggml_cuda_mul_mat_q_impl( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, + const ggml_tensor * src0_pair, + const ggml_tensor * src1, + const ggml_tensor * ids, + ggml_tensor * dst, + ggml_tensor * dst_pair) { GGML_ASSERT( src1->type == GGML_TYPE_F32); GGML_ASSERT( dst->type == GGML_TYPE_F32); GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. + GGML_ASSERT((src0_pair == nullptr) == (dst_pair == nullptr)); + if (src0_pair) { + GGML_ASSERT(src0_pair->type == src0->type); + GGML_ASSERT(dst_pair->type == dst->type); + GGML_ASSERT(ggml_are_same_shape(src0_pair, src0)); + GGML_ASSERT(ggml_are_same_stride(src0_pair, src0)); + GGML_ASSERT(ggml_are_same_shape(dst_pair, dst)); + GGML_ASSERT(ggml_are_same_stride(dst_pair, dst)); + } GGML_TENSOR_BINARY_OP_LOCALS; @@ -101,14 +125,23 @@ void ggml_cuda_mul_mat_q( const float * src1_d = (const float *) src1->data; float * dst_d = (float *) dst->data; - // If src0 is a temporary compute buffer, clear any potential padding. - if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { - const size_t size_data = ggml_nbytes(src0); - const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0); + // Temporary weight tensors may have an allocation tail read by tiled MMQ. + // Clear it once for each projection before launching either multiply. + const ggml_tensor * weights[] = {src0, src0_pair}; + for (const ggml_tensor * weight : weights) { + if (!weight || + ggml_backend_buffer_get_usage(weight->buffer) != + GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + continue; + } + const size_t size_data = ggml_nbytes(weight); + const size_t size_alloc = + ggml_backend_buffer_get_alloc_size(weight->buffer, weight); if (size_alloc > size_data) { - GGML_ASSERT(ggml_is_contiguously_allocated(src0)); - GGML_ASSERT(!src0->view_src); - CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream)); + GGML_ASSERT(ggml_is_contiguously_allocated(weight)); + GGML_ASSERT(!weight->view_src); + CUDA_CHECK(cudaMemsetAsync((char *) weight->data + size_data, 0, + size_alloc - size_data, stream)); } } @@ -121,21 +154,24 @@ void ggml_cuda_mul_mat_q( const int64_t s03 = src0->nb[3] / ts_src0; const int64_t s3 = dst->nb[3] / ts_dst; - bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) - || GGML_CUDA_CC_IS_CDNA(cc); - // LUCE_MMQ_DP_MAX_NE1: for ne1 at or below this, skip stream-k and use - // data-parallel tiles - the stream-k fixup pass costs ~1ms/step at - // spec-decode verify widths (measured sm_86, w6 chain). 0 = always stream-k. - static const int luce_mmq_dp_max_ne1 = []() { - const char * e = getenv("LUCE_MMQ_DP_MAX_NE1"); - return e ? atoi(e) : 0; - }(); - if (use_stream_k && ne11 <= luce_mmq_dp_max_ne1) { - use_stream_k = false; - } + const bool use_stream_k = + (GGML_CUDA_CC_IS_NVIDIA(cc) && + ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) || + GGML_CUDA_CC_IS_CDNA(cc); // TODO: tighter pool buffer size vs q8 path const bool use_native_mxfp4 = blackwell_mma_available(cc) && src0->type == GGML_TYPE_MXFP4; + const bool grouped_src = !ids && ggml_mul_mat_is_grouped_src(dst); + const ggml_tensor * grouped_physical = grouped_src ? src1->view_src : nullptr; + if (grouped_src) { + GGML_ASSERT(grouped_physical && grouped_physical->type == GGML_TYPE_F32); + GGML_ASSERT(grouped_physical->ne[2] == + ggml_mul_mat_grouped_src_groups(dst)); + GGML_ASSERT(grouped_physical->ne[0] * grouped_physical->ne[2] == ne10); + GGML_ASSERT(grouped_physical->ne[1] == ne11); + GGML_ASSERT(grouped_physical->ne[3] == 1); + GGML_ASSERT(!use_native_mxfp4); + } if (!ids) { const size_t nbytes_src1_q8_1 = ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1 + @@ -146,7 +182,14 @@ void ggml_cuda_mul_mat_q( const int64_t s11 = src1->nb[1] / ts_src1; const int64_t s12 = src1->nb[2] / ts_src1; const int64_t s13 = src1->nb[3] / ts_src1; - if (use_native_mxfp4) { + if (grouped_src) { + quantize_mmq_q8_1_grouped_cuda( + src1_d, src1_q8_1.get(), src0->type, + ne10, grouped_physical->ne[0], + grouped_physical->nb[1] / ts_src1, + grouped_physical->nb[2] / ts_src1, + ne10_padded, ne11, stream); + } else if (use_native_mxfp4) { static_assert(sizeof(block_fp4_mmq) == 4 * sizeof(block_q8_1)); quantize_mmq_mxfp4_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); @@ -173,6 +216,12 @@ void ggml_cuda_mul_mat_q( ne03, ne13, s03, s13, s3, use_stream_k, ne1}; ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); + if (src0_pair) { + mmq_args pair_args = args; + pair_args.x = (const char *) src0_pair->data; + pair_args.dst = (float *) dst_pair->data; + ggml_cuda_mul_mat_q_switch_type(ctx, pair_args, stream); + } return; } @@ -238,6 +287,33 @@ void ggml_cuda_mul_mat_q( use_stream_k, ne12}; ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); + if (src0_pair) { + mmq_args pair_args = args; + pair_args.x = (const char *) src0_pair->data; + pair_args.dst = (float *) dst_pair->data; + ggml_cuda_mul_mat_q_switch_type(ctx, pair_args, stream); + } +} + +void ggml_cuda_mul_mat_q( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, + const ggml_tensor * src1, + const ggml_tensor * ids, + ggml_tensor * dst) { + ggml_cuda_mul_mat_q_impl(ctx, src0, nullptr, src1, ids, dst, nullptr); +} + +void ggml_cuda_mul_mat_q_pair( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0_a, + const ggml_tensor * src0_b, + const ggml_tensor * src1, + const ggml_tensor * ids, + ggml_tensor * dst_a, + ggml_tensor * dst_b) { + ggml_cuda_mul_mat_q_impl( + ctx, src0_a, src0_b, src1, ids, dst_a, dst_b); } void ggml_cuda_op_mul_mat_q( @@ -316,6 +392,12 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t case GGML_TYPE_IQ4_NL: mmq_supported = true; break; + case GGML_TYPE_Q2_0_ROCMFP2: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + // ROCmFPX MMQ variants are implemented for gfx1151 only. + mmq_supported = GGML_CUDA_CC_IS_RDNA3_5(cc); + break; default: mmq_supported = false; break; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh index d8df400fb..ad26af456 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmq.cuh @@ -66,6 +66,10 @@ static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) { return MMQ_Q8_1_DS_LAYOUT_DS4; case GGML_TYPE_Q8_0: return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q2_0_ROCMFP2: + case GGML_TYPE_Q3_0_ROCMFPX: + return MMQ_Q8_1_DS_LAYOUT_D4; case GGML_TYPE_MXFP4: return MMQ_Q8_1_DS_LAYOUT_D4; case GGML_TYPE_NVFP4: @@ -101,24 +105,9 @@ struct tile_x_sizes { int sc; }; -// Lucebox: tuned MMQ tile for RDNA3/RDNA4. -// DFlash issues many small mul_mat_q calls (ne[1] ~= ddtree_budget+1, ~23 at -// the default budget=22) where the stock 128x128 / 8-warp tile under-occupies. -// Current settings (mmq_y=128, nwarps=8, mmq_x_max=128): -// - mmq_y=128/nwarps=8: stock value for non-RDNA1 AMD — occupancy for decode -// batches (ne[1]~23) measured at +8.3% vs stock on gfx1201; also the right -// size for prefill (ne[1]=512) where halving block count improves L2 reuse -// (+6.9% prefill speedup on Q4_K m=4096,n=512,k=14336 on gfx1201). -// - mmq_x_max=128: prefill (ne[1]=512) fits in 4 x-tiles instead of 11. -// For decode (ne[1]~23), mmq_x selects 24 regardless of the cap — no impact. -// Original decode measurements (Qwen3.6-27B Q4_K_M, --ddtree-budget=22): -// gfx1100 (RX 7900 XTX): 56.78 -> 60.18 tok/s (+6.0%) -// gfx1201 (R9700): 54.65 -> 59.20 tok/s (+8.3%) -// gfx1151 (Strix Halo): 11.53 -> 12.00 tok/s (+4.1%, 256-token smoke) -// Constraint: nwarps * granularity(mmq_x) / ntx(mmq_x) == mmq_y. -// mmq_x< 128: granularity=16, ntx=1 → 8*16/1=128 ✓ -// mmq_x>=128: granularity=32, ntx=2 → 8*32/2=128 ✓ -// Define LUCEBOX_RDNA_MMQ_TILE_OVERRIDE=0 to disable. +// RDNA uses 128x128, eight-warp MMQ tiles by default. ROCmFPX template +// instances use 64x64, four-warp tiles: their unpacking pressure makes the +// smaller tile faster on gfx1151 without changing other quant formats. #ifndef LUCEBOX_RDNA_MMQ_TILE_OVERRIDE #define LUCEBOX_RDNA_MMQ_TILE_OVERRIDE 1 #endif @@ -130,7 +119,13 @@ struct tile_x_sizes { #endif static int get_mmq_x_max_host(const int cc) { - if (LUCEBOX_RDNA_TILE_HOST(cc)) return 128; + if (LUCEBOX_RDNA_TILE_HOST(cc)) { +#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) + return 64; +#else + return 128; +#endif + } return (amd_mfma_available(cc) || turing_mma_available(cc) || amd_wmma_available(cc)) ? 128 : GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA ? #ifdef GGML_CUDA_FORCE_MMQ @@ -142,8 +137,12 @@ static int get_mmq_x_max_host(const int cc) { static constexpr __device__ int get_mmq_x_max_device() { #if LUCEBOX_RDNA_TILE_DEVICE +#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) + return 64; +#else return 128; #endif +#endif #if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) return 128; #else // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) @@ -167,7 +166,13 @@ static constexpr __device__ int get_mmq_x_max_device() { } static int get_mmq_y_host(const int cc) { - if (LUCEBOX_RDNA_TILE_HOST(cc)) return 128; + if (LUCEBOX_RDNA_TILE_HOST(cc)) { +#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) + return 64; +#else + return 128; +#endif + } return GGML_CUDA_CC_IS_AMD(cc) ? (GGML_CUDA_CC_IS_RDNA1(cc) ? 64 : 128) : ((GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64); } @@ -182,8 +187,12 @@ static constexpr __device__ int get_iter_k([[maybe_unused]] const ggml_type type static constexpr __device__ int get_mmq_y_device() { #if LUCEBOX_RDNA_TILE_DEVICE +#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) + return 64; +#else return 128; #endif +#endif #if defined(GGML_USE_HIP) #if defined(RDNA1) return 64; @@ -226,6 +235,9 @@ static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml case GGML_TYPE_Q5_0: return MMQ_DP4A_TXS_Q8_0; case GGML_TYPE_Q5_1: return MMQ_DP4A_TXS_Q8_1; case GGML_TYPE_Q8_0: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_Q2_0_ROCMFP2: return MMQ_DP4A_TXS_Q8_0_16; + case GGML_TYPE_Q3_0_ROCMFPX: return MMQ_DP4A_TXS_Q8_0_16; case GGML_TYPE_MXFP4: return MMQ_DP4A_TXS_Q8_1; case GGML_TYPE_NVFP4: return MMQ_DP4A_TXS_Q8_0_16; case GGML_TYPE_Q2_K: return MMQ_DP4A_TXS_Q2_K; @@ -270,6 +282,9 @@ static constexpr __host__ __device__ int mmq_get_mma_tile_x_k(ggml_type type) { case GGML_TYPE_Q5_0: return MMQ_MMA_TILE_X_K_Q8_0; case GGML_TYPE_Q5_1: return MMQ_MMA_TILE_X_K_Q8_1; case GGML_TYPE_Q8_0: return MMQ_MMA_TILE_X_K_Q8_0; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: return MMQ_MMA_TILE_X_K_Q8_0; + case GGML_TYPE_Q2_0_ROCMFP2: return MMQ_MMA_TILE_X_K_Q3_K; + case GGML_TYPE_Q3_0_ROCMFPX: return MMQ_MMA_TILE_X_K_Q3_K; // tile sizes are the same for Q8_1 and FP4 for blackwell case GGML_TYPE_MXFP4: return MMQ_MMA_TILE_X_K_Q8_1; case GGML_TYPE_NVFP4: return MMQ_MMA_TILE_X_K_NVFP4; @@ -320,7 +335,13 @@ static constexpr __device__ int mmq_get_granularity_device(const int /*mmq_x*/) #if defined(GGML_USE_HIP) static int mmq_get_nwarps_host(const int cc, const int warp_size) { - if (LUCEBOX_RDNA_TILE_HOST(cc)) return 8; + if (LUCEBOX_RDNA_TILE_HOST(cc)) { +#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) + return 4; +#else + return 8; +#endif + } return amd_mfma_available(cc) ? 8 : 256/warp_size; } #else @@ -331,8 +352,12 @@ static int mmq_get_nwarps_host(const int /*cc*/, const int warp_size) { static constexpr __device__ int mmq_get_nwarps_device() { #if LUCEBOX_RDNA_TILE_DEVICE +#if defined(GGML_CUDA_ROCMFPX_MMQ_TILE) + return 4; +#else return 8; #endif +#endif #if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) return 8; #else @@ -845,6 +870,178 @@ template static __device__ __forceinline__ void loa } } +// Expand packed ROCmFP4 Codebook10 weights into the signed int8 layout used by +// the existing Q8_0 x Q8_1 MMQ kernels. Q4_0_ROCMFP4_FAST has one UE4M3 +// scale per 32-weight block, so after expansion its shared-memory shape is +// identical to Q8_0: 256 int8 weights and eight float scales per row/tile. +template static __device__ __forceinline__ void load_tiles_rocmfp4_fast( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0_ROCMFP4_FAST, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR_ROCMFP4); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI_ROCMFP4; + const int kqsx = txi % QI_ROCMFP4; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_rocmfp4_fast * bxi = (const block_rocmfp4_fast *) x + kbx0 + i*stride + kbx; + const int aux_q4 = rocmfp4_get_qs_i32(bxi->qs, kqsx); + const int2 v = rocmfp4_get_int_from_codebook_16(aux_q4, kvalues_rocmfp4); + const int k0 = kbx * (2 * QI_ROCMFP4) + kqsx; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + k0 + 0] = v.x; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + k0 + QI_ROCMFP4] = v.y; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI_ROCMFP4] = v.y; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = 2*MMQ_TILE_NE_K / QI8_0; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) { + int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_rocmfp4_fast * bxi = (const block_rocmfp4_fast *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kbxd] = rocmfp4_ue4m3_to_fp32_half_finite(bxi->e); +#else + x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + kbxd] = rocmfp4_ue4m3_to_fp32_half_finite(bxi->e); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template +struct rocmfpx_dual_mmq_traits; + +template <> +struct rocmfpx_dual_mmq_traits { + using block_t = block_rocmfp2; + + static __device__ __forceinline__ int pack4(const block_t * block, const int base) { + return rocmfpx_pack4_fp2_vec_cuda(block->qs, base); + } + + static __device__ __forceinline__ float scale(const block_t * block, const int half) { + return rocmfpx_ue4m3_to_fp32_finite(block->e[half]); + } +}; + +template <> +struct rocmfpx_dual_mmq_traits { + using block_t = block_rocmfp3; + + static __device__ __forceinline__ int pack4(const block_t * block, const int base) { + return rocmfpx_pack4_fp3_vec_cuda(block->qs, base); + } + + static __device__ __forceinline__ float scale(const block_t * block, const int half) { + return rocmfpx_ue4m3_to_fp32_finite(block->e[half]); + } +}; + +// Q2/FP3 store one UE4M3 scale per 16 weights. Expand their packed values +// into int8 and reuse the existing scale-per-16 Q8_0 x Q8_1 MMQ kernels. One +// wave loads a full 256-weight row: each lane expands two groups of four. +template +static __device__ __forceinline__ void load_tiles_rocmfpx_dual( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + using traits = rocmfpx_dual_mmq_traits; + using block_t = typename traits::block_t; + + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int groups_per_block = QK_ROCMFPX / 4; + constexpr int blocks_per_tile = MMQ_ITER_K / QK_ROCMFPX; + constexpr int threads_per_row = blocks_per_tile * groups_per_block / 2; + static_assert(threads_per_row == 32, "ROCmFPX MMQ loader expects 32 lanes per row"); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(type, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / (groups_per_block / 2); + const int group = txi % (groups_per_block / 2); + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + if (need_check) { + i = min(i, i_max); + } + + const block_t * block = (const block_t *) x + kbx0 + i*stride + kbx; + const int k0 = kbx*groups_per_block + group; + const int q0 = traits::pack4(block, 4*group); + const int q1 = traits::pack4( + block, 4*(group + groups_per_block/2)); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + k0] = q0; + x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + k0 + groups_per_block/2] = q1; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0] = q0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + groups_per_block/2] = q1; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + } + + constexpr int scales_per_tile = 2*blocks_per_tile; + constexpr int scale_rows_per_warp = warp_size / scales_per_tile; + const int kscale = threadIdx.x % scales_per_tile; + const int scale_block = kscale / 2; + const int scale_half = kscale % 2; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*scale_rows_per_warp) { + int i = i0 + threadIdx.y*scale_rows_per_warp + threadIdx.x/scales_per_tile; + if (need_check) { + i = min(i, i_max); + } + + const block_t * block = (const block_t *) x + kbx0 + i*stride + scale_block; +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q3_K + kscale] = traits::scale(block, scale_half); +#else + x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + kscale] = traits::scale(block, scale_half); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + template static __device__ __forceinline__ void load_tiles_mxfp4_fp4(const char * __restrict__ x, int * __restrict__ x_tile, @@ -3392,6 +3589,30 @@ struct mmq_type_traits { static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; }; +template +struct mmq_type_traits { + static constexpr int vdr = VDR_ROCMFP4_FAST_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_rocmfp4_fast; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_ROCMFP2_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_rocmfpx_dual; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_16_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_16_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_ROCMFP3_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_rocmfpx_dual; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_16_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_16_q8_1_dp4a; +}; + template struct mmq_type_traits { static constexpr int vdr = VDR_MXFP4_Q8_1_MMQ; @@ -3623,7 +3844,7 @@ template #if defined(GGML_USE_HIP) // RDNA4 is compute-bound on MMQ (WMMA path); allow compiler to use more VGPRs // (minBlocks=1 matches NVIDIA Volta+ behavior and reduces register spilling). -#if defined(RDNA4) +#if defined(RDNA4) && !defined(GGML_CUDA_ROCMFPX_MMQ_TILE) __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 1) #elif defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN) __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 2) @@ -4242,6 +4463,9 @@ extern DECL_MMQ_CASE(GGML_TYPE_Q4_1); extern DECL_MMQ_CASE(GGML_TYPE_Q5_0); extern DECL_MMQ_CASE(GGML_TYPE_Q5_1); extern DECL_MMQ_CASE(GGML_TYPE_Q8_0); +extern DECL_MMQ_CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST); +extern DECL_MMQ_CASE(GGML_TYPE_Q2_0_ROCMFP2); +extern DECL_MMQ_CASE(GGML_TYPE_Q3_0_ROCMFPX); extern DECL_MMQ_CASE(GGML_TYPE_MXFP4); extern DECL_MMQ_CASE(GGML_TYPE_NVFP4); extern DECL_MMQ_CASE(GGML_TYPE_Q2_K); @@ -4263,6 +4487,19 @@ extern DECL_MMQ_CASE(GGML_TYPE_IQ4_XS); void ggml_cuda_mul_mat_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst); +// Execute two same-shape MUL_MAT operations with one activation quantization. +// When ids is non-null, also share the MUL_MAT_ID sort. The two MMQ launches +// and their output order remain unchanged, so callers can apply the ordinary +// GLU kernel byte-identically. +void ggml_cuda_mul_mat_q_pair( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0_a, + const ggml_tensor * src0_b, + const ggml_tensor * src1, + const ggml_tensor * ids, + ggml_tensor * dst_a, + ggml_tensor * dst_b); + void ggml_cuda_op_mul_mat_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/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..b181b6fe6 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; @@ -403,7 +415,32 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int return 1; } -template +static bool rocmfp2_fixed_k_enabled() { + static const bool enabled = []() { + const char * value = std::getenv("DFLASH_ROCMFP2_FIXED_K"); + return value && value[0] == '1' && value[1] == '\0'; + }(); + return enabled; +} + +static bool rocmfp3_fixed_k_enabled() { + static const bool enabled = []() { + const char * value = std::getenv("DFLASH_ROCMFP3_FIXED_K"); + return value && value[0] == '1' && value[1] == '\0'; + }(); + return enabled; +} + +static bool rocmfp4_unroll2_enabled() { + static const bool enabled = []() { + const char * value = std::getenv("DFLASH_ROCMFP4_UNROLL2"); + return value && value[0] == '1' && value[1] == '\0'; + }(); + return enabled; +} + +template __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id())*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q( const void * __restrict__ vx, const void * __restrict__ vy, const int32_t * __restrict__ ids, const ggml_cuda_mm_fusion_args_device fusion, float * __restrict__ dst, @@ -425,9 +462,16 @@ static __global__ void mul_mat_vec_q( const int tid = warp_size*threadIdx.y + threadIdx.x; const int row0 = rows_per_cuda_block*blockIdx.x; - const int blocks_per_row_x = ncols_x / qk; + const int blocks_per_row_x = (fixed_ncols_x > 0 ? fixed_ncols_x : ncols_x) / qk; constexpr int blocks_per_iter = vdr * nwarps*warp_size / qi; + static_assert(fixed_ncols_x == 0 || + type == GGML_TYPE_Q3_0_ROCMFPX || + type == GGML_TYPE_Q2_0_ROCMFP2, + "fixed-K decode specialization is limited to profiled ROCmFP2/FP3 types"); + static_assert(!unroll_k_loop_2 || type == GGML_TYPE_Q4_0_ROCMFP4_FAST, + "partial K-loop unrolling is limited to the profiled FP4FAST path"); + const uint32_t channel_dst = blockIdx.y; const bool has_ids = ids != nullptr; @@ -454,6 +498,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 +509,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; } @@ -503,55 +551,127 @@ static __global__ void mul_mat_vec_q( const block_q8_1 * y = ((const block_q8_1 *) vy) + sample_y*stride_sample_y + channel_y*stride_channel_y; - for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { - const int kby = kbx * (qk/QK8_1); // y block index that aligns with kbx + if constexpr (fixed_ncols_x > 0) { + static_assert(fixed_ncols_x % qk == 0, "fixed K must contain whole quant blocks"); + constexpr int fixed_blocks = fixed_ncols_x/qk; + constexpr int fixed_iters = (fixed_blocks + blocks_per_iter - 1) / blocks_per_iter; +#pragma unroll + for (int iter = 0; iter < fixed_iters; ++iter) { + const int kbx = tid / (qi/vdr) + iter*blocks_per_iter; + if constexpr (fixed_blocks % blocks_per_iter != 0) { + if (kbx >= fixed_blocks) { + continue; + } + } + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (tid % (qi/vdr)); - // x block quant index when casting the quants to int - const int kqs = vdr * (tid % (qi/vdr)); +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { + const int kbx_offset = sample_x*stride_sample_x + channel_xs[j]*stride_channel_x + row0*stride_row_x; +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp[j][i] += vec_dot_q_cuda( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[j][i] += vec_dot_q_cuda( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } + } + } + } + } + } else if constexpr (unroll_k_loop_2) { +#pragma unroll 2 + for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (tid % (qi/vdr)); #pragma unroll - for (int j = 0; j < ncols_dst; ++j) { - const int kbx_offset = sample_x*stride_sample_x + channel_xs[j]*stride_channel_x + row0*stride_row_x; + for (int j = 0; j < ncols_dst; ++j) { + const int kbx_offset = sample_x*stride_sample_x + channel_xs[j]*stride_channel_x + row0*stride_row_x; #pragma unroll - for (int i = 0; i < rows_per_cuda_block; ++i) { - tmp[j][i] += vec_dot_q_cuda( - vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); - if constexpr (has_fusion) { - if (use_gate) { - tmp_gate[j][i] += vec_dot_q_cuda( - vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp[j][i] += vec_dot_q_cuda( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[j][i] += vec_dot_q_cuda( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } } } } } - } + } else { + for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (tid % (qi/vdr)); - __shared__ float tmp_shared[nwarps-1 > 0 ? nwarps-1 : 1][ncols_dst][rows_per_cuda_block][warp_size]; - __shared__ float tmp_shared_gate[(has_fusion && (nwarps-1 > 0)) ? nwarps-1 : 1][ncols_dst][rows_per_cuda_block][warp_size]; - if constexpr (!has_fusion) { - (void) tmp_shared_gate; - } else if (!use_gate) { - (void) tmp_shared_gate; +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { + const int kbx_offset = sample_x*stride_sample_x + channel_xs[j]*stride_channel_x + row0*stride_row_x; +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp[j][i] += vec_dot_q_cuda( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[j][i] += vec_dot_q_cuda( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } + } + } + } + } } - if (threadIdx.y > 0) { + if constexpr (nwarps > 1) { + __shared__ float tmp_shared[nwarps-1][ncols_dst][rows_per_cuda_block][warp_size]; + __shared__ float tmp_shared_gate[has_fusion ? nwarps-1 : 1][ncols_dst][rows_per_cuda_block][warp_size]; + if constexpr (!has_fusion) { + (void) tmp_shared_gate; + } else if (!use_gate) { + (void) tmp_shared_gate; + } + + if (threadIdx.y > 0) { +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp_shared[threadIdx.y-1][j][i][threadIdx.x] = tmp[j][i]; + if constexpr (has_fusion) { + if (use_gate) { + tmp_shared_gate[threadIdx.y-1][j][i][threadIdx.x] = tmp_gate[j][i]; + } + } + } + } + } + + __syncthreads(); + if (threadIdx.y > 0) { + return; + } + #pragma unroll for (int j = 0; j < ncols_dst; ++j) { #pragma unroll for (int i = 0; i < rows_per_cuda_block; ++i) { - tmp_shared[threadIdx.y-1][j][i][threadIdx.x] = tmp[j][i]; - if constexpr (has_fusion) { - if (use_gate) { - tmp_shared_gate[threadIdx.y-1][j][i][threadIdx.x] = tmp_gate[j][i]; +#pragma unroll + for (int l = 0; l < nwarps-1; ++l) { + tmp[j][i] += tmp_shared[l][j][i][threadIdx.x]; + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate[j][i] += tmp_shared_gate[l][j][i][threadIdx.x]; + } } } } } } - __syncthreads(); - if (threadIdx.y > 0) { - return; - } dst += sample_dst*stride_sample_dst + channel_dst*stride_channel_dst + row0; @@ -560,15 +680,6 @@ static __global__ void mul_mat_vec_q( for (int j = 0; j < ncols_dst; ++j) { #pragma unroll for (int i = 0; i < rows_per_cuda_block; ++i) { -#pragma unroll - for (int l = 0; l < nwarps-1; ++l) { - tmp[j][i] += tmp_shared[l][j][i][threadIdx.x]; - if constexpr (has_fusion) { - if (use_gate) { - tmp_gate[j][i] += tmp_shared_gate[l][j][i][threadIdx.x]; - } - } - } tmp[j][i] = warp_reduce_sum(tmp[j][i]); if constexpr (has_fusion) { if (use_gate) { @@ -596,7 +707,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 +725,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 +775,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 +786,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 +840,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 +855,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); } } @@ -749,7 +871,8 @@ static std::pair calc_launch_params( return {block_nums, block_dims}; } -template +template static void mul_mat_vec_q_switch_fusion( const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, @@ -761,19 +884,70 @@ static void mul_mat_vec_q_switch_fusion( const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; if (has_fusion) { - mul_mat_vec_q<<>> + mul_mat_vec_q<<>> (vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, ids_tokenwise_samples); return; } - mul_mat_vec_q<<>> + mul_mat_vec_q<<>> (vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, ids_tokenwise_samples); } +template +static void mul_mat_vec_rocmfpx_fixed_k_launch( + const void * vx, const void * vy, const int32_t * ids, + const ggml_cuda_mm_fusion_args_device fusion, float * dst, + const uint32_t ncols_x, const uint32_t nrows_x, const uint3 nchannels_y, + const uint32_t nchannels_dst, const uint32_t stride_row_x, + const uint32_t stride_col_y, const uint32_t stride_col_dst, + const uint3 channel_ratio, const uint32_t stride_channel_x, + const uint32_t stride_channel_y, const uint32_t stride_channel_dst, + const uint32_t nsamples_dst, const uint3 sample_ratio, + const uint32_t stride_sample_x, const uint32_t stride_sample_y, + const uint32_t stride_sample_dst, const uint32_t ids_stride, + const int warp_size, cudaStream_t stream) { + static_assert(type == GGML_TYPE_Q3_0_ROCMFPX || + type == GGML_TYPE_Q2_0_ROCMFP2, + "fixed-K helper is limited to profiled ROCmFP2/FP3 kernels"); + static_assert(fixed_ncols_x == 2048 || fixed_ncols_x == 4096, + "fixed-K helper compiled an unexpected DS4 shape"); + + const dim3 block_nums(nrows_x, nchannels_dst, nsamples_dst); + const dim3 block_dims(warp_size, 1, 1); + mul_mat_vec_q_switch_fusion( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, + stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, + stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, + stride_sample_y, stride_sample_dst, block_nums, block_dims, 0, + ids_stride, stream); +} + +static void mul_mat_vec_rocmfp4_unroll2_launch( + const void * vx, const void * vy, const int32_t * ids, + const ggml_cuda_mm_fusion_args_device fusion, float * dst, + const uint32_t ncols_x, const uint32_t nrows_x, const uint3 nchannels_y, + const uint32_t nchannels_dst, const uint32_t stride_row_x, + const uint32_t stride_col_y, const uint32_t stride_col_dst, + const uint3 channel_ratio, const uint32_t stride_channel_x, + const uint32_t stride_channel_y, const uint32_t stride_channel_dst, + const uint32_t nsamples_dst, const uint3 sample_ratio, + const uint32_t stride_sample_x, const uint32_t stride_sample_y, + const uint32_t stride_sample_dst, const uint32_t ids_stride, + const int warp_size, cudaStream_t stream) { + const dim3 block_nums(nrows_x, nchannels_dst, nsamples_dst); + const dim3 block_dims(warp_size, 1, 1); + mul_mat_vec_q_switch_fusion( + vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, + stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, + stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, + stride_sample_y, stride_sample_dst, block_nums, block_dims, 0, + ids_stride, stream); +} + template static void mul_mat_vec_q_moe_launch( const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, @@ -944,6 +1118,61 @@ static void mul_mat_vec_q_switch_ncols_dst( return; } + // Decode-only gfx1151 specializations. Each one preserves the original + // per-lane K traversal and accumulation order; the environment flags keep + // the optimized kernels independently bisectable. + if constexpr (type == GGML_TYPE_Q4_0_ROCMFP4_FAST) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && ncols_dst == 1 && rocmfp4_unroll2_enabled()) { + mul_mat_vec_rocmfp4_unroll2_launch( + vx, vy, ids, fusion, dst, ncols_x, nrows_x, nchannels_y_fd, + nchannels_dst, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, + stride_channel_dst, nsamples_dst, sample_ratio_fd, + stride_sample_x, stride_sample_y, stride_sample_dst, + ids_stride, warp_size, stream); + return; + } + } + + if constexpr (type == GGML_TYPE_Q3_0_ROCMFPX) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && ncols_dst == 1 && + ncols_x == 2048 && rocmfp3_fixed_k_enabled()) { + mul_mat_vec_rocmfpx_fixed_k_launch( + vx, vy, ids, fusion, dst, ncols_x, nrows_x, nchannels_y_fd, + nchannels_dst, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, + stride_channel_dst, nsamples_dst, sample_ratio_fd, + stride_sample_x, stride_sample_y, stride_sample_dst, + ids_stride, warp_size, stream); + return; + } + } + + if constexpr (type == GGML_TYPE_Q2_0_ROCMFP2) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && ncols_dst == 1 && rocmfp2_fixed_k_enabled()) { + if (ncols_x == 4096) { + mul_mat_vec_rocmfpx_fixed_k_launch( + vx, vy, ids, fusion, dst, ncols_x, nrows_x, nchannels_y_fd, + nchannels_dst, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, + stride_channel_dst, nsamples_dst, sample_ratio_fd, + stride_sample_x, stride_sample_y, stride_sample_dst, + ids_stride, warp_size, stream); + return; + } + if (ncols_x == 2048) { + mul_mat_vec_rocmfpx_fixed_k_launch( + vx, vy, ids, fusion, dst, ncols_x, nrows_x, nchannels_y_fd, + nchannels_dst, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, + stride_channel_dst, nsamples_dst, sample_ratio_fd, + stride_sample_x, stride_sample_y, stride_sample_dst, + ids_stride, warp_size, stream); + return; + } + } + } + switch (ncols_dst) { case 1: { constexpr int c_ncols_dst = 1; @@ -1082,6 +1311,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 +1484,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/quantize.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cu index 4300ffc14..992e12d85 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cu @@ -176,7 +176,8 @@ template static __global__ void quantize_mmq_q8_1( const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, - const int64_t ne0, const int ne1, const int ne2) { + const int64_t ne0, const int ne1, const int ne2, + const int64_t group_width, const int64_t group_stride) { constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32; constexpr int vals_per_sum = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32; @@ -204,8 +205,21 @@ static __global__ void quantize_mmq_q8_1( const int64_t ib = ib0 + (i0 / (4*QK8_1))*ne1 + blockIdx.x; // block index in channel const int64_t iqs = i0 % (4*QK8_1); // quant index in block - // Load 4 floats per thread and calculate max. abs. value between them: - const float4 xi = i0 < ne00 ? x4[(i03*s03 + i02*s02 + i01*s01 + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f); + // Load 4 floats per thread and calculate max. abs. value between them. + // A grouped source is physically [group_width, token, group] while this + // kernel emits the same flattened-K Q8 layout as a materialized permute. + int64_t src_i00 = i00; + int64_t src_group_offset = 0; + if (group_width > 0) { + const int64_t group = i00 / group_width; + src_i00 = i00 - group * group_width; + src_group_offset = group * group_stride; + } + const int64_t src_index = + i03*s03 + i02*s02 + i01*s01 + src_group_offset + src_i00; + const float4 xi = i0 < ne00 + ? x4[src_index/4] + : make_float4(0.0f, 0.0f, 0.0f, 0.0f); float amax = fabsf(xi.x); amax = fmaxf(amax, fabsf(xi.y)); amax = fmaxf(amax, fabsf(xi.z)); @@ -300,15 +314,15 @@ void quantize_mmq_q8_1_cuda( switch (mmq_get_q8_1_ds_layout(type_src0)) { case MMQ_Q8_1_DS_LAYOUT_D4: quantize_mmq_q8_1 - <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2); + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, 0); break; case MMQ_Q8_1_DS_LAYOUT_DS4: quantize_mmq_q8_1 - <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2); + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, 0); break; case MMQ_Q8_1_DS_LAYOUT_D2S6: quantize_mmq_q8_1 - <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2); + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, 0); break; default: GGML_ABORT("fatal error"); @@ -316,6 +330,45 @@ void quantize_mmq_q8_1_cuda( } } +void quantize_mmq_q8_1_grouped_cuda( + const float * x, void * vy, const ggml_type type_src0, + const int64_t ne00, const int64_t group_width, + const int64_t token_stride, const int64_t group_stride, + const int64_t ne0, const int64_t ne1, cudaStream_t stream) { + GGML_ASSERT(ne00 % 4 == 0); + GGML_ASSERT(group_width > 0 && group_width % 4 == 0); + GGML_ASSERT(ne00 % group_width == 0); + GGML_ASSERT(ne0 % (4*QK8_1) == 0); + + const int64_t block_num_y = + (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / + (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const dim3 num_blocks(ne1, block_num_y, 1); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + switch (mmq_get_q8_1_ds_layout(type_src0)) { + case MMQ_Q8_1_DS_LAYOUT_D4: + quantize_mmq_q8_1 + <<>>( + x, nullptr, vy, ne00, token_stride, 0, 0, + ne0, ne1, 1, group_width, group_stride); + break; + case MMQ_Q8_1_DS_LAYOUT_DS4: + quantize_mmq_q8_1 + <<>>( + x, nullptr, vy, ne00, token_stride, 0, 0, + ne0, ne1, 1, group_width, group_stride); + break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: + quantize_mmq_q8_1 + <<>>( + x, nullptr, vy, ne00, token_stride, 0, 0, + ne0, ne1, 1, group_width, group_stride); + break; + default: + GGML_ABORT("fatal error"); + } +} + void quantize_mmq_mxfp4_cuda(const float * x, const int32_t * ids, void * vy, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cuh index 6a91df635..4768791c4 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cuh @@ -26,6 +26,12 @@ void quantize_mmq_q8_1_cuda( ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); +void quantize_mmq_q8_1_grouped_cuda( + const float * x, void * vy, ggml_type type_src0, + int64_t ne00, int64_t group_width, + int64_t token_stride, int64_t group_stride, + int64_t ne0, int64_t ne1, cudaStream_t stream); + void quantize_mmq_mxfp4_cuda(const float * x, const int32_t * ids, void * vy, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py index 17cee853e..22a7d57f4 100755 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py @@ -33,6 +33,7 @@ TYPES_MMQ = [ "GGML_TYPE_Q4_0", "GGML_TYPE_Q4_1", "GGML_TYPE_Q5_0", "GGML_TYPE_Q5_1", "GGML_TYPE_Q8_0", + "GGML_TYPE_Q4_0_ROCMFP4_FAST", "GGML_TYPE_Q2_0_ROCMFP2", "GGML_TYPE_Q3_0_ROCMFPX", "GGML_TYPE_Q2_K", "GGML_TYPE_Q3_K", "GGML_TYPE_Q4_K", "GGML_TYPE_Q5_K", "GGML_TYPE_Q6_K", "GGML_TYPE_IQ2_XXS", "GGML_TYPE_IQ2_XS", "GGML_TYPE_IQ2_S", "GGML_TYPE_IQ3_XXS", "GGML_TYPE_IQ3_S", "GGML_TYPE_IQ1_S", "GGML_TYPE_IQ4_NL", "GGML_TYPE_IQ4_XS", "GGML_TYPE_MXFP4", "GGML_TYPE_NVFP4" diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu new file mode 100644 index 000000000..8221e1d1e --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q2_0_ROCMFP2); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu new file mode 100644 index 000000000..2380af75c --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q3_0_ROCMFPX); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu new file mode 100644 index 000000000..94a2bb0f5 --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu index 59ce36fb1..adfa25312 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu @@ -35,7 +35,7 @@ static void top_k_cub(ggml_cuda_pool & pool, ncols, k, env)); } -#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE +#elif defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) // CUB_TOP_K_AVAILABLE static int next_power_of_2(int x) { int n = 1; @@ -69,7 +69,7 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { for (int i = 0; i < nrows; i++) { top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream); } -#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE +#elif defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) // CUB_TOP_K_AVAILABLE // Fall back to argsort + copy const int ncols_pad = next_power_of_2(ncols); const size_t shared_mem = ncols_pad * sizeof(int); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/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..5ed92c8cc 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,400 @@ 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_bits12_vec_cuda(const uint32_t bits12) { +#if defined(GGML_USE_HIP) + // Byte tables for codes {0, 1, 2, 4} and {0, -1, -2, -4}. + // v_perm_b32 selects a byte with each 3-bit selector, replacing four + // scalar code decodes without changing the packed int8 values. + constexpr uint32_t values_low = 0x04020100u; + constexpr uint32_t values_high = 0xFCFEFF00u; + const uint32_t selectors = + ((bits12 >> 0) & 7u) | + (((bits12 >> 3) & 7u) << 8) | + (((bits12 >> 6) & 7u) << 16) | + (((bits12 >> 9) & 7u) << 24); + return (int) __builtin_amdgcn_perm(values_high, values_low, selectors); +#else + 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)); + return *((const int *) &v); +#endif +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp3_vec_cuda(const uint8_t * qs, const int base) { + // MMQ requests four-value groups, so the 12 packed bits span at most two + // adjacent bytes. Decode them with the same lookup used by MMVQ. + const int bit_pos = 3 * base; + const int byte_pos = bit_pos >> 3; + const uint32_t packed = (uint32_t) qs[byte_pos] | + ((uint32_t) qs[byte_pos + 1] << 8); + return rocmfpx_pack4_fp3_bits12_vec_cuda( + (packed >> (bit_pos & 7)) & 0x0fffu); +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp2_bits8_vec_cuda(const uint32_t bits8) { +#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 int val_packed = rocmfpx_pack4_fp3_bits12_vec_cuda(bits12); + + 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..7bc3fffcd 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,19 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", "TURBO_WHT", + + "MOE_FUSED", + + "DS4_HC", + + "DS4_INDEXER_QAT", + + "DS4_INDEXER_SCORE", + + "DS4_INDEXER_MASK", }; -static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); +static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1187,9 +1247,19 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", "turbo_wht(a)", + + "moe_fused(x)", + + "ds4_hc(x)", + + "ds4_indexer_qat(x)", + + "ds4_indexer_score(q,w,k)", + + "ds4_indexer_mask(x,topk)", }; -static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); +static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -1225,11 +1295,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 +1499,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 +3110,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( @@ -3253,6 +3348,38 @@ struct ggml_tensor * ggml_mul_mat( return result; } +struct ggml_tensor * ggml_mul_mat_grouped_src( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(b->type == GGML_TYPE_F32); + GGML_ASSERT(b->ne[2] > 1 && b->ne[3] == 1); + GGML_ASSERT(ggml_is_contiguous(b)); + GGML_ASSERT(a->ne[0] == b->ne[0] * b->ne[2]); + + const int64_t flat_k = b->ne[0] * b->ne[2]; + struct ggml_tensor * logical = ggml_view_2d( + ctx, b, flat_k, b->ne[1], + (size_t) flat_k * ggml_type_size(b->type), 0); + struct ggml_tensor * result = ggml_mul_mat(ctx, a, logical); + + // Keep the operation as MUL_MAT so scheduling/allocation stay generic. + // The accessors below keep the private tag out of backend code. + ggml_set_op_params_i32(result, 14, (int32_t) b->ne[2]); + ggml_set_op_params_i32(result, 15, 0x4453474d); // "DSGM" + return result; +} + +bool ggml_mul_mat_is_grouped_src(const struct ggml_tensor * tensor) { + return tensor && tensor->op == GGML_OP_MUL_MAT && + ggml_get_op_params_i32(tensor, 15) == 0x4453474d; +} + +int64_t ggml_mul_mat_grouped_src_groups(const struct ggml_tensor * tensor) { + GGML_ASSERT(ggml_mul_mat_is_grouped_src(tensor)); + return ggml_get_op_params_i32(tensor, 14); +} + void ggml_mul_mat_set_prec( struct ggml_tensor * a, enum ggml_prec prec) { @@ -5369,6 +5496,52 @@ void ggml_flash_attn_ext_set_prec( ggml_set_op_params_i32(a, 3, prec_i32); // scale is on first pos, max_bias on second } +void ggml_flash_attn_ext_set_ds4_sparse( + struct ggml_tensor * a, + int raw_rows, + int raw_window, + int keep_rows, + int block_size) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(raw_rows >= 0); + GGML_ASSERT(raw_window > 0 && raw_window <= 0xffff); + GGML_ASSERT(block_size > 0 && block_size <= 0xffff); + ggml_set_op_params_i32(a, 4, raw_rows); + ggml_set_op_params_i32(a, 5, keep_rows); + // Both values are DS4-local and bounded well below 16 bits. Packing them + // keeps the remaining op-parameter slots available for the exact RoPE + // metadata without growing ggml_tensor. + const uint32_t packed_layout = + ((uint32_t) raw_window << 16) | (uint32_t) block_size; + ggml_set_op_params_i32(a, 6, (int32_t) packed_layout); +} + +void ggml_flash_attn_ext_set_ds4_inverse_rope( + struct ggml_tensor * a, + int kv_start, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + int n_ctx_orig, + bool q_unrotated) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(kv_start >= 0); + // Bit 0: inverse RoPE on attention output. Bit 1: Q still needs forward + // RoPE. Both directions share the position and YaRN parameters below. + ggml_set_op_params_i32(a, 7, 1 | (q_unrotated ? 2 : 0)); + ggml_set_op_params_i32(a, 8, kv_start); + ggml_set_op_params_f32(a, 9, freq_base); + ggml_set_op_params_f32(a, 10, freq_scale); + ggml_set_op_params_f32(a, 11, ext_factor); + ggml_set_op_params_f32(a, 12, attn_factor); + ggml_set_op_params_f32(a, 13, beta_fast); + ggml_set_op_params_f32(a, 14, beta_slow); + ggml_set_op_params_i32(a, 15, n_ctx_orig); +} + enum ggml_prec ggml_flash_attn_ext_get_prec( const struct ggml_tensor * a) { GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); @@ -7991,3 +8164,169 @@ 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; + const int64_t n_tokens = mix->ne[1]; + GGML_ASSERT(mix->ne[0] == mix_dim); + GGML_ASSERT(n_tokens > 0); + GGML_ASSERT(ggml_nelements(base) >= mix_dim); + GGML_ASSERT(hc_state->ne[1] == n_tokens); + GGML_ASSERT(hc_state->ne[0] % n_hc == 0); + const int64_t n_embd = hc_state->ne[0] / n_hc; + + struct ggml_tensor * result = n_tokens == 1 + ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd + mix_dim) + : ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd + mix_dim, n_tokens); + result->op = GGML_OP_DS4_HC; + result->src[0] = mix; + result->src[1] = base; + 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; + const int64_t n_tokens = residual_hc->ne[1]; + GGML_ASSERT(n_tokens > 0); + GGML_ASSERT(split->ne[0] == mix_dim && split->ne[1] == n_tokens); + GGML_ASSERT(residual_hc->ne[0] % n_hc == 0); + const int64_t n_embd = residual_hc->ne[0] / n_hc; + GGML_ASSERT(block_out->ne[0] == n_embd && block_out->ne[1] == n_tokens); + + struct ggml_tensor * result = n_tokens == 1 + ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t) n_embd * n_hc) + : ggml_new_tensor_2d(ctx, GGML_TYPE_F32, (int64_t) n_embd * n_hc, n_tokens); + result->op = GGML_OP_DS4_HC; + result->src[0] = residual_hc; + result->src[1] = block_out; + 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); + const int64_t n_tokens = hc_state->ne[1]; + GGML_ASSERT(n_tokens > 0); + GGML_ASSERT(mix->ne[0] >= n_hc && mix->ne[1] == n_tokens); + GGML_ASSERT(ggml_nelements(base) >= n_hc); + GGML_ASSERT(hc_state->ne[0] % n_hc == 0); + const int64_t n_embd = hc_state->ne[0] / n_hc; + + struct ggml_tensor * result = n_tokens == 1 + ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd) + : ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); + result->op = GGML_OP_DS4_HC; + result->src[0] = mix; + result->src[1] = base; + 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; +} + +struct ggml_tensor * ggml_ds4_indexer_qat( + struct ggml_context * ctx, + struct ggml_tensor * input) { + GGML_ASSERT(input->type == GGML_TYPE_F32); + GGML_ASSERT(input->ne[0] == 128); + GGML_ASSERT(ggml_is_contiguous(input)); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, input); + result->op = GGML_OP_DS4_INDEXER_QAT; + result->src[0] = input; + return result; +} + +struct ggml_tensor * ggml_ds4_indexer_score( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * head_weights, + struct ggml_tensor * index_comp, + int kv_start, + int ratio) { + GGML_ASSERT(q->type == GGML_TYPE_F32 && q->ne[0] == 128); + GGML_ASSERT(head_weights->type == GGML_TYPE_F32); + GGML_ASSERT(index_comp->type == GGML_TYPE_F16 && index_comp->ne[0] == 128); + GGML_ASSERT(ggml_is_contiguous(q)); + GGML_ASSERT(ggml_is_contiguous(head_weights)); + GGML_ASSERT(ggml_is_contiguous(index_comp)); + GGML_ASSERT(head_weights->ne[0] == q->ne[1]); + GGML_ASSERT(head_weights->ne[1] == q->ne[2]); + GGML_ASSERT(kv_start >= 0 && ratio > 0); + + struct ggml_tensor * result = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, index_comp->ne[1], q->ne[2]); + result->op = GGML_OP_DS4_INDEXER_SCORE; + result->src[0] = q; + result->src[1] = head_weights; + result->src[2] = index_comp; + ggml_set_op_params_i32(result, 0, kv_start); + ggml_set_op_params_i32(result, 1, ratio); + return result; +} + +struct ggml_tensor * ggml_ds4_indexer_mask( + struct ggml_context * ctx, + struct ggml_tensor * base_mask, + struct ggml_tensor * selected, + int raw_rows) { + GGML_ASSERT(base_mask->type == GGML_TYPE_F32); + GGML_ASSERT(selected->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(base_mask)); + GGML_ASSERT(ggml_is_contiguous(selected)); + GGML_ASSERT(raw_rows >= 0 && raw_rows <= base_mask->ne[0]); + GGML_ASSERT(selected->ne[1] == base_mask->ne[1]); + GGML_ASSERT(selected->ne[2] == base_mask->ne[2]); + GGML_ASSERT(selected->ne[3] == base_mask->ne[3]); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, base_mask); + result->op = GGML_OP_DS4_INDEXER_MASK; + result->src[0] = base_mask; + result->src[1] = selected; + ggml_set_op_params_i32(result, 0, raw_rows); + return result; +} diff --git a/server/src/common/backend_factory.cpp b/server/src/common/backend_factory.cpp index 1383c2166..1b25d940b 100644 --- a/server/src/common/backend_factory.cpp +++ b/server/src/common/backend_factory.cpp @@ -232,6 +232,7 @@ std::unique_ptr create_backend(const BackendArgs & args) { cfg.stream_fd = args.stream_fd; cfg.max_ctx = args.device.max_ctx; cfg.chunk = args.chunk; + cfg.prefill_mode = args.ds4_prefill_mode; auto backend = std::make_unique(cfg); if (!backend->init()) { diff --git a/server/src/common/backend_factory.h b/server/src/common/backend_factory.h index 73ed8698c..544e9e7d2 100644 --- a/server/src/common/backend_factory.h +++ b/server/src/common/backend_factory.h @@ -15,6 +15,7 @@ #include "placement/placement_config.h" #include "placement/remote_draft_config.h" #include "placement/remote_target_shard_config.h" +#include "prefill_attention_mode.h" #include #include @@ -43,7 +44,8 @@ struct BackendArgs { int stream_fd = -1; // Chunked prefill - int chunk = 512; + int chunk = 512; + PrefillAttentionMode ds4_prefill_mode = PrefillAttentionMode::Exact; // qwen35-specific speculative decode options int fa_window = 0; // 0 = full attention. qwen3.6 full-attn layers must see the whole context; a finite window drops the system prompt/tools -> breaks tool calls. diff --git a/server/src/common/prefill_attention_mode.h b/server/src/common/prefill_attention_mode.h new file mode 100644 index 000000000..4eafdfbfe --- /dev/null +++ b/server/src/common/prefill_attention_mode.h @@ -0,0 +1,23 @@ +#pragma once + +namespace dflash::common { + +// Numerics/performance policy for model-specific prefill attention. Exact is +// the tokenwise reference, Dense selects a model's fused dense kernel, and +// Sparse may prune model-specific cache entries. +enum class PrefillAttentionMode { + Exact, + Dense, + Sparse, +}; + +inline const char * prefill_attention_mode_name(PrefillAttentionMode mode) { + switch (mode) { + case PrefillAttentionMode::Exact: return "exact"; + case PrefillAttentionMode::Dense: return "dense"; + case PrefillAttentionMode::Sparse: return "sparse"; + } + return "unknown"; +} + +} // namespace dflash::common diff --git a/server/src/deepseek4/deepseek4_backend.cpp b/server/src/deepseek4/deepseek4_backend.cpp index 9b4616da1..653f1b099 100644 --- a/server/src/deepseek4/deepseek4_backend.cpp +++ b/server/src/deepseek4/deepseek4_backend.cpp @@ -69,6 +69,11 @@ static void add_step_tel(DeepSeek4StepTelemetry & dst, const DeepSeek4StepTeleme dst.output_us += src.output_us; dst.sample_us += src.sample_us; dst.emit_us += src.emit_us; + dst.full_graph_build_us += src.full_graph_build_us; + dst.full_graph_alloc_us += src.full_graph_alloc_us; + dst.full_graph_set_us += src.full_graph_set_us; + dst.full_graph_compute_us += src.full_graph_compute_us; + dst.full_graph_read_us += src.full_graph_read_us; dst.hot_selected += src.hot_selected; dst.cold_selected += src.cold_selected; } @@ -92,6 +97,8 @@ static void log_step_tel(const char * phase, "ffn_hot_graph_build=%llu ffn_hot_graph_hit=%llu ffn_cold_graph_build=%llu ffn_cold_graph_hit=%llu " "hc_pre=%.1fms hc_pre_build=%.1fms hc_pre_input=%.1fms hc_pre_compute=%.1fms " "hc_post=%.1fms output=%.1fms sample=%.1fms emit=%.1fms " + "graph_build=%.1fms graph_alloc=%.1fms graph_set=%.1fms " + "graph_compute=%.1fms graph_read=%.1fms " "hot_sel=%d cold_sel=%d\n", phase, tokens, steps, wall_s, tok_s, ms(t.total_us), ms(t.embed_us), ms(t.attn_build_us), ms(t.attn_compute_us), ms(t.attn_read_us), @@ -107,6 +114,9 @@ static void log_step_tel(const char * phase, ms(t.hc_pre_compute_us), ms(t.hc_post_attn_us + t.hc_post_ffn_us), ms(t.output_us), ms(t.sample_us), ms(t.emit_us), + ms(t.full_graph_build_us), ms(t.full_graph_alloc_us), + ms(t.full_graph_set_us), ms(t.full_graph_compute_us), + ms(t.full_graph_read_us), t.hot_selected, t.cold_selected); } @@ -347,6 +357,7 @@ bool DeepSeek4Backend::init() { std::fprintf(stderr, "[deepseek4] failed to allocate KV cache (ctx=%d)\n", max_ctx); return false; } + cache_.prefill_mode = cfg_.prefill_mode; if (moe_hybrid_) { // Expert IPC removed — layer split replaces expert split. @@ -523,6 +534,7 @@ bool DeepSeek4Backend::unpark(const std::string & what) { moe_placement_ = {}; return false; } + cache_.prefill_mode = cfg_.prefill_mode; parked_ = false; std::printf("[deepseek4] unparked (VRAM restored)\n"); @@ -533,11 +545,38 @@ bool DeepSeek4Backend::unpark(const std::string & what) { int DeepSeek4Backend::do_prefill(const std::vector & tokens, const DaemonIO & io, 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); + // The all-hot layer-range path supports causal chunked prefill. The + // optimized graph snapshots the previous raw SWA window, attends over + // that snapshot plus the current ubatch, and commits only the final SWA + // tail. Learned compressor boundaries are emitted inside the same graph. + // + // Mixed hot/cold hybrid execution still has single-token HC semantics, so + // retain the reference path there. --chunk 1 is the explicit fallback. + const int requested_chunk = cfg_.chunk > 0 ? cfg_.chunk : w_.n_swa; const int n_total = (int)tokens.size(); + // Bound the layer-major graph to the topology validated by the prefill + // kernels. Smaller tail chunks use the same scheduler or its reference + // fallback. + const int layer_major_cap = DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS; + const int chunk = (moe_hybrid_ || + cfg_.prefill_mode == PrefillAttentionMode::Exact) + ? 1 + : std::max(1, std::min(requested_chunk, + layer_major_cap)); int pos = kv_offset; + // New sequence: clear the cache buffer so compressor state double-buffers + // and compressed-KV rows start from zeros, exactly like a fresh server. + // 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); + cache_.cur_pos = 0; + for (DeepSeek4LayerCache & layer : cache_.layers) { + layer.n_comp = 0; + layer.n_index_comp = 0; + } + } last_logits_.clear(); const bool timing = env_flag_enabled("DFLASH_DS4_TIMING"); const auto phase_t0 = Clock::now(); @@ -556,13 +595,28 @@ 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, + cfg_.prefill_mode != + PrefillAttentionMode::Sparse); + } + if (!ok) { std::fprintf(stderr, "[deepseek4] prefill step failed at pos=%d\n", pos); return -1; } @@ -619,12 +673,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..74419e1bf 100644 --- a/server/src/deepseek4/deepseek4_graph.cpp +++ b/server/src/deepseek4/deepseek4_graph.cpp @@ -22,6 +22,7 @@ #include "ggml-backend.h" #include +#include #include #include #include @@ -29,8 +30,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 +53,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) { @@ -114,6 +221,19 @@ struct DeepSeek4I64ArrayBinding { std::vector values; }; +struct DeepSeek4F32ArrayBinding { + ggml_tensor * tensor = nullptr; + std::vector values; +}; + +// Attention implementation selected by the DS4 prefill scheduler. Decode +// retains the established explicit reduction path. +enum class DeepSeek4AttentionImpl { + Explicit, + DenseFlash, + SparseFlash, +}; + static ggml_tensor * build_rms_norm(ggml_context * ctx, ggml_tensor * x, ggml_tensor * weight, float eps); static ggml_tensor * build_clamped_swiglu(ggml_context * ctx, @@ -176,6 +296,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 +323,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 +342,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 +371,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 +474,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 +498,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 +588,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) ─────────────────────────────── @@ -501,6 +663,320 @@ static ggml_tensor * build_tail_rope_2d(ggml_context * ctx, // ─── KV Compressor Step ──────────────────────────────────────────────── +// Build an exact multi-token compressor update for prefill. Complete windows +// are pooled as one batched tensor, so a 2K ubatch does not create hundreds of +// serial softmax subgraphs. The state is assembled functionally from an +// initial snapshot and written back once, avoiding persistent-buffer races. +static bool build_compressor_prefill( + ggml_context * ctx, + ggml_cgraph * gf, + ggml_tensor * cur_all, + ggml_tensor * ape, + ggml_tensor * kv_proj, + ggml_tensor * gate_proj, + ggml_tensor * norm_weight, + DeepSeek4CompressorState & state, + ggml_tensor * comp_cache, + int ratio, + int head_dim, + int kv_start, + int n_tokens, + int n_rot, + float rms_eps, + float compress_rope_freq_base, + float rope_scale_factor, + float rope_yarn_beta_fast, + float rope_yarn_beta_slow, + int rope_orig_ctx, + std::vector & i64_array_inputs, + std::vector & i32_array_inputs, + ggml_tensor ** comp_cache_source_out, + bool indexer_qat) { + if (!cur_all || n_tokens <= 1 || + n_tokens > DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS || + (ratio != 4 && ratio != 128)) { + return false; + } + + const int coff = ratio == 4 ? 2 : 1; + const int comp_width = coff * head_dim; + const int n_state_rows = ratio == 4 ? 2 * ratio : ratio; + + struct Pair { + ggml_tensor * kv = nullptr; + ggml_tensor * score = nullptr; + }; + + auto view_cols = [&](ggml_tensor * src, int width, int first, int count) { + GGML_ASSERT(src && count > 0); + return ggml_cont(ctx, ggml_view_2d(ctx, src, width, count, src->nb[1], + (size_t) first * src->nb[1])); + }; + auto view_pair_cols = [&](const Pair & src, int first, int count) { + return Pair { view_cols(src.kv, comp_width, first, count), + view_cols(src.score, comp_width, first, count) }; + }; + auto concat_tensors = [&](const std::vector & parts) { + GGML_ASSERT(!parts.empty()); + ggml_tensor * out = parts[0]; + for (size_t i = 1; i < parts.size(); ++i) { + out = ggml_concat(ctx, out, parts[i], 1); + } + return ggml_cont(ctx, out); + }; + auto concat_pairs = [&](const std::vector & parts) { + std::vector kv_parts; + std::vector score_parts; + kv_parts.reserve(parts.size()); + score_parts.reserve(parts.size()); + for (const Pair & part : parts) { + kv_parts.push_back(part.kv); + score_parts.push_back(part.score); + } + return Pair { concat_tensors(kv_parts), concat_tensors(score_parts) }; + }; + auto replace_span = [&](const Pair & base, + int first, + const Pair & replacement, + int count, + int width) { + std::vector parts; + if (first > 0) parts.push_back(view_pair_cols(base, 0, first)); + parts.push_back(replacement); + if (first + count < width) { + parts.push_back(view_pair_cols(base, first + count, width - first - count)); + } + return concat_pairs(parts); + }; + + // Project the entire chunk once and add the position-addressed APE score. + Pair projected; + projected.kv = ggml_mul_mat(ctx, kv_proj, cur_all); + projected.score = ggml_mul_mat(ctx, gate_proj, cur_all); + ggml_tensor * ape_rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); + ggml_set_input(ape_rows); + std::vector ape_values((size_t) n_tokens); + for (int i = 0; i < n_tokens; ++i) { + ape_values[(size_t) i] = (kv_start + i) % ratio; + } + i32_array_inputs.push_back({ape_rows, std::move(ape_values)}); + ggml_tensor * ape_cols = ggml_get_rows(ctx, ape, ape_rows); + projected.score = ggml_add(ctx, projected.score, + ds4_cast_if_needed(ctx, ape_cols, GGML_TYPE_F32)); + + // Snapshot before the single writeback. Both compressor output and final + // state depend on these copies, forcing reads to complete before mutation. + Pair initial { ggml_cont(ctx, state.state_kv), + ggml_cont(ctx, state.state_score) }; + ggml_build_forward_expand(gf, initial.kv); + ggml_build_forward_expand(gf, initial.score); + + ggml_tensor * pooled_batch = nullptr; + std::vector comp_rows; + std::vector comp_positions; + + auto pool_groups = [&](ggml_tensor * values_kv, + ggml_tensor * values_score, + int rows, + int groups) { + GGML_ASSERT(values_kv && values_score && rows > 0 && groups > 0); + ggml_tensor * kv3 = ggml_reshape_3d(ctx, values_kv, + head_dim, rows, groups); + ggml_tensor * score3 = ggml_reshape_3d(ctx, values_score, + head_dim, rows, groups); + ggml_tensor * score_t = ggml_cont( + ctx, ggml_permute(ctx, score3, 1, 0, 2, 3)); + ggml_tensor * kv_t = ggml_cont( + ctx, ggml_permute(ctx, kv3, 1, 0, 2, 3)); + ggml_tensor * probs_t = ggml_soft_max(ctx, score_t); + ggml_tensor * weighted_t = ggml_mul(ctx, probs_t, kv_t); + ggml_tensor * pooled_sum = ggml_sum_rows(ctx, weighted_t); + ggml_tensor * pooled = ggml_reshape_2d( + ctx, ggml_cont(ctx, pooled_sum), head_dim, groups); + pooled = ggml_cont(ctx, pooled); + pooled = build_rms_norm(ctx, pooled, norm_weight, rms_eps); + return ggml_reshape_2d(ctx, pooled, head_dim, groups); + }; + + Pair final_state; + if (ratio == 4) { + const int pos_mod = kv_start % ratio; + const int first_count = std::min(ratio - pos_mod, n_tokens); + int consumed = 0; + std::vector complete_parts; + + if (n_tokens >= ratio - pos_mod) { + Pair current = view_pair_cols(initial, ratio, ratio); + Pair first_span = view_pair_cols(projected, 0, first_count); + complete_parts.push_back(replace_span( + current, pos_mod, first_span, first_count, ratio)); + consumed = first_count; + + const int complete_tail = ((n_tokens - consumed) / ratio) * ratio; + if (complete_tail > 0) { + complete_parts.push_back(view_pair_cols( + projected, consumed, complete_tail)); + consumed += complete_tail; + } + } + + if (complete_parts.empty()) { + Pair prev = view_pair_cols(initial, 0, ratio); + Pair current = view_pair_cols(initial, ratio, ratio); + current = replace_span(current, pos_mod, projected, + n_tokens, ratio); + final_state = concat_pairs({prev, current}); + } else { + Pair complete = concat_pairs(complete_parts); + const int groups = (int) complete.kv->ne[1] / ratio; + GGML_ASSERT(groups > 0); + + Pair previous = view_pair_cols(initial, 0, ratio); + if (groups > 1) { + previous = concat_pairs({ + previous, + view_pair_cols(complete, 0, (groups - 1) * ratio), + }); + } + + auto select_half = [&](ggml_tensor * src, int half) { + ggml_tensor * src3 = ggml_reshape_3d( + ctx, src, comp_width, ratio, groups); + return ggml_cont(ctx, ggml_view_3d( + ctx, src3, head_dim, ratio, groups, + src3->nb[1], src3->nb[2], + (size_t) half * head_dim * src3->nb[0])); + }; + ggml_tensor * selected_kv = ggml_concat( + ctx, select_half(previous.kv, 0), + select_half(complete.kv, 1), 1); + ggml_tensor * selected_score = ggml_concat( + ctx, select_half(previous.score, 0), + select_half(complete.score, 1), 1); + pooled_batch = pool_groups( + ggml_cont(ctx, selected_kv), + ggml_cont(ctx, selected_score), 2 * ratio, groups); + + const int first_boundary = kv_start + first_count - 1; + for (int g = 0; g < groups; ++g) { + const int boundary = first_boundary + g * ratio; + const int64_t comp_row = boundary / ratio; + GGML_ASSERT(comp_row >= 0 && comp_row < comp_cache->ne[1]); + comp_rows.push_back(comp_row); + comp_positions.push_back(boundary + 1 - ratio); + } + + Pair last_complete = view_pair_cols( + complete, (groups - 1) * ratio, ratio); + Pair current = last_complete; + const int tail = n_tokens - consumed; + if (tail > 0) { + current = replace_span( + current, 0, view_pair_cols(projected, consumed, tail), + tail, ratio); + } + final_state = concat_pairs({last_complete, current}); + } + } else { + const int pos_mod = kv_start % ratio; + const int to_boundary = ratio - pos_mod; + if (n_tokens < to_boundary) { + final_state = replace_span(initial, pos_mod, projected, n_tokens, ratio); + } else { + std::vector first_parts; + if (pos_mod > 0) { + first_parts.push_back(view_pair_cols(initial, 0, pos_mod)); + } + first_parts.push_back(view_pair_cols(projected, 0, to_boundary)); + Pair first_complete = concat_pairs(first_parts); + + int consumed = to_boundary; + const int complete_tail = ((n_tokens - consumed) / ratio) * ratio; + Pair complete = first_complete; + if (complete_tail > 0) { + complete = concat_pairs({ + first_complete, + view_pair_cols(projected, consumed, complete_tail), + }); + consumed += complete_tail; + } + const int groups = (int) complete.kv->ne[1] / ratio; + pooled_batch = pool_groups(complete.kv, complete.score, + ratio, groups); + const int first_boundary = kv_start + to_boundary - 1; + for (int g = 0; g < groups; ++g) { + const int boundary = first_boundary + g * ratio; + const int64_t comp_row = boundary / ratio; + GGML_ASSERT(comp_row >= 0 && comp_row < comp_cache->ne[1]); + comp_rows.push_back(comp_row); + comp_positions.push_back(boundary + 1 - ratio); + } + + Pair last_complete = view_pair_cols( + complete, (groups - 1) * ratio, ratio); + const int tail = n_tokens - consumed; + if (tail > 0) { + final_state = replace_span( + last_complete, 0, + view_pair_cols(projected, consumed, tail), tail, ratio); + } else { + final_state = last_complete; + } + } + } + + // Persist the exact sequential state using unique row indices. + ggml_tensor * state_rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_state_rows); + ggml_set_input(state_rows); + std::vector state_row_values((size_t) n_state_rows); + for (int i = 0; i < n_state_rows; ++i) state_row_values[(size_t) i] = i; + i64_array_inputs.push_back({state_rows, std::move(state_row_values)}); + final_state.kv = ggml_cont(ctx, final_state.kv); + final_state.score = ggml_cont(ctx, final_state.score); + ggml_tensor * state_kv_source = ggml_set_rows(ctx, state.state_kv, + final_state.kv, state_rows); + ggml_tensor * state_score_source = ggml_set_rows(ctx, state.state_score, + final_state.score, state_rows); + ggml_build_forward_expand(gf, state_kv_source); + ggml_build_forward_expand(gf, state_score_source); + + ggml_tensor * comp_cache_source = comp_cache; + if (pooled_batch) { + ggml_tensor * pooled = pooled_batch; + const int n_pooled = (int) comp_positions.size(); + ggml_tensor * comp_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, + n_pooled); + ggml_set_input(comp_pos); + i32_array_inputs.push_back({comp_pos, std::move(comp_positions)}); + + const float rope_scale = rope_scale_factor > 0.0f + ? (1.0f / rope_scale_factor) : 1.0f; + float rope_attn = 1.0f; + if (rope_scale > 0.0f) { + rope_attn /= (1.0f + 0.1f * logf(1.0f / rope_scale)); + } + pooled = build_tail_rope_2d(ctx, pooled, comp_pos, n_rot, head_dim, + n_pooled, + compress_rope_freq_base, rope_scale, + 1.0f, rope_attn, + rope_yarn_beta_fast, rope_yarn_beta_slow, + rope_orig_ctx); + pooled = ggml_cont(ctx, pooled); + if (indexer_qat) { + pooled = ggml_ds4_indexer_qat(ctx, pooled); + } + + ggml_tensor * comp_row_tensor = ggml_new_tensor_1d( + ctx, GGML_TYPE_I64, (int64_t) comp_rows.size()); + ggml_set_input(comp_row_tensor); + i64_array_inputs.push_back({comp_row_tensor, std::move(comp_rows)}); + comp_cache_source = ggml_set_rows(ctx, comp_cache, pooled, comp_row_tensor); + ggml_build_forward_expand(gf, comp_cache_source); + } + if (comp_cache_source_out) *comp_cache_source_out = comp_cache_source; + return true; +} + static void build_compressor_step( ggml_context * ctx, ggml_cgraph * gf, @@ -526,12 +1002,30 @@ 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, + ggml_tensor * cur_all = nullptr, + int n_tokens_all = 1, + int kv_start_all = -1, + bool indexer_qat = false) { if (!gf || !cur_last || !ape || !kv_proj || !gate_proj || !norm_weight || !state.state_kv || !state.state_score || !comp_cache || ratio <= 0) { return; } + if (cur_all && n_tokens_all > 1 && !state_rows_inp && kv_start_all >= 0 && + build_compressor_prefill(ctx, gf, cur_all, ape, kv_proj, gate_proj, + norm_weight, state, comp_cache, ratio, head_dim, + kv_start_all, n_tokens_all, n_rot, rms_eps, + compress_rope_freq_base, rope_scale_factor, + rope_yarn_beta_fast, rope_yarn_beta_slow, + rope_orig_ctx, i64_array_inputs, + i32_array_inputs, comp_cache_source_out, + indexer_qat)) { + return; + } + // DS4 compression: internal width = coff * head_dim (2x for ratio-4, 1x for ratio-128) const int coff = (ratio == 4) ? 2 : 1; const int comp_width = coff * head_dim; @@ -541,21 +1035,66 @@ 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; + + 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 (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); + 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); } - 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)); + if (batched_rows) { + // Batched state writes allow one compressor boundary 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); + 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 +1106,15 @@ static void build_compressor_step( ggml_build_forward_expand(gf, ggml_cpy(ctx, sc_cur, sc_slot)); } - if (((token_pos + 1) % ratio) != 0) { + if (batched_rows && batched_b < 0) { + // State rows were written, but this batch did not complete a window. + return; + } + if (!batched_rows && !flush_rows_inp && + ((token_pos + 1) % ratio) != 0) { + // Legacy per-layer graphs only pool at flush boundaries. The fused + // stable-topology graph (flush_rows_inp set) pools every step; the + // partial result lands on the masked running comp row. return; } @@ -586,26 +1133,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)); @@ -635,6 +1182,9 @@ static void build_compressor_step( pooled = build_tail_rope_2d(ctx, pooled, comp_pos, n_rot, head_dim, 1, compress_rope_freq_base, rope_scale, 1.0f, rope_attn, rope_yarn_beta_fast, rope_yarn_beta_slow, rope_orig_ctx); + if (indexer_qat) { + pooled = ggml_ds4_indexer_qat(ctx, ggml_cont(ctx, pooled)); + } ggml_tensor * pooled_f16 = ggml_cast(ctx, pooled, GGML_TYPE_F16); const int comp_row = token_pos / ratio; @@ -643,7 +1193,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 +1202,70 @@ 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 (batched_rows) { + if (ratio == 4) { + // Rotate the completed current window into the previous half. + // Reading through the first span makes the dependency explicit. + for (int r = 0; r < ratio; ++r) { + ggml_tensor * src_kv = ggml_view_2d( + ctx, state_kv_source, comp_width, 1, + state_kv_source->nb[1], + (size_t) (ratio + r) * state_kv_source->nb[1]); + ggml_tensor * dst_kv = ggml_view_2d( + ctx, state.state_kv, comp_width, 1, + state.state_kv->nb[1], + (size_t) r * state.state_kv->nb[1]); + ggml_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 + // 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 +1306,13 @@ 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 ** index_comp_cache_source_out = nullptr, + ggml_tensor * flush_rows_inp = nullptr, + ggml_tensor * cur_all = nullptr, + int n_tokens_all = 1, + int kv_start_all = -1, + bool indexer_qat = false) { build_compressor_step(ctx, gf, cur_last, L.indexer_compressor_ape, L.indexer_compressor_kv, @@ -715,81 +1335,134 @@ static void build_indexer_compressor_step( comp_rows_inp, comp_pos_inp, i64_array_inputs, - i32_array_inputs); + i32_array_inputs, + index_comp_cache_source_out, + flush_rows_inp, + cur_all, + n_tokens_all, + kv_start_all, + indexer_qat); } static int ds4_comp_rows_used(const ggml_tensor * comp_cache, int n_cached, int ratio, int token_pos) { if (!comp_cache || ratio <= 0) { return 0; } - const int grew_this_step = ((token_pos + 1) % ratio) == 0 ? 1 : 0; - return std::min(n_cached + grew_this_step, (int) comp_cache->ne[1]); + // n_cached is the committed count before this graph. A multi-token + // prefill graph may cross several compressor boundaries, so derive the + // live count from the query position rather than adding at most one row. + const int through_position = (token_pos + 1) / ratio; + return std::min(std::max(n_cached, through_position), + (int) comp_cache->ne[1]); +} + +// Round the live compressed-row count up to a fixed stride so the fused decode +// 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( +static ggml_tensor * build_indexer_topk( ggml_context * ctx, - ggml_tensor * qr_norm_last, // [n_lora_q, 1] - ggml_tensor * cur_last, // [n_embd, 1] + ggml_tensor * qr_norm, // [n_lora_q, n_tokens] + ggml_tensor * cur, // [n_embd, n_tokens] const DeepSeek4Weights & w, const DeepSeek4Layer & L, - const DeepSeek4LayerCache & lc, - int token_pos, - std::vector & i32_inputs) { - const int n_comp = ds4_comp_rows_used(lc.index_comp_kv, lc.n_index_comp, 4, token_pos); - if (!qr_norm_last || !cur_last || !L.indexer_attn_q_b || !L.indexer_proj || - !lc.index_comp_kv || n_comp <= 0) { + ggml_tensor * index_comp_source, + int n_comp, + int kv_start, + int n_tokens, + ggml_tensor * rope_pos, + std::vector & i32_array_inputs) { + if (!qr_norm || !cur || !L.indexer_attn_q_b || !L.indexer_proj || + !index_comp_source || !rope_pos || n_tokens <= 0 || + n_comp <= w.n_indexer_top_k) { return nullptr; } const int n_indexer_head = w.n_indexer_head; const int head_dim = w.n_indexer_head_dim; + const int top_k = std::min(n_comp, w.n_indexer_top_k); + // A token with <= top_k visible compressed rows needs no ranking: selecting + // [0,top_k) and retaining the ordinary causal mask is exactly equivalent. + // Score only the suffix beginning with the first token that can see row + // top_k. For a zero-prefix ratio-4 2K request this shrinks 2164 score rows + // to just 113. + const int first_scored = std::max( + 0, std::min(n_tokens, 4 * (top_k + 1) - 1 - kv_start)); + const int n_scored = n_tokens - first_scored; + if (n_scored <= 0) return nullptr; + + auto token_slice = [&](ggml_tensor * input, int width) { + if (first_scored == 0) return input; + return ggml_view_2d( + ctx, input, width, n_scored, input->nb[1], + (size_t) first_scored * input->nb[1]); + }; + qr_norm = token_slice(qr_norm, (int) qr_norm->ne[0]); + cur = token_slice(cur, (int) cur->ne[0]); + if (first_scored > 0) { + rope_pos = ggml_view_1d( + ctx, rope_pos, n_scored, + (size_t) first_scored * rope_pos->nb[0]); + } - // DS4 indexer decode scoring mirrors ds4.c::indexer_allowed_decode_one(): - // 1. Build an indexer query from qr_norm (after q_a + RMSNorm, before q_b). - // 2. Apply full-dim RoPE in indexer head space. - // 3. Project per-head scalar weights from the current hidden state. - // 4. Score every compressed row with ReLU(dot(key_h, query_h)) * weight_h. - // 5. Return the top-k compressed-row indices. - ggml_tensor * index_q = ggml_mul_mat(ctx, L.indexer_attn_q_b, qr_norm_last); - index_q = ggml_reshape_3d(ctx, index_q, head_dim, n_indexer_head, 1); - - // TODO: RoPE on indexer query (same gallocr issue as compressor RoPE) - // Skipping for now — correctness deferred. - index_q = ggml_reshape_2d(ctx, index_q, head_dim, n_indexer_head); + // Official ratio-4 indexer graph: q_a-normalized query projection, tail + // RoPE, Hadamard+FP4 QAT, per-head scalar projection, ReLU dot products, + // weighted head reduction and top-512 selection for every query token. + ggml_tensor * index_q = ggml_mul_mat(ctx, L.indexer_attn_q_b, qr_norm); + index_q = ggml_reshape_3d( + ctx, index_q, head_dim, n_indexer_head, n_scored); - ggml_tensor * head_weights = ggml_mul_mat(ctx, L.indexer_proj, cur_last); + const float rope_scale = w.rope_scale_factor > 0.0f + ? (1.0f / w.rope_scale_factor) : 1.0f; + float rope_attn = 1.0f; + if (rope_scale > 0.0f) { + rope_attn /= 1.0f + 0.1f * logf(1.0f / rope_scale); + } + index_q = build_tail_rope_3d( + ctx, index_q, rope_pos, w.n_rot, head_dim, n_indexer_head, + n_scored, w.compress_rope_freq_base, rope_scale, 1.0f, + rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, + (int) w.rope_orig_ctx); + index_q = ggml_ds4_indexer_qat(ctx, ggml_cont(ctx, index_q)); + + ggml_tensor * head_weights = ggml_mul_mat(ctx, L.indexer_proj, cur); head_weights = ggml_scale(ctx, head_weights, 1.0f / std::sqrt((float) head_dim * (float) n_indexer_head)); - // index_comp_kv: [n_indexer_head_dim, comp_cap] — each row is 128-dim - // Score each compressed row against all query heads via broadcast - ggml_tensor * comp_view = ggml_view_2d(ctx, lc.index_comp_kv, - head_dim, n_comp, - lc.index_comp_kv->nb[1], 0); - comp_view = ggml_cast(ctx, comp_view, GGML_TYPE_F32); - // comp_view: [head_dim, n_comp] → [head_dim, 1, n_comp] for broadcast - comp_view = ggml_reshape_3d(ctx, comp_view, head_dim, 1, n_comp); - - // index_q: [head_dim, n_indexer_head, 1] → repeat to [head_dim, n_indexer_head, n_comp] - // But ggml_mul needs same shapes, so use matmul approach: - // Reshape q: [head_dim, n_indexer_head] → used directly as A in matmul - // comp: [head_dim, n_comp] - // matmul: A^T @ B = [n_indexer_head, n_comp] dot scores - ggml_tensor * comp_2d = ggml_reshape_2d(ctx, comp_view, head_dim, n_comp); - // mul_mat(index_q, comp_2d): A=[head_dim, n_indexer_head], B=[head_dim, n_comp] - // → result=[n_indexer_head, n_comp] - ggml_tensor * dots = ggml_mul_mat(ctx, index_q, comp_2d); - dots = ggml_relu(ctx, dots); - - // Weight each head's contribution: dots[n_indexer_head, n_comp] * weights[n_indexer_head, 1] - ggml_tensor * weight_rep = ggml_repeat(ctx, head_weights, dots); - ggml_tensor * weighted = ggml_mul(ctx, dots, weight_rep); - // Sum across heads (ne[0]) → [1, n_comp] - ggml_tensor * scores = ggml_sum_rows(ctx, weighted); - scores = ggml_cont(ctx, scores); - scores = ggml_reshape_2d(ctx, scores, n_comp, 1); - - return ggml_top_k(ctx, scores, std::min(n_comp, w.n_indexer_top_k)); + ggml_tensor * comp = ggml_view_2d( + ctx, index_comp_source, head_dim, n_comp, + index_comp_source->nb[1], 0); + GGML_ASSERT(comp->type == GGML_TYPE_F16); + GGML_ASSERT(ggml_is_contiguous(comp)); + + // Fuse comp^T@q, ReLU, per-head weighting and head reduction. The generic + // graph would retain [n_comp,64,n_tokens] dots (about 2 GiB at 8K) before + // reducing them; this operation stores only the final score matrix. + ggml_tensor * scores = ggml_ds4_indexer_score( + ctx, index_q, head_weights, comp, kv_start + first_scored, 4); + ggml_tensor * selected = ggml_top_k( + ctx, ggml_cont(ctx, scores), top_k); + if (first_scored == 0) return selected; + + ggml_tensor * identity = ggml_new_tensor_2d( + ctx, GGML_TYPE_I32, top_k, first_scored); + ggml_set_input(identity); + std::vector identity_values((size_t) top_k * first_scored); + for (int t = 0; t < first_scored; ++t) { + for (int k = 0; k < top_k; ++k) { + identity_values[(size_t) t * top_k + k] = k; + } + } + i32_array_inputs.push_back({identity, std::move(identity_values)}); + return ggml_concat(ctx, identity, selected, 1); } static ggml_tensor * build_selected_comp_context( @@ -825,12 +1498,13 @@ 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, + DeepSeek4AttentionImpl attention_impl = DeepSeek4AttentionImpl::Explicit) { const int n_embd = w.n_embd; const int head_dim = w.head_dim; const int n_head = w.n_head; - const int n_lora_q = w.n_lora_q; const int n_rot = w.n_rot; const int n_out_group = w.n_out_group; const int n_lora_o = w.n_lora_o; @@ -878,20 +1552,89 @@ static ggml_tensor * build_mla_attention( // n_ctx_orig is critical for YaRN correction on compressed layers const int rope_n_ctx_orig = (int)w.rope_orig_ctx; // 65536 - q = build_tail_rope_3d(ctx, q, rope_pos, n_rot, head_dim, n_head, n_tokens, - rope_freq, rope_scale, rope_ext, rope_attn, - w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + // D=512 flash prefill can rotate Q's 64-d tail inside the exact attention + // kernel. This avoids materializing cont(nope), cont(tail), rope(tail), + // and concat(nope, tail) while retaining the same F32 rounding boundary. + const bool fuse_q_rope = attention_impl != DeepSeek4AttentionImpl::Explicit && + n_tokens > 1 && head_dim == 512 && n_rot == 64; + if (!fuse_q_rope) { + q = build_tail_rope_3d(ctx, q, rope_pos, n_rot, head_dim, n_head, n_tokens, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + } kv = build_tail_rope_2d(ctx, kv, rope_pos, n_rot, head_dim, n_tokens, rope_freq, rope_scale, rope_ext, rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + // ── 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. + const bool causal_batch = + n_tokens > 1 && !cached_inputs && f32_array_inputs; + ggml_tensor * old_rows_scratch = nullptr; + int n_old_rows = 0; + ggml_tensor * prior_rows_scratch = nullptr; + int n_prior_rows = 0; + const bool fused_causal = cached_inputs && cached_inputs->attn_row_mask && n_tokens > 1; + if (fused_causal) { + // Preserve q rows so the cached graph topology remains stable; + // unwrapped rows are masked by host-provided 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) { + // Snapshot the chronological pre-chunk window before any ring writes. + // Attention then consumes [prior F16 rows | current F32 rows], matching + // the single-token path and avoiding an F16 round-trip for this chunk. + n_prior_rows = std::min(kv_start, w.n_swa); + if (n_prior_rows > 0) { + const int first = kv_start < w.n_swa ? 0 : (kv_start % w.n_swa); + const int tail = std::min(n_prior_rows, w.n_swa - first); + auto snapshot_span = [&](int row, int count) { + ggml_tensor * span = ggml_view_2d( + ctx, lc.raw_kv, head_dim, count, lc.raw_kv->nb[1], + (size_t) row * lc.raw_kv->nb[1]); + return ggml_cont(ctx, span); + }; + prior_rows_scratch = snapshot_span(first, tail); + if (tail < n_prior_rows) { + prior_rows_scratch = ggml_concat( + ctx, prior_rows_scratch, + snapshot_span(0, n_prior_rows - tail), 1); + prior_rows_scratch = ggml_cont(ctx, prior_rows_scratch); + } + ggml_build_forward_expand(gf, prior_rows_scratch); + prior_rows_scratch = ds4_cast_if_needed( + ctx, prior_rows_scratch, GGML_TYPE_F32); + } + } + // ── Store ALL KV rows in the raw SWA ring ───────────────────── // For decode (n_tokens=1): write single row. For prefill: write all rows. - if (cached_inputs && cached_inputs->raw_kv_rows) { + ggml_tensor * raw_kv_source = lc.raw_kv; + ggml_tensor * raw_kv_rows = cached_inputs + ? cached_inputs->raw_kv_rows + : nullptr; + if (raw_kv_rows) { ggml_tensor * 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, raw_kv_rows); + ggml_build_forward_expand(gf, raw_kv_source); } else { - for (int ti = 0; ti < n_tokens; ti++) { + // The attention graph consumes the whole current ubatch directly. + // Persist only its final SWA tail so every physical ring row is written + // once, even when the ubatch is much larger than the 128-row ring. + const int first_write = std::max(0, n_tokens - w.n_swa); + for (int ti = first_write; ti < n_tokens; ti++) { const int pos_ti = kv_start + ti; ggml_tensor * kv_row = ggml_view_2d( ctx, kv, head_dim, 1, kv->nb[1], (size_t)ti * kv->nb[1]); @@ -906,8 +1649,7 @@ static ggml_tensor * build_mla_attention( // ── Learned compression update ────────────────────────────────── ggml_tensor * cur_last = ggml_view_2d( ctx, cur, n_embd, 1, cur->nb[1], (size_t)(n_tokens - 1) * cur->nb[1]); - ggml_tensor * qr_last = ggml_view_2d( - ctx, qr, n_lora_q, 1, qr->nb[1], (size_t)(n_tokens - 1) * qr->nb[1]); + ggml_tensor * comp_kv_source = lc.comp_kv; if (ratio > 0 && L.attn_compressor_kv) { build_compressor_step(ctx, gf, cur_last, L.attn_compressor_ape, @@ -931,9 +1673,15 @@ 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, + (causal_batch || fused_causal) ? cur : nullptr, + n_tokens, + kv_start); } + ggml_tensor * index_comp_kv_source = lc.index_comp_kv; if (ratio == 4 && L.indexer_compressor_kv) { build_indexer_compressor_step(ctx, gf, cur_last, w, L, lc, token_pos, cached_inputs ? cached_inputs->index_ape_row : nullptr, @@ -941,8 +1689,14 @@ 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); - (void)build_indexer_score(ctx, qr_last, cur_last, w, L, lc, token_pos, i32_inputs); + i32_array_inputs, + &index_comp_kv_source, + cached_inputs ? cached_inputs->flush_rows : nullptr, + (causal_batch || fused_causal) ? cur : nullptr, + n_tokens, + kv_start, + attention_impl == + DeepSeek4AttentionImpl::SparseFlash); } // ── MLA Dot-Product Attention (SWA + compressed KV) ──────────── @@ -950,85 +1704,389 @@ 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 int n_attn = n_raw + n_comp_attn; + const bool masked_kv = cached_inputs && cached_inputs->attn_row_mask; + const int n_comp_live = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; + ggml_tensor * indexer_topk = nullptr; + if (attention_impl == DeepSeek4AttentionImpl::SparseFlash && + ratio == 4 && f32_array_inputs) { + const int n_index_comp = ds4_comp_rows_used( + lc.index_comp_kv, lc.n_index_comp, 4, token_pos); + indexer_topk = build_indexer_topk( + ctx, qr, cur, w, L, index_comp_kv_source, + n_index_comp, kv_start, n_tokens, rope_pos, + i32_array_inputs); + } + // Stable path reads the full physical ring (masking not-yet-written slots) + // and a padded compressed-row span; the plain path reads only valid rows. + const int n_raw = masked_kv ? w.n_swa + : causal_batch ? n_prior_rows + n_tokens + : std::min(kv_start + n_tokens, w.n_swa); + const int n_comp_attn = masked_kv ? cached_inputs->padded_comp : n_comp_live; + const int n_valid_raw = std::min(kv_start + n_tokens, w.n_swa); + const int n_attn = n_raw + n_comp_attn + n_old_rows; const float kq_scale = 1.0f / sqrtf((float)head_dim); - // 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); - 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); + // 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 (causal_batch) { + ggml_tensor * current = ds4_cast_if_needed(ctx, kv, GGML_TYPE_F32); + kv_attn = prior_rows_scratch + ? ggml_concat(ctx, prior_rows_scratch, current, 1) + : current; + } else if (n_tokens == 1) { + ggml_tensor * cur_kv = ds4_cast_if_needed(ctx, kv, GGML_TYPE_F32); + if (n_raw == w.n_swa && raw_kv_rows) { + // Once the ring is full, use a stable physical row order. The + // cached q=1 graph is first built at position n_swa-1 and then + // reused across every wrap position, so chronological views baked + // into that first topology become stale. Insert the current F32 + // KV at its runtime row in an F32 snapshot instead. The tokenwise + // prefill helper takes the same branch and row ordering. + ggml_tensor * ring = ggml_view_2d( + ctx, lc.raw_kv, head_dim, w.n_swa, lc.raw_kv->nb[1], 0); + ring = ds4_cast_if_needed(ctx, ring, GGML_TYPE_F32); + kv_attn = ggml_set_rows(ctx, ring, cur_kv, raw_kv_rows); + ggml_build_forward_expand(gf, kv_attn); + } else if (n_raw > 1) { + ggml_tensor * prev = nullptr; + 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); - } - 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); - } - // kv_f32: [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] - // → [n_attn, n_head*n_tokens] - ggml_tensor * scores = ggml_mul_mat(ctx, kv_f32, q_flat); - scores = ggml_scale(ctx, scores, kq_scale); - - // 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); + kv_attn = raw_kv_view(0, n_raw); } - // 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)); - ggml_tensor * context = ggml_mul_mat(ctx, kv_T, probs); - // context: [head_dim, n_head*n_tokens] + 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); + } + if (old_rows_scratch) { + kv_attn = ggml_concat(ctx, kv_attn, old_rows_scratch, 1); + } + // kv_attn: [head_dim, n_attn] + + // Build one additive mask tensor and share it between the explicit and + // flash-attention implementations. ggml flash attention expects + // [n_kv,n_query] F16; the explicit path broadcasts the same values over + // heads in F32. + ggml_tensor * score_mask = nullptr; + const bool exact_two_band = + attention_impl == DeepSeek4AttentionImpl::DenseFlash && + causal_batch && + n_tokens > DS4_NUMERICAL_PREFILL_BAND && + n_tokens <= 2 * DS4_NUMERICAL_PREFILL_BAND; + if (!exact_two_band) { + if (masked_kv && n_tokens > 1) { + score_mask = ggml_reshape_2d(ctx, cached_inputs->attn_row_mask, + n_attn, n_tokens); + } else if (masked_kv) { + score_mask = ggml_reshape_2d(ctx, cached_inputs->attn_row_mask, + n_attn, 1); + } else if (causal_batch) { + // Per-token causal mask over [prior rows | current rows | comp rows]. + ggml_tensor * cmask = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_attn, 1, n_tokens); + ggml_set_input(cmask); + std::vector mvals((size_t) n_attn * n_tokens, 0.0f); + for (int i = 0; i < n_tokens; i++) { + const int pos_i = kv_start + i; + float * col = mvals.data() + (size_t) i * n_attn; + const int min_pos = pos_i - w.n_swa + 1; + for (int r = 0; r < n_prior_rows; ++r) { + const int prior_pos = kv_start - n_prior_rows + r; + if (prior_pos < min_pos) col[r] = -1e30f; + } + for (int t = 0; t < n_tokens; ++t) { + const int current_pos = kv_start + t; + if (t > i || current_pos < min_pos) { + col[n_prior_rows + t] = -1e30f; + } + } + if (n_comp_attn > 0) { + const int vis = ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, pos_i); + for (int c = vis; c < n_comp_attn; c++) col[n_raw + c] = -1e30f; + } + } + f32_array_inputs->push_back({cmask, std::move(mvals)}); + score_mask = ggml_reshape_2d(ctx, cmask, n_attn, n_tokens); + } + } + if (indexer_topk) { + if (!score_mask) { + score_mask = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_attn, n_tokens); + ggml_set_input(score_mask); + f32_array_inputs->push_back({ + score_mask, + std::vector((size_t) n_attn * n_tokens, 0.0f), + }); + } + score_mask = ggml_ds4_indexer_mask( + ctx, ggml_cont(ctx, score_mask), indexer_topk, n_raw); + } + (void) n_valid_raw; + + ggml_tensor * context = nullptr; + bool inverse_rope_fused = false; + const bool use_flash = attention_impl != DeepSeek4AttentionImpl::Explicit && + n_tokens > 1; + if (use_flash) { + if (exact_two_band) { + // A larger scheduling band must retain the numerical topology of + // two 2K requests. Prefix queries use the first band's F32 raw KV; + // suffix queries see its final SWA tail after the same F16 cache + // round-trip. HC, projections and MoE still run once over the full + // token batch, avoiding a second expert-weight sweep. + const int first_count = DS4_NUMERICAL_PREFILL_BAND; + const int second_count = n_tokens - first_count; + const int first_comp = ratio > 0 + ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, + kv_start + first_count - 1) + : 0; + const int second_comp = n_comp_live; + const int second_prior_count = std::min(first_count, w.n_swa); + + auto view_kv = [&](int first, int count) { + return ggml_view_2d( + ctx, kv, head_dim, count, kv->nb[1], + (size_t) first * kv->nb[1]); + }; + auto append_comp = [&](ggml_tensor * raw, int count) { + if (count <= 0 || !comp_kv_source) return raw; + ggml_tensor * comp = ggml_view_2d( + ctx, comp_kv_source, head_dim, count, + comp_kv_source->nb[1], 0); + comp = ds4_cast_if_needed(ctx, comp, GGML_TYPE_F32); + return ggml_concat(ctx, raw, comp, 1); + }; + auto make_band_mask = [&](int start, int count, int prior, + int comp_count) { + const int raw_count = prior + count; + const int attn_count = raw_count + comp_count; + ggml_tensor * mask3 = ggml_new_tensor_3d( + ctx, GGML_TYPE_F32, attn_count, 1, count); + ggml_set_input(mask3); + std::vector values( + (size_t) attn_count * count, 0.0f); + for (int i = 0; i < count; ++i) { + const int pos_i = start + i; + const int min_pos = pos_i - w.n_swa + 1; + float * col = values.data() + (size_t) i * attn_count; + for (int r = 0; r < prior; ++r) { + const int prior_pos = start - prior + r; + if (prior_pos < min_pos) col[r] = -1e30f; + } + for (int t = 0; t < count; ++t) { + const int current_pos = start + t; + if (t > i || current_pos < min_pos) { + col[prior + t] = -1e30f; + } + } + if (comp_count > 0) { + const int visible = ds4_comp_rows_used( + lc.comp_kv, lc.n_comp, ratio, pos_i); + for (int c = visible; c < comp_count; ++c) { + col[raw_count + c] = -1e30f; + } + } + } + f32_array_inputs->push_back({mask3, std::move(values)}); + return ggml_reshape_2d(ctx, mask3, attn_count, count); + }; + auto make_flash = [&](ggml_tensor * q_band, + ggml_tensor * kv_band, + ggml_tensor * mask_band, + int raw_count, + int start) { + const int attn_count = (int) kv_band->ne[1]; + ggml_tensor * k_band = ggml_reshape_3d( + ctx, kv_band, head_dim, attn_count, 1); + ggml_tensor * mask_fa = ds4_cast_if_needed( + ctx, mask_band, GGML_TYPE_F16); + ggml_tensor * result = ggml_flash_attn_ext( + ctx, q_band, k_band, k_band, mask_fa, + kq_scale, 0.0f, 0.0f); + if (L.attn_sinks) { + ggml_flash_attn_ext_add_sinks(result, L.attn_sinks); + } + ggml_flash_attn_ext_set_prec(result, GGML_PREC_F32); + ggml_flash_attn_ext_set_ds4_sparse( + result, raw_count, w.n_swa, 0, 32); + ggml_flash_attn_ext_set_ds4_inverse_rope( + result, start, rope_freq, rope_scale, rope_ext, + rope_attn, w.rope_yarn_beta_fast, + w.rope_yarn_beta_slow, rope_n_ctx_orig, fuse_q_rope); + return result; + }; + + ggml_tensor * q_fa = ggml_permute(ctx, q, 0, 2, 1, 3); + auto view_q = [&](int first, int count) { + return ggml_view_3d( + ctx, q_fa, head_dim, count, n_head, + q_fa->nb[1], q_fa->nb[2], + (size_t) first * q_fa->nb[1]); + }; + + ggml_tensor * first_raw = ds4_cast_if_needed( + ctx, view_kv(0, first_count), GGML_TYPE_F32); + if (prior_rows_scratch) { + first_raw = ggml_concat( + ctx, prior_rows_scratch, first_raw, 1); + } + ggml_tensor * first_kv = append_comp(first_raw, first_comp); + ggml_tensor * first_mask = make_band_mask( + kv_start, first_count, n_prior_rows, first_comp); + ggml_tensor * first_context = make_flash( + view_q(0, first_count), first_kv, first_mask, + n_prior_rows + first_count, kv_start); + + ggml_tensor * rounded_prior = ggml_cast( + ctx, view_kv(first_count - second_prior_count, + second_prior_count), + GGML_TYPE_F16); + rounded_prior = ggml_cast(ctx, rounded_prior, GGML_TYPE_F32); + ggml_tensor * second_raw = ggml_concat( + ctx, rounded_prior, view_kv(first_count, second_count), 1); + ggml_tensor * second_kv = append_comp(second_raw, second_comp); + ggml_tensor * second_mask = make_band_mask( + kv_start + first_count, second_count, + second_prior_count, second_comp); + ggml_tensor * second_context = make_flash( + view_q(first_count, second_count), second_kv, second_mask, + second_prior_count + second_count, + kv_start + first_count); + + context = ggml_concat(ctx, first_context, second_context, 2); + inverse_rope_fused = true; + } else { + // ggml FA convention: Q[D,T,H], K/V[D,K,Hkv]. DS4 MLA has one shared + // latent KV head and uses the same latent vector as both key and value. + // The DS4 D=512 kernel consumes Q strides directly, avoiding a full + // [D,H,T] -> [D,T,H] materialization for every layer. + ggml_tensor * q_fa = ggml_permute(ctx, q, 0, 2, 1, 3); + ggml_tensor * kv_fa = ds4_cast_if_needed(ctx, kv_attn, GGML_TYPE_F32); + ggml_tensor * k_fa = ggml_reshape_3d(ctx, kv_fa, head_dim, n_attn, 1); + ggml_tensor * v_fa = k_fa; + ggml_tensor * mask_fa = score_mask + ? ds4_cast_if_needed(ctx, score_mask, GGML_TYPE_F16) + : nullptr; + context = ggml_flash_attn_ext(ctx, q_fa, k_fa, v_fa, mask_fa, + kq_scale, 0.0f, 0.0f); + if (L.attn_sinks) { + ggml_flash_attn_ext_add_sinks(context, L.attn_sinks); + } + ggml_flash_attn_ext_set_prec(context, GGML_PREC_F32); + // Always publish the raw/compressed boundary. A zero keep count leaves + // dense attention unchanged while allowing the D=512 value pass to + // skip the two masked envelopes without guessing DS4 cache layout. + ggml_flash_attn_ext_set_ds4_sparse( + context, n_raw, w.n_swa, + indexer_topk + ? -w.n_indexer_top_k + : attention_impl == DeepSeek4AttentionImpl::SparseFlash + ? w.n_indexer_top_k : 0, + 32); + if (attention_impl != DeepSeek4AttentionImpl::Explicit && + head_dim == 512 && n_rot == 64) { + ggml_flash_attn_ext_set_ds4_inverse_rope( + context, kv_start, rope_freq, rope_scale, rope_ext, + rope_attn, w.rope_yarn_beta_fast, + w.rope_yarn_beta_slow, rope_n_ctx_orig, fuse_q_rope); + inverse_rope_fused = true; + } + } + } else { + // Flatten q to [head_dim, n_head*n_tokens] for batched matmul. + ggml_tensor * q_flat = ggml_reshape_2d(ctx, q, head_dim, + n_head * n_tokens); + ggml_tensor * scores = ggml_mul_mat(ctx, kv_attn, q_flat); + scores = ggml_scale(ctx, scores, kq_scale); + if (score_mask) { + if (n_tokens > 1) { + ggml_tensor * m3 = ggml_reshape_3d(ctx, score_mask, + n_attn, 1, n_tokens); + ggml_tensor * s3 = ggml_reshape_3d(ctx, scores, + n_attn, n_head, n_tokens); + scores = ggml_reshape_2d(ctx, ggml_add(ctx, s3, m3), + n_attn, n_head * n_tokens); + } else { + scores = ggml_add(ctx, scores, score_mask); + } + } - // Reshape back to [head_dim, n_head, n_tokens] - context = ggml_reshape_3d(ctx, context, head_dim, n_head, n_tokens); + // DS4 adds one learned per-head sink logit to the denominator, but the + // sink contributes no value vector. + ggml_tensor * probs = nullptr; + if (L.attn_sinks) { + ggml_tensor * sink_scores = ggml_reshape_2d(ctx, L.attn_sinks, + 1, n_head); + if (n_tokens > 1) { + ggml_tensor * sink_shape = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 1, n_head * n_tokens); + sink_scores = ggml_repeat(ctx, sink_scores, sink_shape); + } + ggml_tensor * scores_with_sink = ggml_concat(ctx, scores, + sink_scores, 0); + ggml_tensor * probs_with_sink = ggml_soft_max(ctx, + scores_with_sink); + probs = ggml_view_2d(ctx, probs_with_sink, n_attn, + n_head * n_tokens, probs_with_sink->nb[1], 0); + } else { + probs = ggml_soft_max(ctx, scores); + } + ggml_tensor * kv_t = ggml_cont(ctx, ggml_transpose(ctx, kv_attn)); + context = ggml_mul_mat(ctx, kv_t, probs); + context = ggml_reshape_3d(ctx, context, head_dim, n_head, n_tokens); + } // ── Inverse tail RoPE on attention output ─────────────────────── - ggml_tensor * neg_pos = cached_inputs ? cached_inputs->neg_pos : nullptr; - if (!neg_pos) { - neg_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); - ggml_set_input(neg_pos); - std::vector neg_vals(n_tokens); - for (int i = 0; i < n_tokens; i++) neg_vals[i] = -(kv_start + i); - i32_array_inputs.push_back({neg_pos, std::move(neg_vals)}); - } - context = build_tail_rope_3d(ctx, context, neg_pos, n_rot, head_dim, n_head, n_tokens, - rope_freq, rope_scale, rope_ext, rope_attn, - w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + if (!inverse_rope_fused) { + ggml_tensor * neg_pos = cached_inputs ? cached_inputs->neg_pos : nullptr; + if (!neg_pos) { + neg_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); + ggml_set_input(neg_pos); + std::vector neg_vals(n_tokens); + for (int i = 0; i < n_tokens; i++) neg_vals[i] = -(kv_start + i); + i32_array_inputs.push_back({neg_pos, std::move(neg_vals)}); + } + context = build_tail_rope_3d( + ctx, context, neg_pos, n_rot, head_dim, n_head, n_tokens, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, + rope_n_ctx_orig); + } // Flatten to [head_dim*n_head, n_tokens] for output projection ggml_tensor * attn_out = ggml_reshape_2d(ctx, context, head_dim * n_head, n_tokens); @@ -1043,16 +2101,28 @@ static ggml_tensor * build_mla_attention( const int group_dim = head_dim * (n_head / n_out_group); // 512 * 8 = 4096 // Reshape attn_out: [32768, n_tokens] → [4096, 8, n_tokens] → permute to [4096, n_tokens, 8] attn_out = ggml_reshape_3d(ctx, attn_out, group_dim, n_out_group, n_tokens); - attn_out = ggml_cont(ctx, ggml_permute(ctx, attn_out, 0, 2, 1, 3)); + attn_out = ggml_permute(ctx, attn_out, 0, 2, 1, 3); + if (n_tokens == 1) { + attn_out = ggml_cont(ctx, attn_out); + } // attn_out is now [group_dim, n_tokens, n_out_group] ggml_tensor * out_a_3d = ggml_reshape_3d(ctx, L.attn_output_a, group_dim, n_lora_o, n_out_group); // out_a_3d: [group_dim, n_lora_o, n_out_group] — ne[2] matches ggml_tensor * attn_low = ggml_mul_mat(ctx, out_a_3d, attn_out); // attn_low: [n_lora_o, n_tokens, n_out_group] - // Permute back to [n_lora_o, n_out_group, n_tokens] then flatten - attn_low = ggml_cont(ctx, ggml_permute(ctx, attn_low, 0, 2, 1, 3)); - attn_low = ggml_reshape_2d(ctx, attn_low, n_lora_o * n_out_group, n_tokens); - ggml_tensor * out = ggml_mul_mat(ctx, L.attn_output_b, attn_low); + ggml_tensor * out = nullptr; + if (n_tokens > 1) { + // Batched ROCmFPX MMQ consumes src1's channel stride directly. This + // avoids materializing both permutations (~256 MiB/layer at 2K). + out = ggml_mul_mat_grouped_src(ctx, L.attn_output_b, attn_low); + } else { + // Preserve the established single-token graph and its numerical + // behavior. Decode is intentionally outside the prefill fast path. + attn_low = ggml_cont(ctx, ggml_permute(ctx, attn_low, 0, 2, 1, 3)); + attn_low = ggml_reshape_2d( + ctx, attn_low, n_lora_o * n_out_group, n_tokens); + out = ggml_mul_mat(ctx, L.attn_output_b, attn_low); + } return out; } @@ -1109,6 +2179,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 +2296,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 +2315,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 +2324,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 +2370,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 +2475,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 +2485,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 +2499,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 +2511,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 +2530,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; @@ -1671,6 +2874,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, @@ -1689,10 +2920,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); @@ -1712,8 +2948,9 @@ 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); ggml_tensor * routed_out = nullptr; @@ -1721,24 +2958,41 @@ 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); - 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); + 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); - 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 +3045,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 +3067,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,11 +3075,87 @@ static float cpu_dot_f16_row(const uint16_t * row, const float * x, int cols) { return acc; } -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) - for (int r = 0; r < rows; r++) { - const uint16_t * row = mat + (size_t)r * cols; +#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) + for (int r = 0; r < rows; r++) { + const uint16_t * row = mat + (size_t)r * cols; out[r] = cpu_dot_f16_row(row, x, cols); } } @@ -1839,6 +3171,126 @@ 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); +} + +// Reuse the persistent HC worker pool across tokens. HC inner work remains +// serial because the pool is intentionally non-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 token) { + fn(token, token + 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]; @@ -1969,7 +3421,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 +3486,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 +3506,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, @@ -2072,7 +3525,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) { @@ -2087,7 +3540,8 @@ static void hc_pre_batch(std::vector & working, sinkhorn_iters, hc_eps, flat.data(), - mix); + mix, + /*serial_fn=*/n_tokens > 1); } }); } @@ -2116,7 +3570,25 @@ 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); - ds4_parallel_for_tokens(n_tokens, 8, [&](int t0, int t1) { + 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_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, @@ -2138,7 +3610,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); @@ -2163,18 +3635,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 +4247,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 +4257,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 +4370,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 +4412,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,51 +4453,1406 @@ bool deepseek4_step( } cache.cur_pos = next_pos; + if (telemetry) { + telemetry->total_us += ds4_elapsed_us(step_t0, Ds4TimingClock::now()); + } return true; } -bool deepseek4_step_layer_range( +// ─── 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, + int n_tokens) { + ggml_tensor * shared_out = build_shared_ffn(ctx, ffn_normed, w, L); + ggml_tensor * logits = ggml_mul_mat(ctx, L.ffn_gate_inp, ffn_normed); + ggml_tensor * probs = ggml_sqrt(ctx, ggml_softplus(ctx, logits)); + + 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, n_tokens); + ggml_tensor * gate_e = ggml_mul_mat_id(ctx, L.ffn_gate_exps, cur_3d, hash_ids); + ggml_tensor * up_e = ggml_mul_mat_id(ctx, L.ffn_up_exps, cur_3d, hash_ids); + 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, hash_ids); + if (!raw_mmid) { + down_e = ggml_reshape_3d(ctx, down_e, w.n_embd, n_used, n_tokens); + } + + ggml_tensor * probs_3d = ggml_reshape_3d( + ctx, probs, 1, w.n_expert, n_tokens); + ggml_tensor * weights = ggml_get_rows(ctx, probs_3d, hash_ids); + weights = ggml_reshape_2d(ctx, weights, n_used, n_tokens); + ggml_tensor * w_sum = ggml_sum_rows(ctx, weights); + w_sum = ggml_clamp(ctx, w_sum, 6.103515625e-5f, INFINITY); + weights = ggml_div(ctx, weights, w_sum); + 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, n_tokens); + 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, n_tokens); + routed_out = ggml_repeat_back(ctx, routed_out, sum_shape); + routed_out = ggml_reshape_2d( + ctx, routed_out, w.n_embd, n_tokens); + } + return ggml_add(ctx, shared_out, routed_out); +} + +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, - std::vector & hc_state, - const float * embed, - int n_tokens, + const std::vector & hc_weights, + const HcWeightsCpu & hc_out_weights, + const std::vector & hash_tables, int kv_start, - int layer_begin, - int layer_end, - std::vector * out_logits, - const int32_t * token_ids, - DeepSeek4StepTelemetry * telemetry) { - const auto step_t0 = Ds4TimingClock::now(); + 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); - // NOTE: The old deepseek4_step() lacks HC implementation. - // Always use the HC-enabled layer_range path below. - - // ── Partial layer-range forward with HC ───────────────────────────── const int n_embd = w.n_embd; const int n_hc = w.n_hc; - const int hc_dim = n_hc * n_embd; - const bool is_last_shard = (layer_end >= w.n_layer); - - // Initialize HC state. - // First shard (layer_begin=0): embed is token embeddings [n_embd × n_tokens], - // replicate into n_hc streams. - // Later shards: embed is full HC state [hc_dim × n_tokens] from previous shard. - if (hc_state.size() != (size_t)hc_dim * (size_t)n_tokens) { - hc_state.resize((size_t)hc_dim * (size_t)n_tokens); + 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); } - if (layer_begin == 0) { - // First shard: replicate embedding into all HC streams - for (int t = 0; t < n_tokens; t++) { - for (int h = 0; h < n_hc; h++) { - memcpy(hc_state.data() + (size_t)t * hc_dim + (size_t)h * n_embd, - embed + (size_t)t * n_embd, (size_t)n_embd * sizeof(float)); - } + 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]); } - } else { - // Later shard: embed contains full HC state from previous shard - memcpy(hc_state.data(), embed, sizeof(float) * (size_t)hc_dim * (size_t)n_tokens); + 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, 1); + } 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; + 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; +} + +// Exact-order prefill control: retain the layer-major HC/FFN schedule, but run +// the attention subgraph one token at a time. This preserves the q=1 QKV, +// compressor, causal-attention, and output-projection reduction order while +// still allowing the token-independent FFN to use a multi-row ROCMFP graph. +static bool ds4_run_exact_tokenwise_prefill_attention( + ggml_backend_t backend, + const DeepSeek4Weights & w, + const DeepSeek4Layer & L, + DeepSeek4LayerCache & lc, + int il, + const float * cur, + int n_tokens, + int kv_start, + DeepSeek4AttentionImpl attention_impl, + std::vector & attn_out_host, + DeepSeek4CachedLayerAlloc & attn_alloc, + DeepSeek4StepTelemetry * telemetry) { + if (!backend || !cur || n_tokens <= 1 || kv_start < 0) return false; + + const int n_embd = w.n_embd; + for (int ti = 0; ti < n_tokens; ++ti) { + const auto build_t0 = Ds4TimingClock::now(); + ggml_init_params params{}; + params.mem_size = ds4_attn_step_meta_size(1); + params.mem_buffer = nullptr; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) return false; + + ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, 1); + ggml_set_input(inp); + std::vector i32_inputs; + std::vector i32_array_inputs; + std::vector i64_array_inputs; + std::vector f32_array_inputs; + ggml_cgraph * gf = ggml_new_graph_custom( + ctx, ds4_attn_step_graph_size(1), false); + ggml_tensor * normed = build_rms_norm(ctx, inp, L.attn_norm, w.rms_eps); + ggml_tensor * attn_out = build_mla_attention( + ctx, gf, normed, w, L, lc, il, kv_start + ti, 1, nullptr, + i32_inputs, i32_array_inputs, i64_array_inputs, &f32_array_inputs, + attention_impl); + ggml_set_output(attn_out); + ggml_build_forward_expand(gf, attn_out); + + if (!attn_alloc.valid() || attn_alloc.owner_ctx != w.ctx || + attn_alloc.backend != backend) { + attn_alloc.free(); + attn_alloc.alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + attn_alloc.owner_ctx = w.ctx; + attn_alloc.backend = backend; + } + if (!attn_alloc.alloc || !ggml_gallocr_alloc_graph(attn_alloc.alloc, gf)) { + std::fprintf(stderr, + "[deepseek4] exact prefill attn alloc failed layer %d token %d\n", + il, ti); + ggml_free(ctx); + return false; + } + if (telemetry) { + telemetry->attn_build_us += ds4_elapsed_us(build_t0, Ds4TimingClock::now()); + } + + ggml_backend_tensor_set(inp, cur + (size_t) ti * n_embd, 0, + sizeof(float) * (size_t) n_embd); + for (const auto & b : i32_inputs) { + ggml_backend_tensor_set(b.tensor, &b.value, 0, sizeof(b.value)); + } + for (const auto & b : i32_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(int32_t) * b.values.size()); + } + for (const auto & b : i64_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(int64_t) * b.values.size()); + } + for (const auto & b : f32_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(float) * b.values.size()); + } + + const auto compute_t0 = Ds4TimingClock::now(); + const ggml_status status = ggml_backend_graph_compute(backend, gf); + if (telemetry) { + telemetry->attn_compute_us += ds4_elapsed_us( + compute_t0, Ds4TimingClock::now()); + } + if (status != GGML_STATUS_SUCCESS) { + std::fprintf(stderr, + "[deepseek4] exact prefill attn compute failed layer %d token %d\n", + il, ti); + ggml_free(ctx); + return false; + } + + const auto read_t0 = Ds4TimingClock::now(); + ggml_backend_tensor_get(attn_out, + attn_out_host.data() + (size_t) ti * n_embd, + 0, sizeof(float) * (size_t) n_embd); + if (telemetry) { + telemetry->attn_read_us += ds4_elapsed_us(read_t0, Ds4TimingClock::now()); + } + + // Publish compressor rows immediately. The next token in this layer + // must observe a row flushed by the current token, matching the q=1 + // reference when a prefill band crosses a compressor boundary. + const int ratio = (int) w.compress_ratios[il]; + if (ratio > 0) { + const int next_pos = kv_start + ti + 1; + lc.n_comp = std::max(lc.n_comp, next_pos / ratio); + if (ratio == 4) { + lc.n_index_comp = std::max(lc.n_index_comp, + next_pos / ratio); + } + } + ggml_free(ctx); + } + return true; +} + +// Layer-major DS4 prefill. Each layer is one GPU graph containing batched HC, +// attention, MoE and HC post-processing. The HC state is kept in two external +// device tensors and ping-ponged between layers, eliminating the two host +// readbacks per layer in the reference implementation. Attention reads a +// snapshot of the previous SWA window plus the current ubatch; only the final +// SWA tail is committed to the persistent ring. The compressor publishes every +// ratio-4/ratio-128 boundary crossed by the ubatch. +// +struct Ds4LayerMajorCachedLayer { + void * meta_buffer = nullptr; + size_t meta_size = 0; + ggml_context * ctx = nullptr; + ggml_cgraph * gf = nullptr; + std::vector i32_inputs; + std::vector i32_array_inputs; + std::vector i64_array_inputs; + std::vector f32_array_inputs; + std::vector allocated_tensors; + ggml_tensor * hash_ids = nullptr; + ggml_tensor * logits = nullptr; + + void destroy() { + if (ctx) { + ggml_free(ctx); + ctx = nullptr; + } + if (meta_buffer) { + std::free(meta_buffer); + meta_buffer = nullptr; + } + meta_size = 0; + gf = nullptr; + hash_ids = nullptr; + logits = nullptr; + i32_inputs.clear(); + i32_array_inputs.clear(); + i64_array_inputs.clear(); + f32_array_inputs.clear(); + allocated_tensors.clear(); + } +}; + +struct Ds4LayerMajorGraphCache { + const ggml_context * owner_ctx = nullptr; + ggml_backend_t backend = nullptr; + PrefillAttentionMode mode = PrefillAttentionMode::Exact; + int n_tokens = 0; + int kv_start = -1; + uint64_t last_use = 0; + bool ready = false; + ggml_context * state_ctx = nullptr; + ggml_backend_buffer_t state_buf = nullptr; + ggml_tensor * state_a = nullptr; + ggml_tensor * state_b = nullptr; + std::vector layers; + + bool matches(const DeepSeek4Weights & w, ggml_backend_t b, + PrefillAttentionMode m, int tokens, int start) const { + return ready && owner_ctx == w.ctx && backend == b && mode == m && + n_tokens == tokens && kv_start == start && + layers.size() == (size_t) w.n_layer; + } + + void destroy() { + for (auto & layer : layers) layer.destroy(); + layers.clear(); + if (state_buf) { + ggml_backend_buffer_free(state_buf); + state_buf = nullptr; + } + if (state_ctx) { + ggml_free(state_ctx); + state_ctx = nullptr; + } + state_a = nullptr; + state_b = nullptr; + owner_ctx = nullptr; + backend = nullptr; + mode = PrefillAttentionMode::Exact; + n_tokens = 0; + kv_start = -1; + last_use = 0; + ready = false; + } +}; + +static thread_local std::array + ds4_layer_major_graph_caches; +static thread_local uint64_t ds4_layer_major_cache_counter = 0; +// A separate gallocr scratch arena per shape consumes several GiB and forces +// the 97-GiB model into managed-memory paging. Every layer executes serially, +// so cached and uncached shapes safely rebind their transient tensors to one +// arena before execution. Retain only the largest/full-chunk topology; tail +// metadata is rebuilt rather than keeping a second long-context graph resident. +static thread_local ggml_gallocr_t ds4_layer_major_shared_alloc = nullptr; +static thread_local const ggml_context * ds4_layer_major_shared_owner = nullptr; +static thread_local ggml_backend_t ds4_layer_major_shared_backend = nullptr; + +static ggml_gallocr_t ds4_layer_major_get_shared_alloc( + const DeepSeek4Weights & w, + ggml_backend_t backend) { + if (ds4_layer_major_shared_alloc && + (ds4_layer_major_shared_owner != w.ctx || + ds4_layer_major_shared_backend != backend)) { + ggml_gallocr_free(ds4_layer_major_shared_alloc); + ds4_layer_major_shared_alloc = nullptr; + } + if (!ds4_layer_major_shared_alloc) { + ds4_layer_major_shared_alloc = ggml_gallocr_new( + ggml_backend_get_default_buffer_type(backend)); + ds4_layer_major_shared_owner = w.ctx; + ds4_layer_major_shared_backend = backend; + } + return ds4_layer_major_shared_alloc; +} + +// Returns 1 on success, 0 when the optimized path is not applicable, and -1 +// after a hard failure. +static int ds4_try_layer_major_prefill( + DeepSeek4FusedDecodeCache & fc, + ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + const std::vector & hc_weights, + const HcWeightsCpu & hc_out_weights, + const std::vector & hash_tables, + std::vector & hash_scratch, + const float * embed, + int n_tokens, + int kv_start, + std::vector & out_logits, + const int32_t * token_ids, + DeepSeek4StepTelemetry * telemetry) { + if (!backend || !embed || n_tokens <= 4 || + n_tokens > DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS || + kv_start < 0 || w.moe_hybrid) { + return 0; + } + if (cache.prefill_mode == PrefillAttentionMode::Exact) return 0; + if (!ds4_backend_is_gpu(backend) || !hc_out_weights.loaded || + hc_out_weights.scale_data.empty() || !w.output_hc_fn || + !w.output_hc_base) { + return 0; + } + for (int il = 0; il < w.n_layer; ++il) { + const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; + const DeepSeek4Layer & L = w.layers[(size_t) il]; + if (!hlw.attn.loaded || hlw.attn.scale_data.size() < 3 || + !hlw.ffn.loaded || hlw.ffn.scale_data.size() < 3 || + !L.hc_attn_base || !L.hc_ffn_base) { + return 0; + } + } + + if (fc.owner_ctx != w.ctx || fc.backend != backend) { + fc.destroy(); + fc.owner_ctx = w.ctx; + fc.backend = backend; + } + if (!ds4_fused_ensure_fn_mirrors(fc, backend, w, hc_weights, + hc_out_weights)) { + std::fprintf(stderr, + "[deepseek4-prefill] failed to create HC weight mirrors\n"); + return -1; + } + + const int n_embd = w.n_embd; + const int n_hc = w.n_hc; + const int64_t hc_dim = (int64_t) n_embd * n_hc; + const int64_t mix_dim = 2 * (int64_t) n_hc + (int64_t) n_hc * n_hc; + const int next_pos = kv_start + n_tokens; + + Ds4LayerMajorGraphCache * graph_cache = nullptr; + bool cache_hit = false; + bool cache_build = false; + if (token_ids) { + ++ds4_layer_major_cache_counter; + for (auto & candidate : ds4_layer_major_graph_caches) { + if (candidate.matches(w, backend, cache.prefill_mode, + n_tokens, kv_start)) { + graph_cache = &candidate; + cache_hit = true; + break; + } + } + if (!graph_cache) { + auto & candidate = ds4_layer_major_graph_caches.front(); + // Do not evict a full/larger chunk for an equal-size graph at a + // later position or for a short tail. Both execute with the shared + // scratch arena below, but only the dominant topology stays cached. + const bool same_owner = candidate.owner_ctx == w.ctx && + candidate.backend == backend && + candidate.mode == cache.prefill_mode; + if (!candidate.ready || !same_owner || + n_tokens > candidate.n_tokens) { + graph_cache = &candidate; + graph_cache->destroy(); + graph_cache->owner_ctx = w.ctx; + graph_cache->backend = backend; + graph_cache->mode = cache.prefill_mode; + graph_cache->n_tokens = n_tokens; + graph_cache->kv_start = kv_start; + graph_cache->layers.resize((size_t) w.n_layer); + cache_build = true; + } + } + if (graph_cache) { + graph_cache->last_use = ds4_layer_major_cache_counter; + } + } + + // Persistent ping-pong state lives outside the per-layer gallocr arena. + ggml_context * state_ctx = cache_hit ? graph_cache->state_ctx : nullptr; + ggml_tensor * state_a = cache_hit ? graph_cache->state_a : nullptr; + ggml_tensor * state_b = cache_hit ? graph_cache->state_b : nullptr; + ggml_backend_buffer_t state_buf = cache_hit ? graph_cache->state_buf : nullptr; + if (!cache_hit) { + ggml_init_params state_params{}; + state_params.mem_size = 4 * ggml_tensor_overhead() + 4096; + state_params.no_alloc = true; + state_ctx = ggml_init(state_params); + if (!state_ctx) { + if (cache_build) graph_cache->destroy(); + return -1; + } + state_a = ggml_new_tensor_2d(state_ctx, GGML_TYPE_F32, + hc_dim, n_tokens); + state_b = ggml_new_tensor_2d(state_ctx, GGML_TYPE_F32, + hc_dim, n_tokens); + state_buf = ggml_backend_alloc_ctx_tensors(state_ctx, backend); + if (!state_buf) { + ggml_free(state_ctx); + if (cache_build) graph_cache->destroy(); + return -1; + } + if (cache_build) { + graph_cache->state_ctx = state_ctx; + graph_cache->state_a = state_a; + graph_cache->state_b = state_b; + graph_cache->state_buf = state_buf; + } + } + + std::vector initial((size_t) hc_dim * n_tokens); + for (int t = 0; t < n_tokens; ++t) { + for (int h = 0; h < n_hc; ++h) { + std::memcpy(initial.data() + (size_t) t * hc_dim + + (size_t) h * n_embd, + embed + (size_t) t * n_embd, + sizeof(float) * (size_t) n_embd); + } + } + ggml_backend_tensor_set(state_a, initial.data(), 0, + sizeof(float) * initial.size()); + initial.clear(); + initial.shrink_to_fit(); + + ggml_gallocr_t alloc = ds4_layer_major_get_shared_alloc(w, backend); + if (!alloc) { + if (cache_build) { + graph_cache->destroy(); + } else { + ggml_backend_buffer_free(state_buf); + ggml_free(state_ctx); + } + return -1; + } + static thread_local std::vector meta_arena; + const size_t meta_bytes = 160u * 1024 * 1024; + if (!graph_cache && meta_arena.size() < meta_bytes) { + meta_arena.resize(meta_bytes); + } + + auto fail = [&](const char * what, int il) { + std::fprintf(stderr, "[deepseek4-prefill] %s at layer %d\n", what, il); + if (graph_cache) { + graph_cache->destroy(); + } else { + ggml_backend_buffer_free(state_buf); + ggml_free(state_ctx); + } + return -1; + }; + + if (cache_hit) { + for (int il = 0; il < w.n_layer; ++il) { + Ds4LayerMajorCachedLayer & layer = + graph_cache->layers[(size_t) il]; + for (ggml_tensor * tensor : layer.allocated_tensors) { + tensor->data = nullptr; + tensor->buffer = nullptr; + } + const auto alloc_t0 = Ds4TimingClock::now(); + if (!layer.ctx || !layer.gf || + !ggml_gallocr_alloc_graph(alloc, layer.gf)) { + return fail("cached scratch allocation failed", il); + } + if (telemetry) { + telemetry->full_graph_alloc_us += ds4_elapsed_us( + alloc_t0, Ds4TimingClock::now()); + } + for (const auto & b : layer.i32_inputs) { + ggml_backend_tensor_set(b.tensor, &b.value, 0, + sizeof(b.value)); + } + for (const auto & b : layer.i32_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(int32_t) * b.values.size()); + } + for (const auto & b : layer.i64_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(int64_t) * b.values.size()); + } + for (const auto & b : layer.f32_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(float) * b.values.size()); + } + if (layer.hash_ids) { + const int n_used = w.n_expert_used; + hash_scratch.resize((size_t) n_used * n_tokens); + const auto & table = hash_tables[(size_t) il].ids; + for (int t = 0; t < n_tokens; ++t) { + std::memcpy( + hash_scratch.data() + (size_t) t * n_used, + table.data() + (size_t) token_ids[t] * n_used, + sizeof(int32_t) * (size_t) n_used); + } + ggml_backend_tensor_set( + layer.hash_ids, hash_scratch.data(), 0, + sizeof(int32_t) * hash_scratch.size()); + } + + const auto compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, layer.gf) != + GGML_STATUS_SUCCESS) { + return fail("cached compute failed", il); + } + if (telemetry) { + telemetry->full_graph_compute_us += ds4_elapsed_us( + compute_t0, Ds4TimingClock::now()); + } + if (layer.logits) { + out_logits.resize((size_t) w.n_vocab); + ggml_backend_tensor_get( + layer.logits, out_logits.data(), 0, + sizeof(float) * (size_t) w.n_vocab); + } + + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + const int ratio = (int) w.compress_ratios[(size_t) il]; + if (ratio > 0) { + lc.n_comp = std::max(lc.n_comp, next_pos / ratio); + if (ratio == 4) { + lc.n_index_comp = std::max( + lc.n_index_comp, next_pos / ratio); + } + } + } + cache.cur_pos = next_pos; + return out_logits.empty() ? -1 : 1; + } + + ggml_tensor * state_in = state_a; + ggml_tensor * state_out = state_b; + for (int il = 0; il < w.n_layer; ++il) { + const auto build_t0 = Ds4TimingClock::now(); + Ds4LayerMajorCachedLayer * cached_layer = cache_build + ? &graph_cache->layers[(size_t) il] : nullptr; + ggml_init_params params{}; + if (cached_layer) { + cached_layer->meta_size = meta_bytes; + cached_layer->meta_buffer = std::malloc(meta_bytes); + if (!cached_layer->meta_buffer) { + return fail("cached metadata allocation failed", il); + } + params.mem_size = cached_layer->meta_size; + params.mem_buffer = cached_layer->meta_buffer; + } else { + params.mem_size = meta_arena.size(); + params.mem_buffer = meta_arena.data(); + } + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + if (!ctx) return fail("metadata allocation failed", il); + if (cached_layer) cached_layer->ctx = ctx; + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 65536, false); + if (!gf) { + if (!cached_layer) ggml_free(ctx); + return fail("graph allocation failed", il); + } + if (cached_layer) cached_layer->gf = gf; + + const DeepSeek4Layer & L = w.layers[(size_t) il]; + DeepSeek4LayerCache & lc = cache.layers[(size_t) il]; + const HcLayerWeightsCpu & hlw = hc_weights[(size_t) il]; + + // HC pre -> batched attention. + ggml_tensor * norm_hc = ggml_rms_norm(ctx, state_in, w.hc_eps); + ggml_tensor * mix_attn = ggml_mul_mat(ctx, + fc.fn_attn_f16[(size_t) il], norm_hc); + mix_attn = ggml_reshape_2d(ctx, mix_attn, mix_dim, n_tokens); + ggml_tensor * attn_base = ds4_fused_hc_base_f32(ctx, + L.hc_attn_base); + ggml_tensor * pre_attn = ggml_ds4_hc_pre( + ctx, mix_attn, attn_base, state_in, n_hc, + w.n_hc_sinkhorn_iter, hlw.attn.scale_data[0], + hlw.attn.scale_data[1], hlw.attn.scale_data[2]); + ggml_tensor * attn_in = ggml_view_2d(ctx, pre_attn, n_embd, + n_tokens, pre_attn->nb[1], 0); + ggml_tensor * split_attn = ggml_view_2d( + ctx, pre_attn, mix_dim, n_tokens, pre_attn->nb[1], + (size_t) n_embd * sizeof(float)); + + std::vector i32_inputs; + std::vector i32_array_inputs; + std::vector i64_array_inputs; + std::vector f32_array_inputs; + ggml_tensor * attn_normed = build_rms_norm(ctx, attn_in, + L.attn_norm, w.rms_eps); + const DeepSeek4AttentionImpl attention_impl = + cache.prefill_mode == PrefillAttentionMode::Sparse + ? DeepSeek4AttentionImpl::SparseFlash + : DeepSeek4AttentionImpl::DenseFlash; + ggml_tensor * attn_out = build_mla_attention( + ctx, gf, attn_normed, w, L, lc, il, kv_start, n_tokens, + nullptr, i32_inputs, i32_array_inputs, i64_array_inputs, + &f32_array_inputs, attention_impl); + if (!attn_out) { + if (!cached_layer) ggml_free(ctx); + return fail("attention graph build failed", il); + } + ggml_tensor * hc_after_attn = ggml_ds4_hc_post( + ctx, state_in, attn_out, split_attn, n_hc); + + // HC pre -> batched MoE. + norm_hc = ggml_rms_norm(ctx, hc_after_attn, w.hc_eps); + ggml_tensor * mix_ffn = ggml_mul_mat(ctx, + fc.fn_ffn_f16[(size_t) il], norm_hc); + mix_ffn = ggml_reshape_2d(ctx, mix_ffn, mix_dim, n_tokens); + ggml_tensor * ffn_base = ds4_fused_hc_base_f32(ctx, L.hc_ffn_base); + ggml_tensor * pre_ffn = ggml_ds4_hc_pre( + ctx, mix_ffn, ffn_base, hc_after_attn, n_hc, + w.n_hc_sinkhorn_iter, hlw.ffn.scale_data[0], + hlw.ffn.scale_data[1], hlw.ffn.scale_data[2]); + ggml_tensor * ffn_in = ggml_view_2d(ctx, pre_ffn, n_embd, + n_tokens, pre_ffn->nb[1], 0); + ggml_tensor * split_ffn = ggml_view_2d( + ctx, pre_ffn, mix_dim, n_tokens, pre_ffn->nb[1], + (size_t) n_embd * sizeof(float)); + ggml_tensor * ffn_normed = build_rms_norm(ctx, ffn_in, + L.ffn_norm, w.rms_eps); + ggml_tensor * hash_ids = nullptr; + ggml_tensor * ffn_out = nullptr; + const bool hash_routed = il < w.n_hash_layer && L.ffn_gate_tid2eid && + token_ids && hash_tables[(size_t) il].loaded; + if (hash_routed) { + hash_ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, + w.n_expert_used, n_tokens); + ggml_set_input(hash_ids); + ffn_out = ds4_build_hash_routed_ffn( + ctx, w, L, ffn_normed, hash_ids, n_tokens); + } else { + ffn_out = build_moe_ffn(ctx, ffn_normed, w, L, il, n_tokens); + } + if (!ffn_out) { + if (!cached_layer) ggml_free(ctx); + return fail("FFN graph build failed", il); + } + ggml_tensor * hc_next = ggml_ds4_hc_post( + ctx, hc_after_attn, ffn_out, split_ffn, n_hc); + + // Persist HC state for the next layer before this layer's gallocr + // scratch buffer is reused. + ggml_tensor * state_copy = ggml_cpy(ctx, hc_next, state_out); + ggml_set_output(state_copy); + ggml_build_forward_expand(gf, state_copy); + + ggml_tensor * logits = nullptr; + if (il + 1 == w.n_layer) { + ggml_tensor * last_hc = ggml_view_2d( + ctx, hc_next, hc_dim, 1, hc_next->nb[1], + (size_t) (n_tokens - 1) * hc_next->nb[1]); + last_hc = ggml_reshape_1d(ctx, last_hc, hc_dim); + ggml_tensor * out_hc_norm = ggml_rms_norm(ctx, last_hc, w.hc_eps); + ggml_tensor * out_mix = ggml_mul_mat(ctx, fc.fn_out_f16, + out_hc_norm); + out_mix = ggml_reshape_1d(ctx, out_mix, n_hc); + ggml_tensor * out_base = ds4_fused_hc_base_f32(ctx, + w.output_hc_base); + ggml_tensor * final_embd = ggml_ds4_hc_out( + ctx, out_mix, out_base, last_hc, n_hc, + hc_out_weights.scale_data[0]); + ggml_tensor * final_2d = ggml_reshape_2d(ctx, final_embd, + n_embd, 1); + ggml_tensor * out_normed = build_rms_norm(ctx, final_2d, + w.out_norm, w.rms_eps); + logits = ggml_mul_mat(ctx, w.output, out_normed); + ggml_set_output(logits); + ggml_build_forward_expand(gf, logits); + } + + if (cached_layer) { + auto remember_unallocated = [&](ggml_tensor * tensor) { + if (tensor && tensor->data == nullptr && + tensor->buffer == nullptr) { + cached_layer->allocated_tensors.push_back(tensor); + } + }; + const int n_graph_nodes = ggml_graph_n_nodes(gf); + for (int i = 0; i < n_graph_nodes; ++i) { + ggml_tensor * node = ggml_graph_node(gf, i); + remember_unallocated(node); + for (int j = 0; j < GGML_MAX_SRC; ++j) { + remember_unallocated(node->src[j]); + } + } + auto & tensors = cached_layer->allocated_tensors; + std::sort(tensors.begin(), tensors.end()); + tensors.erase(std::unique(tensors.begin(), tensors.end()), + tensors.end()); + } + if (!ggml_gallocr_alloc_graph(alloc, gf)) { + if (!cached_layer) ggml_free(ctx); + return fail("scratch allocation failed", il); + } + for (const auto & b : i32_inputs) { + ggml_backend_tensor_set(b.tensor, &b.value, 0, sizeof(b.value)); + } + for (const auto & b : i32_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(int32_t) * b.values.size()); + } + for (const auto & b : i64_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(int64_t) * b.values.size()); + } + for (const auto & b : f32_array_inputs) { + ggml_backend_tensor_set(b.tensor, b.values.data(), 0, + sizeof(float) * b.values.size()); + } + if (hash_ids) { + const int n_used = w.n_expert_used; + hash_scratch.resize((size_t) n_used * n_tokens); + const auto & table = hash_tables[(size_t) il].ids; + for (int t = 0; t < n_tokens; ++t) { + std::memcpy(hash_scratch.data() + (size_t) t * n_used, + table.data() + (size_t) token_ids[t] * n_used, + sizeof(int32_t) * (size_t) n_used); + } + ggml_backend_tensor_set(hash_ids, hash_scratch.data(), 0, + sizeof(int32_t) * hash_scratch.size()); + } + if (telemetry) { + telemetry->full_graph_build_us += ds4_elapsed_us( + build_t0, Ds4TimingClock::now()); + } + + const auto compute_t0 = Ds4TimingClock::now(); + if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { + if (!cached_layer) ggml_free(ctx); + return fail("compute failed", il); + } + if (telemetry) { + telemetry->full_graph_compute_us += ds4_elapsed_us( + compute_t0, Ds4TimingClock::now()); + } + + if (logits) { + out_logits.resize((size_t) w.n_vocab); + ggml_backend_tensor_get(logits, out_logits.data(), 0, + sizeof(float) * (size_t) w.n_vocab); + } + + const int ratio = (int) w.compress_ratios[(size_t) il]; + if (ratio > 0) { + lc.n_comp = std::max(lc.n_comp, next_pos / ratio); + if (ratio == 4) { + lc.n_index_comp = std::max(lc.n_index_comp, + next_pos / ratio); + } + } + if (cached_layer) { + cached_layer->i32_inputs = std::move(i32_inputs); + cached_layer->i32_array_inputs = std::move(i32_array_inputs); + cached_layer->i64_array_inputs = std::move(i64_array_inputs); + cached_layer->f32_array_inputs = std::move(f32_array_inputs); + cached_layer->hash_ids = hash_ids; + cached_layer->logits = logits; + } else { + ggml_free(ctx); + } + std::swap(state_in, state_out); + } + + if (cache_build) { + graph_cache->ready = true; + } else { + ggml_backend_buffer_free(state_buf); + ggml_free(state_ctx); + } + cache.cur_pos = next_pos; + return out_logits.empty() ? -1 : 1; +} + + +bool deepseek4_step_layer_range( + ggml_backend_t backend, + const DeepSeek4Weights & w, + DeepSeek4Cache & cache, + std::vector & hc_state, + const float * embed, + int n_tokens, + int kv_start, + int layer_begin, + int layer_end, + std::vector * out_logits, + const int32_t * token_ids, + DeepSeek4StepTelemetry * telemetry, + bool allow_decode_graph_reuse) { + const auto step_t0 = Ds4TimingClock::now(); + + // NOTE: The old deepseek4_step() lacks HC implementation. + // Always use the HC-enabled layer_range path below. + + // ── Partial layer-range forward with HC ───────────────────────────── + const int n_embd = w.n_embd; + const int n_hc = w.n_hc; + const int hc_dim = n_hc * n_embd; + const bool is_last_shard = (layer_end >= w.n_layer); + + // Initialize HC state. + // First shard (layer_begin=0): embed is token embeddings [n_embd × n_tokens], + // replicate into n_hc streams. + // Later shards: embed is full HC state [hc_dim × n_tokens] from previous shard. + if (hc_state.size() != (size_t)hc_dim * (size_t)n_tokens) { + hc_state.resize((size_t)hc_dim * (size_t)n_tokens); + } + if (layer_begin == 0) { + // First shard: replicate embedding into all HC streams + for (int t = 0; t < n_tokens; t++) { + for (int h = 0; h < n_hc; h++) { + memcpy(hc_state.data() + (size_t)t * hc_dim + (size_t)h * n_embd, + embed + (size_t)t * n_embd, (size_t)n_embd * sizeof(float)); + } + } + } else { + // Later shard: embed contains full HC state from previous shard + memcpy(hc_state.data(), embed, sizeof(float) * (size_t)hc_dim * (size_t)n_tokens); } // Lazy-load per-layer HC weights on CPU (static to avoid reloading) @@ -2860,10 +5867,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 +5901,73 @@ 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); + + // Large full-model batches use the device-resident layer-major pipeline. + if (n_tokens > 4 && + n_tokens <= DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS && layer_begin == 0 && + is_last_shard && out_logits && ds4_backend_is_gpu(backend)) { + const int prc = ds4_try_layer_major_prefill( + fused_decode_graph_cache, backend, w, cache, + hc_layer_weights_range, hc_output_weights_range, + hash_routing_tables_range, scratch.hash_expert_ids, embed, + n_tokens, kv_start, *out_logits, token_ids, telemetry); + if (prc < 0) return false; + if (prc > 0) { + if (telemetry) { + telemetry->total_us += ds4_elapsed_us(step_t0, + Ds4TimingClock::now()); + } + return true; + } + } + + 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 +5976,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 +6031,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 +6051,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,16 +6105,31 @@ 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; DeepSeek4CachedDecodeAttnGraph * cached_attn = nullptr; - if (reuse_decode_attn) { + const bool exact_tokenwise_prefill = + !reuse_decode_attn && n_tokens > 1; + if (exact_tokenwise_prefill) { + const DeepSeek4AttentionImpl attention_impl = + cache.prefill_mode == PrefillAttentionMode::Sparse + ? DeepSeek4AttentionImpl::SparseFlash + : DeepSeek4AttentionImpl::Explicit; + if (!ds4_run_exact_tokenwise_prefill_attention( + backend, w, L, lc, il, cur.data(), n_tokens, kv_start, + attention_impl, attn_out_host, + cached_attn_allocs[(size_t) il], telemetry)) { + return false; + } + } else if (reuse_decode_attn) { const int n_raw = std::min(kv_start + 1, w.n_swa); const int n_comp_attn = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; const int n_index_comp = (ratio == 4) ? ds4_comp_rows_used(lc.index_comp_kv, lc.n_index_comp, 4, token_pos) : 0; + 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 +6139,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 +6152,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 +6172,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 +6204,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,14 +6220,21 @@ 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; + 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); ggml_tensor * normed = build_rms_norm(ctx, inp, L.attn_norm, w.rms_eps); + const DeepSeek4AttentionImpl attention_impl = + cache.prefill_mode == PrefillAttentionMode::Sparse + ? DeepSeek4AttentionImpl::SparseFlash + : DeepSeek4AttentionImpl::Explicit; attn_out = build_mla_attention(ctx, gf, normed, w, L, lc, il, kv_start, n_tokens, nullptr, i32_inputs, i32_array_inputs, - i64_array_inputs); + i64_array_inputs, + &f32_array_inputs, + attention_impl); ggml_set_output(attn_out); ggml_build_forward_expand(gf, attn_out); @@ -3153,8 +6262,11 @@ 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()); } + if (!exact_tokenwise_prefill) { const auto attn_compute_t0 = Ds4TimingClock::now(); if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { std::fprintf(stderr, "[deepseek4] attn compute failed layer %d\n", il); @@ -3162,7 +6274,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 +6292,13 @@ 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 +6316,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 +6336,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 +6441,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 +6458,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 +6478,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()); } @@ -3419,7 +6521,10 @@ bool deepseek4_step_layer_range( ggml_context * ctx = ggml_init(params); if (!ctx) return false; - ggml_tensor * inp = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); + const bool last_only = n_tokens > 1; + const int output_tokens = last_only ? 1 : n_tokens; + ggml_tensor * inp = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, n_embd, output_tokens); ggml_set_input(inp); ggml_tensor * normed = build_rms_norm(ctx, inp, w.out_norm, w.rms_eps); ggml_tensor * logits = ggml_mul_mat(ctx, w.output, normed); @@ -3440,14 +6545,19 @@ bool deepseek4_step_layer_range( ggml_free(ctx); return false; } - ggml_backend_tensor_set(inp, final_embd.data(), 0, sizeof(float) * final_embd.size()); + const float * output_input = last_only + ? final_embd.data() + (size_t)(n_tokens - 1) * n_embd + : final_embd.data(); + ggml_backend_tensor_set(inp, output_input, 0, + sizeof(float) * (size_t)n_embd * output_tokens); if (ggml_backend_graph_compute(backend, gf) != GGML_STATUS_SUCCESS) { ggml_free(ctx); return false; } out_logits->resize((size_t)w.n_vocab); - const size_t logits_offset = (size_t)(n_tokens - 1) * (size_t)w.n_vocab * sizeof(float); + const size_t logits_offset = last_only ? 0 : + (size_t)(n_tokens - 1) * (size_t)w.n_vocab * sizeof(float); ggml_backend_tensor_get(logits, out_logits->data(), logits_offset, sizeof(float) * (size_t)w.n_vocab); ggml_free(ctx); @@ -3459,11 +6569,12 @@ bool deepseek4_step_layer_range( memcpy(out_logits->data(), hc_state.data(), sizeof(float) * hc_dim * n_tokens); } - // Update compressor state + // Update compressor state. Multi-token prefill may cross one or more + // boundaries even when the chunk itself does not end on a boundary. const int next_pos = kv_start + n_tokens; for (int il = layer_begin; il < layer_end; ++il) { const uint32_t ratio = w.compress_ratios[il]; - if (ratio <= 0 || (next_pos % (int)ratio) != 0) continue; + if (ratio <= 0) continue; cache.layers[il].n_comp = std::max(cache.layers[il].n_comp, next_pos / (int)ratio); if (ratio == 4) { cache.layers[il].n_index_comp = std::max(cache.layers[il].n_index_comp, @@ -3472,6 +6583,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 +6691,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 +6878,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..e77e61848 100644 --- a/server/src/deepseek4/deepseek4_internal.h +++ b/server/src/deepseek4/deepseek4_internal.h @@ -24,9 +24,15 @@ #include "internal.h" #include "common/layer_split_utils.h" +#include "common/prefill_attention_mode.h" namespace dflash::common { +// Layer-major prefill may schedule two 2K numerical bands while preserving +// the raw-cache rounding boundary between them. +inline constexpr int DS4_NUMERICAL_PREFILL_BAND = 2048; +inline constexpr int DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS = 4096; + struct MoeHybridPlacement; struct MoeHybridConfig; struct MoeHybridRoutingStats; @@ -64,6 +70,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; }; @@ -255,6 +266,7 @@ struct DeepSeek4Cache { int n_layer = 0; std::vector layers; + PrefillAttentionMode prefill_mode = PrefillAttentionMode::Exact; // HC residual streams: [n_hc * n_embd] persistent state ggml_tensor * hc_state = nullptr; // [n_hc * n_embd] @@ -272,6 +284,7 @@ struct DeepSeek4BackendConfig { DevicePlacement device; int stream_fd = -1; int chunk = 512; // prefill chunk size + PrefillAttentionMode prefill_mode = PrefillAttentionMode::Exact; int max_ctx = 0; // 0 = auto from SWA + compression capacity }; @@ -330,7 +343,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) { diff --git a/server/src/server/server_main.cpp b/server/src/server/server_main.cpp index 317fdbed7..06b99622c 100644 --- a/server/src/server/server_main.cpp +++ b/server/src/server/server_main.cpp @@ -217,6 +217,8 @@ static void print_usage(const char * prog) { " --target-layer-split Reserved layer-split weights\n" " --peer-access Enable peer access for multi-GPU placement\n" " --chunk Chunked-prefill chunk size (default: 512)\n" + " --ds4-prefill DeepSeek4 prefill: exact, dense, or sparse " + "(default: exact; sparse is experimental)\n" " --fa-window Flash-attention sliding window (default: 0=full).\n" " WARNING: >0 drops system prompt / tool definitions\n" " from attention at long contexts. Use 0 for tools.\n" @@ -420,6 +422,19 @@ int main(int argc, char ** argv) { bargs.device.peer_access = true; } else if (std::strcmp(argv[i], "--chunk") == 0 && i + 1 < argc) { bargs.chunk = std::atoi(argv[++i]); + } else if (std::strcmp(argv[i], "--ds4-prefill") == 0 && i + 1 < argc) { + const char * mode = argv[++i]; + if (std::strcmp(mode, "exact") == 0) { + bargs.ds4_prefill_mode = PrefillAttentionMode::Exact; + } else if (std::strcmp(mode, "dense") == 0) { + bargs.ds4_prefill_mode = PrefillAttentionMode::Dense; + } else if (std::strcmp(mode, "sparse") == 0) { + bargs.ds4_prefill_mode = PrefillAttentionMode::Sparse; + } else { + std::fprintf(stderr, + "[server] --ds4-prefill expects exact, dense, or sparse\n"); + return 2; + } } else if (std::strcmp(argv[i], "--fa-window") == 0 && i + 1 < argc) { bargs.fa_window = std::atoi(argv[++i]); } else if (std::strcmp(argv[i], "--model-name") == 0 && i + 1 < argc) { @@ -1060,6 +1075,8 @@ int main(int argc, char ** argv) { std::fprintf(stderr, "[server] │ peer_access = %s\n", bargs.device.peer_access ? "ON" : "off"); std::fprintf(stderr, "[server] │ chunk = %d\n", bargs.chunk); + std::fprintf(stderr, "[server] │ ds4_prefill = %s\n", + prefill_attention_mode_name(bargs.ds4_prefill_mode)); std::fprintf(stderr, "[server] │ fa_window = %d\n", bargs.fa_window); if (bargs.fa_window > 0) { std::fprintf(stderr, "[server] │ ⚠ fa_window > 0 drops system prompt / "