From 2220177c87e9e02c8093d3715de09f02c70e2492 Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Sun, 14 Jun 2026 13:20:47 +0800 Subject: [PATCH 1/8] [FEAT][kernels]: add fused linear log-prob operator (Triton + chunked backward + SM90 TMA/WGMMA) Implements `linear_logp`: log_softmax(hidden @ W^T + b)[target] without materializing the [N, V] logits, differentiable w.r.t. hidden/weight/bias. - PyTorch `NativeLinearLogpOp`: naive F.linear + log_softmax + gather reference. - Triton `TritonLinearLogpOp`: online-softmax forward (zero [N,V] materialization, tensor cores via native-dtype tl.dot), Liger-style chunked backward (cuBLAS matmuls, sequential grad_weight accumulation -> deterministic, peak mem chunk*V). - CUDA SM90 `FusedLinearLogpSM90Op`: TMA + WGMMA streaming forward + smem online softmax (csrc/cuda/fused_linear_logp_sm90.cu); build-guarded behind KERNEL_ALIGN_FORCE_SM90, registry-gated to cc_major==9. Assembles for sm_90a under CUDA 13.1; numerics pending validation on Hopper. - Registry dispatch, tests, benchmark, operator + design docs. Co-Authored-By: Claude Opus 4.8 --- benchmarks/benchmark_linear_logp.py | 158 ++++++++++ csrc/cuda/fused_linear_logp_sm90.cu | 283 ++++++++++++++++++ csrc/ops.cpp | 6 + docs/.nav.yml | 1 + docs/design/fused-linear-logp.md | 216 +++++++++++++ docs/operators/README.md | 1 + docs/operators/linear-logp.md | 134 +++++++++ .../kernels/ops/cuda/loss/linear_logp.py | 120 ++++++++ .../kernels/ops/pytorch/loss/linear_logp.py | 58 ++++ rl_engine/kernels/ops/triton/loss/__init__.py | 2 + .../kernels/ops/triton/loss/linear_logp.py | 218 ++++++++++++++ rl_engine/kernels/registry.py | 18 ++ setup.py | 12 +- tests/test_linear_logp.py | 157 ++++++++++ 14 files changed, 1380 insertions(+), 4 deletions(-) create mode 100644 benchmarks/benchmark_linear_logp.py create mode 100644 csrc/cuda/fused_linear_logp_sm90.cu create mode 100644 docs/design/fused-linear-logp.md create mode 100644 docs/operators/linear-logp.md create mode 100644 rl_engine/kernels/ops/cuda/loss/linear_logp.py create mode 100644 rl_engine/kernels/ops/pytorch/loss/linear_logp.py create mode 100644 rl_engine/kernels/ops/triton/loss/__init__.py create mode 100644 rl_engine/kernels/ops/triton/loss/linear_logp.py create mode 100644 tests/test_linear_logp.py diff --git a/benchmarks/benchmark_linear_logp.py b/benchmarks/benchmark_linear_logp.py new file mode 100644 index 00000000..f3c863d1 --- /dev/null +++ b/benchmarks/benchmark_linear_logp.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Benchmark NativeLinearLogpOp vs TritonLinearLogpOp. + +The native op materializes the full [N, V] logits (a single F.linear + +log_softmax + gather); the Triton op streams the vocab through an online softmax +and never lands the logits, so its peak forward memory is independent of V. This +is the headline number: peak VRAM vs the materializing baseline, swept over +vocab size. Latency (forward and forward+backward) is reported too -- the win is +memory, not FLOPs (the backward recomputes the logit tiles). + +Usage: + python benchmarks/benchmark_linear_logp.py + python benchmarks/benchmark_linear_logp.py --configs "4096,2048,32768;4096,2048,131072" +""" + +import argparse + +import torch +from tabulate import tabulate + +from rl_engine.kernels.ops.pytorch.loss.linear_logp import NativeLinearLogpOp +from rl_engine.kernels.ops.triton.loss.linear_logp import TritonLinearLogpOp +from rl_engine.platforms.device import device_ctx +from rl_engine.utils.logger import logger + +# (num_tokens, hidden_dim, vocab) +DEFAULT_CONFIGS = [ + (4096, 2048, 32768), + (4096, 2048, 50257), + (4096, 2048, 131072), +] + + +def _make_inputs(num_tokens, hidden_dim, vocab, device, dtype): + hidden = torch.randn(num_tokens, hidden_dim, device=device, dtype=dtype) + weight = torch.randn(vocab, hidden_dim, device=device, dtype=dtype) + target = torch.randint(0, vocab, (num_tokens,), device=device) + return hidden, weight, target + + +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_gb(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**3) + + +def run_benchmark(args): + if device_ctx.device_type != "cuda": + raise RuntimeError("linear_logp benchmark requires a CUDA device (Triton op is CUDA-only).") + + device = device_ctx.device + dtype = torch.bfloat16 + native = NativeLinearLogpOp() + triton_op = TritonLinearLogpOp() + + logger.info(f"linear_logp benchmark on {device} (dtype={dtype})") + + rows = [] + for num_tokens, hidden_dim, vocab in args.configs: + hidden, weight, target = _make_inputs(num_tokens, hidden_dim, vocab, device, dtype) + + def native_fwd(h=hidden, w=weight): + with torch.no_grad(): + native(h, w, target) + + def triton_fwd(h=hidden, w=weight): + with torch.no_grad(): + triton_op(h, w, target) + + def native_fwd_bwd(): + h = hidden.clone().requires_grad_(True) + w = weight.clone().requires_grad_(True) + native(h, w, target).sum().backward() + + def triton_fwd_bwd(): + h = hidden.clone().requires_grad_(True) + w = weight.clone().requires_grad_(True) + triton_op(h, w, target).sum().backward() + + n_fwd = _time_ms(native_fwd, args.warmup, args.iters) + t_fwd = _time_ms(triton_fwd, args.warmup, args.iters) + n_fb = _time_ms(native_fwd_bwd, args.warmup, args.iters) + t_fb = _time_ms(triton_fwd_bwd, args.warmup, args.iters) + n_vram = _peak_vram_gb(native_fwd) + t_vram = _peak_vram_gb(triton_fwd) + + rows.append( + [ + f"{num_tokens}x{hidden_dim}x{vocab}", + 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*1024:.0f}", + f"{t_vram*1024:.0f}", + ] + ) + + headers = [ + "shape (N x H x V)", + "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( + "--configs", + type=str, + default=None, + help="Semicolon-separated 'tokens,hidden,vocab' tuples, " + "e.g. '4096,2048,32768;4096,2048,131072'.", + ) + 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/csrc/cuda/fused_linear_logp_sm90.cu b/csrc/cuda/fused_linear_logp_sm90.cu new file mode 100644 index 00000000..72c11ae4 --- /dev/null +++ b/csrc/cuda/fused_linear_logp_sm90.cu @@ -0,0 +1,283 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +// +// Hopper (SM90) fused linear log-prob: +// +// logp[n] = log_softmax(hidden[n] @ W^T + b)[target[n]] +// +// computed WITHOUT materializing the [N, V] logits. Each CTA owns a block of BM +// tokens and streams the vocabulary in BN-wide tiles: for every vocab tile it +// 1. computes the logit tile logits[BM, BN] = H[BM, D] @ W[vtile, D]^T with +// WGMMA (K-loop over D), staging H/W tiles into shared memory via TMA, +// 2. stages the [BM, BN] fp32 logit tile into shared memory, and +// 3. folds it into a per-row online-softmax state (running max + sum) kept in +// shared memory, capturing the target logit when its column lands in tile. +// After the vocab loop it writes logp = z_target - (max + log(sum)) and the row +// log-sum-exp (lse) for the backward pass. +// +// ========================================================================= +// VALIDATION REQUIRED ON SM90. +// This file is written without an SM90 compiler/GPU in the loop. The numerically +// load-bearing, layout-sensitive pieces are: +// (a) the WGMMA shared-memory matrix descriptors (LBO/SBO/swizzle), and +// (b) the mapping from WGMMA fp32 accumulator registers -> (row, col). +// Both follow the PTX ISA "Asynchronous Warpgroup Level Matrix" section, but +// must be checked on hardware. The matmul is isolated in wgmma_m64n64k16 / +// store_acc_to_smem so it can be swapped for an mma.sync (m16n8k16) path if +// WGMMA needs debugging -- the surrounding TMA streaming + online softmax are +// independent of that choice. +// ========================================================================= + +#include "../utils/tma_utils.cuh" +#include +#include // CUDART_INF_F (not pulled in transitively under CUDA 13) +#include + +namespace { + +constexpr int BM = 64; // tokens per CTA == WGMMA M +constexpr int BN = 64; // vocab per tile == WGMMA N +constexpr int BK = 16; // contraction per WGMMA step == WGMMA K (bf16) +constexpr int WG_THREADS = 128; // one warpgroup (4 warps) +constexpr int ACC_REGS = (BM * BN) / WG_THREADS; // = 32 fp32 accumulators / thread + +// -------------------------------------------------------------------------- +// WGMMA helpers (SM90a). See PTX ISA: wgmma.mma_async / shared-memory matrix +// descriptor. Operands are streamed unswizzled (swizzle = 0); the descriptor +// encodes the smem start address plus the leading/stride byte offsets between +// 8x(BK) "core matrices". +// -------------------------------------------------------------------------- +__device__ __forceinline__ uint64_t make_smem_desc(const void *smem_ptr, uint32_t lbo, + uint32_t sbo) { + uint32_t addr = static_cast(__cvta_generic_to_shared(smem_ptr)); + uint64_t desc = 0; + desc |= (static_cast(addr >> 4) & 0x3FFFULL); // bits 0-13: start >> 4 + desc |= (static_cast(lbo >> 4) & 0x3FFFULL) << 16; // bits 16-29: LBO >> 4 + desc |= (static_cast(sbo >> 4) & 0x3FFFULL) << 32; // bits 32-45: SBO >> 4 + desc |= (static_cast(0) << 62); // bits 62-63: swizzle = none + return desc; +} + +__device__ __forceinline__ void wgmma_fence() { asm volatile("wgmma.fence.sync.aligned;"); } +__device__ __forceinline__ void wgmma_commit() { + asm volatile("wgmma.commit_group.sync.aligned;"); +} +template __device__ __forceinline__ void wgmma_wait() { + asm volatile("wgmma.wait_group.sync.aligned %0;" ::"n"(N)); +} + +// d += A @ B^T for one K=16 step. desc_a/desc_b are shared-memory descriptors. +// scale_d == 0 zeroes the accumulator first; 1 accumulates onto it. +__device__ __forceinline__ void wgmma_m64n64k16(float d[ACC_REGS], uint64_t desc_a, + uint64_t desc_b, int scale_d) { + // PTX form: {d}, desc_a, desc_b, scale-d (predicate), imm-scale-a, imm-scale-b, + // imm-trans-a, imm-trans-b. scale-d is a predicate set from scale_d (0 = write, + // 1 = accumulate). scaleA/scaleB = 1 (no negate); transA/transB = 0. + asm volatile( + "{\n" + ".reg .pred p;\n" + "setp.ne.b32 p, %34, 0;\n" + "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " + "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15," + "%16,%17,%18,%19,%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31}, " + "%32, %33, p, 1, 1, 0, 0;\n" + "}\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]), "+f"(d[4]), "+f"(d[5]), "+f"(d[6]), + "+f"(d[7]), "+f"(d[8]), "+f"(d[9]), "+f"(d[10]), "+f"(d[11]), "+f"(d[12]), "+f"(d[13]), + "+f"(d[14]), "+f"(d[15]), "+f"(d[16]), "+f"(d[17]), "+f"(d[18]), "+f"(d[19]), "+f"(d[20]), + "+f"(d[21]), "+f"(d[22]), "+f"(d[23]), "+f"(d[24]), "+f"(d[25]), "+f"(d[26]), "+f"(d[27]), + "+f"(d[28]), "+f"(d[29]), "+f"(d[30]), "+f"(d[31]) + : "l"(desc_a), "l"(desc_b), "r"(scale_d)); +} + +// Scatter the WGMMA fp32 accumulator to a row-major [BM][BN] shared-memory tile. +// m64nN .f32 layout (per warpgroup of 4 warps): warp w owns rows [16w, 16w+16); +// within a warp the N columns are tiled in 8-wide blocks, 4 regs per block: +// reg[4j+0]->(lane/4, 8j+(lane%4)*2 ) reg[4j+1]->(lane/4, 8j+(lane%4)*2+1) +// reg[4j+2]->(lane/4+8, 8j+(lane%4)*2 ) reg[4j+3]->(lane/4+8, 8j+(lane%4)*2+1) +__device__ __forceinline__ void store_acc_to_smem(const float d[ACC_REGS], float *s_logits) { + const int tid = threadIdx.x; + const int warp = tid / 32; + const int lane = tid % 32; + const int r0 = warp * 16 + lane / 4; + const int r1 = r0 + 8; + const int col = (lane % 4) * 2; +#pragma unroll + for (int j = 0; j < BN / 8; ++j) { + const int c = j * 8 + col; + s_logits[r0 * BN + c + 0] = d[4 * j + 0]; + s_logits[r0 * BN + c + 1] = d[4 * j + 1]; + s_logits[r1 * BN + c + 0] = d[4 * j + 2]; + s_logits[r1 * BN + c + 1] = d[4 * j + 3]; + } +} + +// -------------------------------------------------------------------------- +// Kernel: one CTA (one warpgroup) per BM-token block. +// Shared memory layout (one buffer; double-buffering of A/B is a follow-up): +// [ A tile BM*BK bf16 ][ B tile BN*BK bf16 ][ logits BM*BN f32 ] +// [ row_max BM f32 ][ row_sum BM f32 ][ row_zt BM f32 ][ tma mbar 8B ] +// -------------------------------------------------------------------------- +__global__ void fused_linear_logp_sm90_kernel(const __grid_constant__ CUtensorMap h_tmap, + const __grid_constant__ CUtensorMap w_tmap, + const int *__restrict__ target, + const float *__restrict__ bias, // may be null + float *__restrict__ out_logp, + float *__restrict__ out_lse, int N, int D, int V) { + const int tid = threadIdx.x; + const int row_block = blockIdx.x; + const int row_base = row_block * BM; + const int num_rows = min(BM, N - row_base); + const int kd = D / BK; // assumes D % BK == 0 (caller pads/validates) + + extern __shared__ __align__(1024) char smem[]; + nv_bfloat16 *sA = reinterpret_cast(smem); + nv_bfloat16 *sB = reinterpret_cast(sA + BM * BK); + float *sLogits = reinterpret_cast(sB + BN * BK); + float *sMax = sLogits + BM * BN; + float *sSum = sMax + BM; + float *sZt = sSum + BM; + const int tma_mbar = static_cast(__cvta_generic_to_shared(sZt + BM)); + + if (tid < num_rows) { + sMax[tid] = -CUDART_INF_F; + sSum[tid] = 0.0f; + sZt[tid] = 0.0f; + } + if (tid == 0) { + mbarrier_init(tma_mbar, 1); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + __syncthreads(); + + const int a_smem = static_cast(__cvta_generic_to_shared(sA)); + const int b_smem = static_cast(__cvta_generic_to_shared(sB)); + const uint32_t tile_bytes = (BM * BK + BN * BK) * sizeof(nv_bfloat16); + int phase = 0; + + const int num_vtiles = (V + BN - 1) / BN; + for (int vt = 0; vt < num_vtiles; ++vt) { + const int col_base = vt * BN; + + float d[ACC_REGS]; +#pragma unroll + for (int i = 0; i < ACC_REGS; ++i) + d[i] = 0.0f; + + // K-loop over the hidden dimension: stream H[BM,BK] and W[BN,BK] via TMA, + // then issue one WGMMA per K step accumulating into d[]. + for (int k = 0; k < kd; ++k) { + const int k_off = k * BK; + if (tid == 0) { + tma_2d_g2s(a_smem, &h_tmap, k_off, row_base, tma_mbar); + tma_2d_g2s(b_smem, &w_tmap, k_off, col_base, tma_mbar); + mbarrier_arrive_expect_tx(tma_mbar, tile_bytes); + } + mbarrier_wait(tma_mbar, phase); + phase ^= 1; + + wgmma_fence(); + const uint64_t desc_a = make_smem_desc(sA, BK * sizeof(nv_bfloat16), 0); + const uint64_t desc_b = make_smem_desc(sB, BK * sizeof(nv_bfloat16), 0); + wgmma_m64n64k16(d, desc_a, desc_b, 1); + wgmma_commit(); + wgmma_wait<0>(); + __syncthreads(); + } + + store_acc_to_smem(d, sLogits); + __syncthreads(); + + // Online softmax: one thread per row folds this tile's BN columns into + // the running (max, sum) and captures the target logit if present. + if (tid < num_rows) { + const int r = tid; + const int tgt = target[row_base + r]; + float tmax = -CUDART_INF_F; + for (int c = 0; c < BN; ++c) { + const int col = col_base + c; + if (col >= V) + break; + float val = sLogits[r * BN + c]; + if (bias != nullptr) + val += bias[col]; + tmax = fmaxf(tmax, val); + if (col == tgt) + sZt[r] = val; + } + float tsum = 0.0f; + for (int c = 0; c < BN; ++c) { + const int col = col_base + c; + if (col >= V) + break; + float val = sLogits[r * BN + c]; + if (bias != nullptr) + val += bias[col]; + tsum += __expf(val - tmax); + } + float old_max = sMax[r]; + float new_max = fmaxf(old_max, tmax); + sSum[r] = sSum[r] * __expf(old_max - new_max) + tsum * __expf(tmax - new_max); + sMax[r] = new_max; + } + __syncthreads(); + } + + if (tid < num_rows) { + const int r = tid; + const float lse = sMax[r] + logf(sSum[r]); + out_logp[row_base + r] = sZt[r] - lse; + out_lse[row_base + r] = lse; + } +} + +} // namespace + +// Forward: hidden [N, D] bf16, weight [V, D] bf16, target [N] int32, optional +// bias [V] f32. Returns (logp [N] f32, lse [N] f32). Logits are never +// materialized; peak extra memory is the per-CTA shared-memory tiles. +std::vector fused_linear_logp_sm90_forward(torch::Tensor hidden, + torch::Tensor weight, + torch::Tensor target, + torch::optional bias) { + TORCH_CHECK(hidden.is_cuda() && weight.is_cuda(), "inputs must be CUDA tensors"); + TORCH_CHECK(hidden.scalar_type() == at::kBFloat16, "hidden must be bfloat16"); + TORCH_CHECK(weight.scalar_type() == at::kBFloat16, "weight must be bfloat16"); + TORCH_CHECK(hidden.is_contiguous() && weight.is_contiguous(), "inputs must be contiguous"); + const int N = hidden.size(0); + const int D = hidden.size(1); + const int V = weight.size(0); + TORCH_CHECK(weight.size(1) == D, "hidden/weight hidden-dim mismatch"); + TORCH_CHECK(D % BK == 0, "D must be a multiple of ", BK, " for the SM90 kernel"); + + auto opts_f = hidden.options().dtype(torch::kFloat); + auto logp = torch::empty({N}, opts_f); + auto lse = torch::empty({N}, opts_f); + + // TMA descriptors: box [rows=BM/BN, cols=BK]. cols*sizeof(bf16) = 32B; the + // helper selects 32B swizzle for this stride -- the descriptor swizzle bits + // in make_smem_desc must be kept consistent (see VALIDATION note above). + CUtensorMap h_tmap, w_tmap; + init_tensor_map(&h_tmap, reinterpret_cast(hidden.data_ptr()), + N, D, BM, BK); + init_tensor_map(&w_tmap, reinterpret_cast(weight.data_ptr()), + V, D, BN, BK); + + const float *bias_ptr = nullptr; + torch::Tensor bias_f; + if (bias.has_value()) { + bias_f = bias->to(torch::kFloat).contiguous(); + bias_ptr = bias_f.data_ptr(); + } + + const int smem = (BM * BK + BN * BK) * sizeof(nv_bfloat16) + (BM * BN) * sizeof(float) + + 3 * BM * sizeof(float) + 16; + const int grid = (N + BM - 1) / BM; + auto target_i = target.to(torch::kInt32).contiguous(); + + fused_linear_logp_sm90_kernel<<>>( + h_tmap, w_tmap, target_i.data_ptr(), bias_ptr, logp.data_ptr(), + lse.data_ptr(), N, D, V); + + return {logp, lse}; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index c241bf67..f48e4f95 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -9,6 +9,10 @@ torch::Tensor fused_logp_forward(torch::Tensor logits, torch::Tensor token_ids); #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_SM90) torch::Tensor fused_logp_sm90_forward(torch::Tensor logits, torch::Tensor labels); +std::vector fused_linear_logp_sm90_forward(torch::Tensor hidden, + torch::Tensor weight, + torch::Tensor target, + torch::optional bias); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) @@ -75,6 +79,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_SM90) m.def("fused_logp_sm90", &fused_logp_sm90_forward, "TMA-accelerated Online Softmax Fused LogP"); + m.def("fused_linear_logp_sm90", &fused_linear_logp_sm90_forward, + "TMA+WGMMA fused linear log-prob (hidden @ W^T -> selected-token logp), SM90"); #endif #if defined(__CUDACC__) || defined(KERNEL_ALIGN_WITH_CUDA) diff --git a/docs/.nav.yml b/docs/.nav.yml index 3d9b6291..be61fa3e 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -10,6 +10,7 @@ nav: - Operators: - operators/README.md - operators/fused-logp.md + - operators/linear-logp.md - operators/sampling.md - Developer Guide: - contributing/README.md diff --git a/docs/design/fused-linear-logp.md b/docs/design/fused-linear-logp.md new file mode 100644 index 00000000..d3f66204 --- /dev/null +++ b/docs/design/fused-linear-logp.md @@ -0,0 +1,216 @@ +# Design: Fused Linear LogProb (no logit materialization) + fused backward + +## Motivation + +Every current log-prob path — `logp`, `ratio_kl`, `grpo_loss` — takes the +**logits** `[B, S, V]` as input, so the `[B, S, V]` tensor is already resident in +HBM before any fusion happens. For large-vocabulary models that tensor dominates +memory: at `V≈150k`, `B·S≈8k` in bf16 it is ~2.4 GB for the forward activation +alone, plus the same again for its gradient. + +This op fuses the **LM-head projection itself** into the log-prob reduction: + +``` +logits = hidden @ Wᵀ (+ b) # [N, V] — NEVER materialized +logp[n] = logits[n, t[n]] − logsumexp(logits[n, :]) +``` + +The forward streams the vocab in blocks (block matmul → online softmax), keeping +only per-row scalars. The backward recomputes the logit tiles from `hidden`/`W` +instead of storing them — trading a second matmul for the `[N, V]` storage. The +freed memory buys larger batches / longer CoT. + +This is the "cut cross-entropy" / fused-linear-cross-entropy pattern, specialized +to **selected-token log-prob** (the RL training quantity) rather than mean CE. + +## Scope + +- **New op type** `linear_logp` — does **not** replace `logp` (still useful when + logits already exist). It is the differentiable hidden→logp primitive. +- **In scope**: forward + backward w.r.t. `hidden` and `lm_head_weight` (and bias); + Triton, CUDA generic, CUDA SM90 TMA, ROCm backends; native PyTorch reference; + registry wiring; tests; benchmarks; docs. +- **Out of scope (follow-ups)**: rewiring `ratio_kl`/`grpo_loss` to consume hidden + states directly; tensor/sequence-parallel vocab sharding; fp8 weights. + +## Math + +Per row `n`, target token `t = t[n]`, logits `z = Hₙ Wᵀ (+ b) ∈ ℝ^V`: + +**Forward** +``` +lse = logsumexp(z); logp = z[t] − lse +``` +Computed by streaming vocab blocks `Vblk`: each block does a `[Nblk, D]·[D, Vblk]` +matmul tile, folds the tile into an online-softmax state `(m, s)` (running max, +rescaled sum), and captures `z[t]` when `t` lands in the block. Saved for backward: +**only** `lse` (or `(m,s)`) and the gathered `z[t]` — per-row scalars, no `[N,V]`. + +**Backward** — given `g = dL/dlogp ∈ ℝ^N`: +``` +dL/dz[v] = g · (1[v==t] − p[v]), p = softmax(z) # never stored; p recomputed +dL/dHₙ = g · (W[t] − Σ_v p[v]·W[v]) # = (dL/dz) @ W +dL/dW[v] += g · (1[v==t] − p[v]) · Hₙ (reduce over n) +dL/db[v] += g · (1[v==t] − p[v]) (reduce over n) +``` +The backward re-streams vocab blocks, recomputes `z_blk = Hₙ W_blkᵀ`, then +`p_blk = exp(z_blk − lse)`, and accumulates the three gradients. `dL/dW` and +`dL/db` are reductions **over the token dimension** → handled by atomic-add into +the `[V,D]`/`[V]` grad buffers (or a token-blocked two-pass to avoid atomics). + +All reductions/accumulation in **fp32**; inputs may be bf16/fp16. The `c·(onehot − +softmax)` shape is identical to the existing `_ratio_kl_bwd_kernel` — the only new +piece is the surrounding matmul. + +## Public API + +```python +linear_logp = kernel_registry.get_op("linear_logp") + +logp = linear_logp( + hidden, # [B, S, D] or [N, D] (differentiable) + lm_head_weight,# [V, D] (differentiable) + target_ids, # [B, S] or [N] int + bias=None, # [V] optional (differentiable) +) # -> [B, S] or [N], differentiable w.r.t. hidden, weight, bias +``` + +Wrapped in a `torch.autograd.Function` so it is a drop-in differentiable node. +Mask handling follows the existing convention (caller zeroes masked positions, or +pass an ignore-index for `target_ids`). + +## Tensor contract + +| Arg | Shape | Dtype | Notes | +| --- | --- | --- | --- | +| `hidden` | `[N, D]` | bf16/fp16/fp32 | contiguous; leading dims flattened | +| `lm_head_weight` | `[V, D]` | bf16/fp16/fp32 | contiguous (row-major over V) | +| `target_ids` | `[N]` | int32/int64 | in `[0, V)`; optional ignore-index | +| `bias` | `[V]` | float | optional | +| `logp` (out) | `[N]` | fp32 | `z[t] − logsumexp(z)` | + +## Implementation phases + +Ordering (per request): **Triton → CUDA generic → CUDA SM90 TMA → ROCm.** Each +phase is independently shippable and gated on matching the prior phase's reference +to tolerance. + +### Phase 1 — Triton (portable baseline + tolerance target) + +The semantic reference that CUDA/ROCm must match. + +- `rl_engine/kernels/ops/pytorch/loss/linear_logp.py` — `NativeLinearLogpOp`: + **naive** pure-PyTorch reference — a single `F.linear` (materializes the full + `[N, V]` logits) → `log_softmax` → `gather`, differentiable via autograd. No + chunking: this is the straightforward, obviously-correct oracle that the fused + kernels are validated against (and the baseline the benchmark measures the VRAM + win against). Also the CPU / Triton-less fallback. +- `rl_engine/kernels/ops/triton/loss/linear_logp.py`: + - `_linear_logp_fwd_kernel`: program per token-block; loop vocab blocks, `tl.dot` + for the `[Nblk,D]·[D,Vblk]` tile, online-softmax state in registers; store + `logp`, `lse`, gathered `z[t]`. + - `_linear_logp_bwd_kernel`: re-stream vocab, recompute `p_blk`, accumulate + `dH` (per token-block, local) and `dW`/`db` (atomic-add into global buffers). + - `_LinearLogpFunction(autograd.Function)` + `TritonLinearLogpOp`. +- Tests `tests/test_linear_logp.py`: native-vs-`F.linear→log_softmax→gather`; + Triton fwd vs native; Triton bwd vs autograd (gradcheck-style on `hidden`, + `weight`, `bias`) at `atol=1e-3` bf16 / tighter fp32; ignore-index; large `V` + (50257, 131072) memory-flat check; registry dispatch. +- Benchmark `benchmarks/benchmark_linear_logp.py`: sweep `V`; report fwd / fwd+bwd + latency and **peak VRAM vs the `F.linear`+`log_softmax` baseline** (the headline + number — should be flat in `V`). +- Docs `docs/operators/linear-logp.md` (+ `.nav.yml`, `operators/README.md`). + +**Acceptance**: Triton matches native within tolerance; benchmark shows peak VRAM +independent of `V` and a clear win vs the materializing baseline. + +### Phase 2 — CUDA generic fallback (`fused_linear_logp_kernel.cu`) + +Portable CUDA (no arch intrinsics) for all SM≥70. + +- `csrc/fused_linear_logp_kernel.cu`: forward + backward. Block per token-tile; + `cp.async` double-buffer `W` tiles into shared memory; FMA or `mma.sync` matmul; + online-softmax state in registers (reuse `LogSumExpState` / `merge_logsumexp_state` + and `blockReduce*` from `fused_logp_kernel.cu`); fp32 accumulate. Backward + recomputes tiles; atomic-add into `dW`/`db`. +- Bind in `csrc/ops.cpp` under the existing `KERNEL_ALIGN_WITH_CUDA` guard: + `fused_linear_logp_forward`, `fused_linear_logp_backward`. +- `setup.py`: add `csrc/fused_linear_logp_kernel.cu` to `cuda_sources`; env-var + block-size toggles mirroring the `FUSED_LOGP_*` macros. +- `rl_engine/kernels/ops/cuda/loss/linear_logp.py` — `FusedLinearLogpGenericOp` + wrapping fwd/bwd in an `autograd.Function`. + +**Acceptance**: matches Triton/native to tolerance; ≥ parity with Triton latency +on the generic path; same flat-VRAM property. + +### Phase 3 — CUDA SM90 TMA/WGMMA (`fused_linear_logp_sm90.cu`) + +Advanced streaming for Hopper/Blackwell-class (`cc ∈ {9,10,12}`). + +- `csrc/cuda/fused_linear_logp_sm90.cu`: TMA (`cuTensorMapEncodeTiled`) bulk loads + of `H` and `W` tiles into shared memory; WGMMA (`wgmma.mma_async`) for the + matmul; warp-specialized producer (TMA) / consumer (online softmax) with + `mbarrier`; keep state in registers. Backward symmetric with tile recompute. + **Heed the prior TMA lessons** (see issue #91 / the SM120 work): box inner dim + ≤ 256 elems, CUDA 12.9+ for `cuda::maximum`, arch `compute_90a`/`120a`. +- Gate behind `KERNEL_ALIGN_FORCE_SM90` in `setup.py` (arch-specific gencode), bind + under `KERNEL_ALIGN_WITH_SM90` in `ops.cpp`. +- `FusedLinearLogpSM90Op` in the CUDA wrapper; registry auto-prioritizes it in + `_adjust_priority_for_hardware` when `hasattr(_C, "fused_linear_logp_sm90")` and + `cc_major in (9,10,12)` (same pattern as `fused_logp_sm90`). + +**Acceptance**: matches reference to tolerance on SM90+; measurable speedup over +the generic CUDA path; builds cleanly on SM120/CUDA 13 (the dev box) and is cleanly +disabled where unsupported. + +### Phase 4 — ROCm / CDNA + +- **Free coverage first**: the Phase-1 Triton kernel compiles on ROCm — register + `TRITON_LINEAR_LOGP` in the `rocm` priority map so ROCm works from Phase 1. +- **Native HIP** (`csrc/rocm/...` or HIP-ified): wavefront=64-aware reductions and + LDS layouts; replace TMA with **manual double-buffering** into LDS; CDNA MFMA for + the matmul where available. Build via the ROCm/HIP path in `setup.py`. + +**Acceptance**: ROCm Triton matches reference; native HIP ≥ Triton on CDNA. + +## Registry wiring + +Add to `OpBackend`: `TRITON_LINEAR_LOGP`, `PYTORCH_LINEAR_LOGP`, +`CUDA_FUSED_LINEAR_LOGP_GENERIC`, `CUDA_FUSED_LINEAR_LOGP_SM90`. Add `linear_logp` +to each platform map: + +``` +cuda: [CUDA_FUSED_LINEAR_LOGP_GENERIC, TRITON_LINEAR_LOGP, PYTORCH_LINEAR_LOGP] + (+ CUDA_FUSED_LINEAR_LOGP_SM90 inserted at front by _adjust_priority_for_hardware) +rocm: [TRITON_LINEAR_LOGP, PYTORCH_LINEAR_LOGP] (+ native HIP once landed) +cpu: [PYTORCH_LINEAR_LOGP] +``` + +## Numerics & testing strategy + +- The **Triton kernel is the tolerance target**; native PyTorch is the correctness + oracle. Every native backend is gradchecked against `F.linear → log_softmax → + gather` autograd. +- fp32 accumulation throughout; bf16/fp16 I/O. Tolerance `atol≈1e-3` (bf16), + tighter for fp32. +- Edge cases: `V` not a multiple of the block; ignore-index targets; single-token + rows; very large `V` (memory-flat assertion). + +## Risks / open questions + +- **`dW` atomics**: atomic-add into `[V,D]` may contend at small `V`. Fallback is a + token-blocked two-pass (deterministic, atomic-free) — decide per-backend by + benchmark. Determinism flag may be wanted for reproducible RL runs. +- **bf16 weight gradient precision**: accumulate `dW` in fp32, cast on store. +- **Backward recompute cost**: ~2× forward matmul FLOPs in backward; net win is the + memory, not FLOPs — benchmark must report both so the tradeoff is explicit. +- **TMA portability**: the SM90 path repeats the constraints that bit us before + (box dims, CUDA version, arch flags); keep it default-off and feature-detected. + +## Relevant files (to add / touch) + +- `rl_engine/kernels/ops/{pytorch,triton,cuda}/loss/linear_logp.py` +- `csrc/fused_linear_logp_kernel.cu`, `csrc/cuda/fused_linear_logp_sm90.cu` +- `csrc/ops.cpp`, `setup.py`, `rl_engine/kernels/registry.py` +- `tests/test_linear_logp.py`, `benchmarks/benchmark_linear_logp.py` +- `docs/operators/linear-logp.md`, `docs/.nav.yml`, `docs/operators/README.md` diff --git a/docs/operators/README.md b/docs/operators/README.md index 49ab34b3..4188c5bc 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -19,6 +19,7 @@ Every operator page should include: ## Current Pages - [Fused LogP](fused-logp.md) +- [Fused Linear LogP](linear-logp.md) - [GRPO Loss](grpo-loss.md) - [Sampling](sampling.md) - [Operator Doc Template](../contributing/operator-doc-template.md) diff --git a/docs/operators/linear-logp.md b/docs/operators/linear-logp.md new file mode 100644 index 00000000..1e26c23b --- /dev/null +++ b/docs/operators/linear-logp.md @@ -0,0 +1,134 @@ +# Fused Linear LogP + +Fused Linear LogP computes per-token selected log-probabilities directly from +**hidden states and the LM-head weight** — `log_softmax(hidden @ Wᵀ + b)[target]` — +**without ever materializing the `[N, V]` logits**. It targets large-vocabulary RL +post-training, where the `[B, S, V]` logits activation (and its gradient) dominate +memory. The forward streams the vocab in blocks through an online softmax; the +backward recomputes the logit tiles instead of storing them, trading compute for +memory. It is differentiable w.r.t. `hidden`, `lm_head_weight`, and `bias`. + +This differs from [Fused LogP](fused-logp.md), which takes already-materialized +logits as input. Here the LM-head projection is fused into the reduction, so the +`[N, V]` tensor never lands in HBM. + +## Entry Point + +```python +from rl_engine.kernels.registry import kernel_registry + +linear_logp = kernel_registry.get_op("linear_logp") + +logp = linear_logp( + hidden, # [B, S, D] or [N, D] (differentiable) + lm_head_weight, # [V, D] (differentiable) + target_ids, # [B, S] or [N] int, in [0, V) + bias=None, # [V] optional (differentiable) +) # -> [B, S] or [N], float32 + +logp.sum().backward() # gradients flow into hidden, lm_head_weight, bias +``` + +## Backends + +| Backend | Wrapper | Status | +| --- | --- | --- | +| CUDA SM90 (Hopper) | `FusedLinearLogpSM90Op` | TMA + WGMMA streaming forward, online softmax in smem; chunked backward. Compiles for `sm_90a`; **numerics pending validation on an SM90 GPU.** | +| CUDA / ROCm (Triton) | `TritonLinearLogpOp` | Triton online-softmax forward; Liger-style chunked backward (cuBLAS matmuls, deterministic). Phase 1. | +| PyTorch native | `NativeLinearLogpOp` | Naive `F.linear` + `log_softmax` + `gather` reference; CPU / Triton-less fallback. | + +The SM90 backend (`csrc/cuda/fused_linear_logp_sm90.cu`) streams hidden/weight +tiles via TMA, computes each logit tile with WGMMA (`m64n64k16`), folds it into a +per-row online softmax in shared memory, and gathers the target logit — never +materializing `[N, V]`. It is **build-guarded**: only compiled when the extension +is built with `KERNEL_ALIGN_FORCE_SM90=1` on an SM90 device (WGMMA is Hopper-only, +`sm_90a`), and the registry only selects it when `cc_major == 9` and the symbol is +present. The forward kernel requires bf16 hidden/weight; the backward reuses the +deterministic chunked path. It assembles cleanly under CUDA 13.1 `ptxas`, but the +layout-sensitive pieces (WGMMA descriptors, accumulator→smem mapping, B-operand +transpose) require validation on Hopper hardware. + +The backward chunks over the token dimension: for each chunk it materializes only +`[chunk, V]` logits, recomputes the softmax from scratch, and forms the three +gradients with cuBLAS matmuls (`grad_hidden = dz @ W`, `grad_weight += dzᵀ @ X`). +`grad_weight` is accumulated in sequential loop order, so it is **atomic-free and +bitwise-deterministic** while peak backward memory stays `chunk·V` instead of `N·V`. + +The native op materializes the full `[N, V]` logits and is the correctness oracle; +the Triton op is the portable, fp32-accurate baseline that the future CUDA generic, +CUDA SM90 (TMA/WGMMA), and native ROCm backends are validated against. See the +[design doc](../design/fused-linear-logp.md) for the phased plan. + +## Tensor Contract + +| Argument | Shape | Dtype | Requirements | +| --- | --- | --- | --- | +| `hidden` | `[N, D]` / `[B, S, D]` | bf16 / fp16 / fp32 | Differentiable; contiguous. | +| `lm_head_weight` | `[V, D]` | bf16 / fp16 / fp32 | Differentiable; contiguous (row-major over V). | +| `target_ids` | `[N]` / `[B, S]` | int | Token id per position, in `[0, V)`. | +| `bias` | `[V]` | float | Optional; differentiable. | +| Output | `[N]` / `[B, S]` | float32 | `z[target] − logsumexp(z)` per position. | + +Gradients flow into `hidden`, `lm_head_weight`, and `bias`; `target_ids` is +integer and non-differentiable. + +## Reference Semantics + +```python +logits = torch.nn.functional.linear(hidden.float(), weight.float(), bias) # [N, V] +logp = torch.log_softmax(logits, dim=-1) +out = logp.gather(-1, target_ids.long().unsqueeze(-1)).squeeze(-1) +``` + +The Triton kernel accumulates the matmul and softmax in float32, so it matches the +float32 reference to `atol≈1e-3`. For bf16/fp16 inputs it matches the **fp32-upcast** +reference (it is more accurate than a bf16 `F.linear`, which rounds the logits). + +## Performance + +```bash +python benchmarks/benchmark_linear_logp.py +python benchmarks/benchmark_linear_logp.py --configs "4096,2048,32768;4096,2048,131072" +``` + +Indicative results (RTX PRO 6000, SM120, bf16, N=4096, D=2048; native vs Triton): + +| shape (N×H×V) | fwd | fwd+bwd | peak fwd VRAM (native → Triton) | +| --- | --- | --- | --- | +| 4096×2048×32768 | 0.53× | 0.43× | 1280 MB → ~0 MB | +| 4096×2048×50257 | 0.99× | 0.46× | 1965 MB → ~0 MB | +| 4096×2048×131072 | 0.64× | 0.23× | 5120 MB → ~0 MB | + +The headline is memory: the native path allocates the `[N, V]` logits (forward peak +scales with `V`), while the fused op streams them online — its forward peak is +**independent of `V`**, and the chunked backward only ever holds `chunk·V`. The +forward matmul keeps operands in their native dtype, so bf16/fp16 inputs run on +tensor cores (fp32 accumulation) and the forward lands near cuBLAS parity; the fp32 +path stays full-precision (`input_precision="ieee"`) for its role as the tolerance +reference. The backward runs at ~2–4× native: it recomputes the logits per chunk +(native keeps the `[N, V]` it already materialized), which is the compute-for-memory +trade. Closing that gap and the absolute forward latency are the job of the native +CUDA generic / SM90 (TMA/WGMMA) backends in later phases; Phase 1 delivers the +memory reduction and the correctness/tolerance baseline. + +## Tests + +```bash +python -m pytest tests/test_linear_logp.py -v +``` + +Covers the native reference vs the materialized definition, Triton forward (fp32 and +bf16) vs native, Triton backward vs native autograd (with and without bias), leading- +shape preservation, a large-vocab smoke test, and registry dispatch. Triton tests +skip without CUDA + Triton. + +## Implementation Files + +- `rl_engine/kernels/ops/triton/loss/linear_logp.py` +- `rl_engine/kernels/ops/pytorch/loss/linear_logp.py` +- `rl_engine/kernels/ops/cuda/loss/linear_logp.py` (SM90 wrapper + chunked backward) +- `csrc/cuda/fused_linear_logp_sm90.cu`, `csrc/ops.cpp`, `setup.py` (SM90 kernel + build) +- `rl_engine/kernels/registry.py` +- `tests/test_linear_logp.py` +- `benchmarks/benchmark_linear_logp.py` +- `docs/design/fused-linear-logp.md` diff --git a/rl_engine/kernels/ops/cuda/loss/linear_logp.py b/rl_engine/kernels/ops/cuda/loss/linear_logp.py new file mode 100644 index 00000000..191e9f7c --- /dev/null +++ b/rl_engine/kernels/ops/cuda/loss/linear_logp.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from typing import Optional + +import torch + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.utils.logger import logger + +# Backward token-chunk target (mirrors the Triton op): keep peak backward memory +# at ~chunk*V instead of N*V. +_BWD_CHUNK_ELEMS = 1 << 24 + + +class _FusedLinearLogpSM90Function(torch.autograd.Function): + """SM90 TMA+WGMMA fused forward + Liger-style chunked backward. + + The forward calls the compiled ``_C.fused_linear_logp_sm90`` kernel (logits + never materialized). The backward reuses the deterministic chunked cuBLAS + path so gradients flow into ``hidden``, ``lm_head_weight`` and ``bias``. + """ + + @staticmethod + def forward(ctx, hidden, lm_head_weight, bias, target_ids): + hidden_2d = hidden.reshape(-1, hidden.size(-1)).contiguous() + weight = lm_head_weight.contiguous() + target_1d = ( + target_ids.reshape(-1).to(device=hidden_2d.device, dtype=torch.int32).contiguous() + ) + logp, _lse = _C.fused_linear_logp_sm90(hidden_2d, weight, target_1d, bias) + + ctx.save_for_backward(hidden_2d, weight, bias if bias is not None else hidden_2d, target_1d) + ctx.has_bias = bias is not None + ctx.lead_shape = hidden.shape[:-1] + ctx.hidden_dtype = hidden.dtype + ctx.weight_dtype = lm_head_weight.dtype + ctx.bias_dtype = bias.dtype if bias is not None else None + return logp.reshape(hidden.shape[:-1]) + + @staticmethod + def backward(ctx, grad_logp): + hidden_2d, weight, bias_t, target_1d = ctx.saved_tensors + n, d = hidden_2d.shape + v = weight.shape[0] + dt = weight.dtype + g = grad_logp.reshape(-1).to(torch.float32) + + grad_h = torch.empty_like(hidden_2d, dtype=torch.float32) + grad_w = torch.zeros(v, d, device=weight.device, dtype=torch.float32) + grad_b = torch.zeros(v, device=weight.device, dtype=torch.float32) if ctx.has_bias else None + + chunk = max(1, min(n, _BWD_CHUNK_ELEMS // v)) + for i0 in range(0, n, chunk): + i1 = min(i0 + chunk, n) + x = hidden_2d[i0:i1] + logits = torch.matmul(x, weight.t()) + if ctx.has_bias: + logits = logits + bias_t + dz = torch.softmax(logits.float(), dim=-1).neg_() + rows = torch.arange(i1 - i0, device=dz.device) + dz[rows, target_1d[i0:i1].long()] += 1.0 + dz *= g[i0:i1].unsqueeze(1) + + dz_dt = dz.to(dt) + grad_h[i0:i1] = torch.matmul(dz_dt, weight).float() + grad_w += torch.matmul(dz_dt.t(), x).float() + if ctx.has_bias: + grad_b += dz.sum(0) + + grad_hidden = grad_h.to(ctx.hidden_dtype).reshape(ctx.lead_shape + (d,)) + grad_weight = grad_w.to(ctx.weight_dtype) + grad_bias = grad_b.to(ctx.bias_dtype) if ctx.has_bias else None + return grad_hidden, grad_weight, grad_bias, None + + +class FusedLinearLogpSM90Op: + """SM90 (Hopper) TMA+WGMMA fused linear log-prob. + + Computes ``log_softmax(hidden @ W^T + b)[target]`` without materializing the + ``[N, V]`` logits. Requires the extension built with ``KERNEL_ALIGN_FORCE_SM90=1`` + on an SM90 device; bfloat16 hidden/weight only. + """ + + def __init__(self) -> None: + if not _EXT_AVAILABLE or not hasattr(_C, "fused_linear_logp_sm90"): + raise RuntimeError( + "fused_linear_logp_sm90 is not compiled into the extension. Rebuild with " + "KERNEL_ALIGN_FORCE_SM90=1 on an SM90 (Hopper) device: 'pip install -e .'" + ) + logger.info("Successfully linked to precompiled _C.fused_linear_logp_sm90 kernel.") + + def __call__( + self, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.apply(hidden, lm_head_weight, target_ids, bias) + + def apply( + self, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if not hidden.is_cuda: + raise RuntimeError("FusedLinearLogpSM90Op requires CUDA tensors.") + if hidden.dtype != torch.bfloat16 or lm_head_weight.dtype != torch.bfloat16: + raise RuntimeError("FusedLinearLogpSM90Op requires bfloat16 hidden/weight.") + if lm_head_weight.size(-1) != hidden.size(-1): + raise ValueError( + f"hidden dim {hidden.size(-1)} must match lm_head_weight dim " + f"{lm_head_weight.size(-1)}" + ) + return _FusedLinearLogpSM90Function.apply(hidden, lm_head_weight, bias, target_ids) diff --git a/rl_engine/kernels/ops/pytorch/loss/linear_logp.py b/rl_engine/kernels/ops/pytorch/loss/linear_logp.py new file mode 100644 index 00000000..dca4700e --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/loss/linear_logp.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +from typing import Optional + +import torch + + +class NativeLinearLogpOp: + """Naive PyTorch reference for fused linear log-prob. + + Materializes the full ``[N, V]`` logits with a single ``F.linear`` and runs + ``log_softmax`` + ``gather``. This is the obviously-correct oracle the fused + kernels are validated against (and the baseline the benchmark measures the + VRAM win against); it is also the CPU / Triton-less fallback. Differentiable + w.r.t. ``hidden``, ``lm_head_weight`` and ``bias`` through autograd. + """ + + def __init__(self) -> None: + pass + + def __call__( + self, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.apply(hidden, lm_head_weight, target_ids, bias) + + def apply( + self, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Selected-token log-prob ``z[t] - logsumexp(z)``, returned in float32.""" + if hidden.shape[:-1] != target_ids.shape: + raise ValueError( + f"hidden leading shape {tuple(hidden.shape[:-1])} must match " + f"target_ids shape {tuple(target_ids.shape)}" + ) + if lm_head_weight.size(-1) != hidden.size(-1): + raise ValueError( + f"hidden dim {hidden.size(-1)} must match lm_head_weight dim " + f"{lm_head_weight.size(-1)}" + ) + + lead_shape = hidden.shape[:-1] + hidden_2d = hidden.reshape(-1, hidden.size(-1)) + logits = torch.nn.functional.linear(hidden_2d, lm_head_weight, bias) + log_probs = torch.log_softmax(logits.float(), dim=-1) + target_1d = target_ids.reshape(-1).to(device=logits.device, dtype=torch.long) + selected = torch.gather(log_probs, dim=-1, index=target_1d.unsqueeze(1)).squeeze(-1) + return selected.reshape(lead_shape) diff --git a/rl_engine/kernels/ops/triton/loss/__init__.py b/rl_engine/kernels/ops/triton/loss/__init__.py new file mode 100644 index 00000000..86cf4c9d --- /dev/null +++ b/rl_engine/kernels/ops/triton/loss/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors diff --git a/rl_engine/kernels/ops/triton/loss/linear_logp.py b/rl_engine/kernels/ops/triton/loss/linear_logp.py new file mode 100644 index 00000000..40941bce --- /dev/null +++ b/rl_engine/kernels/ops/triton/loss/linear_logp.py @@ -0,0 +1,218 @@ +# 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 + +# Token / vocab / hidden tile sizes (forward Triton kernel). +_BLOCK_N = 32 +_BLOCK_V = 64 +_BLOCK_D = 64 + +# Backward token-chunk size target: process at most this many [chunk, V] logit +# elements per cuBLAS step so peak backward memory stays ~chunk*V, not N*V. +_BWD_CHUNK_ELEMS = 1 << 24 + + +@triton.jit +def _linear_logp_fwd_kernel( + h_ptr, # hidden [N, D] + w_ptr, # lm_head_weight [V, D] + b_ptr, # bias [V] (or dummy when HAS_BIAS=False) + t_ptr, # target_ids [N] + logp_ptr, # output [N] + lse_ptr, # output [N], saved for backward + N, + D, + V, + stride_hn, + stride_hd, + stride_wv, + stride_wd, + HAS_BIAS: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_V: tl.constexpr, + BLOCK_D: tl.constexpr, +): + """One program per token-block. Streams the vocab in BLOCK_V tiles, folding + each ``hidden @ Wblk^T`` tile into an online-softmax state without ever + materializing the full [N, V] logits. Stores logp and the row log-sum-exp.""" + pid = tl.program_id(0) + rows = pid * BLOCK_N + tl.arange(0, BLOCK_N) + row_mask = rows < N + target = tl.load(t_ptr + rows, mask=row_mask, other=0).to(tl.int32) + + m = tl.full((BLOCK_N,), float("-inf"), tl.float32) + s = tl.zeros((BLOCK_N,), tl.float32) + z_t = tl.zeros((BLOCK_N,), tl.float32) + + for v0 in range(0, V, BLOCK_V): + vcols = v0 + tl.arange(0, BLOCK_V) + vmask = vcols < V + + acc = tl.zeros((BLOCK_N, BLOCK_V), tl.float32) + for d0 in range(0, D, BLOCK_D): + offs_d = d0 + tl.arange(0, BLOCK_D) + d_mask = offs_d < D + h = tl.load( + h_ptr + rows[:, None] * stride_hn + offs_d[None, :] * stride_hd, + mask=row_mask[:, None] & d_mask[None, :], + other=0.0, + ) + w = tl.load( + w_ptr + vcols[:, None] * stride_wv + offs_d[None, :] * stride_wd, + mask=vmask[:, None] & d_mask[None, :], + other=0.0, + ) + acc += tl.dot(h, tl.trans(w), input_precision="ieee") + + if HAS_BIAS: + acc += tl.load(b_ptr + vcols, mask=vmask, other=0.0).to(tl.float32)[None, :] + + is_t = (vcols[None, :] == target[:, None]) & vmask[None, :] + z_t += tl.sum(tl.where(is_t, acc, 0.0), axis=1) + acc = tl.where(vmask[None, :], acc, float("-inf")) + + tile_max = tl.max(acc, axis=1) + new_m = tl.maximum(m, tile_max) + s = s * tl.exp(m - new_m) + tl.sum(tl.exp(acc - new_m[:, None]), axis=1) + m = new_m + + lse = m + tl.log(s) + tl.store(logp_ptr + rows, z_t - lse, mask=row_mask) + tl.store(lse_ptr + rows, lse, mask=row_mask) + + +class _LinearLogpFunction(torch.autograd.Function): + """Autograd wrapper: fused forward + recompute-based backward.""" + + @staticmethod + def forward(ctx, hidden, lm_head_weight, bias, target_ids): + hidden_2d = hidden.reshape(-1, hidden.size(-1)).contiguous() + weight = lm_head_weight.contiguous() + target_1d = ( + target_ids.reshape(-1).to(device=hidden_2d.device, dtype=torch.int32).contiguous() + ) + n, d = hidden_2d.shape + v = weight.shape[0] + + logp = torch.empty(n, device=hidden_2d.device, dtype=torch.float32) + lse = torch.empty(n, device=hidden_2d.device, dtype=torch.float32) + bias_t = bias.contiguous() if bias is not None else hidden_2d # dummy ptr when no bias + + grid = (triton.cdiv(n, _BLOCK_N),) + _linear_logp_fwd_kernel[grid]( + hidden_2d, + weight, + bias_t, + target_1d, + logp, + lse, + n, + d, + v, + hidden_2d.stride(0), + hidden_2d.stride(1), + weight.stride(0), + weight.stride(1), + HAS_BIAS=bias is not None, + BLOCK_N=_BLOCK_N, + BLOCK_V=_BLOCK_V, + BLOCK_D=_BLOCK_D, + ) + + ctx.save_for_backward(hidden_2d, weight, bias_t, target_1d, lse) + ctx.has_bias = bias is not None + ctx.lead_shape = hidden.shape[:-1] + ctx.hidden_dtype = hidden.dtype + ctx.weight_dtype = lm_head_weight.dtype + ctx.bias_dtype = bias.dtype if bias is not None else None + return logp.reshape(hidden.shape[:-1]) + + @staticmethod + def backward(ctx, grad_logp): + hidden_2d, weight, bias_t, target_1d, _lse = ctx.saved_tensors + n, d = hidden_2d.shape + v = weight.shape[0] + dt = weight.dtype + g = grad_logp.reshape(-1).to(torch.float32) + + grad_h = torch.empty_like(hidden_2d, dtype=torch.float32) + grad_w = torch.zeros(v, d, device=weight.device, dtype=torch.float32) + grad_b = torch.zeros(v, device=weight.device, dtype=torch.float32) if ctx.has_bias else None + + # Liger-style chunked materialization: process at most `chunk` tokens at a + # time, materializing only [chunk, V] logits. The projections use cuBLAS + # matmuls (tensor cores for bf16/fp16), and grad_weight is accumulated in + # sequential loop order -> deterministic and atomic-free. Peak extra + # memory is chunk*V instead of N*V. + chunk = max(1, min(n, _BWD_CHUNK_ELEMS // v)) + for i0 in range(0, n, chunk): + i1 = min(i0 + chunk, n) + x = hidden_2d[i0:i1] # [C, D] + logits = torch.matmul(x, weight.t()) # [C, V] + if ctx.has_bias: + logits = logits + bias_t + + # dz = g * (onehot - softmax(logits)), recomputed from scratch so it + # is self-normalizing and independent of the forward's saved lse. + dz = torch.softmax(logits.float(), dim=-1).neg_() # [C, V] fp32 + rows = torch.arange(i1 - i0, device=dz.device) + dz[rows, target_1d[i0:i1].long()] += 1.0 + dz *= g[i0:i1].unsqueeze(1) + + dz_dt = dz.to(dt) + grad_h[i0:i1] = torch.matmul(dz_dt, weight).float() # [C, D] + grad_w += torch.matmul(dz_dt.t(), x).float() # [V, D] + if ctx.has_bias: + grad_b += dz.sum(0) + + grad_hidden = grad_h.to(ctx.hidden_dtype).reshape(ctx.lead_shape + (d,)) + grad_weight = grad_w.to(ctx.weight_dtype) + grad_bias = grad_b.to(ctx.bias_dtype) if ctx.has_bias else None + # Inputs: hidden, lm_head_weight, bias, target_ids. + return grad_hidden, grad_weight, grad_bias, None + + +class TritonLinearLogpOp: + """Triton fused linear log-prob op. + + Computes per-token ``log_softmax(hidden @ W^T + b)[target]`` without + materializing the ``[N, V]`` logits: the forward streams the vocab through an + online softmax, the backward recomputes the logit tiles instead of storing + them. Differentiable w.r.t. ``hidden``, ``lm_head_weight`` and ``bias``. + """ + + def __call__( + self, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + return self.apply(hidden, lm_head_weight, target_ids, bias) + + def apply( + self, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if not hidden.is_cuda: + raise RuntimeError("TritonLinearLogpOp requires CUDA tensors.") + if hidden.shape[:-1] != target_ids.shape: + raise ValueError( + f"hidden leading shape {tuple(hidden.shape[:-1])} must match " + f"target_ids shape {tuple(target_ids.shape)}" + ) + if lm_head_weight.size(-1) != hidden.size(-1): + raise ValueError( + f"hidden dim {hidden.size(-1)} must match lm_head_weight dim " + f"{lm_head_weight.size(-1)}" + ) + return _LinearLogpFunction.apply(hidden, lm_head_weight, bias, target_ids) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 7d85d263..558574e2 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -37,6 +37,13 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): TRITON_GRPO_LOSS = "rl_engine.kernels.ops.triton.triton_grpo_loss.TritonGRPOLossOp" PYTORCH_GRPO_LOSS = "rl_engine.kernels.ops.pytorch.loss.grpo_loss.NativeGRPOLossOp" + # Fused linear log-prob (hidden @ W^T -> selected-token logp, no [N, V] logits) + CUDA_FUSED_LINEAR_LOGP_SM90 = ( + "rl_engine.kernels.ops.cuda.loss.linear_logp.FusedLinearLogpSM90Op" + ) + TRITON_LINEAR_LOGP = "rl_engine.kernels.ops.triton.loss.linear_logp.TritonLinearLogpOp" + PYTORCH_LINEAR_LOGP = "rl_engine.kernels.ops.pytorch.loss.linear_logp.NativeLinearLogpOp" + # Generic fallback TRITON_GENERIC = "rl_engine.kernels.ops.triton.generic.TritonOp" PYTORCH_NATIVE = "rl_engine.kernels.ops.pytorch.loss.logp.NativeLogpOp" @@ -74,17 +81,20 @@ def __init__(self): ], "attn": [OpBackend.FLASH_ATTN, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], + "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], # Default dispatch logic for new operators }, "rocm": { "logp": [OpBackend.ROCM_AITER, OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE], "attn": [OpBackend.TRITON_GENERIC, OpBackend.PYTORCH_NATIVE], "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], + "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], }, "cpu": { "logp": [OpBackend.PYTORCH_NATIVE], "attn": [OpBackend.PYTORCH_NATIVE], "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], + "linear_logp": [OpBackend.PYTORCH_LINEAR_LOGP], }, } logger.info(f"KernelRegistry initialized for {device_ctx.device_type}") @@ -112,6 +122,14 @@ def _adjust_priority_for_hardware(self): logp_list = self._priority_map["cuda"]["logp"] if OpBackend.CUDA_FUSED_LOGP_SM90 not in logp_list: logp_list.insert(0, OpBackend.CUDA_FUSED_LOGP_SM90) + + # The fused linear-logp SM90 kernel uses WGMMA, which is Hopper-only + # (sm_90a) -- gate strictly on cc_major == 9, not >= 9. + linear_logp_compiled = _EXT_AVAILABLE and hasattr(_C, "fused_linear_logp_sm90") + if linear_logp_compiled and cc_major == 9: + ll_list = self._priority_map["cuda"]["linear_logp"] + if OpBackend.CUDA_FUSED_LINEAR_LOGP_SM90 not in ll_list: + ll_list.insert(0, OpBackend.CUDA_FUSED_LINEAR_LOGP_SM90) elif cc >= 90: logger.debug( f"SM{cc}: fused TMA LogP kernel not compiled into _C; " diff --git a/setup.py b/setup.py index 8e608137..d5ddb899 100644 --- a/setup.py +++ b/setup.py @@ -110,11 +110,15 @@ def get_extensions(): cxx_flags = ["-O3", "-std=c++17", "-DKERNEL_ALIGN_WITH_CUDA"] extra_link_args = [] - tma_src = "csrc/cuda/fused_logp_sm90.cu" + sm90_srcs = [ + "csrc/cuda/fused_logp_sm90.cu", + "csrc/cuda/fused_linear_logp_sm90.cu", # TMA + WGMMA fused linear log-prob + ] enable_sm90 = os.environ.get("KERNEL_ALIGN_FORCE_SM90") == "1" - if enable_sm90 and os.path.exists(tma_src): - tma_arch = f"{cc_major}{cc_minor}a" - cuda_sources.append(tma_src) + present_sm90 = [s for s in sm90_srcs if os.path.exists(s)] + if enable_sm90 and present_sm90: + tma_arch = f"{cc_major}{cc_minor}a" # WGMMA/TMA require the arch-native 'a' variant + cuda_sources.extend(present_sm90) nvcc_flags.append(f"-gencode=arch=compute_{tma_arch},code=sm_{tma_arch}") cxx_flags.append("-DKERNEL_ALIGN_WITH_SM90") extra_link_args.append("-lcuda") diff --git a/tests/test_linear_logp.py b/tests/test_linear_logp.py new file mode 100644 index 00000000..998f8cf6 --- /dev/null +++ b/tests/test_linear_logp.py @@ -0,0 +1,157 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +import pytest +import torch + +from rl_engine.kernels.ops.pytorch.loss.linear_logp import NativeLinearLogpOp + +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 linear log-prob requires a CUDA device and Triton.", +) + +# Deliberately non-multiples of the kernel block sizes (32 / 64 / 64). +_N = 40 +_D = 80 +_V = 300 + + +def _inputs(seed, *, device, dtype=torch.float32, bias=True, lead=None): + gen = torch.Generator(device=device).manual_seed(seed) + lead = lead or (_N,) + hidden = torch.randn(*lead, _D, generator=gen, device=device, dtype=dtype) + weight = torch.randn(_V, _D, generator=gen, device=device, dtype=dtype) + bias_t = torch.randn(_V, generator=gen, device=device, dtype=dtype) if bias else None + target = torch.randint(0, _V, lead, generator=gen, device=device) + return hidden, weight, target, bias_t + + +def _manual_reference(hidden, weight, target, bias): + """The semantic definition: materialize logits, log_softmax, gather.""" + logits = torch.nn.functional.linear( + hidden.float(), weight.float(), None if bias is None else bias.float() + ) + logp = torch.log_softmax(logits, dim=-1) + idx = target.reshape(-1).long() + sel = logp.reshape(-1, logp.size(-1)).gather(-1, idx.unsqueeze(1)).squeeze(1) + return sel.reshape(target.shape) + + +def test_native_matches_manual_reference(): + native = NativeLinearLogpOp() + hidden, weight, target, bias = _inputs(0, device="cpu") + out = native(hidden, weight, target, bias) + ref = _manual_reference(hidden, weight, target, bias) + assert out.dtype == torch.float32 + assert torch.allclose(out, ref, atol=1e-5) + + +def test_native_rejects_shape_mismatch(): + native = NativeLinearLogpOp() + hidden, weight, _, bias = _inputs(0, device="cpu") + with pytest.raises(ValueError): + native(hidden, weight, torch.zeros(_N + 1, dtype=torch.long), bias) + + +@requires_triton_cuda +def test_triton_forward_matches_native_fp32(): + from rl_engine.kernels.ops.triton.loss.linear_logp import TritonLinearLogpOp + + native, trit = NativeLinearLogpOp(), TritonLinearLogpOp() + hidden, weight, target, bias = _inputs(1, device="cuda") + ref = native(hidden, weight, target, bias) + out = trit(hidden, weight, target, bias) + assert torch.allclose(out, ref, atol=1e-3) + + +@requires_triton_cuda +def test_triton_forward_matches_native_bf16(): + from rl_engine.kernels.ops.triton.loss.linear_logp import TritonLinearLogpOp + + native, trit = NativeLinearLogpOp(), TritonLinearLogpOp() + hidden, weight, target, bias = _inputs(2, device="cuda", dtype=torch.bfloat16) + # The kernel accumulates in fp32, so the oracle uses the fp32-upcast inputs. + ref = native(hidden.float(), weight.float(), target, bias.float()) + out = trit(hidden, weight, target, bias) + assert torch.allclose(out, ref, atol=2e-2) + + +@requires_triton_cuda +@pytest.mark.parametrize("use_bias", [True, False]) +def test_triton_backward_matches_native(use_bias): + from rl_engine.kernels.ops.triton.loss.linear_logp import TritonLinearLogpOp + + native, trit = NativeLinearLogpOp(), TritonLinearLogpOp() + hidden, weight, target, bias = _inputs(3, device="cuda", bias=use_bias) + grad_out = torch.randn(_N, device="cuda") + + def run(op, h, w, b): + h = h.detach().clone().requires_grad_(True) + w = w.detach().clone().requires_grad_(True) + b = b.detach().clone().requires_grad_(True) if b is not None else None + op(h, w, target, b).backward(grad_out) + return h.grad, w.grad, (b.grad if b is not None else None) + + th, tw, tb = run(trit, hidden, weight, bias) + nh, nw, nb = run(native, hidden, weight, bias) + assert torch.allclose(th, nh, atol=2e-3) + assert torch.allclose(tw, nw, atol=2e-3) + if use_bias: + assert torch.allclose(tb, nb, atol=2e-3) + + +@requires_triton_cuda +def test_triton_gradients_flow_to_inputs_only(): + from rl_engine.kernels.ops.triton.loss.linear_logp import TritonLinearLogpOp + + trit = TritonLinearLogpOp() + hidden, weight, target, bias = _inputs(4, device="cuda") + hidden = hidden.requires_grad_(True) + weight = weight.requires_grad_(True) + bias = bias.requires_grad_(True) + trit(hidden, weight, target, bias).sum().backward() + assert hidden.grad is not None and weight.grad is not None and bias.grad is not None + assert target.grad is None # integer targets are non-differentiable + + +@requires_triton_cuda +def test_triton_preserves_leading_shape(): + from rl_engine.kernels.ops.triton.loss.linear_logp import TritonLinearLogpOp + + native, trit = NativeLinearLogpOp(), TritonLinearLogpOp() + hidden, weight, target, bias = _inputs(5, device="cuda", lead=(4, 7)) + out = trit(hidden, weight, target, bias) + assert out.shape == (4, 7) + assert torch.allclose(out, native(hidden, weight, target, bias), atol=1e-3) + + +@requires_triton_cuda +def test_triton_large_vocab_smoke(): + from rl_engine.kernels.ops.triton.loss.linear_logp import TritonLinearLogpOp + + trit = TritonLinearLogpOp() + hidden = torch.randn(8, 64, device="cuda") + weight = torch.randn(50257, 64, device="cuda") + target = torch.randint(0, 50257, (8,), device="cuda") + out = trit(hidden, weight, target) + assert out.shape == (8,) and torch.isfinite(out).all() + + +def test_registry_dispatch_matches_native(): + from rl_engine.kernels.registry import kernel_registry + from rl_engine.platforms.device import device_ctx + + op = kernel_registry.get_op("linear_logp") + device = device_ctx.device if device_ctx.device_type == "cuda" else "cpu" + hidden, weight, target, bias = _inputs(6, device=device) + out = op(hidden, weight, target, bias) + ref = NativeLinearLogpOp()(hidden, weight, target, bias) + assert torch.allclose(out.cpu(), ref.cpu(), atol=1e-3) From 480388ec94b609fef48c1b62c2979cfb2dc40bfd Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Mon, 15 Jun 2026 07:55:45 +0000 Subject: [PATCH 2/8] [FEAT][kernels] SM90 fused linear log-prob: TMA + tensor-core forward, split-V --- benchmarks/benchmark_linear_logp.py | 99 +++-- csrc/cuda/fused_linear_logp_sm90.cu | 406 +++++++++++------- csrc/cuda/fused_logp_sm90.cu | 5 +- docs/operators/linear-logp.md | 31 +- .../kernels/ops/cuda/loss/linear_logp.py | 36 +- rl_engine/kernels/registry.py | 4 +- tests/test_linear_logp.py | 138 ++++++ 7 files changed, 508 insertions(+), 211 deletions(-) diff --git a/benchmarks/benchmark_linear_logp.py b/benchmarks/benchmark_linear_logp.py index f3c863d1..81d91e66 100644 --- a/benchmarks/benchmark_linear_logp.py +++ b/benchmarks/benchmark_linear_logp.py @@ -25,6 +25,24 @@ from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger + +def _maybe_sm90_op(): + """The Hopper TMA+MMA op, or None when unavailable (non-Hopper / not built).""" + import torch + + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + if not ( + torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] == 9 + and _EXT_AVAILABLE + and hasattr(_C, "fused_linear_logp_sm90") + ): + return None + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + + return FusedLinearLogpSM90Op() + # (num_tokens, hidden_dim, vocab) DEFAULT_CONFIGS = [ (4096, 2048, 32768), @@ -75,51 +93,56 @@ def run_benchmark(args): dtype = torch.bfloat16 native = NativeLinearLogpOp() triton_op = TritonLinearLogpOp() + sm90_op = _maybe_sm90_op() - logger.info(f"linear_logp benchmark on {device} (dtype={dtype})") + logger.info( + f"linear_logp benchmark on {device} (dtype={dtype}); " + f"SM90 TMA+MMA backend {'enabled' if sm90_op is not None else 'unavailable'}" + ) rows = [] for num_tokens, hidden_dim, vocab in args.configs: hidden, weight, target = _make_inputs(num_tokens, hidden_dim, vocab, device, dtype) - def native_fwd(h=hidden, w=weight): + def fwd(op, h=hidden, w=weight): with torch.no_grad(): - native(h, w, target) - - def triton_fwd(h=hidden, w=weight): - with torch.no_grad(): - triton_op(h, w, target) - - def native_fwd_bwd(): - h = hidden.clone().requires_grad_(True) - w = weight.clone().requires_grad_(True) - native(h, w, target).sum().backward() + op(h, w, target) - def triton_fwd_bwd(): + def fwd_bwd(op): h = hidden.clone().requires_grad_(True) w = weight.clone().requires_grad_(True) - triton_op(h, w, target).sum().backward() - - n_fwd = _time_ms(native_fwd, args.warmup, args.iters) - t_fwd = _time_ms(triton_fwd, args.warmup, args.iters) - n_fb = _time_ms(native_fwd_bwd, args.warmup, args.iters) - t_fb = _time_ms(triton_fwd_bwd, args.warmup, args.iters) - n_vram = _peak_vram_gb(native_fwd) - t_vram = _peak_vram_gb(triton_fwd) - - rows.append( - [ - f"{num_tokens}x{hidden_dim}x{vocab}", - 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*1024:.0f}", - f"{t_vram*1024:.0f}", + op(h, w, target).sum().backward() + + 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_gb(lambda: fwd(native)) + t_vram = _peak_vram_gb(lambda: fwd(triton_op)) + + row = [ + f"{num_tokens}x{hidden_dim}x{vocab}", + 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*1024:.0f}", + f"{t_vram*1024:.0f}", + ] + if sm90_op is not None: + s_fwd = _time_ms(lambda: fwd(sm90_op), args.warmup, args.iters) + s_fb = _time_ms(lambda: fwd_bwd(sm90_op), args.warmup, args.iters) + s_vram = _peak_vram_gb(lambda: fwd(sm90_op)) + row += [ + f"{s_fwd:.3f}", + f"{n_fwd/s_fwd:.2f}x", + f"{t_fwd/s_fwd:.2f}x", + f"{s_fb:.3f}", + f"{s_vram*1024:.0f}", ] - ) + rows.append(row) headers = [ "shape (N x H x V)", @@ -132,6 +155,14 @@ def triton_fwd_bwd(): "native fwd MB", "triton fwd MB", ] + if sm90_op is not None: + headers += [ + "sm90 fwd ms", + "sm90 vs native", + "sm90 vs triton", + "sm90 f+b ms", + "sm90 fwd MB", + ] print(tabulate(rows, headers=headers, tablefmt="github")) diff --git a/csrc/cuda/fused_linear_logp_sm90.cu b/csrc/cuda/fused_linear_logp_sm90.cu index 72c11ae4..34a74e89 100644 --- a/csrc/cuda/fused_linear_logp_sm90.cu +++ b/csrc/cuda/fused_linear_logp_sm90.cu @@ -2,196 +2,209 @@ // Copyright (c) 2026 RL-Kernel Contributors // // Hopper (SM90) fused linear log-prob: -// // logp[n] = log_softmax(hidden[n] @ W^T + b)[target[n]] -// -// computed WITHOUT materializing the [N, V] logits. Each CTA owns a block of BM -// tokens and streams the vocabulary in BN-wide tiles: for every vocab tile it -// 1. computes the logit tile logits[BM, BN] = H[BM, D] @ W[vtile, D]^T with -// WGMMA (K-loop over D), staging H/W tiles into shared memory via TMA, -// 2. stages the [BM, BN] fp32 logit tile into shared memory, and -// 3. folds it into a per-row online-softmax state (running max + sum) kept in -// shared memory, capturing the target logit when its column lands in tile. -// After the vocab loop it writes logp = z_target - (max + log(sum)) and the row -// log-sum-exp (lse) for the backward pass. -// -// ========================================================================= -// VALIDATION REQUIRED ON SM90. -// This file is written without an SM90 compiler/GPU in the loop. The numerically -// load-bearing, layout-sensitive pieces are: -// (a) the WGMMA shared-memory matrix descriptors (LBO/SBO/swizzle), and -// (b) the mapping from WGMMA fp32 accumulator registers -> (row, col). -// Both follow the PTX ISA "Asynchronous Warpgroup Level Matrix" section, but -// must be checked on hardware. The matmul is isolated in wgmma_m64n64k16 / -// store_acc_to_smem so it can be swapped for an mma.sync (m16n8k16) path if -// WGMMA needs debugging -- the surrounding TMA streaming + online softmax are -// independent of that choice. -// ========================================================================= #include "../utils/tma_utils.cuh" +#include +#include #include -#include // CUDART_INF_F (not pulled in transitively under CUDA 13) +#include #include namespace { -constexpr int BM = 64; // tokens per CTA == WGMMA M -constexpr int BN = 64; // vocab per tile == WGMMA N -constexpr int BK = 16; // contraction per WGMMA step == WGMMA K (bf16) -constexpr int WG_THREADS = 128; // one warpgroup (4 warps) -constexpr int ACC_REGS = (BM * BN) / WG_THREADS; // = 32 fp32 accumulators / thread - -// -------------------------------------------------------------------------- -// WGMMA helpers (SM90a). See PTX ISA: wgmma.mma_async / shared-memory matrix -// descriptor. Operands are streamed unswizzled (swizzle = 0); the descriptor -// encodes the smem start address plus the leading/stride byte offsets between -// 8x(BK) "core matrices". -// -------------------------------------------------------------------------- -__device__ __forceinline__ uint64_t make_smem_desc(const void *smem_ptr, uint32_t lbo, - uint32_t sbo) { - uint32_t addr = static_cast(__cvta_generic_to_shared(smem_ptr)); - uint64_t desc = 0; - desc |= (static_cast(addr >> 4) & 0x3FFFULL); // bits 0-13: start >> 4 - desc |= (static_cast(lbo >> 4) & 0x3FFFULL) << 16; // bits 16-29: LBO >> 4 - desc |= (static_cast(sbo >> 4) & 0x3FFFULL) << 32; // bits 32-45: SBO >> 4 - desc |= (static_cast(0) << 62); // bits 62-63: swizzle = none - return desc; -} +constexpr int BM = 256; // tokens per CTA +constexpr int BN = 64; // vocab per tile +constexpr int BK = 32; // hidden-dim slice streamed per TMA load +constexpr int WARPS = 4; // one warpgroup +constexpr int WG_THREADS = WARPS * 32; // 128 +constexpr int STAGES = 2; // double-buffering -__device__ __forceinline__ void wgmma_fence() { asm volatile("wgmma.fence.sync.aligned;"); } -__device__ __forceinline__ void wgmma_commit() { - asm volatile("wgmma.commit_group.sync.aligned;"); -} -template __device__ __forceinline__ void wgmma_wait() { - asm volatile("wgmma.wait_group.sync.aligned %0;" ::"n"(N)); -} +constexpr int MMA_M = 16; +constexpr int MMA_N = 8; +constexpr int MMA_K = 16; + +constexpr int WARP_M = BM / WARPS; // logit rows per warp +constexpr int M_TILES = WARP_M / MMA_M; // MMA m-tiles each warp owns +constexpr int N_TILES = BN / MMA_N; // 8 n-tiles per warp +constexpr int K_TILES = BK / MMA_K; // MMA k-steps per TMA tile +constexpr int KK_GROUPS = BK / 32; // 32-wide ldmatrix.x4 groups (2 k-steps each) -// d += A @ B^T for one K=16 step. desc_a/desc_b are shared-memory descriptors. -// scale_d == 0 zeroes the accumulator first; 1 accumulates onto it. -__device__ __forceinline__ void wgmma_m64n64k16(float d[ACC_REGS], uint64_t desc_a, - uint64_t desc_b, int scale_d) { - // PTX form: {d}, desc_a, desc_b, scale-d (predicate), imm-scale-a, imm-scale-b, - // imm-trans-a, imm-trans-b. scale-d is a predicate set from scale_d (0 = write, - // 1 = accumulate). scaleA/scaleB = 1 (no negate); transA/transB = 0. - asm volatile( - "{\n" - ".reg .pred p;\n" - "setp.ne.b32 p, %34, 0;\n" - "wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 " - "{%0,%1,%2,%3,%4,%5,%6,%7,%8,%9,%10,%11,%12,%13,%14,%15," - "%16,%17,%18,%19,%20,%21,%22,%23,%24,%25,%26,%27,%28,%29,%30,%31}, " - "%32, %33, p, 1, 1, 0, 0;\n" - "}\n" - : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]), "+f"(d[4]), "+f"(d[5]), "+f"(d[6]), - "+f"(d[7]), "+f"(d[8]), "+f"(d[9]), "+f"(d[10]), "+f"(d[11]), "+f"(d[12]), "+f"(d[13]), - "+f"(d[14]), "+f"(d[15]), "+f"(d[16]), "+f"(d[17]), "+f"(d[18]), "+f"(d[19]), "+f"(d[20]), - "+f"(d[21]), "+f"(d[22]), "+f"(d[23]), "+f"(d[24]), "+f"(d[25]), "+f"(d[26]), "+f"(d[27]), - "+f"(d[28]), "+f"(d[29]), "+f"(d[30]), "+f"(d[31]) - : "l"(desc_a), "l"(desc_b), "r"(scale_d)); +static_assert(WARP_M % MMA_M == 0, "rows per warp must be a multiple of MMA_M"); +static_assert(BK % 32 == 0, "BK must be a multiple of 32 (ldmatrix.x4 spans 32 cols)"); + +// Tensor-core helpers (Ampere/Hopper warp-level MMA). Same layout as +// prefix_shared_attention.cu, validated on this repo's Hopper GPUs. +__device__ __forceinline__ void ldmatrix_x4(uint32_t regs[4], uint32_t addr) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(regs[0]), "=r"(regs[1]), "=r"(regs[2]), "=r"(regs[3]) + : "r"(addr)); } -// Scatter the WGMMA fp32 accumulator to a row-major [BM][BN] shared-memory tile. -// m64nN .f32 layout (per warpgroup of 4 warps): warp w owns rows [16w, 16w+16); -// within a warp the N columns are tiled in 8-wide blocks, 4 regs per block: -// reg[4j+0]->(lane/4, 8j+(lane%4)*2 ) reg[4j+1]->(lane/4, 8j+(lane%4)*2+1) -// reg[4j+2]->(lane/4+8, 8j+(lane%4)*2 ) reg[4j+3]->(lane/4+8, 8j+(lane%4)*2+1) -__device__ __forceinline__ void store_acc_to_smem(const float d[ACC_REGS], float *s_logits) { - const int tid = threadIdx.x; - const int warp = tid / 32; - const int lane = tid % 32; - const int r0 = warp * 16 + lane / 4; - const int r1 = r0 + 8; - const int col = (lane % 4) * 2; -#pragma unroll - for (int j = 0; j < BN / 8; ++j) { - const int c = j * 8 + col; - s_logits[r0 * BN + c + 0] = d[4 * j + 0]; - s_logits[r0 * BN + c + 1] = d[4 * j + 1]; - s_logits[r1 * BN + c + 0] = d[4 * j + 2]; - s_logits[r1 * BN + c + 1] = d[4 * j + 3]; - } +// D[m16,n8] += A[m16,k16] * B[n8,k16] (A row-major, B col-major; fp32 accum) +__device__ __forceinline__ void mma_m16n8k16(const uint32_t A[4], const uint32_t B[2], + float D[4]) { + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};" + : "=f"(D[0]), "=f"(D[1]), "=f"(D[2]), "=f"(D[3]) + : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), + "f"(D[0]), "f"(D[1]), "f"(D[2]), "f"(D[3])); } -// -------------------------------------------------------------------------- -// Kernel: one CTA (one warpgroup) per BM-token block. -// Shared memory layout (one buffer; double-buffering of A/B is a follow-up): -// [ A tile BM*BK bf16 ][ B tile BN*BK bf16 ][ logits BM*BN f32 ] -// [ row_max BM f32 ][ row_sum BM f32 ][ row_zt BM f32 ][ tma mbar 8B ] -// -------------------------------------------------------------------------- + __global__ void fused_linear_logp_sm90_kernel(const __grid_constant__ CUtensorMap h_tmap, const __grid_constant__ CUtensorMap w_tmap, const int *__restrict__ target, const float *__restrict__ bias, // may be null - float *__restrict__ out_logp, - float *__restrict__ out_lse, int N, int D, int V) { + float *__restrict__ part_max, // [n_split, N] + float *__restrict__ part_sum, // [n_split, N] + float *__restrict__ part_zt, // [n_split, N] + int N, int D, int V, int n_split) { const int tid = threadIdx.x; + const int warp = tid / 32; + const int lane = tid % 32; const int row_block = blockIdx.x; + const int split = blockIdx.y; const int row_base = row_block * BM; const int num_rows = min(BM, N - row_base); - const int kd = D / BK; // assumes D % BK == 0 (caller pads/validates) + const int kd = D / BK; // D is validated to be a multiple of BK on the host + + // This CTA owns a contiguous slice of the vocab tiles (split-V): partitioning + // the V loop across blockIdx.y fills the GPU when N/BM alone is too few CTAs. + const int total_vtiles = (V + BN - 1) / BN; + const int vtiles_per_split = (total_vtiles + n_split - 1) / n_split; + const int vt_begin = split * vtiles_per_split; + const int vt_end = min(vt_begin + vtiles_per_split, total_vtiles); extern __shared__ __align__(1024) char smem[]; - nv_bfloat16 *sA = reinterpret_cast(smem); - nv_bfloat16 *sB = reinterpret_cast(sA + BM * BK); - float *sLogits = reinterpret_cast(sB + BN * BK); + nv_bfloat16 *sH = reinterpret_cast(smem); + nv_bfloat16 *sW = reinterpret_cast(sH + STAGES * BM * BK); + float *sLogits = reinterpret_cast(sW + STAGES * BN * BK); float *sMax = sLogits + BM * BN; float *sSum = sMax + BM; float *sZt = sSum + BM; - const int tma_mbar = static_cast(__cvta_generic_to_shared(sZt + BM)); + int *mbar_base = reinterpret_cast(sZt + BM); // STAGES mbarriers (8B each) - if (tid < num_rows) { - sMax[tid] = -CUDART_INF_F; - sSum[tid] = 0.0f; - sZt[tid] = 0.0f; + const uint32_t sH_base = static_cast(__cvta_generic_to_shared(sH)); + const uint32_t sW_base = static_cast(__cvta_generic_to_shared(sW)); + int mbar[STAGES]; +#pragma unroll + for (int s = 0; s < STAGES; ++s) + mbar[s] = static_cast(__cvta_generic_to_shared(mbar_base + 2 * s)); + + for (int r = tid; r < num_rows; r += WG_THREADS) { + sMax[r] = -CUDART_INF_F; + sSum[r] = 0.0f; + sZt[r] = 0.0f; } if (tid == 0) { - mbarrier_init(tma_mbar, 1); +#pragma unroll + for (int s = 0; s < STAGES; ++s) + mbarrier_init(mbar[s], 1); asm volatile("fence.mbarrier_init.release.cluster;"); } __syncthreads(); - const int a_smem = static_cast(__cvta_generic_to_shared(sA)); - const int b_smem = static_cast(__cvta_generic_to_shared(sB)); const uint32_t tile_bytes = (BM * BK + BN * BK) * sizeof(nv_bfloat16); - int phase = 0; - const int num_vtiles = (V + BN - 1) / BN; - for (int vt = 0; vt < num_vtiles; ++vt) { + // Issue the TMA for D-slice k of vocab tile vt into buffer (k % STAGES). + auto issue_load = [&](int k, int col_base) { + const int buf = k % STAGES; + const int k_off = k * BK; + tma_2d_g2s(static_cast(sH_base + buf * BM * BK * sizeof(nv_bfloat16)), &h_tmap, k_off, + row_base, mbar[buf]); + tma_2d_g2s(static_cast(sW_base + buf * BN * BK * sizeof(nv_bfloat16)), &w_tmap, k_off, + col_base, mbar[buf]); + mbarrier_arrive_expect_tx(mbar[buf], tile_bytes); + }; + + int phase[STAGES]; +#pragma unroll + for (int s = 0; s < STAGES; ++s) + phase[s] = 0; + + for (int vt = vt_begin; vt < vt_end; ++vt) { const int col_base = vt * BN; - float d[ACC_REGS]; + // Per-warp accumulators: this warp's M_TILES*16 rows x N_TILES n-tiles. + float acc[M_TILES][N_TILES][4]; #pragma unroll - for (int i = 0; i < ACC_REGS; ++i) - d[i] = 0.0f; + for (int mi = 0; mi < M_TILES; ++mi) +#pragma unroll + for (int n = 0; n < N_TILES; ++n) + acc[mi][n][0] = acc[mi][n][1] = acc[mi][n][2] = acc[mi][n][3] = 0.0f; - // K-loop over the hidden dimension: stream H[BM,BK] and W[BN,BK] via TMA, - // then issue one WGMMA per K step accumulating into d[]. + // Double Buffering: TMA loads in flight so the + // next H/W slices stream in while the current one feeds tensor-core MMAs. + if (tid == 0) { +#pragma unroll + for (int s = 0; s < STAGES - 1; ++s) + if (s < kd) + issue_load(s, col_base); + } for (int k = 0; k < kd; ++k) { - const int k_off = k * BK; - if (tid == 0) { - tma_2d_g2s(a_smem, &h_tmap, k_off, row_base, tma_mbar); - tma_2d_g2s(b_smem, &w_tmap, k_off, col_base, tma_mbar); - mbarrier_arrive_expect_tx(tma_mbar, tile_bytes); + const int buf = k % STAGES; + if (tid == 0 && k + (STAGES - 1) < kd) + issue_load(k + (STAGES - 1), col_base); // overlaps with the MMAs below + mbarrier_wait(mbar[buf], phase[buf]); + phase[buf] ^= 1; + __syncthreads(); + + const uint32_t sH_buf = sH_base + buf * BM * BK * sizeof(nv_bfloat16); + const uint32_t sW_buf = sW_base + buf * BN * BK * sizeof(nv_bfloat16); + + // Load A (this warp's M_TILES*16 rows) for every MMA k-step. + uint32_t A[M_TILES][K_TILES][4]; +#pragma unroll + for (int mi = 0; mi < M_TILES; ++mi) { + const int row0 = warp * WARP_M + mi * MMA_M + (lane % 16); +#pragma unroll + for (int kt = 0; kt < K_TILES; ++kt) { + const uint32_t a_addr = + sH_buf + (row0 * BK + (lane / 16) * 8 + kt * MMA_K) * sizeof(nv_bfloat16); + ldmatrix_x4(A[mi][kt], a_addr); + } + } + + // Load B (all n-tiles, shared across m-tiles) and contract. +#pragma unroll + for (int n = 0; n < N_TILES; ++n) { +#pragma unroll + for (int kk = 0; kk < KK_GROUPS; ++kk) { + uint32_t b4[4]; + const uint32_t b_addr = + sW_buf + ((n * MMA_N + (lane % 8)) * BK + (lane / 8) * 8 + kk * 32) * + sizeof(nv_bfloat16); + ldmatrix_x4(b4, b_addr); + const uint32_t B0[2] = {b4[0], b4[1]}; + const uint32_t B1[2] = {b4[2], b4[3]}; +#pragma unroll + for (int mi = 0; mi < M_TILES; ++mi) { + mma_m16n8k16(A[mi][2 * kk + 0], B0, acc[mi][n]); + mma_m16n8k16(A[mi][2 * kk + 1], B1, acc[mi][n]); + } + } } - mbarrier_wait(tma_mbar, phase); - phase ^= 1; - - wgmma_fence(); - const uint64_t desc_a = make_smem_desc(sA, BK * sizeof(nv_bfloat16), 0); - const uint64_t desc_b = make_smem_desc(sB, BK * sizeof(nv_bfloat16), 0); - wgmma_m64n64k16(d, desc_a, desc_b, 1); - wgmma_commit(); - wgmma_wait<0>(); __syncthreads(); } - store_acc_to_smem(d, sLogits); +#pragma unroll + for (int mi = 0; mi < M_TILES; ++mi) { + const int row = warp * WARP_M + mi * MMA_M + lane / 4; +#pragma unroll + for (int n = 0; n < N_TILES; ++n) { + const int col = n * MMA_N + (lane % 4) * 2; + sLogits[row * BN + col + 0] = acc[mi][n][0]; + sLogits[row * BN + col + 1] = acc[mi][n][1]; + sLogits[(row + 8) * BN + col + 0] = acc[mi][n][2]; + sLogits[(row + 8) * BN + col + 1] = acc[mi][n][3]; + } + } __syncthreads(); - // Online softmax: one thread per row folds this tile's BN columns into - // the running (max, sum) and captures the target logit if present. - if (tid < num_rows) { - const int r = tid; + // Online softmax: threads stride over rows, each folding this tile's BN + // columns into the running (max, sum) and capturing the target logit. + for (int r = tid; r < num_rows; r += WG_THREADS) { const int tgt = target[row_base + r]; float tmax = -CUDART_INF_F; for (int c = 0; c < BN; ++c) { @@ -223,14 +236,63 @@ __global__ void fused_linear_logp_sm90_kernel(const __grid_constant__ CUtensorMa __syncthreads(); } - if (tid < num_rows) { - const int r = tid; - const float lse = sMax[r] + logf(sSum[r]); - out_logp[row_base + r] = sZt[r] - lse; - out_lse[row_base + r] = lse; + // Emit this split's partial online-softmax state; a combine pass merges the + // per-split (max, sum, target-logit) into the final logp/lse. + for (int r = tid; r < num_rows; r += WG_THREADS) { + const int idx = split * N + row_base + r; + part_max[idx] = sMax[r]; + part_sum[idx] = sSum[r]; + part_zt[idx] = sZt[r]; } } +// Merge per-split partials: M = max_s m_s, S = sum_s s_s*exp(m_s - M), +// zt = sum_s zt_s (exactly one split holds the target column), then +// logp = zt - (M + log S). One thread per token row. +__global__ void fused_linear_logp_sm90_combine_kernel(const float *__restrict__ part_max, + const float *__restrict__ part_sum, + const float *__restrict__ part_zt, + float *__restrict__ out_logp, + float *__restrict__ out_lse, int N, + int n_split) { + const int r = blockIdx.x * blockDim.x + threadIdx.x; + if (r >= N) + return; + + float M = -CUDART_INF_F; + for (int s = 0; s < n_split; ++s) + M = fmaxf(M, part_max[s * N + r]); + + float S = 0.0f; + float zt = 0.0f; + for (int s = 0; s < n_split; ++s) { + const int idx = s * N + r; + S += part_sum[idx] * __expf(part_max[idx] - M); + zt += part_zt[idx]; + } + const float lse = M + logf(S); + out_logp[r] = zt - lse; + out_lse[r] = lse; +} + +// 2D bf16 tensor map with swizzle pinned to NONE. This kernel reads its tiles +// with plain row-major ldmatrix addressing, so the TMA must write them +// unswizzled -- the shared init_tensor_map auto-selects a swizzle from the row +// stride, which would not match. Kept local so the shared helper stays untouched. +inline void init_tensor_map_noswizzle(CUtensorMap *tmap, const nv_bfloat16 *gmem, + uint64_t gmem_height, uint64_t gmem_width, + uint32_t box_height, uint32_t box_width) { + uint64_t size[2] = {gmem_width, gmem_height}; + uint64_t stride[1] = {gmem_width * sizeof(nv_bfloat16)}; + uint32_t box[2] = {box_width, box_height}; + uint32_t elem_stride[2] = {1, 1}; + CUresult res = cuTensorMapEncodeTiled( + tmap, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, (void *)gmem, size, stride, box, elem_stride, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TORCH_CHECK(res == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed for fused_linear_logp_sm90"); +} + } // namespace // Forward: hidden [N, D] bf16, weight [V, D] bf16, target [N] int32, optional @@ -240,7 +302,6 @@ std::vector fused_linear_logp_sm90_forward(torch::Tensor hidden, torch::Tensor weight, torch::Tensor target, torch::optional bias) { - TORCH_CHECK(hidden.is_cuda() && weight.is_cuda(), "inputs must be CUDA tensors"); TORCH_CHECK(hidden.scalar_type() == at::kBFloat16, "hidden must be bfloat16"); TORCH_CHECK(weight.scalar_type() == at::kBFloat16, "weight must be bfloat16"); TORCH_CHECK(hidden.is_contiguous() && weight.is_contiguous(), "inputs must be contiguous"); @@ -254,14 +315,14 @@ std::vector fused_linear_logp_sm90_forward(torch::Tensor hidden, auto logp = torch::empty({N}, opts_f); auto lse = torch::empty({N}, opts_f); - // TMA descriptors: box [rows=BM/BN, cols=BK]. cols*sizeof(bf16) = 32B; the - // helper selects 32B swizzle for this stride -- the descriptor swizzle bits - // in make_smem_desc must be kept consistent (see VALIDATION note above). + // TMA descriptors: box [rows=BM/BN, cols=BK], unswizzled (see helper above). CUtensorMap h_tmap, w_tmap; - init_tensor_map(&h_tmap, reinterpret_cast(hidden.data_ptr()), - N, D, BM, BK); - init_tensor_map(&w_tmap, reinterpret_cast(weight.data_ptr()), - V, D, BN, BK); + init_tensor_map_noswizzle( + &h_tmap, reinterpret_cast(hidden.data_ptr()), N, D, BM, + BK); + init_tensor_map_noswizzle( + &w_tmap, reinterpret_cast(weight.data_ptr()), V, D, BN, + BK); const float *bias_ptr = nullptr; torch::Tensor bias_f; @@ -270,14 +331,37 @@ std::vector fused_linear_logp_sm90_forward(torch::Tensor hidden, bias_ptr = bias_f.data_ptr(); } - const int smem = (BM * BK + BN * BK) * sizeof(nv_bfloat16) + (BM * BN) * sizeof(float) + - 3 * BM * sizeof(float) + 16; - const int grid = (N + BM - 1) / BM; + const int smem = STAGES * (BM * BK + BN * BK) * sizeof(nv_bfloat16) + + (BM * BN) * sizeof(float) + 3 * BM * sizeof(float) + STAGES * 8; + const int row_blocks = (N + BM - 1) / BM; + const int total_vtiles = (V + BN - 1) / BN; auto target_i = target.to(torch::kInt32).contiguous(); + // Split the vocab loop across CTAs so the grid fills the GPU: aim for a few + // CTAs per SM, capped by the number of vocab tiles available to split. + int sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + int target_ctas = sm_count * 4; + int n_split = std::max(1, std::min(target_ctas / std::max(row_blocks, 1), total_vtiles)); + + auto part_max = torch::empty({n_split, N}, opts_f); + auto part_sum = torch::empty({n_split, N}, opts_f); + auto part_zt = torch::empty({n_split, N}, opts_f); + + if (smem > 48 * 1024) { + cudaFuncSetAttribute(fused_linear_logp_sm90_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem); + } + + dim3 grid(row_blocks, n_split); fused_linear_logp_sm90_kernel<<>>( - h_tmap, w_tmap, target_i.data_ptr(), bias_ptr, logp.data_ptr(), - lse.data_ptr(), N, D, V); + h_tmap, w_tmap, target_i.data_ptr(), bias_ptr, part_max.data_ptr(), + part_sum.data_ptr(), part_zt.data_ptr(), N, D, V, n_split); + + const int combine_threads = 256; + const int combine_blocks = (N + combine_threads - 1) / combine_threads; + fused_linear_logp_sm90_combine_kernel<<>>( + part_max.data_ptr(), part_sum.data_ptr(), part_zt.data_ptr(), + logp.data_ptr(), lse.data_ptr(), N, n_split); return {logp, lse}; } diff --git a/csrc/cuda/fused_logp_sm90.cu b/csrc/cuda/fused_logp_sm90.cu index a3212dcd..4762614d 100644 --- a/csrc/cuda/fused_logp_sm90.cu +++ b/csrc/cuda/fused_logp_sm90.cu @@ -2,6 +2,7 @@ // Copyright (c) 2026 RL-Kernel Contributors #include "../utils/tma_utils.cuh" +#include #include #include @@ -112,7 +113,9 @@ torch::Tensor fused_logp_sm90_forward(torch::Tensor logits, torch::Tensor labels auto output = torch::empty({B}, logits.options().dtype(torch::kFloat)); CUtensorMap logits_tmap; - init_tensor_map(&logits_tmap, logits.data_ptr(), B, V, 1, TILE_V); + init_tensor_map(&logits_tmap, + reinterpret_cast(logits.data_ptr()), B, V, 1, + TILE_V); int smem_size = (TILE_V * sizeof(nv_bfloat16)) + 16; fused_logp_online_tma_kernel<4><<>>( diff --git a/docs/operators/linear-logp.md b/docs/operators/linear-logp.md index 1e26c23b..f6baca8c 100644 --- a/docs/operators/linear-logp.md +++ b/docs/operators/linear-logp.md @@ -33,20 +33,33 @@ logp.sum().backward() # gradients flow into hidden, lm_head_weight, bias | Backend | Wrapper | Status | | --- | --- | --- | -| CUDA SM90 (Hopper) | `FusedLinearLogpSM90Op` | TMA + WGMMA streaming forward, online softmax in smem; chunked backward. Compiles for `sm_90a`; **numerics pending validation on an SM90 GPU.** | +| CUDA SM90 (Hopper) | `FusedLinearLogpSM90Op` | TMA-streamed, software-pipelined tensor-core forward (`mma.sync.m16n8k16`), online softmax in smem; chunked backward. Compiles for `sm_90a`; **validated fp32-accurate on H100.** Falls back to Triton/native for fp32/fp16 inputs or hidden dims not divisible by 32. | | CUDA / ROCm (Triton) | `TritonLinearLogpOp` | Triton online-softmax forward; Liger-style chunked backward (cuBLAS matmuls, deterministic). Phase 1. | | PyTorch native | `NativeLinearLogpOp` | Naive `F.linear` + `log_softmax` + `gather` reference; CPU / Triton-less fallback. | The SM90 backend (`csrc/cuda/fused_linear_logp_sm90.cu`) streams hidden/weight -tiles via TMA, computes each logit tile with WGMMA (`m64n64k16`), folds it into a -per-row online softmax in shared memory, and gathers the target logit — never +tiles via TMA (`cp.async.bulk.tensor`, mbarrier-completed, double-buffered), +contracts each `[BM, BN]` logit tile with the warp-level tensor-core MMA path +(`ldmatrix` + `mma.sync.m16n8k16`, fp32 accumulation), folds it into a per-row +online softmax in shared memory, and gathers the target logit — never materializing `[N, V]`. It is **build-guarded**: only compiled when the extension -is built with `KERNEL_ALIGN_FORCE_SM90=1` on an SM90 device (WGMMA is Hopper-only, -`sm_90a`), and the registry only selects it when `cc_major == 9` and the symbol is -present. The forward kernel requires bf16 hidden/weight; the backward reuses the -deterministic chunked path. It assembles cleanly under CUDA 13.1 `ptxas`, but the -layout-sensitive pieces (WGMMA descriptors, accumulator→smem mapping, B-operand -transpose) require validation on Hopper hardware. +is built with `KERNEL_ALIGN_FORCE_SM90=1` on an SM90 device (TMA/`sm_90a`), and +the registry only selects it when `cc_major == 9` and the symbol is present. The +forward kernel requires bf16 hidden/weight with `D % 32 == 0`; for any other input +the op transparently falls back to the Triton (else native) backend. The backward +reuses the deterministic chunked path. + +The forward is fp32-accurate (matches the fp32 reference to ~1e-3, and the Triton +op's forward to ~1e-5 with bitwise-identical gradients) and **throughput-tuned**: +it runs at **~2× the Triton online-softmax forward** on H100 (bf16, N=4096, +D=2048) across `V ∈ {32768, 50257, 131072}`, while keeping the zero-`[N,V]` +headline — peak forward memory is the per-CTA shared-memory tiles, independent of +`V`. Two design choices get it there: **register-blocked M-tiling** (`BM = 256` +rows/CTA, so the weight matrix is re-read only `N/BM` times from HBM) and +**split-V** (the vocab loop is partitioned across `blockIdx.y` so the grid fills +all SMs, each split emitting a partial online-softmax state that a tiny combine +kernel merges). Remaining headroom: a register-resident softmax epilogue (drop the +`[BM,BN]` smem logit round-trip) and warp-specialized TMA producers. The backward chunks over the token dimension: for each chunk it materializes only `[chunk, V]` logits, recomputes the softmax from scratch, and forms the three diff --git a/rl_engine/kernels/ops/cuda/loss/linear_logp.py b/rl_engine/kernels/ops/cuda/loss/linear_logp.py index 191e9f7c..3ce798e1 100644 --- a/rl_engine/kernels/ops/cuda/loss/linear_logp.py +++ b/rl_engine/kernels/ops/cuda/loss/linear_logp.py @@ -14,6 +14,33 @@ # at ~chunk*V instead of N*V. _BWD_CHUNK_ELEMS = 1 << 24 +# Hidden-dim slice the SM90 forward streams per TMA load; D must be a multiple of +# it (mirrors `constexpr int BK` in csrc/cuda/fused_linear_logp_sm90.cu). +_SM90_BK = 32 + + +def _sm90_supported(hidden: torch.Tensor, lm_head_weight: torch.Tensor) -> bool: + """Whether the bf16 TMA+MMA forward can run these inputs directly.""" + return ( + hidden.is_cuda + and hidden.dtype == torch.bfloat16 + and lm_head_weight.dtype == torch.bfloat16 + and hidden.size(-1) % _SM90_BK == 0 + ) + + +def _fallback_op(): + """Portable op for inputs the SM90 forward cannot take (fp32/fp16, or a hidden + dim not divisible by the kernel's K slice). Prefers Triton, else native.""" + try: + from rl_engine.kernels.ops.triton.loss.linear_logp import TritonLinearLogpOp + + return TritonLinearLogpOp() + except Exception: # pragma: no cover - Triton missing + from rl_engine.kernels.ops.pytorch.loss.linear_logp import NativeLinearLogpOp + + return NativeLinearLogpOp() + class _FusedLinearLogpSM90Function(torch.autograd.Function): """SM90 TMA+WGMMA fused forward + Liger-style chunked backward. @@ -108,13 +135,14 @@ def apply( target_ids: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if not hidden.is_cuda: - raise RuntimeError("FusedLinearLogpSM90Op requires CUDA tensors.") - if hidden.dtype != torch.bfloat16 or lm_head_weight.dtype != torch.bfloat16: - raise RuntimeError("FusedLinearLogpSM90Op requires bfloat16 hidden/weight.") if lm_head_weight.size(-1) != hidden.size(-1): raise ValueError( f"hidden dim {hidden.size(-1)} must match lm_head_weight dim " f"{lm_head_weight.size(-1)}" ) + # The compiled forward is bf16-only and needs D % BK == 0. For anything + # else (fp32/fp16 references, awkward hidden dims, CPU tensors) fall back + # to a portable backend so the op stays a drop-in for any input. + if not _sm90_supported(hidden, lm_head_weight): + return _fallback_op()(hidden, lm_head_weight, target_ids, bias) return _FusedLinearLogpSM90Function.apply(hidden, lm_head_weight, bias, target_ids) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 558574e2..b2d75fb2 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -123,8 +123,8 @@ def _adjust_priority_for_hardware(self): if OpBackend.CUDA_FUSED_LOGP_SM90 not in logp_list: logp_list.insert(0, OpBackend.CUDA_FUSED_LOGP_SM90) - # The fused linear-logp SM90 kernel uses WGMMA, which is Hopper-only - # (sm_90a) -- gate strictly on cc_major == 9, not >= 9. + # The fused linear-logp SM90 kernel uses TMA bulk-tensor copies built + # for sm_90a -- gate strictly on cc_major == 9 (Hopper), not >= 9. linear_logp_compiled = _EXT_AVAILABLE and hasattr(_C, "fused_linear_logp_sm90") if linear_logp_compiled and cc_major == 9: ll_list = self._priority_map["cuda"]["linear_logp"] diff --git a/tests/test_linear_logp.py b/tests/test_linear_logp.py index 998f8cf6..367b0fc3 100644 --- a/tests/test_linear_logp.py +++ b/tests/test_linear_logp.py @@ -18,6 +18,43 @@ reason="Triton linear log-prob requires a CUDA device and Triton.", ) + +def _sm90_available(): + """SM90 forward needs a Hopper GPU and the kernel compiled into the extension.""" + if not torch.cuda.is_available(): + return False + try: + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + if not (_EXT_AVAILABLE and hasattr(_C, "fused_linear_logp_sm90")): + return False + except Exception: # pragma: no cover + return False + return torch.cuda.get_device_capability()[0] == 9 + + +requires_sm90 = pytest.mark.skipif( + not _sm90_available(), + reason="Fused linear log-prob SM90 kernel requires a Hopper (sm_90) GPU with the " + "extension built KERNEL_ALIGN_FORCE_SM90=1.", +) + +# SM90 forward needs bf16 and a hidden dim that is a multiple of the kernel's K +# slice (32); N / V are deliberately left unaligned to the 64-wide tiles. +_SM90_N = 96 +_SM90_D = 128 +_SM90_V = 500 + + +def _sm90_inputs(seed, *, bias=True, dtype=torch.bfloat16, lead=None): + gen = torch.Generator(device="cuda").manual_seed(seed) + lead = lead or (_SM90_N,) + hidden = torch.randn(*lead, _SM90_D, generator=gen, device="cuda", dtype=dtype) + weight = torch.randn(_SM90_V, _SM90_D, generator=gen, device="cuda", dtype=dtype) + bias_t = torch.randn(_SM90_V, generator=gen, device="cuda", dtype=dtype) if bias else None + target = torch.randint(0, _SM90_V, lead, generator=gen, device="cuda") + return hidden, weight, target, bias_t + # Deliberately non-multiples of the kernel block sizes (32 / 64 / 64). _N = 40 _D = 80 @@ -145,6 +182,107 @@ def test_triton_large_vocab_smoke(): assert out.shape == (8,) and torch.isfinite(out).all() +@requires_sm90 +def test_sm90_forward_matches_native_bf16(): + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + + sm90 = FusedLinearLogpSM90Op() + hidden, weight, target, bias = _sm90_inputs(11) + # The kernel matmul accumulates in fp32 (tensor cores), so the oracle uses the + # fp32-upcast inputs -- like the Triton bf16 test. + ref = NativeLinearLogpOp()(hidden.float(), weight.float(), target, bias.float()) + out = sm90(hidden, weight, target, bias) + assert out.dtype == torch.float32 + assert torch.allclose(out, ref, atol=2e-2) + + +@requires_sm90 +def test_sm90_forward_no_bias(): + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + + sm90 = FusedLinearLogpSM90Op() + hidden, weight, target, _ = _sm90_inputs(12, bias=False) + ref = NativeLinearLogpOp()(hidden.float(), weight.float(), target, None) + out = sm90(hidden, weight, target) + assert torch.allclose(out, ref, atol=2e-2) + + +@requires_sm90 +@pytest.mark.parametrize("use_bias", [True, False]) +def test_sm90_forward_backward_matches_triton(use_bias): + # The SM90 forward is fp32-accurate and the backward reuses the same + # deterministic chunked path as the Triton op, so both match very tightly. + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + from rl_engine.kernels.ops.triton.loss.linear_logp import TritonLinearLogpOp + + sm90, trit = FusedLinearLogpSM90Op(), TritonLinearLogpOp() + hidden, weight, target, bias = _sm90_inputs(13, bias=use_bias) + grad_out = torch.randn(_SM90_N, device="cuda") + + def run(op): + h = hidden.detach().clone().requires_grad_(True) + w = weight.detach().clone().requires_grad_(True) + b = bias.detach().clone().requires_grad_(True) if bias is not None else None + out = op(h, w, target, b) + out.backward(grad_out) + return out.detach(), h.grad, w.grad, (b.grad if b is not None else None) + + so, sh, sw, sb = run(sm90) + to, th, tw, tb = run(trit) + assert torch.allclose(so, to, atol=1e-3) + assert torch.allclose(sh, th, atol=2e-3) + assert torch.allclose(sw, tw, atol=2e-3) + if use_bias: + assert torch.allclose(sb, tb, atol=2e-3) + + +@requires_sm90 +def test_sm90_preserves_leading_shape(): + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + + sm90 = FusedLinearLogpSM90Op() + hidden, weight, target, bias = _sm90_inputs(14, lead=(6, 5)) + out = sm90(hidden, weight, target, bias) + assert out.shape == (6, 5) + ref = NativeLinearLogpOp()(hidden.float(), weight.float(), target, bias.float()) + assert torch.allclose(out, ref, atol=2e-2) + + +@requires_sm90 +def test_sm90_large_vocab_smoke(): + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + + sm90 = FusedLinearLogpSM90Op() + hidden = torch.randn(40, 256, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(50257, 256, device="cuda", dtype=torch.bfloat16) + target = torch.randint(0, 50257, (40,), device="cuda") + out = sm90(hidden, weight, target) + assert out.shape == (40,) and torch.isfinite(out).all() + + +@requires_sm90 +def test_sm90_falls_back_for_unsupported_inputs(): + # fp32 inputs and a hidden dim not divisible by the kernel's K slice are not + # handled by the compiled forward; the op must fall back instead of erroring. + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + + sm90 = FusedLinearLogpSM90Op() + + fp32 = _sm90_inputs(15, dtype=torch.float32) + out = sm90(*fp32) + ref = NativeLinearLogpOp()(*fp32) + assert torch.allclose(out, ref, atol=1e-3) + + # bf16 but D=80 (not a multiple of 32) -> fallback path. + gen = torch.Generator(device="cuda").manual_seed(16) + hidden = torch.randn(40, 80, device="cuda", dtype=torch.bfloat16, generator=gen) + weight = torch.randn(300, 80, device="cuda", dtype=torch.bfloat16, generator=gen) + target = torch.randint(0, 300, (40,), device="cuda", generator=gen) + out = sm90(hidden, weight, target) + ref = NativeLinearLogpOp()(hidden.float(), weight.float(), target, None) + assert torch.allclose(out, ref, atol=2e-2) + + def test_registry_dispatch_matches_native(): from rl_engine.kernels.registry import kernel_registry from rl_engine.platforms.device import device_ctx From 6c366c6134c886bfc56543b1a6ef4e2d84d7707a Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Mon, 15 Jun 2026 08:12:42 +0000 Subject: [PATCH 3/8] Fix the docs --- docs/design/fused-linear-logp.md | 216 ------------------------------- docs/operators/linear-logp.md | 83 ++++++------ 2 files changed, 38 insertions(+), 261 deletions(-) delete mode 100644 docs/design/fused-linear-logp.md diff --git a/docs/design/fused-linear-logp.md b/docs/design/fused-linear-logp.md deleted file mode 100644 index d3f66204..00000000 --- a/docs/design/fused-linear-logp.md +++ /dev/null @@ -1,216 +0,0 @@ -# Design: Fused Linear LogProb (no logit materialization) + fused backward - -## Motivation - -Every current log-prob path — `logp`, `ratio_kl`, `grpo_loss` — takes the -**logits** `[B, S, V]` as input, so the `[B, S, V]` tensor is already resident in -HBM before any fusion happens. For large-vocabulary models that tensor dominates -memory: at `V≈150k`, `B·S≈8k` in bf16 it is ~2.4 GB for the forward activation -alone, plus the same again for its gradient. - -This op fuses the **LM-head projection itself** into the log-prob reduction: - -``` -logits = hidden @ Wᵀ (+ b) # [N, V] — NEVER materialized -logp[n] = logits[n, t[n]] − logsumexp(logits[n, :]) -``` - -The forward streams the vocab in blocks (block matmul → online softmax), keeping -only per-row scalars. The backward recomputes the logit tiles from `hidden`/`W` -instead of storing them — trading a second matmul for the `[N, V]` storage. The -freed memory buys larger batches / longer CoT. - -This is the "cut cross-entropy" / fused-linear-cross-entropy pattern, specialized -to **selected-token log-prob** (the RL training quantity) rather than mean CE. - -## Scope - -- **New op type** `linear_logp` — does **not** replace `logp` (still useful when - logits already exist). It is the differentiable hidden→logp primitive. -- **In scope**: forward + backward w.r.t. `hidden` and `lm_head_weight` (and bias); - Triton, CUDA generic, CUDA SM90 TMA, ROCm backends; native PyTorch reference; - registry wiring; tests; benchmarks; docs. -- **Out of scope (follow-ups)**: rewiring `ratio_kl`/`grpo_loss` to consume hidden - states directly; tensor/sequence-parallel vocab sharding; fp8 weights. - -## Math - -Per row `n`, target token `t = t[n]`, logits `z = Hₙ Wᵀ (+ b) ∈ ℝ^V`: - -**Forward** -``` -lse = logsumexp(z); logp = z[t] − lse -``` -Computed by streaming vocab blocks `Vblk`: each block does a `[Nblk, D]·[D, Vblk]` -matmul tile, folds the tile into an online-softmax state `(m, s)` (running max, -rescaled sum), and captures `z[t]` when `t` lands in the block. Saved for backward: -**only** `lse` (or `(m,s)`) and the gathered `z[t]` — per-row scalars, no `[N,V]`. - -**Backward** — given `g = dL/dlogp ∈ ℝ^N`: -``` -dL/dz[v] = g · (1[v==t] − p[v]), p = softmax(z) # never stored; p recomputed -dL/dHₙ = g · (W[t] − Σ_v p[v]·W[v]) # = (dL/dz) @ W -dL/dW[v] += g · (1[v==t] − p[v]) · Hₙ (reduce over n) -dL/db[v] += g · (1[v==t] − p[v]) (reduce over n) -``` -The backward re-streams vocab blocks, recomputes `z_blk = Hₙ W_blkᵀ`, then -`p_blk = exp(z_blk − lse)`, and accumulates the three gradients. `dL/dW` and -`dL/db` are reductions **over the token dimension** → handled by atomic-add into -the `[V,D]`/`[V]` grad buffers (or a token-blocked two-pass to avoid atomics). - -All reductions/accumulation in **fp32**; inputs may be bf16/fp16. The `c·(onehot − -softmax)` shape is identical to the existing `_ratio_kl_bwd_kernel` — the only new -piece is the surrounding matmul. - -## Public API - -```python -linear_logp = kernel_registry.get_op("linear_logp") - -logp = linear_logp( - hidden, # [B, S, D] or [N, D] (differentiable) - lm_head_weight,# [V, D] (differentiable) - target_ids, # [B, S] or [N] int - bias=None, # [V] optional (differentiable) -) # -> [B, S] or [N], differentiable w.r.t. hidden, weight, bias -``` - -Wrapped in a `torch.autograd.Function` so it is a drop-in differentiable node. -Mask handling follows the existing convention (caller zeroes masked positions, or -pass an ignore-index for `target_ids`). - -## Tensor contract - -| Arg | Shape | Dtype | Notes | -| --- | --- | --- | --- | -| `hidden` | `[N, D]` | bf16/fp16/fp32 | contiguous; leading dims flattened | -| `lm_head_weight` | `[V, D]` | bf16/fp16/fp32 | contiguous (row-major over V) | -| `target_ids` | `[N]` | int32/int64 | in `[0, V)`; optional ignore-index | -| `bias` | `[V]` | float | optional | -| `logp` (out) | `[N]` | fp32 | `z[t] − logsumexp(z)` | - -## Implementation phases - -Ordering (per request): **Triton → CUDA generic → CUDA SM90 TMA → ROCm.** Each -phase is independently shippable and gated on matching the prior phase's reference -to tolerance. - -### Phase 1 — Triton (portable baseline + tolerance target) - -The semantic reference that CUDA/ROCm must match. - -- `rl_engine/kernels/ops/pytorch/loss/linear_logp.py` — `NativeLinearLogpOp`: - **naive** pure-PyTorch reference — a single `F.linear` (materializes the full - `[N, V]` logits) → `log_softmax` → `gather`, differentiable via autograd. No - chunking: this is the straightforward, obviously-correct oracle that the fused - kernels are validated against (and the baseline the benchmark measures the VRAM - win against). Also the CPU / Triton-less fallback. -- `rl_engine/kernels/ops/triton/loss/linear_logp.py`: - - `_linear_logp_fwd_kernel`: program per token-block; loop vocab blocks, `tl.dot` - for the `[Nblk,D]·[D,Vblk]` tile, online-softmax state in registers; store - `logp`, `lse`, gathered `z[t]`. - - `_linear_logp_bwd_kernel`: re-stream vocab, recompute `p_blk`, accumulate - `dH` (per token-block, local) and `dW`/`db` (atomic-add into global buffers). - - `_LinearLogpFunction(autograd.Function)` + `TritonLinearLogpOp`. -- Tests `tests/test_linear_logp.py`: native-vs-`F.linear→log_softmax→gather`; - Triton fwd vs native; Triton bwd vs autograd (gradcheck-style on `hidden`, - `weight`, `bias`) at `atol=1e-3` bf16 / tighter fp32; ignore-index; large `V` - (50257, 131072) memory-flat check; registry dispatch. -- Benchmark `benchmarks/benchmark_linear_logp.py`: sweep `V`; report fwd / fwd+bwd - latency and **peak VRAM vs the `F.linear`+`log_softmax` baseline** (the headline - number — should be flat in `V`). -- Docs `docs/operators/linear-logp.md` (+ `.nav.yml`, `operators/README.md`). - -**Acceptance**: Triton matches native within tolerance; benchmark shows peak VRAM -independent of `V` and a clear win vs the materializing baseline. - -### Phase 2 — CUDA generic fallback (`fused_linear_logp_kernel.cu`) - -Portable CUDA (no arch intrinsics) for all SM≥70. - -- `csrc/fused_linear_logp_kernel.cu`: forward + backward. Block per token-tile; - `cp.async` double-buffer `W` tiles into shared memory; FMA or `mma.sync` matmul; - online-softmax state in registers (reuse `LogSumExpState` / `merge_logsumexp_state` - and `blockReduce*` from `fused_logp_kernel.cu`); fp32 accumulate. Backward - recomputes tiles; atomic-add into `dW`/`db`. -- Bind in `csrc/ops.cpp` under the existing `KERNEL_ALIGN_WITH_CUDA` guard: - `fused_linear_logp_forward`, `fused_linear_logp_backward`. -- `setup.py`: add `csrc/fused_linear_logp_kernel.cu` to `cuda_sources`; env-var - block-size toggles mirroring the `FUSED_LOGP_*` macros. -- `rl_engine/kernels/ops/cuda/loss/linear_logp.py` — `FusedLinearLogpGenericOp` - wrapping fwd/bwd in an `autograd.Function`. - -**Acceptance**: matches Triton/native to tolerance; ≥ parity with Triton latency -on the generic path; same flat-VRAM property. - -### Phase 3 — CUDA SM90 TMA/WGMMA (`fused_linear_logp_sm90.cu`) - -Advanced streaming for Hopper/Blackwell-class (`cc ∈ {9,10,12}`). - -- `csrc/cuda/fused_linear_logp_sm90.cu`: TMA (`cuTensorMapEncodeTiled`) bulk loads - of `H` and `W` tiles into shared memory; WGMMA (`wgmma.mma_async`) for the - matmul; warp-specialized producer (TMA) / consumer (online softmax) with - `mbarrier`; keep state in registers. Backward symmetric with tile recompute. - **Heed the prior TMA lessons** (see issue #91 / the SM120 work): box inner dim - ≤ 256 elems, CUDA 12.9+ for `cuda::maximum`, arch `compute_90a`/`120a`. -- Gate behind `KERNEL_ALIGN_FORCE_SM90` in `setup.py` (arch-specific gencode), bind - under `KERNEL_ALIGN_WITH_SM90` in `ops.cpp`. -- `FusedLinearLogpSM90Op` in the CUDA wrapper; registry auto-prioritizes it in - `_adjust_priority_for_hardware` when `hasattr(_C, "fused_linear_logp_sm90")` and - `cc_major in (9,10,12)` (same pattern as `fused_logp_sm90`). - -**Acceptance**: matches reference to tolerance on SM90+; measurable speedup over -the generic CUDA path; builds cleanly on SM120/CUDA 13 (the dev box) and is cleanly -disabled where unsupported. - -### Phase 4 — ROCm / CDNA - -- **Free coverage first**: the Phase-1 Triton kernel compiles on ROCm — register - `TRITON_LINEAR_LOGP` in the `rocm` priority map so ROCm works from Phase 1. -- **Native HIP** (`csrc/rocm/...` or HIP-ified): wavefront=64-aware reductions and - LDS layouts; replace TMA with **manual double-buffering** into LDS; CDNA MFMA for - the matmul where available. Build via the ROCm/HIP path in `setup.py`. - -**Acceptance**: ROCm Triton matches reference; native HIP ≥ Triton on CDNA. - -## Registry wiring - -Add to `OpBackend`: `TRITON_LINEAR_LOGP`, `PYTORCH_LINEAR_LOGP`, -`CUDA_FUSED_LINEAR_LOGP_GENERIC`, `CUDA_FUSED_LINEAR_LOGP_SM90`. Add `linear_logp` -to each platform map: - -``` -cuda: [CUDA_FUSED_LINEAR_LOGP_GENERIC, TRITON_LINEAR_LOGP, PYTORCH_LINEAR_LOGP] - (+ CUDA_FUSED_LINEAR_LOGP_SM90 inserted at front by _adjust_priority_for_hardware) -rocm: [TRITON_LINEAR_LOGP, PYTORCH_LINEAR_LOGP] (+ native HIP once landed) -cpu: [PYTORCH_LINEAR_LOGP] -``` - -## Numerics & testing strategy - -- The **Triton kernel is the tolerance target**; native PyTorch is the correctness - oracle. Every native backend is gradchecked against `F.linear → log_softmax → - gather` autograd. -- fp32 accumulation throughout; bf16/fp16 I/O. Tolerance `atol≈1e-3` (bf16), - tighter for fp32. -- Edge cases: `V` not a multiple of the block; ignore-index targets; single-token - rows; very large `V` (memory-flat assertion). - -## Risks / open questions - -- **`dW` atomics**: atomic-add into `[V,D]` may contend at small `V`. Fallback is a - token-blocked two-pass (deterministic, atomic-free) — decide per-backend by - benchmark. Determinism flag may be wanted for reproducible RL runs. -- **bf16 weight gradient precision**: accumulate `dW` in fp32, cast on store. -- **Backward recompute cost**: ~2× forward matmul FLOPs in backward; net win is the - memory, not FLOPs — benchmark must report both so the tradeoff is explicit. -- **TMA portability**: the SM90 path repeats the constraints that bit us before - (box dims, CUDA version, arch flags); keep it default-off and feature-detected. - -## Relevant files (to add / touch) - -- `rl_engine/kernels/ops/{pytorch,triton,cuda}/loss/linear_logp.py` -- `csrc/fused_linear_logp_kernel.cu`, `csrc/cuda/fused_linear_logp_sm90.cu` -- `csrc/ops.cpp`, `setup.py`, `rl_engine/kernels/registry.py` -- `tests/test_linear_logp.py`, `benchmarks/benchmark_linear_logp.py` -- `docs/operators/linear-logp.md`, `docs/.nav.yml`, `docs/operators/README.md` diff --git a/docs/operators/linear-logp.md b/docs/operators/linear-logp.md index f6baca8c..a7267884 100644 --- a/docs/operators/linear-logp.md +++ b/docs/operators/linear-logp.md @@ -1,8 +1,8 @@ # Fused Linear LogP Fused Linear LogP computes per-token selected log-probabilities directly from -**hidden states and the LM-head weight** — `log_softmax(hidden @ Wᵀ + b)[target]` — -**without ever materializing the `[N, V]` logits**. It targets large-vocabulary RL +hidden states and the LM-head weight `log_softmax(hidden @ Wᵀ + b)[target]` +without ever materializing the `[N, V]` logits. It targets large-vocabulary RL post-training, where the `[B, S, V]` logits activation (and its gradient) dominate memory. The forward streams the vocab in blocks through an online softmax; the backward recomputes the logit tiles instead of storing them, trading compute for @@ -33,7 +33,7 @@ logp.sum().backward() # gradients flow into hidden, lm_head_weight, bias | Backend | Wrapper | Status | | --- | --- | --- | -| CUDA SM90 (Hopper) | `FusedLinearLogpSM90Op` | TMA-streamed, software-pipelined tensor-core forward (`mma.sync.m16n8k16`), online softmax in smem; chunked backward. Compiles for `sm_90a`; **validated fp32-accurate on H100.** Falls back to Triton/native for fp32/fp16 inputs or hidden dims not divisible by 32. | +| CUDA SM90 (Hopper) | `FusedLinearLogpSM90Op` | TMA-streamed, Double Buffering, tensor-core forward (`mma.sync.m16n8k16`), online softmax in smem; chunked backward. Compiles for `sm_90a`; validated fp32-accurate on H100. Falls back to Triton/native for fp32/fp16 inputs or hidden dims not divisible by 32. | | CUDA / ROCm (Triton) | `TritonLinearLogpOp` | Triton online-softmax forward; Liger-style chunked backward (cuBLAS matmuls, deterministic). Phase 1. | | PyTorch native | `NativeLinearLogpOp` | Naive `F.linear` + `log_softmax` + `gather` reference; CPU / Triton-less fallback. | @@ -47,37 +47,15 @@ is built with `KERNEL_ALIGN_FORCE_SM90=1` on an SM90 device (TMA/`sm_90a`), and the registry only selects it when `cc_major == 9` and the symbol is present. The forward kernel requires bf16 hidden/weight with `D % 32 == 0`; for any other input the op transparently falls back to the Triton (else native) backend. The backward -reuses the deterministic chunked path. - -The forward is fp32-accurate (matches the fp32 reference to ~1e-3, and the Triton -op's forward to ~1e-5 with bitwise-identical gradients) and **throughput-tuned**: -it runs at **~2× the Triton online-softmax forward** on H100 (bf16, N=4096, -D=2048) across `V ∈ {32768, 50257, 131072}`, while keeping the zero-`[N,V]` -headline — peak forward memory is the per-CTA shared-memory tiles, independent of -`V`. Two design choices get it there: **register-blocked M-tiling** (`BM = 256` -rows/CTA, so the weight matrix is re-read only `N/BM` times from HBM) and -**split-V** (the vocab loop is partitioned across `blockIdx.y` so the grid fills -all SMs, each split emitting a partial online-softmax state that a tiny combine -kernel merges). Remaining headroom: a register-resident softmax epilogue (drop the -`[BM,BN]` smem logit round-trip) and warp-specialized TMA producers. - -The backward chunks over the token dimension: for each chunk it materializes only -`[chunk, V]` logits, recomputes the softmax from scratch, and forms the three -gradients with cuBLAS matmuls (`grad_hidden = dz @ W`, `grad_weight += dzᵀ @ X`). -`grad_weight` is accumulated in sequential loop order, so it is **atomic-free and -bitwise-deterministic** while peak backward memory stays `chunk·V` instead of `N·V`. - -The native op materializes the full `[N, V]` logits and is the correctness oracle; -the Triton op is the portable, fp32-accurate baseline that the future CUDA generic, -CUDA SM90 (TMA/WGMMA), and native ROCm backends are validated against. See the -[design doc](../design/fused-linear-logp.md) for the phased plan. +reuses the deterministic chunked path. The native op materializes the full `[N, V]` +logits and is the correctness oracle. ## Tensor Contract | Argument | Shape | Dtype | Requirements | | --- | --- | --- | --- | | `hidden` | `[N, D]` / `[B, S, D]` | bf16 / fp16 / fp32 | Differentiable; contiguous. | -| `lm_head_weight` | `[V, D]` | bf16 / fp16 / fp32 | Differentiable; contiguous (row-major over V). | +| `lm_head_weight` | `[V, D]` | bf16 / fp16 / fp32 | Differentiable; contiguous. | | `target_ids` | `[N]` / `[B, S]` | int | Token id per position, in `[0, V)`. | | `bias` | `[V]` | float | Optional; differentiable. | | Output | `[N]` / `[B, S]` | float32 | `z[target] − logsumexp(z)` per position. | @@ -104,25 +82,40 @@ python benchmarks/benchmark_linear_logp.py python benchmarks/benchmark_linear_logp.py --configs "4096,2048,32768;4096,2048,131072" ``` -Indicative results (RTX PRO 6000, SM120, bf16, N=4096, D=2048; native vs Triton): +Measured on an **NVIDIA H100 80GB** (SM90), bf16, N=4096, D=2048, CUDA 12.8. -| shape (N×H×V) | fwd | fwd+bwd | peak fwd VRAM (native → Triton) | +**Forward latency (ms) and peak forward VRAM:** + +| shape (N×D×V) | native | Triton | **SM90** | SM90 vs Triton | peak fwd VRAM (native → fused) | +| --- | --- | --- | --- | --- | --- | +| 4096×2048×32768 | 1.79 | 6.42 | **3.41** | **1.88×** | 1280 MB → **~0 MB** | +| 4096×2048×50257 | 9.96 | 9.82 | **4.88** | **2.01×** | 1965 MB → **~0 MB** | +| 4096×2048×131072 | 7.28 | 25.56 | **12.88** | **1.98×** | 5120 MB → **~0 MB** | + +**Forward + backward latency (ms):** + +| shape (N×D×V) | native | Triton | **SM90** | | --- | --- | --- | --- | -| 4096×2048×32768 | 0.53× | 0.43× | 1280 MB → ~0 MB | -| 4096×2048×50257 | 0.99× | 0.46× | 1965 MB → ~0 MB | -| 4096×2048×131072 | 0.64× | 0.23× | 5120 MB → ~0 MB | - -The headline is memory: the native path allocates the `[N, V]` logits (forward peak -scales with `V`), while the fused op streams them online — its forward peak is -**independent of `V`**, and the chunked backward only ever holds `chunk·V`. The -forward matmul keeps operands in their native dtype, so bf16/fp16 inputs run on -tensor cores (fp32 accumulation) and the forward lands near cuBLAS parity; the fp32 -path stays full-precision (`input_precision="ieee"`) for its role as the tolerance -reference. The backward runs at ~2–4× native: it recomputes the logits per chunk -(native keeps the `[N, V]` it already materialized), which is the compute-for-memory -trade. Closing that gap and the absolute forward latency are the job of the native -CUDA generic / SM90 (TMA/WGMMA) backends in later phases; Phase 1 delivers the -memory reduction and the correctness/tolerance baseline. +| 4096×2048×32768 | 4.25 | 15.86 | 12.69 | +| 4096×2048×50257 | 23.29 | 47.20 | 42.20 | +| 4096×2048×131072 | 17.05 | 117.62 | 104.97 | + +**Memory**: the native path allocates the `[N, V]` logits (forward +peak scales with `V`), while the fused op streams them online — its forward peak is +just the per-CTA shared-memory tiles, **independent of `V`** (≈0 MB of activation), +and the chunked backward only ever holds `chunk·V`. That freed memory is what lets +you grow the batch or the CoT length. + +**Latency**: the SM90 forward runs at **~2× the memory-free Triton baseline** +across the vocab range — from TMA double-buffering, register-blocked `mma.sync` +M-tiling (`BM=256`, so the weight matrix is re-read `N/BM` times), and split-V over +the grid. It also beats the materializing native path at moderate vocab (2.0× at +V=50257); at the extremes the native cuBLAS GEMM is still faster in raw ms, but only +by paying the 1.3–5 GB `[N, V]` allocation the fused op avoids. The backward is the +shared deterministic chunked-recompute path (recomputes logit tiles rather than +storing them), so the fused forward+backward already beats Triton's. Closing the +remaining forward gap to native's cuBLAS GEMM (WGMMA, a register-resident softmax +epilogue) and a fully fused CUDA backward are future work. ## Tests From 3138dbbeaec2e17dc5c492f1a67db38e58be2362 Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Mon, 15 Jun 2026 08:36:02 +0000 Subject: [PATCH 4/8] style: black (E305 blank lines) --- benchmarks/benchmark_linear_logp.py | 1 + tests/test_linear_logp.py | 1 + 2 files changed, 2 insertions(+) diff --git a/benchmarks/benchmark_linear_logp.py b/benchmarks/benchmark_linear_logp.py index 81d91e66..55a0cade 100644 --- a/benchmarks/benchmark_linear_logp.py +++ b/benchmarks/benchmark_linear_logp.py @@ -43,6 +43,7 @@ def _maybe_sm90_op(): return FusedLinearLogpSM90Op() + # (num_tokens, hidden_dim, vocab) DEFAULT_CONFIGS = [ (4096, 2048, 32768), diff --git a/tests/test_linear_logp.py b/tests/test_linear_logp.py index 367b0fc3..80310908 100644 --- a/tests/test_linear_logp.py +++ b/tests/test_linear_logp.py @@ -55,6 +55,7 @@ def _sm90_inputs(seed, *, bias=True, dtype=torch.bfloat16, lead=None): target = torch.randint(0, _SM90_V, lead, generator=gen, device="cuda") return hidden, weight, target, bias_t + # Deliberately non-multiples of the kernel block sizes (32 / 64 / 64). _N = 40 _D = 80 From 01c53b2bdf9659144e0a54d007b96475542ed7df Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Mon, 15 Jun 2026 08:44:26 +0000 Subject: [PATCH 5/8] style: trim trailing whitespace --- docs/operators/linear-logp.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/operators/linear-logp.md b/docs/operators/linear-logp.md index a7267884..2d8b8ce9 100644 --- a/docs/operators/linear-logp.md +++ b/docs/operators/linear-logp.md @@ -47,7 +47,7 @@ is built with `KERNEL_ALIGN_FORCE_SM90=1` on an SM90 device (TMA/`sm_90a`), and the registry only selects it when `cc_major == 9` and the symbol is present. The forward kernel requires bf16 hidden/weight with `D % 32 == 0`; for any other input the op transparently falls back to the Triton (else native) backend. The backward -reuses the deterministic chunked path. The native op materializes the full `[N, V]` +reuses the deterministic chunked path. The native op materializes the full `[N, V]` logits and is the correctness oracle. ## Tensor Contract From bc93e3602556cf77f91fd0652997251c66f4284f Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Tue, 16 Jun 2026 13:49:38 +0000 Subject: [PATCH 6/8] (bug) add target/bias validation --- csrc/cuda/fused_linear_logp_sm90.cu | 7 +++++++ rl_engine/kernels/ops/cuda/loss/linear_logp.py | 18 +++++++++++++++--- tests/test_linear_logp.py | 16 ++++++++++++++++ 3 files changed, 38 insertions(+), 3 deletions(-) diff --git a/csrc/cuda/fused_linear_logp_sm90.cu b/csrc/cuda/fused_linear_logp_sm90.cu index 34a74e89..c560f7ea 100644 --- a/csrc/cuda/fused_linear_logp_sm90.cu +++ b/csrc/cuda/fused_linear_logp_sm90.cu @@ -310,6 +310,13 @@ std::vector fused_linear_logp_sm90_forward(torch::Tensor hidden, const int V = weight.size(0); TORCH_CHECK(weight.size(1) == D, "hidden/weight hidden-dim mismatch"); TORCH_CHECK(D % BK == 0, "D must be a multiple of ", BK, " for the SM90 kernel"); + TORCH_CHECK(target.numel() == N, "target must have one id per token: expected ", N, + " (hidden rows), got ", target.numel()); + if (bias.has_value()) { + TORCH_CHECK(bias->device() == hidden.device(), + "bias must be on the same device as hidden"); + TORCH_CHECK(bias->numel() == V, "bias must have V=", V, " elements, got ", bias->numel()); + } auto opts_f = hidden.options().dtype(torch::kFloat); auto logp = torch::empty({N}, opts_f); diff --git a/rl_engine/kernels/ops/cuda/loss/linear_logp.py b/rl_engine/kernels/ops/cuda/loss/linear_logp.py index 3ce798e1..e0efd6e9 100644 --- a/rl_engine/kernels/ops/cuda/loss/linear_logp.py +++ b/rl_engine/kernels/ops/cuda/loss/linear_logp.py @@ -140,9 +140,21 @@ def apply( f"hidden dim {hidden.size(-1)} must match lm_head_weight dim " f"{lm_head_weight.size(-1)}" ) - # The compiled forward is bf16-only and needs D % BK == 0. For anything - # else (fp32/fp16 references, awkward hidden dims, CPU tensors) fall back - # to a portable backend so the op stays a drop-in for any input. + n_tokens = hidden.numel() // hidden.size(-1) + if target_ids.numel() != n_tokens: + raise ValueError( + f"target_ids must have one id per token: expected {n_tokens}, " + f"got {target_ids.numel()}" + ) + if bias is not None: + if bias.numel() != lm_head_weight.size(0): + raise ValueError( + f"bias must have V={lm_head_weight.size(0)} elements, got {bias.numel()}" + ) + if bias.device != hidden.device: + raise ValueError( + f"bias device {bias.device} must match hidden device {hidden.device}" + ) if not _sm90_supported(hidden, lm_head_weight): return _fallback_op()(hidden, lm_head_weight, target_ids, bias) return _FusedLinearLogpSM90Function.apply(hidden, lm_head_weight, bias, target_ids) diff --git a/tests/test_linear_logp.py b/tests/test_linear_logp.py index 80310908..e52d3879 100644 --- a/tests/test_linear_logp.py +++ b/tests/test_linear_logp.py @@ -284,6 +284,22 @@ def test_sm90_falls_back_for_unsupported_inputs(): assert torch.allclose(out, ref, atol=2e-2) +@requires_sm90 +def test_sm90_rejects_bad_target_and_bias(): + # Shape/device mismatches must be a clean error, not a CUDA illegal access. + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + + sm90 = FusedLinearLogpSM90Op() + hidden, weight, target, bias = _sm90_inputs(17) + + with pytest.raises((ValueError, RuntimeError)): # wrong target length + sm90(hidden, weight, target[:-1], bias) + with pytest.raises((ValueError, RuntimeError)): # wrong bias length + sm90(hidden, weight, target, bias[:-1]) + with pytest.raises((ValueError, RuntimeError)): # bias on the wrong device + sm90(hidden, weight, target, bias.cpu()) + + def test_registry_dispatch_matches_native(): from rl_engine.kernels.registry import kernel_registry from rl_engine.platforms.device import device_ctx From dc3cd73521340e9b5e1cf101f976b651fc1fa6e4 Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Wed, 17 Jun 2026 15:00:04 +0000 Subject: [PATCH 7/8] check lm_head_weight and hidden device --- csrc/cuda/fused_linear_logp_sm90.cu | 3 +++ rl_engine/kernels/ops/cuda/loss/linear_logp.py | 5 +++++ tests/test_linear_logp.py | 2 ++ 3 files changed, 10 insertions(+) diff --git a/csrc/cuda/fused_linear_logp_sm90.cu b/csrc/cuda/fused_linear_logp_sm90.cu index c560f7ea..ff1ec741 100644 --- a/csrc/cuda/fused_linear_logp_sm90.cu +++ b/csrc/cuda/fused_linear_logp_sm90.cu @@ -302,6 +302,9 @@ std::vector fused_linear_logp_sm90_forward(torch::Tensor hidden, torch::Tensor weight, torch::Tensor target, torch::optional bias) { + TORCH_CHECK(hidden.is_cuda() && weight.is_cuda(), "hidden and weight must be CUDA tensors"); + TORCH_CHECK(weight.device() == hidden.device(), + "lm_head_weight must be on the same device as hidden"); TORCH_CHECK(hidden.scalar_type() == at::kBFloat16, "hidden must be bfloat16"); TORCH_CHECK(weight.scalar_type() == at::kBFloat16, "weight must be bfloat16"); TORCH_CHECK(hidden.is_contiguous() && weight.is_contiguous(), "inputs must be contiguous"); diff --git a/rl_engine/kernels/ops/cuda/loss/linear_logp.py b/rl_engine/kernels/ops/cuda/loss/linear_logp.py index e0efd6e9..a55ea7e4 100644 --- a/rl_engine/kernels/ops/cuda/loss/linear_logp.py +++ b/rl_engine/kernels/ops/cuda/loss/linear_logp.py @@ -140,6 +140,11 @@ def apply( f"hidden dim {hidden.size(-1)} must match lm_head_weight dim " f"{lm_head_weight.size(-1)}" ) + if lm_head_weight.device != hidden.device: + raise ValueError( + f"lm_head_weight device {lm_head_weight.device} must match hidden " + f"device {hidden.device}" + ) n_tokens = hidden.numel() // hidden.size(-1) if target_ids.numel() != n_tokens: raise ValueError( diff --git a/tests/test_linear_logp.py b/tests/test_linear_logp.py index e52d3879..969b5b50 100644 --- a/tests/test_linear_logp.py +++ b/tests/test_linear_logp.py @@ -292,6 +292,8 @@ def test_sm90_rejects_bad_target_and_bias(): sm90 = FusedLinearLogpSM90Op() hidden, weight, target, bias = _sm90_inputs(17) + with pytest.raises((ValueError, RuntimeError)): # weight on the wrong device + sm90(hidden, weight.cpu(), target, bias) with pytest.raises((ValueError, RuntimeError)): # wrong target length sm90(hidden, weight, target[:-1], bias) with pytest.raises((ValueError, RuntimeError)): # wrong bias length From e7966e4b0c95949e57e4eed25e19dc0e0fcd0c6b Mon Sep 17 00:00:00 2001 From: KJLdefeated Date: Mon, 22 Jun 2026 03:20:44 +0000 Subject: [PATCH 8/8] fix the requests --- benchmarks/benchmark_linear_logp.py | 4 +- .../kernels/ops/cuda/loss/linear_logp.py | 55 ++++++--------- .../kernels/ops/pytorch/loss/linear_logp.py | 56 +++++++++++++++ .../kernels/ops/triton/loss/linear_logp.py | 70 +++++++------------ tests/test_linear_logp.py | 23 ++++++ 5 files changed, 128 insertions(+), 80 deletions(-) diff --git a/benchmarks/benchmark_linear_logp.py b/benchmarks/benchmark_linear_logp.py index 55a0cade..2e8ffabf 100644 --- a/benchmarks/benchmark_linear_logp.py +++ b/benchmarks/benchmark_linear_logp.py @@ -87,8 +87,8 @@ def _peak_vram_gb(fn, warmup=3, iters=5): def run_benchmark(args): - if device_ctx.device_type != "cuda": - raise RuntimeError("linear_logp benchmark requires a CUDA device (Triton op is CUDA-only).") + if device_ctx.device_type not in ["cuda", "xpu", "hip"]: + raise RuntimeError("linear_logp benchmark requires a compatible GPU device.") device = device_ctx.device dtype = torch.bfloat16 diff --git a/rl_engine/kernels/ops/cuda/loss/linear_logp.py b/rl_engine/kernels/ops/cuda/loss/linear_logp.py index a55ea7e4..44aa998b 100644 --- a/rl_engine/kernels/ops/cuda/loss/linear_logp.py +++ b/rl_engine/kernels/ops/cuda/loss/linear_logp.py @@ -8,12 +8,9 @@ import torch from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.pytorch.loss.linear_logp import chunked_linear_logp_backward from rl_engine.utils.logger import logger -# Backward token-chunk target (mirrors the Triton op): keep peak backward memory -# at ~chunk*V instead of N*V. -_BWD_CHUNK_ELEMS = 1 << 24 - # Hidden-dim slice the SM90 forward streams per TMA load; D must be a multiple of # it (mirrors `constexpr int BK` in csrc/cuda/fused_linear_logp_sm90.cu). _SM90_BK = 32 @@ -70,36 +67,19 @@ def forward(ctx, hidden, lm_head_weight, bias, target_ids): @staticmethod def backward(ctx, grad_logp): hidden_2d, weight, bias_t, target_1d = ctx.saved_tensors - n, d = hidden_2d.shape - v = weight.shape[0] - dt = weight.dtype - g = grad_logp.reshape(-1).to(torch.float32) - - grad_h = torch.empty_like(hidden_2d, dtype=torch.float32) - grad_w = torch.zeros(v, d, device=weight.device, dtype=torch.float32) - grad_b = torch.zeros(v, device=weight.device, dtype=torch.float32) if ctx.has_bias else None - - chunk = max(1, min(n, _BWD_CHUNK_ELEMS // v)) - for i0 in range(0, n, chunk): - i1 = min(i0 + chunk, n) - x = hidden_2d[i0:i1] - logits = torch.matmul(x, weight.t()) - if ctx.has_bias: - logits = logits + bias_t - dz = torch.softmax(logits.float(), dim=-1).neg_() - rows = torch.arange(i1 - i0, device=dz.device) - dz[rows, target_1d[i0:i1].long()] += 1.0 - dz *= g[i0:i1].unsqueeze(1) - - dz_dt = dz.to(dt) - grad_h[i0:i1] = torch.matmul(dz_dt, weight).float() - grad_w += torch.matmul(dz_dt.t(), x).float() - if ctx.has_bias: - grad_b += dz.sum(0) - - grad_hidden = grad_h.to(ctx.hidden_dtype).reshape(ctx.lead_shape + (d,)) - grad_weight = grad_w.to(ctx.weight_dtype) - grad_bias = grad_b.to(ctx.bias_dtype) if ctx.has_bias else None + grad_hidden, grad_weight, grad_bias = chunked_linear_logp_backward( + grad_logp, + hidden_2d, + weight, + target_1d, + bias_t, + has_bias=ctx.has_bias, + lead_shape=ctx.lead_shape, + hidden_dtype=ctx.hidden_dtype, + weight_dtype=ctx.weight_dtype, + bias_dtype=ctx.bias_dtype, + ) + # Inputs: hidden, lm_head_weight, bias, target_ids. return grad_hidden, grad_weight, grad_bias, None @@ -162,4 +142,11 @@ def apply( ) if not _sm90_supported(hidden, lm_head_weight): return _fallback_op()(hidden, lm_head_weight, target_ids, bias) + vocab = lm_head_weight.size(0) + if bool(((target_ids < 0) | (target_ids >= vocab)).any()): + t_min, t_max = int(target_ids.min()), int(target_ids.max()) + raise ValueError( + f"target_ids out of range: expected [0, {vocab - 1}], got [{t_min}, {t_max}]. " + "Mask or filter padding / ignore-index values (e.g. -100) before this op." + ) return _FusedLinearLogpSM90Function.apply(hidden, lm_head_weight, bias, target_ids) diff --git a/rl_engine/kernels/ops/pytorch/loss/linear_logp.py b/rl_engine/kernels/ops/pytorch/loss/linear_logp.py index dca4700e..5e649556 100644 --- a/rl_engine/kernels/ops/pytorch/loss/linear_logp.py +++ b/rl_engine/kernels/ops/pytorch/loss/linear_logp.py @@ -7,6 +7,62 @@ import torch +# Backward token-chunk target: process at most this many ``[chunk, V]`` logit +# elements per cuBLAS step so peak backward memory stays ~``chunk*V`` instead of +# ``N*V``. +BWD_CHUNK_ELEMS = 1 << 24 + + +def chunked_linear_logp_backward( + grad_logp: torch.Tensor, + hidden_2d: torch.Tensor, + weight: torch.Tensor, + target_1d: torch.Tensor, + bias_t: torch.Tensor, + *, + has_bias: bool, + lead_shape, + hidden_dtype: torch.dtype, + weight_dtype: torch.dtype, + bias_dtype, + chunk_elems: int = BWD_CHUNK_ELEMS, +): + # Liger-style chunked backward shared by the Triton and CUDA SM90 fused ops. + n, d = hidden_2d.shape + v = weight.shape[0] + dt = weight.dtype + g = grad_logp.reshape(-1).to(torch.float32) + + grad_h = torch.empty_like(hidden_2d, dtype=torch.float32) + grad_w = torch.zeros(v, d, device=weight.device, dtype=torch.float32) + grad_b = torch.zeros(v, device=weight.device, dtype=torch.float32) if has_bias else None + + chunk = max(1, min(n, chunk_elems // v)) + for i0 in range(0, n, chunk): + i1 = min(i0 + chunk, n) + x = hidden_2d[i0:i1] # [C, D] + logits = torch.matmul(x, weight.t()) # [C, V] + if has_bias: + logits = logits + bias_t + + # dz = g * (onehot - softmax(logits)), recomputed from scratch so it is + # self-normalizing and independent of the forward's saved lse. + dz = torch.softmax(logits.float(), dim=-1).neg_() # [C, V] fp32 + rows = torch.arange(i1 - i0, device=dz.device) + dz[rows, target_1d[i0:i1].long()] += 1.0 + dz *= g[i0:i1].unsqueeze(1) + + dz_dt = dz.to(dt) + grad_h[i0:i1] = torch.matmul(dz_dt, weight).float() # [C, D] + grad_w += torch.matmul(dz_dt.t(), x).float() # [V, D] + if grad_b is not None: + grad_b += dz.sum(0) + + grad_hidden = grad_h.to(hidden_dtype).reshape(tuple(lead_shape) + (d,)) + grad_weight = grad_w.to(weight_dtype) + grad_bias = grad_b.to(bias_dtype) if grad_b is not None else None + return grad_hidden, grad_weight, grad_bias + class NativeLinearLogpOp: """Naive PyTorch reference for fused linear log-prob. diff --git a/rl_engine/kernels/ops/triton/loss/linear_logp.py b/rl_engine/kernels/ops/triton/loss/linear_logp.py index 40941bce..4ef9ce14 100644 --- a/rl_engine/kernels/ops/triton/loss/linear_logp.py +++ b/rl_engine/kernels/ops/triton/loss/linear_logp.py @@ -8,15 +8,13 @@ import triton import triton.language as tl +from rl_engine.kernels.ops.pytorch.loss.linear_logp import chunked_linear_logp_backward + # Token / vocab / hidden tile sizes (forward Triton kernel). _BLOCK_N = 32 _BLOCK_V = 64 _BLOCK_D = 64 -# Backward token-chunk size target: process at most this many [chunk, V] logit -# elements per cuBLAS step so peak backward memory stays ~chunk*V, not N*V. -_BWD_CHUNK_ELEMS = 1 << 24 - @triton.jit def _linear_logp_fwd_kernel( @@ -136,44 +134,18 @@ def forward(ctx, hidden, lm_head_weight, bias, target_ids): @staticmethod def backward(ctx, grad_logp): hidden_2d, weight, bias_t, target_1d, _lse = ctx.saved_tensors - n, d = hidden_2d.shape - v = weight.shape[0] - dt = weight.dtype - g = grad_logp.reshape(-1).to(torch.float32) - - grad_h = torch.empty_like(hidden_2d, dtype=torch.float32) - grad_w = torch.zeros(v, d, device=weight.device, dtype=torch.float32) - grad_b = torch.zeros(v, device=weight.device, dtype=torch.float32) if ctx.has_bias else None - - # Liger-style chunked materialization: process at most `chunk` tokens at a - # time, materializing only [chunk, V] logits. The projections use cuBLAS - # matmuls (tensor cores for bf16/fp16), and grad_weight is accumulated in - # sequential loop order -> deterministic and atomic-free. Peak extra - # memory is chunk*V instead of N*V. - chunk = max(1, min(n, _BWD_CHUNK_ELEMS // v)) - for i0 in range(0, n, chunk): - i1 = min(i0 + chunk, n) - x = hidden_2d[i0:i1] # [C, D] - logits = torch.matmul(x, weight.t()) # [C, V] - if ctx.has_bias: - logits = logits + bias_t - - # dz = g * (onehot - softmax(logits)), recomputed from scratch so it - # is self-normalizing and independent of the forward's saved lse. - dz = torch.softmax(logits.float(), dim=-1).neg_() # [C, V] fp32 - rows = torch.arange(i1 - i0, device=dz.device) - dz[rows, target_1d[i0:i1].long()] += 1.0 - dz *= g[i0:i1].unsqueeze(1) - - dz_dt = dz.to(dt) - grad_h[i0:i1] = torch.matmul(dz_dt, weight).float() # [C, D] - grad_w += torch.matmul(dz_dt.t(), x).float() # [V, D] - if ctx.has_bias: - grad_b += dz.sum(0) - - grad_hidden = grad_h.to(ctx.hidden_dtype).reshape(ctx.lead_shape + (d,)) - grad_weight = grad_w.to(ctx.weight_dtype) - grad_bias = grad_b.to(ctx.bias_dtype) if ctx.has_bias else None + grad_hidden, grad_weight, grad_bias = chunked_linear_logp_backward( + grad_logp, + hidden_2d, + weight, + target_1d, + bias_t, + has_bias=ctx.has_bias, + lead_shape=ctx.lead_shape, + hidden_dtype=ctx.hidden_dtype, + weight_dtype=ctx.weight_dtype, + bias_dtype=ctx.bias_dtype, + ) # Inputs: hidden, lm_head_weight, bias, target_ids. return grad_hidden, grad_weight, grad_bias, None @@ -203,8 +175,11 @@ def apply( target_ids: torch.Tensor, bias: Optional[torch.Tensor] = None, ) -> torch.Tensor: - if not hidden.is_cuda: - raise RuntimeError("TritonLinearLogpOp requires CUDA tensors.") + if hidden.device.type not in ("cuda", "xpu", "hip"): + raise RuntimeError( + "TritonLinearLogpOp requires a GPU tensor (CUDA / ROCm / XPU), got " + f"device '{hidden.device}'." + ) if hidden.shape[:-1] != target_ids.shape: raise ValueError( f"hidden leading shape {tuple(hidden.shape[:-1])} must match " @@ -215,4 +190,11 @@ def apply( f"hidden dim {hidden.size(-1)} must match lm_head_weight dim " f"{lm_head_weight.size(-1)}" ) + vocab = lm_head_weight.size(0) + if bool(((target_ids < 0) | (target_ids >= vocab)).any()): + t_min, t_max = int(target_ids.min()), int(target_ids.max()) + raise ValueError( + f"target_ids out of range: expected [0, {vocab - 1}], got [{t_min}, {t_max}]. " + "Mask or filter padding / ignore-index values (e.g. -100) before this op." + ) return _LinearLogpFunction.apply(hidden, lm_head_weight, bias, target_ids) diff --git a/tests/test_linear_logp.py b/tests/test_linear_logp.py index 969b5b50..a6ce1b30 100644 --- a/tests/test_linear_logp.py +++ b/tests/test_linear_logp.py @@ -302,6 +302,29 @@ def test_sm90_rejects_bad_target_and_bias(): sm90(hidden, weight, target, bias.cpu()) +@requires_sm90 +def test_sm90_rejects_out_of_range_target(): + # Padding (-100) / out-of-vocab ids must error, not silently corrupt fwd/bwd. + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + + sm90 = FusedLinearLogpSM90Op() + hidden, weight, target, bias = _sm90_inputs(18) + + pad = target.clone() + pad[0] = -100 # typical ignore_index + with pytest.raises(ValueError): + sm90(hidden, weight, pad, bias) + + oob = target.clone() + oob[1] = _SM90_V # == V, one past the last valid id + with pytest.raises(ValueError): + sm90(hidden, weight, oob, bias) + + # A valid target (all in [0, V)) still works. + out = sm90(hidden, weight, target, bias) + assert out.shape == target.shape and torch.isfinite(out).all() + + def test_registry_dispatch_matches_native(): from rl_engine.kernels.registry import kernel_registry from rl_engine.platforms.device import device_ctx