model: NPU-aware prefill_chunk_size default (512 on RBLN-CR) - #634
Merged
Conversation
- Guard `_get_npu_name()` returning None with `or ""` so compiling on a host without an attached NPU falls back to 128 instead of raising TypeError on `"RBLN-CR" in None`. - Apply the same NPU-based default (512 on RBLN-CR, else 128) and the divisible-by-64 validation to gemma4, which reimplements `_update_rbln_config` without calling super(). The hardcoded 128 default is dropped from RBLNGemma4ForCausalLMConfig so it resolves at compile time. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Comment on lines
+50
to
+52
| device_name = ( | ||
| rebel.get_npu_name(0) or os.environ.get("RBLN_FORCE_NPU_NAME") or os.environ.get("RBLN_TARGET_SOC") or None | ||
| ) |
Collaborator
There was a problem hiding this comment.
This is not needed. just use rebel.get_npu_name()
rebel-jongho
left a comment
rebel-jongho
left a comment
Collaborator
There was a problem hiding this comment.
how about modifying set_default_values to set the default value of pcs?
Address review feedback: - Drop the `_get_npu_name` helper and its env-var fallbacks; restore `_resolve_npu`/`warn_deprecated_npu` to use `rebel.get_npu_name` directly. - Resolve the NPU-based `prefill_chunk_size` default (and the divisible-by-64 validation) inside `set_default_values`, the shared entry point both decoderonly and gemma4 already route through via `_update_attention_config`, instead of duplicating it in each `_update_rbln_config`. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
set_default_values took only the locally attached device (rebel.get_npu_name(0)) into account, so compiling with an explicitly pinned target NPU (rbln_config.npu) — e.g. targeting RBLN-CR from a host without one — would pick the wrong default. Prefer the pinned npu and fall back to the local device only when it is unset. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Resolve conflict in gemma4/configuration_gemma4.py: dev refactored image-prefill handling to the single `image_prefill_chunk_size` (#605) while this branch was based on the older `image_prefill_chunk_sizes` form. Take dev's version wholesale and re-apply only the PCC change (drop the hardcoded prefill_chunk_size=128 default so it resolves via set_default_values, plus the docstring). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Moving the prefill_chunk_size default into set_default_values (called from _update_attention_config) left it unresolved (None) when validate_sliding_window ran first, crashing with `TypeError: unsupported operand type(s) for -: 'int' and 'NoneType'` for sliding-window models (e.g. Mistral). Run _update_attention_config before the sliding-window block so prefill_chunk_size is set in time. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
rebel-jongho
approved these changes
Jul 16, 2026
rebel-jongho
left a comment
rebel-jongho
left a comment
Collaborator
There was a problem hiding this comment.
LGTM
do we need any tests? or it's just a simple functionality?
# Conflicts: # src/optimum/rbln/transformers/modeling_attention_utils.py # src/optimum/rbln/transformers/models/gemma4/configuration_gemma4.py
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Contributor
Author
|
Resolved the merge conflict with |
Pin npu=RBLN-CR03 in rbln_config, compile without a CR device, and assert the resolved prefill_chunk_size (512) survives save/reload. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The CI runner's compiler (dev371) fails cross-compiling for RBLN-CR03 on a CA25 machine. Keep the wiring test but skip it, matching the existing TestLlamaForCausalLM_fp8 pattern; unit tests still cover the resolution logic. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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
2 participants
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.
What
Make the default
prefill_chunk_sizedepend on the NPU type — 512 on RBLN-CR devices and 128 otherwise. CR NPUs use a larger prefill chunk for better prefill performance (PCC).Changes
set_default_values: Resolve the NPU-basedprefill_chunk_sizedefault and its divisible-by-64 validation here, usingrebel.get_npu_name(0)(guarded withor ""so compiling on a host without an attached NPU falls back to 128 instead of raising on"RBLN-CR" in None). This is the shared entry point that both decoderonly and gemma4 already route through via_update_attention_config, so the default is defined in one place.prefill_chunk_sizedefault/validation from the config constructor;_update_attention_confignow captures the value returned byset_default_values.RBLNGemma4ForCausalLMConfigso it resolves at compile time through the shared path.Notes
prefill_chunk_size or 256) is left as-is — existing behavior preserved.Review feedback addressed
_get_npu_namehelper and itsRBLN_FORCE_NPU_NAME/RBLN_TARGET_SOCenv fallbacks;_resolve_npu/warn_deprecated_npuare back to usingrebel.get_npu_namedirectly.prefill_chunk_sizedefault intoset_default_valuesinstead of duplicating it across each_update_rbln_config.Test
py_compile+ruff checkpass.tests/test_config.py::TestPrefillChunkSizeDefault) covering NPU-aware default resolution (512 on RBLN-CR, 128 otherwise), attached-NPU fallback, no-NPU fallback, explicit-value precedence, and divisible-by-64 validation — 8 passed.test_prefill_chunk_size_npu_wiring_e2e): pinsnpu=RBLN-CR03inrbln_config, compiles tiny-llama without a CR device, asserts the resolvedprefill_chunk_sizeis 512 and survives save/reload. Verified locally (9 passed, ~13s).🤖 Generated with Claude Code