Add torch211-cxx11-cu130-aarch64-linux SM110 artifact

Browse files

Files changed (5) hide show

build/torch211-cxx11-cu130-aarch64-linux/__init__.py +168 -0
build/torch211-cxx11-cu130-aarch64-linux/_flashrt_vla_video_cuda_c4d802d.abi3.so +3 -0
build/torch211-cxx11-cu130-aarch64-linux/_ops.py +6 -0
build/torch211-cxx11-cu130-aarch64-linux/flashrt_vla_video/__init__.py +14 -0
build/torch211-cxx11-cu130-aarch64-linux/metadata.json +22 -0

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

	@@ -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-aarch64-linux/_flashrt_vla_video_cuda_c4d802d.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

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

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

	@@ -0,0 +1,6 @@

+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):
+    return f"_flashrt_vla_video_cuda_c4d802d::{op_name}"

build/torch211-cxx11-cu130-aarch64-linux/flashrt_vla_video/__init__.py ADDED Viewed

	@@ -0,0 +1,14 @@

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

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

	@@ -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": [
+      "11.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "4/S2aAbFffcUrOnK9JO9NS944YoogdJSwqdRXWU7VLo=",
+      "_flashrt_vla_video_cuda_c4d802d.abi3.so": "wxmqAJfV3Zv6xtn3T0PFR6A/BWKW6C/+l30iyqH2k5Y=",
+      "_ops.py": "NhuyMSqzjMK5Nt+2BoUvQhZOFwXajZ3PMDiBmyqABA8=",
+      "flashrt_vla_video/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
+    }
+  }
+}