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 3, 17:45 UTC
[Bugfix] V1: support tuple model outputs in ubatch wrapper (dbo + spec decode) (#40789)
Signed-off-by: Yufeng He <40085740+he-yufeng@users.noreply.github.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Yufeng He committed
May 14, 2026 at 06:47 UTC
b2198670b19ac597c630f0b8d4f7a7addb764a62
1 file changed
+19
-2
vllm/v1/worker/gpu_ubatch_wrapper.py
+19
-2
@@ -31,6 +31,23 @@ from vllm.v1.worker.ubatching import UBatchContext, make_ubatch_contexts
31
logger = init_logger(__name__)
32
33
34
+def _cat_ubatch_outputs(
35
+ sorted_results: list,
36
+) -> "torch.Tensor | tuple[torch.Tensor, ...]":
37
+ """Concatenate per-ubatch model outputs along the batch dim.
38
+
39
+ Most models return a single hidden-states tensor per ubatch. Target
40
+ models running with auxiliary output (e.g. EAGLE3 speculative decoding,
41
+ which collects aux hidden states for the drafter) return a tuple of
42
+ tensors instead. Fan out over tuple components so `torch.cat` sees
43
+ matching shapes and the caller receives the same structure the model
44
+ produced for a single ubatch (#40769).
45
+ """
46
+ if sorted_results and isinstance(sorted_results[0], tuple):
47
+ return tuple(torch.cat(parts, dim=0) for parts in zip(*sorted_results))
48
+ return torch.cat(sorted_results, dim=0)
49
+
50
+
51
@dataclass
52
class UbatchMetadata:
53
context: UBatchContext
@@ -266,7 +283,7 @@ class UBatchWrapper:
283
for thread in ubatch_threads:
284
thread.join()
285
sorted_results = [value for position, value in sorted(results)]
269
- result = torch.cat(sorted_results, dim=0)
286
+ result = _cat_ubatch_outputs(sorted_results)
287
cudagraph_metadata.outputs = result
288
# Join offloader's copy stream after forward to avoid unjoined
289
# stream error. The last layer's start_prefetch forks copy_stream,
@@ -310,7 +327,7 @@ class UBatchWrapper:
327
for thread in ubatch_threads:
328
thread.join()
329
sorted_results = [value for position, value in sorted(results)]
313
- result = torch.cat(sorted_results, dim=0)
330
+ result = _cat_ubatch_outputs(sorted_results)
331
return result
332
333
def _make_ubatch_metadata(