performance: post-compilation CLI to get/set kvcache_num_blocks

[rebel-jongho]

Type of Change

• Performance improvement

Stacked on #639 (base = jongho/memory-budget-nonsave); auto-retargets to dev once #639 merges.

Changes Overview

Post-compilation subcommands on optimum-rbln-cli operating on an already-compiled artifact directory (--model-id DIR), without recompilation:

• --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.
• --set-memory-budget BUDGET — resizes to the largest block count that fits BUDGET (a float fraction of the NPU available DRAM like 0.8, a "80%" string, or bytes like "10GB"). Mutually exclusive with --set-kvcache-num-blocks.

set edits in place, or with --output-dir writes a full resized copy there (copying the non-.rbln files too) and leaves the source untouched.

optimum-rbln-cli --get-kvcache-num-blocks --model-id ./my_artifact
optimum-rbln-cli --set-kvcache-num-blocks 128 --model-id ./my_artifact
optimum-rbln-cli --set-memory-budget 0.8 --model-id ./my_artifact
optimum-rbln-cli --set-memory-budget 10GB --model-id ./my_artifact --output-dir ./my_artifact_10gb

Internals:

• RBLNDecoderOnlyFlashAttentionMixin.rescale_kvcache_num_blocks(compiled_models, rbln_config, target): saved 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 source of truth.
• --set-memory-budget computes the target via estimate_num_kvcache_blocks(..., current_blocks=current). current_blocks (added to estimate_num_kvcache_blocks / _search_num_kvcache_blocks, default 1) makes the fit search correct on an already-resized artifact (buffers at current, not 1). The transient budget is not persisted to rbln_config.json (model: memory_budget config semantics (non-save, mutually exclusive with kvcache_num_blocks) #639 marks memory_budget non-save).
• Guards: target < num_min_blocks → ValueError; target above the compile-time count or num_full_blocks → warning; older rebel-compiler without exp_rescale_buffer_size → clear RuntimeError.

Dependency

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

Motivation and Context

Lets users tune kvcache_num_blocks on a compiled artifact without recompiling — either to an explicit count or to fit a memory budget. Split out from #629 per review; the memory_budget estimation change (#637) is already merged into dev.

How to test

On a compiled decoder-only artifact: --get prints the count; --set-kvcache-num-blocks N (N ≥ num_min_blocks) rewrites buffers to per_block * N and updates the config (verified by reloading the .rbln); --set-memory-budget BUDGET computes and sets the fitting count (smaller budget → fewer blocks); infeasible targets (< num_min_blocks) raise ValueError; --output-dir leaves the source untouched.

[rebel-eunji]

1. nit: The --get-kvcache-num-blocks and --set-kvcache-num-blocks options should be mutually exclusive.
2. To ensure atomicity, it would be safer to generate the artifact in a temporary directory and then rename it, rather than reusing an existing destination directory.