Adds bias gradient computation to backward kernel

[LoserCheems]

Implements bias gradient calculation in the backward pass by adding a new column-block kernel that computes DBias alongside existing DQ, DK, and DV gradients.

Updates initialization function to support multiple tensor names and modifies autotuning configurations to initialize both DQ and DBias tensors.

Includes extensive masking logic and memory access patterns to handle various sequence length and head dimension configurations while maintaining numerical stability.

[Copilot]

Pull Request Overview

Implements bias gradient computation in the Triton backward kernel, updates the tensor initialization helper to zero multiple tensors, and adjusts autotuning configurations to include the new DBias tensor.

• Introduces _bwd_kernel_one_col_block to compute DBias alongside existing gradients.
• Replaces the single-name init_to_zero helper with a version that accepts multiple tensor names.
• Updates autotuning pre_hook entries to zero both DQ and DBias.

Comments suppressed due to low confidence (2)

flash_dmattn/flash_dmattn_triton.py:339

• Consider adding a docstring above the _bwd_kernel_one_col_block Triton kernel to describe its purpose, parameters, and outputs, which will improve readability and onboarding for new contributors.

@triton.jit

flash_dmattn/flash_dmattn_triton.py:541

• The new bias gradient computation introduces both atomic and non-atomic code paths; add unit tests covering these scenarios (varying ATOMIC_ADD, block sizes, and mask configurations) to ensure correctness across configurations.

        dbias = (p * (dp - Di[:, None]))

[Copilot]

[nitpick] This kernel function spans 300+ lines and handles multiple responsibilities; consider breaking it into smaller helper functions (e.g., pointer setup, mask application, gradient computation) to improve readability and ease future maintenance.

[Copilot]

[nitpick] There are several commented-out Triton configuration lines that include DBias; once DBias support is verified, consider removing or updating these stale lines to reduce clutter.