@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9484641616

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 22:35 UTC

[Model] Add step3 vl (#21998)

Signed-off-by: oliveryuan <yuansong@step.ai> Co-authored-by: oliveryuan <yuansong@step.ai>

Song committed Jul 31, 2025 at 23:19 UTC 94846416166c731939892350d7ab26dcbcb2982d

12 files changed +2124 -1

docs/models/supported_models.md

@@ -625,6 +625,7 @@ See [this page](generative_models.md) for more information on how to use generat

625 | `Qwen2_5OmniThinkerForConditionalGeneration` | Qwen2.5-Omni | T + IE+ + VE+ + A+ | `Qwen/Qwen2.5-Omni-7B` | | ✅︎ | ✅︎ |

626 | `SkyworkR1VChatModel` | Skywork-R1V-38B | T + I | `Skywork/Skywork-R1V-38B` | | ✅︎ | ✅︎ |

627 | `SmolVLMForConditionalGeneration` | SmolVLM2 | T + I | `SmolVLM2-2.2B-Instruct` | ✅︎ | | ✅︎ |

628 +| `Step3VLForConditionalGeneration` | Step3-VL | T + I+ | `stepfun-ai/step3` | | ✅︎ | ✅︎ |

629 | `TarsierForConditionalGeneration` | Tarsier | T + IE+ | `omni-search/Tarsier-7b`, `omni-search/Tarsier-34b` | | ✅︎ | ✅︎ |

630 | `Tarsier2ForConditionalGeneration`^ | Tarsier2 | T + IE+ + VE+ | `omni-research/Tarsier2-Recap-7b`, `omni-research/Tarsier2-7b-0115` | | ✅︎ | ✅︎ |

631

tests/models/registry.py

@@ -279,6 +279,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

279 "StableLMEpochForCausalLM": _HfExamplesInfo("stabilityai/stablelm-zephyr-3b"), # noqa: E501

280 "StableLmForCausalLM": _HfExamplesInfo("stabilityai/stablelm-3b-4e1t"),

281 "Starcoder2ForCausalLM": _HfExamplesInfo("bigcode/starcoder2-3b"),

282 + "Step3TextForCausalLM": _HfExamplesInfo("stepfun-ai/step3",

283 + trust_remote_code=True,

284 + is_available_online=False),

285 "SolarForCausalLM": _HfExamplesInfo("upstage/solar-pro-preview-instruct",

286 trust_remote_code=True),

287 "TeleChat2ForCausalLM": _HfExamplesInfo("Tele-AI/TeleChat2-3B",

@@ -457,6 +460,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {

460 "SkyworkR1VChatModel": _HfExamplesInfo("Skywork/Skywork-R1V-38B",

461 trust_remote_code=True),

462 "SmolVLMForConditionalGeneration": _HfExamplesInfo("HuggingFaceTB/SmolVLM2-2.2B-Instruct"), # noqa: E501

463 + "Step3VLForConditionalGeneration": _HfExamplesInfo("stepfun-ai/step3",

464 + trust_remote_code=True,

465 + is_available_online=False),

466 "UltravoxModel": _HfExamplesInfo("fixie-ai/ultravox-v0_5-llama-3_2-1b", # noqa: E501

467 trust_remote_code=True),

468 "TarsierForConditionalGeneration": _HfExamplesInfo("omni-research/Tarsier-7b", # noqa: E501

vllm/entrypoints/openai/tool_parsers/__init__.py

@@ -18,6 +18,7 @@ from .mistral_tool_parser import MistralToolParser

18 from .phi4mini_tool_parser import Phi4MiniJsonToolParser

19 from .pythonic_tool_parser import PythonicToolParser

20 from .qwen3coder_tool_parser import Qwen3CoderToolParser

21 +from .step3_tool_parser import Step3ToolParser

22 from .xlam_tool_parser import xLAMToolParser

24 __all__ = [

@@ -40,4 +41,5 @@ __all__ = [

41 "HunyuanA13BToolParser",

42 "Glm4MoeModelToolParser",

43 "Qwen3CoderToolParser",

44 + "Step3ToolParser",

45 ]

vllm/entrypoints/openai/tool_parsers/step3_tool_parser.py new

+296

@@ -0,0 +1,296 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import contextlib

5 +import json

6 +from collections.abc import Sequence

7 +from typing import Any, Optional, Union

8 +

9 +import regex as re

10 +

11 +from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,

12 + DeltaFunctionCall, DeltaMessage,

13 + DeltaToolCall,

14 + ExtractedToolCallInformation,

15 + FunctionCall, ToolCall)

16 +from vllm.entrypoints.openai.tool_parsers.abstract_tool_parser import (

17 + ToolParser, ToolParserManager)

18 +from vllm.logger import init_logger

19 +from vllm.transformers_utils.tokenizer import AnyTokenizer

20 +from vllm.utils import random_uuid

21 +

22 +logger = init_logger(__name__)

23 +

24 +

25 +@ToolParserManager.register_module(["step3"])

26 +class Step3ToolParser(ToolParser):

27 + """

28 + Tool parser for a model that uses a specific XML-like format for tool calls.

29 + This version uses a robust, stateful, cursor-based streaming parser and

30 + consolidates tool arguments into a single message.

31 + """

32 +

33 + TOOL_CALLS_BEGIN = "<｜tool_calls_begin｜>"

34 + TOOL_CALLS_END = "<｜tool_calls_end｜>"

35 + TOOL_CALL_BEGIN = "<｜tool_call_begin｜>"

36 + TOOL_CALL_END = "<｜tool_call_end｜>"

37 + TOOL_SEP = "<｜tool_sep｜>"

38 + SPECIAL_TOKENS = [

39 + TOOL_CALLS_BEGIN, TOOL_CALLS_END, TOOL_CALL_BEGIN, TOOL_CALL_END

40 + ]

41 +

42 + def __init__(self, tokenizer: AnyTokenizer):

43 + super().__init__(tokenizer)

44 + self.position = 0

45 + # Explicit state flags for robust streaming

46 + self.tool_block_started = False

47 + self.tool_block_finished = False

48 +

49 + def adjust_request(

50 + self, request: ChatCompletionRequest) -> ChatCompletionRequest:

51 + if request.tools and request.tool_choice != 'none':

52 + request.skip_special_tokens = False

53 + return request

54 +

55 + @staticmethod

56 + def _parse_steptml_invoke(

57 + action_text: str

58 + ) -> tuple[Optional[str], Optional[dict[str, str]]]:

59 + func_name_match = re.search(r'<steptml:invoke name="([^"]+)">',

60 + action_text)

61 + if not func_name_match:

62 + return None, None

63 + func_name = func_name_match.group(1)

64 +

65 + params: dict[str, str] = {}

66 + param_matches = re.findall(

67 + r'<steptml:parameter name="([^"]+)">([^<]*)</steptml:parameter>',

68 + action_text)

69 + for name, value in param_matches:

70 + params[name] = value.strip()

71 + return func_name, params

72 +

73 + def _cast_arguments(

74 + self,

75 + func_name: str,

76 + params: dict[str, Any],

77 + request: ChatCompletionRequest,

78 + ) -> dict[str, Any]:

79 + for tool in request.tools or []:

80 + if tool.function.name == func_name:

81 + schema = tool.function.parameters or {}

82 + properties = schema.get("properties", {})

83 + for key, value in params.items():

84 + if not isinstance(value, str):

85 + continue

86 + prop = properties.get(key, {})

87 + typ = prop.get("type")

88 + if typ == "string":

89 + params[key] = value.strip()

90 + elif typ == "integer":

91 + with contextlib.suppress(ValueError):

92 + params[key] = int(value)

93 + elif typ == "number":

94 + with contextlib.suppress(ValueError):

95 + params[key] = float(value)

96 + elif typ == "boolean":

97 + lower_val = value.lower()

98 + params[key] = lower_val == "true" if lower_val in (

99 + "true", "false") else value

100 + elif typ == "null":

101 + params[key] = None if value.lower(

102 + ) == "null" else value

103 + break

104 + return params

105 +

106 + def extract_tool_calls_streaming(

107 + self,

108 + previous_text: str,

109 + current_text: str,

110 + delta_text: str,

111 + previous_token_ids: Sequence[int],

112 + current_token_ids: Sequence[int],

113 + delta_token_ids: Sequence[int],

114 + request: ChatCompletionRequest,

115 + ) -> Union[DeltaMessage, None]:

116 +

117 + # The main loop processes the stream from the last known position.

118 + while True:

119 + if self.position >= len(current_text):

120 + return None # We've processed the entire stream.

121 +

122 + unprocessed_text = current_text[self.position:]

123 +

124 + # STATE: After all tools are done, all subsequent text is content.

125 + if self.tool_block_finished:

126 + self.position = len(current_text)

127 + return DeltaMessage(content=unprocessed_text)

128 +

129 + # STATE: Before the tool block has started.

130 + if not self.tool_block_started:

131 + if unprocessed_text.startswith(self.TOOL_CALLS_BEGIN):

132 + self.position += len(self.TOOL_CALLS_BEGIN)

133 + self.tool_block_started = True

134 + continue # Token consumed, re-loop.

135 +

136 + start_pos = unprocessed_text.find(self.TOOL_CALLS_BEGIN)

137 + if start_pos == -1:

138 + if self.TOOL_CALLS_BEGIN.startswith(

139 + unprocessed_text.strip()) and unprocessed_text:

140 + return None # It's a prefix, wait.

141 + self.position = len(current_text)

142 + return DeltaMessage(content=unprocessed_text)

143 + else:

144 + content = unprocessed_text[:start_pos]

145 + self.position += len(content)

146 + return DeltaMessage(content=content)

147 +

148 + # STATE: Inside the main tool block.

149 + offset = len(unprocessed_text) - len(unprocessed_text.lstrip())

150 + unprocessed_text = unprocessed_text.lstrip()

151 + self.position += offset

152 +

153 + if unprocessed_text.startswith(self.TOOL_CALLS_END):

154 + self.position += len(self.TOOL_CALLS_END)

155 + self.tool_block_finished = True

156 + self.current_tool_id = -1

157 + continue

158 +

159 + # Check if we are between tool calls.

160 + tool_finished = (

161 + self.current_tool_id != -1 and

162 + self.prev_tool_call_arr[self.current_tool_id].get("finished"))

163 + if self.current_tool_id == -1 or tool_finished:

164 + if unprocessed_text.startswith(self.TOOL_CALL_BEGIN):

165 + self.position += len(self.TOOL_CALL_BEGIN)

166 + if self.current_tool_id == -1:

167 + self.current_tool_id = 0

168 + else:

169 + self.current_tool_id += 1

170 + self.current_tool_name_sent = False

171 + while len(self.prev_tool_call_arr) <= self.current_tool_id:

172 + self.prev_tool_call_arr.append({})

173 + self.prev_tool_call_arr[

174 + self.current_tool_id]["finished"] = False

175 + continue

176 +

177 + if self.TOOL_CALL_BEGIN.startswith(unprocessed_text):

178 + return None

179 +

180 + # STATE: Parsing an active tool call.

181 + if self.current_tool_id != -1 and not self.prev_tool_call_arr[

182 + self.current_tool_id].get("finished", False):

183 + end_tool_pos = unprocessed_text.find(self.TOOL_CALL_END)

184 + if end_tool_pos == -1:

185 + tool_body = unprocessed_text

186 + else:

187 + tool_body = unprocessed_text[:end_tool_pos]

188 +

189 + if end_tool_pos == -1 and self.TOOL_CALL_END.startswith(

190 + tool_body):

191 + return None

192 +

193 + function_name, arguments = self._parse_steptml_invoke(

194 + tool_body)

195 + if not function_name:

196 + return None

197 +

198 + tool_call_arr = {

199 + "name": function_name,

200 + "parameters": arguments or {}

201 + }

202 +

203 + # Send the function name as soon as it's parsed.

204 + if not self.current_tool_name_sent:

205 + self.current_tool_name_sent = True

206 + self.prev_tool_call_arr[self.current_tool_id].update(

207 + tool_call_arr)

208 + return DeltaMessage(tool_calls=[

209 + DeltaToolCall(index=self.current_tool_id,

210 + type="function",

211 + id=f"chatcmpl-tool-{random_uuid()}",

212 + function=DeltaFunctionCall(

213 + name=function_name))

214 + ])

215 +

216 + # Update our internal state with the latest parsed arguments.

217 + self.prev_tool_call_arr[

218 + self.current_tool_id].update( # noqa: E501

219 + tool_call_arr)

220 +

221 + # Only send arguments when the tool call is complete.

222 + if end_tool_pos != -1:

223 + self.position += end_tool_pos + len(self.TOOL_CALL_END)

224 + self.prev_tool_call_arr[

225 + self.current_tool_id]["finished"] = True

226 +

227 + final_args = self._cast_arguments(

228 + function_name,

229 + tool_call_arr.get("parameters", {}), # type: ignore

230 + request)

231 + if final_args:

232 + final_args_json = json.dumps(final_args,

233 + ensure_ascii=False)

234 + return DeltaMessage(tool_calls=[

235 + DeltaToolCall(index=self.current_tool_id,

236 + function=DeltaFunctionCall(

237 + arguments=final_args_json))

238 + ])

239 +

240 + # If tool is not finished, return None to wait for more tokens.

241 + return None

242 +

243 + return None

244 +

245 + def extract_tool_calls(

246 + self,

247 + model_output: str,

248 + request: ChatCompletionRequest,

249 + ) -> ExtractedToolCallInformation:

250 + if self.TOOL_CALLS_BEGIN not in model_output:

251 + return ExtractedToolCallInformation(tools_called=False,

252 + tool_calls=[],

253 + content=model_output)

254 +

255 + pre_text, rest = model_output.split(self.TOOL_CALLS_BEGIN, 1)

256 + if self.TOOL_CALLS_END not in rest:

257 + return ExtractedToolCallInformation(tools_called=False,

258 + tool_calls=[],

259 + content=model_output)

260 +

261 + tool_block, post_text = rest.split(self.TOOL_CALLS_END, 1)

262 + content = (pre_text + post_text).strip()

263 +

264 + tool_calls: list[ToolCall] = []

265 + call_parts = tool_block.split(self.TOOL_CALL_BEGIN)

266 +

267 + for part in call_parts:

268 + if not part or self.TOOL_CALL_END not in part:

269 + continue

270 +

271 + call_content = part.split(self.TOOL_CALL_END, 1)[0]

272 + if self.TOOL_SEP not in call_content:

273 + continue

274 +

275 + type_part, invoke_part = call_content.split(self.TOOL_SEP, 1)

276 + if type_part.strip() != "function":

277 + continue

278 +

279 + function_name, params_dict = self._parse_steptml_invoke(

280 + invoke_part)

281 +

282 + if function_name and params_dict is not None:

283 + params_dict = self._cast_arguments(function_name, params_dict,

284 + request)

285 + params_str = json.dumps(params_dict, ensure_ascii=False)

286 + tool_calls.append(

287 + ToolCall(function=FunctionCall(name=function_name,

288 + arguments=params_str)))

289 + if tool_calls:

290 + return ExtractedToolCallInformation(

291 + tools_called=True,

292 + tool_calls=tool_calls,

293 + content=content if content else None)

294 + return ExtractedToolCallInformation(tools_called=False,

295 + tool_calls=[],

296 + content=model_output)

vllm/model_executor/models/registry.py

@@ -129,6 +129,7 @@ _TEXT_GENERATION_MODELS = {

129 "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),

130 "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),

131 "RWForCausalLM": ("falcon", "FalconForCausalLM"),

132 + "Step3TextForCausalLM": ("step3_text", "Step3TextForCausalLM"),

133 "StableLMEpochForCausalLM": ("stablelm", "StablelmForCausalLM"),

134 "StableLmForCausalLM": ("stablelm", "StablelmForCausalLM"),

135 "Starcoder2ForCausalLM": ("starcoder2", "Starcoder2ForCausalLM"),

@@ -238,6 +239,7 @@ _MULTIMODAL_MODELS = {

239 "Qwen2_5OmniModel": ("qwen2_5_omni_thinker", "Qwen2_5OmniThinkerForConditionalGeneration"), # noqa: E501

240 "Qwen2_5OmniForConditionalGeneration": ("qwen2_5_omni_thinker", "Qwen2_5OmniThinkerForConditionalGeneration"), # noqa: E501

241 "UltravoxModel": ("ultravox", "UltravoxModel"),

242 + "Step3VLForConditionalGeneration": ("step3_vl", "Step3VLForConditionalGeneration"), # noqa: E501

243 "TarsierForConditionalGeneration": ("tarsier", "TarsierForConditionalGeneration"), # noqa: E501

244 "Tarsier2ForConditionalGeneration": ("qwen2_vl", "Tarsier2ForConditionalGeneration"), # noqa: E501

245 "VoxtralForConditionalGeneration": ("voxtral", "VoxtralForConditionalGeneration"), # noqa: E501

vllm/model_executor/models/step3_text.py new

+521

@@ -0,0 +1,521 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""Inference-only Jurassic model."""

4 +from collections.abc import Iterable

5 +from typing import Any, Optional

6 +

7 +import torch

8 +from torch import nn

9 +

10 +from vllm.attention import Attention

11 +from vllm.compilation.decorators import support_torch_compile

12 +from vllm.config import CacheConfig, ModelConfig, VllmConfig

13 +from vllm.distributed import (get_pp_group,

14 + get_tensor_model_parallel_world_size,

15 + tensor_model_parallel_all_reduce)

16 +from vllm.logger import init_logger

17 +from vllm.model_executor.layers.activation import SiluAndMul

18 +from vllm.model_executor.layers.fused_moe import FusedMoE

19 +from vllm.model_executor.layers.layernorm import RMSNorm

20 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,

21 + MergedColumnParallelLinear,

22 + ReplicatedLinear,

23 + RowParallelLinear)

24 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

25 +from vllm.model_executor.layers.quantization.base_config import (

26 + QuantizationConfig)

27 +from vllm.model_executor.layers.rotary_embedding import get_rope

28 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

29 +from vllm.model_executor.layers.vocab_parallel_embedding import (

30 + DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

31 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

32 +from vllm.model_executor.sampling_metadata import SamplingMetadata

33 +from vllm.sequence import IntermediateTensors

34 +

35 +from .interfaces import SupportsPP

36 +from .utils import (PPMissingLayer, is_pp_missing_parameter,

37 + make_empty_intermediate_tensors_factory, make_layers)

38 +

39 +logger = init_logger(__name__)

40 +

41 +

42 +class FusedMoEBlock(nn.Module):

43 +

44 + def __init__(self,

45 + config: ModelConfig,

46 + quant_config: Optional[QuantizationConfig] = None,

47 + prefix: str = ""):

48 + super().__init__()

49 + self.tp_size = get_tensor_model_parallel_world_size()

50 +

51 + if self.tp_size > config.moe_num_experts:

52 + raise ValueError(

53 + f"Tensor parallel size {self.tp_size} is greater than "

54 + f"the number of experts {config.moe_num_experts}.")

55 +

56 + self.experts = FusedMoE(num_experts=config.moe_num_experts,

57 + top_k=config.moe_top_k,

58 + hidden_size=config.hidden_size,

59 + intermediate_size=config.moe_intermediate_size,

60 + reduce_results=False,

61 + renormalize=config.norm_expert_weight,

62 + quant_config=quant_config,

63 + prefix=f"{prefix}.experts")

64 + self.gate = ReplicatedLinear(config.hidden_size,

65 + config.moe_num_experts,

66 + bias=False,

67 + quant_config=None,

68 + prefix=f"{prefix}.gate")

69 +

70 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

71 + orig_shape = hidden_states.shape

72 + hidden_dim = hidden_states.shape[-1]

73 + hidden_states = hidden_states.view(-1, hidden_dim)

74 +

75 + router_logits, _ = self.gate(hidden_states)

76 +

77 + final_hidden_states = self.experts(hidden_states=hidden_states,

78 + router_logits=router_logits)

79 + if self.tp_size > 1:

80 + final_hidden_states = tensor_model_parallel_all_reduce(

81 + final_hidden_states)

82 +

83 + return final_hidden_states.view(orig_shape)

84 +

85 +

86 +class Step3TextMLP(nn.Module):

87 +

88 + def __init__(

89 + self,

90 + hidden_size: int,

91 + intermediate_size: int,

92 + hidden_act: str,

93 + quant_config: Optional[QuantizationConfig] = None,

94 + prefix: str = "",

95 + ) -> None:

96 + super().__init__()

97 + self.gate_up_proj = MergedColumnParallelLinear(

98 + hidden_size, [intermediate_size] * 2,

99 + bias=False,

100 + quant_config=quant_config,

101 + prefix=f"{prefix}.gate_up_proj")

102 + self.down_proj = RowParallelLinear(intermediate_size,

103 + hidden_size,

104 + bias=False,

105 + quant_config=quant_config,

106 + prefix=f"{prefix}.down_proj")

107 + if hidden_act != "silu":

108 + raise ValueError(f"Unsupported activation: {hidden_act}. "

109 + "Only silu is supported for now.")

110 + self.act_fn = SiluAndMul()

111 + self.hidden_size = hidden_size

112 +

113 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

114 + gate_up, _ = self.gate_up_proj(hidden_states)

115 + intermediate_act = self.act_fn(gate_up)

116 + output, _ = self.down_proj(intermediate_act)

117 + return output

118 +

119 +

120 +class Step3TextAttention(nn.Module):

121 +

122 + def __init__(

123 + self,

124 + hidden_size: int,

125 + num_heads: int,

126 + num_kv_heads: int,

127 + norm_eps: float,

128 + rope_theta: int,

129 + share_q_dim: Optional[int] = None,

130 + rope_scaling: Optional[dict[str, Any]] = None,

131 + max_position_embedding: int = 8192,

132 + head_dim: int = 256,

133 + cache_config: Optional[CacheConfig] = None,

134 + quant_config: Optional[QuantizationConfig] = None,

135 + prefix: str = "",

136 + ):

137 + super().__init__()

138 + self.hidden_size = hidden_size

139 + tp_size = get_tensor_model_parallel_world_size()

140 +

141 + self.total_num_heads = num_heads

142 + assert self.total_num_heads % tp_size == 0

143 + self.num_heads = self.total_num_heads // tp_size

144 +

145 + if num_kv_heads != 1:

146 + raise ValueError(f"Step3TextAttention num_kv_heads must be 1, "

147 + f"but got {num_kv_heads}.")

148 + self.num_kv_heads = num_kv_heads

149 +

150 + self.head_dim = head_dim

151 + self.kv_size = self.num_kv_heads * self.head_dim

152 + self.q_size = share_q_dim if share_q_dim else self.head_dim

153 +

154 + self.qkv_proj = ReplicatedLinear(

155 + hidden_size,

156 + self.q_size + self.kv_size * 2,

157 + bias=False,

158 + quant_config=quant_config,

159 + prefix=f"{prefix}.qkv_proj",

160 + )

161 +

162 + self.o_proj = RowParallelLinear(

163 + self.total_num_heads * self.head_dim,

164 + hidden_size,

165 + bias=False,

166 + quant_config=quant_config,

167 + prefix=f"{prefix}.o_proj",

168 + )

169 + self.inter_norm = RMSNorm(self.q_size, eps=norm_eps)

170 + self.wq = ColumnParallelLinear(

171 + self.q_size,

172 + self.head_dim * self.total_num_heads,

173 + bias=False,

174 + quant_config=quant_config,

175 + prefix=f"{prefix}.wq",

176 + )

177 + self.rotary_emb = get_rope(self.head_dim,

178 + rotary_dim=self.head_dim,

179 + max_position=max_position_embedding,

180 + base=rope_theta,

181 + rope_scaling=rope_scaling)

182 + scaling = self.head_dim**-0.5

183 + self.attn = Attention(self.num_heads,

184 + self.head_dim,

185 + scaling,

186 + self.num_kv_heads,

187 + cache_config=cache_config,

188 + prefix=f"{prefix}.attn")

189 +

190 + def forward(self, positions: torch.Tensor,

191 + hidden_states: torch.Tensor) -> torch.Tensor:

192 + qkv, _ = self.qkv_proj(hidden_states)

193 + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

194 + q = self.inter_norm(q)

195 + q = self.wq(q)[0]

196 + q, k = self.rotary_emb(positions, q, k)

197 + attn_output = self.attn(q, k, v)

198 + residual, _ = self.o_proj(attn_output)

199 + return residual

200 +

201 +

202 +class Step3TextDecoderLayer(nn.Module):

203 +

204 + def __init__(self,

205 + config: ModelConfig,

206 + cache_config: Optional[CacheConfig] = None,

207 + quant_config: Optional[QuantizationConfig] = None,

208 + prefix: str = "") -> None:

209 + super().__init__()

210 + config = config.hf_config

211 + self.hidden_size = config.hidden_size

212 + rope_scaling = getattr(config, "rope_scaling", None)

213 +

214 + self.self_attn = Step3TextAttention(

215 + hidden_size=self.hidden_size,

216 + num_heads=config.num_attention_heads,

217 + num_kv_heads=1,

218 + cache_config=cache_config,

219 + quant_config=quant_config,

220 + norm_eps=config.rms_norm_eps,

221 + max_position_embedding=config.max_position_embedding,

222 + head_dim=config.head_dim,

223 + share_q_dim=config.share_q_dim,

224 + rope_theta=config.rope_theta,

225 + rope_scaling=rope_scaling,

226 + prefix=f"{prefix}.self_attn")

227 +

228 + layer_idx = int(prefix.split("layers.")[1].split(".")[0])

229 + moe_layers_enum = getattr(config, "moe_layers_enum", None)

230 + if moe_layers_enum is not None:

231 + moe_layers_idx = [

232 + int(i) for i in moe_layers_enum.strip().split(',')

233 + ]

234 + else:

235 + # Default to 1dense.

236 + moe_layers_idx = [i for i in range(1, config.num_hidden_layers)]

237 +

238 + if layer_idx in moe_layers_idx:

239 + self.moe = FusedMoEBlock(config=config,

240 + quant_config=quant_config,

241 + prefix=f"{prefix}.moe")

242 + self.share_expert = Step3TextMLP(

243 + hidden_size=self.hidden_size,

244 + intermediate_size=config.share_expert_dim,

245 + hidden_act="silu",

246 + quant_config=quant_config,

247 + prefix=f"{prefix}.share_expert")

248 + self.use_moe = True

249 + else:

250 + self.mlp = Step3TextMLP(hidden_size=config.hidden_size,

251 + intermediate_size=config.intermediate_size,

252 + hidden_act="silu",

253 + quant_config=quant_config,

254 + prefix=f"{prefix}.mlp")

255 + self.use_moe = False

256 + self.input_layernorm = RMSNorm(config.hidden_size,

257 + eps=config.rms_norm_eps)

258 + self.post_attention_layernorm = RMSNorm(config.hidden_size,

259 + eps=config.rms_norm_eps)

260 +

261 + def forward(

262 + self, positions: torch.Tensor, hidden_states: torch.Tensor,

263 + residual: Optional[torch.Tensor]

264 + ) -> tuple[torch.Tensor, torch.Tensor]:

265 + if residual is None:

266 + residual = hidden_states

267 + hidden_states = self.input_layernorm(hidden_states)

268 + else:

269 + hidden_states, residual = self.input_layernorm(

270 + hidden_states, residual)

271 +

272 + hidden_states = self.self_attn(

273 + positions=positions,

274 + hidden_states=hidden_states,

275 + )

276 +

277 + hidden_states, residual = self.post_attention_layernorm(

278 + hidden_states, residual)

279 +

280 + if self.use_moe:

281 + share_output = self.share_expert(hidden_states)

282 + moe_output = self.moe(hidden_states)

283 + hidden_states = share_output + moe_output

284 + else:

285 + hidden_states = self.mlp(hidden_states)

286 +

287 + return hidden_states, residual

288 +

289 +

290 +@support_torch_compile

291 +class Step3TextModel(nn.Module):

292 +

293 + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None:

294 + super().__init__()

295 + config = vllm_config.model_config.hf_config

296 + cache_config = vllm_config.cache_config

297 + quant_config = vllm_config.quant_config

298 + self.vocab_size = config.vocab_size

299 + self.config = config

300 +

301 + if get_pp_group().is_first_rank or (config.tie_word_embeddings

302 + and get_pp_group().is_last_rank):

303 + self.embed_tokens = VocabParallelEmbedding(

304 + self.vocab_size,

305 + config.hidden_size,

306 + )

307 + else:

308 + self.embed_tokens = PPMissingLayer()

309 +

310 + self.start_layer, self.end_layer, self.layers = make_layers(

311 + config.num_hidden_layers,

312 + lambda prefix: Step3TextDecoderLayer(config=vllm_config.

313 + model_config,

314 + cache_config=cache_config,

315 + quant_config=quant_config,

316 + prefix=prefix),

317 + prefix=f"{prefix}.layers",

318 + )

319 + if get_pp_group().is_last_rank:

320 + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

321 + else:

322 + self.norm = PPMissingLayer()

323 +

324 + self.make_empty_intermediate_tensors = (

325 + make_empty_intermediate_tensors_factory(["hidden_states"],

326 + config.hidden_size))

327 +

328 + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

329 + return self.embed_tokens(input_ids)

330 +

331 + def forward(

332 + self,

333 + input_ids: torch.Tensor,

334 + positions: torch.Tensor,

335 + intermediate_tensors: Optional[IntermediateTensors] = None,

336 + inputs_embeds: Optional[torch.Tensor] = None,

337 + ) -> torch.Tensor:

338 + if get_pp_group().is_first_rank:

339 + if inputs_embeds is not None:

340 + hidden_states = inputs_embeds

341 + else:

342 + hidden_states = self.get_input_embeddings(input_ids)

343 + residual = None

344 + else:

345 + assert intermediate_tensors is not None

346 + hidden_states = intermediate_tensors["hidden_states"]

347 + residual = intermediate_tensors["residual"]

348 +

349 + for i in range(self.start_layer, self.end_layer):

350 + layer = self.layers[i]

351 + hidden_states, residual = layer(positions, hidden_states, residual)

352 +

353 + if not get_pp_group().is_last_rank:

354 + return IntermediateTensors({

355 + "hidden_states": hidden_states,

356 + "residual": residual,

357 + })

358 +

359 + hidden_states, _ = self.norm(hidden_states, residual)

360 + return hidden_states

361 +

362 +

363 +class Step3TextForCausalLM(nn.Module, SupportsPP):

364 +

365 + def __init__(

366 + self,

367 + *,

368 + vllm_config: VllmConfig,

369 + prefix: str = "",

370 + ):

371 + super().__init__()

372 + config = vllm_config.model_config.hf_config

373 + lora_config = vllm_config.lora_config

374 + self.config = config

375 + self.vllm_config = vllm_config

376 +

377 + self.model = Step3TextModel(vllm_config=vllm_config, prefix=prefix)

378 +

379 + if get_pp_group().is_last_rank:

380 + self.unpadded_vocab_size = config.vocab_size

381 + if lora_config:

382 + self.unpadded_vocab_size += lora_config.lora_extra_vocab_size

383 + self.lm_head = ParallelLMHead(

384 + self.unpadded_vocab_size,

385 + config.hidden_size,

386 + org_num_embeddings=config.vocab_size,

387 + padding_size=DEFAULT_VOCAB_PADDING_SIZE

388 + if not lora_config else lora_config.lora_vocab_padding_size,

389 + )

390 + self.logits_processor = LogitsProcessor(self.unpadded_vocab_size,

391 + config.vocab_size)

392 + self.sampler = get_sampler()

393 + else:

394 + self.lm_head = PPMissingLayer()

395 +

396 + self.make_empty_intermediate_tensors = (

397 + self.model.make_empty_intermediate_tensors)

398 +

399 + def forward(self,

400 + input_ids: torch.Tensor,

401 + positions: torch.Tensor,

402 + intermediate_tensors: Optional[IntermediateTensors] = None,

403 + inputs_embeds: Optional[torch.Tensor] = None):

404 + hidden_states = self.model(input_ids, positions, intermediate_tensors,

405 + inputs_embeds)

406 + return hidden_states

407 +

408 + def compute_logits(self, hidden_states: torch.Tensor,

409 + sampling_metadata: SamplingMetadata) -> torch.Tensor:

410 + logits = self.logits_processor(self.lm_head, hidden_states,

411 + sampling_metadata)

412 + return logits

413 +

414 + def sample(

415 + self,

416 + logits: Optional[torch.Tensor],

417 + sampling_metadata: SamplingMetadata,

418 + ) -> Optional[SamplerOutput]:

419 + next_tokens = self.sampler(logits, sampling_metadata)

420 + return next_tokens

421 +

422 + def load_weights(self, weights: Iterable[tuple[str,

423 + torch.Tensor]]) -> set[str]:

424 + qkv_params_mapping = [

425 + # (param_name, shard_name, relative_start_idx, relative_end_idx)

426 + (".qkv_proj", ".q_proj", 0, self.config.share_q_dim /

427 + (self.config.share_q_dim + self.config.head_dim * 2)),

428 + (".qkv_proj", ".k_proj", self.config.share_q_dim /

429 + (self.config.share_q_dim + self.config.head_dim * 2),

430 + (self.config.share_q_dim + self.config.head_dim) /

431 + (self.config.share_q_dim + self.config.head_dim * 2)),

432 + (".qkv_proj", ".v_proj",

433 + (self.config.share_q_dim + self.config.head_dim) /

434 + (self.config.share_q_dim + self.config.head_dim * 2),

435 + (self.config.share_q_dim + self.config.head_dim * 2) /

436 + (self.config.share_q_dim + self.config.head_dim * 2)),

437 + ]

438 + stacked_params_mapping = [

439 + # (param_name, shard_name, shard_id)

440 + (".gate_up_proj", ".gate_proj", 0),

441 + (".gate_up_proj", ".up_proj", 1),

442 + ]

443 + params_dict = dict(self.named_parameters())

444 + loaded_params: set[str] = set()

445 +

446 + expert_params_mapping = [

447 + (".moe.experts.w13_weight", ".moe.gate_proj.weight", "w1"),

448 + (".moe.experts.w13_weight", ".moe.up_proj.weight", "w3"),

449 + (".moe.experts.w2_weight", ".moe.down_proj.weight", "w2")

450 + ]

451 +

452 + disable_moe_stacked_params = [

453 + data[1] for data in expert_params_mapping

454 + ]

455 +

456 + for name, loaded_weight in weights:

457 + for (param_name, weight_name, shard_id) in stacked_params_mapping:

458 + if weight_name not in name:

459 + continue

460 + if any(disable_moe_stacked_param in name

461 + for disable_moe_stacked_param in

462 + disable_moe_stacked_params):

463 + continue

464 + name = name.replace(weight_name, param_name)

465 + if is_pp_missing_parameter(name, self):

466 + continue

467 + param = params_dict[name]

468 + weight_loader = param.weight_loader

469 + weight_loader(param, loaded_weight, shard_id)

470 + loaded_params.add(name)

471 + break

472 + else:

473 + for mapping in expert_params_mapping:

474 + param_name, weight_name, shard_id = mapping

475 + if weight_name not in name:

476 + continue

477 + name = name.replace(weight_name, param_name)

478 + # Skip layers on other devices.

479 + if is_pp_missing_parameter(name, self):

480 + continue

481 + # Skip loading extra bias for GPTQ models.

482 + if ((name.endswith(".bias") or name.endswith("_bias"))

483 + and name not in params_dict):

484 + continue

485 + param = params_dict[name]

486 + weight_loader = param.weight_loader

487 + for expert_id in range(loaded_weight.shape[0]):

488 + loaded_weight_expert = loaded_weight[expert_id]

489 + weight_loader(param,

490 + loaded_weight_expert,

491 + name,

492 + shard_id=shard_id,

493 + expert_id=expert_id)

494 + loaded_params.add(name)

495 + break

496 + else:

497 + for (param_name, weight_name, start_idx,

498 + end_idx) in qkv_params_mapping:

499 + if weight_name not in name:

500 + continue

501 + name = name.replace(weight_name, param_name)

502 + if is_pp_missing_parameter(name, self):

503 + continue

504 + param = params_dict[name]

505 + dim = param.shape[param.output_dim]

506 + begin_idx = int(start_idx * dim)

507 + end_idx = int(end_idx * dim)

508 + param_slice = param.narrow(param.output_dim, begin_idx,

509 + end_idx - begin_idx)

510 + param_slice.copy_(loaded_weight)

511 + loaded_params.add(name)

512 + break

513 + else:

514 + if is_pp_missing_parameter(name, self):

515 + continue

516 + param = params_dict[name]

517 + weight_loader = getattr(param, "weight_loader",

518 + default_weight_loader)

519 + weight_loader(param, loaded_weight)

520 + loaded_params.add(name)

521 + return loaded_params

vllm/model_executor/models/step3_vl.py new

+1052

@@ -0,0 +1,1052 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import math

4 +from collections.abc import Iterable, Mapping, Sequence

5 +from functools import cached_property

6 +from itertools import product

7 +from math import ceil, sqrt

8 +from typing import Any, Literal, Optional, TypedDict, Union

9 +

10 +import numpy as np

11 +import torch

12 +import torch.nn as nn

13 +import torch.nn.functional as F

14 +from PIL import Image

15 +from torchvision import transforms

16 +from torchvision.transforms.functional import InterpolationMode

17 +from transformers import BatchFeature, PretrainedConfig, TensorType

18 +

19 +from vllm.config import VllmConfig

20 +from vllm.distributed import get_tensor_model_parallel_world_size

21 +from vllm.model_executor.layers.activation import get_act_fn

22 +from vllm.model_executor.layers.linear import (ColumnParallelLinear,

23 + QKVParallelLinear,

24 + RowParallelLinear)

25 +from vllm.model_executor.layers.quantization import QuantizationConfig

26 +from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler

27 +from vllm.model_executor.sampling_metadata import SamplingMetadata

28 +from vllm.multimodal import MULTIMODAL_REGISTRY

29 +from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,

30 + MultiModalKwargs, NestedTensors)

31 +from vllm.multimodal.parse import ImageSize, MultiModalDataItems

32 +from vllm.multimodal.processing import (BaseMultiModalProcessor,

33 + BaseProcessingInfo, PromptReplacement,

34 + PromptUpdate, PromptUpdateDetails)

35 +from vllm.multimodal.profiling import BaseDummyInputsBuilder

36 +from vllm.sequence import IntermediateTensors

37 +from vllm.transformers_utils.configs import Step3VisionEncoderConfig

38 +from vllm.transformers_utils.tokenizer import AnyTokenizer

39 +

40 +from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

41 +from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

42 + init_vllm_registered_model, maybe_prefix,

43 + merge_multimodal_embeddings)

44 +

45 +

46 +class Step3VLImagePixelInputs(TypedDict):

47 + type: Literal["pixel_values"]

48 + pixel_values: torch.Tensor

49 + patch_pixel_values: Optional[torch.Tensor]

50 + num_patches: list[int]

51 +

52 +

53 +class Step3VLImageEmbeddingInputs(TypedDict):

54 + type: Literal["image_embeds"]

55 + image_embeds: torch.Tensor

56 +

57 +

58 +Step3VLImageInputs = Union[Step3VLImagePixelInputs,

59 + Step3VLImageEmbeddingInputs]

60 +

61 +ImageWithPatches = tuple[Image.Image, list[Image.Image], list[int] | None]

62 +

63 +MAX_IMAGE_SIZE: int = 3024

64 +

65 +

66 +class Step3VisionProcessor:

67 +

68 + def __init__(self, size, interpolation_mode="bicubic", patch_size=None):

69 + mean = [0.48145466, 0.4578275, 0.40821073]

70 + std = [0.26862954, 0.26130258, 0.27577711]

71 + patch_size = patch_size if patch_size is not None else size

72 +

73 + self.transform = transforms.Compose([

74 + transforms.ToTensor(),

75 + transforms.Normalize(mean, std),

76 + transforms.Resize(

77 + (size, size),

78 + interpolation=InterpolationMode.BICUBIC if interpolation_mode

79 + == "bicubic" else InterpolationMode.BILINEAR,

80 + antialias=True),

81 + ])

82 +

83 + self.patch_transform = transforms.Compose([

84 + transforms.ToTensor(),

85 + transforms.Normalize(mean, std),

86 + transforms.Resize(

87 + (patch_size, patch_size),

88 + interpolation=InterpolationMode.BICUBIC if interpolation_mode

89 + == "bicubic" else InterpolationMode.BILINEAR,

90 + antialias=True),

91 + ]) if patch_size is not None else None

92 +

93 + def __call__(self, image, is_patch=False):

94 + if is_patch:

95 + return {"pixel_values": self.patch_transform(image).unsqueeze(0)}

96 + else:

97 + return {"pixel_values": self.transform(image).unsqueeze(0)}

98 +

99 +

100 +class ImagePatcher:

101 +

102 + def determine_window_size(self, long: int, short: int) -> int:

103 + if long <= 728:

104 + return short if long / short > 1.5 else 0

105 + return min(short, 504) if long / short > 4 else 504

106 +

107 + def slide_window(

108 + self,

109 + width: int,

110 + height: int,

111 + sizes: list[tuple[int, int]],

112 + steps: list[tuple[int, int]],

113 + img_rate_thr: float = 0.6,

114 + ) -> tuple[list[tuple[int, int, int, int]], tuple[int, int]]:

115 + assert 1 >= img_rate_thr >= 0, "The `in_rate_thr` should lie in 0~1"

116 + windows = []

117 + # Sliding windows.

118 + for size, step in zip(sizes, steps):

119 + size_w, size_h = size

120 + step_w, step_h = step

121 +

122 + x_num = 1 if width <= size_w else ceil((width - size_w) / step_w +

123 + 1)

124 + x_start = [step_w * i for i in range(x_num)]

125 + if len(x_start) > 1 and x_start[-1] + size_w > width:

126 + x_start[-1] = width - size_w

127 +

128 + y_num = 1 if height <= size_h else ceil((height - size_h) /

129 + step_h + 1)

130 + y_start = [step_h * i for i in range(y_num)]

131 + if len(y_start) > 1 and y_start[-1] + size_h > height:

132 + y_start[-1] = height - size_h

133 +

134 + start = np.array(list(product(y_start, x_start)), dtype=int)

135 + start[:, [0, 1]] = start[:, [1, 0]]

136 + windows.append(np.concatenate([start, start + size], axis=1))

137 + windows = np.concatenate(windows, axis=0)

138 +

139 + return [(int(box[0]), int(box[1]), int(box[2] - box[0]),

140 + int(box[3] - box[1])) for box in windows], (x_num, y_num)

141 +

142 + def square_pad(self, img: Image.Image) -> Image.Image:

143 + w, h = img.size

144 + if w == h:

145 + return img

146 + size = max(w, h)

147 + padded = Image.new(img.mode, (size, size), 0)

148 + padded.paste(img, (0, 0))

149 + return padded

150 +

151 + def get_image_size_for_padding(self, img_width: int,

152 + img_height: int) -> tuple[int, int]:

153 + ratio = img_width / img_height

154 + if min(img_height, img_width) < 32 and (ratio > 4 or ratio < 1 / 4):

155 + new_size = max(img_height, img_width)

156 + return new_size, new_size

157 + return img_width, img_height

158 +

159 + def get_image_size_for_preprocess(self, img_width: int,

160 + img_height: int) -> tuple[int, int]:

161 +

162 + if max(img_height, img_width) > MAX_IMAGE_SIZE:

163 + scale_factor = MAX_IMAGE_SIZE / max(img_height, img_width)

164 + img_width = int(img_width * scale_factor)

165 + img_height = int(img_height * scale_factor)

166 + return img_width, img_height

167 +

168 + def get_image_size_for_crop(self, img_width: int, img_height: int,

169 + window_size: int):

170 + w_ratio = img_width / window_size

171 + h_ratio = img_height / window_size

172 +

173 + if w_ratio < 1:

174 + width_new = img_width

175 + else:

176 + decimal_w = w_ratio - img_width // window_size

177 + w_ratio = int(w_ratio) + 1 if decimal_w > 0.2 else int(w_ratio)

178 + width_new = window_size * w_ratio

179 + if h_ratio < 1:

180 + height_new = img_height

181 + else:

182 + decimal_h = h_ratio - img_height // window_size

183 + h_ratio = int(h_ratio) + 1 if decimal_h > 0.2 else int(h_ratio)

184 + height_new = window_size * h_ratio

185 + return int(width_new), int(height_new)

186 +

187 + def patch_crop(self, img: Image.Image, i: int, j: int, th: int, tw: int):

188 + target = img.crop((j, i, j + tw, i + th))

189 + return target

190 +

191 + def get_num_patches(self, img_width: int,

192 + img_height: int) -> tuple[int, int]:

193 + img_width, img_height = self.get_image_size_for_padding(

194 + img_width, img_height)

195 + img_width, img_height = self.get_image_size_for_preprocess(

196 + img_width, img_height)

197 + window_size = self.determine_window_size(max(img_height, img_width),

198 + min(img_height, img_width))

199 + if window_size == 0:

200 + return 0, 0

201 + else:

202 + img_width, img_height = self.get_image_size_for_crop(

203 + img_width, img_height, window_size)

204 + center_list, (x_num, y_num) = self.slide_window(

205 + img_width, img_height, [(window_size, window_size)],

206 + [(window_size, window_size)])

207 + full_rows = (len(center_list) - 1) // x_num + 1

208 + if len(center_list) > 0 and len(center_list) % x_num == 0:

209 + full_rows -= 1

210 + return len(center_list), full_rows

211 +

212 + def __call__(

213 + self, img: Image.Image

214 + ) -> tuple[Image.Image, list[Image.Image], list[bool] | None]:

215 + img_width, img_height = img.size

216 + new_img_width, new_img_height = self.get_image_size_for_padding(

217 + img_width, img_height)

218 + if new_img_width != img_width or new_img_height != img_height:

219 + img = self.square_pad(img)

220 + img_width, img_height = img.size

221 +

222 + new_img_width, new_img_height = self.get_image_size_for_preprocess(

223 + img_width, img_height)

224 + img = img.resize((new_img_width, new_img_height),

225 + Image.Resampling.BILINEAR)

226 + window_size = self.determine_window_size(

227 + max(new_img_height, new_img_width),

228 + min(new_img_height, new_img_width))

229 +

230 + if window_size == 0:

231 + return img, [], None

232 + else:

233 + new_img_width, new_img_height = self.get_image_size_for_crop(

234 + new_img_width, new_img_height, window_size)

235 + if (new_img_width, new_img_height) != (img_width, img_height):

236 + img_for_crop = img.resize((new_img_width, new_img_height),

237 + Image.Resampling.BILINEAR)

238 + else:

239 + img_for_crop = img

240 +

241 + patches = []

242 + newlines = []

243 + center_list, (x_num, y_num) = self.slide_window(

244 + new_img_width, new_img_height, [(window_size, window_size)],

245 + [(window_size, window_size)])

246 + for patch_id, center_lf_point in enumerate(center_list):

247 + x, y, patch_w, patch_h = center_lf_point

248 + big_patch = self.patch_crop(img_for_crop, y, x, patch_h,

249 + patch_w)

250 + patches.append(big_patch)

251 + if (patch_id + 1) % x_num == 0:

252 + newlines.append(patch_id)

253 +

254 + if newlines and newlines[-1] == len(patches) - 1:

255 + newlines.pop()

256 +

257 + return img, patches, [i in newlines for i in range(len(patches))

258 + ] if len(patches) > 0 else None

259 +

260 +

261 +class Step3VLProcessor:

262 +

263 + def __init__(

264 + self,

265 + config: PretrainedConfig,

266 + tokenizer: AnyTokenizer,

267 + ) -> None:

268 + super().__init__()

269 +

270 + self.config = config

271 + self.tokenizer = tokenizer

272 +

273 + self.image_size = 728

274 + self.patch_size = 504

275 + self.image_preprocessor = Step3VisionProcessor(self.image_size,

276 + "bilinear",

277 + self.patch_size)

278 +

279 + self.num_image_feature_size = 169

280 + self.num_patch_feature_size = 81

281 + self.image_token = "<im_patch>"

282 + self.image_feature_placeholder = (self.image_token *

283 + self.num_image_feature_size)

284 + self.patch_feature_placeholder = (self.image_token *

285 + self.num_patch_feature_size)

286 +

287 + self.patcher = ImagePatcher()

288 +

289 + @property

290 + def image_token_id(self) -> int:

291 + return self.tokenizer.get_vocab()[self.image_token]

292 +

293 + def get_num_image_tokens(self, img_width: int, img_height: int) -> int:

294 + num_patches, num_newlines = self.patcher.get_num_patches(

295 + img_width, img_height)

296 +

297 + return num_patches * (

298 + self.num_patch_feature_size +

299 + 2) + self.num_image_feature_size + 2 + num_newlines

300 +

301 + def _split_images(self,

302 + images: list[Image.Image]) -> list[ImageWithPatches]:

303 + result = []

304 + for img in images:

305 + result.append(self.patcher(img))

306 + return result

307 +

308 + def _convert_images_to_pixel_values(

309 + self,

310 + images: list[Image.Image],

311 + is_patch: bool = False,

312 + ) -> list[torch.Tensor]:

313 + return [

314 + self.image_preprocessor(img, is_patch=is_patch)["pixel_values"]

315 + for img in images

316 + ]

317 +

318 + def _get_patch_repl(

319 + self,

320 + num_patches: int,

321 + patch_newline_mask: list[bool] | None,

322 + ) -> tuple[str, list[int]]:

323 + text = ""

324 + token_ids = []

325 + for i in range(num_patches):

326 + assert len(patch_newline_mask) == num_patches

327 + text += f"<patch_start>{self.patch_feature_placeholder}<patch_end>"

328 + token_ids.extend(

329 + [self.tokenizer.convert_tokens_to_ids("<patch_start>")] +

330 + [self.image_token_id] * self.num_patch_feature_size +

331 + [self.tokenizer.convert_tokens_to_ids("<patch_end>")])

332 + if patch_newline_mask and patch_newline_mask[i]:

333 + text += "<patch_newline>"

334 + token_ids.append(

335 + self.tokenizer.convert_tokens_to_ids("<patch_newline>"))

336 + return text, token_ids

337 +

338 + def _get_image_repl(

339 + self,

340 + num_images: int,

341 + ) -> tuple[str, list[int]]:

342 + text = f"<im_start>{self.image_feature_placeholder}<im_end>"

343 + token_ids = [

344 + self.tokenizer.convert_tokens_to_ids("<im_start>")

345 + ] + [self.image_token_id] * self.num_image_feature_size + [

346 + self.tokenizer.convert_tokens_to_ids("<im_end>")

347 + ]

348 + return text * num_images, token_ids * num_images

349 +

350 + def _get_image_repl_features(

351 + self,

352 + num_images: int,

353 + num_patches: int,

354 + patch_new_line_idx: Optional[list[bool]],

355 + ) -> tuple[str, list[int]]:

356 + if num_patches > 0:

357 + patch_repl, patch_repl_ids = self._get_patch_repl(

358 + num_patches, patch_new_line_idx)

359 + else:

360 + patch_repl = ""

361 + patch_repl_ids = []

362 + image_repl, image_repl_ids = self._get_image_repl(num_images)

363 + return patch_repl + image_repl, patch_repl_ids + image_repl_ids

364 +

365 + def replace_placeholder(self, text: str, placeholder: str,

366 + repls: list[str]) -> str:

367 + parts = text.split(placeholder)

368 +

369 + if len(parts) - 1 != len(repls):

370 + raise ValueError(

371 + "The number of placeholders does not match the number of replacements." # noqa: E501

372 + )

373 +

374 + result = [parts[0]]

375 + for i, repl in enumerate(repls):

376 + result.append(repl)

377 + result.append(parts[i + 1])

378 +

379 + return "".join(result)

380 +

381 + def __call__(

382 + self,

383 + text: Optional[Union[str, list[str]]] = None,

384 + images: Optional[Union[Image.Image, list[Image.Image]]] = None,

385 + return_tensors: Optional[Union[str, TensorType]] = None,

386 + ) -> BatchFeature:

387 + if text is None:

388 + text = []

389 + if not isinstance(text, list):

390 + text = [text]

391 + if images is None:

392 + images = []

393 + if not isinstance(images, list):

394 + images = [images]

395 +

396 + if len(images) == 0:

397 + image_inputs = {}

398 + text_inputs = self.tokenizer(text)

399 + else:

400 + splitted_images_data = self._split_images(images)

401 + pixel_values_lst = []

402 + patch_pixel_values_lst = []

403 + patch_newline_mask_lst = []

404 + image_repl_str_lst = []

405 + image_repl_ids_lst = []

406 + num_patches = []

407 + for raw_img, img_patches, patch_newline_mask in splitted_images_data: # noqa: E501

408 + pixel_values_lst.extend(

409 + self._convert_images_to_pixel_values([raw_img]))

410 +

411 + if len(img_patches) > 0:

412 + patch_pixel_values_lst.extend(

413 + self._convert_images_to_pixel_values(img_patches,

414 + is_patch=True))

415 + num_patches.append(len(img_patches))

416 +

417 + image_repl_str, image_repl_ids = self._get_image_repl_features(

418 + 1, len(img_patches), patch_newline_mask)

419 + image_repl_str_lst.append(image_repl_str)

420 + image_repl_ids_lst.extend(image_repl_ids)

421 +

422 + if patch_newline_mask is not None:

423 + patch_newline_mask_lst.extend(patch_newline_mask)

424 +

425 + image_inputs = {

426 + "pixel_values": torch.cat(pixel_values_lst),

427 + "num_patches": num_patches,

428 + }

429 + if patch_pixel_values_lst:

430 + image_inputs["patch_pixel_values"] = torch.cat(

431 + patch_pixel_values_lst)

432 + if patch_newline_mask_lst:

433 + image_inputs["patch_newline_mask"] = torch.tensor(

434 + patch_newline_mask_lst, dtype=torch.bool)

435 +

436 + text = [

437 + self.replace_placeholder(t, self.image_token,

438 + image_repl_str_lst) for t in text

439 + ]

440 + text_inputs = self.tokenizer(text)

441 +

442 + return BatchFeature(

443 + {

444 + **text_inputs,

445 + **image_inputs,

446 + },

447 + tensor_type=return_tensors,

448 + )

449 +

450 +

451 +class Step3VLProcessingInfo(BaseProcessingInfo):

452 +

453 + def get_hf_processor(self) -> Step3VLProcessor:

454 + return Step3VLProcessor(

455 + self.get_hf_config(),

456 + self.get_tokenizer(),

457 + )

458 +

459 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

460 + return {"image": None}

461 +

462 + def get_max_image_tokens(self) -> int:

463 + hf_processor = self.get_hf_processor()

464 + return hf_processor.get_num_image_tokens(

465 + self.get_image_size_with_most_features().width,

466 + self.get_image_size_with_most_features().height)

467 +

468 + def get_mm_max_tokens_per_item(

469 + self,

470 + seq_len: int,

471 + mm_counts: Mapping[str, int],

472 + ) -> Mapping[str, int]:

473 + return {"image": self.get_max_image_tokens()}

474 +

475 + def get_image_size_with_most_features(self) -> ImageSize:

476 + return ImageSize(3024, 3024)

477 +

478 + def get_num_mm_tokens(self, mm_data: MultiModalDataDict) -> int:

479 + if len(mm_data) != 1 or "image" not in mm_data:

480 + raise ValueError(

481 + "mm_data could only contain one key 'image' for steo1o")

482 +

483 + image_data = mm_data["image"]

484 + if not isinstance(image_data, (list, tuple)):

485 + image_data = [image_data]

486 +

487 + return sum(self.get_hf_processor().get_num_image_tokens(

488 + img.width, img.height) for img in image_data)

489 +

490 +

491 +class Step3VLDummyInputsBuilder(BaseDummyInputsBuilder[Step3VLProcessingInfo]):

492 +

493 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

494 + num_images = mm_counts.get("image", 0)

495 + return "<im_patch>" * num_images

496 +

497 + def get_dummy_mm_data(

498 + self,

499 + seq_len: int,

500 + mm_counts: Mapping[str, int],

501 + ) -> MultiModalDataDict:

502 + target_width, target_height = \

503 + self.info.get_image_size_with_most_features()

504 + num_images = mm_counts.get("image", 0)

505 +

506 + return {

507 + "image":

508 + self._get_dummy_images(width=target_width,

509 + height=target_height,

510 + num_images=num_images)

511 + }

512 +

513 +

514 +class Step3VLMultiModalProcessor(BaseMultiModalProcessor[Step3VLProcessingInfo]

515 + ):

516 +

517 + def _get_prompt_updates(

518 + self,

519 + mm_items: MultiModalDataItems,

520 + hf_processor_mm_kwargs: Mapping[str, Any],

521 + out_mm_kwargs: MultiModalKwargs,

522 + ) -> Sequence[PromptUpdate]:

523 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

524 + image_placeholder_token_id = hf_processor.image_token_id

525 + batch_num_patches = out_mm_kwargs["num_patches"].tolist()

526 +

527 + def get_replacement_step1o(item_idx: int):

528 + img_out = out_mm_kwargs.get_item("image", item_idx)

529 + num_patches = batch_num_patches[item_idx]

530 + if num_patches > 0:

531 + patch_newline_mask = img_out["patch_newline_mask"].data.tolist(

532 + )

533 + image_repl_ids = hf_processor._get_image_repl_features(

534 + 1, num_patches, patch_newline_mask)[1]

535 + else:

536 + image_repl_ids = hf_processor._get_image_repl_features(

537 + 1, 0, None)[1]

538 + return PromptUpdateDetails.select_token_id(

539 + seq=image_repl_ids,

540 + embed_token_id=image_placeholder_token_id,

541 + )

542 +

543 + return [

544 + PromptReplacement(

545 + modality="image",

546 + target=[image_placeholder_token_id],

547 + replacement=get_replacement_step1o,

548 + )

549 + ]

550 +

551 + def _get_mm_fields_config(

552 + self,

553 + hf_inputs: BatchFeature,

554 + hf_processor_mm_kwargs: Mapping[str, object],

555 + ) -> Mapping[str, MultiModalFieldConfig]:

556 + num_patches = hf_inputs.get("num_patches", torch.empty(0))

557 +

558 + return dict(

559 + pixel_values=MultiModalFieldConfig.batched("image"),

560 + patch_pixel_values=MultiModalFieldConfig.flat_from_sizes(

561 + "image", num_patches),

562 + num_patches=MultiModalFieldConfig.batched("image"),

563 + patch_newline_mask=MultiModalFieldConfig.flat_from_sizes(

564 + "image", num_patches),

565 + )

566 +

567 +

568 +def get_abs_pos(abs_pos, tgt_size):

569 + dim = abs_pos.size(-1)

570 + abs_pos_new = abs_pos.squeeze(0)

571 + cls_token, old_pos_embed = abs_pos_new[:1], abs_pos_new[1:]

572 +

573 + src_size = int(math.sqrt(abs_pos_new.shape[0] - 1))

574 + tgt_size = int(math.sqrt(tgt_size))

575 + dtype = abs_pos.dtype

576 +

577 + if src_size != tgt_size:

578 + old_pos_embed = old_pos_embed.view(1, src_size, src_size,

579 + dim).permute(0, 3, 1,

580 + 2).contiguous()

581 + old_pos_embed = old_pos_embed.to(torch.float32)

582 + new_pos_embed = F.interpolate(

583 + old_pos_embed,

584 + size=(tgt_size, tgt_size),

585 + mode='bicubic',

586 + antialias=True,

587 + align_corners=False,

588 + ).to(dtype)

589 + new_pos_embed = new_pos_embed.permute(0, 2, 3, 1)

590 + new_pos_embed = new_pos_embed.view(tgt_size * tgt_size, dim)

591 + vision_pos_embed = torch.cat([cls_token, new_pos_embed], dim=0)

592 + vision_pos_embed = vision_pos_embed.view(1, tgt_size * tgt_size + 1,

593 + dim)

594 + return vision_pos_embed

595 + else:

596 + return abs_pos

597 +

598 +

599 +class Step3VisionEmbeddings(nn.Module):

600 +

601 + def __init__(self, config: Step3VisionEncoderConfig):

602 + super().__init__()

603 + self.config = config

604 + self.embed_dim = config.hidden_size

605 + self.image_size = config.image_size

606 + self.patch_size = config.patch_size

607 +

608 + self.class_embedding = nn.Parameter(torch.randn(1, self.embed_dim))

609 +

610 + self.patch_embedding = nn.Conv2d(

611 + in_channels=config.num_channels,

612 + out_channels=self.embed_dim,

613 + kernel_size=self.patch_size,

614 + stride=self.patch_size,

615 + bias=True,

616 + )

617 +

618 + self.num_patches = (self.image_size // self.patch_size)**2

619 + self.pad_tp_size = 4 # hard code for padding

620 + # To load the pretrained weights, we still use P+1 as the seqlen

621 + self.position_embedding = torch.nn.Embedding(self.num_patches + 1,

622 + self.embed_dim)

623 + self.register_buffer("position_ids",

624 + torch.arange(self.num_patches + 1).expand(

625 + (1, -1)),

626 + persistent=False)

627 +

628 + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:

629 + batch_size = pixel_values.shape[0]

630 + patch_embeds = self.patch_embedding(

631 + pixel_values) # shape = [*, width, grid, grid]

632 + patch_embeds = patch_embeds.flatten(2).transpose(1, 2)

633 +

634 + # pad

635 + class_embeds = self.class_embedding.expand(batch_size, 1, -1)

636 + embeddings = torch.cat([class_embeds, patch_embeds], dim=1)

637 + embeddings = embeddings + get_abs_pos(

638 + self.position_embedding(self.position_ids), patch_embeds.size(1))

639 + embeddings = torch.cat([

640 + embeddings[:, 0, :].unsqueeze(1).repeat(1, self.pad_tp_size - 1,

641 + 1), embeddings

642 + ],

643 + dim=1)

644 + return embeddings

645 +

646 +

647 +class Step3VisionAttention(nn.Module):

648 + """Multi-headed attention from 'Attention Is All You Need' paper"""

649 +

650 + def __init__(self,

651 + config,

652 + quant_config: Optional[QuantizationConfig] = None,

653 + prefix: str = ""):

654 + super().__init__()

655 + self.config = config

656 + self.embed_dim = config.hidden_size

657 + self.total_num_heads = config.num_attention_heads

658 + self.head_dim = self.embed_dim // self.total_num_heads

659 +

660 + self.scale = self.head_dim**-0.5

661 +

662 + tp_size = get_tensor_model_parallel_world_size()

663 + assert self.total_num_heads % tp_size == 0

664 + self.num_heads = self.total_num_heads // tp_size

665 + self.qkv_proj = QKVParallelLinear(self.embed_dim,

666 + self.head_dim,

667 + self.total_num_heads,

668 + bias=True,

669 + quant_config=quant_config,

670 + prefix=prefix)

671 + self.out_proj = RowParallelLinear(self.embed_dim,

672 + self.embed_dim,

673 + bias=True,

674 + quant_config=quant_config,

675 + prefix=prefix)

676 +

677 + def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):

678 + return tensor.view(bsz, seq_len, self.num_heads,

679 + self.head_dim).transpose(1, 2).contiguous()

680 +

681 + def forward(

682 + self,

683 + hidden_states: torch.Tensor,

684 + ):

685 + """Input shape: Batch x Time x Channel"""

686 + bsz, tgt_len, _ = hidden_states.size()

687 +

688 + # get query proj

689 + qkv, _ = self.qkv_proj(hidden_states)

690 + q, k, v = qkv.chunk(chunks=3, dim=-1)

691 + q = q.view(bsz, tgt_len, self.num_heads, self.head_dim)

692 + k = k.view(bsz, tgt_len, self.num_heads, self.head_dim)

693 + v = v.view(bsz, tgt_len, self.num_heads, self.head_dim)

694 + q = q.transpose(1, 2)

695 + k = k.transpose(1, 2)

696 + v = v.transpose(1, 2)

697 + attn_output = F.scaled_dot_product_attention(q,

698 + k,

699 + v,

700 + scale=self.scale,

701 + is_causal=False)

702 + attn_output = attn_output.transpose(1, 2).reshape(

703 + bsz, tgt_len, self.num_heads * self.head_dim)

704 +

705 + attn_output, _ = self.out_proj(attn_output)

706 +

707 + return attn_output

708 +

709 +

710 +class Step3VisionMLP(nn.Module):

711 +

712 + def __init__(self,

713 + config,

714 + quant_config: Optional[QuantizationConfig] = None,

715 + prefix: str = ""):

716 + super().__init__()

717 + self.config = config

718 + self.activation_fn = get_act_fn(config.hidden_act)

719 + self.fc1 = ColumnParallelLinear(config.hidden_size,

720 + config.intermediate_size,

721 + bias=True,

722 + quant_config=quant_config,

723 + prefix=prefix)

724 + self.fc2 = RowParallelLinear(config.intermediate_size,

725 + config.hidden_size,

726 + bias=True,

727 + quant_config=quant_config,

728 + prefix=prefix)

729 +

730 + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

731 + hidden_states, _ = self.fc1(hidden_states)

732 + hidden_states = self.activation_fn(hidden_states)

733 + hidden_states, _ = self.fc2(hidden_states)

734 + return hidden_states

735 +

736 +

737 +class Step3VisionEncoderLayer(nn.Module):

738 +

739 + def __init__(self,

740 + config: Step3VisionEncoderConfig,

741 + quant_config: Optional[QuantizationConfig] = None,

742 + prefix: str = ""):

743 + super().__init__()

744 + self.embed_dim = config.hidden_size

745 + self.self_attn = Step3VisionAttention(config,

746 + quant_config,

747 + prefix=f"{prefix}.self_attn")

748 + self.layer_norm1 = nn.LayerNorm(self.embed_dim,

749 + eps=config.layer_norm_eps)

750 + self.mlp = Step3VisionMLP(config, quant_config, prefix=f"{prefix}.mlp")

751 + self.layer_norm2 = nn.LayerNorm(self.embed_dim,

752 + eps=config.layer_norm_eps)

753 +

754 + def forward(

755 + self,

756 + hidden_states: torch.Tensor,

757 + ) -> torch.FloatTensor:

758 + hidden_states = hidden_states + self.layer_norm1(

759 + self.self_attn(hidden_states))

760 + hidden_states = hidden_states + self.layer_norm2(

761 + self.mlp(hidden_states))

762 + return hidden_states

763 +

764 +

765 +class Step3VisionEncoder(nn.Module):

766 +

767 + def __init__(self,

768 + config: Step3VisionEncoderConfig,

769 + quant_config: Optional[QuantizationConfig] = None,

770 + prefix: str = ""):

771 + super().__init__()

772 + self.config = config

773 + self.layers = nn.ModuleList([

774 + Step3VisionEncoderLayer(config,

775 + quant_config,

776 + prefix=f"{prefix}.layers.{i}")

777 + for i in range(config.num_hidden_layers)

778 + ])

779 +

780 + def forward(

781 + self,

782 + inputs_embeds,

783 + ):

784 + hidden_states = inputs_embeds

785 + for encoder_layer in self.layers:

786 + hidden_states = encoder_layer(hidden_states)

787 + return hidden_states

788 +

789 +

790 +class Step3VisionTransformer(nn.Module):

791 +

792 + def __init__(self,

793 + config: Step3VisionEncoderConfig,

794 + quant_config: Optional[QuantizationConfig] = None,

795 + prefix: str = ""):

796 + super().__init__()

797 + self.config = config

798 + self.image_size = config.image_size

799 + self.embeddings = Step3VisionEmbeddings(config)

800 + self.transformer = Step3VisionEncoder(config,

801 + quant_config,

802 + prefix=f"{prefix}.transformer")

803 +

804 + def forward(

805 + self,

806 + pixel_values: torch.Tensor,

807 + ):

808 + hidden_states = self.embeddings(pixel_values)

809 + hidden_states = self.transformer(inputs_embeds=hidden_states)

810 + return hidden_states

811 +

812 +

813 +@MULTIMODAL_REGISTRY.register_processor(Step3VLMultiModalProcessor,

814 + info=Step3VLProcessingInfo,

815 + dummy_inputs=Step3VLDummyInputsBuilder)

816 +class Step3VLForConditionalGeneration(nn.Module, SupportsMultiModal,

817 + SupportsPP):

818 +

819 + hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={

820 + "model.": "language_model.model.",

821 + "lm_head.": "language_model.lm_head.",

822 + })

823 +

824 + @classmethod

825 + def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

826 + if modality.startswith("image"):

827 + return "<im_patch>"

828 +

829 + raise ValueError("Only image modality is supported")

830 +

831 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

832 + super().__init__()

833 +

834 + config = vllm_config.model_config.hf_config

835 + multimodal_config = vllm_config.model_config.multimodal_config

836 +

837 + self.config = config

838 + self.multimodal_config = multimodal_config

839 +

840 + self.vision_model = Step3VisionTransformer(config.vision_config,

841 + None,

842 + prefix=maybe_prefix(

843 + prefix, "vision_model"))

844 + self.vit_downsampler = nn.Conv2d(

845 + config.vision_config.hidden_size,

846 + config.vision_config.output_hidden_size,

847 + kernel_size=2,

848 + stride=config.understand_projector_stride)

849 + self.vit_downsampler2 = nn.Conv2d(

850 + config.vision_config.output_hidden_size,

851 + config.vision_config.output_hidden_size * 2,

852 + kernel_size=3,

853 + stride=2,

854 + padding=1,

855 + )

856 + self.vit_large_projector = nn.Linear(

857 + config.vision_config.output_hidden_size * 2,

858 + config.hidden_size,

859 + bias=config.projector_bias,

860 + )

861 + self.language_model = init_vllm_registered_model(

862 + vllm_config=vllm_config,

863 + hf_config=config.text_config,

864 + prefix=maybe_prefix(prefix, "language_model"))

865 +

866 + self.make_empty_intermediate_tensors = (

867 + self.language_model.make_empty_intermediate_tensors)

868 +

869 + @cached_property

870 + def sampler(self):

871 + if hasattr(self.language_model, "sampler"):

872 + return self.language_model.sampler

873 +

874 + return get_sampler()

875 +

876 + @property

877 + def device(self):

878 + return next(self.parameters()).device

879 +

880 + @property

881 + def dtype(self):

882 + return next(self.parameters()).dtype

883 +

884 + def _parse_and_validate_image_input(

885 + self, **kwargs: object) -> Optional[Step3VLImageInputs]:

886 + pixel_values = kwargs.pop("pixel_values", None)

887 + patch_pixel_values = kwargs.pop("patch_pixel_values", None)

888 + num_patches = kwargs.pop("num_patches", None)

889 + image_embeds = kwargs.pop("image_embeds", None)

890 +

891 + if pixel_values is None and image_embeds is None:

892 + return None

893 +

894 + if pixel_values is not None:

895 + pixel_values = flatten_bn(pixel_values, concat=True)

896 + if pixel_values.dim() >= 3:

897 + pixel_values = pixel_values.view(-1, *pixel_values.shape[-3:])

898 + if patch_pixel_values is not None:

899 + patch_pixel_values = flatten_bn(patch_pixel_values,

900 + concat=True)

901 + patch_pixel_values = patch_pixel_values.view(

902 + -1, *patch_pixel_values.shape[-3:])

903 + # Handle empty patch_pixel_values by setting to None

904 + if patch_pixel_values.shape[0] == 0:

905 + patch_pixel_values = None

906 + num_patches = flatten_bn(num_patches, concat=True).tolist()

907 +

908 + return Step3VLImagePixelInputs(

909 + type="pixel_values",

910 + pixel_values=pixel_values.to(self.dtype).to(self.device),

911 + patch_pixel_values=patch_pixel_values.to(self.dtype).to(

912 + self.device) if patch_pixel_values is not None else None,

913 + num_patches=num_patches,

914 + )

915 +

916 + if image_embeds is not None:

917 + if image_embeds.dim() == 2 or image_embeds.dim() >= 3:

918 + image_embeds = image_embeds.view(-1, image_embeds.shape[-1])

919 + else:

920 + raise ValueError(

921 + f"Unexpected shape for image_embeds: {image_embeds.shape}")

922 +

923 + return Step3VLImageEmbeddingInputs(

924 + type="image_embeds",

925 + image_embeds=image_embeds.to(self.dtype).to(self.device),

926 + )

927 + return None

928 +

929 + def _process_image_features(self,

930 + image_features: torch.Tensor) -> torch.Tensor:

931 + B, P = image_features.shape[:2]

932 + HW = int(sqrt(P))

933 + image_features = image_features.permute(0, 2, 1).view(B, -1, HW, HW)

934 + image_features = self.vit_downsampler(image_features)

935 + image_features = self.vit_downsampler2(image_features)

936 + n_dim = image_features.size(1)

937 + image_features = image_features.view(B, n_dim, -1).permute(0, 2, 1)

938 + image_features = self.vit_large_projector(image_features)

939 + return image_features

940 +

941 + def _get_vision_model_output(self,

942 + input_tensor: torch.Tensor) -> torch.Tensor:

943 + return self.vision_model(input_tensor)[:, 4:]

944 +

945 + def _process_image_input(

946 + self, image_input: Step3VLImageInputs) -> tuple[torch.Tensor, ...]:

947 +

948 + if image_input["type"] == "image_embeds":

949 + image_features = image_input["image_embeds"]

950 + else:

951 + image_features = self._get_vision_model_output(

952 + image_input["pixel_values"])

953 + patch_image_features = self._get_vision_model_output(

954 + image_input["patch_pixel_values"]

955 + ) if image_input["patch_pixel_values"] is not None else None

956 + num_patches = image_input["num_patches"]

957 +

958 + image_features = self._process_image_features(image_features)

959 + patch_image_features = self._process_image_features(

960 + patch_image_features) if patch_image_features is not None else None

961 +

962 + merged_image_features = []

963 + cur_patch_idx = 0

964 + for i, num_patch in enumerate(num_patches):

965 + cur_feature = []

966 + if num_patch > 0:

967 + patch_slice = patch_image_features[

968 + cur_patch_idx:cur_patch_idx + num_patch]

969 + cur_feature.append(patch_slice.view(-1, patch_slice.shape[-1]))

970 + cur_feature.append(image_features[i].view(

971 + -1, image_features.shape[-1]))

972 + cur_patch_idx += num_patch

973 + merged_image_features.append(

974 + torch.cat(cur_feature) if len(cur_feature) >

975 + 1 else cur_feature[0])

976 + return merged_image_features

977 +

978 + def get_multimodal_embeddings(self, **kwargs) -> Optional[NestedTensors]:

979 + image_input = self._parse_and_validate_image_input(**kwargs)

980 + if image_input is None:

981 + return None

982 + vision_embeddings = self._process_image_input(image_input)

983 + return vision_embeddings

984 +

985 + def get_input_embeddings(

986 + self,

987 + input_ids: torch.Tensor,

988 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

989 + ) -> torch.Tensor:

990 + if multimodal_embeddings is None:

991 + inputs_embeds = self.language_model.model.get_input_embeddings(

992 + input_ids)

993 + else:

994 + is_text = input_ids != self.config.image_token_id

995 + text_ids = input_ids[is_text]

996 + text_embeds = self.language_model.model.get_input_embeddings(

997 + text_ids)

998 + inputs_embeds = torch.empty(input_ids.shape[0],

999 + text_embeds.shape[-1],

1000 + dtype=text_embeds.dtype,

1001 + device=text_embeds.device)

1002 + inputs_embeds[is_text] = text_embeds

1003 + inputs_embeds = merge_multimodal_embeddings(

1004 + input_ids, inputs_embeds, multimodal_embeddings,

1005 + self.config.image_token_id)

1006 + return inputs_embeds

1007 +

1008 + def forward(

1009 + self,

1010 + input_ids: torch.Tensor,

1011 + positions: torch.Tensor,

1012 + intermediate_tensors: Optional[IntermediateTensors] = None,

1013 + inputs_embeds: Optional[torch.Tensor] = None,

1014 + **kwargs: object,

1015 + ) -> Union[torch.Tensor, IntermediateTensors]:

1016 + if intermediate_tensors is not None:

1017 + inputs_embeds = None

1018 + elif inputs_embeds is None:

1019 + vision_embeddings = self.get_multimodal_embeddings(**kwargs)

1020 + # always pass the input via `inputs_embeds`

1021 + # to make sure the computation graph is consistent

1022 + inputs_embeds = self.get_input_embeddings(input_ids,

1023 + vision_embeddings)

1024 + input_ids = None

1025 +

1026 + hidden_states = self.language_model(input_ids,

1027 + positions,

1028 + intermediate_tensors,

1029 + inputs_embeds=inputs_embeds)

1030 +

1031 + return hidden_states

1032 +

1033 + def compute_logits(

1034 + self,

1035 + hidden_states: torch.Tensor,

1036 + sampling_metadata: SamplingMetadata,

1037 + ) -> Optional[torch.Tensor]:

1038 + return self.language_model.compute_logits(hidden_states,

1039 + sampling_metadata)

1040 +

1041 + def sample(

1042 + self,

1043 + logits: torch.Tensor,

1044 + sampling_metadata: SamplingMetadata,

1045 + ) -> Optional[SamplerOutput]:

1046 + return self.language_model.sample(logits, sampling_metadata)

1047 +

1048 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

1049 + loader = AutoWeightsLoader(self)

1050 + loaded_weights = loader.load_weights(weights,

1051 + mapper=self.hf_to_vllm_mapper)

1052 + return loaded_weights

vllm/reasoning/__init__.py

@@ -8,6 +8,7 @@ from .granite_reasoning_parser import GraniteReasoningParser

8 from .hunyuan_a13b_reasoning_parser import HunyuanA13BReasoningParser

9 from .mistral_reasoning_parser import MistralReasoningParser

10 from .qwen3_reasoning_parser import Qwen3ReasoningParser

11 +from .step3_reasoning_parser import Step3ReasoningParser

13 __all__ = [

14 "ReasoningParser",

@@ -18,4 +19,5 @@ __all__ = [

19 "Qwen3ReasoningParser",

20 "Glm4MoeModelReasoningParser",

21 "MistralReasoningParser",

22 + "Step3ReasoningParser",

23 ]

vllm/reasoning/step3_reasoning_parser.py new

+109

@@ -0,0 +1,109 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from collections.abc import Sequence

5 +from typing import Optional, Union

6 +

7 +import regex as re

8 +from transformers import PreTrainedTokenizerBase

9 +

10 +from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,

11 + DeltaMessage)

12 +from vllm.logger import init_logger

13 +from vllm.reasoning import ReasoningParser, ReasoningParserManager

14 +

15 +logger = init_logger(__name__)

16 +

17 +

18 +@ReasoningParserManager.register_module("step3")

19 +class Step3ReasoningParser(ReasoningParser):

20 + """

21 + Reasoning parser for Step3 model.

22 +

23 + The Step3 model uses </think> token to denote the end of reasoning

24 + text. This parser extracts all content before </think> as reasoning content.

25 + """

26 +

27 + def __init__(self, tokenizer: PreTrainedTokenizerBase):

28 + super().__init__(tokenizer)

29 + self.think_end_token = "</think>"

30 +

31 + self.reasoning_regex = re.compile(rf"(.*?){self.think_end_token}",

32 + re.DOTALL)

33 +

34 + if not self.model_tokenizer:

35 + raise ValueError(

36 + "The model tokenizer must be passed to the ReasoningParser "

37 + "constructor during construction.")

38 +

39 + self.think_end_token_id = self.vocab.get(self.think_end_token)

40 + if self.think_end_token_id is None:

41 + raise RuntimeError(

42 + "Step3 reasoning parser could not locate think end "

43 + "token in the tokenizer!")

44 +

45 + def extract_reasoning_content_streaming(

46 + self,

47 + previous_text: str,

48 + current_text: str,

49 + delta_text: str,

50 + previous_token_ids: Sequence[int],

51 + current_token_ids: Sequence[int],

52 + delta_token_ids: Sequence[int],

53 + ) -> Union[DeltaMessage, None]:

54 + """

55 + Extract reasoning content from a delta message.

56 + Handles streaming output where previous + delta = current.

57 + Uses token IDs for faster processing.

58 + For text "abc</think>xyz":

59 + - 'abc' goes to reasoning_content

60 + - 'xyz' goes to content

61 + """

62 + # Skip single special token

63 + if len(delta_token_ids

64 + ) == 1 and delta_token_ids[0] == self.think_end_token_id:

65 + return None

66 +

67 + if self.think_end_token_id in delta_token_ids:

68 + # </think> in delta, extract reasoning content and remaining content

69 + end_index = delta_text.find(self.think_end_token)

70 + reasoning_content = delta_text[:end_index]

71 + content = delta_text[end_index + len(self.think_end_token):]

72 + return DeltaMessage(reasoning_content=reasoning_content,

73 + content=content if content else None)

74 + elif self.think_end_token_id in previous_token_ids:

75 + # </think> already seen in previous text, everything is content

76 + return DeltaMessage(content=delta_text)

77 + else:

78 + # No </think> seen yet, everything is reasoning

79 + return DeltaMessage(reasoning_content=delta_text)

80 +

81 + def extract_reasoning_content(

82 + self, model_output: str, request: ChatCompletionRequest

83 + ) -> tuple[Optional[str], Optional[str]]:

84 +

85 + # Check if the model output contains the </think> token

86 + if self.think_end_token not in model_output:

87 + # If no </think> token, everything is reasoning content

88 + return model_output, None

89 + else:

90 + # Find the first occurrence of </think>

91 + end_index = model_output.find(self.think_end_token)

92 + reasoning_content = model_output[:end_index]

93 +

94 + # Content after </think> token

95 + content = model_output[end_index + len(self.think_end_token):]

96 +

97 + if len(content) == 0:

98 + content = None

99 +

100 + return reasoning_content, content

101 +

102 + def is_reasoning_end(self, input_ids: list[int]) -> bool:

103 + return self.think_end_token_id in input_ids

104 +

105 + def extract_content_ids(self, input_ids: list[int]) -> list[int]:

106 + if self.think_end_token_id not in input_ids[:-1]:

107 + return []

108 + else:

109 + return input_ids[input_ids.index(self.think_end_token_id) + 1:]

vllm/transformers_utils/config.py

+4 -1

@@ -35,7 +35,8 @@ from vllm.transformers_utils.configs import (ChatGLMConfig, DeepseekVLV2Config,

35 MllamaConfig, MLPSpeculatorConfig,

36 Nemotron_Nano_VL_Config,

37 NemotronConfig, NVLM_D_Config,

38 - RWConfig, UltravoxConfig)

38 + RWConfig, Step3TextConfig,

39 + Step3VLConfig, UltravoxConfig)

40 # yapf: enable

41 from vllm.transformers_utils.configs.mistral import adapt_config_dict

42 from vllm.transformers_utils.utils import check_gguf_file

@@ -83,6 +84,8 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = {

84 "nemotron": NemotronConfig,

85 "NVLM_D": NVLM_D_Config,

86 "ultravox": UltravoxConfig,

87 + "step3_vl": Step3VLConfig,

88 + "step3_text": Step3TextConfig,

89 **_CONFIG_REGISTRY_OVERRIDE_HF

90 }

vllm/transformers_utils/configs/__init__.py

@@ -24,6 +24,9 @@ from vllm.transformers_utils.configs.nemotron import NemotronConfig

24 from vllm.transformers_utils.configs.nemotron_h import NemotronHConfig

25 from vllm.transformers_utils.configs.nemotron_vl import Nemotron_Nano_VL_Config

26 from vllm.transformers_utils.configs.nvlm_d import NVLM_D_Config

27 +from vllm.transformers_utils.configs.step3_vl import (Step3TextConfig,

28 + Step3VisionEncoderConfig,

29 + Step3VLConfig)

30 from vllm.transformers_utils.configs.ultravox import UltravoxConfig

32 __all__ = [

@@ -42,4 +45,7 @@ __all__ = [

45 "Nemotron_Nano_VL_Config",

46 "NVLM_D_Config",

47 "UltravoxConfig",

48 + "Step3VLConfig",

49 + "Step3VisionEncoderConfig",

50 + "Step3TextConfig",

51 ]

vllm/transformers_utils/configs/step3_vl.py new

+123

@@ -0,0 +1,123 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from typing import Any, Optional, Union

4 +

5 +from transformers.configuration_utils import PretrainedConfig

6 +

7 +

8 +class Step3VisionEncoderConfig(PretrainedConfig):

9 + model_type = "step3_vision_encoder"

10 +

11 + def __init__(

12 + self,

13 + hidden_size=1792,

14 + intermediate_size=3072,

15 + output_hidden_size=4096,

16 + num_hidden_layers=63,

17 + num_attention_heads=16,

18 + num_channels=3,

19 + image_size=728,

20 + patch_size=14,

21 + hidden_act="quick_gelu",

22 + layer_norm_eps=1e-5,

23 + **kwargs,

24 + ):

25 + self.hidden_size = hidden_size

26 + self.intermediate_size = intermediate_size

27 + self.output_hidden_size = output_hidden_size

28 + self.num_hidden_layers = num_hidden_layers

29 + self.num_attention_heads = num_attention_heads

30 + self.num_channels = num_channels

31 + self.patch_size = patch_size

32 + self.image_size = image_size

33 + self.layer_norm_eps = layer_norm_eps

34 + self.hidden_act = hidden_act

35 + super().__init__(**kwargs)

36 +

37 +

38 +class Step3TextConfig(PretrainedConfig):

39 + model_type = "step3_text"

40 + architectures = ["Step3TextForCausalLM"]

41 +

42 + def __init__(

43 + self,

44 + hidden_size: int = 7168,

45 + intermediate_size: int = 18432,

46 + num_attention_heads: int = 64,

47 + num_attention_groups: int = 1,

48 + num_hidden_layers: int = 61,

49 + max_seq_len: int = 65536,

50 + vocab_size: int = 128815,

51 + rms_norm_eps: float = 1e-5,

52 + moe_intermediate_size: int = 5120,

53 + moe_num_experts: int = 48,

54 + moe_top_k: int = 3,

55 + rope_theta: float = 500000,

56 + rope_scaling: Optional[dict[str, Any]] = None,

57 + max_position_embedding: int = 65536,

58 + share_expert_dim: int = 5120,

59 + share_q_dim: int = 2048,

60 + head_dim: int = 256,

61 + norm_expert_weight: bool = False,

62 + moe_layers_enum: tuple[int,

63 + ...] = (4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14,

64 + 15, 16, 17, 18, 19, 20, 21, 22, 23, 24,

65 + 25, 26, 27, 28, 29, 30, 31, 32, 33, 34,

66 + 35, 36, 37, 38, 39, 40, 41, 42, 43, 44,

67 + 45, 46, 47, 48, 49, 50, 51, 52, 53, 54,

68 + 55, 56, 57, 58, 59),

69 + **kwargs,

70 + ) -> None:

71 + self.hidden_size = hidden_size

72 + self.intermediate_size = intermediate_size

73 + self.num_attention_heads = num_attention_heads

74 + self.num_attention_groups = num_attention_groups

75 + self.num_hidden_layers = num_hidden_layers

76 + self.max_seq_len = max_seq_len

77 + self.vocab_size = vocab_size

78 + self.rms_norm_eps = rms_norm_eps

79 + self.moe_intermediate_size = moe_intermediate_size

80 + self.moe_num_experts = moe_num_experts

81 + self.moe_top_k = moe_top_k

82 + self.rope_theta = rope_theta

83 + self.rope_scaling = rope_scaling

84 + self.max_position_embedding = max_position_embedding

85 + self.share_expert_dim = share_expert_dim

86 + self.share_q_dim = share_q_dim

87 + self.head_dim = head_dim

88 + self.norm_expert_weight = norm_expert_weight

89 + self.moe_layers_enum = moe_layers_enum

90 +

91 + super().__init__(**kwargs)

92 +

93 +

94 +class Step3VLConfig(PretrainedConfig):

95 + model_type = "step3_vl"

96 +

97 + def __init__(

98 + self,

99 + vision_config: Optional[Union[dict, Step3VisionEncoderConfig]] = None,

100 + text_config: Optional[Union[dict, Step3TextConfig]] = None,

101 + understand_projector_stride: int = 1,

102 + projector_bias: bool = True,

103 + image_token_id: int = 128001,

104 + **kwargs,

105 + ) -> None:

106 + if vision_config is None:

107 + vision_config = Step3VisionEncoderConfig()

108 + elif isinstance(vision_config, dict):

109 + vision_config = Step3VisionEncoderConfig(**vision_config)

110 + self.vision_config = vision_config

111 +

112 + if text_config is None:

113 + text_config = Step3TextConfig()

114 + elif isinstance(text_config, dict):

115 + text_config = Step3TextConfig(**text_config)

116 + self.text_config = text_config

117 +

118 + self.understand_projector_stride = understand_projector_stride

119 + self.projector_bias = projector_bias

120 + self.hidden_size = text_config.hidden_size

121 + self.image_token_id = image_token_id

122 +

123 + super().__init__(**kwargs)