Skip to content

feat(fastwam): add structured FlashAttention backend - #210

Draft
XueSongTap wants to merge 5 commits into
masterfrom
feat/fastwam-structured-fa2
Draft

XueSongTap wants to merge 5 commits into
masterfrom
feat/fastwam-structured-fa2

Conversation

@XueSongTap

Copy link
Copy Markdown
Collaborator

Summary

  • Add structured attention execution for FastWAM MoT masks.
  • Support FA4, FA3, FA2, SDPA, and automatic backend fallback.
  • Integrate the backend across base MoT, Joint, and IDM paths.
  • Preserve dense-mask semantics while avoiding dense [B, H, Q, K] masks.

Validation

  • Added forward/backward equivalence tests for structured vs dense attention.
  • Added mask-semantic tests for Base, Joint, and IDM variants.
  • compileall and git diff --check pass.
  • BF16 CUDA equivalence probe: forward max error 0.0, gradient max error 0.015625.

Performance

Microbenchmark (video tokens=588, action tokens=32, heads=24, head_dim=128, mbs16):

  • Forward: SDPA 0.991 ms -> FA2 0.577 ms (1.72x).
  • Forward + backward: SDPA 6.027 ms -> FA2 3.993 ms (1.51x).

End-to-end 8-GPU training loss remains aligned with the original FastWAM baseline; the current workload is communication-bound, so overall throughput is approximately unchanged.

@github-actions github-actions Bot added the bug Something isn't working label Sep 3, 2026
@XueSongTap
XueSongTap marked this pull request as draft September 3, 2026 03:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant