liangsu9988 commited on Jun 9

Commit

1fea443

verified ·

1 Parent(s): 56405b8

Uploaded using `kernel-builder`.

Browse files

Files changed (17) hide show

benchmarks/benchmark.py +154 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +140 -0
build/torch211-cxx11-cu128-x86_64-linux/{_flashrt_qkv_cache_rope_cuda_cf903dd.abi3.so → _flashrt_qkv_cache_rope_cuda_5de4768.abi3.so} +2 -2
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +1 -1
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +140 -0
build/torch211-cxx11-cu130-x86_64-linux/{_flashrt_qkv_cache_rope_cuda_cf903dd.abi3.so → _flashrt_qkv_cache_rope_cuda_5de4768.abi3.so} +2 -2
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +3 -2
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +140 -0
build/torch212-cxx11-cu130-x86_64-linux/{_flashrt_qkv_cache_rope_cuda_cf903dd.abi3.so → _flashrt_qkv_cache_rope_cuda_5de4768.abi3.so} +2 -2
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +3 -2
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +140 -0
build/torch212-cxx11-cu132-x86_64-linux/{_flashrt_qkv_cache_rope_cuda_cf903dd.abi3.so → _flashrt_qkv_cache_rope_cuda_5de4768.abi3.so} +2 -2
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +3 -2

benchmarks/benchmark.py CHANGED Viewed

@@ -148,6 +148,47 @@ class SourceOps:
         )
         return q_cat_out, k_cat_out, v_cat_out
 def _preload_cublaslt() -> None:
     for parent in Path(torch.__file__).resolve().parents:
@@ -202,6 +243,13 @@ def make_freqs(seq_len: int, head_dim: int):
     return torch.cos(theta).contiguous(), torch.sin(theta).contiguous()
 def make_case(batch: int, seq_len: int, heads: int, head_dim: int):
     dim = heads * head_dim
     packed = torch.randn((batch, seq_len, 3 * dim), device="cuda", dtype=torch.bfloat16)
@@ -243,6 +291,20 @@ def apply_pair_rope(x: torch.Tensor, freqs_re: torch.Tensor, freqs_im: torch.Ten
     return out.reshape(batch, seq_len, heads, head_dim).to(torch.bfloat16)
 def apply_rotate_half_rope_128(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
     xf = x.float()
     out = torch.empty_like(xf)
@@ -290,6 +352,16 @@ def torch_ref_decode(x, weight, cos, sin, eps):
     return apply_rotate_half_rope_128(rms_norm(x, weight, eps).to(torch.bfloat16), cos, sin)
 def make_joint3_case(video_len: int, action_len: int, und_len: int, heads: int, head_dim: int):
     packed_v, v_q_w, v_k_w, freqs_re, freqs_im, _, _ = make_case(1, video_len, heads, head_dim)
     packed_a, a_q_w, a_k_w, _, _, _, _ = make_case(1, action_len, heads, head_dim)
@@ -623,6 +695,85 @@ def run_decode_kv(ops, name: str, heads: int, devpos: bool, args) -> Result:
     )
 def write_markdown(path: Path, results: list[Result]) -> None:
     lines = [
         "| Shape | B,L,H,D | FlashRT us | Eager us | vs eager | Q p99 | K p99 | Q cosine | K cosine | Status |",
@@ -658,11 +809,14 @@ def main() -> None:
     results = [run_one(ops, name, SHAPES[name], args) for name in SHAPE_GROUPS[args.shapes]]
     if args.shapes in ("smoke", "all"):
         results.append(run_joint3(ops, "joint3_small", 64, 8, 4, 8, 128, args))
     if args.shapes in ("headline", "all"):
         results.append(run_joint3(ops, "joint3_vla", 2520, 16, 16, 24, 128, args))
         results.append(run_decode_q(ops, "decode_q_stage_h24", 24, args))
         results.append(run_decode_kv(ops, "decode_kvwrite_h8", 8, False, args))
         results.append(run_decode_kv(ops, "decode_kvwrite_devpos_h8", 8, True, args))
     for r in results:
         print(

         )
         return q_cat_out, k_cat_out, v_cat_out
+    def qkv_split_rope_kvcache_bf16(
+        self,
+        packed_qkv,
+        rope,
+        q_heads,
+        kv_heads,
+        head_dim,
+        cache_offset,
+        q_out=None,
+        k_cache=None,
+        v_cache=None,
+        max_seq_len=None,
+    ):
+        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.bfloat16,
+            )
+        if k_cache is None or v_cache is None:
+            if max_seq_len is None:
+                max_seq_len = cache_offset + seq_len
+            cache_shape = (batch, max_seq_len, kv_heads, head_dim)
+            if k_cache is None:
+                k_cache = torch.empty(cache_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+            if v_cache is None:
+                v_cache = torch.empty(cache_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+        self._ops.qkv_split_rope_kvcache_bf16(
+            packed_qkv,
+            rope,
+            int(q_heads),
+            int(kv_heads),
+            int(head_dim),
+            int(cache_offset),
+            q_out,
+            k_cache,
+            v_cache,
+        )
+        return q_out, k_cache, v_cache
 def _preload_cublaslt() -> None:
     for parent in Path(torch.__file__).resolve().parents:
     return torch.cos(theta).contiguous(), torch.sin(theta).contiguous()
+def make_interleaved_rope(seq_len: int, head_dim: int):
+    theta = torch.randn((seq_len, head_dim // 2), device="cuda", dtype=torch.float32)
+    cos = torch.cos(theta).to(torch.bfloat16)
+    sin = torch.sin(theta).to(torch.bfloat16)
+    return torch.stack([cos, sin], dim=-1).reshape(seq_len, head_dim).contiguous()
 def make_case(batch: int, seq_len: int, heads: int, head_dim: int):
     dim = heads * head_dim
     packed = torch.randn((batch, seq_len, 3 * dim), device="cuda", dtype=torch.bfloat16)
     return out.reshape(batch, seq_len, heads, head_dim).to(torch.bfloat16)
+def apply_interleaved_pair_rope(x: torch.Tensor, rope: torch.Tensor):
+    batch, seq_len, heads, head_dim = x.shape
+    pair = x.float().reshape(batch, seq_len, heads, head_dim // 2, 2)
+    re = pair[..., 0]
+    im = pair[..., 1]
+    rope_pair = rope[:seq_len].float().reshape(seq_len, head_dim // 2, 2)
+    cos = rope_pair[..., 0].view(1, seq_len, 1, head_dim // 2)
+    sin = rope_pair[..., 1].view(1, seq_len, 1, head_dim // 2)
+    out = torch.empty_like(pair.float())
+    out[..., 0] = re * cos - im * sin
+    out[..., 1] = re * sin + im * cos
+    return out.reshape(batch, seq_len, heads, head_dim).to(torch.bfloat16)
 def apply_rotate_half_rope_128(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
     xf = x.float()
     out = torch.empty_like(xf)
     return apply_rotate_half_rope_128(rms_norm(x, weight, eps).to(torch.bfloat16), cos, sin)
+def torch_ref_kvcache(packed_qkv, rope, q_heads, kv_heads, head_dim):
+    batch, seq_len, _ = packed_qkv.shape
+    q_dim = q_heads * head_dim
+    kv_dim = kv_heads * head_dim
+    q = packed_qkv[:, :, :q_dim].view(batch, seq_len, q_heads, head_dim)
+    k = packed_qkv[:, :, q_dim : q_dim + kv_dim].view(batch, seq_len, kv_heads, head_dim)
+    v = packed_qkv[:, :, q_dim + kv_dim :].view(batch, seq_len, kv_heads, head_dim)
+    return apply_interleaved_pair_rope(q, rope), apply_interleaved_pair_rope(k, rope), v
 def make_joint3_case(video_len: int, action_len: int, und_len: int, heads: int, head_dim: int):
     packed_v, v_q_w, v_k_w, freqs_re, freqs_im, _, _ = make_case(1, video_len, heads, head_dim)
     packed_a, a_q_w, a_k_w, _, _, _, _ = make_case(1, action_len, heads, head_dim)
     )
+def run_kvcache_gqa(
+    ops,
+    name: str,
+    batch: int,
+    seq_len: int,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    args,
+) -> Result:
+    qkv_dim = (q_heads + 2 * kv_heads) * head_dim
+    packed = torch.randn((batch, seq_len, qkv_dim), device="cuda", dtype=torch.bfloat16)
+    rope = make_interleaved_rope(seq_len, head_dim)
+    cache_offset = 2
+    max_seq_len = cache_offset + seq_len + 2
+    q_out = torch.empty((batch, seq_len, q_heads, head_dim), device="cuda", dtype=torch.bfloat16)
+    k_cache = torch.empty((batch, max_seq_len, kv_heads, head_dim), device="cuda", dtype=torch.bfloat16)
+    v_cache = torch.empty_like(k_cache)
+    got_q, got_k_cache, got_v_cache = ops.qkv_split_rope_kvcache_bf16(
+        packed,
+        rope,
+        q_heads,
+        kv_heads,
+        head_dim,
+        cache_offset,
+        q_out,
+        k_cache,
+        v_cache,
+    )
+    exp_q, exp_k, exp_v = torch_ref_kvcache(packed, rope, q_heads, kv_heads, head_dim)
+    sl = slice(cache_offset, cache_offset + seq_len)
+    q_p99, q_cos = metrics(got_q, exp_q)
+    k_p99, k_cos = metrics(got_k_cache[:, sl], exp_k)
+    v_p99, v_cos = metrics(got_v_cache[:, sl], exp_v)
+    def flashrt_fn():
+        return ops.qkv_split_rope_kvcache_bf16(
+            packed,
+            rope,
+            q_heads,
+            kv_heads,
+            head_dim,
+            cache_offset,
+            q_out,
+            k_cache,
+            v_cache,
+        )
+    def eager_fn():
+        exp_q_local, exp_k_local, exp_v_local = torch_ref_kvcache(packed, rope, q_heads, kv_heads, head_dim)
+        q_out.copy_(exp_q_local)
+        k_cache[:, sl].copy_(exp_k_local)
+        v_cache[:, sl].copy_(exp_v_local)
+        return q_out, k_cache, v_cache
+    flashrt_us = time_us(flashrt_fn, args.warmup, args.iters)
+    eager_us = time_us(eager_fn, args.warmup, args.iters)
+    status = (
+        "PASS"
+        if q_p99 <= args.p99_abs_limit and k_p99 <= args.p99_abs_limit and v_p99 == 0.0
+        else "FAIL"
+    )
+    return Result(
+        shape=name,
+        batch=batch,
+        seq_len=seq_len,
+        heads=q_heads,
+        head_dim=head_dim,
+        flashrt_us=flashrt_us,
+        torch_eager_us=eager_us,
+        speedup_vs_eager=eager_us / flashrt_us,
+        q_p99_abs=max(q_p99, v_p99),
+        k_p99_abs=k_p99,
+        q_cosine=min(q_cos, v_cos),
+        k_cosine=k_cos,
+        status=status,
+    )
 def write_markdown(path: Path, results: list[Result]) -> None:
     lines = [
         "| Shape | B,L,H,D | FlashRT us | Eager us | vs eager | Q p99 | K p99 | Q cosine | K cosine | Status |",
     results = [run_one(ops, name, SHAPES[name], args) for name in SHAPE_GROUPS[args.shapes]]
     if args.shapes in ("smoke", "all"):
         results.append(run_joint3(ops, "joint3_small", 64, 8, 4, 8, 128, args))
+        results.append(run_kvcache_gqa(ops, "pi05_decoder_gqa_kvcache", 1, 10, 8, 1, 256, args))
     if args.shapes in ("headline", "all"):
         results.append(run_joint3(ops, "joint3_vla", 2520, 16, 16, 24, 128, args))
         results.append(run_decode_q(ops, "decode_q_stage_h24", 24, args))
         results.append(run_decode_kv(ops, "decode_kvwrite_h8", 8, False, args))
         results.append(run_decode_kv(ops, "decode_kvwrite_devpos_h8", 8, True, args))
+    if args.shapes == "headline":
+        results.append(run_kvcache_gqa(ops, "pi05_decoder_gqa_kvcache", 1, 10, 8, 1, 256, args))
     for r in results:
         print(

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

@@ -78,6 +78,53 @@ def _decode_k_norm_rope_kvwrite_devpos_bf16_fake(
     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,
@@ -230,6 +277,21 @@ def _check_packed_qkv(
         raise RuntimeError("norm weights must have shape (heads * head_dim,)")
 def _check_freqs(
     freqs_re: torch.Tensor,
     freqs_im: torch.Tensor,
@@ -286,6 +348,82 @@ def qkv_split_norm_rope_bf16(
     return q_out, k_out
 def decode_q_norm_rope_stage_bf16(
     q_pre: torch.Tensor,
     q_norm_weight: torch.Tensor,
@@ -489,6 +627,8 @@ __all__ = [
     "decode_q_norm_rope_stage_bf16",
     "decode_k_norm_rope_kvwrite_bf16",
     "decode_k_norm_rope_kvwrite_devpos_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_bf16"))
+def _qkv_split_rope_kvcache_bf16_fake(
+    packed_qkv: torch.Tensor,
+    rope: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    cache_offset: int,
+    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 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 must have shape (batch, seq_len, q_heads, head_dim)")
+    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 or cache_offset + seq_len > k_cache.shape[1]:
+        raise RuntimeError("cache_offset + seq_len must be within k_cache.shape[1]")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qkv_split_bf16"))
+def _qkv_split_bf16_fake(
+    packed_qkv: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+) -> None:
+    if packed_qkv.dim() != 3:
+        raise RuntimeError("packed_qkv must have shape (batch, seq_len, 3 * heads * head_dim)")
+    batch, seq_len, cols = packed_qkv.shape
+    if cols != 3 * heads * head_dim:
+        raise RuntimeError("packed_qkv.shape[2] must be 3 * heads * head_dim")
+    out_shape = (batch, seq_len, heads, head_dim)
+    if q_out.shape != out_shape or k_out.shape != out_shape or v_out.shape != out_shape:
+        raise RuntimeError("q_out, k_out, and v_out must have shape (batch, seq_len, heads, head_dim)")
+    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,
         raise RuntimeError("norm weights must have shape (heads * head_dim,)")
+def _check_packed_gqa_qkv(
+    packed_qkv: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+) -> None:
+    if packed_qkv.dim() != 3:
+        raise RuntimeError("packed_qkv must have shape (batch, seq_len, (q_heads + 2 * kv_heads) * head_dim)")
+    if q_heads <= 0 or kv_heads <= 0 or head_dim <= 0 or head_dim % 2 != 0:
+        raise RuntimeError("q_heads, kv_heads, and even head_dim must be positive")
+    expected = (q_heads + 2 * kv_heads) * head_dim
+    if packed_qkv.shape[2] != expected:
+        raise RuntimeError("packed_qkv.shape[2] must be (q_heads + 2 * kv_heads) * head_dim")
 def _check_freqs(
     freqs_re: torch.Tensor,
     freqs_im: torch.Tensor,
     return q_out, k_out
+def qkv_split_rope_kvcache_bf16(
+    packed_qkv: torch.Tensor,
+    rope: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    cache_offset: int,
+    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 GQA packed QKV, apply interleaved RoPE, and write K/V cache.
+    ``packed_qkv`` has shape ``(batch, seq_len, (q_heads + 2 * kv_heads) * head_dim)``.
+    ``rope`` has BF16 interleaved ``[cos0, sin0, cos1, sin1, ...]`` rows with
+    shape ``(>= seq_len, head_dim)``. ``q_out`` has shape
+    ``(batch, seq_len, q_heads, head_dim)``. K/V are written in-place into
+    ``(batch, max_seq_len, kv_heads, head_dim)`` caches starting at
+    ``cache_offset``.
+    """
+    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.bfloat16,
+        )
+    if k_cache is None or v_cache is None:
+        if max_seq_len is None:
+            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.bfloat16)
+        if v_cache is None:
+            v_cache = torch.empty(cache_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    ops.qkv_split_rope_kvcache_bf16(
+        packed_qkv,
+        rope,
+        int(q_heads),
+        int(kv_heads),
+        int(head_dim),
+        int(cache_offset),
+        q_out,
+        k_cache,
+        v_cache,
+    )
+    return q_out, k_cache, v_cache
+def qkv_split_bf16(
+    packed_qkv: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    q_out: torch.Tensor | None = None,
+    k_out: torch.Tensor | None = None,
+    v_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Split packed BF16 QKV into Q/K/V tensors.
+    ``packed_qkv`` has shape ``(batch, seq_len, 3 * heads * head_dim)``.
+    Outputs have shape ``(batch, seq_len, heads, head_dim)``.
+    """
+    out_shape = (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim)
+    if q_out is None:
+        q_out = torch.empty(out_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    if k_out is None:
+        k_out = torch.empty(out_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    if v_out is None:
+        v_out = torch.empty(out_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    ops.qkv_split_bf16(packed_qkv, int(heads), int(head_dim), q_out, k_out, v_out)
+    return q_out, k_out, v_out
 def decode_q_norm_rope_stage_bf16(
     q_pre: torch.Tensor,
     q_norm_weight: torch.Tensor,
     "decode_q_norm_rope_stage_bf16",
     "decode_k_norm_rope_kvwrite_bf16",
     "decode_k_norm_rope_kvwrite_devpos_bf16",
+    "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",

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

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:6c52cad3d53d935f3488e0b2e357d5a59065cc7a3bab2d4f22c67e30473e3eb1
+size 1870824

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-qkv-cache-rope",
-  "id": "_flashrt_qkv_cache_rope_cuda_cf903dd",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

 {
   "name": "flashrt-qkv-cache-rope",
+  "id": "_flashrt_qkv_cache_rope_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],

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

@@ -78,6 +78,53 @@ def _decode_k_norm_rope_kvwrite_devpos_bf16_fake(
     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,
@@ -230,6 +277,21 @@ def _check_packed_qkv(
         raise RuntimeError("norm weights must have shape (heads * head_dim,)")
 def _check_freqs(
     freqs_re: torch.Tensor,
     freqs_im: torch.Tensor,
@@ -286,6 +348,82 @@ def qkv_split_norm_rope_bf16(
     return q_out, k_out
 def decode_q_norm_rope_stage_bf16(
     q_pre: torch.Tensor,
     q_norm_weight: torch.Tensor,
@@ -489,6 +627,8 @@ __all__ = [
     "decode_q_norm_rope_stage_bf16",
     "decode_k_norm_rope_kvwrite_bf16",
     "decode_k_norm_rope_kvwrite_devpos_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_bf16"))
+def _qkv_split_rope_kvcache_bf16_fake(
+    packed_qkv: torch.Tensor,
+    rope: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    cache_offset: int,
+    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 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 must have shape (batch, seq_len, q_heads, head_dim)")
+    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 or cache_offset + seq_len > k_cache.shape[1]:
+        raise RuntimeError("cache_offset + seq_len must be within k_cache.shape[1]")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qkv_split_bf16"))
+def _qkv_split_bf16_fake(
+    packed_qkv: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+) -> None:
+    if packed_qkv.dim() != 3:
+        raise RuntimeError("packed_qkv must have shape (batch, seq_len, 3 * heads * head_dim)")
+    batch, seq_len, cols = packed_qkv.shape
+    if cols != 3 * heads * head_dim:
+        raise RuntimeError("packed_qkv.shape[2] must be 3 * heads * head_dim")
+    out_shape = (batch, seq_len, heads, head_dim)
+    if q_out.shape != out_shape or k_out.shape != out_shape or v_out.shape != out_shape:
+        raise RuntimeError("q_out, k_out, and v_out must have shape (batch, seq_len, heads, head_dim)")
+    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,
         raise RuntimeError("norm weights must have shape (heads * head_dim,)")
+def _check_packed_gqa_qkv(
+    packed_qkv: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+) -> None:
+    if packed_qkv.dim() != 3:
+        raise RuntimeError("packed_qkv must have shape (batch, seq_len, (q_heads + 2 * kv_heads) * head_dim)")
+    if q_heads <= 0 or kv_heads <= 0 or head_dim <= 0 or head_dim % 2 != 0:
+        raise RuntimeError("q_heads, kv_heads, and even head_dim must be positive")
+    expected = (q_heads + 2 * kv_heads) * head_dim
+    if packed_qkv.shape[2] != expected:
+        raise RuntimeError("packed_qkv.shape[2] must be (q_heads + 2 * kv_heads) * head_dim")
 def _check_freqs(
     freqs_re: torch.Tensor,
     freqs_im: torch.Tensor,
     return q_out, k_out
+def qkv_split_rope_kvcache_bf16(
+    packed_qkv: torch.Tensor,
+    rope: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    cache_offset: int,
+    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 GQA packed QKV, apply interleaved RoPE, and write K/V cache.
+    ``packed_qkv`` has shape ``(batch, seq_len, (q_heads + 2 * kv_heads) * head_dim)``.
+    ``rope`` has BF16 interleaved ``[cos0, sin0, cos1, sin1, ...]`` rows with
+    shape ``(>= seq_len, head_dim)``. ``q_out`` has shape
+    ``(batch, seq_len, q_heads, head_dim)``. K/V are written in-place into
+    ``(batch, max_seq_len, kv_heads, head_dim)`` caches starting at
+    ``cache_offset``.
+    """
+    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.bfloat16,
+        )
+    if k_cache is None or v_cache is None:
+        if max_seq_len is None:
+            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.bfloat16)
+        if v_cache is None:
+            v_cache = torch.empty(cache_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    ops.qkv_split_rope_kvcache_bf16(
+        packed_qkv,
+        rope,
+        int(q_heads),
+        int(kv_heads),
+        int(head_dim),
+        int(cache_offset),
+        q_out,
+        k_cache,
+        v_cache,
+    )
+    return q_out, k_cache, v_cache
+def qkv_split_bf16(
+    packed_qkv: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    q_out: torch.Tensor | None = None,
+    k_out: torch.Tensor | None = None,
+    v_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Split packed BF16 QKV into Q/K/V tensors.
+    ``packed_qkv`` has shape ``(batch, seq_len, 3 * heads * head_dim)``.
+    Outputs have shape ``(batch, seq_len, heads, head_dim)``.
+    """
+    out_shape = (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim)
+    if q_out is None:
+        q_out = torch.empty(out_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    if k_out is None:
+        k_out = torch.empty(out_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    if v_out is None:
+        v_out = torch.empty(out_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    ops.qkv_split_bf16(packed_qkv, int(heads), int(head_dim), q_out, k_out, v_out)
+    return q_out, k_out, v_out
 def decode_q_norm_rope_stage_bf16(
     q_pre: torch.Tensor,
     q_norm_weight: torch.Tensor,
     "decode_q_norm_rope_stage_bf16",
     "decode_k_norm_rope_kvwrite_bf16",
     "decode_k_norm_rope_kvwrite_devpos_bf16",
+    "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",

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

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

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

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-qkv-cache-rope",
-  "id": "_flashrt_qkv_cache_rope_cuda_cf903dd",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -9,7 +9,8 @@
     "archs": [
       "10.0",
       "11.0",
-      "12.0+PTX",
       "7.5",
       "8.0",
       "8.6",

 {
   "name": "flashrt-qkv-cache-rope",
+  "id": "_flashrt_qkv_cache_rope_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
     "archs": [
       "10.0",
       "11.0",
+      "12.0",
+      "12.1+PTX",
       "7.5",
       "8.0",
       "8.6",

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

@@ -78,6 +78,53 @@ def _decode_k_norm_rope_kvwrite_devpos_bf16_fake(
     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,
@@ -230,6 +277,21 @@ def _check_packed_qkv(
         raise RuntimeError("norm weights must have shape (heads * head_dim,)")
 def _check_freqs(
     freqs_re: torch.Tensor,
     freqs_im: torch.Tensor,
@@ -286,6 +348,82 @@ def qkv_split_norm_rope_bf16(
     return q_out, k_out
 def decode_q_norm_rope_stage_bf16(
     q_pre: torch.Tensor,
     q_norm_weight: torch.Tensor,
@@ -489,6 +627,8 @@ __all__ = [
     "decode_q_norm_rope_stage_bf16",
     "decode_k_norm_rope_kvwrite_bf16",
     "decode_k_norm_rope_kvwrite_devpos_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_bf16"))
+def _qkv_split_rope_kvcache_bf16_fake(
+    packed_qkv: torch.Tensor,
+    rope: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    cache_offset: int,
+    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 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 must have shape (batch, seq_len, q_heads, head_dim)")
+    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 or cache_offset + seq_len > k_cache.shape[1]:
+        raise RuntimeError("cache_offset + seq_len must be within k_cache.shape[1]")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qkv_split_bf16"))
+def _qkv_split_bf16_fake(
+    packed_qkv: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+) -> None:
+    if packed_qkv.dim() != 3:
+        raise RuntimeError("packed_qkv must have shape (batch, seq_len, 3 * heads * head_dim)")
+    batch, seq_len, cols = packed_qkv.shape
+    if cols != 3 * heads * head_dim:
+        raise RuntimeError("packed_qkv.shape[2] must be 3 * heads * head_dim")
+    out_shape = (batch, seq_len, heads, head_dim)
+    if q_out.shape != out_shape or k_out.shape != out_shape or v_out.shape != out_shape:
+        raise RuntimeError("q_out, k_out, and v_out must have shape (batch, seq_len, heads, head_dim)")
+    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,
         raise RuntimeError("norm weights must have shape (heads * head_dim,)")
+def _check_packed_gqa_qkv(
+    packed_qkv: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+) -> None:
+    if packed_qkv.dim() != 3:
+        raise RuntimeError("packed_qkv must have shape (batch, seq_len, (q_heads + 2 * kv_heads) * head_dim)")
+    if q_heads <= 0 or kv_heads <= 0 or head_dim <= 0 or head_dim % 2 != 0:
+        raise RuntimeError("q_heads, kv_heads, and even head_dim must be positive")
+    expected = (q_heads + 2 * kv_heads) * head_dim
+    if packed_qkv.shape[2] != expected:
+        raise RuntimeError("packed_qkv.shape[2] must be (q_heads + 2 * kv_heads) * head_dim")
 def _check_freqs(
     freqs_re: torch.Tensor,
     freqs_im: torch.Tensor,
     return q_out, k_out
+def qkv_split_rope_kvcache_bf16(
+    packed_qkv: torch.Tensor,
+    rope: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    cache_offset: int,
+    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 GQA packed QKV, apply interleaved RoPE, and write K/V cache.
+    ``packed_qkv`` has shape ``(batch, seq_len, (q_heads + 2 * kv_heads) * head_dim)``.
+    ``rope`` has BF16 interleaved ``[cos0, sin0, cos1, sin1, ...]`` rows with
+    shape ``(>= seq_len, head_dim)``. ``q_out`` has shape
+    ``(batch, seq_len, q_heads, head_dim)``. K/V are written in-place into
+    ``(batch, max_seq_len, kv_heads, head_dim)`` caches starting at
+    ``cache_offset``.
+    """
+    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.bfloat16,
+        )
+    if k_cache is None or v_cache is None:
+        if max_seq_len is None:
+            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.bfloat16)
+        if v_cache is None:
+            v_cache = torch.empty(cache_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    ops.qkv_split_rope_kvcache_bf16(
+        packed_qkv,
+        rope,
+        int(q_heads),
+        int(kv_heads),
+        int(head_dim),
+        int(cache_offset),
+        q_out,
+        k_cache,
+        v_cache,
+    )
+    return q_out, k_cache, v_cache
+def qkv_split_bf16(
+    packed_qkv: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    q_out: torch.Tensor | None = None,
+    k_out: torch.Tensor | None = None,
+    v_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Split packed BF16 QKV into Q/K/V tensors.
+    ``packed_qkv`` has shape ``(batch, seq_len, 3 * heads * head_dim)``.
+    Outputs have shape ``(batch, seq_len, heads, head_dim)``.
+    """
+    out_shape = (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim)
+    if q_out is None:
+        q_out = torch.empty(out_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    if k_out is None:
+        k_out = torch.empty(out_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    if v_out is None:
+        v_out = torch.empty(out_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    ops.qkv_split_bf16(packed_qkv, int(heads), int(head_dim), q_out, k_out, v_out)
+    return q_out, k_out, v_out
 def decode_q_norm_rope_stage_bf16(
     q_pre: torch.Tensor,
     q_norm_weight: torch.Tensor,
     "decode_q_norm_rope_stage_bf16",
     "decode_k_norm_rope_kvwrite_bf16",
     "decode_k_norm_rope_kvwrite_devpos_bf16",
+    "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",

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

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

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

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-qkv-cache-rope",
-  "id": "_flashrt_qkv_cache_rope_cuda_cf903dd",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -9,7 +9,8 @@
     "archs": [
       "10.0",
       "11.0",
-      "12.0+PTX",
       "7.5",
       "8.0",
       "8.6",

 {
   "name": "flashrt-qkv-cache-rope",
+  "id": "_flashrt_qkv_cache_rope_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
     "archs": [
       "10.0",
       "11.0",
+      "12.0",
+      "12.1+PTX",
       "7.5",
       "8.0",
       "8.6",

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

@@ -78,6 +78,53 @@ def _decode_k_norm_rope_kvwrite_devpos_bf16_fake(
     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,
@@ -230,6 +277,21 @@ def _check_packed_qkv(
         raise RuntimeError("norm weights must have shape (heads * head_dim,)")
 def _check_freqs(
     freqs_re: torch.Tensor,
     freqs_im: torch.Tensor,
@@ -286,6 +348,82 @@ def qkv_split_norm_rope_bf16(
     return q_out, k_out
 def decode_q_norm_rope_stage_bf16(
     q_pre: torch.Tensor,
     q_norm_weight: torch.Tensor,
@@ -489,6 +627,8 @@ __all__ = [
     "decode_q_norm_rope_stage_bf16",
     "decode_k_norm_rope_kvwrite_bf16",
     "decode_k_norm_rope_kvwrite_devpos_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_bf16"))
+def _qkv_split_rope_kvcache_bf16_fake(
+    packed_qkv: torch.Tensor,
+    rope: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    cache_offset: int,
+    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 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 must have shape (batch, seq_len, q_heads, head_dim)")
+    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 or cache_offset + seq_len > k_cache.shape[1]:
+        raise RuntimeError("cache_offset + seq_len must be within k_cache.shape[1]")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("qkv_split_bf16"))
+def _qkv_split_bf16_fake(
+    packed_qkv: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    v_out: torch.Tensor,
+) -> None:
+    if packed_qkv.dim() != 3:
+        raise RuntimeError("packed_qkv must have shape (batch, seq_len, 3 * heads * head_dim)")
+    batch, seq_len, cols = packed_qkv.shape
+    if cols != 3 * heads * head_dim:
+        raise RuntimeError("packed_qkv.shape[2] must be 3 * heads * head_dim")
+    out_shape = (batch, seq_len, heads, head_dim)
+    if q_out.shape != out_shape or k_out.shape != out_shape or v_out.shape != out_shape:
+        raise RuntimeError("q_out, k_out, and v_out must have shape (batch, seq_len, heads, head_dim)")
+    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,
         raise RuntimeError("norm weights must have shape (heads * head_dim,)")
+def _check_packed_gqa_qkv(
+    packed_qkv: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+) -> None:
+    if packed_qkv.dim() != 3:
+        raise RuntimeError("packed_qkv must have shape (batch, seq_len, (q_heads + 2 * kv_heads) * head_dim)")
+    if q_heads <= 0 or kv_heads <= 0 or head_dim <= 0 or head_dim % 2 != 0:
+        raise RuntimeError("q_heads, kv_heads, and even head_dim must be positive")
+    expected = (q_heads + 2 * kv_heads) * head_dim
+    if packed_qkv.shape[2] != expected:
+        raise RuntimeError("packed_qkv.shape[2] must be (q_heads + 2 * kv_heads) * head_dim")
 def _check_freqs(
     freqs_re: torch.Tensor,
     freqs_im: torch.Tensor,
     return q_out, k_out
+def qkv_split_rope_kvcache_bf16(
+    packed_qkv: torch.Tensor,
+    rope: torch.Tensor,
+    q_heads: int,
+    kv_heads: int,
+    head_dim: int,
+    cache_offset: int,
+    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 GQA packed QKV, apply interleaved RoPE, and write K/V cache.
+    ``packed_qkv`` has shape ``(batch, seq_len, (q_heads + 2 * kv_heads) * head_dim)``.
+    ``rope`` has BF16 interleaved ``[cos0, sin0, cos1, sin1, ...]`` rows with
+    shape ``(>= seq_len, head_dim)``. ``q_out`` has shape
+    ``(batch, seq_len, q_heads, head_dim)``. K/V are written in-place into
+    ``(batch, max_seq_len, kv_heads, head_dim)`` caches starting at
+    ``cache_offset``.
+    """
+    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.bfloat16,
+        )
+    if k_cache is None or v_cache is None:
+        if max_seq_len is None:
+            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.bfloat16)
+        if v_cache is None:
+            v_cache = torch.empty(cache_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    ops.qkv_split_rope_kvcache_bf16(
+        packed_qkv,
+        rope,
+        int(q_heads),
+        int(kv_heads),
+        int(head_dim),
+        int(cache_offset),
+        q_out,
+        k_cache,
+        v_cache,
+    )
+    return q_out, k_cache, v_cache
+def qkv_split_bf16(
+    packed_qkv: torch.Tensor,
+    heads: int,
+    head_dim: int,
+    q_out: torch.Tensor | None = None,
+    k_out: torch.Tensor | None = None,
+    v_out: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Split packed BF16 QKV into Q/K/V tensors.
+    ``packed_qkv`` has shape ``(batch, seq_len, 3 * heads * head_dim)``.
+    Outputs have shape ``(batch, seq_len, heads, head_dim)``.
+    """
+    out_shape = (packed_qkv.shape[0], packed_qkv.shape[1], heads, head_dim)
+    if q_out is None:
+        q_out = torch.empty(out_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    if k_out is None:
+        k_out = torch.empty(out_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    if v_out is None:
+        v_out = torch.empty(out_shape, device=packed_qkv.device, dtype=torch.bfloat16)
+    ops.qkv_split_bf16(packed_qkv, int(heads), int(head_dim), q_out, k_out, v_out)
+    return q_out, k_out, v_out
 def decode_q_norm_rope_stage_bf16(
     q_pre: torch.Tensor,
     q_norm_weight: torch.Tensor,
     "decode_q_norm_rope_stage_bf16",
     "decode_k_norm_rope_kvwrite_bf16",
     "decode_k_norm_rope_kvwrite_devpos_bf16",
+    "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",

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

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:6f9b056bfcac039776c37f83e30e348dde8d7ec3f2ab1d7f7062e8a57f2bb39f
+size 1828568

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "flashrt-qkv-cache-rope",
-  "id": "_flashrt_qkv_cache_rope_cuda_cf903dd",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -9,7 +9,8 @@
     "archs": [
       "10.0",
       "11.0",
-      "12.0+PTX",
       "7.5",
       "8.0",
       "8.6",

 {
   "name": "flashrt-qkv-cache-rope",
+  "id": "_flashrt_qkv_cache_rope_cuda_5de4768",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
     "archs": [
       "10.0",
       "11.0",
+      "12.0",
+      "12.1+PTX",
       "7.5",
       "8.0",
       "8.6",