Repository navigation
Conversation
This was referenced Oct 15, 2026
![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 addresses a silent data corruption bug in the short-sequence dot-product fallback path (_apply_attention_dot) by correctly transposing 4-D inputs instead of reshaping them, and adds support for Grouped-Query Attention (GQA) via kv_heads. It also fixes an issue in the tile-size grid-search utility when bkv_compute is None and adds corresponding unit tests. The review feedback identifies a critical issue where kv_heads is missing from the signature and context of _apply_attention, making the GQA fallback path unreachable in practice. Additionally, it is recommended to add defensive checks to ensure heads and effective_kv_heads are positive integers to prevent potential division-by-zero errors.
Perseus14
added this pull request to stack #512
October 15, 2026 16:00
…_grid_search iters/bkv_compute guards
Two small fixes peeled out of the Ulysses x Ring PR so that they can be
reviewed on their own.
Dot-product fallback (`_apply_attention_dot`)
- Sequences below `flash_min_seq_length` bypass the flash/ulysses kernels and
run the dot-product fallback. Callers that apply rotary embeddings hand the
dispatcher `[B, H, S, D]` (the dispatcher reads S from axis 2 when
`ndim == 4`), but the `split_head_dim` path reshaped those tensors as if
they were `[B, S, H*D]`. Both layouts hold the same number of elements, so
the reshape succeeded and returned a wrong answer silently. 4-D inputs are
now transposed to `[B, S, H, D]` (`_to_bshd`); the flat 3-D contract is
unchanged.
- `kv_heads` (optional, default None) is accepted so that K/V with fewer heads
than Q are repeated on the head axis after the transpose, on both the
`split_head_dim` and batched-heads paths; `heads % kv_heads != 0` raises
ValueError instead of silently truncating the repeat factor.
`dot_product_kernel` forwards `context.get("kv_heads")`.
tile_size_grid_search
- `time_callable` rejects `iters < 1` instead of dividing by zero.
- `_broadcast_winner` encodes `bkv_compute=None` as -1 in the broadcast
payload and decodes it back, so a winner without an explicit compute block
is found again on the non-source hosts.
Tests
- dot_fallback_layout_test.py (6, new, CPU): the reviewer's counterexample
(uniform softmax must give per-head token means), random 4-D inputs vs a
dense reference and GQA repeat on 4-D (both over split_head_dim in
{True, False}), 3-D/4-D agreement, the heads % kv_heads guard, and the
dispatcher seq-axis contract: `_apply_attention` on a [B, H=128, S=32, D]
query with flash_min_seq_length=64 must route to a spied dot_product kernel
(S read from axis 2) that then returns the dense reference.
- tile_size_grid_search_test.py (16, +2): `iters=0` raises; `_broadcast_winner`
round-trips `bkv_compute=None`.
- Both files are added to end_to_end/tpu/run_wan_stack_tests.sh.
Verified (this head): TPU v6e-8 (jax 0.11.2), every file of
end_to_end/tpu/run_wan_stack_tests.sh at this head run directly with unittest:
custom_splash_fixed_m 31, ring_fixed_m 18, dot_fallback_layout 6,
tile_size_grid_search 16 -- 71 tests, all OK.
Perseus14
force-pushed
the
feat/attn-dot-fallback-fixes
branch
from
October 15, 2026 19:20
167b48c to
ee90fa1
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #477 (
feat/fixed-m-kernel). Two small fixes peeled out of #478 so they can be reviewed on their own:_apply_attention_dot: sequences belowflash_min_seq_lengthtake the dot-product fallback, and callers that apply RoPE hand it[B, H, S, D], but thesplit_head_dimpath reshaped that as if it were[B, S, H*D]— same element count, so it silently returned wrong values. 4-D inputs are now transposed (_to_bshd); the flat 3-D contract is unchanged, an optionalkv_headsrepeats K/V on the head axis after the transpose (for bothsplit_head_dim=TrueandFalse), andheads % kv_heads != 0raisesValueError.tile_size_grid_search:time_callablerejectsiters < 1, and_broadcast_winnerround-tripsbkv_compute=Nonethrough the multihost broadcast.Tests:
dot_fallback_layout_test.py(6, new, CPU — includes the counterexample,split_head_dimin{True, False}, theheads % kv_headsguard, and a spied_apply_attentiondispatcher test) andtile_size_grid_search_test.py(+2); both added toend_to_end/tpu/run_wan_stack_tests.sh.E2E latency (TPU v6e-8, Wan 2.2 T2V-A14B, 720p / 81 frames / 40 steps,
CP=4, DP=2, 1D UlyssesU=4, DVFS unpinned): 138.5 s generate / 136.4 s denoise (neutral vs #477's 139.3 s / 136.6 s; correctness & guard PR).Stack: #477 → #507 → #478 → #508 → #479 → #488 → #509 → #510 → #511 → #491