diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py index 03d09c9ca..ef0ad8b5a 100644 --- a/python/sglang/srt/layers/logits_processor.py +++ b/python/sglang/srt/layers/logits_processor.py @@ -273,8 +273,8 @@ class LogitsMetadata: mm_input_embeds: Optional[torch.Tensor] = None - # DRAFT_EXTEND_V2: when set, lm_head runs only on these rows (see - # EagleDraftExtendInput.select_index). + # DRAFT_EXTEND_V2: when set, lm_head and LAST hidden capture use only these + # rows (see EagleDraftExtendInput.select_index). draft_extend_select_index: Optional[torch.Tensor] = None @classmethod @@ -539,19 +539,37 @@ class LogitsProcessor(nn.Module): or logits_metadata.forward_mode.is_target_verify() or logits_metadata.forward_mode.is_draft_extend_v2() ): - if logits_metadata.draft_extend_select_index is not None: - # Only next_token_logits narrows to [bs, vocab]; the - # FULL-capture hidden stays unpruned. - pruned_states = hidden_states[logits_metadata.draft_extend_select_index] + draft_extend_select_index = logits_metadata.draft_extend_select_index + if draft_extend_select_index is not None: + # The draft-extend graph returns LAST hidden states alongside + # selected logits. Build selected variants for every hidden-state + # representation; FULL capture below still uses the original + # unpruned tensors. + pruned_states = hidden_states[draft_extend_select_index] + pruned_states_before_norm = ( + hidden_states_before_norm[draft_extend_select_index] + if hidden_states_before_norm is not None + else None + ) else: pruned_states = hidden_states - pruned_states_before_norm = hidden_states_before_norm + pruned_states_before_norm = hidden_states_before_norm if aux_hidden_states is not None: - aux_pruned_states = ( - aux_hidden_states - if isinstance(aux_hidden_states, torch.Tensor) - else [hidden for hidden in aux_hidden_states] - ) + if draft_extend_select_index is not None: + aux_pruned_states = ( + aux_hidden_states[draft_extend_select_index] + if isinstance(aux_hidden_states, torch.Tensor) + else [ + hidden[draft_extend_select_index] + for hidden in aux_hidden_states + ] + ) + else: + aux_pruned_states = ( + aux_hidden_states + if isinstance(aux_hidden_states, torch.Tensor) + else [hidden for hidden in aux_hidden_states] + ) sample_indices = None input_logprob_indices = None diff --git a/python/sglang/srt/model_executor/input_buffers.py b/python/sglang/srt/model_executor/input_buffers.py index 15d74385f..6d1a2d31c 100644 --- a/python/sglang/srt/model_executor/input_buffers.py +++ b/python/sglang/srt/model_executor/input_buffers.py @@ -2,7 +2,7 @@ from __future__ import annotations import dataclasses from dataclasses import dataclass, fields -from typing import Dict, Tuple +from typing import Collection, Dict, Tuple import torch @@ -68,13 +68,15 @@ class ForwardInputBuffers: if buffer is not None: buffer.zero_() - def share_buffers(self): + def share_buffers(self, *, exclude: Collection[str] = ()): # disable share input buffer on npu due to accuracy issue if is_npu(): return for f in fields(self): name = f.name + if name in exclude: + continue buffer = getattr(self, name) if buffer is None: diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 919b78aa4..0ef10e2b3 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -68,6 +68,7 @@ class EagleDraftExtendInputBuffers(ForwardInputBuffers): extend_seq_lens: torch.Tensor num_correct_drafts: torch.Tensor num_accept_tokens: torch.Tensor + select_index: torch.Tensor next_token_logits_buffer: torch.Tensor global_num_tokens_gpu: Optional[torch.Tensor] global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] @@ -190,6 +191,11 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): num_accept_tokens = torch.full( (self.max_bs,), self.captured_req_width, dtype=torch.int32 ) + select_index = ( + torch.arange(self.max_bs, dtype=torch.int64) * self.captured_req_width + + self.captured_req_width + - 1 + ) if self.require_gathered_buffer: if self.require_mlp_tp_gather: @@ -227,7 +233,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): next_token_logits_buffer = ( self.model_runner.graph_shared_output.get_logits_buffer( - vocab_size, rows=self.max_bs * self.captured_req_width + vocab_size, rows=self.max_bs ) ) @@ -258,12 +264,17 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): extend_seq_lens=extend_seq_lens, num_correct_drafts=num_correct_drafts, num_accept_tokens=num_accept_tokens, + select_index=select_index, next_token_logits_buffer=next_token_logits_buffer, global_num_tokens_gpu=global_num_tokens_gpu, global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu, dsa_seed_topk_capture=dsa_seed_topk_capture, ) - self.buffers.share_buffers() + # The values depend on captured_req_width, while adaptive speculative + # decoding owns multiple runners whose select_index buffers all have + # shape [max_bs]. Sharing by field name and shape would alias those + # width-specific indices and can make a narrower graph gather OOB. + self.buffers.share_buffers(exclude={"select_index"}) self.backend = resolve_decode_backend(self) @@ -343,17 +354,23 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): ) num_correct_drafts = buffers.num_correct_drafts[:bs] num_accept_tokens = buffers.num_accept_tokens[:bs] - next_token_logits_buffer = buffers.next_token_logits_buffer[:num_tokens] + select_index = buffers.select_index[:bs] + next_token_logits_buffer = buffers.next_token_logits_buffer[:bs] - # pruned_states = num_tokens (all tokens) - num_tokens_for_logprob = num_tokens + # The worker samples only the last accepted row from each request. + # Keep the full tree width for the draft forward, but run the lm_head + # and logits path on those selected rows only. + num_tokens_for_logprob = bs if self.require_mlp_tp_gather: global_num_tokens_cpu = [num_tokens] * self.attn_dp_size + global_num_tokens_for_logprob_cpu = [bs] * self.attn_dp_size elif self.require_attn_tp_gather: global_num_tokens_cpu = [num_tokens] + global_num_tokens_for_logprob_cpu = [bs] else: global_num_tokens_cpu = None + global_num_tokens_for_logprob_cpu = None if global_num_tokens_cpu is not None: global_dp_buffer_len = sum(global_num_tokens_cpu) @@ -380,6 +397,8 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): num_accept_tokens=num_accept_tokens, # Padded tree width per req; drives the constant qo layout. num_tokens_per_req=self.captured_req_width, + num_tokens_for_logprob_per_req=1, + select_index=select_index, ) forward_batch = ForwardBatch( @@ -398,6 +417,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): positions=positions, mrope_positions=mrope_positions, global_num_tokens_gpu=buffers.global_num_tokens_gpu, + global_num_tokens_for_logprob_cpu=global_num_tokens_for_logprob_cpu, global_num_tokens_for_logprob_gpu=buffers.global_num_tokens_for_logprob_gpu, dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(), global_dp_buffer_len=global_dp_buffer_len, @@ -469,7 +489,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): post_warmup_hook=post_warmup_hook, ) - def execute(self, forward_batch: ForwardBatch): + def execute(self, forward_batch: ForwardBatch, select_index: torch.Tensor): assert forward_batch.out_cache_loc is not None self.deepep_adapter.replay() buffers = self.buffers @@ -526,6 +546,8 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): copy_srcs.append(forward_batch.spec_info.num_correct_drafts) copy_dsts.append(buffers.num_accept_tokens[:raw_bs]) copy_srcs.append(forward_batch.spec_info.num_accept_tokens) + copy_dsts.append(buffers.select_index[:raw_bs]) + copy_srcs.append(select_index) _grouped_foreach_copy_(copy_dsts, copy_srcs) # hidden_states is large + contiguous: copy_() uses the cudaMemcpyAsync @@ -543,9 +565,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): # TODO(ch-wan): support num_token_non_padded if self.require_gathered_buffer: buffers.global_num_tokens_gpu.fill_(bs * self.captured_req_width) - buffers.global_num_tokens_for_logprob_gpu.fill_( - bs * self.captured_req_width - ) + buffers.global_num_tokens_for_logprob_gpu.fill_(bs) if forward_batch.seq_lens_cpu is not None: if bs != raw_bs: @@ -607,7 +627,9 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): out = self._replay_graph(shape_key, forward_batch) out = LogitsProcessorOutput( - next_token_logits=out.next_token_logits[:num_tokens], - hidden_states=out.hidden_states[:num_tokens], + next_token_logits=out.next_token_logits[:raw_bs], + # CUDA graph replay reuses its captured output storage. These states + # survive into the next draft step, so detach them from that buffer. + hidden_states=out.hidden_states[:raw_bs].clone(), ) return out diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index a927efc8b..e51f3f901 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -315,8 +315,8 @@ class EagleDraftExtendInput(SpecInput): # Flat per-req index of each request's last accepted window row # (i * window + front + num_correct_drafts[i]). When set, the logits - # processor runs lm_head only on these rows. None under gathered-buffer - # (DP) modes, whose logprob buffer sizing assumes all-row logits. + # processor runs lm_head and LAST hidden capture only on these rows; FULL + # hidden capture remains unpruned. select_index: Optional[torch.Tensor] = None # None for draft-extend's idle batch; attention backends fall back to diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 875c81688..48e46f8a2 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -978,7 +978,7 @@ class EagleDraftWorker(EagleDraftWorkerBase): with canary_ctx: if can_run_decode_cuda_graph: draft_logits_output = self.cuda_graph_runner_for_draft_extend.execute( - forward_batch + forward_batch, select_index ) else: draft_logits_output = self.draft_runner.forward( @@ -1008,15 +1008,16 @@ class EagleDraftWorker(EagleDraftWorkerBase): dsa_seed_topk_indices = dsa_extend_topk_capture[select_index] # Reorganize the spec info for the next batch - draft_logits_output.next_token_logits = draft_logits_output.next_token_logits[ - select_index - ] - if draft_logits_output.hidden_states is not None: - draft_logits_output.hidden_states = draft_logits_output.hidden_states[ - select_index - ] - # The draft-extend graph only anchors full logits; selected-row topk is - # owned by the worker for both graph and eager paths. + if not can_run_decode_cuda_graph: + draft_logits_output.next_token_logits = ( + draft_logits_output.next_token_logits[select_index] + ) + if draft_logits_output.hidden_states is not None: + draft_logits_output.hidden_states = draft_logits_output.hidden_states[ + select_index + ] + # Selected-row top-k remains worker-owned for both graph and eager + # paths; the graph runner only moves the row selection before lm_head. if get_spec().speculative_use_rejection_sampling: ret_draft_probs, ret_topk_p, ret_topk_index = sample_draft_proposal( draft_logits_output.next_token_logits, diff --git a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py index e0df0d4b2..d37693b0c 100644 --- a/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py +++ b/python/sglang/test/kits/attention_unittest/runner_modes/speculative_draft_extend_runner.py @@ -388,12 +388,9 @@ def run_mla_draft_extend_v2_cuda_graph_case( def _assert_draft_extend_v2_outputs_close(actual, expected, settings) -> None: - # DRAFT_EXTEND_V2 graph runner only anchors the full-row - # `next_token_logits` / `hidden_states`; the selected-row `topk_p` / - # `topk_index` are owned by EAGLEWorkerV2 and computed *after* replay (see - # `eagle_worker_v2._draft_extend_for_decode`). The graph computes no topk - # and `EAGLEDraftExtendCudaGraphRunner.replay` returns no topk fields, so - # the runner-mode reference must only compare what the graph anchors. + # DRAFT_EXTEND_V2 graph runner anchors selected-row next_token_logits and + # hidden_states. Selected-row topk_p / topk_index remain worker-owned and + # are computed after replay, so compare only what the graph anchors. torch.testing.assert_close( actual.next_token_logits, expected.next_token_logits, @@ -582,6 +579,21 @@ def _run_eagle_draft_extend_eager( return ret +def _draft_extend_select_index( + batch: ForwardBatch, settings: EagleDraftRunnerSettings +) -> torch.Tensor: + return ( + torch.arange( + batch.batch_size, + dtype=torch.int64, + device=batch.input_ids.device, + ) + * settings.speculative_num_draft_tokens + + batch.spec_info.num_accept_tokens + - 1 + ) + + def run_eagle_draft_extend_cuda_graph_runner_case( testcase, case, @@ -613,6 +625,11 @@ def run_eagle_draft_extend_cuda_graph_runner_case( settings, ) expected = _run_eagle_draft_extend_eager(eager_worker, eager_batch, settings) + select_index = _draft_extend_select_index(eager_batch, settings) + expected = LogitsProcessorOutput( + next_token_logits=expected.next_token_logits[select_index], + hidden_states=expected.hidden_states[select_index], + ) graph_fixture, graph_worker, graph_backend = _build_eagle_draft_extend_fixture( testcase, @@ -636,7 +653,7 @@ def run_eagle_draft_extend_cuda_graph_runner_case( adapter.prepare_replay_state(graph_fixture, case, draft_inputs, settings) testcase.assertTrue(graph_runner.can_run_graph(graph_batch)) - actual = graph_runner.execute(graph_batch) + actual = graph_runner.execute(graph_batch, select_index) adapter.assert_outputs_close(actual, expected, settings) finally: _reset_cuda_graph_test_buffers() @@ -663,6 +680,8 @@ class _EagleDraftExtendForward(nn.Module): def _select_logits_positions(self, forward_batch: ForwardBatch) -> torch.Tensor: if forward_batch.forward_mode.is_draft_extend_v2(): + if forward_batch.spec_info.select_index is not None: + return forward_batch.spec_info.select_index return torch.arange( forward_batch.input_ids.shape[0], dtype=torch.int64, @@ -689,11 +708,12 @@ class _EagleDraftExtendForward(nn.Module): hidden_states = hidden_states + self.token_embed(input_ids) hidden_states = self.module(hidden_states, forward_batch) - logits = self.lm_head(hidden_states).float() select_index = self._select_logits_positions(forward_batch) + hidden_states = hidden_states[select_index] + logits = self.lm_head(hidden_states).float() return LogitsProcessorOutput( - next_token_logits=logits[select_index], - hidden_states=hidden_states[select_index], + next_token_logits=logits, + hidden_states=hidden_states, ) 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 1cf9d42a7..2bee42d48 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 @@ -26,6 +26,7 @@ from sglang.srt.model_executor.cuda_graph_buffer_registry import ( GraphSlot, PaddingPolicy, ) +from sglang.srt.model_executor.input_buffers import ForwardInputBuffers from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="base-a-test-cpu") @@ -59,6 +60,12 @@ class _MiniForwardBatch: spec_info: Optional[object] = None +@dataclasses.dataclass +class _PoolInputBuffers(ForwardInputBuffers): + input_ids: torch.Tensor + select_index: torch.Tensor + + def _make_registry(max_bs: int = 8, max_num_tokens: int = 16): return CudaGraphBufferRegistry( device=torch.device("cpu"), @@ -642,6 +649,27 @@ class TestPoolBackedAlloc(unittest.TestCase): small_first, big_after = _ptrs(16, 32) self.assertNotEqual(small_first.data_ptr(), big_after.data_ptr()) + def test_forward_input_buffers_can_exclude_width_specific_fields(self): + first = _PoolInputBuffers( + input_ids=torch.zeros(4, dtype=torch.int64), + select_index=torch.tensor([1, 3], dtype=torch.int64), + ) + second = _PoolInputBuffers( + input_ids=torch.ones(4, dtype=torch.int64), + select_index=torch.tensor([3, 7], dtype=torch.int64), + ) + + first.share_buffers() + second.share_buffers(exclude={"select_index"}) + + self.assertEqual(first.input_ids.data_ptr(), second.input_ids.data_ptr()) + self.assertNotEqual( + first.select_index.data_ptr(), second.select_index.data_ptr() + ) + torch.testing.assert_close( + second.select_index, torch.tensor([3, 7], dtype=torch.int64) + ) + class TestBuildDecodeRegistry(unittest.TestCase): """``build_decode_registry`` registers the always-on FB-shared decode diff --git a/test/registered/unit/spec/test_eagle_draft_extend_logits.py b/test/registered/unit/spec/test_eagle_draft_extend_logits.py new file mode 100644 index 000000000..00193ca49 --- /dev/null +++ b/test/registered/unit/spec/test_eagle_draft_extend_logits.py @@ -0,0 +1,98 @@ +import unittest + +import torch + +from sglang.srt.layers.aux_hidden_states import pack_aux_hidden_states +from sglang.srt.layers.logits_processor import LogitsMetadata, LogitsProcessor +from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + ForwardMode, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestEagleDraftExtendLogitsPruning(unittest.TestCase): + def setUp(self): + self.hidden_states = torch.arange(12, dtype=torch.float32).reshape(6, 2) + self.select_index = torch.tensor([2, 4], dtype=torch.int64) + + def _metadata(self, capture_hidden_mode=CaptureHiddenMode.LAST): + return LogitsMetadata( + forward_mode=ForwardMode.DRAFT_EXTEND_V2, + capture_hidden_mode=capture_hidden_mode, + draft_extend_select_index=self.select_index, + ) + + def _stored_hidden_states( + self, + *, + hidden_states_before_norm=None, + aux_hidden_states=None, + capture_hidden_mode=CaptureHiddenMode.LAST, + ): + metadata = self._metadata(capture_hidden_mode) + ( + pruned_states, + pruned_states_before_norm, + aux_pruned_states, + sample_indices, + _, + _, + ) = LogitsProcessor._get_pruned_states( + None, + self.hidden_states, + hidden_states_before_norm, + aux_hidden_states, + metadata, + ) + return LogitsProcessor._get_hidden_states_to_store( + None, + self.hidden_states, + hidden_states_before_norm, + aux_hidden_states, + pruned_states, + pruned_states_before_norm, + aux_pruned_states, + sample_indices, + metadata, + ) + + def test_last_hidden_states_use_selected_rows(self): + actual = self._stored_hidden_states() + torch.testing.assert_close(actual, self.hidden_states[self.select_index]) + + def test_last_pre_norm_hidden_states_use_selected_rows(self): + hidden_states_before_norm = self.hidden_states + 100 + actual = self._stored_hidden_states( + hidden_states_before_norm=hidden_states_before_norm + ) + torch.testing.assert_close(actual, hidden_states_before_norm[self.select_index]) + + def test_last_aux_hidden_states_use_selected_rows(self): + aux_hidden_states = [self.hidden_states + 100, self.hidden_states + 200] + actual = self._stored_hidden_states(aux_hidden_states=aux_hidden_states) + expected = pack_aux_hidden_states( + [hidden[self.select_index] for hidden in aux_hidden_states] + ) + torch.testing.assert_close(actual, expected) + + def test_last_packed_aux_hidden_states_use_selected_rows(self): + aux_hidden_states = torch.cat( + [self.hidden_states + 100, self.hidden_states + 200], dim=-1 + ) + actual = self._stored_hidden_states(aux_hidden_states=aux_hidden_states) + torch.testing.assert_close(actual, aux_hidden_states[self.select_index]) + + def test_full_hidden_capture_stays_unpruned(self): + aux_hidden_states = [self.hidden_states + 100, self.hidden_states + 200] + actual = self._stored_hidden_states( + aux_hidden_states=aux_hidden_states, + capture_hidden_mode=CaptureHiddenMode.FULL, + ) + torch.testing.assert_close(actual, pack_aux_hidden_states(aux_hidden_states)) + + +if __name__ == "__main__": + unittest.main()