diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 8eaa4168f..7239ccf23 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -763,6 +763,10 @@ class Envs: # CUDA graph SGLANG_PREP_IN_CUDA_GRAPH = EnvBool(True) + # Eager forward wraps the ForwardBatch's own tensors instead of copying them + # into the CUDA graph buffer registry (no per-iter device-to-device copy). + SGLANG_EAGER_INPUT_NO_COPY = EnvBool(False) + # Distributed SGLANG_DSV4_FIX_TP_ATTN_A2A_SCATTER = EnvBool(True) SGLANG_SHARED_EXPERT_TP1 = EnvBool(False) diff --git a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py index 6c189d0ca..0224cea3d 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py @@ -462,20 +462,11 @@ class CudaGraphBufferRegistry: padded_num_tokens: int, forward_batch_template: "ForwardBatch", ) -> "ForwardBatch": - """Return a FB view backed by registry slot buffers. - - ``forward_batch_template`` provides the non-slot fields - (``forward_mode`` / ``spec_info`` / ``sampling_info`` / - ``capture_hidden_mode`` / ``dp_*`` / ``lora_ids`` / ...). Slot - fields are replaced with views into the registry buffers via - ``dataclasses.replace`` — the template itself is not mutated. - - NOTE: currently parked / unused. It is NOT a drop-in for the decode - replay path's ``build_replay_fb_view``: it returns the *padded* - out_cache_loc slot slice (vs the raw ``fb.out_cache_loc`` that path - keeps), does not recompute ``seq_lens_sum`` for the padded tail, and - does not split ``forward_mode`` vs ``actual_forward_mode``. Reconcile - those before wiring it into any replay path. + """Return a FB view (``dataclasses.replace`` of ``forward_batch_template``) + whose slot fields are buffer views and whose non-slot fields are carried + from the template. A plain copy slot whose FB field is ``None`` this iter + is carried (not exposed as a stale buffer); computed slots are always + exposed. """ import dataclasses @@ -488,6 +479,14 @@ class CudaGraphBufferRegistry: # adopted backing object, not re-attached to the FB view here. if "." in slot.name: continue + is_computed = slot.post_fill is not None or not slot.copy_from_fb + if ( + not is_computed + and slot.source_fn is None + and getattr(forward_batch_template, slot.name, None) is None + ): + # Absent this iter (fill_from skipped it): carry the template. + continue replace_kwargs[slot.name] = slot.slice_for(padded_bs, padded_num_tokens) return dataclasses.replace(forward_batch_template, **replace_kwargs) @@ -507,6 +506,7 @@ def build_decode_registry( enable_prefill_cp: bool = False, require_mlp_tp_gather: bool = False, dp_size: int = 1, + register_global_num_tokens: bool = True, share_pool: bool = True, source: Optional[Any] = None, ) -> CudaGraphBufferRegistry: @@ -643,28 +643,34 @@ def build_decode_registry( ) ) - def _global_num_tokens_post_fill(buf, fb, ctx): - # Filled with the padded token count on the gathered (DP) path; left - # untouched otherwise. Not an FB copy (copy_from_fb=False). - if require_gathered_buffer: - buf.fill_(ctx.padded_num_tokens) + # Computed slots, always exposed by extract_buffer; callers that already set + # global_num_tokens_* on the batch pass register_global_num_tokens=False. + if register_global_num_tokens: - _global_shape = ( - (lambda _bs, _mt: (dp_size,)) - if require_mlp_tp_gather - else (lambda _bs, _mt: (1,)) - ) - for _global_name in ("global_num_tokens_gpu", "global_num_tokens_for_logprob_gpu"): - slots.append( - GraphSlot( - _global_name, - _global_shape, - torch.int32, - axis="none", - copy_from_fb=False, - post_fill=_global_num_tokens_post_fill, - ) + def _global_num_tokens_post_fill(buf, fb, ctx): + # Only the gathered (DP) path writes a value; otherwise left as init. + if require_gathered_buffer: + buf.fill_(ctx.padded_num_tokens) + + _global_shape = ( + (lambda _bs, _mt: (dp_size,)) + if require_mlp_tp_gather + else (lambda _bs, _mt: (1,)) ) + for _global_name in ( + "global_num_tokens_gpu", + "global_num_tokens_for_logprob_gpu", + ): + slots.append( + GraphSlot( + _global_name, + _global_shape, + torch.int32, + axis="none", + copy_from_fb=False, + post_fill=_global_num_tokens_post_fill, + ) + ) for slot in slots: bind = None @@ -760,12 +766,17 @@ def build_prefill_registry( hidden_size: int = 0, embed_dtype: Optional[torch.dtype] = None, enable_mamba_track: bool = False, + register_input_embeds: bool = True, share_pool: bool = True, source: Optional[Any] = None, ) -> CudaGraphBufferRegistry: """Registry mirroring the **token-axis** FB-shared buffers for the piecewise / breakable (prefill) cuda-graph runners. + ``register_input_embeds`` (default ``True``) registers the multimodal + ``input_embeds`` slot; the eager extend path passes ``False`` so it is + carried from the batch (a read input) rather than written in-graph. + Padding policies match the inline copy/zero in ``PiecewiseCudaGraphRunner.replay_prepare``: ``input_ids`` / ``positions`` / ``out_cache_loc`` / ``mrope_positions`` / ``input_embeds`` reset their @@ -828,16 +839,17 @@ def build_prefill_registry( slice_fn=lambda buf, n: buf[:, :n], ) ) - slots.append( - GraphSlot( - "input_embeds", - lambda _bs2, mt: (mt, hidden_size), - embed_dtype, - axis="tokens", - padding_policy=PaddingPolicy.ZERO, - copy_from_fb=False, + if register_input_embeds: + slots.append( + GraphSlot( + "input_embeds", + lambda _bs2, mt: (mt, hidden_size), + embed_dtype, + axis="tokens", + padding_policy=PaddingPolicy.ZERO, + copy_from_fb=False, + ) ) - ) if enable_mamba_track: slots.append(GraphSlot("mamba_track_indices", _bs, torch.int64, axis="bs")) slots.append(GraphSlot("mamba_track_mask", _bs, torch.bool, axis="bs")) diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 5c9dc6429..c46405511 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -26,7 +26,7 @@ import socket import threading import time from collections import defaultdict -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from typing import Any, Callable, List, Optional, Tuple, Union @@ -140,6 +140,11 @@ from sglang.srt.model_executor.breakable_cuda_graph_runner import ( BreakableCudaGraphRunner, ) from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner +from sglang.srt.model_executor.cuda_graph_buffer_registry import ( + CudaGraphBufferRegistry, + build_decode_registry, + build_prefill_registry, +) from sglang.srt.model_executor.cuda_graph_runner import ( CudaGraphRunner, _allocate_decode_buffers, @@ -214,7 +219,7 @@ from sglang.srt.utils import ( set_cuda_arch, slow_rank_detector, ) -from sglang.srt.utils.common import ceil_align, require_mlp_sync +from sglang.srt.utils.common import ceil_align, next_power_of_2, require_mlp_sync from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto from sglang.srt.utils.nvtx_pytorch_hooks import PytHooks from sglang.srt.utils.offloader import ( @@ -338,6 +343,14 @@ class ModelRunnerOutput: indexer_topk_output: Optional[TopkCaptureOutput] = None +@dataclass +class _EagerBufferRegistry: + # Lazily-built eager input-buffer registry plus the capacity it was sized to. + registry: Optional["CudaGraphBufferRegistry"] = None + max_bs: int = 0 + max_num_tokens: int = 0 + + class ModelRunner(ModelRunnerKVCacheMixin): """ModelRunner runs the forward passes of the models.""" @@ -414,6 +427,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.enable_elastic_ep = server_args.elastic_ep_backend is not None self.forward_pass_id = 0 self.init_new_workspace = False + self._eager_decode_registry = _EagerBufferRegistry() + self._eager_prefill_registry = _EagerBufferRegistry() self.draft_model_idx = draft_model_idx self.enable_hisparse = server_args.enable_hisparse @@ -3077,11 +3092,118 @@ class ModelRunner(ModelRunnerKVCacheMixin): def update_decode_attn_backend(self, stream_idx: int): self.decode_attn_backend = self.decode_attn_backend_group[stream_idx] + def _ensure_eager_registry( + self, + cache: _EagerBufferRegistry, + raw_bs: int, + raw_num_tokens: int, + build: Callable[[int, int], "CudaGraphBufferRegistry"], + ) -> "CudaGraphBufferRegistry": + # Built on first use and grown (next power of two) when a batch exceeds + # the current capacity. + if ( + cache.registry is not None + and raw_bs <= cache.max_bs + and raw_num_tokens <= cache.max_num_tokens + ): + return cache.registry + cache.max_bs = next_power_of_2(max(raw_bs, cache.max_bs)) + cache.max_num_tokens = next_power_of_2( + max(raw_num_tokens, cache.max_num_tokens) + ) + cache.registry = build(cache.max_bs, cache.max_num_tokens) + return cache.registry + + def _ensure_eager_decode_registry( + self, raw_bs: int, raw_num_tokens: int + ) -> "CudaGraphBufferRegistry": + is_encoder_decoder = self.model_config.is_encoder_decoder + return self._ensure_eager_registry( + self._eager_decode_registry, + raw_bs, + raw_num_tokens, + lambda bs, num_tokens: build_decode_registry( + device=self.device, + max_bs=bs, + max_num_token=num_tokens, + # Eager has no padding so this sentinel is never read; 0 avoids the + # cuda-graph-only fill-value method that some backends lack. + seq_len_fill_value=0, + cache_loc_dtype=torch.int64, + enable_mamba_track=( + self.server_args.enable_mamba_extra_buffer() + and self.spec_algorithm.is_none() + ), + is_encoder_decoder=is_encoder_decoder, + encoder_len_fill_value=( + getattr(self.model_config.hf_config, "max_source_positions", 0) + if is_encoder_decoder + else 0 + ), + enable_num_token_non_padded=False, + register_global_num_tokens=False, + require_gathered_buffer=False, + require_mlp_tp_gather=False, + dp_size=self.server_args.dp_size, + share_pool=False, + source=None, + ), + ) + + def _ensure_eager_prefill_registry( + self, raw_bs: int, raw_num_tokens: int + ) -> "CudaGraphBufferRegistry": + return self._ensure_eager_registry( + self._eager_prefill_registry, + raw_bs, + raw_num_tokens, + lambda bs, num_tokens: build_prefill_registry( + device=self.device, + max_bs=bs, + max_num_token=num_tokens, + cache_loc_dtype=torch.int64, + is_multimodal=self.is_multimodal, + enable_mamba_track=False, + register_input_embeds=False, + share_pool=False, + source=None, + ), + ) + + def _eager_fb_view( + self, forward_batch: ForwardBatch, pp_proxy_tensors=None + ) -> ForwardBatch: + if envs.SGLANG_EAGER_INPUT_NO_COPY.get(): + return replace(forward_batch) + raw_bs = forward_batch.batch_size + raw_num_tokens = forward_batch.input_ids.shape[0] + ensure = ( + self._ensure_eager_prefill_registry + if forward_batch.forward_mode.is_extend(include_draft_extend_v2=True) + else self._ensure_eager_decode_registry + ) + registry = ensure(raw_bs, raw_num_tokens) + registry.fill_from( + forward_batch, + raw_bs=raw_bs, + padded_bs=raw_bs, + raw_num_tokens=raw_num_tokens, + padded_num_tokens=raw_num_tokens, + pp_proxy_tensors=pp_proxy_tensors, + ) + return registry.extract_buffer( + padded_bs=raw_bs, + padded_num_tokens=raw_num_tokens, + forward_batch_template=forward_batch, + ) + def forward_decode( self, forward_batch: ForwardBatch, pp_proxy_tensors=None, ) -> Union[LogitsProcessorOutput, PPProxyTensors]: + if not self.server_args.enable_pdmux: + forward_batch = self._eager_fb_view(forward_batch, pp_proxy_tensors) # Set extra arguments pdmux_override = False if forward_batch.needs_forward_metadata_init(): @@ -3170,6 +3292,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): ret = self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs) return (ret, can_run_graph) + if not self.server_args.enable_pdmux: + forward_batch = self._eager_fb_view(forward_batch, pp_proxy_tensors) + # Launch model forward if forward_batch.needs_forward_metadata_init(): if hasattr(self.model, "prepare_forward_batch"): @@ -3203,6 +3328,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): # called from the idle path can re-read a prior batch's req_pool # indices and trigger SWA mapping use-after-free. if forward_batch.batch_size > 0: + if not self.server_args.enable_pdmux: + forward_batch = self._eager_fb_view(forward_batch, pp_proxy_tensors) self.attn_backend.init_forward_metadata(forward_batch) else: self.attn_backend.forward_metadata = None diff --git a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py index 766891f5a..179152e2f 100644 --- a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py +++ b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py @@ -67,5 +67,74 @@ class TestTboFilterBatchMarker(CustomTestCase): self.assertFalse(child.forward_metadata_replan_equivalent) +def _make_valued_batch(bs: int) -> ForwardBatch: + # Distinct per-position values so a filtered slice is unambiguous. + return ForwardBatch( + forward_mode=ForwardMode.TARGET_VERIFY, + batch_size=bs, + input_ids=torch.arange(bs, dtype=torch.long), + positions=torch.arange(bs, dtype=torch.long), + out_cache_loc=torch.arange(bs, dtype=torch.long) + 100, + req_pool_indices=torch.arange(bs, dtype=torch.long), + seq_lens=torch.arange(1, bs + 1, dtype=torch.int32), + seq_lens_cpu=torch.arange(1, bs + 1, dtype=torch.int32), + seq_lens_sum=int(torch.arange(1, bs + 1).sum()), + spec_info=None, + ) + + +class TestTboFilterBatchOnRegistryView(CustomTestCase): + """TBO runs one forward on the parent and splits via filter_batch, so an + eager registry-backed view must filter into children identically to the raw + batch.""" + + def _registry_view(self, batch: ForwardBatch) -> ForwardBatch: + from sglang.srt.model_executor.cuda_graph_buffer_registry import ( + build_decode_registry, + ) + + bs = batch.batch_size + reg = build_decode_registry( + device=torch.device("cpu"), + max_bs=bs, + max_num_token=bs, + seq_len_fill_value=1, + cache_loc_dtype=torch.int64, + register_global_num_tokens=False, # eager decode config + share_pool=False, + source=None, + ) + reg.fill_from( + batch, raw_bs=bs, padded_bs=bs, raw_num_tokens=bs, padded_num_tokens=bs + ) + view = reg.extract_buffer( + padded_bs=bs, padded_num_tokens=bs, forward_batch_template=batch + ) + # The buffered fields are registry buffers (distinct storage), same values. + self.assertNotEqual(view.input_ids.data_ptr(), batch.input_ids.data_ptr()) + self.assertTrue(torch.equal(view.input_ids, batch.input_ids)) + return view + + def test_filter_registry_view_matches_raw_batch(self): + raw_child = _filter(_make_valued_batch(8), lo=0, hi=4) + view_child = _filter(self._registry_view(_make_valued_batch(8)), lo=0, hi=4) + + self.assertEqual(view_child.batch_size, raw_child.batch_size) + for f in ( + "input_ids", + "positions", + "out_cache_loc", + "seq_lens", + "seq_lens_cpu", + "req_pool_indices", + ): + self.assertTrue( + torch.equal(getattr(view_child, f), getattr(raw_child, f)), + f"{f} differs between raw-batch and registry-view filtering", + ) + # filter_batch still resets the plan marker on the child. + self.assertFalse(view_child.forward_metadata_ready) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py index 0b1322589..a9cce1013 100644 --- a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py +++ b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py @@ -437,6 +437,65 @@ class TestMissingAndOptionalSlots(unittest.TestCase): ) ) + def test_extract_carries_none_for_absent_plain_slot(self): + # A plain copy slot absent this iter (mrope on a non-multimodal batch) + # must be carried as None, not exposed as the stale/zero buffer. + r = _make_registry(max_bs=4, max_num_tokens=8) + r.register_slot( + GraphSlot("input_ids", lambda bs, mt: (mt,), torch.int64, axis="tokens") + ) + r.register_slot( + GraphSlot( + "mrope_positions", + lambda bs, mt: (3, mt), + torch.int64, + axis="tokens", + slice_fn=lambda buf, n: buf[:, :n], + ) + ) + fb = _MiniForwardBatch( + batch_size=2, + input_ids=torch.arange(2, dtype=torch.int64), + mrope_positions=None, # non-multimodal: FB doesn't carry it + ) + r.fill_from(fb, raw_bs=2, padded_bs=2, raw_num_tokens=2, padded_num_tokens=2) + fb_view = r.extract_buffer( + padded_bs=2, padded_num_tokens=2, forward_batch_template=fb + ) + # input_ids was present -> buffer-backed; mrope absent -> carried None. + self.assertEqual( + fb_view.input_ids.data_ptr(), r.get_slot("input_ids").buffer.data_ptr() + ) + self.assertIsNone(fb_view.mrope_positions) + + def test_extract_exposes_computed_slot_even_when_fb_field_none(self): + # A computed slot (copy_from_fb=False) is always exposed, even when its + # FB field is None — the None-skip carry applies only to plain copies. + def _fill_two(buf, fb, ctx): + buf.fill_(2) + + r = _make_registry(max_bs=4, max_num_tokens=8) + r.register_slot( + GraphSlot( + "global_num_tokens_gpu", + lambda bs, mt: (1,), + torch.int32, + axis="none", + copy_from_fb=False, + post_fill=_fill_two, + ) + ) + fb = _MiniForwardBatch(batch_size=2, global_num_tokens_gpu=None) + r.fill_from(fb, raw_bs=2, padded_bs=2, raw_num_tokens=2, padded_num_tokens=2) + fb_view = r.extract_buffer( + padded_bs=2, padded_num_tokens=2, forward_batch_template=fb + ) + self.assertIsNotNone(fb_view.global_num_tokens_gpu) + self.assertEqual( + fb_view.global_num_tokens_gpu.data_ptr(), + r.get_slot("global_num_tokens_gpu").buffer.data_ptr(), + ) + class TestPostFillHook(unittest.TestCase): def test_post_fill_runs_after_copy(self): @@ -888,6 +947,61 @@ class TestBuildDecodeRegistry(unittest.TestCase): # local = clamp(100 - rank*4, 0, 4) = 4 (NOT the raw FB copy of 100). self.assertEqual(int(src.num_token_non_padded.item()), 4) + def test_register_global_num_tokens_false_carries_fb_values(self): + # register_global_num_tokens=False (eager) excludes the computed + # global_num_tokens_* slots so the batch's DP values are carried, not + # clobbered by the zero buffer extract_buffer would otherwise expose. + from sglang.srt.model_executor.cuda_graph_buffer_registry import ( + build_decode_registry, + ) + + reg = build_decode_registry( + device=torch.device("cpu"), + max_bs=4, + max_num_token=8, + seq_len_fill_value=5, + cache_loc_dtype=torch.int64, + register_global_num_tokens=False, + share_pool=False, + source=None, + ) + self.assertFalse(reg.has_slot("global_num_tokens_gpu")) + self.assertFalse(reg.has_slot("global_num_tokens_for_logprob_gpu")) + + gnt = torch.tensor([37], dtype=torch.int32) + gntlp = torch.tensor([41], dtype=torch.int32) + fb = _MiniForwardBatch( + batch_size=2, + input_ids=torch.arange(2, dtype=torch.int64), + positions=torch.arange(2, dtype=torch.int64), + out_cache_loc=torch.arange(2, dtype=torch.int64), + req_pool_indices=torch.zeros(2, dtype=torch.int64), + seq_lens=torch.full((2,), 5, dtype=torch.int32), + seq_lens_cpu=torch.full((2,), 5, dtype=torch.int32), + global_num_tokens_gpu=gnt, + global_num_tokens_for_logprob_gpu=gntlp, + ) + reg.fill_from(fb, raw_bs=2, padded_bs=2, raw_num_tokens=2, padded_num_tokens=2) + fb_view = reg.extract_buffer( + padded_bs=2, padded_num_tokens=2, forward_batch_template=fb + ) + # Carried from the batch (same tensors), not a zero registry buffer. + self.assertIs(fb_view.global_num_tokens_gpu, gnt) + self.assertIs(fb_view.global_num_tokens_for_logprob_gpu, gntlp) + + # Default (graph path) still registers the computed slots. + reg2 = build_decode_registry( + device=torch.device("cpu"), + max_bs=4, + max_num_token=8, + seq_len_fill_value=5, + cache_loc_dtype=torch.int64, + share_pool=False, + source=None, + ) + self.assertTrue(reg2.has_slot("global_num_tokens_gpu")) + self.assertTrue(reg2.has_slot("global_num_tokens_for_logprob_gpu")) + def test_source_with_ngram_registers_structured_slots(self): from sglang.srt.model_executor.cuda_graph_buffer_registry import ( build_decode_registry, @@ -1188,6 +1302,42 @@ class TestBuildPrefillRegistry(unittest.TestCase): ) self.assertTrue(torch.all(ids[3:8] == 0)) + def test_register_input_embeds_false_keeps_mrope_carries_embeds(self): + # register_input_embeds=False (eager): mrope stays registered but + # input_embeds is carried from the FB (a read input), not a zero buffer. + from sglang.srt.model_executor.cuda_graph_buffer_registry import ( + build_prefill_registry, + ) + + reg = build_prefill_registry( + device=torch.device("cpu"), + max_bs=2, + max_num_token=8, + cache_loc_dtype=torch.int64, + is_multimodal=True, + hidden_size=4, + embed_dtype=torch.float32, + register_input_embeds=False, + share_pool=False, + source=None, + ) + self.assertTrue(reg.has_slot("mrope_positions")) + self.assertFalse(reg.has_slot("input_embeds")) + # extract_buffer carries the FB's real input_embeds (not a zero buffer). + embeds = torch.randn(3, 4) + fb = _MiniForwardBatch( + batch_size=1, + input_ids=torch.tensor([1, 2, 3], dtype=torch.int64), + positions=torch.tensor([0, 1, 2], dtype=torch.int64), + out_cache_loc=torch.tensor([7, 8, 9], dtype=torch.int64), + input_embeds=embeds, + ) + reg.fill_from(fb, raw_bs=1, padded_bs=1, raw_num_tokens=3, padded_num_tokens=3) + fb_view = reg.extract_buffer( + padded_bs=1, padded_num_tokens=3, forward_batch_template=fb + ) + self.assertIs(fb_view.input_embeds, embeds) + class TestFillOncePolicy(unittest.TestCase): """FILL_ONCE initializes the whole buffer at alloc and never resets the