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