Skip to content

Gemma-4-31B (Bucketing)

Overview

This guide is intended for users already familiar with the basics of optimum-rbln. It first compiles and runs the dense Gemma4 model (google/gemma-4-31B-it), a Vision-Language Model (VLM), on a single Rebel100™ (RBLN-CR13) NPU, then shows how to enable bucketing on the same example.

Bucketing compiles a single model that supports several predefined input shapes, so one runtime can handle inputs of different sizes and switch between them automatically at inference without building a separate model for each size. This is especially useful for a VLM. This guide covers only the Gemma4-specific configuration; for the general concept and trade-offs, see the Bucketing guide.

Serving on Rebel100™

On Rebel100™, use RBLN vLLM (vllm serve) as the main way to serve models such as Gemma4. It compiles the model and exposes an OpenAI-compatible API in one step; see the gemma-4-31B-it recipe. This tutorial covers the optimum-rbln Python API, which is useful for checking compilation and running inference directly.

Note

Rebel100™ does not support RSD (see the version matrix), and every Rebel100™ example in the Model Zoo compiles for one NPU, so both the language model and the vision encoder set num_devices to 1 in this tutorial.

Note

Gemma4 model published on the HuggingFace Hub has restricted access under Google's Gemma license. Accept the license and authenticate before downloading.

Setup & Installation

Before you begin, ensure that your system environment is properly configured and that all required packages are installed. This includes:

Note

  • rebel-compiler requires an RBLN Portal account.
  • The commands above are intended for a default pip install on Debian-based Linux such as Ubuntu. For all other configurations, refer to the Installation Guide for the supported install matrix and the applicable commands.

Note

The google/gemma-4-31B-it model on HuggingFace has restricted access. Once access is granted, you can log in using the hf (huggingface-cli) command as shown below:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
$ hf auth login

    _|    _|  _|    _|    _|_|_|    _|_|_|  _|_|_|  _|      _|    _|_|_|      _|_|_|_|    _|_|      _|_|_|  _|_|_|_|
    _|    _|  _|    _|  _|        _|          _|    _|_|    _|  _|            _|        _|    _|  _|        _|
    _|_|_|_|  _|    _|  _|  _|_|  _|  _|_|    _|    _|  _|  _|  _|  _|_|      _|_|_|    _|_|_|_|  _|        _|_|_|
    _|    _|  _|    _|  _|    _|  _|    _|    _|    _|    _|_|  _|    _|      _|        _|    _|  _|        _|
    _|    _|    _|_|      _|_|_|    _|_|_|  _|_|_|  _|      _|    _|_|_|      _|        _|    _|    _|_|_|  _|_|_|_|

    To login, `huggingface_hub` requires a token generated from [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens) .
Enter your token (input will not be visible):

Using RBLN Optimum

This section compiles google/gemma-4-31B-it, loads the compiled model, and runs image-text inference without bucketing. Gemma4 is a VLM whose top-level RBLNGemma4ForConditionalGeneration wraps two submodules, a vision_tower (image encoder) and a language_model (causal LM), each configured through rbln_config.

Model Compilation

To begin, import the RBLNGemma4ForConditionalGeneration class from optimum-rbln. This class's from_pretrained() method downloads the Gemma4 model from the HuggingFace Hub and compiles it using the RBLN Compiler. When exporting the model, set export=True and provide rbln_config, which sets the batch size, each submodule's device count (1 on Rebel100™), and the language model's context length. The language model uses flash attention with a max_seq_len of 262_144 and a kvcache_partition_len of 8_192. On Rebel100™, kvcache_partition_len must be between 1_024 and 16_384, and max_seq_len must be a multiple of it and at least twice it. After compilation, save the model artifacts to disk using the save_pretrained() method. This creates a directory (e.g., gemma-4-31B-it) containing the compiled model.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
from optimum.rbln import RBLNGemma4ForConditionalGeneration

# Define the HuggingFace model ID
model_id = "google/gemma-4-31B-it"

# Compile the model for a single Rebel100™ NPU
model = RBLNGemma4ForConditionalGeneration.from_pretrained(
    model_id=model_id,
    export=True,
    rbln_config={
        "batch_size": 1,
        "language_model": {
            "num_devices": 1,
            "max_seq_len": 262_144,
            "kvcache_partition_len": 8_192,
            "prefill_chunk_size": 128,
            "attn_impl": "flash_attn",
        },
        "vision_tower": {
            "num_devices": 1,
        },
    },
)

model.save_pretrained("gemma-4-31B-it")

Model Inference

Load the compiled model on NPU 0 (rbln_device=0), build a chat message with an image and a text prompt, and generate a response. apply_chat_template() returns model-ready tensors; only the newly generated tokens are decoded.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
from optimum.rbln import RBLNGemma4ForConditionalGeneration
from transformers import AutoProcessor

# Load the compiled model
model = RBLNGemma4ForConditionalGeneration.from_pretrained(
    model_id="gemma-4-31B-it",
    export=False,
    rbln_device=0,
)

# Build a chat message with an image and a text prompt
processor = AutoProcessor.from_pretrained("google/gemma-4-31B-it")
messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "url": "https://raw.githubusercontent.com/google-gemma/cookbook/refs/heads/main/apps/sample-data/GoldenGate.png"},
            {"type": "text", "text": "What is shown in this image?"},
        ],
    }
]
inputs = processor.apply_chat_template(
    messages,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
    add_generation_prompt=True,
)
input_len = inputs["input_ids"].shape[-1]

# Generate and decode only the newly generated tokens
outputs = model.generate(**inputs, max_new_tokens=512)
response = processor.decode(outputs[0][input_len:], skip_special_tokens=False)
print(processor.parse_response(response, prefix=inputs["input_ids"][0][:input_len])['content'])

Example Output: 

1
This image shows the Golden Gate Bridge in San Francisco, California. The bridge's iconic international orange suspension structure is the central focus, spanning across the water. In the foreground, a portion of Fort Point is visible on the left, and a small rocky outcrop with a bird on it sits in the water. Mountains can be seen in the background under a clear blue sky.

Enabling bucketing

Gemma4's image processor compresses each image into a fixed number of image tokens set by max_soft_tokens (one of 70, 140, 280, 560, 1120), and the model is compiled for one such value, called a bucket. With a single bucket, every image uses that one setting, and the token count cannot be lowered for a simple image or raised for a detail-heavy one without recompiling. Bucketing compiles several buckets together so the max_soft_tokens value can be chosen per image from one model.

Bucketing applies to two linked submodules: the vision_tower is sized to the chosen max_soft_tokens value, and the language_model sizes the separate bidirectional graph that prefills those image tokens (image_prefill_chunk_size, distinct from the causal text prefill_chunk_size; each image or video frame is one prefilled run). The chosen max_soft_tokens value flows into the LM prefill, so optimum-rbln derives and validates the LM buckets against the vision buckets at compile time. Set both through rbln_config, per submodule.

Field Submodule Accepted values Effect
max_soft_tokens vision_tower List[int] drawn from {70, 140, 280, 560, 1120} Image-token count buckets for the vision encoder. The processor's max_soft_tokens sets the count per image and must match one of these compiled values.
image_prefill_chunk_size language_model List[int], each a multiple of 128 Image-prefill chunk-size buckets for the LM. At runtime the smallest bucket ≥ the max_soft_tokens used for the image is selected automatically, so the largest value must be ≥ the largest max_soft_tokens.

Bucket Configuration

Starting from the baseline compile above, pass max_soft_tokens (on vision_tower) and image_prefill_chunk_size (on language_model) as lists so their buckets are compiled together. The full compile call is shown below, with the added bucketing lines marked with +. For the 280-token value that the baseline also compiles, the bucketed model runs the same computation; bucketing adds precompiled buckets for the same weights without changing the computation. The generated text can still differ from run to run, because Gemma 4's generation_config samples (do_sample=True, top_k=64, top_p=0.95). Bucketing lets you choose max_soft_tokens per image: a smaller value spends fewer tokens and less compute on an image, and a larger value keeps more detail. The single-bucket baseline requires recompilation to change this value.

Caution

More buckets mean longer compilation and higher device-memory usage. Choose buckets to match the max_soft_tokens values your workload needs; two to four values are typically sufficient.

Note

Use only max_soft_tokens values the Gemma4 image processor supports (70, 140, 280, 560, 1120); optimum-rbln does not check these values itself. It does check image_prefill_chunk_size against max_soft_tokens at compile time: every value must be a positive multiple of 128, the largest value must be at least the largest max_soft_tokens, and every value must be the smallest bucket ≥ at least one max_soft_tokens (a bucket that no image would use fails the check). Keep the two axes matched:

  • Leaving image_prefill_chunk_size unset derives one bucket per max_soft_tokens value (each rounded up to a multiple of 128; [128, 256, 384, 640, 1152] for the values below).
  • A single value below the largest max_soft_tokens fails to compile.
  • A single large value prefills every image at that size, wasting compute on small ones.

In the example below, [384, 640, 1152] serves 70, 140, and 280 with 384, 560 with 640, and 1120 with 1152, so every bucket is used.

At inference, set the processor's max_soft_tokens to any compiled bucket (for example, AutoProcessor.from_pretrained(model_id, max_soft_tokens=560)), and the matching image_prefill_chunk_size is selected automatically as the smallest bucket ≥ that value.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
 from optimum.rbln import RBLNGemma4ForConditionalGeneration

 # Define the HuggingFace model ID
 model_id = "google/gemma-4-31B-it"

 # Compile the model for a single Rebel100™ NPU
 model = RBLNGemma4ForConditionalGeneration.from_pretrained(
     model_id=model_id,
     export=True,
     rbln_config={
         "batch_size": 1,
         "language_model": {
             "num_devices": 1,
             "max_seq_len": 262_144,
             "kvcache_partition_len": 8_192,
             "prefill_chunk_size": 128,
             "attn_impl": "flash_attn",
+            "image_prefill_chunk_size": [384, 640, 1152],
         },
         "vision_tower": {
             "num_devices": 1,
+            "max_soft_tokens": [70, 140, 280, 560, 1120],
         },
     },
 )

 model.save_pretrained("gemma-4-31B-it")

Bucket Selection

Select a compiled bucket at inference by setting max_soft_tokens on the processor. This is the only change from the baseline inference above. A smaller value (70) spends the fewest tokens per image: the highest throughput and lowest latency, at the coarsest detail, so it suits simple or low-information images. A larger value (1120) preserves the most detail, so it suits detail-heavy images such as documents or fine textures, at the cost of the most compute and memory. Any compiled bucket works; reuse the same model and messages from the baseline.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
# Select any compiled bucket via the processor (70, 140, 280, 560, 1120);
# the rest is the baseline inference flow, reusing `model` and `messages`.
processor = AutoProcessor.from_pretrained(
    "google/gemma-4-31B-it",
    max_soft_tokens=70,
)

inputs = processor.apply_chat_template(
    messages,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
    add_generation_prompt=True,
)
input_len = inputs["input_ids"].shape[-1]

outputs = model.generate(**inputs, max_new_tokens=512)
response = processor.decode(outputs[0][input_len:], skip_special_tokens=False)
print(processor.parse_response(response, prefix=inputs["input_ids"][0][:input_len])['content'])

Example Output: 

1
2
3
4
5
6
7
max_soft_tokens=70:
The image shows the Golden Gate Bridge in San Francisco, California. The red suspension bridge spans across the water, with one of its large towers and a portion of the roadway visible. To the left, there is a brick building and a shoreline. In the foreground, a small rock with a bird on it sits in the water.

max_soft_tokens=1120:
This image shows the Golden Gate Bridge in San Francisco, California. The iconic orange-red suspension bridge spans across the water, connecting the city to Marin County.

In the foreground, there is a rocky shoreline with Fort Point, a historic masonry fortification, situated directly beneath the bridge's southern end. A small rock with a bird perched on it is visible in the water, and hills can be seen in the background under a clear blue sky.

References