fix(sglang): request hidden states for every trainable drafter algorithm

tests/integration/test_sglang_adapter_contract.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
+import os
 import sys
 from types import SimpleNamespace
 
@@ -183,3 +184,113 @@ def test_eagle3_last_hidden_collection_does_not_request_raw_topk_without_logits(
     assert DFLASH_RETURN_AUX_HIDDEN_PARAM not in custom_params
     assert DRAFTER_RAW_TOP_LOGPROBS_PARAM not in custom_params
     assert sampling_params["custom_params"] == custom_params
+
+
+@pytest.mark.parametrize(
+    ("algorithm", "expected_param", "other_param"),
+    [
+        ("EAGLE1", DRAFTER_RETURN_LAST_HIDDEN_PARAM, DFLASH_RETURN_AUX_HIDDEN_PARAM),
+        ("EAGLE2", DRAFTER_RETURN_LAST_HIDDEN_PARAM, DFLASH_RETURN_AUX_HIDDEN_PARAM),
+        ("EAGLE3", DRAFTER_RETURN_LAST_HIDDEN_PARAM, DFLASH_RETURN_AUX_HIDDEN_PARAM),
+        ("PEAGLE", DRAFTER_RETURN_LAST_HIDDEN_PARAM, DFLASH_RETURN_AUX_HIDDEN_PARAM),
+        ("DFLASH", DFLASH_RETURN_AUX_HIDDEN_PARAM, DRAFTER_RETURN_LAST_HIDDEN_PARAM),
+        ("DSPARK", DFLASH_RETURN_AUX_HIDDEN_PARAM, DRAFTER_RETURN_LAST_HIDDEN_PARAM),
+        ("DOMINO", DFLASH_RETURN_AUX_HIDDEN_PARAM, DRAFTER_RETURN_LAST_HIDDEN_PARAM),
+    ],
+)
+def test_every_trainable_drafter_requests_its_hidden_state_layout(
+    algorithm, expected_param, other_param
+) -> None:
+    """Every algorithm speco_worker can train must get its hidden states.
+
+    EAGLE-1/2 and P-EAGLE hard-raise in compute_loss without last_hidden_states,
+    and Domino consumes the same aux context layers as DFlash, so none of them
+    can be left out of the request gate.
+    """
+    server = _SpecoSGLangHttpServerMixin()
+    server.replica_rank = 0
+    server._drafter_collection_step = None
+    server._drafter_collection_samples = 0
+    server._drafter_collection_tokens = 0
+    server._speco_drafter_config = {
+        "enable": True,
+        "enable_drafter_training": True,
+        "speculative_algorithm": algorithm,
+        "training": {
+            "collect_hidden_states_from_sgl": True,
+            "collect_interval_steps": 5,
+            "use_logits": False,
+            "hidden_state_window_mode": "front",
+        },
+    }
+
+    sampling_params: dict = {}
+    should_collect, _, custom_params = server._speco_request_hidden_state_params(
+        sampling_params,
+        prompt_len=8,
+        request_id=f"req-{algorithm.lower()}",
+        collection_global_steps=5,
+        max_new_tokens=16,
+    )
+
+    assert should_collect is True
+    assert custom_params[expected_param] is True
+    assert other_param not in custom_params
+
+
+@pytest.mark.parametrize(
+    ("algorithm", "expected"),
+    [
+        ("EAGLE1", "1"),
+        ("EAGLE2", "1"),
+        ("EAGLE3", "1"),
+        ("PEAGLE", "1"),
+        ("DFLASH", "0"),
+        ("DSPARK", "0"),
+        ("DOMINO", "0"),
+    ],
+)
+def test_last_hidden_env_follows_the_shared_layout_resolver(
+    monkeypatch, algorithm, expected
+) -> None:
+    from verl_speco.integration.sglang_runtime import _set_sglang_patch_envs
+
+    monkeypatch.delenv("VERL_SGLANG_DRAFTER_RETURN_LAST_HIDDEN", raising=False)
+    _set_sglang_patch_envs(
+        {
+            "enable": True,
+            "enable_drafter_training": True,
+            "speculative_algorithm": algorithm,
+            "training": {
+                "collect_hidden_states_from_sgl": True,
+                "use_logits": False,
+            },
+        }
+    )
+
+    assert os.environ["VERL_SGLANG_DRAFTER_RETURN_LAST_HIDDEN"] == expected
+
+
+def test_hidden_state_gates_stay_off_without_drafter_collection() -> None:
+    from verl_speco.integration.sglang_runtime import (
+        _drafter_uses_dflash_aux_hidden,
+        _drafter_uses_eagle_last_hidden,
+    )
+
+    for drafter_cfg in (
+        {"enable": False, "speculative_algorithm": "EAGLE3", "training": {}},
+        {
+            "enable": True,
+            "enable_drafter_training": False,
+            "speculative_algorithm": "DFLASH",
+            "training": {"collect_hidden_states_from_sgl": True},
+        },
+        {
+            "enable": True,
+            "enable_drafter_training": True,
+            "speculative_algorithm": "EAGLE3",
+            "training": {"collect_hidden_states_from_sgl": True, "use_logits": True},
+        },
+    ):
+        assert _drafter_uses_eagle_last_hidden(drafter_cfg) is False
+        assert _drafter_uses_dflash_aux_hidden(drafter_cfg) is False

verl_speco/integration/sglang_runtime.py

@@ -42,6 +42,9 @@
     sglang_needs_qwen3_rope_compat_patch,
     speco_step_matches_interval,
 )
+from verl_speco.integration.oldlogprob_layer_ids import (
+    resolve_drafter_hidden_states_layout,
+)
 
 try:
     import torch as _torch
@@ -158,30 +161,36 @@ def clear_sglang_runtime_config() -> None:
     os.environ.pop(SPECO_SGLANG_DRAFTER_CONFIG_ENV, None)
 
 
-def _drafter_uses_eagle_last_hidden(drafter_cfg: dict[str, Any]) -> bool:
-    algorithm = str(drafter_cfg.get("speculative_algorithm", "") or "").upper()
+def _drafter_collects_hidden_from_sgl(drafter_cfg: dict[str, Any]) -> bool:
+    """Whether drafter training reads hidden states off the SGLang rollout."""
     training_cfg = drafter_cfg.get("training") or {}
     return bool(
-        algorithm == "EAGLE3"
-        and drafter_cfg.get("enable")
+        drafter_cfg.get("enable")
         and drafter_cfg.get("enable_drafter_training")
         and training_cfg.get("collect_hidden_states_from_sgl")
         and not training_cfg.get("use_logits")
     )
 
 
-def _drafter_uses_dflash_aux_hidden(drafter_cfg: dict[str, Any]) -> bool:
-    algorithm = str(drafter_cfg.get("speculative_algorithm", "") or "").upper()
-    training_cfg = drafter_cfg.get("training") or {}
-    return bool(
-        algorithm in {"DFLASH", "DSPARK"}
-        and drafter_cfg.get("enable")
-        and drafter_cfg.get("enable_drafter_training")
-        and training_cfg.get("collect_hidden_states_from_sgl")
-        and not training_cfg.get("use_logits")
+def _drafter_hidden_states_layout(drafter_cfg: dict[str, Any]) -> str:
+    return resolve_drafter_hidden_states_layout(
+        drafter_cfg.get("speculative_algorithm"),
+        drafter_cfg.get("training") or {},
     )
 
 
+def _drafter_uses_eagle_last_hidden(drafter_cfg: dict[str, Any]) -> bool:
+    return _drafter_collects_hidden_from_sgl(
+        drafter_cfg
+    ) and _drafter_hidden_states_layout(drafter_cfg).startswith("eagle3_")
+
+
+def _drafter_uses_dflash_aux_hidden(drafter_cfg: dict[str, Any]) -> bool:
+    return _drafter_collects_hidden_from_sgl(
+        drafter_cfg
+    ) and _drafter_hidden_states_layout(drafter_cfg).startswith("dflash_")
+
+
 def _positive_int_or_none(value: Any) -> Optional[int]:
     if value is None:
         return None