PR2: entropy-coded cold pool (window policy)

apps/cli/main.cpp

@@ -283,6 +283,9 @@ int main(int argc, char** argv) {
         engine_options.speculative    = cli.speculative;
         engine_options.enable_vision  = cli.enable_vision;
         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;
         // 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

@@ -85,6 +85,8 @@ std::string usage_text(const char* argv0) {
            "       [--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"
+           "       [--cold-policy none|off|window|host] [--cold-keep-tokens N]\n"
+           "       [--cold-host-bytes N[g|m|k]]\n"
            "       [--no-cuda-graph]\n"
            "\n"
            "Streams answer content to stdout and reasoning plus diagnostics to stderr.\n"
@@ -152,6 +154,17 @@ 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 == "--no-cuda-graph") {
             options.use_cuda_graph = false;
         } else if (arg == "--stop-token-id") {

apps/cli/options.h

@@ -27,6 +27,9 @@ 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 raw_output      = false;
     bool print_token_ids = false;

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, 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,
+                              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

include/ninfer/types.h

@@ -37,6 +37,15 @@ enum class EnginePurpose : std::uint8_t {
     CausalScoring,
 };
 
+// Entropy-coded cold KV pool. Window compresses fully-written pages that
+// are at least cold_keep_tokens behind the decode frontier; attention
+// producers decode those pages inline from fixed-size slots.
+enum class ColdPolicy : std::uint8_t {
+    None,
+    Window,
+    Host,
+};
+
 enum class KvCapacityMode : std::uint8_t {
     Explicit,
     Automatic,
@@ -124,6 +133,10 @@ struct EngineOptions {
     bool use_cuda_graph                    = true;
     ContextCacheOptions context_cache;
     ContextCostOptions context_cost;
+    ColdPolicy cold_policy            = ColdPolicy::None;
+    std::uint32_t cold_keep_tokens    = 128;
+    // Pinned host-memory budget for ColdPolicy::Host offload. Default 4 GiB.
+    std::uint64_t cold_host_bytes     = 4ULL << 30;
     LoadProgress load_progress;
 };
 

src/CMakeLists.txt

@@ -70,6 +70,8 @@ add_library(ninfer_ops STATIC
   ops/launcher/embed_gather.cu
   ops/launcher/gdn_gating.cu
   ops/launcher/gelu.cu
+  ops/launcher/cold_i8.cu
+  ops/launcher/entropy_cold_requant.cu
   ops/launcher/l2norm.cu
   ops/launcher/layer_norm.cu
   ops/launcher/mtp_pack.cu
@@ -241,6 +243,8 @@ add_library(ninfer_ops STATIC
   ops/wrapper/cast.cpp
   ops/wrapper/causal_conv1d_silu.cpp
   ops/wrapper/embedding.cpp
+  ops/wrapper/cold_i8.cpp
+  ops/wrapper/entropy_cold_requant.cpp
   ops/wrapper/gdn_gating.cpp
   ops/wrapper/gdn_gating_proj.cpp
   ops/wrapper/gdn_input_proj.cpp

src/core/paged_kv_cache.h

@@ -16,13 +16,33 @@ namespace ninfer {
 
 inline constexpr std::int32_t kPagedKVPageSize = 64;
 
+// Cold-slot sentinel encoding in block tables: entries <= -2 address the
+// entropy pool (slot base = -2 - entry). The pool itself is per-layer and
+// owned by the decoder state; the table row only carries the slot base.
+inline constexpr std::int32_t kPagedKVColdSentinelBase = -2;
+
+[[nodiscard]] inline bool paged_kv_is_cold(std::int32_t entry) noexcept {
+    return entry <= kPagedKVColdSentinelBase;
+}
+[[nodiscard]] inline std::int32_t paged_kv_cold_slot_base(std::int32_t entry) noexcept {
+    return kPagedKVColdSentinelBase - entry;
+}
+[[nodiscard]] inline std::int32_t paged_kv_cold_entry(std::int32_t slot_base) noexcept {
+    return kPagedKVColdSentinelBase - slot_base;
+}
+
 /** Non-owning, single-sequence view consumed by growing-cache Ops. */
 struct PagedKVLayerView {
     Tensor k_pages;
     Tensor v_pages;
     Tensor k_scale_pages;
     Tensor v_scale_pages;
     Tensor block_table;
+    // Entropy-coded cold pool (single-sequence window view): fixed raw slots +
+    // validity flags. Empty when the cache has no cold pool.
+    Tensor cold_slots;
+    Tensor cold_slot_valid;
+    std::int32_t cold_slot_bytes = 0;
     std::int32_t head_dim     = 0;
     std::int32_t num_kv_heads = 0;
     DType dtype               = DType::BF16;
@@ -36,6 +56,10 @@ struct PagedKVBatchLayerView {
     Tensor k_scale_pages;
     Tensor v_scale_pages;
     Tensor block_tables;
+    // Entropy-coded cold pool: fixed raw slots + validity flags per layer.
+    Tensor cold_slots;
+    Tensor cold_slot_valid;
+    std::int32_t cold_slot_bytes = 0;
     std::int32_t head_dim     = 0;
     std::int32_t num_kv_heads = 0;
     DType dtype               = DType::BF16;
@@ -112,6 +136,10 @@ class DeviceKVPageHandle {
 
     [[nodiscard]] bool valid() const noexcept { return owner_ != nullptr; }
 
+    // Physical page index within the owning pool (public read access for the
+    // cold-compression pass).
+    [[nodiscard]] std::int32_t index() const noexcept { return index_; }
+
 private:
     friend class DeviceKVPagePool;
     friend class DeviceKVPageLease;
@@ -351,13 +379,14 @@ class KVExecutionTablePool {
 
     [[nodiscard]] const Tensor& matrix() const noexcept { return block_tables_; }
 
+    void publish_indices(KVExecutionRowHandle row, std::uint32_t logical_begin,
+                         std::span<const std::int32_t> indices, cudaStream_t stream);
 private:
     friend class KVExecutionRowLease;
 
     [[nodiscard]] bool valid_handle(KVExecutionRowHandle handle) const noexcept;
     bool release_row(std::int32_t row, std::uint32_t generation) noexcept;
-    void publish_indices(KVExecutionRowHandle row, std::uint32_t logical_begin,
-                         std::span<const std::int32_t> indices, cudaStream_t stream);
+
 
     KVExecutionTableSpec spec_;
     const DeviceKVPagePool* pages_ = nullptr;

src/ops/kernel/cold_i8_kernels.cuh

@@ -0,0 +1,105 @@
+#pragma once
+
+// ninfer::ops::detail - raw cold-slot codec for INT8-tier KV pages (v2b).
+//
+// The INT8 tier's requantized E2M1 codes are near-uniform (measured H
+// 3.6-3.9 on real planes), so rANS gains nothing; this slot format stores
+// the requant output verbatim with a fixed layout and no overflow path:
+//
+//   [ 16 B header | 8192 B packed E2M1 nibbles | 1024 B E4M3 g16 scales ]
+//
+// The nibble/scale planes use the same page-major geometry the NVFP4 tier
+// stores (codes [128, 64, kv_heads, pages], scales [16, 64, kv_heads,
+// pages]), produced by entropy_cold_requant's Int8G64 mode. Restore
+// converts a slot back into the INT8 tier's native planes (int8 codes +
+// fp16 group-64 scales) with the upper-bound group scale, adding <0.4%
+// quantization noise on top of the requant's measured 0.012 NMSE.
+//
+// Kernel DEFINITIONS live only in ops/launcher/cold_i8.cu; this header
+// declares them plus the shared device helpers so attention kernels can
+// include it without duplicate device-link definitions.
+
+#include "ninfer/ops/cold_i8.h"
+#include "ops/kernel/gqa_attention_kv_nvfp4.cuh"
+
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+namespace ninfer::ops::detail {
+
+inline constexpr int kColdI8SlotHeaderBytes = 16;
+inline constexpr int kColdI8SlotCodeBytes   = 8192;   // 64 rows x 128 B nibbles
+inline constexpr int kColdI8SlotScaleBytes  = 1024;   // 64 rows x 16 B E4M3
+static_assert(kColdI8SlotHeaderBytes + kColdI8SlotCodeBytes + kColdI8SlotScaleBytes ==
+              ninfer::ops::kColdI8SlotBytes);
+inline constexpr std::uint32_t kColdI8SlotMagic = 0x49384352u; // "RC8I"
+
+// Pack one requantized (page, head, plane) into a raw slot. src uses the
+// nvfp4 page-major geometry entropy_cold_requant emits; slots layout is
+// [slot_bytes, kv_heads, 2, pages] with V one nb[2] step past K.
+__global__ void cold_i8_slot_pack_kernel(const std::uint8_t* __restrict__ src_codes,
+                                         const std::uint8_t* __restrict__ src_scales,
+                                         int kv_heads,
+                                         std::uint8_t* __restrict__ slots,
+                                         std::int32_t* __restrict__ slot_valid);
+
+// Warm restore: unpack one (page, head, plane) slot into the INT8 tier's
+// native planes (int8 codes [256,64,kv_heads,pages], fp16 scales
+// [4,64,kv_heads,pages]). One block per (head, page); 256 threads split
+// the 64 rows.
+__global__ void cold_i8_slot_restore_kernel(const std::uint8_t* __restrict__ slots,
+                                            int kv_heads, std::int8_t* __restrict__ dst_codes,
+                                            __half* __restrict__ dst_scales);
+
+// Slot region accessors for producers.
+__device__ __forceinline__ const std::uint8_t*
+cold_i8_slot_scales(const std::uint8_t* slot) {
+    return slot + kColdI8SlotHeaderBytes + kColdI8SlotCodeBytes;
+}
+
+__device__ __forceinline__ const std::uint8_t*
+cold_i8_slot_codes(const std::uint8_t* slot) {
+    return slot + kColdI8SlotHeaderBytes;
+}
+
+// Decode one key row of a raw slot into INT8-tier native form: 256 int8
+// codes plus one fp16 scale per 64-channel group. The group scale is the
+// upper bound 6*max(e4m3 sub-scales)/127 so no amax scan of the decoded
+// values is needed; codes clamp at 127 so fp16 rounding-down is safe.
+// Used by the warm-restore kernel and the attention producers' cold staging.
+__device__ __forceinline__ void cold_i8_decode_row(const std::uint8_t* slot, int row,
+                                                   std::int8_t* codes_out,   // 256, d-major
+                                                   __half* scales_out) {     // 4 groups
+    const std::uint8_t* row_codes  = cold_i8_slot_codes(slot) + row * 128;
+    const std::uint8_t* row_scales = cold_i8_slot_scales(slot) + row * 16;
+#pragma unroll
+    for (int g = 0; g < 4; ++g) {
+        float mx = 0.0f;
+#pragma unroll
+        for (int s = 0; s < 4; ++s) {
+            mx = fmaxf(mx, gqa_kv_nvfp4_e4m3_to_f32(row_scales[g * 4 + s]));
+        }
+        const float scale = mx * 6.0f / 127.0f;
+        scales_out[g]     = __float2half(scale);
+        const float inv   = scale > 0.0f ? 1.0f / scale : 0.0f;
+#pragma unroll
+        for (int i = 0; i < 64; i += 2) {
+            const int d          = g * 64 + i;
+            const std::uint8_t b = row_codes[d >> 1];
+            const float v0 = gqa_kv_nvfp4_e2m1_to_f32(b & 0x0F) *
+                             gqa_kv_nvfp4_e4m3_to_f32(row_scales[d >> 4]);
+            const float v1 = gqa_kv_nvfp4_e2m1_to_f32(b >> 4) *
+                             gqa_kv_nvfp4_e4m3_to_f32(row_scales[(d + 1) >> 4]);
+            int c0         = __float2int_rn(v0 * inv);
+            int c1         = __float2int_rn(v1 * inv);
+            c0             = max(-127, min(127, c0));
+            c1             = max(-127, min(127, c1));
+            codes_out[d]     = static_cast<std::int8_t>(c0);
+            codes_out[d + 1] = static_cast<std::int8_t>(c1);
+        }
+    }
+}
+
+} // namespace ninfer::ops::detail