Qwen image latent pack unpack

csrc/cuda/latent_pack_unpack.cu

@@ -0,0 +1,127 @@
+// SPDX-License-Identifier: Apache-2.0
+// Copyright (c) 2026 RL-Kernel Contributors
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
+
+#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 <typename Bits> struct Pair { Bits lo, hi; };
+
+template <typename Bits, bool Unpack>
+__global__ void latent_permute(const Pair<Bits>* __restrict__ input, Pair<Bits>* __restrict__ output, int64_t num_channels, int64_t latent_height, int64_t latent_width) {
+  __shared__ Pair<Bits> 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<int64_t>(blockIdx.z) * num_channels * latent_height * token_cols;
+  const int64_t token_col_start = static_cast<int64_t>(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 <typename Bits, bool Unpack>
+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<Bits, Unpack><<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(reinterpret_cast<const Pair<Bits>*>(x.data_ptr()), reinterpret_cast<Pair<Bits>*>(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<uint32_t, true>(x, y, B, C, H, W);
+    } else {
+      launch<uint32_t, false>(x, y, B, C, H, W);
+    }
+  } else {
+    if (unpack) {
+      launch<uint16_t, true>(x, y, B, C, H, W);
+    } else {
+      launch<uint16_t, false>(x, y, B, C, H, W);
+    }
+  }
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return y;
+}

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<std::vector<int64_t>, 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");

rl_engine/kernels/ops/cuda/packing/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

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

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)

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

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

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

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

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"] = {

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")