diff --git a/specforge/offline_capture/sglang.py b/specforge/offline_capture/sglang.py index 0fd29f8d3..d76a71652 100644 --- a/specforge/offline_capture/sglang.py +++ b/specforge/offline_capture/sglang.py @@ -83,6 +83,10 @@ def capture( loss_mask=torch.cat([row[2] for row in data], dim=0), ) + def capture_rows(self, input_ids: List[List[int]]): + """Capture variable-length rows without padding target compute.""" + return self._backend.capture_rows(input_ids) + def load_offline_capture( pretrained_model_name_or_path: str, diff --git a/specforge/offline_capture/sglang_backend/capture.py b/specforge/offline_capture/sglang_backend/capture.py index 306bab3fe..b4af47584 100644 --- a/specforge/offline_capture/sglang_backend/capture.py +++ b/specforge/offline_capture/sglang_backend/capture.py @@ -10,6 +10,7 @@ from __future__ import annotations +import logging from array import array from typing import List, Optional @@ -31,6 +32,8 @@ from .model_runner import SGLangRunner from .utils import wrap_offline_eagle3_logits_processors +logger = logging.getLogger(__name__) + class OfflineSGLangCaptureBackend: """Frozen local target used only to materialize offline features.""" @@ -94,23 +97,38 @@ def set_capture_layers( ) -> None: """Set auxiliary layers through the strategy's SGLang capture API.""" - setter_name = { - "eagle3": "set_eagle3_layers_to_capture", - "dflash": "set_dflash_layers_to_capture", - "dspark": "set_dspark_layers_to_capture", + setter_names = { + "eagle3": ("set_eagle3_layers_to_capture",), + "dflash": ("set_dflash_layers_to_capture",), + # K3-class targets expose a native DSpark hook. Dense SGLang + # targets serve the same auxiliary-hidden-state layout through the + # DFlash hook. + "dspark": ( + "set_dspark_layers_to_capture", + "set_dflash_layers_to_capture", + ), }.get(capture_method) - if setter_name is None: + if setter_names is None: raise ValueError( "offline SGLang capture method must be 'eagle3', 'dflash', or " "'dspark', " f"got {capture_method!r}" ) - setter = getattr(self.model_runner.model, setter_name, None) - if not callable(setter): - raise RuntimeError( - f"target model does not expose SGLang capture hook {setter_name!r}" - ) - setter(layer_ids) + for setter_name in setter_names: + setter = getattr(self.model_runner.model, setter_name, None) + if callable(setter): + if setter_name != setter_names[0]: + logger.info( + "capture method %r resolved through compatibility hook %s", + capture_method, + setter_name, + ) + setter(layer_ids) + return + raise RuntimeError( + "target model does not expose a compatible SGLang capture hook; " + f"tried {setter_names!r}" + ) def _maybe_prepare_mlp_sync_batch(self, batch: ScheduleBatch) -> None: if require_mlp_sync(self.model_runner.server_args): @@ -164,29 +182,20 @@ def _clear_pools(self) -> None: self.model_runner.token_to_kv_pool_allocator.clear() @torch.no_grad() - def capture_eagle3( - self, - *, - input_ids: torch.Tensor, - attention_mask: torch.Tensor, - loss_mask: torch.Tensor, - ): - """Capture per-request auxiliary and final hidden states without logits.""" + def capture_rows(self, input_ids: list[list[int]]): + """Capture variable-length request rows in one packed prefill.""" + if not input_ids: + return (), () + if any(not row for row in input_ids): + raise ValueError("SGLang capture rows must contain at least one token") sampling_params = SamplingParams(temperature=0, max_new_tokens=1, top_k=1) reqs: list[Req] = [] - data = [] - input_rows = torch.split(input_ids, 1, dim=0) - attention_rows = torch.split(attention_mask, 1, dim=0) - loss_rows = torch.split(loss_mask, 1, dim=0) - - for idx, (input_row, attention_row, loss_row) in enumerate( - zip(input_rows, attention_rows, loss_rows) - ): + for idx, input_row in enumerate(input_ids): req = Req( rid=str(idx), origin_input_text="", - origin_input_ids=input_row.view(-1).tolist(), + origin_input_ids=list(input_row), sampling_params=sampling_params, ) req.full_untruncated_fill_ids = array("q", req.origin_input_ids) @@ -194,7 +203,6 @@ def capture_eagle3( req.extend_input_len = req.fill_len - len(req.prefix_indices) req.logprob_start_len = len(req.origin_input_ids) - 1 reqs.append(req) - data.append((input_row, attention_row, loss_row)) input_lens = [len(req.origin_input_ids) for req in reqs] try: @@ -203,13 +211,36 @@ def capture_eagle3( last_hidden_states = getattr(output, "last_hidden_states", None) if aux_hidden_states is None or last_hidden_states is None: raise RuntimeError( - "SGLang did not return the hidden states required for " - "offline feature preparation" + "SGLang did not return the hidden states required for capture" ) aux_rows = torch.split(aux_hidden_states, input_lens, dim=0) last_rows = torch.split(last_hidden_states, input_lens, dim=0) finally: self._clear_pools() + return aux_rows, last_rows + + @torch.no_grad() + def capture_eagle3( + self, + *, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + loss_mask: torch.Tensor, + ): + """Capture per-request auxiliary and final hidden states without logits.""" + + data = [] + input_rows = torch.split(input_ids, 1, dim=0) + attention_rows = torch.split(attention_mask, 1, dim=0) + loss_rows = torch.split(loss_mask, 1, dim=0) + for input_row, attention_row, loss_row in zip( + input_rows, attention_rows, loss_rows + ): + data.append((input_row, attention_row, loss_row)) + + aux_rows, last_rows = self.capture_rows( + [input_row.view(-1).tolist() for input_row, _, _ in data] + ) return data, aux_rows, last_rows diff --git a/tests/test_runtime/test_sglang_0514_compat.py b/tests/test_runtime/test_sglang_0514_compat.py index 92250d438..107902803 100644 --- a/tests/test_runtime/test_sglang_0514_compat.py +++ b/tests/test_runtime/test_sglang_0514_compat.py @@ -9,9 +9,39 @@ from specforge.offline_capture.sglang_backend import patch as sglang_patch from specforge.offline_capture.sglang_backend import utils as sglang_utils +from specforge.offline_capture.sglang_backend.capture import OfflineSGLangCaptureBackend class SGLang0514CompatibilityTest(unittest.TestCase): + def test_dspark_prefers_its_native_capture_hook(self): + model = mock.Mock( + spec=["set_dspark_layers_to_capture", "set_dflash_layers_to_capture"] + ) + backend = object.__new__(OfflineSGLangCaptureBackend) + backend.model_runner = types.SimpleNamespace(model=model) + + backend.set_capture_layers([1, 9, 17], capture_method="dspark") + + model.set_dspark_layers_to_capture.assert_called_once_with([1, 9, 17]) + model.set_dflash_layers_to_capture.assert_not_called() + + def test_dspark_uses_dense_dflash_capture_hook_as_compatibility_fallback(self): + model = mock.Mock(spec=["set_dflash_layers_to_capture"]) + backend = object.__new__(OfflineSGLangCaptureBackend) + backend.model_runner = types.SimpleNamespace(model=model) + + backend.set_capture_layers([1, 9, 17], capture_method="dspark") + + model.set_dflash_layers_to_capture.assert_called_once_with([1, 9, 17]) + + def test_missing_capture_hooks_fail_with_the_tried_names(self): + model = mock.Mock(spec=[]) + backend = object.__new__(OfflineSGLangCaptureBackend) + backend.model_runner = types.SimpleNamespace(model=model) + + with self.assertRaisesRegex(RuntimeError, "set_dspark_layers_to_capture"): + backend.set_capture_layers([1], capture_method="dspark") + def test_tp_and_pdmux_calls_omit_removed_keywords(self): tree = ast.parse( textwrap.dedent(inspect.getsource(sglang_patch.initialize_model_parallel))