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.
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 for the complete contract. Source provenance is FlashRT
commit 24df793f4fa2d50780aea03b644208c6e0cb4162.