Add optional Triton fast path for categorical TD-target construction

configs/agent/flashSAC.yaml

@@ -40,6 +40,7 @@ critic_num_bins: 101
 critic_min_v: ${eval:'-${agent.normalized_G_max}'}
 critic_max_v: ${eval:'${agent.normalized_G_max}'}
 critic_target_update_tau: 0.01
+categorical_target_backend: 'pytorch'  # 'pytorch' or optional 'triton'
 
 temp_initial_value: 0.01
 temp_target_sigma: 0.15
@@ -53,4 +54,4 @@ compile_mode: 'auto'     # 'max-autotune' if torch>=2.9.0 else 'reduce-overhead'
 use_amp: true
 
 load_optimizer: true
-load_reward_normalizer: true
+load_reward_normalizer: true

flash_rl/agents/flashSAC/_triton/__init__.py

@@ -0,0 +1,127 @@
+"""Optional Triton accelerators for FlashSAC."""
+
+from __future__ import annotations
+
+import warnings
+from collections.abc import Callable
+from functools import cache
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as package_version
+from typing import Optional
+
+import torch
+
+_NUM_BINS = 101
+_MEASURED_TORCH_VERSION = (2, 9)
+_MEASURED_TRITON_VERSION = (3, 5)
+_kernel: Optional[Callable[..., torch.Tensor]] = None
+_disabled = False
+
+
+def _major_minor(version: str) -> Optional[tuple[int, int]]:
+    try:
+        major, minor = version.split(".")[:2]
+        return int(major), int(minor)
+    except (ValueError, IndexError):
+        return None
+
+
+@cache
+def _has_measured_dependency_versions() -> bool:
+    """Use only the Torch/Triton line where the kernel measured faster."""
+    try:
+        triton_version = package_version("triton")
+    except PackageNotFoundError:
+        return False
+    return (
+        _major_minor(str(torch.__version__)) == _MEASURED_TORCH_VERSION
+        and _major_minor(triton_version) == _MEASURED_TRITON_VERSION
+    )
+
+
+def _can_use_kernel(
+    next_qs: torch.Tensor,
+    next_q_log_probs: torch.Tensor,
+    reward: torch.Tensor,
+    done: torch.Tensor,
+    actor_entropy: torch.Tensor,
+    num_bins: int,
+    min_v: float,
+    max_v: float,
+) -> bool:
+    if _disabled or not _has_measured_dependency_versions():
+        return False
+    if num_bins != _NUM_BINS or max_v <= min_v:
+        return False
+    if next_qs.ndim != 2 or next_qs.shape[0] != 2:
+        return False
+    if next_q_log_probs.ndim != 3 or next_q_log_probs.shape != (*next_qs.shape, num_bins):
+        return False
+
+    batch_size = next_qs.shape[1]
+    if batch_size == 0 or any(tensor.shape != (batch_size,) for tensor in (reward, done, actor_entropy)):
+        return False
+
+    tensors = (next_qs, next_q_log_probs, reward, done, actor_entropy)
+    if torch.is_grad_enabled() and any(tensor.requires_grad for tensor in tensors):
+        return False
+    if not all(
+        tensor.is_cuda and tensor.dtype == torch.float32 and tensor.device == next_qs.device for tensor in tensors
+    ):
+        return False
+    return bool(torch.cuda.get_device_capability(next_qs.device) >= (8, 0))
+
+
+def try_fused_categorical_td_target(
+    next_qs: torch.Tensor,
+    next_q_log_probs: torch.Tensor,
+    reward: torch.Tensor,
+    done: torch.Tensor,
+    actor_entropy: torch.Tensor,
+    gamma: float,
+    num_bins: int,
+    min_v: float,
+    max_v: float,
+) -> Optional[torch.Tensor]:
+    """Return a fused target, or ``None`` so the caller uses PyTorch."""
+    global _disabled, _kernel
+
+    try:
+        if not _can_use_kernel(
+            next_qs,
+            next_q_log_probs,
+            reward,
+            done,
+            actor_entropy,
+            num_bins,
+            min_v,
+            max_v,
+        ):
+            return None
+
+        if _kernel is None:
+            from .categorical_target import fused_categorical_td_target
+
+            _kernel = fused_categorical_td_target
+        return _kernel(
+            next_qs,
+            next_q_log_probs,
+            reward,
+            done,
+            actor_entropy,
+            gamma,
+            min_v,
+            max_v,
+        )
+    except Exception as error:
+        _disabled = True
+        try:
+            warnings.warn(
+                "Disabling the optional Triton categorical target; "
+                f"using the original PyTorch path instead ({type(error).__name__}: {error}).",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+        except RuntimeWarning:
+            pass
+        return None

flash_rl/agents/flashSAC/_triton/categorical_target.py

@@ -0,0 +1,148 @@
+"""Triton fast path for FlashSAC's categorical TD target."""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+_FUSED_CATEGORICAL_TARGET_BLOCK_SIZE = 128
+
+
+@triton.jit  # type: ignore[misc]
+def _fused_categorical_td_target_kernel(  # type: ignore[no-untyped-def]
+    next_qs_ptr,
+    next_q_log_probs_ptr,
+    reward_ptr,
+    done_ptr,
+    actor_entropy_ptr,
+    target_probs_ptr,
+    num_bins,
+    gamma,
+    min_v,
+    max_v,
+    next_qs_stride_critic,
+    next_qs_stride_batch,
+    log_probs_stride_critic,
+    log_probs_stride_batch,
+    log_probs_stride_bin,
+    reward_stride,
+    done_stride,
+    actor_entropy_stride,
+    target_probs_stride_batch,
+    target_probs_stride_bin,
+    BLOCK: tl.constexpr,
+) -> None:
+    """Select the minimum-Q critic and project one batch row without atomics."""
+    batch_idx = tl.program_id(0)
+    offsets = tl.arange(0, BLOCK)
+    valid_source = offsets < num_bins
+
+    q0 = tl.load(next_qs_ptr + batch_idx * next_qs_stride_batch)
+    q1 = tl.load(next_qs_ptr + next_qs_stride_critic + batch_idx * next_qs_stride_batch)
+
+    # Match torch.argmin: select the first value on ties and the first NaN.
+    selected_critic = tl.where((q0 <= q1) | (q0 != q0), 0, 1)
+    selected_log_probs_ptr = (
+        next_q_log_probs_ptr
+        + selected_critic * log_probs_stride_critic
+        + batch_idx * log_probs_stride_batch
+        + offsets * log_probs_stride_bin
+    )
+    log_probs = tl.load(selected_log_probs_ptr, mask=valid_source, other=-float("inf")).to(tl.float32)
+
+    reward = tl.load(reward_ptr + batch_idx * reward_stride).to(tl.float32)
+    done = tl.load(done_ptr + batch_idx * done_stride).to(tl.float32)
+    actor_entropy = tl.load(actor_entropy_ptr + batch_idx * actor_entropy_stride).to(tl.float32)
+
+    bin_width = (max_v - min_v) / (num_bins - 1)
+    bin_values = min_v + offsets.to(tl.float32) * bin_width
+    target_bin_values = reward + gamma * (bin_values - actor_entropy) * (1.0 - done)
+    target_bin_values = tl.minimum(
+        tl.maximum(target_bin_values, min_v, propagate_nan=tl.PropagateNan.ALL),
+        max_v,
+        propagate_nan=tl.PropagateNan.ALL,
+    )
+
+    projected = (target_bin_values - min_v) / bin_width
+    projected = tl.minimum(
+        tl.maximum(projected, 0.0, propagate_nan=tl.PropagateNan.ALL),
+        num_bins - 1.0,
+        propagate_nan=tl.PropagateNan.ALL,
+    )
+    projection_is_nan = projected != projected
+    # Triton's NaN-to-integer conversion is undefined. Use an in-range index
+    # for the reduction, then explicitly surface the numerical failure below.
+    projected_for_index = tl.where(projection_is_nan, 0.0, projected)
+    lower = tl.floor(projected_for_index).to(tl.int32)
+    upper = tl.minimum(lower + 1, num_bins - 1)
+    fraction = projected_for_index - lower.to(tl.float32)
+
+    probabilities = tl.where(valid_source, tl.exp(log_probs), 0.0)
+    lower_mass = probabilities * (1.0 - fraction)
+    upper_mass = probabilities * fraction
+
+    # Dense one-hot reduction avoids nondeterministic atomic accumulation. Rows
+    # are destination bins and columns are source bins.
+    destination = offsets[:, None]
+    valid_columns = valid_source[None, :]
+    lower_contribution = tl.where(
+        valid_columns & (lower[None, :] == destination),
+        lower_mass[None, :],
+        0.0,
+    )
+    upper_contribution = tl.where(
+        valid_columns & (upper[None, :] == destination),
+        upper_mass[None, :],
+        0.0,
+    )
+    accumulated_mass = tl.sum(lower_contribution + upper_contribution, axis=1)
+    accumulated_mass = tl.where(projection_is_nan, float("nan"), accumulated_mass)
+
+    target_offsets = target_probs_ptr + batch_idx * target_probs_stride_batch + offsets * target_probs_stride_bin
+    tl.store(target_offsets, accumulated_mass, mask=offsets < num_bins)
+
+
+def fused_categorical_td_target(
+    next_qs: torch.Tensor,
+    next_q_log_probs: torch.Tensor,
+    reward: torch.Tensor,
+    done: torch.Tensor,
+    actor_entropy: torch.Tensor,
+    gamma: float,
+    min_v: float,
+    max_v: float,
+) -> torch.Tensor:
+    """Launch the kernel after the optional dispatcher validates its inputs."""
+    batch_size = next_qs.shape[1]
+    num_bins = next_q_log_probs.shape[2]
+    target_probs = torch.empty(
+        (batch_size, num_bins),
+        dtype=torch.float32,
+        device=next_qs.device,
+    )
+    _fused_categorical_td_target_kernel[(batch_size,)](
+        next_qs,
+        next_q_log_probs,
+        reward,
+        done,
+        actor_entropy,
+        target_probs,
+        num_bins,
+        gamma,
+        min_v,
+        max_v,
+        next_qs.stride(0),
+        next_qs.stride(1),
+        next_q_log_probs.stride(0),
+        next_q_log_probs.stride(1),
+        next_q_log_probs.stride(2),
+        reward.stride(0),
+        done.stride(0),
+        actor_entropy.stride(0),
+        target_probs.stride(0),
+        target_probs.stride(1),
+        BLOCK=_FUSED_CATEGORICAL_TARGET_BLOCK_SIZE,
+        num_warps=1,
+    )
+    return target_probs

flash_rl/agents/flashSAC/agent.py

@@ -76,6 +76,7 @@ class FlashSACConfig:
 
     load_optimizer: bool
     load_reward_normalizer: bool
+    categorical_target_backend: str = "pytorch"
 
 
 def _init_flashsac_networks(
@@ -306,6 +307,7 @@ def _update_networks(
         device=device,
         use_amp=cfg.use_amp,
         grad_scaler=grad_scaler,
+        categorical_target_backend=cfg.categorical_target_backend,
     )
 
     target_critic_info = update_target_network(
@@ -354,6 +356,10 @@ def __init__(
 
         temp_target_entropy = 0.5 * self._action_dim * math.log(2 * math.pi * math.e * cfg.temp_target_sigma**2)
         compile_mode = _resolve_compile_mode(cfg.compile_mode)
+        if cfg.categorical_target_backend not in ("pytorch", "triton"):
+            raise ValueError(
+                f"categorical_target_backend must be 'pytorch' or 'triton', got {cfg.categorical_target_backend!r}"
+            )
         cfg = replace(cfg, temp_target_entropy=temp_target_entropy, compile_mode=compile_mode)
 
         super().__init__(

flash_rl/agents/flashSAC/update.py

@@ -178,6 +178,7 @@ def update_critic(
     device: torch.device,
     use_amp: bool,
     grad_scaler: Optional[GradScaler],
+    categorical_target_backend: str = "pytorch",
 ) -> dict[str, torch.Tensor]:
     """Update critic network.
 
@@ -195,6 +196,7 @@ def update_critic(
         device: Device to use.
         use_amp: Whether to use automatic mixed precision.
         grad_scaler: GradScaler for FP16 AMP.
+        categorical_target_backend: Backend for categorical TD-target construction.
     """
 
     with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=use_amp):
@@ -223,19 +225,36 @@ def update_critic(
             )
             next_qs = qs_all.chunk(2, dim=1)[1]
             next_q_log_probs = q_infos_all["log_prob"].chunk(2, dim=1)[1]
-            next_q_log_probs = _select_min_q_log_probs(next_qs, next_q_log_probs)
 
             # Compute target probs
-            target_probs = _compute_categorical_td_target(
-                target_log_probs=next_q_log_probs,
-                reward=batch["reward"],  # type: ignore
-                done=batch["terminated"],  # type: ignore
-                actor_entropy=next_actor_entropy,
-                gamma=gamma**n_step,
-                num_bins=num_bins,
-                min_v=min_v,
-                max_v=max_v,
-            )
+            target_probs = None
+            if categorical_target_backend == "triton":
+                from flash_rl.agents.flashSAC._triton import try_fused_categorical_td_target
+
+                target_probs = try_fused_categorical_td_target(
+                    next_qs,
+                    next_q_log_probs,
+                    batch["reward"],  # type: ignore
+                    batch["terminated"],  # type: ignore
+                    next_actor_entropy,
+                    gamma**n_step,
+                    num_bins,
+                    min_v,
+                    max_v,
+                )
+
+            if target_probs is None:
+                next_q_log_probs = _select_min_q_log_probs(next_qs, next_q_log_probs)
+                target_probs = _compute_categorical_td_target(
+                    target_log_probs=next_q_log_probs,
+                    reward=batch["reward"],  # type: ignore
+                    done=batch["terminated"],  # type: ignore
+                    actor_entropy=next_actor_entropy,
+                    gamma=gamma**n_step,
+                    num_bins=num_bins,
+                    min_v=min_v,
+                    max_v=max_v,
+                )
             max_entropy_bonus = next_actor_entropy.max()
 
         # Compute predicted q-value