[Bug] flex_attention BlockMask incorrectly created with H=num_heads

Checklist

• 1. I have searched related issues but cannot get the expected help.
• 2. The bug has not been fixed in the latest version.
• 3. Please note that if the bug-related issue you submitted lacks corresponding environment info and a minimal reproducible demo, it will be challenging for us to reproduce and resolve the issue, reducing the likelihood of receiving feedback.
• 4. If the issue you raised is not a bug but a question, please raise a discussion at https://github.com/sgl-project/SpecForge/discussions/new/choose Otherwise, it will be closed.
• 5. Please use English, otherwise it will be closed.

Describe the bug

Description:

In the current implementation of DFlash training with flex_attention backend, the BlockMask is created with H=num_heads (line 87), which is unnecessary and potentially problematic. When we train the dflash model with long context, it will cause OOM:

[rank7]: Traceback (most recent call last):
[rank7]: File "/ossfs/workspace/SpecForge/scripts/train_dflash.py", line 587, in
[rank7]: main()
[rank7]: File "/ossfs/workspace/SpecForge/scripts/train_dflash.py", line 531, in main
[rank7]: loss, accuracy = dflash_model(
[rank7]: ^^^^^^^^^^^^^
[rank7]: File "/usr/local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank7]: return self._call_impl(*args, **kwargs)
[rank7]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]: File "/usr/local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank7]: return forward_call(*args, **kwargs)
[rank7]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]: File "/usr/local/lib/python3.11/site-packages/torch/distributed/fsdp/fully_sharded_data_parallel.py", line 851, in forward
[rank7]: output = self._fsdp_wrapped_module(*args, **kwargs)
[rank7]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]: File "/usr/local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
[rank7]: return self._call_impl(*args, **kwargs)
[rank7]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]: File "/usr/local/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1786, in _call_impl
[rank7]: return forward_call(*args, **kwargs)
[rank7]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]: File "/ossfs/workspace/SpecForge/specforge/core/dflash.py", line 158, in forward
[rank7]: dflash_attn_mask = self._get_or_create_block_mask(
[rank7]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]: File "/ossfs/workspace/SpecForge/specforge/core/dflash.py", line 85, in _get_or_create_block_mask
[rank7]: block_mask = create_block_mask(
[rank7]: ^^^^^^^^^^^^^^^^^^
[rank7]: File "/usr/local/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 1092, in create_block_mask
[rank7]: partial_block_mask, full_block_mask = _convert_mask_to_block_mask(
[rank7]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank7]: File "/usr/local/lib/python3.11/site-packages/torch/nn/attention/flex_attention.py", line 892, in _convert_mask_to_block_mask
[rank7]: mask_block_sum = mask.sum(
[rank7]: ^^^^^^^^^
[rank7]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 256.00 GiB. GPU 7 has a total capacity of 139.81 GiB of which 98.64 GiB is free. Including non-PyTorch memory, this process has 41.17 GiB memory in use. Of the allocated memory 39.96 GiB is allocated by PyTorch, and 43.03 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting.

Problem: The DFlash attention mask pattern (dflash_mask_fn) does not vary across attention heads - all heads use the same visibility pattern (context blocks before current block + noise tokens from same block). Creating separate masks for each head:

1. Wastes memory by duplicating identical mask data num_heads times
2. Adds unnecessary computation overhead during mask creation
3. Requires caching num_heads as part of the cache key, complicating the caching logic.

Expected Behavior: Since the mask pattern is head-independent, BlockMask should be created with H=1, allowing the mask to be broadcast across all attention heads automatically.

Reproduction

bash run_qwen3_8b_dflash_offline.sh

Environment

.

Start in specforge/core/dflash.py at _get_or_create_block_mask around line 85, then trace its call from forward. Reproduce with run_qwen3_8b_dflash_offline.sh and inspect the flex_attention BlockMask creation. Done when the head-independent mask no longer allocates redundant per-head data and the long-context run avoids the reported CUDA out-of-memory failure.