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

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][2/N] Dual-Batch Overlap add DeepEP High Throughput support and Prefill support (#24845)

Signed-off-by: Sage Moore <sage@neuralmagic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: yewentao256 <zhyanwentao@126.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Sage Moore <sage@neuralmagic.com> Co-authored-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>

Lucas Wilkinson committed Sep 23, 2025 at 12:02 UTC cc1dc7ed6d2d901df46847724ad90147872a9d79

19 files changed +604 -238

tests/v1/attention/test_attention_splitting.py

+82 -1

@@ -5,11 +5,12 @@ import pytest

5 import torch

7 from tests.v1.attention.test_attention_backends import BATCH_SPECS

8 -from tests.v1.attention.utils import create_common_attn_metadata

8 +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata

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

10 _make_metadata_with_slice,

11 slice_query_start_locs,

12 split_attn_metadata)

13 +from vllm.v1.worker.ubatch_utils import create_ubatch_slices

16 @pytest.fixture

@@ -155,3 +156,83 @@ def test_split_attn_metadata_decode_batch(large_decode_metadata):

156 assert results[1].num_reqs == mid_point

157 assert results[1].num_actual_tokens == mid_point

158 assert torch.equal(results[1].seq_lens, torch.tensor([2048] * mid_point))

159 +

160 +

161 +@pytest.mark.parametrize(

162 + "seq_lens,query_lens,split_point,expected_first_reqs,expected_second_reqs",

163 + [

164 + # Split in the middle of request 1

165 + ([32, 40], [8, 8], 12, 2, 1),

166 + # Split inside the first request

167 + ([32, 40], [8, 8], 4, 1, 2),

168 + ],

169 +)

170 +def test_prefill_split_across_ubatches(seq_lens, query_lens, split_point,

171 + expected_first_reqs,

172 + expected_second_reqs):

173 + """Test splitting a prefill across ubatches"""

174 + import numpy as np

175 +

176 + device = torch.device("cpu")

177 + batch_spec = BatchSpec(seq_lens=seq_lens, query_lens=query_lens)

178 + common = create_common_attn_metadata(batch_spec,

179 + block_size=16,

180 + device=device)

181 +

182 + num_scheduled_tokens = np.array(query_lens, dtype=np.int32)

183 + qsl_np = common.query_start_loc_cpu.numpy()

184 + num_tokens = common.num_actual_tokens

185 +

186 + ubatch_slices = create_ubatch_slices(num_scheduled_tokens, split_point)

187 + assert len(ubatch_slices) == 2

188 +

189 + first_meta = _make_metadata_with_slice(ubatch_slices[0], common)

190 + second_meta = _make_metadata_with_slice(ubatch_slices[1], common)

191 +

192 + # Token counts match the split

193 + assert first_meta.num_actual_tokens == split_point

194 + assert second_meta.num_actual_tokens == num_tokens - split_point

195 +

196 + # Number of requests per ubatch

197 + assert first_meta.num_reqs == expected_first_reqs

198 + assert second_meta.num_reqs == expected_second_reqs

199 +

200 + # Identify which request is split and how many tokens are in the first chunk

201 + split_req_idx = int(np.searchsorted(qsl_np, split_point, side="right") - 1)

202 + tokens_in_first_chunk = split_point - int(qsl_np[split_req_idx])

203 + orig_q_lens = (common.query_start_loc_cpu[1:] -

204 + common.query_start_loc_cpu[:-1])

205 +

206 + # Check query length continuity: first-chunk + second-chunk == original qlen

207 + # First ubatch last request query length

208 + qlen_first_last = int(first_meta.query_start_loc_cpu[-1] -

209 + first_meta.query_start_loc_cpu[-2])

210 + # Second ubatch first request query length

211 + qlen_second_first = int(second_meta.query_start_loc_cpu[1] -

212 + second_meta.query_start_loc_cpu[0])

213 + assert qlen_first_last == tokens_in_first_chunk

214 + assert qlen_first_last + qlen_second_first == int(

215 + orig_q_lens[split_req_idx])

216 +

217 + # Check seq_lens adjustments

218 + # Context lengths per original request

219 + context_lens = [s - q for s, q in zip(seq_lens, query_lens)]

220 +

221 + # First ubatch: last request's seq_len should be

222 + # context + tokens_in_first_chunk

223 + expected_seqlen = context_lens[split_req_idx] + tokens_in_first_chunk

224 + assert int(first_meta.seq_lens[-1]) == expected_seqlen

225 +

226 + # For full preceding requests in first ubatch, seq_lens should match

227 + # originals

228 + for i in range(first_meta.num_reqs - 1):

229 + assert int(first_meta.seq_lens[i]) == seq_lens[i]

230 +

231 + # Second ubatch: first request (continuation) seq_len should be full

232 + # original

233 + assert int(second_meta.seq_lens[0]) == seq_lens[split_req_idx]

234 + # Any following full requests in second ubatch should match originals

235 + for j in range(1, second_meta.num_reqs):

236 + # Map to original request index

237 + orig_idx = split_req_idx + j

238 + assert int(second_meta.seq_lens[j]) == seq_lens[orig_idx]

tests/v1/spec_decode/test_eagle.py

+2 -4

@@ -532,9 +532,8 @@ def test_propose(method, attn_backend, num_speculative_tokens, monkeypatch):

532 # Mock runner for attention metadata building

533 proposer.runner = mock.MagicMock()

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

535 - proposer.runner.attn_groups[0][0].metadata_builders = [

535 + proposer.runner.attn_groups[0][0].get_metadata_builder.return_value = \

536 attn_metadata_builder

537 - ]

537

538 result = proposer.propose(target_token_ids=target_token_ids,

539 target_positions=target_positions,

@@ -659,9 +658,8 @@ def test_propose_tree(spec_token_tree):

658 # Mock runner for attention metadata building.

659 proposer.runner = mock.MagicMock()

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

662 - proposer.runner.attn_groups[0][0].metadata_builders = [

661 + proposer.runner.attn_groups[0][0].get_metadata_builder.return_value = \

662 attn_metadata_builder

664 - ]

663

664 # Setup inputs for the proposer.

665 target_token_ids = torch.randint(0,

vllm/config/__init__.py

+7 -5

@@ -638,11 +638,13 @@ class VllmConfig:

638

639 if self.parallel_config.enable_dbo:

640 a2a_backend = envs.VLLM_ALL2ALL_BACKEND

641 - assert a2a_backend == "deepep_low_latency", \

642 - "Microbatching currently only supports the deepep_low_latency "\

643 - f"all2all backend. {a2a_backend} is not supported. To fix set "\

644 - "the VLLM_ALL2ALL_BACKEND environment variable to "\

645 - "deepep_low_latency and install the DeepEP kerenls."

641 + assert a2a_backend in \

642 + ["deepep_low_latency", "deepep_high_throughput"], \

643 + "Microbatching currently only supports the deepep_low_latency and "\

644 + f"deepep_high_throughput all2all backend. {a2a_backend} is not "\

645 + "supported. To fix set the VLLM_ALL2ALL_BACKEND environment "\

646 + "variable to deepep_low_latency or deepep_high_throughput and "\

647 + "install the DeepEP kernels."

648

649 if not self.instance_id:

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

vllm/config/parallel.py

+10 -4

@@ -139,12 +139,18 @@ class ParallelConfig:

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

140

141 enable_dbo: bool = False

142 - """Enable microbatching for the model executor."""

142 + """Enable dual batch overlap for the model executor."""

143

144 dbo_decode_token_threshold: int = 32

145 - """The threshold for microbatching. If the number of tokens in the

146 - request is greater than this threshold, microbatching will be used.

147 - Otherwise, the request will be processed in a single batch."""

145 + """The threshold for dual batch overlap for batches only containing decodes.

146 + If the number of tokens in the request is greater than this threshold,

147 + microbatching will be used. Otherwise, the request will be processed in a

148 + single batch."""

149 + dbo_prefill_token_threshold: int = 512 # TODO(lucas): tune

150 + """The threshold for dual batch overlap for batches that contain one or more

151 + prefills. If the number of tokens in the request is greater than this

152 + threshold, microbatching will be used. Otherwise, the request will be

153 + processed in a single batch."""

154

155 ray_workers_use_nsight: bool = False

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

vllm/distributed/device_communicators/all2all.py

+19 -9

@@ -1,10 +1,11 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 -from typing import Any

3 +from typing import Any, Optional

5 import torch

6 import torch.distributed as dist

8 +import vllm.envs as envs

9 from vllm.distributed import get_dp_group

10 from vllm.forward_context import get_forward_context

11 from vllm.logger import init_logger

@@ -200,12 +201,12 @@ class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):

201

202 def _make_all2all_kwargs(self) -> dict[Any, Any]:

203 # Defaults for internode and intranode are taken from DeepEP tests.

203 - num_nvl_bytes = 1024 * 1024 * 1024

204 + num_nvl_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024

205 num_rdma_bytes = None

206 num_qps_per_rank = None

207

208 if self.internode:

208 - num_rdma_bytes = 1024 * 1024 * 1024

209 + num_rdma_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024

210 num_qps_per_rank = self.num_sms // 2

211 else:

212 num_rdma_bytes = 0

@@ -230,13 +231,18 @@ class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):

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

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

233 buffer_kwargs, deep_ep.Buffer)

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

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

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

236 - # in get_or_create must be updated.

237 - handle.set_num_sms(self.num_sms)

234 return handle

235

236 + def set_num_sms(self, num_sms: int):

237 + import deep_ep

238 +

239 + # Right now the buffers are sized for only what the kernels were

240 + # created with. So we can only reduce the number of SMS used

241 + # but not increase it.

242 + if num_sms > self.num_sms:

243 + num_sms = self.num_sms

244 + deep_ep.Buffer.set_num_sms(num_sms)

245 +

246

247 class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):

248 """

@@ -265,7 +271,7 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):

271 import deep_ep

272

273 # Defaults for internode and intranode are taken from DeepEP tests.

268 - num_nvl_bytes = 1024 * 1024 * 1024

274 + num_nvl_bytes = envs.VLLM_DEEPEP_BUFFER_SIZE_MB * 1024 * 1024

275 num_qps_per_rank = num_local_experts

276 num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(

277 num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank,

@@ -291,3 +297,7 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):

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

298 buffer_kwargs, deep_ep.Buffer)

299 return handle

300 +

301 + # DeepEP LL uses RDMA so no SMs are used for communication

302 + def max_sms_used(self) -> Optional[int]:

303 + return 0

\ No newline at end of file

vllm/distributed/device_communicators/base_device_communicator.py

@@ -60,6 +60,12 @@ class All2AllManagerBase:

60 # and reuse it for the same config.

61 raise NotImplementedError

63 + def set_num_sms(self, num_sms: int):

64 + pass

65 +

66 + def max_sms_used(self) -> Optional[int]:

67 + return None # None means it could use the whole GPU

68 +

69 def dispatch(self, hidden_states: torch.Tensor,

70 router_logits: torch.Tensor):

71 raise NotImplementedError

vllm/engine/arg_utils.py

@@ -330,6 +330,8 @@ class EngineArgs:

330 enable_dbo: bool = ParallelConfig.enable_dbo

331 dbo_decode_token_threshold: int = \

332 ParallelConfig.dbo_decode_token_threshold

333 + dbo_prefill_token_threshold: int = \

334 + ParallelConfig.dbo_prefill_token_threshold

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

336 enable_eplb: bool = ParallelConfig.enable_eplb

337 expert_placement_strategy: ExpertPlacementStrategy = \

@@ -698,6 +700,9 @@ class EngineArgs:

700 parallel_group.add_argument(

701 "--dbo-decode-token-threshold",

702 **parallel_kwargs["dbo_decode_token_threshold"])

703 + parallel_group.add_argument(

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

705 + **parallel_kwargs["dbo_prefill_token_threshold"])

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

707 **parallel_kwargs["enable_eplb"])

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

@@ -1316,6 +1321,7 @@ class EngineArgs:

1321 enable_expert_parallel=self.enable_expert_parallel,

1322 enable_dbo=self.enable_dbo,

1323 dbo_decode_token_threshold=self.dbo_decode_token_threshold,

1324 + dbo_prefill_token_threshold=self.dbo_prefill_token_threshold,

1325 enable_eplb=self.enable_eplb,

1326 eplb_config=self.eplb_config,

1327 expert_placement_strategy=self.expert_placement_strategy,

vllm/envs.py

+11

@@ -189,6 +189,8 @@ if TYPE_CHECKING:

189 VLLM_CUSTOM_SCOPES_FOR_PROFILING: bool = False

190 VLLM_KV_EVENTS_USE_INT_BLOCK_HASHES: bool = True

191 VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME: str = "VLLM_OBJECT_STORAGE_SHM_BUFFER"

192 + VLLM_DEEPEP_BUFFER_SIZE_MB: int = 1024

193 + VLLM_DBO_COMM_SMS: int = 20

194 GPT_OSS_SYSTEM_TOOL_MCP_LABELS: list[str] = []

195 VLLM_PATTERN_MATCH_DEBUG: Optional[str] = None

196

@@ -1392,6 +1394,15 @@ environment_variables: dict[str, Callable[[], Any]] = {

1394 lambda: os.getenv("VLLM_OBJECT_STORAGE_SHM_BUFFER_NAME",

1395 "VLLM_OBJECT_STORAGE_SHM_BUFFER"),

1396

1397 + # The size in MB of the buffers (NVL and RDMA) used by DeepEP

1398 + "VLLM_DEEPEP_BUFFER_SIZE_MB":

1399 + lambda: int(os.getenv("VLLM_DEEPEP_BUFFER_SIZE_MB", "1024")),

1400 +

1401 + # The number of SMs to allocate for communication kernels when running DBO

1402 + # the rest of the SMs on the device will be allocated to compute

1403 + "VLLM_DBO_COMM_SMS":

1404 + lambda: int(os.getenv("VLLM_DBO_COMM_SMS", "20")),

1405 +

1406 # Valid values are container,code_interpreter,web_search_preview

1407 # ex GPT_OSS_SYSTEM_TOOL_MCP_LABELS=container,code_interpreter

1408 "GPT_OSS_SYSTEM_TOOL_MCP_LABELS":

vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py

+51 -25

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

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

13 moe_kernel_quantize_input)

14 from vllm.utils import round_up

15 +from vllm.v1.worker.ubatching import (

16 + dbo_current_ubatch_id, dbo_enabled, dbo_switch_to_comm,

17 + dbo_switch_to_compute, dbo_switch_to_compute_sync,

18 + dbo_yield_and_switch_from_comm_to_compute,

19 + dbo_yield_and_switch_from_compute_to_comm)

22 class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

@@ -46,9 +51,9 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

51 self.async_prepare = True

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

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

50 - # combine function.

51 - self.handle = None

54 + # requires. Under DBO microbatching we must track one handle per

55 + # micro-batch to avoid races between threads.

56 + self.handles = [None, None]

58 # From https://github.com/deepseek-ai/DeepEP/blob/9fe9021f29c9083cd1808ab36b740208524d9f63/deep_ep/buffer.py#L164

59 self.available_rank_configs = [2, 4, 8, 16, 24, 32, 64, 128, 144, 160]

@@ -89,6 +94,11 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

95 has_scales = token_scales is not None

97 + # We yield before launching the dispatch kernel since the dispatch

98 + # kernel will block the CPU so we want to queue up all the compute

99 + # for the other ubatch before the dispatch kernel starts.

100 + dbo_yield_and_switch_from_compute_to_comm()

101 +

102 (num_tokens_per_rank, num_tokens_per_rdma_rank,

103 dispatch_expert_num_tokens, is_token_in_rank,

104 event) = self.buffer.get_dispatch_layout(

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

114

115 (

116 token_data, expert_topk_ids, expert_topk_weights,

107 - expert_num_tokens_per_expert_list, self.handle, event

117 + expert_num_tokens_per_expert_list, handle, event

118 ) = self.buffer.dispatch(

119 x=token_data,

120 handle=None,

@@ -119,9 +129,15 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

129 expert_alignment=1,

130 config=self._get_dispatch_config(),

131 previous_event=None,

122 - async_finish=self.async_prepare,

132 + async_finish=self.async_prepare and not dbo_enabled(),

133 allocate_on_comm_stream=False)

134

135 + # record the handle for this ubatch

136 + a2a_idx = dbo_current_ubatch_id()

137 + self.handles[a2a_idx] = handle

138 +

139 + dbo_switch_to_compute_sync()

140 +

141 return lambda: self._receiver(

142 event,

143 has_scales,

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

162 a1_scale: Optional[torch.Tensor],

163 quant_config: FusedMoEQuantConfig,

164 ) -> mk.PrepareResultType:

149 - if self.async_prepare:

165 + if event.event is not None:

166 event.current_stream_wait()

167

168 if has_scales:

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

223 expert_map: Optional[torch.Tensor],

224 apply_router_weight_on_input: bool,

225 quant_config: FusedMoEQuantConfig,

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

226 + ) -> mk.ReceiverType:

227

228 if apply_router_weight_on_input:

229 topk = topk_ids.size(1)

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

249 a1q_scale = None

250 a1_post_scale = quant_config.a1_scale

251

236 - return (lambda *args: None,

237 - self._do_dispatch(tokens=a1q,

238 - token_scales=a1q_scale,

239 - rank_topk_ids=topk_ids,

240 - rank_topk_weights=topk_weights,

241 - num_experts=num_experts,

242 - a1_scale=a1_post_scale,

243 - quant_config=quant_config))

252 + return self._do_dispatch(tokens=a1q,

253 + token_scales=a1q_scale,

254 + rank_topk_ids=topk_ids,

255 + rank_topk_weights=topk_weights,

256 + num_experts=num_experts,

257 + a1_scale=a1_post_scale,

258 + quant_config=quant_config)

259

260 def prepare(

261 self,

@@ -252,10 +267,9 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

267 apply_router_weight_on_input: bool,

268 quant_config: FusedMoEQuantConfig,

269 ) -> mk.PrepareResultType:

255 - (_, receiver) = self.prepare_async(a1, topk_weights, topk_ids,

256 - num_experts, expert_map,

257 - apply_router_weight_on_input,

258 - quant_config)

270 + receiver = self.prepare_async(a1, topk_weights, topk_ids, num_experts,

271 + expert_map, apply_router_weight_on_input,

272 + quant_config)

273 return receiver()

274

275 def _finalize(

@@ -269,7 +283,9 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

283 do_async: bool,

284 ) -> Optional[Callable]:

285

272 - assert self.handle is not None

286 + a2a_idx = dbo_current_ubatch_id()

287 + handle = self.handles[a2a_idx]

288 + assert handle is not None

289

290 # fused_expert_output can have 0 tokens - This happens when none of the

291 # tokens from the all2all reach this EP rank.

@@ -283,25 +299,35 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

299 topk_ids=topk_ids,

300 apply_router_weight_on_input=apply_router_weight_on_input,

301 )

286 -

302 + dbo_yield_and_switch_from_compute_to_comm()

303 combined_x, _, event = self.buffer.combine(

304 x=fused_expert_output,

289 - handle=self.handle,

305 + handle=handle,

306 topk_weights=None,

307 config=self._get_combine_config(),

308 previous_event=None,

293 - async_finish=do_async,

309 + async_finish=do_async and not dbo_enabled(),

310 allocate_on_comm_stream=False)

311

312 + dbo_switch_to_compute()

313 +

314 if do_async:

315

316 def _receiver():

299 - event.current_stream_wait()

317 + if event.event is not None:

318 + event.current_stream_wait()

319 + dbo_switch_to_comm()

320 # Respect inplace outputs.

321 output.copy_(combined_x, non_blocking=True)

322

303 - return lambda: _receiver()

323 + # TODO(lucas): refactor the modular kernel so this will be

324 + # handled there

325 + dbo_yield_and_switch_from_comm_to_compute()

326 +

327 + return _receiver

328 else:

329 + # TODO(lucas): support this case with the refactored modular kernel

330 + assert not dbo_enabled()

331 # Respect inplace outputs.

332 output.copy_(combined_x, non_blocking=True)

333 return None

vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py

+4 -6

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

206 apply_router_weight_on_input: bool,

207 weight_and_reduce_impl: mk.TopKWeightAndReduce,

208 do_async: bool,

209 - ) -> Optional[Callable]:

209 + ) -> tuple[Callable, Callable]:

210 assert isinstance(

211 weight_and_reduce_impl, TopKWeightAndReduceDelegate

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

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

233 return_recv_hook=do_recv_hook,

234 out=output)

235

236 - return recv_hook

236 + return recv_hook, lambda: None

237

238 def finalize_async(

239 self,

@@ -243,8 +243,8 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

243 topk_ids: torch.Tensor,

244 apply_router_weight_on_input: bool,

245 weight_and_reduce_impl: mk.TopKWeightAndReduce,

246 - ) -> Callable:

247 - recv_hook = self._finalize(

246 + ) -> tuple[Callable, Callable]:

247 + return self._finalize(

248 output,

249 fused_expert_output,

250 topk_weights,

@@ -253,8 +253,6 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):

253 weight_and_reduce_impl,

254 do_async=True,

255 )

256 - assert recv_hook is not None

257 - return recv_hook

256

257 def finalize(

258 self,

vllm/model_executor/layers/fused_moe/modular_kernel.py

+95 -36

@@ -13,7 +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,

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

17 + dbo_maybe_run_recv_hook,

18 dbo_register_recv_hook, dbo_yield)

20 #

@@ -223,7 +224,7 @@ class FusedMoEPrepareAndFinalize(ABC):

224 expert_map: Optional[torch.Tensor],

225 apply_router_weight_on_input: bool,

226 quant_config: FusedMoEQuantConfig,

226 - ) -> tuple[Callable, ReceiverType]:

227 + ) -> Union[tuple[Callable, ReceiverType], ReceiverType]:

228 """

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

230 but do not wait for results from other workers.

@@ -239,10 +240,21 @@ class FusedMoEPrepareAndFinalize(ABC):

240 - apply_router_weight_on_input: When True, apply the weights to the

241 activations, before quantization + dispatching.

242

242 - Returns a callback that when invoked waits for results from other

243 - workers and has the same return signature as `prepare`, e.g.

244 -

245 - receiver = obj.prepare_async(...)

243 + Returns a callback or a hook callback pair that when invoked waits for

244 + results from other workers and has the same return signature as

245 + `prepare`, if a hook is returned this is more lightweight check that

246 + the recv is complete without doing extra work (used by DBO, will be

247 + refactored in the very near future)

248 +

249 + e.g.

250 +

251 + ret = obj.prepare_async(...)

252 +

253 + if isinstance(ret, tuple):

254 + hook, receiver = ret

255 + hook()

256 +

257 + if hook is not None:

258 a, a_scales, expert_meta, topk_ids, topk_weights = receiver()

259

260 is equivalent to:

@@ -284,7 +296,7 @@ class FusedMoEPrepareAndFinalize(ABC):

296 topk_ids: torch.Tensor,

297 apply_router_weight_on_input: bool,

298 weight_and_reduce_impl: TopKWeightAndReduce,

287 - ) -> Callable:

299 + ) -> Union[tuple[Callable, Callable], Callable]:

300 """

301 Perform any combine plus apply weights and perform a reduction on the

302 fused experts output but do not wait for results from other workers.

@@ -298,11 +310,17 @@ class FusedMoEPrepareAndFinalize(ABC):

310 - weight_and_reduce_impl: An optional TopKWeightAndReduce

311 implementation.

312

301 - Returns a callback that when invoked waits for results from other

302 - workers and has the same return signature as `finalize`, e.g.

313 + Returns a callback or a hook callback pair that when invoked waits for

314 + results from other workers and has the same return signature as

315 + `finalize`, if a hook is returned this is more lightweight check that

316 + the recv is complete without doing extra work (used by DBO, will be

317 + refactored in the very near future)

318

304 - receiver = obj.finalize_async(output, ...)

319 + ret = obj.finalize_async(output, ...)

320 ... output not valid yet ...

321 + if isinstance(ret, tuple):

322 + hook, receiver = ret

323 + hook()

324 receiver()

325 ... output valid here ...

326

@@ -600,9 +618,23 @@ class FusedMoEModularKernel(torch.nn.Module):

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

619 objects.

620 """

603 - fused_out_buffer = SharedResizableBuffer()

604 - workspace13_buffer = SharedResizableBuffer()

605 - workspace2_buffer = SharedResizableBuffer()

621 +

622 + class SharedBuffers:

623 +

624 + def __init__(self) -> None:

625 + self.fused_out = SharedResizableBuffer()

626 + self.workspace13 = SharedResizableBuffer()

627 + self.workspace2 = SharedResizableBuffer()

628 +

629 + # Persistent buffers that are shared across `FusedMoEModularKernel`

630 + # instances (layers), to save memory and allocattions.

631 + #

632 + # We have two sets of buffers to support dual batch overlap (DBO) where each

633 + # microbatch (ubatch) should use its own set of buffers to avoid

634 + # cross-ubatch contimination.

635 + # NOTE that memory is lazily allocated for these buffers, meaning that if

636 + # DBO isn't being used, the second SharedBuffers will be empty.

637 + shared_buffers: list[SharedBuffers] = [SharedBuffers(), SharedBuffers()]

638

639 def __init__(

640 self,

@@ -647,14 +679,18 @@ class FusedMoEModularKernel(torch.nn.Module):

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

680 expert_tokens_meta)

681

682 + # select per-ubatch buffers to avoid cross-ubatch reuse under DBO

683 + ubatch_idx = dbo_current_ubatch_id()

684 + buffers = self.shared_buffers[ubatch_idx]

685 +

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

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

652 - workspace13 = self.workspace13_buffer.get(workspace13_shape,

653 - device=a1.device,

654 - dtype=workspace_dtype)

655 - workspace2 = self.workspace2_buffer.get(workspace2_shape,

656 - device=a1.device,

657 - dtype=workspace_dtype)

688 + workspace13 = buffers.workspace13.get(workspace13_shape,

689 + device=a1.device,

690 + dtype=workspace_dtype)

691 + workspace2 = buffers.workspace2.get(workspace2_shape,

692 + device=a1.device,

693 + dtype=workspace_dtype)

694

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

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

@@ -733,9 +769,11 @@ class FusedMoEModularKernel(torch.nn.Module):

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

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

771 expert_tokens_meta)

736 - fused_out = self.fused_out_buffer.get(fused_out_shape,

737 - device=a1q.device,

738 - dtype=a1.dtype)

772 + ubatch_idx = dbo_current_ubatch_id()

773 + buffers = self.shared_buffers[ubatch_idx]

774 + fused_out = buffers.fused_out.get(fused_out_shape,

775 + device=a1q.device,

776 + dtype=a1.dtype)

777

778 def slice_input_tensors(

779 chunk_idx: int

@@ -868,6 +906,7 @@ class FusedMoEModularKernel(torch.nn.Module):

906 if not self.prepare_finalize.supports_async():

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

908 # support async prepare/finalize

909 + # TODO(lucas): enable in follow-up

910 assert not dbo_enabled()

911

912 (a1q, a1q_scale, expert_tokens_meta, _expert_topk_ids,

@@ -883,7 +922,7 @@ class FusedMoEModularKernel(torch.nn.Module):

922 else:

923 # Overlap shared expert compute with all2all dispatch.

924 dbo_maybe_run_recv_hook()

886 - hook, receiver = self.prepare_finalize.prepare_async(

925 + prepare_ret = self.prepare_finalize.prepare_async(

926 a1,

927 topk_weights,

928 topk_ids,

@@ -893,13 +932,21 @@ class FusedMoEModularKernel(torch.nn.Module):

932 self.fused_experts.quant_config,

933 )

934

896 - # If DBO is being used, register the hook with the ubatch context

897 - # and call it in dbo_maybe_run_recv_hook instead of passing it to

898 - # the receiver.

899 - dbo_register_recv_hook(hook)

900 - dbo_yield()

901 - if not dbo_enabled():

902 - hook()

935 + # TODO(lucas): refactor this in the alternative schedules followup

936 + # currently unpack if we have hook + receiver pair or just

937 + # receiver (see finalize_async docstring)

938 + hook, receiver = prepare_ret \

939 + if isinstance(prepare_ret, tuple) else (None, prepare_ret)

940 +

941 + if hook is not None:

942 + if dbo_enabled():

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

944 + # context and call it in dbo_maybe_run_recv_hook instead of

945 + # passing it to the receiver.

946 + dbo_register_recv_hook(hook)

947 + dbo_yield()

948 + else:

949 + hook()

950

951 (a1q, a1q_scale, expert_tokens_meta, _expert_topk_ids,

952 _expert_topk_weights) = receiver()

@@ -952,7 +999,7 @@ class FusedMoEModularKernel(torch.nn.Module):

999 if self.shared_experts is not None:

1000 shared_output = self.shared_experts(a1)

1001 else:

955 - recv_hook = self.prepare_finalize.finalize_async(

1002 + finalize_ret = self.prepare_finalize.finalize_async(

1003 output,

1004 fused_out,

1005 topk_weights,

@@ -964,11 +1011,23 @@ class FusedMoEModularKernel(torch.nn.Module):

1011 if self.shared_experts is not None:

1012 shared_output = self.shared_experts(a1)

1013

967 - assert recv_hook is not None

968 - dbo_register_recv_hook(recv_hook)

969 - dbo_yield()

970 - if not dbo_enabled():

971 - recv_hook()

1014 + # TODO(lucas): refactor this in the alternative schedules followup

1015 + # currently unpack if we have hook + receiver pair or just

1016 + # receiver (see finalize_async docstring)

1017 + hook, receiver = finalize_ret \

1018 + if isinstance(finalize_ret, tuple) else (None, finalize_ret)

1019 +

1020 + if hook is not None:

1021 + if dbo_enabled():

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

1023 + # context and call it in dbo_maybe_run_recv_hook instead of

1024 + # passing it to the receiver.

1025 + dbo_register_recv_hook(hook)

1026 + dbo_yield()

1027 + else:

1028 + hook()

1029 +

1030 + receiver()

1031

1032 if self.shared_experts is None:

1033 return output

vllm/v1/attention/backends/utils.py

+41 -3

@@ -107,19 +107,57 @@ def _make_metadata_with_slice(

107 the requests included in ubatch_slice

108 """

109

110 + assert not ubatch_slice.is_empty(), (

111 + f"Ubatch slice {ubatch_slice} is empty")

112 +

113 request_slice = ubatch_slice.request_slice

114 token_slice = ubatch_slice.token_slice

115

116 + start_locs = attn_metadata.query_start_loc_cpu

117 + first_req = request_slice.start

118 + first_tok = token_slice.start

119 + last_req = request_slice.stop - 1

120 + last_tok = token_slice.stop - 1

121 +

122 + assert start_locs[first_req] <= first_tok < start_locs[first_req + 1], \

123 + "Token slice start outside of first request"

124 + assert start_locs[last_req] <= last_tok < start_locs[last_req+1], \

125 + "Token slice end outside of last request"

126 +

127 + # If the "middle" request has tokens in both ubatches, we have to split it.

128 + # If ubatch_slice is the first ubatch then we will be splitting the last

129 + # request. If it's the second microbatch, then we will be splitting the

130 + # first request

131 + splits_first_request = first_tok > start_locs[first_req]

132 + splits_last_request = last_tok < start_locs[last_req + 1] - 1

133 +

134 + query_start_loc_cpu = slice_query_start_locs(start_locs, request_slice)

135 query_start_loc = slice_query_start_locs(attn_metadata.query_start_loc,

136 request_slice)

137 +

138 assert len(query_start_loc) >= 2, (

139 f"query_start_loc must have at least 2 elements, "

140 f"got {len(query_start_loc)}")

118 - query_start_loc_cpu = slice_query_start_locs(

119 - attn_metadata.query_start_loc_cpu, request_slice)

141

142 + if splits_first_request:

143 + tokens_skipped = first_tok - start_locs[first_req]

144 + query_start_loc[1:] -= tokens_skipped

145 + query_start_loc_cpu[1:] -= tokens_skipped

146 seq_lens = attn_metadata.seq_lens[request_slice]

147 seq_lens_cpu = attn_metadata.seq_lens_cpu[request_slice]

148 +

149 + if splits_last_request:

150 + tokens_skipped = query_start_loc_cpu[-1] - token_slice.stop

151 + query_start_loc[-1] -= tokens_skipped

152 + query_start_loc_cpu[-1] -= tokens_skipped

153 +

154 + # Make sure we don't modify the seq_lens tensors

155 + # (not cudagraph compatible)

156 + seq_lens = seq_lens.clone()

157 + seq_lens_cpu = seq_lens_cpu.clone()

158 + seq_lens[-1] -= tokens_skipped

159 + seq_lens_cpu[-1] -= tokens_skipped

160 +

161 max_seq_len = int(seq_lens_cpu.max())

162 num_computed_tokens_cpu = attn_metadata.num_computed_tokens_cpu[

163 request_slice]

@@ -167,6 +205,7 @@ def split_attn_metadata(

205 for ubatch_slice in ubatch_slices:

206 results.append(

207 _make_metadata_with_slice(ubatch_slice, common_attn_metadata))

208 +

209 return results

210

211

@@ -696,7 +735,6 @@ def split_decodes_and_prefills(

735 return num_reqs, 0, num_tokens, 0

736

737 first_prefill = is_prefill.int().argmax(dim=-1).item()

699 - assert torch.all(query_lens[first_prefill:] > decode_threshold)

738 assert torch.all(query_lens[:first_prefill] <= decode_threshold)

739 num_decodes = first_prefill

740 num_prefills = num_reqs - num_decodes

vllm/v1/spec_decode/eagle.py

+3 -6

@@ -30,7 +30,6 @@ from vllm.v1.sample.metadata import SamplingMetadata

30 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata

31 from vllm.v1.utils import CpuGpuBuffer

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

33 -from vllm.v1.worker.ubatching import dbo_current_ubatch_id

34 logger = init_logger(__name__)

@@ -192,9 +191,8 @@ class EagleProposer:

191 assert self.runner is not None

192

193 # FIXME: need to consider multiple kv_cache_groups

195 - ubatch_id = dbo_current_ubatch_id()

194 attn_metadata_builder = \

197 - self.runner.attn_groups[0][0].metadata_builders[ubatch_id]

195 + self.runner.attn_groups[0][0].get_metadata_builder()

196 attn_metadata = attn_metadata_builder.build_for_drafting(

197 common_attn_metadata=common_attn_metadata, draft_index=0)

198

@@ -330,7 +328,7 @@ class EagleProposer:

328

329 # Rebuild attention metadata

330 attn_metadata_builder = \

333 - self.runner.attn_groups[0][0].metadata_builders[ubatch_id]

331 + self.runner.attn_groups[0][0].get_metadata_builder()

332 attn_metadata = attn_metadata_builder\

333 .build_for_drafting(common_attn_metadata=common_attn_metadata,

334 draft_index=token_index + 1)

@@ -538,9 +536,8 @@ class EagleProposer:

536 hidden_states: torch.Tensor,

537 common_attn_metadata: CommonAttentionMetadata,

538 ) -> list[torch.Tensor]:

541 - ubatch_id = dbo_current_ubatch_id()

539 tree_attn_metadata_builder = \

543 - self.runner.attn_groups[0][0].metadata_builders[ubatch_id]

540 + self.runner.attn_groups[0][0].get_metadata_builder()

541 assert isinstance(tree_attn_metadata_builder,

542 TreeAttentionMetadataBuilder)

543

vllm/v1/worker/gpu_model_runner.py

+76 -97

@@ -96,7 +96,8 @@ from vllm.v1.worker.gpu_ubatch_wrapper import UBatchWrapper

96 from vllm.v1.worker.kv_connector_model_runner_mixin import (

97 KVConnectorModelRunnerMixin)

98 from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin

99 -from vllm.v1.worker.ubatch_splitting import get_dp_padding_ubatch, ubatch_split

99 +from vllm.v1.worker.ubatch_splitting import (check_ubatch_thresholds,

100 + ubatch_split)

101 from vllm.v1.worker.ubatch_utils import UBatchSlice, UBatchSlices

102 from vllm.v1.worker.utils import is_residual_scattered_for_sp

103

@@ -1032,7 +1033,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1033 num_tokens_padded = num_tokens_unpadded + self.get_local_padding(

1034 num_tokens_unpadded)

1035 ubatch_slices, num_tokens_after_padding = \

1035 - ubatch_split(max_num_scheduled_tokens,

1036 + ubatch_split(num_scheduled_tokens,

1037 num_tokens_unpadded,

1038 num_tokens_padded,

1039 self.vllm_config)

@@ -1206,7 +1207,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1207 ubatch_slices, common_attn_metadata)

1208 for ubid, common_attn_metadata in enumerate(

1209 common_attn_metadata_list):

1209 - assert common_attn_metadata.max_query_len == 1

1210 attn_metadata_i = (attn_group.get_metadata_builder(

1211 ubatch_id=ubid).build(

1212 common_prefix_len=common_prefix_len,

@@ -2182,9 +2182,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2182 ) = self._preprocess(scheduler_output, intermediate_tensors,

2183 ubatch_slices, num_tokens_after_padding)

2184

2185 - if ubatch_slices is not None:

2186 - num_input_tokens = num_input_tokens // 2

2187 -

2185 uniform_decode = (max_query_len

2186 == self.uniform_decode_query_len) and (

2187 num_scheduled_tokens

@@ -2194,6 +2191,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2191 cudagraph_runtime_mode, batch_descriptor = \

2192 self.cudagraph_dispatcher.dispatch(batch_descriptor)

2193

2194 + # This is currently to get around the assert in the DPMetadata

2195 + # where it wants `num_tokens_across_dp` to align with `num_tokens`

2196 + if ubatch_slices is not None:

2197 + num_input_tokens = ubatch_slices[0].num_tokens

2198 +

2199 # Run the model.

2200 # Use persistent buffers for CUDA graphs.

2201 with (set_forward_context(

@@ -2821,7 +2823,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2823 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

2824 force_attention: bool = False,

2825 uniform_decode: bool = False,

2824 - allow_microbatching: bool = False,

2826 + allow_microbatching: bool = True,

2827 skip_eplb: bool = False,

2828 is_profile: bool = False,

2829 create_mixed_batch: bool = False,

@@ -2847,32 +2849,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

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

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

2851 """

2850 - ubatch_enabled = self.parallel_config.enable_dbo

2851 - num_tokens_across_dp = None

2852 - num_pad = 0

2853 - should_ubatch = False

2854 - if ubatch_enabled:

2855 - should_ubatch = num_tokens >= \

2856 - self.parallel_config.dbo_decode_token_threshold and \

2857 - allow_microbatching

2858 -

2859 - (should_ubatch, num_tokens_across_dp) = get_dp_padding_ubatch(

2860 - num_tokens, num_tokens, should_ubatch, self.vllm_config)

2861 -

2862 - # Currently the dummy run should only be ubatching during

2863 - # cuda graph capture, meaning all DP ranks should already

2864 - # have the same batch size

2865 - if num_tokens_across_dp is not None:

2866 - assert int(num_tokens_across_dp[0]) == num_tokens // 2

2867 -

2852 assert cudagraph_runtime_mode in {

2853 CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL

2854 }

2855

2872 - if not should_ubatch:

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

2874 - num_tokens += num_pad

2875 -

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

2857 # cudagraph_mode.separate_routine(). This means that we are using

2858 # different graphs and/or modes for mixed prefill-decode batches vs.

@@ -2888,10 +2868,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2868 # for GQA/MQA.

2869 max_query_len = self.uniform_decode_query_len if uniform_decode else \

2870 num_tokens

2891 - if allow_microbatching:

2892 - assert self.uniform_decode_query_len == 1

2893 - assert uniform_decode is True

2894 - assert max_query_len == 1

2871

2872 # Set num_scheduled_tokens based on num_tokens and max_num_seqs

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

@@ -2930,20 +2906,31 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2906 assert len(num_scheduled_tokens_list) == num_reqs

2907 num_scheduled_tokens = np.array(num_scheduled_tokens_list,

2908 dtype=np.int32)

2909 + total_num_scheduled_tokens = int(num_scheduled_tokens.sum())

2910

2911 ubatch_slices = None

2912 + num_tokens_after_padding = None

2913 +

2914 # We currently only microbatch if the number of tokens is

2915 # over a certain threshold.

2937 - if should_ubatch:

2938 - # We only support decode-only cudagraphs

2939 - assert num_reqs == num_tokens

2940 - assert num_tokens % 2 == 0

2941 - ubatch_slices = [

2942 - UBatchSlice(slice(0, num_reqs // 2), slice(0,

2943 - num_tokens // 2)),

2944 - UBatchSlice(slice(num_reqs // 2, num_reqs),

2945 - slice(num_tokens // 2, num_tokens))

2946 - ]

2916 + if self.parallel_config.enable_dbo and allow_microbatching:

2917 + ubatch_slices, num_tokens_after_padding = ubatch_split(

2918 + num_scheduled_tokens,

2919 + total_num_scheduled_tokens,

2920 + total_num_scheduled_tokens,

2921 + self.vllm_config,

2922 + )

2923 +

2924 + # If we failed to microbatch, currently need to resynchronize

2925 + # TODO(lucas,sage): we should be able to avoid this second sync by

2926 + # refactoring `get_dp_padding_ubatch` and `get_dp_padding` into

2927 + # a single `coordinate_batch_across_dp` function.

2928 + if num_tokens_after_padding is None:

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

2930 + num_tokens_after_padding = num_tokens + num_pad

2931 + else:

2932 + num_tokens_across_dp = num_tokens_after_padding

2933 + num_tokens_after_padding = int(num_tokens_after_padding[0].item())

2934

2935 attn_metadata: Optional[PerLayerAttnMetadata] = None

2936

@@ -2966,6 +2953,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2953 self.seq_lens.np[num_reqs:] = 0

2954 self.seq_lens.copy_to_gpu()

2955

2956 + cum_num_tokens, _ = self._get_cumsum_and_arange(

2957 + num_scheduled_tokens)

2958 + self.query_start_loc.np[1:num_reqs + 1] = cum_num_tokens

2959 + self.query_start_loc.copy_to_gpu()

2960 +

2961 for kv_cache_group_id, kv_cache_group_spec in enumerate(

2962 self.kv_cache_config.kv_cache_groups):

2963 common_attn_metadata = CommonAttentionMetadata(

@@ -3060,7 +3052,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3052 with self.maybe_randomize_inputs(input_ids), set_forward_context(

3053 attn_metadata,

3054 self.vllm_config,

3063 - num_tokens=num_tokens,

3055 + num_tokens=num_tokens_after_padding,

3056 num_tokens_across_dp=num_tokens_across_dp,

3057 cudagraph_runtime_mode=cudagraph_runtime_mode,

3058 batch_descriptor=batch_descriptor,

@@ -3395,56 +3387,51 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

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

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

3389 cudagraph_runtime_mode.name))

3398 - enable_dbo = self.parallel_config.enable_dbo

3399 - # DBO Only supports running Full cudagraphs with uniform

3400 - # decode lengths

3401 - if enable_dbo and uniform_decode:

3402 - for num_tokens in compilation_cases:

3403 - # If the number of tokens is greater than the microbatching

3404 - # threshold, don't generate a microbatched cudagraph

3405 - if (num_tokens

3406 - < self.parallel_config.dbo_decode_token_threshold):

3407 - continue

3390

3409 - # Warmup

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

3392 + for num_tokens in compilation_cases:

3393 + # We currently only capture ubatched graphs when its a FULL

3394 + # cudagraph and for uniform decode batches.

3395 + capture_ubatched_graph = self.parallel_config.enable_dbo \

3396 + and cudagraph_runtime_mode == CUDAGraphMode.FULL \

3397 + and uniform_decode \

3398 + and check_ubatch_thresholds(

3399 + config=self.vllm_config.parallel_config,

3400 + num_tokens=num_tokens,

3401 + uniform_decode=uniform_decode,

3402 + )

3403 +

3404 + # Currently we capture both microbatched and non-microbatched

3405 + # graphs when capture_ubatched_graph is True, this is because

3406 + # occasionally we will be forced out of microbatching due to other

3407 + # DP ranks not microbatching (usually caused by an empty second

3408 + # microbatch; once we resolve this, we can remove the

3409 + # non-microbatched graph capture).

3410 + allow_microbatching_options = [True, False] if \

3411 + capture_ubatched_graph else [False]

3412 + for allow_microbatching in allow_microbatching_options:

3413 for _ in range(

3414 self.compilation_config.cudagraph_num_of_warmups):

3415 + # Use CUDAGraphRuntimeStyle.NONE (default) for warmup.

3416 + # But be careful, warm up with `NONE`is orthogonal to

3417 + # if we want to warm up attention or not. This is

3418 + # different from the case where `FULL` implies capture

3419 + # attention while `PIECEWISE` implies no attention.

3420 force_attention = (

3421 cudagraph_runtime_mode == CUDAGraphMode.FULL)

3422 self._dummy_run(num_tokens,

3423 cudagraph_runtime_mode=CUDAGraphMode.NONE,

3424 force_attention=force_attention,

3417 - uniform_decode=True,

3418 - allow_microbatching=True,

3419 - skip_eplb=True)

3420 -

3421 - # Graph Capture

3425 + uniform_decode=uniform_decode,

3426 + allow_microbatching=allow_microbatching,

3427 + skip_eplb=True,

3428 + remove_lora=False)

3429 self._dummy_run(num_tokens,

3423 - cudagraph_runtime_mode=CUDAGraphMode.FULL,

3424 - uniform_decode=True,

3425 - allow_microbatching=True,

3426 - skip_eplb=True)

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

3428 - for num_tokens in compilation_cases:

3429 - for _ in range(self.compilation_config.cudagraph_num_of_warmups):

3430 - # Use CUDAGraphRuntimeStyle.NONE (default) for warmup.

3431 - # But be careful, warm up with `NONE`is orthogonal to

3432 - # if we want to warm up attention or not. This is

3433 - # different from the case where `FULL` implies capture

3434 - # attention while `PIECEWISE` implies no attention.

3435 - force_attention = (

3436 - cudagraph_runtime_mode == CUDAGraphMode.FULL)

3437 - self._dummy_run(num_tokens,

3438 - cudagraph_runtime_mode=CUDAGraphMode.NONE,

3439 - force_attention=force_attention,

3430 + cudagraph_runtime_mode=cudagraph_runtime_mode,

3431 uniform_decode=uniform_decode,

3432 + allow_microbatching=allow_microbatching,

3433 skip_eplb=True,

3434 remove_lora=False)

3443 - self._dummy_run(num_tokens,

3444 - cudagraph_runtime_mode=cudagraph_runtime_mode,

3445 - uniform_decode=uniform_decode,

3446 - skip_eplb=True,

3447 - remove_lora=False)

3435 self.maybe_remove_all_loras(self.lora_config)

3436

3437 def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:

@@ -3500,24 +3487,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

3487 attn_groups: list[AttentionGroup] = []

3488 for (attn_backend,

3489 kv_cache_spec), layer_names in attn_backends_map.items():

3503 - attn_metadata_builders = []

3504 - attn_metadata_builders.append(attn_backend.get_builder_cls()(

3505 - kv_cache_spec,

3490 + attn_group = AttentionGroup.create_with_metadata_builders(

3491 + attn_backend,

3492 layer_names,

3493 + kv_cache_spec,

3494 self.vllm_config,

3495 self.device,

3509 - ))

3510 - if self.parallel_config.enable_dbo:

3511 - attn_metadata_builders.append(

3512 - attn_backend.get_builder_cls()(

3513 - kv_cache_spec,

3514 - layer_names,

3515 - self.vllm_config,

3516 - self.device,

3517 - ))

3518 - attn_group = AttentionGroup(attn_backend,

3519 - attn_metadata_builders,

3520 - layer_names, kv_cache_spec)

3496 + num_metadata_builders=1

3497 + if not self.parallel_config.enable_dbo else 2,

3498 + )

3499 +

3500 attn_groups.append(attn_group)

3501 return attn_groups

3502

vllm/v1/worker/gpu_ubatch_wrapper.py

+81 -6

@@ -1,25 +1,28 @@

1 # SPDX-License-Identifier: Apache-2.0

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

4 -import dataclasses

4 import threading

5 +from dataclasses import dataclass

6 from typing import Any, Callable, Optional

8 import torch

10 +import vllm.envs as envs

11 from vllm.compilation.cuda_graph import CUDAGraphWrapper

12 from vllm.config import CUDAGraphMode, VllmConfig

13 +from vllm.distributed import get_ep_group

14 from vllm.forward_context import (create_forward_context, get_forward_context,

15 override_forward_context)

16 from vllm.logger import init_logger

17 from vllm.platforms import current_platform

18 from vllm.sequence import IntermediateTensors

19 +from vllm.utils import has_deep_gemm

20 from vllm.v1.worker.ubatching import UBatchContext, make_ubatch_contexts

22 logger = init_logger(__name__)

22 -@dataclasses.dataclass

25 +@dataclass

26 class UbatchMetadata:

27 context: UBatchContext

28 input_ids: torch.Tensor

@@ -29,13 +32,55 @@ class UbatchMetadata:

32 num_tokens: int

32 -@dataclasses.dataclass

35 +@dataclass

36 class CUDAGraphMetaData:

37 cudagraph: torch.cuda.CUDAGraph

38 ubatch_metadata: UbatchMetadata

39 outputs: Optional[Any] = None

42 +class SMControlContextManager:

43 +

44 + def __init__(self, comm_sms: int, set_comm_sms: Callable[[int], None],

45 + set_compute_sms: Callable[[int], None]):

46 + """

47 + Context manager for controlling SM (Streaming Multiprocessor)

48 + allocation. Upon entering the context, it sets the number of SMs

49 + allocated for communication and computation to comm_sms and

50 + total_sms - comm_sms respectively. Upon exiting, it restores the

51 + allocation to use all available SMs (i.e. total_sms).

52 +

53 + Args:

54 + comm_sms (int): The number of SMs to allocate for communication.

55 + (The remainder will be used for computation.)

56 + set_comm_sms (Callable[[int], None]):

57 + A function that sets the number of SMs for communication.

58 + set_compute_sms (Callable[[int], None]):

59 + A function that sets the number of SMs for computation.

60 + """

61 +

62 + assert current_platform.is_cuda(), \

63 + "SM control is currently only supported on CUDA"

64 +

65 + props = torch.cuda.get_device_properties(torch.cuda.current_device())

66 + total_sms = props.multi_processor_count

67 +

68 + assert comm_sms < total_sms

69 + self.total_sms = total_sms

70 + self.compute_sms = total_sms - comm_sms

71 + self.comm_sms = comm_sms

72 + self.set_comm_sms = set_comm_sms

73 + self.set_compute_sms = set_compute_sms

74 +

75 + def __enter__(self):

76 + self.set_comm_sms(self.comm_sms)

77 + self.set_compute_sms(self.compute_sms)

78 +

79 + def __exit__(self, exc_type, exc_value, traceback):

80 + self.set_comm_sms(self.total_sms)

81 + self.set_compute_sms(self.total_sms)

82 +

83 +

84 class UBatchWrapper:

86 def __init__(self, runnable: Callable, vllm_config: VllmConfig,

@@ -56,6 +101,35 @@ class UBatchWrapper:

101 runnable, vllm_config, runtime_mode=runtime_mode)

102 self.graph_pool = current_platform.get_global_graph_pool()

103

104 + self.sm_control = self._create_sm_control_context(vllm_config)

105 +

106 + @staticmethod

107 + def _create_sm_control_context(vllm_config: VllmConfig):

108 + comm_sms = envs.VLLM_DBO_COMM_SMS

109 +

110 + set_comm_sms = lambda sms: None

111 + if vllm_config.parallel_config.enable_expert_parallel:

112 + # Currently only DeepEP highthroughput supports SM control so this

113 + # only affects that case.

114 + all2all_manager = get_ep_group(

115 + ).device_communicator.all2all_manager

116 +

117 + if all2all_manager.max_sms_used() is not None:

118 + comm_sms = min(comm_sms, all2all_manager.max_sms_used())

119 +

120 + if comm_sms > 0:

121 + set_comm_sms = lambda sms: all2all_manager.set_num_sms(sms)

122 +

123 + # TODO(lucas): support other kernels besides DeepGEMM

124 + set_compute_sms = lambda sms: None

125 + if has_deep_gemm() and comm_sms > 0:

126 + import deep_gemm as dg

127 + set_compute_sms = lambda sms: dg.set_num_sms(sms)

128 +

129 + return SMControlContextManager(comm_sms=comm_sms,

130 + set_comm_sms=set_comm_sms,

131 + set_compute_sms=set_compute_sms)

132 +

133 def __getattr__(self, key: str):

134 # allow accessing the attributes of the runnable.

135 if hasattr(self.runnable, key):

@@ -282,8 +356,8 @@ class UBatchWrapper:

356 dp_metadata=dp_metadata,

357 batch_descriptor=batch_descriptor,

358 cudagraph_runtime_mode=CUDAGraphMode.NONE)

285 -

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

359 + with self.sm_control:

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

361 elif num_tokens in self.cudagraphs:

362 cudagraph_metadata = self.cudagraphs[num_tokens]

363 cudagraph_metadata.cudagraph.replay()

@@ -300,4 +374,5 @@ class UBatchWrapper:

374 dp_metadata=dp_metadata,

375 batch_descriptor=batch_descriptor,

376 cudagraph_runtime_mode=CUDAGraphMode.NONE)

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

377 + with self.sm_control:

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

vllm/v1/worker/ubatch_splitting.py

+57 -21

@@ -3,9 +3,10 @@

4 from typing import Optional

6 +import numpy as np

7 import torch

8 -from vllm.config import VllmConfig

9 +from vllm.config import ParallelConfig, VllmConfig

10 from vllm.forward_context import DPMetadata

11 from vllm.logger import init_logger

12 from vllm.utils import round_up

@@ -29,6 +30,16 @@ def should_ubatch_with_num_tokens(

30 dp_size, dp_rank)

33 +def check_ubatch_thresholds(config: ParallelConfig, num_tokens: int,

34 + uniform_decode: bool) -> bool:

35 + if not config.enable_dbo:

36 + return False

37 + if uniform_decode:

38 + return num_tokens >= config.dbo_decode_token_threshold

39 + else:

40 + return num_tokens >= config.dbo_prefill_token_threshold

41 +

42 +

43 def get_dp_padding_ubatch(

44 num_tokens_unpadded: int, num_tokens_padded: int,

45 should_attempt_ubatching: bool,

@@ -95,9 +106,37 @@ def get_dp_padding_ubatch(

106 dtype=torch.int32)

107 return should_ubatch, num_tokens_after_padding

108

109 +def create_ubatch_slices(num_scheduled_tokens: np.ndarray, split_point: int) \

110 + -> UBatchSlices:

111 + # TODO(lucas): Refactor the gpu_model_runner.py so we can pass

112 + # in cu_num_tokens directly (i.e. query_start_loc)

113 + cu_num_tokens = np.zeros(len(num_scheduled_tokens) + 1, dtype=np.int32)

114 + np.cumsum(num_scheduled_tokens, dtype=np.int32, out=cu_num_tokens[1:])

115 +

116 + first_ubatch_token_slice = slice(0, split_point)

117 + second_ubatch_token_slice = slice(split_point, cu_num_tokens[-1])

118 +

119 + # Determine request slices using exclusive stop semantics

120 + # First ubatch includes requests whose tokens overlap [0, split_point)

121 + first_ubatch_req_stop = int(

122 + np.searchsorted(cu_num_tokens, split_point, side="left"))

123 + first_ubatch_req_slice = slice(0, first_ubatch_req_stop)

124 +

125 + # Second ubatch starts at the request that contains the split_point

126 + # or the request starting exactly at split_point (if on boundary)

127 + second_ubatch_req_start = int(

128 + np.searchsorted(cu_num_tokens, split_point, side="right") - 1)

129 + second_ubatch_req_slice = slice(second_ubatch_req_start,

130 + len(cu_num_tokens) - 1)

131 +

132 + return [

133 + UBatchSlice(first_ubatch_req_slice, first_ubatch_token_slice),

134 + UBatchSlice(second_ubatch_req_slice, second_ubatch_token_slice)

135 + ]

136 +

137

138 def ubatch_split(

100 - max_num_scheduled_tokens: int,

139 + num_scheduled_tokens_per_request: np.ndarray,

140 num_tokens_unpadded: int,

141 num_tokens_padded: int,

142 vllm_config: VllmConfig,

@@ -122,17 +161,20 @@ def ubatch_split(

161 return (None, None)

162

163 # 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

164 + should_attempt_ubatching = check_ubatch_thresholds(

165 + parallel_config,

166 + num_tokens_unpadded,

167 + vllm_config,

168 + )

169

170 # 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)

171 + should_ubatch, num_tokens_after_padding = get_dp_padding_ubatch(

172 + num_tokens_unpadded,

173 + num_tokens_padded,

174 + should_attempt_ubatching,

175 + vllm_config,

176 + )

177 +

178 if not should_ubatch:

179 return (None, None)

180

@@ -141,15 +183,9 @@ def ubatch_split(

183 # to the second ubatch in pad_out_ubatch_slice after attention

184 # metadata creation

185 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 - ]

186 + token_split_point = int(num_tokens_after_padding[0].item())

187 +

188 + ubatch_slices = create_ubatch_slices(num_scheduled_tokens_per_request,

189 + token_split_point)

190

191 return (ubatch_slices, num_tokens_after_padding)

vllm/v1/worker/ubatch_utils.py

@@ -10,6 +10,14 @@ class UBatchSlice:

10 request_slice: slice

11 token_slice: slice

13 + def is_empty(self) -> bool:

14 + return self.request_slice.start == self.request_slice.stop \

15 + or self.token_slice.start == self.token_slice.stop

16 +

17 + @property

18 + def num_tokens(self) -> int:

19 + return self.token_slice.stop - self.token_slice.start

20 +

22 UBatchSlices: TypeAlias = list[UBatchSlice]

vllm/v1/worker/ubatching.py

+24 -11

@@ -51,8 +51,8 @@ class UBatchContext:

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

54 + # Assume we want to start on the compute stream

55 + self.update_stream(self.compute_stream)

56 return self

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

@@ -62,17 +62,15 @@ class UBatchContext:

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)

65 return False

67 def _restore_context(self):

68 forward_context._forward_context = self.forward_context

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

70 def update_stream(self, stream):

71 self.current_stream = stream

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

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

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

75 def _signal_comm_done(self):

76 self.gpu_comm_done_event.record(self.comm_stream)

@@ -99,9 +97,20 @@ class UBatchContext:

97 self.cpu_wait_event.clear()

98 self._restore_context()

100 + def switch_to_comm(self):

101 + self.update_stream(self.comm_stream)

102 +

103 + def switch_to_compute(self):

104 + self.update_stream(self.compute_stream)

105 +

106 def switch_to_comm_sync(self):

107 self._signal_compute_done()

108 self.update_stream(self.comm_stream)

109 + self._wait_compute_done()

110 +

111 + def switch_to_compute_sync(self):

112 + self._signal_comm_done()

113 + self.update_stream(self.compute_stream)

114 self._wait_comm_done()

115

116 def maybe_run_recv_hook(self):

@@ -112,8 +121,7 @@ class UBatchContext:

121 def yield_(self):

122 self.current_stream = current_stream()

123 self._cpu_yield()

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

116 - self.update_stream(self.current_stream)

124 + self.update_stream(self.current_stream)

125

126 def yield_and_switch_from_compute_to_comm(self):

127 assert current_stream() == self.compute_stream

@@ -153,15 +161,20 @@ def _register_ubatch_function(func):

161 return wrapper

162

163

164 +dbo_maybe_run_recv_hook = _register_ubatch_function(

165 + UBatchContext.maybe_run_recv_hook)

166 +dbo_yield = _register_ubatch_function(UBatchContext.yield_)

167 dbo_yield_and_switch_from_compute_to_comm = _register_ubatch_function(

168 UBatchContext.yield_and_switch_from_compute_to_comm)

169 dbo_yield_and_switch_from_comm_to_compute = _register_ubatch_function(

170 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)

171 +dbo_switch_to_comm = _register_ubatch_function(UBatchContext.switch_to_comm)

172 +dbo_switch_to_compute = _register_ubatch_function(

173 + UBatchContext.switch_to_compute)

174 dbo_switch_to_comm_sync = _register_ubatch_function(

175 UBatchContext.switch_to_comm_sync)

176 +dbo_switch_to_compute_sync = _register_ubatch_function(

177 + UBatchContext.switch_to_compute_sync)

178

179

180 def dbo_register_recv_hook(recv_hook):

vllm/v1/worker/utils.py

+21 -4

@@ -130,15 +130,32 @@ class MultiModalBudget:

130 @dataclass

131 class AttentionGroup:

132 backend: type[AttentionBackend]

133 + # When ubatching is enabled we will have a metadata builder for each ubatch

134 + # so that if they use internal persistant buffers for cudagraphs, and they

135 + # won't have to worry about conflicting with the other ubatches.

136 metadata_builders: list[AttentionMetadataBuilder]

137 layer_names: list[str]

138 kv_cache_spec: KVCacheSpec

139

140 + @staticmethod

141 + def create_with_metadata_builders(

142 + backend: type[AttentionBackend],

143 + layer_names: list[str],

144 + kv_cache_spec: KVCacheSpec,

145 + vllm_config: VllmConfig,

146 + device: torch.device,

147 + num_metadata_builders: int = 1,

148 + ) -> 'AttentionGroup':

149 + metadata_builders = [

150 + backend.get_builder_cls()(kv_cache_spec, layer_names, vllm_config,

151 + device)

152 + for _ in range(num_metadata_builders)

153 + ]

154 + return AttentionGroup(backend, metadata_builders, layer_names,

155 + kv_cache_spec)

156 +

157 def get_metadata_builder(self,

138 - ubatch_id: Optional[int] = None

139 - ) -> AttentionMetadataBuilder:

140 - if ubatch_id is None:

141 - return self.metadata_builders[0]

158 + ubatch_id: int = 0) -> AttentionMetadataBuilder:

159 assert len(self.metadata_builders) > ubatch_id

160 return self.metadata_builders[ubatch_id]

161