Repository navigation
Conversation
dbd6ace to
7a6e0e6
Compare
![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 a fused RMSNorm + RoPE + head-transposition Pallas kernel for TPU, integrates it into the Flax attention pipeline, and implements several optimizations such as in-kernel VMEM transpose and Classifier-Free Guidance (CFG) before unpatchify. Feedback on these changes highlights critical runtime issues, including an AttributeError from using jax.shard_map instead of jax.experimental.shard_map.shard_map and a potential NameError due to an undefined _warn_once helper. Additionally, the reviewer recommends adding defensive checks to validate the sequence length of freqs_cis and removing redundant JAX compatibility shims that are already handled in the package initialization.
5b75348 to
82681ec
Compare
82681ec to
b33e35e
Compare
b33e35e to
7740dcf
Compare
7740dcf to
5c0f8bb
Compare
5c0f8bb to
aee8c4a
Compare
2f8b92d to
a507d0d
Compare
6ac9049 to
96dd50f
Compare
96dd50f to
359249b
Compare
593e337 to
4ebc6fd
Compare
4ebc6fd to
5a1fa32
Compare
a75047a to
0b1315d
Compare
0b1315d to
77afa86
Compare
77afa86 to
2e70bc8
Compare
2e70bc8 to
973594c
Compare
mbohlool
left a comment
mbohlool
left a comment
There was a problem hiding this comment.
Approving once the inline items are addressed; the Pallas producer matches the unfused reference op for op, and tiling, tails and VMEM look right.
| def resolve_from_config(config: Any, name: str) -> Any: | ||
| """Coerced value of `name`: `config` first, then the legacy env var, then the default.""" | ||
| env, dflt, _ = _OPTIONS[name] | ||
| keys = _extract_keys(config) | ||
| legacy = os.environ.get(env) | ||
| if name in keys and keys[name] is not None: | ||
| value = _coerce(name, keys[name]) | ||
| if legacy is not None and env not in _WARNED_ENV: | ||
| try: | ||
| legacy_coerced = _coerce(name, legacy) | ||
| except ValueError: | ||
| legacy_coerced = legacy | ||
| if legacy_coerced != value: | ||
| _WARNED_ENV.add(env) | ||
| max_logging.log( | ||
| f"[wan_runtime_options] Ignoring legacy env {env}={legacy!r}: config sets {name}={value!r}. " | ||
| f"Set {name} in the YAML or on the command line instead." | ||
| ) | ||
| return value | ||
| return _coerce(name, legacy) if legacy is not None else dflt |
There was a problem hiding this comment.
Could we drop the legacy WAN_* env tier and resolve config, then default, only? None of these variables exist on main or in #477/#478/#479, so there is nothing to stay compatible with, and today a stale shell variable silently changes the compiled graph whenever the config lacks a key. Please re-point the tests and remove the "used to be WAN_* environment variables" YAML comments and the WAN_ROPE_ACCUM mention in resolve_rope_accum's docstring.
| # Its output is equivalent but not bit-identical to the XLA | ||
| # producer: on v6e the bf16 trajectory diverges (PSNR ~15 dB | ||
| # against the XLA-path video, visually clean); on tpu7x the |
There was a problem hiding this comment.
Which stack and seed are the ~15 dB here and the 34.1 dB in the fused_rmsnorm_rope_pallas.py docstring (L48-49) each measured on? The v6e profile turns fused RoPE on for about 4% denoise time, so a VBench or CLIP score would help justify that default, or the kernel could default off on v6e.
| IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" | ||
| # Kernel numerics grids and multi-device tests are skipped in CI; see | ||
| # end_to_end/tpu/run_wan_stack_tests.sh. | ||
| _SKIP_IN_GITHUB_ACTIONS = unittest.skipIf( | ||
| IN_GITHUB_ACTIONS, "TPU kernel / multi-device test, skipped in GitHub Actions; run end_to_end/tpu/run_wan_stack_tests.sh" |
There was a problem hiding this comment.
_SKIP_IN_GITHUB_ACTIONS also skips the tiny interpret-mode parity, pair-swap and backward tests, so CI never checks the kernel numerics. Could we gate only the TPU-shape and multi-device classes (e.g. on _on_tpu() / jax.device_count()) and let the rest run, and fail loudly, in CPU CI?
| can_prescale = attn_kernel in supported_custom_kernels and not use_mem_eff and q.shape[1] >= min_seq_len | ||
| norm_mode_env = self.wan_rope_norm_mode | ||
| fuse_qk_prescale = bool(self.wan_fuse_qk_prescale and can_prescale) | ||
| q_prescale_val = float(LOG2E) if (fuse_qk_prescale and getattr(self.attention_op, "use_base2_exp", True)) else 1.0 | ||
| k_prescale_val = float(self.attention_op.scale) if fuse_qk_prescale else 1.0 |
There was a problem hiding this comment.
Could we add a test for the qk_prescaled producer to Ulysses hand-off? The existing mesh test uses dot_product below flash_min_seq_length, so can_prescale is False there. Mocking _ulysses_attention and _ulysses_ring_custom_attention, running FlaxWanAttention(attention_kernel="ulysses_custom", flash_min_seq_length=0, use_fused_rope_kernel=True) on the fake-TPU mesh and asserting qk_prescaled plus k == xla_k * bf16(scale) would catch a dropped not qk_prescaled guard, which today would silently double-scale.
|
|
||
| self.use_fused_rope_kernel = attention_config["use_fused_rope_kernel"] | ||
| self.fused_rope_block_s = attention_config["fused_rope_block_s"] | ||
| self.fused_rope_head_block = attention_config["fused_rope_head_block"] |
There was a problem hiding this comment.
Could we check here in __init__ that fused_rope_head_block divides heads and raise ValueError? Today a value that doesn't divide it (e.g. 16 with Wan-14B's 40 heads) only raises inside shard_map at first trace, after the weights are loaded.
| if k_prescaled or qk_prescaled: | ||
| key_states = key_states / jnp.asarray(scale, dtype=key_states.dtype) | ||
| if qk_prescaled and use_base2_exp: | ||
| query_states = query_states / jnp.asarray(LOG2E, dtype=query_states.dtype) | ||
| if not split_head_dim: | ||
| query_states = query_states.transpose(1, 0, 2) | ||
| key_states = key_states.transpose(1, 0, 2) | ||
| value_states = value_states.transpose(1, 0, 2) |
There was a problem hiding this comment.
These _apply_attention_dot edits also run on the SDXL/Flux dot_product path (the memory-efficient branch now handles split_head_dim=True differently). Could we add an SDXL-shaped parity test for this branch and call the change out in the PR description?
…, not env
Fused producer (kernels/fused_rmsnorm_rope_pallas.py)
- A Pallas kernel computing FP32 RMSNorm + RoPE + the [B,S,H*D] -> [B,H,S,D]
head transpose for self-attention Q/K. It optionally folds log2(e) into Q
and 1/sqrt(d) into K.
- norm_mode="exact" leaves the feature reduction to XLA. It is 0 ULP
against the separately jitted XLA producer (`fused_rmsnorm_rope`).
- norm_mode="fused" does the reduction in-kernel. It stays within 2x
pair-relative bf16 eps.
- rope_accum ("dtype" | "f32") matches how the local compiler rounds the
RoPE multiply-add. "auto" uses the per-platform measured mode: tpu7x
"dtype", v6e "f32".
- Off by default in the YAML (`use_fused_rope_kernel: False`) and in the
launcher's generic profile. Once inlined into the 40-layer graph, XLA rounds
the unfused producer differently, so the output is equivalent but not
identical (720p/81f, 40 steps: pre-decode bf16 latent PSNR is 53.8 dB after
1 step and 34.1 dB after 40 steps on v6e seed 12345; post-VAE decoded uint8
video PSNR / SSIM with peak=255 at stack tip is 25.71 dB / 0.9125 on v6e
seed 777, 17.92 dB / 0.6556 on v6e seed 12345, and 17.69 dB / 0.6918 on
tpu7x seed 12345).
- Fails closed (`ValueError`) when `use_fused_rope_kernel=True` is explicitly
requested on non-TPU meshes, unmeasured TPU generations, sharded feature
axes, `dim_head % 128 != 0`, or uneven sequence splits.
- `run_wan_fast_inference.sh` turns it ON in the v6e and v7 profiles. Measured
at the stack tip (Wan 2.2 T2V-A14B, 720p/81f/40 steps, warm AOT, CP=4, DP=2),
kernel on vs off:
v6e-8 (DVFS unpinned): denoise 124.5s - 125.0s (generate 126.6s - 127.5s) vs 131.2s off
tpu7x-8 (DVFS unpinned): denoise 104.9s (generate 107.8s) vs 106.6s (generate 109.6s) off
tpu7x-8 (DVFS pinned): denoise 91.7s (generate 94.6s) vs 93.4s (generate 96.4s) off
- Isolated speed, one v6e chip, jax 0.11.2, 40 heads, d=128, bf16, exact
mode, block_s=1024, median of 50 runs: per-shard seq 9450 takes 0.964 ms
vs 1.789 ms for XLA (1.86x); seq 18900 takes 1.953 ms vs 4.018 ms (2.06x).
- VMEM budget: 64 MiB is requested only when every TPU in the mesh is a v6e or
tpu7x (the validated generations).
Wan graph-changing switches
- `wan_rope_norm_mode`, `wan_fuse_qk_prescale`, `wan_splash_transpose_out`,
`wan_cfg_before_unpatchify`, `wan_cross_attn_prescale_kv` and `wan_rope_accum`
are pure YAML/config keys declared in every `base_wan*.yml` with the code
default and resolved at model-build time via `from_config` /
`attention_config_entries(config)` (no environment-variable fallback or
process-global store). Every value lives on the GraphDef and in the AOT key.
- `wan_cfg_before_unpatchify` (default True) applies classifier-free guidance
on packed tokens before unpatchify instead of after unpatchify.
- splash `transpose_out`: an in-kernel output layout [H,S,D] instead of
[H,D,S], for the custom splash kernel and the non-ring Ulysses wrapper. Off
by default.
Other changes
- `split_head_dim` is plumbed from the config into `WanModel` and its
attention layers (`False` in the six `base_wan*.yml` files).
- `_apply_attention_dot`: handles prescaled Q/K (`qk_prescaled` /
`k_prescaled` / `use_base2_exp`) and refuses prescaled Q/K for unsupported
kernels and head-local SVG attention.
Tests
- `fused_rmsnorm_rope_pallas_test.py`: parity grid, fail-closed guards,
backward, VMEM resolver, `FlaxWanAttention` sequence-sharded integration, and
`transpose_out` tests; runs in GitHub Actions on CPU and unvalidated TPUs via `interpret=not _on_validated_tpu()`.
- `dot_fallback_layout_test.py`, `wan_runtime_options_test.py`, `aot_cache_test.py`.
3 participants
Summary
Adds a fused RMSNorm + RoPE + head-transpose Pallas producer for Wan self-attention Q/K, and moves Wan's graph-changing switches out of environment variables into config.
What's in it
kernels/fused_rmsnorm_rope_pallas.py): computes FP32 RMSNorm + RoPE + the[B,S,H*D] → [B,H,S,D]transpose, with optional log2(e) / 1/√d prescale folding.norm_mode="exact"is 0 ULP against the separately jitted XLA producer (fused_rmsnorm_ropefromfeat/fused-rmsnorm-rope-xla).Off by default in the yml (
use_fused_rope_kernel: False) and in the launcher's generic profile: once inlined into the 40-layer graph, XLA rounds the unfused producer differently, so the 40-step v6e trajectory diverges (~15 dB PSNR against the XLA-path video, visually clean; bit-identical on tpu7x).Turned on in the launcher's v6e and v7 profiles. End-to-end latency (Wan 2.2 T2V-A14B, 720p / 81 frames / 40 steps,
CP=4, DP=2, warm AOT, DVFS unpinned):Set
USE_FUSED_ROPE_KERNEL=falseto turn it off.wan_rope_norm_mode,wan_fuse_qk_prescale,wan_splash_transpose_out,wan_cfg_before_unpatchify,wan_cross_attn_prescale_kvandwan_rope_accumare yml keys passed viaattention_config_entries(config)acrossWanPipeline, VACE, Animate, andwan_block_benchmark.wan_cfg_before_unpatchify(on by default): applies CFG before unpatchify. A compiled parity test shows bit-identical output in f32 everywhere and in bf16 on CPU/v6e, and within 1 bf16 ULP on tpu7x.Tests
fused_rmsnorm_rope_pallas_test.py(38 test methods / 60 parameterized cases),wan_runtime_options_test.py(11), plus updates todot_fallback_layout_test.py(12) andaot_cache_test.py(45).rope_accumskip on unvalidated TPUs (vmem_limit_is_validated,rope_accum_is_measured).GITHUB_ACTIONS=true(88 of 283 cumulative); guards, VMEM resolver, runtime options, AOT cache, and dot-fallback layout tests run in CI.run_wan_stack_tests.sh, 575.2 s).Stack: #477 → #507 → #478 → #508 → #479 → #488 → #509 → #510 → #511 → #491