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

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, 05:30 UTC

[Core/DBO][1/N] Add Dual-Batch Overlap mechanism to VLLM (#23693)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Sage Moore <sage@neuralmagic.com> Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Lucas Wilkinson <lwilkinson@neuralmagic.com> Co-authored-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

Sage Moore committed Sep 16, 2025 at 09:21 UTC 567939953b7a9cb0ded6bf0bb21a76917b8fed97

22 files changed +1255 -170

examples/offline_inference/data_parallel.py

@@ -87,6 +87,11 @@ def parse_args():

87 default=0.8,

88 help=("Fraction of GPU memory vLLM is allowed to allocate (0.0, 1.0]."),

89 )

90 + parser.add_argument(

91 + "--enable-dbo",

92 + action="store_true",

93 + help=("Enable microbatched execution"),

94 + )

95 parser.add_argument(

96 "--compilation-config",

97 type=int,

@@ -113,6 +118,7 @@ def main(

118 max_model_len,

119 compilation_config,

120 gpu_memory_utilization,

121 + enable_dbo,

122 quantization,

123 ):

124 os.environ["VLLM_DP_RANK"] = str(global_dp_rank)

@@ -167,6 +173,7 @@ def main(

173 max_num_seqs=max_num_seqs,

174 max_model_len=max_model_len,

175 gpu_memory_utilization=gpu_memory_utilization,

176 + enable_dbo=enable_dbo,

177 quantization=quantization,

178 compilation_config=compilation_config,

179 )

@@ -227,6 +234,7 @@ if __name__ == "__main__":

234 args.max_model_len,

235 args.compilation_config,

236 args.gpu_memory_utilization,

237 + args.enable_dbo,

238 args.quantization,

239 ),

240 )

tests/v1/attention/test_attention_splitting.py

+5 -5

@@ -6,7 +6,7 @@ import torch

7 from tests.v1.attention.test_attention_backends import BATCH_SPECS

8 from tests.v1.attention.utils import create_common_attn_metadata

9 -from vllm.v1.attention.backends.utils import (UbatchSlice,

9 +from vllm.v1.attention.backends.utils import (UBatchSlice,

10 _make_metadata_with_slice,

11 slice_query_start_locs,

12 split_attn_metadata)

@@ -106,7 +106,7 @@ def mixed_small_metadata():

106 def test_make_metadata_with_slice_decode_batch(small_decode_metadata):

107 """Test slicing decode batch metadata"""

108 # Split first request only

109 - ubatch_slice = UbatchSlice(slice(0, 1), slice(0, 1))

109 + ubatch_slice = UBatchSlice(slice(0, 1), slice(0, 1))

110

111 result = _make_metadata_with_slice(ubatch_slice, small_decode_metadata)

112

@@ -120,7 +120,7 @@ def test_make_metadata_with_slice_decode_batch(small_decode_metadata):

120

121 def test_make_metadata_with_slice_mixed_batch(mixed_small_metadata):

122 """Test slicing mixed batch metadata"""

123 - ubatch_slice = UbatchSlice(slice(1, 3),

123 + ubatch_slice = UBatchSlice(slice(1, 3),

124 slice(1, 7)) # Requests 1-3, tokens 1-7

125

126 result = _make_metadata_with_slice(ubatch_slice, mixed_small_metadata)

@@ -137,8 +137,8 @@ def test_split_attn_metadata_decode_batch(large_decode_metadata):

137 num_tokens = large_decode_metadata.num_reqs

138 mid_point = num_tokens // 2

139 ubatch_slices = [

140 - UbatchSlice(slice(0, mid_point), slice(0, mid_point)),

141 - UbatchSlice(slice(mid_point, num_tokens), slice(mid_point,

140 + UBatchSlice(slice(0, mid_point), slice(0, mid_point)),

141 + UBatchSlice(slice(mid_point, num_tokens), slice(mid_point,

142 num_tokens)),

143 ]

144

tests/v1/spec_decode/test_eagle.py

+6 -2

@@ -365,7 +365,9 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):

365 # Mock runner for attention metadata building

366 proposer.runner = mock.MagicMock()

367 proposer.runner.attn_groups.append([mock.MagicMock()])

368 - proposer.runner.attn_groups[0][0].metadata_builder = attn_metadata_builder

368 + proposer.runner.attn_groups[0][0].metadata_builders = [

369 + attn_metadata_builder

370 + ]

371

372 result = proposer.propose(target_token_ids=target_token_ids,

373 target_positions=target_positions,

@@ -489,7 +491,9 @@ def test_propose_tree(spec_token_tree):

491 # Mock runner for attention metadata building.

492 proposer.runner = mock.MagicMock()

493 proposer.runner.attn_groups.append([mock.MagicMock()])

492 - proposer.runner.attn_groups[0][0].metadata_builder = attn_metadata_builder

494 + proposer.runner.attn_groups[0][0].metadata_builders = [

495 + attn_metadata_builder

496 + ]

497

498 # Setup inputs for the proposer.

499 target_token_ids = torch.randint(0,

vllm/config/__init__.py

@@ -2848,6 +2848,14 @@ class VllmConfig:

2848 "when cudagraph_mode piecewise cudagraphs is used, "\

2849 f"cudagraph_mode={self.compilation_config.cudagraph_mode}"

2850

2851 + if self.parallel_config.enable_dbo:

2852 + a2a_backend = envs.VLLM_ALL2ALL_BACKEND

2853 + assert a2a_backend == "deepep_low_latency", \

2854 + "Microbatching currently only supports the deepep_low_latency "\

2855 + f"all2all backend. {a2a_backend} is not supported. To fix set "\

2856 + "the VLLM_ALL2ALL_BACKEND environment variable to "\

2857 + "deepep_low_latency and install the DeepEP kerenls."

2858 +

2859 if not self.instance_id:

2860 self.instance_id = random_uuid()[:5]

2861

vllm/config/parallel.py

@@ -137,6 +137,14 @@ class ParallelConfig:

137 disable_custom_all_reduce: bool = False

138 """Disable the custom all-reduce kernel and fall back to NCCL."""

139

140 + enable_dbo: bool = False

141 + """Enable microbatching for the model executor."""

142 +

143 + dbo_decode_token_threshold: int = 32

144 + """The threshold for microbatching. If the number of tokens in the

145 + request is greater than this threshold, microbatching will be used.

146 + Otherwise, the request will be processed in a single batch."""

147 +

148 ray_workers_use_nsight: bool = False

149 """Whether to profile Ray workers with nsight, see https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html#profiling-nsight-profiler."""

150

vllm/distributed/device_communicators/all2all.py

-5

@@ -251,9 +251,4 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):

251 logger.debug("DeepEP all2all args %s", buffer_kwargs)

252 handle: deep_ep.Buffer = self.handle_cache.get_or_create(

253 buffer_kwargs, deep_ep.Buffer)

254 - # It is dangerous to set num sms outside this function. num_sms is not

255 - # a part of the hash-key that identifies this object. If we are in a

256 - # situation where we make objects with different num_sms, the hash key

257 - # in get_or_create must be updated.

258 - handle.set_num_sms(self.num_sms)

254 return handle

vllm/engine/arg_utils.py

+10

@@ -327,6 +327,9 @@ class EngineArgs:

327 data_parallel_hybrid_lb: bool = False

328 data_parallel_backend: str = ParallelConfig.data_parallel_backend

329 enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel

330 + enable_dbo: bool = ParallelConfig.enable_dbo

331 + dbo_decode_token_threshold: int = \

332 + ParallelConfig.dbo_decode_token_threshold

333 eplb_config: EPLBConfig = get_field(ParallelConfig, "eplb_config")

334 enable_eplb: bool = ParallelConfig.enable_eplb

335 expert_placement_strategy: ExpertPlacementStrategy = \

@@ -695,6 +698,11 @@ class EngineArgs:

698 parallel_group.add_argument(

699 "--enable-expert-parallel",

700 **parallel_kwargs["enable_expert_parallel"])

701 + parallel_group.add_argument("--enable-dbo",

702 + **parallel_kwargs["enable_dbo"])

703 + parallel_group.add_argument(

704 + "--dbo-decode-token-threshold",

705 + **parallel_kwargs["dbo_decode_token_threshold"])

706 parallel_group.add_argument("--enable-eplb",

707 **parallel_kwargs["enable_eplb"])

708 parallel_group.add_argument("--eplb-config",

@@ -1339,6 +1347,8 @@ class EngineArgs:

1347 data_parallel_backend=self.data_parallel_backend,

1348 data_parallel_hybrid_lb=self.data_parallel_hybrid_lb,

1349 enable_expert_parallel=self.enable_expert_parallel,

1350 + enable_dbo=self.enable_dbo,

1351 + dbo_decode_token_threshold=self.dbo_decode_token_threshold,

1352 enable_eplb=self.enable_eplb,

1353 eplb_config=self.eplb_config,

1354 expert_placement_strategy=self.expert_placement_strategy,

vllm/forward_context.py

+101 -20

@@ -14,6 +14,7 @@ import vllm.envs as envs

14 from vllm.config import CUDAGraphMode, ParallelConfig, VllmConfig

15 from vllm.logger import init_logger

16 from vllm.platforms import current_platform

17 +from vllm.v1.worker.ubatch_utils import UBatchSlices, is_second_ubatch_empty

19 if TYPE_CHECKING:

20 from vllm.attention.backends.abstract import AttentionMetadata

@@ -97,6 +98,53 @@ class DPMetadata:

98 dist.all_reduce(num_tokens_tensor, group=group)

99 return num_tokens_tensor.cpu()

100

101 + @staticmethod

102 + def should_ubatch_across_dp(

103 + should_ubatch: bool, orig_num_tokens_per_ubatch: int,

104 + padded_num_tokens_per_ubatch: int, dp_size: int,

105 + dp_rank: int) -> tuple[bool, Optional[torch.Tensor]]:

106 + """

107 + 1. Decides if each DP rank is going to microbatch. Either all ranks

108 + run with microbatching or none of them do. If this function decides

109 + not to run with microbatching. It will "abort" meaning that no padding

110 + information will be returned to the caller. It will return (False, None)

111 +

112 + 2. Determines the total number of tokens that each rank will run.

113 + All ranks will be padded out so that the run with the same number

114 + of tokens

115 +

116 + Returns: tuple[

117 + should_ubatch: Are all DP ranks going to microbatch

118 + num_tokens_after_padding: A tensor containing the total number of

119 + tokens per-microbatch for each DP rank including padding. Will be

120 + None if should_ubatch if False

121 + ]

122 + """

123 +

124 + device = current_platform.device_type

125 + tensor = torch.zeros(3, dp_size, device=device, dtype=torch.int32)

126 + tensor[0][dp_rank] = orig_num_tokens_per_ubatch

127 + tensor[1][dp_rank] = padded_num_tokens_per_ubatch

128 + tensor[2][dp_rank] = 1 if should_ubatch else 0

129 +

130 + from vllm.distributed.parallel_state import get_dp_group

131 + dist.all_reduce(tensor, group=get_dp_group().device_group)

132 +

133 + result: bool = bool(torch.all(tensor[2] == 1).item())

134 + if not result:

135 + return result, None

136 +

137 + orig_num_tokens_tensor = tensor[0, :]

138 + padded_num_tokens_tensor = tensor[1, :]

139 +

140 + orig_min_num_tokens = int(orig_num_tokens_tensor.min().item())

141 + padded_max_num_tokens = int(padded_num_tokens_tensor.max().item())

142 + if is_second_ubatch_empty(orig_min_num_tokens, padded_max_num_tokens):

143 + logger.debug("Aborting ubatching %s %s", orig_min_num_tokens,

144 + padded_max_num_tokens)

145 + return False, None

146 + return result, padded_num_tokens_tensor.cpu()

147 +

148 @staticmethod

149 def make(

150 parallel_config: ParallelConfig,

@@ -119,14 +167,15 @@ class DPMetadata:

167

168 # If num_tokens_across_dp is None, it will be computed by all_reduce

169 # Otherwise, num_tokens_across_dp[dp_rank] should be equal to batchsize

122 - assert (num_tokens_across_dp is None

123 - or num_tokens_across_dp[dp_rank] == batchsize)

170 + assert (num_tokens_across_dp is None or num_tokens_across_dp[dp_rank]

171 + == batchsize), f"{num_tokens_across_dp[dp_rank]} {batchsize}"

172 if num_tokens_across_dp is None:

173 num_tokens_across_dp = DPMetadata.num_tokens_across_dp(

174 batchsize, dp_size, dp_rank)

175 max_tokens_across_dp_cpu = torch.max(num_tokens_across_dp)

176 cu_tokens_across_dp_cpu = torch.cumsum(num_tokens_across_dp, dim=0)

129 - return DPMetadata(max_tokens_across_dp_cpu, cu_tokens_across_dp_cpu)

177 + return DPMetadata(max_tokens_across_dp_cpu, cu_tokens_across_dp_cpu,

178 + num_tokens_across_dp)

179

180 @contextmanager

181 def chunked_sizes(self, max_chunk_size_per_rank: int, chunk_idx: int):

@@ -179,9 +228,12 @@ class ForwardContext:

228 Type AttentionMetadata for v0,

229 Type Dict[str, AttentionMetadata] for v1, map from layer_name of each

230 attention layer to its attention metadata

182 - set dynamically for each forward pass

231 + Type List[Dict[str, AttentionMetadata]] for DBO. List of size two, one

232 + for each microbatch.

233 + Set dynamically for each forward pass

234 """

184 - attn_metadata: Union["AttentionMetadata", dict[str, "AttentionMetadata"]]

235 + attn_metadata: Union["AttentionMetadata", dict[str, "AttentionMetadata"],

236 + list[dict[str, "AttentionMetadata"]]]

237 # TODO: remove after making all virtual_engines share the same kv cache

238 virtual_engine: int # set dynamically for each forward pass

239 # set dynamically for each forward pass

@@ -191,6 +243,8 @@ class ForwardContext:

243 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE

244 batch_descriptor: Optional[BatchDescriptor] = None

245

246 + ubatch_slices: Optional[UBatchSlices] = None

247 +

248 def __post_init__(self):

249 assert self.cudagraph_runtime_mode in [

250 CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL], \

@@ -208,6 +262,39 @@ def get_forward_context() -> ForwardContext:

262 return _forward_context

263

264

265 +def create_forward_context(

266 + attn_metadata: Any,

267 + vllm_config: VllmConfig,

268 + virtual_engine: int = 0,

269 + dp_metadata: Optional[DPMetadata] = None,

270 + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

271 + batch_descriptor: Optional[BatchDescriptor] = None,

272 + ubatch_slices: Optional[UBatchSlices] = None):

273 + return ForwardContext(no_compile_layers=vllm_config.compilation_config.

274 + static_forward_context,

275 + virtual_engine=virtual_engine,

276 + attn_metadata=attn_metadata,

277 + dp_metadata=dp_metadata,

278 + cudagraph_runtime_mode=cudagraph_runtime_mode,

279 + batch_descriptor=batch_descriptor,

280 + ubatch_slices=ubatch_slices)

281 +

282 +

283 +@contextmanager

284 +def override_forward_context(forward_context: Optional[ForwardContext]):

285 + """A context manager that overrides the current forward context.

286 + This is used to override the forward context for a specific

287 + forward pass.

288 + """

289 + global _forward_context

290 + prev_context = _forward_context

291 + _forward_context = forward_context

292 + try:

293 + yield

294 + finally:

295 + _forward_context = prev_context

296 +

297 +

298 @contextmanager

299 def set_forward_context(

300 attn_metadata: Any,

@@ -216,7 +303,8 @@ def set_forward_context(

303 num_tokens: Optional[int] = None,

304 num_tokens_across_dp: Optional[torch.Tensor] = None,

305 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

219 - batch_descriptor: Optional[BatchDescriptor] = None):

306 + batch_descriptor: Optional[BatchDescriptor] = None,

307 + ubatch_slices: Optional[UBatchSlices] = None):

308 """A context manager that stores the current forward context,

309 can be attention metadata, etc.

310 Here we can inject common logic for every model forward pass.

@@ -225,6 +313,7 @@ def set_forward_context(

313 need_to_track_batchsize = track_batchsize and attn_metadata is not None

314 if need_to_track_batchsize:

315 forward_start_time = time.perf_counter()

316 +

317 dp_metadata: Optional[DPMetadata] = None

318 if vllm_config.parallel_config.data_parallel_size > 1 and (

319 attn_metadata is not None or num_tokens is not None):

@@ -232,20 +321,14 @@ def set_forward_context(

321 attn_metadata, num_tokens or 0,

322 num_tokens_across_dp)

323

235 - global _forward_context

236 - prev_context = _forward_context

237 - _forward_context = ForwardContext(

238 - no_compile_layers=vllm_config.compilation_config.

239 - static_forward_context,

240 - virtual_engine=virtual_engine,

241 - attn_metadata=attn_metadata,

242 - dp_metadata=dp_metadata,

243 - cudagraph_runtime_mode=cudagraph_runtime_mode,

244 - batch_descriptor=batch_descriptor,

245 - )

324 + forward_context = create_forward_context(attn_metadata, vllm_config,

325 + virtual_engine, dp_metadata,

326 + cudagraph_runtime_mode,

327 + batch_descriptor, ubatch_slices)

328

329 try:

248 - yield

330 + with override_forward_context(forward_context):

331 + yield

332 finally:

333 global last_logging_time, batchsize_logging_interval

334 if need_to_track_batchsize:

@@ -282,5 +365,3 @@ def set_forward_context(

365 logger.info(("Batchsize forward time stats "

366 "(batchsize, count, median_time(ms)): %s"),

367 forward_stats)

285 -

286 - _forward_context = prev_context

vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py

+14 -12

@@ -191,7 +191,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

191 expert_map: Optional[torch.Tensor],

192 apply_router_weight_on_input: bool,

193 quant_config: FusedMoEQuantConfig,

194 - ) -> Callable:

194 + ) -> tuple[Callable, mk.ReceiverType]:

195

196 if apply_router_weight_on_input:

197 topk = topk_ids.size(1)

@@ -217,13 +217,14 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

217 a1q_scale = None

218 a1_post_scale = a1_scale

219

220 - return self._do_dispatch(tokens=a1q,

221 - token_scales=a1q_scale,

222 - rank_topk_ids=topk_ids,

223 - rank_topk_weights=topk_weights,

224 - num_experts=num_experts,

225 - a1_scale=a1_post_scale,

226 - quant_config=quant_config)

220 + return (lambda *args: None,

221 + self._do_dispatch(tokens=a1q,

222 + token_scales=a1q_scale,

223 + rank_topk_ids=topk_ids,

224 + rank_topk_weights=topk_weights,

225 + num_experts=num_experts,

226 + a1_scale=a1_post_scale,

227 + quant_config=quant_config))

228

229 def prepare(

230 self,

@@ -237,10 +238,11 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

238 apply_router_weight_on_input: bool,

239 quant_config: FusedMoEQuantConfig,

240 ) -> mk.PrepareResultType:

240 - receiver = self.prepare_async(a1, a1_scale, a2_scale, topk_weights,

241 - topk_ids, num_experts, expert_map,

242 - apply_router_weight_on_input,

243 - quant_config)

241 + (_, receiver) = self.prepare_async(a1, a1_scale, a2_scale,

242 + topk_weights, topk_ids, num_experts,

243 + expert_map,

244 + apply_router_weight_on_input,

245 + quant_config)

246 return receiver()

247

248 def finalize(

vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py

+29 -16

@@ -11,6 +11,9 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

11 TopKWeightAndReduceDelegate)

12 from vllm.model_executor.layers.fused_moe.utils import (

13 moe_kernel_quantize_input, normalize_batched_scales_shape)

14 +from vllm.v1.worker.ubatching import (dbo_current_ubatch_id, dbo_enabled,

15 + dbo_maybe_run_recv_hook,

16 + dbo_register_recv_hook, dbo_yield)

18 # DeepEP kernels quantize dispatch inputs in 128 element chunks.

19 DEEPEP_QUANT_BLOCK_SIZE = 128

@@ -55,7 +58,7 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

58 # The dispatch function returns a handle that the combine function

59 # requires. We store the handle here so it is available to the

60 # combine function.

58 - self.handle = None

61 + self.handles: list[Optional[tuple]] = [None, None]

62 self.num_dispatchers_ = num_dispatchers

64 def num_dispatchers(self) -> int:

@@ -123,13 +126,15 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

126 expert_map: Optional[torch.Tensor],

127 apply_router_weight_on_input: bool,

128 quant_config: FusedMoEQuantConfig,

126 - ) -> mk.ReceiverType:

129 + ) -> tuple[Callable, mk.ReceiverType]:

130

131 hidden_size = a1.size(1)

132 assert hidden_size in self.SUPPORTED_HIDDEN_SIZES, \

133 (f"Hidden Size {hidden_size} not in supported list of hidden sizes"

134 f"{self.SUPPORTED_HIDDEN_SIZES}")

135

136 + a2a_idx = dbo_current_ubatch_id()

137 +

138 if self.use_fp8_dispatch:

139 assert hidden_size % 128 == 0, \

140 "DeepEP kernels quantize the inputs in blocks of shape 128"

@@ -148,7 +153,7 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

153 a1 = a1 * topk_weights.to(a1.dtype)

154

155 # Dispatch

151 - expert_x, expert_num_tokens, self.handle, event, hook = \

156 + expert_x, expert_num_tokens, handle, _, hook= \

157 self.buffer.low_latency_dispatch(a1,

158 topk_ids,

159 self.max_tokens_per_rank,

@@ -156,21 +161,19 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

161 use_fp8=self.use_fp8_dispatch,

162 async_finish=False,

163 return_recv_hook=True)

164 + self.handles[a2a_idx] = handle

165

160 - return lambda: self._receiver(hook, expert_x, expert_num_tokens,

161 - a1_scale, a1.dtype, quant_config)

166 + return (hook, lambda: self._receiver(expert_x, expert_num_tokens,

167 + a1_scale, a1.dtype, quant_config))

168

169 def _receiver(

170 self,

165 - hook: Callable,

171 expert_x: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],

172 expert_num_tokens: torch.Tensor,

173 a1_scale,

174 a1_dtype,

175 quant_config: FusedMoEQuantConfig,

176 ) -> mk.PrepareResultType:

172 - hook()

173 -

177 expert_x, expert_x_scale = self._do_quant(

178 expert_x, a1_scale, a1_dtype, quant_config.quant_dtype,

179 quant_config.per_act_token_quant, quant_config.block_shape)

@@ -192,10 +195,12 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

195 apply_router_weight_on_input: bool,

196 quant_config: FusedMoEQuantConfig,

197 ) -> mk.PrepareResultType:

195 - receiver = self.prepare_async(a1, a1_scale, a2_scale, topk_weights,

196 - topk_ids, num_experts, expert_map,

197 - apply_router_weight_on_input,

198 - quant_config)

198 + hook, receiver = self.prepare_async(a1, a1_scale, a2_scale,

199 + topk_weights, topk_ids,

200 + num_experts, expert_map,

201 + apply_router_weight_on_input,

202 + quant_config)

203 + hook()

204 return receiver()

205

206 def finalize(

@@ -210,7 +215,11 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

215 assert isinstance(

216 weight_and_reduce_impl, TopKWeightAndReduceDelegate

217 ), ("Weight application and reduction happens in the combine kernel.")

213 - assert self.handle is not None

218 +

219 + a2a_idx = dbo_current_ubatch_id()

220 + do_recv_hook = dbo_enabled()

221 + handle = self.handles[a2a_idx]

222 + assert handle is not None

223

224 combine_topk_weights = topk_weights

225 if apply_router_weight_on_input:

@@ -218,12 +227,16 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

227 combine_topk_weights = torch.ones_like(topk_weights)

228

229 # TODO (varun) : Enable zero copy mode

221 - _, event, hook = self.buffer.low_latency_combine(

230 + dbo_maybe_run_recv_hook()

231 + _, _, recv_hook = self.buffer.low_latency_combine(

232 fused_expert_output,

233 topk_ids,

234 combine_topk_weights,

225 - self.handle,

235 + handle,

236 async_finish=False,

237 zero_copy=False,

228 - return_recv_hook=False,

238 + return_recv_hook=do_recv_hook,

239 out=output)

240 + if recv_hook is not None:

241 + dbo_register_recv_hook(recv_hook)

242 + dbo_yield()

vllm/model_executor/layers/fused_moe/layer.py

+43 -15

@@ -38,6 +38,7 @@ from vllm.platforms import current_platform

38 from vllm.platforms.interface import CpuArchEnum

39 from vllm.utils import (cdiv, direct_register_custom_op, has_deep_ep, has_pplx,

40 round_up)

41 +from vllm.v1.worker.ubatching import dbo_current_ubatch_id

43 if current_platform.is_cuda_alike():

44 from .fused_batched_moe import BatchedTritonExperts

@@ -992,16 +993,28 @@ class FusedMoE(CustomOp):

993 if (self.moe_parallel_config.use_pplx_kernels

994 or self.moe_parallel_config.use_deepep_ll_kernels

995 or self.moe_config.use_flashinfer_cutlass_kernels):

995 - self.batched_hidden_states = torch.zeros(

996 - (moe.max_num_tokens, self.hidden_size),

997 - dtype=moe.in_dtype,

998 - device=torch.cuda.current_device())

996 + if vllm_config.parallel_config.enable_dbo:

997 + self.batched_hidden_states = torch.zeros(

998 + (2, moe.max_num_tokens, self.hidden_size),

999 + dtype=moe.in_dtype,

1000 + device=torch.cuda.current_device())

1001 +

1002 + # Note here we use `num_experts` which is logical expert count

1003 + self.batched_router_logits = torch.zeros(

1004 + (2, moe.max_num_tokens, num_experts),

1005 + dtype=moe.in_dtype,

1006 + device=torch.cuda.current_device())

1007 + else:

1008 + self.batched_hidden_states = torch.zeros(

1009 + (moe.max_num_tokens, self.hidden_size),

1010 + dtype=moe.in_dtype,

1011 + device=torch.cuda.current_device())

1012

1000 - # Note here we use `num_experts` which is logical expert count

1001 - self.batched_router_logits = torch.zeros(

1002 - (moe.max_num_tokens, num_experts),

1003 - dtype=moe.in_dtype,

1004 - device=torch.cuda.current_device())

1013 + # Note here we use `num_experts` which is logical expert count

1014 + self.batched_router_logits = torch.zeros(

1015 + (moe.max_num_tokens, num_experts),

1016 + dtype=moe.in_dtype,

1017 + device=torch.cuda.current_device())

1018

1019 @property

1020 def shared_experts(self) -> Optional[torch.nn.Module]:

@@ -1708,14 +1721,29 @@ class FusedMoE(CustomOp):

1721 hidden_states = full_hidden_states[chunk_start:chunk_end, :]

1722 router_logits = full_router_logits[chunk_start:chunk_end, :]

1723

1711 - assert (self.batched_hidden_states.size(0) # type: ignore

1724 + assert self.batched_hidden_states is not None

1725 + assert self.batched_router_logits is not None

1726 + # This is only true when DBO has been enabled in the config.

1727 + # Both tensors will have an outer dimension for the ubatch id

1728 + if self.batched_hidden_states.dim() == 3:

1729 + assert self.batched_router_logits.dim() == 3

1730 + batch_buffer_idx = dbo_current_ubatch_id()

1731 + batched_hidden_states = self.batched_hidden_states[

1732 + batch_buffer_idx, :]

1733 + batched_router_logits = self.batched_router_logits[

1734 + batch_buffer_idx, :]

1735 + else:

1736 + batched_hidden_states = self.batched_hidden_states

1737 + batched_router_logits = self.batched_router_logits

1738 +

1739 + assert (batched_hidden_states.size(0) # type: ignore

1740 >= chunk_size)

1713 - assert (self.batched_router_logits.size(0) # type: ignore

1741 + assert (batched_router_logits.size(0) # type: ignore

1742 >= chunk_size)

1715 - staged_hidden_states = self.batched_hidden_states[:

1716 - chunk_size, :] # type: ignore

1717 - staged_router_logits = self.batched_router_logits[:

1718 - chunk_size, :] # type: ignore

1743 + staged_hidden_states = batched_hidden_states[:

1744 + chunk_size, :] # type: ignore

1745 + staged_router_logits = batched_router_logits[:

1746 + chunk_size, :] # type: ignore

1747 staged_hidden_states.copy_(hidden_states, non_blocking=True)

1748 staged_router_logits.copy_(router_logits, non_blocking=True)

1749

vllm/model_executor/layers/fused_moe/modular_kernel.py

+48 -15

@@ -13,6 +13,8 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig

13 from vllm.model_executor.layers.fused_moe.utils import ( # yapf: disable

14 _resize_cache, count_expert_num_tokens)

15 from vllm.utils import cdiv

16 +from vllm.v1.worker.ubatching import (dbo_enabled, dbo_maybe_run_recv_hook,

17 + dbo_register_recv_hook, dbo_yield)

19 #

20 # This file defines a set of base classes used to make MoE kernels more modular.

@@ -226,7 +228,7 @@ class FusedMoEPrepareAndFinalize(ABC):

228 expert_map: Optional[torch.Tensor],

229 apply_router_weight_on_input: bool,

230 quant_config: FusedMoEQuantConfig,

229 - ) -> ReceiverType:

231 + ) -> tuple[Callable, ReceiverType]:

232 """

233 Perform any quantization (and/or) dispatching needed for this kernel

234 but do not wait for results from other workers.

@@ -496,6 +498,23 @@ def _chunk_scales(scales: Optional[torch.Tensor], start: int,

498 return None

499

500

501 +class SharedResizableBuffer:

502 +

503 + def __init__(self):

504 + self.buffer = None

505 +

506 + def get(self, shape: tuple[int, ...], device: torch.device,

507 + dtype: torch.dtype):

508 + shape_numel = prod(shape)

509 + if self.buffer is None or self.buffer.numel() < shape_numel:

510 + self.buffer = torch.empty(shape_numel, device=device, dtype=dtype)

511 + assert self.buffer.device == device, \

512 + f"Buffer device mismatch: {self.buffer.device} != {device}"

513 + assert self.buffer.dtype == dtype, \

514 + f"Buffer dtype mismatch: {self.buffer.dtype} != {dtype}"

515 + return self.buffer[:shape_numel].view(*shape)

516 +

517 +

518 @final

519 class FusedMoEModularKernel(torch.nn.Module):

520 """

@@ -509,6 +528,9 @@ class FusedMoEModularKernel(torch.nn.Module):

528 layer due to any layer specific state that may be used by the component

529 objects.

530 """

531 + fused_out_buffer = SharedResizableBuffer()

532 + workspace13_buffer = SharedResizableBuffer()

533 + workspace2_buffer = SharedResizableBuffer()

534

535 def __init__(

536 self,

@@ -559,12 +581,12 @@ class FusedMoEModularKernel(torch.nn.Module):

581

582 # We can reuse the memory between cache1 and cache3 because by the

583 # time we need cache3, we're done with cache1.

562 - workspace13 = torch.empty(prod(workspace13_shape),

563 - device=a1.device,

564 - dtype=workspace_dtype)

565 - workspace2 = torch.empty(prod(workspace2_shape),

566 - device=a1.device,

567 - dtype=workspace_dtype)

584 + workspace13 = self.workspace13_buffer.get(workspace13_shape,

585 + device=a1.device,

586 + dtype=workspace_dtype)

587 + workspace2 = self.workspace2_buffer.get(workspace2_shape,

588 + device=a1.device,

589 + dtype=workspace_dtype)

590

591 assert fused_out is None or fused_out.shape == fused_out_shape, (

592 f"fused_out {fused_out.shape} but expected {fused_out_shape}")

@@ -656,9 +678,9 @@ class FusedMoEModularKernel(torch.nn.Module):

678 (_, _, fused_out_shape, _) = self.fused_experts.workspace_shapes(

679 a1, a1q, M, N, K, top_k, global_num_experts, local_num_experts,

680 expert_tokens_meta)

659 - fused_out = torch.empty(fused_out_shape,

660 - device=a1q.device,

661 - dtype=a1.dtype)

681 + fused_out = self.fused_out_buffer.get(fused_out_shape,

682 + device=a1q.device,

683 + dtype=a1.dtype)

684

685 def slice_input_tensors(

686 chunk_idx: int

@@ -801,8 +823,10 @@ class FusedMoEModularKernel(torch.nn.Module):

823

824 shared_output: torch.Tensor

825

804 - if (not self.prepare_finalize.supports_async()

805 - or self.shared_experts is None):

826 + if not self.prepare_finalize.supports_async():

827 + # We shouldn't be running an a2a kernel that doesn't

828 + # support async prepare/finalize

829 + assert not dbo_enabled()

830

831 # Run shared experts serially with dispatch.

832 if self.shared_experts is not None:

@@ -822,7 +846,8 @@ class FusedMoEModularKernel(torch.nn.Module):

846 )

847 else:

848 # Overlap shared expert compute with all2all dispatch.

825 - receiver = self.prepare_finalize.prepare_async(

849 + dbo_maybe_run_recv_hook()

850 + hook, receiver = self.prepare_finalize.prepare_async(

851 a1,

852 a1_scale,

853 a2_scale,

@@ -834,8 +859,16 @@ class FusedMoEModularKernel(torch.nn.Module):

859 self.fused_experts.quant_config,

860 )

861

837 - assert self.shared_experts is not None

838 - shared_output = self.shared_experts(a1)

862 + if self.shared_experts is not None:

863 + shared_output = self.shared_experts(a1)

864 +

865 + # If DBO is being used, register the hook with the ubatch context

866 + # and call it in dbo_maybe_run_recv_hook instead of passing it to

867 + # the receiver.

868 + dbo_register_recv_hook(hook)

869 + dbo_yield()

870 + if not dbo_enabled():

871 + hook()

872

873 (a1q, a1q_scale, expert_tokens_meta, _expert_topk_ids,

874 _expert_topk_weights) = receiver()

vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py

+18 -25

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Optional, Union

3 +from typing import Callable, Optional, Union

5 import pplx_kernels as pplx

6 import torch

@@ -103,7 +103,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

103 expert_map: Optional[torch.Tensor],

104 apply_router_weight_on_input: bool,

105 quant_config: FusedMoEQuantConfig,

106 - ) -> mk.ReceiverType:

106 + ) -> tuple[Callable, mk.ReceiverType]:

107 num_tokens = a1.size(0) # M

108 hidden_dim = a1.size(-1) # K

109

@@ -214,41 +214,33 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

214 do_recv=False,

215 )

216

217 - return lambda: self._receiver(

217 + hook = lambda: self.a2a.dispatch(

218 + out_expert_num_tokens=expert_num_tokens,

219 + out_expert_x=expert_x,

220 + out_expert_x_scale=expert_x_scale,

221 + dp_x=a1q,

222 + dp_x_scale=a1q_scale,

223 + indices=topk_ids,

224 + bound_m=bound_m,

225 + do_send=False,

226 + do_recv=True,

227 + )

228 +

229 + return (hook, lambda: self._receiver(

230 expert_num_tokens,

231 expert_x,

232 expert_x_scale,

221 - a1q,

222 - a1q_scale,

223 - topk_ids,

224 - bound_m,

233 orig_a_scale_block_shape,

226 - )

234 + ))

235

236 def _receiver(

237 self,

238 expert_num_tokens: torch.Tensor,

239 expert_x: torch.Tensor,

240 expert_x_scale: Optional[torch.Tensor],

233 - a1q: torch.Tensor,

234 - a1q_scale: Optional[torch.Tensor],

235 - topk_ids: torch.Tensor,

236 - bound_m: Optional[torch.Tensor],

241 orig_a_scale_block_shape: Optional[int],

242 ) -> mk.PrepareResultType:

243

240 - self.a2a.dispatch(

241 - out_expert_num_tokens=expert_num_tokens,

242 - out_expert_x=expert_x,

243 - out_expert_x_scale=expert_x_scale,

244 - dp_x=a1q,

245 - dp_x_scale=a1q_scale,

246 - indices=topk_ids,

247 - bound_m=bound_m,

248 - do_send=False,

249 - do_recv=True,

250 - )

251 -

244 if expert_x_scale is not None:

245 expert_x_scale = expert_x_scale[:, :, :orig_a_scale_block_shape]

246 assert expert_x_scale.ndim == 3

@@ -270,7 +262,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

262 apply_router_weight_on_input: bool,

263 quant_config: FusedMoEQuantConfig,

264 ) -> mk.PrepareResultType:

273 - receiver = self.prepare_async(

265 + hook, receiver = self.prepare_async(

266 a1,

267 a1_scale,

268 a2_scale,

@@ -281,6 +273,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

273 apply_router_weight_on_input,

274 quant_config,

275 )

276 + hook()

277 return receiver()

278

279 def finalize(

vllm/v1/attention/backends/utils.py

+9 -9

@@ -28,6 +28,7 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (

28 get_kv_connector_cache_layout)

29 from vllm.logger import init_logger

30 from vllm.v1.kv_cache_interface import AttentionSpec

31 +from vllm.v1.worker.ubatch_utils import UBatchSlice

33 logger = init_logger(__name__)

34 KVCacheLayoutType = Literal["NHD", "HND"]

@@ -81,12 +82,6 @@ class CommonAttentionMetadata:

82 encoder_seq_lens: Optional[np.ndarray] = None

84 -@dataclass

85 -class UbatchSlice:

86 - request_slice: slice

87 - token_slice: slice

88 -

89 -

85 def slice_query_start_locs(

86 query_start_loc: torch.Tensor,

87 request_slice: slice,

@@ -103,7 +98,7 @@ def slice_query_start_locs(

100 def _make_metadata_with_slice(

106 - ubatch_slice: UbatchSlice,

101 + ubatch_slice: UBatchSlice,

102 attn_metadata: CommonAttentionMetadata) -> CommonAttentionMetadata:

103 """

104 This function creates a new CommonAttentionMetadata that corresponds to

@@ -133,6 +128,11 @@ def _make_metadata_with_slice(

128 torch.max(torch.abs(query_start_loc_cpu[1:] -

129 query_start_loc_cpu[:-1])).item())

130

131 + # This is to account for the case where we are in a dummy

132 + # run and query_start_loc_cpu is full of 0s

133 + if max_query_len == 0:

134 + max_query_len = attn_metadata.max_query_len

135 +

136 block_table_tensor = attn_metadata.block_table_tensor[request_slice]

137 slot_mapping = attn_metadata.slot_mapping[token_slice]

138

@@ -152,12 +152,12 @@ def _make_metadata_with_slice(

152

153

154 def split_attn_metadata(

155 - ubatch_slices: list[UbatchSlice],

155 + ubatch_slices: list[UBatchSlice],

156 common_attn_metadata: CommonAttentionMetadata,

157 ) -> list[CommonAttentionMetadata]:

158 """

159 Creates a new CommonAttentionMetadata instance that corresponds to the

160 - requests for each UbatchSlice in ubatch_slices.

160 + requests for each UBatchSlice in ubatch_slices.

161

162 Note: This function does not modify common_attn_metadata

163 """

vllm/v1/spec_decode/eagle.py

+8 -4

@@ -27,6 +27,7 @@ from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata

27 from vllm.v1.attention.backends.utils import CommonAttentionMetadata

28 from vllm.v1.kv_cache_interface import KVCacheConfig

29 from vllm.v1.sample.metadata import SamplingMetadata

30 +from vllm.v1.worker.ubatching import dbo_current_ubatch_id

32 logger = init_logger(__name__)

@@ -179,9 +180,11 @@ class EagleProposer:

180 assert self.runner is not None

181

182 # FIXME: need to consider multiple kv_cache_groups

182 - attn_metadata = self.runner.attn_groups[0][0].metadata_builder\

183 - .build_for_drafting(common_attn_metadata=common_attn_metadata,

184 - draft_index=0)

183 + ubatch_id = dbo_current_ubatch_id()

184 + attn_metadata_builder = \

185 + self.runner.attn_groups[0][0].metadata_builders[ubatch_id]

186 + attn_metadata = attn_metadata_builder.build_for_drafting(

187 + common_attn_metadata=common_attn_metadata, draft_index=0)

188

189 # At this moment, we assume all eagle layers belong to the same KV

190 # cache group, thus using the same attention metadata.

@@ -355,8 +358,9 @@ class EagleProposer:

358 hidden_states: torch.Tensor,

359 common_attn_metadata: CommonAttentionMetadata,

360 ) -> list[torch.Tensor]:

361 + ubatch_id = dbo_current_ubatch_id()

362 tree_attn_metadata_builder = \

359 - self.runner.attn_groups[0][0].metadata_builder

363 + self.runner.attn_groups[0][0].metadata_builders[ubatch_id]

364 assert isinstance(tree_attn_metadata_builder,

365 TreeAttentionMetadataBuilder)

366

vllm/v1/worker/cpu_model_runner.py

+7 -2

@@ -64,8 +64,13 @@ class CPUModelRunner(GPUModelRunner):

64 if not self.attn_groups[0]:

65 return

67 - mb = getattr(self.attn_groups[0][0], "metadata_builder", None)

68 - if not isinstance(mb, TorchSDPAMetadataBuilderV1):

67 + mb = getattr(self.attn_groups[0][0], "metadata_builders", None)

68 + if isinstance(mb, list):

69 + if not isinstance(mb[0], TorchSDPAMetadataBuilderV1):

70 + return

71 + mb[0].reorder_batch(self.input_batch, scheduler_output)

72 + return

73 + elif not isinstance(mb, TorchSDPAMetadataBuilderV1):

74 # Encoder-only / rerank models do not benefit from reordering,

75 # so we safely skip here.

76 return

vllm/v1/worker/gpu_model_runner.py

+236 -39

@@ -15,6 +15,7 @@ import torch

15 import torch.distributed

16 import torch.nn as nn

17 from tqdm import tqdm

18 +from typing_extensions import TypeAlias

20 import vllm.envs as envs

21 from vllm.attention import Attention, AttentionType

@@ -55,11 +56,12 @@ from vllm.tasks import GenerationTask, PoolingTask, SupportedTask

56 from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, DeviceMemoryProfiler,

57 GiB_bytes, LazyLoader, check_use_alibi, get_dtype_size,

58 is_pin_memory_available, round_up, supports_dynamo)

59 +from vllm.v1.attention.backends.flash_attn import AttentionMetadata

60 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder

61 from vllm.v1.attention.backends.utils import (

62 AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata,

63 create_fast_prefill_custom_backend,

62 - reorder_batch_to_split_decodes_and_prefills)

64 + reorder_batch_to_split_decodes_and_prefills, split_attn_metadata)

65 from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher

66 # yapf conflicts with isort for this block

67 # yapf: disable

@@ -85,9 +87,12 @@ from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

87 from vllm.v1.spec_decode.ngram_proposer import NgramProposer

88 from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext

89 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch

90 +from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper

91 from vllm.v1.worker.kv_connector_model_runner_mixin import (

92 KVConnectorModelRunnerMixin)

93 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

94 +from vllm.v1.worker.ubatch_splitting import get_dp_padding_ubatch, ubatch_split

95 +from vllm.v1.worker.ubatch_utils import UBatchSlice, UBatchSlices

96 from vllm.v1.worker.utils import is_residual_scattered_for_sp

98 from .utils import (AttentionGroup, MultiModalBudget,

@@ -105,6 +110,11 @@ else:

110

111 logger = init_logger(__name__)

112

113 +AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata]

114 +# list when ubatching is enabled

115 +PerLayerAttnMetadata: TypeAlias = Union[list[AttnMetadataDict],

116 + AttnMetadataDict]

117 +

118

119 # Wrapper for ModelRunnerOutput to support overlapped execution.

120 class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

@@ -274,6 +284,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

284

285 # Request states.

286 self.requests: dict[str, CachedRequestState] = {}

287 + self.comm_stream = torch.cuda.Stream()

288

289 # Input Batch

290 # NOTE(Chen): Ideally, we should initialize the input batch inside

@@ -872,10 +883,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

883 return encoder_seq_lens

884

885 def _prepare_inputs(

875 - self,

876 - scheduler_output: "SchedulerOutput",

877 - ) -> tuple[dict[str, Any], torch.Tensor, Optional[SpecDecodeMetadata],

878 - np.ndarray, Optional[CommonAttentionMetadata], int]:

886 + self, scheduler_output: "SchedulerOutput"

887 + ) -> tuple[PerLayerAttnMetadata, torch.Tensor,

888 + Optional[SpecDecodeMetadata], np.ndarray,

889 + Optional[CommonAttentionMetadata], int, Optional[UBatchSlices],

890 + Optional[torch.Tensor]]:

891 """

892 :return: tuple[

893 attn_metadata: layer-to-attention_metadata mapping,

@@ -947,6 +959,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

959 self.query_start_loc.copy_to_gpu()

960 query_start_loc = self.query_start_loc.gpu[:num_reqs + 1]

961

962 + num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens

963 + num_tokens_padded = num_tokens_unpadded + self.get_local_padding(

964 + num_tokens_unpadded)

965 + ubatch_slices, num_tokens_after_padding = \

966 + ubatch_split(max_num_scheduled_tokens,

967 + num_tokens_unpadded,

968 + num_tokens_padded,

969 + self.vllm_config)

970 +

971 self.seq_lens.np[:num_reqs] = (

972 self.input_batch.num_computed_tokens_cpu[:num_reqs] +

973 num_scheduled_tokens)

@@ -1001,7 +1022,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1022 logits_indices_padded = self._prepare_kv_sharing_fast_prefill(

1023 logits_indices)

1024

1004 - attn_metadata: dict[str, Any] = {}

1025 + attn_metadata: PerLayerAttnMetadata = {}

1026 + if ubatch_slices is not None:

1027 + attn_metadata = [dict() for _ in range(len(ubatch_slices))]

1028

1029 # Used in the below loop.

1030 query_start_loc_cpu = self.query_start_loc.cpu[:num_reqs + 1]

@@ -1075,7 +1098,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1098 for attn_group in self.attn_groups[kv_cache_group_id]:

1099 # Prepare for cascade attention if enabled & beneficial.

1100 common_prefix_len = 0

1078 - builder = attn_group.metadata_builder

1101 + builder = attn_group.get_metadata_builder()

1102 if self.cascade_attn_enabled:

1103 common_prefix_len = self._compute_cascade_attn_prefix_len(

1104 num_scheduled_tokens,

@@ -1093,13 +1116,27 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1116 num_draft_tokens=self.num_draft_tokens.gpu[:num_reqs],

1117 )

1118

1096 - attn_metadata_i = builder.build(

1097 - common_prefix_len=common_prefix_len,

1098 - common_attn_metadata=common_attn_metadata,

1099 - **extra_attn_metadata_args)

1100 -

1101 - for layer_name in attn_group.layer_names:

1102 - attn_metadata[layer_name] = attn_metadata_i

1119 + if ubatch_slices is not None:

1120 + common_attn_metadata_list = split_attn_metadata(

1121 + ubatch_slices, common_attn_metadata)

1122 + for ubid, common_attn_metadata in enumerate(

1123 + common_attn_metadata_list):

1124 + assert common_attn_metadata.max_query_len == 1

1125 + attn_metadata_i = (attn_group.get_metadata_builder(

1126 + ubatch_id=ubid).build(

1127 + common_prefix_len=common_prefix_len,

1128 + common_attn_metadata=common_attn_metadata))

1129 + for layer_name in kv_cache_group_spec.layer_names:

1130 + assert type(attn_metadata) is list

1131 + attn_metadata[ubid][layer_name] = attn_metadata_i

1132 + else:

1133 + assert isinstance(attn_metadata, dict)

1134 + attn_metadata_i = builder.build(

1135 + common_prefix_len=common_prefix_len,

1136 + common_attn_metadata=common_attn_metadata,

1137 + **extra_attn_metadata_args)

1138 + for layer_name in attn_group.layer_names:

1139 + attn_metadata[layer_name] = attn_metadata_i

1140

1141 # Hot-Swap lora model

1142 if self.lora_config:

@@ -1107,7 +1144,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1144

1145 return (attn_metadata, logits_indices, spec_decode_metadata,

1146 num_scheduled_tokens, spec_decode_common_attn_metadata,

1110 - max_num_scheduled_tokens)

1147 + max_num_scheduled_tokens, ubatch_slices,

1148 + num_tokens_after_padding)

1149

1150 def _compute_cascade_attn_prefix_len(

1151 self,

@@ -1508,7 +1546,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1546

1547 def get_model(self) -> nn.Module:

1548 # get raw model out of the cudagraph wrapper.

1511 - if isinstance(self.model, CUDAGraphWrapper):

1549 + if isinstance(self.model, (CUDAGraphWrapper, UBatchWrapper)):

1550 return self.model.unwrap()

1551 return self.model

1552

@@ -1675,6 +1713,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1713

1714 def get_dp_padding(self,

1715 num_tokens: int) -> tuple[int, Optional[torch.Tensor]]:

1716 + """

1717 + Determines the total number of tokens that each rank will run.

1718 + All ranks will be padded out so that they run with the same number

1719 + of tokens

1720 +

1721 + Returns: tuple[

1722 + num_pad_tokens: The number of tokens that will be added to the batch

1723 + num_tokens_after_padding: A tensor containing the total number of

1724 + tokens for each DP rank including padding.

1725 + ]

1726 + """

1727 dp_size = self.vllm_config.parallel_config.data_parallel_size

1728 dp_rank = self.vllm_config.parallel_config.data_parallel_rank

1729

@@ -1698,6 +1747,39 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1747 dtype=torch.int32)

1748 return max_tokens_across_dp_cpu - num_tokens, num_tokens_after_padding

1749

1750 + def get_local_padding(self, num_tokens_unpadded: int) -> int:

1751 +

1752 + num_tokens_padded = num_tokens_unpadded

1753 +

1754 + if (self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

1755 + and num_tokens_unpadded <= self.cudagraph_batch_sizes[-1]):

1756 + # Use piecewise CUDA graphs.

1757 + # Add padding to the batch size.

1758 + num_tokens_padded = self.vllm_config.pad_for_cudagraph(

1759 + num_tokens_unpadded)

1760 + else:

1761 + # Eager mode.

1762 + # Pad tokens to multiple of tensor_parallel_size when

1763 + # enabled collective fusion for SP

1764 + tp_size = self.vllm_config.parallel_config.tensor_parallel_size

1765 + if self.vllm_config.compilation_config.pass_config. \

1766 + enable_sequence_parallelism and tp_size > 1:

1767 + num_tokens_padded = round_up(num_tokens_unpadded, tp_size)

1768 +

1769 + num_pad_tokens = num_tokens_padded - num_tokens_unpadded

1770 + return num_pad_tokens

1771 +

1772 + # This is where the second ubatch is adjusted to account for the padding.

1773 + # Should be called after attention metadata creation. This just pads

1774 + # the second ubatch slice out to the total number of tokens

1775 + # (num_tokens + padding)

1776 + def pad_out_ubatch_slice(self, ubatch_slices: UBatchSlices,

1777 + num_total_tokens: int):

1778 + padded_second_ubatch_slice = slice(ubatch_slices[1].token_slice.start,

1779 + num_total_tokens)

1780 + ubatch_slices[1] = UBatchSlice(padded_second_ubatch_slice,

1781 + padded_second_ubatch_slice)

1782 +

1783 def _pool(

1784 self,

1785 hidden_states: torch.Tensor,

@@ -1758,15 +1840,22 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1840 self,

1841 scheduler_output: "SchedulerOutput",

1842 intermediate_tensors: Optional[IntermediateTensors] = None,

1843 + ubatch_slices: Optional[UBatchSlices] = None,

1844 + num_tokens_after_padding: Optional[torch.Tensor] = None,

1845 ) -> tuple[int, int, Optional[torch.Tensor], Optional[torch.Tensor],

1846 Optional[torch.Tensor], torch.Tensor,

1847 Optional[IntermediateTensors], dict[str, Any]]:

1848

1849 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

1766 - num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)

1767 - # Padding for DP

1768 - num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)

1769 - num_input_tokens += num_pad

1850 + if ubatch_slices:

1851 + assert num_tokens_after_padding is not None

1852 + num_input_tokens = int(num_tokens_after_padding[0].item() * 2)

1853 + self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens)

1854 + elif ubatch_slices is None:

1855 + num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens)

1856 + num_pad, num_tokens_after_padding = self.get_dp_padding(

1857 + num_input_tokens)

1858 + num_input_tokens += num_pad

1859

1860 # _prepare_inputs may reorder the batch, so we must gather multi

1861 # modal outputs after that to ensure the correct order

@@ -1821,7 +1910,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1910 return (

1911 num_scheduled_tokens,

1912 num_input_tokens,

1824 - num_tokens_across_dp,

1913 + num_tokens_after_padding,

1914 input_ids,

1915 inputs_embeds,

1916 positions,

@@ -2027,7 +2116,8 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2116 # Prepare the decoder inputs.

2117 (attn_metadata, logits_indices, spec_decode_metadata,

2118 num_scheduled_tokens_np, spec_decode_common_attn_metadata,

2030 - max_query_len) = self._prepare_inputs(scheduler_output)

2119 + max_query_len, ubatch_slices, num_tokens_after_padding

2120 + ) = self._prepare_inputs(scheduler_output)

2121

2122 finally:

2123 if self.prepare_inputs_event is not None:

@@ -2042,7 +2132,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2132 positions,

2133 intermediate_tensors,

2134 model_kwargs,

2045 - ) = self._preprocess(scheduler_output, intermediate_tensors)

2135 + ) = self._preprocess(scheduler_output, intermediate_tensors,

2136 + ubatch_slices, num_tokens_after_padding)

2137 +

2138 + if ubatch_slices is not None:

2139 + num_input_tokens = num_input_tokens // 2

2140

2141 uniform_decode = (max_query_len

2142 == self.uniform_decode_query_len) and (

@@ -2062,6 +2156,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2156 num_tokens_across_dp=num_tokens_across_dp,

2157 cudagraph_runtime_mode=cudagraph_runtime_mode,

2158 batch_descriptor=batch_descriptor,

2159 + ubatch_slices=ubatch_slices,

2160 ), record_function_or_nullcontext("Forward"),

2161 self.maybe_get_kv_connector_output(scheduler_output) as

2162 kv_connector_output):

@@ -2441,10 +2536,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2536 # CudagraphWraper and CudagraphDispatcher of vllm.

2537

2538 # wrap the model with full cudagraph wrapper if needed.

2444 - if self.compilation_config.cudagraph_mode.has_full_cudagraphs():

2539 + if self.compilation_config.cudagraph_mode.has_full_cudagraphs() \

2540 + and not self.parallel_config.enable_dbo:

2541 self.model = CUDAGraphWrapper(self.model,

2542 self.vllm_config,

2543 runtime_mode=CUDAGraphMode.FULL)

2544 + elif self.parallel_config.enable_dbo:

2545 + if self.compilation_config.cudagraph_mode.has_full_cudagraphs():

2546 + self.model = UBatchWrapper(self.model, self.vllm_config,

2547 + CUDAGraphMode.FULL, self.device)

2548 + else:

2549 + self.model = UBatchWrapper(self.model, self.vllm_config,

2550 + CUDAGraphMode.NONE, self.device)

2551

2552 def reload_weights(self) -> None:

2553 assert getattr(self, "model", None) is not None, \

@@ -2642,6 +2745,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2745 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

2746 force_attention: bool = False,

2747 uniform_decode: bool = False,

2748 + allow_microbatching: bool = False,

2749 skip_eplb: bool = False,

2750 is_profile: bool = False,

2751 create_mixed_batch: bool = False,

@@ -2667,12 +2771,30 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2771 (1 token) and prefill (multiple tokens) requests.

2772 remove_lora: If False, dummy LoRAs are not destroyed after the run

2773 """

2774 + ubatch_enabled = self.parallel_config.enable_dbo

2775 + num_tokens_across_dp = None

2776 + num_pad = 0

2777 + should_ubatch = False

2778 + if ubatch_enabled:

2779 + should_ubatch = num_tokens >= \

2780 + self.parallel_config.dbo_decode_token_threshold and \

2781 + allow_microbatching

2782 +

2783 + (should_ubatch, num_tokens_across_dp) = get_dp_padding_ubatch(

2784 + num_tokens, num_tokens, should_ubatch, self.vllm_config)

2785 +

2786 + # Currently the dummy run should only be ubatching during

2787 + # cuda graph capture, meaning all DP ranks should already

2788 + # have the same batch size

2789 + if num_tokens_across_dp is not None:

2790 + assert int(num_tokens_across_dp[0]) == num_tokens // 2

2791 +

2792 assert cudagraph_runtime_mode in {

2793 CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL

2794 }

2795

2674 - # Padding for DP

2675 - num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)

2796 + if not should_ubatch:

2797 + num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)

2798 num_tokens += num_pad

2799

2800 # If cudagraph_mode.decode_mode() == FULL and

@@ -2690,6 +2812,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2812 # for GQA/MQA.

2813 max_query_len = self.uniform_decode_query_len if uniform_decode else \

2814 num_tokens

2815 + if allow_microbatching:

2816 + assert self.uniform_decode_query_len == 1

2817 + assert uniform_decode is True

2818 + assert max_query_len == 1

2819

2820 # Set num_scheduled_tokens based on num_tokens and max_num_seqs

2821 # for dummy run with LoRA so that the num_reqs collectively

@@ -2728,12 +2854,28 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2854 num_scheduled_tokens = np.array(num_scheduled_tokens_list,

2855 dtype=np.int32)

2856

2731 - attn_metadata: Optional[dict[str, Any]] = None

2857 + ubatch_slices = None

2858 + # We currently only microbatch if the number of tokens is

2859 + # over a certain threshold.

2860 + if should_ubatch:

2861 + # We only support decode-only cudagraphs

2862 + assert num_reqs == num_tokens

2863 + assert num_tokens % 2 == 0

2864 + ubatch_slices = [

2865 + UBatchSlice(slice(0, num_reqs // 2), slice(0,

2866 + num_tokens // 2)),

2867 + UBatchSlice(slice(num_reqs // 2, num_reqs),

2868 + slice(num_tokens // 2, num_tokens))

2869 + ]

2870 +

2871 + attn_metadata: Optional[PerLayerAttnMetadata] = None

2872

2873 # If force_attention is True, we always capture attention. Otherwise,

2874 # it only happens for cudagraph_runtime_mode=FULL.

2875 if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:

2876 attn_metadata = {}

2877 + if ubatch_slices is not None:

2878 + attn_metadata = [dict() for _ in range(len(ubatch_slices))]

2879

2880 if create_mixed_batch:

2881 # In the mixed batch mode (used for FI warmup), we use

@@ -2766,12 +2908,26 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2908 slot_mapping=self.input_batch.

2909 block_table[kv_cache_group_id].slot_mapping[:num_tokens],

2910 causal=True)

2769 -

2911 for attn_group in self.attn_groups[kv_cache_group_id]:

2771 - attn_metadata_i = attn_group.metadata_builder\

2772 - .build_for_cudagraph_capture(common_attn_metadata)

2773 - for layer_name in kv_cache_group_spec.layer_names:

2774 - attn_metadata[layer_name] = attn_metadata_i

2912 + if ubatch_slices is not None:

2913 + common_attn_metadata_list = split_attn_metadata(

2914 + ubatch_slices, common_attn_metadata)

2915 + for ubid, common_attn_metadata in enumerate(

2916 + common_attn_metadata_list):

2917 + assert common_attn_metadata.max_query_len == 1

2918 + attn_metadata_i = (attn_group\

2919 + .get_metadata_builder(ubatch_id=ubid)\

2920 + .build_for_cudagraph_capture(common_attn_metadata))

2921 + for layer_name in kv_cache_group_spec.layer_names:

2922 + assert type(attn_metadata) is list

2923 + attn_metadata[ubid][

2924 + layer_name] = attn_metadata_i

2925 + else:

2926 + assert type(attn_metadata) is dict

2927 + attn_metadata_i = attn_group.get_metadata_builder()\

2928 + .build_for_cudagraph_capture(common_attn_metadata)

2929 + for layer_name in kv_cache_group_spec.layer_names:

2930 + attn_metadata[layer_name] = attn_metadata_i

2931

2932 with self.maybe_dummy_run_with_lora(self.lora_config,

2933 num_scheduled_tokens, remove_lora):

@@ -2818,13 +2974,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2974 f"Cudagraph runtime mode mismatch at dummy_run. "

2975 f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}.")

2976

2977 + if ubatch_slices is not None:

2978 + num_tokens = num_tokens // 2

2979 with self.maybe_randomize_inputs(input_ids), set_forward_context(

2980 attn_metadata,

2981 self.vllm_config,

2982 num_tokens=num_tokens,

2983 num_tokens_across_dp=num_tokens_across_dp,

2984 cudagraph_runtime_mode=cudagraph_runtime_mode,

2827 - batch_descriptor=batch_descriptor):

2985 + batch_descriptor=batch_descriptor,

2986 + ubatch_slices=ubatch_slices):

2987 outputs = self.model(

2988 input_ids=input_ids,

2989 positions=positions,

@@ -3096,6 +3255,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3255 set_cudagraph_capturing_enabled(True)

3256 with freeze_gc(), graph_capture(device=self.device):

3257 cudagraph_mode = self.compilation_config.cudagraph_mode

3258 + assert cudagraph_mode is not None

3259 if cudagraph_mode.mixed_mode() != CUDAGraphMode.NONE:

3260 cudagraph_runtime_mode = cudagraph_mode.mixed_mode()

3261

@@ -3153,6 +3313,35 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3313 desc="Capturing CUDA graphs ({}, {})".format(

3314 "decode" if uniform_decode else "mixed prefill-decode",

3315 cudagraph_runtime_mode.name))

3316 + enable_dbo = self.parallel_config.enable_dbo

3317 + # DBO Only supports running Full cudagraphs with uniform

3318 + # decode lengths

3319 + if enable_dbo and uniform_decode:

3320 + for num_tokens in compilation_cases:

3321 + # If the number of tokens is greater than the microbatching

3322 + # threshold, don't generate a microbatched cudagraph

3323 + if (num_tokens

3324 + < self.parallel_config.dbo_decode_token_threshold):

3325 + continue

3326 +

3327 + # Warmup

3328 + for _ in range(

3329 + self.compilation_config.cudagraph_num_of_warmups):

3330 + force_attention = (

3331 + cudagraph_runtime_mode == CUDAGraphMode.FULL)

3332 + self._dummy_run(num_tokens,

3333 + cudagraph_runtime_mode=CUDAGraphMode.NONE,

3334 + force_attention=force_attention,

3335 + uniform_decode=True,

3336 + allow_microbatching=True,

3337 + skip_eplb=True)

3338 +

3339 + # Graph Capture

3340 + self._dummy_run(num_tokens,

3341 + cudagraph_runtime_mode=CUDAGraphMode.FULL,

3342 + uniform_decode=True,

3343 + allow_microbatching=True,

3344 + skip_eplb=True)

3345 # We skip EPLB here since we don't want to record dummy metrics

3346 for num_tokens in compilation_cases:

3347 for _ in range(self.compilation_config.cudagraph_num_of_warmups):

@@ -3219,14 +3408,23 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3408 ) -> list[AttentionGroup]:

3409 attn_groups: list[AttentionGroup] = []

3410 for attn_backend, layer_names in attn_backends_map.items():

3222 - attn_metadata_builder_i = attn_backend.get_builder_cls()(

3411 + attn_metadata_builders = []

3412 + attn_metadata_builders.append(attn_backend.get_builder_cls()(

3413 kv_cache_spec,

3414 layer_names,

3415 self.vllm_config,

3416 self.device,

3227 - )

3417 + ))

3418 + if self.parallel_config.enable_dbo:

3419 + attn_metadata_builders.append(

3420 + attn_backend.get_builder_cls()(

3421 + kv_cache_spec,

3422 + layer_names,

3423 + self.vllm_config,

3424 + self.device,

3425 + ))

3426 attn_group = AttentionGroup(attn_backend,

3229 - attn_metadata_builder_i,

3427 + attn_metadata_builders,

3428 layer_names)

3429 attn_groups.append(attn_group)

3430 return attn_groups

@@ -3246,11 +3444,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3444 min_cg_builder_name = None

3445

3446 for attn_group in self._attn_group_iterator():

3249 - builder = attn_group.metadata_builder

3447 + builder = attn_group.get_metadata_builder()

3448 if builder.cudagraph_support.value < min_cg_support.value:

3449 min_cg_support = builder.cudagraph_support

3450 min_cg_builder_name = builder.__class__.__name__

3253 -

3451 # Flexible resolve the cudagraph mode

3452 cudagraph_mode = self.compilation_config.cudagraph_mode

3453 # check cudagraph for mixed batch is supported

@@ -3316,7 +3513,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3513 is compatible (e.g., decode threshold is the same)

3514 """

3515 for group in self._attn_group_iterator():

3319 - attn_metadata_builder_i = group.metadata_builder

3516 + attn_metadata_builder_i = group.get_metadata_builder()

3517

3518 # check that if any backends reorder batches; that the reordering

3519 # is compatible (e.g., decode threshold is the same)

vllm/v1/worker/gpu_ubatch_wrapper.py new

+303

@@ -0,0 +1,303 @@

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

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

3 +

4 +import dataclasses

5 +import threading

6 +from typing import Any, Callable, Optional

7 +

8 +import torch

9 +

10 +from vllm.compilation.cuda_graph import CUDAGraphWrapper

11 +from vllm.config import CUDAGraphMode, VllmConfig

12 +from vllm.forward_context import (create_forward_context, get_forward_context,

13 + override_forward_context)

14 +from vllm.logger import init_logger

15 +from vllm.platforms import current_platform

16 +from vllm.sequence import IntermediateTensors

17 +from vllm.v1.worker.ubatching import UBatchContext, make_ubatch_contexts

18 +

19 +logger = init_logger(__name__)

20 +

21 +

22 +@dataclasses.dataclass

23 +class UbatchMetadata:

24 + context: UBatchContext

25 + input_ids: torch.Tensor

26 + positions: torch.Tensor

27 + inputs_embeds: Optional[torch.Tensor]

28 + intermediate_tensors: Optional[IntermediateTensors]

29 + num_tokens: int

30 +

31 +

32 +@dataclasses.dataclass

33 +class CUDAGraphMetaData:

34 + cudagraph: torch.cuda.CUDAGraph

35 + ubatch_metadata: UbatchMetadata

36 + outputs: Optional[Any] = None

37 +

38 +

39 +class UBatchWrapper:

40 +

41 + def __init__(self, runnable: Callable, vllm_config: VllmConfig,

42 + runtime_mode: CUDAGraphMode, device: torch.cuda.device):

43 + self.runnable = runnable

44 + self.vllm_config = vllm_config

45 + self.compilation_config = vllm_config.compilation_config

46 + self.comm_stream = torch.cuda.Stream(device=device)

47 + # Two ubatch threads plus the main thread

48 + self.ready_barrier = threading.Barrier(3)

49 +

50 + self.cudagraphs: dict[int, CUDAGraphMetaData] = {}

51 +

52 + self.cudagraph_wrapper = None

53 + self.graph_pool = None

54 + if runtime_mode is not CUDAGraphMode.NONE:

55 + self.cudagraph_wrapper = CUDAGraphWrapper(

56 + runnable, vllm_config, runtime_mode=runtime_mode)

57 + self.graph_pool = current_platform.get_global_graph_pool()

58 +

59 + def __getattr__(self, key: str):

60 + # allow accessing the attributes of the runnable.

61 + if hasattr(self.runnable, key):

62 + return getattr(self.runnable, key)

63 + raise AttributeError(f"Attribute {key} not exists in the runnable of "

64 + f"cudagraph wrapper: {self.runnable}")

65 +

66 + def unwrap(self) -> Callable:

67 + # in case we need to access the original runnable.

68 + return self.runnable

69 +

70 + def _capture_ubatches(self, ubatch_metadata, model) -> torch.Tensor:

71 + """

72 + Capture a cudagraph for a microbatched run.

73 +

74 + The logic here is somewhat complicated because we need to make sure that

75 + each of the ubatch threads initialize the cuda context before we start

76 + the graph capture.

77 +

78 + The flow is as follows:

79 + 1. The main thread starts up each ubatch thread. Each thread will

80 + initialize its cuda context (torch.cuda.current_blas_handle())

81 + before going to sleep upon entering the ubatch_context.

82 +

83 + 2. The main thread starts the graph capture and wakes up the first

84 + ubatch thread.

85 +

86 + 3. Each ubatch thread runs the model to completion and returns the

87 + completed output tensors back to the main thread.

88 +

89 + 4. The main thread stores the captured cudagraph along with its metadata

90 + and returns

91 + """

92 +

93 + @torch.inference_mode()

94 + def _capture_ubatch_thread(results, ubatch_metadata):

95 + ubatch_context = ubatch_metadata.context

96 + with torch.cuda.stream(ubatch_context.compute_stream):

97 + _ = torch.cuda.current_blas_handle()

98 + with torch.cuda.stream(ubatch_context.comm_stream):

99 + _ = torch.cuda.current_blas_handle()

100 + with ubatch_context:

101 + model_output = model(

102 + input_ids=ubatch_metadata.input_ids,

103 + positions=ubatch_metadata.positions,

104 + intermediate_tensors=ubatch_metadata.intermediate_tensors,

105 + inputs_embeds=ubatch_metadata.inputs_embeds,

106 + )

107 +

108 + results.append((ubatch_metadata.context.id, model_output))

109 +

110 + results: list[tuple[int, torch.Tensor]] = []

111 + compute_stream = ubatch_metadata[0].context.compute_stream

112 + num_tokens = ubatch_metadata[0].num_tokens + \

113 + ubatch_metadata[1].num_tokens

114 +

115 + # Ubatches will manually manage the forward context, so we override

116 + # it to None here so we can have it restored correctly later

117 + with override_forward_context(None):

118 + ubatch_threads = []

119 + for metadata in ubatch_metadata:

120 + thread = threading.Thread(target=_capture_ubatch_thread,

121 + args=(

122 + results,

123 + metadata,

124 + ))

125 + ubatch_threads.append(thread)

126 + thread.start()

127 + self.ready_barrier.wait() # Wait for both threads to be ready

128 +

129 + # Capture the cudagraph

130 + cudagraph_metadata = \

131 + CUDAGraphMetaData(

132 + cudagraph=torch.cuda.CUDAGraph(),

133 + ubatch_metadata=ubatch_metadata,

134 + )

135 + with torch.cuda.graph(cudagraph_metadata.cudagraph,

136 + stream=compute_stream,

137 + pool=self.graph_pool):

138 + ubatch_metadata[0].context.cpu_wait_event.set()

139 + for thread in ubatch_threads:

140 + thread.join()

141 + sorted_results = [value for position, value in sorted(results)]

142 + result = torch.cat(sorted_results, dim=0)

143 + cudagraph_metadata.outputs = result

144 + self.cudagraphs[num_tokens] = cudagraph_metadata

145 + return cudagraph_metadata.outputs

146 +

147 + def _run_ubatches(self, ubatch_metadata, model) -> torch.Tensor:

148 +

149 + @torch.inference_mode()

150 + def _ubatch_thread(results, model, ubatch_metadata):

151 + with ubatch_metadata.context:

152 + model_output = model(

153 + input_ids=ubatch_metadata.input_ids,

154 + positions=ubatch_metadata.positions,

155 + intermediate_tensors=ubatch_metadata.intermediate_tensors,

156 + inputs_embeds=ubatch_metadata.inputs_embeds,

157 + )

158 + results.append((ubatch_metadata.context.id, model_output))

159 +

160 + results: list[tuple[int, torch.Tensor]] = []

161 +

162 + # Ubatch threads will manually manage the forward context, so we

163 + # override it to None here so we can have it restored correctly

164 + # after both threads have finished

165 + with override_forward_context(None):

166 + ubatch_threads = []

167 + for metadata in ubatch_metadata:

168 + thread = threading.Thread(target=_ubatch_thread,

169 + args=(

170 + results,

171 + model,

172 + metadata,

173 + ))

174 + ubatch_threads.append(thread)

175 + thread.start()

176 + self.ready_barrier.wait() # Wait for both threads to be ready

177 + ubatch_metadata[0].context.cpu_wait_event.set()

178 + for thread in ubatch_threads:

179 + thread.join()

180 + sorted_results = [value for position, value in sorted(results)]

181 + result = torch.cat(sorted_results, dim=0)

182 + return result

183 +

184 + def _make_ubatch_metadata(self, ubatch_slices, attn_metadata, input_ids,

185 + positions, inputs_embeds, intermediate_tensors,

186 + compute_stream, dp_metadata, batch_descriptor,

187 + cudagraph_runtime_mode) -> list[UbatchMetadata]:

188 +

189 + # Create one forward context per ubatch

190 + forward_contexts = []

191 + for i, ubatch_slice in enumerate(ubatch_slices):

192 + forward_contexts.append(

193 + create_forward_context(

194 + attn_metadata[i] if attn_metadata is not None else None,

195 + self.vllm_config,

196 + dp_metadata=dp_metadata,

197 + batch_descriptor=batch_descriptor,

198 + cudagraph_runtime_mode=cudagraph_runtime_mode))

199 +

200 + ubatch_ctxs = make_ubatch_contexts(

201 + num_micro_batches=len(ubatch_slices),

202 + comm_stream=self.comm_stream,

203 + compute_stream=compute_stream,

204 + forward_contexts=forward_contexts,

205 + ready_barrier=self.ready_barrier)

206 +

207 + ubatch_metadata: list[UbatchMetadata] = []

208 + for i, ubatch_slice in enumerate(ubatch_slices):

209 + sliced_input_ids, sliced_positions, sliced_inputs_embeds, \

210 + sliced_intermediate_tensors = \

211 + self._slice_model_inputs(

212 + ubatch_slice.token_slice, input_ids, positions,

213 + inputs_embeds, intermediate_tensors)

214 + ubatch_metadata.append(

215 + UbatchMetadata(

216 + context=ubatch_ctxs[i],

217 + input_ids=sliced_input_ids,

218 + positions=sliced_positions,

219 + inputs_embeds=sliced_inputs_embeds,

220 + intermediate_tensors=sliced_intermediate_tensors,

221 + num_tokens=ubatch_slice.token_slice.stop -

222 + ubatch_slice.token_slice.start))

223 +

224 + return ubatch_metadata

225 +

226 + def _slice_model_inputs(self, tokens_slice: slice, input_ids, positions,

227 + inputs_embeds, intermediate_tensors):

228 + sliced_input_ids = input_ids[tokens_slice]

229 + # if we are using mrope. Mrope adds an additional dimension to the

230 + # positions tensor

231 + if positions.ndim == 2:

232 + sliced_positions = positions[:, tokens_slice]

233 + else:

234 + sliced_positions = positions[tokens_slice]

235 + sliced_inputs_embeds = inputs_embeds[

236 + tokens_slice] if inputs_embeds else None

237 + sliced_intermediate_tensors = intermediate_tensors[

238 + tokens_slice] if intermediate_tensors else None

239 +

240 + return (sliced_input_ids, sliced_positions, sliced_inputs_embeds,

241 + sliced_intermediate_tensors)

242 +

243 + def __call__(self, *args, **kwargs):

244 + forward_context = get_forward_context()

245 + batch_descriptor = forward_context.batch_descriptor

246 + ubatch_slices = forward_context.ubatch_slices

247 + cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode

248 +

249 + # If there's no ubatching, just run the runnable object

250 + if ubatch_slices is None:

251 + if cudagraph_runtime_mode in (CUDAGraphMode.NONE,

252 + CUDAGraphMode.PIECEWISE):

253 + return self.runnable(*args, **kwargs)

254 + else:

255 + assert self.cudagraph_wrapper is not None

256 + return self.cudagraph_wrapper(*args, **kwargs)

257 +

258 + attn_metadata = forward_context.attn_metadata

259 + num_tokens = (ubatch_slices[0].token_slice.stop -

260 + ubatch_slices[0].token_slice.start) * 2

261 + input_ids = kwargs['input_ids']

262 + positions = kwargs['positions']

263 + intermediate_tensors = kwargs['intermediate_tensors']

264 + inputs_embeds = kwargs['inputs_embeds']

265 + compute_stream = torch.cuda.current_stream()

266 +

267 + dp_metadata = forward_context.dp_metadata

268 +

269 + # We shouldn't be here unless we are running with multiple DP ranks

270 + assert dp_metadata is not None

271 +

272 + if num_tokens not in self.cudagraphs \

273 + and cudagraph_runtime_mode is CUDAGraphMode.FULL:

274 + ubatch_metadata = self._make_ubatch_metadata(

275 + ubatch_slices=ubatch_slices,

276 + attn_metadata=attn_metadata,

277 + input_ids=input_ids,

278 + positions=positions,

279 + intermediate_tensors=intermediate_tensors,

280 + inputs_embeds=inputs_embeds,

281 + compute_stream=compute_stream,

282 + dp_metadata=dp_metadata,

283 + batch_descriptor=batch_descriptor,

284 + cudagraph_runtime_mode=CUDAGraphMode.NONE)

285 +

286 + return self._capture_ubatches(ubatch_metadata, self.model)

287 + elif num_tokens in self.cudagraphs:

288 + cudagraph_metadata = self.cudagraphs[num_tokens]

289 + cudagraph_metadata.cudagraph.replay()

290 + return cudagraph_metadata.outputs

291 + else:

292 + ubatch_metadata = self._make_ubatch_metadata(

293 + ubatch_slices=ubatch_slices,

294 + attn_metadata=attn_metadata,

295 + input_ids=input_ids,

296 + positions=positions,

297 + intermediate_tensors=intermediate_tensors,

298 + inputs_embeds=inputs_embeds,

299 + compute_stream=compute_stream,

300 + dp_metadata=dp_metadata,

301 + batch_descriptor=batch_descriptor,

302 + cudagraph_runtime_mode=CUDAGraphMode.NONE)

303 + return self._run_ubatches(ubatch_metadata, self.model)

vllm/v1/worker/ubatch_splitting.py new

+155

@@ -0,0 +1,155 @@

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

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

3 +

4 +from typing import Optional

5 +

6 +import torch

7 +

8 +from vllm.config import VllmConfig

9 +from vllm.forward_context import DPMetadata

10 +from vllm.logger import init_logger

11 +from vllm.utils import round_up

12 +from vllm.v1.worker.ubatch_utils import (UBatchSlice, UBatchSlices,

13 + is_second_ubatch_empty)

14 +

15 +logger = init_logger(__name__)

16 +

17 +

18 +def should_ubatch_with_num_tokens(

19 + should_ubatch: bool,

20 + orig_num_tokens_per_ubatch: int,

21 + padded_num_tokens_per_ubatch: int,

22 + vllm_config: VllmConfig,

23 +) -> tuple[bool, Optional[torch.Tensor]]:

24 + dp_size = vllm_config.parallel_config.data_parallel_size

25 + dp_rank = vllm_config.parallel_config.data_parallel_rank

26 + return DPMetadata.should_ubatch_across_dp(should_ubatch,

27 + orig_num_tokens_per_ubatch,

28 + padded_num_tokens_per_ubatch,

29 + dp_size, dp_rank)

30 +

31 +

32 +def get_dp_padding_ubatch(

33 + num_tokens_unpadded: int, num_tokens_padded: int,

34 + should_attempt_ubatching: bool,

35 + vllm_config: VllmConfig) -> tuple[bool, Optional[torch.Tensor]]:

36 + """

37 + 1. Decides if each DP rank is going to microbatch. Either all ranks

38 + run with microbatching or none of them do. If this function decides

39 + not to run with microbatching. It will "abort" meaning that no padding

40 + information will be returned to the caller. It will return (False, None)

41 +

42 + 2. Determines the total number of tokens that each rank will run.

43 + All ranks will be padded out so that the run with the same number

44 + of tokens

45 +

46 + Returns: tuple[

47 + should_ubatch: Are all DP ranks going to microbatch

48 + num_tokens_after_padding: A tensor containing the total number of

49 + tokens per-microbatch for each DP rank including padding. Will be

50 + None if should_ubatch if False

51 + ]

52 +

53 + """

54 + assert num_tokens_padded >= num_tokens_unpadded

55 + dp_size = vllm_config.parallel_config.data_parallel_size

56 + if dp_size == 1:

57 + # Early exit.

58 + return False, None

59 +

60 + # If this DP rank doesn't want to attempt microbatching

61 + if not should_attempt_ubatching:

62 + (should_ubatch, num_tokens_across_dp) = should_ubatch_with_num_tokens(

63 + False, 0, 0, vllm_config)

64 + assert should_ubatch is False

65 + assert num_tokens_across_dp is None

66 + return should_ubatch, num_tokens_across_dp

67 +

68 + # Round up to the next multiple of two for even divisibility

69 + num_tokens_padded = round_up(num_tokens_padded, 2)

70 + num_tokens_per_ubatch = num_tokens_padded // 2

71 + should_ubatch = True

72 +

73 + # Sanity Check that the existing padding isn't giving us an empty second

74 + # ubatch. Abort if so

75 + if is_second_ubatch_empty(num_tokens_unpadded, num_tokens_padded):

76 + logger.debug(

77 + "Empty second µbatch detected: unpadded tokens: %s, padded "

78 + "tokens: %s", num_tokens_unpadded, num_tokens_padded)

79 + should_ubatch = False

80 +

81 + # Note that we compute the number of padded tokens per ubatch

82 + (should_ubatch, num_tokens_across_dp) = should_ubatch_with_num_tokens(

83 + should_ubatch, num_tokens_unpadded // 2, num_tokens_per_ubatch,

84 + vllm_config)

85 + if not should_ubatch:

86 + assert num_tokens_across_dp is None

87 + return should_ubatch, num_tokens_across_dp

88 +

89 + assert num_tokens_across_dp is not None

90 +

91 + max_tokens_across_dp_cpu = int(torch.max(num_tokens_across_dp).item())

92 + num_tokens_after_padding = torch.tensor([max_tokens_across_dp_cpu] *

93 + dp_size,

94 + device="cpu",

95 + dtype=torch.int32)

96 + return should_ubatch, num_tokens_after_padding

97 +

98 +

99 +def ubatch_split(

100 + max_num_scheduled_tokens: int,

101 + num_tokens_unpadded: int,

102 + num_tokens_padded: int,

103 + vllm_config: VllmConfig,

104 +) -> tuple[Optional[UBatchSlices], Optional[torch.Tensor]]:

105 + """

106 + Coordinates amongst all DP ranks to determine if and how the full batch

107 + should be split into microbatches.

108 +

109 + Returns: tuple[

110 + ubatch_slices: if this is set then all DP ranks have agreed to

111 + microbatch

112 + num_tokens_after_padding: A tensor containing the total number of

113 + tokens per-microbatch for each DP rank including padding. Will be

114 + None if ubatch_slices is None

115 + ]

116 +

117 + """

118 + parallel_config = vllm_config.parallel_config

119 + # Don't bother with the should_ubatch handshaking unless microbatching

120 + # is enabled

121 + if not parallel_config.enable_dbo:

122 + return (None, None)

123 +

124 + # Check preconditions for microbatching

125 + should_attempt_ubatching = \

126 + parallel_config.enable_dbo and \

127 + num_tokens_unpadded >= \

128 + parallel_config.dbo_decode_token_threshold \

129 + and max_num_scheduled_tokens == 1

130 +

131 + # Don't microbatch unless every other DP worker is also microbatching

132 + num_tokens_after_padding = None

133 + (should_ubatch, num_tokens_after_padding) = get_dp_padding_ubatch(

134 + num_tokens_unpadded, num_tokens_padded, should_attempt_ubatching,

135 + vllm_config)

136 + if not should_ubatch:

137 + return (None, None)

138 +

139 + # This doesn't actually pad the ubatch slices. It just initializes the

140 + # split point to the padded value so that padding can be applied

141 + # to the second ubatch in pad_out_ubatch_slice after attention

142 + # metadata creation

143 + assert num_tokens_after_padding is not None

144 + total_num_tokens_per_ubatch = int(num_tokens_after_padding[0].item())

145 + padded_first_ubatch_slice = slice(0, total_num_tokens_per_ubatch)

146 + padded_second_ubatch_slice = slice(total_num_tokens_per_ubatch,

147 + num_tokens_unpadded)

148 +

149 + # Note there's an assumption here that there's 1 token per request

150 + ubatch_slices = [

151 + UBatchSlice(padded_first_ubatch_slice, padded_first_ubatch_slice),

152 + UBatchSlice(padded_second_ubatch_slice, padded_second_ubatch_slice)

153 + ]

154 +

155 + return (ubatch_slices, num_tokens_after_padding)

vllm/v1/worker/ubatch_utils.py new

+19

@@ -0,0 +1,19 @@

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

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

3 +from dataclasses import dataclass

4 +

5 +from typing_extensions import TypeAlias

6 +

7 +

8 +@dataclass

9 +class UBatchSlice:

10 + request_slice: slice

11 + token_slice: slice

12 +

13 +

14 +UBatchSlices: TypeAlias = list[UBatchSlice]

15 +

16 +

17 +def is_second_ubatch_empty(orig_num_tokens_per_ubatch: int,

18 + padded_num_tokens_per_ubatch: int) -> bool:

19 + return padded_num_tokens_per_ubatch >= 2 * orig_num_tokens_per_ubatch

vllm/v1/worker/ubatching.py new

+211

@@ -0,0 +1,211 @@

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

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

3 +import threading

4 +from typing import Optional

5 +

6 +import torch

7 +

8 +from vllm import forward_context

9 +from vllm.forward_context import ForwardContext

10 +from vllm.utils import current_stream

11 +

12 +_THREAD_ID_TO_CONTEXT: dict = {}

13 +_CURRENT_CONTEXTS: list[Optional['UBatchContext']] = [None, None]

14 +

15 +

16 +class UBatchContext:

17 + """

18 + Context manager for micro-batching synchronization using threading events.

19 + """

20 +

21 + def __init__(self,

22 + id: int,

23 + comm_stream: torch.cuda.Stream,

24 + compute_stream: torch.cuda.Stream,

25 + forward_context: ForwardContext,

26 + ready_barrier: threading.Barrier,

27 + cpu_wait_event: threading.Event,

28 + cpu_signal_event: threading.Event,

29 + gpu_comm_done_event: torch.cuda.Event,

30 + gpu_compute_done_event: torch.cuda.Event,

31 + schedule: str = "default"):

32 + self.id = id

33 + self.comm_stream = comm_stream

34 + self.compute_stream = compute_stream

35 + self.forward_context = forward_context

36 + self.ready_barrier = ready_barrier

37 + self.cpu_wait_event = cpu_wait_event

38 + self.cpu_signal_event = cpu_signal_event

39 + self.current_stream = compute_stream

40 + self.gpu_comm_done_event = gpu_comm_done_event

41 + self.gpu_compute_done_event = gpu_compute_done_event

42 + self.schedule = schedule

43 + self.recv_hook = None

44 +

45 + def __enter__(self):

46 + global _CURRENT_CONTEXTS, _THREAD_ID_TO_CONTEXT

47 + _THREAD_ID_TO_CONTEXT[threading.get_ident()] = self.id

48 + _CURRENT_CONTEXTS[self.id] = self

49 + self.ready_barrier.wait()

50 +

51 + self.cpu_wait_event.wait()

52 + self.cpu_wait_event.clear()

53 + self._restore_context()

54 + # Assume we start on the compute stream

55 + assert current_stream() == self.compute_stream

56 + return self

57 +

58 + def __exit__(self, exc_type, exc_val, exc_tb):

59 + global _CURRENT_CONTEXTS, _THREAD_ID_TO_CONTEXT

60 + _CURRENT_CONTEXTS[self.id] = None

61 + del _THREAD_ID_TO_CONTEXT[threading.get_ident()]

62 + self.maybe_run_recv_hook()

63 + self.cpu_signal_event.set()

64 + self.cpu_wait_event.clear()

65 + self.current_stream = self.compute_stream

66 + torch.cuda.set_stream(self.current_stream)

67 + return False

68 +

69 + def _restore_context(self):

70 + forward_context._forward_context = self.forward_context

71 + torch.cuda.set_stream(self.current_stream)

72 +

73 + def update_stream(self, stream):

74 + self.current_stream = stream

75 + torch.cuda.set_stream(self.current_stream)

76 +

77 + def _signal_comm_done(self):

78 + self.gpu_comm_done_event.record(self.comm_stream)

79 +

80 + def _signal_compute_done(self):

81 + self.gpu_compute_done_event.record(self.compute_stream)

82 +

83 + def _wait_compute_done(self):

84 + self.comm_stream.wait_event(self.gpu_compute_done_event)

85 +

86 + def _wait_comm_done(self):

87 + self.compute_stream.wait_event(self.gpu_comm_done_event)

88 +

89 + def _cpu_yield(self):

90 + # It is critical for correctness that only one thread is running

91 + # at a time. These asserts just make sure that this is the only

92 + # thread running before waking the other one up and going to sleep

93 + assert forward_context._forward_context == self.forward_context

94 + assert current_stream() == self.current_stream

95 + assert not self.cpu_wait_event.is_set()

96 +

97 + self.cpu_signal_event.set()

98 + self.cpu_wait_event.wait()

99 + self.cpu_wait_event.clear()

100 + self._restore_context()

101 +

102 + def switch_to_comm_sync(self):

103 + self._signal_compute_done()

104 + self.update_stream(self.comm_stream)

105 + self._wait_comm_done()

106 +

107 + def maybe_run_recv_hook(self):

108 + if self.recv_hook is not None:

109 + self.recv_hook()

110 + self.recv_hook = None

111 +

112 + def yield_(self):

113 + self.current_stream = current_stream()

114 + self._cpu_yield()

115 + if self.current_stream != current_stream():

116 + self.update_stream(self.current_stream)

117 +

118 + def yield_and_switch_from_compute_to_comm(self):

119 + assert current_stream() == self.compute_stream

120 + self._signal_compute_done()

121 + self._cpu_yield()

122 + assert self.current_stream == self.compute_stream

123 + self.update_stream(self.comm_stream)

124 + self._wait_compute_done()

125 +

126 + def yield_and_switch_from_comm_to_compute(self):

127 + assert current_stream() == self.comm_stream

128 + self._signal_comm_done()

129 + self._cpu_yield()

130 + assert self.current_stream == self.comm_stream

131 + self.update_stream(self.compute_stream)

132 + self._wait_comm_done()

133 +

134 +

135 +def dbo_enabled() -> bool:

136 + return len(_THREAD_ID_TO_CONTEXT) > 0

137 +

138 +

139 +def dbo_current_ubatch_id() -> int:

140 + if len(_THREAD_ID_TO_CONTEXT) == 0:

141 + return 0

142 + return _THREAD_ID_TO_CONTEXT[threading.get_ident()]

143 +

144 +

145 +def _register_ubatch_function(func):

146 +

147 + def wrapper(*args, **kwargs):

148 + if len(_THREAD_ID_TO_CONTEXT) > 0:

149 + ctx_idx = _THREAD_ID_TO_CONTEXT[threading.get_ident()]

150 + ctx = _CURRENT_CONTEXTS[ctx_idx]

151 + func(ctx, *args, **kwargs)

152 +

153 + return wrapper

154 +

155 +

156 +dbo_yield_and_switch_from_compute_to_comm = _register_ubatch_function(

157 + UBatchContext.yield_and_switch_from_compute_to_comm)

158 +dbo_yield_and_switch_from_comm_to_compute = _register_ubatch_function(

159 + UBatchContext.yield_and_switch_from_comm_to_compute)

160 +dbo_yield = _register_ubatch_function(UBatchContext.yield_)

161 +dbo_maybe_run_recv_hook = _register_ubatch_function(

162 + UBatchContext.maybe_run_recv_hook)

163 +dbo_switch_to_comm_sync = _register_ubatch_function(

164 + UBatchContext.switch_to_comm_sync)

165 +

166 +

167 +def dbo_register_recv_hook(recv_hook):

168 + if len(_THREAD_ID_TO_CONTEXT) > 0:

169 + ctx_idx = _THREAD_ID_TO_CONTEXT[threading.get_ident()]

170 + next_ctx = _CURRENT_CONTEXTS[(ctx_idx + 1) % 2]

171 + next_ctx.recv_hook = recv_hook

172 +

173 +

174 +def make_ubatch_contexts(

175 + num_micro_batches: int,

176 + compute_stream: torch.cuda.Stream,

177 + comm_stream: torch.cuda.Stream,

178 + forward_contexts: list[ForwardContext],

179 + ready_barrier: threading.Barrier,

180 + schedule: str = "default",

181 +) -> list[UBatchContext]:

182 + assert num_micro_batches == 2, "only been tested with 2 micro-batches"

183 + """

184 + Create a context manager for micro-batching synchronization.

185 + """

186 + cpu_events = [threading.Event() for _ in range(num_micro_batches)]

187 + gpu_comm_done_events = [

188 + torch.cuda.Event() for _ in range(num_micro_batches)

189 + ]

190 + gpu_compute_done_events = [

191 + torch.cuda.Event() for _ in range(num_micro_batches)

192 + ]

193 +

194 + assert len(forward_contexts) == 2

195 +

196 + ctxs = []

197 + for i in range(num_micro_batches):

198 + ctx = UBatchContext(id=i,

199 + compute_stream=compute_stream,

200 + comm_stream=comm_stream,

201 + forward_context=forward_contexts[i],

202 + ready_barrier=ready_barrier,

203 + cpu_wait_event=cpu_events[i],

204 + cpu_signal_event=cpu_events[(i + 1) %

205 + num_micro_batches],

206 + gpu_comm_done_event=gpu_comm_done_events[i],

207 + gpu_compute_done_event=gpu_compute_done_events[i],

208 + schedule=schedule)

209 + ctxs.append(ctx)

210 +

211 + return ctxs

vllm/v1/worker/utils.py

+9 -1

@@ -130,9 +130,17 @@ class MultiModalBudget:

130 @dataclass

131 class AttentionGroup:

132 backend: type[AttentionBackend]

133 - metadata_builder: AttentionMetadataBuilder

133 + metadata_builders: list[AttentionMetadataBuilder]

134 layer_names: list[str]

135

136 + def get_metadata_builder(self,

137 + ubatch_id: Optional[int] = None

138 + ) -> AttentionMetadataBuilder:

139 + if ubatch_id is None:

140 + return self.metadata_builders[0]

141 + assert len(self.metadata_builders) > ubatch_id

142 + return self.metadata_builders[ubatch_id]

143 +

144

145 def sanity_check_mm_encoder_outputs(

146 mm_embeddings: MultiModalEmbeddings,