[bingdongni]

Purpose

Keep masked FlashAttention consistent with the SDPA path when a tokenizer supplies an integer padding mask or a caller sets a non-default softmax scale.

Changes

• Interpret integer masks as 0/1 validity masks, while retaining bool and additive floating-point mask handling.
• Pass the configured scale to the equal-length masked self-attention path.

Test Plan

• pre-commit run --all-files --hook-stage manual (upstream CI)
• CUDA comparisons against PyTorch SDPA for bool, integer and additive masks with equal and unequal query/key lengths and 2D, 3D and 4D masks.

Test Results

On head 4731dfa, the GitHub Actions run passed all 18 CUDA adapter checks in the FA2 simulation on an RTX 4060; 12 cases failed against the original backend. The FA3 simulation passed nine equal-length cases and skipped nine unequal-length cases. These checks substitute PyTorch CUDA SDPA for the FlashAttention no-padding calls; they do not exercise a native FlashAttention kernel.

The upstream GPU unit-test, kernel-test, and documentation jobs passed on the current head. The remaining CI jobs are running.

Checklist

• pre-commit run --all-files --hook-stage manual passed in upstream CI on the current head
• I added or updated tests for my changes
• I updated documentation if needed (not applicable)
• I considered GPU memory impact of my changes

[mergify]

Merge Protections

🔴 1 of 1 protections blocking · waiting on 👀 reviews and 🤖 CI

	Protection	Waiting on
🔴	PR merge requirements	👀 reviews and 🤖 CI

🔴 PR merge requirements

Waiting for

• #approved-reviews-by>=1
• check-success=full-suite-passed

This rule is failing.

• #approved-reviews-by>=1
• check-success=full-suite-passed
• check-success=fastcheck-passed
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[greptile-apps]

Confidence Score: 4/5

[Medium impact] The change is limited to masked calls through the FlashAttention backend.

The code changes appear safe, but the added inline comment must be removed to meet the repository's rule.

Findings

1.  Inline comment violates repository rule ▶

Fix with agent prompt

### Issue 1
fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py:316
The comment on `scale` explains the test value, which falls outside the repository's allowed comment types. Remove it before merging to meet the requirement to remove comments outside that keep-list.

```suggestion
    scale = 0.37
```

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Summary

This PR fixes integer padding masks and preserves the configured attention scale.

• Masked FlashAttention honors integer padding masks and custom scales.

_{Reviews (3) · Last reviewed commit: "[test] Cover supported padding mask rank..." · Reviewed by Greptile}

[fv-superagent]

fv-superagent score: 3/5 — the source fix is correct and CI is green, but the new test errors on FA3.

Welcome, and thanks for the precise problem statement and the SDPA comparison — that made this quick to check.

• _key_padding_mask_from_attn_mask now matches the SDPA contract: int → != 0, bool unchanged, additive float unchanged. SDPAMetadataBuilder stores tokenizer-style int64 2D masks and the layer selector can route them here, so the integer half fixes a real cross-backend mismatch.
• softmax_scale=self.softmax_scale on the equal-length masked path is right: every other branch already passes it, and layer.py always resolves self.softmax_scale (default head_size**-0.5), so default runs are unchanged. The Buildkite unit lane confirms it — all six test_backend_padding_mask_and_custom_scale[bool|integer|additive-48|64] params PASSED.

Blocking (agrees with Greptile's P1 on this file):

• fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py:291 — the query_len=48 params take the flash_attn_varlen_qk_no_pad route, which every sibling test here guards with _fa2_only because FA3 rejects its dropout_p call. On an FA3/FA4 checkout this test will raise instead of skipping. Guard just the unequal-length case, e.g. if query_len != key_len: _fa2_only(no_pad_impls) after the seed — that keeps the equal-length int-mask/scale coverage running everywhere.

Optional, not blocking

• This test exercises the backend, not the custom ops the file is named for; a docstring line or a move beside test_sdpa_metadata_mask_contract.py would keep that contract honest.
• Only 2D masks are covered; the 3D/4D normalization branches are untouched by tests.
• Pre-existing, not from this PR: the float branch hard-masks on >= 0, unlike SDPA's additive pass-through for finite negative biases. No current producer emits those.

— fv-superagent · integer masks no longer attend to everyone · run by aryan5v

[fv-superagent]

fv-superagent score: 4/5 (was 3/5) — the FA3 guard is in and Greptile's follow-up is clean.

• 23e89e9 adds if query_len != key_len: _fa2_only(no_pad_impls), so the unequal-length params skip cleanly off FA2 while the equal-length integer-mask/scale coverage still runs everywhere. That closes the P1.
• Source fix unchanged and correct; the Buildkite unit lane re-ran all six test_backend_padding_mask_and_custom_scale params green on this head.
• Remaining points are optional, carried over: this test exercises the backend rather than the custom op the file is named for, and the 3D/4D mask-normalization branches still have no coverage.

— fv-superagent · one guard, three skips, zero TypeErrors · run by aryan5v

[greptile-apps]

Inline comment violates repository rule

The comment on scale explains the test value, which falls outside the repository's allowed comment types. Remove it before merging to meet the requirement to remove comments outside that keep-list.

Suggested change

	scale = 0.37 # Deliberately different from head_dim**-0.5.
	scale = 0.37

Rule Used: What: Remove all unwanted code comments in the scoped diff except for a small keep-list (license headers, dependency/protocol-required non-obvious behavior, // prettier-ignore, public API doc comments, and issue/RFC constraint links). Why: Comment... (source)

Prompt To Fix With AI

This is a comment left during a code review.
Path: fastvideo/tests/attention/test_flash_attn_no_pad_custom_op.py
Line: 316

Comment:
**Inline comment violates repository rule**

The comment on `scale` explains the test value, which falls outside the repository's allowed comment types. Remove it before merging to meet the requirement to remove comments outside that keep-list.

```suggestion
    scale = 0.37
```

**Rule Used:** What: Remove all unwanted code comments in the scoped diff except for a small keep-list (license headers, dependency/protocol-required non-obvious behavior, `// prettier-ignore`, public API doc comments, and issue/RFC constraint links).  Why: Comment... ([source](https://app.greptile.com/-/switch-tenant?tenant=hao-ai-lab&next=%2Fcustom-context%3Fmemory%3D856788fb-b8e1-47f8-b01b-d83ea205fc6b))

---

For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

[aryan5v]

@fv-superagent check again

[fv-superagent]

Checked again: #2080 is still green and mergeable at 4731dfa - pre-commit, the docs build and all six fastcheck lanes pass (build #2272), the new test's 18 params pass, and Greptile's latest review is 4/5. Since the earlier check the author added three test-only commits: the _fa2_only guard, 3D/4D mask-rank coverage and the P2 comment removal. Nothing new is wrong; the only thing still showing on the PR is my own first-pass CHANGES_REQUESTED, which those fixes superseded.