Overview

All kernels in rl-triton share a single architectural idea: express the RL recurrence as a linear recurrence of the form

\[A_t = a_t + b_t \cdot A_{t+1}\]

and solve it in \(O(\log N)\) parallel steps using an associative scan on the GPU, rather than the \(O(N)\) sequential loop a naïve implementation requires. The scan runs entirely inside Streaming Multiprocessor SRAM, avoiding repeated round-trips to High Bandwidth Memory (HBM).

For sequences that fit in a single thread block (seq_len ≤ 131072), a fully-fused kernel computes every intermediate quantity - TD errors, IS ratios, decay products - without materialising any intermediate tensors. Longer sequences fall back to a chunked scan that stitches results across blocks.

Choosing the Right Kernel

Situation	Kernel
On-policy PPO / A2C advantage estimation	`compute_gae`
Off-policy actor-critic (IMPALA, APPO)	`compute_vtrace`
Off-policy Q-learning, discrete actions	`compute_retrace`
Critic targets with bias-variance control	`compute_lambda_returns`
Simple reward-to-go baseline	`compute_discounted_returns`
TD(λ) parameter updates with traces	`compute_eligibility_traces`
Episode-scoped cumulative statistics	`compute_episodic_prefix_sum`

API Reference

`rl_triton.ops.gae.compute_gae(rewards, values, terminateds, truncateds=None, gamma=0.99, lambda_=0.95, bootstrap_values=None, last_value=None)`

Compute Generalized Advantage Estimation via a backward associative scan.

Recurrence:

A[t] = δ[t] + β[t] * A[t+1], A[T] = 0
δ[t] = r[t] + γ·(1-terminated[t]) * v_next[t] - V(s_t)
β[t] = γ·λ·(1-done[t]) (scan decay coefficient)
done[t] = terminated[t] | truncated[t]

terminated[t] gates the one-step bootstrap inside δ[t]: set to 1 for true episode ends (s_{t+1} is a reset state with no meaningful value). truncated[t] stops trace propagation and injects the true continuation value bootstrap_values[env, t] = V(s_{t+1}^true) into δ[t].

bootstrap_values supplies the true continuation value V(s_{t+1}) wherever the stored values[t+1] is invalid. Two situations require this, under one rule: (a) truncated steps, where values[t+1] belongs to the next episode; and (b) the final column t=T-1, where values[t+1] lies past the buffer. In both cases set bootstrap_values to the true continuation value; leave it zero everywhere else.

The final column feeds delta[T-1] as the next-state value ONLY. The scan's additive boundary carry A[T] is always 0: an advantage carry represents trace mass from steps past the buffer, and there is none there. (β[T-1], the multiplicative decay coefficient, is unaffected by this -- only the additive A[T] term changes.) A single entry in bootstrap_values[:, -1] is sufficient -- set it to V(s_T) if the episode continues past the window, or 0 if it terminated at T-1.

Most users have no interior truncations and should use the last_value argument (shape [num_envs]) instead, which populates the boundary column automatically.

Parameters:

Name	Type	Description	Default
`rewards`	`Tensor`	Per-step rewards, [num_envs, seq_len], float32, CUDA.	required
`values`	`Tensor`	V(s_t), [num_envs, seq_len], float32, CUDA.	required
`terminateds`	`Tensor`	True termination flags (1.0=terminated), [num_envs, seq_len], float32, CUDA.	required
`truncateds`	`Tensor \| None`	Time-limit truncation flags (1.0=truncated), [num_envs, seq_len], float32, CUDA. If None, terminateds is used for both gating roles (conservative: treats all boundaries as terminations).	`None`
`gamma`	`float`	Discount factor (default 0.99).	`0.99`
`lambda_`	`float`	GAE trace parameter in [0, 1] (default 0.95).	`0.95`
`bootstrap_values`	`Tensor \| None`	True continuation values V(s_{t+1}^true), [num_envs, seq_len], float32, CUDA. Set bootstrap_values[env, t] = V(s_{t+1}^true) at every truncated step and at t=T-1 if the window ends mid-episode. Zero elsewhere. If None, defaults to all zeros. Mutually exclusive with last_value.	`None`
`last_value`	`Tensor \| None`	Convenience arg for the common case of no interior truncations: V(s_T) per environment, shape [num_envs], float32, CUDA. Populates bootstrap_values[:, -1] automatically. Mutually exclusive with bootstrap_values.	`None`

Returns:

Name	Type	Description
`advantages`	`Tensor`	A[t], shape [num_envs, seq_len], float32.

Source code in src/rl_triton/ops/gae.py

def compute_gae(
    rewards: torch.Tensor,
    values: torch.Tensor,
    terminateds: torch.Tensor,
    truncateds: torch.Tensor | None = None,
    gamma: float = 0.99,
    lambda_: float = 0.95,
    bootstrap_values: torch.Tensor | None = None,
    last_value: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute Generalized Advantage Estimation via a backward associative scan.

    Recurrence:

    - A[t] = δ[t] + β[t] * A[t+1],  A[T] = 0
    - δ[t] = r[t] + γ·(1-terminated[t]) * v_next[t] - V(s_t)
    - β[t] = γ·λ·(1-done[t])       (scan decay coefficient)
    - done[t] = terminated[t] | truncated[t]

    `terminated[t]` gates the one-step bootstrap inside δ[t]: set to 1 for true
    episode ends (s_{t+1} is a reset state with no meaningful value).
    `truncated[t]` stops trace propagation and injects the true continuation
    value bootstrap_values[env, t] = V(s_{t+1}^true) into δ[t].

    bootstrap_values supplies the true continuation value V(s_{t+1}) wherever
    the stored values[t+1] is invalid.  Two situations require this, under one
    rule: (a) truncated steps, where values[t+1] belongs to the next episode;
    and (b) the final column t=T-1, where values[t+1] lies past the buffer.
    In both cases set bootstrap_values to the true continuation value; leave
    it zero everywhere else.

    The final column feeds delta[T-1] as the next-state value ONLY.  The scan's
    additive boundary carry A[T] is always 0: an advantage carry represents
    trace mass from steps past the buffer, and there is none there.  (β[T-1],
    the multiplicative decay coefficient, is unaffected by this -- only the
    additive A[T] term changes.)  A single entry in bootstrap_values[:, -1] is
    sufficient -- set it to V(s_T) if the episode continues past the window,
    or 0 if it terminated at T-1.

    Most users have no interior truncations and should use the `last_value`
    argument (shape [num_envs]) instead, which populates the boundary column
    automatically.

    Args:
        rewards:          Per-step rewards, [num_envs, seq_len], float32, CUDA.
        values:           V(s_t), [num_envs, seq_len], float32, CUDA.
        terminateds:      True termination flags (1.0=terminated),
                          [num_envs, seq_len], float32, CUDA.
        truncateds:       Time-limit truncation flags (1.0=truncated),
                          [num_envs, seq_len], float32, CUDA.
                          If None, terminateds is used for both gating roles
                          (conservative: treats all boundaries as terminations).
        gamma:            Discount factor (default 0.99).
        lambda_:          GAE trace parameter in [0, 1] (default 0.95).
        bootstrap_values: True continuation values V(s_{t+1}^true),
                          [num_envs, seq_len], float32, CUDA.
                          Set bootstrap_values[env, t] = V(s_{t+1}^true) at every
                          truncated step and at t=T-1 if the window ends mid-episode.
                          Zero elsewhere.  If None, defaults to all zeros.
                          Mutually exclusive with last_value.
        last_value:       Convenience arg for the common case of no interior
                          truncations: V(s_T) per environment, shape [num_envs],
                          float32, CUDA.  Populates bootstrap_values[:, -1]
                          automatically.  Mutually exclusive with bootstrap_values.

    Returns:
        advantages: A[t], shape [num_envs, seq_len], float32.
    """
    num_envs, seq_len = rewards.shape
    has_truncations   = truncateds is not None

    # Cheap structural checks -- always-on.
    for name, t in [("rewards", rewards), ("values", values), ("terminateds", terminateds)]:
        assert t.is_cuda,                f"{name} must be on CUDA"
        assert t.dtype == torch.float32, f"{name}: expected float32, got {t.dtype}"
        assert t.shape == rewards.shape, f"{name} shape {t.shape} != rewards shape {rewards.shape}"
    if has_truncations:
        assert truncateds.is_cuda,                "truncateds must be on CUDA"
        assert truncateds.dtype == torch.float32, "truncateds: expected float32"
        assert truncateds.shape == rewards.shape, \
            f"truncateds shape {truncateds.shape} != rewards shape {rewards.shape}"
    if last_value is not None:
        assert bootstrap_values is None, \
            "pass either last_value (shape [num_envs], convenience for the " \
            "window boundary) or bootstrap_values (shape [num_envs, seq_len], " \
            "full per-step control), not both."
        assert last_value.shape == (num_envs,), \
            f"last_value must have shape [{num_envs}], got {last_value.shape}"
        assert not has_truncations, \
            "last_value cannot be combined with truncateds; use bootstrap_values instead."
    if bootstrap_values is not None:
        assert bootstrap_values.is_cuda,                "bootstrap_values must be on CUDA"
        assert bootstrap_values.dtype == torch.float32, "bootstrap_values: expected float32"
        assert bootstrap_values.shape == rewards.shape, \
            f"bootstrap_values shape {bootstrap_values.shape} != rewards shape {rewards.shape}"

    # Expensive tensor scans -- correctness-warning path only (not in benchmark hot loop).
    if _CORRECTNESS_WARNINGS():
        if has_truncations:
            assert not (terminateds.bool() & truncateds.bool()).any(), \
                "terminated and truncated are mutually exclusive: a step cannot be both"
        if has_truncations and bootstrap_values is not None:
            interior = torch.ones_like(truncateds, dtype=torch.bool)
            interior[:, -1] = False
            stray = (bootstrap_values != 0) & (truncateds == 0) & interior
            assert not stray.any(), (
                "bootstrap_values must be zero at non-truncated interior steps. "
                "Nonzero entries there double-count via the additive v_next trick "
                "and corrupt delta. Populate bootstrap_values only at truncated "
                "steps and at the final column (window boundary)."
            )

    rewards     = rewards.contiguous()
    values      = values.contiguous()
    terminateds = terminateds.contiguous()
    if has_truncations:
        truncateds = truncateds.contiguous()

    if last_value is not None:
        scalar_bootstrap = last_value.contiguous()
        bootstrap_values = None
    elif bootstrap_values is not None:
        bootstrap_values = bootstrap_values.contiguous()
        scalar_bootstrap = bootstrap_values[:, -1].contiguous()
    else:
        scalar_bootstrap = None

    out = torch.empty_like(rewards)

    if seq_len <= _FLAT_MAX_SEQ_LEN:
        BLOCK_SIZE = triton.next_power_of_2(seq_len)
        num_warps  = _WARPS.get(BLOCK_SIZE, 16)
        num_stages = 2 if BLOCK_SIZE >= 2048 else 1

        if has_truncations:
            if bootstrap_values is None:
                bootstrap_values = torch.zeros_like(rewards)
            gae_fused_kernel[(num_envs,)](
                rewards, values, terminateds, truncateds,
                out, bootstrap_values,
                seq_len, rewards.stride(0),
                gamma=gamma, lambda_=lambda_,
                BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages,
                HAS_TRUNCATIONS=True, HAS_BOOTSTRAP=True,
            )
        else:
            has_bootstrap = scalar_bootstrap is not None
            gae_fused_kernel[(num_envs,)](
                rewards, values, terminateds, None,
                out, scalar_bootstrap,
                seq_len, rewards.stride(0),
                gamma=gamma, lambda_=lambda_,
                BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages,
                HAS_TRUNCATIONS=False, HAS_BOOTSTRAP=has_bootstrap,
            )
        return out

    # Chunked path for seq_len > 131072.
    if not has_truncations:
        truncateds = torch.zeros_like(terminateds)
    if bootstrap_values is None:
        bootstrap_values = torch.zeros_like(rewards)
        if scalar_bootstrap is not None:
            bootstrap_values[:, -1] = scalar_bootstrap

    not_terminated = 1.0 - terminateds
    not_done       = 1.0 - (terminateds + truncateds).clamp(max=1.0)
    next_values    = torch.empty_like(values)
    next_values[:, :-1] = values[:, 1:]
    next_values[:, -1]  = 0.0
    next_values = next_values * (1.0 - truncateds)
    next_values[:, -1]  = 0.0
    next_values = next_values + bootstrap_values

    deltas = rewards + gamma * not_terminated * next_values - values
    decays = gamma * lambda_ * not_done   # beta[t], the scan decay coefficient
    # Additive boundary carry A[T] = 0: bootstrap_values[:, -1] already entered
    # deltas[:, -1] via next_values above (weight 1). Seeding the scan's
    # boundary carry with it too would double-count it -- see kernels/gae.py's
    # module docstring. beta (decays) itself is unaffected by this fix.
    return _run_scan(deltas, decays)

`rl_triton.ops.vtrace.compute_vtrace(log_pi_target, log_pi_behavior, values, rewards, terminateds, truncateds=None, gamma=0.99, rho_bar=1.0, c_bar=1.0, bootstrap_values=None, last_value=None)`

Compute V-Trace targets and advantages via a backward associative scan.

V-Trace (Espeholt et al. 2018, IMPALA) corrects for off-policy data using clipped importance sampling ratios ρ and c.

Recurrence:

Δ[t] = δ[t] + β[t] * Δ[t+1], Δ[T] = 0
δ[t] = ρ[t] * (r[t] + γ * V(s_{t+1}) * (1 - terminated[t]) - V(s_t))
β[t] = γ * c[t] * (1 - done[t]), done[t] = terminated[t] | truncated[t] (scan decay coefficient)
ρ[t] = min(ρ_bar, π_target[t] / π_behavior[t])
c[t] = min(c_bar, π_target[t] / π_behavior[t])

Outputs:

vs[t] = Δ[t] + V(s_t) (critic targets)
A[t] = ρ[t] * (r[t] + γ * vs[t+1] * (1 - terminated[t]) - V(s_t)) (actor advantages)

bootstrap_values supplies the true continuation value V(s_{t+1}) wherever the stored values[t+1] is invalid. Two situations require this, under one rule: (a) truncated steps, where values[t+1] belongs to the next episode; and (b) the final column t=T-1, where values[t+1] lies past the buffer. In both cases set bootstrap_values to the true continuation value; leave it zero everywhere else.

The final column feeds delta[T-1] as the next-state value ONLY. The scan's additive boundary carry Δ[T] is always 0: an advantage carry represents trace mass from steps past the buffer, and there is none there. (β[T-1], the multiplicative decay coefficient above, is unaffected by this -- only the additive Δ[T] term changes.) A single entry in bootstrap_values[:, -1] is sufficient -- set it to V(s_T) if the episode continues past the window, or 0 if it terminated at T-1. (The same bootstrap_values[:, -1] is also used directly, not as a carry, when computing next_vtrace_targets[:, -1] for the advantage formula -- that use is single-counted and correct as-is.)

Most users have no interior truncations and should use the last_value argument (shape [num_envs]) instead, which populates the boundary column automatically.

terminated[t] gates the one-step bootstrap in δ[t]: set to 1 at true episode ends. truncated[t] stops trace propagation and injects bootstrap_values[env, t] = V(s_{t+1}^true), correcting for stale reset observations in the values buffer under Gymnasium next-step autoreset.

Parameters:

Name	Type	Description	Default
`log_pi_target`	`Tensor`	Log probabilities under target policy, [num_envs, seq_len], float32, CUDA.	required
`log_pi_behavior`	`Tensor`	Log probabilities under behavior policy, [num_envs, seq_len], float32, CUDA.	required
`values`	`Tensor`	V(s_t), [num_envs, seq_len], float32, CUDA.	required
`rewards`	`Tensor`	Per-step rewards, [num_envs, seq_len], float32, CUDA.	required
`terminateds`	`Tensor`	True termination flags (1.0=terminated), [num_envs, seq_len], float32, CUDA.	required
`truncateds`	`Tensor \| None`	Time-limit truncation flags (1.0=truncated), [num_envs, seq_len], float32, CUDA. If None, terminateds is used for both gating roles.	`None`
`gamma`	`float`	Discount factor (default 0.99).	`0.99`
`rho_bar`	`float`	IS ratio clip for δ (default 1.0).	`1.0`
`c_bar`	`float`	IS ratio clip for decay (default 1.0).	`1.0`
`bootstrap_values`	`Tensor \| None`	True continuation values V(s_{t+1}^true), [num_envs, seq_len], float32, CUDA. Set bootstrap_values[env, t] = V(s_{t+1}^true) at every truncated step and at t=T-1 if the window ends mid-episode. Zero elsewhere. If None, defaults to all zeros. Mutually exclusive with last_value.	`None`
`last_value`	`Tensor \| None`	Convenience arg for the common case of no interior truncations: V(s_T) per environment, shape [num_envs], float32, CUDA. Populates bootstrap_values[:, -1] automatically. Mutually exclusive with bootstrap_values.	`None`

Returns:

Name	Type	Description
`vtrace_targets`	`Tensor`	vs[t], shape [num_envs, seq_len], float32.
`vtrace_advantages`	`Tensor`	A[t], shape [num_envs, seq_len], float32.

Source code in src/rl_triton/ops/vtrace.py

def compute_vtrace(
    log_pi_target: torch.Tensor,
    log_pi_behavior: torch.Tensor,
    values: torch.Tensor,
    rewards: torch.Tensor,
    terminateds: torch.Tensor,
    truncateds: torch.Tensor | None = None,
    gamma: float = 0.99,
    rho_bar: float = 1.0,
    c_bar: float = 1.0,
    bootstrap_values: torch.Tensor | None = None,
    last_value: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Compute V-Trace targets and advantages via a backward associative scan.

    V-Trace (Espeholt et al. 2018, IMPALA) corrects for off-policy data using
    clipped importance sampling ratios ρ and c.

    Recurrence:

    - Δ[t] = δ[t] + β[t] * Δ[t+1],   Δ[T] = 0
    - δ[t] = ρ[t] * (r[t] + γ * V(s_{t+1}) * (1 - terminated[t]) - V(s_t))
    - β[t] = γ * c[t] * (1 - done[t]),  done[t] = terminated[t] | truncated[t]
             (scan decay coefficient)
    - ρ[t] = min(ρ_bar, π_target[t] / π_behavior[t])
    - c[t] = min(c_bar, π_target[t] / π_behavior[t])

    Outputs:

    - vs[t] = Δ[t] + V(s_t)   (critic targets)
    - A[t]  = ρ[t] * (r[t] + γ * vs[t+1] * (1 - terminated[t]) - V(s_t))   (actor advantages)

    bootstrap_values supplies the true continuation value V(s_{t+1}) wherever
    the stored values[t+1] is invalid.  Two situations require this, under one
    rule: (a) truncated steps, where values[t+1] belongs to the next episode;
    and (b) the final column t=T-1, where values[t+1] lies past the buffer.
    In both cases set bootstrap_values to the true continuation value; leave
    it zero everywhere else.

    The final column feeds delta[T-1] as the next-state value ONLY.  The
    scan's additive boundary carry Δ[T] is always 0: an advantage carry
    represents trace mass from steps past the buffer, and there is none
    there.  (β[T-1], the multiplicative decay coefficient above, is
    unaffected by this -- only the additive Δ[T] term changes.)  A single
    entry in bootstrap_values[:, -1] is sufficient -- set it to V(s_T) if the
    episode continues past the window, or 0 if it terminated at T-1.  (The
    same bootstrap_values[:, -1] is also used directly, not as a carry, when
    computing next_vtrace_targets[:, -1] for the advantage formula -- that use
    is single-counted and correct as-is.)

    Most users have no interior truncations and should use the `last_value`
    argument (shape [num_envs]) instead, which populates the boundary column
    automatically.

    `terminated[t]` gates the one-step bootstrap in δ[t]: set to 1 at true episode
    ends.  `truncated[t]` stops trace propagation and injects bootstrap_values[env, t]
    = V(s_{t+1}^true), correcting for stale reset observations in the values buffer
    under Gymnasium next-step autoreset.

    Args:
        log_pi_target:    Log probabilities under target policy, [num_envs, seq_len], float32, CUDA.
        log_pi_behavior:  Log probabilities under behavior policy, [num_envs, seq_len], float32, CUDA.
        values:           V(s_t), [num_envs, seq_len], float32, CUDA.
        rewards:          Per-step rewards, [num_envs, seq_len], float32, CUDA.
        terminateds:      True termination flags (1.0=terminated),
                          [num_envs, seq_len], float32, CUDA.
        truncateds:       Time-limit truncation flags (1.0=truncated),
                          [num_envs, seq_len], float32, CUDA.
                          If None, terminateds is used for both gating roles.
        gamma:            Discount factor (default 0.99).
        rho_bar:          IS ratio clip for δ (default 1.0).
        c_bar:            IS ratio clip for decay (default 1.0).
        bootstrap_values: True continuation values V(s_{t+1}^true),
                          [num_envs, seq_len], float32, CUDA.
                          Set bootstrap_values[env, t] = V(s_{t+1}^true) at every
                          truncated step and at t=T-1 if the window ends mid-episode.
                          Zero elsewhere.  If None, defaults to all zeros.
                          Mutually exclusive with last_value.
        last_value:       Convenience arg for the common case of no interior
                          truncations: V(s_T) per environment, shape [num_envs],
                          float32, CUDA.  Populates bootstrap_values[:, -1]
                          automatically.  Mutually exclusive with bootstrap_values.

    Returns:
        vtrace_targets:    vs[t], shape [num_envs, seq_len], float32.
        vtrace_advantages: A[t],  shape [num_envs, seq_len], float32.
    """
    num_envs, seq_len = rewards.shape
    has_truncations   = truncateds is not None

    if _CORRECTNESS_WARNINGS():
        for name, t in [
            ("log_pi_target",   log_pi_target),
            ("log_pi_behavior", log_pi_behavior),
            ("values",          values),
            ("rewards",         rewards),
            ("terminateds",     terminateds),
        ]:
            assert t.is_cuda,                f"{name} must be on CUDA"
            assert t.dtype == torch.float32, f"{name}: expected float32, got {t.dtype}"
            assert t.shape == rewards.shape, f"{name} shape {t.shape} != rewards shape {rewards.shape}"
        if has_truncations:
            assert truncateds.is_cuda,                "truncateds must be on CUDA"
            assert truncateds.dtype == torch.float32, "truncateds: expected float32"
            assert truncateds.shape == rewards.shape, \
                f"truncateds shape {truncateds.shape} != rewards shape {rewards.shape}"
            assert not (terminateds.bool() & truncateds.bool()).any(), \
                "terminated and truncated are mutually exclusive: a step cannot be both"
        if last_value is not None:
            assert bootstrap_values is None, \
                "pass either last_value (shape [num_envs], convenience for the " \
                "window boundary) or bootstrap_values (shape [num_envs, seq_len], " \
                "full per-step control), not both."
            assert last_value.shape == (num_envs,), \
                f"last_value must have shape [{num_envs}], got {last_value.shape}"
        if bootstrap_values is not None:
            assert bootstrap_values.is_cuda,                "bootstrap_values must be on CUDA"
            assert bootstrap_values.dtype == torch.float32, "bootstrap_values: expected float32"
            assert bootstrap_values.shape == rewards.shape, \
                f"bootstrap_values shape {bootstrap_values.shape} != rewards shape {rewards.shape}"
        if has_truncations and bootstrap_values is not None:
            interior = torch.ones_like(truncateds, dtype=torch.bool)
            interior[:, -1] = False
            stray = (bootstrap_values != 0) & (truncateds == 0) & interior
            assert not stray.any(), (
                "bootstrap_values must be zero at non-truncated interior steps. "
                "Nonzero entries there double-count via the additive v_next trick "
                "and corrupt delta. Populate bootstrap_values only at truncated "
                "steps and at the final column (window boundary)."
            )

    log_pi_target   = log_pi_target.contiguous()
    log_pi_behavior = log_pi_behavior.contiguous()
    values          = values.contiguous()
    rewards         = rewards.contiguous()
    terminateds     = terminateds.contiguous()
    if has_truncations:
        truncateds = truncateds.contiguous()

    # Fused kernel for seq_len <= 131072.
    # Pass truncateds=None for the no-truncation path so compute_vtrace_fused
    # dispatches HAS_TRUNCATIONS=False -- no zero-tensor allocations for truncateds
    # or 2D bootstrap_values in that path.
    if seq_len <= _FLAT_MAX_SEQ_LEN:
        return compute_vtrace_fused(
            log_pi_target, log_pi_behavior,
            values, rewards, terminateds,
            truncateds=truncateds if has_truncations else None,
            gamma=gamma, rho_bar=rho_bar, c_bar=c_bar,
            bootstrap_values=bootstrap_values,
            last_value=last_value,
        )

    # Chunked path for seq_len > 131072.  Materialize None inputs here only.
    if not has_truncations:
        truncateds = torch.zeros_like(terminateds)
    if last_value is not None:
        bootstrap_values = torch.zeros_like(rewards)
        bootstrap_values[:, -1] = last_value
    elif bootstrap_values is None:
        bootstrap_values = torch.zeros_like(rewards)

    not_terminated = 1.0 - terminateds
    not_done       = 1.0 - (terminateds + truncateds).clamp(max=1.0)
    next_values    = torch.empty_like(values)
    next_values[:, :-1] = values[:, 1:]
    next_values[:, -1]  = 0.0
    next_values = next_values * (1.0 - truncateds)
    next_values[:, -1]  = 0.0
    next_values = next_values + bootstrap_values

    is_ratios      = torch.exp(log_pi_target - log_pi_behavior)
    rho            = torch.clamp(is_ratios, max=rho_bar)
    c              = torch.clamp(is_ratios, max=c_bar)
    u = rho * (rewards + gamma * next_values * not_terminated - values)
    v = gamma * c * not_done   # beta[t], the scan decay coefficient

    # Additive boundary carry Delta[T] = 0: bootstrap_values[:, -1] already
    # entered u[:, -1] via next_values above (weight 1). Seeding the scan's
    # boundary carry with it too would double-count it -- see kernels/gae.py's
    # module docstring for the same bug class in GAE. beta (v) itself is
    # unaffected by this fix.
    value_deltas   = _run_scan(u, v)
    vtrace_targets = value_deltas + values

    next_vtrace_targets = torch.empty_like(vtrace_targets)
    next_vtrace_targets[:, :-1] = vtrace_targets[:, 1:]
    next_vtrace_targets[:, -1]  = 0.0
    next_vtrace_targets = next_vtrace_targets * (1.0 - truncateds)
    next_vtrace_targets[:, -1]  = 0.0
    next_vtrace_targets = next_vtrace_targets + bootstrap_values

    vtrace_advantages = rho * (rewards + gamma * next_vtrace_targets * not_terminated - values)
    return vtrace_targets, vtrace_advantages

`rl_triton.ops.retrace.compute_retrace(action_probs_target, action_probs_behavior, q_values, next_q_values_all, actions, rewards, terminateds, truncateds, gamma, lambda_=1.0, c_bar=1.0, rho_bar=1.0)`

Compute Retrace(λ) Q-value targets and advantages via a backward associative scan.

Retrace(λ) (Munos et al. 2016) corrects off-policy Q-value estimates using truncated IS traces applied only to the decay factor, not the TD error. Discrete actions only -- use compute_vtrace for continuous action spaces.

Recurrence:

Δ[t] = δ[t] + decay[t] * Δ[t+1], Δ[T] = 0
δ[t] = r[t] + γ · E_π[Q(s_{t+1},a)] · (1-terminated[t]) - Q(s_t,a_t)
decay[t] = γ · c[t+1] · (1-done[t]), done[t] = terminated[t] | truncated[t]
c[t] = λ · min(c_bar, π(a_t|s_t) / μ(a_t|s_t))

Q-value targets: Q_ret[t] = Q(s_t, a_t) + Δ[t] Advantages: A[t] = ρ[t] · (r[t] + γ · Q_ret[t+1] · (1-terminated[t]) - Q(s_t,a_t)) ρ[t] = min(rho_bar, π(a_t|s_t) / μ(a_t|s_t))

The Q-bootstrap γ·E_π[Q(s_{t+1},·)] is folded into δ[t] via next_q_values_all, so no separate bootstrap_values argument is needed. terminated gates the one-step bootstrap in δ[t]; terminated | truncated gates trace decay in β[t].

Parameters:

Name	Type	Description	Default
`action_probs_target`	`Tensor`	Target policy probabilities over all actions, [num_envs, seq_len, num_actions], float32, CUDA.	required
`action_probs_behavior`	`Tensor`	Behavior policy probability of the taken action, [num_envs, seq_len], float32, CUDA.	required
`q_values`	`Tensor`	Q(s_t, a_t) for the taken action, [num_envs, seq_len], float32, CUDA.	required
`next_q_values_all`	`Tensor`	Q(s_{t+1}, a) for all actions, [num_envs, seq_len, num_actions], float32, CUDA.	required
`actions`	`Tensor`	Indices of the taken action, [num_envs, seq_len], int64, CUDA.	required
`rewards`	`Tensor`	Per-step rewards, [num_envs, seq_len], float32, CUDA.	required
`terminateds`	`Tensor`	True termination flags (1.0 = terminated). Zeros the bootstrap γ·E_π[Q(s_{t+1},·)] in δ[t]. [num_envs, seq_len], float32, CUDA.	required
`truncateds`	`Tensor`	Time-limit truncation flags (1.0 = truncated). Keeps the bootstrap in δ[t] but severs the trace. [num_envs, seq_len], float32, CUDA.	required
`gamma`	`float`	Discount factor.	required
`lambda_`	`float`	Trace decay parameter (default 1.0).	`1.0`
`c_bar`	`float`	IS ratio clip for trace weights (default 1.0).	`1.0`
`rho_bar`	`float`	IS ratio clip for advantage scaling (default 1.0).	`1.0`

Returns:

Name	Type	Description
`retrace_targets`	`Tensor`	Q_ret[t], shape [num_envs, seq_len], float32.
`advantages`	`Tensor`	A[t], shape [num_envs, seq_len], float32.

Source code in src/rl_triton/ops/retrace.py

def compute_retrace(
    action_probs_target: torch.Tensor,
    action_probs_behavior: torch.Tensor,
    q_values: torch.Tensor,
    next_q_values_all: torch.Tensor,
    actions: torch.Tensor,
    rewards: torch.Tensor,
    terminateds: torch.Tensor,
    truncateds: torch.Tensor,
    gamma: float,
    lambda_: float = 1.0,
    c_bar: float = 1.0,
    rho_bar: float = 1.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Compute Retrace(λ) Q-value targets and advantages via a backward associative scan.

    Retrace(λ) (Munos et al. 2016) corrects off-policy Q-value estimates using
    truncated IS traces applied only to the decay factor, not the TD error.
    Discrete actions only -- use compute_vtrace for continuous action spaces.

    Recurrence:

    - Δ[t] = δ[t] + decay[t] * Δ[t+1],   Δ[T] = 0
    - δ[t]     = r[t] + γ · E_π[Q(s_{t+1},a)] · (1-terminated[t]) - Q(s_t,a_t)
    - decay[t] = γ · c[t+1] · (1-done[t]),  done[t] = terminated[t] | truncated[t]
    - c[t]     = λ · min(c_bar, π(a_t|s_t) / μ(a_t|s_t))

    Q-value targets: Q_ret[t] = Q(s_t, a_t) + Δ[t]
    Advantages:      A[t]     = ρ[t] · (r[t] + γ · Q_ret[t+1] · (1-terminated[t]) - Q(s_t,a_t))
                     ρ[t]     = min(rho_bar, π(a_t|s_t) / μ(a_t|s_t))

    The Q-bootstrap γ·E_π[Q(s_{t+1},·)] is folded into δ[t] via
    `next_q_values_all`, so no separate bootstrap_values argument is needed.
    `terminated` gates the one-step bootstrap in δ[t]; `terminated | truncated`
    gates trace decay in β[t].

    Args:
        action_probs_target:   Target policy probabilities over all actions,
                               [num_envs, seq_len, num_actions], float32, CUDA.
        action_probs_behavior: Behavior policy probability of the taken action,
                               [num_envs, seq_len], float32, CUDA.
        q_values:              Q(s_t, a_t) for the taken action,
                               [num_envs, seq_len], float32, CUDA.
        next_q_values_all:     Q(s_{t+1}, a) for all actions,
                               [num_envs, seq_len, num_actions], float32, CUDA.
        actions:               Indices of the taken action,
                               [num_envs, seq_len], int64, CUDA.
        rewards:               Per-step rewards, [num_envs, seq_len], float32, CUDA.
        terminateds:           True termination flags (1.0 = terminated).
                               Zeros the bootstrap γ·E_π[Q(s_{t+1},·)] in δ[t].
                               [num_envs, seq_len], float32, CUDA.
        truncateds:            Time-limit truncation flags (1.0 = truncated).
                               Keeps the bootstrap in δ[t] but severs the trace.
                               [num_envs, seq_len], float32, CUDA.
        gamma:                 Discount factor.
        lambda_:               Trace decay parameter (default 1.0).
        c_bar:                 IS ratio clip for trace weights (default 1.0).
        rho_bar:               IS ratio clip for advantage scaling (default 1.0).

    Returns:
        retrace_targets: Q_ret[t], shape [num_envs, seq_len], float32.
        advantages:      A[t],     shape [num_envs, seq_len], float32.
    """
    # Cheap structural checks -- always-on.
    for name, t in [
        ("action_probs_behavior", action_probs_behavior),
        ("q_values",              q_values),
        ("rewards",               rewards),
        ("terminateds",           terminateds),
        ("truncateds",            truncateds),
    ]:
        assert t.is_cuda,                 f"{name} must be on CUDA"
        assert t.dtype == torch.float32,  f"{name}: expected float32, got {t.dtype}"
        assert t.shape == rewards.shape,  f"{name} shape {t.shape} != rewards shape {rewards.shape}"

    assert action_probs_target.is_cuda,                "action_probs_target must be on CUDA"
    assert action_probs_target.dtype == torch.float32, "action_probs_target: expected float32"
    assert action_probs_target.shape[:2] == rewards.shape, (
        f"action_probs_target shape {action_probs_target.shape} incompatible with rewards {rewards.shape}"
    )
    assert next_q_values_all.is_cuda,                "next_q_values_all must be on CUDA"
    assert next_q_values_all.dtype == torch.float32, "next_q_values_all: expected float32"
    assert next_q_values_all.shape == action_probs_target.shape, (
        f"next_q_values_all {next_q_values_all.shape} != action_probs_target {action_probs_target.shape}"
    )
    assert actions.is_cuda,              "actions must be on CUDA"
    assert actions.dtype == torch.int64, f"actions: expected int64, got {actions.dtype}"
    assert actions.shape == rewards.shape, \
        f"actions shape {actions.shape} != rewards shape {rewards.shape}"

    # Expensive tensor scan -- correctness-warning path only.
    if _CORRECTNESS_WARNINGS():
        assert not (terminateds.bool() & truncateds.bool()).any(), \
            "terminated and truncated are mutually exclusive: a step cannot be both"

    action_probs_target   = action_probs_target.contiguous()
    action_probs_behavior = action_probs_behavior.contiguous()
    q_values              = q_values.contiguous()
    next_q_values_all     = next_q_values_all.contiguous()
    actions               = actions.contiguous()
    rewards               = rewards.contiguous()
    terminateds           = terminateds.contiguous()
    truncateds            = truncateds.contiguous()

    num_envs, seq_len = rewards.shape

    # Fused Triton kernel for seq_len <= _TRITON_SEQ_LEN_CEILING (confirmed win);
    # generic materialize-u/v + _run_scan path above that (fused kernel is a
    # confirmed loss there -- see _TRITON_SEQ_LEN_CEILING comment; this reroute
    # is a smaller loss too, not a win -- see the same comment).
    # done[t] = terminated[t] | truncated[t] is computed in-kernel from the two
    # raw flags for the fused path -- no separate PyTorch combine here.
    if seq_len <= _TRITON_SEQ_LEN_CEILING:
        return compute_retrace_fused(
            action_probs_target, action_probs_behavior,
            q_values, next_q_values_all, actions,
            rewards, truncateds, terminateds,
            gamma=gamma, lambda_=lambda_, c_bar=c_bar, rho_bar=rho_bar,
        )

    # seq_len > _TRITON_SEQ_LEN_CEILING: materialize u/v with plain PyTorch ops
    # (bandwidth-bound, no in-kernel 3D reads) and hand off to the generic
    # scan kernel via _run_scan, which internally picks the flat associative
    # scan for seq_len <= 131072 or the chunked kernel above that.
    dones = (terminateds + truncateds).clamp(max=1.0)
    expected_next_q = (action_probs_target * next_q_values_all).sum(dim=-1)
    u = rewards + gamma * expected_next_q * (1.0 - terminateds) - q_values

    pi_a   = action_probs_target.gather(-1, actions.unsqueeze(-1)).squeeze(-1)
    c      = lambda_ * torch.clamp(pi_a / action_probs_behavior, max=c_bar)
    c_next = torch.empty_like(c)
    c_next[:, :-1] = c[:, 1:]
    c_next[:, -1]  = 0.0
    v = gamma * c_next * (1.0 - dones)

    retrace_targets = _run_scan(u, v) + q_values

    rho = torch.clamp(pi_a / action_probs_behavior, max=rho_bar)
    next_q_ret = torch.empty_like(retrace_targets)
    next_q_ret[:, :-1] = retrace_targets[:, 1:]
    next_q_ret[:, -1]  = 0.0
    advantages = rho * (rewards + gamma * next_q_ret * (1.0 - terminateds) - q_values)

    return retrace_targets, advantages

`rl_triton.ops.returns.compute_lambda_returns(rewards, next_values, terminateds, truncateds=None, gamma=0.99, lambda_=0.95, bootstrap_values=None)`

Compute TD(λ) targets (λ-returns) via a backward associative scan.

Recurrence:

G[t] = r[t] + γ(1-done[t]) * [(1-λ)V(s_{t+1}) + λG[t+1]] + γtruncated[t]*bootstrap[t]
done[t] = terminated[t] | truncated[t]

Special cases: λ=0 reduces to one-step TD; λ=1 reduces to discounted returns.

terminated[t] stops both the value bootstrap and trace propagation. truncated[t] stops trace propagation and injects bootstrap_values[env, t] = V(s_{t+1}^true) as the continuation value. The caller must zero next_values[env, t] at truncated steps to avoid double-counting.

Parameters:

Name	Type	Description	Default
`rewards`	`Tensor`	Per-step rewards, [num_envs, seq_len], float32, CUDA.	required
`next_values`	`Tensor`	V(s_{t+1}), [num_envs, seq_len], float32, CUDA. Must be zeroed by caller at truncated steps.	required
`terminateds`	`Tensor`	True termination flags (1.0=terminated), [num_envs, seq_len], float32, CUDA.	required
`truncateds`	`Tensor \| None`	Time-limit truncation flags (1.0=truncated), [num_envs, seq_len], float32, CUDA. If None, terminateds is used for both gating roles.	`None`
`gamma`	`float`	Discount factor (default 0.99).	`0.99`
`lambda_`	`float`	Trace parameter in [0, 1] (default 0.95).	`0.95`
`bootstrap_values`	`Tensor \| None`	True continuation values V(s_{t+1}^true), [num_envs, seq_len], float32, CUDA. Nonzero at truncated steps and at t=T-1 when the window ends mid-episode; zero elsewhere. If None, defaults to all zeros.	`None`

Returns:

Name	Type	Description
`lambda_returns`	`Tensor`	G[t], shape [num_envs, seq_len], float32.

Source code in src/rl_triton/ops/returns.py

def compute_lambda_returns(
    rewards: torch.Tensor,
    next_values: torch.Tensor,
    terminateds: torch.Tensor,
    truncateds: torch.Tensor | None = None,
    gamma: float = 0.99,
    lambda_: float = 0.95,
    bootstrap_values: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute TD(λ) targets (λ-returns) via a backward associative scan.

    Recurrence:

    - G[t] = r[t] + γ*(1-done[t]) * [(1-λ)*V(s_{t+1}) + λ*G[t+1]]
             + γ*truncated[t]*bootstrap[t]
    - done[t] = terminated[t] | truncated[t]

    Special cases: λ=0 reduces to one-step TD; λ=1 reduces to discounted returns.

    `terminated[t]` stops both the value bootstrap and trace propagation.
    `truncated[t]` stops trace propagation and injects bootstrap_values[env, t]
    = V(s_{t+1}^true) as the continuation value.  The caller must zero
    next_values[env, t] at truncated steps to avoid double-counting.

    Args:
        rewards:          Per-step rewards, [num_envs, seq_len], float32, CUDA.
        next_values:      V(s_{t+1}), [num_envs, seq_len], float32, CUDA.
                          Must be zeroed by caller at truncated steps.
        terminateds:      True termination flags (1.0=terminated),
                          [num_envs, seq_len], float32, CUDA.
        truncateds:       Time-limit truncation flags (1.0=truncated),
                          [num_envs, seq_len], float32, CUDA.
                          If None, terminateds is used for both gating roles.
        gamma:            Discount factor (default 0.99).
        lambda_:          Trace parameter in [0, 1] (default 0.95).
        bootstrap_values: True continuation values V(s_{t+1}^true),
                          [num_envs, seq_len], float32, CUDA.
                          Nonzero at truncated steps and at t=T-1 when the window
                          ends mid-episode; zero elsewhere.
                          If None, defaults to all zeros.

    Returns:
        lambda_returns: G[t], shape [num_envs, seq_len], float32.
    """
    num_envs, seq_len = rewards.shape
    has_truncations   = truncateds is not None

    # Cheap structural checks -- always-on.
    for name, t in [("rewards", rewards), ("next_values", next_values), ("terminateds", terminateds)]:
        assert t.is_cuda,                f"{name} must be on CUDA"
        assert t.dtype == torch.float32, f"{name}: expected float32, got {t.dtype}"
        assert t.shape == rewards.shape, f"{name} shape {t.shape} != rewards shape {rewards.shape}"
    if has_truncations:
        assert truncateds.is_cuda,                "truncateds must be on CUDA"
        assert truncateds.dtype == torch.float32, "truncateds: expected float32"
        assert truncateds.shape == rewards.shape, \
            f"truncateds shape {truncateds.shape} != rewards shape {rewards.shape}"
    if bootstrap_values is not None:
        assert bootstrap_values.is_cuda,                "bootstrap_values must be on CUDA"
        assert bootstrap_values.dtype == torch.float32, "bootstrap_values: expected float32"
        assert bootstrap_values.shape == rewards.shape, \
            f"bootstrap_values shape {bootstrap_values.shape} != rewards shape {rewards.shape}"

    # Expensive tensor scans -- correctness-warning path only.
    if _CORRECTNESS_WARNINGS():
        if has_truncations:
            assert not (terminateds.bool() & truncateds.bool()).any(), \
                "terminated and truncated are mutually exclusive: a step cannot be both"

    rewards     = rewards.contiguous()
    next_values = next_values.contiguous()
    terminateds = terminateds.contiguous()
    if has_truncations:
        truncateds = truncateds.contiguous()
    if bootstrap_values is not None:
        bootstrap_values = bootstrap_values.contiguous()

    out = torch.empty_like(rewards)

    if seq_len <= _FLAT_MAX_SEQ_LEN:
        BLOCK_SIZE = triton.next_power_of_2(seq_len)
        num_warps  = _WARPS_LAMBDA.get(BLOCK_SIZE, 16)
        num_stages = 2 if BLOCK_SIZE >= 2048 else 1
        if has_truncations:
            if bootstrap_values is None:
                bootstrap_values = torch.zeros_like(rewards)
            lambda_returns_fused_kernel[(num_envs,)](
                rewards, next_values, terminateds, truncateds,
                out, bootstrap_values,
                seq_len, rewards.stride(0),
                gamma=gamma, lambda_=lambda_,
                BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages,
                HAS_TRUNCATIONS=True, HAS_BOOTSTRAP=True,
            )
        else:
            scalar_bootstrap = bootstrap_values[:, -1].contiguous() \
                               if bootstrap_values is not None else None
            has_bootstrap = scalar_bootstrap is not None
            lambda_returns_fused_kernel[(num_envs,)](
                rewards, next_values, terminateds, None,
                out, scalar_bootstrap,
                seq_len, rewards.stride(0),
                gamma=gamma, lambda_=lambda_,
                BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages,
                HAS_TRUNCATIONS=False, HAS_BOOTSTRAP=has_bootstrap,
            )
        return out

    # Chunked fallback for seq_len > 131072.
    if not has_truncations:
        truncateds = torch.zeros_like(terminateds)
    if bootstrap_values is None:
        bootstrap_values = torch.zeros_like(rewards)
    not_done = 1.0 - (terminateds + truncateds).clamp(max=1.0)
    carry    = bootstrap_values[:, -1]
    u = rewards + gamma * (1.0 - lambda_) * not_done * next_values \
        + gamma * truncateds * bootstrap_values
    v = gamma * lambda_ * not_done
    return _run_scan(u, v, carry)

`rl_triton.ops.returns.compute_discounted_returns(rewards, terminateds, truncateds=None, gamma=0.99, bootstrap_values=None)`

Compute discounted returns (reward-to-go) via a backward associative scan.

Recurrence:

G[t] = r[t] + γ(1-done[t])G[t+1] + γtruncated[t]bootstrap[t]
done[t] = terminated[t] | truncated[t]

terminated[t] stops the return propagation (episode ends, G[t+1]=0). truncated[t] also stops propagation but injects bootstrap_values[env, t] = V(s_{t+1}^true) as the continuation value instead of zero.

Parameters:

Name	Type	Description	Default
`rewards`	`Tensor`	Per-step rewards, [num_envs, seq_len], float32, CUDA.	required
`terminateds`	`Tensor`	True termination flags (1.0=terminated), [num_envs, seq_len], float32, CUDA.	required
`truncateds`	`Tensor \| None`	Time-limit truncation flags (1.0=truncated), [num_envs, seq_len], float32, CUDA. If None, terminateds is used for both gating roles.	`None`
`gamma`	`float`	Discount factor (default 0.99).	`0.99`
`bootstrap_values`	`Tensor \| None`	True continuation values V(s_{t+1}^true), [num_envs, seq_len], float32, CUDA. Nonzero at truncated steps and at t=T-1 when the window ends mid-episode; zero elsewhere. If None, defaults to all zeros.	`None`

Returns:

Name	Type	Description
`returns`	`Tensor`	G[t], shape [num_envs, seq_len], float32.

Source code in src/rl_triton/ops/returns.py

def compute_discounted_returns(
    rewards: torch.Tensor,
    terminateds: torch.Tensor,
    truncateds: torch.Tensor | None = None,
    gamma: float = 0.99,
    bootstrap_values: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute discounted returns (reward-to-go) via a backward associative scan.

    Recurrence:

    - G[t] = r[t] + γ*(1-done[t])*G[t+1] + γ*truncated[t]*bootstrap[t]
    - done[t] = terminated[t] | truncated[t]

    `terminated[t]` stops the return propagation (episode ends, G[t+1]=0).
    `truncated[t]` also stops propagation but injects bootstrap_values[env, t]
    = V(s_{t+1}^true) as the continuation value instead of zero.

    Args:
        rewards:          Per-step rewards, [num_envs, seq_len], float32, CUDA.
        terminateds:      True termination flags (1.0=terminated),
                          [num_envs, seq_len], float32, CUDA.
        truncateds:       Time-limit truncation flags (1.0=truncated),
                          [num_envs, seq_len], float32, CUDA.
                          If None, terminateds is used for both gating roles.
        gamma:            Discount factor (default 0.99).
        bootstrap_values: True continuation values V(s_{t+1}^true),
                          [num_envs, seq_len], float32, CUDA.
                          Nonzero at truncated steps and at t=T-1 when the window
                          ends mid-episode; zero elsewhere.
                          If None, defaults to all zeros.

    Returns:
        returns: G[t], shape [num_envs, seq_len], float32.
    """
    num_envs, seq_len = rewards.shape
    has_truncations   = truncateds is not None

    # Cheap structural checks -- always-on.
    for name, t in [("rewards", rewards), ("terminateds", terminateds)]:
        assert t.is_cuda,                f"{name} must be on CUDA"
        assert t.dtype == torch.float32, f"{name}: expected float32, got {t.dtype}"
        assert t.shape == rewards.shape, f"{name} shape {t.shape} != rewards shape {rewards.shape}"
    if has_truncations:
        assert truncateds.is_cuda,                "truncateds must be on CUDA"
        assert truncateds.dtype == torch.float32, "truncateds: expected float32"
        assert truncateds.shape == rewards.shape, \
            f"truncateds shape {truncateds.shape} != rewards shape {rewards.shape}"
    if bootstrap_values is not None:
        assert bootstrap_values.is_cuda,                "bootstrap_values must be on CUDA"
        assert bootstrap_values.dtype == torch.float32, "bootstrap_values: expected float32"
        assert bootstrap_values.shape == rewards.shape, \
            f"bootstrap_values shape {bootstrap_values.shape} != rewards shape {rewards.shape}"

    # Expensive tensor scans -- correctness-warning path only.
    if _CORRECTNESS_WARNINGS():
        if has_truncations:
            assert not (terminateds.bool() & truncateds.bool()).any(), \
                "terminated and truncated are mutually exclusive: a step cannot be both"

    rewards     = rewards.contiguous()
    terminateds = terminateds.contiguous()
    if has_truncations:
        truncateds = truncateds.contiguous()
    if bootstrap_values is not None:
        bootstrap_values = bootstrap_values.contiguous()

    out = torch.empty_like(rewards)

    if seq_len <= _FLAT_MAX_SEQ_LEN:
        BLOCK_SIZE = triton.next_power_of_2(seq_len)
        num_warps  = _WARPS_LIGHT.get(BLOCK_SIZE, 16)
        num_stages = 2 if BLOCK_SIZE >= 1024 else 1
        if has_truncations:
            if bootstrap_values is None:
                bootstrap_values = torch.zeros_like(rewards)
            discounted_returns_fused_kernel[(num_envs,)](
                rewards, terminateds, truncateds,
                out, bootstrap_values,
                seq_len, rewards.stride(0),
                gamma=gamma,
                BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages,
                HAS_TRUNCATIONS=True, HAS_BOOTSTRAP=True,
            )
        else:
            scalar_bootstrap = bootstrap_values[:, -1].contiguous() \
                               if bootstrap_values is not None else None
            has_bootstrap = scalar_bootstrap is not None
            discounted_returns_fused_kernel[(num_envs,)](
                rewards, terminateds, None,
                out, scalar_bootstrap,
                seq_len, rewards.stride(0),
                gamma=gamma,
                BLOCK_SIZE=BLOCK_SIZE, num_warps=num_warps, num_stages=num_stages,
                HAS_TRUNCATIONS=False, HAS_BOOTSTRAP=has_bootstrap,
            )
        return out

    # Chunked fallback for seq_len > 131072.
    if not has_truncations:
        truncateds = torch.zeros_like(terminateds)
    if bootstrap_values is None:
        bootstrap_values = torch.zeros_like(rewards)
    done  = (terminateds + truncateds).clamp(max=1.0)
    carry = bootstrap_values[:, -1]
    u = rewards + gamma * truncateds * bootstrap_values
    v = gamma * (1.0 - done)
    return _run_scan(u, v, carry)

`rl_triton.ops.returns.compute_eligibility_traces(gradients, dones, gamma, lambda_, seed_values=None)`

Compute accumulating eligibility traces via a forward associative scan.

Recurrence:

z[t] = g[t] + γ·λ·(1 - done[t-1]) * z[t-1], z[-1] = seed, done[-1] := 0

g[t] is the per-step input: the value-function gradient ∇_w V̂(s_t) for general function approximation, or the feature vector x(s_t) in the linear case. Unlike all other kernels, this scan runs forward in time.

done uses the same convention as every other function in this package (done[t]=1 means the episode ends AT t) -- callers pass the same raw done array used elsewhere, unshifted. Internally, the trace carried into t is severed whenever the PRECEDING step ended an episode (done[t-1]=1), i.e. whenever t is the first step of a new episode; gating on done[t] instead would sever the carry at the old episode's own last step (still that episode's data) and fail to sever it at the new episode's first step. Limited to seq_len <= 131072.

Parameters:

Name	Type	Description	Default
`gradients`	`Tensor`	g[t] -- value-function gradients or feature vectors, [num_envs, seq_len], float32, CUDA.	required
`dones`	`Tensor`	Episode termination flags (1.0=done), [num_envs, seq_len], float32, CUDA. done[t]=1 means episode ends at t (same convention as compute_gae etc.); do not pre-shift this array.	required
`gamma`	`float`	Discount factor.	required
`lambda_`	`float`	Trace decay parameter.	required
`seed_values`	`Tensor \| None`	Initial trace z[-1] per environment, shape [num_envs]. Defaults to zeros.	`None`

Returns:

Name	Type	Description
`traces`	`Tensor`	z[t], shape [num_envs, seq_len], float32.

Source code in src/rl_triton/ops/returns.py

def compute_eligibility_traces(
    gradients: torch.Tensor,
    dones: torch.Tensor,
    gamma: float,
    lambda_: float,
    seed_values: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute accumulating eligibility traces via a forward associative scan.

    Recurrence:

    - z[t] = g[t] + γ·λ·(1 - done[t-1]) * z[t-1],  z[-1] = seed,  done[-1] := 0

    g[t] is the per-step input: the value-function gradient ∇_w V̂(s_t) for
    general function approximation, or the feature vector x(s_t) in the linear
    case.  Unlike all other kernels, this scan runs forward in time.

    `done` uses the same convention as every other function in this package
    (done[t]=1 means the episode ends AT t) -- callers pass the same raw done
    array used elsewhere, unshifted.  Internally, the trace carried into t is
    severed whenever the PRECEDING step ended an episode (done[t-1]=1), i.e.
    whenever t is the first step of a new episode; gating on done[t] instead
    would sever the carry at the old episode's own last step (still that
    episode's data) and fail to sever it at the new episode's first step.
    Limited to seq_len <= 131072.

    Args:
        gradients:   g[t] -- value-function gradients or feature vectors,
                     [num_envs, seq_len], float32, CUDA.
        dones:       Episode termination flags (1.0=done), [num_envs, seq_len], float32, CUDA.
                     done[t]=1 means episode ends at t (same convention as
                     compute_gae etc.); do not pre-shift this array.
        gamma:       Discount factor.
        lambda_:     Trace decay parameter.
        seed_values: Initial trace z[-1] per environment, shape [num_envs].
                     Defaults to zeros.

    Returns:
        traces: z[t], shape [num_envs, seq_len], float32.
    """
    num_envs, seq_len = gradients.shape

    # Cheap structural checks -- always-on.
    assert gradients.is_cuda and dones.is_cuda, "gradients and dones must be on CUDA"
    assert gradients.dtype == torch.float32, f"gradients: expected float32, got {gradients.dtype}"
    assert dones.dtype == torch.float32,     f"dones: expected float32, got {dones.dtype}"
    assert gradients.shape == dones.shape,   "gradients and dones must have the same shape"
    if seed_values is not None:
        assert seed_values.shape == (num_envs,), \
            f"seed_values must have shape [{num_envs}], got {seed_values.shape}"
        assert seed_values.is_cuda, "seed_values must be on CUDA"

    assert seq_len <= _FLAT_MAX_SEQ_LEN, (
        f"seq_len={seq_len} exceeds the flat kernel limit {_FLAT_MAX_SEQ_LEN}. "
        "A chunked forward scan kernel has not been implemented yet."
    )

    gradients = gradients.contiguous()
    dones     = dones.contiguous()

    has_seed = seed_values is not None
    if has_seed:
        seed_values = seed_values.contiguous()

    out = torch.empty_like(gradients)

    BLOCK_SIZE = triton.next_power_of_2(seq_len)
    num_warps  = _WARPS_LIGHT.get(BLOCK_SIZE, 16)
    num_stages = 2 if BLOCK_SIZE >= 1024 else 1

    eligibility_traces_fused_kernel[(num_envs,)](
        gradients, dones,
        out, seed_values,
        seq_len, gradients.stride(0),
        gamma=gamma, lambda_=lambda_,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=num_warps,
        num_stages=num_stages,
        HAS_SEED=has_seed,
    )
    return out

`rl_triton.ops.prefix_sum.compute_episodic_prefix_sum(inputs, dones, seed_values=None, boundary='ends_at')`

Cumulative sum that resets to zero at episode/segment boundaries.

Two mutually-exclusive boundary conventions, selected by boundary:

"ends_at" (default): done[t]=1 means the segment ENDS at t; the reset lands at t+1. C[t] = x[t] + (1 - done[t-1]) * C[t-1], C[-1] = seed, done[-1] := 0 This is the Gymnasium-next-step / GAE-canonical convention used identically by every other kernel in this library (compute_gae, compute_lambda_returns, compute_discounted_returns, compute_vtrace, compute_retrace, and compute_eligibility_traces). An RL user computing advantages, returns, and a reset-aware timestep counter from one rollout buffer's terminated/truncated flags gets boundary-consistent results across every kernel by default.
"starts_at": done[t]=1 means the segment STARTS at t; the reset lands at t itself, immediately. C[t] = x[t] + (1 - done[t]) * C[t-1], C[-1] = seed This is the convention sequence-packing callers want: a document- boundary flag marks a new document's first token, and a RoPE local position counter must reset exactly there, not one token later. Pass this explicitly for that use case -- it is not the default.

In both modes: done[t]=1 accumulator resets (C[t]=x[t], modulo which index the reset is keyed to); done[t]=0 accumulator continues.

Limited to seq_len <= 131072.

Parameters:

Name	Type	Description	Default
`inputs`	`Tensor`	Values to accumulate x[t], [num_envs, seq_len], float32, CUDA.	required
`dones`	`Tensor`	Episode/segment boundary flags (1.0=flagged), [num_envs, seq_len], float32, CUDA. Meaning depends on `boundary`; never pre-shift this array yourself -- the kernel handles the index internally.	required
`seed_values`	`Tensor \| None`	Initial carry C[-1] per environment, shape [num_envs]. Defaults to zeros.	`None`
`boundary`	`str`	"ends_at" (default) or "starts_at" -- see above.	`'ends_at'`

Returns:

Name	Type	Description
`prefix_sums`	`Tensor`	C[t], shape [num_envs, seq_len], float32.

Source code in src/rl_triton/ops/prefix_sum.py

def compute_episodic_prefix_sum(
    inputs: torch.Tensor,
    dones: torch.Tensor,
    seed_values: torch.Tensor | None = None,
    boundary: str = "ends_at",
) -> torch.Tensor:
    """
    Cumulative sum that resets to zero at episode/segment boundaries.

    Two mutually-exclusive boundary conventions, selected by `boundary`:

    - "ends_at" (default): done[t]=1 means the segment ENDS at t; the reset
      lands at t+1.
        C[t] = x[t] + (1 - done[t-1]) * C[t-1],  C[-1] = seed,  done[-1] := 0
      This is the Gymnasium-next-step / GAE-canonical convention used
      identically by every other kernel in this library (compute_gae,
      compute_lambda_returns, compute_discounted_returns, compute_vtrace,
      compute_retrace, and compute_eligibility_traces). An RL user computing
      advantages, returns, and a reset-aware timestep counter from one
      rollout buffer's terminated/truncated flags gets boundary-consistent
      results across every kernel by default.
    - "starts_at": done[t]=1 means the segment STARTS at t; the reset lands
      at t itself, immediately.
        C[t] = x[t] + (1 - done[t]) * C[t-1],  C[-1] = seed
      This is the convention sequence-packing callers want: a document-
      boundary flag marks a new document's first token, and a RoPE local
      position counter must reset exactly there, not one token later.
      Pass this explicitly for that use case -- it is not the default.

    In both modes: done[t]=1 accumulator resets (C[t]=x[t], modulo which
    index the reset is keyed to); done[t]=0 accumulator continues.

    Limited to seq_len <= 131072.

    Args:
        inputs:      Values to accumulate x[t], [num_envs, seq_len], float32, CUDA.
        dones:       Episode/segment boundary flags (1.0=flagged), [num_envs, seq_len],
                     float32, CUDA. Meaning depends on `boundary`; never pre-shift
                     this array yourself -- the kernel handles the index internally.
        seed_values: Initial carry C[-1] per environment, shape [num_envs].
                     Defaults to zeros.
        boundary:    "ends_at" (default) or "starts_at" -- see above.

    Returns:
        prefix_sums: C[t], shape [num_envs, seq_len], float32.
    """
    num_envs, seq_len = inputs.shape

    if _CORRECTNESS_WARNINGS():
        assert inputs.is_cuda and dones.is_cuda, "inputs and dones must be on CUDA"
        assert inputs.dtype == torch.float32, f"inputs: expected float32, got {inputs.dtype}"
        assert dones.dtype == torch.float32,  f"dones: expected float32, got {dones.dtype}"
        assert inputs.shape == dones.shape,   "inputs and dones must have the same shape"
        if seed_values is not None:
            assert seed_values.shape == (num_envs,), \
                f"seed_values must have shape [{num_envs}], got {seed_values.shape}"
            assert seed_values.is_cuda, "seed_values must be on CUDA"
        assert boundary in ("ends_at", "starts_at"), \
            f"boundary must be 'ends_at' or 'starts_at', got {boundary!r}"

    assert seq_len <= _FLAT_MAX_SEQ_LEN, (
        f"seq_len={seq_len} exceeds the flat kernel limit {_FLAT_MAX_SEQ_LEN}. "
        "A chunked forward scan kernel has not been implemented yet."
    )

    inputs = inputs.contiguous()
    dones  = dones.contiguous()

    has_seed = seed_values is not None
    if has_seed:
        seed_values = seed_values.contiguous()

    out = torch.empty_like(inputs)

    BLOCK_SIZE = triton.next_power_of_2(seq_len)
    num_warps  = _WARPS.get(BLOCK_SIZE, 16)
    num_stages = 2 if BLOCK_SIZE >= 1024 else 1

    prefix_sum_fused_kernel[(num_envs,)](
        inputs, dones,
        out, seed_values,
        seq_len, inputs.stride(0),
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=num_warps,
        num_stages=num_stages,
        HAS_SEED=has_seed,
        BOUNDARY_ENDS_AT=(boundary == "ends_at"),
    )
    return out