Skip to content

perf(ds4): batch exact prefill through width-4 vector kernels #666

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dusterbloom wants to merge 1 commit into
base: main
Choose a base branch
from
+21 −6
Open

perf(ds4): batch exact prefill through width-4 vector kernels #666

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2787,7 +2787,9 @@ static bool ggml_cuda_try_fuse_mul_mat_glu(
const int64_t ncols = ids ? src1->ne[2] : src1->ne[1];
if (ggml_cuda_should_use_mmq(
src0->type, cc, ncols,
ids ? src0->ne[2] : /*n_experts=*/0)) {
ids ? src0->ne[2] : /*n_experts=*/0) &&
!(ggml_cuda_mmvq_max_ncols_override > 0 &&
ncols <= ggml_cuda_mmvq_max_ncols_override)) {

@cubic-dev-ai cubic-dev-ai Bot Aug 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: When the DS4 fused MMQ-pair path is skipped because ncols<=override, the fallback only lands on the quantized matrix-vector path if the grouped MMVQ batch ceiling (get_mmvq_mmid_max_batch) is >= ncols for the qtype. For an ids/grouped DS4 feed-forward, ncols is src1->ne[2], but the override ceiling is designed for plain mul_mat's src1->ne[1]; the two limits are not the same constant. If a qtype's mmid MMVQ ceiling is below ncols, this gate silently drops the fused path and the node falls back to the slower unfused MMQ/dequant path instead of the intended MMVQ win.

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At server/deps/llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu, line 2792:

<comment>When the DS4 fused MMQ-pair path is skipped because ncols<=override, the fallback only lands on the quantized matrix-vector path if the grouped MMVQ batch ceiling (get_mmvq_mmid_max_batch) is >= ncols for the qtype. For an ids/grouped DS4 feed-forward, ncols is src1->ne[2], but the override ceiling is designed for plain mul_mat's src1->ne[1]; the two limits are not the same constant. If a qtype's mmid MMVQ ceiling is below ncols, this gate silently drops the fused path and the node falls back to the slower unfused MMQ/dequant path instead of the intended MMVQ win.</comment>

<file context>
@@ -2787,7 +2787,9 @@ static bool ggml_cuda_try_fuse_mul_mat_glu(
-                ids ? src0->ne[2] : /*n_experts=*/0)) {
+                ids ? src0->ne[2] : /*n_experts=*/0) &&
+            !(ggml_cuda_mmvq_max_ncols_override > 0 &&
+              ncols <= ggml_cuda_mmvq_max_ncols_override)) {
             ggml_cuda_mul_mat_q_pair(
                 ctx, up->src[0], gate->src[0], src1, ids, up, gate);
</file context>

ggml_cuda_mul_mat_q_pair(
ctx, up->src[0], gate->src[0], src1, ids, up, gate);
ggml_cuda_op_swiglu_ds4(ctx, glu);
Expand Down
5 changes: 3 additions & 2 deletions server/src/deepseek4/deepseek4_backend.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1896,8 +1896,9 @@ int DeepSeek4Backend::do_prefill(const std::vector<int32_t> & tokens,
const bool hybrid_batch_supported =
!moe_hybrid_ || cfg_.prefill_mode == PrefillAttentionMode::Sparse;
const int base_chunk =
!prefill_attention_mode_is_approximate(cfg_.prefill_mode) ||
!hybrid_batch_supported
!hybrid_batch_supported ||
(cfg_.prefill_mode == PrefillAttentionMode::Exact &&
spec_drafter_ != nullptr)
? 1
: std::max(1, std::min(requested_chunk,
layer_major_cap));
Expand Down
18 changes: 15 additions & 3 deletions server/src/deepseek4/deepseek4_graph.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6965,7 +6965,17 @@ bool deepseek4_step_layer_range(
// each sub-forward then writes at most one window and, if present, its
// boundary is the final token. This preserves the same pool/rotate order
// as sequential execution while retaining safe batched prefixes.
const int first_chunk = deepseek4_safe_compressor_batch_tokens(w, kv_start, n_tokens);
const bool exact_prefill_band =
cache.prefill_mode == PrefillAttentionMode::Exact &&
allow_decode_graph_reuse && !fused_verify_candidate;
const int first_chunk = std::min(
deepseek4_safe_compressor_batch_tokens(w, kv_start, n_tokens),
exact_prefill_band ? 4 : n_tokens);
const bool exact_multi_token_band =
exact_prefill_band && n_tokens > 1 && n_tokens <= 4;
ScopedCudaGraphOverrides exact_mmvq_scope(
/*disable_graphs=*/false,
/*mmvq_max_ncols=*/exact_multi_token_band ? 4 : 0);
if (first_chunk > 0 && first_chunk < n_tokens &&
Comment on lines +6976 to 6979

@cubic-dev-ai cubic-dev-ai Bot Aug 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2: When LUCE_MMVQ_MAX_NCOLS is set below 4, this scope ignores the operator's explicit dispatch threshold and forces MMVQ for every quantized matmul in the exact batch. Preserve an explicit environment threshold, or scope the selection to the qualified feed-forward operations.

(Based on your team's feedback about respecting MMVQ dispatch overrides.)

View Feedback

Prompt for AI agents 
Check if this issue is valid — if so, understand the root cause and fix it. At server/src/deepseek4/deepseek4_graph.cpp, line 6976:

<comment>When `LUCE_MMVQ_MAX_NCOLS` is set below 4, this scope ignores the operator's explicit dispatch threshold and forces MMVQ for every quantized matmul in the exact batch. Preserve an explicit environment threshold, or scope the selection to the qualified feed-forward operations.

(Based on your team's feedback about respecting MMVQ dispatch overrides.) </comment>

<file context>
@@ -6965,7 +6965,17 @@ bool deepseek4_step_layer_range(
+        exact_prefill_band ? 4 : n_tokens);
+    const bool exact_multi_token_band =
+        exact_prefill_band && n_tokens > 1 && n_tokens <= 4;
+    ScopedCudaGraphOverrides exact_mmvq_scope(
+        /*disable_graphs=*/false,
+        /*mmvq_max_ncols=*/exact_multi_token_band ? 4 : 0);
</file context>
Suggested change
ScopedCudaGraphOverrides exact_mmvq_scope(
/*disable_graphs=*/false,
/*mmvq_max_ncols=*/exact_multi_token_band ? 4 : 0);
if (first_chunk > 0 && first_chunk < n_tokens &&
const int exact_mmvq_max_ncols =
exact_multi_token_band && std::getenv("LUCE_MMVQ_MAX_NCOLS") == nullptr
? 4 : 0;
ScopedCudaGraphOverrides exact_mmvq_scope(
/*disable_graphs=*/false,
/*mmvq_max_ncols=*/exact_mmvq_max_ncols);

!fused_verify_candidate && !heterogeneous_sparse_prefill &&
!standard_layer_major_prefill) {
Expand All @@ -6989,8 +6999,10 @@ bool deepseek4_step_layer_range(
}

for (int off = 0; off < n_tokens;) {
const int chunk = deepseek4_safe_compressor_batch_tokens(
w, kv_start + off, n_tokens - off);
const int remaining = n_tokens - off;
const int chunk = std::min(
deepseek4_safe_compressor_batch_tokens(w, kv_start + off, remaining),
exact_prefill_band ? 4 : remaining);
std::vector<float> chunk_hc;
std::vector<float> chunk_out;
std::vector<float> chunk_capture;
Expand Down
Loading