Summary

#163

• Converted the attention mask tensor storage in global memory from float16/bfloat16 to bool.
• Motivation: cut mask memory footprint and HBM bandwidth by ~50%, while preserving numerical behavior by converting to compute dtype only when loading into SMEM.

Design

• Store mask in global memory as bool (1 byte/element).
• On gmem→smem transfer, convert per-tile from bool to compute Element (fp16/bf16) and proceed with the same softmax/masking pipeline.
• Keep smem/register representations and downstream math unchanged (Element(0/1) for mask).
• Preserve bias path and cp.async usage; only the mask path uses a convert-on-copy code path.

Alternatives considered:

• Full float mask: simpler, but doubles mask memory/bandwidth.
• Bitpacked mask (1 bit/element): more memory efficient but substantially more complex (alignment, unpack cost, divergence), deferred.

Changes

• CUDA kernels:
  • Forward: read mask as const bool* and use copy_MN<..., Bool_to_Element=true>(...) to convert to Element in smem.
  • Backward: same bool gmem representation and convert-on-copy in all mask loads.
• Utils:
  • Extended copy_MN with a Bool_to_Element compile-time path (defaults keep old behavior). Maintains signature compatibility via defaults.
• API (C++):
  • Forward/Varlen forward: enforce mask.dtype() == torch::kBool.
  • Backward/Varlen backward: enforce mask.dtype() == torch::kBool (aligned with kernel behavior).
  • Strides remain element-based; no change needed.
• Python interface:
  • Normalize input mask to torch.bool if not already.
  • When padding K/V to multiples of 128, pad mask with False (and bias with -inf as before).

Implementation Notes

• gmem bool → smem Element conversion uses thread-sliced loops over tile fragments; preserves existing OOB clearing logic and identity/predicate layouts.
• cp.async remains for Q/K/V and bias; mask conversion path does not use cp.async (impact negligible as mask bandwidth is small relative to Q/K/V).
• OR-reduce on mask for tile activity detection unchanged, operates on Element(0/1) in smem.
• Kept bias copy paths unchanged for performance (cp.async friendly).
• Kept occupancy unaffected; smem usage unchanged (mask still Element in smem).

Tests

• Correctness:
  • Forward parity vs. previous float-mask path on random tensors across:
    • head_dim ∈ {64, 96, 128, 192, 256}, seqlen_q/k ∈ {128, 512, 2k, 8k}
    • causal / non-causal
    • GQA/MQA
  • Backward parity: gradcheck on (q, k, v, bias) with the same masks.
• Integration:
  • Varlen forward/backward and split-k dispatch paths.
  • OR-reduce skip path: ensure inactive tiles produce zero outputs/gradients.
• Performance:
  • Microbench: end-to-end latency and HBM reads; verify ~50% reduction in mask bytes and neutral to ±3% runtime delta overall.

Docs

• Updated user docs and examples:
  • attn_mask must be torch.bool.
  • Padding behavior: mask padded with False.
  • Notes on memory reduction and expected performance characteristics.

Checklist

• Linked issue provided
• API stable (mask now bool)
• Tests added or updated (fwd/bwd parity, varlen, causal, GQA/MQA)
• Docs added or updated (dtype and padding semantics)
• No known performance regressions (microbench within ±3%)