From 8b9d9af050c990c48940445e3cd57d820b451606 Mon Sep 17 00:00:00 2001 From: anthony0727 Date: Mon, 10 Aug 2026 01:14:27 +0900 Subject: [PATCH] Add optional Triton categorical target kernel --- configs/agent/flashSAC.yaml | 3 +- flash_rl/agents/flashSAC/_triton/__init__.py | 127 +++++++++++++++ .../flashSAC/_triton/categorical_target.py | 148 ++++++++++++++++++ flash_rl/agents/flashSAC/agent.py | 6 + flash_rl/agents/flashSAC/update.py | 41 +++-- 5 files changed, 313 insertions(+), 12 deletions(-) create mode 100644 flash_rl/agents/flashSAC/_triton/__init__.py create mode 100644 flash_rl/agents/flashSAC/_triton/categorical_target.py diff --git a/configs/agent/flashSAC.yaml b/configs/agent/flashSAC.yaml index c584cf2..09378fc 100755 --- a/configs/agent/flashSAC.yaml +++ b/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 \ No newline at end of file +load_reward_normalizer: true diff --git a/flash_rl/agents/flashSAC/_triton/__init__.py b/flash_rl/agents/flashSAC/_triton/__init__.py new file mode 100644 index 0000000..1e26c8e --- /dev/null +++ b/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 diff --git a/flash_rl/agents/flashSAC/_triton/categorical_target.py b/flash_rl/agents/flashSAC/_triton/categorical_target.py new file mode 100644 index 0000000..590ab68 --- /dev/null +++ b/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 diff --git a/flash_rl/agents/flashSAC/agent.py b/flash_rl/agents/flashSAC/agent.py index 254e6f5..2522157 100644 --- a/flash_rl/agents/flashSAC/agent.py +++ b/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__( diff --git a/flash_rl/agents/flashSAC/update.py b/flash_rl/agents/flashSAC/update.py index 2686b06..c214f6f 100644 --- a/flash_rl/agents/flashSAC/update.py +++ b/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