Promote latest kernel artifacts to main

README.md

@@ -1,9 +0,0 @@
-# flashrt/flashrt-spatiotemporal-layout
-
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
-
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/flashrt-spatiotemporal-layout
-
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.

benchmarks/benchmark_native_parity.py

@@ -0,0 +1,433 @@
+#!/usr/bin/env python3
+"""Benchmark Tensor wrappers against native FlashRT and torch baselines."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import os
+import sys
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import torch
+
+
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "flashrt-spatiotemporal-layout"
+REGISTRATION = (
+    ROOT.parent / "kernels/kernel-builder/src/pyproject/templates/torch"
+)
+SHAPES = {
+    "latent-small": (1, 64, 4, 32, 32),
+    "vae-channel320": (1, 320, 17, 32, 32),
+    "vae-channel512": (1, 512, 4, 64, 64),
+}
+
+
+@dataclass
+class Result:
+    workload: str
+    shape: str
+    dtype: str
+    native_us: float
+    wrapper_us: float
+    native_parity: float
+    graph_native_us: float
+    graph_wrapper_us: float
+    graph_native_parity: float
+    eager_us: float
+    compile_us: float
+    strong_library_us: str
+    max_abs: float
+    mean_abs: float
+    p99_abs: float
+    cosine: float
+    accepted: bool
+
+
+def time_us(fn, warmup: int, iters: int) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    begin = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    begin.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return begin.elapsed_time(end) * 1000.0 / iters
+
+
+def graph_time_us(fn, warmup: int, iters: int) -> float:
+    side = torch.cuda.Stream()
+    side.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(side):
+        fn()
+    torch.cuda.current_stream().wait_stream(side)
+    torch.cuda.synchronize()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        fn()
+    return time_us(graph.replay, warmup, iters)
+
+
+def build_source():
+    from torch.utils.cpp_extension import load
+
+    major, minor = torch.cuda.get_device_capability()
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", f"{major}.{minor}")
+    namespace = "flashrt_spatiotemporal_layout_native_parity"
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext/torch_binding.cpp"),
+            str(PACKAGE / "csrc/spatiotemporal_layout.cu"),
+            str(PACKAGE / "csrc/bf16_ndhwc_to_ncdhw_transpose.cu"),
+            str(PACKAGE / "csrc/bf16_quant_fp8_ncdhw_to_ndhwc.cu"),
+        ],
+        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION)],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=["-O3", "-DCUDA_KERNEL"],
+        is_python_module=False,
+        verbose=False,
+    )
+    return getattr(torch.ops, namespace)
+
+
+def load_wrapper(backend: str, artifact: str | None):
+    if backend == "source":
+        return build_source()
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("flashrt_spatiotemporal_layout")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
+def build_native():
+    from torch.utils.cpp_extension import load
+
+    return load(
+        name="flashrt_spatiotemporal_layout_raw_native",
+        sources=[
+            str(PACKAGE / "benchmarks/native_binding.cpp"),
+            str(PACKAGE / "csrc/spatiotemporal_layout.cu"),
+            str(PACKAGE / "csrc/bf16_ndhwc_to_ncdhw_transpose.cu"),
+            str(PACKAGE / "csrc/bf16_quant_fp8_ncdhw_to_ndhwc.cu"),
+        ],
+        extra_include_paths=[str(PACKAGE / "csrc")],
+        extra_cflags=["-O3"],
+        extra_cuda_cflags=["-O3"],
+        verbose=False,
+    )
+
+
+def metrics(got: torch.Tensor, ref: torch.Tensor):
+    if got.dtype == torch.float8_e4m3fn:
+        exact = torch.equal(got.view(torch.uint8), ref.view(torch.uint8))
+        return (0.0, 0.0, 0.0, 1.0) if exact else (float("inf"),) * 4
+    diff = (got.float() - ref.float()).abs().flatten()
+    cosine = torch.nn.functional.cosine_similarity(
+        got.float().flatten(), ref.float().flatten(), dim=0
+    ).item()
+    return (
+        diff.max().item(),
+        diff.mean().item(),
+        torch.quantile(diff, 0.99).item(),
+        cosine,
+    )
+
+
+def add_result(
+    rows,
+    workload,
+    shape,
+    wrapper_fn,
+    native_fn,
+    eager_fn,
+    compiled_fn,
+    got,
+    ref,
+    args,
+):
+    native_us = time_us(native_fn, args.warmup, args.iters)
+    wrapper_us = time_us(wrapper_fn, args.warmup, args.iters)
+    graph_native_us = graph_time_us(native_fn, args.warmup, args.iters)
+    graph_wrapper_us = graph_time_us(wrapper_fn, args.warmup, args.iters)
+    eager_us = time_us(eager_fn, args.warmup, args.iters)
+    compile_us = time_us(compiled_fn, args.warmup, args.iters)
+    max_abs, mean_abs, p99_abs, cosine = metrics(got, ref)
+    parity = wrapper_us / native_us
+    rows.append(
+        Result(
+            workload,
+            str(shape),
+            str(got.dtype),
+            native_us,
+            wrapper_us,
+            parity,
+            graph_native_us,
+            graph_wrapper_us,
+            graph_wrapper_us / graph_native_us,
+            eager_us,
+            compile_us,
+            "N/A (no equivalent single library op)",
+            max_abs,
+            mean_abs,
+            p99_abs,
+            cosine,
+            (
+                wrapper_us - native_us <= max(0.75, native_us * 0.05)
+                or graph_wrapper_us - graph_native_us
+                <= max(0.5, graph_native_us * 0.05)
+            )
+            and min(wrapper_us, graph_wrapper_us)
+            <= min(eager_us, compile_us) * 0.98
+            and max_abs == 0.0,
+        )
+    )
+
+
+def run_shape(wrapper, native, name, shape, args):
+    b, c, t, h, w = shape
+    x = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+    ndhwc = x.permute(0, 2, 3, 4, 1).contiguous()
+    bias = torch.randn(c, device="cuda", dtype=torch.bfloat16)
+    residual = torch.randn_like(x)
+    scale = 0.03125
+    rows = []
+
+    def wrapper_layout():
+        wrapper.ndhwc_to_ncdhw_bf16(ndhwc, out_layout)
+
+    def eager_layout():
+        out_layout.copy_(ndhwc.permute(0, 4, 1, 2, 3))
+
+    compiled_layout = torch.compile(eager_layout, fullgraph=True)
+    out_layout = torch.empty_like(x)
+    native_layout = lambda: native.ndhwc_to_ncdhw(ndhwc, out_layout)
+    wrapper_layout()
+    ref = ndhwc.permute(0, 4, 1, 2, 3).contiguous()
+    add_result(
+        rows, f"{name}/ndhwc_to_ncdhw", shape, wrapper_layout, native_layout,
+        eager_layout, compiled_layout, out_layout, ref, args
+    )
+
+    out_bias = torch.empty_like(x)
+    wrapper_bias = lambda: wrapper.ndhwc_to_ncdhw_bias_bf16(
+        ndhwc, bias, out_bias
+    )
+    native_bias = lambda: native.ndhwc_to_ncdhw_bias(ndhwc, bias, out_bias)
+
+    def eager_bias():
+        out_bias.copy_(
+            (ndhwc.permute(0, 4, 1, 2, 3).float()
+             + bias.float().view(1, c, 1, 1, 1)).to(torch.bfloat16)
+        )
+
+    compiled_bias = torch.compile(eager_bias, fullgraph=True)
+    wrapper_bias()
+    ref_bias = (
+        x.float() + bias.float().view(1, c, 1, 1, 1)
+    ).to(torch.bfloat16)
+    add_result(
+        rows, f"{name}/ndhwc_to_ncdhw_bias", shape, wrapper_bias, native_bias,
+        eager_bias, compiled_bias, out_bias, ref_bias, args
+    )
+
+    out_add = torch.empty_like(x)
+    wrapper_add = lambda: wrapper.ndhwc_to_ncdhw_add_bf16(
+        ndhwc, residual, out_add
+    )
+    native_add = lambda: native.ndhwc_to_ncdhw_add(
+        ndhwc, residual, out_add
+    )
+
+    def eager_add():
+        out_add.copy_(
+            (ndhwc.permute(0, 4, 1, 2, 3).float()
+             + residual.float()).to(torch.bfloat16)
+        )
+
+    compiled_add = torch.compile(eager_add, fullgraph=True)
+    wrapper_add()
+    ref_add = (x.float() + residual.float()).to(torch.bfloat16)
+    add_result(
+        rows, f"{name}/ndhwc_to_ncdhw_add", shape, wrapper_add, native_add,
+        eager_add, compiled_add, out_add, ref_add, args
+    )
+
+    out_fp8 = torch.empty(
+        (b, t, h, w, c), device="cuda", dtype=torch.float8_e4m3fn
+    )
+    wrapper_quant = lambda: wrapper.ncdhw_quantize_fp8_static_ndhwc_bf16(
+        x, scale, out_fp8
+    )
+    native_quant = lambda: native.ncdhw_quantize(x, scale, out_fp8)
+
+    def eager_quant():
+        out_fp8.copy_(
+            (x.float() / scale).clamp(-448.0, 448.0)
+            .to(torch.float8_e4m3fn).permute(0, 2, 3, 4, 1)
+        )
+
+    compiled_quant = torch.compile(eager_quant, fullgraph=True)
+    wrapper_quant()
+    ref_fp8 = (
+        (x.float() / scale).clamp(-448.0, 448.0)
+        .to(torch.float8_e4m3fn).permute(0, 2, 3, 4, 1).contiguous()
+    )
+    add_result(
+        rows, f"{name}/ncdhw_quantize_fp8_ndhwc", shape, wrapper_quant,
+        native_quant, eager_quant, compiled_quant, out_fp8, ref_fp8, args
+    )
+    return rows
+
+
+def run_cache_and_upsample(wrapper, native, args):
+    rows = []
+    previous = torch.randn(
+        (1, 64, 2, 32, 32), device="cuda", dtype=torch.bfloat16
+    )
+    current = torch.randn(
+        (1, 64, 1, 32, 32), device="cuda", dtype=torch.bfloat16
+    )
+    packed = torch.empty(
+        (1, 32, 32, 192), device="cuda", dtype=torch.bfloat16
+    )
+    wrapper_pack = lambda: wrapper.pack_causal_cache3_nhwc_bf16(
+        previous, current, packed
+    )
+    native_pack = lambda: native.pack_causal_cache3_nhwc(
+        previous, current, packed
+    )
+
+    def eager_pack():
+        packed.copy_(
+            torch.cat(
+                (previous[:, :, 0], previous[:, :, 1], current[:, :, 0]),
+                dim=1,
+            ).permute(0, 2, 3, 1)
+        )
+
+    compiled_pack = torch.compile(eager_pack, fullgraph=True)
+    wrapper_pack()
+    got_pack = packed.clone()
+    ref_pack = torch.cat(
+        (previous[:, :, 0], previous[:, :, 1], current[:, :, 0]), dim=1
+    ).permute(0, 2, 3, 1).contiguous()
+    add_result(
+        rows,
+        "vae-t1/pack_causal_cache3_nhwc",
+        tuple(current.shape),
+        wrapper_pack,
+        native_pack,
+        eager_pack,
+        compiled_pack,
+        got_pack,
+        ref_pack,
+        args,
+    )
+
+    for temporal_factor, first_chunk in ((1, False), (2, True)):
+        spatial_factor = 2
+        out_channels = 16 if temporal_factor == 1 else 8
+        input = current if temporal_factor == 1 else torch.randn(
+            (1, 64, 4, 16, 16), device="cuda", dtype=torch.bfloat16
+        )
+        out_t = input.shape[2] * temporal_factor - (
+            temporal_factor - 1 if first_chunk else 0
+        )
+        out = torch.empty(
+            (
+                input.shape[0], out_channels, out_t,
+                input.shape[3] * spatial_factor,
+                input.shape[4] * spatial_factor,
+            ),
+            device="cuda",
+            dtype=torch.bfloat16,
+        )
+        wrapper_up = lambda: wrapper.channel_to_space3d_bf16(
+            input, out_channels, temporal_factor, spatial_factor, 1,
+            first_chunk, out
+        )
+        native_up = lambda: native.channel_to_space3d(
+            input, out_channels, temporal_factor, spatial_factor, 1,
+            first_chunk, out
+        )
+
+        def eager_up():
+            expanded = input[:, : out_channels * temporal_factor * 4]
+            expanded = expanded.view(
+                input.shape[0], out_channels, temporal_factor, 2, 2,
+                input.shape[2], input.shape[3], input.shape[4],
+            ).permute(0, 1, 5, 2, 6, 3, 7, 4)
+            value = expanded.reshape(
+                input.shape[0], out_channels,
+                input.shape[2] * temporal_factor,
+                input.shape[3] * 2, input.shape[4] * 2,
+            )
+            if first_chunk:
+                value = value[:, :, temporal_factor - 1 :]
+            out.copy_(value)
+
+        compiled_up = torch.compile(eager_up, fullgraph=True)
+        wrapper_up()
+        got = out.clone()
+        eager_up()
+        ref = out.clone()
+        out.copy_(got)
+        add_result(
+            rows,
+            f"vae/channel_to_space3d-ft{temporal_factor}"
+            f"-first{int(first_chunk)}",
+            tuple(input.shape),
+            wrapper_up,
+            native_up,
+            eager_up,
+            compiled_up,
+            got,
+            ref,
+            args,
+        )
+    return rows
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact")
+    parser.add_argument("--warmup", type=int, default=20)
+    parser.add_argument("--iters", type=int, default=100)
+    parser.add_argument("--output")
+    args = parser.parse_args()
+    wrapper = load_wrapper(args.backend, args.artifact)
+    native = build_native()
+    rows = []
+    for name, shape in SHAPES.items():
+        rows.extend(run_shape(wrapper, native, name, shape, args))
+    rows.extend(run_cache_and_upsample(wrapper, native, args))
+    for row in rows:
+        print(
+            f"{row.workload}: native={row.native_us:.3f}us "
+            f"wrapper={row.wrapper_us:.3f}us parity={row.native_parity:.3f} "
+            f"graph={row.graph_wrapper_us:.3f}/"
+            f"{row.graph_native_us:.3f}us "
+            f"eager={row.eager_us:.3f}us compile={row.compile_us:.3f}us "
+            f"accepted={row.accepted}"
+        )
+    if args.output:
+        path = Path(args.output)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(json.dumps([asdict(row) for row in rows], indent=2) + "\n")
+    if not all(row.accepted for row in rows):
+        raise SystemExit("performance/correctness acceptance failed")
+
+
+if __name__ == "__main__":
+    main()

build/torch211-cxx11-cu128-x86_64-linux/__init__.py

@@ -58,6 +58,137 @@ def _update_cache2_ncdhw_bf16_fake(cur: torch.Tensor, prev: torch.Tensor, out: t
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("channel_to_space3d_bf16"))
+def _channel_to_space3d_bf16_fake(
+    x: torch.Tensor,
+    out_channels: int,
+    temporal_factor: int,
+    spatial_factor: int,
+    repeats: int,
+    first_chunk: bool,
+    out: torch.Tensor,
+) -> None:
+    _check_ncdhw(x, "x")
+    out_t = x.shape[2] * temporal_factor - (
+        temporal_factor - 1 if first_chunk else 0
+    )
+    expected = (
+        x.shape[0],
+        out_channels,
+        out_t,
+        x.shape[3] * spatial_factor,
+        x.shape[4] * spatial_factor,
+    )
+    if (
+        out_channels <= 0
+        or temporal_factor <= 0
+        or spatial_factor <= 0
+        or repeats <= 0
+        or x.shape[1] * repeats
+        < out_channels * temporal_factor * spatial_factor * spatial_factor
+        or out.shape != expected
+    ):
+        raise RuntimeError("channel_to_space3d_bf16 shape contract failed")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("pack_causal_cache3_nhwc_bf16")
+)
+def _pack_causal_cache3_nhwc_bf16_fake(
+    previous: torch.Tensor, current: torch.Tensor, out: torch.Tensor
+) -> None:
+    _check_ncdhw(previous, "previous")
+    _check_ncdhw(current, "current")
+    b, c, t, h, w = current.shape
+    if (
+        t != 1
+        or previous.shape != (b, c, 2, h, w)
+        or out.shape != (b, h, w, 3 * c)
+    ):
+        raise RuntimeError("causal cache pack shape contract failed")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("avg_pool3d_channels_bf16"))
+def _avg_pool3d_channels_bf16_fake(
+    x: torch.Tensor,
+    out_channels: int,
+    factor_t: int,
+    factor_s: int,
+    group_size: int,
+    out: torch.Tensor,
+) -> None:
+    _check_ncdhw(x, "x")
+    expected = (
+        x.shape[0],
+        out_channels,
+        (x.shape[2] + factor_t - 1) // factor_t,
+        x.shape[3] // factor_s,
+        x.shape[4] // factor_s,
+    )
+    if out.shape != expected:
+        raise RuntimeError("out has the wrong pooled NCDHW shape")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_bf16"))
+def _ndhwc_to_ncdhw_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B, T, H, W, C)")
+    expected = (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (B, C, T, H, W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_bias_bf16"))
+def _ndhwc_to_ncdhw_bias_bf16_fake(
+    x: torch.Tensor, bias: torch.Tensor, out: torch.Tensor
+) -> None:
+    _ndhwc_to_ncdhw_bf16_fake(x, out)
+    if bias.shape != (x.shape[4],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_add_bf16"))
+def _ndhwc_to_ncdhw_add_bf16_fake(
+    x: torch.Tensor, residual: torch.Tensor, out: torch.Tensor
+) -> None:
+    _ndhwc_to_ncdhw_bf16_fake(x, out)
+    if residual.shape != out.shape:
+        raise RuntimeError("residual must match the NCDHW output shape")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("ncdhw_quantize_fp8_static_ndhwc_bf16")
+)
+def _ncdhw_quantize_fp8_static_ndhwc_bf16_fake(
+    x: torch.Tensor, scale: float, out: torch.Tensor
+) -> None:
+    _check_ncdhw(x, "x")
+    expected = (x.shape[0], x.shape[2], x.shape[3], x.shape[4], x.shape[1])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (B, T, H, W, C)")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("upsample2x_quantize_fp8_static_nhwc_bf16")
+)
+def _upsample2x_quantize_fp8_static_nhwc_bf16_fake(
+    x: torch.Tensor, scale: float, out: torch.Tensor
+) -> None:
+    if x.dim() != 4:
+        raise RuntimeError("x must have shape (N, C, H, W)")
+    expected = (x.shape[0], 2 * x.shape[2], 2 * x.shape[3], x.shape[1])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (N, 2H, 2W, C)")
+    return None
+
+
 def ncdhw_to_blc_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
     """Convert BF16 NCDHW tensor to BLC where L = T * H * W."""
 
@@ -101,10 +232,183 @@ def update_cache2_ncdhw_bf16(cur: torch.Tensor, prev: torch.Tensor, out: torch.T
     return out
 
 
+def channel_to_space3d_bf16(
+    x: torch.Tensor,
+    out_channels: int,
+    temporal_factor: int,
+    spatial_factor: int,
+    repeats: int = 1,
+    first_chunk: bool = False,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Move expanded channels into temporal/spatial dimensions."""
+
+    out_t = x.shape[2] * temporal_factor - (
+        temporal_factor - 1 if first_chunk else 0
+    )
+    if out is None:
+        out = torch.empty(
+            (
+                x.shape[0],
+                out_channels,
+                out_t,
+                x.shape[3] * spatial_factor,
+                x.shape[4] * spatial_factor,
+            ),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.channel_to_space3d_bf16(
+        x,
+        int(out_channels),
+        int(temporal_factor),
+        int(spatial_factor),
+        int(repeats),
+        bool(first_chunk),
+        out,
+    )
+    return out
+
+
+def pack_causal_cache3_nhwc_bf16(
+    previous: torch.Tensor,
+    current: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Pack two cached and one current NCDHW frame into NHWC with 3C."""
+
+    if out is None:
+        out = torch.empty(
+            (
+                current.shape[0],
+                current.shape[3],
+                current.shape[4],
+                3 * current.shape[1],
+            ),
+            device=current.device,
+            dtype=current.dtype,
+        )
+    ops.pack_causal_cache3_nhwc_bf16(previous, current, out)
+    return out
+
+
+def avg_pool3d_channels_bf16(
+    x: torch.Tensor,
+    out_channels: int,
+    factor_t: int,
+    factor_s: int,
+    group_size: int,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Pool NCDHW blocks while folding spatiotemporal factors into channels."""
+
+    if out is None:
+        out = torch.empty(
+            (
+                x.shape[0],
+                out_channels,
+                (x.shape[2] + factor_t - 1) // factor_t,
+                x.shape[3] // factor_s,
+                x.shape[4] // factor_s,
+            ),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.avg_pool3d_channels_bf16(
+        x, out_channels, factor_t, factor_s, group_size, out
+    )
+    return out
+
+
+def ndhwc_to_ncdhw_bf16(
+    x: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Convert contiguous BF16 NDHWC to contiguous BF16 NCDHW."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3]),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.ndhwc_to_ncdhw_bf16(x, out)
+    return out
+
+
+def ndhwc_to_ncdhw_bias_bf16(
+    x: torch.Tensor, bias: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Convert NDHWC to NCDHW and add a BF16 channel bias."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3]),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.ndhwc_to_ncdhw_bias_bf16(x, bias, out)
+    return out
+
+
+def ndhwc_to_ncdhw_add_bf16(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Convert NDHWC to NCDHW and add a BF16 NCDHW residual."""
+
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.ndhwc_to_ncdhw_add_bf16(x, residual, out)
+    return out
+
+
+def ncdhw_quantize_fp8_static_ndhwc_bf16(
+    x: torch.Tensor,
+    scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Quantize BF16 NCDHW to FP8 E4M3 NDHWC using a static scale."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[2], x.shape[3], x.shape[4], x.shape[1]),
+            device=x.device,
+            dtype=torch.float8_e4m3fn,
+        )
+    ops.ncdhw_quantize_fp8_static_ndhwc_bf16(x, scale, out)
+    return out
+
+
+def upsample2x_quantize_fp8_static_nhwc_bf16(
+    x: torch.Tensor,
+    scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Nearest-neighbor 2x upsample NCHW and emit static-scaled FP8 NHWC."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], 2 * x.shape[2], 2 * x.shape[3], x.shape[1]),
+            device=x.device,
+            dtype=torch.float8_e4m3fn,
+        )
+    ops.upsample2x_quantize_fp8_static_nhwc_bf16(x, scale, out)
+    return out
+
+
 __all__ = [
     "ncdhw_to_blc_bf16",
     "patch_im2col_bf16",
     "time_unshuffle2_bf16",
     "add_bias_ncdhw_bf16",
     "update_cache2_ncdhw_bf16",
+    "channel_to_space3d_bf16",
+    "pack_causal_cache3_nhwc_bf16",
+    "avg_pool3d_channels_bf16",
+    "ndhwc_to_ncdhw_bf16",
+    "ndhwc_to_ncdhw_bias_bf16",
+    "ndhwc_to_ncdhw_add_bf16",
+    "ncdhw_quantize_fp8_static_ndhwc_bf16",
+    "upsample2x_quantize_fp8_static_nhwc_bf16",
 ]

build/{torch211-cxx11-cu130-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_5de4768.abi3.so → torch211-cxx11-cu128-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_7781728.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b925abfb8b8b8dadffaeba05b23a0059612b6ccc83ddcd75a620bcb606daad6f
-size 696272
+oid sha256:54170a1e1ac254a6ba0e14432e3fe9b5ec1f09f3ff218a9e4fbf0463134653ab
+size 1887872

build/torch211-cxx11-cu128-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _flashrt_spatiotemporal_layout_cuda_5de4768
-ops = torch.ops._flashrt_spatiotemporal_layout_cuda_5de4768
+from . import _flashrt_spatiotemporal_layout_cuda_7781728
+ops = torch.ops._flashrt_spatiotemporal_layout_cuda_7781728
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_flashrt_spatiotemporal_layout_cuda_5de4768::{op_name}"
+    return f"_flashrt_spatiotemporal_layout_cuda_7781728::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-spatiotemporal-layout",
-  "id": "_flashrt_spatiotemporal_layout_cuda_5de4768",
+  "id": "_flashrt_spatiotemporal_layout_cuda_7781728",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch211-cxx11-cu130-aarch64-linux/__init__.py

@@ -0,0 +1,414 @@
+"""FlashRT spatiotemporal layout kernels."""
+
+from __future__ import annotations
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def _check_ncdhw(x: torch.Tensor, name: str) -> None:
+    if x.dim() != 5:
+        raise RuntimeError(f"{name} must have shape (B, C, T, H, W)")
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ncdhw_to_blc_bf16"))
+def _ncdhw_to_blc_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c, t, h, w = x.shape
+    if out.shape != (b, t * h * w, c):
+        raise RuntimeError("out must have shape (B, T * H * W, C)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("patch_im2col_bf16"))
+def _patch_im2col_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    if x.dim() != 4 or x.shape[1:] != (224, 224, 3):
+        raise RuntimeError("x must have shape (num_views, 224, 224, 3)")
+    if out.shape != (x.shape[0] * 256, 588):
+        raise RuntimeError("out must have shape (num_views * 256, 588)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("time_unshuffle2_bf16"))
+def _time_unshuffle2_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c2, t, h, w = x.shape
+    if c2 % 2 != 0:
+        raise RuntimeError("x.shape[1] must be even")
+    if out.shape != (b, c2 // 2, 2 * t, h, w):
+        raise RuntimeError("out must have shape (B, C / 2, 2 * T, H, W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("add_bias_ncdhw_bf16"))
+def _add_bias_ncdhw_bf16_fake(x: torch.Tensor, bias: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    if bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("update_cache2_ncdhw_bf16"))
+def _update_cache2_ncdhw_bf16_fake(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(cur, "cur")
+    b, c, _, h, w = cur.shape
+    if prev.shape != (b, c, 2, h, w) or out.shape != (b, c, 2, h, w):
+        raise RuntimeError("prev and out must have shape (B, C, 2, H, W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("channel_to_space3d_bf16"))
+def _channel_to_space3d_bf16_fake(
+    x: torch.Tensor,
+    out_channels: int,
+    temporal_factor: int,
+    spatial_factor: int,
+    repeats: int,
+    first_chunk: bool,
+    out: torch.Tensor,
+) -> None:
+    _check_ncdhw(x, "x")
+    out_t = x.shape[2] * temporal_factor - (
+        temporal_factor - 1 if first_chunk else 0
+    )
+    expected = (
+        x.shape[0],
+        out_channels,
+        out_t,
+        x.shape[3] * spatial_factor,
+        x.shape[4] * spatial_factor,
+    )
+    if (
+        out_channels <= 0
+        or temporal_factor <= 0
+        or spatial_factor <= 0
+        or repeats <= 0
+        or x.shape[1] * repeats
+        < out_channels * temporal_factor * spatial_factor * spatial_factor
+        or out.shape != expected
+    ):
+        raise RuntimeError("channel_to_space3d_bf16 shape contract failed")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("pack_causal_cache3_nhwc_bf16")
+)
+def _pack_causal_cache3_nhwc_bf16_fake(
+    previous: torch.Tensor, current: torch.Tensor, out: torch.Tensor
+) -> None:
+    _check_ncdhw(previous, "previous")
+    _check_ncdhw(current, "current")
+    b, c, t, h, w = current.shape
+    if (
+        t != 1
+        or previous.shape != (b, c, 2, h, w)
+        or out.shape != (b, h, w, 3 * c)
+    ):
+        raise RuntimeError("causal cache pack shape contract failed")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("avg_pool3d_channels_bf16"))
+def _avg_pool3d_channels_bf16_fake(
+    x: torch.Tensor,
+    out_channels: int,
+    factor_t: int,
+    factor_s: int,
+    group_size: int,
+    out: torch.Tensor,
+) -> None:
+    _check_ncdhw(x, "x")
+    expected = (
+        x.shape[0],
+        out_channels,
+        (x.shape[2] + factor_t - 1) // factor_t,
+        x.shape[3] // factor_s,
+        x.shape[4] // factor_s,
+    )
+    if out.shape != expected:
+        raise RuntimeError("out has the wrong pooled NCDHW shape")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_bf16"))
+def _ndhwc_to_ncdhw_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B, T, H, W, C)")
+    expected = (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (B, C, T, H, W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_bias_bf16"))
+def _ndhwc_to_ncdhw_bias_bf16_fake(
+    x: torch.Tensor, bias: torch.Tensor, out: torch.Tensor
+) -> None:
+    _ndhwc_to_ncdhw_bf16_fake(x, out)
+    if bias.shape != (x.shape[4],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_add_bf16"))
+def _ndhwc_to_ncdhw_add_bf16_fake(
+    x: torch.Tensor, residual: torch.Tensor, out: torch.Tensor
+) -> None:
+    _ndhwc_to_ncdhw_bf16_fake(x, out)
+    if residual.shape != out.shape:
+        raise RuntimeError("residual must match the NCDHW output shape")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("ncdhw_quantize_fp8_static_ndhwc_bf16")
+)
+def _ncdhw_quantize_fp8_static_ndhwc_bf16_fake(
+    x: torch.Tensor, scale: float, out: torch.Tensor
+) -> None:
+    _check_ncdhw(x, "x")
+    expected = (x.shape[0], x.shape[2], x.shape[3], x.shape[4], x.shape[1])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (B, T, H, W, C)")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("upsample2x_quantize_fp8_static_nhwc_bf16")
+)
+def _upsample2x_quantize_fp8_static_nhwc_bf16_fake(
+    x: torch.Tensor, scale: float, out: torch.Tensor
+) -> None:
+    if x.dim() != 4:
+        raise RuntimeError("x must have shape (N, C, H, W)")
+    expected = (x.shape[0], 2 * x.shape[2], 2 * x.shape[3], x.shape[1])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (N, 2H, 2W, C)")
+    return None
+
+
+def ncdhw_to_blc_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 NCDHW tensor to BLC where L = T * H * W."""
+
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[2] * x.shape[3] * x.shape[4], x.shape[1]), device=x.device, dtype=x.dtype)
+    ops.ncdhw_to_blc_bf16(x, out)
+    return out
+
+
+def patch_im2col_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 NHWC images into flattened 14x14 patch rows."""
+
+    if out is None:
+        out = torch.empty((x.shape[0] * 256, 588), device=x.device, dtype=x.dtype)
+    ops.patch_im2col_bf16(x, out)
+    return out
+
+
+def time_unshuffle2_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 (B, 2C, T, H, W) into (B, C, 2T, H, W)."""
+
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[1] // 2, 2 * x.shape[2], x.shape[3], x.shape[4]), device=x.device, dtype=x.dtype)
+    ops.time_unshuffle2_bf16(x, out)
+    return out
+
+
+def add_bias_ncdhw_bf16(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    """Add channel bias to an NCDHW tensor in place."""
+
+    ops.add_bias_ncdhw_bf16(x, bias)
+    return x
+
+
+def update_cache2_ncdhw_bf16(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Update two-frame latent cache from current and previous NCDHW tensors."""
+
+    if out is None:
+        out = torch.empty((cur.shape[0], cur.shape[1], 2, cur.shape[3], cur.shape[4]), device=cur.device, dtype=cur.dtype)
+    ops.update_cache2_ncdhw_bf16(cur, prev, out)
+    return out
+
+
+def channel_to_space3d_bf16(
+    x: torch.Tensor,
+    out_channels: int,
+    temporal_factor: int,
+    spatial_factor: int,
+    repeats: int = 1,
+    first_chunk: bool = False,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Move expanded channels into temporal/spatial dimensions."""
+
+    out_t = x.shape[2] * temporal_factor - (
+        temporal_factor - 1 if first_chunk else 0
+    )
+    if out is None:
+        out = torch.empty(
+            (
+                x.shape[0],
+                out_channels,
+                out_t,
+                x.shape[3] * spatial_factor,
+                x.shape[4] * spatial_factor,
+            ),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.channel_to_space3d_bf16(
+        x,
+        int(out_channels),
+        int(temporal_factor),
+        int(spatial_factor),
+        int(repeats),
+        bool(first_chunk),
+        out,
+    )
+    return out
+
+
+def pack_causal_cache3_nhwc_bf16(
+    previous: torch.Tensor,
+    current: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Pack two cached and one current NCDHW frame into NHWC with 3C."""
+
+    if out is None:
+        out = torch.empty(
+            (
+                current.shape[0],
+                current.shape[3],
+                current.shape[4],
+                3 * current.shape[1],
+            ),
+            device=current.device,
+            dtype=current.dtype,
+        )
+    ops.pack_causal_cache3_nhwc_bf16(previous, current, out)
+    return out
+
+
+def avg_pool3d_channels_bf16(
+    x: torch.Tensor,
+    out_channels: int,
+    factor_t: int,
+    factor_s: int,
+    group_size: int,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Pool NCDHW blocks while folding spatiotemporal factors into channels."""
+
+    if out is None:
+        out = torch.empty(
+            (
+                x.shape[0],
+                out_channels,
+                (x.shape[2] + factor_t - 1) // factor_t,
+                x.shape[3] // factor_s,
+                x.shape[4] // factor_s,
+            ),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.avg_pool3d_channels_bf16(
+        x, out_channels, factor_t, factor_s, group_size, out
+    )
+    return out
+
+
+def ndhwc_to_ncdhw_bf16(
+    x: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Convert contiguous BF16 NDHWC to contiguous BF16 NCDHW."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3]),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.ndhwc_to_ncdhw_bf16(x, out)
+    return out
+
+
+def ndhwc_to_ncdhw_bias_bf16(
+    x: torch.Tensor, bias: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Convert NDHWC to NCDHW and add a BF16 channel bias."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3]),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.ndhwc_to_ncdhw_bias_bf16(x, bias, out)
+    return out
+
+
+def ndhwc_to_ncdhw_add_bf16(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Convert NDHWC to NCDHW and add a BF16 NCDHW residual."""
+
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.ndhwc_to_ncdhw_add_bf16(x, residual, out)
+    return out
+
+
+def ncdhw_quantize_fp8_static_ndhwc_bf16(
+    x: torch.Tensor,
+    scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Quantize BF16 NCDHW to FP8 E4M3 NDHWC using a static scale."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[2], x.shape[3], x.shape[4], x.shape[1]),
+            device=x.device,
+            dtype=torch.float8_e4m3fn,
+        )
+    ops.ncdhw_quantize_fp8_static_ndhwc_bf16(x, scale, out)
+    return out
+
+
+def upsample2x_quantize_fp8_static_nhwc_bf16(
+    x: torch.Tensor,
+    scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Nearest-neighbor 2x upsample NCHW and emit static-scaled FP8 NHWC."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], 2 * x.shape[2], 2 * x.shape[3], x.shape[1]),
+            device=x.device,
+            dtype=torch.float8_e4m3fn,
+        )
+    ops.upsample2x_quantize_fp8_static_nhwc_bf16(x, scale, out)
+    return out
+
+
+__all__ = [
+    "ncdhw_to_blc_bf16",
+    "patch_im2col_bf16",
+    "time_unshuffle2_bf16",
+    "add_bias_ncdhw_bf16",
+    "update_cache2_ncdhw_bf16",
+    "channel_to_space3d_bf16",
+    "pack_causal_cache3_nhwc_bf16",
+    "avg_pool3d_channels_bf16",
+    "ndhwc_to_ncdhw_bf16",
+    "ndhwc_to_ncdhw_bias_bf16",
+    "ndhwc_to_ncdhw_add_bf16",
+    "ncdhw_quantize_fp8_static_ndhwc_bf16",
+    "upsample2x_quantize_fp8_static_nhwc_bf16",
+]

build/{torch212-cxx11-cu132-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_5de4768.abi3.so → torch211-cxx11-cu130-aarch64-linux/_flashrt_spatiotemporal_layout_cuda_7781728.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5df2034b32ce3972e68eebd56c25981208b778c3fe6f19f1fce64a47695fa3d8
-size 702464
+oid sha256:ec2a08c6333c68c9b6d126414c005a8fd83a6499a8141b90c8c96a842c107df4
+size 458904

build/torch211-cxx11-cu130-aarch64-linux/_ops.py

@@ -0,0 +1,6 @@
+import torch
+from . import _flashrt_spatiotemporal_layout_cuda_7781728
+ops = torch.ops._flashrt_spatiotemporal_layout_cuda_7781728
+
+def add_op_namespace_prefix(op_name: str):
+    return f"_flashrt_spatiotemporal_layout_cuda_7781728::{op_name}"

build/torch211-cxx11-cu130-aarch64-linux/flashrt_spatiotemporal_layout/__init__.py

@@ -0,0 +1,14 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+
+def _import_from_path(file_path: Path):
+    path_hash = '{:x}'.format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    spec = importlib.util.spec_from_file_location(path_hash, file_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[path_hash] = module
+    spec.loader.exec_module(module)
+    return module
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / '__init__.py')))

build/torch211-cxx11-cu130-aarch64-linux/metadata.json

@@ -0,0 +1,22 @@
+{
+  "name": "flashrt-spatiotemporal-layout",
+  "id": "_flashrt_spatiotemporal_layout_cuda_7781728",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "4WRzvPydgNf2NK7ZljUlhO8NWc3qbONgjAzNb7YayKE=",
+      "_flashrt_spatiotemporal_layout_cuda_7781728.abi3.so": "7CoIxjM8aMm20SZBTABaj9g6ZJmoFBuQyMlqhCwQffQ=",
+      "_ops.py": "jbNHl3Z/6eFjOcwAfHyy7aALNepozLB2kHUYItLSAaE=",
+      "flashrt_spatiotemporal_layout/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  }
+}

build/torch211-cxx11-cu130-x86_64-linux/__init__.py

@@ -58,6 +58,137 @@ def _update_cache2_ncdhw_bf16_fake(cur: torch.Tensor, prev: torch.Tensor, out: t
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("channel_to_space3d_bf16"))
+def _channel_to_space3d_bf16_fake(
+    x: torch.Tensor,
+    out_channels: int,
+    temporal_factor: int,
+    spatial_factor: int,
+    repeats: int,
+    first_chunk: bool,
+    out: torch.Tensor,
+) -> None:
+    _check_ncdhw(x, "x")
+    out_t = x.shape[2] * temporal_factor - (
+        temporal_factor - 1 if first_chunk else 0
+    )
+    expected = (
+        x.shape[0],
+        out_channels,
+        out_t,
+        x.shape[3] * spatial_factor,
+        x.shape[4] * spatial_factor,
+    )
+    if (
+        out_channels <= 0
+        or temporal_factor <= 0
+        or spatial_factor <= 0
+        or repeats <= 0
+        or x.shape[1] * repeats
+        < out_channels * temporal_factor * spatial_factor * spatial_factor
+        or out.shape != expected
+    ):
+        raise RuntimeError("channel_to_space3d_bf16 shape contract failed")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("pack_causal_cache3_nhwc_bf16")
+)
+def _pack_causal_cache3_nhwc_bf16_fake(
+    previous: torch.Tensor, current: torch.Tensor, out: torch.Tensor
+) -> None:
+    _check_ncdhw(previous, "previous")
+    _check_ncdhw(current, "current")
+    b, c, t, h, w = current.shape
+    if (
+        t != 1
+        or previous.shape != (b, c, 2, h, w)
+        or out.shape != (b, h, w, 3 * c)
+    ):
+        raise RuntimeError("causal cache pack shape contract failed")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("avg_pool3d_channels_bf16"))
+def _avg_pool3d_channels_bf16_fake(
+    x: torch.Tensor,
+    out_channels: int,
+    factor_t: int,
+    factor_s: int,
+    group_size: int,
+    out: torch.Tensor,
+) -> None:
+    _check_ncdhw(x, "x")
+    expected = (
+        x.shape[0],
+        out_channels,
+        (x.shape[2] + factor_t - 1) // factor_t,
+        x.shape[3] // factor_s,
+        x.shape[4] // factor_s,
+    )
+    if out.shape != expected:
+        raise RuntimeError("out has the wrong pooled NCDHW shape")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_bf16"))
+def _ndhwc_to_ncdhw_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B, T, H, W, C)")
+    expected = (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (B, C, T, H, W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_bias_bf16"))
+def _ndhwc_to_ncdhw_bias_bf16_fake(
+    x: torch.Tensor, bias: torch.Tensor, out: torch.Tensor
+) -> None:
+    _ndhwc_to_ncdhw_bf16_fake(x, out)
+    if bias.shape != (x.shape[4],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_add_bf16"))
+def _ndhwc_to_ncdhw_add_bf16_fake(
+    x: torch.Tensor, residual: torch.Tensor, out: torch.Tensor
+) -> None:
+    _ndhwc_to_ncdhw_bf16_fake(x, out)
+    if residual.shape != out.shape:
+        raise RuntimeError("residual must match the NCDHW output shape")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("ncdhw_quantize_fp8_static_ndhwc_bf16")
+)
+def _ncdhw_quantize_fp8_static_ndhwc_bf16_fake(
+    x: torch.Tensor, scale: float, out: torch.Tensor
+) -> None:
+    _check_ncdhw(x, "x")
+    expected = (x.shape[0], x.shape[2], x.shape[3], x.shape[4], x.shape[1])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (B, T, H, W, C)")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("upsample2x_quantize_fp8_static_nhwc_bf16")
+)
+def _upsample2x_quantize_fp8_static_nhwc_bf16_fake(
+    x: torch.Tensor, scale: float, out: torch.Tensor
+) -> None:
+    if x.dim() != 4:
+        raise RuntimeError("x must have shape (N, C, H, W)")
+    expected = (x.shape[0], 2 * x.shape[2], 2 * x.shape[3], x.shape[1])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (N, 2H, 2W, C)")
+    return None
+
+
 def ncdhw_to_blc_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
     """Convert BF16 NCDHW tensor to BLC where L = T * H * W."""
 
@@ -101,10 +232,183 @@ def update_cache2_ncdhw_bf16(cur: torch.Tensor, prev: torch.Tensor, out: torch.T
     return out
 
 
+def channel_to_space3d_bf16(
+    x: torch.Tensor,
+    out_channels: int,
+    temporal_factor: int,
+    spatial_factor: int,
+    repeats: int = 1,
+    first_chunk: bool = False,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Move expanded channels into temporal/spatial dimensions."""
+
+    out_t = x.shape[2] * temporal_factor - (
+        temporal_factor - 1 if first_chunk else 0
+    )
+    if out is None:
+        out = torch.empty(
+            (
+                x.shape[0],
+                out_channels,
+                out_t,
+                x.shape[3] * spatial_factor,
+                x.shape[4] * spatial_factor,
+            ),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.channel_to_space3d_bf16(
+        x,
+        int(out_channels),
+        int(temporal_factor),
+        int(spatial_factor),
+        int(repeats),
+        bool(first_chunk),
+        out,
+    )
+    return out
+
+
+def pack_causal_cache3_nhwc_bf16(
+    previous: torch.Tensor,
+    current: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Pack two cached and one current NCDHW frame into NHWC with 3C."""
+
+    if out is None:
+        out = torch.empty(
+            (
+                current.shape[0],
+                current.shape[3],
+                current.shape[4],
+                3 * current.shape[1],
+            ),
+            device=current.device,
+            dtype=current.dtype,
+        )
+    ops.pack_causal_cache3_nhwc_bf16(previous, current, out)
+    return out
+
+
+def avg_pool3d_channels_bf16(
+    x: torch.Tensor,
+    out_channels: int,
+    factor_t: int,
+    factor_s: int,
+    group_size: int,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Pool NCDHW blocks while folding spatiotemporal factors into channels."""
+
+    if out is None:
+        out = torch.empty(
+            (
+                x.shape[0],
+                out_channels,
+                (x.shape[2] + factor_t - 1) // factor_t,
+                x.shape[3] // factor_s,
+                x.shape[4] // factor_s,
+            ),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.avg_pool3d_channels_bf16(
+        x, out_channels, factor_t, factor_s, group_size, out
+    )
+    return out
+
+
+def ndhwc_to_ncdhw_bf16(
+    x: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Convert contiguous BF16 NDHWC to contiguous BF16 NCDHW."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3]),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.ndhwc_to_ncdhw_bf16(x, out)
+    return out
+
+
+def ndhwc_to_ncdhw_bias_bf16(
+    x: torch.Tensor, bias: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Convert NDHWC to NCDHW and add a BF16 channel bias."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3]),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.ndhwc_to_ncdhw_bias_bf16(x, bias, out)
+    return out
+
+
+def ndhwc_to_ncdhw_add_bf16(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Convert NDHWC to NCDHW and add a BF16 NCDHW residual."""
+
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.ndhwc_to_ncdhw_add_bf16(x, residual, out)
+    return out
+
+
+def ncdhw_quantize_fp8_static_ndhwc_bf16(
+    x: torch.Tensor,
+    scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Quantize BF16 NCDHW to FP8 E4M3 NDHWC using a static scale."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[2], x.shape[3], x.shape[4], x.shape[1]),
+            device=x.device,
+            dtype=torch.float8_e4m3fn,
+        )
+    ops.ncdhw_quantize_fp8_static_ndhwc_bf16(x, scale, out)
+    return out
+
+
+def upsample2x_quantize_fp8_static_nhwc_bf16(
+    x: torch.Tensor,
+    scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Nearest-neighbor 2x upsample NCHW and emit static-scaled FP8 NHWC."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], 2 * x.shape[2], 2 * x.shape[3], x.shape[1]),
+            device=x.device,
+            dtype=torch.float8_e4m3fn,
+        )
+    ops.upsample2x_quantize_fp8_static_nhwc_bf16(x, scale, out)
+    return out
+
+
 __all__ = [
     "ncdhw_to_blc_bf16",
     "patch_im2col_bf16",
     "time_unshuffle2_bf16",
     "add_bias_ncdhw_bf16",
     "update_cache2_ncdhw_bf16",
+    "channel_to_space3d_bf16",
+    "pack_causal_cache3_nhwc_bf16",
+    "avg_pool3d_channels_bf16",
+    "ndhwc_to_ncdhw_bf16",
+    "ndhwc_to_ncdhw_bias_bf16",
+    "ndhwc_to_ncdhw_add_bf16",
+    "ncdhw_quantize_fp8_static_ndhwc_bf16",
+    "upsample2x_quantize_fp8_static_nhwc_bf16",
 ]

build/{torch212-cxx11-cu130-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_5de4768.abi3.so → torch211-cxx11-cu130-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_7781728.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c470dfb5ff53452ad0eddb0e5a4204dcdd151c6058e3b7b1ef7b77c765aaa91e
-size 702464
+oid sha256:4849a4c21396b8d1ca4173e837ea16491d7962e7720820d191008890de99ccf1
+size 1804176

build/torch211-cxx11-cu130-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _flashrt_spatiotemporal_layout_cuda_5de4768
-ops = torch.ops._flashrt_spatiotemporal_layout_cuda_5de4768
+from . import _flashrt_spatiotemporal_layout_cuda_7781728
+ops = torch.ops._flashrt_spatiotemporal_layout_cuda_7781728
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_flashrt_spatiotemporal_layout_cuda_5de4768::{op_name}"
+    return f"_flashrt_spatiotemporal_layout_cuda_7781728::{op_name}"

build/torch211-cxx11-cu130-x86_64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-spatiotemporal-layout",
-  "id": "_flashrt_spatiotemporal_layout_cuda_5de4768",
+  "id": "_flashrt_spatiotemporal_layout_cuda_7781728",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch212-cxx11-cu130-x86_64-linux/__init__.py

@@ -58,6 +58,137 @@ def _update_cache2_ncdhw_bf16_fake(cur: torch.Tensor, prev: torch.Tensor, out: t
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("channel_to_space3d_bf16"))
+def _channel_to_space3d_bf16_fake(
+    x: torch.Tensor,
+    out_channels: int,
+    temporal_factor: int,
+    spatial_factor: int,
+    repeats: int,
+    first_chunk: bool,
+    out: torch.Tensor,
+) -> None:
+    _check_ncdhw(x, "x")
+    out_t = x.shape[2] * temporal_factor - (
+        temporal_factor - 1 if first_chunk else 0
+    )
+    expected = (
+        x.shape[0],
+        out_channels,
+        out_t,
+        x.shape[3] * spatial_factor,
+        x.shape[4] * spatial_factor,
+    )
+    if (
+        out_channels <= 0
+        or temporal_factor <= 0
+        or spatial_factor <= 0
+        or repeats <= 0
+        or x.shape[1] * repeats
+        < out_channels * temporal_factor * spatial_factor * spatial_factor
+        or out.shape != expected
+    ):
+        raise RuntimeError("channel_to_space3d_bf16 shape contract failed")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("pack_causal_cache3_nhwc_bf16")
+)
+def _pack_causal_cache3_nhwc_bf16_fake(
+    previous: torch.Tensor, current: torch.Tensor, out: torch.Tensor
+) -> None:
+    _check_ncdhw(previous, "previous")
+    _check_ncdhw(current, "current")
+    b, c, t, h, w = current.shape
+    if (
+        t != 1
+        or previous.shape != (b, c, 2, h, w)
+        or out.shape != (b, h, w, 3 * c)
+    ):
+        raise RuntimeError("causal cache pack shape contract failed")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("avg_pool3d_channels_bf16"))
+def _avg_pool3d_channels_bf16_fake(
+    x: torch.Tensor,
+    out_channels: int,
+    factor_t: int,
+    factor_s: int,
+    group_size: int,
+    out: torch.Tensor,
+) -> None:
+    _check_ncdhw(x, "x")
+    expected = (
+        x.shape[0],
+        out_channels,
+        (x.shape[2] + factor_t - 1) // factor_t,
+        x.shape[3] // factor_s,
+        x.shape[4] // factor_s,
+    )
+    if out.shape != expected:
+        raise RuntimeError("out has the wrong pooled NCDHW shape")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_bf16"))
+def _ndhwc_to_ncdhw_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B, T, H, W, C)")
+    expected = (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (B, C, T, H, W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_bias_bf16"))
+def _ndhwc_to_ncdhw_bias_bf16_fake(
+    x: torch.Tensor, bias: torch.Tensor, out: torch.Tensor
+) -> None:
+    _ndhwc_to_ncdhw_bf16_fake(x, out)
+    if bias.shape != (x.shape[4],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_add_bf16"))
+def _ndhwc_to_ncdhw_add_bf16_fake(
+    x: torch.Tensor, residual: torch.Tensor, out: torch.Tensor
+) -> None:
+    _ndhwc_to_ncdhw_bf16_fake(x, out)
+    if residual.shape != out.shape:
+        raise RuntimeError("residual must match the NCDHW output shape")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("ncdhw_quantize_fp8_static_ndhwc_bf16")
+)
+def _ncdhw_quantize_fp8_static_ndhwc_bf16_fake(
+    x: torch.Tensor, scale: float, out: torch.Tensor
+) -> None:
+    _check_ncdhw(x, "x")
+    expected = (x.shape[0], x.shape[2], x.shape[3], x.shape[4], x.shape[1])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (B, T, H, W, C)")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("upsample2x_quantize_fp8_static_nhwc_bf16")
+)
+def _upsample2x_quantize_fp8_static_nhwc_bf16_fake(
+    x: torch.Tensor, scale: float, out: torch.Tensor
+) -> None:
+    if x.dim() != 4:
+        raise RuntimeError("x must have shape (N, C, H, W)")
+    expected = (x.shape[0], 2 * x.shape[2], 2 * x.shape[3], x.shape[1])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (N, 2H, 2W, C)")
+    return None
+
+
 def ncdhw_to_blc_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
     """Convert BF16 NCDHW tensor to BLC where L = T * H * W."""
 
@@ -101,10 +232,183 @@ def update_cache2_ncdhw_bf16(cur: torch.Tensor, prev: torch.Tensor, out: torch.T
     return out
 
 
+def channel_to_space3d_bf16(
+    x: torch.Tensor,
+    out_channels: int,
+    temporal_factor: int,
+    spatial_factor: int,
+    repeats: int = 1,
+    first_chunk: bool = False,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Move expanded channels into temporal/spatial dimensions."""
+
+    out_t = x.shape[2] * temporal_factor - (
+        temporal_factor - 1 if first_chunk else 0
+    )
+    if out is None:
+        out = torch.empty(
+            (
+                x.shape[0],
+                out_channels,
+                out_t,
+                x.shape[3] * spatial_factor,
+                x.shape[4] * spatial_factor,
+            ),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.channel_to_space3d_bf16(
+        x,
+        int(out_channels),
+        int(temporal_factor),
+        int(spatial_factor),
+        int(repeats),
+        bool(first_chunk),
+        out,
+    )
+    return out
+
+
+def pack_causal_cache3_nhwc_bf16(
+    previous: torch.Tensor,
+    current: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Pack two cached and one current NCDHW frame into NHWC with 3C."""
+
+    if out is None:
+        out = torch.empty(
+            (
+                current.shape[0],
+                current.shape[3],
+                current.shape[4],
+                3 * current.shape[1],
+            ),
+            device=current.device,
+            dtype=current.dtype,
+        )
+    ops.pack_causal_cache3_nhwc_bf16(previous, current, out)
+    return out
+
+
+def avg_pool3d_channels_bf16(
+    x: torch.Tensor,
+    out_channels: int,
+    factor_t: int,
+    factor_s: int,
+    group_size: int,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Pool NCDHW blocks while folding spatiotemporal factors into channels."""
+
+    if out is None:
+        out = torch.empty(
+            (
+                x.shape[0],
+                out_channels,
+                (x.shape[2] + factor_t - 1) // factor_t,
+                x.shape[3] // factor_s,
+                x.shape[4] // factor_s,
+            ),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.avg_pool3d_channels_bf16(
+        x, out_channels, factor_t, factor_s, group_size, out
+    )
+    return out
+
+
+def ndhwc_to_ncdhw_bf16(
+    x: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Convert contiguous BF16 NDHWC to contiguous BF16 NCDHW."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3]),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.ndhwc_to_ncdhw_bf16(x, out)
+    return out
+
+
+def ndhwc_to_ncdhw_bias_bf16(
+    x: torch.Tensor, bias: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Convert NDHWC to NCDHW and add a BF16 channel bias."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3]),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.ndhwc_to_ncdhw_bias_bf16(x, bias, out)
+    return out
+
+
+def ndhwc_to_ncdhw_add_bf16(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Convert NDHWC to NCDHW and add a BF16 NCDHW residual."""
+
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.ndhwc_to_ncdhw_add_bf16(x, residual, out)
+    return out
+
+
+def ncdhw_quantize_fp8_static_ndhwc_bf16(
+    x: torch.Tensor,
+    scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Quantize BF16 NCDHW to FP8 E4M3 NDHWC using a static scale."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[2], x.shape[3], x.shape[4], x.shape[1]),
+            device=x.device,
+            dtype=torch.float8_e4m3fn,
+        )
+    ops.ncdhw_quantize_fp8_static_ndhwc_bf16(x, scale, out)
+    return out
+
+
+def upsample2x_quantize_fp8_static_nhwc_bf16(
+    x: torch.Tensor,
+    scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Nearest-neighbor 2x upsample NCHW and emit static-scaled FP8 NHWC."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], 2 * x.shape[2], 2 * x.shape[3], x.shape[1]),
+            device=x.device,
+            dtype=torch.float8_e4m3fn,
+        )
+    ops.upsample2x_quantize_fp8_static_nhwc_bf16(x, scale, out)
+    return out
+
+
 __all__ = [
     "ncdhw_to_blc_bf16",
     "patch_im2col_bf16",
     "time_unshuffle2_bf16",
     "add_bias_ncdhw_bf16",
     "update_cache2_ncdhw_bf16",
+    "channel_to_space3d_bf16",
+    "pack_causal_cache3_nhwc_bf16",
+    "avg_pool3d_channels_bf16",
+    "ndhwc_to_ncdhw_bf16",
+    "ndhwc_to_ncdhw_bias_bf16",
+    "ndhwc_to_ncdhw_add_bf16",
+    "ncdhw_quantize_fp8_static_ndhwc_bf16",
+    "upsample2x_quantize_fp8_static_nhwc_bf16",
 ]

build/torch212-cxx11-cu130-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_7781728.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:917fb550fc3481af63a6c63aa6ad864c76f76c61ce59e9236769c3f08b2e54b6
+size 1810440

build/torch212-cxx11-cu130-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _flashrt_spatiotemporal_layout_cuda_5de4768
-ops = torch.ops._flashrt_spatiotemporal_layout_cuda_5de4768
+from . import _flashrt_spatiotemporal_layout_cuda_7781728
+ops = torch.ops._flashrt_spatiotemporal_layout_cuda_7781728
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_flashrt_spatiotemporal_layout_cuda_5de4768::{op_name}"
+    return f"_flashrt_spatiotemporal_layout_cuda_7781728::{op_name}"

build/torch212-cxx11-cu130-x86_64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-spatiotemporal-layout",
-  "id": "_flashrt_spatiotemporal_layout_cuda_5de4768",
+  "id": "_flashrt_spatiotemporal_layout_cuda_7781728",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch212-cxx11-cu132-x86_64-linux/__init__.py

@@ -58,6 +58,137 @@ def _update_cache2_ncdhw_bf16_fake(cur: torch.Tensor, prev: torch.Tensor, out: t
     return None
 
 
+@torch.library.register_fake(add_op_namespace_prefix("channel_to_space3d_bf16"))
+def _channel_to_space3d_bf16_fake(
+    x: torch.Tensor,
+    out_channels: int,
+    temporal_factor: int,
+    spatial_factor: int,
+    repeats: int,
+    first_chunk: bool,
+    out: torch.Tensor,
+) -> None:
+    _check_ncdhw(x, "x")
+    out_t = x.shape[2] * temporal_factor - (
+        temporal_factor - 1 if first_chunk else 0
+    )
+    expected = (
+        x.shape[0],
+        out_channels,
+        out_t,
+        x.shape[3] * spatial_factor,
+        x.shape[4] * spatial_factor,
+    )
+    if (
+        out_channels <= 0
+        or temporal_factor <= 0
+        or spatial_factor <= 0
+        or repeats <= 0
+        or x.shape[1] * repeats
+        < out_channels * temporal_factor * spatial_factor * spatial_factor
+        or out.shape != expected
+    ):
+        raise RuntimeError("channel_to_space3d_bf16 shape contract failed")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("pack_causal_cache3_nhwc_bf16")
+)
+def _pack_causal_cache3_nhwc_bf16_fake(
+    previous: torch.Tensor, current: torch.Tensor, out: torch.Tensor
+) -> None:
+    _check_ncdhw(previous, "previous")
+    _check_ncdhw(current, "current")
+    b, c, t, h, w = current.shape
+    if (
+        t != 1
+        or previous.shape != (b, c, 2, h, w)
+        or out.shape != (b, h, w, 3 * c)
+    ):
+        raise RuntimeError("causal cache pack shape contract failed")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("avg_pool3d_channels_bf16"))
+def _avg_pool3d_channels_bf16_fake(
+    x: torch.Tensor,
+    out_channels: int,
+    factor_t: int,
+    factor_s: int,
+    group_size: int,
+    out: torch.Tensor,
+) -> None:
+    _check_ncdhw(x, "x")
+    expected = (
+        x.shape[0],
+        out_channels,
+        (x.shape[2] + factor_t - 1) // factor_t,
+        x.shape[3] // factor_s,
+        x.shape[4] // factor_s,
+    )
+    if out.shape != expected:
+        raise RuntimeError("out has the wrong pooled NCDHW shape")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_bf16"))
+def _ndhwc_to_ncdhw_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B, T, H, W, C)")
+    expected = (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (B, C, T, H, W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_bias_bf16"))
+def _ndhwc_to_ncdhw_bias_bf16_fake(
+    x: torch.Tensor, bias: torch.Tensor, out: torch.Tensor
+) -> None:
+    _ndhwc_to_ncdhw_bf16_fake(x, out)
+    if bias.shape != (x.shape[4],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_add_bf16"))
+def _ndhwc_to_ncdhw_add_bf16_fake(
+    x: torch.Tensor, residual: torch.Tensor, out: torch.Tensor
+) -> None:
+    _ndhwc_to_ncdhw_bf16_fake(x, out)
+    if residual.shape != out.shape:
+        raise RuntimeError("residual must match the NCDHW output shape")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("ncdhw_quantize_fp8_static_ndhwc_bf16")
+)
+def _ncdhw_quantize_fp8_static_ndhwc_bf16_fake(
+    x: torch.Tensor, scale: float, out: torch.Tensor
+) -> None:
+    _check_ncdhw(x, "x")
+    expected = (x.shape[0], x.shape[2], x.shape[3], x.shape[4], x.shape[1])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (B, T, H, W, C)")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("upsample2x_quantize_fp8_static_nhwc_bf16")
+)
+def _upsample2x_quantize_fp8_static_nhwc_bf16_fake(
+    x: torch.Tensor, scale: float, out: torch.Tensor
+) -> None:
+    if x.dim() != 4:
+        raise RuntimeError("x must have shape (N, C, H, W)")
+    expected = (x.shape[0], 2 * x.shape[2], 2 * x.shape[3], x.shape[1])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (N, 2H, 2W, C)")
+    return None
+
+
 def ncdhw_to_blc_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
     """Convert BF16 NCDHW tensor to BLC where L = T * H * W."""
 
@@ -101,10 +232,183 @@ def update_cache2_ncdhw_bf16(cur: torch.Tensor, prev: torch.Tensor, out: torch.T
     return out
 
 
+def channel_to_space3d_bf16(
+    x: torch.Tensor,
+    out_channels: int,
+    temporal_factor: int,
+    spatial_factor: int,
+    repeats: int = 1,
+    first_chunk: bool = False,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Move expanded channels into temporal/spatial dimensions."""
+
+    out_t = x.shape[2] * temporal_factor - (
+        temporal_factor - 1 if first_chunk else 0
+    )
+    if out is None:
+        out = torch.empty(
+            (
+                x.shape[0],
+                out_channels,
+                out_t,
+                x.shape[3] * spatial_factor,
+                x.shape[4] * spatial_factor,
+            ),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.channel_to_space3d_bf16(
+        x,
+        int(out_channels),
+        int(temporal_factor),
+        int(spatial_factor),
+        int(repeats),
+        bool(first_chunk),
+        out,
+    )
+    return out
+
+
+def pack_causal_cache3_nhwc_bf16(
+    previous: torch.Tensor,
+    current: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Pack two cached and one current NCDHW frame into NHWC with 3C."""
+
+    if out is None:
+        out = torch.empty(
+            (
+                current.shape[0],
+                current.shape[3],
+                current.shape[4],
+                3 * current.shape[1],
+            ),
+            device=current.device,
+            dtype=current.dtype,
+        )
+    ops.pack_causal_cache3_nhwc_bf16(previous, current, out)
+    return out
+
+
+def avg_pool3d_channels_bf16(
+    x: torch.Tensor,
+    out_channels: int,
+    factor_t: int,
+    factor_s: int,
+    group_size: int,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Pool NCDHW blocks while folding spatiotemporal factors into channels."""
+
+    if out is None:
+        out = torch.empty(
+            (
+                x.shape[0],
+                out_channels,
+                (x.shape[2] + factor_t - 1) // factor_t,
+                x.shape[3] // factor_s,
+                x.shape[4] // factor_s,
+            ),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.avg_pool3d_channels_bf16(
+        x, out_channels, factor_t, factor_s, group_size, out
+    )
+    return out
+
+
+def ndhwc_to_ncdhw_bf16(
+    x: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Convert contiguous BF16 NDHWC to contiguous BF16 NCDHW."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3]),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.ndhwc_to_ncdhw_bf16(x, out)
+    return out
+
+
+def ndhwc_to_ncdhw_bias_bf16(
+    x: torch.Tensor, bias: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Convert NDHWC to NCDHW and add a BF16 channel bias."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3]),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.ndhwc_to_ncdhw_bias_bf16(x, bias, out)
+    return out
+
+
+def ndhwc_to_ncdhw_add_bf16(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Convert NDHWC to NCDHW and add a BF16 NCDHW residual."""
+
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.ndhwc_to_ncdhw_add_bf16(x, residual, out)
+    return out
+
+
+def ncdhw_quantize_fp8_static_ndhwc_bf16(
+    x: torch.Tensor,
+    scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Quantize BF16 NCDHW to FP8 E4M3 NDHWC using a static scale."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[2], x.shape[3], x.shape[4], x.shape[1]),
+            device=x.device,
+            dtype=torch.float8_e4m3fn,
+        )
+    ops.ncdhw_quantize_fp8_static_ndhwc_bf16(x, scale, out)
+    return out
+
+
+def upsample2x_quantize_fp8_static_nhwc_bf16(
+    x: torch.Tensor,
+    scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Nearest-neighbor 2x upsample NCHW and emit static-scaled FP8 NHWC."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], 2 * x.shape[2], 2 * x.shape[3], x.shape[1]),
+            device=x.device,
+            dtype=torch.float8_e4m3fn,
+        )
+    ops.upsample2x_quantize_fp8_static_nhwc_bf16(x, scale, out)
+    return out
+
+
 __all__ = [
     "ncdhw_to_blc_bf16",
     "patch_im2col_bf16",
     "time_unshuffle2_bf16",
     "add_bias_ncdhw_bf16",
     "update_cache2_ncdhw_bf16",
+    "channel_to_space3d_bf16",
+    "pack_causal_cache3_nhwc_bf16",
+    "avg_pool3d_channels_bf16",
+    "ndhwc_to_ncdhw_bf16",
+    "ndhwc_to_ncdhw_bias_bf16",
+    "ndhwc_to_ncdhw_add_bf16",
+    "ncdhw_quantize_fp8_static_ndhwc_bf16",
+    "upsample2x_quantize_fp8_static_nhwc_bf16",
 ]

build/torch212-cxx11-cu132-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_7781728.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e671b827971af10d928dcfc1cfea4676931152975e33829a1e6cf3743cb746e
+size 1818632

build/torch212-cxx11-cu132-x86_64-linux/_ops.py

@@ -1,9 +1,9 @@
 import torch
-from . import _flashrt_spatiotemporal_layout_cuda_5de4768
-ops = torch.ops._flashrt_spatiotemporal_layout_cuda_5de4768
+from . import _flashrt_spatiotemporal_layout_cuda_7781728
+ops = torch.ops._flashrt_spatiotemporal_layout_cuda_7781728
 
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_flashrt_spatiotemporal_layout_cuda_5de4768::{op_name}"
+    return f"_flashrt_spatiotemporal_layout_cuda_7781728::{op_name}"

build/torch212-cxx11-cu132-x86_64-linux/metadata.json

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-spatiotemporal-layout",
-  "id": "_flashrt_spatiotemporal_layout_cuda_5de4768",
+  "id": "_flashrt_spatiotemporal_layout_cuda_7781728",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

build/torch213-cxx11-cu130-aarch64-linux/__init__.py

@@ -0,0 +1,414 @@
+"""FlashRT spatiotemporal layout kernels."""
+
+from __future__ import annotations
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+def _check_ncdhw(x: torch.Tensor, name: str) -> None:
+    if x.dim() != 5:
+        raise RuntimeError(f"{name} must have shape (B, C, T, H, W)")
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ncdhw_to_blc_bf16"))
+def _ncdhw_to_blc_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c, t, h, w = x.shape
+    if out.shape != (b, t * h * w, c):
+        raise RuntimeError("out must have shape (B, T * H * W, C)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("patch_im2col_bf16"))
+def _patch_im2col_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    if x.dim() != 4 or x.shape[1:] != (224, 224, 3):
+        raise RuntimeError("x must have shape (num_views, 224, 224, 3)")
+    if out.shape != (x.shape[0] * 256, 588):
+        raise RuntimeError("out must have shape (num_views * 256, 588)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("time_unshuffle2_bf16"))
+def _time_unshuffle2_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    b, c2, t, h, w = x.shape
+    if c2 % 2 != 0:
+        raise RuntimeError("x.shape[1] must be even")
+    if out.shape != (b, c2 // 2, 2 * t, h, w):
+        raise RuntimeError("out must have shape (B, C / 2, 2 * T, H, W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("add_bias_ncdhw_bf16"))
+def _add_bias_ncdhw_bf16_fake(x: torch.Tensor, bias: torch.Tensor) -> None:
+    _check_ncdhw(x, "x")
+    if bias.shape != (x.shape[1],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("update_cache2_ncdhw_bf16"))
+def _update_cache2_ncdhw_bf16_fake(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor) -> None:
+    _check_ncdhw(cur, "cur")
+    b, c, _, h, w = cur.shape
+    if prev.shape != (b, c, 2, h, w) or out.shape != (b, c, 2, h, w):
+        raise RuntimeError("prev and out must have shape (B, C, 2, H, W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("channel_to_space3d_bf16"))
+def _channel_to_space3d_bf16_fake(
+    x: torch.Tensor,
+    out_channels: int,
+    temporal_factor: int,
+    spatial_factor: int,
+    repeats: int,
+    first_chunk: bool,
+    out: torch.Tensor,
+) -> None:
+    _check_ncdhw(x, "x")
+    out_t = x.shape[2] * temporal_factor - (
+        temporal_factor - 1 if first_chunk else 0
+    )
+    expected = (
+        x.shape[0],
+        out_channels,
+        out_t,
+        x.shape[3] * spatial_factor,
+        x.shape[4] * spatial_factor,
+    )
+    if (
+        out_channels <= 0
+        or temporal_factor <= 0
+        or spatial_factor <= 0
+        or repeats <= 0
+        or x.shape[1] * repeats
+        < out_channels * temporal_factor * spatial_factor * spatial_factor
+        or out.shape != expected
+    ):
+        raise RuntimeError("channel_to_space3d_bf16 shape contract failed")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("pack_causal_cache3_nhwc_bf16")
+)
+def _pack_causal_cache3_nhwc_bf16_fake(
+    previous: torch.Tensor, current: torch.Tensor, out: torch.Tensor
+) -> None:
+    _check_ncdhw(previous, "previous")
+    _check_ncdhw(current, "current")
+    b, c, t, h, w = current.shape
+    if (
+        t != 1
+        or previous.shape != (b, c, 2, h, w)
+        or out.shape != (b, h, w, 3 * c)
+    ):
+        raise RuntimeError("causal cache pack shape contract failed")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("avg_pool3d_channels_bf16"))
+def _avg_pool3d_channels_bf16_fake(
+    x: torch.Tensor,
+    out_channels: int,
+    factor_t: int,
+    factor_s: int,
+    group_size: int,
+    out: torch.Tensor,
+) -> None:
+    _check_ncdhw(x, "x")
+    expected = (
+        x.shape[0],
+        out_channels,
+        (x.shape[2] + factor_t - 1) // factor_t,
+        x.shape[3] // factor_s,
+        x.shape[4] // factor_s,
+    )
+    if out.shape != expected:
+        raise RuntimeError("out has the wrong pooled NCDHW shape")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_bf16"))
+def _ndhwc_to_ncdhw_bf16_fake(x: torch.Tensor, out: torch.Tensor) -> None:
+    if x.dim() != 5:
+        raise RuntimeError("x must have shape (B, T, H, W, C)")
+    expected = (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (B, C, T, H, W)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_bias_bf16"))
+def _ndhwc_to_ncdhw_bias_bf16_fake(
+    x: torch.Tensor, bias: torch.Tensor, out: torch.Tensor
+) -> None:
+    _ndhwc_to_ncdhw_bf16_fake(x, out)
+    if bias.shape != (x.shape[4],):
+        raise RuntimeError("bias must have shape (C,)")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("ndhwc_to_ncdhw_add_bf16"))
+def _ndhwc_to_ncdhw_add_bf16_fake(
+    x: torch.Tensor, residual: torch.Tensor, out: torch.Tensor
+) -> None:
+    _ndhwc_to_ncdhw_bf16_fake(x, out)
+    if residual.shape != out.shape:
+        raise RuntimeError("residual must match the NCDHW output shape")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("ncdhw_quantize_fp8_static_ndhwc_bf16")
+)
+def _ncdhw_quantize_fp8_static_ndhwc_bf16_fake(
+    x: torch.Tensor, scale: float, out: torch.Tensor
+) -> None:
+    _check_ncdhw(x, "x")
+    expected = (x.shape[0], x.shape[2], x.shape[3], x.shape[4], x.shape[1])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (B, T, H, W, C)")
+    return None
+
+
+@torch.library.register_fake(
+    add_op_namespace_prefix("upsample2x_quantize_fp8_static_nhwc_bf16")
+)
+def _upsample2x_quantize_fp8_static_nhwc_bf16_fake(
+    x: torch.Tensor, scale: float, out: torch.Tensor
+) -> None:
+    if x.dim() != 4:
+        raise RuntimeError("x must have shape (N, C, H, W)")
+    expected = (x.shape[0], 2 * x.shape[2], 2 * x.shape[3], x.shape[1])
+    if out.shape != expected:
+        raise RuntimeError("out must have shape (N, 2H, 2W, C)")
+    return None
+
+
+def ncdhw_to_blc_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 NCDHW tensor to BLC where L = T * H * W."""
+
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[2] * x.shape[3] * x.shape[4], x.shape[1]), device=x.device, dtype=x.dtype)
+    ops.ncdhw_to_blc_bf16(x, out)
+    return out
+
+
+def patch_im2col_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 NHWC images into flattened 14x14 patch rows."""
+
+    if out is None:
+        out = torch.empty((x.shape[0] * 256, 588), device=x.device, dtype=x.dtype)
+    ops.patch_im2col_bf16(x, out)
+    return out
+
+
+def time_unshuffle2_bf16(x: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Convert BF16 (B, 2C, T, H, W) into (B, C, 2T, H, W)."""
+
+    if out is None:
+        out = torch.empty((x.shape[0], x.shape[1] // 2, 2 * x.shape[2], x.shape[3], x.shape[4]), device=x.device, dtype=x.dtype)
+    ops.time_unshuffle2_bf16(x, out)
+    return out
+
+
+def add_bias_ncdhw_bf16(x: torch.Tensor, bias: torch.Tensor) -> torch.Tensor:
+    """Add channel bias to an NCDHW tensor in place."""
+
+    ops.add_bias_ncdhw_bf16(x, bias)
+    return x
+
+
+def update_cache2_ncdhw_bf16(cur: torch.Tensor, prev: torch.Tensor, out: torch.Tensor | None = None) -> torch.Tensor:
+    """Update two-frame latent cache from current and previous NCDHW tensors."""
+
+    if out is None:
+        out = torch.empty((cur.shape[0], cur.shape[1], 2, cur.shape[3], cur.shape[4]), device=cur.device, dtype=cur.dtype)
+    ops.update_cache2_ncdhw_bf16(cur, prev, out)
+    return out
+
+
+def channel_to_space3d_bf16(
+    x: torch.Tensor,
+    out_channels: int,
+    temporal_factor: int,
+    spatial_factor: int,
+    repeats: int = 1,
+    first_chunk: bool = False,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Move expanded channels into temporal/spatial dimensions."""
+
+    out_t = x.shape[2] * temporal_factor - (
+        temporal_factor - 1 if first_chunk else 0
+    )
+    if out is None:
+        out = torch.empty(
+            (
+                x.shape[0],
+                out_channels,
+                out_t,
+                x.shape[3] * spatial_factor,
+                x.shape[4] * spatial_factor,
+            ),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.channel_to_space3d_bf16(
+        x,
+        int(out_channels),
+        int(temporal_factor),
+        int(spatial_factor),
+        int(repeats),
+        bool(first_chunk),
+        out,
+    )
+    return out
+
+
+def pack_causal_cache3_nhwc_bf16(
+    previous: torch.Tensor,
+    current: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Pack two cached and one current NCDHW frame into NHWC with 3C."""
+
+    if out is None:
+        out = torch.empty(
+            (
+                current.shape[0],
+                current.shape[3],
+                current.shape[4],
+                3 * current.shape[1],
+            ),
+            device=current.device,
+            dtype=current.dtype,
+        )
+    ops.pack_causal_cache3_nhwc_bf16(previous, current, out)
+    return out
+
+
+def avg_pool3d_channels_bf16(
+    x: torch.Tensor,
+    out_channels: int,
+    factor_t: int,
+    factor_s: int,
+    group_size: int,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Pool NCDHW blocks while folding spatiotemporal factors into channels."""
+
+    if out is None:
+        out = torch.empty(
+            (
+                x.shape[0],
+                out_channels,
+                (x.shape[2] + factor_t - 1) // factor_t,
+                x.shape[3] // factor_s,
+                x.shape[4] // factor_s,
+            ),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.avg_pool3d_channels_bf16(
+        x, out_channels, factor_t, factor_s, group_size, out
+    )
+    return out
+
+
+def ndhwc_to_ncdhw_bf16(
+    x: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Convert contiguous BF16 NDHWC to contiguous BF16 NCDHW."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3]),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.ndhwc_to_ncdhw_bf16(x, out)
+    return out
+
+
+def ndhwc_to_ncdhw_bias_bf16(
+    x: torch.Tensor, bias: torch.Tensor, out: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Convert NDHWC to NCDHW and add a BF16 channel bias."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[4], x.shape[1], x.shape[2], x.shape[3]),
+            device=x.device,
+            dtype=x.dtype,
+        )
+    ops.ndhwc_to_ncdhw_bias_bf16(x, bias, out)
+    return out
+
+
+def ndhwc_to_ncdhw_add_bf16(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Convert NDHWC to NCDHW and add a BF16 NCDHW residual."""
+
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.ndhwc_to_ncdhw_add_bf16(x, residual, out)
+    return out
+
+
+def ncdhw_quantize_fp8_static_ndhwc_bf16(
+    x: torch.Tensor,
+    scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Quantize BF16 NCDHW to FP8 E4M3 NDHWC using a static scale."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], x.shape[2], x.shape[3], x.shape[4], x.shape[1]),
+            device=x.device,
+            dtype=torch.float8_e4m3fn,
+        )
+    ops.ncdhw_quantize_fp8_static_ndhwc_bf16(x, scale, out)
+    return out
+
+
+def upsample2x_quantize_fp8_static_nhwc_bf16(
+    x: torch.Tensor,
+    scale: float,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Nearest-neighbor 2x upsample NCHW and emit static-scaled FP8 NHWC."""
+
+    if out is None:
+        out = torch.empty(
+            (x.shape[0], 2 * x.shape[2], 2 * x.shape[3], x.shape[1]),
+            device=x.device,
+            dtype=torch.float8_e4m3fn,
+        )
+    ops.upsample2x_quantize_fp8_static_nhwc_bf16(x, scale, out)
+    return out
+
+
+__all__ = [
+    "ncdhw_to_blc_bf16",
+    "patch_im2col_bf16",
+    "time_unshuffle2_bf16",
+    "add_bias_ncdhw_bf16",
+    "update_cache2_ncdhw_bf16",
+    "channel_to_space3d_bf16",
+    "pack_causal_cache3_nhwc_bf16",
+    "avg_pool3d_channels_bf16",
+    "ndhwc_to_ncdhw_bf16",
+    "ndhwc_to_ncdhw_bias_bf16",
+    "ndhwc_to_ncdhw_add_bf16",
+    "ncdhw_quantize_fp8_static_ndhwc_bf16",
+    "upsample2x_quantize_fp8_static_nhwc_bf16",
+]

build/torch213-cxx11-cu130-aarch64-linux/_ops.py

@@ -0,0 +1,6 @@
+import torch
+from . import flashrt_spatiotemporal_layout_test
+ops = torch.ops.flashrt_spatiotemporal_layout_test
+
+def add_op_namespace_prefix(op_name: str):
+    return f"flashrt_spatiotemporal_layout_test::{op_name}"

build/torch213-cxx11-cu130-aarch64-linux/flashrt_spatiotemporal_layout/__init__.py

@@ -0,0 +1,14 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+
+def _import_from_path(file_path: Path):
+    path_hash = '{:x}'.format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    spec = importlib.util.spec_from_file_location(path_hash, file_path)
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[path_hash] = module
+    spec.loader.exec_module(module)
+    return module
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / '__init__.py')))

build/{torch211-cxx11-cu128-x86_64-linux/_flashrt_spatiotemporal_layout_cuda_5de4768.abi3.so → torch213-cxx11-cu130-aarch64-linux/flashrt_spatiotemporal_layout_test.abi3.so}

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:280efdb5e795d088a9d44e031b64b3d652a70cb070325579c2623690b8953b3d
-size 727528
+oid sha256:16f0b72634927e102b0ce7c17a8addd5eaed791906b19d034f2729cf909e895a
+size 461408

build/torch213-cxx11-cu130-aarch64-linux/metadata.json

@@ -0,0 +1,32 @@
+{
+  "name": "flashrt-spatiotemporal-layout",
+  "id": "flashrt_spatiotemporal_layout_test",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "4WRzvPydgNf2NK7ZljUlhO8NWc3qbONgjAzNb7YayKE=",
+      "flashrt_spatiotemporal_layout_test.abi3.so": "FvC3JjSSfhArDOfBeord1erteRkGsZ0DTycpz5CeiVo=",
+      "_ops.py": "P29ghFUoONBA3NS6cY+kzbwr/RczjL8AdMewXrgZDCA=",
+      "flashrt_spatiotemporal_layout/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "456d297",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.13.0+cu130",
+      "cuda": "13.0"
+    }
+  }
+}