diff --git a/benchmarks/gsm8k.py b/benchmarks/gsm8k.py index b753a55d..f06db051 100644 --- a/benchmarks/gsm8k.py +++ b/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") diff --git a/benchmarks/gsm8k_backends.py b/benchmarks/gsm8k_backends.py index 2931f9d5..974b9cf5 100644 --- a/benchmarks/gsm8k_backends.py +++ b/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: diff --git a/tests/test_engine.py b/tests/test_engine.py index 1ad00979..66c66370 100644 --- a/tests/test_engine.py +++ b/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 diff --git a/tests/test_gsm8k_accuracy.py b/tests/test_gsm8k_accuracy.py index c7da98b0..8c90c3a6 100644 --- a/tests/test_gsm8k_accuracy.py +++ b/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), diff --git a/vllm_rlt/config.py b/vllm_rlt/config.py index 8872d018..dead3d9c 100644 --- a/vllm_rlt/config.py +++ b/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", diff --git a/vllm_rlt/engine/stage_engine.py b/vllm_rlt/engine/stage_engine.py index 40a6c9a8..6e199276 100644 --- a/vllm_rlt/engine/stage_engine.py +++ b/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, diff --git a/vllm_rlt/entrypoints/runtime_args.py b/vllm_rlt/entrypoints/runtime_args.py index a166705b..7d19882e 100644 --- a/vllm_rlt/entrypoints/runtime_args.py +++ b/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",