Enhance NPU optimizations with KV cache slicing and buffer management - #314
Open
zhaixuejun1993 wants to merge 15 commits into
Open
Enhance NPU optimizations with KV cache slicing and buffer management#314zhaixuejun1993 wants to merge 15 commits into
zhaixuejun1993 wants to merge 15 commits into
Conversation
Add GGML_OPENVINO_NPU_KV_SLICE as an opt-in NPU prefill optimization. When enabled, static prefill keeps the graph's attention dimension at the active prompt/KV length instead of expanding it to the full context capacity, and the KV-cache input/output tensors are bound with a shortened context axis over the same backing storage. This reduces the amount of KV cache and attention-mask data imported by the NPU plugin and lets the prefill graph plan smaller attention buffers for early prompt chunks. For example, with a 4096-token context and a 256-token prefill chunk, the first chunk can expose KV/mask dimensions at 256 instead of 4096, avoiding work over padded cache rows. The trade-off is graph stability: NPU static shapes now depend on the aligned active attention length for prefill, so a cached prefill model is reused only when the attention sizes match. Decode keeps a fixed full-context graph to avoid recompiling on every generated token, and SWA/ring-cache cases still fall back to full-size binding when the live prefix is not a simple leading slice.
Add GGML_OPENVINO_NPU_L0_HOST_TENSORS as an opt-in NPU buffer allocation mode. When enabled, non-remote OpenVINO backend buffers on NPU are allocated through the NPU default context with create_host_tensor(), and ggml tensor data points directly at that OpenVINO-owned Level Zero host allocation. The optimization targets the data path rather than the graph math: KV-cache and compute tensors can be passed back to the NPU plugin as OpenVINO tensors backed by importable Level Zero host memory. That avoids an extra ordinary-host-memory-to-plugin-staging copy on each infer boundary, which is especially useful for prefill where large KV/cache buffers dominate input and output traffic. If Level Zero host allocation fails, the backend logs a warning and falls back to the existing allocation path. The destructor tracks OpenVINO-owned storage so it does not free memory owned by the host tensor. The trade-off is that this is NPU-specific, consumes Level Zero host allocation resources, and improves transfer/import overhead without changing attention compute complexity or tensor shapes.
Add GGML_OPENVINO_NPU_FAST_MASK as an opt-in NPU static prefill input-packing optimization. When enabled, the prefill attention mask is built directly in the OpenVINO input tensor instead of first materializing a padded std::vector and then copying that whole buffer into the tensor. The new fill_prefill_mask() helper writes each row once: it copies the valid mask prefix, fills padding columns with -inf, handles padded rows, and restores the causal diagonal to zero. This keeps the final mask numerically identical to the existing pad_input() plus set_zero_diagonal() path while removing one intermediate allocation and one chunk_size-by-context_size host copy per prefill chunk. The optimization targets CPU-side staging overhead, not attention math. It composes with GGML_OPENVINO_NPU_KV_SLICE: KV slice reduces the mask dimensions, while FAST_MASK reduces the work needed to construct whatever mask shape remains. The trade-off is a second mask construction path to maintain, with limited benefit for very small prompts or contexts.
Add GGML_OPENVINO_NPU_REQUANT_POLICY as the NPU-side selector for the default weight requantization layout. The existing NPU behavior remains the default: unset or group-128 maps quantized weights to Q4_0_128, while channel-wise selects the experimental Q4_0_C layout. Unknown policy values fail fast during model loading instead of silently compiling a graph with unintended weight packing. The optimization target is NPU weight bandwidth and compressed-weight code generation. Q4_0_128 keeps weights in a 4-bit symmetric layout with one scale per 128 values, reducing both payload bytes and scale metadata versus smaller groups while preserving a more conservative quantization granularity than channel-wise. Q4_0_C uses one scale over the innermost dimension, which can further shrink metadata and simplify dequant broadcast, but may cost accuracy because the quantization block is much larger. Include the selected NPU requant policy in the OpenVINO model cache discriminator. The same cgraph compiled with group-128 and channel-wise has different weight constants and dequant shapes, so reusing a cached decoder across policies would be incorrect. This keeps cache reuse stable within a policy while forcing recompilation when the layout policy changes. Trade-off: policy selection gives a performance/accuracy knob for NPU experiments, but all options still pay the load/compile-time requantization cost. The default remains group-128 for a safer balance of bandwidth, metadata, and accuracy; channel-wise is intentionally documented as experimental.
Add GGML_OPENVINO_NPU_CONFIG as the recommended catch-all override for OpenVINO NPU plugin and compiler properties. The value is parsed as comma-separated KEY=VALUE pairs and applied after the built-in defaults and named NPU aliases, so advanced users can override any exposed property without adding a new llama.cpp environment variable for every plugin tuning knob. Keep the common knobs as named aliases for discoverability and compatibility: GGML_OPENVINO_NPU_COMPILER_TYPE maps to NPU_COMPILER_TYPE, GGML_OPENVINO_NPUW_FUNCALL_FOR_ALL maps to NPUW_FUNCALL_FOR_ALL, GGML_OPENVINO_NPUW_UNFOLD_IREQS maps to NPUW_UNFOLD_IREQS, and GGML_OPENVINO_COMPILATION_NUM_THREADS maps to COMPILATION_NUM_THREADS. GGML_OPENVINO_NPU_COMPILE_CONFIG is retained as a legacy alias for NPU_COMPILATION_MODE_PARAMS, while the documentation now recommends expressing it through GGML_OPENVINO_NPU_CONFIG. This keeps the public tuning surface compact while preserving high-signal aliases for settings with known performance or stability implications. NPU_COMPILER_TYPE can select the driver compiler, which currently generates faster prefill kernels on the tested NPU stack. NPUW_FUNCALL_FOR_ALL and NPUW_UNFOLD_IREQS expose NPUW graph lowering trade-offs around dispatch overhead, memory use, and long-context stability. COMPILATION_NUM_THREADS gives a way to cap compiler worker parallelism to reduce peak host memory during large graph compilation. Trade-off: a free-form KEY=VALUE string is less type-safe than one env var per option and invalid plugin properties will only be diagnosed by the OpenVINO stack. Applying it last is intentional: it gives one escape hatch for bisecting driver/compiler behavior and for testing new plugin options without rebuilding llama.cpp, while the named aliases continue to document the most important performance knobs.
Add GGML_OPENVINO_KV_SCATTER_ELEMENTS as an opt-in lowering for non-stateful single-row KV cache writes. When the flag is enabled and the set_rows indices are the simple decode case, translate_set_rows uses ScatterElementsUpdate with broadcasted rank-4 indices instead of ScatterUpdate. The motivation is NPU decode performance. The default ScatterUpdate path updates one KV row but the NPU plugin does not perform that update in place; it effectively copies the full destination tensor to write a small slice. As context length grows, that full-tensor movement dominates the per-token KV update cost. ScatterElementsUpdate expresses the same single-row write at element granularity and is measurably cheaper for this shape. Keep the optimization guarded. Stateful execution keeps its existing concat/update path, and multidimensional indices already use the broader ScatterElementsUpdate lowering. The new flag only changes the common non-stateful decode write, so the default behavior and graph shape remain unchanged unless the user opts in. Trade-off: ScatterElementsUpdate introduces extra ShapeOf, Reshape, and Broadcast nodes to build indices matching the update tensor. That overhead is worthwhile only when it avoids the NPU plugin's full-destination ScatterUpdate copy, so the path is controlled by an environment variable rather than becoming the unconditional lowering.
Add NPU token-embedding controls for models whose token_embd.weight would otherwise be widened from Q6_K to FP16. GGML_OPENVINO_TOKEN_EMBD_I8 keeps the embedding table on the normal int8 compressed-weight path, while GGML_OPENVINO_TOKEN_EMBD_I4 requantizes it to the group-128 int4 layout used by the NPU weight path. The performance target is tied embeddings and large vocabulary tables. When token_embd.weight is also consumed by the lm_head matmul, a normal Gather from the dequantized table gives the dequantization subgraph a second consumer and can force the full table to be materialized. The new gather_compressed_rows helper recognizes the dequantization chain emitted for int4/int8 constants, gathers only the selected packed rows from the leaf constants, and rebuilds the convert/subtract/multiply/reshape chain on those rows. This reduces memory pressure and avoids dequantizing or materializing a full embedding table when inference only needs the token rows for the current input. The helper is deliberately conservative: it only rewrites static 2D tables above a size threshold, walks only the known compressed-weight node patterns, limits recursion depth, and falls back to the original Gather when the subgraph is not recognized. Trade-off: int8 and especially int4 token embeddings can change accuracy and may not be faster for small untied embedding tables, because the rewrite adds extra Gather/Reshape nodes and only pays off when it avoids a large full-table dequantization. The default behavior is unchanged; the NPU-specific I8/I4 paths are opt-in through environment variables.
Extend the GGML_OPENVINO_REDUCE_COMPILE_MEM streaming requantization path to Q4_0/u4 targets. The previous streaming path intentionally skipped u4 because Q4_0 packs two weights per byte and the zero-point tensor stores two 4-bit values per byte, so the original quantize_q4_0 helper wrote as if every call started at block zero. Teach quantize_q4_0 to accept a destination block offset and use the absolute block index for weights, scales, and zero-point writes. Streaming requantization emits whole-row chunks in increasing order, and the target block size is required to divide the row width, so no target block straddles a chunk boundary. That keeps packed weight bytes aligned and preserves the even-block assignment before the odd-block zero-point OR for each shared zero-point byte. With that indexing in place, reduce-compile-memory mode can dequantize a chunk of rows into the scratch buffer and immediately quantize it into the final u4 buffers instead of materializing the entire tensor as temporary F32. This extends the peak-memory reduction already used by Q8/F16 targets to the NPU group-128/channel u4 layouts used for compressed weights. Trade-off: the u4 path is more sensitive to block alignment than Q8/F16 because of nibble packing, so streaming remains guarded by the same row-divisibility check and only runs when GGML_OPENVINO_REDUCE_COMPILE_MEM is enabled. When the flag is off, the full-materialization path and quantized output layout remain unchanged.
Extend the GGML_OPENVINO_RELEASE_WEIGHTS safety check to the static graph cache-miss path. Once host weight buffers have been released, any later compile would read invalid/zeroed host pages and bake the wrong constants into a new OpenVINO model. The dynamic path already rejected this situation. Static execution can still miss the decoder cache when graph shapes or model parameters change, so it needs the same guard before erasing cached infer requests and rebuilding prefill/decode models. This is a correctness and diagnosability change rather than a throughput optimization. It preserves the memory-saving release-weights mode for stable graph shapes, but turns unsupported recompilation into an explicit abort instead of silently producing a corrupted compiled model.
Add GGML_OPENVINO_NPU_KEEP_Q4_0 as an opt-in escape hatch for NPU weight handling. When enabled for native Q4_0 tensors, the NPU requant path returns nullopt and keeps the original Q4_0 extraction instead of regrouping the tensor through the selected NPU requant policy. The motivation is to avoid unnecessary work for weights that are already stored as 4-bit values. The default NPU policy maps quantized weights to layouts such as Q4_0_128, which can reduce scale metadata but requires round-tripping Q4_0 through the requantization path. Keeping Q4_0 native can save that conversion work and preserve the original group-32 packing. Keep the option experimental and disabled by default. Current NPU driver/plugin stacks can compile the resulting group-32 graph, but inference may hang or fail with ZE_RESULT_ERROR_DEVICE_LOST. The default behavior therefore continues to use the safer NPU requant policy; this flag is mainly useful for debugging driver behavior, measuring the true cost of regrouping, and testing future NPU stacks. Trade-off: native Q4_0 avoids regrouping overhead and keeps the original quantization blocks, but it gives up the more regular group-128 layout expected by the default NPU compressed-weight path and may be unstable on current hardware/software combinations.
Compile the static prefill and decode models with separate OpenVINO config maps so the decode graph can opt into NPUW_UNFOLD_IREQS without forcing the same lowering on prefill. The build_static_model helper now receives the config map explicitly, and the decode path gets a copy of the base config with NPUW_UNFOLD_IREQS=YES when the user has not already set it. The performance target is NPU token generation. Decode executes a long stream of single-token infers, where folded NPUW function-call dispatch overhead becomes visible in every token. Unfolding the NPUW calls into separate infer requests removes that repeated dispatch cost and measured about +20% token-generation throughput on the tested stack, matching the behavior seen in OpenVINO GenAI. Prefill keeps the original folded config. Prompt processing is a batched workload dominated by larger matmuls, and the folded form is faster there while also using less memory. Splitting the config lets prefill and decode use the lowering that matches their workload shape instead of treating NPUW_UNFOLD_IREQS as a single global compromise. Trade-off: unfolding can increase memory and request-management overhead, so the automatic default is limited to the NPU decode model. If a user explicitly provides NPUW_UNFOLD_IREQS through the named env alias or GGML_OPENVINO_NPU_CONFIG, that value is respected because the decode override is only applied when the property is absent.
Collect the remaining small cleanups from the NPU tuning series. This reorders the NPU KV-slice documentation and environment-variable cache entry next to the other NPU opt-in toggles, adds a short header comment for the centralized NPU helper declarations, and makes the static prefill attention-size assignment in the decoder a little more explicit. There is no intended behavior change in this commit. The purpose is to make the final source layout match the split commits that introduced the NPU tuning knobs, so future readers can find the related flags and helpers in one place without mixing that mechanical cleanup into the functional performance commits.
Move the default NPU_COMPILER_TYPE=DRIVER setting into the NPU compile_config initializer so the default NPU plugin options are declared in one place. Route GGML_OPENVINO_NPU_COMPILE_CONFIG and GGML_OPENVINO_NPU_COMPILER_TYPE through the same set_compile_option_from_env helper used by the other named NPU compile option aliases. This keeps the legacy NPU_COMPILE_CONFIG mapping to NPU_COMPILATION_MODE_PARAMS intact while making all optional env-to-compile_config overrides follow the same non-empty-value rule. GGML_OPENVINO_NPU_CONFIG is still parsed last, so explicit comma-separated KEY=VALUE overrides continue to have the highest precedence over both defaults and named aliases.
Default the streaming requantization path on for NPU so the compile-time F32 dequant transient (~1-2 GB for token_embd) is capped without an env flag. Output is byte-identical to full materialization; GGML_OPENVINO_REDUCE_COMPILE_MEM / MEMORY_OPTIMIZE still override (set to 0 to force off). Kept the g_nonov_weight_cache opt-in to avoid holding token_embd resident (which regressed decode steady memory). pp1024_d0: PrefillPeakPrivWS 5974->4664 MiB (-22%), decode steady and pp/tg throughput unchanged.
Wire the frontend compiled-model cache (GGML_OPENVINO_COMPILED_MODEL_CACHE_DIR) into ov_graph_compute_static so the NPU path can skip the plugin compile on a warm start, like the dynamic/GPU path. Prefill and decode each export a blob keyed by the shared graph fingerprint mixed with per-model bits (is_prefill, chunk size, attention sizes, decode NPUW_UNFOLD_IREQS). NPUW blobs are weightless, so import still builds the ov::Model (weight requant + convert) and hands it over via MODEL_PTR, but skips compile. Opt-in; inactive when the env var is unset. Verified on phi-4-mini NPU: compile 15158 ms -> 268 ms on a cache hit (~15s faster start), identical coherent output.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces a set of advanced configuration options and performance optimizations for OpenVINO NPU (Neural Processing Unit) support, along with corresponding documentation updates and code refactoring. The changes focus on enabling fine-grained NPU tuning, new memory and compilation strategies, and additional quantization controls, all exposed via new environment variables.
Key changes include:
New NPU configuration and tuning options:
GGML_OPENVINO_NPU_CONFIGand related environment variables, allowing users to override NPU plugin/compiler options with comma-separatedKEY=VALUEpairs for advanced tuning and debugging. This includes options for compiler type, function call folding, thread count, and more. (docs/backend/OPENVINO.md,ggml-openvino-extra.cpp,ggml-openvino-extra.h) [1] [2] [3]docs/backend/OPENVINO.md,ggml-openvino-extra.cpp,ggml-openvino-extra.h) [1] [2] [3]Memory and buffer management enhancements:
ggml-openvino.cpp,ggml-openvino-extra.cpp,ggml-openvino-extra.h) [1] [2] [3]NPU quantization and weight handling improvements:
GGML_OPENVINO_NPU_REQUANT_POLICYto control NPU weight requantization layout (e.g.,group-128or experimentalchannel-wise), and new toggles for keeping native Q4_0 weights or using int8/int4 compressed embeddings. (docs/backend/OPENVINO.md,ggml-openvino-extra.cpp,ggml-openvino-extra.h) [1] [2] [3]KV-cache and attention optimizations:
ggml-decoder.cpp,ggml-openvino-extra.cpp,ggml-openvino-extra.h) [1] [2] [3] [4] [5]Documentation and environment variable registration:
OPENVINO.mdwith detailed descriptions and recommendations for all new environment variables and options, clarifying legacy aliases and preferred usage. Registered all new variables in the device config initialization. (docs/backend/OPENVINO.md,ggml-openvino-extra.cpp) [1] [2] [3] [4]These changes provide advanced users and developers with much greater control over OpenVINO NPU execution, facilitating performance tuning, troubleshooting, and experimentation with new hardware features.## Overview
Additional information
Requirements