diff --git a/csrc/cuda/latent_pack_unpack.cu b/csrc/cuda/latent_pack_unpack.cu new file mode 100644 index 00000000..f95cb06a --- /dev/null +++ b/csrc/cuda/latent_pack_unpack.cu @@ -0,0 +1,127 @@ +// SPDX-License-Identifier: Apache-2.0 +// Copyright (c) 2026 RL-Kernel Contributors +#include +#include +#include +#include + +#ifndef LATENT_TILE_Y +#define LATENT_TILE_Y 32 +#endif +#ifndef LATENT_THREADS_Y +#define LATENT_THREADS_Y 8 +#endif + +namespace { +constexpr int tile_y = LATENT_TILE_Y; +constexpr int tile_x = tile_y; +constexpr int threads_x = tile_x; +constexpr int threads_y = LATENT_THREADS_Y; +static_assert(tile_y % threads_y == 0 && threads_x * threads_y <= 1024); +template struct Pair { Bits lo, hi; }; + +template +__global__ void latent_permute(const Pair* __restrict__ input, Pair* __restrict__ output, int64_t num_channels, int64_t latent_height, int64_t latent_width) { + __shared__ Pair shared_pairs[tile_y][tile_x + 1]; + const int thread_x = threadIdx.x, thread_y = threadIdx.y; + const int64_t token_cols = latent_width / 2, pairs_per_token = num_channels * 2; + const int64_t batch_pair_offset = static_cast(blockIdx.z) * num_channels * latent_height * token_cols; + const int64_t token_col_start = static_cast(blockIdx.x) * tile_x; + const int64_t token_row = blockIdx.y; + for (int64_t token_pair_start = 0; token_pair_start < pairs_per_token; token_pair_start += tile_y) { + #pragma unroll + for (int thread_y_offset = 0; thread_y_offset < tile_y; thread_y_offset += threads_y) { + if constexpr (Unpack) { + const int64_t token_col = token_col_start + thread_y + thread_y_offset; + const int64_t pair_in_token = token_pair_start + thread_x; + if (token_col < token_cols && pair_in_token < pairs_per_token) { + const int64_t input_offset = batch_pair_offset + (token_row * token_cols + token_col) * pairs_per_token + pair_in_token; + shared_pairs[thread_y + thread_y_offset][thread_x] = input[input_offset]; + } + } else { + const int64_t token_col = token_col_start + thread_x; + const int64_t pair_in_token = token_pair_start + thread_y + thread_y_offset; + if (token_col < token_cols && pair_in_token < pairs_per_token) { + const int64_t channel = pair_in_token / 2; + const int64_t spatial_row = token_row * 2 + pair_in_token % 2; + const int64_t input_offset = batch_pair_offset + (channel * latent_height + spatial_row) * token_cols + token_col; + shared_pairs[thread_y + thread_y_offset][thread_x] = input[input_offset]; + } + } + } + __syncthreads(); + #pragma unroll + for (int thread_y_offset = 0; thread_y_offset < tile_x; thread_y_offset += threads_y) { + if constexpr (Unpack) { + const int64_t token_col = token_col_start + thread_x; + const int64_t pair_in_token = token_pair_start + thread_y + thread_y_offset; + if (token_col < token_cols && pair_in_token < pairs_per_token) { + const int64_t channel = pair_in_token / 2; + const int64_t spatial_row = token_row * 2 + pair_in_token % 2; + const int64_t output_offset = batch_pair_offset + (channel * latent_height + spatial_row) * token_cols + token_col; + output[output_offset] = shared_pairs[thread_x][thread_y + thread_y_offset]; + } + } else { + const int64_t token_col = token_col_start + thread_y + thread_y_offset; + const int64_t pair_in_token = token_pair_start + thread_x; + if (token_col < token_cols && pair_in_token < pairs_per_token) { + const int64_t output_offset = batch_pair_offset + (token_row * token_cols + token_col) * pairs_per_token + pair_in_token; + output[output_offset] = shared_pairs[thread_x][thread_y + thread_y_offset]; + } + } + } + if (token_pair_start + tile_y < pairs_per_token) { + __syncthreads(); + } + } +} + +template +void launch(const torch::Tensor& x, torch::Tensor& y, int64_t B, int64_t C, int64_t H, int64_t W) { + const dim3 grid((W / 2 + tile_x - 1) / tile_x, H / 2, B); + const dim3 block(threads_x, threads_y); + latent_permute<<>>(reinterpret_cast*>(x.data_ptr()), reinterpret_cast*>(y.data_ptr()), C, H, W); +} +} + +torch::Tensor latent_pack_unpack_cuda(torch::Tensor x, int64_t B, int64_t C, int64_t H, int64_t W, bool unpack) { + TORCH_CHECK(x.is_cuda() && x.is_contiguous(), "expected contiguous CUDA input"); + TORCH_CHECK(x.scalar_type() == at::kFloat || x.scalar_type() == at::kHalf || x.scalar_type() == at::kBFloat16, "expected fp32, fp16, or bf16"); + TORCH_CHECK(B >= 0 && B <= 65535 && C > 0 && H > 0 && W > 0 && H % 2 == 0 && W % 2 == 0, "invalid latent dimensions"); + TORCH_CHECK(C <= INT32_MAX / 4 && H <= INT32_MAX && W <= INT32_MAX, "latent dimensions exceed indexing limits"); + const int64_t P = (H / 2) * (W / 2); + TORCH_CHECK(P <= INT32_MAX && H / 2 <= 65535, "latent dimensions exceed launch limits"); + TORCH_CHECK(x.numel() / (C * 4) == B * P && x.numel() % (C * 4) == 0, "input element count does not match latent dimensions"); + if (unpack) { + TORCH_CHECK(x.dim() == 3 && x.size(0) == B && x.size(1) == P && x.size(2) == C * 4, "expected packed [B,P,4*C] input"); + } else { + const bool nchw = x.dim() == 4 && x.size(1) == C; + const bool singleton = x.dim() == 5 && ((x.size(1) == C && x.size(2) == 1) || (x.size(1) == 1 && x.size(2) == C)); + TORCH_CHECK((nchw || singleton) && x.size(0) == B && x.size(-2) == H && x.size(-1) == W, "expected spatial latent input"); + } + const c10::cuda::CUDAGuard guard(x.device()); + torch::Tensor y; + if (unpack) { + y = torch::empty({B, C, 1, H, W}, x.options()); + } else { + y = torch::empty({B, P, C * 4}, x.options()); + } + if (B == 0) { + return y; + } + if (x.scalar_type() == at::kFloat) { + if (unpack) { + launch(x, y, B, C, H, W); + } else { + launch(x, y, B, C, H, W); + } + } else { + if (unpack) { + launch(x, y, B, C, H, W); + } else { + launch(x, y, B, C, H, W); + } + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return y; +} diff --git a/csrc/ops.cpp b/csrc/ops.cpp index 026ea23d..a9d1ed04 100644 --- a/csrc/ops.cpp +++ b/csrc/ops.cpp @@ -126,6 +126,8 @@ torch::Tensor hip_deterministic_logp_backward( #endif #if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) +torch::Tensor latent_pack_unpack_cuda(torch::Tensor x, int64_t B, int64_t C, + int64_t H, int64_t W, bool unpack); // Single-node TP=8 deterministic CUDA IPC collectives. ROCm uses the // rank-ordered RCCL transport in rl_engine.distributed.collectives. std::tuple, int64_t> deterministic_collective_ipc_meta( @@ -474,6 +476,10 @@ at::Tensor prefix_shared_attention( // PyBind11 Module Registration PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { +#if !defined(USE_ROCM) && !defined(KERNEL_ALIGN_WITH_ROCM) + m.def("latent_pack_unpack", &latent_pack_unpack_cuda, + "Qwen-Image bitwise latent pack/unpack CUDA"); +#endif m.doc() = "RL-Kernel High-Performance Operator Extension Library"; m.def("fused_logp", &fused_logp_forward, "Fused logp forward fallback"); diff --git a/rl_engine/kernels/ops/cuda/packing/__init__.py b/rl_engine/kernels/ops/cuda/packing/__init__.py new file mode 100644 index 00000000..98813136 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/packing/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/rl_engine/kernels/ops/cuda/packing/latent_pack_unpack.py b/rl_engine/kernels/ops/cuda/packing/latent_pack_unpack.py new file mode 100644 index 00000000..a44aeca2 --- /dev/null +++ b/rl_engine/kernels/ops/cuda/packing/latent_pack_unpack.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +import torch + +from rl_engine.kernels.ops import base +from rl_engine.kernels.ops.latent_layout import LatentKernelOp + + +class CudaLatentPackOp(LatentKernelOp): + def __init__(self): + if torch.version.hip is not None or not torch.cuda.is_available(): + raise RuntimeError("CUDA latent permutation requires an NVIDIA GPU") + if not base._EXT_AVAILABLE or not hasattr(base._C, "latent_pack_unpack"): + raise RuntimeError("Rebuild rl_engine._C with latent_pack_unpack.cu") + self.kernel = base._C.latent_pack_unpack + + +class CudaLatentUnpackOp(CudaLatentPackOp): + unpack = True diff --git a/rl_engine/kernels/ops/latent_layout.py b/rl_engine/kernels/ops/latent_layout.py new file mode 100644 index 00000000..a88db0d1 --- /dev/null +++ b/rl_engine/kernels/ops/latent_layout.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Shared Qwen-Image shape contract and inverse-permutation autograd.""" + +import operator + +import torch + + +def dimensions(x, batch_size, num_channels_latents, height, width, *, unpack=False): + b, c, h, w = map(operator.index, (batch_size, num_channels_latents, height, width)) + if x.dtype not in (torch.float32, torch.float16, torch.bfloat16): + raise TypeError("latent input must be fp32, fp16, or bf16") + if not x.is_contiguous(): + raise ValueError("latent input must be contiguous") + if not (0 <= b <= 65535 and c > 0 and h > 0 and w > 0 and h % 2 == w % 2 == 0): + raise ValueError("expected B in [0,65535], positive C and positive even H,W") + if c > (2**31 - 1) // 4 or max(h, w, h // 2 * (w // 2)) > 2**31 - 1: + raise ValueError("latent dimensions exceed indexing limits") + if h // 2 > 65535: + raise ValueError("latent height exceeds launch limits") + shapes = ( + ((b, h // 2 * (w // 2), 4 * c),) + if unpack + else ((b, c, h, w), (b, c, 1, h, w), (b, 1, c, h, w)) + ) + if tuple(x.shape) not in shapes: + raise ValueError(f"latent shape {tuple(x.shape)} does not match {shapes}") + return b, c, h, w + + +def unpack_dimensions(x, height, width, vae_scale_factor): + scale = operator.index(vae_scale_factor) + if scale <= 0 or x.ndim != 3 or x.shape[-1] % 4: + raise ValueError("expected positive VAE scale and packed [B,P,4*C] input") + h, w = (2 * (int(size) // (2 * scale)) for size in (height, width)) + return dimensions(x, x.shape[0], x.shape[-1] // 4, h, w, unpack=True) + + +class _LatentPermutation(torch.autograd.Function): + @staticmethod + def forward(ctx, x, dims, unpack, kernel): + ctx.dims, ctx.unpack, ctx.kernel, ctx.shape = dims, unpack, kernel, x.shape + return kernel(x, *dims, unpack) + + @staticmethod + def backward(ctx, grad): + # Re-enter Function.apply so higher-order gradients are also permutations. + result = _LatentPermutation.apply(grad.contiguous(), ctx.dims, not ctx.unpack, ctx.kernel) + return result.reshape(ctx.shape), None, None, None + + +class LatentKernelOp: + unpack = False + op_class = "permutation" + + def __call__(self, x, *args, **kwargs): + return self.forward(x, *args, **kwargs) + + def forward(self, x, *args, **kwargs): + if x.device.type != "cuda": + raise ValueError("GPU latent backend requires a CUDA tensor") + dims = ( + unpack_dimensions(x, *args, **kwargs) if self.unpack else dimensions(x, *args, **kwargs) + ) + return _LatentPermutation.apply(x, dims, self.unpack, self.kernel) diff --git a/rl_engine/kernels/ops/pytorch/packing/latent_pack_unpack.py b/rl_engine/kernels/ops/pytorch/packing/latent_pack_unpack.py new file mode 100644 index 00000000..6286cc1d --- /dev/null +++ b/rl_engine/kernels/ops/pytorch/packing/latent_pack_unpack.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Qwen-Image reference: reshape/permute only, with no dtype conversion.""" + +import torch + +from rl_engine.kernels.ops.latent_layout import dimensions, unpack_dimensions + + +class NativeLatentPackOp(torch.nn.Module): + op_class = "permutation" + + def forward(self, x, batch_size, num_channels_latents, height, width): + b, c, h, w = dimensions(x, batch_size, num_channels_latents, height, width) + return ( + x.view(b, c, h // 2, 2, w // 2, 2) + .permute(0, 2, 4, 1, 3, 5) + .reshape(b, h // 2 * (w // 2), c * 4) + ) + + +class NativeLatentUnpackOp(torch.nn.Module): + op_class = "permutation" + + def forward(self, x, height, width, vae_scale_factor): + b, c, h, w = unpack_dimensions(x, height, width, vae_scale_factor) + return x.view(b, h // 2, w // 2, c, 2, 2).permute(0, 3, 1, 4, 2, 5).reshape(b, c, 1, h, w) # 1 is frame, just to adapt to the interface diff --git a/rl_engine/kernels/ops/triton/packing/__init__.py b/rl_engine/kernels/ops/triton/packing/__init__.py new file mode 100644 index 00000000..98813136 --- /dev/null +++ b/rl_engine/kernels/ops/triton/packing/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/rl_engine/kernels/ops/triton/packing/latent_pack_unpack.py b/rl_engine/kernels/ops/triton/packing/latent_pack_unpack.py new file mode 100644 index 00000000..e34312ad --- /dev/null +++ b/rl_engine/kernels/ops/triton/packing/latent_pack_unpack.py @@ -0,0 +1,53 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +import torch +import triton +import triton.language as tl + +from rl_engine.kernels.ops.latent_layout import LatentKernelOp + + +@triton.jit +def _permute( + X, + Y, + C: tl.constexpr, + H: tl.constexpr, + W: tl.constexpr, + UNPACK: tl.constexpr, + BLOCK: tl.constexpr, +): + # Integer pointers preserve every input bit, including signaling NaNs. + i = tl.program_id(0).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + n: tl.constexpr = C * H * W + channel = i // (H * W) + row, col = i // W % H, i % W + packed = ((row // 2 * (W // 2) + col // 2) * C + channel) * 4 + row % 2 * 2 + col % 2 # pack * channel * 4pos + base = tl.program_id(1).to(tl.int64) * n + src, dst = (packed, i) if UNPACK else (i, packed) + value = tl.load(X + base + src, i < n, other=0) + tl.store(Y + base + dst, value, i < n) + + +def _launch(x, b, c, h, w, unpack): + shape = (b, c, 1, h, w) if unpack else (b, h // 2 * (w // 2), c * 4) + y = torch.empty(shape, dtype=x.dtype, device=x.device) + if b: + bits = torch.int32 if x.element_size() == 4 else torch.int16 + with torch.cuda.device(x.device): + _permute[(triton.cdiv(c * h * w, 256), b)]( + x.view(bits), y.view(bits), c, h, w, unpack, 256 + ) + return y + + +class TritonLatentPackOp(LatentKernelOp): + kernel = staticmethod(_launch) + + def __init__(self): + if not torch.cuda.is_available(): + raise RuntimeError("Triton latent permutation requires a GPU") + + +class TritonLatentUnpackOp(TritonLatentPackOp): + unpack = True diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index d272fdb9..88c19b36 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -106,6 +106,20 @@ class OpBackend(Enum, metaclass=_KernelEnumMeta): # Variable-length packing (pack-and-pad), [B,S,...] -> [Total_Active,...] PYTORCH_PACK = "rl_engine.kernels.ops.pytorch.packing.pack.NativePackOp" + PYTORCH_LATENT_PACK = ( + "rl_engine.kernels.ops.pytorch.packing.latent_pack_unpack.NativeLatentPackOp" + ) + PYTORCH_LATENT_UNPACK = ( + "rl_engine.kernels.ops.pytorch.packing.latent_pack_unpack.NativeLatentUnpackOp" + ) + CUDA_LATENT_PACK = "rl_engine.kernels.ops.cuda.packing.latent_pack_unpack.CudaLatentPackOp" + CUDA_LATENT_UNPACK = "rl_engine.kernels.ops.cuda.packing.latent_pack_unpack.CudaLatentUnpackOp" + TRITON_LATENT_PACK = ( + "rl_engine.kernels.ops.triton.packing.latent_pack_unpack.TritonLatentPackOp" + ) + TRITON_LATENT_UNPACK = ( + "rl_engine.kernels.ops.triton.packing.latent_pack_unpack.TritonLatentUnpackOp" + ) # Batch-invariant deterministic GEMM (WS1 #146) CUDA_DET_GEMM = "rl_engine.kernels.ops.cuda.matmul.det_gemm.DetGemmOp" TRITON_DET_GEMM = "rl_engine.kernels.ops.triton.matmul.det_gemm.TritonDetGemmOp" @@ -709,6 +723,14 @@ def __init__(self): ], }, } + for platform, ops in self._priority_map.items(): + for direction in ("PACK", "UNPACK"): + backends = [OpBackend[f"PYTORCH_LATENT_{direction}"]] + if platform in ("cuda", "rocm"): + backends.insert(0, OpBackend[f"TRITON_LATENT_{direction}"]) + if platform == "cuda": + backends.insert(0, OpBackend[f"CUDA_LATENT_{direction}"]) + ops[f"latent_{direction.lower()}"] = backends # Preserve the former CPU fallback behavior for every operator on NPU, # then override only the operators with an Ascend-specific backend. self._priority_map["npu"] = { diff --git a/setup.py b/setup.py index aef36c9a..b4ac3bd0 100644 --- a/setup.py +++ b/setup.py @@ -153,6 +153,7 @@ def get_extensions(): ] ) else: + cuda_sources.append("csrc/cuda/latent_pack_unpack.cu") # CUDA IPC and the fixed-tree collective implementation are not # part of the ROCm extension. cuda_sources.append("csrc/cuda/distributed/deterministic_collective.cu") diff --git a/tests/test_latent_pack_unpack.py b/tests/test_latent_pack_unpack.py new file mode 100644 index 00000000..da934aac --- /dev/null +++ b/tests/test_latent_pack_unpack.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors +"""Exact-value and raw-bit acceptance for Qwen-Image layout permutations.""" + +from importlib import import_module +from pathlib import Path + +import pytest +import torch + +from rl_engine.kernels.ops import base +from rl_engine.kernels.ops.pytorch.packing.latent_pack_unpack import ( + NativeLatentPackOp, + NativeLatentUnpackOp, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + +DTYPES = (torch.float32, torch.float16, torch.bfloat16) +SIZES = ((1024, 1024), (1328, 1328), (1664, 928)) + + +def assert_bits(actual, expected): + assert actual.shape == expected.shape and actual.dtype == expected.dtype + assert torch.equal( + actual.detach().contiguous().cpu().view(torch.uint8), + expected.detach().contiguous().cpu().view(torch.uint8), + ) + + +@pytest.fixture(params=("PYTORCH_CPU", "PYTORCH", "CUDA", "TRITON")) +def ops(request): + backend = request.param.split("_")[0] + device = "cpu" if request.param == "PYTORCH_CPU" else "cuda" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("GPU unavailable") + if backend == "CUDA" and (torch.version.hip or not hasattr(base._C, "latent_pack_unpack")): + pytest.skip("CUDA latent extension unavailable") + if backend == "TRITON": + pytest.importorskip("triton") + return tuple( + getattr(import_module(module), cls)() + for direction in ("PACK", "UNPACK") + for module, cls in [OpBackend[f"{backend}_LATENT_{direction}"].value.rsplit(".", 1)] + ) + (device,) + + +# Check exact forward results, gradients, round trips, and batch invariance across backends. +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize("image_size", SIZES) +@pytest.mark.parametrize("batch", (1, 2, 4, 8)) +def test_forward_backward_exact(ops, dtype, image_size, batch): + pack, unpack, device = ops + ih, iw = image_size + h, w, c = ih // 8, iw // 8, 16 + generator = torch.Generator().manual_seed(386) + cpu = torch.randn(batch, c, 1, h, w, generator=generator).to(dtype).requires_grad_() + x = cpu.detach().to(device).requires_grad_() + expected = NativeLatentPackOp()(cpu, batch, c, h, w) + packed = pack(x, batch, c, h, w) + torch.testing.assert_close(packed.cpu(), expected, rtol=0, atol=0) + assert_bits(packed, expected) + restored = unpack(packed, ih, iw, 8) + assert_bits(restored, cpu) + # Independent unpack input prevents a paired pack/unpack mistake cancelling out. + tokens = torch.randn(expected.shape, generator=generator).to(dtype).requires_grad_() + gpu_tokens = tokens.detach().to(device).requires_grad_() + spatial = unpack(gpu_tokens, ih, iw, 8) + ref_spatial = NativeLatentUnpackOp()(tokens, ih, iw, 8) + assert_bits(spatial, ref_spatial) + assert_bits(pack(spatial, batch, c, h, w), tokens) + # Non-contiguous gradients exercise the explicit backward materialization path. + dp = torch.randn(batch, 4 * c, expected.shape[1], generator=generator).to(dtype) + dp = dp.transpose(1, 2) + ds = torch.randn(batch, c, 1, w, h, generator=generator).to(dtype).transpose(-1, -2) + dx = torch.autograd.grad(packed, x, dp.to(device))[0] + ref_dx = torch.autograd.grad(expected, cpu, dp)[0] + dt = torch.autograd.grad(spatial, gpu_tokens, ds.to(device))[0] + ref_dt = torch.autograd.grad(ref_spatial, tokens, ds)[0] + assert_bits(dx, ref_dx) + assert_bits(dt, ref_dt) + assert_bits(pack(x[:1], 1, c, h, w), packed[:1]) + assert_bits(unpack(gpu_tokens[:1], ih, iw, 8), spatial[:1]) + + +# Check empty and edge shapes, coordinate order, singleton dimensions, and size rounding. +@pytest.mark.parametrize("dtype", DTYPES) +@pytest.mark.parametrize( + "shape", ((0, 16, 2, 2), (1, 1, 2, 2), (2, 3, 6, 10), (2, 17, 10, 14), (1, 16, 4, 66)) +) +def test_edges_and_coordinate_order(ops, dtype, shape): + pack, unpack, device = ops + b, c, h, w = shape + x = torch.arange(b * c * h * w, dtype=torch.float32).reshape(shape).to(dtype).to(device) + actual = pack(x, b, c, h, w) + expected = NativeLatentPackOp()(x.cpu(), b, c, h, w) + assert_bits(actual, expected) + assert_bits(unpack(actual, h * 8 + 7, w * 8 + 15, 8), x.unsqueeze(2)) + for view in (x.unsqueeze(1), x.unsqueeze(2)): + assert_bits(pack(view, b, c, h, w), expected) + if b and c: + assert_bits(actual[0, 0, :4], x[0, 0, :2, :2].flatten()) + + +# Check bit preservation for special values and gradients with odd storage offsets. +@pytest.mark.parametrize("dtype", DTYPES) +def test_special_bits_and_unaligned_storage(ops, dtype): + pack, unpack, device = ops + patterns = { + torch.float32: ( + 0, + 0x80000000, + 0x7F800000, + 0xFF800000, + 0x7FC01234, + 0x7F801234, + 1, + 0x80000001, + 0x7F7FFFFF, + ), + torch.float16: (0, 0x8000, 0x7C00, 0xFC00, 0x7E55, 0x7C55, 1, 0x8001, 0x7BFF), + torch.bfloat16: (0, 0x8000, 0x7F80, 0xFF80, 0x7FC5, 0x7F85, 1, 0x8001, 0x7F7F), + } + bits = torch.uint32 if dtype == torch.float32 else torch.uint16 + storage = torch.tensor(patterns[dtype], dtype=bits).view(dtype).repeat(100).to(device) + x = storage[1 : 1 + 2 * 3 * 6 * 10].view(2, 3, 6, 10).requires_grad_() + y = pack(x, 2, 3, 6, 10) + assert_bits(y, NativeLatentPackOp()(x.cpu(), 2, 3, 6, 10)) + assert_bits(unpack(y, 48, 80, 8), x.unsqueeze(2)) + tokens = storage[1:361].view(2, 15, 12).detach().requires_grad_() + z = unpack(tokens, 48, 80, 8) + assert_bits(z, NativeLatentUnpackOp()(tokens.cpu(), 48, 80, 8)) + assert_bits(torch.autograd.grad(y, x, tokens)[0], z.squeeze(2)) + assert_bits(torch.autograd.grad(z, tokens, x.unsqueeze(2))[0], y) + + +# Check forward/backward consistency, gradients through backward, and CUDA stream execution. +def test_forward_backward_and_stream(ops): + pack, unpack, device = ops + x = torch.randn(2, 3, 6, 10, device=device, requires_grad=True) + y = pack(x, 2, 3, 6, 10) + grad = torch.randn_like(y, requires_grad=True) + dx = torch.autograd.grad(y, x, grad, create_graph=True)[0] + probe = torch.randn_like(x) + assert_bits(torch.autograd.grad(dx, grad, probe)[0], pack(probe, 2, 3, 6, 10)) + if device == "cuda": + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + result = unpack(pack(x, 2, 3, 6, 10), 48, 80, 8) + torch.cuda.current_stream().wait_stream(stream) + assert_bits(result, x.unsqueeze(2)) + + +# Check invalid input rejection and equivalent positional and keyword arguments. +def test_validation(ops): + pack, unpack, device = ops + x = torch.empty(1, 3, 6, 10, device=device) + for args in ((1, 3, 5, 10), (1, 4, 6, 10), (1, 3, 0, 10)): + with pytest.raises(ValueError): + pack(x, *args) + with pytest.raises(ValueError, match="contiguous"): + pack(x.transpose(-1, -2), 1, 3, 10, 6) + with pytest.raises(TypeError): + pack(x.to(torch.float64), 1, 3, 6, 10) + y = pack(x, 1, 3, 6, 10) + assert_bits(pack(x, batch_size=1, num_channels_latents=3, height=6, width=10), y) + assert_bits(unpack(y, height=48, width=80, vae_scale_factor=8), x.unsqueeze(2)) + for args in ((48, 80, 0), (64, 80, 8), (1, 1, 8)): + with pytest.raises(ValueError): + unpack(y, *args) + + +# Check backend fallback order, cached instance reuse, and failure when no backend works. +def test_registry_fallback(monkeypatch): + registry = KernelRegistry() + original = registry._get_or_create_backend + attempted = [] + + def native_only(backend): + attempted.append(backend) + return original(backend) if backend.name.startswith("PYTORCH") else None + + monkeypatch.setattr(torch.version, "hip", None) + monkeypatch.setattr(registry, "_get_or_create_backend", native_only) + for direction, op_class in (("pack", NativeLatentPackOp), ("unpack", NativeLatentUnpackOp)): + attempted.clear() + op = registry.get_op(f"latent_{direction}", "cuda") + assert isinstance(op, op_class) + assert attempted == [ + OpBackend[f"{backend}_LATENT_{direction.upper()}"] + for backend in ("CUDA", "TRITON", "PYTORCH") + ] + assert registry.get_op(f"latent_{direction}", "cpu") is op + + monkeypatch.setattr(registry, "_get_or_create_backend", lambda backend: None) + for direction in ("pack", "unpack"): + with pytest.raises(RuntimeError, match="No functional backend"): + registry.get_op(f"latent_{direction}", "cuda") + + +# Check that CUDA initialization fails when the extension lacks the latent kernel. +def test_cuda_missing_symbol_gate(monkeypatch): + from rl_engine.kernels.ops.cuda.packing.latent_pack_unpack import CudaLatentPackOp + + monkeypatch.setattr(base, "_C", object()) + monkeypatch.setattr(base, "_EXT_AVAILABLE", True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.version, "hip", None) + with pytest.raises(RuntimeError, match="Rebuild"): + CudaLatentPackOp() + + +# Check that the CUDA binding rejects invalid devices, dtypes, shapes, and dimensions. +@pytest.mark.cuda_only +def test_native_binding_validation(): + if not torch.cuda.is_available() or not hasattr(base._C, "latent_pack_unpack"): + pytest.skip("CUDA latent extension unavailable") + x = torch.empty(1, 3, 6, 10, device="cuda") + for invalid in (x.cpu(), x.double(), x.flatten(), x.transpose(-1, -2)): + with pytest.raises(RuntimeError): + base._C.latent_pack_unpack(invalid, 1, 3, 6, 10, False) + for dims in ((-1, 3, 6, 10), (1, 4, 6, 10), (1, 3, 5, 10), (1, 2**62, 6, 10)): + with pytest.raises(RuntimeError): + base._C.latent_pack_unpack(x, *dims, False)