[Test][Qwen3-VL] Cover compiled DeepStack input contract (#53529)
Signed-off-by: maithilijoshi20 <97733343+maithilijoshi20@users.noreply.github.com>
maithilijoshi20 committed
Aug 31, 2026 at 19:46 UTC
9debcd5990bf275db5ac8ed7bea653a1fa8a23a0
1 file changed
+71
tests/compile/test_deepstack_input_contract.py
new
+71
@@ -0,0 +1,71 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+from torch import nn
6
+
7
+from vllm.compilation.decorators import support_torch_compile
8
+from vllm.config import (
9
+ CompilationConfig,
10
+ CompilationMode,
11
+ VllmConfig,
12
+ set_current_vllm_config,
13
+)
14
+from vllm.forward_context import set_forward_context
15
+from vllm.model_executor.models.qwen3_vl import (
16
+ Qwen3VLForConditionalGeneration,
17
+)
18
+from vllm.sequence import IntermediateTensors
19
+
20
+
21
+@support_torch_compile(
22
+ dynamic_arg_dims={"x": 0, "deepstack_input_embeds": 0},
23
+)
24
+class _DeepStackDecoder(nn.Module):
25
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
26
+ super().__init__()
27
+
28
+ def forward(
29
+ self,
30
+ x: torch.Tensor,
31
+ deepstack_input_embeds: IntermediateTensors | None,
32
+ ) -> torch.Tensor:
33
+ if deepstack_input_embeds is not None:
34
+ x = x + deepstack_input_embeds["deepstack_input_embeds_0"]
35
+ return x
36
+
37
+
38
+def _run_contract(
39
+ deepstack_buffer: Qwen3VLForConditionalGeneration,
40
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
41
+ config = VllmConfig(
42
+ compilation_config=CompilationConfig(mode=CompilationMode.VLLM_COMPILE)
43
+ )
44
+ with set_current_vllm_config(config):
45
+ model = _DeepStackDecoder(vllm_config=config).eval().cuda()
46
+
47
+ x = torch.arange(8, dtype=torch.float32, device="cuda").reshape(2, 4)
48
+ with torch.inference_mode(), set_forward_context(None, vllm_config=config):
49
+ empty_payload = deepstack_buffer._get_deepstack_input_embeds(num_tokens=2)
50
+ first = model(x, empty_payload)
51
+
52
+ payload = torch.full((1, 2, 4), 7.0, device="cuda")
53
+ deepstack_buffer._set_deepstack_input_embeds(payload)
54
+ active_payload = deepstack_buffer._get_deepstack_input_embeds(num_tokens=2)
55
+ assert active_payload is not None
56
+ second = model(x, active_payload)
57
+
58
+ return x.cpu(), first.cpu(), second.cpu()
59
+
60
+
61
+def test_qwen3_vl_deepstack_payload_after_empty_compile(monkeypatch):
62
+ monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
63
+
64
+ deepstack_buffer = object.__new__(Qwen3VLForConditionalGeneration)
65
+ deepstack_buffer.deepstack_input_embeds = [torch.zeros(2, 4, device="cuda")]
66
+ deepstack_buffer.deepstack_input_embeds_num_tokens = 0
67
+ deepstack_buffer.deepstack_num_level = 1
68
+
69
+ x, first, second = _run_contract(deepstack_buffer)
70
+ torch.testing.assert_close(first, x)
71
+ torch.testing.assert_close(second, x + 7)