feat(dflash): add flex_attention with BlockMask support

[xiaomin-D]

• Add attention_backend parameter to OnlineDFlashModel
• Implement BlockMask creation for flex_attention optimization
• Default to flex_attention backend (~30% speedup)
• Add iter_time display in training progress bar
• Clean up comments and simplify code

Motivation

Modifications

Related Issues

Accuracy Test

Benchmark & Profiling

eager

flex

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://sgl-fru7574.slack.com/archives/C09784E3EN6 to discuss your PR.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[Ximingwang-09]

Thank you for your PR. While training with long context(40K), I found that the flex attention backend causes OOM, whereas SDPA works fine. Upon investigation, the extra H = num_heads setting in create_block_mask seems to be the cause.
Is there any special reason why H must be set to num_heads in create_block_mask? In the implementation of eagle3, H is set to 1.
Related Issue：#452

[xiaomin-D]

Thank you for your PR. While training with long context(40K), I found that the flex attention backend causes OOM, whereas SDPA works fine. Upon investigation, the extra H = num_heads setting in create_block_mask seems to be the cause. Is there any special reason why H must be set to num_heads in create_block_mask? In the implementation of eagle3, H is set to 1. Related Issue：#452

I think you’re right — this was an oversight on my part.