Enable unified KV for KVarN

Allow KVarN runtime validation and server auto-slot mode to use kv-unified, and gate prompt-cache state restore on actual memory stream safety instead of disabling unified KV globally. Add CUDA KVarN irregular-index coverage and update docs/tests.

Author: Anbeeld

docs/beellama-args.md

@@ -205,7 +205,9 @@ For current long-context Qwen and Gemma DFlash serving, the common asymmetric ch
 
 Do not assume enum compatibility with TheTom's public TurboQuant fork. Bee uses the buun enum order for Turbo/TCQ cache types while keeping a 128-value `turbo3` block. Bee keeps TCQ cache IDs at `45` and `46`; Tom's `TQ3_1S` and `TQ4_1S` model weight formats use new GGML type IDs `47` and `48`.
 
-KVarN pseudo types are accepted only on the target `--cache-type-k` / `--cache-type-v` path. Draft cache arguments intentionally reject them. If only one side is specified, Bee mirrors that KVarN bit width to the other side; all K/V pairs over `2`, `3`, `4`, `5`, `6`, and `8` bits are selectable. KVarN still requires a supported attention/cache path, non-unified KV streams, 128-slice-compatible key/value head dimensions, and backend native KVarN ops unless `--kv-kvarn-fallback` is enabled.
+KVarN pseudo types are accepted only on the target `--cache-type-k` / `--cache-type-v` path. Draft cache arguments intentionally reject them. If only one side is specified, Bee mirrors that KVarN bit width to the other side; all K/V pairs over `2`, `3`, `4`, `5`, `6`, and `8` bits are selectable. KVarN supports both normal per-sequence KV streams and `--kv-unified`, and still requires a supported attention/cache path, 128-slice-compatible key/value head dimensions, and backend native KVarN ops unless `--kv-kvarn-fallback` is enabled.
+
+When KVarN runs with `--kv-unified`, it uses one physical structured stream and standard KV metadata for sequence isolation. Per-sequence prompt-cache state restore is used only when that shared stream is exclusive; otherwise the server keeps correctness by skipping the RAM state restore and processing the prompt normally.
 
 Layers that cannot use structured KVarN records (for example iSWA sliding-window layers, whose rolling eviction is incompatible with KVarN 128-token tile groups) fall back to a normal KV cache whose type matches the requested bit width: `kvarn2` → `q2_0`, `kvarn3` → `q3_0`, `kvarn4` → `q4_0`, `kvarn5` → `q5_0`, `kvarn6` → `q6_0`, `kvarn8` → `q8_0`. This keeps total KV memory at the requested bit level instead of silently storing the fallback layers in `f16`.
 

include/llama.h

@@ -911,6 +911,17 @@ extern "C" {
     // Check if the memory supports shifting
     LLAMA_API bool llama_memory_can_shift(llama_memory_t mem);
 
+    // Returns the actual number of physical KV attention streams used by the memory object.
+    // A value of 1 means all sequences share one physical KV stream.
+    LLAMA_API uint32_t llama_memory_kv_n_stream(llama_memory_t mem);
+
+    // Convenience query for a context. Returns 0 if the context has no KV attention memory.
+    LLAMA_API uint32_t llama_context_kv_n_stream(const struct llama_context * ctx);
+
+    // Returns true when restoring a per-sequence state can overwrite physical KV data
+    // that may also contain cells for other sequences in the same stream.
+    LLAMA_API bool llama_memory_state_seq_restore_requires_exclusive_kv_stream(llama_memory_t mem);
+
     // Expand the recurrent state to new_n_seq_max cells (for deferred backup allocation).
     // Returns true on success. No-op if the memory is already large enough or has no recurrent component.
     LLAMA_API bool llama_memory_recurrent_expand(llama_memory_t mem, uint32_t new_n_seq_max);

src/llama-context.cpp

@@ -8393,7 +8393,7 @@ llama_context * llama_init_from_model(
                 /*.attention_supported =*/ attention_supported,
                 /*.head_dims_supported =*/ head_dims_supported,
                 /*.n_seq_max           =*/ std::max(1u, params.n_seq_max),
-                /*.kv_unified          =*/ false,
+                /*.kv_unified          =*/ params.kv_unified,
             };
 
             if (const char * reason = llama_kvarn_validate_runtime(params.kvarn, requirements)) {
@@ -8406,10 +8406,6 @@ llama_context * llama_init_from_model(
                         __func__, llama_kvarn_type_name(params.kvarn.type), reason);
                 params.kvarn = llama_kvarn_default_params();
             } else {
-                if (params.kv_unified) {
-                    LLAMA_LOG_WARN("%s: KVarN requires non-unified KV streams; forcing kv_unified=false\n", __func__);
-                    params.kv_unified = false;
-                }
                 if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_ENABLED) {
                     LLAMA_LOG_WARN("%s: KVarN requires Flash Attention; enabling it\n", __func__);
                     params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
@@ -9442,6 +9438,18 @@ bool llama_memory_can_shift(llama_memory_t mem) {
     return mem->get_can_shift();
 }
 
+uint32_t llama_memory_kv_n_stream(llama_memory_t mem) {
+    return mem ? mem->get_kv_n_stream() : 0;
+}
+
+uint32_t llama_context_kv_n_stream(const llama_context * ctx) {
+    return ctx ? llama_memory_kv_n_stream(llama_get_memory(ctx)) : 0;
+}
+
+bool llama_memory_state_seq_restore_requires_exclusive_kv_stream(llama_memory_t mem) {
+    return mem && mem->state_seq_restore_requires_exclusive_kv_stream();
+}
+
 static llama_memory_recurrent * get_recurrent_mem(llama_memory_t mem) {
     if (auto * h = dynamic_cast<llama_memory_hybrid *>(mem))      return h->get_mem_recr();
     if (auto * h = dynamic_cast<llama_memory_hybrid_iswa *>(mem)) return h->get_mem_recr();

src/llama-kv-cache-iswa.cpp

@@ -304,6 +304,11 @@ bool llama_kv_cache_iswa::get_can_shift() const {
            kv_base->get_kv_size() == kv_swa->get_kv_size();
 }
 
+bool llama_kv_cache_iswa::state_seq_restore_requires_exclusive_kv_stream() const {
+    return kv_base->state_seq_restore_requires_exclusive_kv_stream() ||
+           kv_swa->state_seq_restore_requires_exclusive_kv_stream();
+}
+
 void llama_kv_cache_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
     const bool include_base = (flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0 ||
                               kv_base->requires_state_for_partial_restore();

src/llama-kv-cache-iswa.h

@@ -73,6 +73,7 @@ class llama_kv_cache_iswa : public llama_memory_i {
     // state write/load
 
     bool requires_state_for_partial_restore() const override;
+    bool state_seq_restore_requires_exclusive_kv_stream() const override;
     void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
     void state_read (llama_io_read_i  & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
 

src/llama-kv-cache-kvarn.cpp

@@ -651,6 +651,10 @@ bool llama_kv_cache_kvarn::requires_state_for_partial_restore() const {
     return true;
 }
 
+bool llama_kv_cache_kvarn::state_seq_restore_requires_exclusive_kv_stream() const {
+    return get_kv_n_stream() == 1;
+}
+
 bool llama_kv_cache_kvarn::has_pending_stream_copies() const {
     return !pending_stream_copies.empty();
 }

src/llama-kv-cache-kvarn.h

@@ -116,6 +116,7 @@ class llama_kv_cache_kvarn : public llama_memory_i {
     std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const override;
 
     bool requires_state_for_partial_restore() const override;
+    bool state_seq_restore_requires_exclusive_kv_stream() const override;
     void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
     void state_read(llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;
 

src/llama-kvarn.cpp

@@ -165,10 +165,6 @@ const char * llama_kvarn_validate_runtime(
     if (!requirements.head_dims_supported) {
         return "KVarN requires key and value head dimensions to be 128-slice-compatible";
     }
-    if (requirements.kv_unified) {
-        return "KVarN requires non-unified KV streams";
-    }
-
     return nullptr;
 }
 

src/llama-memory-hybrid-iswa.cpp

@@ -226,6 +226,11 @@ bool llama_memory_hybrid_iswa::requires_state_for_partial_restore() const {
            mem_recr->requires_state_for_partial_restore();
 }
 
+bool llama_memory_hybrid_iswa::state_seq_restore_requires_exclusive_kv_stream() const {
+    return mem_attn->state_seq_restore_requires_exclusive_kv_stream() ||
+           mem_recr->state_seq_restore_requires_exclusive_kv_stream();
+}
+
 void llama_memory_hybrid_iswa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
     mem_attn->state_write(io, seq_id, flags);
     mem_recr->state_write(io, seq_id, flags);

src/llama-memory-hybrid-iswa.h

@@ -84,6 +84,7 @@ class llama_memory_hybrid_iswa : public llama_memory_i {
     // state write/load
 
     bool requires_state_for_partial_restore() const override;
+    bool state_seq_restore_requires_exclusive_kv_stream() const override;
     void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;
     void state_read (llama_io_read_i  & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0)       override;