[BugFix][VLM] keep Qwen3-VL MoE inference deepstack order (#34690)

This commit is contained in:
Zheng Wengang
2026-08-28 17:02:16 +08:00
committed by GitHub
parent 3785b2d20f
commit 1e6d041f78
3 changed files with 156 additions and 36 deletions
+41 -18
View File
@@ -1152,6 +1152,11 @@ class Qwen3LLMModel(Qwen3Model):
self.deepstack_embed_to_decoder_layer = range(
len(config.vision_config.deepstack_visual_indexes)
)
# Use HF deepstack order only if rl_on_policy_target is set;
# otherwise, retain original order for inference accuracy.
self.use_hf_deepstack_order = (
get_exec().deterministic.rl_on_policy_target is not None
)
def get_deepstack_embeds(
self, layer_idx: int, input_deepstack_embeds: Optional[torch.Tensor]
@@ -1196,25 +1201,43 @@ class Qwen3LLMModel(Qwen3Model):
hidden_states + residual if residual is not None else hidden_states
)
# SGLang applies residual at the START of the next layer, not at the END like HuggingFace.
# See: https://github.com/huggingface/transformers/blob/v5.0.0rc0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L549
# To match HF behavior, deepstack must be added AFTER residual: (hidden_states + residual) + deepstack
# The order matters because addition with different tensors is not associative in practice.
# Deepstack for prev_layer is applied at the start of current layer via post_residual_addition.
deepstack_embeds = self.get_deepstack_embeds(
layer_idx - 1, input_deepstack_embeds
)
hidden_states, residual = layer(
positions,
hidden_states,
forward_batch,
residual,
post_residual_addition=deepstack_embeds,
)
if self.use_hf_deepstack_order:
# HF-order path (RL on-policy / FSDP). SGLang applies residual at the START of the
# next layer, so to match HF's (hidden_states + residual) + deepstack, deepstack for
# the previous layer is added after residual via post_residual_addition.
deepstack_embeds = self.get_deepstack_embeds(
layer_idx - 1, input_deepstack_embeds
)
hidden_states, residual = layer(
positions,
hidden_states,
forward_batch,
residual,
post_residual_addition=deepstack_embeds,
)
else:
# Inference path: add deepstack directly to hidden_states at the end of the layer
# (original, grounding-correct order).
hidden_states, residual = layer(
positions,
hidden_states,
forward_batch,
residual,
)
if (
input_deepstack_embeds is not None
and layer_idx in self.deepstack_embed_to_decoder_layer
):
sep = self.hidden_size * layer_idx
hidden_states.add_(
input_deepstack_embeds[:, sep : sep + self.hidden_size]
)
# Handle deepstack for the last processed layer if it exists.
last_deepstack = self.get_deepstack_embeds(
self.end_layer - 1, input_deepstack_embeds
# Handle deepstack for the last processed layer (HF-order path only).
last_deepstack = (
self.get_deepstack_embeds(self.end_layer - 1, input_deepstack_embeds)
if self.use_hf_deepstack_order
else None
)
if not self.pp_group.is_last_rank:
+42 -18
View File
@@ -31,6 +31,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTe
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeModel
from sglang.srt.models.qwen3_vl import Qwen3VLForConditionalGeneration
from sglang.srt.runtime_context import get_exec
from sglang.srt.utils.hf_transformers_utils import get_processor
logger = logging.getLogger(__name__)
@@ -58,6 +59,11 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
# This approach follows the original implementation.
# TODO: make config of type Qwen3VLMoeConfig, so that we can directly obtain deepstack_visual_indexes.
self.deepstack_embed_to_decoder_layer = range(3)
# Use HF deepstack order only if rl_on_policy_target is set;
# otherwise, retain original order for inference accuracy.
self.use_hf_deepstack_order = (
get_exec().deterministic.rl_on_policy_target is not None
)
def get_input_embeddings(self) -> nn.Embedding:
return self.embed_tokens
@@ -104,25 +110,43 @@ class Qwen3MoeLLMModel(Qwen3MoeModel):
hidden_states + residual if residual is not None else hidden_states
)
# SGLang applies residual at the START of the next layer, not at the END like HuggingFace.
# See: https://github.com/huggingface/transformers/blob/v5.0.0rc0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L549
# To match HF behavior, deepstack must be added AFTER residual: (hidden_states + residual) + deepstack
# The order matters because addition with different tensors is not associative in practice.
# Deepstack for prev_layer is applied at the start of current layer via post_residual_addition.
deepstack_embeds = self.get_deepstack_embeds(
layer_idx - 1, input_deepstack_embeds
)
hidden_states, residual = layer(
positions,
hidden_states,
forward_batch,
residual,
post_residual_addition=deepstack_embeds,
)
if self.use_hf_deepstack_order:
# HF-order path (RL on-policy / FSDP). SGLang applies residual at the START of the
# next layer, so to match HF's (hidden_states + residual) + deepstack, deepstack for
# the previous layer is added after residual via post_residual_addition.
deepstack_embeds = self.get_deepstack_embeds(
layer_idx - 1, input_deepstack_embeds
)
hidden_states, residual = layer(
positions,
hidden_states,
forward_batch,
residual,
post_residual_addition=deepstack_embeds,
)
else:
# Inference path: add deepstack directly to hidden_states at the end of the layer
# (original, grounding-correct order).
hidden_states, residual = layer(
positions,
hidden_states,
forward_batch,
residual,
)
if (
input_deepstack_embeds is not None
and layer_idx in self.deepstack_embed_to_decoder_layer
):
sep = self.hidden_size * layer_idx
hidden_states.add_(
input_deepstack_embeds[:, sep : sep + self.hidden_size]
)
# Handle deepstack for the last processed layer if it exists.
last_deepstack = self.get_deepstack_embeds(
self.end_layer - 1, input_deepstack_embeds
# Handle deepstack for the last processed layer (HF-order path only).
last_deepstack = (
self.get_deepstack_embeds(self.end_layer - 1, input_deepstack_embeds)
if self.use_hf_deepstack_order
else None
)
if not self.pp_group.is_last_rank:
@@ -4,9 +4,13 @@ python3 -m unittest test_vision_openai_server.TestOpenAIVisionServer.test_mixed_
python3 -m unittest test_vision_openai_server.TestOpenAIVisionServer.test_multi_images_chat_completion
"""
import base64
import io
import re
import unittest
import openai
from PIL import Image, ImageDraw
from sglang.srt.environ import envs
from sglang.test.ci.ci_register import register_cuda_ci
@@ -28,6 +32,34 @@ from sglang.test.vlm_utils import (
register_cuda_ci(est_time=560, stage="base-b", runner_config="1-gpu-large")
# --- Qwen3-VL grounding regression (deepstack fusion) --------------------------
# Guards Qwen3MoeLLMModel.forward: deepstack (multi-scale ViT features) injection
# must keep its original inference order. PR #14636 rerouted it through
# post_residual_addition (for RL on-policy / FSDP), which is FP-order-sensitive
# and regresses FP8 visual grounding (the predicted point drifts by ~150+ px).
_GROUNDING_IMG_SIZE = 1000
# Target box in pixels; on a 1000x1000 canvas this equals the 0-1000 normalized
# coordinate, so the check is robust to normalized-vs-pixel conventions.
_GROUNDING_BOX = (620, 180, 880, 360) # (x0, y0, x1, y1), center (750, 270)
_GROUNDING_MARGIN = 60
_GROUNDING_SYSTEM = (
"You are a UI grounding model. Treat the image as a 1000x1000 normalized "
"coordinate system with the top-left at (0,0) and the bottom-right at "
"(1000,1000). Return the geometric center of the requested element. "
"Output ONLY one coordinate in the form (x, y) and nothing else."
)
_GROUNDING_COORD_RE = re.compile(r"\(?\s*(\d{1,4})\s*,\s*(\d{1,4})\s*\)?")
def _make_grounding_image() -> str:
"""White canvas with a single red box at _GROUNDING_BOX; base64 data URI."""
img = Image.new("RGB", (_GROUNDING_IMG_SIZE, _GROUNDING_IMG_SIZE), (255, 255, 255))
ImageDraw.Draw(img).rectangle(_GROUNDING_BOX, fill=(220, 30, 30))
buf = io.BytesIO()
img.save(buf, format="PNG")
return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode("utf-8")
class TestLlavaServer(ImageOpenAITestMixin):
model = "lmms-lab/llava-onevision-qwen2-0.5b-ov"
@@ -52,6 +84,47 @@ class TestQwen3VLServer(ImageOpenAITestMixin, VideoOpenAITestMixin):
with envs.SGLANG_MM_FEATURE_CACHE_MB.override(512):
super().setUpClass()
def test_deepstack_grounding_hits_target_box(self):
# Regression guard for the Qwen3-VL MoE deepstack fusion order: the
# predicted point must land inside the target box; a deepstack corruption
# drifts it out (see PR #14636).
client = openai.Client(api_key=self.api_key, base_url=self.base_url)
response = client.chat.completions.create(
model="default",
messages=[
{"role": "system", "content": _GROUNDING_SYSTEM},
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {"url": _make_grounding_image()},
},
{
"type": "text",
"text": "Point at the center of the red rectangle.",
},
],
},
],
temperature=0,
**(self.get_vision_request_kwargs()),
)
out = response.choices[0].message.content
match = _GROUNDING_COORD_RE.search(out or "")
self.assertIsNotNone(match, f"could not parse a coordinate from: {out!r}")
x, y = int(match.group(1)), int(match.group(2))
x0, y0, x1, y1 = _GROUNDING_BOX
inside = (x0 - _GROUNDING_MARGIN <= x <= x1 + _GROUNDING_MARGIN) and (
y0 - _GROUNDING_MARGIN <= y <= y1 + _GROUNDING_MARGIN
)
self.assertTrue(
inside,
f"grounding output {out!r} -> ({x}, {y}) fell outside target box "
f"{_GROUNDING_BOX} (margin {_GROUNDING_MARGIN}); deepstack fusion "
f"likely regressed grounding.",
)
class TestQwen2VLContextLengthServer(CustomTestCase):
# --context-length 300 is calibrated to this model's mm-token expansion: