PR1: per-layer KV storage (mixed-precision layer tables)

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"
 
@@ -283,6 +284,13 @@ 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;
+        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;
+        }
         // 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,7 +79,7 @@ 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"
@@ -134,7 +135,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");

apps/cli/options.h

@@ -27,6 +27,8 @@ struct Options {
     SpeculativeOptions speculative;
     bool enable_vision  = false;
     bool use_cuda_graph = true;
+    std::string kv_layer_storage_spec;
+    bool kv_layer_storage_explicit = false;
 
     bool raw_output      = false;
     bool print_token_ids = false;

include/ninfer/types.h

@@ -30,8 +30,15 @@ enum class KvCacheStorage : std::uint8_t {
     BFloat16,
     Int8Group64,
     Fp8E4M3Row256,
+    Nvfp4Group16,
+    Fp8Group16,
+    Iso3Group16,
 };
 
+// Per-layer table width shared by targets that publish per-layer KV storage.
+// Entries are indexed by full-attention layer order, not physical model layer.
+inline constexpr std::size_t kKvLayerStorageSlots = 16;
+
 enum class EnginePurpose : std::uint8_t {
     Generation,
     CausalScoring,
@@ -115,6 +122,12 @@ struct EngineOptions {
     std::uint32_t pending_timeout_ms   = 30000;
     std::uint32_t prefill_chunk        = 1024;
     KvCacheStorage kv_cache            = KvCacheStorage::BFloat16;
+    // Per-layer KV storage override, indexed by full-attention layer order.
+    // BFloat16 entries inherit kv_cache. Any non-BFloat16 entry replaces the
+    // target's registered per-layer default table wholesale; entries outside
+    // the target's full-attention layer count are rejected.
+    std::array<KvCacheStorage, kKvLayerStorageSlots> kv_layer_storage{};
+    bool kv_layer_storage_explicit = false;
     SpeculativeOptions speculative;
     std::size_t media_cache_bytes = kDefaultMediaCacheBytes;
     std::size_t media_live_bytes  = kDefaultMediaLiveBytes;

src/core/dtype.h

@@ -14,6 +14,9 @@ enum class DType : std::uint8_t {
     I8         = 5,
     FP16       = 6,
     FP8_E4M3FN = 7,
+    // Packed E2M1 nibble plane (two codes per byte) with per-16-channel
+    // E4M3FN scales; see the per-layer KV storage table.
+    NVFP4      = 8,
 };
 
 std::size_t dtype_size(DType dtype);

src/core/paged_kv_cache.h

@@ -27,6 +27,7 @@ struct PagedKVLayerView {
     std::int32_t num_kv_heads = 0;
     DType dtype               = DType::BF16;
     std::int32_t quant_group  = 0;
+    std::array<DType, 16> layer_dtypes{};
 };
 
 /** Non-owning multi-sequence view consumed by batched growing-cache Ops. */

src/product/kv_options.h

@@ -0,0 +1,87 @@
+#pragma once
+
+#include "ninfer/types.h"
+
+#include <array>
+#include <cstddef>
+#include <optional>
+#include <stdexcept>
+#include <string>
+#include <string_view>
+
+namespace ninfer::product {
+
+// Per-layer KV storage spec parsing for CLI and serving.
+//
+// Spec grammar (comma separated):
+//   all:<type>     every registered full-attention layer
+//   <type>         shorthand for all:<type>
+//   A:<type>       one layer, A in [0, 15]
+//   A-B:<type>     inclusive layer range A..B, A <= B in [0, 15]
+// where <type> is bf16, int8, fp8, or nvfp4. Unlisted slots stay BFloat16,
+// which means "inherit the global --kv-dtype". A slot may be written exactly
+// once. Cold-policy parsing lives with the cold-pool change, not here.
+
+[[nodiscard]] inline std::optional<KvCacheStorage> parse_kv_storage(std::string_view text) {
+    if (text == "bf16") { return KvCacheStorage::BFloat16; }
+    if (text == "int8") { return KvCacheStorage::Int8Group64; }
+    if (text == "nvfp4") { return KvCacheStorage::Nvfp4Group16; }
+    if (text == "fp8") { return KvCacheStorage::Fp8Group16; }
+    if (text == "iso3") { return KvCacheStorage::Iso3Group16; }
+    return std::nullopt;
+}
+
+[[nodiscard]] inline std::array<KvCacheStorage, kKvLayerStorageSlots>
+parse_kv_layer_storage(std::string_view spec) {
+    std::array<KvCacheStorage, kKvLayerStorageSlots> table{};
+    if (spec.empty()) { return table; }
+
+    std::size_t begin = 0;
+    while (begin < spec.size()) {
+        const std::size_t comma = spec.find(',', begin);
+        const std::string_view item =
+            spec.substr(begin, comma == std::string_view::npos ? spec.size() - begin
+                                                               : comma - begin);
+        if (item.empty()) { throw std::invalid_argument("kv-layer-storage has an empty entry"); }
+
+        const std::size_t colon = item.find(':');
+        std::size_t first = 0;
+        std::size_t last = kKvLayerStorageSlots - 1;
+        std::string_view type = item;
+        if (colon != std::string_view::npos) {
+            const std::string_view layers = item.substr(0, colon);
+            type = item.substr(colon + 1);
+            if (layers == "all") {
+                first = 0;
+                last = kKvLayerStorageSlots - 1;
+            } else {
+                const std::size_t dash = layers.find('-');
+                if (dash == std::string_view::npos) {
+                    first = last =
+                        static_cast<std::size_t>(std::stoul(std::string(layers)));
+                } else {
+                    first = static_cast<std::size_t>(std::stoul(std::string(
+                        layers.substr(0, dash))));
+                    last = static_cast<std::size_t>(std::stoul(std::string(
+                        layers.substr(dash + 1))));
+                }
+                if (first > last || last >= kKvLayerStorageSlots) {
+                    throw std::invalid_argument("kv-layer-storage layer index out of range");
+                }
+            }
+        }
+        const auto value = parse_kv_storage(type);
+        if (!value) { throw std::invalid_argument("kv-layer-storage has an invalid type"); }
+        for (std::size_t slot = first; slot <= last; ++slot) {
+            if (table[slot] != KvCacheStorage::BFloat16) {
+                throw std::invalid_argument("kv-layer-storage slot written twice");
+            }
+            table[slot] = *value;
+        }
+        if (comma == std::string_view::npos) { break; }
+        begin = comma + 1;
+    }
+    return table;
+}
+
+} // namespace ninfer::product

src/targets/qwen3_6/export/ninfer/targets/qwen3_6/decoder_state.h

@@ -11,6 +11,7 @@ namespace ninfer::targets::qwen3_6 {
 
 inline constexpr std::int32_t kKvInt8QuantGroup = 64;
 inline constexpr std::int32_t kKvFp8QuantGroup  = 256;
+inline constexpr std::int32_t kNvfp4KvQuantGroup = 16;
 
 struct DecoderStateSpec {
     std::uint32_t full_attention_layers     = 0;
@@ -20,6 +21,9 @@ struct DecoderStateSpec {
     std::int32_t attention_head_dim         = 0;
     DType kv_dtype                          = DType::BF16;
     std::int32_t kv_quant_group             = 0;
+    // Per-layer storage table indexed by full-attention layer order.
+    // BFloat16 entries inherit kv_dtype; empty (all-BF16) inherits wholesale.
+    std::array<DType, 16> layer_kv_dtypes{};
     bool enable_mtp                         = false;
     std::int32_t kv_table_rows              = 1;
     std::uint32_t text_physical_page_groups = 0;
@@ -35,6 +39,12 @@ struct PagedKVCacheLayout {
     std::int32_t head_dim     = 0;
     DType dtype               = DType::BF16;
     std::int32_t quant_group  = 0;
+    // Resolved per-layer storage (one entry per full-attention layer).
+    std::array<DType, 16> layer_dtypes{};
+    // Plane offset of each layer in the page geometry (prefix sums over
+    // per-layer plane counts; mixed BF16/quantized tables have unequal
+    // strides).
+    std::array<std::uint32_t, 16> layer_plane_base{};
 
     [[nodiscard]] std::size_t payload_bytes() const noexcept { return pages.payload_bytes(); }
 };
@@ -96,6 +106,8 @@ class PagedKVCache {
     std::int32_t kv_heads_     = 0;
     std::int32_t head_dim_     = 0;
     DType dtype_               = DType::BF16;
+    std::array<DType, 16> layer_dtypes_{};
+    std::array<std::uint32_t, 16> layer_plane_base_{};
     std::int32_t quant_group_  = 0;
 };
 

src/targets/qwen3_6/impl/runtime/layouts.h

@@ -76,6 +76,7 @@ struct SequencePlanningInputs {
     SpeculativeBackend speculative_backend = SpeculativeBackend::None;
     DType kv_dtype                         = DType::BF16;
     std::int32_t kv_quant_group            = 0;
+    std::array<DType, 16> layer_kv_dtypes{};
     ProposalHead proposal_head             = ProposalHead::Full;
     StartupFeatures features;
     bool use_cuda_graph = true;
@@ -100,6 +101,7 @@ struct SequencePlanImpl<NINFER_QWEN36_VARIANT> {
     SpeculativeBackend speculative_backend = SpeculativeBackend::None;
     DType kv_dtype                         = DType::BF16;
     std::int32_t kv_quant_group            = 0;
+    std::array<DType, 16> layer_kv_dtypes{};
     ProposalHead proposal_head             = ProposalHead::Full;
     StartupFeatures features;
     bool use_cuda_graph = true;

src/targets/qwen3_6/impl/runtime/layouts_impl.h

@@ -136,6 +136,7 @@ PersistentLayout persistent_layout(const SequencePlanImpl& plan) {
                      .attention_head_dim        = TextConfig::head_dim,
                      .kv_dtype                  = plan.kv_dtype,
                      .kv_quant_group            = plan.kv_quant_group,
+                     .layer_kv_dtypes           = plan.layer_kv_dtypes,
                      .enable_mtp                = plan.features.mtp(),
                      .kv_table_rows             = static_cast<std::int32_t>(plan.max_concurrency),
                      .text_physical_page_groups = physical_pages,
@@ -675,6 +676,7 @@ std::unique_ptr<SequencePlanImpl> build_sequence_candidate(const SequencePlannin
     impl->context_cache       = inputs.context_cache;
     impl->kv_dtype            = inputs.kv_dtype;
     impl->kv_quant_group      = inputs.kv_quant_group;
+    impl->layer_kv_dtypes     = inputs.layer_kv_dtypes;
     impl->persistent          = persistent_layout(*impl);
     impl->workspace           = build_workspace_plan(*impl);
     if (impl->use_cuda_graph) {
@@ -733,6 +735,25 @@ make_sequence_planner_impl(DeviceContext& device, const EngineOptions& options,
     validate_target_options(device, options);
     const TargetKVCacheProfile kv_profile = target_kv_cache_profile(options.kv_cache);
 
+    std::array<DType, 16> layer_overrides{};
+    const bool has_override = options.kv_layer_storage_explicit;
+    if (has_override) {
+        for (std::size_t i = 0; i < layer_overrides.size(); ++i) {
+            const auto v = options.kv_layer_storage[i];
+            layer_overrides[i] = v == KvCacheStorage::BFloat16
+                                     ? DType::BF16
+                                     : (v == KvCacheStorage::Int8Group64
+                                            ? DType::I8
+                                            : (v == KvCacheStorage::Nvfp4Group16
+                                                   ? DType::NVFP4
+                                                   : (v == KvCacheStorage::Fp8Group16
+                                                          ? DType::FP8_E4M3FN
+                                                          : DType::BF16)));
+        }
+    } else if constexpr (Variant::supports_per_layer_kv_defaults) {
+        layer_overrides = Variant::default_layer_kv_dtypes(
+            weights_profile);
+    }
     SequencePlanningInputs inputs{
         .weights_profile     = weights_profile,
         .capacity            = options.max_context,
@@ -742,6 +763,7 @@ make_sequence_planner_impl(DeviceContext& device, const EngineOptions& options,
         .speculative_backend = options.speculative.backend,
         .kv_dtype            = kv_profile.dtype,
         .kv_quant_group      = kv_profile.quant_group,
+        .layer_kv_dtypes     = layer_overrides,
         .proposal_head       = options.speculative.proposal_head,
         .features            = qwen3_6::startup_features(options),
         .use_cuda_graph      = options.use_cuda_graph,