Publish GROOT N1.7 Thor kernel sources

Browse files

Files changed (14) hide show

CARD.md +23 -0
README.md +23 -0
SYNC.md +6 -0
VALIDATION.md +5 -0
benchmarks/README.md +4 -0
benchmarks/RESULTS.md +5 -0
build.toml +24 -0
csrc/attention_mha_masked.cu +243 -0
csrc/attention_mha_masked.cuh +22 -0
examples/basic_usage.py +10 -0
flake.nix +6 -0
tests/test_masked_mha_runtime.py +173 -0
torch-ext/masked_mha_runtime/__init__.py +57 -0
torch-ext/torch_binding.cpp +104 -0

CARD.md ADDED Viewed

	@@ -0,0 +1,23 @@

+---
+tags: [kernel, cuda, attention, inference, cuda-graphs]
+library_name: kernels
+license: apache-2.0
+---
+# Masked MHA Runtime
+Allocation-free FP16/BF16 attention that masks padded logits inside softmax,
+removing the per-call `-inf` pre-fill. BF16 accepts fused-QKV token strides.
+## API
+- `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.
+Rows wider than 1024 keys use a deterministic multi-pass softmax.
+This package contains the native masked-MHA execution path validated in
+FlashRT's GROOT N1.7 Thor runtime. It is separate from FlashAttention-4.

README.md ADDED Viewed

	@@ -0,0 +1,23 @@

+# masked-mha-runtime
+FlashRT native SM110 masked FP16/BF16 MHA for fixed-shape CUDA Graph runtimes.
+It masks padded logits during softmax, supports fused-QKV token strides, and
+keeps caller-owned logits/output buffers stable across graph replay.
+```python
+from kernels import get_kernel
+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`.

SYNC.md ADDED Viewed

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

+# Source sync
+- Upstream: `flashrt-project/FlashRT`
+- Commit: `24df793f4fa2d50780aea03b644208c6e0cb4162`
+- Source: `csrc/kernels/attention_mha_masked.cu`
+- Local changes: Tensor validation and `torch.library` binding only.

VALIDATION.md ADDED Viewed

	@@ -0,0 +1,5 @@

+# Validation
+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.

benchmarks/README.md ADDED Viewed

	@@ -0,0 +1,4 @@

+# Benchmark
+Benchmark the installed artifact against PyTorch SDPA and the matching
+FlashRT native symbol on the same Thor host. Report both static-buffer paths.

benchmarks/RESULTS.md ADDED Viewed

	@@ -0,0 +1,5 @@

+# 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.

build.toml ADDED Viewed

	@@ -0,0 +1,24 @@

+[general]
+name = "masked-mha-runtime"
+license = "Apache-2.0"
+version = 1
+backends = ["cuda"]
+[general.cuda]
+minver = "13"
+[general.hub]
+repo-id = "flashrt/masked-mha-runtime"
+[torch]
+include = ["csrc"]
+src = ["torch-ext/torch_binding.cpp"]
+[kernel.masked_mha_runtime]
+backend = "cuda"
+depends = ["torch"]
+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"]

csrc/attention_mha_masked.cu ADDED Viewed

	@@ -0,0 +1,243 @@

+// ============================================================================
+//  FlashRT — MHA attention without the -inf logits pre-fill.
+//
+//  The plain ``attention_mha_{fp16,bf16}`` kernels softmax over the padded
+//  logits width, so callers must pre-fill the whole (NH, max_q, max_kv)
+//  scratch with -inf every invocation — a full DRAM sweep per layer. These
+//  variants run a column-masked softmax that reads and writes only the
+//  valid S_kv columns (row stride = padded width); the PV GEMM already
+//  uses k = S_kv, so the padding is never read anywhere and the pre-fill
+//  disappears.
+//
+//  Any S_kv is supported: rows up to SMM_MAX_COLS use a register-tiled
+//  softmax sized to the row, wider rows fall back to a multi-pass kernel
+//  that holds no per-column registers.
+//
+//  Additive: new symbols only.
+// ============================================================================
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+#include <cublas_v2.h>
+#define SMM_WARP 32
+#define SMM_MAX_COLS 1024
+#define SMM_ITERS (SMM_MAX_COLS / SMM_WARP)
+namespace {
+template <typename T>
+__device__ __forceinline__ float to_f(T v);
+template <>
+__device__ __forceinline__ float to_f<__half>(__half v) { return __half2float(v); }
+template <>
+__device__ __forceinline__ float to_f<__nv_bfloat16>(__nv_bfloat16 v) { return __bfloat162float(v); }
+template <typename T>
+__device__ __forceinline__ T from_f(float v);
+template <>
+__device__ __forceinline__ __half from_f<__half>(float v) { return __float2half(v); }
+template <>
+__device__ __forceinline__ __nv_bfloat16 from_f<__nv_bfloat16>(float v) { return __float2bfloat16(v); }
+// Row-wise softmax over the first ``cols_valid`` of each ``cols_pad``-wide
+// row. Padding columns are neither read nor written.
+//
+// ITERS is the per-thread register tile, dispatched from the actual valid
+// column count: a 41-key DiT self-attention row needs 2 registers, not the
+// 32 a worst-case 1024-key row would. Sizing it per call keeps the loops
+// fully unrolled without paying occupancy for columns that do not exist.
+template <typename T, int ITERS>
+__global__ void softmax_masked_kernel(T* data, int rows, int cols_pad,
+                                      int cols_valid) {
+    const int lane = threadIdx.x % SMM_WARP;
+    const int row = blockIdx.x;
+    if (row >= rows) return;
+    T* src = data + (long)row * cols_pad;
+    float reg[ITERS];
+    float mx = -1e30f;
+    #pragma unroll
+    for (int it = 0; it < ITERS; ++it) {
+        const int c = it * SMM_WARP + lane;
+        if (c < cols_valid) {
+            reg[it] = to_f<T>(src[c]);
+            mx = fmaxf(mx, reg[it]);
+        } else {
+            reg[it] = -1e30f;
+        }
+    }
+    #pragma unroll
+    for (int o = 16; o > 0; o >>= 1)
+        mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, o));
+    float sum = 0.f;
+    #pragma unroll
+    for (int it = 0; it < ITERS; ++it) {
+        const int c = it * SMM_WARP + lane;
+        if (c < cols_valid) {
+            reg[it] = __expf(reg[it] - mx);
+            sum += reg[it];
+        }
+    }
+    #pragma unroll
+    for (int o = 16; o > 0; o >>= 1)
+        sum += __shfl_xor_sync(0xffffffffu, sum, o);
+    const float inv = 1.0f / sum;
+    #pragma unroll
+    for (int it = 0; it < ITERS; ++it) {
+        const int c = it * SMM_WARP + lane;
+        if (c < cols_valid)
+            src[c] = from_f<T>(reg[it] * inv);
+    }
+}
+// Register-tiled variants above cap at SMM_MAX_COLS columns. Rows wider
+// than that go through this multi-pass kernel instead: it keeps no
+// per-column registers, so it is correct for any width. Three passes over
+// the row (max, exp+sum, scale) make it slower than the tiled path, which
+// is why it only runs past the tiled path's reach.
+template <typename T>
+__global__ void softmax_masked_wide_kernel(T* data, int rows, int cols_pad,
+                                           int cols_valid) {
+    const int lane = threadIdx.x % SMM_WARP;
+    const int row = blockIdx.x;
+    if (row >= rows) return;
+    T* src = data + (long)row * cols_pad;
+    float mx = -1e30f;
+    for (int c = lane; c < cols_valid; c += SMM_WARP)
+        mx = fmaxf(mx, to_f<T>(src[c]));
+    #pragma unroll
+    for (int o = 16; o > 0; o >>= 1)
+        mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, o));
+    float sum = 0.f;
+    for (int c = lane; c < cols_valid; c += SMM_WARP) {
+        const float e = __expf(to_f<T>(src[c]) - mx);
+        src[c] = from_f<T>(e);
+        sum += e;
+    }
+    #pragma unroll
+    for (int o = 16; o > 0; o >>= 1)
+        sum += __shfl_xor_sync(0xffffffffu, sum, o);
+    const float inv = 1.0f / sum;
+    for (int c = lane; c < cols_valid; c += SMM_WARP)
+        src[c] = from_f<T>(to_f<T>(src[c]) * inv);
+}
+template <typename T>
+inline void launch_softmax_masked(T* data, int rows, int cols_pad,
+                                  int cols_valid, cudaStream_t stream) {
+    const int iters = (cols_valid + SMM_WARP - 1) / SMM_WARP;
+    if (iters > SMM_ITERS) {
+        softmax_masked_wide_kernel<T><<<rows, SMM_WARP, 0, stream>>>(
+            data, rows, cols_pad, cols_valid);
+        return;
+    }
+    if (iters <= 2) {
+        softmax_masked_kernel<T, 2><<<rows, SMM_WARP, 0, stream>>>(
+            data, rows, cols_pad, cols_valid);
+    } else if (iters <= 4) {
+        softmax_masked_kernel<T, 4><<<rows, SMM_WARP, 0, stream>>>(
+            data, rows, cols_pad, cols_valid);
+    } else if (iters <= 8) {
+        softmax_masked_kernel<T, 8><<<rows, SMM_WARP, 0, stream>>>(
+            data, rows, cols_pad, cols_valid);
+    } else if (iters <= 16) {
+        softmax_masked_kernel<T, 16><<<rows, SMM_WARP, 0, stream>>>(
+            data, rows, cols_pad, cols_valid);
+    } else {
+        softmax_masked_kernel<T, SMM_ITERS><<<rows, SMM_WARP, 0, stream>>>(
+            data, rows, cols_pad, cols_valid);
+    }
+}
+}  // namespace
+extern "C" {
+void attention_mha_fp16_masked(
+    cublasHandle_t handle,
+    const __half* Q, const __half* K, const __half* V,
+    __half* logits, __half* out,
+    int S_q, int S_kv, int NH, int HD,
+    float attn_scale, cudaStream_t stream) {
+    cublasSetStream(handle, stream);
+    const int S_kv_pad = ((S_kv + 7) / 8) * 8;
+    float zero = 0.0f, one = 1.0f;
+    const long long strideC = (long long)S_q * S_kv_pad;
+    cublasGemmStridedBatchedEx(handle,
+        CUBLAS_OP_T, CUBLAS_OP_N,
+        S_kv, S_q, HD,
+        &attn_scale,
+        K, CUDA_R_16F, NH * HD, (long long)HD,
+        Q, CUDA_R_16F, NH * HD, (long long)HD,
+        &zero,
+        logits, CUDA_R_16F, S_kv_pad, strideC,
+        NH,
+        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
+    launch_softmax_masked<__half>(logits, NH * S_q, S_kv_pad, S_kv, stream);
+    cublasGemmStridedBatchedEx(handle,
+        CUBLAS_OP_N, CUBLAS_OP_N,
+        HD, S_q, S_kv,
+        &one,
+        V, CUDA_R_16F, NH * HD, (long long)HD,
+        logits, CUDA_R_16F, S_kv_pad, strideC,
+        &zero,
+        out, CUDA_R_16F, NH * HD, (long long)HD,
+        NH,
+        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
+}
+void attention_mha_bf16_masked(
+    cublasHandle_t handle,
+    const __nv_bfloat16* Q, const __nv_bfloat16* K, const __nv_bfloat16* V,
+    __nv_bfloat16* logits, __nv_bfloat16* out,
+    int S_q, int S_kv, int NH, int HD,
+    float attn_scale, int logits_kv_stride, int qkv_token_stride,
+    cudaStream_t stream) {
+    cublasSetStream(handle, stream);
+    const int S_kv_pad = ((S_kv + 7) / 8) * 8;
+    const int kv_stride = (logits_kv_stride > 0) ? logits_kv_stride : S_kv_pad;
+    // Token stride (elements) of the Q/K/V sources. NH*HD for packed
+    // per-site buffers; 3*NH*HD lets the fused-QKV GEMM output be read in
+    // place (no split copies).
+    const int tstride = (qkv_token_stride > 0) ? qkv_token_stride : NH * HD;
+    float zero = 0.0f, one = 1.0f;
+    const long long strideC = (long long)S_q * kv_stride;
+    cublasGemmStridedBatchedEx(handle,
+        CUBLAS_OP_T, CUBLAS_OP_N,
+        S_kv, S_q, HD,
+        &attn_scale,
+        K, CUDA_R_16BF, tstride, (long long)HD,
+        Q, CUDA_R_16BF, tstride, (long long)HD,
+        &zero,
+        logits, CUDA_R_16BF, kv_stride, strideC,
+        NH,
+        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
+    launch_softmax_masked<__nv_bfloat16>(logits, NH * S_q, kv_stride, S_kv,
+                                         stream);
+    cublasGemmStridedBatchedEx(handle,
+        CUBLAS_OP_N, CUBLAS_OP_N,
+        HD, S_q, S_kv,
+        &one,
+        V, CUDA_R_16BF, tstride, (long long)HD,
+        logits, CUDA_R_16BF, kv_stride, strideC,
+        &zero,
+        out, CUDA_R_16BF, NH * HD, (long long)HD,
+        NH,
+        CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT);
+}
+}  // extern "C"

csrc/attention_mha_masked.cuh ADDED Viewed

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

+#pragma once
+#include <cublas_v2.h>
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+extern "C" {
+void attention_mha_fp16_masked(
+    cublasHandle_t handle, const __half* q, const __half* k,
+    const __half* v, __half* logits, __half* out, int sequence_q,
+    int sequence_kv, int heads, int head_dim, float scale,
+    cudaStream_t stream);
+void attention_mha_bf16_masked(
+    cublasHandle_t handle, const __nv_bfloat16* q,
+    const __nv_bfloat16* k, const __nv_bfloat16* v,
+    __nv_bfloat16* logits, __nv_bfloat16* out, int sequence_q,
+    int sequence_kv, int heads, int head_dim, float scale,
+    int logits_kv_stride, int qkv_token_stride, cudaStream_t stream);
+}

examples/basic_usage.py ADDED Viewed

	@@ -0,0 +1,10 @@

+import torch
+from kernels import get_kernel
+mha = get_kernel("flashrt/masked-mha-runtime", version=1)
+q = torch.randn((41, 32, 48), device="cuda", dtype=torch.bfloat16)
+k = torch.randn_like(q)
+v = torch.randn_like(q)
+logits = mha.allocate_workspace(q, k)
+out = torch.empty_like(q)
+mha.forward_static(q, k, v, logits=logits, out=out)

flake.nix ADDED Viewed

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

+{
+  description = "Flake for FlashRT masked MHA runtime kernels";
+  inputs.kernel-builder.url = "github:huggingface/kernels/870e825d881664e39f9287a27a74ef63ff3c545e";
+  outputs = { self, kernel-builder }:
+    kernel-builder.lib.genKernelFlakeOutputs { inherit self; path = ./.; };
+}

tests/test_masked_mha_runtime.py ADDED Viewed

	@@ -0,0 +1,173 @@

+#!/usr/bin/env python3
+"""Strict source and installed-artifact tests for masked-mha-runtime."""
+from __future__ import annotations
+import argparse
+import importlib
+import os
+import re
+import subprocess
+import sys
+from pathlib import Path
+import torch
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "masked-mha-runtime"
+REGISTRATION_INCLUDE = (
+    ROOT.parent / "kernels" / "kernel-builder" / "src" / "pyproject"
+    / "templates" / "torch"
+)
+class SourceOps:
+    def __init__(self, namespace: str):
+        self.ops = getattr(torch.ops, namespace)
+    @staticmethod
+    def allocate_workspace(q, k):
+        stride = (k.shape[0] + 7) // 8 * 8
+        return torch.empty(
+            (q.shape[1], q.shape[0], stride), device=q.device, dtype=q.dtype
+        )
+    def forward_static(self, q, k, v, *, logits, out, scale=None):
+        scale = q.shape[-1] ** -0.5 if scale is None else scale
+        self.ops.forward_static(q, k, v, logits, out, float(scale))
+        return out
+def load_source_ops():
+    from torch.utils.cpp_extension import load
+    nvcc = subprocess.check_output(
+        ["nvcc", "--version"], text=True
+    )
+    match = re.search(r"release\s+(\d+)\.", nvcc)
+    torch_cuda_major = int(torch.version.cuda.split(".", 1)[0])
+    if match and int(match.group(1)) != torch_cuda_major:
+        raise RuntimeError(
+            "source test requires PyTorch and nvcc from the same CUDA major; "
+            f"torch={torch.version.cuda}, nvcc={match.group(1)}.x. Use the "
+            "installed artifact or a matching isolated build environment."
+        )
+    major, minor = torch.cuda.get_device_capability(0)
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", f"{major}.{minor}")
+    namespace = "masked_mha_runtime_source_test"
+    load(
+        name=namespace,
+        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"],
+        extra_cuda_cflags=["-O3", "--use_fast_math", "-DCUDA_KERNEL"],
+        extra_ldflags=["-lcublas"],
+        is_python_module=False,
+        verbose=False,
+    )
+    return SourceOps(namespace)
+def load_installed_ops(artifact):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("masked_mha_runtime")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+def metrics(got, ref):
+    diff = (got.float() - ref.float()).abs()
+    cosine = torch.nn.functional.cosine_similarity(
+        got.float().flatten(), ref.float().flatten(), dim=0
+    ).item()
+    return float(diff.max()), float(torch.quantile(diff.flatten(), 0.99)), float(cosine)
+def run_case(ops, dtype, sq, sk, heads, dim, fused_stride=False):
+    torch.manual_seed(1000 + sq + sk + dim)
+    if fused_stride:
+        packed = torch.randn((sk, 3, heads, dim), device="cuda", dtype=dtype)
+        q = packed[:sq, 0]
+        k = packed[:, 1]
+        v = packed[:, 2]
+    else:
+        q = torch.randn((sq, heads, dim), device="cuda", dtype=dtype)
+        k = torch.randn((sk, heads, dim), device="cuda", dtype=dtype)
+        v = torch.randn_like(k)
+    logits = ops.allocate_workspace(q, k)
+    logits.fill_(float("nan"))
+    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(),
+        v.permute(1, 0, 2).unsqueeze(0).float(),
+    ).squeeze(0).permute(1, 0, 2).to(dtype)
+    max_abs, p99_abs, cosine = metrics(got, ref)
+    atol = 0.00390625 if dtype is torch.float16 else 0.015625
+    if not torch.isfinite(got.float()).all() or cosine < 0.999 or p99_abs > atol:
+        raise AssertionError(
+            f"dtype={dtype} sq={sq} sk={sk} h={heads} d={dim}: "
+            f"max={max_abs} p99={p99_abs} cos={cosine}"
+        )
+    # Preserve fused-QKV token strides. Cloning each view separately can
+    # normalize a size-one query dimension and destroy the shared stride.
+    static_q = q
+    static_k = k
+    static_v = v
+    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()
+    torch.cuda.synchronize()
+    if not torch.equal(out, first):
+        raise AssertionError("CUDA Graph replay is not bitwise deterministic")
+    print(
+        f"PASS {dtype} sq={sq} sk={sk} h={heads} d={dim} "
+        f"fused_stride={fused_stride} max={max_abs:.6f} "
+        f"p99={p99_abs:.6f} cos={cosine:.8f}"
+    )
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact")
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    args = parser.parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required")
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    cases = [
+        (torch.float16, 41, 41, 32, 48, False),
+        (torch.bfloat16, 41, 41, 32, 48, True),
+    ]
+    if args.mode == "full":
+        cases.extend([
+            (torch.float16, 1, 277, 16, 128, False),
+            (torch.bfloat16, 1, 1024, 1, 16, True),
+            (torch.bfloat16, 1, 1025, 1, 16, True),
+            (torch.bfloat16, 1, 2048, 1, 16, True),
+        ])
+    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__":
+    main()

torch-ext/masked_mha_runtime/__init__.py ADDED Viewed

	@@ -0,0 +1,57 @@

+"""Allocation-free masked MHA runtime operators."""
+from __future__ import annotations
+import math
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+@torch.library.register_fake(add_op_namespace_prefix("forward_static"))
+def _forward_static_fake(q, k, v, logits, out, scale: float) -> None:
+    del scale
+    if q.dim() != 3 or k.dim() != 3 or v.dim() != 3:
+        raise RuntimeError("q/k/v must have shape (sequence, heads, head_dim)")
+    if out.shape != q.shape:
+        raise RuntimeError("out must match q")
+    if logits.dim() != 3 or logits.shape[:2] != (q.shape[1], q.shape[0]):
+        raise RuntimeError("logits must have shape (heads, sequence_q, stride)")
+    if logits.shape[2] < k.shape[0]:
+        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
+    return torch.empty(
+        (q.shape[1], q.shape[0], stride), device=q.device, dtype=q.dtype
+    )
+def forward_static(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    logits: torch.Tensor,
+    out: torch.Tensor,
+    scale: Optional[float] = None,
+) -> torch.Tensor:
+    """Run MHA without pre-filling padded logits; all buffers are caller-owned."""
+    if scale is None:
+        scale = 1.0 / math.sqrt(q.shape[-1])
+    ops.forward_static(q, k, v, logits, out, float(scale))
+    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)
+    out = torch.empty_like(q, memory_format=torch.contiguous_format)
+    return forward_static(q, k, v, logits=logits, out=out, scale=scale)
+__all__ = ["allocate_workspace", "forward", "forward_static"]

torch-ext/torch_binding.cpp ADDED Viewed

	@@ -0,0 +1,104 @@

+#include <torch/all.h>
+#include <torch/library.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAException.h>
+#include <limits>
+#include "attention_mha_masked.cuh"
+#include "registration.h"
+namespace {
+int checked_int(int64_t value, const char* name) {
+  TORCH_CHECK(value > 0 && value <= std::numeric_limits<int>::max(),
+              name, " must fit in a positive int");
+  return static_cast<int>(value);
+}
+void check_qkv(torch::Tensor const& tensor, const char* name,
+               c10::ScalarType dtype) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be CUDA");
+  TORCH_CHECK(tensor.scalar_type() == dtype, name, " has the wrong dtype");
+  TORCH_CHECK(tensor.dim() == 3, name, " must have shape (S, H, D)");
+  TORCH_CHECK(tensor.stride(2) == 1 && tensor.stride(1) == tensor.size(2),
+              name, " must be contiguous within each token");
+}
+void masked_mha_forward_static(
+    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 ||
+                  q.scalar_type() == torch::kBFloat16,
+              "q must be FP16 or BF16");
+  check_qkv(q, "q", q.scalar_type());
+  check_qkv(k, "k", q.scalar_type());
+  check_qkv(v, "v", q.scalar_type());
+  TORCH_CHECK(q.size(1) == k.size(1) && q.size(1) == v.size(1) &&
+                  q.size(2) == k.size(2) && q.size(2) == v.size(2) &&
+                  k.size(0) == v.size(0),
+              "q/k/v head shapes must match");
+  TORCH_CHECK(q.get_device() == k.get_device() &&
+                  q.get_device() == v.get_device(),
+              "q/k/v must be on the same device");
+  TORCH_CHECK(out.is_cuda() && out.is_contiguous() &&
+                  out.scalar_type() == q.scalar_type() &&
+                  out.sizes() == q.sizes(),
+              "out must be contiguous and match q");
+  TORCH_CHECK(logits.is_cuda() && logits.scalar_type() == q.scalar_type() &&
+                  logits.dim() == 3 && logits.size(0) == q.size(1) &&
+                  logits.size(1) == q.size(0) &&
+                  logits.size(2) >= k.size(0) && logits.stride(2) == 1,
+              "logits must have shape (H, S_q, stride >= S_kv)");
+  TORCH_CHECK(logits.get_device() == q.get_device() &&
+                  out.get_device() == q.get_device(),
+              "outputs must be on the q device");
+  TORCH_CHECK(logits.stride(1) == logits.size(2) &&
+                  logits.stride(0) == logits.size(1) * logits.size(2),
+              "logits must use a dense padded row stride");
+  c10::cuda::CUDAGuard guard(q.device());
+  auto stream = at::cuda::getCurrentCUDAStream(q.get_device()).stream();
+  auto handle = at::cuda::getCurrentCUDABlasHandle();
+  const int sq = checked_int(q.size(0), "S_q");
+  const int sk = checked_int(k.size(0), "S_kv");
+  const int heads = checked_int(q.size(1), "heads");
+  const int dim = checked_int(q.size(2), "head_dim");
+  if (q.scalar_type() == torch::kFloat16) {
+    TORCH_CHECK(q.stride(0) == heads * dim &&
+                    k.stride(0) == heads * dim &&
+                    v.stride(0) == heads * dim,
+                "FP16 q/k/v must be contiguous across tokens");
+    attention_mha_fp16_masked(
+        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()), sq, sk, heads, dim,
+        static_cast<float>(scale), stream);
+  } else {
+    TORCH_CHECK(q.stride(0) == k.stride(0) && q.stride(0) == v.stride(0),
+                "BF16 q/k/v must share one token stride");
+    attention_mha_bf16_masked(
+        handle, static_cast<const __nv_bfloat16*>(q.data_ptr()),
+        static_cast<const __nv_bfloat16*>(k.data_ptr()),
+        static_cast<const __nv_bfloat16*>(v.data_ptr()),
+        static_cast<__nv_bfloat16*>(logits.data_ptr()),
+        static_cast<__nv_bfloat16*>(out.data_ptr()), sq, sk, heads, dim,
+        static_cast<float>(scale), checked_int(logits.size(2), "logits stride"),
+        checked_int(q.stride(0), "qkv token stride"), 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.impl("forward_static", torch::kCUDA, &masked_mha_forward_static);
+}
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)