diff --git a/apps/cli/main.cpp b/apps/cli/main.cpp index 933192da9b..cbb0038468 100644 --- a/apps/cli/main.cpp +++ b/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; diff --git a/apps/cli/options.cpp b/apps/cli/options.cpp index b5c5798d9a..1c79d5b6ee 100644 --- a/apps/cli/options.cpp +++ b/apps/cli/options.cpp @@ -1,5 +1,6 @@ #include "options.h" #include "product/speculative_options.h" +#include "product/kv_options.h" #include #include @@ -78,7 +79,7 @@ std::string usage_text(const char* argv0) { " (--prompt |--messages )\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"); diff --git a/apps/cli/options.h b/apps/cli/options.h index 3c0c2e0960..fca7acf2cc 100644 --- a/apps/cli/options.h +++ b/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; diff --git a/include/ninfer/types.h b/include/ninfer/types.h index 677c381e88..99dbb1810b 100644 --- a/include/ninfer/types.h +++ b/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 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; diff --git a/src/core/dtype.h b/src/core/dtype.h index 9d68741181..e860abdb5a 100644 --- a/src/core/dtype.h +++ b/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); diff --git a/src/core/paged_kv_cache.h b/src/core/paged_kv_cache.h index 9a4e079148..d09c92fa77 100644 --- a/src/core/paged_kv_cache.h +++ b/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 layer_dtypes{}; }; /** Non-owning multi-sequence view consumed by batched growing-cache Ops. */ diff --git a/src/product/kv_options.h b/src/product/kv_options.h new file mode 100644 index 0000000000..652936ba01 --- /dev/null +++ b/src/product/kv_options.h @@ -0,0 +1,87 @@ +#pragma once + +#include "ninfer/types.h" + +#include +#include +#include +#include +#include +#include + +namespace ninfer::product { + +// Per-layer KV storage spec parsing for CLI and serving. +// +// Spec grammar (comma separated): +// all: every registered full-attention layer +// shorthand for all: +// A: one layer, A in [0, 15] +// A-B: inclusive layer range A..B, A <= B in [0, 15] +// where 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 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 +parse_kv_layer_storage(std::string_view spec) { + std::array 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::stoul(std::string(layers))); + } else { + first = static_cast(std::stoul(std::string( + layers.substr(0, dash)))); + last = static_cast(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 diff --git a/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/decoder_state.h b/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/decoder_state.h index f1193f588a..121a4f8f59 100644 --- a/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/decoder_state.h +++ b/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 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 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 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 layer_dtypes_{}; + std::array layer_plane_base_{}; std::int32_t quant_group_ = 0; }; diff --git a/src/targets/qwen3_6/impl/runtime/layouts.h b/src/targets/qwen3_6/impl/runtime/layouts.h index 859f4f5d7e..c232877886 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts.h +++ b/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 layer_kv_dtypes{}; ProposalHead proposal_head = ProposalHead::Full; StartupFeatures features; bool use_cuda_graph = true; @@ -100,6 +101,7 @@ struct SequencePlanImpl { SpeculativeBackend speculative_backend = SpeculativeBackend::None; DType kv_dtype = DType::BF16; std::int32_t kv_quant_group = 0; + std::array layer_kv_dtypes{}; ProposalHead proposal_head = ProposalHead::Full; StartupFeatures features; bool use_cuda_graph = true; diff --git a/src/targets/qwen3_6/impl/runtime/layouts_impl.h b/src/targets/qwen3_6/impl/runtime/layouts_impl.h index b21165ac86..87f2c8e462 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts_impl.h +++ b/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(plan.max_concurrency), .text_physical_page_groups = physical_pages, @@ -675,6 +676,7 @@ std::unique_ptr 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 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, diff --git a/src/targets/qwen3_6/impl/state/decoder_state.cpp b/src/targets/qwen3_6/impl/state/decoder_state.cpp index 5e7372df45..7450ac6e70 100644 --- a/src/targets/qwen3_6/impl/state/decoder_state.cpp +++ b/src/targets/qwen3_6/impl/state/decoder_state.cpp @@ -1,6 +1,7 @@ #include #include +#include #include #include @@ -14,21 +15,44 @@ std::uint32_t page_count(std::uint32_t capacity) { PagedKVCacheLayout plan_cache(LayoutBuilder& builder, std::uint32_t layers, std::uint32_t capacity, std::int32_t kv_heads, std::int32_t head_dim, DType dtype, - std::int32_t quant_group, std::int32_t table_rows, - std::uint32_t physical_page_groups) { + std::int32_t quant_group, std::span layer_dtypes, + std::int32_t table_rows, std::uint32_t physical_page_groups) { if (layers == 0 || layers > static_cast(std::numeric_limits::max()) || kv_heads <= 0 || head_dim <= 0 || table_rows <= 0) { throw std::invalid_argument("Paged KV cache geometry is invalid"); } - const bool scaled = dtype == DType::I8 || dtype == DType::FP8_E4M3FN; - const bool valid_profile = - (dtype == DType::BF16 && quant_group == 0) || - (dtype == DType::I8 && quant_group == kKvInt8QuantGroup && head_dim % quant_group == 0) || - (dtype == DType::FP8_E4M3FN && head_dim == kKvFp8QuantGroup && - quant_group == kKvFp8QuantGroup); - if (!valid_profile) { - throw std::invalid_argument("Paged KV cache dtype or quantization is invalid"); + if (!layer_dtypes.empty() && layer_dtypes.size() < layers) { + throw std::invalid_argument("Paged KV per-layer dtype table is shorter than the layer count"); + } + // Per-layer resolution: BF16 entries inherit the global dtype. Accepted + // per-layer storages are the quantized codecs with their native group. + const auto layer_dtype = [&](std::uint32_t layer) { + const DType override_dtype = layer_dtypes.empty() ? DType::BF16 : layer_dtypes[layer]; + const DType selected = override_dtype == DType::BF16 ? dtype : override_dtype; + if (selected != DType::BF16 && selected != DType::I8 && selected != DType::NVFP4 && + selected != DType::FP8_E4M3FN) { + throw std::invalid_argument("Paged KV per-layer dtype is invalid"); + } + return selected; + }; + const auto layer_quant_group = [&](DType selected) { + return selected == DType::I8 + ? kKvInt8QuantGroup + : (selected == DType::BF16 + ? 0 + : (selected == DType::FP8_E4M3FN ? kKvFp8QuantGroup + : kNvfp4KvQuantGroup)); + }; + (void)quant_group; + for (std::uint32_t layer = 0; layer < layers; ++layer) { + const DType selected = layer_dtype(layer); + if (selected != DType::BF16) { + const std::int32_t group = layer_quant_group(selected); + if (head_dim % group != 0) { + throw std::invalid_argument("Paged KV per-layer quantization is invalid"); + } + } } const std::uint32_t logical_pages = page_count(capacity); @@ -37,14 +61,40 @@ PagedKVCacheLayout plan_cache(LayoutBuilder& builder, std::uint32_t layers, std: } KVPageGeometry geometry; - geometry.planes.reserve(static_cast(layers) * (scaled ? 4ULL : 2ULL)); + geometry.planes.reserve(static_cast(layers) * 4ULL); + std::array stored{}; + std::array plane_base{}; + std::uint32_t plane_cursor = 0; for (std::uint32_t layer = 0; layer < layers; ++layer) { - geometry.planes.push_back({dtype, head_dim, kv_heads, 256}); - geometry.planes.push_back({dtype, head_dim, kv_heads, 256}); - if (scaled) { - geometry.planes.push_back({DType::FP16, head_dim / quant_group, kv_heads, 256}); - geometry.planes.push_back({DType::FP16, head_dim / quant_group, kv_heads, 256}); + const DType selected = layer_dtype(layer); + stored[layer] = selected; + plane_base[layer] = plane_cursor; + const std::int32_t group = layer_quant_group(selected); + if (selected == DType::BF16) { + geometry.planes.push_back({DType::BF16, head_dim, kv_heads, 256}); + geometry.planes.push_back({DType::BF16, head_dim, kv_heads, 256}); + } else if (selected == DType::I8) { + geometry.planes.push_back({DType::I8, head_dim, kv_heads, 256}); + geometry.planes.push_back({DType::I8, head_dim, kv_heads, 256}); + geometry.planes.push_back({DType::FP16, head_dim / group, kv_heads, 256}); + geometry.planes.push_back({DType::FP16, head_dim / group, kv_heads, 256}); + } else if (selected == DType::FP8_E4M3FN) { + geometry.planes.push_back({DType::FP8_E4M3FN, head_dim, kv_heads, 256}); + geometry.planes.push_back({DType::FP8_E4M3FN, head_dim, kv_heads, 256}); + // FP8 per-group scales are FP16 in the production codecs; the + // attention kernels require FP16 scale planes. + geometry.planes.push_back({DType::FP16, head_dim / group, kv_heads, 256}); + geometry.planes.push_back({DType::FP16, head_dim / group, kv_heads, 256}); + } else { + // NVFP4 tier: K keeps E2M1 packed codes with per-16 E4M3FN scales; + // V stores ISO3 sign-magnitude nibbles in the same plane geometry + // (semantic split via v_dtype, no extra payload). + geometry.planes.push_back({DType::U8, head_dim / 2, kv_heads, 256}); + geometry.planes.push_back({DType::U8, head_dim / 2, kv_heads, 256}); + geometry.planes.push_back({DType::FP8_E4M3FN, head_dim / group, kv_heads, 256}); + geometry.planes.push_back({DType::FP8_E4M3FN, head_dim / group, kv_heads, 256}); } + plane_cursor += selected == DType::BF16 ? 2U : 4U; } return PagedKVCacheLayout{ .pages = plan_device_kv_page_pool( @@ -59,6 +109,8 @@ PagedKVCacheLayout plan_cache(LayoutBuilder& builder, std::uint32_t layers, std: .head_dim = head_dim, .dtype = dtype, .quant_group = quant_group, + .layer_dtypes = stored, + .layer_plane_base = plane_base, }; } @@ -66,13 +118,16 @@ PagedKVCacheLayout plan_cache(LayoutBuilder& builder, std::uint32_t layers, std: DecoderStateLayout plan_decoder_state(LayoutBuilder& builder, const DecoderStateSpec& spec) { DecoderStateLayout layout; + const std::span layer_dtypes(spec.layer_kv_dtypes.data(), + spec.full_attention_layers); layout.text_kv = plan_cache(builder, spec.full_attention_layers, spec.capacity, spec.kv_heads, spec.attention_head_dim, spec.kv_dtype, spec.kv_quant_group, - spec.kv_table_rows, spec.text_physical_page_groups); + layer_dtypes, spec.kv_table_rows, + spec.text_physical_page_groups); if (spec.enable_mtp) { layout.mtp_kv = plan_cache(builder, spec.mtp_layers, spec.capacity, spec.kv_heads, spec.attention_head_dim, spec.kv_dtype, spec.kv_quant_group, - spec.kv_table_rows, spec.mtp_physical_page_groups); + {}, spec.kv_table_rows, spec.mtp_physical_page_groups); } return layout; } @@ -80,7 +135,17 @@ DecoderStateLayout plan_decoder_state(LayoutBuilder& builder, const DecoderState PagedKVCache::PagedKVCache(DeviceSpan backing, const PagedKVCacheLayout& layout) : pages_(backing, layout.pages), execution_tables_(backing, layout.execution_tables, pages_), layers_(layout.layers), max_context_(layout.max_context), kv_heads_(layout.kv_heads), - head_dim_(layout.head_dim), dtype_(layout.dtype), quant_group_(layout.quant_group) {} + head_dim_(layout.head_dim), dtype_(layout.dtype), quant_group_(layout.quant_group), + cold_slot_bytes_(layout.cold_slot_bytes), max_cold_pages_(layout.max_cold_pages), + layer_dtypes_(layout.layer_dtypes), layer_plane_base_(layout.layer_plane_base) { + cold_slot_used_.assign(max_cold_pages_, 0); + for (std::uint32_t layer = 0; layer < layers_; ++layer) { + if (layout.cold_slots[layer].region.bytes != 0) { + cold_slots_[layer] = layout.cold_slots[layer].bind(backing); + cold_slot_valid_[layer] = layout.cold_slot_valid[layer].bind(backing); + } + } +} PagedKVCacheView::PagedKVCacheView(const PagedKVCache& cache, Tensor block_table) noexcept : cache_(&cache), block_table_(block_table) {} @@ -103,9 +168,14 @@ PagedKVCacheView PagedKVCache::execution_view(const KVExecutionRowLease& row) co PagedKVLayerView PagedKVCache::layer_view(std::uint32_t layer, Tensor block_table) const { if (layer >= layers_) { throw std::out_of_range("Paged KV layer is out of range"); } - const bool scaled = dtype_ == DType::I8 || dtype_ == DType::FP8_E4M3FN; - const std::size_t stride = scaled ? 4ULL : 2ULL; - const std::size_t base = static_cast(layer) * stride; + // The scaled/stride decision follows the layer's resolved dtype so a + // per-layer table (PR1) can mix quantized and BF16 layers in one pool. + const DType layer_dtype = + layer_dtypes_.empty() ? dtype_ : layer_dtypes_[layer]; + const bool scaled = layer_dtype == DType::I8 || layer_dtype == DType::FP8_E4M3FN; + const std::size_t base = layer_plane_base_.empty() + ? static_cast(layer) * (scaled ? 4ULL : 2ULL) + : layer_plane_base_[layer]; return PagedKVLayerView{ .k_pages = pages_.plane(base), .v_pages = pages_.plane(base + 1), @@ -114,16 +184,27 @@ PagedKVLayerView PagedKVCache::layer_view(std::uint32_t layer, Tensor block_tabl .block_table = block_table, .head_dim = head_dim_, .num_kv_heads = kv_heads_, - .dtype = dtype_, - .quant_group = quant_group_, + .dtype = layer_dtypes_.empty() ? dtype_ : layer_dtypes_[layer], + .quant_group = layer_dtypes_.empty() + ? quant_group_ + : (layer_dtypes_[layer] == DType::I8 + ? kKvInt8QuantGroup + : (layer_dtypes_[layer] == DType::FP8_E4M3FN + ? kKvFp8QuantGroup + : 0)), }; } PagedKVBatchLayerView PagedKVCache::batch_layer_view(std::uint32_t layer) const { if (layer >= layers_) { throw std::out_of_range("Paged KV layer is out of range"); } - const bool scaled = dtype_ == DType::I8 || dtype_ == DType::FP8_E4M3FN; - const std::size_t stride = scaled ? 4ULL : 2ULL; - const std::size_t base = static_cast(layer) * stride; + // The scaled/stride decision follows the layer's resolved dtype so a + // per-layer table (PR1) can mix quantized and BF16 layers in one pool. + const DType layer_dtype = + layer_dtypes_.empty() ? dtype_ : layer_dtypes_[layer]; + const bool scaled = layer_dtype == DType::I8 || layer_dtype == DType::FP8_E4M3FN; + const std::size_t base = layer_plane_base_.empty() + ? static_cast(layer) * (scaled ? 4ULL : 2ULL) + : layer_plane_base_[layer]; return PagedKVBatchLayerView{ .k_pages = pages_.plane(base), .v_pages = pages_.plane(base + 1), @@ -132,8 +213,14 @@ PagedKVBatchLayerView PagedKVCache::batch_layer_view(std::uint32_t layer) const .block_tables = execution_tables_.matrix(), .head_dim = head_dim_, .num_kv_heads = kv_heads_, - .dtype = dtype_, - .quant_group = quant_group_, + .dtype = layer_dtypes_.empty() ? dtype_ : layer_dtypes_[layer], + .quant_group = layer_dtypes_.empty() + ? quant_group_ + : (layer_dtypes_[layer] == DType::I8 + ? kKvInt8QuantGroup + : (layer_dtypes_[layer] == DType::FP8_E4M3FN + ? kKvFp8QuantGroup + : 0)), }; } diff --git a/src/targets/qwen3_6_27b/impl/variant.cpp b/src/targets/qwen3_6_27b/impl/variant.cpp index c2036d9145..4644c94064 100644 --- a/src/targets/qwen3_6_27b/impl/variant.cpp +++ b/src/targets/qwen3_6_27b/impl/variant.cpp @@ -19,6 +19,20 @@ #include "targets/qwen3_6/impl/runtime/instantiate.h" namespace ninfer::targets::qwen3_6_27b::detail { + +std::array Variant::default_layer_kv_dtypes(WeightsProfile) { + // Data-driven prior from the offline calibration history: layer 14 is an + // extreme outlier (uniform-precision K NMSE ~30x the next layer), and the + // next five layers dominate the remaining error. Upgrading those six to + // INT8 keeps the long-generation error budget bounded at a modest byte + // cost. The rest of the table is BF16 = inherit the global --kv-dtype. + std::array table{}; + for (const int layer : {2, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}) { + table[static_cast(layer)] = DType::I8; + } + return table; +} + namespace { std::vector diff --git a/src/targets/qwen3_6_27b/impl/variant.h b/src/targets/qwen3_6_27b/impl/variant.h index 75332671ad..6868dd999f 100644 --- a/src/targets/qwen3_6_27b/impl/variant.h +++ b/src/targets/qwen3_6_27b/impl/variant.h @@ -15,6 +15,8 @@ using GraphExecutionProfile = qwen3_6::GraphExecutionProfile; // Compile-time data and the three closed execution leaves supplied to the Qwen3.6 family runtime. // It owns no request state, execution phase, graph object, or schedule callback. struct Variant { + static constexpr bool supports_per_layer_kv_defaults = true; + [[nodiscard]] static std::array default_layer_kv_dtypes(WeightsProfile profile); using WeightsProfile = detail::WeightsProfile; using TextConfig = detail::TextConfig; using VisionConfig = detail::VisionConfig; diff --git a/src/targets/qwen3_6_35b_a3b/impl/variant.cpp b/src/targets/qwen3_6_35b_a3b/impl/variant.cpp index 96ab3ac8e6..d3ce82588d 100644 --- a/src/targets/qwen3_6_35b_a3b/impl/variant.cpp +++ b/src/targets/qwen3_6_35b_a3b/impl/variant.cpp @@ -14,6 +14,10 @@ #include "targets/qwen3_6/impl/runtime/instantiate.h" namespace ninfer::targets::qwen3_6_35b_a3b::detail { +std::array Variant::default_layer_kv_dtypes(WeightsProfile) { + return {}; // no per-layer calibration prior for this target +} + namespace { std::vector diff --git a/src/targets/qwen3_6_35b_a3b/impl/variant.h b/src/targets/qwen3_6_35b_a3b/impl/variant.h index c6802e43f4..21f7c4acfe 100644 --- a/src/targets/qwen3_6_35b_a3b/impl/variant.h +++ b/src/targets/qwen3_6_35b_a3b/impl/variant.h @@ -33,6 +33,8 @@ struct Variant { static constexpr std::uint32_t maximum_dflash_draft_tokens = kMaximumDFlashDraftTokens; static constexpr std::uint32_t maximum_context = kNativeContext; static constexpr bool supports_dflash = DFlashConfig::supported; + static constexpr bool supports_per_layer_kv_defaults = false; + [[nodiscard]] static std::array default_layer_kv_dtypes(WeightsProfile profile); static constexpr std::int32_t draft_head_rows = 131072; [[nodiscard]] static std::vector