Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements the backward pass for custom splash attention and custom ring attention on TPU using Pallas, including the necessary VJP registrations, backward kernels (for dQ, dK, and dV), and comprehensive unit tests. The reviewer identified several critical correctness issues and race conditions across the backward kernels. Specifically, the kernels rely on specific execution orders of grid blocks (e.g., for initialization, accumulation, and writing results) despite the grid dimensions being configured with 'arbitrary' semantics, which allows the compiler to schedule them in any order. Refactoring these grid-level accumulations into sequential loops inside the kernels is required to prevent silent numerical corruption.
| if use_dq_aliasing: | ||
|
|
||
| @pl.when(j < 3) | ||
| def write_dq_first(): | ||
| dq_ref[...] = dq_val | ||
|
|
||
| @pl.when(j >= 3) | ||
| def write_dq_acc(): | ||
| dq_ref[...] = dq_alias_ref[...] + dq_val | ||
|
|
||
| else: | ||
| dq_ref[...] = dq_val |
There was a problem hiding this comment.
Critical Correctness Bug / Race Condition in DQ Aliasing:
When use_dq_aliasing is enabled, the kernel relies on the grid dimension j executing in strict sequential order (specifically, that j < 3 executes before j >= 3 for any given alias index j % 3).
However, because j is a grid dimension with "arbitrary" semantics, the compiler is free to schedule the blocks in any order.
If a block with j >= 3 (e.g., j = 3, which maps to alias index 0) is scheduled and executed before j = 0 (which also maps to alias index 0):
write_dq_accwill read fromdq_alias_ref[...], which at that point contains uninitialized garbage memory fromlax.empty.- Later, when
j = 0executes,write_dq_firstwill overwritedq_ref[...]withdq_val, completely wiping out the accumulated gradient fromj = 3.
This will result in silent numerical corruption and incorrect gradients. To fix this, you must avoid relying on grid-level scheduling order for accumulation and aliasing. All accumulation across KV blocks must be handled via sequential loops inside the kernel or via deterministic reduction patterns.
| @pl.when(j == 0) | ||
| def init(): | ||
| dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref) |
There was a problem hiding this comment.
Critical Correctness Issue / Race Condition:
The grid dimension j (KV block index) is defined with "arbitrary" semantics in dimension_semantics. In the Pallas programming model, "arbitrary" semantics means the compiler is free to schedule and execute the grid blocks in any order, or even concurrently.
Assuming that j == 0 executes first to initialize the scratchpad and j == grid_width - 1 executes last to write the results is unsafe. If the compiler schedules these blocks out of order, the scratchpad will be cleared mid-computation, or incomplete gradients will be written to dq_ref.
To fix this, you should perform the accumulation loop over KV blocks sequentially inside the kernel (e.g., using lax.fori_loop) rather than mapping it to a grid dimension.
| @pl.when(step_q == 0) | ||
| def init(): | ||
| dk_scratch_ref[...] = jnp.zeros_like(dk_scratch_ref) | ||
| dv_scratch_ref[...] = jnp.zeros_like(dv_scratch_ref) | ||
|
|
There was a problem hiding this comment.
Critical Correctness Issue / Race Condition:
The grid dimension step_q is defined with "arbitrary" semantics in dimension_semantics. There is no guarantee that step_q == 0 will execute first to initialize the scratchpads or that step_q == total_q_steps - 1 will execute last to write the results.
If the compiler schedules these blocks out of order, the scratchpads will be cleared mid-computation or incomplete gradients will be written to dk_ref and dv_ref.
Consider refactoring the accumulation over Q steps to be a sequential loop inside the kernel instead of a grid dimension.
| should_init_dkv = jnp.logical_and(i == 0, q_head_in_group == 0) | ||
| should_write_dkv = jnp.logical_and( | ||
| i == grid_height - 1, q_head_in_group == q_heads_per_kv_head - 1 | ||
| ) |
There was a problem hiding this comment.
Critical Correctness Issue / Race Condition:
In _flash_attention_bwd_fused_kernel, both i (Q block index) and h_q (Q head index) are grid dimensions with "arbitrary" semantics.
The conditions should_init_dkv and should_write_dkv assume a specific sequential execution order across the grid blocks (i.e., that i == 0 and q_head_in_group == 0 runs first, and i == grid_height - 1 and q_head_in_group == q_heads_per_kv_head - 1 runs last).
Since the compiler is free to schedule "arbitrary" grid dimensions in any order, this assumption is unsafe and will lead to silent data corruption or incorrect gradients if the blocks are executed out of order.
To resolve this, you should perform the accumulation over Q blocks and GQA heads sequentially inside the kernel using loops rather than relying on grid-level scheduling.
1 participant
Summary
Implements custom mask-free backward attention kernels (
_flash_attention_bwd_fused_kernel,_flash_attention_dq_kernel,_flash_attention_dkv_kernel) and custom ring backward attention (_custom_ring_attention_backward) with dynamic KV tail slicing and zero segment masks.Key Features
orig_kv_seq_len):SegmentIdsand mask overhead by dynamically slicing the last KV block (slice_k_len = kv_seq_len % bkv_compute) in both forward and backward passes.active_q_len/active_kv_lenand slices outputs back to exactorig_q_seq_len/orig_kv_seq_len, supporting arbitrary sequence lengths not divisible by block sizes.use_fused_bwd_kernel=True):grid_width), Q inner (num_q_heads, grid_height), computingdQ,dK, anddVin a single Pallas kernel pass.dKanddVin VMEM scratch across Q blocks and GQA heads (q_heads_per_kv_head) with zero transposes.dq_reduction_steps=3) to reducedQacross KV blocks with minimal HBM memory footprint._custom_ring_attention_backward):ring_axis), rotatingK, V, dK, dVacross devices vialax.ppermuteand accumulating localdQ, dK, dV.Performance Benchmarks (TPU v6e, 128 MB VMEM)
Workload: WAN 14B 76k ($\implies$
75,600total tokens,40heads,head_dim = 128, Context ParallelismCP = 418,900local tokens/shard).fwd/bwd)fwd=1024x1024, bwd=2048x2048, SegMask)20,480/20,480+1,580(+8.4%)150.29 ms212.87 ms363.16 ms1.00x(Baseline)fwd=1024x1024, bwd=2048x2048, NoMask)20,480/20,480+1,580(+8.4%)140.22 ms212.12 ms352.33 ms1.03xfwd=4096x1024, bwd=4096x1024)20,480/20,480+1,580(+8.4%)98.48 ms201.41 ms299.89 ms1.21xfwd=4736x1024, bwd=4736x1024, in=512)18,944/18,944+44(+0.23%)94.89 ms185.76 ms280.65 ms1.29xfwd=9472x1024, bwd=4736x2048, in=512)18,944/18,944+44(+0.23%)88.46 ms182.72 ms271.19 ms1.34xfwd=18944x1024, bwd=4736x2048, in=512)18,944/18,944+44(+0.23%)86.56 ms182.80 ms269.35 ms1.35x(-93.8 ms)86.56 msvs150.29 ms) usingblock_q = 18944(grid_height = 1, only 44 padded tokens), keeping the entire shard's Q resident in VMEM across the ring.182.80 msvs212.87 ms) using fused backward withblock_q_dkv = 4736, block_kv_dkv = 2048.Testing
Added
src/maxdiffusion/tests/custom_splash_backward_test.pycovering:dQ, dK, dV).orig_kv_seq_len = 1350) withvmapand GQA.dq_reduction_steps=3andNone) vs unfused backward (use_fused_bwd_kernel=False).sq=1350, skv=1777withblock_q=512, block_kv=512).1000.0garbage in the KV tail verifying exact gradient isolation.sq_local=650, skv_local=733).