Repository navigation
model : deepseek41 quantized K cache, shared prompt tokens, Engram host placement - #293
danielhanchen wants to merge 5 commits into
Conversation
llama-server --embeddings --pooling none toggles llama_set_embeddings per request. With pooling NONE the toggle did not trigger a re-reserve, so a model whose graph depends on cparams.embeddings (DeepSeek-V4.1 returns the flat HC state as t_embd) could run on a stale reservation, and ggml_gallocr segfaulted on the first /completion. Re-reserve on every toggle.
…ngram host placement - Tokens shared by several sequences (llama-batched decodes the prompt once with n_seq_id = n_parallel) read the n-gram history of their first sequence. llama_kv_cache_dsv4::init_batch now checks that all of the token's sequences have the same predecessors (batch first, then the cache via the new llama_kv_cache::seq_pos_tok_le) and refuses the batch otherwise (decode returns 1, caches untouched). - Mixed token/embd ubatches: embedding (image) positions hash as DEAD, get Engram gate 0, and their KV cells record LLAMA_TOKEN_NULL. - The Engram table (per_layer_token_embd) stays in host memory under tensor overrides and split modes when it is not read lazily.
… engram_gate - The compressed-attention path did not apply the Hadamard rotation that build_raw_attention applies to a quantized K cache. It now rotates q, the window K and the output like the window-only path, and rotates the compressed rows in F32 before storing them in the cache type for the concat. - Quantized compressed K / index K views are dequantized directly to F32 (CUDA has no same-type copy for strided quantized views); with LLAMA_ATTN_ROT_DISABLE=1 they go through F32 as well. - engram_gate is an elementwise weight: dequantize a quantized copy in the graph, and keep llama-quantize from quantizing it. F16/BF16/F32 caches build the same graph as before.
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 920f23c13a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| std::map<std::pair<llama_seq_id, llama_pos>, llama_token> in_batch; | ||
| for (int32_t i = 0; i < batch.n_tokens; ++i) { | ||
| for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { | ||
| in_batch[{ batch.seq_id[i][s], batch.pos[i] }] = batch.token[i]; |
There was a problem hiding this comment.
For a mixed token/embedding batch, llama_batch_allocr stores token ID 0 as the placeholder for every embedding row, so this map treats an embedding predecessor as a real token 0 and also collapses rows that share a multimodal position. The graph later identifies those rows through ubatch.type, stores LLAMA_TOKEN_NULL, and resolves mixed histories by ubatch order. Consequently, a shared token following an embedding in one sequence and token 0 in another can pass this check even though their Engram hashes differ, causing the first sequence's hash to be reused silently for the other sequence. Build this validation from the typed ubatches, or otherwise retain each row's token/embedding kind and order.
Useful? React with 👍 / 👎.
| ggml_tensor * engram_keep = nullptr; | ||
| if (!model.engram_layers.empty() && ubatch.token && ubatch.embd) { | ||
| engram_keep = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, n_tokens); |
There was a problem hiding this comment.
When a caller supplies token-only and embedding-only entries in one DeepSeek-V4.1 batch, this new mixed-input path is unreachable: llm_arch_supports_mixed_batch() still explicitly returns false for LLM_ARCH_DEEPSEEK41, that value initializes llama_batch_allocr::allow_mixed, and llama_batch_allocr::init() rejects the request before this graph is built. Thus the mixed token/embedding support added here still returns an input-validation failure; remove DeepSeek-V4.1 from that denylist or otherwise enable mixed allocation for this architecture.
Useful? React with 👍 / 👎.
…ntiation) # Conflicts: # src/models/deepseek41.cpp
Runtime support for DeepSeek-V4.1 beyond the reference path: a quantized K cache, batches that share prompt tokens or mix tokens and embeddings, host placement of the Engram table, and a scheduler re-reserve when embeddings are toggled.
Stacked on the CUDA PR (#292,
dsv41-cuda-exact). The diff is only this change: 9 files, +152/-12.What it adds
-ctk q8_0/q4_0,src/models/deepseek41.cpp,src/llama-kv-cache-dsv4.cpp):build_raw_attentionapplies to a quantized K cache. It now rotates q, the window K and the output like the window-only path, and rotates the compressed rows in F32 before storing them in the cache type.LLAMA_ATTN_ROT_DISABLE=1they go through F32 as well.engram_gate: it is an elementwise weight, so a quantized copy is dequantized in the graph, andllama-quantizeno longer quantizes it.llama-batcheddecodes the prompt once withn_seq_id = n_parallel) read the n-gram history of their first sequence.init_batchchecks that all of the token's sequences have the same predecessors (batch first, then the cache via the newllama_kv_cache::seq_pos_tok_le) and refuses the batch otherwise: decode returns 1 and the caches are untouched.LLAMA_TOKEN_NULL.per_layer_token_embdstays in host memory under tensor overrides and split modes when it is not read lazily.src/llama-context.cpp):llama-server --embeddings --pooling nonetogglesllama_set_embeddingsper request. With pooling NONE the toggle did not re-reserve, so a model whose graph depends oncparams.embeddingscould run on a stale reservation and segfault inggml_gallocron the first/completion. The scheduler now re-reserves on every toggle.Verification
All results are local, on B200 (CUDA
sm_100a) and x86 CPU.q8_0and 2.0e-2 to 2.2e-2 withq4_0; the quantized-cache gate passes.q8_00.02136,q4_00.03261 (q4_0without rotation 0.03580).q8_00.008094,q4_00.015061.q8_0is safe;q4_0costs accuracy. The KV cache at 32k is 190 MiB (F16), 101 MiB (q8_0) and 53 MiB (q4_0).llama-batchedwith shared prompts and parallel sequences, and 43/43 server checks.llama-server --embeddings --pooling nonefollowed by/completionno longer segfaults.Known gaps
-ctk q4_0is supported but measurably less accurate (see above);q8_0is the recommended quantized cache.