Skip to content

Adding jit support to Qwen3VL, readme fixes #170

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
coder0143 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Adding jit support to Qwen3VL, readme fixes #170

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
529356a
adding jit to qwen3vl and readme fix
coder0143 Mar 3, 2026
914d864
Merge branch 'main' into main
coder0143 Apr 14, 2026
31fff6b
pre-commit fix
coder0143 Apr 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion bonsai/models/qwen3_vl/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@ This directory contains a pure JAX implementation of the [Qwen3-VL SOTA Vision L
| [Qwen3-VL-32B-Thinking](https://huggingface.co/Qwen/Qwen3-VL-32B-Thinking) | **✅ Supported** |
| [NVIDIA-Cosmos-Reason2-2B](https://huggingface.co/nvidia/Cosmos-Reason2-2B) | **✅ Supported** |
| [NVIDIA-Cosmos-Reason2-8B](https://huggingface.co/nvidia/Cosmos-Reason2-8B) | **✅ Supported** |

| **MoE Models** | |
| [Qwen3-VL-30B-A3B](https://huggingface.co/Qwen/Qwen3-30B-A3B) | **🟡 Not started** |
| [Qwen3-VL-235B-A22B](https://huggingface.co/Qwen/Qwen3-235B-A22B) | **🟡 Not started** |
Expand Down
85 changes: 41 additions & 44 deletions bonsai/models/qwen3_vl/modeling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,20 +1,10 @@
# Copyright 2025 The JAX Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import functools
import math
from dataclasses import dataclass
from typing import Optional, Tuple, TypeAlias

# One entry per image/video: ((T1, H1, W1), (T2, H2, W2), ...)
GridTHW: TypeAlias = tuple[tuple[int, int, int], ...]
from enum import Enum

import jax
Expand Down Expand Up @@ -599,10 +589,8 @@ def __init__(self, config: Qwen3VLVisionConfig, *, rngs: nnx.Rngs):
]
)

def _fast_pos_embed_interpolate(self, grid_thw: Array) -> Array:
"""Bilinear interpolation for position embeddings, matching PyTorch."""
grid_h, grid_w = int(grid_thw[0, 1]), int(grid_thw[0, 2])

def _fast_pos_embed_interpolate_single(self, grid_t: int, grid_h: int, grid_w: int) -> Array:
"""Bilinear interpolation for position embeddings for a single image."""
# Create interpolation indices
h_idxs = jnp.linspace(0, self.num_grid_per_side - 1, grid_h)
w_idxs = jnp.linspace(0, self.num_grid_per_side - 1, grid_w)
Expand Down Expand Up @@ -640,7 +628,6 @@ def _fast_pos_embed_interpolate(self, grid_thw: Array) -> Array:

# Apply spatial merge permutation
merge_size = self.config.spatial_merge_size
grid_t = int(grid_thw[0, 0])

# Reshape: (H*W, D) -> (H, W, D)
pos_embeds = pos_embeds.reshape(grid_h, grid_w, -1)
Expand All @@ -659,11 +646,16 @@ def _fast_pos_embed_interpolate(self, grid_thw: Array) -> Array:

return pos_embeds

def _rot_pos_emb(self, grid_thw: Array) -> Tuple[Array, Array]:
"""Compute rotary position embeddings matching PyTorch rot_pos_emb."""
def _fast_pos_embed_interpolate(self, grid_thw: GridTHW) -> Array:
"""Bilinear interpolation for position embeddings, handles multiple images."""
results = [
self._fast_pos_embed_interpolate_single(grid_t, grid_h, grid_w) for grid_t, grid_h, grid_w in grid_thw
]
return jnp.concatenate(results, axis=0)

def _rot_pos_emb_single(self, grid_t: int, grid_h: int, grid_w: int) -> Array:
"""Compute rotary position embeddings for a single image."""
merge_size = self.config.spatial_merge_size
grid_h, grid_w = int(grid_thw[0, 1]), int(grid_thw[0, 2])
grid_t = int(grid_thw[0, 0])

# Compute merged dimensions
merged_h, merged_w = grid_h // merge_size, grid_w // merge_size
Expand All @@ -678,7 +670,7 @@ def _rot_pos_emb(self, grid_thw: Array) -> Tuple[Array, Array]:
row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None]
col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :]

# Expand and reshape to match spatial merge order: (merged_h, merged_w, m, m) -> (merged_h, merged_w, m, m)
# Expand and reshape to match spatial merge order
row_idx = jnp.broadcast_to(row_idx, (merged_h, merged_w, merge_size, merge_size)).reshape(-1)
col_idx = jnp.broadcast_to(col_idx, (merged_h, merged_w, merge_size, merge_size)).reshape(-1)

Expand All @@ -687,32 +679,33 @@ def _rot_pos_emb(self, grid_thw: Array) -> Tuple[Array, Array]:
row_idx = jnp.tile(row_idx, grid_t)
col_idx = jnp.tile(col_idx, grid_t)

# Create frequency table - PyTorch uses rotary_dim = head_dim // 2
# And inv_freq has length rotary_dim // 2 = head_dim // 4
# Create frequency table
max_hw = max(grid_h, grid_w)
head_dim = self.config.head_dim
rotary_dim = head_dim // 2 # = 32 for head_dim=64
inv_freq_dim = rotary_dim // 2 # = 16
rotary_dim = head_dim // 2
inv_freq = 1.0 / (self.config.rope_theta ** (jnp.arange(0, rotary_dim, 2, dtype=jnp.float32) / rotary_dim))
seq_positions = jnp.arange(max_hw, dtype=jnp.float32)
freq_table = jnp.outer(seq_positions, inv_freq) # (max_hw, rotary_dim//2) = (max_hw, 16)
freq_table = jnp.outer(seq_positions, inv_freq)

# Lookup embeddings for row and col: (seq, rotary_dim//2) each
row_emb = freq_table[row_idx] # (seq, 16)
col_emb = freq_table[col_idx] # (seq, 16)
# Lookup embeddings for row and col
row_emb = freq_table[row_idx]
col_emb = freq_table[col_idx]

# Concatenate row and col: (seq, 32)
# Concatenate row and col, then double
emb = jnp.concatenate([row_emb, col_emb], axis=-1)

# Double the embedding (matching PyTorch cat): (seq, 64)
emb = jnp.concatenate([emb, emb], axis=-1)

# Apply cos/sin
return emb

def _rot_pos_emb(self, grid_thw: GridTHW) -> Tuple[Array, Array]:
"""Compute rotary position embeddings, handles multiple images."""
embs = [self._rot_pos_emb_single(grid_t, grid_h, grid_w) for grid_t, grid_h, grid_w in grid_thw]
emb = jnp.concatenate(embs, axis=0)
cos = jnp.cos(emb)
sin = jnp.sin(emb)
return cos, sin

def __call__(self, hidden_states: Array, grid_thw: Array) -> Tuple[Array, list[Array]]:
def __call__(self, hidden_states: Array, grid_thw: GridTHW) -> Tuple[Array, list[Array]]:
hidden_states = self.patch_embed(hidden_states)
seq_len = hidden_states.shape[0]

Expand Down Expand Up @@ -1012,10 +1005,9 @@ def __init__(self, config: ModelConfig, *, rngs: nnx.Rngs):
def __call__(
self,
input_ids: Array,
pixel_values: Optional[Array] = None,
image_grid_thw: Optional[Array] = None,
*,
cache: Cache,
pixel_values: Optional[Array] = None,
image_grid_thw: Optional[GridTHW] = None,
token_type_ids: Optional[Array] = None,
) -> Array:
"""Forward pass with KV-cache."""
Expand Down Expand Up @@ -1099,15 +1091,20 @@ def forward(model: Qwen3VLForConditionalGeneration, cache: Cache, input_ids: Arr
return logits[:, -1, :], cache


# TODO: Add jit compilation support by fixing image size
@functools.partial(jax.jit, static_argnums=(4,))
def forward_vision(
model: Qwen3VLForConditionalGeneration,
cache: Cache,
input_ids: Array,
pixel_values: Array,
image_grid_thw: Array,
image_grid_thw: GridTHW,
token_type_ids: Array,
) -> Tuple[Array, Cache]:
"""Forward pass with vision inputs (not JIT - vision has data-dependent shapes)."""
logits = model(input_ids, pixel_values, image_grid_thw, cache=cache, token_type_ids=token_type_ids)
"""JIT-compiled forward pass with vision inputs.

image_grid_thw is marked static since its values control array shapes
(linspace, reshape, arange) in the vision encoder. Different image
resolutions will trigger recompilation (one-time cost per resolution).
"""
logits = model(input_ids, cache, pixel_values, image_grid_thw, token_type_ids)
return logits[:, -1, :], cache
2 changes: 1 addition & 1 deletion bonsai/models/qwen3_vl/tests/run_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ def main():
# Check for vision inputs
if "pixel_values" in inputs_vision:
pixel_values = jnp.array(inputs_vision["pixel_values"].numpy())
image_grid_thw = jnp.array(inputs_vision["image_grid_thw"].numpy())
image_grid_thw = tuple(tuple(row) for row in inputs_vision["image_grid_thw"].long().tolist())

@vfdev-5 vfdev-5 Mar 3, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see that finally we could not make the original code simple as we still need to iterate over its structure. I wonder if we can accept the type of image_grid_thw as list[list[int]] instead of tuple[tuple[int, int, int]] ? Can this work with jit and the type checker etc?

@coder0143 coder0143 Mar 5, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, list[list[int]] won't work directly with static_argnums because lists are not hashable — JAX needs to hash static arguments for its compilation cache, and list raises TypeError: unhashable type: 'list'.

To make it simple we can do something like:

@functools.partial(jax.jit, static_argnums=(4,))
def _forward_vision_jit(model, cache, input_ids, pixel_values, image_grid_thw, token_type_ids):
    logits = model(input_ids, cache, pixel_values, image_grid_thw, token_type_ids)
    return logits[:, -1, :], cache

def forward_vision(model, cache, input_ids, pixel_values, image_grid_thw, token_type_ids):
    """Accepts list[list[int]] or tuple — converts to hashable tuple for JIT."""
    grid_thw = tuple(tuple(row) for row in image_grid_thw)
    return _forward_vision_jit(model, cache, input_ids, pixel_values, grid_thw, token_type_ids)

Which makes the callers simple:

image_grid_thw = inputs["image_grid_thw"].long().tolist()  # [[1, 64, 64]]
logits, cache = forward_vision(model, cache, ids, pixels, image_grid_thw, tti)

@vfdev-5 Do tell if you want me to make the changes.

@vfdev-5 vfdev-5 Mar 9, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Finally, I think it's fine to keep it everywhere as tuple of tuples. We have to do the conversion in tests only so it's ok


# Create token_type_ids (1 for image tokens, 0 for text)
# Image token ID is 151655
Expand Down
38 changes: 18 additions & 20 deletions bonsai/models/qwen3_vl/tests/test_outputs_qwen3vl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -374,17 +374,16 @@ def test_rope_embedding(self):
flax_visual = self.flax_model.model.visual

# Create grid_thw for a small image: 1 frame, 16x16 grid
grid_thw_np = np.array([[1, 16, 16]], dtype=np.int64)
grid_thw_pt = torch.tensor(grid_thw_np)
grid_thw_jax = jnp.array(grid_thw_np)
grid_thw_tuple = ((1, 16, 16),)
grid_thw_pt = torch.tensor(grid_thw_tuple, dtype=torch.long)

with torch.inference_mode():
pt_rope = pt_visual.rot_pos_emb(grid_thw_pt)
pt_emb = torch.cat([pt_rope, pt_rope], dim=-1)
pt_cos = pt_emb.cos().numpy()
pt_sin = pt_emb.sin().numpy()

flax_cos, flax_sin = flax_visual._rot_pos_emb(grid_thw_jax)
flax_cos, flax_sin = flax_visual._rot_pos_emb(grid_thw_tuple)

np.testing.assert_allclose(np.array(flax_cos), pt_cos, rtol=1e-5, atol=1e-5)
np.testing.assert_allclose(np.array(flax_sin), pt_sin, rtol=1e-5, atol=1e-5)
Expand All @@ -401,9 +400,8 @@ def test_full_vision_encoder(self):
grid_t, grid_h, grid_w = 1, 16, 16
num_patches = grid_t * grid_h * grid_w # 256 patches before merge

grid_thw_np = np.array([[grid_t, grid_h, grid_w]], dtype=np.int64)
grid_thw_pt = torch.tensor(grid_thw_np)
grid_thw_jax = jnp.array(grid_thw_np)
grid_thw_tuple = ((grid_t, grid_h, grid_w),)
grid_thw_pt = torch.tensor(grid_thw_tuple, dtype=torch.long)

key = jax.random.PRNGKey(42)
jx = jax.random.normal(key, (num_patches, per_patch_size), dtype=jnp.float32)
Expand All @@ -420,7 +418,7 @@ def test_full_vision_encoder(self):
else:
pt_out = pt_result[0].numpy()

flax_out, flax_deepstack = flax_visual(jx, grid_thw_jax)
flax_out, flax_deepstack = flax_visual(jx, grid_thw_tuple)
flax_out = np.array(flax_out)

# After spatial merge (2x2), 256 patches become 64
Expand Down Expand Up @@ -621,27 +619,27 @@ def test_position_embedding_output(self):
"""Check position embedding interpolation output matches."""
inputs = self._create_dummy_image_input()
grid_thw_pt = inputs["image_grid_thw"]
grid_thw_jax = jnp.array(grid_thw_pt.numpy())
grid_thw_tuple = tuple(tuple(row) for row in grid_thw_pt.long().tolist())

with torch.inference_mode():
pt_pos = self.pt_model.model.visual.fast_pos_embed_interpolate(grid_thw_pt).numpy()
flax_pos = np.array(self.flax_model.model.visual._fast_pos_embed_interpolate(grid_thw_jax))
flax_pos = np.array(self.flax_model.model.visual._fast_pos_embed_interpolate(grid_thw_tuple))

np.testing.assert_allclose(flax_pos, pt_pos, rtol=1e-5, atol=3e-5)

def test_rope_embedding(self):
"""Check RoPE embedding output matches."""
inputs = self._create_dummy_image_input()
grid_thw_pt = inputs["image_grid_thw"]
grid_thw_jax = jnp.array(grid_thw_pt.numpy())
grid_thw_tuple = tuple(tuple(row) for row in grid_thw_pt.long().tolist())

with torch.inference_mode():
pt_rope = self.pt_model.model.visual.rot_pos_emb(grid_thw_pt)
pt_emb = torch.cat([pt_rope, pt_rope], dim=-1)
pt_cos = pt_emb.cos().numpy()
pt_sin = pt_emb.sin().numpy()

flax_cos, flax_sin = self.flax_model.model.visual._rot_pos_emb(grid_thw_jax)
flax_cos, flax_sin = self.flax_model.model.visual._rot_pos_emb(grid_thw_tuple)

np.testing.assert_allclose(np.array(flax_cos), pt_cos, rtol=1e-6, atol=2e-5)
np.testing.assert_allclose(np.array(flax_sin), pt_sin, rtol=1e-6, atol=2e-5)
Expand All @@ -652,7 +650,7 @@ def test_vision_patch_plus_pos_output(self):
pixel_values_pt = inputs["pixel_values"]
grid_thw_pt = inputs["image_grid_thw"]
pixel_values_jax = jnp.array(pixel_values_pt.numpy())
grid_thw_jax = jnp.array(grid_thw_pt.numpy())
grid_thw_tuple = tuple(tuple(row) for row in grid_thw_pt.long().tolist())

# Get patch + pos embeddings
with torch.inference_mode():
Expand All @@ -661,7 +659,7 @@ def test_vision_patch_plus_pos_output(self):
pt_hidden = (pt_patches + pt_pos).numpy()

flax_patches = self.flax_model.model.visual.patch_embed(pixel_values_jax)
flax_pos = self.flax_model.model.visual._fast_pos_embed_interpolate(grid_thw_jax)
flax_pos = self.flax_model.model.visual._fast_pos_embed_interpolate(grid_thw_tuple)
flax_hidden = np.array(flax_patches + flax_pos)

np.testing.assert_allclose(flax_hidden, pt_hidden, rtol=1e-5, atol=1e-5)
Expand All @@ -672,7 +670,7 @@ def test_full_vision_output(self):
pixel_values_pt = inputs["pixel_values"]
grid_thw_pt = inputs["image_grid_thw"]
pixel_values_jax = jnp.array(pixel_values_pt.numpy())
grid_thw_jax = jnp.array(grid_thw_pt.numpy())
grid_thw_tuple = tuple(tuple(row) for row in grid_thw_pt.long().tolist())

with torch.inference_mode():
pt_result = self.pt_model.model.visual(pixel_values_pt, grid_thw_pt)
Expand All @@ -683,7 +681,7 @@ def test_full_vision_output(self):
pt_out = self.pt_model.model.visual.merger(pt_hidden).numpy()
else:
pt_out = pt_result[0].numpy()
flax_out, _ = self.flax_model.model.visual(pixel_values_jax, grid_thw_jax)
flax_out, _ = self.flax_model.model.visual(pixel_values_jax, grid_thw_tuple)

self.assertEqual(
np.array(flax_out).shape,
Expand Down Expand Up @@ -722,7 +720,7 @@ def test_vision_forward_with_numeric_check(self):
input_ids_pt = inputs["input_ids"]

pixel_values_jax = jnp.array(pixel_values_pt.numpy())
grid_thw_jax = jnp.array(grid_thw_pt.numpy())
grid_thw_tuple = tuple(tuple(row) for row in grid_thw_pt.long().tolist())
input_ids_jax = jnp.array(input_ids_pt.numpy())

# Create token_type_ids: 1 for image tokens, 0 for text
Expand All @@ -743,7 +741,7 @@ def test_vision_forward_with_numeric_check(self):
cache = model_lib.init_cache(self.flax_config, batch, seq_len, 20)

flax_logits, cache = model_lib.forward_vision(
self.flax_model, cache, input_ids_jax, pixel_values_jax, grid_thw_jax, token_type_ids_jax
self.flax_model, cache, input_ids_jax, pixel_values_jax, grid_thw_tuple, token_type_ids_jax
)

# Vision+text forward has larger tolerance due to accumulated numerical diffs
Expand Down Expand Up @@ -777,7 +775,7 @@ def test_generation_with_vision_input(self):
)

pixel_values_jax = jnp.array(inputs["pixel_values"].numpy())
grid_thw_jax = jnp.array(inputs["image_grid_thw"].numpy())
grid_thw_tuple = tuple(tuple(row) for row in inputs["image_grid_thw"].long().tolist())
input_ids_jax = jnp.array(inputs["input_ids"].numpy())

# Create token_type_ids: 1 for image tokens, 0 for text
Expand All @@ -789,7 +787,7 @@ def test_generation_with_vision_input(self):

# Prefill with vision
logits, cache = model_lib.forward_vision(
self.flax_model, cache, input_ids_jax, pixel_values_jax, grid_thw_jax, token_type_ids_jax
self.flax_model, cache, input_ids_jax, pixel_values_jax, grid_thw_tuple, token_type_ids_jax
)

# Verify cache position
Expand Down
Loading