performance: allow user-adjustable kvcache_num_blocks estimation

[rebel-jongho]

Pull Request Description

⚠️ Important: Branch Target

• New features, enhancements, and non-critical fixes: Merge to dev branch
• Critical hotfixes only: Merge to main branch (must also merge to dev)

Please ensure you've selected the correct base branch before submitting!

Type of Change

• Release (dev → main merge for production release)
• New Model Support
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Performance improvement
• Code refactoring
• Other (please describe):

Changes Overview

Add an optional kvcache_num_blocks_adjuster callable that lets users override the auto-estimated kvcache_num_blocks at compile time.

• RBLNDecoderOnlyModelConfig gains a kvcache_num_blocks_adjuster: Optional[Callable[[int, RBLNModelConfig], int]] field (defaults to None). It is registered in subclass_non_save_attributes so it is not serialized into rbln_config.json.
• RBLNDecoderOnlyFlashAttentionMixin.set_kvcache_num_blocks_after_compilation accepts an optional adjuster argument. When provided, it is called as adjuster(estimated, rbln_config) and the return value replaces the estimate.
• The decoder-only and Gemma4 compile paths forward rbln_config.kvcache_num_blocks_adjuster into that argument.
• RBLNQwen3VLForConditionalGenerationConfig and RBLNQwen3VLModelConfig override subclass_non_save_attributes; their lists are extended with "kvcache_num_blocks_adjuster" so the inherited exclusion is not shadowed.

Validation policy for the adjusted value:

• A non-positive or non-int return value raises ValueError.
• A value below num_min_blocks raises the existing "Memory is not enough" ValueError.
• A value above the estimate emits a warning but is honored (the upper bound is the caller's responsibility).

Motivation and Context

The current estimate_num_kvcache_blocks logic assumes the model owns all available DRAM on the NPU. In deployments where some memory must be reserved for other processes or components, the auto-estimated block count can be too high and lead to OOM at compile or runtime. Users currently have no way to influence the auto-estimated value other than fully overriding kvcache_num_blocks with a manual constant, which loses the benefit of DRAM-based estimation.

The adjuster gives users a hook to scale the estimate to their actual memory budget while still relying on the compiler's DRAM-based estimate as the starting point. It is a compile-time-only knob: when loading a precompiled artifact, the estimation path does not run, so the adjuster is ignored and the persisted kvcache_num_blocks is used.

Note for reviewers / future maintainers

subclass_non_save_attributes is a plain class attribute that is shadowed (not merged) when a subclass redefines it. Any future subclass of RBLNDecoderOnlyModelConfig that overrides subclass_non_save_attributes must also include "kvcache_num_blocks_adjuster", otherwise a callable passed by the user would be serialized into rbln_config.json and crash json.dump.

Related Issues

N/A

Usage Example

def reserve_half(estimated: int, rbln_config) -> int:
    return estimated // 2

model = RBLNLlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-3-8B",
    rbln_config={"kvcache_num_blocks": 0, "kvcache_num_blocks_adjuster": reserve_half},
)

Test Plan

• CPU-side: construct RBLNDecoderOnlyModelConfig(kvcache_num_blocks_adjuster=fn) and confirm the field is stored but absent from _prepare_for_serialization().
• NPU: compile a decoder-only model with kvcache_num_blocks=0 and a reducing adjuster; confirm the resulting rbln_config.kvcache_num_blocks matches the adjuster output and the model runs.
• NPU: compile with an adjuster returning a value above the estimate; confirm a warning is logged and compilation still succeeds when memory allows.
• NPU: compile with an adjuster returning a value below num_min_blocks; confirm the expected ValueError is raised.

Made with Cursor

[rebel-jonghewk]

Can this be adjusted after compilation? For example, say I've already compiled the model and the rbln_compiled_dir exists — can I modify it at the start of inference (e.g., during vLLM startup)?

[rebel-jongho]

Can this be adjusted after compilation? For example, say I've already compiled the model and the rbln_compiled_dir exists — can I modify it at the start of inference (e.g., during vLLM startup)?

No. with this PR, It can't be done after compilation. This PR is just for pre-compilation.

We need further refactoring to support that fuctionality.

[rebel-jonghewk]

We need further refactoring to support that fuctionality.

@rebel-jongho Are you planning to handle it in a follow-up PR, or include it here? The latter is actually what I need.

[rebel-jongho]

Superseded by #629, which takes a different approach: instead of a compile-time kvcache_num_blocks_adjuster callback, it adds a post-compilation optimum-rbln-cli --get/--set-kvcache-num-blocks that resizes an already-compiled artifact in place (backed by rebel-compiler's new exp_rescale_buffer_size, rebellions-sw/rebel_compiler#12067). Closing in favor of that direction.