[Engine] Packed prefill resolves automatically, for Huginn too

benchmarks/gsm8k.py

@@ -123,7 +123,12 @@ def select_doc_ids(population, *, limit, seed, split):
 
 
 def checkpoint_loops(config):
-    """Full recurrent depth of an Ouro checkpoint, from its config.json."""
+    """Full recurrent depth: Ouro ``total_ut_steps`` or Huginn ``mean_recurrence``."""
+    if "RavenForCausalLM" in (config.get("architectures") or []):
+        loops = config.get("mean_recurrence")
+        if type(loops) is not int or loops < 1:
+            raise ValueError("A Huginn checkpoint needs an integer mean_recurrence >= 1")
+        return loops
     loops = config.get("total_ut_steps")
     if config.get("model_type") != "ouro" or type(loops) is not int or loops < 1:
         raise ValueError("This evaluation requires an Ouro checkpoint with total_ut_steps >= 1")

benchmarks/gsm8k_backends.py

@@ -28,10 +28,11 @@ def __call__(self, input_ids, scores):
 
 
 def check_depth(config, loops):
-    if config.total_ut_steps != loops:
-        raise ValueError(
-            f"Model has {config.total_ut_steps} loops but the protocol expects {loops}"
-        )
+    depth = getattr(config, "total_ut_steps", None)
+    if depth is None:
+        depth = getattr(config, "mean_recurrence", None)  # Huginn
+    if depth != loops:
+        raise ValueError(f"Model has {depth} loops but the protocol expects {loops}")
 
 
 class ReleaseExitHooks:

tests/test_engine.py

@@ -415,3 +415,35 @@ def test_finish_reason_preserves_public_string_values(reason):
     output = RequestOutput.from_request(request)
     assert output.finished and type(output.finish_reason) is str
     assert output.finish_reason == reason.value
+
+
+def test_packed_prefill_resolves_automatically_when_unset(monkeypatch):
+    from vllm_rlt.config import ExecutionConfig
+    from vllm_rlt.engine.stage_engine import packed_prefill_eligible
+
+    common = dict(layout="last_exited", boundary_kv=False, staged_prefill=False)
+    # Unset: on for CUDA, fully recurrent, last_exited, ordinary prefill; off otherwise.
+    assert packed_prefill_eligible(None, device_type="cuda", **common)
+    assert not packed_prefill_eligible(None, device_type="cpu", **common)
+    assert not packed_prefill_eligible(None, device_type="cuda", **{**common, "layout": "shared"})
+    # Huginn packs by default (GSM8K-87 paired gate: no regression); =0 keeps per-row prefill.
+    assert packed_prefill_eligible(None, device_type="cuda", **{**common, "boundary_kv": True})
+    monkeypatch.setenv("VLLM_RLT_HUGINN_PACKED_PREFILL", "0")
+    assert not packed_prefill_eligible(None, device_type="cuda", **{**common, "boundary_kv": True})
+    monkeypatch.delenv("VLLM_RLT_HUGINN_PACKED_PREFILL")
+    assert not packed_prefill_eligible(
+        None, device_type="cuda", **{**common, "staged_prefill": True}
+    )
+    # Explicit settings win regardless of eligibility.
+    assert not packed_prefill_eligible(False, device_type="cuda", **common)
+    assert packed_prefill_eligible(True, device_type="cuda", **{**common, "boundary_kv": True})
+    assert ExecutionConfig().prefill_uva is None
+    with pytest.raises(ValueError, match="prefill_uva"):
+        ExecutionConfig(prefill_uva=1)
+    # A CPU engine resolves the unset value to the per-row path and records the decision.
+    engine = LLMEngine(tiny_model())
+    assert engine.execution_config.prefill_uva is False
+    with pytest.raises(ValueError, match="prefill_uva requires CUDA"):
+        LLMEngine(tiny_model(), execution_config=ExecutionConfig(prefill_uva=True))
+    engine = LLMEngine(tiny_model(), execution_config=ExecutionConfig(prefill_uva=False))
+    assert engine.execution_config.prefill_uva is False

tests/test_gsm8k_accuracy.py

@@ -522,6 +522,7 @@ def test_min_loops_is_validated_against_the_checkpoint_depth():
         ({"model_type": "ouro", "total_ut_steps": 4}, 4),
         ({"model_type": "ouro", "total_ut_steps": 6}, 6),
         ({"model_type": "llama", "total_ut_steps": 4}, None),
+        ({"architectures": ["RavenForCausalLM"], "mean_recurrence": 32}, 32),
         ({"model_type": "ouro"}, None),
         ({"model_type": "ouro", "total_ut_steps": 0}, None),
         ({"model_type": "ouro", "total_ut_steps": 4.0}, None),

vllm_rlt/config.py

@@ -118,11 +118,17 @@ class ExecutionConfig:
     use_uva: bool = True
     # Select packed prefill metadata staging (and require packed attention).
     # The historical name is retained; use_uva controls its transfer mechanism.
-    prefill_uva: bool = False
+    # None resolves at engine construction: on when the attention backend packs
+    # prefill (Triton opts in, FlashAttention-4 natively) with last_exited KV and
+    # ordinary (non-staged) prefill, including Huginn unless
+    # VLLM_RLT_HUGINN_PACKED_PREFILL=0; off otherwise. Measured 2.0-2.4x on
+    # 2048-token prompts (Ouro and Huginn, refill mode).
+    prefill_uva: bool | None = None
 
     def __post_init__(self):
+        if self.prefill_uva is not None and type(self.prefill_uva) is not bool:
+            raise ValueError("prefill_uva must be a boolean or None")
         for name in (
-            "prefill_uva",
             "use_uva",
             "async_scheduling",
             "multi_stream",

vllm_rlt/engine/stage_engine.py

@@ -29,6 +29,33 @@ class SignalHandle:
     row: int
 
 
+def _huginn_packed_prefill_disabled() -> bool:
+    """Huginn packs prefill like Ouro; VLLM_RLT_HUGINN_PACKED_PREFILL=0 keeps per-row prefill.
+
+    Measured on GH200 (Huginn-0125, 32 depths): 1.2-2.1x throughput and -39..-64%
+    TTFT at 512-2048-token prompts; GSM8K-87 (70 paired questions) 21 vs 20 correct
+    with 47/70 identical token sequences, i.e. no regression.
+    """
+    import os
+
+    return os.environ.get("VLLM_RLT_HUGINN_PACKED_PREFILL", "1") == "0"
+
+
+def packed_prefill_eligible(prefill_uva, *, device_type, layout, boundary_kv, staged_prefill):
+    """Whether an unset ``prefill_uva`` may resolve to packed prefill.
+
+    Explicit settings win. Packed prefill needs CUDA packed attention and
+    last_exited KV; staged prefill (wavefront chunking, mixed mode) keeps the
+    per-row path, which its schedulers were written against. Boundary-KV models
+    (Huginn) pack unless VLLM_RLT_HUGINN_PACKED_PREFILL=0.
+    """
+    if prefill_uva is not None:
+        return bool(prefill_uva)
+    if boundary_kv and _huginn_packed_prefill_disabled():
+        return False
+    return device_type == "cuda" and layout == "last_exited" and not staged_prefill
+
+
 class StageEngine:
     """Single-device loop-level engine with synchronous or pipelined scheduling."""
 
@@ -67,8 +94,12 @@ def __init__(
                 raise ValueError(
                     "Huginn does not support staged prefill required by this scheduler"
                 )
-            if self.execution_config.async_scheduling or self.execution_config.prefill_uva:
-                raise ValueError("Huginn requires synchronous scheduling without prefill UVA")
+            if self.execution_config.async_scheduling:
+                raise ValueError("Huginn requires synchronous scheduling")
+            if self.execution_config.prefill_uva is True and _huginn_packed_prefill_disabled():
+                raise ValueError(
+                    "Huginn packed prefill is disabled by VLLM_RLT_HUGINN_PACKED_PREFILL=0"
+                )
             if speculative_config is not None:
                 raise ValueError("Huginn speculative decoding is not yet supported")
         self.speculative_config = speculative_config
@@ -122,21 +153,40 @@ def __init__(
             raise ValueError("prefill_uva requires CUDA packed prefill with last_exited KV")
         if staged_prefill and cache_config.layout != "last_exited":
             raise ValueError("staged prefill requires last_exited KV planes")
-        attention = create_backend(
-            attention_backend,
-            parameter.device,
-            parameter.dtype,
-            config.head_dim,
-            cache_config.block_size,
-            required_capabilities=BackendCapabilities(
-                packed_prefill=self.execution_config.prefill_uva,
-                cuda_graphs=self.execution_config.cuda_graphs,
-                async_scheduling=(
-                    parameter.device.type == "cuda" and self.execution_config.async_scheduling
-                ),
-                device_context_lengths=self.async_speculative,
-            ),
+        auto_packed = packed_prefill_eligible(
+            self.execution_config.prefill_uva,
+            device_type=parameter.device.type,
+            layout=cache_config.layout,
+            boundary_kv=getattr(model, "requires_boundary_kv", False),
+            staged_prefill=staged_prefill,
         )
+
+        def backend(packed_prefill):
+            return create_backend(
+                attention_backend,
+                parameter.device,
+                parameter.dtype,
+                config.head_dim,
+                cache_config.block_size,
+                required_capabilities=BackendCapabilities(
+                    packed_prefill=packed_prefill,
+                    cuda_graphs=self.execution_config.cuda_graphs,
+                    async_scheduling=(
+                        parameter.device.type == "cuda" and self.execution_config.async_scheduling
+                    ),
+                    device_context_lengths=self.async_speculative,
+                ),
+            )
+
+        attention = backend(bool(self.execution_config.prefill_uva))
+        if auto_packed:
+            # Keep the backend the user would get anyway; only Triton needs the
+            # explicit opt-in to pack prefill, FlashAttention-4 packs natively.
+            if attention.info["selected_backend"] == "triton":
+                attention = backend(True)
+            auto_packed = bool(attention.capabilities.packed_prefill)
+        if self.execution_config.prefill_uva is None:
+            self.execution_config = replace(self.execution_config, prefill_uva=auto_packed)
         num_blocks, self.memory_plan = plan_cache(
             model,
             cache_config,

vllm_rlt/entrypoints/runtime_args.py

@@ -26,7 +26,15 @@ def add_runtime_args(parser):
     parser.add_argument(
         "--prefill-uva",
         action="store_true",
-        help="Use packed prefill metadata staging (requires packed attention)",
+        default=None,
+        help="Force packed prefill metadata staging (requires packed attention); "
+        "default resolves automatically (on for CUDA with last_exited KV)",
+    )
+    parser.add_argument(
+        "--no-prefill-uva",
+        action="store_false",
+        dest="prefill_uva",
+        help="Keep per-row prefill even when packed attention is available",
     )
     parser.add_argument(
         "--no-uva",