Repository navigation
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 introduces an optimized fused_rmsnorm_rope producer for Wan Attention to perform FP32 RMSNorm, cast to input dtype, and apply RoPE rotation and head transposition, along with corresponding unit tests. Feedback on the changes highlights a critical issue where accessing self.norm_q.scale[...] directly will raise a TypeError since Flax NNX Param objects are not subscriptable, recommending the use of .value instead. Additionally, it is suggested to cast the sliced rotary embedding tensors (cos and sin) to their respective query and key dtypes rather than globally casting to raw_q.dtype to avoid potential precision loss in mixed-precision configurations.
This was referenced Oct 15, 2026
Perseus14
added this pull request to stack #512
October 15, 2026 16:00
Perseus14
force-pushed
the
feat/fused-rmsnorm-rope-xla
branch
from
October 15, 2026 17:53
9f73261 to
519699f
Compare
Adds `fused_rmsnorm_rope` (kernels/fused_producers.py) and uses it in `FlaxWanAttention.__call__` for the self-attention Q/K producer chain. What - `fused_rmsnorm_rope(raw_q, raw_k, q_scale, k_scale, freqs_cis, ...)` does FP32 RMSNorm, casts back to the input dtype, then applies RoPE and the head transposition in input dtype, so no FP32 RoPE intermediates are materialised at long sequence lengths. It mirrors nnx.RMSNorm's association (`x * (rsqrt(var + eps) * scale)`): bit-for-bit in f32, and in bf16 when the norm scale is bf16. With bf16 activations and an f32 scale, nnx.RMSNorm's `promote_dtype` rounds the scale to bf16 before folding it into the rsqrt, while the fused producer keeps the scale in f32 and rounds only the final product; the two then differ by at most one bf16 ulp. It takes separate q/k epsilons and accepts an optional `kv_heads` for GQA (defaults to `q_heads`). - It is `jax.jit`-wrapped with heads/dim/eps static: under the pipeline's outer jit it is simply inlined, but when a full-size block is called eagerly (as wan_vace_transformer_test does at 75,600 tokens) XLA fuses the chain instead of materialising ~1.5 GB of extra intermediates, which OOMed that test on the v4-8 CI runner. - `FlaxWanAttention` takes the fused path when `rotary_emb` is given, `qk_norm` is on and the call is self-attention; cross-attention and the no-RoPE path keep the separate norm_q / norm_k / `_apply_rope` steps. `self.eps` is kept so the fused call can fall back to the layer epsilon when a norm module does not expose `epsilon`. Tests - fused_producers_test.py (8, new, CPU): bit-for-bit match with nnx.RMSNorm in f32 and in bf16 (bf16 scale), the bf16-activation / f32-scale case pinned to <= 1 bf16 ulp of nnx.RMSNorm, the reassociated form is shown to differ, fused Q matches flax RMSNorm, separate q/k epsilons are applied, non-trivial RoPE in bfloat16 matches the unfused reference, and GQA shapes/values. - 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 32, ring_fixed_m 24, dot_fallback_layout 6, tile_size_grid_search 16, attention_config_guards 20, custom_splash_unpadded 12, fused_producers 8 -- 118 tests, all OK.
Perseus14
force-pushed
the
feat/fused-rmsnorm-rope-xla
branch
from
October 15, 2026 19:20
519699f to
36b073d
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 #478 (
feat/ring-attention). Addsfused_rmsnorm_rope(kernels/fused_producers.py), an XLA-level fused FP32-RMSNorm + RoPE + head-transpose producer for the self-attention Q/K chain, and uses it inFlaxWanAttention.__call__whenrotary_embis given,qk_normis on, and the call is self-attention (cross-attention and the no-RoPE path keep the separatenorm_q/norm_k/_apply_ropesteps). It mirrorsnnx.RMSNorm's association (x * (rsqrt(var + eps) * scale)) — bit-for-bit infloat32and forbfloat16activations with abfloat16scale, and within 1bfloat16ULP whenraw_qisbfloat16andscaleisfloat32(wherennx.RMSNorm.promote_dtyperounds the scale tobfloat16first while the fused producer keeps it infloat32) — takes separate q/k epsilons, and acceptskv_headsfor GQA. It isjax.jit-wrapped (heads/dim/eps static): inlined under the pipeline's outer jit, but when a full-size block is called eagerly (aswan_vace_transformer_testdoes at 75,600 tokens) XLA fuses the chain instead of materialising ~1.5 GB of intermediates, which OOMed that test on the v4-8 CI runner. #488's Pallas_fused_rope_producerbuilds on this XLA producer.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.2 s generate / 136.0 s denoise, 67.1 s compile (neutral step time under full-block XLA fusion vs #478's 138.3 s / 136.1 s; eliminates the ~1.5 GB eager intermediate allocation).Tests:
fused_producers_test.py(8, new, CPU), added torun_wan_stack_tests.sh.Stack: #477 → #507 → #478 → #508 → #479 → #488 → #509 → #510 → #511 → #491