Promote latest kernel artifacts to main

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text
+*.pyd filter=lfs diff=lfs merge=lfs -text
+*.dylib filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,9 +0,0 @@
-# flashrt/flashrt-vla-video
-
-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-vla-video
-
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.

benchmarks/benchmark_q_norm_rope.py

@@ -0,0 +1,398 @@
+import torch
+
+from kernels.benchmark import Benchmark
+
+
+_original_allclose = torch.allclose
+
+
+def _flashrt_allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False):
+    if input.dtype == torch.bfloat16 and other.dtype == torch.bfloat16:
+        abs_err = (input.float() - other.float()).abs()
+        rel_err = abs_err / other.float().abs().clamp_min(1.0)
+        return bool(abs_err.max().item() <= 0.03125 and rel_err.max().item() <= 0.05)
+    return _original_allclose(input, other, rtol=rtol, atol=atol, equal_nan=equal_nan)
+
+
+torch.allclose = _flashrt_allclose
+
+
+def _reference_norm_rope(x, weight, cos, sin, eps=1e-6):
+    half = x.shape[-1] // 2
+    rstd = torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps)
+    normed = x.float() * rstd * weight.float()
+    lo = normed[..., :half]
+    hi = normed[..., half:]
+    out_lo = lo * cos.float() - hi * sin.float()
+    out_hi = hi * cos.float() + lo * sin.float()
+    return torch.cat([out_lo, out_hi], dim=-1).to(torch.bfloat16)
+
+
+def _reference_qkv_split_norm_rope(
+    packed_qkv, norm_q_weight, norm_k_weight, freqs_re, freqs_im, heads, head_dim, eps=1e-6
+):
+    batch, tokens, _ = packed_qkv.shape
+    dim = heads * head_dim
+    q = packed_qkv[..., :dim].reshape(batch, tokens, heads, head_dim)
+    k = packed_qkv[..., dim : 2 * dim].reshape(batch, tokens, heads, head_dim)
+    qf = q.float()
+    kf = k.float()
+    qn = qf * torch.rsqrt((qf * qf).mean(dim=(-2, -1), keepdim=True) + eps)
+    kn = kf * torch.rsqrt((kf * kf).mean(dim=(-2, -1), keepdim=True) + eps)
+    qn = qn * norm_q_weight.reshape(1, 1, heads, head_dim).float()
+    kn = kn * norm_k_weight.reshape(1, 1, heads, head_dim).float()
+
+    def rope(x):
+        xr = x[..., 0::2].float()
+        xi = x[..., 1::2].float()
+        fr = freqs_re[:tokens][None, :, None, :]
+        fi = freqs_im[:tokens][None, :, None, :]
+        out = torch.empty_like(x, dtype=torch.float32)
+        out[..., 0::2] = xr * fr - xi * fi
+        out[..., 1::2] = xr * fi + xi * fr
+        return out.to(torch.bfloat16)
+
+    return rope(qn), rope(kn)
+
+
+class QNormRopeBenchmark(Benchmark):
+    seed = 11
+
+    def _setup_heads(self, n_heads: int) -> None:
+        self.q = torch.randn(
+            (n_heads, 128), device=self.device, dtype=torch.bfloat16
+        ).contiguous()
+        self.weight = torch.randn(
+            (128,), device=self.device, dtype=torch.bfloat16
+        ).contiguous()
+        self.cos = torch.randn((64,), device=self.device, dtype=torch.bfloat16)
+        self.sin = torch.randn((64,), device=self.device, dtype=torch.bfloat16)
+        self.out = torch.empty_like(self.q)
+
+    def _reference(self):
+        return _reference_norm_rope(self.q, self.weight, self.cos, self.sin)
+
+    def setup_heads1(self):
+        self._setup_heads(1)
+
+    def benchmark_heads1(self):
+        self.kernel.q_norm_rope_bf16(
+            self.q, self.weight, self.cos, self.sin, out=self.out
+        )
+
+    def verify_heads1(self):
+        return self._reference()
+
+    def setup_heads4(self):
+        self._setup_heads(4)
+
+    def benchmark_heads4(self):
+        self.kernel.q_norm_rope_bf16(
+            self.q, self.weight, self.cos, self.sin, out=self.out
+        )
+
+    def verify_heads4(self):
+        return self._reference()
+
+    def setup_heads8(self):
+        self._setup_heads(8)
+
+    def benchmark_heads8(self):
+        self.kernel.q_norm_rope_bf16(
+            self.q, self.weight, self.cos, self.sin, out=self.out
+        )
+
+    def verify_heads8(self):
+        return self._reference()
+
+    def setup_heads16(self):
+        self._setup_heads(16)
+
+    def benchmark_heads16(self):
+        self.kernel.q_norm_rope_bf16(
+            self.q, self.weight, self.cos, self.sin, out=self.out
+        )
+
+    def verify_heads16(self):
+        return self._reference()
+
+    def setup_heads24(self):
+        self._setup_heads(24)
+
+    def benchmark_heads24(self):
+        self.kernel.q_norm_rope_bf16(
+            self.q, self.weight, self.cos, self.sin, out=self.out
+        )
+
+    def verify_heads24(self):
+        return self._reference()
+
+    def setup_heads32(self):
+        self._setup_heads(32)
+
+    def benchmark_heads32(self):
+        self.kernel.q_norm_rope_bf16(
+            self.q, self.weight, self.cos, self.sin, out=self.out
+        )
+
+    def verify_heads32(self):
+        return self._reference()
+
+    def setup_heads48(self):
+        self._setup_heads(48)
+
+    def benchmark_heads48(self):
+        self.kernel.q_norm_rope_bf16(
+            self.q, self.weight, self.cos, self.sin, out=self.out
+        )
+
+    def verify_heads48(self):
+        return self._reference()
+
+
+class QKVSplitNormRopeBenchmark(Benchmark):
+    seed = 13
+
+    def _setup_tokens(self, tokens: int) -> None:
+        self.heads = 24
+        self.head_dim = 128
+        dim = self.heads * self.head_dim
+        self.packed_qkv = torch.randn(
+            (1, tokens, 3 * dim), device=self.device, dtype=torch.bfloat16
+        ).contiguous()
+        self.norm_q_weight = torch.randn(
+            (dim,), device=self.device, dtype=torch.bfloat16
+        ).contiguous()
+        self.norm_k_weight = torch.randn(
+            (dim,), device=self.device, dtype=torch.bfloat16
+        ).contiguous()
+        self.freqs_re = torch.randn(
+            (4096, self.head_dim // 2), device=self.device, dtype=torch.float32
+        ).contiguous()
+        self.freqs_im = torch.randn(
+            (4096, self.head_dim // 2), device=self.device, dtype=torch.float32
+        ).contiguous()
+        self.q_out = torch.empty(
+            (1, tokens, self.heads, self.head_dim),
+            device=self.device,
+            dtype=torch.bfloat16,
+        )
+        self.k_out = torch.empty_like(self.q_out)
+        self.out = self.q_out
+
+    def _reference_pair(self):
+        return _reference_qkv_split_norm_rope(
+            self.packed_qkv,
+            self.norm_q_weight,
+            self.norm_k_weight,
+            self.freqs_re,
+            self.freqs_im,
+            self.heads,
+            self.head_dim,
+        )
+
+    def _reference(self):
+        return self._reference_pair()[0]
+
+    def _benchmark(self):
+        self.kernel.qkv_split_norm_rope_bf16(
+            self.packed_qkv,
+            self.norm_q_weight,
+            self.norm_k_weight,
+            self.freqs_re,
+            self.freqs_im,
+            heads=self.heads,
+            head_dim=self.head_dim,
+            q_out=self.q_out,
+            k_out=self.k_out,
+        )
+
+    def setup_tokens1(self):
+        self._setup_tokens(1)
+
+    def benchmark_tokens1(self):
+        self._benchmark()
+
+    def verify_tokens1(self):
+        return self._reference()
+
+    def setup_tokens4(self):
+        self._setup_tokens(4)
+
+    def benchmark_tokens4(self):
+        self._benchmark()
+
+    def verify_tokens4(self):
+        return self._reference()
+
+    def setup_tokens16(self):
+        self._setup_tokens(16)
+
+    def benchmark_tokens16(self):
+        self._benchmark()
+
+    def verify_tokens16(self):
+        return self._reference()
+
+    def setup_tokens64(self):
+        self._setup_tokens(64)
+
+    def benchmark_tokens64(self):
+        self._benchmark()
+
+    def verify_tokens64(self):
+        return self._reference()
+
+    def setup_tokens256(self):
+        self._setup_tokens(256)
+
+    def benchmark_tokens256(self):
+        self._benchmark()
+
+    def verify_tokens256(self):
+        return self._reference()
+
+    def setup_tokens1024(self):
+        self._setup_tokens(1024)
+
+    def benchmark_tokens1024(self):
+        self._benchmark()
+
+    def verify_tokens1024(self):
+        return self._reference()
+
+    def setup_tokens2520(self):
+        self._setup_tokens(2520)
+
+    def benchmark_tokens2520(self):
+        self._benchmark()
+
+    def verify_tokens2520(self):
+        return self._reference()
+
+    def setup_tokens4096(self):
+        self._setup_tokens(4096)
+
+    def benchmark_tokens4096(self):
+        self._benchmark()
+
+    def verify_tokens4096(self):
+        return self._reference()
+
+
+class QKVSplitNormRopeKBenchmark(QKVSplitNormRopeBenchmark):
+    seed = 14
+
+    def _setup_tokens(self, tokens: int) -> None:
+        super()._setup_tokens(tokens)
+        self.out = self.k_out
+
+    def _reference(self):
+        return self._reference_pair()[1]
+
+
+class KNormRopeVCacheBenchmark(Benchmark):
+    seed = 12
+
+    def _setup_heads(self, n_heads: int) -> None:
+        self.k = torch.randn(
+            (n_heads, 128), device=self.device, dtype=torch.bfloat16
+        ).contiguous()
+        self.v = torch.randn(
+            (n_heads, 128), device=self.device, dtype=torch.bfloat16
+        ).contiguous()
+        self.weight = torch.randn(
+            (128,), device=self.device, dtype=torch.bfloat16
+        ).contiguous()
+        self.cos = torch.randn((64,), device=self.device, dtype=torch.bfloat16)
+        self.sin = torch.randn((64,), device=self.device, dtype=torch.bfloat16)
+        self.k_out = torch.empty_like(self.k)
+        self.v_out = torch.empty_like(self.v)
+        self.out = self.k_out
+
+    def _reference(self):
+        return _reference_norm_rope(self.k, self.weight, self.cos, self.sin)
+
+    def setup_heads1(self):
+        self._setup_heads(1)
+
+    def benchmark_heads1(self):
+        self.kernel.k_norm_rope_v_cache_bf16(
+            self.k, self.v, self.weight, self.cos, self.sin,
+            k_out=self.k_out, v_out=self.v_out
+        )
+
+    def verify_heads1(self):
+        return self._reference()
+
+    def setup_heads4(self):
+        self._setup_heads(4)
+
+    def benchmark_heads4(self):
+        self.kernel.k_norm_rope_v_cache_bf16(
+            self.k, self.v, self.weight, self.cos, self.sin,
+            k_out=self.k_out, v_out=self.v_out
+        )
+
+    def verify_heads4(self):
+        return self._reference()
+
+    def setup_heads8(self):
+        self._setup_heads(8)
+
+    def benchmark_heads8(self):
+        self.kernel.k_norm_rope_v_cache_bf16(
+            self.k, self.v, self.weight, self.cos, self.sin,
+            k_out=self.k_out, v_out=self.v_out
+        )
+
+    def verify_heads8(self):
+        return self._reference()
+
+    def setup_heads16(self):
+        self._setup_heads(16)
+
+    def benchmark_heads16(self):
+        self.kernel.k_norm_rope_v_cache_bf16(
+            self.k, self.v, self.weight, self.cos, self.sin,
+            k_out=self.k_out, v_out=self.v_out
+        )
+
+    def verify_heads16(self):
+        return self._reference()
+
+    def setup_heads24(self):
+        self._setup_heads(24)
+
+    def benchmark_heads24(self):
+        self.kernel.k_norm_rope_v_cache_bf16(
+            self.k, self.v, self.weight, self.cos, self.sin,
+            k_out=self.k_out, v_out=self.v_out
+        )
+
+    def verify_heads24(self):
+        return self._reference()
+
+    def setup_heads32(self):
+        self._setup_heads(32)
+
+    def benchmark_heads32(self):
+        self.kernel.k_norm_rope_v_cache_bf16(
+            self.k, self.v, self.weight, self.cos, self.sin,
+            k_out=self.k_out, v_out=self.v_out
+        )
+
+    def verify_heads32(self):
+        return self._reference()
+
+    def setup_heads48(self):
+        self._setup_heads(48)
+
+    def benchmark_heads48(self):
+        self.kernel.k_norm_rope_v_cache_bf16(
+            self.k, self.v, self.weight, self.cos, self.sin,
+            k_out=self.k_out, v_out=self.v_out
+        )
+
+    def verify_heads48(self):
+        return self._reference()

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

@@ -0,0 +1,168 @@
+"""FlashRT VLA, vision, video, and diffusion kernels."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.register_fake(add_op_namespace_prefix("q_norm_rope_bf16"))
+def _q_norm_rope_bf16_fake(
+    q: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if out.shape != q.shape:
+        raise RuntimeError("out shape must match q shape")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("k_norm_rope_v_cache_bf16"))
+def _k_norm_rope_v_cache_bf16_fake(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if k_out.shape != k.shape or v_out.shape != v.shape:
+        raise RuntimeError("k_out and v_out shapes must match k and v")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qkv_split_norm_rope_bf16"))
+def _qkv_split_norm_rope_bf16_fake(
+    packed_qkv: torch.Tensor,
+    norm_q_weight: torch.Tensor,
+    norm_k_weight: torch.Tensor,
+    freqs_re: torch.Tensor,
+    freqs_im: torch.Tensor,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    seq_len: int,
+    eps: float = 1e-6,
+) -> None:
+    expected = (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim)
+    if q_out.shape != expected or k_out.shape != expected:
+        raise RuntimeError(
+            "q_out and k_out shapes must be (batch, tokens, heads, head_dim)"
+        )
+    return None
+
+
+def q_norm_rope_bf16(
+    q: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    *,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Compute per-row RMSNorm plus rotate-half RoPE for BF16 Q heads.
+
+    ``q`` must be contiguous BF16 with shape ``(..., 128)``. ``weight`` must
+    have shape ``(128,)``. ``cos`` and ``sin`` must have shape ``(64,)`` for
+    the current decode position. If ``out`` is omitted, a BF16 tensor with the
+    same shape as ``q`` is allocated.
+    """
+
+    if out is None:
+        out = torch.empty_like(q)
+    ops.q_norm_rope_bf16(q, weight, cos, sin, out, eps)
+    return out
+
+
+def k_norm_rope_v_cache_bf16(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    k_out: Optional[torch.Tensor] = None,
+    v_out: Optional[torch.Tensor] = None,
+    *,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute K RMSNorm plus RoPE and copy V into cache-shaped outputs.
+
+    ``k`` and ``v`` must be contiguous BF16 tensors with shape ``(..., 128)``.
+    ``weight`` must have shape ``(128,)``. ``cos`` and ``sin`` must have shape
+    ``(64,)`` for the current decode position. If outputs are omitted, BF16
+    tensors with the same shape as ``k`` are allocated.
+    """
+
+    if k_out is None:
+        k_out = torch.empty_like(k)
+    if v_out is None:
+        v_out = torch.empty_like(v)
+    ops.k_norm_rope_v_cache_bf16(k, v, weight, cos, sin, k_out, v_out, eps)
+    return k_out, v_out
+
+
+def qkv_split_norm_rope_bf16(
+    packed_qkv: torch.Tensor,
+    norm_q_weight: torch.Tensor,
+    norm_k_weight: torch.Tensor,
+    freqs_re: torch.Tensor,
+    freqs_im: torch.Tensor,
+    *,
+    heads: int,
+    head_dim: int,
+    seq_len: Optional[int] = None,
+    q_out: Optional[torch.Tensor] = None,
+    k_out: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Split packed QKV, RMS-normalize Q/K, and apply interleaved RoPE.
+
+    ``packed_qkv`` must be contiguous BF16 with shape
+    ``(batch, tokens, 3 * heads * head_dim)``. ``norm_q_weight`` and
+    ``norm_k_weight`` must have shape ``(heads * head_dim,)``. ``freqs_re`` and
+    ``freqs_im`` must be contiguous FP32 tensors with shape
+    ``(seq_len_table, head_dim // 2)``. Outputs have shape
+    ``(batch, tokens, heads, head_dim)``.
+    """
+
+    if seq_len is None:
+        seq_len = packed_qkv.shape[1]
+    if q_out is None:
+        q_out = torch.empty(
+            (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim),
+            device=packed_qkv.device,
+            dtype=packed_qkv.dtype,
+        )
+    if k_out is None:
+        k_out = torch.empty_like(q_out)
+    ops.qkv_split_norm_rope_bf16(
+        packed_qkv,
+        norm_q_weight,
+        norm_k_weight,
+        freqs_re,
+        freqs_im,
+        q_out,
+        k_out,
+        heads,
+        head_dim,
+        seq_len,
+        eps,
+    )
+    return q_out, k_out
+
+
+__all__ = [
+    "q_norm_rope_bf16",
+    "k_norm_rope_v_cache_bf16",
+    "qkv_split_norm_rope_bf16",
+]

build/torch211-cxx11-cu126-x86_64-linux/_flashrt_vla_video_cuda_c4d802d.abi3.so

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

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

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

build/torch211-cxx11-cu126-x86_64-linux/flashrt_vla_video/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

@@ -0,0 +1,20 @@
+{
+  "name": "flashrt-vla-video",
+  "id": "_flashrt_vla_video_cuda_c4d802d",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "7.0",
+      "7.2",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0+PTX"
+    ]
+  }
+}

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

@@ -0,0 +1,168 @@
+"""FlashRT VLA, vision, video, and diffusion kernels."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.register_fake(add_op_namespace_prefix("q_norm_rope_bf16"))
+def _q_norm_rope_bf16_fake(
+    q: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if out.shape != q.shape:
+        raise RuntimeError("out shape must match q shape")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("k_norm_rope_v_cache_bf16"))
+def _k_norm_rope_v_cache_bf16_fake(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if k_out.shape != k.shape or v_out.shape != v.shape:
+        raise RuntimeError("k_out and v_out shapes must match k and v")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qkv_split_norm_rope_bf16"))
+def _qkv_split_norm_rope_bf16_fake(
+    packed_qkv: torch.Tensor,
+    norm_q_weight: torch.Tensor,
+    norm_k_weight: torch.Tensor,
+    freqs_re: torch.Tensor,
+    freqs_im: torch.Tensor,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    seq_len: int,
+    eps: float = 1e-6,
+) -> None:
+    expected = (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim)
+    if q_out.shape != expected or k_out.shape != expected:
+        raise RuntimeError(
+            "q_out and k_out shapes must be (batch, tokens, heads, head_dim)"
+        )
+    return None
+
+
+def q_norm_rope_bf16(
+    q: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    *,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Compute per-row RMSNorm plus rotate-half RoPE for BF16 Q heads.
+
+    ``q`` must be contiguous BF16 with shape ``(..., 128)``. ``weight`` must
+    have shape ``(128,)``. ``cos`` and ``sin`` must have shape ``(64,)`` for
+    the current decode position. If ``out`` is omitted, a BF16 tensor with the
+    same shape as ``q`` is allocated.
+    """
+
+    if out is None:
+        out = torch.empty_like(q)
+    ops.q_norm_rope_bf16(q, weight, cos, sin, out, eps)
+    return out
+
+
+def k_norm_rope_v_cache_bf16(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    k_out: Optional[torch.Tensor] = None,
+    v_out: Optional[torch.Tensor] = None,
+    *,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute K RMSNorm plus RoPE and copy V into cache-shaped outputs.
+
+    ``k`` and ``v`` must be contiguous BF16 tensors with shape ``(..., 128)``.
+    ``weight`` must have shape ``(128,)``. ``cos`` and ``sin`` must have shape
+    ``(64,)`` for the current decode position. If outputs are omitted, BF16
+    tensors with the same shape as ``k`` are allocated.
+    """
+
+    if k_out is None:
+        k_out = torch.empty_like(k)
+    if v_out is None:
+        v_out = torch.empty_like(v)
+    ops.k_norm_rope_v_cache_bf16(k, v, weight, cos, sin, k_out, v_out, eps)
+    return k_out, v_out
+
+
+def qkv_split_norm_rope_bf16(
+    packed_qkv: torch.Tensor,
+    norm_q_weight: torch.Tensor,
+    norm_k_weight: torch.Tensor,
+    freqs_re: torch.Tensor,
+    freqs_im: torch.Tensor,
+    *,
+    heads: int,
+    head_dim: int,
+    seq_len: Optional[int] = None,
+    q_out: Optional[torch.Tensor] = None,
+    k_out: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Split packed QKV, RMS-normalize Q/K, and apply interleaved RoPE.
+
+    ``packed_qkv`` must be contiguous BF16 with shape
+    ``(batch, tokens, 3 * heads * head_dim)``. ``norm_q_weight`` and
+    ``norm_k_weight`` must have shape ``(heads * head_dim,)``. ``freqs_re`` and
+    ``freqs_im`` must be contiguous FP32 tensors with shape
+    ``(seq_len_table, head_dim // 2)``. Outputs have shape
+    ``(batch, tokens, heads, head_dim)``.
+    """
+
+    if seq_len is None:
+        seq_len = packed_qkv.shape[1]
+    if q_out is None:
+        q_out = torch.empty(
+            (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim),
+            device=packed_qkv.device,
+            dtype=packed_qkv.dtype,
+        )
+    if k_out is None:
+        k_out = torch.empty_like(q_out)
+    ops.qkv_split_norm_rope_bf16(
+        packed_qkv,
+        norm_q_weight,
+        norm_k_weight,
+        freqs_re,
+        freqs_im,
+        q_out,
+        k_out,
+        heads,
+        head_dim,
+        seq_len,
+        eps,
+    )
+    return q_out, k_out
+
+
+__all__ = [
+    "q_norm_rope_bf16",
+    "k_norm_rope_v_cache_bf16",
+    "qkv_split_norm_rope_bf16",
+]

build/torch211-cxx11-cu128-x86_64-linux/_flashrt_vla_video_cuda_c4d802d.abi3.so

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

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

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

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

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

@@ -0,0 +1,23 @@
+{
+  "name": "flashrt-vla-video",
+  "id": "_flashrt_vla_video_cuda_c4d802d",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "10.1",
+      "12.0+PTX",
+      "7.0",
+      "7.2",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

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

@@ -0,0 +1,168 @@
+"""FlashRT VLA, vision, video, and diffusion kernels."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.register_fake(add_op_namespace_prefix("q_norm_rope_bf16"))
+def _q_norm_rope_bf16_fake(
+    q: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if out.shape != q.shape:
+        raise RuntimeError("out shape must match q shape")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("k_norm_rope_v_cache_bf16"))
+def _k_norm_rope_v_cache_bf16_fake(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if k_out.shape != k.shape or v_out.shape != v.shape:
+        raise RuntimeError("k_out and v_out shapes must match k and v")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qkv_split_norm_rope_bf16"))
+def _qkv_split_norm_rope_bf16_fake(
+    packed_qkv: torch.Tensor,
+    norm_q_weight: torch.Tensor,
+    norm_k_weight: torch.Tensor,
+    freqs_re: torch.Tensor,
+    freqs_im: torch.Tensor,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    seq_len: int,
+    eps: float = 1e-6,
+) -> None:
+    expected = (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim)
+    if q_out.shape != expected or k_out.shape != expected:
+        raise RuntimeError(
+            "q_out and k_out shapes must be (batch, tokens, heads, head_dim)"
+        )
+    return None
+
+
+def q_norm_rope_bf16(
+    q: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    *,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Compute per-row RMSNorm plus rotate-half RoPE for BF16 Q heads.
+
+    ``q`` must be contiguous BF16 with shape ``(..., 128)``. ``weight`` must
+    have shape ``(128,)``. ``cos`` and ``sin`` must have shape ``(64,)`` for
+    the current decode position. If ``out`` is omitted, a BF16 tensor with the
+    same shape as ``q`` is allocated.
+    """
+
+    if out is None:
+        out = torch.empty_like(q)
+    ops.q_norm_rope_bf16(q, weight, cos, sin, out, eps)
+    return out
+
+
+def k_norm_rope_v_cache_bf16(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    k_out: Optional[torch.Tensor] = None,
+    v_out: Optional[torch.Tensor] = None,
+    *,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute K RMSNorm plus RoPE and copy V into cache-shaped outputs.
+
+    ``k`` and ``v`` must be contiguous BF16 tensors with shape ``(..., 128)``.
+    ``weight`` must have shape ``(128,)``. ``cos`` and ``sin`` must have shape
+    ``(64,)`` for the current decode position. If outputs are omitted, BF16
+    tensors with the same shape as ``k`` are allocated.
+    """
+
+    if k_out is None:
+        k_out = torch.empty_like(k)
+    if v_out is None:
+        v_out = torch.empty_like(v)
+    ops.k_norm_rope_v_cache_bf16(k, v, weight, cos, sin, k_out, v_out, eps)
+    return k_out, v_out
+
+
+def qkv_split_norm_rope_bf16(
+    packed_qkv: torch.Tensor,
+    norm_q_weight: torch.Tensor,
+    norm_k_weight: torch.Tensor,
+    freqs_re: torch.Tensor,
+    freqs_im: torch.Tensor,
+    *,
+    heads: int,
+    head_dim: int,
+    seq_len: Optional[int] = None,
+    q_out: Optional[torch.Tensor] = None,
+    k_out: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Split packed QKV, RMS-normalize Q/K, and apply interleaved RoPE.
+
+    ``packed_qkv`` must be contiguous BF16 with shape
+    ``(batch, tokens, 3 * heads * head_dim)``. ``norm_q_weight`` and
+    ``norm_k_weight`` must have shape ``(heads * head_dim,)``. ``freqs_re`` and
+    ``freqs_im`` must be contiguous FP32 tensors with shape
+    ``(seq_len_table, head_dim // 2)``. Outputs have shape
+    ``(batch, tokens, heads, head_dim)``.
+    """
+
+    if seq_len is None:
+        seq_len = packed_qkv.shape[1]
+    if q_out is None:
+        q_out = torch.empty(
+            (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim),
+            device=packed_qkv.device,
+            dtype=packed_qkv.dtype,
+        )
+    if k_out is None:
+        k_out = torch.empty_like(q_out)
+    ops.qkv_split_norm_rope_bf16(
+        packed_qkv,
+        norm_q_weight,
+        norm_k_weight,
+        freqs_re,
+        freqs_im,
+        q_out,
+        k_out,
+        heads,
+        head_dim,
+        seq_len,
+        eps,
+    )
+    return q_out, k_out
+
+
+__all__ = [
+    "q_norm_rope_bf16",
+    "k_norm_rope_v_cache_bf16",
+    "qkv_split_norm_rope_bf16",
+]

build/torch211-cxx11-cu130-x86_64-linux/_flashrt_vla_video_cuda_c4d802d.abi3.so

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

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

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

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

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

@@ -0,0 +1,22 @@
+{
+  "name": "flashrt-vla-video",
+  "id": "_flashrt_vla_video_cuda_c4d802d",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

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

@@ -0,0 +1,168 @@
+"""FlashRT VLA, vision, video, and diffusion kernels."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.register_fake(add_op_namespace_prefix("q_norm_rope_bf16"))
+def _q_norm_rope_bf16_fake(
+    q: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if out.shape != q.shape:
+        raise RuntimeError("out shape must match q shape")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("k_norm_rope_v_cache_bf16"))
+def _k_norm_rope_v_cache_bf16_fake(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if k_out.shape != k.shape or v_out.shape != v.shape:
+        raise RuntimeError("k_out and v_out shapes must match k and v")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qkv_split_norm_rope_bf16"))
+def _qkv_split_norm_rope_bf16_fake(
+    packed_qkv: torch.Tensor,
+    norm_q_weight: torch.Tensor,
+    norm_k_weight: torch.Tensor,
+    freqs_re: torch.Tensor,
+    freqs_im: torch.Tensor,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    seq_len: int,
+    eps: float = 1e-6,
+) -> None:
+    expected = (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim)
+    if q_out.shape != expected or k_out.shape != expected:
+        raise RuntimeError(
+            "q_out and k_out shapes must be (batch, tokens, heads, head_dim)"
+        )
+    return None
+
+
+def q_norm_rope_bf16(
+    q: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    *,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Compute per-row RMSNorm plus rotate-half RoPE for BF16 Q heads.
+
+    ``q`` must be contiguous BF16 with shape ``(..., 128)``. ``weight`` must
+    have shape ``(128,)``. ``cos`` and ``sin`` must have shape ``(64,)`` for
+    the current decode position. If ``out`` is omitted, a BF16 tensor with the
+    same shape as ``q`` is allocated.
+    """
+
+    if out is None:
+        out = torch.empty_like(q)
+    ops.q_norm_rope_bf16(q, weight, cos, sin, out, eps)
+    return out
+
+
+def k_norm_rope_v_cache_bf16(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    k_out: Optional[torch.Tensor] = None,
+    v_out: Optional[torch.Tensor] = None,
+    *,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute K RMSNorm plus RoPE and copy V into cache-shaped outputs.
+
+    ``k`` and ``v`` must be contiguous BF16 tensors with shape ``(..., 128)``.
+    ``weight`` must have shape ``(128,)``. ``cos`` and ``sin`` must have shape
+    ``(64,)`` for the current decode position. If outputs are omitted, BF16
+    tensors with the same shape as ``k`` are allocated.
+    """
+
+    if k_out is None:
+        k_out = torch.empty_like(k)
+    if v_out is None:
+        v_out = torch.empty_like(v)
+    ops.k_norm_rope_v_cache_bf16(k, v, weight, cos, sin, k_out, v_out, eps)
+    return k_out, v_out
+
+
+def qkv_split_norm_rope_bf16(
+    packed_qkv: torch.Tensor,
+    norm_q_weight: torch.Tensor,
+    norm_k_weight: torch.Tensor,
+    freqs_re: torch.Tensor,
+    freqs_im: torch.Tensor,
+    *,
+    heads: int,
+    head_dim: int,
+    seq_len: Optional[int] = None,
+    q_out: Optional[torch.Tensor] = None,
+    k_out: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Split packed QKV, RMS-normalize Q/K, and apply interleaved RoPE.
+
+    ``packed_qkv`` must be contiguous BF16 with shape
+    ``(batch, tokens, 3 * heads * head_dim)``. ``norm_q_weight`` and
+    ``norm_k_weight`` must have shape ``(heads * head_dim,)``. ``freqs_re`` and
+    ``freqs_im`` must be contiguous FP32 tensors with shape
+    ``(seq_len_table, head_dim // 2)``. Outputs have shape
+    ``(batch, tokens, heads, head_dim)``.
+    """
+
+    if seq_len is None:
+        seq_len = packed_qkv.shape[1]
+    if q_out is None:
+        q_out = torch.empty(
+            (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim),
+            device=packed_qkv.device,
+            dtype=packed_qkv.dtype,
+        )
+    if k_out is None:
+        k_out = torch.empty_like(q_out)
+    ops.qkv_split_norm_rope_bf16(
+        packed_qkv,
+        norm_q_weight,
+        norm_k_weight,
+        freqs_re,
+        freqs_im,
+        q_out,
+        k_out,
+        heads,
+        head_dim,
+        seq_len,
+        eps,
+    )
+    return q_out, k_out
+
+
+__all__ = [
+    "q_norm_rope_bf16",
+    "k_norm_rope_v_cache_bf16",
+    "qkv_split_norm_rope_bf16",
+]

build/torch212-cxx11-cu126-x86_64-linux/_flashrt_vla_video_cuda_c4d802d.abi3.so

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

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

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

build/torch212-cxx11-cu126-x86_64-linux/flashrt_vla_video/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

@@ -0,0 +1,20 @@
+{
+  "name": "flashrt-vla-video",
+  "id": "_flashrt_vla_video_cuda_c4d802d",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "7.0",
+      "7.2",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0+PTX"
+    ]
+  }
+}

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

@@ -0,0 +1,168 @@
+"""FlashRT VLA, vision, video, and diffusion kernels."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.register_fake(add_op_namespace_prefix("q_norm_rope_bf16"))
+def _q_norm_rope_bf16_fake(
+    q: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if out.shape != q.shape:
+        raise RuntimeError("out shape must match q shape")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("k_norm_rope_v_cache_bf16"))
+def _k_norm_rope_v_cache_bf16_fake(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if k_out.shape != k.shape or v_out.shape != v.shape:
+        raise RuntimeError("k_out and v_out shapes must match k and v")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qkv_split_norm_rope_bf16"))
+def _qkv_split_norm_rope_bf16_fake(
+    packed_qkv: torch.Tensor,
+    norm_q_weight: torch.Tensor,
+    norm_k_weight: torch.Tensor,
+    freqs_re: torch.Tensor,
+    freqs_im: torch.Tensor,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    seq_len: int,
+    eps: float = 1e-6,
+) -> None:
+    expected = (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim)
+    if q_out.shape != expected or k_out.shape != expected:
+        raise RuntimeError(
+            "q_out and k_out shapes must be (batch, tokens, heads, head_dim)"
+        )
+    return None
+
+
+def q_norm_rope_bf16(
+    q: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    *,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Compute per-row RMSNorm plus rotate-half RoPE for BF16 Q heads.
+
+    ``q`` must be contiguous BF16 with shape ``(..., 128)``. ``weight`` must
+    have shape ``(128,)``. ``cos`` and ``sin`` must have shape ``(64,)`` for
+    the current decode position. If ``out`` is omitted, a BF16 tensor with the
+    same shape as ``q`` is allocated.
+    """
+
+    if out is None:
+        out = torch.empty_like(q)
+    ops.q_norm_rope_bf16(q, weight, cos, sin, out, eps)
+    return out
+
+
+def k_norm_rope_v_cache_bf16(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    k_out: Optional[torch.Tensor] = None,
+    v_out: Optional[torch.Tensor] = None,
+    *,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute K RMSNorm plus RoPE and copy V into cache-shaped outputs.
+
+    ``k`` and ``v`` must be contiguous BF16 tensors with shape ``(..., 128)``.
+    ``weight`` must have shape ``(128,)``. ``cos`` and ``sin`` must have shape
+    ``(64,)`` for the current decode position. If outputs are omitted, BF16
+    tensors with the same shape as ``k`` are allocated.
+    """
+
+    if k_out is None:
+        k_out = torch.empty_like(k)
+    if v_out is None:
+        v_out = torch.empty_like(v)
+    ops.k_norm_rope_v_cache_bf16(k, v, weight, cos, sin, k_out, v_out, eps)
+    return k_out, v_out
+
+
+def qkv_split_norm_rope_bf16(
+    packed_qkv: torch.Tensor,
+    norm_q_weight: torch.Tensor,
+    norm_k_weight: torch.Tensor,
+    freqs_re: torch.Tensor,
+    freqs_im: torch.Tensor,
+    *,
+    heads: int,
+    head_dim: int,
+    seq_len: Optional[int] = None,
+    q_out: Optional[torch.Tensor] = None,
+    k_out: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Split packed QKV, RMS-normalize Q/K, and apply interleaved RoPE.
+
+    ``packed_qkv`` must be contiguous BF16 with shape
+    ``(batch, tokens, 3 * heads * head_dim)``. ``norm_q_weight`` and
+    ``norm_k_weight`` must have shape ``(heads * head_dim,)``. ``freqs_re`` and
+    ``freqs_im`` must be contiguous FP32 tensors with shape
+    ``(seq_len_table, head_dim // 2)``. Outputs have shape
+    ``(batch, tokens, heads, head_dim)``.
+    """
+
+    if seq_len is None:
+        seq_len = packed_qkv.shape[1]
+    if q_out is None:
+        q_out = torch.empty(
+            (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim),
+            device=packed_qkv.device,
+            dtype=packed_qkv.dtype,
+        )
+    if k_out is None:
+        k_out = torch.empty_like(q_out)
+    ops.qkv_split_norm_rope_bf16(
+        packed_qkv,
+        norm_q_weight,
+        norm_k_weight,
+        freqs_re,
+        freqs_im,
+        q_out,
+        k_out,
+        heads,
+        head_dim,
+        seq_len,
+        eps,
+    )
+    return q_out, k_out
+
+
+__all__ = [
+    "q_norm_rope_bf16",
+    "k_norm_rope_v_cache_bf16",
+    "qkv_split_norm_rope_bf16",
+]

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

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

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

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

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

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

@@ -0,0 +1,22 @@
+{
+  "name": "flashrt-vla-video",
+  "id": "_flashrt_vla_video_cuda_c4d802d",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}

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

@@ -0,0 +1,168 @@
+"""FlashRT VLA, vision, video, and diffusion kernels."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+from ._ops import add_op_namespace_prefix, ops
+
+
+@torch.library.register_fake(add_op_namespace_prefix("q_norm_rope_bf16"))
+def _q_norm_rope_bf16_fake(
+    q: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if out.shape != q.shape:
+        raise RuntimeError("out shape must match q shape")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("k_norm_rope_v_cache_bf16"))
+def _k_norm_rope_v_cache_bf16_fake(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    if k_out.shape != k.shape or v_out.shape != v.shape:
+        raise RuntimeError("k_out and v_out shapes must match k and v")
+    return None
+
+
+@torch.library.register_fake(add_op_namespace_prefix("qkv_split_norm_rope_bf16"))
+def _qkv_split_norm_rope_bf16_fake(
+    packed_qkv: torch.Tensor,
+    norm_q_weight: torch.Tensor,
+    norm_k_weight: torch.Tensor,
+    freqs_re: torch.Tensor,
+    freqs_im: torch.Tensor,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    seq_len: int,
+    eps: float = 1e-6,
+) -> None:
+    expected = (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim)
+    if q_out.shape != expected or k_out.shape != expected:
+        raise RuntimeError(
+            "q_out and k_out shapes must be (batch, tokens, heads, head_dim)"
+        )
+    return None
+
+
+def q_norm_rope_bf16(
+    q: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    out: Optional[torch.Tensor] = None,
+    *,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Compute per-row RMSNorm plus rotate-half RoPE for BF16 Q heads.
+
+    ``q`` must be contiguous BF16 with shape ``(..., 128)``. ``weight`` must
+    have shape ``(128,)``. ``cos`` and ``sin`` must have shape ``(64,)`` for
+    the current decode position. If ``out`` is omitted, a BF16 tensor with the
+    same shape as ``q`` is allocated.
+    """
+
+    if out is None:
+        out = torch.empty_like(q)
+    ops.q_norm_rope_bf16(q, weight, cos, sin, out, eps)
+    return out
+
+
+def k_norm_rope_v_cache_bf16(
+    k: torch.Tensor,
+    v: torch.Tensor,
+    weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    k_out: Optional[torch.Tensor] = None,
+    v_out: Optional[torch.Tensor] = None,
+    *,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute K RMSNorm plus RoPE and copy V into cache-shaped outputs.
+
+    ``k`` and ``v`` must be contiguous BF16 tensors with shape ``(..., 128)``.
+    ``weight`` must have shape ``(128,)``. ``cos`` and ``sin`` must have shape
+    ``(64,)`` for the current decode position. If outputs are omitted, BF16
+    tensors with the same shape as ``k`` are allocated.
+    """
+
+    if k_out is None:
+        k_out = torch.empty_like(k)
+    if v_out is None:
+        v_out = torch.empty_like(v)
+    ops.k_norm_rope_v_cache_bf16(k, v, weight, cos, sin, k_out, v_out, eps)
+    return k_out, v_out
+
+
+def qkv_split_norm_rope_bf16(
+    packed_qkv: torch.Tensor,
+    norm_q_weight: torch.Tensor,
+    norm_k_weight: torch.Tensor,
+    freqs_re: torch.Tensor,
+    freqs_im: torch.Tensor,
+    *,
+    heads: int,
+    head_dim: int,
+    seq_len: Optional[int] = None,
+    q_out: Optional[torch.Tensor] = None,
+    k_out: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Split packed QKV, RMS-normalize Q/K, and apply interleaved RoPE.
+
+    ``packed_qkv`` must be contiguous BF16 with shape
+    ``(batch, tokens, 3 * heads * head_dim)``. ``norm_q_weight`` and
+    ``norm_k_weight`` must have shape ``(heads * head_dim,)``. ``freqs_re`` and
+    ``freqs_im`` must be contiguous FP32 tensors with shape
+    ``(seq_len_table, head_dim // 2)``. Outputs have shape
+    ``(batch, tokens, heads, head_dim)``.
+    """
+
+    if seq_len is None:
+        seq_len = packed_qkv.shape[1]
+    if q_out is None:
+        q_out = torch.empty(
+            (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim),
+            device=packed_qkv.device,
+            dtype=packed_qkv.dtype,
+        )
+    if k_out is None:
+        k_out = torch.empty_like(q_out)
+    ops.qkv_split_norm_rope_bf16(
+        packed_qkv,
+        norm_q_weight,
+        norm_k_weight,
+        freqs_re,
+        freqs_im,
+        q_out,
+        k_out,
+        heads,
+        head_dim,
+        seq_len,
+        eps,
+    )
+    return q_out, k_out
+
+
+__all__ = [
+    "q_norm_rope_bf16",
+    "k_norm_rope_v_cache_bf16",
+    "qkv_split_norm_rope_bf16",
+]

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

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

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

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

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

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

@@ -0,0 +1,22 @@
+{
+  "name": "flashrt-vla-video",
+  "id": "_flashrt_vla_video_cuda_c4d802d",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "11.0",
+      "12.0",
+      "12.1+PTX",
+      "7.5",
+      "8.0",
+      "8.6",
+      "8.7",
+      "8.9",
+      "9.0"
+    ]
+  }
+}