refactor: extract packed capture_rows from the offline SGLang backend

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,

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,37 +182,27 @@ 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)
             req.fill_len = len(req.full_untruncated_fill_ids)
             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
 

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