[FEATURE SUPPORT] Broadcastable 4D mask/bias, 128‑rounded key length, stride‑0 broadcasting, and dbias reductions

[LoserCheems]

Summary

• Implements a robust contract for attention mask/bias with strict 4D shapes and broadcast over batch, heads, and query dims.
• Rounds key length to a multiple of 128 for mask/bias to match kernel access, preventing OOB and NaNs.
• Uses stride=0 broadcasting in CUDA for size‑1 dims to avoid materialization along B/H/S_q.
• Fixes swapped seqlen_q==1 path and aligns dbias reductions with broadcast semantics.
• Resolves issue [FEATURE REQUEST] Support {num_heads, num_kv_heads, 1} shaped bias in attention functions #189.

Design

• Contract for aux tensors:
  • Shapes: ({batch_size|1}, {num_heads|num_kv_heads|1}, {seqlen_q|1}, round_multiple(seqlen_k, 128)).
  • Boolean mask (True=keep), floating bias (additive).
  • Last dimension must be contiguous; only K-dim is materialized/padded to 128 alignment. Q/K/V remain unpadded; no tail copies for K/V to reduce memory overhead.
• Broadcasting:
  • CUDA sets zero strides for singleton B/H/S_q to broadcast without expanding.
  • Only the last K-dim is expanded/padded (S_k==1 → expand; S_k==seqlen_k → pad to 128).
• Swapped path (seqlen_q==1 and H>H_k):
  • Reshape heads to (H_k, ngroups) via views; preserve original mask/bias batch and head semantics.
  • Pass seqlenq_ngroups_swapped into params so batch/head strides and LSE writeback are correct.

Changes

• Python (flash_dmattn_interface.py):
  • Pads/expands mask/bias last dim to seqlen_k_rounded=round_multiple(seqlen_k, 128).
  • Enforces 4D mask/bias; avoids materializing broadcast dims; keeps last dim contiguous.
  • Backward slices returned dbias to the true seqlen_k so callers see trimmed gradients.
• CUDA C++ (csrc/flash_dmattn/flash_api.cpp):
  • set_params_fprop/dgrad:
    • mask/bias batch/head/row strides set to 0 for size‑1 dims (broadcast-safe).
    • Batch stride handling honors seqlenq_ngroups_swapped.
  • mha_fwd:
    • Enforces 4D mask/bias with K-dim = seqlen_k_rounded and contiguous last dim.
    • Validates B∈{B,1}, H∈{1,H_k,H}, S_q∈{1,L_q}; no implicit 3D unsqueeze/expand.
    • Swapped path uses reshape-only views for mask/bias and restores original views after compute.
  • mha_bwd:
    • Enforces the same 4D contract; allocates/accepts dbias with K-dim rounded to 128.
    • Correctly reduces dbias over groups, batch, and S_q when broadcast, then copies back into user-provided dbias.
• Documentation:
  • README/README_zh updated to reflect broadcastable 4D shapes and removal of “dbias” jargon in prose.

Implementation Notes

• No physical expand on B/H/S_q; broadcasting via stride=0 avoids large temporaries and keeps ABI stable.
• Only the K-dim is materialized/rounded, which the kernels iterate over; last dim is enforced contiguous.
• Swapped path:
  • Save original sizes for mask/bias, reshape to (H_k, ngroups) forms during compute, and restore after.
  • Pass seqlenq_ngroups_swapped into params; missing this led to NaNs in multi-batch/multi-head settings.
• Stability fix:
  • With mask/bias present, kernels iterate N blocks against rounded length. Aligning mask/bias K-dim to 128 prevents OOB; K/V are not padded to reduce memory.

Tests

• Benchmarks:
  • forward_equivalence.py: Passed for various B, H in {1,H_k,H}, S_q in {1,L_q}, causal/non‑causal, bf16/fp16.
  • backward_equivalence.py: dQ/dK/dV matched; dbias now non‑zero and matches reference; gradients sliced back to original seqlen_k on return.
• Regressions covered:
  • seqlen_q==1 swapped path with H>H_k across B>1.
  • mask/bias None vs provided; broadcast over batch/head/S_q dims.

Docs

• Updated shape conventions in README and README_zh:
  • Now documented as ({1|batch_size}, {1|num_kv_heads|num_heads}, {1|query_len}, {1|key_len}) with K rounded to 128.
  • Clarified broadcast semantics and gradient wording.

Checklist

• Linked issue provided: [FEATURE REQUEST] Support {num_heads, num_kv_heads, 1} shaped bias in attention functions #189
• API stable: mask/bias 4D contract with broadcast preserved
• Tests updated: forward/backward equivalence scripts pass
• Docs updated: README and Chinese README synced on shapes/semantics
• No known performance regressions; swapped-path correctness fix and stride‑0 broadcasting reduce overhead; memory impact limited to K-dim rounding for mask/bias only

[Copilot]

Pull Request Overview

This PR introduces broadcastable 4D mask/bias with strict shape contracts, 128-rounded key length for kernel access alignment, stride-0 broadcasting to avoid materialization, and corrected dbias reduction semantics. It resolves issue #189 by implementing robust attention mask/bias support with proper broadcasting and memory efficiency optimizations.

• Implements strict 4D shape contract for mask/bias with broadcast support over batch, heads, and query dimensions
• Rounds key length to 128 multiples for kernel access patterns, preventing out-of-bounds access and NaNs
• Uses stride-0 broadcasting in CUDA to avoid memory expansion along broadcast dimensions

Reviewed Changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
flash_dmattn/flash_dmattn_interface.py	Implements 128-rounding for mask/bias K-dim and backward gradient slicing
csrc/flash_dmattn/flash_api.cpp	Adds stride-0 broadcasting, 4D shape validation, and corrected swapped path handling
benchmarks/forward_performance.py	Updates test functions to use simplified mask preparation with proper broadcasting
benchmarks/forward_equivalence.py	Refactors test code to remove complex zoh_states calculations and use direct bias inputs
benchmarks/backward_performance.py	Simplifies backward test functions with direct bias inputs and proper mask handling
benchmarks/backward_equivalence.py	Updates backward equivalence tests to use simplified bias approach
README_zh.md	Updates Chinese documentation to reflect new broadcastable 4D shape requirements
README.md	Updates English documentation to reflect new broadcastable 4D shape requirements

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[Copilot]

The round_multiple lambda function is defined but there's already a global function with the same name and logic. Consider using the existing global function to avoid code duplication.

Suggested change

auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; };

[Copilot]

Pull Request Overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

csrc/flash_dmattn/flash_api.cpp:96

• The multiplication by seqlen_q for mask and bias batch strides is missing when seqlenq_ngroups_swapped is true, unlike the q_batch_stride which is correctly updated on line 94. This could lead to incorrect memory access patterns in the swapped path.

            params.mask_batch_stride *= seqlen_q;
            params.bias_batch_stride *= seqlen_q;

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[Copilot]

[nitpick] Commented out test cases should either be removed if no longer needed or have a comment explaining why they are disabled and when they might be re-enabled.

Suggested change

	# Not support head_dim > 128 in triton yet
	# The following test cases are disabled because Triton does not currently support head_dim > 128.
	# Once Triton adds support for head_dim > 128, these test cases should be re-enabled.