[quantization] Implement build_static_inputs for Gemma4 static runtime#828
Conversation
| first_true = int(true_indices[0].item()) | ||
| if first_true > 0 and torch.any(row[:first_true]): | ||
| raise ValueError("Left padding expected but not found") |
There was a problem hiding this comment.
The condition torch.any(row[:first_true]) is always False because first_true is defined as the index of the first True in the row. By definition, everything before first_true is False.
If I understand correctly, for 'left' padding we need to check that rest values after first_true are valid.
We can do it with using following condition:
if not torch.all(row[first_true:]):
raise ValueError("Left padding expected but not found")
There was a problem hiding this comment.
cfg.padding_side is currently dead: StaticGemma4Layout has no such field, the runtime falls back to "right", and manual padding always writes to [:seq_len]. #822 keeps the same behavior. You can support right-padding-only in this step and reject/remove the left option.
if cfg.padding_side != "right":
raise ValueError(
"StaticGemma4Runtime currently supports right padding only"
)
| device: cpu | ||
| prompt: "<|image|>Describe the image." | ||
| verify_steps: 4 | ||
| gen_steps: 16 No newline at end of file |
There was a problem hiding this comment.
| gen_steps: 16 | |
| gen_steps: 16 | |
|
Since this script validates a single fixed prompt/image against a fixed static export profile, keeping Without this check, a stale or incorrect Could we add a validation like the following before replacing the placeholder tokens? raw_input_ids = inputs["input_ids"] # (1, seq_len)
image_mask, video_mask, audio_mask = self.model.model.get_placeholder_mask(
input_ids=raw_input_ids
)
image_positions = torch.nonzero(image_mask[0], as_tuple=True)[0]
if image_positions.numel() == 0:
raise ValueError("No image placeholder tokens were found in processor output")
actual_start = int(image_positions[0].item())
actual_count = int(image_positions.numel())
expected_positions = torch.arange(
actual_start,
actual_start + actual_count,
device=image_positions.device,
)
if not torch.equal(image_positions, expected_positions):
raise ValueError(
"Image placeholder tokens must form one contiguous span. "
f"positions={image_positions.tolist()}"
)
if actual_start != self.layout.visual_start_idx:
raise ValueError(
"Processor output does not match the static export profile: "
f"expected visual_start_idx={self.layout.visual_start_idx}, "
f"actual={actual_start}"
)
if actual_count != self.layout.num_visual_tokens:
raise ValueError(
"Processor output does not match the static export profile: "
f"expected num_visual_tokens={self.layout.num_visual_tokens}, "
f"actual={actual_count}"
)
input_ids = raw_input_ids.squeeze(0) |
|
LGTM |
Adds static input building, side-by-side validation, and helper tests. Co-authored-by: Cline TICO-DCO-1.0-Signed-off-by: d.savchenkov <d.savchenkov@partner.samsung.com>
3 participants
Related issue: #768
Related PR: #822
What
This PR implements
StaticGemma4Runtime.build_static_inputs— the first concrete piece of the Gemma4 E2B static-shape NPU inference runtime. It processes a prompt+image through the HF processor, pads to a fixedmax_seq, and replaces image placeholder tokens withpad_token_idto producellm_input_idssuitable for static NPU export. A side-by-side validation function (verify_step_build_static_inputs) re-derives each sub-step from raw HF processor output and compares against the runtime's output. The runtime entry point (run_static_gemma4_runtime) is wired to run only this validation step, with clear SKIP messages for prefill/decode/generation which are not yet implemented.Why
The Gemma4 E2B static runtime follows the Llama static runtime pattern: CPU owns dynamic/control-heavy orchestration (tokenization, layout validation, padding, mask generation, cache management), while NPU executes fixed-shape quantized tensor-compute subgraphs.
build_static_inputsis the CPU-side entry point that transforms variable-length processor output into the fixed-shape tensor contract the NPU subgraphs expect. Without correctllm_input_idsconstruction (especially the image placeholder replacement), downstream NPU prefill would receive image-token IDs in the embedding lookup, producing silently wrong results. This PR establishes the input-building foundation and its validation harness so that subsequent prefill/decode PRs can build on a verified input contract.Key Design Decisions
image_token_idread fromself.config, notself.text_config: Theimage_token_idattribute lives on the top-levelGemma4Config, not onGemma4TextConfig. Reading it fromself.text_configreturnsNone(the attribute doesn't exist there), which silently skips the image placeholder replacement — a critical bug that was caught and fixed during development.Side-by-side validation rather than golden-file comparison:
verify_step_build_static_inputsre-derives the expected output from raw HF processor output at validation time, rather than comparing against pre-recorded golden tensors. This makes the validation self-contained and robust to processor version changes.Helper functions as pure module-level functions:
_normalize_valid_token_maskand_validate_padding_layoutare extracted as standalone functions (not methods) so they can be unit-tested without instantiating a fullStaticGemma4Runtime(which requires a real model and processor).padding_sideaccessed defensively viahasattr:StaticGemma4Layoutdoes not currently have apadding_sidefield. The code useshasattr(self.layout, "padding_side")with a"right"fallback, so the layout dataclass doesn't need modification in this PR.Validation-only entry point:
run_static_gemma4_runtimeruns onlybuild_static_inputsvalidation and prints explicit SKIP messages for prefill, decode, and generation. This keeps the PR focused and reviewable while establishing the validation harness for future PRs.Changes
tico/quantization/recipes/debug/static_gemma4_runtime.py— Replacedbuild_static_inputsstub with full implementation (processor call, padding, image token replacement, valid_token_mask normalization, padding layout validation). Added_normalize_valid_token_maskand_validate_padding_layoutmodule-level helper functions. Addedverify_step_build_static_inputsside-by-side validation function (6 sub-step checks). Replacedrun_static_gemma4_runtimestub with validation-only entry point that loads model/processor, creates runtime, runs validation, and prints SKIP messages for prefill/decode/generation.tico/quantization/examples/inspector.py— Addedstatic-gemma4-runtimemode to the inspector CLI, importing and dispatching torun_static_gemma4_runtime/StaticGemma4RuntimeConfig.tico/quantization/examples/configs/static_gemma4_runtime.yaml— New example config for the static Gemma4 runtime inspector mode.test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.py— New test file with 12 unit tests for_normalize_valid_token_mask(5 tests: with attention mask, without mask using pad_token_id, without mask and no pad_token_id, shape mismatch raises, batched input) and_validate_padding_layout(7 tests: right padding valid, right padding no padding, right padding invalid raises, left padding valid, left padding no padding, batched right padding valid, batched right padding invalid raises).Tests
The 12 unit tests in
test/quantization/recipes/integration/test_static_gemma4_runtime_helpers.pyexercise the pure-Python helper functions without requiring a real Gemma4 model or processor:_normalize_valid_token_mask(5 tests): attention mask conversion to bool, pad_token_id derivation when mask is absent, all-valid fallback when both mask and pad_token_id are None, shape mismatch error, batched 2D input._validate_padding_layout(7 tests): valid right padding, no-padding case, invalid right padding raises, valid left padding, left no-padding case, batched valid, batched invalid raises.All 12 tests pass:
Example Scripts
The static runtime can be invoked via the inspector CLI:
This loads the model, creates the runtime, runs
build_static_inputsvalidation, and prints SKIP messages for the not-yet-implemented prefill/decode/generation steps.