[FEATURE SUPPORT] Flexible head dims for mask/bias with in-kernel conversion path

[LoserCheems]

Summary

#163

• Add support for mask and bias head dimensions: (B, H, Q, K), (B, H_k, Q, K), and (B, 1, Q, K).

Design

• Flexible head indexing using h_mask/h_bias to select:
  • H-aligned: head = h
  • H_k-aligned: head = h % num_kv_heads
  • Broadcast: head = 0

Changes

• Public/API:
  • Allow mask/bias shapes with head dim ∈ {1, num_kv_heads, num_heads}.
  • Params expose h_mask and h_bias to drive kernel-side head selection.
• Kernel:
  • Unified head-index selection across fwd/bwd/splitKV/varlen.
  • Mask loads use the conversion copy path during
• Python:
  • Interface accepts flexible head dims and applies broadcast semantics.
  • Padding path uses consistent, non-influential defaults for masked-out positions.
• Misc:
  • Resolved signature/argument ordering in set_params_fprop/set_params_dgrad.

Implementation Notes

• Conversion is fused with tiled copy; no extra passes.
• Maintains vectorization where safe; clean tail handling.
• SplitKV uses correct per-kv-head index for mask/bias.
• Guarded to avoid undue register/shared memory pressure.

Tests

• Shapes: head dims H, H_k, and 1; uniform and varlen sequences.
• Dtypes: fp16/bf16 for Q/K/V and outputs.
• Correctness: forward/backward equivalence and gradient checks.
• Performance: microbenchmarks (latency/TFLOPs) across typical configs.
• Integration: end-to-end for splitKV and varlen.

Docs

• Document supported mask/bias shapes and broadcast semantics.
• Note that conversion occurs in-kernel; no user changes required.
• Update examples to show head-dim flexibility.

Checklist

• Linked issue provided
• API stable
• Tests added or updated
• Docs added or updated
• No known performance regressions

[Copilot]

Pull Request Overview

This PR adds support for flexible head dimensions in attention masks and biases, allowing for shapes (B, H, Q, K), (B, H_k, Q, K), and (B, 1, Q, K) instead of being restricted to only (B, H_k, Q, K).

• Introduces kernel-side head indexing logic to handle different head dimension configurations
• Updates API and parameter structures to track mask and bias head counts separately
• Implements proper broadcast semantics and conversion paths for different head arrangements

Reviewed Changes

Copilot reviewed 6 out of 7 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
flash_dmattn/integrations/flash_dynamic_mask_attention.py	Updates documentation to reflect flexible mask/bias shape support
flash_dmattn/flash_dmattn_interface.py	Updates API documentation for flexible head dimensions
csrc/flash_dmattn/src/flash_fwd_kernel.h	Implements head indexing logic for mask/bias in forward kernel
csrc/flash_dmattn/src/flash_bwd_kernel.h	Implements head indexing logic for mask/bias in backward kernel
csrc/flash_dmattn/src/flash.h	Adds h_mask and h_bias fields to parameter structures
csrc/flash_dmattn/flash_api.cpp	Updates function signatures and implements flexible head dimension handling

---

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

[Copilot]

The comment format is inconsistent between mask and bias parameters. The bias comment has an extra comma after the first shape specification that should be removed to match the mask comment format.

Suggested change

	at::Tensor &bias, // batch_size x num_heads x seqlen_q x seqlen_k, or batch_size x num_heads_k x seqlen_q x seqlen_k or batch_size x 1 x seqlen_q x seqlen_k
	at::Tensor &bias, // batch_size x num_heads x seqlen_q x seqlen_k or batch_size x num_heads_k x seqlen_q x seqlen_k or batch_size x 1 x seqlen_q x seqlen_k

[Copilot]

The order of shape specifications in the bias comment differs from the mask comment above it. For consistency, the bias comment should list shapes in the same order: num_heads first, then num_heads_k, then 1.

Suggested change

	const at::Tensor &bias, // batch_size x num_heads_k x seqlen_q x seqlen_k or batch_size x num_heads x seqlen_q x seqlen_k or batch_size x 1 x seqlen_q x seqlen_k
	const at::Tensor &bias, // batch_size x num_heads x seqlen_q x seqlen_k or batch_size x num_heads_k x seqlen_q x seqlen_k or batch_size x 1 x seqlen_q x seqlen_k