[FEAT][kernels] Fused Gumbel-Softmax sampler — Triton forward, ST hard path, fused backward

benchmarks/benchmark_gumbel_softmax.py

@@ -0,0 +1,159 @@
+# 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, internal_gumbels):
+    logits = torch.randn(batch, seq, vocab, device=device, dtype=dtype)
+    gumbels = None
+    if not internal_gumbels:
+        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}, "
+        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, 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, 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, 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))
+
+        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(
+        "--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,
+        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())

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

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)

docs/operators/gumbel-softmax.md

@@ -0,0 +1,68 @@
+# 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. |
+| `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.
+
+## Backends
+
+| Backend | Status | Notes |
+| --- | --- | --- |
+| PyTorch | Reference | Uses PyTorch autograd end to end. |
+| 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
+
+```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.

rl_engine/kernels/ops/pytorch/sampling/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors

rl_engine/kernels/ops/pytorch/sampling/gumbel_softmax.py

@@ -0,0 +1,84 @@
+# 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)

rl_engine/kernels/ops/triton/sampling/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# Copyright (c) 2026 RL-Kernel Contributors