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.
9f73261 to
519699f
Compare
519699f to
36b073d
Compare
mbohlool
left a comment
mbohlool
left a comment
There was a problem hiding this comment.
LGTM after addressing comments.
| q_eps = getattr(self.norm_q, "epsilon", self.eps) | ||
| k_eps = getattr(self.norm_k, "epsilon", self.eps) |
There was a problem hiding this comment.
Could these read self.norm_q.epsilon / self.norm_k.epsilon directly? norm_q and norm_k are always nnx.RMSNorm on this branch, and nnx.RMSNorm always has epsilon, so the getattr default never fires and only hides a real breakage if the norm type changes. That also makes self.eps = eps unnecessary.
| with self.conditional_named_scope("attn_rope"): | ||
| query_proj = _unflatten_heads(query_proj, self.heads) | ||
| key_proj = _unflatten_heads(key_proj, self.heads) | ||
| if rotary_emb is not None and self.qk_norm and is_self_attention: |
There was a problem hiding this comment.
Could we add one module-level test for this branch? The new tests call fused_rmsnorm_rope against a copy of the RoPE math, so the wiring here (scale and epsilon reads, _unflatten_heads on V) is untested. A tiny FlaxWanAttention with qk_norm and random rotary_emb, compared against norm_q/norm_k + _unflatten_heads + self._apply_rope, would pin it.
| cos_q, sin_q = cos[:, :, :Sq, :].astype(raw_q.dtype), sin[:, :, :Sq, :].astype(raw_q.dtype) | ||
| cos_k, sin_k = cos[:, :, :Sk, :].astype(raw_k.dtype), sin[:, :, :Sk, :].astype(raw_k.dtype) |
There was a problem hiding this comment.
nit: since the only caller is self-attention, could this check freqs_cis.shape[2] == Sq == Sk and raise instead of slicing? _apply_rope fails loudly on a length mismatch, while the slice silently uses the first Sq positions of a longer table.
800418f to
f0232b9
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, accepts an optional `kv_heads` for GQA (defaults to `q_heads`), and validates `freqs_cis.shape[2] == Sq == Sk`. - 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, reading `self.norm_q.epsilon` and `self.norm_k.epsilon` directly; cross-attention and the no-RoPE path keep the separate norm_q / norm_k / `_apply_rope` steps. Tests - fused_producers_test.py (10, 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, GQA shapes/values, `freqs_cis` sequence-length mismatch guard, VJP parity against unfused grads, and end-to-end `FlaxWanAttention` module wiring + unfused forward parity. - Added to end_to_end/tpu/run_wan_stack_tests.sh (139 tests cumulative).
f0232b9 to
11bf69d
Compare
2 participants
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