Skip to content

Allow configurable NPU KV cache initialization value via model metadata. #2776

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Jul 10, 2026
+104 −40
Merged

Allow configurable NPU KV cache initialization value via model metadata. #2776

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
123 changes: 89 additions & 34 deletions runtime/executor/llm_litert_npu_compiled_model_executor.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,28 @@ using ::litert::TensorBuffer;

constexpr int kInvalidTokenId = -1;

absl::Status FillKVCacheBuffer(TensorBuffer& buffer, int64_t init_value) {
LITERT_ASSIGN_OR_RETURN(RankedTensorType tensor_type, buffer.TensorType());
LITERT_ASSIGN_OR_RETURN(auto size, buffer.PackedSize());
LITERT_ASSIGN_OR_RETURN(
auto lock, ::litert::TensorBufferScopedLock::Create(
buffer, ::litert::TensorBuffer::LockMode::kWrite));

auto element_type = tensor_type.ElementType();
if (element_type == ::litert::ElementType::Int16) {
auto* ptr = static_cast<int16_t*>(lock.second);
std::fill(ptr, ptr + size / sizeof(int16_t),
static_cast<int16_t>(init_value));
} else if (element_type == ::litert::ElementType::UInt16) {
auto* ptr = static_cast<uint16_t*>(lock.second);
std::fill(ptr, ptr + size / sizeof(uint16_t),
static_cast<uint16_t>(init_value));
} else {
std::memset(lock.second, 0, size);
}
return absl::OkStatus();
}

constexpr char kPrefillSignature[] = "prefill_128";
constexpr int kPrefillSize = 128;
constexpr char kDecodeSignature[] = "decode";
Expand Down Expand Up @@ -292,6 +314,37 @@ litert::Expected<bool> HasPerLayerEmbedder(
return false;
}

int64_t GetKvCacheInitValue(ModelResources& resources) {
int64_t kv_cache_init_value = 0;
if (auto metadata_status = resources.GetLlmMetadata(); metadata_status.ok()) {
const proto::LlmMetadata* metadata = *metadata_status;
if (metadata && metadata->has_kv_cache_init_value()) {
kv_cache_init_value = metadata->kv_cache_init_value();
}
}
return kv_cache_init_value;
}

absl::Status ClearKVCacheToZero(
absl::flat_hash_map<absl::string_view, ::litert::TensorBuffer>& buffers) {
for (auto& [buffer_name, buffer] : buffers) {
if (buffer_name.starts_with(kv_cache_k_root_name) ||
buffer_name.starts_with(kv_cache_v_root_name) ||
buffer_name.starts_with(kv_cache_c_root_name)) {
auto status = buffer.Clear();
if (!status) {
LITERT_ASSIGN_OR_RETURN(
auto lock_and_addr,
::litert::TensorBufferScopedLock::Create(
buffer, ::litert::TensorBuffer::LockMode::kWrite));
LITERT_ASSIGN_OR_RETURN(size_t size, buffer.Size());
std::memset(lock_and_addr.second, 0, size);
}
}
}
return absl::OkStatus();
}

} // namespace

std::ostream& operator<<(
Expand Down Expand Up @@ -888,7 +941,8 @@ absl::Status LlmLiteRtNpuCompiledModelExecutor::AllocateTransformerBuffers(
decode_output_kv_cache_slice_buffers,
absl::flat_hash_map<absl::string_view, ::litert::TensorBuffer>&
verify_output_kv_cache_slice_buffers,
absl::flat_hash_map<absl::string_view, HWQuantParams>& kv_quant_params) {
absl::flat_hash_map<absl::string_view, HWQuantParams>& kv_quant_params,
int64_t kv_cache_init_value) {
auto prefill_signature = transformer_model->FindSignature(kPrefillSignature);

if (prefill_signature.HasValue()) {
Expand Down Expand Up @@ -917,7 +971,8 @@ absl::Status LlmLiteRtNpuCompiledModelExecutor::AllocateTransformerBuffers(
LITERT_ASSIGN_OR_RETURN(
input_kv_cache_buffers[input_name],
llm_compiled_model.CreateInputBuffer(kPrefillSignature, input_name));
input_kv_cache_buffers[input_name].Clear();
LITERT_RETURN_IF_ERROR(FillKVCacheBuffer(
input_kv_cache_buffers[input_name], kv_cache_init_value));
} else {
LITERT_ASSIGN_OR_RETURN(
gemma_prefill_input_buffers[input_name],
Expand All @@ -938,7 +993,8 @@ absl::Status LlmLiteRtNpuCompiledModelExecutor::AllocateTransformerBuffers(
LITERT_ASSIGN_OR_RETURN(
input_kv_cache_buffers[input_name],
llm_compiled_model.CreateInputBuffer(kDecodeSignature, input_name));
input_kv_cache_buffers[input_name].Clear();
LITERT_RETURN_IF_ERROR(FillKVCacheBuffer(
input_kv_cache_buffers[input_name], kv_cache_init_value));
}
continue;
}
Expand Down Expand Up @@ -1473,10 +1529,11 @@ absl::Status LlmLiteRtNpuCompiledModelExecutor::WarmupInference(

// Clear the KV cache buffers after warmup.
ABSL_RETURN_IF_ERROR(
ClearKVCache(llm_inference_context.prefill_input_buffers));
ClearKVCacheToZero(llm_inference_context.prefill_input_buffers));
return absl::OkStatus();
}


absl::Status LlmLiteRtNpuCompiledModelExecutor::WarmupDrafterInference(
const DrafterContext& drafter_context,
const DrafterAuxContext& drafter_aux_context) {
Expand Down Expand Up @@ -3399,6 +3456,7 @@ LlmLiteRtNpuCompiledModelExecutor::CreateForModelHasPerLayerEmbedding(
const LlmExecutorSettings& executor_settings, ModelResources& resources,
litert::Environment& env, const litert::Model* transformer_model,
LogitsQuantizationParams quantization_params) {
int64_t kv_cache_init_value = GetKvCacheInitValue(resources);
// If the model is fully AOT compiled for NPU, NPU accelerator is used
// automatically.
LITERT_ASSIGN_OR_RETURN(auto options,
Expand Down Expand Up @@ -3426,25 +3484,26 @@ LlmLiteRtNpuCompiledModelExecutor::CreateForModelHasPerLayerEmbedding(
verify_output_kv_cache_slice_buffers;

absl::flat_hash_map<absl::string_view, HWQuantParams> kv_quant_params;
ABSL_RETURN_IF_ERROR(AllocateTransformerBuffers(
LITERT_RETURN_IF_ERROR(AllocateTransformerBuffers(
env, transformer_model, llm_compiled_model, gemma_prefill_input_buffers,
gemma_decode_input_buffers, gemma_verify_input_buffers,
input_kv_cache_buffers, prefill_output_kv_cache_slice_buffers,
decode_output_kv_cache_slice_buffers,
verify_output_kv_cache_slice_buffers, kv_quant_params));
verify_output_kv_cache_slice_buffers, kv_quant_params,
kv_cache_init_value));

// Gemma3n specific fix: KV cache buffer 19 of *prefill* is not connected
// to any OPs in the model, making the LiteRT runtime allocate host memory
// for it. This is incompatible when running the transformer model on the NPU.
if (input_kv_cache_buffers.contains(cache_k19)) {
LITERT_ASSIGN_OR_RETURN(auto buffer_k, llm_compiled_model.CreateInputBuffer(
kDecodeSignature, cache_k19));
buffer_k.Clear();
LITERT_RETURN_IF_ERROR(FillKVCacheBuffer(buffer_k, kv_cache_init_value));
input_kv_cache_buffers[cache_k19] = std::move(buffer_k);

LITERT_ASSIGN_OR_RETURN(auto buffer_v, llm_compiled_model.CreateInputBuffer(
kDecodeSignature, cache_v19));
buffer_v.Clear();
LITERT_RETURN_IF_ERROR(FillKVCacheBuffer(buffer_v, kv_cache_init_value));
input_kv_cache_buffers[cache_v19] = std::move(buffer_v);
}
LITERT_ASSIGN_OR_RETURN(
Expand Down Expand Up @@ -3537,11 +3596,6 @@ LlmLiteRtNpuCompiledModelExecutor::CreateForModelHasPerLayerEmbedding(
std::move(prefill_valid_mask), std::move(decode_valid_mask),
std::move(verify_valid_mask)));

ABSL_RETURN_IF_ERROR(WarmupInference(
llm_compiled_model, llm_inference_context,
npu_auxiliary_context.npu_auxiliary_compiled_model, rope_context,
mask_context, cache_update_inference_context));

// For now we only support one prefill length in the model.
SortedPrefillSignatureMap prefill_runner_set;
prefill_runner_set[kPrefillSize] = kPrefillSignature;
Expand Down Expand Up @@ -3749,6 +3803,11 @@ LlmLiteRtNpuCompiledModelExecutor::CreateForModelHasPerLayerEmbedding(
}
}

ABSL_RETURN_IF_ERROR(WarmupInference(
llm_compiled_model, llm_inference_context,
npu_auxiliary_context.npu_auxiliary_compiled_model, rope_context,
mask_context, cache_update_inference_context));

auto executor = absl::WrapUnique(new LlmLiteRtNpuCompiledModelExecutor(
executor_settings, env, std::move(embedder_context),
std::move(npu_auxiliary_context), std::move(mask_context),
Expand All @@ -3761,9 +3820,9 @@ LlmLiteRtNpuCompiledModelExecutor::CreateForModelHasPerLayerEmbedding(
std::move(ple_table_ptrs), std::move(ple_quant_params),
std::move(ple_per_tensor_scales), table_count, ple_embedding_dim_val,
output_type, ple_table_element_type, mul_scale, output_scale,
final_zero_point, std::move(kv_quant_params), speculative_decoding_type,
std::move(drafter_context), std::move(drafter_aux_context),
embedder_per_layer_model));
final_zero_point, std::move(kv_quant_params), kv_cache_init_value,
speculative_decoding_type, std::move(drafter_context),
std::move(drafter_aux_context), embedder_per_layer_model));
return executor;
}

Expand All @@ -3772,6 +3831,7 @@ LlmLiteRtNpuCompiledModelExecutor::CreateForModelWithoutPerLayerEmbedding(
const LlmExecutorSettings& executor_settings, ModelResources& resources,
litert::Environment& env, const litert::Model* transformer_model,
LogitsQuantizationParams quantization_params) {
int64_t kv_cache_init_value = GetKvCacheInitValue(resources);
// Set up LiteRt options.
LITERT_ASSIGN_OR_RETURN(auto options,
CreateLiteRtNpuOptions(executor_settings));
Expand Down Expand Up @@ -3803,7 +3863,8 @@ LlmLiteRtNpuCompiledModelExecutor::CreateForModelWithoutPerLayerEmbedding(
gemma_decode_input_buffers, gemma_verify_input_buffers,
input_kv_cache_buffers, prefill_output_kv_cache_slice_buffers,
decode_output_kv_cache_slice_buffers,
verify_output_kv_cache_slice_buffers, kv_quant_params));
verify_output_kv_cache_slice_buffers, kv_quant_params,
kv_cache_init_value));
LITERT_ASSIGN_OR_RETURN(
auto llm_inference_context,
CreateLlmInferenceContextWithBufferSharing(
Expand Down Expand Up @@ -3924,11 +3985,6 @@ LlmLiteRtNpuCompiledModelExecutor::CreateForModelWithoutPerLayerEmbedding(
std::move(prefill_valid_mask), std::move(decode_valid_mask),
std::move(verify_valid_mask)));

ABSL_RETURN_IF_ERROR(WarmupInference(
llm_compiled_model, llm_inference_context,
npu_auxiliary_context.npu_auxiliary_compiled_model, rope_context,
mask_context, cache_update_inference_context));

// For now we only support one prefill length in the model.
SortedPrefillSignatureMap prefill_runner_set;
prefill_runner_set[kPrefillSize] = kPrefillSignature;
Expand Down Expand Up @@ -3983,6 +4039,11 @@ LlmLiteRtNpuCompiledModelExecutor::CreateForModelWithoutPerLayerEmbedding(
}
}

ABSL_RETURN_IF_ERROR(WarmupInference(
llm_compiled_model, llm_inference_context,
npu_auxiliary_context.npu_auxiliary_compiled_model, rope_context,
mask_context, cache_update_inference_context));

auto executor = absl::WrapUnique(new LlmLiteRtNpuCompiledModelExecutor(
executor_settings, env, std::move(embedder_context),
std::move(npu_auxiliary_context), std::move(mask_context),
Expand All @@ -3993,26 +4054,20 @@ LlmLiteRtNpuCompiledModelExecutor::CreateForModelWithoutPerLayerEmbedding(
/*per_layer_embedding_lookup_manager=*/nullptr,
/*embedder_per_layer_context=*/std::nullopt, quantization_params, {}, {},
{}, 0, 0, litert::ElementType::None, litert::ElementType::None, 1.0f,
1.0f, 0, std::move(kv_quant_params), speculative_decoding_type,
std::move(drafter_context), std::move(drafter_aux_context)));
1.0f, 0, std::move(kv_quant_params), kv_cache_init_value,
speculative_decoding_type, std::move(drafter_context),
std::move(drafter_aux_context)));
return executor;
}

absl::Status LlmLiteRtNpuCompiledModelExecutor::ClearKVCache(
absl::flat_hash_map<absl::string_view, ::litert::TensorBuffer>& buffers) {
absl::flat_hash_map<absl::string_view, ::litert::TensorBuffer>& buffers)
const {
for (auto& [buffer_name, buffer] : buffers) {
if (buffer_name.starts_with(kv_cache_k_root_name) ||
buffer_name.starts_with(kv_cache_v_root_name) ||
buffer_name.starts_with(kv_cache_c_root_name)) {
auto status = buffer.Clear();
if (!status) {
LITERT_ASSIGN_OR_RETURN(
auto lock_and_addr,
::litert::TensorBufferScopedLock::Create(
buffer, ::litert::TensorBuffer::LockMode::kWrite));
LITERT_ASSIGN_OR_RETURN(size_t size, buffer.Size());
std::memset(lock_and_addr.second, 0, size);
}
LITERT_RETURN_IF_ERROR(FillKVCacheBuffer(buffer, kv_cache_init_value_));
}
}
return absl::OkStatus();
Expand Down
15 changes: 9 additions & 6 deletions runtime/executor/llm_litert_npu_compiled_model_executor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -343,6 +343,7 @@ class LlmLiteRtNpuCompiledModelExecutor : public LlmExecutor {
int32_t final_zero_point = 0,
absl::flat_hash_map<absl::string_view, HWQuantParams> kv_quant_params =
{},
int64_t kv_cache_init_value = 0,
SpeculativeDecodingType speculative_decoding_type =
SpeculativeDecodingType::kNone,
std::optional<DrafterContext> drafter_context = std::nullopt,
Expand Down Expand Up @@ -379,7 +380,8 @@ class LlmLiteRtNpuCompiledModelExecutor : public LlmExecutor {
drafter_aux_context_(std::move(drafter_aux_context)),
embedder_per_layer_model_(embedder_per_layer_model),
per_tensor_logits_scale_(quantization_params.scale),
per_tensor_logits_zero_point_(quantization_params.zero_point) {
per_tensor_logits_zero_point_(quantization_params.zero_point),
kv_cache_init_value_(kv_cache_init_value) {
auto npu_config_status = executor_settings_.GetBackendConfig<NpuConfig>();
if (npu_config_status.ok()) {
npu_config_ = *npu_config_status;
Expand Down Expand Up @@ -583,8 +585,6 @@ class LlmLiteRtNpuCompiledModelExecutor : public LlmExecutor {
::litert::TensorBuffer prefill_valid_mask,
::litert::TensorBuffer decode_valid_mask,
::litert::TensorBuffer verify_valid_mask);
// Run a 'warmup' inference on every model (prefill and decode). This is
// intended to be called before the first actual inference.
static absl::Status WarmupInference(
::litert::CompiledModel& compiled_model_llm,
InferenceContext& llm_inference_context,
Expand All @@ -601,8 +601,9 @@ class LlmLiteRtNpuCompiledModelExecutor : public LlmExecutor {

// Clears all buffers in the provided 'buffers' map that belong to the KV
// cache.
static absl::Status ClearKVCache(
absl::flat_hash_map<absl::string_view, ::litert::TensorBuffer>& buffers);
absl::Status ClearKVCache(
absl::flat_hash_map<absl::string_view, ::litert::TensorBuffer>& buffers)
const;

bool UseEmbeddingLookupManager() const {
return embedding_lookup_manager_ != nullptr;
Expand Down Expand Up @@ -640,7 +641,8 @@ class LlmLiteRtNpuCompiledModelExecutor : public LlmExecutor {
decode_output_kv_cache_slice_buffers,
absl::flat_hash_map<absl::string_view, ::litert::TensorBuffer>&
verify_output_kv_cache_slice_buffers,
absl::flat_hash_map<absl::string_view, HWQuantParams>& kv_quant_params);
absl::flat_hash_map<absl::string_view, HWQuantParams>& kv_quant_params,
int64_t kv_cache_init_value);

// Create the executor for Gemma3n, with multi-modality support.
static absl::StatusOr<std::unique_ptr<LlmLiteRtNpuCompiledModelExecutor>>
Expand Down Expand Up @@ -726,6 +728,7 @@ class LlmLiteRtNpuCompiledModelExecutor : public LlmExecutor {
// Tracks whether a decode step was run so we know how to update the logits
// processor state.
bool ran_decode_ = false;
int64_t kv_cache_init_value_ = 0;
};

std::ostream& operator<<(
Expand Down
6 changes: 6 additions & 0 deletions runtime/proto/llm_metadata.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,4 +90,10 @@ message LlmMetadata {

// The suppress tokens for the model.
TokenIds suppress_tokens = 9;

// The value to initialize the KV cache with, to respect underlying backend
// optimizations. Defaults to 0.
// Note: This field is curently only used by NPU backends for certain models,
// proceed with caution.
optional int64 kv_cache_init_value = 10;
}
Loading