liangsu9988 commited on about 6 hours ago

Commit

287dee9

verified ·

1 Parent(s): 6e5f010

Uploaded using `kernel-builder`.

Browse files

Files changed (17) hide show

benchmarks/benchmark_native_parity.py +520 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +331 -1
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu128-x86_64-linux/{_world_model_conv_cuda_f14c443.abi3.so → _world_model_conv_cuda_7781728.abi3.so} +2 -2
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +1 -1
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +331 -1
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu130-x86_64-linux/{_world_model_conv_cuda_f14c443.abi3.so → _world_model_conv_cuda_7781728.abi3.so} +2 -2
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +1 -1
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +331 -1
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu130-x86_64-linux/{_world_model_conv_cuda_f14c443.abi3.so → _world_model_conv_cuda_7781728.abi3.so} +2 -2
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +1 -1
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +331 -1
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu132-x86_64-linux/{_world_model_conv_cuda_f14c443.abi3.so → _world_model_conv_cuda_7781728.abi3.so} +2 -2
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +1 -1

benchmarks/benchmark_native_parity.py ADDED Viewed

	@@ -0,0 +1,520 @@

+#!/usr/bin/env python3
+"""World-model Conv benchmark with native, wrapper, compile, and cuDNN paths."""
+from __future__ import annotations
+import argparse
+import json
+import os
+import sys
+from dataclasses import asdict, dataclass
+from pathlib import Path
+import torch
+import torch.nn.functional as F
+PACKAGE = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PACKAGE / "tests"))
+from test_world_model_conv import (  # noqa: E402
+    dequantize_linear_nvfp4,
+    load_installed_ops,
+    load_source_ops,
+    quantize_conv_tensor,
+)
+CONV3D_SHAPES = {
+    "causal-c32": (1, 2, 4, 16, 16, 32, 32),
+    "causal-small": (1, 2, 4, 16, 16, 64, 64),
+}
+NVFP4_CONV3D_SHAPES = {
+    "nvfp4-c64": (1, 2, 4, 16, 16, 64, 64),
+    "nvfp4-c128": (1, 2, 4, 16, 16, 128, 128),
+    "nvfp4-c512": (1, 2, 4, 16, 16, 512, 512),
+}
+CONV2D_SHAPES = {
+    "resample-c64": (4, 32, 32, 64, 64),
+    "resample-c320": (17, 32, 32, 320, 320),
+}
+@dataclass
+class Result:
+    workload: str
+    shape: str
+    native_us: float
+    wrapper_us: float
+    wrapper_native: float
+    eager_cudnn_us: float
+    compile_cudnn_us: float
+    diagnostic_predequant_cudnn_us: float | None
+    diagnostic_predequant_compile_us: float | None
+    max_abs: float
+    mean_abs: float
+    p99_abs: float
+    cosine: float
+    accepted: bool
+def bench(fn, warmup, iters):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iters):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) * 1000.0 / iters
+def build_native():
+    from torch.utils.cpp_extension import load
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0a")
+    return load(
+        name="world_model_conv_raw_native",
+        sources=[
+            str(PACKAGE / "benchmarks/native_binding.cpp"),
+            str(PACKAGE / "csrc/fp8_conv3d_sm120_v18.cu"),
+            str(PACKAGE / "csrc/fp8_causal_conv3d_sm120.cu"),
+            str(PACKAGE / "csrc/fp8_conv2d_3x3_sm120.cu"),
+            str(PACKAGE / "csrc/nvfp4_causal_conv3d_sm120.cu"),
+            str(PACKAGE / "csrc/nvfp4_causal_conv3d_residual_sm120.cu"),
+            str(PACKAGE / "csrc/nvfp4_causal_conv3d_residual_k128_sm120.cu"),
+        ],
+        extra_include_paths=[str(PACKAGE / "csrc")],
+        extra_cflags=["-O3"],
+        extra_cuda_cflags=["-O3"],
+        verbose=False,
+    )
+def metrics(got, ref):
+    diff = (got.float() - ref.float()).abs().flatten()
+    cosine = F.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 source_call(ops, name, *args, out):
+    if hasattr(ops, "_ops"):
+        getattr(ops._ops, name)(*args, out)
+    else:
+        getattr(ops, name)(*args, out=out)
+def run_conv3d(ops, native, label, shape, args):
+    n, tc, tn, h, w, ci, co = shape
+    cache = (torch.randn((n, tc, h, w, ci), device="cuda") * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    new = (torch.randn((n, tn, h, w, ci), device="cuda") * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    weight = (torch.randn((co, 3, 3, 3, ci), device="cuda") * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    bias = (torch.randn(co, device="cuda") * 0.01).to(torch.bfloat16)
+    out = torch.empty((n, tn, h, w, co), device="cuda", dtype=torch.bfloat16)
+    alpha = 0.75
+    wrapper = lambda: source_call(
+        ops,
+        "fp8_causal_conv3d_ndhwc_bf16",
+        cache,
+        new,
+        weight,
+        bias,
+        alpha,
+        out=out,
+    )
+    raw = lambda: native.causal_conv3d(
+        cache, new, weight, bias, alpha, out
+    )
+    def cudnn_ref():
+        x = torch.cat((cache, new), dim=1).float().permute(0, 4, 1, 2, 3)
+        wt = weight.float().permute(0, 4, 1, 2, 3)
+        y = F.conv3d(x, wt, padding=(0, 1, 1))
+        y = y.mul(alpha).add(bias.float().view(1, -1, 1, 1, 1))
+        out.copy_(y[:, :, :tn].permute(0, 2, 3, 4, 1).to(torch.bfloat16))
+    compiled = torch.compile(cudnn_ref, fullgraph=True)
+    wrapper()
+    got = out.clone()
+    cudnn_ref()
+    ref = out.clone()
+    max_abs, mean_abs, p99_abs, cosine = metrics(got, ref)
+    native_us = bench(raw, args.warmup, args.iters)
+    wrapper_us = bench(wrapper, args.warmup, args.iters)
+    eager_us = bench(cudnn_ref, args.warmup, args.iters)
+    compile_us = bench(compiled, args.warmup, args.iters)
+    return Result(
+        label,
+        str(shape),
+        native_us,
+        wrapper_us,
+        wrapper_us / native_us,
+        eager_us,
+        compile_us,
+        None,
+        None,
+        max_abs,
+        mean_abs,
+        p99_abs,
+        cosine,
+        wrapper_us - native_us <= max(0.5, native_us * 0.05)
+        and wrapper_us <= min(eager_us, compile_us) * 0.98
+        and cosine >= 0.999
+        and mean_abs <= 0.01,
+    )
+def run_conv3d_residual(ops, native, label, shape, args):
+    n, tc, tn, h, w, ci, co = shape
+    cache = (torch.randn((n, tc, h, w, ci), device="cuda") * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    new = (torch.randn((n, tn, h, w, ci), device="cuda") * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    weight = (torch.randn((co, 3, 3, 3, ci), device="cuda") * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    bias = (torch.randn(co, device="cuda") * 0.01).to(torch.bfloat16)
+    residual = torch.randn(
+        (n, co, tn, h, w), device="cuda", dtype=torch.bfloat16
+    )
+    out = torch.empty_like(residual)
+    alpha = 0.75
+    wrapper = lambda: source_call(
+        ops,
+        "fp8_conv3d_v18_ncdhw_res_bf16out",
+        cache,
+        new,
+        weight,
+        bias,
+        residual,
+        alpha,
+        out=out,
+    )
+    raw = lambda: native.causal_conv3d_residual(
+        cache, new, weight, bias, residual, alpha, out
+    )
+    def cudnn_ref():
+        x = torch.cat((cache, new), dim=1).float().permute(0, 4, 1, 2, 3)
+        wt = weight.float().permute(0, 4, 1, 2, 3)
+        y = F.conv3d(x, wt, padding=(0, 1, 1))
+        y = y.mul(alpha).add(bias.float().view(1, -1, 1, 1, 1))
+        y = (
+            y[:, :, :tn].to(torch.bfloat16).float()
+            + residual.float()
+        ).to(torch.bfloat16)
+        out.copy_(y)
+    compiled = torch.compile(cudnn_ref, fullgraph=True)
+    wrapper()
+    got = out.clone()
+    cudnn_ref()
+    ref = out.clone()
+    max_abs, mean_abs, p99_abs, cosine = metrics(got, ref)
+    native_us = bench(raw, args.warmup, args.iters)
+    wrapper_us = bench(wrapper, args.warmup, args.iters)
+    eager_us = bench(cudnn_ref, args.warmup, args.iters)
+    compile_us = bench(compiled, args.warmup, args.iters)
+    return Result(
+        f"{label}-residual",
+        str(shape),
+        native_us,
+        wrapper_us,
+        wrapper_us / native_us,
+        eager_us,
+        compile_us,
+        None,
+        None,
+        max_abs,
+        mean_abs,
+        p99_abs,
+        cosine,
+        wrapper_us - native_us <= max(0.5, native_us * 0.05)
+        and wrapper_us <= min(eager_us, compile_us) * 0.98
+        and cosine >= 0.999
+        and mean_abs <= 0.01,
+    )
+def run_conv2d(ops, native, label, shape, args):
+    n, h, w, ci, co = shape
+    input = (torch.randn((n, h, w, ci), device="cuda") * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    weight = (torch.randn((co, 3, 3, ci), device="cuda") * 0.1).to(
+        torch.float8_e4m3fn
+    )
+    bias = (torch.randn(co, device="cuda") * 0.01).to(torch.bfloat16)
+    out = torch.empty((n, h, w, co), device="cuda", dtype=torch.bfloat16)
+    alpha = 0.75
+    wrapper = lambda: source_call(
+        ops,
+        "fp8_conv2d_3x3_nhwc_bf16",
+        input,
+        weight,
+        bias,
+        alpha,
+        out=out,
+    )
+    raw = lambda: native.conv2d(input, weight, bias, alpha, out)
+    def cudnn_ref():
+        x = input.float().permute(0, 3, 1, 2)
+        wt = weight.float().permute(0, 3, 1, 2)
+        y = F.conv2d(x, wt, padding=1).mul(alpha)
+        y = y.add(bias.float().view(1, -1, 1, 1))
+        out.copy_(y.permute(0, 2, 3, 1).to(torch.bfloat16))
+    compiled = torch.compile(cudnn_ref, fullgraph=True)
+    wrapper()
+    got = out.clone()
+    cudnn_ref()
+    ref = out.clone()
+    max_abs, mean_abs, p99_abs, cosine = metrics(got, ref)
+    native_us = bench(raw, args.warmup, args.iters)
+    wrapper_us = bench(wrapper, args.warmup, args.iters)
+    eager_us = bench(cudnn_ref, args.warmup, args.iters)
+    compile_us = bench(compiled, args.warmup, args.iters)
+    return Result(
+        label,
+        str(shape),
+        native_us,
+        wrapper_us,
+        wrapper_us / native_us,
+        eager_us,
+        compile_us,
+        None,
+        None,
+        max_abs,
+        mean_abs,
+        p99_abs,
+        cosine,
+        wrapper_us - native_us <= max(0.5, native_us * 0.05)
+        and wrapper_us <= min(eager_us, compile_us) * 0.98
+        and cosine >= 0.999
+        and mean_abs <= 0.01,
+    )
+def run_nvfp4_conv3d(ops, native, label, shape, args, residual_path):
+    n, tc, tn, h, w, ci, co = shape
+    cache_bf16 = (
+        torch.randn((n, tc, h, w, ci), device="cuda") * 0.1
+    ).to(torch.bfloat16)
+    input_bf16 = (
+        torch.randn((n, tn, h, w, ci), device="cuda") * 0.1
+    ).to(torch.bfloat16)
+    weight_bf16 = (
+        torch.randn((co, 3, 3, 3, ci), device="cuda") * 0.1
+    ).to(torch.bfloat16)
+    cache, cache_sf = quantize_conv_tensor(cache_bf16)
+    input, input_sf = quantize_conv_tensor(input_bf16)
+    weight, weight_sf = quantize_conv_tensor(weight_bf16)
+    cache_dequant = dequantize_linear_nvfp4(
+        cache.reshape(-1, ci // 2), cache_sf.reshape(-1, ci // 16)
+    ).reshape_as(cache_bf16)
+    input_dequant = dequantize_linear_nvfp4(
+        input.reshape(-1, ci // 2), input_sf.reshape(-1, ci // 16)
+    ).reshape_as(input_bf16)
+    weight_dequant = dequantize_linear_nvfp4(
+        weight.reshape(-1, ci // 2), weight_sf.reshape(-1, ci // 16)
+    ).reshape_as(weight_bf16)
+    bias = (torch.randn(co, device="cuda") * 0.01).to(torch.bfloat16)
+    alpha = 0.75
+    if residual_path:
+        residual = torch.randn(
+            (n, co, tn, h, w), device="cuda", dtype=torch.bfloat16
+        )
+        out = torch.empty_like(residual)
+        wrapper = lambda: source_call(
+            ops,
+            "nvfp4_causal_conv3d_residual_ncdhw_bf16",
+            cache, input, weight, cache_sf, input_sf, weight_sf, bias,
+            residual, None, alpha, out=out,
+        )
+        raw = lambda: native.nvfp4_causal_conv3d_residual(
+            cache, input, weight, cache_sf, input_sf, weight_sf, bias,
+            residual, alpha, out,
+        )
+    else:
+        residual = None
+        out = torch.empty(
+            (n, tn, h, w, co), device="cuda", dtype=torch.bfloat16
+        )
+        wrapper = lambda: source_call(
+            ops,
+            "nvfp4_causal_conv3d_ndhwc_bf16",
+            cache, input, weight, cache_sf, input_sf, weight_sf, bias,
+            None, alpha, out=out,
+        )
+        raw = lambda: native.nvfp4_causal_conv3d(
+            cache, input, weight, cache_sf, input_sf, weight_sf, bias,
+            alpha, out,
+        )
+    def store_cudnn_result(cache_value, input_value):
+        x = torch.cat((cache_value, input_value), dim=1).permute(
+            0, 4, 1, 2, 3
+        )
+        wt = weight_dequant.permute(0, 4, 1, 2, 3)
+        y = F.conv3d(x, wt, padding=(0, 1, 1))[:, :, :tn]
+        y = y.mul(alpha).add(bias.float().view(1, -1, 1, 1, 1))
+        if residual_path:
+            out.copy_(
+                (y.to(torch.bfloat16).float() + residual.float()).to(
+                    torch.bfloat16
+                )
+            )
+        else:
+            out.copy_(y.permute(0, 2, 3, 4, 1).to(torch.bfloat16))
+    def predequant_cudnn_ref():
+        store_cudnn_result(cache_dequant, input_dequant)
+    magnitude = torch.tensor(
+        [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], device="cuda"
+    )
+    scale_values, scale_bytes = [], []
+    for byte in list(range(0x78)) + [0xFE]:
+        exponent = (byte >> 3) & 0xF
+        mantissa = byte & 0x7
+        value = (
+            (mantissa / 8.0) * (2.0 ** -6)
+            if exponent == 0
+            else (1.0 + mantissa / 8.0) * (2.0 ** (exponent - 7))
+        )
+        scale_values.append(value)
+        scale_bytes.append(byte)
+    scale_lookup = torch.zeros(256, device="cuda")
+    scale_lookup[
+        torch.tensor(scale_bytes, device="cuda", dtype=torch.long)
+    ] = torch.tensor(scale_values, device="cuda")
+    def unpack(packed_value, scale_value):
+        low = packed_value & 0xF
+        high = packed_value >> 4
+        low_value = magnitude[(low & 0x7).long()] * torch.where(
+            low & 0x8 != 0, -1.0, 1.0
+        )
+        high_value = magnitude[(high & 0x7).long()] * torch.where(
+            high & 0x8 != 0, -1.0, 1.0
+        )
+        values = torch.stack((low_value, high_value), dim=-1).flatten(-2)
+        scales_value = scale_lookup[scale_value.long()].repeat_interleave(
+            16, dim=-1
+        )
+        return values * scales_value
+    def cudnn_ref():
+        cache_value = unpack(cache, cache_sf).reshape_as(cache_bf16)
+        input_value = unpack(input, input_sf).reshape_as(input_bf16)
+        store_cudnn_result(cache_value, input_value)
+    compiled = torch.compile(cudnn_ref, fullgraph=True)
+    predequant_compiled = torch.compile(predequant_cudnn_ref, fullgraph=True)
+    wrapper()
+    got = out.clone()
+    cudnn_ref()
+    ref = out.clone()
+    max_abs, mean_abs, p99_abs, cosine = metrics(got, ref)
+    native_us = bench(raw, args.warmup, args.iters)
+    wrapper_us = bench(wrapper, args.warmup, args.iters)
+    eager_us = bench(cudnn_ref, args.warmup, args.iters)
+    compile_us = bench(compiled, args.warmup, args.iters)
+    diagnostic_eager_us = bench(
+        predequant_cudnn_ref, args.warmup, args.iters
+    )
+    diagnostic_compile_us = bench(
+        predequant_compiled, args.warmup, args.iters
+    )
+    return Result(
+        f"{label}{'-residual' if residual_path else ''}",
+        str(shape),
+        native_us,
+        wrapper_us,
+        wrapper_us / native_us,
+        eager_us,
+        compile_us,
+        diagnostic_eager_us,
+        diagnostic_compile_us,
+        max_abs,
+        mean_abs,
+        p99_abs,
+        cosine,
+        wrapper_us - native_us <= max(0.5, native_us * 0.05)
+        and wrapper_us <= min(eager_us, compile_us) * 0.98
+        and cosine >= 0.998
+        and mean_abs <= 0.02,
+    )
+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=10)
+    parser.add_argument("--iters", type=int, default=30)
+    parser.add_argument("--output")
+    args = parser.parse_args()
+    ops = (
+        load_source_ops()
+        if args.backend == "source"
+        else load_installed_ops(args.artifact)
+    )
+    native = build_native()
+    rows = [
+        *(run_conv3d(ops, native, name, shape, args)
+          for name, shape in CONV3D_SHAPES.items()),
+        *(run_conv3d_residual(ops, native, name, shape, args)
+          for name, shape in CONV3D_SHAPES.items()
+          if shape[-1] % 8 == 0),
+        *(run_conv2d(ops, native, name, shape, args)
+          for name, shape in CONV2D_SHAPES.items()),
+        *(run_nvfp4_conv3d(ops, native, name, shape, args, False)
+          for name, shape in NVFP4_CONV3D_SHAPES.items()),
+        *(run_nvfp4_conv3d(ops, native, name, shape, args, True)
+          for name, shape in NVFP4_CONV3D_SHAPES.items()),
+    ]
+    for row in rows:
+        print(
+            f"{row.workload}: native={row.native_us:.3f}us "
+            f"wrapper={row.wrapper_us:.3f}us ({row.wrapper_native:.3f}) "
+            f"cuDNN-eager={row.eager_cudnn_us:.3f}us "
+            f"cuDNN-compile={row.compile_cudnn_us:.3f}us "
+            + (
+                f"predequant-compile="
+                f"{row.diagnostic_predequant_compile_us:.3f}us "
+                if row.diagnostic_predequant_compile_us is not None
+                else ""
+            )
+            + f"cos={row.cosine:.7f} 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("world-model Conv acceptance failed")
+if __name__ == "__main__":
+    main()

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

@@ -27,6 +27,11 @@ def _fp8_conv3d_fake(
         raise RuntimeError("cache_x must have shape (N,2,H,W,Ci)")
     if weight.shape != (co, 3, 3, 3, ci):
         raise RuntimeError("weight must have shape (Co,3,3,3,Ci)")
     if residual.shape != (n, co, t_new, h, w) or out.shape != residual.shape:
         raise RuntimeError("residual/out must be NCDHW")
     if bias.shape != (co,):
@@ -34,6 +39,179 @@ def _fp8_conv3d_fake(
     return None
 def fp8_conv3d_v18_ncdhw_res_bf16out(
     cache_x: torch.Tensor,
     new_x: torch.Tensor,
@@ -54,4 +232,156 @@ def fp8_conv3d_v18_ncdhw_res_bf16out(
     return out
-__all__ = ["fp8_conv3d_v18_ncdhw_res_bf16out"]

         raise RuntimeError("cache_x must have shape (N,2,H,W,Ci)")
     if weight.shape != (co, 3, 3, 3, ci):
         raise RuntimeError("weight must have shape (Co,3,3,3,Ci)")
+    if ci not in (32, 64):
+        raise RuntimeError(
+            "FP8 Conv3D is accepted only for Ci=32/64; use the "
+            "strong-library or NVFP4 path for larger channels"
+        )
     if residual.shape != (n, co, t_new, h, w) or out.shape != residual.shape:
         raise RuntimeError("residual/out must be NCDHW")
     if bias.shape != (co,):
     return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_causal_conv3d_ndhwc_bf16")
+)
+def _fp8_causal_conv3d_fake(
+    cache_x: torch.Tensor,
+    new_x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    if cache_x.dim() != 5 or new_x.dim() != 5:
+        raise RuntimeError("cache_x/new_x must be NDHWC")
+    n, t, h, w, ci = new_x.shape
+    co = weight.shape[0]
+    if cache_x.shape != (n, 2, h, w, ci):
+        raise RuntimeError("cache_x must have shape (N,2,H,W,Ci)")
+    if weight.shape != (co, 3, 3, 3, ci):
+        raise RuntimeError("weight must have shape (Co,3,3,3,Ci)")
+    if ci not in (32, 64):
+        raise RuntimeError(
+            "FP8 Conv3D is accepted only for Ci=32/64; use the "
+            "strong-library or NVFP4 path for larger channels"
+        )
+    if bias.shape != (co,) or out.shape != (n, t, h, w, co):
+        raise RuntimeError("bias/out shape mismatch")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_conv2d_3x3_nhwc_bf16")
+)
+def _fp8_conv2d_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 4:
+        raise RuntimeError("input must have shape (N,H,W,Ci)")
+    n, h, w, ci = input.shape
+    co = weight.shape[0]
+    if weight.shape != (co, 3, 3, ci):
+        raise RuntimeError("weight must have shape (Co,3,3,Ci)")
+    if bias.shape != (co,) or out.shape != (n, h, w, co):
+        raise RuntimeError("bias/out shape mismatch")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_conv2d_3x3_ncdhw_bf16")
+)
+def _fp8_conv2d_ncdhw_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 5:
+        raise RuntimeError("input must have shape (B,T,H,W,Ci)")
+    b, t, h, w, ci = input.shape
+    co = weight.shape[0]
+    if weight.shape != (co, 3, 3, ci):
+        raise RuntimeError("weight must have shape (Co,3,3,Ci)")
+    if bias.shape != (co,) or out.shape != (b, co, t, h, w):
+        raise RuntimeError("bias/out shape mismatch")
+    return None
+def _check_nvfp4_conv_shapes(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    outer_weight: torch.Tensor | None,
+) -> tuple[int, int, int, int, int]:
+    if (
+        cache_packed.dim() != 5
+        or new_packed.dim() != 5
+        or weight_packed.dim() != 5
+    ):
+        raise RuntimeError("packed inputs must be five-dimensional")
+    n, t, h, w, ci_half = new_packed.shape
+    ci = ci_half * 2
+    co = weight_packed.shape[0]
+    if (
+        (ci != 64 and ci % 128 != 0)
+        or co % 8 != 0
+        or cache_packed.shape != (n, 2, h, w, ci // 2)
+        or weight_packed.shape != (co, 3, 3, 3, ci // 2)
+        or cache_sf.shape != (n, 2, h, w, ci // 16)
+        or new_sf.shape != (n, t, h, w, ci // 16)
+        or weight_sf.shape != (co, 3, 3, 3, ci // 16)
+        or bias.shape != (co,)
+        or (outer_weight is not None and outer_weight.shape != (co,))
+    ):
+        raise RuntimeError(
+            "NVFP4 Conv3D accepts Ci=64 or multiples of 128; "
+            "other shapes must use the strong-library path"
+        )
+    return n, t, h, w, co
+@torch.library.register_fake(
+    add_op_namespace_prefix("nvfp4_causal_conv3d_ndhwc_bf16")
+)
+def _nvfp4_causal_conv3d_fake(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    outer_weight: torch.Tensor | None,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    del alpha
+    n, t, h, w, co = _check_nvfp4_conv_shapes(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        outer_weight,
+    )
+    if out.shape != (n, t, h, w, co):
+        raise RuntimeError("out must have shape (N,T,H,W,Co)")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix(
+        "nvfp4_causal_conv3d_residual_ncdhw_bf16"
+    )
+)
+def _nvfp4_causal_conv3d_residual_fake(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    outer_weight: torch.Tensor | None,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    del alpha
+    n, t, h, w, co = _check_nvfp4_conv_shapes(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        outer_weight,
+    )
+    if residual.shape != (n, co, t, h, w) or out.shape != residual.shape:
+        raise RuntimeError("residual/out must have shape (N,Co,T,H,W)")
+    return None
 def fp8_conv3d_v18_ncdhw_res_bf16out(
     cache_x: torch.Tensor,
     new_x: torch.Tensor,
     return out
+def fp8_causal_conv3d_ndhwc_bf16(
+    cache_x: torch.Tensor,
+    new_x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 causal 3D convolution with virtual two-frame cache concat."""
+    n, t, h, w, _ = new_x.shape
+    if out is None:
+        out = torch.empty(
+            (n, t, h, w, weight.shape[0]),
+            device=new_x.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_causal_conv3d_ndhwc_bf16(
+        cache_x, new_x, weight, bias, float(alpha), out
+    )
+    return out
+def fp8_conv2d_3x3_nhwc_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 3x3 Conv2D with NHWC input/output and BF16 epilogue."""
+    if out is None:
+        out = torch.empty(
+            (*input.shape[:3], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_conv2d_3x3_nhwc_bf16(
+        input, weight, bias, float(alpha), out
+    )
+    return out
+def fp8_conv2d_3x3_ncdhw_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 3x3 Conv2D over B*T frames with direct BF16 NCDHW output."""
+    if out is None:
+        out = torch.empty(
+            (
+                input.shape[0],
+                weight.shape[0],
+                input.shape[1],
+                input.shape[2],
+                input.shape[3],
+            ),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_conv2d_3x3_ncdhw_bf16(
+        input, weight, bias, float(alpha), out
+    )
+    return out
+def nvfp4_causal_conv3d_ndhwc_bf16(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    outer_weight: torch.Tensor | None = None,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """NVFP4 causal Conv3D with linear UE4M3 scale factors and BF16 NDHWC output."""
+    n, t, h, w, _ = new_packed.shape
+    if out is None:
+        out = torch.empty(
+            (n, t, h, w, weight_packed.shape[0]),
+            device=new_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_causal_conv3d_ndhwc_bf16(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        outer_weight,
+        float(alpha),
+        out,
+    )
+    return out
+def nvfp4_causal_conv3d_residual_ncdhw_bf16(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    outer_weight: torch.Tensor | None = None,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """NVFP4 causal Conv3D with fused bias, residual, and BF16 NCDHW output."""
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_causal_conv3d_residual_ncdhw_bf16(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        residual,
+        outer_weight,
+        float(alpha),
+        out,
+    )
+    return out
+__all__ = [
+    "fp8_conv3d_v18_ncdhw_res_bf16out",
+    "fp8_causal_conv3d_ndhwc_bf16",
+    "fp8_conv2d_3x3_nhwc_bf16",
+    "fp8_conv2d_3x3_ncdhw_bf16",
+    "nvfp4_causal_conv3d_ndhwc_bf16",
+    "nvfp4_causal_conv3d_residual_ncdhw_bf16",
+]

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

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

 import torch
+from . import _world_model_conv_cuda_7781728
+ops = torch.ops._world_model_conv_cuda_7781728
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_world_model_conv_cuda_7781728::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/{_world_model_conv_cuda_f14c443.abi3.so → _world_model_conv_cuda_7781728.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:724ae0ca8e8d4638c2d8aae226836993124662348462c1d83e6f558109bdf221
-size 296248

 version https://git-lfs.github.com/spec/v1
+oid sha256:d963296bec8c011fc81e22ef3d66ae0df353d36550f0844bbdf5541a4f23cd2a
+size 1271616

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

@@ -1,6 +1,6 @@
 {
   "name": "world-model-conv",
-  "id": "_world_model_conv_cuda_f14c443",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "world-model-conv",
+  "id": "_world_model_conv_cuda_7781728",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -27,6 +27,11 @@ def _fp8_conv3d_fake(
         raise RuntimeError("cache_x must have shape (N,2,H,W,Ci)")
     if weight.shape != (co, 3, 3, 3, ci):
         raise RuntimeError("weight must have shape (Co,3,3,3,Ci)")
     if residual.shape != (n, co, t_new, h, w) or out.shape != residual.shape:
         raise RuntimeError("residual/out must be NCDHW")
     if bias.shape != (co,):
@@ -34,6 +39,179 @@ def _fp8_conv3d_fake(
     return None
 def fp8_conv3d_v18_ncdhw_res_bf16out(
     cache_x: torch.Tensor,
     new_x: torch.Tensor,
@@ -54,4 +232,156 @@ def fp8_conv3d_v18_ncdhw_res_bf16out(
     return out
-__all__ = ["fp8_conv3d_v18_ncdhw_res_bf16out"]

         raise RuntimeError("cache_x must have shape (N,2,H,W,Ci)")
     if weight.shape != (co, 3, 3, 3, ci):
         raise RuntimeError("weight must have shape (Co,3,3,3,Ci)")
+    if ci not in (32, 64):
+        raise RuntimeError(
+            "FP8 Conv3D is accepted only for Ci=32/64; use the "
+            "strong-library or NVFP4 path for larger channels"
+        )
     if residual.shape != (n, co, t_new, h, w) or out.shape != residual.shape:
         raise RuntimeError("residual/out must be NCDHW")
     if bias.shape != (co,):
     return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_causal_conv3d_ndhwc_bf16")
+)
+def _fp8_causal_conv3d_fake(
+    cache_x: torch.Tensor,
+    new_x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    if cache_x.dim() != 5 or new_x.dim() != 5:
+        raise RuntimeError("cache_x/new_x must be NDHWC")
+    n, t, h, w, ci = new_x.shape
+    co = weight.shape[0]
+    if cache_x.shape != (n, 2, h, w, ci):
+        raise RuntimeError("cache_x must have shape (N,2,H,W,Ci)")
+    if weight.shape != (co, 3, 3, 3, ci):
+        raise RuntimeError("weight must have shape (Co,3,3,3,Ci)")
+    if ci not in (32, 64):
+        raise RuntimeError(
+            "FP8 Conv3D is accepted only for Ci=32/64; use the "
+            "strong-library or NVFP4 path for larger channels"
+        )
+    if bias.shape != (co,) or out.shape != (n, t, h, w, co):
+        raise RuntimeError("bias/out shape mismatch")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_conv2d_3x3_nhwc_bf16")
+)
+def _fp8_conv2d_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 4:
+        raise RuntimeError("input must have shape (N,H,W,Ci)")
+    n, h, w, ci = input.shape
+    co = weight.shape[0]
+    if weight.shape != (co, 3, 3, ci):
+        raise RuntimeError("weight must have shape (Co,3,3,Ci)")
+    if bias.shape != (co,) or out.shape != (n, h, w, co):
+        raise RuntimeError("bias/out shape mismatch")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_conv2d_3x3_ncdhw_bf16")
+)
+def _fp8_conv2d_ncdhw_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 5:
+        raise RuntimeError("input must have shape (B,T,H,W,Ci)")
+    b, t, h, w, ci = input.shape
+    co = weight.shape[0]
+    if weight.shape != (co, 3, 3, ci):
+        raise RuntimeError("weight must have shape (Co,3,3,Ci)")
+    if bias.shape != (co,) or out.shape != (b, co, t, h, w):
+        raise RuntimeError("bias/out shape mismatch")
+    return None
+def _check_nvfp4_conv_shapes(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    outer_weight: torch.Tensor | None,
+) -> tuple[int, int, int, int, int]:
+    if (
+        cache_packed.dim() != 5
+        or new_packed.dim() != 5
+        or weight_packed.dim() != 5
+    ):
+        raise RuntimeError("packed inputs must be five-dimensional")
+    n, t, h, w, ci_half = new_packed.shape
+    ci = ci_half * 2
+    co = weight_packed.shape[0]
+    if (
+        (ci != 64 and ci % 128 != 0)
+        or co % 8 != 0
+        or cache_packed.shape != (n, 2, h, w, ci // 2)
+        or weight_packed.shape != (co, 3, 3, 3, ci // 2)
+        or cache_sf.shape != (n, 2, h, w, ci // 16)
+        or new_sf.shape != (n, t, h, w, ci // 16)
+        or weight_sf.shape != (co, 3, 3, 3, ci // 16)
+        or bias.shape != (co,)
+        or (outer_weight is not None and outer_weight.shape != (co,))
+    ):
+        raise RuntimeError(
+            "NVFP4 Conv3D accepts Ci=64 or multiples of 128; "
+            "other shapes must use the strong-library path"
+        )
+    return n, t, h, w, co
+@torch.library.register_fake(
+    add_op_namespace_prefix("nvfp4_causal_conv3d_ndhwc_bf16")
+)
+def _nvfp4_causal_conv3d_fake(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    outer_weight: torch.Tensor | None,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    del alpha
+    n, t, h, w, co = _check_nvfp4_conv_shapes(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        outer_weight,
+    )
+    if out.shape != (n, t, h, w, co):
+        raise RuntimeError("out must have shape (N,T,H,W,Co)")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix(
+        "nvfp4_causal_conv3d_residual_ncdhw_bf16"
+    )
+)
+def _nvfp4_causal_conv3d_residual_fake(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    outer_weight: torch.Tensor | None,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    del alpha
+    n, t, h, w, co = _check_nvfp4_conv_shapes(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        outer_weight,
+    )
+    if residual.shape != (n, co, t, h, w) or out.shape != residual.shape:
+        raise RuntimeError("residual/out must have shape (N,Co,T,H,W)")
+    return None
 def fp8_conv3d_v18_ncdhw_res_bf16out(
     cache_x: torch.Tensor,
     new_x: torch.Tensor,
     return out
+def fp8_causal_conv3d_ndhwc_bf16(
+    cache_x: torch.Tensor,
+    new_x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 causal 3D convolution with virtual two-frame cache concat."""
+    n, t, h, w, _ = new_x.shape
+    if out is None:
+        out = torch.empty(
+            (n, t, h, w, weight.shape[0]),
+            device=new_x.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_causal_conv3d_ndhwc_bf16(
+        cache_x, new_x, weight, bias, float(alpha), out
+    )
+    return out
+def fp8_conv2d_3x3_nhwc_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 3x3 Conv2D with NHWC input/output and BF16 epilogue."""
+    if out is None:
+        out = torch.empty(
+            (*input.shape[:3], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_conv2d_3x3_nhwc_bf16(
+        input, weight, bias, float(alpha), out
+    )
+    return out
+def fp8_conv2d_3x3_ncdhw_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 3x3 Conv2D over B*T frames with direct BF16 NCDHW output."""
+    if out is None:
+        out = torch.empty(
+            (
+                input.shape[0],
+                weight.shape[0],
+                input.shape[1],
+                input.shape[2],
+                input.shape[3],
+            ),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_conv2d_3x3_ncdhw_bf16(
+        input, weight, bias, float(alpha), out
+    )
+    return out
+def nvfp4_causal_conv3d_ndhwc_bf16(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    outer_weight: torch.Tensor | None = None,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """NVFP4 causal Conv3D with linear UE4M3 scale factors and BF16 NDHWC output."""
+    n, t, h, w, _ = new_packed.shape
+    if out is None:
+        out = torch.empty(
+            (n, t, h, w, weight_packed.shape[0]),
+            device=new_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_causal_conv3d_ndhwc_bf16(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        outer_weight,
+        float(alpha),
+        out,
+    )
+    return out
+def nvfp4_causal_conv3d_residual_ncdhw_bf16(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    outer_weight: torch.Tensor | None = None,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """NVFP4 causal Conv3D with fused bias, residual, and BF16 NCDHW output."""
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_causal_conv3d_residual_ncdhw_bf16(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        residual,
+        outer_weight,
+        float(alpha),
+        out,
+    )
+    return out
+__all__ = [
+    "fp8_conv3d_v18_ncdhw_res_bf16out",
+    "fp8_causal_conv3d_ndhwc_bf16",
+    "fp8_conv2d_3x3_nhwc_bf16",
+    "fp8_conv2d_3x3_ncdhw_bf16",
+    "nvfp4_causal_conv3d_ndhwc_bf16",
+    "nvfp4_causal_conv3d_residual_ncdhw_bf16",
+]

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

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

 import torch
+from . import _world_model_conv_cuda_7781728
+ops = torch.ops._world_model_conv_cuda_7781728
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_world_model_conv_cuda_7781728::{op_name}"

build/torch211-cxx11-cu130-x86_64-linux/{_world_model_conv_cuda_f14c443.abi3.so → _world_model_conv_cuda_7781728.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5bf4f3051e364b68f0fa536f7d9a7199a71c14a64f4863ab0485ef3cde7ffd7c
-size 285088

 version https://git-lfs.github.com/spec/v1
+oid sha256:8dfb73bc6c3dd31b73b180a9781a09a736607d33fb7846cf3b5d6d9dd99804cd
+size 1242552

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

@@ -1,6 +1,6 @@
 {
   "name": "world-model-conv",
-  "id": "_world_model_conv_cuda_f14c443",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "world-model-conv",
+  "id": "_world_model_conv_cuda_7781728",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -27,6 +27,11 @@ def _fp8_conv3d_fake(
         raise RuntimeError("cache_x must have shape (N,2,H,W,Ci)")
     if weight.shape != (co, 3, 3, 3, ci):
         raise RuntimeError("weight must have shape (Co,3,3,3,Ci)")
     if residual.shape != (n, co, t_new, h, w) or out.shape != residual.shape:
         raise RuntimeError("residual/out must be NCDHW")
     if bias.shape != (co,):
@@ -34,6 +39,179 @@ def _fp8_conv3d_fake(
     return None
 def fp8_conv3d_v18_ncdhw_res_bf16out(
     cache_x: torch.Tensor,
     new_x: torch.Tensor,
@@ -54,4 +232,156 @@ def fp8_conv3d_v18_ncdhw_res_bf16out(
     return out
-__all__ = ["fp8_conv3d_v18_ncdhw_res_bf16out"]

         raise RuntimeError("cache_x must have shape (N,2,H,W,Ci)")
     if weight.shape != (co, 3, 3, 3, ci):
         raise RuntimeError("weight must have shape (Co,3,3,3,Ci)")
+    if ci not in (32, 64):
+        raise RuntimeError(
+            "FP8 Conv3D is accepted only for Ci=32/64; use the "
+            "strong-library or NVFP4 path for larger channels"
+        )
     if residual.shape != (n, co, t_new, h, w) or out.shape != residual.shape:
         raise RuntimeError("residual/out must be NCDHW")
     if bias.shape != (co,):
     return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_causal_conv3d_ndhwc_bf16")
+)
+def _fp8_causal_conv3d_fake(
+    cache_x: torch.Tensor,
+    new_x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    if cache_x.dim() != 5 or new_x.dim() != 5:
+        raise RuntimeError("cache_x/new_x must be NDHWC")
+    n, t, h, w, ci = new_x.shape
+    co = weight.shape[0]
+    if cache_x.shape != (n, 2, h, w, ci):
+        raise RuntimeError("cache_x must have shape (N,2,H,W,Ci)")
+    if weight.shape != (co, 3, 3, 3, ci):
+        raise RuntimeError("weight must have shape (Co,3,3,3,Ci)")
+    if ci not in (32, 64):
+        raise RuntimeError(
+            "FP8 Conv3D is accepted only for Ci=32/64; use the "
+            "strong-library or NVFP4 path for larger channels"
+        )
+    if bias.shape != (co,) or out.shape != (n, t, h, w, co):
+        raise RuntimeError("bias/out shape mismatch")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_conv2d_3x3_nhwc_bf16")
+)
+def _fp8_conv2d_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 4:
+        raise RuntimeError("input must have shape (N,H,W,Ci)")
+    n, h, w, ci = input.shape
+    co = weight.shape[0]
+    if weight.shape != (co, 3, 3, ci):
+        raise RuntimeError("weight must have shape (Co,3,3,Ci)")
+    if bias.shape != (co,) or out.shape != (n, h, w, co):
+        raise RuntimeError("bias/out shape mismatch")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_conv2d_3x3_ncdhw_bf16")
+)
+def _fp8_conv2d_ncdhw_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 5:
+        raise RuntimeError("input must have shape (B,T,H,W,Ci)")
+    b, t, h, w, ci = input.shape
+    co = weight.shape[0]
+    if weight.shape != (co, 3, 3, ci):
+        raise RuntimeError("weight must have shape (Co,3,3,Ci)")
+    if bias.shape != (co,) or out.shape != (b, co, t, h, w):
+        raise RuntimeError("bias/out shape mismatch")
+    return None
+def _check_nvfp4_conv_shapes(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    outer_weight: torch.Tensor | None,
+) -> tuple[int, int, int, int, int]:
+    if (
+        cache_packed.dim() != 5
+        or new_packed.dim() != 5
+        or weight_packed.dim() != 5
+    ):
+        raise RuntimeError("packed inputs must be five-dimensional")
+    n, t, h, w, ci_half = new_packed.shape
+    ci = ci_half * 2
+    co = weight_packed.shape[0]
+    if (
+        (ci != 64 and ci % 128 != 0)
+        or co % 8 != 0
+        or cache_packed.shape != (n, 2, h, w, ci // 2)
+        or weight_packed.shape != (co, 3, 3, 3, ci // 2)
+        or cache_sf.shape != (n, 2, h, w, ci // 16)
+        or new_sf.shape != (n, t, h, w, ci // 16)
+        or weight_sf.shape != (co, 3, 3, 3, ci // 16)
+        or bias.shape != (co,)
+        or (outer_weight is not None and outer_weight.shape != (co,))
+    ):
+        raise RuntimeError(
+            "NVFP4 Conv3D accepts Ci=64 or multiples of 128; "
+            "other shapes must use the strong-library path"
+        )
+    return n, t, h, w, co
+@torch.library.register_fake(
+    add_op_namespace_prefix("nvfp4_causal_conv3d_ndhwc_bf16")
+)
+def _nvfp4_causal_conv3d_fake(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    outer_weight: torch.Tensor | None,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    del alpha
+    n, t, h, w, co = _check_nvfp4_conv_shapes(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        outer_weight,
+    )
+    if out.shape != (n, t, h, w, co):
+        raise RuntimeError("out must have shape (N,T,H,W,Co)")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix(
+        "nvfp4_causal_conv3d_residual_ncdhw_bf16"
+    )
+)
+def _nvfp4_causal_conv3d_residual_fake(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    outer_weight: torch.Tensor | None,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    del alpha
+    n, t, h, w, co = _check_nvfp4_conv_shapes(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        outer_weight,
+    )
+    if residual.shape != (n, co, t, h, w) or out.shape != residual.shape:
+        raise RuntimeError("residual/out must have shape (N,Co,T,H,W)")
+    return None
 def fp8_conv3d_v18_ncdhw_res_bf16out(
     cache_x: torch.Tensor,
     new_x: torch.Tensor,
     return out
+def fp8_causal_conv3d_ndhwc_bf16(
+    cache_x: torch.Tensor,
+    new_x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 causal 3D convolution with virtual two-frame cache concat."""
+    n, t, h, w, _ = new_x.shape
+    if out is None:
+        out = torch.empty(
+            (n, t, h, w, weight.shape[0]),
+            device=new_x.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_causal_conv3d_ndhwc_bf16(
+        cache_x, new_x, weight, bias, float(alpha), out
+    )
+    return out
+def fp8_conv2d_3x3_nhwc_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 3x3 Conv2D with NHWC input/output and BF16 epilogue."""
+    if out is None:
+        out = torch.empty(
+            (*input.shape[:3], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_conv2d_3x3_nhwc_bf16(
+        input, weight, bias, float(alpha), out
+    )
+    return out
+def fp8_conv2d_3x3_ncdhw_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 3x3 Conv2D over B*T frames with direct BF16 NCDHW output."""
+    if out is None:
+        out = torch.empty(
+            (
+                input.shape[0],
+                weight.shape[0],
+                input.shape[1],
+                input.shape[2],
+                input.shape[3],
+            ),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_conv2d_3x3_ncdhw_bf16(
+        input, weight, bias, float(alpha), out
+    )
+    return out
+def nvfp4_causal_conv3d_ndhwc_bf16(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    outer_weight: torch.Tensor | None = None,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """NVFP4 causal Conv3D with linear UE4M3 scale factors and BF16 NDHWC output."""
+    n, t, h, w, _ = new_packed.shape
+    if out is None:
+        out = torch.empty(
+            (n, t, h, w, weight_packed.shape[0]),
+            device=new_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_causal_conv3d_ndhwc_bf16(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        outer_weight,
+        float(alpha),
+        out,
+    )
+    return out
+def nvfp4_causal_conv3d_residual_ncdhw_bf16(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    outer_weight: torch.Tensor | None = None,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """NVFP4 causal Conv3D with fused bias, residual, and BF16 NCDHW output."""
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_causal_conv3d_residual_ncdhw_bf16(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        residual,
+        outer_weight,
+        float(alpha),
+        out,
+    )
+    return out
+__all__ = [
+    "fp8_conv3d_v18_ncdhw_res_bf16out",
+    "fp8_causal_conv3d_ndhwc_bf16",
+    "fp8_conv2d_3x3_nhwc_bf16",
+    "fp8_conv2d_3x3_ncdhw_bf16",
+    "nvfp4_causal_conv3d_ndhwc_bf16",
+    "nvfp4_causal_conv3d_residual_ncdhw_bf16",
+]

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

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

 import torch
+from . import _world_model_conv_cuda_7781728
+ops = torch.ops._world_model_conv_cuda_7781728
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_world_model_conv_cuda_7781728::{op_name}"

build/torch212-cxx11-cu130-x86_64-linux/{_world_model_conv_cuda_f14c443.abi3.so → _world_model_conv_cuda_7781728.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:21963c9f0459cc516906960624490bd34898ac669436bf3a4e6cfd4f6f2bc9c7
+size 1248816

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

@@ -1,6 +1,6 @@
 {
   "name": "world-model-conv",
-  "id": "_world_model_conv_cuda_f14c443",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "world-model-conv",
+  "id": "_world_model_conv_cuda_7781728",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -27,6 +27,11 @@ def _fp8_conv3d_fake(
         raise RuntimeError("cache_x must have shape (N,2,H,W,Ci)")
     if weight.shape != (co, 3, 3, 3, ci):
         raise RuntimeError("weight must have shape (Co,3,3,3,Ci)")
     if residual.shape != (n, co, t_new, h, w) or out.shape != residual.shape:
         raise RuntimeError("residual/out must be NCDHW")
     if bias.shape != (co,):
@@ -34,6 +39,179 @@ def _fp8_conv3d_fake(
     return None
 def fp8_conv3d_v18_ncdhw_res_bf16out(
     cache_x: torch.Tensor,
     new_x: torch.Tensor,
@@ -54,4 +232,156 @@ def fp8_conv3d_v18_ncdhw_res_bf16out(
     return out
-__all__ = ["fp8_conv3d_v18_ncdhw_res_bf16out"]

         raise RuntimeError("cache_x must have shape (N,2,H,W,Ci)")
     if weight.shape != (co, 3, 3, 3, ci):
         raise RuntimeError("weight must have shape (Co,3,3,3,Ci)")
+    if ci not in (32, 64):
+        raise RuntimeError(
+            "FP8 Conv3D is accepted only for Ci=32/64; use the "
+            "strong-library or NVFP4 path for larger channels"
+        )
     if residual.shape != (n, co, t_new, h, w) or out.shape != residual.shape:
         raise RuntimeError("residual/out must be NCDHW")
     if bias.shape != (co,):
     return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_causal_conv3d_ndhwc_bf16")
+)
+def _fp8_causal_conv3d_fake(
+    cache_x: torch.Tensor,
+    new_x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    if cache_x.dim() != 5 or new_x.dim() != 5:
+        raise RuntimeError("cache_x/new_x must be NDHWC")
+    n, t, h, w, ci = new_x.shape
+    co = weight.shape[0]
+    if cache_x.shape != (n, 2, h, w, ci):
+        raise RuntimeError("cache_x must have shape (N,2,H,W,Ci)")
+    if weight.shape != (co, 3, 3, 3, ci):
+        raise RuntimeError("weight must have shape (Co,3,3,3,Ci)")
+    if ci not in (32, 64):
+        raise RuntimeError(
+            "FP8 Conv3D is accepted only for Ci=32/64; use the "
+            "strong-library or NVFP4 path for larger channels"
+        )
+    if bias.shape != (co,) or out.shape != (n, t, h, w, co):
+        raise RuntimeError("bias/out shape mismatch")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_conv2d_3x3_nhwc_bf16")
+)
+def _fp8_conv2d_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 4:
+        raise RuntimeError("input must have shape (N,H,W,Ci)")
+    n, h, w, ci = input.shape
+    co = weight.shape[0]
+    if weight.shape != (co, 3, 3, ci):
+        raise RuntimeError("weight must have shape (Co,3,3,Ci)")
+    if bias.shape != (co,) or out.shape != (n, h, w, co):
+        raise RuntimeError("bias/out shape mismatch")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix("fp8_conv2d_3x3_ncdhw_bf16")
+)
+def _fp8_conv2d_ncdhw_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    if input.dim() != 5:
+        raise RuntimeError("input must have shape (B,T,H,W,Ci)")
+    b, t, h, w, ci = input.shape
+    co = weight.shape[0]
+    if weight.shape != (co, 3, 3, ci):
+        raise RuntimeError("weight must have shape (Co,3,3,Ci)")
+    if bias.shape != (co,) or out.shape != (b, co, t, h, w):
+        raise RuntimeError("bias/out shape mismatch")
+    return None
+def _check_nvfp4_conv_shapes(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    outer_weight: torch.Tensor | None,
+) -> tuple[int, int, int, int, int]:
+    if (
+        cache_packed.dim() != 5
+        or new_packed.dim() != 5
+        or weight_packed.dim() != 5
+    ):
+        raise RuntimeError("packed inputs must be five-dimensional")
+    n, t, h, w, ci_half = new_packed.shape
+    ci = ci_half * 2
+    co = weight_packed.shape[0]
+    if (
+        (ci != 64 and ci % 128 != 0)
+        or co % 8 != 0
+        or cache_packed.shape != (n, 2, h, w, ci // 2)
+        or weight_packed.shape != (co, 3, 3, 3, ci // 2)
+        or cache_sf.shape != (n, 2, h, w, ci // 16)
+        or new_sf.shape != (n, t, h, w, ci // 16)
+        or weight_sf.shape != (co, 3, 3, 3, ci // 16)
+        or bias.shape != (co,)
+        or (outer_weight is not None and outer_weight.shape != (co,))
+    ):
+        raise RuntimeError(
+            "NVFP4 Conv3D accepts Ci=64 or multiples of 128; "
+            "other shapes must use the strong-library path"
+        )
+    return n, t, h, w, co
+@torch.library.register_fake(
+    add_op_namespace_prefix("nvfp4_causal_conv3d_ndhwc_bf16")
+)
+def _nvfp4_causal_conv3d_fake(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    outer_weight: torch.Tensor | None,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    del alpha
+    n, t, h, w, co = _check_nvfp4_conv_shapes(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        outer_weight,
+    )
+    if out.shape != (n, t, h, w, co):
+        raise RuntimeError("out must have shape (N,T,H,W,Co)")
+    return None
+@torch.library.register_fake(
+    add_op_namespace_prefix(
+        "nvfp4_causal_conv3d_residual_ncdhw_bf16"
+    )
+)
+def _nvfp4_causal_conv3d_residual_fake(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    outer_weight: torch.Tensor | None,
+    alpha: float,
+    out: torch.Tensor,
+) -> None:
+    del alpha
+    n, t, h, w, co = _check_nvfp4_conv_shapes(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        outer_weight,
+    )
+    if residual.shape != (n, co, t, h, w) or out.shape != residual.shape:
+        raise RuntimeError("residual/out must have shape (N,Co,T,H,W)")
+    return None
 def fp8_conv3d_v18_ncdhw_res_bf16out(
     cache_x: torch.Tensor,
     new_x: torch.Tensor,
     return out
+def fp8_causal_conv3d_ndhwc_bf16(
+    cache_x: torch.Tensor,
+    new_x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 causal 3D convolution with virtual two-frame cache concat."""
+    n, t, h, w, _ = new_x.shape
+    if out is None:
+        out = torch.empty(
+            (n, t, h, w, weight.shape[0]),
+            device=new_x.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_causal_conv3d_ndhwc_bf16(
+        cache_x, new_x, weight, bias, float(alpha), out
+    )
+    return out
+def fp8_conv2d_3x3_nhwc_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 3x3 Conv2D with NHWC input/output and BF16 epilogue."""
+    if out is None:
+        out = torch.empty(
+            (*input.shape[:3], weight.shape[0]),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_conv2d_3x3_nhwc_bf16(
+        input, weight, bias, float(alpha), out
+    )
+    return out
+def fp8_conv2d_3x3_ncdhw_bf16(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """FP8 3x3 Conv2D over B*T frames with direct BF16 NCDHW output."""
+    if out is None:
+        out = torch.empty(
+            (
+                input.shape[0],
+                weight.shape[0],
+                input.shape[1],
+                input.shape[2],
+                input.shape[3],
+            ),
+            device=input.device,
+            dtype=torch.bfloat16,
+        )
+    ops.fp8_conv2d_3x3_ncdhw_bf16(
+        input, weight, bias, float(alpha), out
+    )
+    return out
+def nvfp4_causal_conv3d_ndhwc_bf16(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    outer_weight: torch.Tensor | None = None,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """NVFP4 causal Conv3D with linear UE4M3 scale factors and BF16 NDHWC output."""
+    n, t, h, w, _ = new_packed.shape
+    if out is None:
+        out = torch.empty(
+            (n, t, h, w, weight_packed.shape[0]),
+            device=new_packed.device,
+            dtype=torch.bfloat16,
+        )
+    ops.nvfp4_causal_conv3d_ndhwc_bf16(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        outer_weight,
+        float(alpha),
+        out,
+    )
+    return out
+def nvfp4_causal_conv3d_residual_ncdhw_bf16(
+    cache_packed: torch.Tensor,
+    new_packed: torch.Tensor,
+    weight_packed: torch.Tensor,
+    cache_sf: torch.Tensor,
+    new_sf: torch.Tensor,
+    weight_sf: torch.Tensor,
+    bias: torch.Tensor,
+    residual: torch.Tensor,
+    outer_weight: torch.Tensor | None = None,
+    alpha: float = 1.0,
+    *,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """NVFP4 causal Conv3D with fused bias, residual, and BF16 NCDHW output."""
+    if out is None:
+        out = torch.empty_like(residual)
+    ops.nvfp4_causal_conv3d_residual_ncdhw_bf16(
+        cache_packed,
+        new_packed,
+        weight_packed,
+        cache_sf,
+        new_sf,
+        weight_sf,
+        bias,
+        residual,
+        outer_weight,
+        float(alpha),
+        out,
+    )
+    return out
+__all__ = [
+    "fp8_conv3d_v18_ncdhw_res_bf16out",
+    "fp8_causal_conv3d_ndhwc_bf16",
+    "fp8_conv2d_3x3_nhwc_bf16",
+    "fp8_conv2d_3x3_ncdhw_bf16",
+    "nvfp4_causal_conv3d_ndhwc_bf16",
+    "nvfp4_causal_conv3d_residual_ncdhw_bf16",
+]

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

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

 import torch
+from . import _world_model_conv_cuda_7781728
+ops = torch.ops._world_model_conv_cuda_7781728
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_world_model_conv_cuda_7781728::{op_name}"

build/torch212-cxx11-cu132-x86_64-linux/{_world_model_conv_cuda_f14c443.abi3.so → _world_model_conv_cuda_7781728.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:cf595248e19065ecdb2fe638be1a705647fad52912c88d0aa9c614195d09bf8f
+size 1285680

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

@@ -1,6 +1,6 @@
 {
   "name": "world-model-conv",
-  "id": "_world_model_conv_cuda_f14c443",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "world-model-conv",
+  "id": "_world_model_conv_cuda_7781728",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],