Class StatelessForwardExecutor

rl_engine/executors/stateless_executor.py:93–209

Lightweight Reference/Reward scoring wrapper. The executor runs one full-sequence forward pass with ``use_cache=False`` whenever the wrapped model accepts that argument. It never calls a generation loop and it does not instantiate a paged-KV serving runtime.

class StatelessForwardExecutor:
    """
    Lightweight Reference/Reward scoring wrapper.

    The executor runs one full-sequence forward pass with ``use_cache=False``
    whenever the wrapped model accepts that argument. It never calls a
    generation loop and it does not instantiate a paged-KV serving runtime.
    """

    def __init__(
        self,
        model: torch.nn.Module,
        config: Optional[StatelessForwardConfig] = None,
        *,
        reward_adapter: Optional[RewardAdapter] = None,
    ):
        self.model = model
        self.config = config or StatelessForwardConfig()
        self.reward_adapter = reward_adapter or default_reward_adapter

    def score(self, inputs: StatelessForwardInputs) -> StatelessForwardResult:
        _validate_inputs(inputs, self.config)

        device = inputs.input_ids.device
        cuda_tracking = device.type == "cuda" and torch.cuda.is_available()
        if cuda_tracking:
            torch.cuda.reset_peak_memory_stats(device)
            torch.cuda.synchronize(device)

        started_at = time.perf_counter()
        with _temporarily_configure_stateless_model(self.model, self.config) as no_cache_policy:
            with torch.no_grad():
                try:
                    raw_outputs, use_cache_passed = _run_no_cache_forward(self.model, inputs)
                    no_cache_policy["attention_backend_fallback"] = False
                except Exception as exc:
                    if not _should_fallback_attention_backend(exc, self.config):
                        raise
                    fallback_config = replace(self.config, attention_backend="eager")
                    with _temporarily_configure_stateless_model(
                        self.model,
                        fallback_config,
                    ) as fallback_policy:
                        raw_outputs, use_cache_passed = _run_no_cache_forward(self.model, inputs)
                    no_cache_policy.update(fallback_policy)
                    no_cache_policy.update(
                        {
                            "attention_backend_requested": self.config.attention_backend,
                            "attention_backend_fallback": True,
                            "attention_backend_fallback_reason": (
                                f"{type(exc).__name__}: {str(exc)[:200]}"
                            ),
                        }
                    )
                kv_cache = extract_kv_cache_outputs(raw_outputs)
                kv_cache_summary = summarize_tensor_tree(kv_cache)
                if self.config.reject_kv_cache_outputs and kv_cache is not None:
                    raise ValueError(