From 6bf6d15767708406f39f3e0eba9cfe2e08c68e31 Mon Sep 17 00:00:00 2001 From: JLiu4Coding Date: Wed, 8 Jul 2026 14:45:13 +0800 Subject: [PATCH 1/7] Add fused Gumbel-Softmax sampler --- benchmarks/benchmark_gumbel_softmax.py | 149 ++++++++++++++++ docs/.nav.yml | 1 + docs/operators/README.md | 1 + docs/operators/gumbel-softmax.md | 66 +++++++ .../kernels/ops/pytorch/sampling/__init__.py | 3 + .../ops/pytorch/sampling/gumbel_softmax.py | 86 +++++++++ .../kernels/ops/triton/sampling/__init__.py | 3 + .../ops/triton/sampling/gumbel_softmax.py | 166 ++++++++++++++++++ rl_engine/kernels/registry.py | 16 ++ tests/test_gumbel_softmax.py | 152 ++++++++++++++++ 10 files changed, 643 insertions(+) create mode 100644 benchmarks/benchmark_gumbel_softmax.py create mode 100644 docs/operators/gumbel-softmax.md create mode 100644 rl_engine/kernels/ops/pytorch/sampling/__init__.py create mode 100644 rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py create mode 100644 rl_engine/kernels/ops/triton/sampling/__init__.py create mode 100644 rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py create mode 100644 tests/test_gumbel_softmax.py diff --git a/benchmarks/benchmark_gumbel_softmax.py b/benchmarks/benchmark_gumbel_softmax.py new file mode 100644 index 00000000..dbd2a34d --- /dev/null +++ b/benchmarks/benchmark_gumbel_softmax.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Benchmark NativeGumbelSoftmaxOp vs TritonGumbelSoftmaxOp. + +Usage: + python benchmarks/benchmark_gumbel_softmax.py + python benchmarks/benchmark_gumbel_softmax.py --configs "1,512,32000;4,512,50257" +""" + +import argparse + +import torch +from tabulate import tabulate + +from rl_engine.kernels.ops.pytorch.sampling.gumbel_softmax import NativeGumbelSoftmaxOp +from rl_engine.kernels.ops.triton.sampling.gumbel_softmax import TritonGumbelSoftmaxOp +from rl_engine.platforms.device import device_ctx +from rl_engine.utils.logger import logger + + +DEFAULT_CONFIGS = [ + (1, 512, 32000), + (4, 512, 32000), + (4, 1024, 50257), +] + + +def _make_inputs(batch, seq, vocab, device, dtype): + logits = torch.randn(batch, seq, vocab, device=device, dtype=dtype) + gumbels = ( + -torch.empty(batch, seq, vocab, device=device, dtype=torch.float32).exponential_().log() + ) + return logits, gumbels + + +def _time_ms(fn, warmup, iters): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iters + + +def _peak_vram_mb(fn, warmup=3, iters=5): + for _ in range(warmup): + fn() + torch.cuda.synchronize() + torch.cuda.empty_cache() + baseline = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + for _ in range(iters): + fn() + torch.cuda.synchronize() + return (torch.cuda.max_memory_allocated() - baseline) / (1024**2) + + +def run_benchmark(args): + if device_ctx.device_type not in ["cuda", "xpu", "hip"]: + raise RuntimeError("gumbel_softmax benchmark requires a compatible GPU device.") + + device = device_ctx.device + dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16 + native = NativeGumbelSoftmaxOp() + triton_op = TritonGumbelSoftmaxOp() + + logger.info( + f"gumbel_softmax benchmark on {device} (dtype={dtype}, hard={args.hard}, tau={args.tau})" + ) + + rows = [] + for batch, seq, vocab in args.configs: + logits, gumbels = _make_inputs(batch, seq, vocab, device, dtype) + upstream = torch.randn_like(logits) + + def fwd(op, x=logits, g=gumbels): + with torch.no_grad(): + op(x, tau=args.tau, hard=args.hard, gumbels=g) + + def fwd_bwd(op, x_src=logits, g=gumbels, grad=upstream): + x = x_src.detach().clone().requires_grad_(True) + op(x, tau=args.tau, hard=args.hard, gumbels=g).backward(grad) + + n_fwd = _time_ms(lambda: fwd(native), args.warmup, args.iters) + t_fwd = _time_ms(lambda: fwd(triton_op), args.warmup, args.iters) + n_fb = _time_ms(lambda: fwd_bwd(native), args.warmup, args.iters) + t_fb = _time_ms(lambda: fwd_bwd(triton_op), args.warmup, args.iters) + n_vram = _peak_vram_mb(lambda: fwd(native)) + t_vram = _peak_vram_mb(lambda: fwd(triton_op)) + + rows.append( + [ + f"{batch}x{seq}x{vocab}", + str(dtype).replace("torch.", ""), + f"{n_fwd:.3f}", + f"{t_fwd:.3f}", + f"{n_fwd / t_fwd:.2f}x", + f"{n_fb:.3f}", + f"{t_fb:.3f}", + f"{n_fb / t_fb:.2f}x", + f"{n_vram:.0f}", + f"{t_vram:.0f}", + ] + ) + + headers = [ + "shape (B x S x V)", + "dtype", + "native fwd ms", + "triton fwd ms", + "fwd speedup", + "native f+b ms", + "triton f+b ms", + "f+b speedup", + "native fwd MB", + "triton fwd MB", + ] + print(tabulate(rows, headers=headers, tablefmt="github")) + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--iters", type=int, default=20) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument("--tau", type=float, default=1.0) + parser.add_argument("--hard", action="store_true") + parser.add_argument("--dtype", choices=("float16", "bfloat16"), default="float16") + parser.add_argument( + "--configs", + type=str, + default=None, + help="Semicolon-separated 'batch,seq,vocab' tuples, e.g. '1,512,32000;4,512,50257'.", + ) + args = parser.parse_args() + if args.configs: + args.configs = [tuple(int(x) for x in tup.split(",")) for tup in args.configs.split(";")] + else: + args.configs = DEFAULT_CONFIGS + return args + + +if __name__ == "__main__": + run_benchmark(parse_args()) diff --git a/docs/.nav.yml b/docs/.nav.yml index f2178984..78f29d6d 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -15,6 +15,7 @@ nav: - operators/attention.md - operators/fused-logp.md - operators/linear-logp.md + - operators/gumbel-softmax.md - operators/linear-logp-tp-test.md - operators/grpo-loss.md - operators/lm_head.md diff --git a/docs/operators/README.md b/docs/operators/README.md index bf38f368..1b076e0d 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -22,6 +22,7 @@ Every operator page should include: - [Standard Attention](attention.md) - [Fused LogP](fused-logp.md) - [Fused Linear LogP](linear-logp.md) +- [Gumbel-Softmax](gumbel-softmax.md) - [Fused Linear LogP TP Test Runbook](linear-logp-tp-test.md) - [GRPO Loss](grpo-loss.md) - [RoPE](rope.md) diff --git a/docs/operators/gumbel-softmax.md b/docs/operators/gumbel-softmax.md new file mode 100644 index 00000000..3490c13e --- /dev/null +++ b/docs/operators/gumbel-softmax.md @@ -0,0 +1,66 @@ +# Gumbel-Softmax + +Gumbel-Softmax provides differentiable sampling from token logits for rollout and +alignment experiments. It returns either soft probabilities or straight-through +one-hot samples while preserving gradients through the softmax path. + +## Entry Point + +```python +from rl_engine.kernels.registry import kernel_registry + +op = kernel_registry.get_op("gumbel_softmax") +samples = op(logits, tau=1.0, hard=True) +``` + +Backend classes: + +- `NativeGumbelSoftmaxOp` +- `TritonGumbelSoftmaxOp` + +## Definition + +For logits `x`, Gumbel noise `g`, and temperature `tau`: + +```text +z = (x + g) / tau +y_soft = softmax(z) +``` + +When `hard=False`, the output is `y_soft`. When `hard=True`, the forward output +is one-hot at `argmax(y_soft)` and the backward pass uses the straight-through +softmax gradient. + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `logits` | `[..., V]` | fp32/fp16/bf16 | Floating-point tensor, `V` is vocab size. | +| `tau` | scalar | Python float | Must be positive. | +| `hard` | scalar | Python bool | Enables straight-through one-hot output. | +| `gumbels` | `[..., V]` or `None` | Floating point | Optional deterministic noise for tests/repro. | +| Output | `[..., V]` | Same as `logits` | Probabilities or one-hot rows over vocab. | + +The Triton backend flattens leading dimensions to `[N, V]` and launches one +program per row. It supports vocab sizes up to 131072 in this implementation. + +## Backends + +| Backend | Status | Notes | +| --- | --- | --- | +| PyTorch | Reference | Uses PyTorch autograd end to end. | +| Triton | GPU optimized forward | Uses Triton forward and PyTorch math in custom autograd backward. | + +For deterministic correctness tests, pass the same precomputed `gumbels` tensor +to both backends. If `gumbels=None`, each backend samples its own Gumbel noise. + +## Tests and Benchmarks + +```bash +python -m pytest tests/test_gumbel_softmax.py +python benchmarks/benchmark_gumbel_softmax.py +``` + +The benchmark reports forward latency, forward+backward latency, speedup, and +peak forward VRAM for PyTorch and Triton on a single GPU. + diff --git a/rl_engine/kernels/ops/pytorch/sampling/__init__.py b/rl_engine/kernels/ops/pytorch/sampling/__init__.py new file mode 100644 index 00000000..4f7d09a0 --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/sampling/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + diff --git a/rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py b/rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py new file mode 100644 index 00000000..443dc22d --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from typing import Optional + +import torch + + +def _validate_gumbel_softmax_inputs( + logits: torch.Tensor, + tau: float, + gumbels: Optional[torch.Tensor], +) -> None: + if logits.ndim < 2: + raise ValueError(f"logits must have at least 2 dimensions, got shape {tuple(logits.shape)}") + if logits.size(-1) <= 0: + raise ValueError("logits vocab dimension must be non-empty") + if not logits.is_floating_point(): + raise TypeError(f"logits must be a floating-point tensor, got dtype {logits.dtype}") + if tau <= 0: + raise ValueError(f"tau must be positive, got {tau}") + if gumbels is not None: + if gumbels.shape != logits.shape: + raise ValueError( + f"gumbels shape {tuple(gumbels.shape)} must match logits shape " + f"{tuple(logits.shape)}" + ) + if gumbels.device != logits.device: + raise ValueError( + f"gumbels device {gumbels.device} must match logits device {logits.device}" + ) + if not gumbels.is_floating_point(): + raise TypeError(f"gumbels must be floating-point, got dtype {gumbels.dtype}") + + +def _sample_gumbels_like(logits: torch.Tensor) -> torch.Tensor: + # Matches torch.nn.functional.gumbel_softmax's exponential sampling path. + return ( + -torch.empty_like(logits, memory_format=torch.legacy_contiguous_format) + .exponential_() + .log() + ) + + +def gumbel_softmax_reference( + logits: torch.Tensor, + *, + tau: float = 1.0, + hard: bool = False, + gumbels: Optional[torch.Tensor] = None, +) -> torch.Tensor: + _validate_gumbel_softmax_inputs(logits, tau, gumbels) + noise = _sample_gumbels_like(logits) if gumbels is None else gumbels.to(dtype=logits.dtype) + y_soft = torch.softmax((logits + noise) / tau, dim=-1) + if not hard: + return y_soft + + index = y_soft.argmax(dim=-1, keepdim=True) + y_hard = torch.zeros_like(y_soft).scatter_(-1, index, 1.0) + return y_hard - y_soft.detach() + y_soft + + +class NativeGumbelSoftmaxOp: + """PyTorch reference implementation for differentiable Gumbel-Softmax sampling.""" + + def __call__( + self, + logits: torch.Tensor, + *, + tau: float = 1.0, + hard: bool = False, + gumbels: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.apply(logits, tau=tau, hard=hard, gumbels=gumbels) + + def apply( + self, + logits: torch.Tensor, + *, + tau: float = 1.0, + hard: bool = False, + gumbels: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return gumbel_softmax_reference(logits, tau=float(tau), hard=bool(hard), gumbels=gumbels) diff --git a/rl_engine/kernels/ops/triton/sampling/__init__.py b/rl_engine/kernels/ops/triton/sampling/__init__.py new file mode 100644 index 00000000..4f7d09a0 --- /dev/null +++ b/rl_engine/kernels/ops/triton/sampling/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + diff --git a/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py b/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py new file mode 100644 index 00000000..2350c068 --- /dev/null +++ b/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from rl_engine.kernels.ops.pytorch.sampling.gumbel_softmax import ( + _validate_gumbel_softmax_inputs, +) + +_BLOCK_V = 1024 + + +@triton.jit +def _gumbel_softmax_fwd_kernel( + logits_ptr, + gumbels_ptr, + y_soft_ptr, + out_ptr, + seed, + n_rows, + V: tl.constexpr, + tau: tl.constexpr, + HAS_GUMBELS: tl.constexpr, + HARD: tl.constexpr, + BLOCK_V: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_V) + row_off = row.to(tl.int64) * V + + z_max = -float("inf") + for start in range(0, V, BLOCK_V): + offs = start + cols + mask = offs < V + logits = tl.load(logits_ptr + row_off + offs, mask=mask, other=-float("inf")).to( + tl.float32 + ) + if HAS_GUMBELS: + gumbels = tl.load(gumbels_ptr + row_off + offs, mask=mask, other=0.0).to(tl.float32) + else: + u = tl.rand(seed, row_off + offs) + u = tl.minimum(tl.maximum(u, 1.0e-20), 1.0 - 1.0e-7) + gumbels = -tl.log(-tl.log(u)) + z = tl.where(mask, (logits + gumbels) / tau, -float("inf")) + z_max = tl.maximum(z_max, tl.max(z, axis=0)) + + denom = 0.0 + for start in range(0, V, BLOCK_V): + offs = start + cols + mask = offs < V + logits = tl.load(logits_ptr + row_off + offs, mask=mask, other=-float("inf")).to( + tl.float32 + ) + if HAS_GUMBELS: + gumbels = tl.load(gumbels_ptr + row_off + offs, mask=mask, other=0.0).to(tl.float32) + else: + u = tl.rand(seed, row_off + offs) + u = tl.minimum(tl.maximum(u, 1.0e-20), 1.0 - 1.0e-7) + gumbels = -tl.log(-tl.log(u)) + z = tl.where(mask, (logits + gumbels) / tau, -float("inf")) + denom += tl.sum(tl.exp(z - z_max), axis=0) + + for start in range(0, V, BLOCK_V): + offs = start + cols + mask = offs < V + logits = tl.load(logits_ptr + row_off + offs, mask=mask, other=-float("inf")).to( + tl.float32 + ) + if HAS_GUMBELS: + gumbels = tl.load(gumbels_ptr + row_off + offs, mask=mask, other=0.0).to(tl.float32) + else: + u = tl.rand(seed, row_off + offs) + u = tl.minimum(tl.maximum(u, 1.0e-20), 1.0 - 1.0e-7) + gumbels = -tl.log(-tl.log(u)) + z = tl.where(mask, (logits + gumbels) / tau, -float("inf")) + y = tl.exp(z - z_max) / denom + y = tl.where(mask, y, 0.0) + tl.store(y_soft_ptr + row_off + offs, y, mask=mask) + if HARD: + out = tl.where(z == z_max, 1.0, 0.0) + tl.store(out_ptr + row_off + offs, out, mask=mask) + else: + tl.store(out_ptr + row_off + offs, y, mask=mask) + + +class _GumbelSoftmaxFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, logits, gumbels, tau: float, hard: bool, seed: int): + V = logits.shape[-1] + lead_shape = logits.shape[:-1] + logits_2d = logits.contiguous().view(-1, V) + gumbels_2d = ( + gumbels.contiguous().view(-1, V) if gumbels is not None else logits_2d + ) + n_rows = logits_2d.shape[0] + block_v = min(_BLOCK_V, triton.next_power_of_2(V)) + + y_soft = torch.empty_like(logits_2d, dtype=torch.float32) + out = torch.empty_like(logits_2d, dtype=torch.float32) + _gumbel_softmax_fwd_kernel[(n_rows,)]( + logits_2d, + gumbels_2d, + y_soft, + out, + int(seed), + n_rows, + V, + float(tau), + HAS_GUMBELS=gumbels is not None, + HARD=bool(hard), + BLOCK_V=block_v, + ) + + ctx.save_for_backward(y_soft) + ctx.tau = float(tau) + ctx.logits_shape = tuple(logits.shape) + ctx.logits_dtype = logits.dtype + return out.view(*lead_shape, V).to(logits.dtype) + + @staticmethod + def backward(ctx, grad_output): + (y_soft,) = ctx.saved_tensors + grad_2d = grad_output.contiguous().view_as(y_soft).to(torch.float32) + dot = (grad_2d * y_soft).sum(dim=-1, keepdim=True) + grad_logits = y_soft * (grad_2d - dot) / ctx.tau + return grad_logits.view(ctx.logits_shape).to(ctx.logits_dtype), None, None, None, None + + +class TritonGumbelSoftmaxOp: + """Triton Gumbel-Softmax sampler with straight-through hard samples.""" + + def __call__( + self, + logits: torch.Tensor, + *, + tau: float = 1.0, + hard: bool = False, + gumbels: Optional[torch.Tensor] = None, + seed: Optional[int] = None, + ) -> torch.Tensor: + return self.apply(logits, tau=tau, hard=hard, gumbels=gumbels, seed=seed) + + def apply( + self, + logits: torch.Tensor, + *, + tau: float = 1.0, + hard: bool = False, + gumbels: Optional[torch.Tensor] = None, + seed: Optional[int] = None, + ) -> torch.Tensor: + _validate_gumbel_softmax_inputs(logits, float(tau), gumbels) + if logits.device.type not in ("cuda", "xpu", "hip"): + raise RuntimeError( + "TritonGumbelSoftmaxOp requires a GPU tensor (CUDA / ROCm / XPU), got " + f"device '{logits.device}'." + ) + if seed is None: + seed = int(torch.randint(0, 2**31 - 1, (), device="cpu").item()) + return _GumbelSoftmaxFunction.apply(logits, gumbels, float(tau), bool(hard), int(seed)) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 4c3c64e3..88e3eaa2 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -48,6 +48,13 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): # Fused policy-ratio + KL-penalty front-end (PPO/GRPO), logits -> (ratio, kl) TRITON_RATIO_KL = "rl_engine.kernels.ops.triton.loss.ratio_kl.TritonRatioKLOp" PYTORCH_RATIO_KL = "rl_engine.kernels.ops.pytorch.loss.ratio_kl.NativeRatioKLOp" + # Differentiable rollout sampler + TRITON_GUMBEL_SOFTMAX = ( + "rl_engine.kernels.ops.triton.sampling.gumbel_softmax.TritonGumbelSoftmaxOp" + ) + PYTORCH_GUMBEL_SOFTMAX = ( + "rl_engine.kernels.ops.pytorch.sampling.gumbel_softmax.NativeGumbelSoftmaxOp" + ) # RMSNorm(pre-norm / QK-Norm) - pure Pytorch reference(ws1 ground-truth) PYTORCH_NATIVE_RMS_NORM = "rl_engine.kernels.ops.pytorch.norm.rms_norm.NativeRMSNormOp" @@ -113,6 +120,10 @@ def __init__(self): "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], + "gumbel_softmax": [ + OpBackend.TRITON_GUMBEL_SOFTMAX, + OpBackend.PYTORCH_GUMBEL_SOFTMAX, + ], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], "embedding": [OpBackend.PYTORCH_NATIVE_EMBEDDING], @@ -135,6 +146,10 @@ def __init__(self): "rope": [OpBackend.PYTORCH_NATIVE_ROPE], "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], + "gumbel_softmax": [ + OpBackend.TRITON_GUMBEL_SOFTMAX, + OpBackend.PYTORCH_GUMBEL_SOFTMAX, + ], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], @@ -151,6 +166,7 @@ def __init__(self): "rope": [OpBackend.PYTORCH_NATIVE_ROPE], "linear_logp": [OpBackend.PYTORCH_LINEAR_LOGP], "ratio_kl": [OpBackend.PYTORCH_RATIO_KL], + "gumbel_softmax": [OpBackend.PYTORCH_GUMBEL_SOFTMAX], "matmul": [OpBackend.PYTORCH_NATIVE_MATMUL], "rms_norm": [OpBackend.PYTORCH_NATIVE_RMS_NORM], "lm_head": [OpBackend.PYTORCH_NATIVE_LM_HEAD], diff --git a/tests/test_gumbel_softmax.py b/tests/test_gumbel_softmax.py new file mode 100644 index 00000000..20c9f66b --- /dev/null +++ b/tests/test_gumbel_softmax.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.sampling.gumbel_softmax import NativeGumbelSoftmaxOp + +try: + import triton # noqa: F401 + + _HAS_TRITON = True +except ImportError: # pragma: no cover + _HAS_TRITON = False + +requires_triton_cuda = pytest.mark.skipif( + not (_HAS_TRITON and torch.cuda.is_available()), + reason="Triton Gumbel-Softmax requires a CUDA device and Triton.", +) + + +def _inputs(seed, *, device="cpu", dtype=torch.float32, shape=(4, 17)): + gen = torch.Generator(device=device).manual_seed(seed) + logits = torch.randn(*shape, generator=gen, device=device, dtype=dtype) + # Keep deterministic test noise in fp32, then let each backend cast as needed. + gumbels = -torch.empty(*shape, device=device, dtype=torch.float32).exponential_( + generator=gen + ).log() + return logits, gumbels + + +def _run_grad(op, logits, gumbels, grad_out, *, hard=False, tau=0.7): + x = logits.detach().clone().requires_grad_(True) + out = op(x, tau=tau, hard=hard, gumbels=gumbels) + out.backward(grad_out) + return out.detach(), x.grad + + +def test_native_matches_torch_reference_with_supplied_noise(): + logits, gumbels = _inputs(0) + op = NativeGumbelSoftmaxOp() + actual = op(logits, tau=0.9, hard=False, gumbels=gumbels) + expected = torch.softmax((logits + gumbels.to(logits.dtype)) / 0.9, dim=-1) + assert torch.allclose(actual, expected, atol=1e-6) + + +def test_native_hard_is_one_hot_with_straight_through_gradient(): + logits, gumbels = _inputs(1) + logits = logits.requires_grad_(True) + out = NativeGumbelSoftmaxOp()(logits, tau=0.8, hard=True, gumbels=gumbels) + assert torch.allclose(out.sum(dim=-1), torch.ones_like(out[..., 0])) + assert torch.all((out == 0.0) | (out == 1.0)) + + out[..., 0].sum().backward() + assert logits.grad is not None + assert torch.isfinite(logits.grad).all() + assert logits.grad.abs().sum() > 0 + + +def test_native_preserves_leading_shape(): + logits, gumbels = _inputs(2, shape=(2, 3, 19)) + out = NativeGumbelSoftmaxOp()(logits, tau=1.2, hard=False, gumbels=gumbels) + assert out.shape == (2, 3, 19) + assert torch.allclose(out.sum(dim=-1), torch.ones(2, 3), atol=1e-6) + + +def test_native_temperature_changes_distribution_sharpness(): + logits, gumbels = _inputs(3, shape=(5, 23)) + op = NativeGumbelSoftmaxOp() + cool = op(logits, tau=0.25, gumbels=gumbels) + warm = op(logits, tau=2.0, gumbels=gumbels) + assert cool.max(dim=-1).values.mean() > warm.max(dim=-1).values.mean() + + +def test_native_rejects_invalid_inputs(): + op = NativeGumbelSoftmaxOp() + logits, gumbels = _inputs(4) + with pytest.raises(ValueError, match="tau must be positive"): + op(logits, tau=0.0, gumbels=gumbels) + with pytest.raises(ValueError, match="gumbels shape"): + op(logits, gumbels=gumbels[:, :-1]) + with pytest.raises(ValueError, match="at least 2 dimensions"): + op(torch.randn(8)) + + +def test_registry_dispatch_matches_native_cpu(): + from rl_engine.kernels.registry import kernel_registry + + logits, gumbels = _inputs(5) + op = kernel_registry.get_op("gumbel_softmax") + out = op(logits, tau=0.7, hard=False, gumbels=gumbels) + ref = NativeGumbelSoftmaxOp()(logits, tau=0.7, hard=False, gumbels=gumbels) + assert torch.allclose(out, ref, atol=1e-6) + + +@requires_triton_cuda +def test_triton_forward_matches_native_soft_fp32(): + from rl_engine.kernels.ops.triton.sampling.gumbel_softmax import TritonGumbelSoftmaxOp + + logits, gumbels = _inputs(6, device="cuda", shape=(7, 257)) + out = TritonGumbelSoftmaxOp()(logits, tau=0.8, hard=False, gumbels=gumbels) + ref = NativeGumbelSoftmaxOp()(logits, tau=0.8, hard=False, gumbels=gumbels) + assert torch.allclose(out, ref, atol=1e-5, rtol=1e-5) + + +@requires_triton_cuda +def test_triton_forward_matches_native_hard_fp32(): + from rl_engine.kernels.ops.triton.sampling.gumbel_softmax import TritonGumbelSoftmaxOp + + logits, gumbels = _inputs(7, device="cuda", shape=(9, 251)) + out = TritonGumbelSoftmaxOp()(logits, tau=0.6, hard=True, gumbels=gumbels) + ref = NativeGumbelSoftmaxOp()(logits, tau=0.6, hard=True, gumbels=gumbels) + assert torch.allclose(out, ref, atol=0.0) + assert torch.allclose(out.sum(dim=-1), torch.ones_like(out[..., 0])) + + +@requires_triton_cuda +def test_triton_backward_matches_native(): + from rl_engine.kernels.ops.triton.sampling.gumbel_softmax import TritonGumbelSoftmaxOp + + logits, gumbels = _inputs(8, device="cuda", shape=(4, 3, 127)) + grad_out = torch.randn_like(logits) + triton_out, triton_grad = _run_grad( + TritonGumbelSoftmaxOp(), logits, gumbels, grad_out, hard=True, tau=0.9 + ) + native_out, native_grad = _run_grad( + NativeGumbelSoftmaxOp(), logits, gumbels, grad_out, hard=True, tau=0.9 + ) + assert torch.allclose(triton_out, native_out, atol=0.0) + assert torch.allclose(triton_grad, native_grad, atol=2e-5, rtol=2e-5) + + +@requires_triton_cuda +def test_triton_preserves_leading_shape_and_dtype(): + from rl_engine.kernels.ops.triton.sampling.gumbel_softmax import TritonGumbelSoftmaxOp + + logits, gumbels = _inputs(9, device="cuda", dtype=torch.float16, shape=(2, 5, 313)) + out = TritonGumbelSoftmaxOp()(logits, tau=1.0, hard=False, gumbels=gumbels) + assert out.shape == (2, 5, 313) + assert out.dtype == torch.float16 + assert torch.allclose(out.float().sum(dim=-1), torch.ones(2, 5, device="cuda"), atol=1e-3) + + +@requires_triton_cuda +def test_triton_generated_noise_smoke(): + from rl_engine.kernels.ops.triton.sampling.gumbel_softmax import TritonGumbelSoftmaxOp + + logits = torch.randn(8, 1024, device="cuda") + out = TritonGumbelSoftmaxOp()(logits, tau=1.0, hard=False, seed=2026) + assert out.shape == logits.shape + assert torch.isfinite(out).all() + assert torch.allclose(out.sum(dim=-1), torch.ones(8, device="cuda"), atol=1e-5) From 72424de60013abb5b126322ae3604724827f8da7 Mon Sep 17 00:00:00 2001 From: JLiu4Coding Date: Wed, 8 Jul 2026 15:26:13 +0800 Subject: [PATCH 2/7] Optimize Triton Gumbel-Softmax sampler --- benchmarks/benchmark_gumbel_softmax.py | 23 ++- .../ops/triton/sampling/gumbel_softmax.py | 161 +++++++++++------- tests/test_gumbel_softmax.py | 8 +- 3 files changed, 122 insertions(+), 70 deletions(-) diff --git a/benchmarks/benchmark_gumbel_softmax.py b/benchmarks/benchmark_gumbel_softmax.py index dbd2a34d..b86e9f4d 100644 --- a/benchmarks/benchmark_gumbel_softmax.py +++ b/benchmarks/benchmark_gumbel_softmax.py @@ -26,11 +26,15 @@ ] -def _make_inputs(batch, seq, vocab, device, dtype): +def _make_inputs(batch, seq, vocab, device, dtype, internal_gumbels): logits = torch.randn(batch, seq, vocab, device=device, dtype=dtype) - gumbels = ( - -torch.empty(batch, seq, vocab, device=device, dtype=torch.float32).exponential_().log() - ) + gumbels = None + if not internal_gumbels: + gumbels = ( + -torch.empty(batch, seq, vocab, device=device, dtype=torch.float32) + .exponential_() + .log() + ) return logits, gumbels @@ -71,12 +75,13 @@ def run_benchmark(args): triton_op = TritonGumbelSoftmaxOp() logger.info( - f"gumbel_softmax benchmark on {device} (dtype={dtype}, hard={args.hard}, tau={args.tau})" + f"gumbel_softmax benchmark on {device} (dtype={dtype}, hard={args.hard}, " + f"tau={args.tau}, internal_gumbels={args.internal_gumbels})" ) rows = [] for batch, seq, vocab in args.configs: - logits, gumbels = _make_inputs(batch, seq, vocab, device, dtype) + logits, gumbels = _make_inputs(batch, seq, vocab, device, dtype, args.internal_gumbels) upstream = torch.randn_like(logits) def fwd(op, x=logits, g=gumbels): @@ -131,6 +136,12 @@ def parse_args(): parser.add_argument("--tau", type=float, default=1.0) parser.add_argument("--hard", action="store_true") parser.add_argument("--dtype", choices=("float16", "bfloat16"), default="float16") + parser.add_argument( + "--internal-gumbels", + action="store_true", + help="Let each backend sample Gumbel noise internally. By default the benchmark " + "passes a materialized Gumbel tensor for deterministic apples-to-apples timing.", + ) parser.add_argument( "--configs", type=str, diff --git a/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py b/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py index 2350c068..c0b2642a 100644 --- a/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py +++ b/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py @@ -13,7 +13,7 @@ _validate_gumbel_softmax_inputs, ) -_BLOCK_V = 1024 +_MAX_BLOCK_V = 131072 @triton.jit @@ -32,61 +32,59 @@ def _gumbel_softmax_fwd_kernel( ): row = tl.program_id(0) cols = tl.arange(0, BLOCK_V) + mask = cols < V row_off = row.to(tl.int64) * V - z_max = -float("inf") - for start in range(0, V, BLOCK_V): - offs = start + cols - mask = offs < V - logits = tl.load(logits_ptr + row_off + offs, mask=mask, other=-float("inf")).to( - tl.float32 - ) - if HAS_GUMBELS: - gumbels = tl.load(gumbels_ptr + row_off + offs, mask=mask, other=0.0).to(tl.float32) - else: - u = tl.rand(seed, row_off + offs) - u = tl.minimum(tl.maximum(u, 1.0e-20), 1.0 - 1.0e-7) - gumbels = -tl.log(-tl.log(u)) - z = tl.where(mask, (logits + gumbels) / tau, -float("inf")) - z_max = tl.maximum(z_max, tl.max(z, axis=0)) - - denom = 0.0 - for start in range(0, V, BLOCK_V): - offs = start + cols - mask = offs < V - logits = tl.load(logits_ptr + row_off + offs, mask=mask, other=-float("inf")).to( - tl.float32 - ) - if HAS_GUMBELS: - gumbels = tl.load(gumbels_ptr + row_off + offs, mask=mask, other=0.0).to(tl.float32) - else: - u = tl.rand(seed, row_off + offs) - u = tl.minimum(tl.maximum(u, 1.0e-20), 1.0 - 1.0e-7) - gumbels = -tl.log(-tl.log(u)) - z = tl.where(mask, (logits + gumbels) / tau, -float("inf")) - denom += tl.sum(tl.exp(z - z_max), axis=0) - - for start in range(0, V, BLOCK_V): - offs = start + cols - mask = offs < V - logits = tl.load(logits_ptr + row_off + offs, mask=mask, other=-float("inf")).to( - tl.float32 - ) - if HAS_GUMBELS: - gumbels = tl.load(gumbels_ptr + row_off + offs, mask=mask, other=0.0).to(tl.float32) - else: - u = tl.rand(seed, row_off + offs) - u = tl.minimum(tl.maximum(u, 1.0e-20), 1.0 - 1.0e-7) - gumbels = -tl.log(-tl.log(u)) - z = tl.where(mask, (logits + gumbels) / tau, -float("inf")) - y = tl.exp(z - z_max) / denom - y = tl.where(mask, y, 0.0) - tl.store(y_soft_ptr + row_off + offs, y, mask=mask) - if HARD: - out = tl.where(z == z_max, 1.0, 0.0) - tl.store(out_ptr + row_off + offs, out, mask=mask) - else: - tl.store(out_ptr + row_off + offs, y, mask=mask) + logits = tl.load(logits_ptr + row_off + cols, mask=mask, other=-float("inf")).to(tl.float32) + if HAS_GUMBELS: + gumbels = tl.load(gumbels_ptr + row_off + cols, mask=mask, other=0.0).to(tl.float32) + else: + u = tl.rand(seed, row_off + cols) + u = tl.minimum(tl.maximum(u, 1.0e-20), 1.0 - 1.0e-7) + gumbels = -tl.log(-tl.log(u)) + + z = tl.where(mask, (logits + gumbels) / tau, -float("inf")) + z_max = tl.max(z, axis=0) + exp_z = tl.exp(z - z_max) + denom = tl.sum(exp_z, axis=0) + y = tl.where(mask, exp_z / denom, 0.0) + + tl.store(y_soft_ptr + row_off + cols, y, mask=mask) + if HARD: + out = tl.where(z == z_max, 1.0, 0.0) + tl.store(out_ptr + row_off + cols, out, mask=mask) + else: + tl.store(out_ptr + row_off + cols, y, mask=mask) + + +@triton.jit +def _gumbel_softmax_hard_nograd_kernel( + logits_ptr, + gumbels_ptr, + out_ptr, + seed, + n_rows, + V: tl.constexpr, + HAS_GUMBELS: tl.constexpr, + BLOCK_V: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_V) + mask = cols < V + row_off = row.to(tl.int64) * V + + logits = tl.load(logits_ptr + row_off + cols, mask=mask, other=-float("inf")).to(tl.float32) + if HAS_GUMBELS: + gumbels = tl.load(gumbels_ptr + row_off + cols, mask=mask, other=0.0).to(tl.float32) + else: + u = tl.rand(seed, row_off + cols) + u = tl.minimum(tl.maximum(u, 1.0e-20), 1.0 - 1.0e-7) + gumbels = -tl.log(-tl.log(u)) + + z = tl.where(mask, logits + gumbels, -float("inf")) + z_max = tl.max(z, axis=0) + out = tl.where(z == z_max, 1.0, 0.0) + tl.store(out_ptr + row_off + cols, out, mask=mask) class _GumbelSoftmaxFunction(torch.autograd.Function): @@ -99,10 +97,13 @@ def forward(ctx, logits, gumbels, tau: float, hard: bool, seed: int): gumbels.contiguous().view(-1, V) if gumbels is not None else logits_2d ) n_rows = logits_2d.shape[0] - block_v = min(_BLOCK_V, triton.next_power_of_2(V)) + block_v = triton.next_power_of_2(V) + if block_v > _MAX_BLOCK_V: + raise ValueError(f"vocab size {V} exceeds TritonGumbelSoftmaxOp limit {_MAX_BLOCK_V}") + num_warps = 32 if block_v >= 65536 else 16 if block_v >= 32768 else 8 - y_soft = torch.empty_like(logits_2d, dtype=torch.float32) - out = torch.empty_like(logits_2d, dtype=torch.float32) + y_soft = torch.empty_like(logits_2d) + out = torch.empty_like(logits_2d) if hard else y_soft _gumbel_softmax_fwd_kernel[(n_rows,)]( logits_2d, gumbels_2d, @@ -115,23 +116,59 @@ def forward(ctx, logits, gumbels, tau: float, hard: bool, seed: int): HAS_GUMBELS=gumbels is not None, HARD=bool(hard), BLOCK_V=block_v, + num_warps=num_warps, ) ctx.save_for_backward(y_soft) ctx.tau = float(tau) ctx.logits_shape = tuple(logits.shape) ctx.logits_dtype = logits.dtype - return out.view(*lead_shape, V).to(logits.dtype) + return out.view(*lead_shape, V) @staticmethod def backward(ctx, grad_output): (y_soft,) = ctx.saved_tensors - grad_2d = grad_output.contiguous().view_as(y_soft).to(torch.float32) - dot = (grad_2d * y_soft).sum(dim=-1, keepdim=True) - grad_logits = y_soft * (grad_2d - dot) / ctx.tau + grad_2d = grad_output.contiguous().view_as(y_soft) + grad_logits = torch._softmax_backward_data( + grad_2d, + y_soft, + -1, + ctx.logits_dtype, + ) + if ctx.tau != 1.0: + grad_logits = grad_logits / ctx.tau return grad_logits.view(ctx.logits_shape).to(ctx.logits_dtype), None, None, None, None +def _hard_gumbel_softmax_nograd( + logits: torch.Tensor, + gumbels: Optional[torch.Tensor], + seed: int, +) -> torch.Tensor: + V = logits.shape[-1] + lead_shape = logits.shape[:-1] + logits_2d = logits.contiguous().view(-1, V) + gumbels_2d = gumbels.contiguous().view(-1, V) if gumbels is not None else logits_2d + block_v = triton.next_power_of_2(V) + if block_v > _MAX_BLOCK_V: + raise ValueError(f"vocab size {V} exceeds TritonGumbelSoftmaxOp limit {_MAX_BLOCK_V}") + num_warps = 32 if block_v >= 65536 else 16 if block_v >= 32768 else 8 + + out = torch.empty_like(logits_2d) + _gumbel_softmax_hard_nograd_kernel[(logits_2d.shape[0],)]( + logits_2d, + gumbels_2d, + out, + int(seed), + logits_2d.shape[0], + V, + HAS_GUMBELS=gumbels is not None, + BLOCK_V=block_v, + num_warps=num_warps, + ) + return out.view(*lead_shape, V) + + class TritonGumbelSoftmaxOp: """Triton Gumbel-Softmax sampler with straight-through hard samples.""" @@ -163,4 +200,6 @@ def apply( ) if seed is None: seed = int(torch.randint(0, 2**31 - 1, (), device="cpu").item()) + if hard and (not torch.is_grad_enabled() or not logits.requires_grad): + return _hard_gumbel_softmax_nograd(logits, gumbels, int(seed)) return _GumbelSoftmaxFunction.apply(logits, gumbels, float(tau), bool(hard), int(seed)) diff --git a/tests/test_gumbel_softmax.py b/tests/test_gumbel_softmax.py index 20c9f66b..925b8755 100644 --- a/tests/test_gumbel_softmax.py +++ b/tests/test_gumbel_softmax.py @@ -83,14 +83,16 @@ def test_native_rejects_invalid_inputs(): op(torch.randn(8)) -def test_registry_dispatch_matches_native_cpu(): +def test_registry_dispatch_matches_native(): from rl_engine.kernels.registry import kernel_registry + from rl_engine.platforms.device import device_ctx - logits, gumbels = _inputs(5) + device = device_ctx.device if device_ctx.device_type == "cuda" else "cpu" + logits, gumbels = _inputs(5, device=device) op = kernel_registry.get_op("gumbel_softmax") out = op(logits, tau=0.7, hard=False, gumbels=gumbels) ref = NativeGumbelSoftmaxOp()(logits, tau=0.7, hard=False, gumbels=gumbels) - assert torch.allclose(out, ref, atol=1e-6) + assert torch.allclose(out.cpu(), ref.cpu(), atol=1e-5, rtol=1e-5) @requires_triton_cuda From cfa5333811ed57d060b3ab5871ef720e480aa737 Mon Sep 17 00:00:00 2001 From: JLiu4Coding Date: Wed, 8 Jul 2026 15:26:13 +0800 Subject: [PATCH 3/7] Optimize Triton Gumbel-Softmax sampler --- benchmarks/benchmark_gumbel_softmax.py | 15 ++++---- docs/operators/gumbel-softmax.md | 6 ++- .../ops/triton/sampling/gumbel_softmax.py | 37 +++++++++++-------- tests/test_gumbel_softmax.py | 19 ++++++++-- 4 files changed, 48 insertions(+), 29 deletions(-) diff --git a/benchmarks/benchmark_gumbel_softmax.py b/benchmarks/benchmark_gumbel_softmax.py index b86e9f4d..99b3b1ee 100644 --- a/benchmarks/benchmark_gumbel_softmax.py +++ b/benchmarks/benchmark_gumbel_softmax.py @@ -18,7 +18,6 @@ from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger - DEFAULT_CONFIGS = [ (1, 512, 32000), (4, 512, 32000), @@ -31,9 +30,7 @@ def _make_inputs(batch, seq, vocab, device, dtype, internal_gumbels): gumbels = None if not internal_gumbels: gumbels = ( - -torch.empty(batch, seq, vocab, device=device, dtype=torch.float32) - .exponential_() - .log() + -torch.empty(batch, seq, vocab, device=device, dtype=torch.float32).exponential_().log() ) return logits, gumbels @@ -83,19 +80,21 @@ def run_benchmark(args): for batch, seq, vocab in args.configs: logits, gumbels = _make_inputs(batch, seq, vocab, device, dtype, args.internal_gumbels) upstream = torch.randn_like(logits) + native_leaf = logits.detach().clone().requires_grad_(True) + triton_leaf = logits.detach().clone().requires_grad_(True) def fwd(op, x=logits, g=gumbels): with torch.no_grad(): op(x, tau=args.tau, hard=args.hard, gumbels=g) - def fwd_bwd(op, x_src=logits, g=gumbels, grad=upstream): - x = x_src.detach().clone().requires_grad_(True) + def fwd_bwd(op, x, g=gumbels, grad=upstream): + x.grad = None op(x, tau=args.tau, hard=args.hard, gumbels=g).backward(grad) n_fwd = _time_ms(lambda: fwd(native), args.warmup, args.iters) t_fwd = _time_ms(lambda: fwd(triton_op), args.warmup, args.iters) - n_fb = _time_ms(lambda: fwd_bwd(native), args.warmup, args.iters) - t_fb = _time_ms(lambda: fwd_bwd(triton_op), args.warmup, args.iters) + n_fb = _time_ms(lambda: fwd_bwd(native, native_leaf), args.warmup, args.iters) + t_fb = _time_ms(lambda: fwd_bwd(triton_op, triton_leaf), args.warmup, args.iters) n_vram = _peak_vram_mb(lambda: fwd(native)) t_vram = _peak_vram_mb(lambda: fwd(triton_op)) diff --git a/docs/operators/gumbel-softmax.md b/docs/operators/gumbel-softmax.md index 3490c13e..3ed784f4 100644 --- a/docs/operators/gumbel-softmax.md +++ b/docs/operators/gumbel-softmax.md @@ -39,6 +39,7 @@ softmax gradient. | `tau` | scalar | Python float | Must be positive. | | `hard` | scalar | Python bool | Enables straight-through one-hot output. | | `gumbels` | `[..., V]` or `None` | Floating point | Optional deterministic noise for tests/repro. | +| `seed` | scalar or `None` | Python int | Triton-only seed for backend-internal Gumbel noise when `gumbels=None`. | | Output | `[..., V]` | Same as `logits` | Probabilities or one-hot rows over vocab. | The Triton backend flattens leading dimensions to `[N, V]` and launches one @@ -49,10 +50,12 @@ program per row. It supports vocab sizes up to 131072 in this implementation. | Backend | Status | Notes | | --- | --- | --- | | PyTorch | Reference | Uses PyTorch autograd end to end. | -| Triton | GPU optimized forward | Uses Triton forward and PyTorch math in custom autograd backward. | +| Triton | GPU optimized | Uses Triton forward, a no-grad hard-sampling fast path, and fused softmax backward. | For deterministic correctness tests, pass the same precomputed `gumbels` tensor to both backends. If `gumbels=None`, each backend samples its own Gumbel noise. +The Triton backend accepts `seed` to make backend-internal Gumbel generation +reproducible for smoke tests and benchmarks. ## Tests and Benchmarks @@ -63,4 +66,3 @@ python benchmarks/benchmark_gumbel_softmax.py The benchmark reports forward latency, forward+backward latency, speedup, and peak forward VRAM for PyTorch and Triton on a single GPU. - diff --git a/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py b/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py index c0b2642a..2bab3659 100644 --- a/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py +++ b/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py @@ -9,13 +9,21 @@ import triton import triton.language as tl -from rl_engine.kernels.ops.pytorch.sampling.gumbel_softmax import ( - _validate_gumbel_softmax_inputs, -) +from rl_engine.kernels.ops.pytorch.sampling.gumbel_softmax import _validate_gumbel_softmax_inputs _MAX_BLOCK_V = 131072 +def _launch_config(vocab_size: int) -> tuple[int, int]: + block_v = triton.next_power_of_2(vocab_size) + if block_v > _MAX_BLOCK_V: + raise ValueError( + f"vocab size {vocab_size} exceeds TritonGumbelSoftmaxOp limit {_MAX_BLOCK_V}" + ) + num_warps = 32 if block_v >= 65536 else 16 if block_v >= 32768 else 8 + return block_v, num_warps + + @triton.jit def _gumbel_softmax_fwd_kernel( logits_ptr, @@ -51,7 +59,10 @@ def _gumbel_softmax_fwd_kernel( tl.store(y_soft_ptr + row_off + cols, y, mask=mask) if HARD: - out = tl.where(z == z_max, 1.0, 0.0) + # Tie-break to the first maximum so each row is exactly one-hot. + is_max = z == z_max + first_max = tl.min(tl.where(is_max, cols, BLOCK_V), axis=0) + out = tl.where(cols == first_max, 1.0, 0.0) tl.store(out_ptr + row_off + cols, out, mask=mask) else: tl.store(out_ptr + row_off + cols, y, mask=mask) @@ -83,7 +94,9 @@ def _gumbel_softmax_hard_nograd_kernel( z = tl.where(mask, logits + gumbels, -float("inf")) z_max = tl.max(z, axis=0) - out = tl.where(z == z_max, 1.0, 0.0) + is_max = z == z_max + first_max = tl.min(tl.where(is_max, cols, BLOCK_V), axis=0) + out = tl.where(cols == first_max, 1.0, 0.0) tl.store(out_ptr + row_off + cols, out, mask=mask) @@ -93,14 +106,9 @@ def forward(ctx, logits, gumbels, tau: float, hard: bool, seed: int): V = logits.shape[-1] lead_shape = logits.shape[:-1] logits_2d = logits.contiguous().view(-1, V) - gumbels_2d = ( - gumbels.contiguous().view(-1, V) if gumbels is not None else logits_2d - ) + gumbels_2d = gumbels.contiguous().view(-1, V) if gumbels is not None else logits_2d n_rows = logits_2d.shape[0] - block_v = triton.next_power_of_2(V) - if block_v > _MAX_BLOCK_V: - raise ValueError(f"vocab size {V} exceeds TritonGumbelSoftmaxOp limit {_MAX_BLOCK_V}") - num_warps = 32 if block_v >= 65536 else 16 if block_v >= 32768 else 8 + block_v, num_warps = _launch_config(V) y_soft = torch.empty_like(logits_2d) out = torch.empty_like(logits_2d) if hard else y_soft @@ -149,10 +157,7 @@ def _hard_gumbel_softmax_nograd( lead_shape = logits.shape[:-1] logits_2d = logits.contiguous().view(-1, V) gumbels_2d = gumbels.contiguous().view(-1, V) if gumbels is not None else logits_2d - block_v = triton.next_power_of_2(V) - if block_v > _MAX_BLOCK_V: - raise ValueError(f"vocab size {V} exceeds TritonGumbelSoftmaxOp limit {_MAX_BLOCK_V}") - num_warps = 32 if block_v >= 65536 else 16 if block_v >= 32768 else 8 + block_v, num_warps = _launch_config(V) out = torch.empty_like(logits_2d) _gumbel_softmax_hard_nograd_kernel[(logits_2d.shape[0],)]( diff --git a/tests/test_gumbel_softmax.py b/tests/test_gumbel_softmax.py index 925b8755..f7c3982e 100644 --- a/tests/test_gumbel_softmax.py +++ b/tests/test_gumbel_softmax.py @@ -23,9 +23,9 @@ def _inputs(seed, *, device="cpu", dtype=torch.float32, shape=(4, 17)): gen = torch.Generator(device=device).manual_seed(seed) logits = torch.randn(*shape, generator=gen, device=device, dtype=dtype) # Keep deterministic test noise in fp32, then let each backend cast as needed. - gumbels = -torch.empty(*shape, device=device, dtype=torch.float32).exponential_( - generator=gen - ).log() + gumbels = ( + -torch.empty(*shape, device=device, dtype=torch.float32).exponential_(generator=gen).log() + ) return logits, gumbels @@ -152,3 +152,16 @@ def test_triton_generated_noise_smoke(): assert out.shape == logits.shape assert torch.isfinite(out).all() assert torch.allclose(out.sum(dim=-1), torch.ones(8, device="cuda"), atol=1e-5) + + +@requires_triton_cuda +def test_triton_hard_nograd_generated_noise_fast_path(): + from rl_engine.kernels.ops.triton.sampling.gumbel_softmax import TritonGumbelSoftmaxOp + + logits = torch.randn(3, 5, 1024, device="cuda") + with torch.no_grad(): + out = TritonGumbelSoftmaxOp()(logits, tau=0.7, hard=True, seed=2027) + assert out.shape == logits.shape + assert out.dtype == logits.dtype + assert torch.allclose(out.sum(dim=-1), torch.ones(3, 5, device="cuda")) + assert torch.all((out == 0.0) | (out == 1.0)) From 64210be4a1b7f8cf233fb9e7ab31353d3e3ca29d Mon Sep 17 00:00:00 2001 From: JLiu4Coding Date: Mon, 13 Jul 2026 22:48:26 +0800 Subject: [PATCH 4/7] Fix Gumbel-Softmax lint formatting --- rl_engine/kernels/ops/pytorch/sampling/__init__.py | 1 - rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py | 4 +--- rl_engine/kernels/ops/triton/sampling/__init__.py | 1 - 3 files changed, 1 insertion(+), 5 deletions(-) diff --git a/rl_engine/kernels/ops/pytorch/sampling/__init__.py b/rl_engine/kernels/ops/pytorch/sampling/__init__.py index 4f7d09a0..86cf4c9d 100644 --- a/rl_engine/kernels/ops/pytorch/sampling/__init__.py +++ b/rl_engine/kernels/ops/pytorch/sampling/__init__.py @@ -1,3 +1,2 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors - diff --git a/rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py b/rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py index 443dc22d..4684bce1 100644 --- a/rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py +++ b/rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py @@ -38,9 +38,7 @@ def _validate_gumbel_softmax_inputs( def _sample_gumbels_like(logits: torch.Tensor) -> torch.Tensor: # Matches torch.nn.functional.gumbel_softmax's exponential sampling path. return ( - -torch.empty_like(logits, memory_format=torch.legacy_contiguous_format) - .exponential_() - .log() + -torch.empty_like(logits, memory_format=torch.legacy_contiguous_format).exponential_().log() ) diff --git a/rl_engine/kernels/ops/triton/sampling/__init__.py b/rl_engine/kernels/ops/triton/sampling/__init__.py index 4f7d09a0..86cf4c9d 100644 --- a/rl_engine/kernels/ops/triton/sampling/__init__.py +++ b/rl_engine/kernels/ops/triton/sampling/__init__.py @@ -1,3 +1,2 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026 RL-Kernel Contributors - From 4011e632f20277d65725d3ea4ac63c4d471f2877 Mon Sep 17 00:00:00 2001 From: JLiu4Coding Date: Tue, 21 Jul 2026 21:48:44 +0800 Subject: [PATCH 5/7] Address Gumbel-Softmax review feedback --- docs/operators/gumbel-softmax.md | 4 ++- .../ops/pytorch/sampling/gumbel_softmax.py | 4 ++- .../ops/triton/sampling/gumbel_softmax.py | 7 +++++- tests/test_gumbel_softmax.py | 25 ++++++++++++++++++- 4 files changed, 36 insertions(+), 4 deletions(-) diff --git a/docs/operators/gumbel-softmax.md b/docs/operators/gumbel-softmax.md index 3ed784f4..7462a01d 100644 --- a/docs/operators/gumbel-softmax.md +++ b/docs/operators/gumbel-softmax.md @@ -38,12 +38,14 @@ softmax gradient. | `logits` | `[..., V]` | fp32/fp16/bf16 | Floating-point tensor, `V` is vocab size. | | `tau` | scalar | Python float | Must be positive. | | `hard` | scalar | Python bool | Enables straight-through one-hot output. | -| `gumbels` | `[..., V]` or `None` | Floating point | Optional deterministic noise for tests/repro. | +| `gumbels` | `[..., V]` or `None` | Floating point | Optional deterministic fixed noise for tests/repro; gradients are not propagated to this tensor. | | `seed` | scalar or `None` | Python int | Triton-only seed for backend-internal Gumbel noise when `gumbels=None`. | | Output | `[..., V]` | Same as `logits` | Probabilities or one-hot rows over vocab. | The Triton backend flattens leading dimensions to `[N, V]` and launches one program per row. It supports vocab sizes up to 131072 in this implementation. +For larger vocab sizes, `TritonGumbelSoftmaxOp` falls back to the PyTorch +reference implementation instead of raising. ## Backends diff --git a/rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py b/rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py index 4684bce1..58ed9762 100644 --- a/rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py +++ b/rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py @@ -50,7 +50,9 @@ def gumbel_softmax_reference( gumbels: Optional[torch.Tensor] = None, ) -> torch.Tensor: _validate_gumbel_softmax_inputs(logits, tau, gumbels) - noise = _sample_gumbels_like(logits) if gumbels is None else gumbels.to(dtype=logits.dtype) + noise = ( + _sample_gumbels_like(logits) if gumbels is None else gumbels.detach().to(dtype=logits.dtype) + ) y_soft = torch.softmax((logits + noise) / tau, dim=-1) if not hard: return y_soft diff --git a/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py b/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py index 2bab3659..7e6e7bc6 100644 --- a/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py +++ b/rl_engine/kernels/ops/triton/sampling/gumbel_softmax.py @@ -9,7 +9,10 @@ import triton import triton.language as tl -from rl_engine.kernels.ops.pytorch.sampling.gumbel_softmax import _validate_gumbel_softmax_inputs +from rl_engine.kernels.ops.pytorch.sampling.gumbel_softmax import ( + NativeGumbelSoftmaxOp, + _validate_gumbel_softmax_inputs, +) _MAX_BLOCK_V = 131072 @@ -198,6 +201,8 @@ def apply( seed: Optional[int] = None, ) -> torch.Tensor: _validate_gumbel_softmax_inputs(logits, float(tau), gumbels) + if triton.next_power_of_2(logits.shape[-1]) > _MAX_BLOCK_V: + return NativeGumbelSoftmaxOp()(logits, tau=tau, hard=hard, gumbels=gumbels) if logits.device.type not in ("cuda", "xpu", "hip"): raise RuntimeError( "TritonGumbelSoftmaxOp requires a GPU tensor (CUDA / ROCm / XPU), got " diff --git a/tests/test_gumbel_softmax.py b/tests/test_gumbel_softmax.py index f7c3982e..615d925a 100644 --- a/tests/test_gumbel_softmax.py +++ b/tests/test_gumbel_softmax.py @@ -72,6 +72,18 @@ def test_native_temperature_changes_distribution_sharpness(): assert cool.max(dim=-1).values.mean() > warm.max(dim=-1).values.mean() +def test_native_treats_supplied_gumbels_as_fixed_noise(): + logits, gumbels = _inputs(11) + logits = logits.requires_grad_(True) + gumbels = gumbels.requires_grad_(True) + + out = NativeGumbelSoftmaxOp()(logits, tau=0.8, hard=False, gumbels=gumbels) + out[..., 0].sum().backward() + + assert logits.grad is not None + assert gumbels.grad is None + + def test_native_rejects_invalid_inputs(): op = NativeGumbelSoftmaxOp() logits, gumbels = _inputs(4) @@ -87,7 +99,7 @@ def test_registry_dispatch_matches_native(): from rl_engine.kernels.registry import kernel_registry from rl_engine.platforms.device import device_ctx - device = device_ctx.device if device_ctx.device_type == "cuda" else "cpu" + device = device_ctx.device if device_ctx.device_type == "cuda" or device_ctx.is_rocm else "cpu" logits, gumbels = _inputs(5, device=device) op = kernel_registry.get_op("gumbel_softmax") out = op(logits, tau=0.7, hard=False, gumbels=gumbels) @@ -95,6 +107,17 @@ def test_registry_dispatch_matches_native(): assert torch.allclose(out.cpu(), ref.cpu(), atol=1e-5, rtol=1e-5) +@pytest.mark.skipif(not _HAS_TRITON, reason="Triton import is required for fallback wrapper test.") +def test_triton_wrapper_falls_back_to_native_for_unsupported_vocab(monkeypatch): + from rl_engine.kernels.ops.triton.sampling import gumbel_softmax as triton_gs + + monkeypatch.setattr(triton_gs, "_MAX_BLOCK_V", 8) + logits, gumbels = _inputs(12, shape=(2, 9)) + out = triton_gs.TritonGumbelSoftmaxOp()(logits, tau=0.9, hard=True, gumbels=gumbels) + ref = NativeGumbelSoftmaxOp()(logits, tau=0.9, hard=True, gumbels=gumbels) + assert torch.allclose(out, ref, atol=0.0) + + @requires_triton_cuda def test_triton_forward_matches_native_soft_fp32(): from rl_engine.kernels.ops.triton.sampling.gumbel_softmax import TritonGumbelSoftmaxOp From c5e33ccc1ec24a2ba2186b98133e086d3581dd96 Mon Sep 17 00:00:00 2001 From: JLiu4Coding Date: Tue, 21 Jul 2026 23:21:18 +0800 Subject: [PATCH 6/7] Add GPU CI smoke script --- scripts/ci_smoke.py | 78 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 scripts/ci_smoke.py diff --git a/scripts/ci_smoke.py b/scripts/ci_smoke.py new file mode 100644 index 00000000..0cbe1cda --- /dev/null +++ b/scripts/ci_smoke.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""CI fail-fast smoke check for the compiled CUDA extension (``rl_engine._C``). + +Exits non-zero with a clear message when either failure mode from issue #191 +occurs, so GPU CI cannot pass while silently running only native fallbacks: + + * Bug A - the extension did not build at all (e.g. PEP 517 build isolation hid + torch from ``setup.py``), so ``from rl_engine import _C`` raises ImportError. + * Bug B - the extension built for the wrong GPU architecture; the import + succeeds (``dlopen`` does not check arch) but the first kernel launch raises + ``cudaErrorNoKernelImageForDevice`` once the stream is synchronized. + +Uses only ``fused_logp`` - the op registered unconditionally in ``csrc/ops.cpp`` - +so it does not require ``KERNEL_ALIGN_FORCE_SM90=1`` / a Hopper build. +""" +import sys + +import torch + + +def main() -> int: + if not torch.cuda.is_available(): + print("[smoke] FATAL: CUDA is not available in this CI environment", file=sys.stderr) + return 2 + + print(f"[smoke] torch: {torch.__version__} (cuda {torch.version.cuda})") + print(f"[smoke] device: {torch.cuda.get_device_name()}") + cc = torch.cuda.get_device_capability() + print(f"[smoke] capability: sm_{cc[0]}{cc[1]}") + + # (1) Import check -> catches Bug A (no .so was compiled at all). + try: + from rl_engine import _C + except ImportError as exc: + print( + "[smoke] FATAL: compiled extension rl_engine._C is missing - the CUDA " + "kernels were not built.\n" + " Likely PEP 517 build isolation hid torch from setup.py; install " + "with `pip install --no-build-isolation -e .`.\n" + f" Underlying error: {exc}", + file=sys.stderr, + ) + return 1 + print(f"[smoke] _C file: {getattr(_C, '__file__', None)}") + + # (2) Real launch + synchronize -> catches Bug B (arch mismatch only surfaces + # on launch, asynchronously, so the sync is required to raise it here). + try: + logits = torch.randn(4, 32, device="cuda", dtype=torch.float32) + token_ids = torch.randint(0, 32, (4,), device="cuda", dtype=torch.long) + out = _C.fused_logp(logits, token_ids) + torch.cuda.synchronize() + # Broad on purpose: an arch mismatch raises a CUDA RuntimeError (not ImportError), + # and any launch failure whatsoever must fail the smoke check loudly. + except Exception as exc: + print( + "[smoke] FATAL: rl_engine._C built but fused_logp failed to launch on " + f"sm_{cc[0]}{cc[1]}.\n" + " The extension was likely compiled for a different architecture; " + "set TORCH_CUDA_ARCH_LIST / TARGET_SM to match this GPU.\n" + f" Underlying error: {type(exc).__name__}: {exc}", + file=sys.stderr, + ) + return 1 + + if tuple(out.shape) != (4,): + print( + f"[smoke] FATAL: unexpected fused_logp output shape {tuple(out.shape)}", file=sys.stderr + ) + return 1 + + print(f"[smoke] OK: rl_engine._C built and fused_logp ran on sm_{cc[0]}{cc[1]}.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From fe8c56c8e8311505b6e7620f40a0b446c4e65075 Mon Sep 17 00:00:00 2001 From: JLiu4Coding Date: Tue, 21 Jul 2026 23:45:17 +0800 Subject: [PATCH 7/7] Silence CI smoke broad-exception lint --- scripts/ci_smoke.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/ci_smoke.py b/scripts/ci_smoke.py index 0cbe1cda..d42d31d3 100644 --- a/scripts/ci_smoke.py +++ b/scripts/ci_smoke.py @@ -53,7 +53,7 @@ def main() -> int: torch.cuda.synchronize() # Broad on purpose: an arch mismatch raises a CUDA RuntimeError (not ImportError), # and any launch failure whatsoever must fail the smoke check loudly. - except Exception as exc: + except Exception as exc: # noqa: BLE001 print( "[smoke] FATAL: rl_engine._C built but fused_logp failed to launch on " f"sm_{cc[0]}{cc[1]}.\n"