diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index a83c090b7..05c71d00d 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -702,6 +702,39 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): if skip_attn_backend_init: self.mark_forward_metadata_ready() + def init_mlp_sync_metadata( + self, batch: ScheduleBatch, device: Union[str, torch.device] + ) -> None: + """Populate per-rank token counts for DP-attention MLP synchronization.""" + if batch.global_num_tokens is None: + return + + assert batch.global_num_tokens_for_logprob is not None + if self.spec_info is not None: + from sglang.srt.speculative.spec_info import spec_scale_global_num_tokens + + global_num_tokens, global_num_tokens_for_logprob = ( + spec_scale_global_num_tokens( + self.spec_info, + batch.global_num_tokens, + batch.global_num_tokens_for_logprob, + ) + ) + else: + global_num_tokens = batch.global_num_tokens + global_num_tokens_for_logprob = batch.global_num_tokens_for_logprob + + self.original_global_num_tokens_cpu = batch.global_num_tokens + self.global_num_tokens_cpu = global_num_tokens + self.global_num_tokens_gpu = torch.tensor( + global_num_tokens, dtype=torch.int64 + ).to(device, non_blocking=True) + self.global_num_tokens_for_logprob_cpu = global_num_tokens_for_logprob + self.global_num_tokens_for_logprob_gpu = torch.tensor( + global_num_tokens_for_logprob, dtype=torch.int64 + ).to(device, non_blocking=True) + self.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph + @classmethod def init_new( cls, @@ -841,37 +874,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): ) ret.num_token_non_padded_cpu = num_tokens - # For MLP sync - if batch.global_num_tokens is not None: - assert batch.global_num_tokens_for_logprob is not None - - # process global_num_tokens and global_num_tokens_for_logprob - if batch.spec_info is not None: - from sglang.srt.speculative.spec_info import ( - spec_scale_global_num_tokens, - ) - - global_num_tokens, global_num_tokens_for_logprob = ( - spec_scale_global_num_tokens( - batch.spec_info, - batch.global_num_tokens, - batch.global_num_tokens_for_logprob, - ) - ) - else: - global_num_tokens = batch.global_num_tokens - global_num_tokens_for_logprob = batch.global_num_tokens_for_logprob - - ret.original_global_num_tokens_cpu = batch.global_num_tokens - ret.global_num_tokens_cpu = global_num_tokens - ret.global_num_tokens_gpu = torch.tensor( - global_num_tokens, dtype=torch.int64 - ).to(device, non_blocking=True) - - ret.global_num_tokens_for_logprob_cpu = global_num_tokens_for_logprob - ret.global_num_tokens_for_logprob_gpu = torch.tensor( - global_num_tokens_for_logprob, dtype=torch.int64 - ).to(device, non_blocking=True) + ret.init_mlp_sync_metadata(batch, device) if ret.forward_mode.is_idle(): ret.positions = torch.empty((0,), dtype=torch.int64, device=device) @@ -1520,9 +1523,10 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): spec_info.num_accept_tokens = self._pad_tensor_to_size( spec_info.num_accept_tokens, bs ) - spec_info.hidden_states = self._pad_tensor_to_size( - spec_info.hidden_states, num_tokens - ) + if spec_info.hidden_states is not None: + spec_info.hidden_states = self._pad_tensor_to_size( + spec_info.hidden_states, num_tokens + ) def prepare_attn_tp_scatter_input(self, model_runner: ModelRunner): from sglang.srt.layers.communicator import get_attn_tp_context @@ -1569,12 +1573,15 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): ] logits_output.hidden_states = logits_output.hidden_states[:num_tokens] elif self.forward_mode.is_target_verify(): # verify - num_tokens = bs * self.spec_info.draft_token_num + num_tokens = bs * self.spec_info.num_tokens_per_req if logits_output.next_token_logits is not None: logits_output.next_token_logits = logits_output.next_token_logits[ :num_tokens ] - logits_output.hidden_states = logits_output.hidden_states[:num_tokens] + if logits_output.hidden_states is not None: + logits_output.hidden_states = logits_output.hidden_states[ + :num_tokens + ] elif self.forward_mode.is_draft_extend_v2(): # draft extend_v2 bs = bs * self.spec_info.num_tokens_per_req if logits_output.next_token_logits is not None: diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index c7a20c8ef..91e717f21 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -523,6 +523,15 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): cap_layout.verify_lens.copy_(live.verify_lens) cap_layout.qo_indptr_device.copy_(live.qo_indptr_device) + @staticmethod + def _max_dp_batch_size(forward_batch: ForwardBatch) -> int: + request_counts = forward_batch.original_global_num_tokens_cpu + if request_counts is None: + raise RuntimeError( + "DP CUDA graph replay requires raw per-rank request counts" + ) + return max(request_counts) + def can_run_graph(self, forward_batch: ForwardBatch): # Disable for token embedding overrides (dynamic per-request) if forward_batch.replace_embeds is not None: @@ -550,9 +559,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): return False if self.require_mlp_tp_gather: - # Raw sync values are per-rank request counts on decode-family - # rounds -- no width division, no per-algorithm enumeration. - cuda_graph_bs = max(forward_batch.original_global_num_tokens_cpu) + cuda_graph_bs = self._max_dp_batch_size(forward_batch) else: cuda_graph_bs = forward_batch.batch_size @@ -1103,15 +1110,8 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): else: raw_num_token = raw_bs * self.captured_req_width if self.require_mlp_tp_gather: - max_num_tokens = max(forward_batch.global_num_tokens_cpu) - max_batch_size = ( - max_num_tokens / self.captured_req_width - if self.model_runner.spec_algorithm.is_eagle() - or self.model_runner.spec_algorithm.is_standalone() - or self.model_runner.spec_algorithm.is_dflash_family() - else max_num_tokens - ) - bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs) + max_batch_size = self._max_dp_batch_size(forward_batch) + bs = self._pad_to_bucket(max_batch_size, self.capture_bs) else: bs = self._pad_to_bucket(raw_bs, self.capture_bs) padded_num_tokens = bs * self.captured_req_width diff --git a/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py b/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py index b89b8d77e..55cf84243 100644 --- a/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py +++ b/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py @@ -16,6 +16,9 @@ from unittest.mock import MagicMock import torch from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode +from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( + DecodeCudaGraphRunner, +) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -37,6 +40,70 @@ def _logits_output(num_rows: int) -> SimpleNamespace: class TestMlpSyncPadUnpad(CustomTestCase): + def test_init_mlp_sync_metadata_scales_speculative_request_width(self): + spec_info = SimpleNamespace( + num_tokens_per_req=4, + num_tokens_for_logprob_per_req=2, + ) + fb = ForwardBatch( + forward_mode=ForwardMode.TARGET_VERIFY, + batch_size=2, + input_ids=torch.arange(8), + req_pool_indices=torch.tensor([0, 1]), + seq_lens=torch.tensor([5, 6]), + out_cache_loc=torch.arange(8), + seq_lens_sum=11, + positions=torch.arange(8), + spec_info=spec_info, + ) + batch = SimpleNamespace( + global_num_tokens=[2, 0, 3], + global_num_tokens_for_logprob=[2, 0, 3], + can_run_dp_cuda_graph=True, + ) + + fb.init_mlp_sync_metadata(batch, torch.device("cpu")) + + self.assertEqual(fb.original_global_num_tokens_cpu, [2, 0, 3]) + self.assertEqual(fb.global_num_tokens_cpu, [8, 0, 12]) + self.assertEqual(fb.global_num_tokens_for_logprob_cpu, [4, 0, 6]) + torch.testing.assert_close(fb.global_num_tokens_gpu, torch.tensor([8, 0, 12])) + torch.testing.assert_close( + fb.global_num_tokens_for_logprob_gpu, torch.tensor([4, 0, 6]) + ) + self.assertTrue(fb.can_run_dp_cuda_graph) + + def test_draft_input_without_hidden_states_can_be_padded(self): + spec_info = SimpleNamespace( + is_draft_input=lambda: True, + hidden_states=None, + ) + fb = ForwardBatch( + forward_mode=ForwardMode.TARGET_VERIFY, + batch_size=1, + input_ids=torch.tensor([11]), + req_pool_indices=torch.tensor([5]), + seq_lens=torch.tensor([7]), + out_cache_loc=torch.tensor([0]), + seq_lens_sum=7, + positions=torch.tensor([6]), + seq_lens_cpu=torch.tensor([7]), + lora_ids=[None], + spec_info=spec_info, + ) + + fb._pad_inputs_to_size(_mock_model_runner(), num_tokens=2, bs=1) + + self.assertIsNone(spec_info.hidden_states) + + def test_dp_cuda_graph_batch_size_uses_raw_request_counts(self): + fb = SimpleNamespace(original_global_num_tokens_cpu=[3, 11, 7]) + self.assertEqual(DecodeCudaGraphRunner._max_dp_batch_size(fb), 11) + + fb.original_global_num_tokens_cpu = None + with self.assertRaisesRegex(RuntimeError, "raw per-rank request counts"): + DecodeCudaGraphRunner._max_dp_batch_size(fb) + def test_decode_post_forward_unpads_per_request_tensors(self): fb = ForwardBatch( forward_mode=ForwardMode.DECODE,