PR6: NVFP4-tier KV + entropy cold pool (10L default table, +33% context ceiling)

apps/cli/main.cpp

@@ -1,4 +1,5 @@
 #include "options.h"
+#include "product/kv_options.h"
 #include "product/load_progress/load_progress.h"
 #include "product/prompt_input/prompt_input.h"
 
@@ -282,7 +283,19 @@ int main(int argc, char** argv) {
         engine_options.kv_cache       = cli.kv_cache;
         engine_options.speculative    = cli.speculative;
         engine_options.enable_vision  = cli.enable_vision;
+        engine_options.yarn_enabled  = cli.yarn_enabled;
         engine_options.use_cuda_graph = cli.use_cuda_graph;
+        engine_options.cold_policy           = cli.cold_policy;
+        engine_options.cold_keep_tokens      = cli.cold_keep_tokens;
+        engine_options.cold_host_bytes       = cli.cold_host_bytes;
+        if (cli.kv_layer_storage_explicit) {
+            const auto table = ninfer::product::parse_kv_layer_storage(cli.kv_layer_storage_spec);
+            for (std::size_t i = 0; i < table.size(); ++i) {
+                engine_options.kv_layer_storage[i] = table[i];
+            }
+            engine_options.kv_layer_storage_explicit = true;
+        }
+        engine_options.graph_capture_ceiling = cli.graph_capture_ceiling;
         // One CLI invocation owns exactly one request, so retained cross-request context has no
         // consumer and must not reserve an extra Device StateImage or run terminal capture.
         engine_options.context_cache.enabled                = false;

apps/cli/options.cpp

@@ -1,5 +1,6 @@
 #include "options.h"
 #include "product/speculative_options.h"
+#include "product/kv_options.h"
 
 #include <cerrno>
 #include <cmath>
@@ -78,14 +79,16 @@ std::string usage_text(const char* argv0) {
            " <model.ninfer> (--prompt <text>|--messages <messages.json>)\n"
            "       [--max-context N] [--kv-capacity N|auto] [--prefill-chunk N] [--max-new N]\n"
            "       [--device N]\n"
-           "       [--kv-dtype bf16|int8|fp8] [--spec mtp|dflash --draft-tokens N]\n"
+           "       [--kv-dtype bf16|int8|fp8] [--kv-layer-storage SPEC] [--spec mtp|dflash --draft-tokens N]\n"
            "       [--lm-head-draft]\n"
            "       [--temperature F] [--top-p F] [--top-k N] [--min-p F]\n"
            "       [--presence-penalty F] [--frequency-penalty F] [--seed N] [--greedy]\n"
            "       [--stop-token-id N]... [--stop <text>]... [--reasoning-stop <text>]...\n"
            "       [--raw-output] [--print-token-ids] [--no-thinking] [--thinking-budget N]\n"
            "       [--reasoning-effort low|medium|xhigh] [--vision]\n"
-           "       [--no-cuda-graph]\n"
+           "       [--cold-policy none|off|window|host] [--cold-keep-tokens N]\n"
+           "       [--cold-host-bytes N[g|m|k]]\n"
+           "       [--no-cuda-graph] [--graph-capture-ceiling N]\n"
            "\n"
            "Streams answer content to stdout and reasoning plus diagnostics to stderr.\n"
            "Structured message content accepts text, image/image_url, and video/video_url parts;\n"
@@ -134,7 +137,10 @@ Options parse_options(int argc, char** argv) {
             options.device = parse_device(value(arg));
         } else if (arg == "--kv-dtype") {
             options.kv_cache = parse_kv_cache(value(arg));
-        } else if (arg == "--spec") {
+                } else if (arg == "--kv-layer-storage") {
+            options.kv_layer_storage_spec = value(arg);
+            options.kv_layer_storage_explicit = true;
+} else if (arg == "--spec") {
             options.speculative.backend = product::parse_speculative_backend(value(arg));
         } else if (arg == "--draft-tokens") {
             options.speculative.draft_tokens = parse_u32(value(arg), "draft-tokens");
@@ -152,8 +158,23 @@ Options parse_options(int argc, char** argv) {
             options.reasoning_effort = parse_reasoning_effort(value(arg));
         } else if (arg == "--vision") {
             options.enable_vision = true;
+        } else if (arg == "--cold-policy") {
+            const std::string v = value(arg);
+            if (v == "none" || v == "off") { options.cold_policy = ColdPolicy::None; }
+            else if (v == "window") { options.cold_policy = ColdPolicy::Window; }
+            else if (v == "host") { options.cold_policy = ColdPolicy::Host; }
+            else { throw std::invalid_argument("invalid cold-policy: " + v); }
+            options.cold_keep_tokens = 128;
+        } else if (arg == "--cold-keep-tokens") {
+            options.cold_keep_tokens = parse_u32(value(arg), "cold-keep-tokens");
+        } else if (arg == "--cold-host-bytes") {
+            options.cold_host_bytes = parse_u32(value(arg), "cold-host-bytes");
+        } else if (arg == "--graph-capture-ceiling") {
+            options.graph_capture_ceiling = parse_u32(value(arg), "graph-capture-ceiling");
         } else if (arg == "--no-cuda-graph") {
             options.use_cuda_graph = false;
+        } else if (arg == "--yarn") {
+            options.yarn_enabled = true;
         } else if (arg == "--stop-token-id") {
             const std::uint32_t token = parse_u32(value(arg), "stop-token-id", true);
             if (token > static_cast<std::uint32_t>(std::numeric_limits<TokenId>::max())) {
@@ -207,8 +228,8 @@ Options parse_options(int argc, char** argv) {
         throw std::invalid_argument("--prefill-chunk must be a multiple of 128");
     }
     if (options.kv_capacity.mode == KvCapacityMode::Explicit &&
-        options.kv_capacity.explicit_tokens < options.max_context) {
-        throw std::invalid_argument("--kv-capacity must be at least --max-context");
+        options.kv_capacity.explicit_tokens == 0) {
+        throw std::invalid_argument("--kv-capacity must be positive");
     }
     product::validate_speculative_cli_options(options.speculative);
     if (options.speculative.backend == SpeculativeBackend::DFlash && options.enable_vision) {

apps/cli/options.h

@@ -27,6 +27,13 @@ struct Options {
     SpeculativeOptions speculative;
     bool enable_vision  = false;
     bool use_cuda_graph = true;
+    ColdPolicy cold_policy         = ColdPolicy::None;
+    std::uint32_t cold_keep_tokens = 128;
+    std::uint64_t cold_host_bytes  = 4ULL << 30;
+    bool yarn_enabled              = false;
+    std::string kv_layer_storage_spec;
+    bool kv_layer_storage_explicit = false;
+    std::uint32_t graph_capture_ceiling = 0;
 
     bool raw_output      = false;
     bool print_token_ids = false;

apps/perplexity/main.cpp

@@ -2,6 +2,7 @@
 #include "evaluation.h"
 
 #include "ninfer/engine.h"
+#include "product/kv_options.h"
 
 #include <nlohmann/json.hpp>
 
@@ -44,6 +45,8 @@ struct Options {
     std::uint32_t stride      = 2048;
     int device                = 0;
     ninfer::KvCacheStorage kv = ninfer::KvCacheStorage::Fp8E4M3Row256;
+    std::array<ninfer::KvCacheStorage, ninfer::kKvLayerStorageSlots> kv_layer_storage{};
+    bool kv_layer_storage_explicit = false;
     bool quick                = false;
 };
 
@@ -52,7 +55,7 @@ struct Options {
                                 "\nusage: ninfer-perplexity <model.ninfer> "
                                 "(--corpus <manifest.json> [--quick] | --text <utf8-file>) "
                                 "[--context N] [--stride N] [--device N] "
-                                "[--kv-dtype bf16|int8|fp8] [--output <directory>]");
+                                "[--kv-dtype bf16|int8|fp8] [--kv-layer-storage SPEC] [--output <directory>]");
 }
 
 template <class Integer>
@@ -104,9 +107,15 @@ Options parse_options(int argc, char** argv) {
                 out.kv = ninfer::KvCacheStorage::Int8Group64;
             } else if (dtype == "fp8") {
                 out.kv = ninfer::KvCacheStorage::Fp8E4M3Row256;
+            } else if (dtype == "nvfp4") {
+                out.kv = ninfer::KvCacheStorage::Nvfp4Group16;
             } else {
-                usage_error("--kv-dtype must be bf16, int8, or fp8");
+                usage_error("--kv-dtype must be bf16, int8, fp8, or nvfp4");
             }
+        } else if (option == "--kv-layer-storage") {
+            const auto table = ninfer::product::parse_kv_layer_storage(value("--kv-layer-storage"));
+            out.kv_layer_storage = table;
+            out.kv_layer_storage_explicit = true;
         } else if (option == "--output") {
             out.output = std::filesystem::path(value("--output"));
         } else {
@@ -201,6 +210,8 @@ int run(const Options& options) {
     engine_options.device                 = options.device;
     engine_options.max_context            = options.context;
     engine_options.kv_cache               = options.kv;
+    engine_options.kv_layer_storage      = options.kv_layer_storage;
+    engine_options.kv_layer_storage_explicit = options.kv_layer_storage_explicit;
     engine_options.load_progress.callback = [&](std::string_view phase, std::uint64_t done,
                                                 std::uint64_t total) {
         const std::uint64_t bucket =

include/ninfer/ops/bidirectional_gqa_attention.h

@@ -0,0 +1,68 @@
+#pragma once
+
+#include "core/arena.h"
+#include "core/paged_kv_cache.h"
+#include "core/tensor.h"
+
+#include <cuda_runtime.h>
+
+#include <cstddef>
+#include <cstdint>
+
+namespace ninfer::ops {
+
+/**
+ * Host execution-resource promise for bidirectional_gqa_attention.
+ *
+ * Device context_lengths define the exact per-row mathematical contexts. This envelope bounds
+ * every row so a fixed launch can be captured and replayed without a host read.
+ */
+struct GqaContextExecutionEnvelope {
+    std::uint32_t min_context = 0;
+    std::uint32_t max_context = 0;
+};
+
+/**
+ * Op: bidirectional grouped-query attention over persistent context and one query block
+ *
+ * For D=128, Hq=32, Hkv=8, group=4, query row i, query head h, and kvh=floor(h/4):
+ *
+ *   keys   = context K rows [0,L) followed logically by every live query K row [0,V)
+ *   score  = scale * dot(q[:,h,i], key[:,kvh,j])
+ *   prob   = softmax over the complete logical key set
+ *   ideal[:,h,i] = sum_j prob[j] * value[:,kvh,j]
+ *
+ * q/out are contiguous BF16 [128,32,W,B]. query_k/query_v are contiguous BF16 [128,8,W,B].
+ * context_lengths, valid_columns, and table_rows are contiguous device I32 [B]. Row b has
+ * V=valid_columns[b] live query columns and reads logical context [0,context_lengths[b]) through
+ * table row table_rows[b]. Columns i>=V are an inert physical tail and produce zero output.
+ * context is a read-only paged BF16 cache with head-major page planes [128,64,Nphysical,8]. scale
+ * is 1/sqrt(128).
+ *
+ * There is no causal triangle: every live query row attends every other live query K/V row in the
+ * same batch row. Context and query K/V remain separate physical segments and every input/cache
+ * byte is unchanged. The oracle evaluates `ideal` naively in FP64 from represented inputs. The
+ * BF16 out is promoted and compared directly with that result; output storage rounding belongs to
+ * the Op's numerical criterion, not the oracle. out is the only observable mutation and is
+ * completely overwritten. The current optimized implementation domain is W=1..16 on sm_120a.
+ *
+ * The caller guarantees min_context <= L <= max_context and that every logical page intersecting
+ * [0,L) is materialized. The execution envelope may affect finite launch selection and workspace
+ * capacity, never the admitted key set or numerical result.
+ */
+void bidirectional_gqa_attention(const Tensor& q, const Tensor& query_k, const Tensor& query_v,
+                                 const Tensor& context_lengths, const Tensor& valid_columns,
+                                 const Tensor& table_rows, float scale,
+                                 const PagedKVBatchLayerView& context,
+                                 GqaContextExecutionEnvelope envelope, WorkspaceArena& workspace,
+                                 Tensor& out, cudaStream_t stream);
+
+/**
+ * Returns the transient arena capacity required for every T in the inclusive optimized interval.
+ * The execution envelope is the fixed profile; invalid profiles or intervals throw.
+ */
+[[nodiscard]] std::size_t bidirectional_gqa_attention_workspace_capacity_bytes(
+    GqaContextExecutionEnvelope envelope, std::int32_t min_tokens, std::int32_t max_tokens,
+    std::int32_t batch_size);
+
+} // namespace ninfer::ops

include/ninfer/ops/cold_i8.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+namespace ninfer::ops {
+
+// Fixed raw-slot size: 16 B header + 8192 B E2M1 nibbles + 1024 B E4M3 g16
+// scales (see ops/kernel/cold_i8_kernels.cuh for the composition).
+inline constexpr std::int32_t kColdI8SlotBytes = 9232;
+
+// Raw cold-slot codec for INT8-tier KV pages (entropy-cold revision 2b).
+//
+// Slot layout (9232 B, fixed, no overflow): 16 B header | 8192 B packed
+// E2M1 nibbles | 1024 B E4M3 group-16 scales, both in the nvfp4 page-major
+// geometry produced by entropy_cold_requant_raw(Int8G64). Measured on real
+// Qwen3.8 INT8 planes: NMSE 0.012-0.014 (inside the accepted NVFP4-layer
+// envelope) at 1.83x per head-page vs the raw int8 plane (16896 B).
+void cold_i8_slot_pack_raw(const std::uint8_t* src_codes, const std::uint8_t* src_scales,
+                           int kv_heads, int page_count, std::uint8_t* slots,
+                           std::int32_t* slot_valid, int slot_bytes, cudaStream_t stream);
+
+// Inverse: unpack slots into the INT8 tier's native planes (int8 codes +
+// fp16 group-64 scales). Used by the warm-restore path.
+void cold_i8_slot_restore_raw(const std::uint8_t* slots, int kv_heads, int page_count,
+                              std::int8_t* dst_codes, void* dst_scales_fp16,
+                              int slot_bytes, cudaStream_t stream);
+
+} // namespace ninfer::ops

include/ninfer/ops/dflash2_grouped_conv.h

@@ -0,0 +1,38 @@
+#pragma once
+
+#include "core/tensor.h"
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+namespace ninfer::ops {
+
+/**
+ * Op: DFlash2 grouped convolution pass
+ *
+ * `hidden` is contiguous BF16 [H,T]: the token-major block produced by the
+ * draft projections. `delta` is contiguous BF16 [N,T] with N=2*taps*G and is
+ * the exact token-major output of the kernel projection (N fastest). `side`
+ * selects the pass's tap bank. `base` is a BF16 [taps,H] Weight (one side of
+ * the stored [2,taps,H] base kernel). `out` is contiguous BF16 [H,T] and is
+ * completely overwritten.
+ *
+ * For token t with in-block position p = t & (block_size-1) and channel h in
+ * group g = h / group_size,
+ *
+ *   out[h,t] = hidden[h,t] * (base[0,h] + delta[G*(side*taps)+g,t])
+ *            + sum_{tap>=1, p>=tap} hidden[h,t-tap]
+ *              * (base[tap,h] + delta[G*(side*taps+tap)+g,t]).
+ *
+ * Cross-block reads introduced by t-tap are multiplied by the zero in-block
+ * position mask, so rows never leak into a neighbour block. The registered
+ * domain is H=5120, T=1..64, taps=2, G=320, group_size=16, block_size=8,
+ * side in {0,1}. Inputs and base are unchanged and the Op owns no workspace or
+ * persistent state. Intermediate arithmetic is FP32 and out is rounded to BF16.
+ */
+void dflash2_grouped_conv(const Tensor& hidden, const Tensor& delta, const Weight& base,
+                          std::int32_t block_size, std::int32_t group_size, std::int32_t taps,
+                          std::int32_t side, Tensor& out, cudaStream_t stream);
+
+} // namespace ninfer::ops

include/ninfer/ops/dflash2_selector.h

@@ -0,0 +1,47 @@
+#pragma once
+
+#include "core/tensor.h"
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+namespace ninfer::ops {
+
+/**
+ * Op: DFlash2 candidate selector with a greedy path walk
+ *
+ * `unary_logits` is contiguous BF16 [V,S*B]: full-vocab base logits of S
+ * proposal positions for B batch rows, token-major (each column holds one
+ * token's V logits). `projected_hidden` is contiguous BF16 [R,S*B]: the
+ * candidate selector's hidden projection of the same proposal columns.
+ * `predecessor_codebook` and `successor_codebook` are BF16 [V,R] Weights.
+ * `anchors` is contiguous I32 [B] and holds the previous target token per row.
+ *
+ * `candidates` (I32), `unary` (F32), and `scores` (F32) are caller-owned
+ * scratch with shapes [B,S,K], [B,S,K], and [B,S,K,K] respectively; all are
+ * completely overwritten. `drafts` is contiguous I32 [S*B] and receives the
+ * selected token for every proposal position.
+ *
+ * Selection: per row/step the top-K base logits (higher value first, lower
+ * token id breaking ties) form the candidate set. Each candidate c at step s
+ * receives the pair score
+ *
+ *   score(s,p,c) = unary[s,c]
+ *                + sum_r projected[r,s] * predecessor_codebook[pred,r]
+ *                                      * successor_codebook[candidate[s,c],r],
+ *
+ * where pred is the anchor token at s=0 and candidate[s-1,p] afterwards; only
+ * p=0 is admitted at s=0. The walk starts at previous candidate index 0 and
+ * greedily picks the highest-score current candidate at every step (lowest
+ * candidate index breaking ties).
+ *
+ * The registered domain is V=248320, R=256, S=7, B=1..8, K=16. Inputs and
+ * weights are unchanged and the Op owns no workspace or persistent state.
+ */
+void dflash2_selector(const Tensor& unary_logits, const Tensor& projected_hidden,
+                      const Weight& predecessor_codebook, const Weight& successor_codebook,
+                      const Tensor& anchors, Tensor& candidates, Tensor& unary, Tensor& scores,
+                      Tensor& drafts, std::int32_t steps, std::int32_t top_k, cudaStream_t stream);
+
+} // namespace ninfer::ops

include/ninfer/ops/entropy_cold_requant.h

@@ -0,0 +1,32 @@
+#pragma once
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+namespace ninfer::ops {
+
+// Cold-page requantization for the entropy slot codec (revision 2).
+// Requantizes one or more stored KV page planes (E2M1 g16 NVFP4 or int8 g64)
+// into fresh NVFP4 planes whose codes carry one E4M3FN scale per 64 channels
+// (replicated into the four g16 scale slots). The output layout matches the
+// page-major nvfp4 planes entropy_nvfp4_slot_encode_raw consumes, so the slot
+// codec, decode path, scale scatter, and attention producers stay unchanged;
+// the g64 requant is what makes the rANS streams compressible (measured
+// 2.0-2.6 bits/code vs ~4.0 for g16 storage codes on Qwen3.8 frames).
+enum class EntropyColdRequantMode : int {
+    Nvfp4G16 = 0,
+    Int8G64  = 1,
+    Iso3VG16 = 2,
+};
+
+// src planes use page-major strides: codes kv_heads*8192 (nvfp4) or
+// kv_heads*16384 (int8) bytes per page, scales kv_heads*1024 (nvfp4) or
+// kv_heads*512 (int8). dst planes are nvfp4 page-major: codes
+// [128, 64, kv_heads, page_count], scales [16, 64, kv_heads, page_count].
+void entropy_cold_requant_raw(const std::uint8_t* src_codes, const std::uint8_t* src_scales,
+                              EntropyColdRequantMode mode, int kv_heads, int page_count,
+                              std::uint8_t* dst_codes, std::uint8_t* dst_scales,
+                              cudaStream_t stream);
+
+} // namespace ninfer::ops