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

Browse files

Files changed (4) hide show

build/torch211-cxx11-cu130-aarch64-linux/__init__.py +130 -0
build/torch211-cxx11-cu130-aarch64-linux/_ops.py +3 -3
build/torch211-cxx11-cu130-aarch64-linux/flashrt_qkv_cache_rope_test.abi3.so +3 -0
build/torch211-cxx11-cu130-aarch64-linux/metadata.json +14 -4

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

@@ -206,6 +206,57 @@ def _qkv_split_rope_kvcache_bf16_fake(
     return None
 @torch.library.register_fake(add_op_namespace_prefix("qkv_split_bf16"))
 def _qkv_split_bf16_fake(
     packed_qkv: torch.Tensor,
@@ -500,6 +551,83 @@ def qkv_split_rope_kvcache_bf16(
     return q_out, k_cache, v_cache
 def qkv_split_per_head_norm_rope_bf16(
     packed_qkv: torch.Tensor,
     q_norm_weight: torch.Tensor,
@@ -867,6 +995,8 @@ __all__ = [
     "qkv_split_bias_rope_fp16",
     "qkv_split_bf16",
     "qkv_split_rope_kvcache_bf16",
     "qkv_split_norm_rope_bf16",
     "qkv_split_bias_norm_rope_v_bf16",
     "qkv_split_bias_norm_rope_v_cat_bf16",

     return None
+@torch.library.register_fake(add_op_namespace_prefix("qkv_split_rope_kvcache_fp16"))
+def _qkv_split_rope_kvcache_fp16_fake(
+    packed_qkv: torch.Tensor,
+    rope: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    cache_offset: int,
+    device_position: torch.Tensor | None,
+    q_out: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+) -> None:
+    _check_packed_gqa_qkv(packed_qkv, q_heads, kv_heads, head_dim)
+    batch, seq_len, _ = packed_qkv.shape
+    if packed_qkv.dtype != torch.float16 or rope.dtype != torch.float16:
+        raise RuntimeError("packed_qkv and rope must have dtype torch.float16")
+    if rope.dim() != 2 or rope.shape[0] < seq_len or rope.shape[1] != head_dim:
+        raise RuntimeError("rope must have shape (>= seq_len, head_dim)")
+    if q_out.shape != (batch, seq_len, q_heads, head_dim):
+        raise RuntimeError("q_out has an invalid shape")
+    if k_cache.dim() != 4 or k_cache.shape[0] != batch or k_cache.shape[2:] != (kv_heads, head_dim):
+        raise RuntimeError("k_cache must have shape (batch, max_seq_len, kv_heads, head_dim)")
+    if v_cache.shape != k_cache.shape:
+        raise RuntimeError("v_cache must have the same shape as k_cache")
+    if cache_offset < 0:
+        raise RuntimeError("cache_offset must be non-negative")
+    if device_position is not None and (device_position.dtype != torch.int32 or device_position.numel() != 1):
+        raise RuntimeError("device_position must contain one int32 element")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qk_norm_rope_strided_bf16"))
+def _qk_norm_rope_strided_bf16_fake(
+    q_in, k_in, q_weight, k_weight, cos, sin,
+    q_heads: int, k_heads: int, eps: float, q_out, k_out,
+) -> None:
+    rows = q_in.shape[0]
+    if q_in.dim() != 2 or k_in.dim() != 2 or k_in.shape[0] != rows:
+        raise RuntimeError("q_in and k_in must be rank-2 with the same row count")
+    if q_in.shape[1] < q_heads * 128 or k_in.shape[1] < k_heads * 128:
+        raise RuntimeError("input row widths must cover heads * 128")
+    if q_weight.shape != (128,) or k_weight.shape != (128,):
+        raise RuntimeError("Q/K weights must have shape (128,)")
+    if cos.shape != (rows, 128) or sin.shape != cos.shape:
+        raise RuntimeError("cos and sin must have shape (rows, 128)")
+    if q_out.shape != (rows, q_heads, 128) or k_out.shape != (rows, k_heads, 128):
+        raise RuntimeError("output shapes must be (rows, heads, 128)")
+    return None
 @torch.library.register_fake(add_op_namespace_prefix("qkv_split_bf16"))
 def _qkv_split_bf16_fake(
     packed_qkv: torch.Tensor,
     return q_out, k_cache, v_cache
+def qkv_split_rope_kvcache_fp16(
+    packed_qkv: torch.Tensor,
+    rope: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    cache_offset: int = 0,
+    device_position: torch.Tensor | None = None,
+    q_out: torch.Tensor | None = None,
+    k_cache: torch.Tensor | None = None,
+    v_cache: torch.Tensor | None = None,
+    max_seq_len: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Split FP16 GQA QKV, apply interleaved RoPE, and update K/V caches.
+    ``device_position`` is an optional CUDA int32 scalar used by CUDA Graph
+    decode loops. When provided, cache rows start at
+    ``cache_offset + device_position[0]`` without a host synchronization.
+    """
+    batch, seq_len, _ = packed_qkv.shape
+    if q_out is None:
+        q_out = torch.empty(
+            (batch, seq_len, q_heads, head_dim),
+            device=packed_qkv.device,
+            dtype=torch.float16,
+        )
+    if k_cache is None or v_cache is None:
+        if max_seq_len is None:
+            if device_position is not None:
+                raise ValueError("max_seq_len is required when device_position is used")
+            max_seq_len = cache_offset + seq_len
+        cache_shape = (batch, int(max_seq_len), kv_heads, head_dim)
+        if k_cache is None:
+            k_cache = torch.empty(cache_shape, device=packed_qkv.device, dtype=torch.float16)
+        if v_cache is None:
+            v_cache = torch.empty(cache_shape, device=packed_qkv.device, dtype=torch.float16)
+    ops.qkv_split_rope_kvcache_fp16(
+        packed_qkv,
+        rope,
+        int(q_heads),
+        int(kv_heads),
+        int(head_dim),
+        int(cache_offset),
+        device_position,
+        q_out,
+        k_cache,
+        v_cache,
+    )
+    return q_out, k_cache, v_cache
+def qk_norm_rope_strided_bf16(
+    q_in: torch.Tensor,
+    k_in: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    q_heads: int,
+    k_heads: int,
+    eps: float = 1e-6,
+    q_out: torch.Tensor | None = None,
+    k_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Per-head RMSNorm and rotate-half RoPE from padded Q/K row strides."""
+    rows = q_in.shape[0]
+    if q_out is None:
+        q_out = torch.empty((rows, q_heads, 128), device=q_in.device, dtype=torch.bfloat16)
+    if k_out is None:
+        k_out = torch.empty((rows, k_heads, 128), device=k_in.device, dtype=torch.bfloat16)
+    ops.qk_norm_rope_strided_bf16(
+        q_in, k_in, q_weight, k_weight, cos, sin,
+        int(q_heads), int(k_heads), float(eps), q_out, k_out,
+    )
+    return q_out, k_out
 def qkv_split_per_head_norm_rope_bf16(
     packed_qkv: torch.Tensor,
     q_norm_weight: torch.Tensor,
     "qkv_split_bias_rope_fp16",
     "qkv_split_bf16",
     "qkv_split_rope_kvcache_bf16",
+    "qkv_split_rope_kvcache_fp16",
+    "qk_norm_rope_strided_bf16",
     "qkv_split_norm_rope_bf16",
     "qkv_split_bias_norm_rope_v_bf16",
     "qkv_split_bias_norm_rope_v_cat_bf16",

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

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

 import torch
+from . import flashrt_qkv_cache_rope_test
+ops = torch.ops.flashrt_qkv_cache_rope_test
 def add_op_namespace_prefix(op_name: str):
+    return f"flashrt_qkv_cache_rope_test::{op_name}"

build/torch211-cxx11-cu130-aarch64-linux/flashrt_qkv_cache_rope_test.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:0511d29a589f1fb1380c7555f54b673d4e055350a98a6b6b3e6a24cc40d4046b
+size 764184

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-qkv-cache-rope",
-  "id": "_flashrt_qkv_cache_rope_cuda_7781728",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,10 +13,20 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "A5t1ELQb0pbOoFMjZR7VzYvIds3qUO94UZ/4ZbqHd6U=",
-      "_flashrt_qkv_cache_rope_cuda_7781728.abi3.so": "Lwrt2ZZRrbXY/3OleQGifvVdwC402AyVLXzfUhG1Qls=",
-      "_ops.py": "IZnrJOS3RgtEaCdzR1d8T8OPo0IQT2EeJ3J00gc4lrw=",
       "flashrt_qkv_cache_rope/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   }
 }

 {
   "name": "flashrt-qkv-cache-rope",
+  "id": "flashrt_qkv_cache_rope_test",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "Izk+vnt9b4h2bryy/G+JK29rnXtvBhAbTLxsiNQGCSk=",
+      "flashrt_qkv_cache_rope_test.abi3.so": "BRHSmlifH7E4DHVV9UtnPU4FU1CpimtrPmokzEDUBGs=",
+      "_ops.py": "ZSIUGrzgBBffx37jx/znR0nnvtuaR3+G3FSvPKa9UHQ=",
       "flashrt_qkv_cache_rope/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
+  },
+  "provenance": {
+    "kernel": {
+      "sha": "7db5ebe",
+      "dirty": false
+    },
+    "validation": {
+      "torch": "2.11.0+cu130",
+      "cuda": "13.0"
+    }
   }
 }