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 RBLN NPUs, 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.

Info

Gemma4 runs across multiple RBLN NPUs and requires a Rebellions Scalable Design (RSD)-capable configuration: in this example the language model compiles across 16 NPUs and the vision encoder across 8. RSD is available on ATOM™+ (RBLN-CA22) and ATOM™-Max (RBLN-CA25). You can check your RBLN NPU type using the rbln-stat command.

Note

The 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 yet. 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, and the language model's context length. After compilation, save the model artifacts to disk using the save_pretrained() method. This will create 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 RBLN NPUs
model = RBLNGemma4ForConditionalGeneration.from_pretrained(
    model_id=model_id,
    export=True,
    rbln_config={
        "batch_size": 1,
        "language_model": {
            "num_devices": 16,
            "max_seq_len": 262_144,
            "kvcache_partition_len": 16384,
            "prefill_chunk_size": 128,
            "attn_impl": "flash_attn",
        },
        "vision_tower": {
            "num_devices": 8,
        },
    },
)

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

Model Inference

Load the compiled model, 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
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,
)

# 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)['content'])

Example Output: 

1
This image shows the Golden Gate Bridge in San Francisco, California. The red suspension bridge spans across the water, with Fort Point visible in the lower-left foreground. In the center foreground, a single bird is perched on a rock in the water.

Enabling bucketing

Gemma4's image processor compresses each image into a fixed number of image tokens - a chosen max_soft_tokens value (one of 70, 140, 280, 560, 1120) - and the model is compiled for one such value, called a bucket. A single bucket locks every image to that one setting: without recompiling, the token count cannot be lowered for a simple image or raised for a detail-heavy one. 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 returns the same output - it adds pre-compiled buckets for the same weights rather than changing the computation. Its advantage is per-image choice: a smaller max_soft_tokens spends fewer tokens and less compute on an image, and a larger one keeps more detail - choices the single-bucket baseline cannot make without recompiling.

Caution

More buckets mean longer compilation and higher device-memory usage. Choose buckets to match the max_soft_tokens values your workload needs - 2-4 values is a reasonable default.

Note

max_soft_tokens accepts only the values the Gemma4 image processor supports (70, 140, 280, 560, 1120), and the largest image_prefill_chunk_size must be at least the largest max_soft_tokens; optimum-rbln validates this at compile time. Keep the two axes matched:

  • Leaving image_prefill_chunk_size unset derives one bucket per max_soft_tokens value.
  • 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.

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 RBLN NPUs
 model = RBLNGemma4ForConditionalGeneration.from_pretrained(
     model_id=model_id,
     export=True,
     rbln_config={
         "batch_size": 1,
         "language_model": {
             "num_devices": 16,
             "max_seq_len": 262_144,
             "kvcache_partition_len": 16384,
             "prefill_chunk_size": 128,
             "attn_impl": "flash_attn",
+            "image_prefill_chunk_size": [384, 640, 1152],
         },
         "vision_tower": {
             "num_devices": 8,
+            "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 - 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)['content'])

Example Output: 

1
2
3
4
5
max_soft_tokens=70:
This image shows the Golden Gate Bridge in San Francisco, California. The iconic orange-red suspension bridge stretches across the water, with the city's coastline and hills visible in the background. In the foreground, there is a rocky shoreline with a building and a small rock in the water with a bird perched on it.

max_soft_tokens=1120:
The image shows the Golden Gate Bridge in San Francisco, California. It is a large, red suspension bridge spanning a body of water. In the foreground, there is a rocky shoreline and a large rock in the water with a bird perched on it. To the left, there is a historic stone building (Fort Point) and a parking area with several cars. In the background, rolling hills are visible under a clear blue sky.

References