diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 92cd0433..7167ea28 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -68,6 +68,15 @@ jobs: python -m pytest rl_engine/tests/test_dispatch.py -v PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest tests/test_attention_correctness.py -q -rs + - name: Run Linear Log-Prob Ground-Truth Tests (CPU-safe) + run: | + python -m pytest \ + tests/test_batch_invariant_linear_logp.py \ + tests/test_kernel_registry.py \ + tests/test_linear_logp.py \ + tests/test_operator_inputs.py \ + -q -rs + - name: Run Attention Ground-Truth Tests (CPU-safe) run: | python -m pytest tests/test_attention.py -v -k "not large and not gpu" diff --git a/.github/workflows/gpu-ci.yml b/.github/workflows/gpu-ci.yml index 46caf703..22fe4675 100644 --- a/.github/workflows/gpu-ci.yml +++ b/.github/workflows/gpu-ci.yml @@ -12,6 +12,7 @@ on: - 'requirements*.txt' - 'docker/Dockerfile.cuda' - 'ci/run_gpu_ci.sh' + - 'scripts/ci_smoke.py' - '.github/workflows/gpu-ci.yml' types: [ opened, synchronize, reopened, labeled ] diff --git a/benchmarks/benchmark_batch_invariant_linear_logp.py b/benchmarks/benchmark_batch_invariant_linear_logp.py new file mode 100644 index 00000000..f6b14cd8 --- /dev/null +++ b/benchmarks/benchmark_batch_invariant_linear_logp.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Benchmark single-card batch-invariant fused linear log-probability. + +The benchmark compares three SM90 paths: + +1. the new batch-invariant fused forward; +2. the existing throughput-oriented fused ``linear_logp`` forward; +3. a batch-invariant materialized composition (LM head, then logp). + +The two fused timings call their compiled symbols directly with the same prepared +inputs, so neither side includes Python validation or a host synchronization. +FP32 inputs for the materialized composition are allocated only after the fused +measurements and outside its timed region. Memory columns report incremental +operator allocations above those prepared inputs. +""" + +from __future__ import annotations + +import argparse +from collections.abc import Callable +from functools import partial + +import torch +from tabulate import tabulate + +from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE +from rl_engine.kernels.ops.cuda.linear.lm_head import SM90LMHeadOp +from rl_engine.kernels.ops.cuda.loss.batch_invariant_logp import BatchInvariantLogpSM90Op + +DEFAULT_CONFIGS = [ + (1, 4096, 151936), + (32, 4096, 151936), + (256, 4096, 151936), + (4096, 2048, 32768), + (4096, 4096, 151936), +] +_HIDDEN_TILE = 32 +_FP32_LOGP_VOCAB_ALIGNMENT = 4 + + +def _make_inputs(num_tokens: int, hidden_dim: int, vocab_size: int): + hidden = torch.randn(num_tokens, hidden_dim, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(vocab_size, hidden_dim, device="cuda", dtype=torch.bfloat16) + target_ids = torch.randint(0, vocab_size, (num_tokens,), device="cuda") + return hidden, weight, target_ids + + +def _time_ms(fn: Callable[[], torch.Tensor], warmup: int, iterations: int) -> float: + 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(iterations): + fn() + end.record() + torch.cuda.synchronize() + return start.elapsed_time(end) / iterations + + +def _peak_memory_mb(fn: Callable[[], torch.Tensor], warmup: int = 2) -> float: + for _ in range(warmup): + fn() + torch.cuda.synchronize() + torch.cuda.empty_cache() + baseline = torch.cuda.memory_allocated() + torch.cuda.reset_peak_memory_stats() + fn() + torch.cuda.synchronize() + peak = torch.cuda.max_memory_allocated() - baseline + return peak / (1024**2) + + +def _validate_run_counts(warmup: int, iterations: int) -> None: + if warmup < 0: + raise ValueError(f"warmup must be non-negative, got {warmup}") + if iterations <= 0: + raise ValueError(f"iterations must be positive, got {iterations}") + + +def _validate_config(config: tuple[int, ...]) -> tuple[int, int, int]: + if len(config) != 3: + raise ValueError(f"each config must contain N,D,V, got {config}") + + num_tokens, hidden_dim, vocab_size = config + if num_tokens <= 0 or hidden_dim <= 0 or vocab_size <= 0: + raise ValueError(f"N, D, and V must be positive, got {config}") + if hidden_dim % _HIDDEN_TILE != 0: + raise ValueError(f"D must be divisible by {_HIDDEN_TILE}, got {hidden_dim}") + if vocab_size % _FP32_LOGP_VOCAB_ALIGNMENT != 0: + raise ValueError( + "V must be divisible by 4 so the materialized FP32 logp comparator " + f"stays on its SM90 TMA backend, got {vocab_size}" + ) + return num_tokens, hidden_dim, vocab_size + + +def _parse_configs(value: str | None) -> list[tuple[int, int, int]]: + if value is None: + return list(DEFAULT_CONFIGS) + try: + configs = [tuple(int(item) for item in group.split(",")) for group in value.split(";")] + except ValueError as error: + raise ValueError( + "--configs must contain semicolon-separated integer N,D,V triples" + ) from error + return [_validate_config(config) for config in configs] + + +def _run_materialized(lm_head, logp, hidden, weight, target_ids): + logits = lm_head(hidden, weight) + return logp(logits, target_ids) + + +def _run_fused(symbol, hidden, weight, target_ids): + logp, _lse = symbol(hidden, weight, target_ids, None) + return logp + + +def run(args: argparse.Namespace) -> None: + _validate_run_counts(args.warmup, args.iterations) + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9: + raise RuntimeError("benchmark requires an SM90 Hopper GPU") + required_symbols = ( + "batch_invariant_linear_logp_sm90", + "fused_linear_logp_sm90", + "lm_head_sm90_forward", + "lm_head_sm90_forward_fp32", + "batch_invariant_logp_sm90", + ) + if not _EXT_AVAILABLE or any(not hasattr(_C, name) for name in required_symbols): + raise RuntimeError("benchmark requires all compiled SM90 comparison symbols") + + invariant_lm_head = SM90LMHeadOp() + invariant_logp = BatchInvariantLogpSM90Op() + rows = [] + + with torch.inference_mode(): + for config in args.configs: + num_tokens, hidden_dim, vocab_size = _validate_config(tuple(config)) + hidden, weight, target_ids = _make_inputs(num_tokens, hidden_dim, vocab_size) + target_ids_i32 = target_ids.to(torch.int32) + + run_invariant_fused = partial( + _run_fused, + _C.batch_invariant_linear_logp_sm90, + hidden, + weight, + target_ids_i32, + ) + run_throughput_fused = partial( + _run_fused, + _C.fused_linear_logp_sm90, + hidden, + weight, + target_ids_i32, + ) + invariant_ms = _time_ms(run_invariant_fused, args.warmup, args.iterations) + throughput_ms = _time_ms(run_throughput_fused, args.warmup, args.iterations) + invariant_memory = _peak_memory_mb(run_invariant_fused) + throughput_memory = _peak_memory_mb(run_throughput_fused) + invariant_output = run_invariant_fused() + throughput_output = run_throughput_fused() + + hidden_fp32 = hidden.float() + weight_fp32 = weight.float() + del ( + run_invariant_fused, + run_throughput_fused, + hidden, + weight, + target_ids_i32, + ) + + run_materialized = partial( + _run_materialized, + invariant_lm_head.forward_fp32, + invariant_logp, + hidden_fp32, + weight_fp32, + target_ids, + ) + materialized_ms = _time_ms(run_materialized, args.warmup, args.iterations) + materialized_memory = _peak_memory_mb(run_materialized) + materialized_output = run_materialized() + invariant_max_abs = float((invariant_output - materialized_output).abs().max()) + throughput_max_abs = float((throughput_output - materialized_output).abs().max()) + + rows.append( + [ + f"{num_tokens}x{hidden_dim}x{vocab_size}", + f"{invariant_ms:.3f}", + f"{throughput_ms:.3f}", + f"{invariant_ms / throughput_ms:.2f}x", + f"{materialized_ms:.3f}", + f"{invariant_memory:.0f}", + f"{throughput_memory:.0f}", + f"{materialized_memory:.0f}", + f"{invariant_max_abs:.3e}", + f"{throughput_max_abs:.3e}", + ] + ) + del ( + run_materialized, + hidden_fp32, + weight_fp32, + target_ids, + invariant_output, + throughput_output, + materialized_output, + ) + torch.cuda.empty_cache() + + print( + tabulate( + rows, + headers=[ + "shape (N x D x V)", + "invariant fused ms", + "throughput fused ms", + "invariance overhead", + "materialized ms", + "invariant incremental MB", + "throughput incremental MB", + "materialized incremental MB", + "invariant max abs vs materialized", + "throughput max abs vs materialized", + ], + tablefmt="github", + ) + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--iterations", type=int, default=20) + parser.add_argument("--warmup", type=int, default=5) + parser.add_argument( + "--configs", + type=str, + default=None, + help=( + "Semicolon-separated N,D,V triples with D divisible by 32 and V divisible by 4, " + "for example '4096,4096,151936'" + ), + ) + args = parser.parse_args() + try: + _validate_run_counts(args.warmup, args.iterations) + args.configs = _parse_configs(args.configs) + except ValueError as error: + parser.error(str(error)) + return args + + +if __name__ == "__main__": + run(parse_args()) diff --git a/csrc/cuda/fused_linear_logp_sm90.cu b/csrc/cuda/fused_linear_logp_sm90.cu index e1788ad9..0059b960 100644 --- a/csrc/cuda/fused_linear_logp_sm90.cu +++ b/csrc/cuda/fused_linear_logp_sm90.cu @@ -7,10 +7,13 @@ #include "../utils/tma_utils.cuh" #include #include +#include #include #include +#include #include #include +#include #include #include #include @@ -54,6 +57,16 @@ constexpr int GRAD_W_WMMA_K = 16; constexpr int GRAD_W_WMMA_WARPS = 8; constexpr int GRAD_W_WMMA_THREADS = GRAD_W_WMMA_WARPS * 32; +// The batch-invariant entry point uses a split-V schedule derived only from V. +// In particular, it must never use N (the batch/token dimension) to choose a +// different floating-point reduction topology. +constexpr int MAX_BATCH_INVARIANT_VOCAB_SPLITS = 32; + +enum class VocabSplitPolicy { + kThroughput, + kBatchInvariant, +}; + 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)"); @@ -61,6 +74,64 @@ 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); +constexpr int select_batch_invariant_vocab_splits(int total_vtiles) { + // Keep at most MAX_BATCH_INVARIANT_VOCAB_SPLITS non-empty, contiguous + // ranges. Both the partition and the combine order depend on V only. + const int vtiles_per_split = + (total_vtiles + MAX_BATCH_INVARIANT_VOCAB_SPLITS - 1) / + MAX_BATCH_INVARIANT_VOCAB_SPLITS; + return (total_vtiles + vtiles_per_split - 1) / vtiles_per_split; +} + +static_assert( + select_batch_invariant_vocab_splits(MAX_BATCH_INVARIANT_VOCAB_SPLITS + 1) <= + MAX_BATCH_INVARIANT_VOCAB_SPLITS, + "batch-invariant split selection must respect its workspace cap"); + +int select_throughput_vocab_splits(int row_blocks, int total_vtiles) { + // The performance path intentionally adapts split-V to occupancy. + const int sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + const int target_ctas = sm_count * 4; + return std::max(1, std::min(target_ctas / std::max(row_blocks, 1), total_vtiles)); +} + +bool is_aligned(const torch::Tensor &tensor, std::uintptr_t alignment) { + return reinterpret_cast(tensor.data_ptr()) % alignment == 0; +} + +torch::Tensor ensure_aligned(torch::Tensor tensor, std::uintptr_t alignment, + const char *name, const char *consumer) { + if (is_aligned(tensor, alignment)) + return tensor; + + auto aligned = tensor.clone(); + TORCH_CHECK(is_aligned(aligned, alignment), name, " could not be copied to ", alignment, + "-byte-aligned storage for ", consumer); + return aligned; +} + +torch::Tensor ensure_tma_aligned(torch::Tensor tensor, const char *name) { + return ensure_aligned(tensor, 16, name, "TMA"); +} + +torch::Tensor ensure_wmma_aligned(torch::Tensor tensor, const char *name) { + // wmma::load_matrix_sync requires a 256-bit-aligned base pointer. + return ensure_aligned(tensor, 32, name, "WMMA"); +} + +torch::Tensor prepare_int32_targets(torch::Tensor target) { + auto target_i = target.to(torch::kInt32).contiguous(); + if (target.scalar_type() == at::kLong) { + const auto invalid = target.lt(std::numeric_limits::min()) | + target.gt(std::numeric_limits::max()); + // All supported vocab ranges are non-negative, so -1 can never alias a + // real class ID. Saturating to INT_MAX would be wrong for a TP shard + // whose final legitimate class is exactly INT_MAX. + target_i.masked_fill_(invalid, -1); + } + return target_i; +} + // 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) { @@ -87,7 +158,8 @@ __global__ void fused_linear_logp_sm90_kernel(const __grid_constant__ CUtensorMa float *__restrict__ part_sum, // [n_split, N] float *__restrict__ part_zt, // [n_split, N] int N, int D, int V, int n_split, - int vocab_start_index) { + int vocab_start_index, + bool target_must_be_local) { const int tid = threadIdx.x; const int warp = tid / 32; const int lane = tid % 32; @@ -99,7 +171,8 @@ __global__ void fused_linear_logp_sm90_kernel(const __grid_constant__ CUtensorMa // 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 total_vtiles = + static_cast((static_cast(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); @@ -127,7 +200,15 @@ __global__ void fused_linear_logp_sm90_kernel(const __grid_constant__ CUtensorMa for (int r = tid; r < num_rows; r += WG_THREADS) { sMax[r] = -CUDART_INF_F; sSum[r] = 0.0f; - sZt[r] = 0.0f; + const int tgt = target[row_base + r]; + const int64_t local_begin = static_cast(vocab_start_index); + const int64_t local_end = local_begin + V; + const bool target_is_local = static_cast(tgt) >= local_begin && + static_cast(tgt) < local_end; + // The public single-card path remains asynchronous by default. Mark an + // invalid target as NaN instead of silently treating its logit as zero; + // the global-target TP path intentionally permits out-of-shard IDs. + sZt[r] = target_must_be_local && !target_is_local ? CUDART_NAN_F : 0.0f; } if (tid == 0) { #pragma unroll @@ -247,7 +328,8 @@ __global__ void fused_linear_logp_sm90_kernel(const __grid_constant__ CUtensorMa if (bias != nullptr) val += bias[col]; tmax = fmaxf(tmax, val); - if (vocab_start_index + col == tgt) + if (static_cast(vocab_start_index) + col == + static_cast(tgt)) sZt[r] = val; } float tsum = 0.0f; @@ -271,7 +353,7 @@ __global__ void fused_linear_logp_sm90_kernel(const __grid_constant__ CUtensorMa // 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; + const int64_t idx = static_cast(split) * N + row_base + r; part_max[idx] = sMax[r]; part_sum[idx] = sSum[r]; part_zt[idx] = sZt[r]; @@ -450,12 +532,12 @@ __global__ void fused_linear_logp_sm90_combine_kernel(const float *__restrict__ float M = -CUDART_INF_F; for (int s = 0; s < n_split; ++s) - M = fmaxf(M, part_max[s * N + r]); + M = fmaxf(M, part_max[static_cast(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; + const int64_t idx = static_cast(s) * N + r; S += part_sum[idx] * __expf(part_max[idx] - M); zt += part_zt[idx]; } @@ -1198,6 +1280,8 @@ void launch_logits_tile_bf16_mma(torch::Tensor hidden, "bf16 MMA logits tile requires bf16 hidden and weight"); TORCH_CHECK(hidden.is_contiguous() && weight.is_contiguous(), "bf16 MMA logits tile requires contiguous hidden and weight"); + TORCH_CHECK(is_aligned(hidden, 16) && is_aligned(weight, 16), + "bf16 MMA logits tile requires 16-byte-aligned hidden and weight"); const int N = hidden.size(0); const int D = hidden.size(1); const int V = weight.size(0); @@ -1241,6 +1325,8 @@ void launch_dlogits_tile_bf16_mma(torch::Tensor hidden, "bf16 MMA dlogits tile requires bf16 dlogits output"); TORCH_CHECK(hidden.is_contiguous() && weight.is_contiguous(), "bf16 MMA dlogits tile requires contiguous hidden and weight"); + TORCH_CHECK(is_aligned(hidden, 16) && is_aligned(weight, 16), + "bf16 MMA dlogits tile requires 16-byte-aligned hidden and weight"); TORCH_CHECK(target.is_contiguous() && grad_logp.is_contiguous() && lse.is_contiguous(), "bf16 MMA dlogits tile requires contiguous target, grad_logp, and lse"); const int N = hidden.size(0); @@ -1285,6 +1371,8 @@ void launch_grad_weight_tile_wmma(torch::Tensor dlogits, grad_weight.scalar_type() == at::kBFloat16, "WMMA grad_weight tile requires bf16 dlogits, hidden, and grad_weight"); TORCH_CHECK(hidden.is_contiguous(), "WMMA grad_weight tile requires contiguous hidden"); + TORCH_CHECK(is_aligned(hidden, 32), + "WMMA grad_weight tile requires 32-byte-aligned hidden"); TORCH_CHECK(grad_weight.dim() == 2 && dlogits.dim() == 2 && hidden.dim() == 2, "WMMA grad_weight tile expects 2-D tensors"); const int N = dlogits.size(0); @@ -1324,6 +1412,8 @@ void launch_grad_weight_from_logits_tile_wmma(torch::Tensor logits_or_probs, grad_weight.scalar_type() == at::kBFloat16, "fused tile grad_weight requires bf16 hidden and grad_weight"); TORCH_CHECK(hidden.is_contiguous(), "fused tile grad_weight requires contiguous hidden"); + TORCH_CHECK(is_aligned(hidden, 32), + "fused tile grad_weight requires 32-byte-aligned hidden"); TORCH_CHECK(logits_or_probs.dim() == 2 && hidden.dim() == 2 && grad_weight.dim() == 2, "fused tile grad_weight expects 2-D tensors"); TORCH_CHECK(logits_or_probs.stride(1) == 1 && hidden.stride(1) == 1 && @@ -1772,26 +1862,75 @@ torch::Tensor linear_logp_logits_bf16_to_dlogits(torch::Tensor logits, std::vector fused_linear_logp_sm90_forward_impl( torch::Tensor hidden, torch::Tensor weight, torch::Tensor target, - torch::optional bias, int64_t vocab_start_index, bool return_target_logit) { - TORCH_CHECK(hidden.is_cuda() && weight.is_cuda(), "hidden and weight must be CUDA tensors"); + torch::optional bias, int64_t vocab_start_index, bool return_target_logit, + VocabSplitPolicy split_policy) { + TORCH_CHECK(hidden.dim() == 2, "hidden must be 2-D [N, D]"); + TORCH_CHECK(weight.dim() == 2, "weight must be 2-D [V, D]"); + TORCH_CHECK(target.dim() == 1, "target must be 1-D [N]"); + TORCH_CHECK(hidden.is_cuda() && weight.is_cuda() && target.is_cuda(), + "hidden, weight, and target must be CUDA tensors"); TORCH_CHECK(weight.device() == hidden.device(), "lm_head_weight must be on the same device as hidden"); + TORCH_CHECK(target.device() == hidden.device(), + "target 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"); - const int N = hidden.size(0); - const int D = hidden.size(1); - const int V = weight.size(0); + TORCH_CHECK(hidden.size(0) <= std::numeric_limits::max(), + "hidden row count exceeds the SM90 kernel's int32 indexing limit"); + TORCH_CHECK(hidden.size(1) <= std::numeric_limits::max(), + "hidden dimension exceeds the SM90 kernel's int32 indexing limit"); + TORCH_CHECK(weight.size(0) <= std::numeric_limits::max(), + "vocabulary size exceeds the SM90 kernel's int32 indexing limit"); + const int N = static_cast(hidden.size(0)); + const int D = static_cast(hidden.size(1)); + const int V = static_cast(weight.size(0)); + TORCH_CHECK(N > 0, "hidden must contain at least one token row"); + TORCH_CHECK(D > 0, "hidden dimension must be positive"); + TORCH_CHECK(V > 0, "weight must contain at least one vocabulary row"); 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()); + TORCH_CHECK(vocab_start_index >= 0 && + vocab_start_index <= std::numeric_limits::max(), + "vocab_start_index must be representable as a non-negative int32 value"); + TORCH_CHECK(vocab_start_index + static_cast(V) <= + static_cast(std::numeric_limits::max()) + 1, + "local vocabulary range exceeds the SM90 kernel's int32 target limit"); + + const auto target_type = target.scalar_type(); + TORCH_CHECK(target_type == at::kByte || target_type == at::kChar || + target_type == at::kShort || target_type == at::kInt || + target_type == at::kLong, + "target must have dtype uint8, int8, int16, int32, or int64"); + + const bool batch_invariant = split_policy == VocabSplitPolicy::kBatchInvariant; if (bias.has_value()) { + TORCH_CHECK(bias->dim() == 1, "bias must be 1-D [V]"); 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()); + if (batch_invariant) { + const auto bias_type = bias->scalar_type(); + TORCH_CHECK(bias_type == at::kHalf || bias_type == at::kBFloat16 || + bias_type == at::kFloat, + "bias must have dtype float16, bfloat16, or float32"); + } } + c10::cuda::CUDAGuard device_guard(hidden.device()); + const auto *device_properties = at::cuda::getCurrentDeviceProperties(); + TORCH_CHECK(device_properties->major == 9, + "fused_linear_logp_sm90 requires an SM90 Hopper GPU, got compute capability ", + device_properties->major, ".", device_properties->minor); + + // A contiguous view may still start at a non-zero storage offset. TMA + // requires a 16-byte-aligned global base address, so copy only those rare + // misaligned views into allocator-aligned storage. + hidden = ensure_tma_aligned(hidden, "hidden"); + weight = ensure_tma_aligned(weight, "weight"); + auto opts_f = hidden.options().dtype(torch::kFloat); auto out_value = torch::empty({N}, opts_f); auto lse = torch::empty({N}, opts_f); @@ -1814,48 +1953,65 @@ std::vector fused_linear_logp_sm90_forward_impl( 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(); + const int row_blocks = static_cast((static_cast(N) + BM - 1) / BM); + const int total_vtiles = + static_cast((static_cast(V) + BN - 1) / BN); + auto target_i = prepare_int32_targets(target); - // 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)); + const int n_split = + split_policy == VocabSplitPolicy::kBatchInvariant + ? select_batch_invariant_vocab_splits(total_vtiles) + : select_throughput_vocab_splits(row_blocks, 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); + C10_CUDA_CHECK(cudaFuncSetAttribute(fused_linear_logp_sm90_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); } dim3 grid(row_blocks, n_split); - fused_linear_logp_sm90_kernel<<>>( + fused_linear_logp_sm90_kernel<<>>( 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, - static_cast(vocab_start_index)); + static_cast(vocab_start_index), !return_target_logit); + C10_CUDA_KERNEL_LAUNCH_CHECK(); const int combine_threads = 256; - const int combine_blocks = (N + combine_threads - 1) / combine_threads; - fused_linear_logp_sm90_combine_kernel<<>>( + const int combine_blocks = + static_cast((static_cast(N) + combine_threads - 1) / combine_threads); + fused_linear_logp_sm90_combine_kernel<<>>( part_max.data_ptr(), part_sum.data_ptr(), part_zt.data_ptr(), out_value.data_ptr(), lse.data_ptr(), N, n_split, return_target_logit); + C10_CUDA_KERNEL_LAUNCH_CHECK(); return {out_value, lse}; } // 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. +// materialized. Global scratch grows with n_split (which depends on V up to a +// fixed cap), so its worst-case bound is O(N). std::vector fused_linear_logp_sm90_forward(torch::Tensor hidden, torch::Tensor weight, torch::Tensor target, torch::optional bias) { - return fused_linear_logp_sm90_forward_impl(hidden, weight, target, bias, 0, false); + return fused_linear_logp_sm90_forward_impl(hidden, weight, target, bias, 0, false, + VocabSplitPolicy::kThroughput); +} + +// Single-card batch-invariant forward. It deliberately shares the validated +// TMA+MMA device kernel with the throughput path, but pins split-V independently +// of N so each row sees the same D and V reduction topology across batch layouts. +std::vector batch_invariant_linear_logp_sm90_forward( + torch::Tensor hidden, torch::Tensor weight, torch::Tensor target, + torch::optional bias) { + return fused_linear_logp_sm90_forward_impl(hidden, weight, target, bias, 0, false, + VocabSplitPolicy::kBatchInvariant); } // Vocab-parallel local-shard forward. Target ids stay in global-vocab @@ -1865,7 +2021,8 @@ std::vector fused_linear_logp_sm90_global_target_forward( torch::Tensor hidden, torch::Tensor weight, torch::Tensor target, torch::optional bias, int64_t vocab_start_index) { return fused_linear_logp_sm90_forward_impl( - hidden, weight, target, bias, vocab_start_index, true); + hidden, weight, target, bias, vocab_start_index, true, + VocabSplitPolicy::kThroughput); } // Backward fast path: compute local dlogits in one CUDA kernel, then use GEMMs @@ -1907,6 +2064,7 @@ std::vector fused_linear_logp_sm90_backward(torch::Tensor grad_lo c10::cuda::CUDAGuard device_guard(hidden.device()); auto opts_f = hidden.options().dtype(torch::kFloat); auto empty = torch::empty({0}, opts_f); + auto target_i = prepare_int32_targets(target.reshape({N})); if (streaming_output_backward_enabled() && !compute_grad_hidden && compute_grad_weight && hidden.scalar_type() == at::kBFloat16 && weight.scalar_type() == at::kBFloat16 && @@ -1915,8 +2073,6 @@ std::vector fused_linear_logp_sm90_backward(torch::Tensor grad_lo auto grad_weight = torch::empty_like(weight); auto grad_f = grad_logp.reshape({N}).to(torch::kFloat).contiguous(); auto lse_f = lse.reshape({N}).to(torch::kFloat).contiguous(); - auto target_i = target.reshape({N}).to(torch::kInt32).contiguous(); - auto hidden_gemm = hidden.contiguous(); const int tile_v = streaming_output_backward_vocab_tile(V); const int threads = 256; const bool use_cuda_tf32_logits = streaming_output_backward_cuda_logits_enabled(); @@ -1936,6 +2092,17 @@ std::vector fused_linear_logp_sm90_backward(torch::Tensor grad_lo const bool use_bf16_mma_tile = use_bf16_mma_logits || use_bf16_mma_dz; const bool use_grad_weight_wmma = !use_fused_tile_dw && streaming_output_backward_grad_weight_wmma_enabled(); + if (use_fused_tile_dw) + hidden = ensure_wmma_aligned(hidden, "hidden"); + if (use_bf16_mma_tile) { + // Forward accepts misaligned contiguous views by copying them + // before TMA setup. Autograd saves the original views, so do + // the same only when this backward actually creates tensor maps. + hidden = ensure_tma_aligned(hidden, "hidden"); + weight = ensure_tma_aligned(weight, "weight"); + } + auto hidden_gemm = hidden.contiguous(); + bool grad_weight_wmma_inputs_aligned = false; torch::Tensor hidden_logits; if (!use_bf16_mma_tile) hidden_logits = hidden.to(torch::kFloat).contiguous(); @@ -1997,6 +2164,10 @@ std::vector fused_linear_logp_sm90_backward(torch::Tensor grad_lo auto grad_weight_view = grad_weight.narrow(0, v0, vc); if (use_grad_weight_wmma && can_use_grad_weight_tile_wmma(N, D, vc)) { + if (!grad_weight_wmma_inputs_aligned) { + hidden_gemm = ensure_wmma_aligned(hidden_gemm, "hidden"); + grad_weight_wmma_inputs_aligned = true; + } launch_grad_weight_tile_wmma(dlogits_gemm, hidden_gemm, grad_weight_view); } else { auto grad_weight_tile = @@ -2014,7 +2185,6 @@ std::vector fused_linear_logp_sm90_backward(torch::Tensor grad_lo auto grad_f = grad_logp.reshape({N}).to(torch::kFloat).contiguous(); auto lse_f = lse.reshape({N}).to(torch::kFloat).contiguous(); - auto target_i = target.reshape({N}).to(torch::kInt32).contiguous(); torch::Tensor bias_f; const float *bias_ptr = nullptr; if (bias.has_value()) { @@ -2040,9 +2210,11 @@ std::vector fused_linear_logp_sm90_backward(torch::Tensor grad_lo (compute_grad_hidden || compute_grad_weight) && !compute_grad_bias && !bias.has_value() && hidden.scalar_type() == at::kBFloat16 && weight.scalar_type() == at::kBFloat16 && D % BK == 0) { + // This branch always uses the TMA-backed dlogits tile kernel. + hidden = ensure_tma_aligned(hidden, "hidden"); + weight = ensure_tma_aligned(weight, "weight"); auto grad_f = grad_logp.reshape({N}).to(torch::kFloat).contiguous(); auto lse_f = lse.reshape({N}).to(torch::kFloat).contiguous(); - auto target_i = target.reshape({N}).to(torch::kInt32).contiguous(); auto hidden_gemm = hidden.contiguous(); const int tile_v = streaming_output_backward_vocab_tile(V); auto dlogits_workspace = torch::empty({N, tile_v}, hidden.options()); @@ -2121,7 +2293,6 @@ std::vector fused_linear_logp_sm90_backward(torch::Tensor grad_lo auto grad_f = grad_logp.reshape({N}).to(torch::kFloat).contiguous(); auto lse_f = lse.reshape({N}).to(torch::kFloat).contiguous(); - auto target_i = target.reshape({N}).to(torch::kInt32).contiguous(); const int threads = 256; const int64_t total = static_cast(N) * V; const int blocks = static_cast((total + threads - 1) / threads); @@ -2195,6 +2366,7 @@ std::vector fused_linear_logp_sm90_backward(torch::Tensor grad_lo hidden_grad_gemm.scalar_type() == at::kBFloat16 && can_use_grad_weight_tile_wmma(N, D, V); if (use_full_fused_tile_dw) { + hidden_grad_gemm = ensure_wmma_aligned(hidden_grad_gemm, "hidden"); grad_weight = torch::empty_like(weight); if (full_fused_tile_mode == "from_logits" || full_fused_tile_mode == "logits" || diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 97b753d3..ac5d97d7 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -13,6 +13,11 @@ std::vector fused_linear_logp_sm90_forward(torch::Tensor hidden, torch::Tensor weight, torch::Tensor target, torch::optional bias); +std::vector batch_invariant_linear_logp_sm90_forward( + torch::Tensor hidden, + torch::Tensor weight, + torch::Tensor target, + torch::optional bias); std::vector batch_invariant_logp_sm90_forward(torch::Tensor logits, torch::Tensor target, int64_t ignore_index); @@ -309,6 +314,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { 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"); + m.def("batch_invariant_linear_logp_sm90", &batch_invariant_linear_logp_sm90_forward, + "Batch-invariant fused linear log-prob from hidden states, single-card SM90"); m.def("batch_invariant_logp_sm90", &batch_invariant_logp_sm90_forward, "TMA online-softmax batch-invariant selected-token log-prob from logits, SM90"); m.def("fused_linear_logp_sm90_global_target", &fused_linear_logp_sm90_global_target_forward, diff --git a/docs/.nav.yml b/docs/.nav.yml index 73281687..70a5b633 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -18,6 +18,7 @@ nav: - operators/attention.md - operators/fused-logp.md - operators/linear-logp.md + - operators/batch-invariant-linear-logp.md - operators/batch-invariant-logp.md - operators/linear-logp-tp-test.md - operators/grpo-loss.md diff --git a/docs/benchmarking/README.md b/docs/benchmarking/README.md index 99da15c0..41332984 100644 --- a/docs/benchmarking/README.md +++ b/docs/benchmarking/README.md @@ -14,9 +14,16 @@ python benchmarks/profiler.py --format json --output reports/profile.json python benchmarks/benchmark_sampling.py python benchmarks/benchmark_grpo_op.py python benchmarks/benchmark_pack.py --smoke +python benchmarks/benchmark_batch_invariant_linear_logp.py python scripts/run_perf.py ``` +The batch-invariant linear-logp entry point is a manual three-path comparison, +not a single profiler workload; it reports invariant-versus-throughput overhead +over the same logical values. The fused paths share BF16 inputs; the materialized +path uses FP32 copies, and each memory result is relative to its own prepared-input +baseline. + The automated profiler records one row per workload shape with: - `tokens_per_sec`: active tokens divided by median latency. diff --git a/docs/operators/README.md b/docs/operators/README.md index 00f4cbb4..5c171628 100644 --- a/docs/operators/README.md +++ b/docs/operators/README.md @@ -22,6 +22,7 @@ Every operator page should include: - [Standard Attention](attention.md) - [Fused LogP](fused-logp.md) - [Fused Linear LogP](linear-logp.md) +- [Batch-Invariant Fused Linear LogP](batch-invariant-linear-logp.md) - [Batch-Invariant LogP](batch-invariant-logp.md) - [Fused Linear LogP TP Test Runbook](linear-logp-tp-test.md) - [GRPO Loss](grpo-loss.md) diff --git a/docs/operators/batch-invariant-linear-logp.md b/docs/operators/batch-invariant-linear-logp.md new file mode 100644 index 00000000..09794f15 --- /dev/null +++ b/docs/operators/batch-invariant-linear-logp.md @@ -0,0 +1,211 @@ +# Batch-Invariant Fused Linear LogP + +Batch-Invariant Fused Linear LogP computes the selected token log-probability +directly from hidden states and LM-head weights: + +```text +logp[row] = log_softmax(hidden[row] @ weight.T + bias)[target_ids[row]] +``` + +The operator does not materialize the full `[N, V]` logits tensor. Unlike the +throughput-oriented [`linear_logp`](linear-logp.md), its floating-point reduction +topology is independent of the token-row count `N`. This gives a bitwise +single-card batch-invariance contract for rollout and train-inference alignment. + +## Prerequisite + +Build the extension on a Hopper host after installing CUDA-enabled PyTorch: + +```bash +RL_KERNEL_REQUIRE_EXT=1 KERNEL_ALIGN_FORCE_SM90=1 \ + python -m pip install --no-build-isolation -e . +``` + +## Entry Point + +```python +import torch + +from rl_engine.kernels.registry import kernel_registry + +op = kernel_registry.get_op("batch_invariant_linear_logp", device="cuda") + +with torch.no_grad(): + logp = op( + hidden, # [B, S, D] or [N, D], bf16 + lm_head_weight, # [V, D], bf16 + target_ids, # [B, S] or [N], integer + bias=None, # optional [V] + validate=False, # async; invalid targets produce NaN + ) # [B, S] or [N], float32 +``` + +## Contract + +For fixed `hidden[row]`, `weight`, `target_ids[row]`, and optional `bias`, the +output bits do not depend on: + +- batch size or flattened token-row count; +- the row's position in the batch; +- neighboring row contents; +- splitting the batch dimension into chunks and concatenating the results; +- repeated execution on the same SM90 device and software stack. + +The contract is stronger than repeated-run determinism. It does not promise +bitwise equality across GPU architectures, CUDA versions, compiler versions, +different vocabulary sizes, or different hidden dimensions. + +## Why the Existing Fused Path Is Separate + +The regular SM90 `linear_logp` path chooses the number of split-V CTAs using +the number of row blocks and the device SM count. That occupancy heuristic is +useful for throughput, but changing `N` can change both vocab partitioning and +the final log-sum-exp merge order. + +The batch-invariant entry point instead uses: + +- fixed `BM=256`, `BN=64`, and `BK=32` tiles; +- fixed ascending hidden-dimension traversal with FP32 tensor-core accumulation; +- no split-K reduction; +- a split-V schedule derived only from `V`, capped at 32 contiguous ranges; +- an ascending, fixed-order merge of per-split online-softmax states. + +Both entry points reuse the same validated TMA + `mma.sync` device kernel. Their +host launch policies are intentionally distinct so performance tuning cannot +silently weaken the batch-invariance contract. + +## Tensor Contract + +| Value | Shape | Dtype | Device and layout | +| --- | --- | --- | --- | +| `hidden` | `[*lead, D]`, at least 2-D, non-empty | BF16 | Hopper CUDA device; non-contiguous layouts are copied | +| `lm_head_weight` | `[V, D]`, `V > 0` | BF16 | Same device; non-contiguous layouts are copied | +| `target_ids` | `[*lead]` | `uint8`, `int8`, `int16`, `int32`, or `int64` | Same device; converted to contiguous int32 without wrapping extreme int64 values | +| `bias` | Optional `[V]` | FP16, BF16, or FP32 | Same device; consumed as FP32 | +| output | `[*lead]` | FP32 | Same device | + +`D` must be positive and divisible by 32. A contiguous tensor can still have a +misaligned non-zero storage offset; the CUDA owner detects this and conditionally +clones only the affected TMA input to a 16-byte-aligned allocation. + +Target IDs must be in `[0, V)`. The default `validate=False` stays asynchronous +and writes `NaN` for an invalid row instead of silently returning a finite but +wrong log-probability. Pass `validate=True` at trust boundaries to raise a +detailed `ValueError`; that diagnostic mode performs a GPU-to-CPU synchronization. +Padding and ignore-index values should normally be filtered before the call. + +Backward, tensor parallelism, and fallback are deliberately unsupported. The +wrapper rejects differentiable execution while autograd is enabled, and registry +dispatch fails closed outside Hopper or when the compiled symbol is absent. + +Model parameters may still have `requires_grad=True` during rollout; call the +operator inside `torch.no_grad()` or `torch.inference_mode()`. + +## Memory + +The forward never stores `[N, V]` logits or probabilities. Its partial-reduction +workspace is: + +```text +3 * num_vocab_splits * N * sizeof(float) +``` + +for the partial max, exponential sum, and selected target logit. The workspace +grows with the number of vocab splits until the 32-split cap; its worst-case bound +is therefore `O(N)` and independent of `V`. The entry point also allocates FP32 +`logp` and `lse` outputs (`2 * N` values), `O(N)` target conversion/sentinel +buffers for non-int32 or non-contiguous IDs, and up to `V` FP32 values when a +non-FP32 or non-contiguous bias must be converted. Layout or TMA-alignment fixes +can additionally create conditional input copies; aligned contiguous inputs stay +zero-copy. + +## Accuracy + +The hidden projection uses BF16 tensor-core inputs with FP32 accumulation. The +online softmax and split-state merge are FP32. Correctness tests compare against: + +```python +from rl_engine.kernels.ops.pytorch.loss.linear_logp import NativeLinearLogpOp + +expected = NativeLinearLogpOp().forward_fp32(hidden, weight, target_ids, bias) +``` + +Reference parity is tolerance-based because the reference uses a different GEMM +and log-sum-exp reduction. Batch-invariance checks within the SM90 backend use +exact `torch.equal` comparisons. + +## Tests + +```bash +RL_KERNEL_REQUIRE_EXT=1 \ + python -m pytest tests/test_batch_invariant_linear_logp.py -q -rs +``` + +The explicit extension gate prevents a Hopper build that omitted the SM90 +symbol from reporting these feature checks as hardware skips. Ordinary test +runs without that flag still skip the optional suite when the symbol is absent. + +GPU coverage includes: + +- correctness with and without bias; +- repeated-run bitwise determinism; +- a fixed probe row at multiple batch positions; +- unrelated neighboring-row noise; +- batch size 1 versus 4096 rows; +- full-batch versus chunk sizes around `BM=256`; +- hidden-tile traversals from `D=32` through production-style `D=4096`; +- vocab-tile and split-cap boundaries at `V=63/64/65` and `V=2048/2049`; +- a production-style `V=50257` with edge and partial vocab tiles; +- strided public inputs and contiguous inputs with misaligned storage offsets; +- adversarial mutation, temporary copies, launch, and dependent consumption on a non-default stream; +- supported bias dtypes and rejection of unsupported bias dtypes; +- every supported target dtype, extreme int64 sentinels, asynchronous invalid-target + `NaN`, opt-in range errors, and the forward-only boundary. + +The batch-size test uses a vocabulary large enough to cross the adaptive split-V +boundary in the regular fused path. + +The shared operator checker can run forward correctness on Hopper: + +```bash +python scripts/check_operator.py \ + --op batch_invariant_linear_logp \ + --candidate cuda-sm90 \ + --device cuda \ + --arch-key sm90 \ + --dtype bf16 \ + --batch 4 \ + --seq 32 \ + --normalized-dim 128 \ + --vocab 4096 +``` + +Do not pass `--check-grad`; backward is outside this operator's contract. + +## Benchmark + +```bash +python benchmarks/benchmark_batch_invariant_linear_logp.py +python benchmarks/benchmark_batch_invariant_linear_logp.py \ + --configs "4096,4096,151936" +``` + +This is a manual comparative microbenchmark over the same logical values, so it +is intentionally separate from the profiler's single-workload registry. The two +fused symbols share prepared BF16 tensors and exclude Python validation from +timing. The materialized batch-invariant LM-head + logp reference uses FP32 copies +created after the fused measurements. Each fused path reports maximum absolute +error against that reference; memory columns are incremental allocations above +each path's prepared-input baseline. Default cases include `N=1`, `32`, and `256` +to expose the fixed split-V schedule's small-batch occupancy tradeoff. Custom +configs require `D % 32 == 0` for the fused kernels and `V % 4 == 0` so the +materialized FP32 logp comparator remains on its SM90 TMA backend. + +## Implementation Files + +- `csrc/cuda/fused_linear_logp_sm90.cu` +- `csrc/ops.cpp` +- `rl_engine/kernels/ops/cuda/loss/batch_invariant_linear_logp.py` +- `rl_engine/kernels/registry.py` +- `tests/test_batch_invariant_linear_logp.py` +- `benchmarks/benchmark_batch_invariant_linear_logp.py` diff --git a/docs/operators/linear-logp.md b/docs/operators/linear-logp.md index 4b5231ef..b82cec9d 100644 --- a/docs/operators/linear-logp.md +++ b/docs/operators/linear-logp.md @@ -13,6 +13,11 @@ This differs from [Fused LogP](fused-logp.md), which takes already-materialized logits as input. Here the LM-head projection is fused into the forward reduction, so the forward `[N, V]` tensor never lands in HBM. +The default operator is throughput-oriented and does not promise bitwise batch +invariance: its SM90 split-V schedule may adapt to the token-row count. Use +[Batch-Invariant Fused Linear LogP](batch-invariant-linear-logp.md) when the same +row must produce identical bits across batch layouts. + ## Entry Point ```python diff --git a/rl_engine/_C.pyi b/rl_engine/_C.pyi index 20f17461..62f0844c 100644 --- a/rl_engine/_C.pyi +++ b/rl_engine/_C.pyi @@ -15,6 +15,12 @@ def fused_linear_logp_sm90( target: torch.Tensor, bias: torch.Tensor | None, ) -> list[torch.Tensor]: ... +def batch_invariant_linear_logp_sm90( + hidden: torch.Tensor, + weight: torch.Tensor, + target: torch.Tensor, + bias: torch.Tensor | None, +) -> list[torch.Tensor]: ... def fused_linear_logp_sm90_global_target( hidden: torch.Tensor, weight: torch.Tensor, diff --git a/rl_engine/kernels/gtest/operator_inputs.py b/rl_engine/kernels/gtest/operator_inputs.py index 835ee0e4..4bb0504a 100644 --- a/rl_engine/kernels/gtest/operator_inputs.py +++ b/rl_engine/kernels/gtest/operator_inputs.py @@ -31,6 +31,7 @@ def make_operator_inputs( "attention": _make_attention_inputs, "logp": _make_logp_inputs, "linear_logp": _make_linear_logp_inputs, + "batch_invariant_linear_logp": _make_linear_logp_inputs, "batch_invariant_logp": _make_batch_invariant_logp_inputs, "rope": _make_rope_inputs, "silu": _make_silu_inputs, @@ -55,6 +56,7 @@ def operator_shape_name(op_name: str, args: argparse.Namespace) -> str: "attention": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", "logp": f"{batch}x{seq}x{vocab}", "linear_logp": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", + "batch_invariant_linear_logp": f"{batch}x{seq}x{_normalized_dim(args)}x{vocab}", "batch_invariant_logp": f"{batch}x{seq}x{vocab}", "rope": f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}", "silu": f"{batch}x{seq}x{DEFAULT_INTERMEDIATE}", diff --git a/rl_engine/kernels/gtest/operator_specs.py b/rl_engine/kernels/gtest/operator_specs.py index 08925021..ad5809bb 100644 --- a/rl_engine/kernels/gtest/operator_specs.py +++ b/rl_engine/kernels/gtest/operator_specs.py @@ -5,7 +5,7 @@ import argparse import importlib -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any import torch @@ -22,6 +22,7 @@ class OperatorSpec: gold_method: str candidate_paths: dict[str, str] grad_input_names: tuple[str, ...] = () + candidate_methods: dict[str, str] = field(default_factory=dict) def _load_object(path: str) -> Any: @@ -88,6 +89,20 @@ def _load_object(path: str) -> Any: }, grad_input_names=("hidden", "lm_head_weight"), ), + "batch_invariant_linear_logp": OperatorSpec( + name="batch_invariant_linear_logp", + op_class="logprob", + gold_path="rl_engine.kernels.ops.pytorch.loss.linear_logp.NativeLinearLogpOp", + gold_method="forward_fp32", + candidate_paths={ + "pytorch": "rl_engine.kernels.ops.pytorch.loss.linear_logp.NativeLinearLogpOp", + "cuda-sm90": ( + "rl_engine.kernels.ops.cuda.loss.batch_invariant_linear_logp." + "BatchInvariantLinearLogpSM90Op" + ), + }, + candidate_methods={"pytorch": "forward_fp32"}, + ), "embedding": OperatorSpec( name="embedding", op_class="elementwise", @@ -214,6 +229,9 @@ def make_candidate(args: argparse.Namespace) -> CandidateSpec: if candidate_name in spec.candidate_paths: candidate_op = _load_object(spec.candidate_paths[candidate_name])() + candidate_method = spec.candidate_methods.get(candidate_name) + if candidate_method is not None: + candidate_op = getattr(candidate_op, candidate_method) if args.op == "logp" and candidate_name == "cuda-sm90": candidate_op = _LogpSM90CandidateAdapter(candidate_op) return CandidateSpec( diff --git a/rl_engine/kernels/ops/cuda/loss/batch_invariant_linear_logp.py b/rl_engine/kernels/ops/cuda/loss/batch_invariant_linear_logp.py new file mode 100644 index 00000000..06e7b3fc --- /dev/null +++ b/rl_engine/kernels/ops/cuda/loss/batch_invariant_linear_logp.py @@ -0,0 +1,166 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Single-card batch-invariant fused linear log-probability on SM90.""" + +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 + +_INTEGER_DTYPES = { + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, +} +_BIAS_DTYPES = {torch.float16, torch.bfloat16, torch.float32} +_HIDDEN_TILE = 32 + + +def _validate_target_range(target_ids: torch.Tensor, vocab_size: int) -> None: + # Compare in int64 so Python scalars such as a large vocabulary size do not + # overflow narrow integer tensors during validation. + target_i64 = target_ids.to(dtype=torch.int64) + invalid = (target_i64 < 0) | (target_i64 >= vocab_size) + if bool(invalid.any()): + invalid_targets = target_i64[invalid] + target_min = int(invalid_targets.min()) + target_max = int(invalid_targets.max()) + raise ValueError( + f"target_ids must be in [0, {vocab_size}), got invalid range " + f"[{target_min}, {target_max}]" + ) + + +def _validate_inputs( + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor], + *, + validate_targets: bool = False, +) -> None: + if hidden.dim() < 2: + raise ValueError(f"hidden must be at least 2-D [*lead, D], got {tuple(hidden.shape)}") + if lm_head_weight.dim() != 2: + raise ValueError(f"lm_head_weight must be 2-D [V, D], got {tuple(lm_head_weight.shape)}") + 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 weight dim {lm_head_weight.size(1)}" + ) + if hidden.size(-1) == 0: + raise ValueError("hidden dimension must be positive") + if target_ids.numel() == 0: + raise ValueError("hidden must contain at least one token row") + if lm_head_weight.size(0) == 0: + raise ValueError("lm_head_weight must contain at least one vocabulary row") + if not hidden.is_cuda or not lm_head_weight.is_cuda or not target_ids.is_cuda: + raise ValueError("hidden, lm_head_weight, and target_ids must be CUDA tensors") + if lm_head_weight.device != hidden.device or target_ids.device != hidden.device: + raise ValueError("hidden, lm_head_weight, and target_ids must be on the same CUDA device") + if hidden.dtype != torch.bfloat16 or lm_head_weight.dtype != torch.bfloat16: + raise TypeError("batch_invariant_linear_logp_sm90 supports bf16 hidden and weight only") + if target_ids.dtype not in _INTEGER_DTYPES: + raise TypeError(f"target_ids must have an integer dtype, got {target_ids.dtype}") + if hidden.size(-1) % _HIDDEN_TILE != 0: + raise ValueError(f"hidden dim must be a multiple of {_HIDDEN_TILE}, got {hidden.size(-1)}") + if torch.cuda.get_device_capability(hidden.device)[0] != 9: + raise RuntimeError("batch_invariant_linear_logp_sm90 requires an SM90 Hopper GPU") + + if bias is not None: + if bias.dim() != 1 or bias.numel() != lm_head_weight.size(0): + raise ValueError( + f"bias must have shape ({lm_head_weight.size(0)},), got {tuple(bias.shape)}" + ) + if not bias.is_cuda or bias.device != hidden.device: + raise ValueError("bias must be a CUDA tensor on the same device as hidden") + if bias.dtype not in _BIAS_DTYPES: + raise TypeError(f"bias must have dtype fp16, bf16, or fp32, got {bias.dtype}") + + if torch.is_grad_enabled() and ( + hidden.requires_grad + or lm_head_weight.requires_grad + or (bias is not None and bias.requires_grad) + ): + raise RuntimeError( + "batch_invariant_linear_logp_sm90 is forward-only; call it under " + "torch.no_grad() or use a differentiable linear_logp backend" + ) + + if validate_targets: + _validate_target_range(target_ids, lm_head_weight.size(0)) + + +class BatchInvariantLinearLogpSM90Op: + """Fused ``log_softmax(hidden @ weight.T + bias)[target]`` for Hopper. + + The operator never materializes ``[N, V]`` logits. Its SM90 launch uses a + fixed D traversal and a split-V schedule derived only from V, so a row's + float32 output is bitwise invariant to batch size, row position, neighboring + rows, and batch-dimension chunking. Backward and tensor parallelism are + intentionally outside this first single-card contract. + """ + + op_class = "logprob" + is_batch_invariant = True + supports_backward = False + + def __init__(self) -> None: + if not _EXT_AVAILABLE or not hasattr(_C, "batch_invariant_linear_logp_sm90"): + raise RuntimeError( + "batch_invariant_linear_logp_sm90 is not compiled into the extension. " + "Rebuild on Hopper with KERNEL_ALIGN_FORCE_SM90=1." + ) + logger.info("Linked _C.batch_invariant_linear_logp_sm90.") + + def __call__( + self, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + validate: bool = False, + ) -> torch.Tensor: + return self.forward(hidden, lm_head_weight, target_ids, bias, validate=validate) + + def forward( + self, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + validate: bool = False, + ) -> torch.Tensor: + _validate_inputs( + hidden, + lm_head_weight, + target_ids, + bias, + validate_targets=validate, + ) + + lead_shape = hidden.shape[:-1] + hidden_2d = hidden.reshape(-1, hidden.size(-1)).contiguous() + weight_2d = lm_head_weight.contiguous() + target_1d = target_ids.reshape(-1) + + logp, _lse = _C.batch_invariant_linear_logp_sm90( + hidden_2d, + weight_2d, + target_1d, + bias, + ) + return logp.reshape(lead_shape) diff --git a/rl_engine/kernels/ops/pytorch/loss/linear_logp.py b/rl_engine/kernels/ops/pytorch/loss/linear_logp.py index 620b3c24..6c46c345 100644 --- a/rl_engine/kernels/ops/pytorch/loss/linear_logp.py +++ b/rl_engine/kernels/ops/pytorch/loss/linear_logp.py @@ -8,6 +8,8 @@ import torch +from rl_engine.kernels.ops.pytorch.linear.lm_head import NativeLMHeadOp + # 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``. @@ -662,3 +664,51 @@ def apply( 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) + + def forward_fp32( + self, + hidden: torch.Tensor, + lm_head_weight: torch.Tensor, + target_ids: torch.Tensor, + bias: Optional[torch.Tensor] = None, + *, + tp_group: Any = None, + vocab_start_index: int = 0, + global_vocab_size: Optional[int] = None, + ) -> torch.Tensor: + """FP32 oracle without low-precision logit rounding.""" + 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)}" + ) + if should_use_tensor_parallel_linear_logp( + tp_group, + int(vocab_start_index), + global_vocab_size, + lm_head_weight.size(0), + ): + hidden_fp32 = hidden.float() + weight_fp32 = lm_head_weight.float() + bias_fp32 = None if bias is None else bias.float() + with NativeLMHeadOp._strict_fp32_matmul(hidden.device.type): + return tensor_parallel_linear_logp( + hidden_fp32, + weight_fp32, + target_ids, + bias_fp32, + tp_group=tp_group, + vocab_start_index=vocab_start_index, + global_vocab_size=global_vocab_size, + ) + + logits = NativeLMHeadOp().forward_fp32(hidden, lm_head_weight, bias=bias) + with torch.autocast(device_type=logits.device.type, enabled=False): + log_probs = torch.log_softmax(logits, dim=-1) + target = target_ids.to(device=logits.device, dtype=torch.long).unsqueeze(-1) + return torch.gather(log_probs, dim=-1, index=target).squeeze(-1) diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 89a4c29b..50b2b338 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -50,6 +50,10 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): CUDA_FUSED_LINEAR_LOGP_SM90 = ( "rl_engine.kernels.ops.cuda.loss.linear_logp.FusedLinearLogpSM90Op" ) + CUDA_BATCH_INVARIANT_LINEAR_LOGP_SM90 = ( + "rl_engine.kernels.ops.cuda.loss.batch_invariant_linear_logp." + "BatchInvariantLinearLogpSM90Op" + ) TRITON_LINEAR_LOGP = "rl_engine.kernels.ops.triton.loss.linear_logp.TritonLinearLogpOp" PYTORCH_LINEAR_LOGP = "rl_engine.kernels.ops.pytorch.loss.linear_logp.NativeLinearLogpOp" # Fused policy-ratio + KL-penalty front-end (PPO/GRPO), logits -> (ratio, kl) @@ -207,6 +211,10 @@ def __init__(self): OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP, ], + # Strict single-card contract: no non-invariant fallback. + "batch_invariant_linear_logp": [ + OpBackend.CUDA_BATCH_INVARIANT_LINEAR_LOGP_SM90, + ], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], "pack": [OpBackend.PYTORCH_PACK], "det_gemm": [OpBackend.CUDA_DET_GEMM, OpBackend.TRITON_DET_GEMM], @@ -253,6 +261,7 @@ def __init__(self): "grpo_loss": [OpBackend.TRITON_GRPO_LOSS, OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.TRITON_ROPE, OpBackend.PYTORCH_NATIVE_ROPE], "linear_logp": [OpBackend.TRITON_LINEAR_LOGP, OpBackend.PYTORCH_LINEAR_LOGP], + "batch_invariant_linear_logp": [], "ratio_kl": [OpBackend.TRITON_RATIO_KL, OpBackend.PYTORCH_RATIO_KL], "pack": [OpBackend.PYTORCH_PACK], "det_gemm": [OpBackend.TRITON_DET_GEMM], @@ -280,6 +289,7 @@ def __init__(self): "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], "linear_logp": [OpBackend.PYTORCH_LINEAR_LOGP], + "batch_invariant_linear_logp": [], "ratio_kl": [OpBackend.PYTORCH_RATIO_KL], "pack": [OpBackend.PYTORCH_PACK], "det_gemm": [], @@ -301,6 +311,7 @@ def __init__(self): "grpo_loss": [OpBackend.PYTORCH_GRPO_LOSS], "rope": [OpBackend.PYTORCH_NATIVE_ROPE], "linear_logp": [OpBackend.PYTORCH_LINEAR_LOGP], + "batch_invariant_linear_logp": [], "ratio_kl": [OpBackend.PYTORCH_RATIO_KL], "pack": [OpBackend.PYTORCH_PACK], "batch_invariant_logp": [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP], @@ -405,6 +416,8 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: + if not self._backend_supports_device(backend, device): + continue if backend.name in self._instance_cache: return self._instance_cache[backend.name] @@ -425,6 +438,34 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + @staticmethod + def _backend_supports_device( + backend: OpBackend, + device: torch.device | str | None, + ) -> bool: + if backend is not OpBackend.CUDA_BATCH_INVARIANT_LINEAR_LOGP_SM90: + return True + if torch.version.hip is not None: + return False + + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + if not _EXT_AVAILABLE or not hasattr(_C, "batch_invariant_linear_logp_sm90"): + return False + + resolved = torch.device("cuda" if device is None else device) + if resolved.type != "cuda": + return False + try: + return torch.cuda.get_device_capability(resolved)[0] == 9 + except Exception as error: + logger.warning( + "Failed to probe %s for batch-invariant linear_logp: %s", + resolved, + error, + ) + return False + def _platform_for_device(self, device: torch.device | str | None) -> str: if device is None: if device_ctx.is_rocm: diff --git a/rl_engine/tests/test_dispatch.py b/rl_engine/tests/test_dispatch.py index eba66da8..093a4f40 100644 --- a/rl_engine/tests/test_dispatch.py +++ b/rl_engine/tests/test_dispatch.py @@ -118,6 +118,13 @@ def test_musa_det_gemm_fails_closed(self, monkeypatch): with pytest.raises(RuntimeError, match="No functional backend"): registry.get_op("det_gemm", device="musa") + def test_musa_batch_invariant_linear_logp_fails_closed(self, monkeypatch): + self._mock_musa_device(monkeypatch) + registry = KernelRegistry() + + with pytest.raises(RuntimeError, match="No functional backend"): + registry.get_op("batch_invariant_linear_logp", device="musa") + def test_registry_explicit_device_selects_device_platform(monkeypatch): registry = KernelRegistry() diff --git a/scripts/ci_smoke.py b/scripts/ci_smoke.py index 0cbe1cda..47341c0c 100644 --- a/scripts/ci_smoke.py +++ b/scripts/ci_smoke.py @@ -11,8 +11,9 @@ succeeds (``dlopen`` does not check arch) but the first kernel launch raises ``cudaErrorNoKernelImageForDevice`` once the stream is synchronized. -Uses only ``fused_logp`` - the op registered unconditionally in ``csrc/ops.cpp`` - -so it does not require ``KERNEL_ALIGN_FORCE_SM90=1`` / a Hopper build. +Every GPU runs the generic ``fused_logp`` launch. Hopper additionally requires +and launches the batch-invariant fused linear-logp symbol, so an SM90 build +cannot pass CI while silently omitting the feature under test. """ import sys @@ -70,7 +71,49 @@ def main() -> int: ) return 1 - print(f"[smoke] OK: rl_engine._C built and fused_logp ran on sm_{cc[0]}{cc[1]}.") + if cc[0] == 9: + symbol = "batch_invariant_linear_logp_sm90" + if not hasattr(_C, symbol): + print( + f"[smoke] FATAL: Hopper extension is missing _C.{symbol}.\n" + " Rebuild with KERNEL_ALIGN_FORCE_SM90=1 and TARGET_SM=9.0.", + file=sys.stderr, + ) + return 1 + try: + hidden = torch.randn(1, 32, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(64, 32, device="cuda", dtype=torch.bfloat16) + target = torch.tensor([63], device="cuda", dtype=torch.int32) + logp, lse = _C.batch_invariant_linear_logp_sm90( + hidden, + weight, + target, + None, + ) + torch.cuda.synchronize() + except Exception as exc: + print( + f"[smoke] FATAL: _C.{symbol} failed to launch on sm_{cc[0]}{cc[1]}.\n" + f" Underlying error: {type(exc).__name__}: {exc}", + file=sys.stderr, + ) + return 1 + + if tuple(logp.shape) != (1,) or tuple(lse.shape) != (1,): + print( + f"[smoke] FATAL: _C.{symbol} returned shapes " + f"{tuple(logp.shape)}, {tuple(lse.shape)}", + file=sys.stderr, + ) + return 1 + if not bool(torch.isfinite(logp).all() and torch.isfinite(lse).all()): + print( + f"[smoke] FATAL: _C.{symbol} returned non-finite values", + file=sys.stderr, + ) + return 1 + + print(f"[smoke] OK: required extension kernels ran on sm_{cc[0]}{cc[1]}.") return 0 diff --git a/tests/test_batch_invariant_linear_logp.py b/tests/test_batch_invariant_linear_logp.py new file mode 100644 index 00000000..93b97ec2 --- /dev/null +++ b/tests/test_batch_invariant_linear_logp.py @@ -0,0 +1,768 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +from __future__ import annotations + +import os +import re +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from rl_engine.kernels.gtest.op_checks import run_operator_suite +from rl_engine.kernels.gtest.operator_specs import OP_SPECS, make_candidate, make_operator_case +from rl_engine.kernels.ops.cuda.loss import batch_invariant_linear_logp as op_module +from rl_engine.kernels.ops.pytorch.loss.linear_logp import NativeLinearLogpOp + + +def _sm90_op_available() -> bool: + if not torch.cuda.is_available(): + return False + try: + if torch.cuda.get_device_capability()[0] != 9: + return False + # GPU CI and explicit feature verification must not turn a missing + # symbol into a skip. Let the op constructor fail with its build hint. + if os.environ.get("RL_KERNEL_REQUIRE_EXT") == "1": + return True + + from rl_engine.kernels.ops.base import _C, _EXT_AVAILABLE + + return bool(_EXT_AVAILABLE and hasattr(_C, "batch_invariant_linear_logp_sm90")) + except Exception: + return False + + +requires_sm90_op = pytest.mark.skipif( + not _sm90_op_available(), + reason="batch-invariant linear_logp requires Hopper and the compiled SM90 symbol", +) + + +def test_sm90_availability_hard_mode_does_not_skip_a_missing_symbol(monkeypatch): + from rl_engine.kernels.ops import base as base_module + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (9, 0)) + monkeypatch.setattr(base_module, "_EXT_AVAILABLE", False) + monkeypatch.setattr(base_module, "_C", object()) + + monkeypatch.delenv("RL_KERNEL_REQUIRE_EXT", raising=False) + assert not _sm90_op_available() + + monkeypatch.setenv("RL_KERNEL_REQUIRE_EXT", "1") + assert _sm90_op_available() + + +def _make_inputs( + seed: int, + *, + num_tokens: int, + hidden_dim: int, + vocab_size: int, + bias: bool, +): + generator = torch.Generator(device="cuda").manual_seed(seed) + hidden = torch.randn( + num_tokens, + hidden_dim, + device="cuda", + dtype=torch.bfloat16, + generator=generator, + ) + weight = torch.randn( + vocab_size, + hidden_dim, + device="cuda", + dtype=torch.bfloat16, + generator=generator, + ) + target_ids = torch.randint( + 0, + vocab_size, + (num_tokens,), + device="cuda", + generator=generator, + ) + bias_tensor = ( + torch.randn( + vocab_size, + device="cuda", + dtype=torch.bfloat16, + generator=generator, + ) + if bias + else None + ) + return hidden, weight, target_ids, bias_tensor + + +def _reference(hidden, weight, target_ids, bias): + return NativeLinearLogpOp().forward_fp32( + hidden, + weight, + target_ids, + bias, + ) + + +def test_cuda_source_pins_batch_invariant_split_to_vocab_only(): + source_path = ( + Path(__file__).resolve().parents[1] / "csrc" / "cuda" / "fused_linear_logp_sm90.cu" + ) + source = source_path.read_text(encoding="utf-8") + + match = re.search( + r"int select_batch_invariant_vocab_splits\(int total_vtiles\) " r"\{(?P.*?)\n\}", + source, + flags=re.DOTALL, + ) + assert match is not None + invariant_branch = match.group("body") + assert "total_vtiles" in invariant_branch + assert "MAX_BATCH_INVARIANT_VOCAB_SPLITS" in invariant_branch + assert "row_blocks" not in invariant_branch + assert "multiProcessorCount" not in invariant_branch + assert "MAX_BATCH_INVARIANT_VOCAB_SPLITS = 32" in source + + entry_start = source.index( + "std::vector batch_invariant_linear_logp_sm90_forward(" + ) + entry_end = source.index("// Vocab-parallel local-shard forward", entry_start) + invariant_entry = source[entry_start:entry_end] + assert "VocabSplitPolicy::kBatchInvariant" in invariant_entry + assert "VocabSplitPolicy::kThroughput" not in invariant_entry + + +def _operator_checker_args(candidate: str, *, seq: int = 16) -> SimpleNamespace: + return SimpleNamespace( + op="batch_invariant_linear_logp", + candidate=candidate, + arch_key="sm90" if candidate == "cuda-sm90" else None, + batch=2, + seq=seq, + vocab=257, + seed=123, + input_mode="random", + normalized_dim=128, + ) + + +def test_operator_checker_uses_strict_fp32_gold_and_supports_default_candidate(): + invariant_spec = OP_SPECS["batch_invariant_linear_logp"] + args = _operator_checker_args("pytorch") + case = make_operator_case(args, torch.float32, torch.device("cpu")) + + assert invariant_spec.gold_method == "forward_fp32" + assert OP_SPECS["linear_logp"].gold_method == "apply" + for candidate_name in ("pytorch", "native"): + args.candidate = candidate_name + candidate = make_candidate(args) + report = run_operator_suite( + "batch_invariant_linear_logp", + candidates=[candidate], + cases=[case], + ) + + assert candidate.backend == "pytorch" + assert getattr(candidate.fn, "__name__", None) == "forward_fp32" + assert report.passed + + +@requires_sm90_op +def test_operator_checker_runs_cuda_sm90_candidate_end_to_end(): + args = _operator_checker_args("cuda-sm90", seq=4) + case = make_operator_case(args, torch.bfloat16, torch.device("cuda")) + candidate = make_candidate(args) + + with torch.no_grad(): + report = run_operator_suite( + "batch_invariant_linear_logp", + candidates=[candidate], + cases=[case], + ) + + assert candidate.backend == "cuda-sm90" + assert report.passed + + +@pytest.mark.parametrize( + ("config", "error"), + [ + ("1,31,64", "D must be divisible by 32"), + ("1,32,65", "V must be divisible by 4"), + ("1,32", "must contain N,D,V"), + ], +) +def test_benchmark_rejects_configs_that_leave_the_sm90_comparison_path(config, error): + from benchmarks.benchmark_batch_invariant_linear_logp import _parse_configs + + with pytest.raises(ValueError, match=error): + _parse_configs(config) + + +@pytest.mark.parametrize( + ("warmup", "iterations", "error"), + [ + (-1, 1, "warmup must be non-negative"), + (0, 0, "iterations must be positive"), + (0, -1, "iterations must be positive"), + ], +) +def test_benchmark_rejects_invalid_run_counts(warmup, iterations, error): + from benchmarks.benchmark_batch_invariant_linear_logp import _validate_run_counts + + with pytest.raises(ValueError, match=error): + _validate_run_counts(warmup, iterations) + + +@pytest.mark.parametrize( + ("arguments", "error"), + [ + (["--configs", "1,32"], "must contain N,D,V"), + (["--configs", "one,32,64"], "integer N,D,V triples"), + (["--warmup", "-1"], "warmup must be non-negative"), + (["--iterations", "0"], "iterations must be positive"), + ], +) +def test_benchmark_cli_reports_invalid_arguments(monkeypatch, capsys, arguments, error): + from benchmarks.benchmark_batch_invariant_linear_logp import parse_args + + monkeypatch.setattr(sys, "argv", ["benchmark_batch_invariant_linear_logp.py", *arguments]) + with pytest.raises(SystemExit) as caught: + parse_args() + + assert caught.value.code == 2 + assert error in capsys.readouterr().err + + +def test_wrapper_uses_dedicated_extension_symbol_and_preserves_leading_shape( + monkeypatch, +): + calls = [] + validation_options = [] + + class FakeExtension: + @staticmethod + def batch_invariant_linear_logp_sm90(hidden, weight, target_ids, bias): + calls.append((hidden, weight, target_ids, bias)) + num_tokens = hidden.size(0) + return torch.arange(num_tokens, dtype=torch.float32), torch.zeros(num_tokens) + + monkeypatch.setattr(op_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(op_module, "_C", FakeExtension()) + + def fake_validate(*args, **kwargs): + validation_options.append(kwargs["validate_targets"]) + + monkeypatch.setattr(op_module, "_validate_inputs", fake_validate) + + hidden = torch.randn(2, 3, 32) + weight = torch.randn(17, 32) + target_ids = torch.randint(0, 17, (2, 3)) + op = op_module.BatchInvariantLinearLogpSM90Op() + output = op(hidden, weight, target_ids, validate=True) + + assert output.shape == (2, 3) + assert output.dtype == torch.float32 + assert validation_options == [True] + assert len(calls) == 1 + called_hidden, called_weight, called_targets, called_bias = calls[0] + assert called_hidden.shape == (6, 32) and called_hidden.is_contiguous() + assert called_weight.shape == (17, 32) and called_weight.is_contiguous() + assert called_targets.shape == (6,) and called_targets.dtype == target_ids.dtype + assert called_bias is None + + +def test_wrapper_fails_closed_when_extension_symbol_is_missing(monkeypatch): + monkeypatch.setattr(op_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(op_module, "_C", object()) + + with pytest.raises(RuntimeError, match="not compiled"): + op_module.BatchInvariantLinearLogpSM90Op() + + +def test_validation_rejects_shape_mismatch_before_device_probe(): + hidden = torch.randn(2, 3, 32) + weight = torch.randn(17, 32) + target_ids = torch.randint(0, 17, (2, 4)) + + with pytest.raises(ValueError, match="leading shape"): + op_module._validate_inputs(hidden, weight, target_ids, None) + + +@pytest.mark.parametrize( + "target_ids", + [ + torch.tensor([0, 100], dtype=torch.int8), + torch.tensor([0, 250], dtype=torch.uint8), + torch.tensor([0, 499], dtype=torch.int16), + ], +) +def test_target_range_validation_avoids_narrow_integer_overflow(target_ids): + op_module._validate_target_range(target_ids, vocab_size=500) + + +@requires_sm90_op +@pytest.mark.parametrize("bias_dtype", [None, torch.bfloat16, torch.float16, torch.float32]) +def test_sm90_matches_fp32_reference(bias_dtype): + op = op_module.BatchInvariantLinearLogpSM90Op() + hidden, weight, target_ids, bias = _make_inputs( + 310, + num_tokens=97, + hidden_dim=128, + vocab_size=500, + bias=bias_dtype is not None, + ) + if bias is not None: + bias = bias.to(dtype=bias_dtype) + + output = op(hidden, weight, target_ids, bias) + expected = _reference(hidden, weight, target_ids, bias) + + assert output.dtype == torch.float32 + assert torch.allclose(output, expected, atol=2e-2, rtol=0.0) + + +@requires_sm90_op +@pytest.mark.parametrize("bias_dtype", [torch.float16, torch.float32]) +def test_sm90_preserves_bias_precision_above_bfloat16(bias_dtype): + op = op_module.BatchInvariantLinearLogpSM90Op() + hidden = torch.zeros(1, 32, device="cuda", dtype=torch.bfloat16) + weight = torch.zeros(64, 32, device="cuda", dtype=torch.bfloat16) + target_ids = torch.zeros(1, device="cuda", dtype=torch.int64) + bias = torch.zeros(64, device="cuda", dtype=bias_dtype) + bias[0] = -8.03 + + actual = op(hidden, weight, target_ids, bias) + expected = _reference(hidden, weight, target_ids, bias) + rounded_bf16 = _reference(hidden, weight, target_ids, bias.to(torch.bfloat16)) + + assert torch.allclose(actual, expected, atol=2e-3, rtol=0.0) + assert (expected - rounded_bf16).abs().item() > 2e-2 + + +@requires_sm90_op +@pytest.mark.parametrize( + "hidden_dim,vocab_size", + [ + (32, 63), + (64, 64), + (96, 65), + (128, 2048), + (160, 2049), + (4096, 4097), + ], +) +def test_sm90_matches_fp32_reference_across_hidden_and_vocab_tiles( + hidden_dim, + vocab_size, +): + op = op_module.BatchInvariantLinearLogpSM90Op() + num_tokens = 3 if hidden_dim >= 4096 else 17 + hidden, weight, target_ids, _ = _make_inputs( + 322 + hidden_dim + vocab_size, + num_tokens=num_tokens, + hidden_dim=hidden_dim, + vocab_size=vocab_size, + bias=False, + ) + + output = op(hidden, weight, target_ids) + expected = _reference(hidden, weight, target_ids, None) + probe = op(hidden[:1], weight, target_ids[:1]) + + assert torch.allclose(output, expected, atol=2e-2, rtol=0.0) + assert torch.equal(probe[0], output[0]) + + +@requires_sm90_op +def test_sm90_is_bitwise_invariant_across_batch_size_position_and_noise(): + op = op_module.BatchInvariantLinearLogpSM90Op() + hidden, weight, target_ids, _ = _make_inputs( + 311, + num_tokens=4096, + hidden_dim=128, + vocab_size=4096, + bias=False, + ) + probe_hidden = hidden[137].clone() + probe_target = target_ids[137].clone() + baseline = op(probe_hidden.unsqueeze(0), weight, probe_target.unsqueeze(0))[0] + + probe_positions = (0, 127, 128, 255, 256, 2048, 4095) + for position in probe_positions: + hidden[position].copy_(probe_hidden) + target_ids[position].copy_(probe_target) + + packed = op(hidden, weight, target_ids) + for position in probe_positions: + assert torch.equal(packed[position], baseline) + + repeated = op(hidden, weight, target_ids) + assert torch.equal(repeated, packed) + + +@requires_sm90_op +@pytest.mark.parametrize( + "num_tokens,vocab_size,chunk_sizes", + [ + (33, 500, (1, 7)), + (4096, 4096, (255, 256, 257, 1024)), + ], +) +def test_sm90_full_batch_matches_batch_dimension_chunks_bitwise( + num_tokens, vocab_size, chunk_sizes +): + op = op_module.BatchInvariantLinearLogpSM90Op() + hidden, weight, target_ids, bias = _make_inputs( + 312, + num_tokens=num_tokens, + hidden_dim=128, + vocab_size=vocab_size, + bias=True, + ) + + full = op(hidden, weight, target_ids, bias) + for chunk_size in chunk_sizes: + chunked = torch.cat( + [ + op( + hidden[start : start + chunk_size], + weight, + target_ids[start : start + chunk_size], + bias, + ) + for start in range(0, hidden.size(0), chunk_size) + ] + ) + assert torch.equal(chunked, full), f"batch chunk size {chunk_size} changed output bits" + + +@requires_sm90_op +def test_sm90_production_vocab_matches_reference_and_is_batch_invariant(): + op = op_module.BatchInvariantLinearLogpSM90Op() + hidden, weight, target_ids, _ = _make_inputs( + 315, + num_tokens=17, + hidden_dim=128, + vocab_size=50257, + bias=False, + ) + target_ids[:8] = torch.tensor( + [0, 63, 64, 2047, 2048, 4095, 50255, 50256], + device="cuda", + ) + + full = op(hidden, weight, target_ids) + expected = _reference(hidden, weight, target_ids, None) + repeated = op(hidden, weight, target_ids) + chunked = torch.cat( + [ + op(hidden[start : start + 7], weight, target_ids[start : start + 7]) + for start in range(0, 17, 7) + ] + ) + probe = op(hidden[8:9], weight, target_ids[8:9]) + + assert torch.allclose(full, expected, atol=2e-2, rtol=0.0) + assert torch.equal(repeated, full) + assert torch.equal(chunked, full) + assert torch.equal(probe[0], full[8]) + + +@requires_sm90_op +def test_sm90_accepts_contiguous_tma_inputs_with_misaligned_storage_offsets(): + op = op_module.BatchInvariantLinearLogpSM90Op() + num_tokens, hidden_dim, vocab_size = 9, 128, 500 + hidden_storage = torch.randn( + num_tokens * hidden_dim + 1, + device="cuda", + dtype=torch.bfloat16, + ) + weight_storage = torch.randn( + vocab_size * hidden_dim + 1, + device="cuda", + dtype=torch.bfloat16, + ) + hidden = hidden_storage[1:].view(num_tokens, hidden_dim) + weight = weight_storage[1:].view(vocab_size, hidden_dim) + target_ids = torch.arange(num_tokens, device="cuda") % vocab_size + + assert hidden.is_contiguous() and hidden.data_ptr() % 16 != 0 + assert weight.is_contiguous() and weight.data_ptr() % 16 != 0 + + actual = op(hidden, weight, target_ids) + aligned = op(hidden.clone(), weight.clone(), target_ids) + expected = _reference(hidden, weight, target_ids, None) + + assert torch.equal(actual, aligned) + assert torch.allclose(actual, expected, atol=2e-2, rtol=0.0) + + +@requires_sm90_op +def test_sm90_wrapper_copies_strided_inputs_and_preserves_leading_shape(): + op = op_module.BatchInvariantLinearLogpSM90Op() + generator = torch.Generator(device="cuda").manual_seed(324) + hidden = torch.randn( + 2, + 3, + 192, + device="cuda", + dtype=torch.bfloat16, + generator=generator, + )[..., ::2] + weight = torch.randn( + 96, + 257, + device="cuda", + dtype=torch.bfloat16, + generator=generator, + ).t() + target_ids = torch.randint( + 0, + 257, + (2, 3, 2), + device="cuda", + generator=generator, + )[..., 0] + bias = torch.randn( + 257, + 2, + device="cuda", + dtype=torch.float16, + generator=generator, + )[:, 0] + + assert not hidden.is_contiguous() + assert not weight.is_contiguous() + assert not target_ids.is_contiguous() + assert not bias.is_contiguous() + + actual = op(hidden, weight, target_ids, bias) + expected = _reference(hidden, weight, target_ids, bias) + + assert actual.shape == (2, 3) + assert torch.allclose(actual, expected, atol=2e-2, rtol=0.0) + + +@requires_sm90_op +def test_sm90_uses_the_current_non_default_stream(): + op = op_module.BatchInvariantLinearLogpSM90Op() + desired_hidden, desired_weight, target_ids, desired_bias = _make_inputs( + 316, + num_tokens=33, + hidden_dim=128, + vocab_size=4096, + bias=True, + ) + hidden_storage = torch.zeros( + desired_hidden.numel() + 1, + device="cuda", + dtype=torch.bfloat16, + ) + weight_storage = torch.zeros( + desired_weight.numel() + 1, + device="cuda", + dtype=torch.bfloat16, + ) + bias_storage = torch.zeros( + desired_bias.numel() + 1, + device="cuda", + dtype=torch.bfloat16, + ) + hidden = hidden_storage[1:].view_as(desired_hidden) + weight = weight_storage[1:].view_as(desired_weight) + bias = bias_storage[1:] + + assert hidden.is_contiguous() and hidden.data_ptr() % 16 != 0 + assert weight.is_contiguous() and weight.data_ptr() % 16 != 0 + assert bias.is_contiguous() and bias.data_ptr() % 16 != 0 + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + + with torch.cuda.stream(stream): + # Delay the mutations on this stream. A kernel accidentally launched on + # the default stream would race ahead and consume the initial zeros. + torch.cuda._sleep(50_000_000) + hidden.copy_(desired_hidden) + weight.copy_(desired_weight) + bias.copy_(desired_bias) + output = op(hidden, weight, target_ids, bias, validate=False) + dependent = output.clone() + done = torch.cuda.Event() + done.record() + + done.synchronize() + expected = _reference(desired_hidden, desired_weight, target_ids, desired_bias) + assert torch.allclose(dependent, expected, atol=2e-2, rtol=0.0) + + +@requires_sm90_op +@pytest.mark.parametrize("requires_grad_input", ["hidden", "weight", "bias"]) +def test_sm90_forward_only_contract_rejects_autograd(requires_grad_input): + op = op_module.BatchInvariantLinearLogpSM90Op() + hidden, weight, target_ids, bias = _make_inputs( + 313, + num_tokens=8, + hidden_dim=128, + vocab_size=500, + bias=True, + ) + assert bias is not None + tensors = {"hidden": hidden, "weight": weight, "bias": bias} + tensors[requires_grad_input].requires_grad_(True) + + with pytest.raises(RuntimeError, match="forward-only"): + op(hidden, weight, target_ids, bias) + + with torch.no_grad(): + output = op(hidden, weight, target_ids, bias) + assert output.shape == (8,) + + +@requires_sm90_op +def test_sm90_invalid_targets_are_nan_or_raise_when_validation_is_requested(): + op = op_module.BatchInvariantLinearLogpSM90Op() + hidden, weight, target_ids, _ = _make_inputs( + 314, + num_tokens=8, + hidden_dim=128, + vocab_size=500, + bias=False, + ) + target_ids[3] = -100 + target_ids[4] = 2**40 + + unchecked = op(hidden, weight, target_ids) + assert torch.isnan(unchecked[3]) + assert torch.isnan(unchecked[4]) + assert torch.isfinite(torch.cat((unchecked[:3], unchecked[5:]))).all() + + with pytest.raises(ValueError, match="must be in"): + op(hidden, weight, target_ids, validate=True) + + +@requires_sm90_op +@pytest.mark.parametrize( + "target_dtype,invalid_value", + [ + (torch.uint8, 250), + (torch.int8, -1), + (torch.int16, 500), + (torch.int32, 97), + (torch.int64, 2**40), + ], +) +def test_sm90_supported_target_dtypes_preserve_valid_rows_and_mark_invalid( + target_dtype, + invalid_value, +): + op = op_module.BatchInvariantLinearLogpSM90Op() + hidden, weight, valid_targets, _ = _make_inputs( + 323, + num_tokens=8, + hidden_dim=128, + vocab_size=97, + bias=False, + ) + expected = _reference(hidden, weight, valid_targets, None) + target_ids = valid_targets.to(dtype=target_dtype) + target_ids[0] = invalid_value + + unchecked = op(hidden, weight, target_ids) + + assert torch.isnan(unchecked[0]) + assert torch.allclose(unchecked[1:], expected[1:], atol=2e-2, rtol=0.0) + with pytest.raises(ValueError, match="must be in"): + op(hidden, weight, target_ids, validate=True) + + +@requires_sm90_op +@pytest.mark.parametrize( + "target_dtype,target_value", + [ + (torch.uint8, 250), + (torch.int16, 499), + ], +) +def test_sm90_preserves_valid_high_narrow_integer_targets(target_dtype, target_value): + op = op_module.BatchInvariantLinearLogpSM90Op() + hidden, weight, target_ids, _ = _make_inputs( + 325, + num_tokens=4, + hidden_dim=128, + vocab_size=500, + bias=False, + ) + target_ids = target_ids.to(dtype=target_dtype) + target_ids[0] = target_value + + actual = op(hidden, weight, target_ids) + expected = _reference(hidden, weight, target_ids, None) + + assert torch.isfinite(actual).all() + assert torch.allclose(actual, expected, atol=2e-2, rtol=0.0) + + +@requires_sm90_op +def test_sm90_rejects_unsupported_bias_dtype(): + op = op_module.BatchInvariantLinearLogpSM90Op() + hidden, weight, target_ids, _ = _make_inputs( + 317, + num_tokens=8, + hidden_dim=128, + vocab_size=500, + bias=False, + ) + bias = torch.randn(500, device="cuda", dtype=torch.float64) + + with pytest.raises(TypeError, match="bias must have dtype"): + op(hidden, weight, target_ids, bias) + + +@requires_sm90_op +@pytest.mark.parametrize("target_dtype", [torch.float32, torch.bool]) +def test_sm90_raw_symbol_rejects_non_integer_targets(target_dtype): + from rl_engine.kernels.ops.base import _C + + hidden, weight, target_ids, _ = _make_inputs( + 318, + num_tokens=8, + hidden_dim=128, + vocab_size=500, + bias=False, + ) + invalid_dtype_targets = target_ids.to(dtype=target_dtype) + + with pytest.raises(RuntimeError, match="target must have dtype"): + _C.batch_invariant_linear_logp_sm90( + hidden, + weight, + invalid_dtype_targets, + None, + ) + + +@requires_sm90_op +def test_sm90_raw_symbol_rejects_long_target_metadata_before_conversion(): + from rl_engine.kernels.ops.base import _C + + hidden, weight, _target_ids, _ = _make_inputs( + 320, + num_tokens=8, + hidden_dim=128, + vocab_size=500, + bias=False, + ) + cpu_target = torch.zeros(8, dtype=torch.int64) + expanded_target = torch.zeros(1, device="cuda", dtype=torch.int64).expand(9) + + with pytest.raises(RuntimeError, match="must be CUDA tensors"): + _C.batch_invariant_linear_logp_sm90(hidden, weight, cpu_target, None) + with pytest.raises(RuntimeError, match="one id per token"): + _C.batch_invariant_linear_logp_sm90(hidden, weight, expanded_target, None) diff --git a/tests/test_kernel_registry.py b/tests/test_kernel_registry.py index 66a09297..91c661ef 100644 --- a/tests/test_kernel_registry.py +++ b/tests/test_kernel_registry.py @@ -2,6 +2,7 @@ # Copyright (c) 2026 RL-Kernel Contributors import pytest +import torch from rl_engine.kernels import registry as registry_module from rl_engine.kernels.registry import KernelRegistry, OpBackend @@ -89,11 +90,16 @@ def test_sm90_linear_ops_prioritize_cuda_when_extension_symbols_exist(monkeypatc class FakeExtension: fused_linear_logp_sm90 = object() + batch_invariant_linear_logp_sm90 = object() embedding_sm90_forward = object() lm_head_sm90_forward = object() monkeypatch.setattr(registry_module.device_ctx, "device_type", "cuda") - monkeypatch.setattr(registry_module.torch.cuda, "get_device_capability", lambda: (9, 0)) + monkeypatch.setattr( + registry_module.torch.cuda, + "get_device_capability", + lambda device=None: (9, 0), + ) monkeypatch.setattr(base_module, "_EXT_AVAILABLE", True) monkeypatch.setattr(base_module, "_C", FakeExtension()) @@ -101,17 +107,26 @@ class FakeExtension: assert registry._priority_map["cuda"]["embedding"][0] is OpBackend.CUDA_SM90_EMBEDDING assert registry._priority_map["cuda"]["lm_head"][0] is OpBackend.CUDA_SM90_LM_HEAD + assert registry._priority_map["cuda"]["batch_invariant_linear_logp"] == [ + OpBackend.CUDA_BATCH_INVARIANT_LINEAR_LOGP_SM90 + ] def test_sm90_linear_ops_do_not_prioritize_cuda_on_non_hopper(monkeypatch): from rl_engine.kernels.ops import base as base_module class FakeExtension: + fused_linear_logp_sm90 = object() + batch_invariant_linear_logp_sm90 = object() embedding_sm90_forward = object() lm_head_sm90_forward = object() monkeypatch.setattr(registry_module.device_ctx, "device_type", "cuda") - monkeypatch.setattr(registry_module.torch.cuda, "get_device_capability", lambda: (8, 0)) + monkeypatch.setattr( + registry_module.torch.cuda, + "get_device_capability", + lambda device=None: (8, 0), + ) monkeypatch.setattr(base_module, "_EXT_AVAILABLE", True) monkeypatch.setattr(base_module, "_C", FakeExtension()) @@ -119,3 +134,98 @@ class FakeExtension: assert OpBackend.CUDA_SM90_EMBEDDING not in registry._priority_map["cuda"]["embedding"] assert OpBackend.CUDA_SM90_LM_HEAD not in registry._priority_map["cuda"]["lm_head"] + assert ( + OpBackend.CUDA_FUSED_LINEAR_LOGP_SM90 not in registry._priority_map["cuda"]["linear_logp"] + ) + assert registry._priority_map["cuda"]["batch_invariant_linear_logp"] == [ + OpBackend.CUDA_BATCH_INVARIANT_LINEAR_LOGP_SM90 + ] + + with pytest.raises(RuntimeError, match="No functional backend"): + registry.get_op("batch_invariant_linear_logp", device="cuda") + + +def test_batch_invariant_linear_logp_dispatch_uses_requested_cuda_device(monkeypatch): + from rl_engine.kernels.ops import base as base_module + + class FakeExtension: + batch_invariant_linear_logp_sm90 = object() + + class FakeOp: + pass + + probed_devices = [] + + def capability(device=None): + resolved = torch.device("cuda:0" if device is None else device) + probed_devices.append(resolved) + return (9, 0) if resolved.index == 1 else (8, 0) + + monkeypatch.setattr(registry_module.device_ctx, "device_type", "cuda") + monkeypatch.setattr(registry_module.torch.cuda, "get_device_capability", capability) + monkeypatch.setattr(base_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(base_module, "_C", FakeExtension()) + monkeypatch.setattr(KernelRegistry, "_load_backend", lambda self, backend: FakeOp) + + registry = KernelRegistry() + probed_devices.clear() + + with pytest.raises(RuntimeError, match="No functional backend"): + registry.get_op("batch_invariant_linear_logp", device="cuda:0") + assert OpBackend.CUDA_BATCH_INVARIANT_LINEAR_LOGP_SM90.name not in registry._failed_backends + + hopper_op = registry.get_op("batch_invariant_linear_logp", device="cuda:1") + assert isinstance(hopper_op, FakeOp) + assert registry.get_op("batch_invariant_linear_logp", device="cuda:1") is hopper_op + + with pytest.raises(RuntimeError, match="No functional backend"): + registry.get_op("batch_invariant_linear_logp", device="cuda:0") + + assert torch.device("cuda:0") in probed_devices + assert torch.device("cuda:1") in probed_devices + + +def test_batch_invariant_linear_logp_explicit_device_survives_init_probe_failure(monkeypatch): + from rl_engine.kernels.ops import base as base_module + + class FakeExtension: + batch_invariant_linear_logp_sm90 = object() + + class FakeOp: + pass + + def capability(device=None): + if device is None: + raise RuntimeError("current device is unavailable") + return (9, 0) if torch.device(device).index == 1 else (8, 0) + + monkeypatch.setattr(registry_module.device_ctx, "device_type", "cuda") + monkeypatch.setattr(registry_module.torch.cuda, "get_device_capability", capability) + monkeypatch.setattr(base_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(base_module, "_C", FakeExtension()) + monkeypatch.setattr(KernelRegistry, "_load_backend", lambda self, backend: FakeOp) + + registry = KernelRegistry() + + assert isinstance( + registry.get_op("batch_invariant_linear_logp", device="cuda:1"), + FakeOp, + ) + + +def test_batch_invariant_linear_logp_hopper_fails_closed_without_symbol(monkeypatch): + from rl_engine.kernels.ops import base as base_module + + monkeypatch.setattr(registry_module.device_ctx, "device_type", "cuda") + monkeypatch.setattr( + registry_module.torch.cuda, + "get_device_capability", + lambda device=None: (9, 0), + ) + monkeypatch.setattr(base_module, "_EXT_AVAILABLE", True) + monkeypatch.setattr(base_module, "_C", object()) + + registry = KernelRegistry() + + with pytest.raises(RuntimeError, match="No functional backend"): + registry.get_op("batch_invariant_linear_logp", device="cuda:0") diff --git a/tests/test_linear_logp.py b/tests/test_linear_logp.py index 659127ce..eeaecc14 100644 --- a/tests/test_linear_logp.py +++ b/tests/test_linear_logp.py @@ -143,6 +143,41 @@ def _sm90_inputs(seed, *, bias=True, dtype=torch.bfloat16, lead=None): return hidden, weight, target, bias_t +def _select_sm90_fused_backward(monkeypatch): + """Disable Python alternatives so the public op reaches the fused C++ backward.""" + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_FUSED_BACKWARD", "1") + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_FUSED_BWD_PRECISION", "auto") + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_FUSED_BWD_BF16_DLOGITS", "1") + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_SAVE_PROBS_BF16", "0") + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL", "0") + + +def _select_sm90_streaming_backward(monkeypatch, *, logits_mode, grad_weight_mode=""): + """Pin the public op to one streaming CUDA-backward implementation.""" + _select_sm90_fused_backward(monkeypatch) + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_STREAMING_BWD", "1") + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_STREAMING_BWD_MODE", "tiled") + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_STREAMING_BWD_VOCAB_TILE", "128") + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_STREAMING_BWD_LOGITS", logits_mode) + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_STREAMING_BWD_GW", grad_weight_mode) + monkeypatch.setenv( + "RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD", + "1" if grad_weight_mode == "fused_tile" else "0", + ) + + +def _copy_with_storage_offset(tensor, offset): + storage = torch.empty( + tensor.numel() + offset, + device=tensor.device, + dtype=tensor.dtype, + ) + result = storage[offset:].view_as(tensor) + result.copy_(tensor) + assert result.is_contiguous() + return result + + # Deliberately non-multiples of the kernel block sizes (32 / 64 / 64). _N = 40 _D = 80 @@ -221,6 +256,57 @@ def test_native_matches_manual_reference(): assert torch.allclose(out, ref, atol=1e-5) +def test_native_forward_fp32_avoids_bf16_logit_rounding(): + generator = torch.Generator().manual_seed(2027) + hidden = torch.randn(2, 3, 32, generator=generator).bfloat16() + weight = torch.randn(67, 32, generator=generator).bfloat16() + bias = torch.randn(67, generator=generator).bfloat16() + target = torch.randint(0, 67, (2, 3), generator=generator) + + actual = NativeLinearLogpOp().forward_fp32( + hidden, + weight, + target, + bias, + tp_group=None, + vocab_start_index=0, + global_vocab_size=None, + ) + logits = torch.nn.functional.linear(hidden.float(), weight.float(), bias.float()) + expected = ( + torch.log_softmax(logits, dim=-1) + .gather( + -1, + target.unsqueeze(-1), + ) + .squeeze(-1) + ) + + assert actual.dtype == torch.float32 + assert torch.allclose(actual, expected, atol=1e-6, rtol=0.0) + + +def test_native_forward_fp32_ignores_ambient_autocast_and_restores_tf32(): + generator = torch.Generator().manual_seed(2028) + hidden = torch.randn(2, 3, 32, generator=generator).bfloat16() + weight = torch.randn(67, 32, generator=generator).bfloat16() + target = torch.randint(0, 67, (2, 3), generator=generator) + op = NativeLinearLogpOp() + + previous_tf32 = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = True + try: + expected = op.forward_fp32(hidden, weight, target) + with torch.autocast(device_type="cpu", dtype=torch.bfloat16): + actual = op.forward_fp32(hidden, weight, target) + + assert actual.dtype == torch.float32 + assert torch.equal(actual, expected) + assert torch.backends.cuda.matmul.allow_tf32 is True + finally: + torch.backends.cuda.matmul.allow_tf32 = previous_tf32 + + def test_linear_logp_handoff_matches_masked_reference_across_layouts(): torch.manual_seed(2026) op = NativeLinearLogpOp() @@ -399,7 +485,10 @@ def test_tied_embedding_lm_head_shared_gradient_is_layout_invariant(): model.zero_grad(set_to_none=True) hidden = model(input_ids) logps = op( - hidden, model.lm_head.weight, _safe_token_ids(masked_target, mask), model.lm_head.bias + hidden, + model.lm_head.weight, + _safe_token_ids(masked_target, mask), + model.lm_head.bias, ) logps = logps.masked_fill(~mask, 0.0) logits = torch.nn.functional.linear(hidden.float(), model.lm_head.weight.float(), None) @@ -829,6 +918,144 @@ def run(logits_mode=None, grad_weight_mode=None): assert torch.allclose(fused_tile_w, tiled_w, atol=8e-2, rtol=8e-2) +@requires_sm90 +@pytest.mark.parametrize("logits_mode", ["sm90_mma", "sm90_mma_dz"]) +@pytest.mark.parametrize("misaligned_input", ["hidden", "weight"]) +def test_sm90_streaming_mma_backward_realigns_tma_inputs( + monkeypatch, misaligned_input, logits_mode +): + from rl_engine.kernels.ops.base import _C + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + + assert hasattr(_C, "fused_linear_logp_sm90_backward") + op = FusedLinearLogpSM90Op() + hidden, weight, target, _ = _sm90_inputs(138, bias=False) + grad_out = torch.randn(_SM90_N, device="cuda") + + misaligned_hidden = ( + _copy_with_storage_offset(hidden, 1) if misaligned_input == "hidden" else hidden + ) + misaligned_weight = ( + _copy_with_storage_offset(weight, 1) if misaligned_input == "weight" else weight + ) + misaligned_tensor = misaligned_hidden if misaligned_input == "hidden" else misaligned_weight + assert misaligned_tensor.data_ptr() % 16 != 0 + misaligned_weight = misaligned_weight.detach().requires_grad_(True) + + aligned_hidden = misaligned_hidden.clone() + aligned_weight = misaligned_weight.detach().clone().requires_grad_(True) + assert aligned_hidden.data_ptr() % 16 == 0 + assert aligned_weight.data_ptr() % 16 == 0 + + _select_sm90_streaming_backward(monkeypatch, logits_mode=logits_mode) + + misaligned_out = op(misaligned_hidden, misaligned_weight, target, None) + misaligned_out.backward(grad_out) + aligned_out = op(aligned_hidden, aligned_weight, target, None) + aligned_out.backward(grad_out) + + assert torch.equal(misaligned_out, aligned_out) + assert misaligned_weight.grad is not None + assert torch.isfinite(misaligned_weight.grad).all() + assert torch.allclose(misaligned_weight.grad, aligned_weight.grad, atol=8e-2, rtol=8e-2) + + +@requires_sm90 +@pytest.mark.parametrize("grad_weight_mode", ["wmma", "fused_tile"]) +def test_sm90_streaming_wmma_backward_realigns_hidden(monkeypatch, grad_weight_mode): + from rl_engine.kernels.ops.base import _C + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + + assert hasattr(_C, "fused_linear_logp_sm90_backward") + gen = torch.Generator(device="cuda").manual_seed(139) + hidden = torch.randn(_SM90_N, _SM90_D, device="cuda", dtype=torch.bfloat16, generator=gen) + weight = torch.randn(512, _SM90_D, device="cuda", dtype=torch.bfloat16, generator=gen) + target = torch.randint(0, 512, (_SM90_N,), device="cuda", generator=gen) + grad_out = torch.randn(_SM90_N, device="cuda", generator=gen) + + misaligned_hidden = _copy_with_storage_offset(hidden, 8) + assert misaligned_hidden.data_ptr() % 16 == 0 + assert misaligned_hidden.data_ptr() % 32 != 0 + + misaligned_weight = weight.detach().clone().requires_grad_(True) + aligned_hidden = misaligned_hidden.clone() + aligned_weight = weight.detach().clone().requires_grad_(True) + assert aligned_hidden.data_ptr() % 32 == 0 + _select_sm90_streaming_backward( + monkeypatch, + logits_mode="", + grad_weight_mode=grad_weight_mode, + ) + + op = FusedLinearLogpSM90Op() + misaligned_out = op(misaligned_hidden, misaligned_weight, target, None) + misaligned_out.backward(grad_out) + aligned_out = op(aligned_hidden, aligned_weight, target, None) + aligned_out.backward(grad_out) + + assert torch.equal(misaligned_out, aligned_out) + assert misaligned_weight.grad is not None + assert torch.isfinite(misaligned_weight.grad).all() + assert torch.allclose(misaligned_weight.grad, aligned_weight.grad, atol=8e-2, rtol=8e-2) + + +@requires_sm90 +@pytest.mark.parametrize( + ("full_mode", "misaligned_input", "offset"), + [ + ("tile_cublas", "hidden", 1), + ("tile_cublas", "weight", 1), + ("from_logits", "hidden", 8), + ], +) +def test_sm90_full_backward_realigns_inputs(monkeypatch, full_mode, misaligned_input, offset): + from rl_engine.kernels.ops.base import _C + from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op + + assert hasattr(_C, "fused_linear_logp_sm90_backward") + gen = torch.Generator(device="cuda").manual_seed(140) + hidden = torch.randn(_SM90_N, _SM90_D, device="cuda", dtype=torch.bfloat16, generator=gen) + weight = torch.randn(512, _SM90_D, device="cuda", dtype=torch.bfloat16, generator=gen) + target = torch.randint(0, 512, (_SM90_N,), device="cuda", generator=gen) + grad_out = torch.randn(_SM90_N, device="cuda", generator=gen) + + misaligned_hidden = ( + _copy_with_storage_offset(hidden, offset) if misaligned_input == "hidden" else hidden + ) + misaligned_weight = ( + _copy_with_storage_offset(weight, offset) if misaligned_input == "weight" else weight + ) + misaligned_tensor = misaligned_hidden if misaligned_input == "hidden" else misaligned_weight + required_alignment = 32 if full_mode == "from_logits" else 16 + assert misaligned_tensor.data_ptr() % required_alignment != 0 + + misaligned_weight = misaligned_weight.detach().requires_grad_(True) + aligned_hidden = misaligned_hidden.clone() + aligned_weight = misaligned_weight.detach().clone().requires_grad_(True) + aligned_tensor = aligned_hidden if misaligned_input == "hidden" else aligned_weight + assert aligned_tensor.data_ptr() % required_alignment == 0 + _select_sm90_fused_backward(monkeypatch) + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_STREAMING_BWD", "0") + op = FusedLinearLogpSM90Op() + + def run(hidden_input, weight_input): + # Keep the regular autograd wrapper for forward, then select the raw + # full-backward branch only after it has saved the original views. + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL", "0") + output = op(hidden_input, weight_input, target, None) + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_FUSED_TILE_BWD_FULL", full_mode) + output.backward(grad_out) + return output.detach(), weight_input.grad + + misaligned_out, misaligned_grad = run(misaligned_hidden, misaligned_weight) + aligned_out, aligned_grad = run(aligned_hidden, aligned_weight) + + assert torch.equal(misaligned_out, aligned_out) + assert misaligned_grad is not None + assert torch.isfinite(misaligned_grad).all() + assert torch.allclose(misaligned_grad, aligned_grad, atol=8e-2, rtol=8e-2) + + @requires_sm90 def test_sm90_preserves_leading_shape(): from rl_engine.kernels.ops.cuda.loss.linear_logp import FusedLinearLogpSM90Op @@ -917,6 +1144,119 @@ def test_sm90_rejects_out_of_range_target(): assert out.shape == target.shape and torch.isfinite(out).all() +@requires_sm90 +@pytest.mark.parametrize( + "vocab_start_index,error", + [ + (-1, "vocab_start_index"), + (2**32, "vocab_start_index"), + (2**31 - 250, "local vocabulary range"), + ], +) +def test_sm90_raw_tp_symbol_rejects_unrepresentable_vocab_ranges( + vocab_start_index, + error, +): + from rl_engine.kernels.ops.base import _C + + hidden, weight, target, _ = _sm90_inputs(319, bias=False) + with pytest.raises(RuntimeError, match=error): + _C.fused_linear_logp_sm90_global_target( + hidden[:8], + weight, + target[:8], + None, + vocab_start_index, + ) + + +@requires_sm90 +def test_sm90_raw_tp_symbol_does_not_wrap_extreme_int64_target(): + from rl_engine.kernels.ops.base import _C + + hidden, weight, target, _ = _sm90_inputs(321, bias=False) + target[0] = 2**40 + + local_target_logit, local_lse = _C.fused_linear_logp_sm90_global_target( + hidden[:8], + weight, + target[:8], + None, + 0, + ) + + assert local_target_logit[0].item() == 0.0 + assert torch.isfinite(local_lse).all() + + +@requires_sm90 +def test_sm90_raw_tp_extreme_target_cannot_alias_legal_int32_max_class(): + from rl_engine.kernels.ops.base import _C + + int32_max = torch.iinfo(torch.int32).max + hidden = torch.ones(1, 32, device="cuda", dtype=torch.bfloat16) + weight = torch.ones(1, 32, device="cuda", dtype=torch.bfloat16) + invalid_target = torch.tensor([2**40], device="cuda", dtype=torch.int64) + valid_target = torch.tensor([int32_max], device="cuda", dtype=torch.int64) + + invalid_logit, invalid_lse = _C.fused_linear_logp_sm90_global_target( + hidden, + weight, + invalid_target, + None, + int32_max, + ) + valid_logit, valid_lse = _C.fused_linear_logp_sm90_global_target( + hidden, + weight, + valid_target, + None, + int32_max, + ) + + assert invalid_logit.item() == 0.0 + assert valid_logit.item() == 32.0 + assert torch.isfinite(invalid_lse).all() and torch.isfinite(valid_lse).all() + + +@requires_sm90 +def test_sm90_backward_extreme_target_matches_out_of_shard_sentinel(monkeypatch): + from rl_engine.kernels.ops.base import _C + + monkeypatch.setenv("RL_KERNEL_LINEAR_LOGP_FUSED_BWD_PRECISION", "fp32") + hidden, weight, _target, _ = _sm90_inputs(322, bias=False) + grad_logp = torch.randn(_SM90_N, device="cuda") + logits = hidden.float().matmul(weight.float().t()) + lse = torch.logsumexp(logits, dim=-1) + + def grad_weight_for(target_value): + target = torch.full( + (_SM90_N,), + target_value, + device="cuda", + dtype=torch.int64, + ) + _grad_hidden, grad_weight, _grad_bias = _C.fused_linear_logp_sm90_backward( + grad_logp, + hidden, + weight, + target, + lse, + None, + 0, + False, + True, + False, + True, + ) + return grad_weight + + sentinel_grad = grad_weight_for(-1) + extreme_grad = grad_weight_for(2**40) + + assert torch.equal(extreme_grad, sentinel_grad) + + def test_sm90_tp_metadata_prefers_sm90_tp_helper(monkeypatch): from rl_engine.kernels.ops.cuda.loss import linear_logp as cuda_linear_logp @@ -987,7 +1327,14 @@ def test_sm90_tp_metadata_prefers_save_probs_for_output_only(monkeypatch): monkeypatch.setattr(cuda_linear_logp, "_sm90_save_probs_bf16_tp_available", lambda: True) def fake_save_probs_apply(hidden_arg, weight_arg, target_arg, vocab_start, global_vocab, group): - calls["save_probs"] = (hidden_arg, weight_arg, target_arg, vocab_start, global_vocab, group) + calls["save_probs"] = ( + hidden_arg, + weight_arg, + target_arg, + vocab_start, + global_vocab, + group, + ) return sentinel def forbidden_sm90_tp(*args, **kwargs): @@ -1125,7 +1472,11 @@ class FakeC: @staticmethod def fused_linear_logp_sm90_backward(*args): calls["backward"] = args - return torch.ones_like(hidden), torch.ones_like(weight), torch.ones_like(bias) + return ( + torch.ones_like(hidden), + torch.ones_like(weight), + torch.ones_like(bias), + ) monkeypatch.setattr(cuda_linear_logp, "_EXT_AVAILABLE", True) monkeypatch.setattr(cuda_linear_logp, "_C", FakeC) diff --git a/tests/test_operator_inputs.py b/tests/test_operator_inputs.py index 4f742734..4bd21271 100644 --- a/tests/test_operator_inputs.py +++ b/tests/test_operator_inputs.py @@ -39,6 +39,7 @@ def _args(**overrides): "attention", "logp", "linear_logp", + "batch_invariant_linear_logp", "batch_invariant_logp", "rope", "silu",