benchmark: add block_sparse_attention (BSA fwd+bwd vs FA4) - #674
benchmark: add block_sparse_attention (BSA fwd+bwd vs FA4)#674vedaanta wants to merge 9 commits into
Conversation
Forward-only block-sparse attention benchmark on video-diffusion workload shapes (Wan2.1 1.3B/14B at 480P/720P), sweeping 64/128/256-token block granularities at top-k sparsity with a dense peak-reference bar. Arms: cudnn.block_sparse_attention_forward and an optional flash_attn.cute FA4 arm on identical masks, with fp32 parity checks. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Open-weights attention config (56 heads, head_dim 128) at its published sequence scales (~31k tokens per 5 s clip, ~91k for 15 s). H3 ships full attention with sparse support planned, so the dense bar reflects current cost and the sparse bars the headroom. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Both 24 heads x head_dim 128 per open-weights configs: HunyuanVideo 720p 129-frame (~118.8k tokens, the standard shape in the sparse-attention literature) and Mochi-1 480p 163-frame (~44.5k tokens). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This reverts commit 1b64264.
Structural mask for autoregressive video generation: each query attends its own frame, a window of previous frames (default 1, ~90% sparsity at these sequence lengths), and an anchor frame. Identical across heads, contiguous KV runs, variable per-row block counts (exercises the variable-count metadata contract), and expressed exactly at every granularity since the frame size is a multiple of all of them. TFLOPS now derive from the actual mask population, and index rows are count-prefixed to support variable counts. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
--direction fwd|bwd|both (default both). The backward times the public cudnn.block_sparse_attention_backward against flash_attn.cute's block-sparse backward, which consumes Q-direction (per-KV-block) index lists and requires 128-token KV sparse blocks; at 64-token granularity its whole fwd+bwd pipeline runs the KV-aggregated mask so gradients stay exact. Parity now also checks dq/dk/dv of both arms against an fp32 autograd reference with a deterministic shared dout. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
vedaanta
changed the title
The benchmark covers both inference and training (fwd+bwd), so the name follows the API it drives. results/ holds one full run of both mask families (plots + CSVs) on a clock-locked B200. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Updated results after the rename and restructure: one case per model, results organized as Backward bars live inside the topk/frame-causal charts: every configuration is an fwd/bwd pair of bar groups (second line of each x-label). Decode charts are fwd-only. All numbers on a B200 (SM100) clock-locked at 847 MHz — absolute TFLOPS scale with clocks; the relative comparisons are the point. TFLOPS count selected blocks only; both arms attend identical token sets. Wan2.1-14B 720P (40 heads x d128, S=92160)top-k 90% (VSA-style learned masks) — cuDNN leads fwd and bwd at 64/128 granularity (bwd by ~2–3x), ties at 256/dense:
frame-causal ~93% (same token set expressed exactly at every granularity):
decode shape (s_q=2048, KV history swept 22k -> 90k):
MiniMax-H3 15s (56 heads x d128, S=91392)
CSVs sit next to each plot. Parity: 43 checks (fwd + dq/dk/dv + decode, both arms x both families x all granularities) against fp32 references, all passing. |
--decode-q N runs the chunked autoregressive-generation shape (query = last N tokens over the full KV history, forward only); both mask families support it and decode parity is checked against the fp32 reference. Cases trim to one per model (wan14b-720p, minimax-h3-15s; other shapes remain reachable as custom HxS cases) and results move to results/<model>/<arch>/ matching the other benchmarks, with full B200 runs of both families plus decode for each model. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
--decode-kv-fracs (default 25/50/75/100% of the case seqlen) sweeps the KV history the decode chunk attends, modeling per-step cost as generated frames accumulate; chart width now scales with the group count. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Please target develop not main branch |
|
Please add dockerfile and review steps |
2 participants
What
New
benchmark/block_sparse_attentionmicrobenchmark: block-sparse attention forward and backward through the publiccudnn.block_sparse_attention_forward/_backwardAPIs, with an optional FA4-lineage CuTe DSL arm (flash_attn.cute) on identical block masks.topk(VSA-style learned, per-head data-dependent) andframe_causal(autoregressive video: frame window + anchor, variable per-row counts, ~90% sparse defaults)--checkvalidates fwd + dq/dk/dv of every arm/family/granularity against an fp32 autograd reference (40 checks)--csv/--plotoutputs following the repo chart paletteTesting
All 40 parity checks pass on B200 (SM100). Full grids (both families, fwd+bwd) run end-to-end; example numbers in the README-described format were collected on a clock-locked B200.
🤖 Generated with Claude Code