Skip to content

Commit 6eb37d4

Browse files
jglee-sqbitsjglee-sqbits
committed
[Models] Add QwenImage encoder (rebased onto main)
1 parent 62c6aea commit 6eb37d4

6 files changed

Lines changed: 628 additions & 0 deletions

File tree

‎max/python/max/pipelines/architectures/qwen2_5vl/encoder/__init__.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
# ===----------------------------------------------------------------------=== #
2+
# Copyright (c) 2026, Modular Inc. All rights reserved.
3+
#
4+
# Licensed under the Apache License v2.0 with LLVM Exceptions:
5+
# https://llvm.org/LICENSE.txt
6+
#
7+
# Unless required by applicable law or agreed to in writing, software
8+
# distributed under the License is distributed on an "AS IS" BASIS,
9+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
# See the License for the specific language governing permissions and
11+
# limitations under the License.
12+
# ===----------------------------------------------------------------------=== #
13+
14+
from .model import Qwen25VLEncoderModel
15+
from .multimodal_encoder import Qwen25VLMultimodalEncoderModel
16+
17+
__all__ = ["Qwen25VLEncoderModel", "Qwen25VLMultimodalEncoderModel"]

‎max/python/max/pipelines/architectures/qwen2_5vl/encoder/layers/__init__.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# ===----------------------------------------------------------------------=== #
2+
# Copyright (c) 2026, Modular Inc. All rights reserved.
3+
#
4+
# Licensed under the Apache License v2.0 with LLVM Exceptions:
5+
# https://llvm.org/LICENSE.txt
6+
#
7+
# Unless required by applicable law or agreed to in writing, software
8+
# distributed under the License is distributed on an "AS IS" BASIS,
9+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
# See the License for the specific language governing permissions and
11+
# limitations under the License.
12+
# ===----------------------------------------------------------------------=== #
13+
14+
from .attention import Qwen25VLEncoderAttention
15+
16+
__all__ = ["Qwen25VLEncoderAttention"]

‎max/python/max/pipelines/architectures/qwen2_5vl/encoder/layers/attention.py‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
# ===----------------------------------------------------------------------=== #
2+
# Copyright (c) 2026, Modular Inc. All rights reserved.
3+
#
4+
# Licensed under the Apache License v2.0 with LLVM Exceptions:
5+
# https://llvm.org/LICENSE.txt
6+
#
7+
# Unless required by applicable law or agreed to in writing, software
8+
# distributed under the License is distributed on an "AS IS" BASIS,
9+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
# See the License for the specific language governing permissions and
11+
# limitations under the License.
12+
# ===----------------------------------------------------------------------=== #
13+
14+
"""Qwen2.5-VL encoder-only attention with bias support (module v2)."""
15+
16+
from __future__ import annotations
17+
18+
from max.dtype import DType
19+
from max.graph import DeviceRef, TensorValue, ops
20+
from max.nn.attention.mask_config import MHAMaskVariant
21+
from max.nn.kernels import flash_attention_gpu
22+
from max.nn.layer import Module
23+
from max.nn.linear import Linear
24+
from max.nn.rotary_embedding import RotaryEmbedding
25+
26+
27+
class Qwen25VLEncoderAttention(Module):
28+
"""Encoder-only attention with bias for Qwen2.5-VL (module v2)."""
29+
30+
def __init__(
31+
self,
32+
num_attention_heads: int,
33+
num_key_value_heads: int,
34+
hidden_size: int,
35+
head_dim: int,
36+
scale: float,
37+
attention_bias: bool = True,
38+
*,
39+
dtype: DType,
40+
device: DeviceRef,
41+
) -> None:
42+
super().__init__()
43+
self.n_heads = num_attention_heads
44+
self.n_kv_heads = num_key_value_heads
45+
self.head_dim = head_dim
46+
self.scale = scale
47+
48+
q_dim = head_dim * num_attention_heads
49+
kv_dim = head_dim * num_key_value_heads
50+
51+
self.q_proj = Linear(
52+
hidden_size,
53+
q_dim,
54+
dtype=dtype,
55+
device=device,
56+
has_bias=attention_bias,
57+
)
58+
self.k_proj = Linear(
59+
hidden_size,
60+
kv_dim,
61+
dtype=dtype,
62+
device=device,
63+
has_bias=attention_bias,
64+
)
65+
self.v_proj = Linear(
66+
hidden_size,
67+
kv_dim,
68+
dtype=dtype,
69+
device=device,
70+
has_bias=attention_bias,
71+
)
72+
self.o_proj = Linear(
73+
q_dim,
74+
hidden_size,
75+
dtype=dtype,
76+
device=device,
77+
has_bias=False,
78+
)
79+
80+
def _repeat_kv(self, x: TensorValue, n_rep: int) -> TensorValue:
81+
if n_rep == 1:
82+
return x
83+
seq_len = x.shape[0]
84+
n_kv_heads = x.shape[1]
85+
head_dim = x.shape[2]
86+
x = ops.unsqueeze(x, 2)
87+
x = ops.broadcast_to(x, (seq_len, n_kv_heads, n_rep, head_dim))
88+
return ops.reshape(x, (seq_len, n_kv_heads * n_rep, head_dim))
89+
90+
def __call__(
91+
self,
92+
x: TensorValue,
93+
rope: RotaryEmbedding,
94+
) -> TensorValue:
95+
total_seq_len = x.shape[0]
96+
97+
q = self.q_proj(x)
98+
k = self.k_proj(x)
99+
v = self.v_proj(x)
100+
101+
q = ops.reshape(q, (total_seq_len, self.n_heads, self.head_dim))
102+
k = ops.reshape(k, (total_seq_len, self.n_kv_heads, self.head_dim))
103+
v = ops.reshape(v, (total_seq_len, self.n_kv_heads, self.head_dim))
104+
105+
q = ops.squeeze(rope(ops.unsqueeze(q, 0)), 0)
106+
k = ops.squeeze(rope(ops.unsqueeze(k, 0)), 0)
107+
108+
if self.n_kv_heads != self.n_heads:
109+
n_rep = self.n_heads // self.n_kv_heads
110+
k = self._repeat_kv(k, n_rep)
111+
v = self._repeat_kv(v, n_rep)
112+
113+
q = ops.unsqueeze(q, 0)
114+
k = ops.unsqueeze(k, 0)
115+
v = ops.unsqueeze(v, 0)
116+
117+
attn_out = flash_attention_gpu(
118+
q,
119+
k,
120+
v,
121+
mask_variant=MHAMaskVariant.CAUSAL_MASK,
122+
scale=self.scale,
123+
)
124+
125+
attn_out = ops.squeeze(attn_out, 0)
126+
attn_out = ops.reshape(attn_out, (total_seq_len, -1))
127+
return self.o_proj(attn_out)

‎max/python/max/pipelines/architectures/qwen2_5vl/encoder/model.py‎

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
# ===----------------------------------------------------------------------=== #
2+
# Copyright (c) 2026, Modular Inc. All rights reserved.
3+
#
4+
# Licensed under the Apache License v2.0 with LLVM Exceptions:
5+
# https://llvm.org/LICENSE.txt
6+
#
7+
# Unless required by applicable law or agreed to in writing, software
8+
# distributed under the License is distributed on an "AS IS" BASIS,
9+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
# See the License for the specific language governing permissions and
11+
# limitations under the License.
12+
# ===----------------------------------------------------------------------=== #
13+
14+
"""Qwen2.5-VL encoder ComponentModel wrapper (module v2)."""
15+
16+
from __future__ import annotations
17+
18+
from collections.abc import Callable
19+
from typing import Any
20+
21+
from max.driver import Device
22+
from max.dtype import DType
23+
from max.engine import InferenceSession, Model
24+
from max.graph import DeviceRef, Graph, TensorType
25+
from max.graph.weights import Weights
26+
from max.nn.embedding import Embedding
27+
from max.nn.layer import Module
28+
from max.pipelines.architectures.llama3.weight_adapters import (
29+
LLAMA_SAFETENSOR_MAPPING as QWEN_SAFETENSOR_MAP,
30+
)
31+
from max.pipelines.lib import SupportedEncoding
32+
from max.pipelines.lib.interfaces.component_model import ComponentModel
33+
34+
from .model_config import Qwen25VLTextEncoderConfig
35+
from .qwen25vl import Qwen25VLTextEncoderTransformer
36+
37+
38+
class _EmbedOnly(Module):
39+
"""Token embedding only (module v2)."""
40+
41+
def __init__(
42+
self,
43+
vocab_size: int,
44+
hidden_size: int,
45+
*,
46+
dtype: DType,
47+
device: DeviceRef,
48+
) -> None:
49+
super().__init__()
50+
self.embed_tokens = Embedding(
51+
vocab_size,
52+
hidden_size,
53+
dtype=dtype,
54+
device=device,
55+
)
56+
57+
def __call__(self, tokens: Any) -> Any:
58+
return self.embed_tokens(tokens)
59+
60+
61+
class Qwen25VLEncoderModel(ComponentModel):
62+
"""Qwen2.5-VL language-side encoder ComponentModel wrapper (module v2)."""
63+
64+
def __init__(
65+
self,
66+
config: dict[str, Any],
67+
encoding: SupportedEncoding,
68+
devices: list[Device],
69+
weights: Weights,
70+
session: InferenceSession | None = None,
71+
) -> None:
72+
super().__init__(config, encoding, devices, weights)
73+
self.config = Qwen25VLTextEncoderConfig.generate(
74+
config,
75+
encoding,
76+
devices,
77+
)
78+
self.session = session
79+
self.load_model()
80+
81+
def load_model(self) -> Callable[..., Any]:
82+
embed_state: dict[str, Any] = {}
83+
transform_state: dict[str, Any] = {}
84+
85+
for key, value in self.weights.items():
86+
wd = value.data()
87+
88+
# Normalize floating-point weights to bf16
89+
if wd.dtype.is_float() and not wd.dtype.is_float8():
90+
is_scale = key.endswith(".weight_scale") or key.endswith(
91+
".input_scale"
92+
)
93+
if not is_scale:
94+
wd = wd.astype(DType.bfloat16)
95+
96+
# Key mapping
97+
adapted_key = key
98+
if adapted_key.startswith("model.language_model."):
99+
adapted_key = adapted_key[len("model.language_model.") :]
100+
else:
101+
for before, after in QWEN_SAFETENSOR_MAP.items():
102+
adapted_key = adapted_key.replace(before, after)
103+
104+
# Skip vision weights
105+
if adapted_key.startswith("visual.") or adapted_key.startswith(
106+
"vision_encoder."
107+
):
108+
continue
109+
110+
# Strip "model." prefix
111+
adapted_key = adapted_key.removeprefix("model.")
112+
113+
if adapted_key.startswith("embed_tokens."):
114+
embed_state[adapted_key] = wd
115+
elif (
116+
adapted_key.startswith("layers.")
117+
or adapted_key.startswith("norm.")
118+
or adapted_key.startswith("rope.")
119+
):
120+
transform_state[adapted_key] = wd
121+
122+
lc = self.config
123+
device_ref = DeviceRef.from_device(self.devices[0])
124+
125+
# --- Compile embed_tokens ---
126+
embed_model = _EmbedOnly(
127+
lc.vocab_size,
128+
lc.hidden_size,
129+
dtype=lc.dtype,
130+
device=device_ref,
131+
)
132+
embed_model.load_state_dict(
133+
embed_state, weight_alignment=1, strict=True
134+
)
135+
embed_input_types = [
136+
TensorType(DType.int64, shape=["total_seq_len"], device=device_ref),
137+
]
138+
with Graph("qwen_te_embed", input_types=embed_input_types) as g:
139+
out = embed_model(*(v.tensor for v in g.inputs))
140+
g.output(out)
141+
142+
session = self.session
143+
if session is None:
144+
session = InferenceSession(devices=self.devices)
145+
146+
self._embed_model: Model = session.load(
147+
g,
148+
weights_registry=embed_model.state_dict(),
149+
)
150+
151+
# --- Compile transformer layers + norm ---
152+
transform_model = Qwen25VLTextEncoderTransformer(lc)
153+
transform_model.load_state_dict(
154+
transform_state,
155+
weight_alignment=1,
156+
strict=True,
157+
)
158+
transform_input_types = [
159+
TensorType(
160+
lc.dtype,
161+
shape=["total_seq_len", lc.hidden_size],
162+
device=device_ref,
163+
),
164+
]
165+
with Graph("qwen_te_transform", input_types=transform_input_types) as g:
166+
out = transform_model(*(v.tensor for v in g.inputs))
167+
g.output(out)
168+
self._transform_model: Model = session.load(
169+
g,
170+
weights_registry=transform_model.state_dict(),
171+
)
172+
173+
return self._embed_model
174+
175+
def __call__(self, token_input: Any) -> tuple[Any]:
176+
"""Run text encoder: embed_tokens → transformer → normed output.
177+
178+
Accepts both Buffer (v2) and experimental Tensor (v3 compat).
179+
Returns a tuple wrapping the result in the same type as input.
180+
"""
181+
# Extract Buffer from _Tensor if needed
182+
is_tensor = hasattr(token_input, "driver_tensor")
183+
buf = token_input.driver_tensor if is_tensor else token_input
184+
185+
embed_result = self._embed_model.execute(buf)
186+
transform_result = self._transform_model.execute(embed_result[0])
187+
result_buf = transform_result[0]
188+
189+
if is_tensor:
190+
from max.experimental.tensor import Tensor as _Tensor
191+
192+
return (_Tensor(storage=result_buf),)
193+
194+
return (result_buf,)

0 commit comments

Comments
 (0)