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 1d8f22bb8..e83cfaa93 100644 --- a/server/deps/llama.cpp/ggml/include/ggml.h +++ b/server/deps/llama.cpp/ggml/include/ggml.h @@ -603,6 +603,12 @@ extern "C" { GGML_OP_DS4_HC, // Fused DeepSeek4 hyper-connection pre/post/out mixing + GGML_OP_DS4_INDEXER_QAT, // DS4 indexer Hadamard + FP4 activation simulation + + GGML_OP_DS4_INDEXER_SCORE, // Fused DS4 indexer dot/ReLU/head reduction + + GGML_OP_DS4_INDEXER_MASK, // Apply per-token DS4 indexer top-k to an attention mask + GGML_OP_COUNT, }; @@ -1446,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( @@ -2377,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); @@ -2426,9 +2475,12 @@ extern "C" { struct ggml_tensor * experts, struct ggml_tensor * expert_weights); - // Fused DeepSeek4 hyper-connection helpers (decode, n_tokens == 1). - // ggml_ds4_hc_pre: mix[2*n_hc+n_hc^2] + base + hc_state[n_embd*n_hc] -> - // dst[n_embd + 2*n_hc + n_hc^2] = { working, split(pre,post,comb) } + // 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, @@ -2440,7 +2492,8 @@ extern "C" { float post_scale, float comb_scale); - // ggml_ds4_hc_post: residual hc_state + block_out + split -> new hc_state + // ggml_ds4_hc_post: residual hc_state + block_out + split -> new hc_state; + // all non-base tensors may carry the same n_tokens second dimension. GGML_API struct ggml_tensor * ggml_ds4_hc_post( struct ggml_context * ctx, struct ggml_tensor * residual_hc, @@ -2448,7 +2501,8 @@ extern "C" { struct ggml_tensor * split, int n_hc); - // ggml_ds4_hc_out: output-stage merge of hc streams into one embedding + // ggml_ds4_hc_out: output-stage merge of hc streams into one embedding per + // token. mix and hc_state may carry the same n_tokens second dimension. GGML_API struct ggml_tensor * ggml_ds4_hc_out( struct ggml_context * ctx, struct ggml_tensor * mix, @@ -2457,6 +2511,36 @@ extern "C" { int n_hc, float pre_scale); + // Official DS4 ratio-4 indexer transform. Each contiguous 128-wide F32 + // row is Hadamard-rotated and passed through the model's blockwise FP4 + // activation-simulation round trip. The operation is out-of-place so the + // pre-QAT query remains available to the main attention path. + GGML_API struct ggml_tensor * ggml_ds4_indexer_qat( + struct ggml_context * ctx, + struct ggml_tensor * input); + + // Compute the official ratio-4 indexer score matrix without materializing + // [n_comp,n_head,n_tokens] per-head dots. q is F32 + // [128,n_head,n_tokens], head_weights is F32 [n_head,n_tokens], and + // index_comp is F16 [128,n_comp]. head_weights already includes the + // model's 1/sqrt(128*n_head) scale. + GGML_API struct ggml_tensor * ggml_ds4_indexer_score( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * head_weights, + struct ggml_tensor * index_comp, + int kv_start, + int ratio); + + // Preserve the raw rows of base_mask and retain only selected compressed + // rows. selected is I32 [top_k,n_tokens], indexing the compressed span; + // base_mask is F32 [raw_rows+n_comp,n_tokens]. + GGML_API struct ggml_tensor * ggml_ds4_indexer_mask( + struct ggml_context * ctx, + struct ggml_tensor * base_mask, + struct ggml_tensor * selected, + int raw_rows); + // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, diff --git a/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp b/server/deps/llama.cpp/ggml/src/ggml-backend-meta.cpp index 8e4273212..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-cpu/ggml-cpu.c b/server/deps/llama.cpp/ggml/src/ggml-cpu/ggml-cpu.c index 170f1e4a9..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; @@ -1830,6 +1992,18 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { GGML_ABORT("GGML_OP_DS4_HC is only implemented for CUDA"); } + case GGML_OP_DS4_INDEXER_QAT: + { + ggml_compute_forward_ds4_indexer_qat(params, tensor); + } break; + case GGML_OP_DS4_INDEXER_SCORE: + { + ggml_compute_forward_ds4_indexer_score(params, tensor); + } break; + case GGML_OP_DS4_INDEXER_MASK: + { + ggml_compute_forward_ds4_indexer_mask(params, tensor); + } break; case GGML_OP_OUT_PROD: { ggml_compute_forward_out_prod(params, tensor); @@ -2373,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-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 eb7d125a1..ddcb37c4b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/common.cuh @@ -129,6 +129,13 @@ using block_rocmfp6_device = block_rocmfp6; # define GGML_CUDA_USE_CUB #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +// rocPRIM ships hipCUB, a CUB-compatible device-sort API. The ordinary +// single-block bitonic argsort cannot represent DS4's long-context indexer once +// the number of compressed rows exceeds 1024, so use hipCUB for that case. +#if defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +# define GGML_CUDA_USE_HIPCUB +#endif // defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + #ifdef __CUDA_ARCH_LIST__ constexpr bool ggml_cuda_has_arch_impl(int) { return false; diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/concat.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/concat.cu index 51a3efa53..774b6e099 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/concat.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/concat.cu @@ -236,6 +236,9 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { case GGML_TYPE_I8: concat_cuda_typed(ctx, src0, src1, dst, dim); break; + case GGML_TYPE_I32: + concat_cuda_typed(ctx, src0, src1, dst, dim); + break; default: GGML_ABORT("unsupported concat type %s", ggml_type_name(src0->type)); } diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu index 1c6b86972..77ad8005a 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-hc.cu @@ -2,26 +2,26 @@ // Fused DeepSeek4 hyper-connection ops. // -// mode 0 (pre): src0 = mix [mix_dim] (f32, from fn @ rms_norm(hc_state)) +// mode 0 (pre): src0 = mix [mix_dim,n_tokens] (f32, from fn @ rms_norm(hc_state)) // src1 = base [mix_dim] (f32) -// src2 = hc_state [n_embd*n_hc] (f32, raw residual streams) -// dst = [n_embd + mix_dim]: +// src2 = hc_state [n_embd*n_hc,n_tokens] (f32, raw residual streams) +// dst = [n_embd + mix_dim,n_tokens]: // dst[0..n_embd) = working vector (pre-mixed input) // dst[n_embd..n_embd+mix) = split = {pre[n_hc], post[n_hc], comb[n_hc*n_hc]} // Math matches cpu_hc_sinkhorn + finish_hc_pre_from_mix_into in // deepseek4_graph.cpp (sigmoid gates + Sinkhorn-normalized combine). // -// mode 1 (post): src0 = residual hc_state [n_embd*n_hc] -// src1 = block_out [n_embd] -// src2 = split [mix_dim] (view of a mode-0 dst tail) -// dst = new hc_state [n_embd*n_hc]: +// mode 1 (post): src0 = residual hc_state [n_embd*n_hc,n_tokens] +// src1 = block_out [n_embd,n_tokens] +// src2 = split [mix_dim,n_tokens] (view of a mode-0 dst tail) +// dst = new hc_state [n_embd*n_hc,n_tokens]: // dst[h*n_embd+d] = post[h]*block_out[d] // + sum_src comb[h + src*n_hc] * residual[src*n_embd+d] // -// mode 2 (out): src0 = mix [n_hc] +// mode 2 (out): src0 = mix [n_hc,n_tokens] // src1 = base [n_hc] -// src2 = hc_state [n_embd*n_hc] -// dst = [n_embd]: weights[h] = sigmoid(mix[h]*s0+base[h]) + 1e-6; +// src2 = hc_state [n_embd*n_hc,n_tokens] +// dst = [n_embd,n_tokens]: weights[h] = sigmoid(mix[h]*s0+base[h]) + 1e-6; // dst[d] = sum_h weights[h]*hc_state[h*n_embd+d] #define DS4_HC_SINKHORN_EPS 1.0e-6f @@ -183,7 +183,15 @@ static __global__ void ds4_hc_pre_kernel_t( int iters, float pre_scale, float post_scale, - float comb_scale) { + float comb_scale, + size_t mix_stride, + size_t hc_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + mix += (size_t) token * mix_stride; + hc_state += (size_t) token * hc_stride; + dst += (size_t) token * dst_stride; + __shared__ float split[DS4_HC_MAX_MIX]; __shared__ float s_mix[DS4_HC_MAX_MIX]; __shared__ float s_base[DS4_HC_MAX_MIX]; @@ -218,6 +226,78 @@ static __global__ void ds4_hc_pre_kernel_t( } } +// Large prefill batches already expose thousands of token blocks, so they do +// not need every embedding tile to recompute the same serial Sinkhorn. Split +// it into one deterministic solve per token followed by the parallel mixing +// pass. The small-token path keeps the fused kernel above to avoid an extra +// launch during decode and speculative verification. +template +static __global__ void ds4_hc_pre_split_kernel_t( + const float * __restrict__ mix, + const float * __restrict__ base, + float * __restrict__ dst, + int n_embd, + int iters, + float pre_scale, + float post_scale, + float comb_scale, + size_t mix_stride, + size_t dst_stride) { + const int token = (int) blockIdx.x; + mix += (size_t) token * mix_stride; + dst += (size_t) token * dst_stride; + + __shared__ float split[DS4_HC_MAX_MIX]; + __shared__ float s_mix[DS4_HC_MAX_MIX]; + __shared__ float s_base[DS4_HC_MAX_MIX]; + constexpr int mix_dim = 2 * NHC + NHC * NHC; + const int tid = (int) threadIdx.x; + + if (tid < mix_dim) { + s_mix[tid] = mix[tid]; + s_base[tid] = base[tid]; + } + __syncthreads(); + + if (tid == 0) { + ds4_hc_sinkhorn_split_t( + s_mix, s_base, pre_scale, post_scale, comb_scale, iters, split); +#pragma unroll + for (int i = 0; i < mix_dim; ++i) { + dst[n_embd + i] = split[i]; + } + } +} + +template +static __global__ void ds4_hc_pre_mix_kernel_t( + const float * __restrict__ hc_state, + float * __restrict__ dst, + int n_embd, + size_t hc_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + hc_state += (size_t) token * hc_stride; + dst += (size_t) token * dst_stride; + + __shared__ float pre[NHC]; + const int tid = (int) threadIdx.x; + if (tid < NHC) { + pre[tid] = dst[n_embd + tid]; + } + __syncthreads(); + + const int d = (int) blockIdx.x * (int) blockDim.x + tid; + if (d < n_embd) { + float acc = 0.0f; +#pragma unroll + for (int h = 0; h < NHC; ++h) { + acc += pre[h] * hc_state[(size_t) h * n_embd + d]; + } + dst[d] = acc; + } +} + static __global__ void ds4_hc_pre_kernel( const float * __restrict__ mix, const float * __restrict__ base, @@ -228,7 +308,15 @@ static __global__ void ds4_hc_pre_kernel( int iters, float pre_scale, float post_scale, - float comb_scale) { + float comb_scale, + size_t mix_stride, + size_t hc_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + mix += (size_t) token * mix_stride; + hc_state += (size_t) token * hc_stride; + dst += (size_t) token * dst_stride; + __shared__ float split[DS4_HC_MAX_MIX]; __shared__ float s_mix[DS4_HC_MAX_MIX]; __shared__ float s_base[DS4_HC_MAX_MIX]; @@ -271,7 +359,17 @@ static __global__ void ds4_hc_post_kernel( const float * __restrict__ split, float * __restrict__ dst, int n_embd, - int n_hc) { + int n_hc, + size_t residual_stride, + size_t block_out_stride, + size_t split_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + residual += (size_t) token * residual_stride; + block_out += (size_t) token * block_out_stride; + split += (size_t) token * split_stride; + dst += (size_t) token * dst_stride; + const int i = blockIdx.x * blockDim.x + threadIdx.x; const int total = n_embd * n_hc; if (i >= total) { @@ -295,7 +393,15 @@ static __global__ void ds4_hc_out_kernel( float * __restrict__ dst, int n_embd, int n_hc, - float pre_scale) { + float pre_scale, + size_t mix_stride, + size_t hc_stride, + size_t dst_stride) { + const int token = (int) blockIdx.y; + mix += (size_t) token * mix_stride; + hc_state += (size_t) token * hc_stride; + dst += (size_t) token * dst_stride; + const int d = blockIdx.x * blockDim.x + threadIdx.x; if (d >= n_embd) { return; @@ -321,8 +427,14 @@ void ggml_cuda_op_ds4_hc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int mode = ggml_get_op_params_i32(dst, 0); const int n_embd = ggml_get_op_params_i32(dst, 1); const int n_hc = ggml_get_op_params_i32(dst, 2); + const int n_tokens = (int) dst->ne[1]; GGML_ASSERT(n_hc > 0 && n_hc <= DS4_HC_MAX_HC); + GGML_ASSERT(n_tokens > 0); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(src1->nb[0] == sizeof(float)); + GGML_ASSERT(src2->nb[0] == sizeof(float)); + GGML_ASSERT(dst->nb[0] == sizeof(float)); cudaStream_t stream = ctx.stream(); @@ -333,33 +445,54 @@ void ggml_cuda_op_ds4_hc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const float post_scale = ggml_get_op_params_f32(dst, 5); const float comb_scale = ggml_get_op_params_f32(dst, 6); const int pre_blocks = (n_embd + 255) / 256; - if (n_hc == 4) { - ds4_hc_pre_kernel_t<4><<>>( + const dim3 grid(pre_blocks, n_tokens, 1); + if (n_hc == 4 && n_tokens >= 64) { + ds4_hc_pre_split_kernel_t<4><<>>( + (const float *) src0->data, (const float *) src1->data, + (float *) dst->data, + n_embd, iters, pre_scale, post_scale, comb_scale, + src0->nb[1] / sizeof(float), dst->nb[1] / sizeof(float)); + ds4_hc_pre_mix_kernel_t<4><<>>( + (const float *) src2->data, (float *) dst->data, + n_embd, src2->nb[1] / sizeof(float), + dst->nb[1] / sizeof(float)); + } else if (n_hc == 4) { + ds4_hc_pre_kernel_t<4><<>>( (const float *) src0->data, (const float *) src1->data, (const float *) src2->data, (float *) dst->data, - n_embd, iters, pre_scale, post_scale, comb_scale); + n_embd, iters, pre_scale, post_scale, comb_scale, + src0->nb[1] / sizeof(float), src2->nb[1] / sizeof(float), + dst->nb[1] / sizeof(float)); } else { - ds4_hc_pre_kernel<<>>( + ds4_hc_pre_kernel<<>>( (const float *) src0->data, (const float *) src1->data, (const float *) src2->data, (float *) dst->data, - n_embd, n_hc, iters, pre_scale, post_scale, comb_scale); + n_embd, n_hc, iters, pre_scale, post_scale, comb_scale, + src0->nb[1] / sizeof(float), src2->nb[1] / sizeof(float), + dst->nb[1] / sizeof(float)); } } break; case 1: { const int total = n_embd * n_hc; const int blocks = (total + 255) / 256; - ds4_hc_post_kernel<<>>( + const dim3 grid(blocks, n_tokens, 1); + ds4_hc_post_kernel<<>>( (const float *) src0->data, (const float *) src1->data, (const float *) src2->data, (float *) dst->data, - n_embd, n_hc); + n_embd, n_hc, + src0->nb[1] / sizeof(float), src1->nb[1] / sizeof(float), + src2->nb[1] / sizeof(float), dst->nb[1] / sizeof(float)); } break; case 2: { const float pre_scale = ggml_get_op_params_f32(dst, 4); const int blocks = (n_embd + 255) / 256; - ds4_hc_out_kernel<<>>( + const dim3 grid(blocks, n_tokens, 1); + ds4_hc_out_kernel<<>>( (const float *) src0->data, (const float *) src1->data, (const float *) src2->data, (float *) dst->data, - n_embd, n_hc, pre_scale); + n_embd, n_hc, pre_scale, + src0->nb[1] / sizeof(float), src2->nb[1] / sizeof(float), + dst->nb[1] / sizeof(float)); } break; default: GGML_ABORT("ds4_hc: unknown mode"); 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 e991bf4a3..5def8828b 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu @@ -31,34 +31,464 @@ __device__ static float ds4_fa_block_max(float v) { return smem[0]; } -__global__ static void ds4_flash_attn_d512_f32_shared_kv_kernel( +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, - const float * k, - const float * v, + 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) { + 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 scores[]; - const float * qh = q + ((size_t) h * (size_t) n_tokens + (size_t) t) * D; - const float scale = rsqrtf((float) D); + 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) { - const float * kr = k + (size_t) r * D; - float dot = 0.0f; + 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) { - dot += qh[d] * kr[d]; + 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; } - const float s = dot * scale; scores[r] = s; local_max = fmaxf(local_max, s); } @@ -76,38 +506,1032 @@ __global__ static void ds4_flash_attn_d512_f32_shared_kv_kernel( 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 = 0; r < n_kv; ++r) { - acc += scores[r] * v[(size_t) r * D + d]; + 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; } - dst[((size_t) t * (size_t) n_heads + (size_t) h) * D + d] = acc * inv_denom; } } -static bool ggml_cuda_ds4_flash_attn_d512_f32(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - const char * e = getenv("DFLASH_DS4_FLASH_ATTN"); - if (!e || !e[0] || e[0] == (char)48) { +// 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]; - if (!Q || !K || !V || mask || - Q->type != GGML_TYPE_F32 || K->type != GGML_TYPE_F32 || V->type != GGML_TYPE_F32 || + 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) || - K->nb[0] != (int64_t) sizeof(float) || - V->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; } @@ -115,22 +1539,258 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32(ggml_backend_cuda_context & ctx, g 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 size_t shmem = (size_t) n_kv * sizeof(float); - ds4_flash_attn_d512_f32_shared_kv_kernel<<>>( - (float *) dst->data, - (const float *) Q->data, - (const float *) K->data, - (const float *) V->data, - sinks ? (const float *) sinks->data : nullptr, - n_tokens, - n_heads, - n_kv); + 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; } 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 e5463faf9..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 @@ -65,6 +65,7 @@ #include "ggml-cuda/fill.cuh" #include "ggml-cuda/moe-fused.cuh" #include "ggml-cuda/ds4-hc.cuh" +#include "ggml-cuda/ds4-indexer.cuh" #include "ggml.h" #include @@ -2302,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; } @@ -2344,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; } } @@ -2469,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); @@ -2887,6 +2955,15 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_DS4_HC: ggml_cuda_op_ds4_hc(ctx, dst); break; + case GGML_OP_DS4_INDEXER_QAT: + ggml_cuda_op_ds4_indexer_qat(ctx, dst); + break; + case GGML_OP_DS4_INDEXER_SCORE: + ggml_cuda_op_ds4_indexer_score(ctx, dst); + break; + case GGML_OP_DS4_INDEXER_MASK: + ggml_cuda_op_ds4_indexer_mask(ctx, dst); + break; case GGML_OP_GROUP_NORM: ggml_cuda_op_group_norm(ctx, dst); break; @@ -3584,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 })) { @@ -3611,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 })) { @@ -4109,6 +4211,32 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud fused_node_count = 5; break; } + } else if (ggml_cuda_can_fuse( + cgraph, i, + { op, GGML_OP_RESHAPE, op, GGML_OP_RESHAPE, GGML_OP_GLU }, {})) { + ggml_tensor * glu = cgraph->nodes[i + 4]; + ggml_tensor * gate_view = glu->src[0]; + ggml_tensor * up_view = glu->src[1]; + ggml_tensor * gate = nullptr; + ggml_tensor * up = nullptr; + + if (gate_view == cgraph->nodes[i + 1] && + up_view == cgraph->nodes[i + 3]) { + gate = cgraph->nodes[i]; + up = cgraph->nodes[i + 2]; + } else if (gate_view == cgraph->nodes[i + 3] && + up_view == cgraph->nodes[i + 1]) { + gate = cgraph->nodes[i + 2]; + up = cgraph->nodes[i]; + } else { + continue; + } + + if (ggml_cuda_try_fuse_mul_mat_glu(*cuda_ctx, gate, up, glu)) { + fused_mul_mat_vec = true; + fused_node_count = 5; + break; + } } else if (ggml_cuda_can_fuse(cgraph, i, { op, op, GGML_OP_GLU }, {})) { ggml_tensor * glu = cgraph->nodes[i + 2]; ggml_tensor * gate = glu->src[0]; @@ -4119,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]; - ggml_cuda_set_fusion_glu_params(fusion_data, glu); - - ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 3; - break; - } - - if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate->src[0]; - ggml_cuda_set_fusion_glu_params(fusion_data, glu); - - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + if (ggml_cuda_try_fuse_mul_mat_glu(*cuda_ctx, gate, up, glu)) { fused_mul_mat_vec = true; fused_node_count = 3; break; @@ -5021,6 +5129,24 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return true; case GGML_OP_DS4_HC: return true; + case GGML_OP_DS4_INDEXER_QAT: + return op->src[0]->type == GGML_TYPE_F32 && + op->src[0]->ne[0] == 128 && + ggml_is_contiguous(op->src[0]); + case GGML_OP_DS4_INDEXER_SCORE: + return op->src[0]->type == GGML_TYPE_F32 && + op->src[0]->ne[0] == 128 && + op->src[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: { @@ -5222,7 +5348,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_F16 || src0_type == GGML_TYPE_BF16 || - src0_type == GGML_TYPE_I8; + src0_type == GGML_TYPE_I8 || + src0_type == GGML_TYPE_I32; } break; case GGML_OP_CONV_TRANSPOSE_1D: { @@ -5312,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/quantize.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cu index 4300ffc14..992e12d85 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cu @@ -176,7 +176,8 @@ template static __global__ void quantize_mmq_q8_1( const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, - const int64_t ne0, const int ne1, const int ne2) { + const int64_t ne0, const int ne1, const int ne2, + const int64_t group_width, const int64_t group_stride) { constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32; constexpr int vals_per_sum = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32; @@ -204,8 +205,21 @@ static __global__ void quantize_mmq_q8_1( const int64_t ib = ib0 + (i0 / (4*QK8_1))*ne1 + blockIdx.x; // block index in channel const int64_t iqs = i0 % (4*QK8_1); // quant index in block - // Load 4 floats per thread and calculate max. abs. value between them: - const float4 xi = i0 < ne00 ? x4[(i03*s03 + i02*s02 + i01*s01 + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f); + // Load 4 floats per thread and calculate max. abs. value between them. + // A grouped source is physically [group_width, token, group] while this + // kernel emits the same flattened-K Q8 layout as a materialized permute. + int64_t src_i00 = i00; + int64_t src_group_offset = 0; + if (group_width > 0) { + const int64_t group = i00 / group_width; + src_i00 = i00 - group * group_width; + src_group_offset = group * group_stride; + } + const int64_t src_index = + i03*s03 + i02*s02 + i01*s01 + src_group_offset + src_i00; + const float4 xi = i0 < ne00 + ? x4[src_index/4] + : make_float4(0.0f, 0.0f, 0.0f, 0.0f); float amax = fabsf(xi.x); amax = fmaxf(amax, fabsf(xi.y)); amax = fmaxf(amax, fabsf(xi.z)); @@ -300,15 +314,15 @@ void quantize_mmq_q8_1_cuda( switch (mmq_get_q8_1_ds_layout(type_src0)) { case MMQ_Q8_1_DS_LAYOUT_D4: quantize_mmq_q8_1 - <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2); + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, 0); break; case MMQ_Q8_1_DS_LAYOUT_DS4: quantize_mmq_q8_1 - <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2); + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, 0); break; case MMQ_Q8_1_DS_LAYOUT_D2S6: quantize_mmq_q8_1 - <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2); + <<>>(x, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2, 0, 0); break; default: GGML_ABORT("fatal error"); @@ -316,6 +330,45 @@ void quantize_mmq_q8_1_cuda( } } +void quantize_mmq_q8_1_grouped_cuda( + const float * x, void * vy, const ggml_type type_src0, + const int64_t ne00, const int64_t group_width, + const int64_t token_stride, const int64_t group_stride, + const int64_t ne0, const int64_t ne1, cudaStream_t stream) { + GGML_ASSERT(ne00 % 4 == 0); + GGML_ASSERT(group_width > 0 && group_width % 4 == 0); + GGML_ASSERT(ne00 % group_width == 0); + GGML_ASSERT(ne0 % (4*QK8_1) == 0); + + const int64_t block_num_y = + (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / + (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const dim3 num_blocks(ne1, block_num_y, 1); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + switch (mmq_get_q8_1_ds_layout(type_src0)) { + case MMQ_Q8_1_DS_LAYOUT_D4: + quantize_mmq_q8_1 + <<>>( + x, nullptr, vy, ne00, token_stride, 0, 0, + ne0, ne1, 1, group_width, group_stride); + break; + case MMQ_Q8_1_DS_LAYOUT_DS4: + quantize_mmq_q8_1 + <<>>( + x, nullptr, vy, ne00, token_stride, 0, 0, + ne0, ne1, 1, group_width, group_stride); + break; + case MMQ_Q8_1_DS_LAYOUT_D2S6: + quantize_mmq_q8_1 + <<>>( + x, nullptr, vy, ne00, token_stride, 0, 0, + ne0, ne1, 1, group_width, group_stride); + break; + default: + GGML_ABORT("fatal error"); + } +} + void quantize_mmq_mxfp4_cuda(const float * x, const int32_t * ids, void * vy, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cuh index 6a91df635..4768791c4 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/quantize.cuh @@ -26,6 +26,12 @@ void quantize_mmq_q8_1_cuda( ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); +void quantize_mmq_q8_1_grouped_cuda( + const float * x, void * vy, ggml_type type_src0, + int64_t ne00, int64_t group_width, + int64_t token_stride, int64_t group_stride, + int64_t ne0, int64_t ne1, cudaStream_t stream); + void quantize_mmq_mxfp4_cuda(const float * x, const int32_t * ids, void * vy, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py index 17cee853e..22a7d57f4 100755 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/generate_cu_files.py @@ -33,6 +33,7 @@ TYPES_MMQ = [ "GGML_TYPE_Q4_0", "GGML_TYPE_Q4_1", "GGML_TYPE_Q5_0", "GGML_TYPE_Q5_1", "GGML_TYPE_Q8_0", + "GGML_TYPE_Q4_0_ROCMFP4_FAST", "GGML_TYPE_Q2_0_ROCMFP2", "GGML_TYPE_Q3_0_ROCMFPX", "GGML_TYPE_Q2_K", "GGML_TYPE_Q3_K", "GGML_TYPE_Q4_K", "GGML_TYPE_Q5_K", "GGML_TYPE_Q6_K", "GGML_TYPE_IQ2_XXS", "GGML_TYPE_IQ2_XS", "GGML_TYPE_IQ2_S", "GGML_TYPE_IQ3_XXS", "GGML_TYPE_IQ3_S", "GGML_TYPE_IQ1_S", "GGML_TYPE_IQ4_NL", "GGML_TYPE_IQ4_XS", "GGML_TYPE_MXFP4", "GGML_TYPE_NVFP4" diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu new file mode 100644 index 000000000..8221e1d1e --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q2_0_rocmfp2.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q2_0_ROCMFP2); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu new file mode 100644 index 000000000..2380af75c --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q3_0_rocmfpx.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q3_0_ROCMFPX); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu new file mode 100644 index 000000000..94a2bb0f5 --- /dev/null +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/template-instances/mmq-instance-q4_0_rocmfp4_fast.cu @@ -0,0 +1,6 @@ +// This file has been autogenerated by generate_cu_files.py, do not edit manually. + +#define GGML_CUDA_ROCMFPX_MMQ_TILE 1 +#include "../mmq.cuh" + +DECL_MMQ_CASE(GGML_TYPE_Q4_0_ROCMFP4_FAST); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu index 59ce36fb1..adfa25312 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/top-k.cu @@ -35,7 +35,7 @@ static void top_k_cub(ggml_cuda_pool & pool, ncols, k, env)); } -#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE +#elif defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) // CUB_TOP_K_AVAILABLE static int next_power_of_2(int x) { int n = 1; @@ -69,7 +69,7 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { for (int i = 0; i < nrows; i++) { top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream); } -#elif defined(GGML_CUDA_USE_CUB) // CUB_TOP_K_AVAILABLE +#elif defined(GGML_CUDA_USE_CUB) || defined(GGML_CUDA_USE_HIPCUB) // CUB_TOP_K_AVAILABLE // Fall back to argsort + copy const int ncols_pad = next_power_of_2(ncols); const size_t shared_mem = ncols_pad * sizeof(int); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh index e00f3c9ec..5ed92c8cc 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh @@ -441,12 +441,14 @@ static __device__ __forceinline__ int rocmfpx_pack4_fp3_bits12_vec_cuda(const ui } static __device__ __forceinline__ int rocmfpx_pack4_fp3_vec_cuda(const uint8_t * qs, const int base) { - const char4 v = make_char4( - (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 0)*3, 3)), - (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 1)*3, 3)), - (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 2)*3, 3)), - (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 3)*3, 3))); - return *((const int *) &v); + // MMQ requests four-value groups, so the 12 packed bits span at most two + // adjacent bytes. Decode them with the same lookup used by MMVQ. + const int bit_pos = 3 * base; + const int byte_pos = bit_pos >> 3; + const uint32_t packed = (uint32_t) qs[byte_pos] | + ((uint32_t) qs[byte_pos + 1] << 8); + return rocmfpx_pack4_fp3_bits12_vec_cuda( + (packed >> (bit_pos & 7)) & 0x0fffu); } static __device__ __forceinline__ int rocmfpx_pack4_fp2_bits8_vec_cuda(const uint32_t bits8) { diff --git a/server/deps/llama.cpp/ggml/src/ggml.c b/server/deps/llama.cpp/ggml/src/ggml.c index 591ab9bda..7bc3fffcd 100644 --- a/server/deps/llama.cpp/ggml/src/ggml.c +++ b/server/deps/llama.cpp/ggml/src/ggml.c @@ -1128,9 +1128,15 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "MOE_FUSED", "DS4_HC", + + "DS4_INDEXER_QAT", + + "DS4_INDEXER_SCORE", + + "DS4_INDEXER_MASK", }; -static_assert(GGML_OP_COUNT == 100, "GGML_OP_COUNT != 100"); +static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1245,9 +1251,15 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "moe_fused(x)", "ds4_hc(x)", + + "ds4_indexer_qat(x)", + + "ds4_indexer_score(q,w,k)", + + "ds4_indexer_mask(x,topk)", }; -static_assert(GGML_OP_COUNT == 100, "GGML_OP_COUNT != 100"); +static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -3336,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) { @@ -5452,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); @@ -8090,12 +8180,17 @@ struct ggml_tensor * ggml_ds4_hc_pre( GGML_ASSERT(hc_state->type == GGML_TYPE_F32); GGML_ASSERT(n_hc > 0 && n_hc <= 8); const int64_t mix_dim = 2*(int64_t)n_hc + (int64_t)n_hc*n_hc; - GGML_ASSERT(ggml_nelements(mix) == mix_dim); + 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(ggml_nelements(hc_state) % n_hc == 0); - const int64_t n_embd = ggml_nelements(hc_state) / n_hc; + 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 = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd + mix_dim); + 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; @@ -8121,12 +8216,16 @@ struct ggml_tensor * ggml_ds4_hc_post( GGML_ASSERT(split->type == GGML_TYPE_F32); GGML_ASSERT(n_hc > 0 && n_hc <= 8); const int64_t mix_dim = 2*(int64_t)n_hc + (int64_t)n_hc*n_hc; - GGML_ASSERT(ggml_nelements(split) == mix_dim); - GGML_ASSERT(ggml_nelements(residual_hc) % n_hc == 0); - const int64_t n_embd = ggml_nelements(residual_hc) / n_hc; - GGML_ASSERT(ggml_nelements(block_out) == n_embd); - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t) n_embd * n_hc); + 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; @@ -8148,12 +8247,16 @@ struct ggml_tensor * ggml_ds4_hc_out( GGML_ASSERT(base->type == GGML_TYPE_F32); GGML_ASSERT(hc_state->type == GGML_TYPE_F32); GGML_ASSERT(n_hc > 0 && n_hc <= 8); - GGML_ASSERT(ggml_nelements(mix) >= n_hc); + 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(ggml_nelements(hc_state) % n_hc == 0); - const int64_t n_embd = ggml_nelements(hc_state) / 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 = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); + 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; @@ -8164,3 +8267,66 @@ struct ggml_tensor * ggml_ds4_hc_out( ggml_set_op_params_f32(result, 4, pre_scale); return result; } + +struct ggml_tensor * ggml_ds4_indexer_qat( + struct ggml_context * ctx, + struct ggml_tensor * input) { + GGML_ASSERT(input->type == GGML_TYPE_F32); + GGML_ASSERT(input->ne[0] == 128); + GGML_ASSERT(ggml_is_contiguous(input)); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, input); + result->op = GGML_OP_DS4_INDEXER_QAT; + result->src[0] = input; + return result; +} + +struct ggml_tensor * ggml_ds4_indexer_score( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * head_weights, + struct ggml_tensor * index_comp, + int kv_start, + int ratio) { + GGML_ASSERT(q->type == GGML_TYPE_F32 && q->ne[0] == 128); + GGML_ASSERT(head_weights->type == GGML_TYPE_F32); + GGML_ASSERT(index_comp->type == GGML_TYPE_F16 && index_comp->ne[0] == 128); + GGML_ASSERT(ggml_is_contiguous(q)); + GGML_ASSERT(ggml_is_contiguous(head_weights)); + GGML_ASSERT(ggml_is_contiguous(index_comp)); + GGML_ASSERT(head_weights->ne[0] == q->ne[1]); + GGML_ASSERT(head_weights->ne[1] == q->ne[2]); + GGML_ASSERT(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 5a3edbf9b..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,19 +545,24 @@ 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. + // 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. // - // Chunked prefill (n_tokens > 1) is only exact while the whole prompt fits - // inside the raw SWA ring: the chunk graph writes ring slots pos % n_swa - // while attention reads the ring in the same graph, and the learned - // compressor only runs for the chunk's last token. Prompts beyond n_swa - // degrade into incoherence. Token-by-token prefill matches the reference - // semantics at any length; chunked stays available via - // DFLASH_DS4_CHUNKED_PREFILL=1 for short-prompt benchmarking only. - const bool unsafe_chunked = env_flag_enabled("DFLASH_DS4_CHUNKED_PREFILL"); - const int chunk = (moe_hybrid_ || !unsafe_chunked) ? 1 : (cfg_.chunk > 0 ? cfg_.chunk : 512); + // 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. @@ -554,6 +571,11 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, // request on can drift by a token or two. if (kv_offset == 0 && cache_.buf) { ggml_backend_buffer_clear(cache_.buf, 0); + cache_.cur_pos = 0; + for (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"); @@ -590,7 +612,9 @@ int DeepSeek4Backend::do_prefill(const std::vector & tokens, embed.data(), n_tok, pos, 0, w_.n_layer, &logits, tokens.data() + i, - timing ? &step_tel : nullptr); + timing ? &step_tel : nullptr, + cfg_.prefill_mode != + PrefillAttentionMode::Sparse); } if (!ok) { std::fprintf(stderr, "[deepseek4] prefill step failed at pos=%d\n", pos); diff --git a/server/src/deepseek4/deepseek4_graph.cpp b/server/src/deepseek4/deepseek4_graph.cpp index 3f81ba9aa..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 @@ -220,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, @@ -649,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, @@ -676,12 +1004,28 @@ static void build_compressor_step( std::vector & i64_array_inputs, std::vector & i32_array_inputs, ggml_tensor ** comp_cache_source_out = nullptr, - ggml_tensor * flush_rows_inp = nullptr) { + ggml_tensor * flush_rows_inp = nullptr, + ggml_tensor * cur_all = nullptr, + int n_tokens_all = 1, + int kv_start_all = -1, + 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; @@ -695,18 +1039,58 @@ static void build_compressor_step( 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); - } - sc_cur = ggml_add(ctx, sc_cur, ape_col); - - if (state_rows_inp) { + if (!batched_rows) { + if (ape_row_inp) { + ape_col = ggml_get_rows(ctx, ape, ape_row_inp); + ape_col = ggml_reshape_2d(ctx, ape_col, comp_width, 1); + } else { + ape_col = ggml_view_2d( + ctx, ape, comp_width, 1, ape->nb[1], (size_t)pos_mod * ape->nb[1]); + ape_col = ggml_cast(ctx, ape_col, GGML_TYPE_F32); + } + sc_cur = ggml_add(ctx, sc_cur, ape_col); + } + + if (batched_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); @@ -722,7 +1106,12 @@ static void build_compressor_step( ggml_build_forward_expand(gf, ggml_cpy(ctx, sc_cur, sc_slot)); } - if (!flush_rows_inp && ((token_pos + 1) % ratio) != 0) { + if (batched_rows && batched_b < 0) { + // 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. @@ -793,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; @@ -814,6 +1206,52 @@ static void build_compressor_step( *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 @@ -869,7 +1307,12 @@ static void build_indexer_compressor_step( ggml_tensor * comp_pos_inp, std::vector & i64_array_inputs, std::vector & i32_array_inputs, - ggml_tensor * flush_rows_inp = nullptr) { + ggml_tensor ** 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, @@ -893,16 +1336,24 @@ static void build_indexer_compressor_step( comp_pos_inp, i64_array_inputs, i32_array_inputs, - nullptr, - flush_rows_inp); + 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 @@ -917,70 +1368,101 @@ static int ds4_padded_comp_rows(int n_comp, int cap) { return padded < cap ? padded : cap; } -static ggml_tensor * build_indexer_score( +static ggml_tensor * build_indexer_topk( ggml_context * ctx, - ggml_tensor * qr_norm_last, // [n_lora_q, 1] - ggml_tensor * cur_last, // [n_embd, 1] + ggml_tensor * qr_norm, // [n_lora_q, n_tokens] + ggml_tensor * cur, // [n_embd, n_tokens] const DeepSeek4Weights & w, const DeepSeek4Layer & L, - const DeepSeek4LayerCache & lc, - int token_pos, - std::vector & i32_inputs) { - const int n_comp = ds4_comp_rows_used(lc.index_comp_kv, lc.n_index_comp, 4, token_pos); - if (!qr_norm_last || !cur_last || !L.indexer_attn_q_b || !L.indexer_proj || - !lc.index_comp_kv || n_comp <= 0) { + ggml_tensor * index_comp_source, + int n_comp, + int kv_start, + int n_tokens, + ggml_tensor * rope_pos, + std::vector & i32_array_inputs) { + if (!qr_norm || !cur || !L.indexer_attn_q_b || !L.indexer_proj || + !index_comp_source || !rope_pos || n_tokens <= 0 || + n_comp <= w.n_indexer_top_k) { return nullptr; } const int n_indexer_head = w.n_indexer_head; const int head_dim = w.n_indexer_head_dim; + const int top_k = std::min(n_comp, w.n_indexer_top_k); + // A token with <= top_k visible compressed rows needs no ranking: selecting + // [0,top_k) and retaining the ordinary causal mask is exactly equivalent. + // Score only the suffix beginning with the first token that can see row + // top_k. For a zero-prefix ratio-4 2K request this shrinks 2164 score rows + // to just 113. + const int first_scored = std::max( + 0, std::min(n_tokens, 4 * (top_k + 1) - 1 - kv_start)); + const int n_scored = n_tokens - first_scored; + if (n_scored <= 0) return nullptr; + + auto token_slice = [&](ggml_tensor * input, int width) { + if (first_scored == 0) return input; + return ggml_view_2d( + ctx, input, width, n_scored, input->nb[1], + (size_t) first_scored * input->nb[1]); + }; + qr_norm = token_slice(qr_norm, (int) qr_norm->ne[0]); + cur = token_slice(cur, (int) cur->ne[0]); + if (first_scored > 0) { + rope_pos = ggml_view_1d( + ctx, rope_pos, n_scored, + (size_t) first_scored * rope_pos->nb[0]); + } + + // Official ratio-4 indexer graph: q_a-normalized query projection, tail + // RoPE, Hadamard+FP4 QAT, per-head scalar projection, ReLU dot products, + // weighted head reduction and top-512 selection for every query token. + ggml_tensor * index_q = ggml_mul_mat(ctx, L.indexer_attn_q_b, qr_norm); + index_q = ggml_reshape_3d( + ctx, index_q, head_dim, n_indexer_head, n_scored); + + const float rope_scale = w.rope_scale_factor > 0.0f + ? (1.0f / w.rope_scale_factor) : 1.0f; + float rope_attn = 1.0f; + if (rope_scale > 0.0f) { + rope_attn /= 1.0f + 0.1f * logf(1.0f / rope_scale); + } + index_q = build_tail_rope_3d( + ctx, index_q, rope_pos, w.n_rot, head_dim, n_indexer_head, + n_scored, w.compress_rope_freq_base, rope_scale, 1.0f, + rope_attn, w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, + (int) w.rope_orig_ctx); + index_q = ggml_ds4_indexer_qat(ctx, ggml_cont(ctx, index_q)); - // DS4 indexer decode scoring mirrors ds4.c::indexer_allowed_decode_one(): - // 1. Build an indexer query from qr_norm (after q_a + RMSNorm, before q_b). - // 2. Apply full-dim RoPE in indexer head space. - // 3. Project per-head scalar weights from the current hidden state. - // 4. Score every compressed row with ReLU(dot(key_h, query_h)) * weight_h. - // 5. Return the top-k compressed-row indices. - ggml_tensor * index_q = ggml_mul_mat(ctx, L.indexer_attn_q_b, qr_norm_last); - index_q = ggml_reshape_3d(ctx, index_q, head_dim, n_indexer_head, 1); - - // TODO: RoPE on indexer query (same gallocr issue as compressor RoPE) - // Skipping for now — correctness deferred. - index_q = ggml_reshape_2d(ctx, index_q, head_dim, n_indexer_head); - - ggml_tensor * head_weights = ggml_mul_mat(ctx, L.indexer_proj, cur_last); + ggml_tensor * head_weights = ggml_mul_mat(ctx, L.indexer_proj, cur); head_weights = ggml_scale(ctx, head_weights, 1.0f / std::sqrt((float) head_dim * (float) n_indexer_head)); - // index_comp_kv: [n_indexer_head_dim, comp_cap] — each row is 128-dim - // Score each compressed row against all query heads via broadcast - ggml_tensor * comp_view = ggml_view_2d(ctx, lc.index_comp_kv, - head_dim, n_comp, - lc.index_comp_kv->nb[1], 0); - comp_view = ggml_cast(ctx, comp_view, GGML_TYPE_F32); - // comp_view: [head_dim, n_comp] → [head_dim, 1, n_comp] for broadcast - comp_view = ggml_reshape_3d(ctx, comp_view, head_dim, 1, n_comp); - - // index_q: [head_dim, n_indexer_head, 1] → repeat to [head_dim, n_indexer_head, n_comp] - // But ggml_mul needs same shapes, so use matmul approach: - // Reshape q: [head_dim, n_indexer_head] → used directly as A in matmul - // comp: [head_dim, n_comp] - // matmul: A^T @ B = [n_indexer_head, n_comp] dot scores - ggml_tensor * comp_2d = ggml_reshape_2d(ctx, comp_view, head_dim, n_comp); - // mul_mat(index_q, comp_2d): A=[head_dim, n_indexer_head], B=[head_dim, n_comp] - // → result=[n_indexer_head, n_comp] - ggml_tensor * dots = ggml_mul_mat(ctx, index_q, comp_2d); - dots = ggml_relu(ctx, dots); - - // Weight each head's contribution: dots[n_indexer_head, n_comp] * weights[n_indexer_head, 1] - ggml_tensor * weight_rep = ggml_repeat(ctx, head_weights, dots); - ggml_tensor * weighted = ggml_mul(ctx, dots, weight_rep); - // Sum across heads (ne[0]) → [1, n_comp] - ggml_tensor * scores = ggml_sum_rows(ctx, weighted); - scores = ggml_cont(ctx, scores); - scores = ggml_reshape_2d(ctx, scores, n_comp, 1); - - return ggml_top_k(ctx, scores, std::min(n_comp, w.n_indexer_top_k)); + 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( @@ -1016,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; @@ -1069,22 +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. ggml_tensor * raw_kv_source = lc.raw_kv; - if (cached_inputs && cached_inputs->raw_kv_rows) { + ggml_tensor * raw_kv_rows = cached_inputs + ? cached_inputs->raw_kv_rows + : nullptr; + if (raw_kv_rows) { ggml_tensor * kv_f32 = ggml_is_contiguous(kv) ? kv : ggml_cont(ctx, kv); - raw_kv_source = ggml_set_rows(ctx, lc.raw_kv, kv_f32, cached_inputs->raw_kv_rows); + raw_kv_source = ggml_set_rows(ctx, lc.raw_kv, kv_f32, raw_kv_rows); ggml_build_forward_expand(gf, raw_kv_source); } else { - for (int ti = 0; ti < n_tokens; ti++) { + // The attention graph consumes the whole current ubatch directly. + // Persist only its final SWA tail so every physical ring row is written + // once, even when the ubatch is much larger than the 128-row ring. + const int first_write = std::max(0, n_tokens - w.n_swa); + for (int ti = first_write; ti < n_tokens; ti++) { const int pos_ti = kv_start + ti; ggml_tensor * kv_row = ggml_view_2d( ctx, kv, head_dim, 1, kv->nb[1], (size_t)ti * kv->nb[1]); @@ -1099,8 +1649,6 @@ static ggml_tensor * build_mla_attention( // ── Learned compression update ────────────────────────────────── ggml_tensor * cur_last = ggml_view_2d( ctx, cur, n_embd, 1, cur->nb[1], (size_t)(n_tokens - 1) * cur->nb[1]); - ggml_tensor * qr_last = ggml_view_2d( - ctx, qr, n_lora_q, 1, qr->nb[1], (size_t)(n_tokens - 1) * qr->nb[1]); ggml_tensor * comp_kv_source = lc.comp_kv; if (ratio > 0 && L.attn_compressor_kv) { build_compressor_step(ctx, gf, cur_last, @@ -1127,9 +1675,13 @@ static ggml_tensor * build_mla_attention( i64_array_inputs, i32_array_inputs, &comp_kv_source, - cached_inputs ? cached_inputs->flush_rows : nullptr); + cached_inputs ? cached_inputs->flush_rows : nullptr, + (causal_batch || fused_causal) ? cur : nullptr, + n_tokens, + kv_start); } + 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, @@ -1138,8 +1690,13 @@ static ggml_tensor * build_mla_attention( cached_inputs ? cached_inputs->index_comp_pos : nullptr, i64_array_inputs, i32_array_inputs, - cached_inputs ? cached_inputs->flush_rows : nullptr); - (void)build_indexer_score(ctx, qr_last, cur_last, w, L, lc, token_pos, i32_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) ──────────── @@ -1147,14 +1704,26 @@ static ggml_tensor * build_mla_attention( // raw_kv: [head_dim, n_swa] F16 persistent ring buffer (single KV head, shared) // comp_kv: [head_dim, comp_cap] F16 compressed rows. // n_raw = min(kv_start + n_tokens, n_swa) - const bool masked_kv = (n_tokens == 1) && cached_inputs && cached_inputs->attn_row_mask; + const bool masked_kv = cached_inputs && cached_inputs->attn_row_mask; const int n_comp_live = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; + ggml_tensor * indexer_topk = nullptr; + if (attention_impl == DeepSeek4AttentionImpl::SparseFlash && + ratio == 4 && f32_array_inputs) { + const int n_index_comp = ds4_comp_rows_used( + lc.index_comp_kv, lc.n_index_comp, 4, token_pos); + indexer_topk = build_indexer_topk( + ctx, qr, cur, w, L, index_comp_kv_source, + n_index_comp, kv_start, n_tokens, rope_pos, + i32_array_inputs); + } // Stable path reads the full physical ring (masking not-yet-written slots) // and a padded compressed-row span; the plain path reads only valid rows. - const int n_raw = masked_kv ? w.n_swa : std::min(kv_start + n_tokens, w.n_swa); + const int n_raw = masked_kv ? w.n_swa + : 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; + const int n_attn = n_raw + n_comp_attn + n_old_rows; const float kq_scale = 1.0f / sqrtf((float)head_dim); // Get valid KV rows. For single-token decode, include the current in-graph @@ -1177,9 +1746,26 @@ static ggml_tensor * build_mla_attention( ggml_tensor * ring = ggml_view_2d( ctx, raw_kv_source, head_dim, w.n_swa, raw_kv_source->nb[1], 0); kv_attn = ds4_cast_if_needed(ctx, ring, GGML_TYPE_F32); + } else if (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 > 1) { + if (n_raw == w.n_swa && raw_kv_rows) { + // Once the ring is full, use a stable physical row order. The + // cached q=1 graph is first built at position n_swa-1 and then + // reused across every wrap position, so chronological views baked + // into that first topology become stale. Insert the current F32 + // KV at its runtime row in an F32 snapshot instead. The tokenwise + // prefill helper takes the same branch and row ordering. + ggml_tensor * ring = ggml_view_2d( + ctx, lc.raw_kv, head_dim, w.n_swa, lc.raw_kv->nb[1], 0); + ring = ds4_cast_if_needed(ctx, ring, GGML_TYPE_F32); + kv_attn = ggml_set_rows(ctx, ring, cur_kv, raw_kv_rows); + ggml_build_forward_expand(gf, kv_attn); + } else if (n_raw > 1) { ggml_tensor * prev = nullptr; if (kv_start < w.n_swa) { prev = raw_kv_view(0, n_raw - 1); @@ -1207,60 +1793,300 @@ static ggml_tensor * build_mla_attention( comp = ds4_cast_if_needed(ctx, comp, GGML_TYPE_F32); kv_attn = ggml_concat(ctx, kv_attn, comp, 1); } + if (old_rows_scratch) { + kv_attn = ggml_concat(ctx, kv_attn, old_rows_scratch, 1); + } // kv_attn: [head_dim, n_attn] - // Flatten q to [head_dim, n_head*n_tokens] for batched matmul - ggml_tensor * q_flat = ggml_reshape_2d(ctx, q, head_dim, n_head * n_tokens); - - // Scores: mul_mat(kv_attn, q_flat) = kv_attn^T[n_attn, head_dim] @ q_flat[head_dim, n_head*n_tokens] - // → [n_attn, n_head*n_tokens] - ggml_tensor * scores = ggml_mul_mat(ctx, kv_attn, q_flat); - scores = ggml_scale(ctx, scores, kq_scale); - if (masked_kv) { - // Broadcast-add the [n_attn,1] additive mask across all query columns. - scores = ggml_add(ctx, scores, cached_inputs->attn_row_mask); + // 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; - // 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); + 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 { - probs = ggml_soft_max(ctx, scores); - } - // probs: [n_attn, n_head*n_tokens] - - // Context: kv_T^T[head_dim, n_attn] @ probs[n_attn, n_head*n_tokens] → [head_dim, n_head*n_tokens] - // i.e. mul_mat(kv_T, probs) where kv_T = cont(transpose(kv_attn)) = [n_raw, head_dim] - ggml_tensor * kv_T = ggml_cont(ctx, ggml_transpose(ctx, kv_attn)); - ggml_tensor * context = ggml_mul_mat(ctx, kv_T, probs); - // context: [head_dim, n_head*n_tokens] + // Flatten q to [head_dim, n_head*n_tokens] for batched matmul. + ggml_tensor * q_flat = ggml_reshape_2d(ctx, q, head_dim, + n_head * n_tokens); + ggml_tensor * scores = ggml_mul_mat(ctx, kv_attn, q_flat); + scores = ggml_scale(ctx, scores, kq_scale); + if (score_mask) { + if (n_tokens > 1) { + ggml_tensor * m3 = ggml_reshape_3d(ctx, score_mask, + n_attn, 1, n_tokens); + ggml_tensor * s3 = ggml_reshape_3d(ctx, scores, + n_attn, n_head, n_tokens); + scores = ggml_reshape_2d(ctx, ggml_add(ctx, s3, m3), + n_attn, n_head * n_tokens); + } else { + scores = ggml_add(ctx, scores, score_mask); + } + } - // Reshape back to [head_dim, n_head, n_tokens] - context = ggml_reshape_3d(ctx, context, head_dim, n_head, n_tokens); + // DS4 adds one learned per-head sink logit to the denominator, but the + // sink contributes no value vector. + ggml_tensor * probs = nullptr; + if (L.attn_sinks) { + ggml_tensor * sink_scores = ggml_reshape_2d(ctx, L.attn_sinks, + 1, n_head); + if (n_tokens > 1) { + ggml_tensor * sink_shape = ggml_new_tensor_2d( + ctx, GGML_TYPE_F32, 1, n_head * n_tokens); + sink_scores = ggml_repeat(ctx, sink_scores, sink_shape); + } + ggml_tensor * scores_with_sink = ggml_concat(ctx, scores, + sink_scores, 0); + ggml_tensor * probs_with_sink = ggml_soft_max(ctx, + scores_with_sink); + probs = ggml_view_2d(ctx, probs_with_sink, n_attn, + n_head * n_tokens, probs_with_sink->nb[1], 0); + } else { + probs = ggml_soft_max(ctx, scores); + } + ggml_tensor * kv_t = ggml_cont(ctx, ggml_transpose(ctx, kv_attn)); + context = ggml_mul_mat(ctx, kv_t, probs); + context = ggml_reshape_3d(ctx, context, head_dim, n_head, n_tokens); + } // ── Inverse tail RoPE on attention output ─────────────────────── - ggml_tensor * neg_pos = cached_inputs ? cached_inputs->neg_pos : nullptr; - if (!neg_pos) { - neg_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); - ggml_set_input(neg_pos); - std::vector neg_vals(n_tokens); - for (int i = 0; i < n_tokens; i++) neg_vals[i] = -(kv_start + i); - i32_array_inputs.push_back({neg_pos, std::move(neg_vals)}); - } - context = build_tail_rope_3d(ctx, context, neg_pos, n_rot, head_dim, n_head, n_tokens, - rope_freq, rope_scale, rope_ext, rope_attn, - w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, rope_n_ctx_orig); + if (!inverse_rope_fused) { + ggml_tensor * neg_pos = cached_inputs ? cached_inputs->neg_pos : nullptr; + if (!neg_pos) { + neg_pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_tokens); + ggml_set_input(neg_pos); + std::vector neg_vals(n_tokens); + for (int i = 0; i < n_tokens; i++) neg_vals[i] = -(kv_start + i); + i32_array_inputs.push_back({neg_pos, std::move(neg_vals)}); + } + context = build_tail_rope_3d( + ctx, context, neg_pos, n_rot, head_dim, n_head, n_tokens, + rope_freq, rope_scale, rope_ext, rope_attn, + w.rope_yarn_beta_fast, w.rope_yarn_beta_slow, + rope_n_ctx_orig); + } // Flatten to [head_dim*n_head, n_tokens] for output projection ggml_tensor * attn_out = ggml_reshape_2d(ctx, context, head_dim * n_head, n_tokens); @@ -1275,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; } @@ -2438,6 +3276,21 @@ static void cpu_matvec_f16_pooled(float * out, const uint16_t * mat, const float 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]; @@ -2672,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) { @@ -2735,7 +3588,7 @@ static void hc_post_batch(std::vector & out_hc, }); return; } - ds4_parallel_for_tokens(n_tokens, 8, [&](int t0, int t1) { + ds4_pool_for_tokens(n_tokens, [&](int t0, int t1) { for (int t = t0; t < t1; ++t) { cpu_hc_post(out_hc.data() + (size_t)t * hc_dim, block_out + (size_t)t * n_embd, @@ -2757,7 +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); @@ -3742,7 +4595,8 @@ static ggml_tensor * ds4_build_hash_routed_ffn( const DeepSeek4Weights & w, const DeepSeek4Layer & L, ggml_tensor * ffn_normed, - ggml_tensor * hash_ids) { + ggml_tensor * hash_ids, + int n_tokens) { ggml_tensor * shared_out = build_shared_ffn(ctx, ffn_normed, w, L); ggml_tensor * logits = ggml_mul_mat(ctx, L.ffn_gate_inp, ffn_normed); ggml_tensor * probs = ggml_sqrt(ctx, ggml_softplus(ctx, logits)); @@ -3750,22 +4604,24 @@ static ggml_tensor * ds4_build_hash_routed_ffn( const int n_used = w.n_expert_used; const int n_ff_exp = w.n_ff_exp; const bool raw_mmid = ds4_ffn_raw_mmid_enabled(); - ggml_tensor * cur_3d = ggml_reshape_3d(ctx, ffn_normed, w.n_embd, 1, 1); + ggml_tensor * 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, 1); - up_e = ggml_reshape_3d(ctx, up_e, n_ff_exp, n_used, 1); + gate_e = ggml_reshape_3d(ctx, gate_e, n_ff_exp, n_used, n_tokens); + up_e = ggml_reshape_3d(ctx, up_e, n_ff_exp, n_used, n_tokens); } ggml_tensor * mid_e = build_clamped_swiglu(ctx, gate_e, up_e, w.swiglu_clamp_exp); ggml_tensor * down_e = ggml_mul_mat_id(ctx, L.ffn_down_exps, mid_e, hash_ids); if (!raw_mmid) { - down_e = ggml_reshape_3d(ctx, down_e, w.n_embd, n_used, 1); + down_e = ggml_reshape_3d(ctx, down_e, w.n_embd, n_used, n_tokens); } - ggml_tensor * probs_3d = ggml_reshape_3d(ctx, probs, 1, w.n_expert, 1); + ggml_tensor * probs_3d = ggml_reshape_3d( + ctx, probs, 1, w.n_expert, n_tokens); ggml_tensor * weights = ggml_get_rows(ctx, probs_3d, hash_ids); - weights = ggml_reshape_2d(ctx, weights, n_used, 1); + weights = ggml_reshape_2d(ctx, weights, n_used, n_tokens); ggml_tensor * w_sum = ggml_sum_rows(ctx, weights); w_sum = ggml_clamp(ctx, w_sum, 6.103515625e-5f, INFINITY); weights = ggml_div(ctx, weights, w_sum); @@ -3777,11 +4633,14 @@ static ggml_tensor * ds4_build_hash_routed_ffn( if (ds4_ffn_fused_combine_enabled()) { routed_out = ggml_laguna_moe_combine(ctx, down_e, weights); } else { - ggml_tensor * weights_3d = ggml_reshape_3d(ctx, weights, 1, n_used, 1); + 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, 1); + 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, 1); + routed_out = ggml_reshape_2d( + ctx, routed_out, w.n_embd, n_tokens); } return ggml_add(ctx, shared_out, routed_out); } @@ -3853,7 +4712,6 @@ static bool ds4_build_fused_decode_graph( const int n_embd = w.n_embd; const int n_hc = w.n_hc; - const int n_raw = std::min(kv_start + 1, w.n_swa); const int token_pos = kv_start; const size_t arena_size = 192u * 1024 * 1024; if (fg.sg.meta_arena.size() < arena_size) { @@ -3991,7 +4849,8 @@ static bool ds4_build_fused_decode_graph( 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); + 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); } @@ -4075,7 +4934,6 @@ static int ds4_try_fused_decode_step( } const int token_pos = kv_start; - const int n_raw = std::min(kv_start + 1, w.n_swa); std::vector key; key.reserve((size_t) w.n_layer + 2); key.push_back(w.n_swa); @@ -4226,6 +5084,731 @@ static int ds4_try_fused_decode_step( 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, @@ -4341,6 +5924,25 @@ bool deepseek4_step_layer_range( 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()) { @@ -4509,7 +6111,20 @@ bool deepseek4_step_layer_range( ggml_context * ctx = nullptr; DeepSeek4CachedDecodeAttnGraph * cached_attn = nullptr; - if (reuse_decode_attn) { + const bool exact_tokenwise_prefill = + !reuse_decode_attn && n_tokens > 1; + if (exact_tokenwise_prefill) { + const DeepSeek4AttentionImpl attention_impl = + cache.prefill_mode == PrefillAttentionMode::Sparse + ? DeepSeek4AttentionImpl::SparseFlash + : DeepSeek4AttentionImpl::Explicit; + if (!ds4_run_exact_tokenwise_prefill_attention( + backend, w, L, lc, il, cur.data(), n_tokens, kv_start, + attention_impl, attn_out_host, + cached_attn_allocs[(size_t) il], telemetry)) { + return false; + } + } else if (reuse_decode_attn) { const int n_raw = std::min(kv_start + 1, w.n_swa); const int n_comp_attn = (ratio > 0) ? ds4_comp_rows_used(lc.comp_kv, lc.n_comp, ratio, token_pos) : 0; const int n_index_comp = (ratio == 4) ? ds4_comp_rows_used(lc.index_comp_kv, lc.n_index_comp, 4, token_pos) : 0; @@ -4605,14 +6220,21 @@ bool deepseek4_step_layer_range( std::vector i32_inputs; std::vector i32_array_inputs; std::vector i64_array_inputs; + std::vector f32_array_inputs; const size_t graph_size = ds4_attn_step_graph_size(n_tokens); gf = ggml_new_graph_custom(ctx, graph_size, false); 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); @@ -4640,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); @@ -4670,6 +6295,7 @@ bool deepseek4_step_layer_range( if (telemetry) telemetry->attn_read_us += ds4_elapsed_us(attn_read_t0, Ds4TimingClock::now()); } if (ctx) ggml_free(ctx); + } // ── HC post (attention) ───────────────────────────────── if (!(use_backend_decode_hc_graph || use_backend_decode_hc_direct)) { @@ -4895,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); @@ -4916,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); @@ -4935,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, diff --git a/server/src/deepseek4/deepseek4_internal.h b/server/src/deepseek4/deepseek4_internal.h index a5d28d23c..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; @@ -260,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] @@ -277,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 }; 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 / "