Views
No views yet
1from kernels import get_kernel
2import torch
3
4ops = get_kernel("flashrt/masked-mha-runtime", version=1)
5logits = ops.allocate_workspace(q, k)
6out = torch.empty_like(q, memory_format=torch.contiguous_format)
7ops.attention_mha_fp16_masked(q, k, v, logits=logits, out=out)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.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.24df793f4fa2d50780aea03b644208c6e0cb4162.