performance: post-compilation kvcache_num_blocks CLI and memory_budget for estimation

[rebel-jongho]

Type of Change

• New feature (non-breaking change which adds functionality)

Changes Overview

Two related capabilities for controlling the decoder-only kv-cache block count:

1. Post-compilation CLI to inspect/resize kvcache_num_blocks on an already-compiled artifact, without recompiling.
2. Compile-time memory_budget to bound the DRAM the automatic block estimation may assume.

1. Post-compilation CLI

Two subcommands on optimum-rbln-cli operating on an artifact directory (--model-id DIR):

• --get-kvcache-num-blocks — prints kvcache_num_blocks from rbln_config.json.
• --set-kvcache-num-blocks N — resizes the kv-cache of every *.rbln to N blocks and writes the .rbln files and rbln_config.json. Edits in place, or with --output-dir writes a full resized copy there (copying the non-.rbln files too) and leaves the source untouched.

# Read the current kv-cache block count
optimum-rbln-cli --get-kvcache-num-blocks --model-id ./my_artifact

# Resize in place to 128 blocks
optimum-rbln-cli --set-kvcache-num-blocks 128 --model-id ./my_artifact

# Resize into a new directory, leaving the source untouched
optimum-rbln-cli --set-kvcache-num-blocks 128 --model-id ./my_artifact --output-dir ./my_artifact_128

Internals — new RBLNDecoderOnlyFlashAttentionMixin.rescale_kvcache_num_blocks(compiled_models, rbln_config, target):

• Saved kv-cache buffers hold per_block * current bytes (current = rbln_config.kvcache_num_blocks). Rescaling by the rational ratio target / current (rebel-compiler exp_rescale_buffer_size) yields exactly per_block * target.
• Stateless: rbln_config.json is the single source of truth for the current block count.
• target < num_min_blocks → ValueError (config only). Non-fatal warnings (config only): target above the compile-time count (may not fit device DRAM), and target above num_full_blocks (excess blocks never used). Older rebel-compiler without the API → clear RuntimeError via hasattr capability detection.

Depends on rebel-compiler's new exp_rescale_buffer_size (rebellions-sw/rebel_compiler#12067).

2. Compile-time memory_budget

RBLNDecoderOnlyModelConfig gains a memory_budget member that caps the device DRAM the automatic kvcache_num_blocks estimation may assume:

• Prefer a float in (0, 1] as a fraction of the NPU available DRAM (e.g. 0.8 for 80%, matching vllm's gpu_memory_utilization). A "80%" string is also accepted (handy for passing through the compile CLI), as are bytes given as an int or string ("10GB", "512MB"). The fraction is of usable memory (after the per-chiplet system reserve), not the raw total.
• Defaults to the full available DRAM and must not exceed it (ValueError otherwise). The budget is split evenly per chiplet with no further reserve subtraction (already excluded), so the default reproduces current estimation exactly.
• Flows through the compile CLI as a forwarded rbln_config key:

optimum-rbln-cli --model-id meta-llama/Llama-2-7b-chat-hf --memory-budget 0.8
optimum-rbln-cli --model-id meta-llama/Llama-2-7b-chat-hf --memory-budget 80%
optimum-rbln-cli --model-id meta-llama/Llama-2-7b-chat-hf --memory-budget 100GB

Also adds RBLNDecoderOnlyFlashAttentionMixin._required_memory_at(compiled_models, rbln_config, num_blocks), returning the device-wide kv-cache DRAM at a given block count with 2MB alignment applied (non-linear in the block count). It works for any current block size — resizable tensors are normalized to per-block bytes using rbln_config.kvcache_num_blocks. The aligned per-chiplet math is factored into _kvcache_bytes_per_chiplet, shared with the block search (unchanged behavior).

Motivation and Context

• Users need to tune kvcache_num_blocks on a compiled artifact without recompiling the whole model (CLI).
• Users need to constrain kv-cache block estimation to a memory budget smaller than the full device (e.g. to leave headroom), expressed intuitively in bytes or a percentage (memory_budget).

How to test

• CLI: on a compiled decoder-only artifact, --get prints the stored count; --set N (N ≥ num_min_blocks) rewrites the kv-cache buffers to per_block * N and updates the config (verified by reloading the .rbln); --set below num_min_blocks raises ValueError; --output-dir leaves the source untouched.
• memory_budget: estimation with a smaller budget yields fewer blocks (monotonic), None reproduces the current result, and a budget above the NPU total raises ValueError. _required_memory_at(N) is monotonic in N, equals the current footprint at N = kvcache_num_blocks, and stays correct for an already-resized artifact.

[rebel-eunji]

How about using a decimal to express a percentage? e.g. using 0.8 to express 80%.
Upstream vllm is using decimal.
https://github.com/vllm-project/vllm/blob/dcfebf93f4eccf30f71872283331eee757915daf/vllm/config/cache.py#L68

[rebel-eunji]

It can also be tricky to validate strictly when we stick with the % representation.

[rebel-jongho]

Thanks for the pointer — I made the float fraction the recommended form (e.g. 0.8), matching vllm's gpu_memory_utilization, and the docstring now steers users to it.

I did keep the "80%" string as an accepted alias, for one reason: it's convenient to pass through the compile CLI, where values arrive as strings (e.g. optimum-rbln-cli --memory-budget 80%). The float form works there too (--memory-budget 0.8).

[rebel-jonghewk]

@rebel-jongho Since Post-compilation CLI depends on the compiler fix, could you split Post-compilation CLI and Compile-time memory_budget into separate PRs? That would make Compile-time memory_budget easier to review.

[rebel-jongho]

Split as requested, thanks @rebel-jonghewk:

• performance: add memory_budget for kv-cache block estimation #637 — Compile-time memory_budget (first PR, no compiler dependency)
• performance: post-compilation CLI to get/set kvcache_num_blocks #638 — Post-compilation CLI (follow-up, depends on rebellions-sw/rebel_compiler#12067)

Closing this in favor of those two.