Publish GROOT Thor Batch 4 (torch213-cxx11-cu130-aarch64-linux, 6ab0803)

Browse files

Files changed (16) hide show

CARD.md +3 -0
README.md +14 -2
SYNC.md +7 -0
VALIDATION.md +21 -0
benchmarks/RESULTS.md +41 -3
build.toml +6 -1
build/torch213-cxx11-cu130-aarch64-linux/__init__.py +53 -0
build/torch213-cxx11-cu130-aarch64-linux/{masked_mha_runtime_source_test.abi3.so → _masked_mha_runtime_cuda_6ab0803.abi3.so} +2 -2
build/torch213-cxx11-cu130-aarch64-linux/_ops.py +3 -3
build/torch213-cxx11-cu130-aarch64-linux/metadata.json +6 -6
csrc/attention_seqused_fused.cu +137 -0
csrc/attention_seqused_fused.cuh +22 -0
flake.lock +118 -0
tests/test_masked_mha_runtime.py +118 -4
torch-ext/masked_mha_runtime/__init__.py +87 -1
torch-ext/torch_binding.cpp +74 -0

CARD.md CHANGED Viewed

@@ -13,7 +13,10 @@ removing the per-call `-inf` pre-fill. BF16 accepts fused-QKV token strides.
 - `forward(q, k, v, *, scale=None)`
 - `forward_static(q, k, v, *, logits, out, scale=None)`
 - `allocate_workspace(q, k)`
 Inputs use `(sequence, heads, head_dim)`. `forward_static` is the CUDA Graph
 hot-path API: allocate `logits` and `out` once and reuse their addresses.

 - `forward(q, k, v, *, scale=None)`
 - `forward_static(q, k, v, *, logits, out, scale=None)`
+- `attention_mha_fp16_masked(q, k, v, *, logits, out, scale=None)`
+- `attention_mha_bf16_masked(q, k, v, *, logits, out, qkv_token_stride=None, scale=None)`
 - `allocate_workspace(q, k)`
+- `forward_seqused_static(q, k, v, valid_k, *, logits, out, scale=None)`
 Inputs use `(sequence, heads, head_dim)`. `forward_static` is the CUDA Graph
 hot-path API: allocate `logits` and `out` once and reuse their addresses.

README.md CHANGED Viewed

@@ -11,13 +11,25 @@ import torch
 ops = get_kernel("flashrt/masked-mha-runtime", version=1)
 logits = ops.allocate_workspace(q, k)
 out = torch.empty_like(q, memory_format=torch.contiguous_format)
-ops.forward_static(q, k, v, logits=logits, out=out)
 ```
-Public functions are `allocate_workspace`, `forward_static`, and `forward`.
 Inputs use `(sequence, heads, head_dim)`. The production GROOT gate covers
 DiT `(41, 32, 48)`, ViT/LLM sequence lengths `277/1024`, padded boundaries
 `1025/2048`, FP16 and BF16, fused strides, and bitwise CUDA Graph replay.
 See [CARD.md](CARD.md) for the complete contract. Source provenance is FlashRT
 commit `24df793f4fa2d50780aea03b644208c6e0cb4162`.

 ops = get_kernel("flashrt/masked-mha-runtime", version=1)
 logits = ops.allocate_workspace(q, k)
 out = torch.empty_like(q, memory_format=torch.contiguous_format)
+ops.attention_mha_fp16_masked(q, k, v, logits=logits, out=out)
 ```
+Public functions are `allocate_workspace`, `attention_mha_fp16_masked`,
+`attention_mha_bf16_masked`, `forward_static`, `forward`, and
+`forward_seqused_static`. The explicit masked names are additive; the legacy
+`forward_static` entry dispatches the same kernels.
 Inputs use `(sequence, heads, head_dim)`. The production GROOT gate covers
 DiT `(41, 32, 48)`, ViT/LLM sequence lengths `277/1024`, padded boundaries
 `1025/2048`, FP16 and BF16, fused strides, and bitwise CUDA Graph replay.
+`forward_seqused_static` is the FP16 shared-KV path used by static VLA
+decoders: Q is `(Sq,H,D)`, K/V are `(Sk_max,D)`, and a CUDA int32 `valid_k`
+limits softmax without a standalone `-inf` fill. It supports `Sk_max <= 1024`.
+The BF16 masked entry accepts `qkv_token_stride=` and validates it against the
+actual Tensor token stride. This permits direct views into fused QKV GEMM
+output while rejecting a mismatched layout before launch. Both masked entries
+only read and write the valid `S_kv` columns, so padded workspace bytes never
+participate in softmax or CUDA Graph replay.
 See [CARD.md](CARD.md) for the complete contract. Source provenance is FlashRT
 commit `24df793f4fa2d50780aea03b644208c6e0cb4162`.

SYNC.md CHANGED Viewed

@@ -4,3 +4,10 @@
 - Commit: `24df793f4fa2d50780aea03b644208c6e0cb4162`
 - Source: `csrc/kernels/attention_mha_masked.cu`
 - Local changes: Tensor validation and `torch.library` binding only.

 - Commit: `24df793f4fa2d50780aea03b644208c6e0cb4162`
 - Source: `csrc/kernels/attention_mha_masked.cu`
 - Local changes: Tensor validation and `torch.library` binding only.
+The CUDA source was reverified byte-for-byte against FlashRT revision
+`ff7795a88c257ccf7195c3c066ce69b7fe84be95`. SHA256 is
+`751b1de5e227bfff00f57c30c0a5bfb8efefaeedd7230fb3b2d0c67a28773ff9`.
+The explicit FP16/BF16 public entries are additive binding surfaces over this
+same native implementation; the established `forward_static` entry remains
+unchanged.

VALIDATION.md CHANGED Viewed

@@ -3,3 +3,24 @@
 Release requires FP16 and BF16 parity against PyTorch SDPA at key lengths
 41, 277, 1024, 1025, and 2048; poisoned padded scratch; CUDA Graph replay;
 source and installed-artifact runs on SM110; and native-vs-package timing.

 Release requires FP16 and BF16 parity against PyTorch SDPA at key lengths
 41, 277, 1024, 1025, and 2048; poisoned padded scratch; CUDA Graph replay;
 source and installed-artifact runs on SM110; and native-vs-package timing.
+The additive `attention_mha_fp16_masked` and
+`attention_mha_bf16_masked` entries must remain bitwise-equal to
+`forward_static`. The BF16 entry additionally checks `qkv_token_stride`
+against the actual Q tensor stride so fused-QKV views fail early instead of
+reading the wrong token row. Release gating includes poisoned padded columns,
+bitwise graph replay, and an explicit/static latency ratio no greater than
+1.05x.
+On 2026-08-06, source full passed `10/10` on NVIDIA Thor (SM110), PyTorch
+`2.13.0+cu130`, and CUDA 13.0. The additive `forward_seqused_static` rows cover
+PI0.5 `(Sq,H,D)=(10,8,256)`, `Sk_max=456/968`, and device-resident
+`valid_k=456/712/968`. Worst p99 absolute error was `0.000122`, worst cosine
+was `0.99999982`, and graph replay was bitwise deterministic.
+On 2026-08-08, the source and clean installed-artifact gates both passed
+`10/10` on the same Thor class. The explicit FP16/BF16 entries measured
+`1.005x/1.000x` of `forward_static`; BF16 fused-stride execution also passed
+`torch.compile(fullgraph=True)` exactly. For the GROOT row
+`Sq=Sk=41,H=32,D=48`, worst p99 absolute error was `0.003906` and cosine was
+`0.99999213`; sequence lengths 1024, 1025, and 2048 also passed.

benchmarks/RESULTS.md CHANGED Viewed

@@ -1,5 +1,43 @@
 # Results
-Pending the clean SM110 installed-artifact benchmark. No speedup claim is
-published until the package is measured against both PyTorch SDPA and the
-FlashRT native symbol on the same Thor host.

 # Results
+## NVIDIA Thor installed artifact
+- Device: NVIDIA Thor, SM110, aarch64
+- Runtime: PyTorch 2.13.0 + CUDA 13.0
+- Variant: `torch213-cxx11-cu130-aarch64-linux`
+- Native source: FlashRT `24df793f4fa2d50780aea03b644208c6e0cb4162`
+The installed-artifact full gate passed 6/6. Representative BF16 rows:
+| Shape | p99 absolute error | Cosine | CUDA Graph replay |
+|---|---:|---:|---|
+| `sk=1024` | 0.000452 | 0.99999362 | bitwise |
+| `sk=1025` | 0.000452 | 0.99999356 | bitwise |
+| `sk=2048` | 0.000174 | 0.99999666 | bitwise |
+The Hub Tensor wrapper and the raw registered native op were measured in the
+same process with CUDA events. Direct wrapper/raw latency was
+`24.7932/24.7544 us` (`1.0016x`). With the caller-owned static buffers used by
+the production CUDA Graph path, wrapper/raw latency was
+`19.0034/19.0034 us` (`1.0000x`).
+These numbers establish package parity with the FlashRT native symbol. They
+are not presented as a universal speedup over PyTorch SDPA; backend selection
+must still be profiled for the target sequence and padding profile.
+## GROOT Thor explicit masked-attention entries
+Clean installed artifact built from source commit
+`6ab0803f010b2af49474577dddef866938c3bdcd` on NVIDIA Thor (SM110), PyTorch
+2.13.0+cu130. Measurements use caller-owned buffers, CUDA Graph replay, and
+A-B-B-A ordering.
+| Entry | Shape | `forward_static` us | Explicit us | Explicit/static |
+|---|---:|---:|---:|---:|
+| `attention_mha_fp16_masked` | `Sq=41, Sk=277, H=32, D=48` | 22.342 | 22.454 | 1.005x |
+| `attention_mha_bf16_masked` | `Sq=41, Sk=277, H=32, D=48` | 22.649 | 22.644 | 1.000x |
+Both entries are bitwise-equal to the established wrapper. The BF16 entry
+also validates the fused-QKV token stride explicitly. Source and clean
+installed-artifact gates passed `10/10`; BF16 fused-stride execution passed
+`torch.compile(fullgraph=True)` exactly.

build.toml CHANGED Viewed

@@ -21,4 +21,9 @@ include = ["csrc"]
 cuda-minver = "13"
 cuda-capabilities = ["11.0a"]
 cuda-flags = ["-O3", "--use_fast_math"]
-src = ["csrc/attention_mha_masked.cu", "csrc/attention_mha_masked.cuh"]

 cuda-minver = "13"
 cuda-capabilities = ["11.0a"]
 cuda-flags = ["-O3", "--use_fast_math"]
+src = [
+  "csrc/attention_mha_masked.cu",
+  "csrc/attention_mha_masked.cuh",
+  "csrc/attention_seqused_fused.cu",
+  "csrc/attention_seqused_fused.cuh",
+]

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

@@ -23,6 +23,20 @@ def _forward_static_fake(q, k, v, logits, out, scale: float) -> None:
         raise RuntimeError("logits stride must cover sequence_kv")
 @torch.library.register_fake(add_op_namespace_prefix("forward_seqused_static"))
 def _forward_seqused_static_fake(q, k, v, valid_k, logits, out, scale: float) -> None:
     del scale
@@ -58,6 +72,43 @@ def forward_static(
     return out
 def forward(q, k, v, *, scale: Optional[float] = None):
     """Convenience allocation wrapper; use ``forward_static`` in hot paths."""
     logits = allocate_workspace(q, k)
@@ -83,6 +134,8 @@ def forward_seqused_static(
 __all__ = [
     "allocate_workspace",
     "forward",
     "forward_static",

         raise RuntimeError("logits stride must cover sequence_kv")
+@torch.library.register_fake(add_op_namespace_prefix("attention_mha_fp16_masked"))
+def _attention_mha_fp16_masked_fake(q, k, v, logits, out, scale: float) -> None:
+    _forward_static_fake(q, k, v, logits, out, scale)
+@torch.library.register_fake(add_op_namespace_prefix("attention_mha_bf16_masked"))
+def _attention_mha_bf16_masked_fake(
+    q, k, v, logits, out, scale: float, qkv_token_stride: int
+) -> None:
+    _forward_static_fake(q, k, v, logits, out, scale)
+    if qkv_token_stride != q.stride(0):
+        raise RuntimeError("qkv_token_stride must match q.stride(0)")
 @torch.library.register_fake(add_op_namespace_prefix("forward_seqused_static"))
 def _forward_seqused_static_fake(q, k, v, valid_k, logits, out, scale: float) -> None:
     del scale
     return out
+def attention_mha_fp16_masked(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    logits: torch.Tensor,
+    out: torch.Tensor,
+    scale: Optional[float] = None,
+) -> torch.Tensor:
+    """Run the deterministic valid-column FP16 masked MHA kernel."""
+    if scale is None:
+        scale = 1.0 / math.sqrt(q.shape[-1])
+    ops.attention_mha_fp16_masked(q, k, v, logits, out, float(scale))
+    return out
+def attention_mha_bf16_masked(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    logits: torch.Tensor,
+    out: torch.Tensor,
+    qkv_token_stride: Optional[int] = None,
+    scale: Optional[float] = None,
+) -> torch.Tensor:
+    """Run BF16 masked MHA directly from contiguous or fused-QKV views."""
+    if scale is None:
+        scale = 1.0 / math.sqrt(q.shape[-1])
+    if qkv_token_stride is None:
+        qkv_token_stride = q.stride(0)
+    ops.attention_mha_bf16_masked(
+        q, k, v, logits, out, float(scale), int(qkv_token_stride)
+    )
+    return out
 def forward(q, k, v, *, scale: Optional[float] = None):
     """Convenience allocation wrapper; use ``forward_static`` in hot paths."""
     logits = allocate_workspace(q, k)
 __all__ = [
+    "attention_mha_bf16_masked",
+    "attention_mha_fp16_masked",
     "allocate_workspace",
     "forward",
     "forward_static",

build/torch213-cxx11-cu130-aarch64-linux/{masked_mha_runtime_source_test.abi3.so → _masked_mha_runtime_cuda_6ab0803.abi3.so} RENAMED Viewed

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:4a28073db408aaea10e601f842cb523aecfead2965dd28d5928028552dfda4cc
+size 382720

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

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

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

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

@@ -1,27 +1,27 @@
 {
   "name": "masked-mha-runtime",
-  "id": "masked_mha_runtime_source_test",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
-      "11.0"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "sZLa3RWzOCoLFNml3p4SGOBbeTUtLYDp2tmIL0FlToo=",
-      "masked_mha_runtime_source_test.abi3.so": "O1RY70UCFV3HYcLmdC59LX7ikE+Fw+Cm5uCVrMa2StA=",
-      "_ops.py": "dXHlvnzts0eDL5fgM/I0tMNAsCB3uhZT03SIVpj4bvU=",
       "masked_mha_runtime/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
-      "sha": "3a505c0",
       "dirty": false
     },
     "validation": {

 {
   "name": "masked-mha-runtime",
+  "id": "_masked_mha_runtime_cuda_6ab0803",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda",
     "archs": [
+      "11.0a"
     ]
   },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "w4T07fNWmll9bTURpoHVj+llGGGmT6f7hRpFFKZjBoQ=",
+      "_masked_mha_runtime_cuda_6ab0803.abi3.so": "SigHPbQIquoQ5gH4QstSOuz+rSll3SjVkoAoVS39pMw=",
+      "_ops.py": "tWBn/Gj4KYWoPtEY8bDBbJ0/A61ndt9MyvId4N+G0SE=",
       "masked_mha_runtime/__init__.py": "v6p5XMfQzddhi1fLSAw4HX9CyS0rQsidvu9VsT01xi4="
     }
   },
   "provenance": {
     "kernel": {
+      "sha": "6ab0803",
       "dirty": false
     },
     "validation": {

csrc/attention_seqused_fused.cu ADDED Viewed

	@@ -0,0 +1,137 @@

+// ================================================================
+// FlashRT — seqused attention with the mask folded into softmax.
+//
+// attention_qkv_fp16_seqused runs QK^T, a standalone -inf mask
+// kernel over rows [valid, S_kv_max), softmax, then AV. The v2 chain
+// folds the seqused bound into the softmax kernel itself: max/sum
+// run over [0, valid) and positions beyond valid are written as
+// zero probabilities (identical to exp(-65504 - max) ~ 0 in the
+// reference), removing one kernel per call.
+//
+// Additive: attention_cublas.cu is untouched.
+// ================================================================
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cublas_v2.h>
+#include "attention_seqused_fused.cuh"
+namespace {
+#define SMV2_WARP_SIZE 32
+#define SMV2_MAX_COLS 1024
+#define SMV2_ITERS (SMV2_MAX_COLS / SMV2_WARP_SIZE)
+// One warp per logits row. The register-to-column mapping, the reduction
+// order, and the arithmetic are all identical to softmax_fp16_kernel; the
+// only difference is that out-of-range columns take -1e30 at load time
+// instead of being read back as the -65504 a separate mask kernel wrote.
+// Both exponentiate to exactly zero, so the stored probabilities are
+// bit-identical to the mask + softmax chain.
+__global__ void softmax_fp16_seqused_kernel(
+    __half* data, int rows, int cols, const int* __restrict__ seqused_k) {
+    int lane = threadIdx.x % SMV2_WARP_SIZE;
+    int row = blockIdx.x;
+    if (row >= rows) return;
+    int valid = seqused_k[0];
+    if (valid < 0) valid = 0;
+    if (valid > cols) valid = cols;
+    __half* src = data + row * cols;
+    int cols2 = cols / 2;
+    __half2* src2 = reinterpret_cast<__half2*>(src);
+    float reg[SMV2_ITERS];
+    float mx = -1e30f;
+    #pragma unroll
+    for (int it = 0; it < SMV2_ITERS / 2; it++) {
+        int c2 = it * SMV2_WARP_SIZE + lane;
+        if (c2 < cols2) {
+            __half2 v2 = src2[c2];
+            reg[it*2]   = (2 * c2     < valid) ? __half2float(v2.x) : -1e30f;
+            reg[it*2+1] = (2 * c2 + 1 < valid) ? __half2float(v2.y) : -1e30f;
+            mx = fmaxf(mx, fmaxf(reg[it*2], reg[it*2+1]));
+        } else {
+            reg[it*2] = -1e30f;
+            reg[it*2+1] = -1e30f;
+        }
+    }
+    if ((cols & 1) && lane == 0) {
+        float v = (cols - 1 < valid) ? __half2float(src[cols-1]) : -1e30f;
+        reg[SMV2_ITERS-1] = v;
+        mx = fmaxf(mx, v);
+    }
+    #pragma unroll
+    for (int o = 16; o > 0; o >>= 1)
+        mx = fmaxf(mx, __shfl_xor_sync(0xffffffff, mx, o));
+    float sm = 0.f;
+    #pragma unroll
+    for (int it = 0; it < SMV2_ITERS; it++) {
+        reg[it] = __expf(reg[it] - mx);
+        sm += reg[it];
+    }
+    #pragma unroll
+    for (int o = 16; o > 0; o >>= 1)
+        sm += __shfl_xor_sync(0xffffffff, sm, o);
+    float inv = 1.f / (sm + 1e-8f);
+    #pragma unroll
+    for (int it = 0; it < SMV2_ITERS / 2; it++) {
+        int c2 = it * SMV2_WARP_SIZE + lane;
+        if (c2 < cols2) {
+            __half2 v2;
+            v2.x = __float2half(reg[it*2] * inv);
+            v2.y = __float2half(reg[it*2+1] * inv);
+            src2[c2] = v2;
+        }
+    }
+    if ((cols & 1) && lane == 0) {
+        src[cols-1] = __float2half(reg[SMV2_ITERS-1] * inv);
+    }
+}
+}  // namespace
+void attention_qkv_fp16_seqused_v2(
+    cublasHandle_t handle,
+    const __half* Q,
+    const __half* K,
+    const __half* V,
+    __half* logits,
+    __half* out,
+    int S, int S_kv_max, int NH, int HD,
+    const int* seqused_k,
+    float attn_scale,
+    cudaStream_t stream)
+{
+    cublasSetStream(handle, stream);
+    float zero = 0.0f;
+    cublasGemmEx(handle,
+        CUBLAS_OP_T, CUBLAS_OP_N,
+        S_kv_max, S * NH, HD,
+        &attn_scale,
+        K, CUDA_R_16F, HD,
+        Q, CUDA_R_16F, HD,
+        &zero,
+        logits, CUDA_R_16F, S_kv_max,
+        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
+    softmax_fp16_seqused_kernel<<<S * NH, SMV2_WARP_SIZE, 0, stream>>>(
+        logits, S * NH, S_kv_max, seqused_k);
+    float one = 1.0f;
+    cublasGemmEx(handle,
+        CUBLAS_OP_N, CUBLAS_OP_N,
+        HD, S * NH, S_kv_max,
+        &one,
+        V, CUDA_R_16F, HD,
+        logits, CUDA_R_16F, S_kv_max,
+        &zero,
+        out, CUDA_R_16F, HD,
+        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
+}

csrc/attention_seqused_fused.cuh ADDED Viewed

	@@ -0,0 +1,22 @@

+// ================================================================
+// FlashRT — seqused attention with the mask folded into softmax.
+// Drop-in for attention_qkv_fp16_seqused with one fewer kernel;
+// requires S_kv_max <= 1024 (softmax per-warp register budget).
+// ================================================================
+#pragma once
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cublas_v2.h>
+void attention_qkv_fp16_seqused_v2(
+    cublasHandle_t handle,
+    const __half* Q,
+    const __half* K,
+    const __half* V,
+    __half* logits,
+    __half* out,
+    int S, int S_kv_max, int NH, int HD,
+    const int* seqused_k,
+    float attn_scale,
+    cudaStream_t stream);

flake.lock ADDED Viewed

	@@ -0,0 +1,118 @@

+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1767039857,
+        "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs",
+        "rust-overlay": "rust-overlay"
+      },
+      "locked": {
+        "lastModified": 1785740631,
+        "narHash": "sha256-bK6AXiZSJYaCYK2iTT6Xy+24bab6k0IURV80PB1LmZI=",
+        "owner": "huggingface",
+        "repo": "kernels",
+        "rev": "870e825d881664e39f9287a27a74ef63ff3c545e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernels",
+        "rev": "870e825d881664e39f9287a27a74ef63ff3c545e",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1776927958,
+        "narHash": "sha256-XOzEtft7E0P6TgQViLUOQeGHlEYiQ0+FY24BPEksj6s=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "fec2c46cca5bf9767486a290abae51200b656d69",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "ref": "nixos-unstable-small",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "rust-overlay": {
+      "inputs": {
+        "nixpkgs": [
+          "kernel-builder",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1776914043,
+        "narHash": "sha256-qug5r56yW1qOsjSI99l3Jm15JNT9CvS2otkXNRNtrPI=",
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "rev": "2d35c4358d7de3a0e606a6e8b27925d981c01cc3",
+        "type": "github"
+      },
+      "original": {
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "type": "github"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

tests/test_masked_mha_runtime.py CHANGED Viewed

@@ -38,6 +38,33 @@ class SourceOps:
         self.ops.forward_static(q, k, v, logits, out, float(scale))
         return out
 def load_source_ops():
     from torch.utils.cpp_extension import load
@@ -62,6 +89,7 @@ def load_source_ops():
         sources=[
             str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
             str(PACKAGE / "csrc" / "attention_mha_masked.cu"),
         ],
         extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
         extra_cflags=["-O3", "-DCUDA_KERNEL"],
@@ -107,6 +135,32 @@ def run_case(ops, dtype, sq, sk, heads, dim, fused_stride=False):
     out = torch.empty_like(q, memory_format=torch.contiguous_format)
     got = ops.forward_static(q, k, v, logits=logits, out=out)
     torch.cuda.synchronize()
     ref = torch.nn.functional.scaled_dot_product_attention(
         q.permute(1, 0, 2).unsqueeze(0).float(),
         k.permute(1, 0, 2).unsqueeze(0).float(),
@@ -128,9 +182,15 @@ def run_case(ops, dtype, sq, sk, heads, dim, fused_stride=False):
     graph = torch.cuda.CUDAGraph()
     torch.cuda.synchronize()
     with torch.cuda.graph(graph):
-        ops.forward_static(
-            static_q, static_k, static_v, logits=logits, out=out
-        )
     graph.replay()
     first = out.clone()
     graph.replay()
@@ -144,6 +204,54 @@ def run_case(ops, dtype, sq, sk, heads, dim, fused_stride=False):
     )
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--backend", choices=["source", "installed"], default="source")
@@ -166,7 +274,13 @@ def main():
         ])
     for case in cases:
         run_case(ops, *case)
-    print(f"masked-mha-runtime {args.backend} {args.mode}: passed {len(cases)}/{len(cases)}")
 if __name__ == "__main__":

         self.ops.forward_static(q, k, v, logits, out, float(scale))
         return out
+    def attention_mha_fp16_masked(self, q, k, v, *, logits, out, scale=None):
+        scale = q.shape[-1] ** -0.5 if scale is None else scale
+        self.ops.attention_mha_fp16_masked(
+            q, k, v, logits, out, float(scale)
+        )
+        return out
+    def attention_mha_bf16_masked(
+        self, q, k, v, *, logits, out, qkv_token_stride=None, scale=None
+    ):
+        scale = q.shape[-1] ** -0.5 if scale is None else scale
+        if qkv_token_stride is None:
+            qkv_token_stride = q.stride(0)
+        self.ops.attention_mha_bf16_masked(
+            q, k, v, logits, out, float(scale), int(qkv_token_stride)
+        )
+        return out
+    def forward_seqused_static(
+        self, q, k, v, valid_k, *, logits, out, scale=None
+    ):
+        scale = q.shape[-1] ** -0.5 if scale is None else scale
+        self.ops.forward_seqused_static(
+            q, k, v, valid_k, logits, out, float(scale)
+        )
+        return out
 def load_source_ops():
     from torch.utils.cpp_extension import load
         sources=[
             str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
             str(PACKAGE / "csrc" / "attention_mha_masked.cu"),
+            str(PACKAGE / "csrc" / "attention_seqused_fused.cu"),
         ],
         extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
         extra_cflags=["-O3", "-DCUDA_KERNEL"],
     out = torch.empty_like(q, memory_format=torch.contiguous_format)
     got = ops.forward_static(q, k, v, logits=logits, out=out)
     torch.cuda.synchronize()
+    alias_logits = ops.allocate_workspace(q, k)
+    alias_logits.fill_(float("nan"))
+    alias_out = torch.empty_like(q, memory_format=torch.contiguous_format)
+    if dtype is torch.float16:
+        alias_got = ops.attention_mha_fp16_masked(
+            q, k, v, logits=alias_logits, out=alias_out
+        )
+    else:
+        alias_got = ops.attention_mha_bf16_masked(
+            q, k, v, logits=alias_logits, out=alias_out,
+            qkv_token_stride=q.stride(0),
+        )
+    torch.cuda.synchronize()
+    if not torch.equal(alias_got, got):
+        raise AssertionError("explicit masked MHA entry differs from forward_static")
+    if dtype is torch.bfloat16:
+        try:
+            ops.attention_mha_bf16_masked(
+                q, k, v, logits=alias_logits, out=alias_out,
+                qkv_token_stride=q.stride(0) + 1,
+            )
+        except RuntimeError as exc:
+            if "qkv_token_stride" not in str(exc):
+                raise
+        else:
+            raise AssertionError("invalid qkv_token_stride was not rejected")
     ref = torch.nn.functional.scaled_dot_product_attention(
         q.permute(1, 0, 2).unsqueeze(0).float(),
         k.permute(1, 0, 2).unsqueeze(0).float(),
     graph = torch.cuda.CUDAGraph()
     torch.cuda.synchronize()
     with torch.cuda.graph(graph):
+        if dtype is torch.float16:
+            ops.attention_mha_fp16_masked(
+                static_q, static_k, static_v, logits=logits, out=out
+            )
+        else:
+            ops.attention_mha_bf16_masked(
+                static_q, static_k, static_v, logits=logits, out=out,
+                qkv_token_stride=static_q.stride(0),
+            )
     graph.replay()
     first = out.clone()
     graph.replay()
     )
+def run_seqused_case(ops, sq, sk, valid, heads=8, dim=256):
+    torch.manual_seed(164000 + sq + sk + valid)
+    q = torch.randn((sq, heads, dim), device="cuda", dtype=torch.float16)
+    k = torch.randn((sk, dim), device="cuda", dtype=torch.float16)
+    v = torch.randn_like(k)
+    valid_k = torch.tensor([valid], device="cuda", dtype=torch.int32)
+    logits = torch.full(
+        (sq * heads, sk), float("nan"), device="cuda", dtype=torch.float16
+    )
+    out = torch.empty_like(q)
+    got = ops.forward_seqused_static(
+        q, k, v, valid_k, logits=logits, out=out
+    )
+    torch.cuda.synchronize()
+    ref = torch.nn.functional.scaled_dot_product_attention(
+        q.permute(1, 0, 2).unsqueeze(0).float(),
+        k[:valid].unsqueeze(0).unsqueeze(0).float(),
+        v[:valid].unsqueeze(0).unsqueeze(0).float(),
+        enable_gqa=True,
+    ).squeeze(0).permute(1, 0, 2).half()
+    max_abs, p99_abs, cosine = metrics(got, ref)
+    if not torch.isfinite(got.float()).all() or cosine < 0.999 or p99_abs > 0.00390625:
+        raise AssertionError(
+            f"seqused sq={sq} sk={sk} valid={valid} h={heads} d={dim}: "
+            f"max={max_abs} p99={p99_abs} cos={cosine}"
+        )
+    if not torch.equal(logits[:, valid:], torch.zeros_like(logits[:, valid:])):
+        raise AssertionError("seqused probabilities beyond valid_k are not zero")
+    residual = torch.randn_like(out)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        ops.forward_seqused_static(
+            q, k, v, valid_k, logits=logits, out=out
+        )
+    graph.replay()
+    first = out.clone()
+    graph.replay()
+    torch.cuda.synchronize()
+    if not torch.equal(first, out):
+        raise AssertionError("seqused CUDA Graph replay is not bitwise deterministic")
+    del residual
+    print(
+        f"PASS seqused sq={sq} sk={sk} valid={valid} h={heads} d={dim} "
+        f"max={max_abs:.6f} p99={p99_abs:.6f} cos={cosine:.8f}"
+    )
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--backend", choices=["source", "installed"], default="source")
         ])
     for case in cases:
         run_case(ops, *case)
+    seqused_cases = [(10, 456, 456)]
+    if args.mode == "full":
+        seqused_cases.extend([(10, 968, 456), (10, 968, 712), (10, 968, 968)])
+    for case in seqused_cases:
+        run_seqused_case(ops, *case)
+    total = len(cases) + len(seqused_cases)
+    print(f"masked-mha-runtime {args.backend} {args.mode}: passed {total}/{total}")
 if __name__ == "__main__":

torch-ext/masked_mha_runtime/__init__.py CHANGED Viewed

@@ -23,6 +23,31 @@ def _forward_static_fake(q, k, v, logits, out, scale: float) -> None:
         raise RuntimeError("logits stride must cover sequence_kv")
 def allocate_workspace(q: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
     """Allocate padded logits scratch once, outside the hot path."""
     stride = (k.shape[0] + 7) // 8 * 8
@@ -47,6 +72,43 @@ def forward_static(
     return out
 def forward(q, k, v, *, scale: Optional[float] = None):
     """Convenience allocation wrapper; use ``forward_static`` in hot paths."""
     logits = allocate_workspace(q, k)
@@ -54,4 +116,28 @@ def forward(q, k, v, *, scale: Optional[float] = None):
     return forward_static(q, k, v, logits=logits, out=out, scale=scale)
-__all__ = ["allocate_workspace", "forward", "forward_static"]

         raise RuntimeError("logits stride must cover sequence_kv")
+@torch.library.register_fake(add_op_namespace_prefix("attention_mha_fp16_masked"))
+def _attention_mha_fp16_masked_fake(q, k, v, logits, out, scale: float) -> None:
+    _forward_static_fake(q, k, v, logits, out, scale)
+@torch.library.register_fake(add_op_namespace_prefix("attention_mha_bf16_masked"))
+def _attention_mha_bf16_masked_fake(
+    q, k, v, logits, out, scale: float, qkv_token_stride: int
+) -> None:
+    _forward_static_fake(q, k, v, logits, out, scale)
+    if qkv_token_stride != q.stride(0):
+        raise RuntimeError("qkv_token_stride must match q.stride(0)")
+@torch.library.register_fake(add_op_namespace_prefix("forward_seqused_static"))
+def _forward_seqused_static_fake(q, k, v, valid_k, logits, out, scale: float) -> None:
+    del scale
+    if q.dim() != 3 or k.dim() != 2 or v.dim() != 2:
+        raise RuntimeError("q must be (sequence, heads, head_dim); k/v must be (max_sequence, head_dim)")
+    if valid_k.numel() != 1 or valid_k.dtype != torch.int32:
+        raise RuntimeError("valid_k must be an int32 scalar tensor")
+    if out.shape != q.shape or logits.shape != (q.shape[0] * q.shape[1], k.shape[0]):
+        raise RuntimeError("caller-owned output shapes do not match q/k")
 def allocate_workspace(q: torch.Tensor, k: torch.Tensor) -> torch.Tensor:
     """Allocate padded logits scratch once, outside the hot path."""
     stride = (k.shape[0] + 7) // 8 * 8
     return out
+def attention_mha_fp16_masked(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    logits: torch.Tensor,
+    out: torch.Tensor,
+    scale: Optional[float] = None,
+) -> torch.Tensor:
+    """Run the deterministic valid-column FP16 masked MHA kernel."""
+    if scale is None:
+        scale = 1.0 / math.sqrt(q.shape[-1])
+    ops.attention_mha_fp16_masked(q, k, v, logits, out, float(scale))
+    return out
+def attention_mha_bf16_masked(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    logits: torch.Tensor,
+    out: torch.Tensor,
+    qkv_token_stride: Optional[int] = None,
+    scale: Optional[float] = None,
+) -> torch.Tensor:
+    """Run BF16 masked MHA directly from contiguous or fused-QKV views."""
+    if scale is None:
+        scale = 1.0 / math.sqrt(q.shape[-1])
+    if qkv_token_stride is None:
+        qkv_token_stride = q.stride(0)
+    ops.attention_mha_bf16_masked(
+        q, k, v, logits, out, float(scale), int(qkv_token_stride)
+    )
+    return out
 def forward(q, k, v, *, scale: Optional[float] = None):
     """Convenience allocation wrapper; use ``forward_static`` in hot paths."""
     logits = allocate_workspace(q, k)
     return forward_static(q, k, v, logits=logits, out=out, scale=scale)
+def forward_seqused_static(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    valid_k: torch.Tensor,
+    *,
+    logits: torch.Tensor,
+    out: torch.Tensor,
+    scale: Optional[float] = None,
+) -> torch.Tensor:
+    """Run FP16 shared-KV attention with a graph-resident valid-key count."""
+    if scale is None:
+        scale = 1.0 / math.sqrt(q.shape[-1])
+    ops.forward_seqused_static(q, k, v, valid_k, logits, out, float(scale))
+    return out
+__all__ = [
+    "attention_mha_bf16_masked",
+    "attention_mha_fp16_masked",
+    "allocate_workspace",
+    "forward",
+    "forward_static",
+    "forward_seqused_static",
+]

torch-ext/torch_binding.cpp CHANGED Viewed

@@ -8,6 +8,7 @@
 #include <limits>
 #include "attention_mha_masked.cuh"
 #include "registration.h"
 namespace {
@@ -94,11 +95,84 @@ void masked_mha_forward_static(
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
 }  // namespace
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("forward_static(Tensor q, Tensor k, Tensor v, Tensor! logits, Tensor! out, float scale) -> ()");
   ops.impl("forward_static", torch::kCUDA, &masked_mha_forward_static);
 }
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

 #include <limits>
 #include "attention_mha_masked.cuh"
+#include "attention_seqused_fused.cuh"
 #include "registration.h"
 namespace {
   C10_CUDA_KERNEL_LAUNCH_CHECK();
 }
+void attention_mha_fp16_masked_op(
+    torch::Tensor const& q, torch::Tensor const& k, torch::Tensor const& v,
+    torch::Tensor& logits, torch::Tensor& out, double scale) {
+  TORCH_CHECK(q.scalar_type() == torch::kFloat16,
+              "attention_mha_fp16_masked requires FP16 q/k/v");
+  masked_mha_forward_static(q, k, v, logits, out, scale);
+}
+void attention_mha_bf16_masked_op(
+    torch::Tensor const& q, torch::Tensor const& k, torch::Tensor const& v,
+    torch::Tensor& logits, torch::Tensor& out, double scale,
+    int64_t qkv_token_stride) {
+  TORCH_CHECK(q.scalar_type() == torch::kBFloat16,
+              "attention_mha_bf16_masked requires BF16 q/k/v");
+  TORCH_CHECK(qkv_token_stride == q.stride(0),
+              "qkv_token_stride must match q.stride(0)");
+  masked_mha_forward_static(q, k, v, logits, out, scale);
+}
+void masked_mha_forward_seqused_static(
+    torch::Tensor const& q, torch::Tensor const& k, torch::Tensor const& v,
+    torch::Tensor const& valid_k, torch::Tensor& logits, torch::Tensor& out,
+    double scale) {
+  check_qkv(q, "q", torch::kFloat16);
+  TORCH_CHECK(k.is_cuda() && v.is_cuda() && k.scalar_type() == torch::kFloat16 &&
+                  v.scalar_type() == torch::kFloat16 && k.dim() == 2 &&
+                  v.dim() == 2 && k.is_contiguous() && v.is_contiguous() &&
+                  k.sizes() == v.sizes(),
+              "k/v must be contiguous FP16 tensors with shape (S_kv_max, D)");
+  TORCH_CHECK(q.is_contiguous() && q.size(2) == k.size(1),
+              "q must be contiguous and share head_dim with k/v");
+  TORCH_CHECK(k.size(0) <= 1024,
+              "forward_seqused_static supports S_kv_max <= 1024");
+  TORCH_CHECK(valid_k.is_cuda() && valid_k.scalar_type() == torch::kInt &&
+                  valid_k.numel() == 1 && valid_k.is_contiguous(),
+              "valid_k must be a contiguous CUDA int32 scalar tensor");
+  TORCH_CHECK(q.get_device() == k.get_device() && q.get_device() == v.get_device() &&
+                  q.get_device() == valid_k.get_device(),
+              "q/k/v/valid_k must be on the same device");
+  TORCH_CHECK(logits.is_cuda() && logits.scalar_type() == torch::kFloat16 &&
+                  logits.is_contiguous() && logits.dim() == 2 &&
+                  logits.size(0) == q.size(0) * q.size(1) &&
+                  logits.size(1) == k.size(0),
+              "logits must be contiguous FP16 with shape (S_q * H, S_kv_max)");
+  TORCH_CHECK(out.is_cuda() && out.scalar_type() == torch::kFloat16 &&
+                  out.is_contiguous() && out.sizes() == q.sizes(),
+              "out must be contiguous FP16 and match q");
+  c10::cuda::CUDAGuard guard(q.device());
+  auto stream = at::cuda::getCurrentCUDAStream(q.get_device()).stream();
+  auto handle = at::cuda::getCurrentCUDABlasHandle();
+  attention_qkv_fp16_seqused_v2(
+      handle, static_cast<const __half*>(q.data_ptr()),
+      static_cast<const __half*>(k.data_ptr()),
+      static_cast<const __half*>(v.data_ptr()),
+      static_cast<__half*>(logits.data_ptr()),
+      static_cast<__half*>(out.data_ptr()), checked_int(q.size(0), "S_q"),
+      checked_int(k.size(0), "S_kv_max"), checked_int(q.size(1), "heads"),
+      checked_int(q.size(2), "head_dim"),
+      static_cast<const int*>(valid_k.data_ptr()), static_cast<float>(scale),
+      stream);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
 }  // namespace
 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("forward_static(Tensor q, Tensor k, Tensor v, Tensor! logits, Tensor! out, float scale) -> ()");
+  ops.def("attention_mha_fp16_masked(Tensor q, Tensor k, Tensor v, Tensor! logits, Tensor! out, float scale) -> ()");
+  ops.def("attention_mha_bf16_masked(Tensor q, Tensor k, Tensor v, Tensor! logits, Tensor! out, float scale, int qkv_token_stride) -> ()");
+  ops.def("forward_seqused_static(Tensor q, Tensor k, Tensor v, Tensor valid_k, Tensor! logits, Tensor! out, float scale) -> ()");
   ops.impl("forward_static", torch::kCUDA, &masked_mha_forward_static);
+  ops.impl("attention_mha_fp16_masked", torch::kCUDA,
+           &attention_mha_fp16_masked_op);
+  ops.impl("attention_mha_bf16_masked", torch::kCUDA,
+           &attention_mha_bf16_masked_op);
+  ops.impl("forward_seqused_static", torch::kCUDA,
+           &masked_mha_forward_seqused_static);
 }
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)