From 8dd630549ce1cb55e167a66102dcc3aef4cc1926 Mon Sep 17 00:00:00 2001 From: James Huang Date: Mon, 27 Jul 2026 08:18:12 +0000 Subject: [PATCH] Fixed-m softmax for 2D-ring x Ulysses attention (WAN 2.2) Removes the online-softmax running-max and rescale chain from the ring attention kernel, replacing it with a per-query Cauchy-Schwarz bound that is known before streaming starts. Heads whose bound cannot be proven safe keep online softmax, so the result is exact either way. 720p 81f T2V, 40 steps, v7x-8, dp2 x (ring2 x uly2), denoise total: attention=ulysses_ring_custom 9472/1024 109.3s 2.733 s/step attention=ulysses_ring_custom_fixed_m 6400/2048 105.3s 2.633 s/step The tile differs per mode because the optimum moves. Plain is MXU-bound and peaks at bq=9472, degrading to ~3.14 s/step at 6400/2048; fixed-m is flat across both (2.63-2.69), which is the mechanism working -- with the softmax VPU chain gone the kernel stops caring which tile it was given. In-kernel the ring hop goes 22.65 -> 21.15 ms, i.e. 85% -> 91% of the K=128 MXU roofline. Ring specifics, over the flat fixed-m design: - No k-smoothing is possible on a ring (no rank holds the full K, and RoPE does not preserve a zero mean), so the gate uses the two-sided halved bound 106.5 rather than the flat 213. - The bound is global over every K shard, so a head's pinned m is identical on every hop and the partials combine by accumulation. When any head fails the gate the call falls back to an LSE-space merge, which is invariant to the bound's overshoot; a naive (m, l) merge would flush the other partials. - The all-heads-fixed branch compiles a branch-free kernel. Per-head dispatch in the ragged last KV block degrades the whole grid's instruction schedule (measured 3x), so exactly one body is emitted there. Also in this change: - Warmup now runs one real step per weight set. A 2-step warmup puts both steps on the high-noise transformer, leaving the low-noise weights untouched until mid-generation, where their first call cost 26.4s against a 2.6s step. The first generation now reports steady-state latency instead of needing a throwaway run (denoise 118.5s -> 105.3s on a cold process). - _ulysses_attention no longer folds the CFG batch into heads when data/fsdp already shard it; the folded size-1 batch is not divisible by those axes. - Cross-attention demotes to ulysses_custom for the fixed-m ring variant too, matching the other two ring variants. - run_wan_fast_inference.sh takes FIXEDM=0/1 and picks the matching tile. Tests: src/maxdiffusion/tests/ring_fixed_m_test.py covers the online path, an all-eligible call, a sink head that falls back everywhere, a ragged KV tail, and a caller-supplied local max||k|| being promoted to the global bound. --- end_to_end/tpu/run_wan_fast_inference.sh | 21 +- src/maxdiffusion/aot_cache.py | 20 ++ .../kernels/custom_splash_attention.py | 115 +++++++-- .../splash_attention/ring_attention_kernel.py | 228 +++++++++++++++--- src/maxdiffusion/models/attention_flax.py | 189 ++++++++++++--- .../pipelines/wan/wan_pipeline_2_2.py | 87 ++++++- src/maxdiffusion/tests/ring_fixed_m_test.py | 193 +++++++++++++++ 7 files changed, 767 insertions(+), 86 deletions(-) create mode 100644 src/maxdiffusion/tests/ring_fixed_m_test.py diff --git a/end_to_end/tpu/run_wan_fast_inference.sh b/end_to_end/tpu/run_wan_fast_inference.sh index 39e63705f..89680760b 100755 --- a/end_to_end/tpu/run_wan_fast_inference.sh +++ b/end_to_end/tpu/run_wan_fast_inference.sh @@ -26,6 +26,11 @@ # OUTPUT_DIR video/metrics output (default /tmp/wan_out) # COMPILE_TE=true torch.compile the text encoder (adds ~30s to load, # saves ~10s/encode; worth it for long-lived processes) +# FIXEDM=0 plain online softmax instead of fixed-m +# EXTRA_LIBTPU extra libtpu flags, appended to the tuned set +# +# 720p 81f / 40 steps denoise: fixed-m 105.3s, plain 109.6s. Each mode gets its +# own optimal tile below; plain degrades badly on fixed-m's. set -u MODEL=${1:-22} STEPS=${2:-40} @@ -45,6 +50,17 @@ mkdir -p "$CACHE_ROOT/jax" "$CACHE_ROOT/aot_wan$MODEL" "$CACHE_ROOT/converted" " # Tuned collective/scheduler flag set for v7 (from the PR #430 2D-ring # baseline). One line: libtpu stops parsing at a literal backslash. export LIBTPU_INIT_ARGS="--xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_dot_strength_reduction=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_enable_async_collective_permute=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true --xla_tpu_scoped_vmem_limit_kib=65536 --xla_tpu_enable_async_all_to_all=true --xla_tpu_enable_all_experimental_scheduler_features=true --xla_tpu_enable_scheduler_memory_pressure_tracking=true --xla_tpu_host_transfer_overlap_limit=24 --xla_tpu_aggressive_opt_barrier_removal=ENABLED --xla_lhs_prioritize_async_depth_over_stall=ENABLED --xla_should_allow_loop_variant_parameter_in_chain=ENABLED --xla_should_add_loop_invariant_op_in_chain=ENABLED --xla_tpu_enable_ici_ag_pipelining=true --xla_max_concurrent_host_send_recv=100 --xla_tpu_scheduler_percent_shared_memory_limit=100 --xla_latency_hiding_scheduler_rerun=2 --xla_tpu_use_minor_sharding_for_major_trivial_input=true --xla_tpu_relayout_group_size_threshold_for_reduce_scatter=1 --xla_tpu_enable_latency_hiding_scheduler=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_megacore_fusion=true --xla_tpu_megacore_fusion_allow_ags=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true --xla_tpu_sparse_core_all_gather_latency_multiplier=1 --xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 --xla_tpu_enable_sparse_core_collective_aggregator=true --xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true --xla_tpu_enable_sparse_core_reduce_scatter_v2=true --xla_tpu_enable_sparse_core_collective_offload_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true --xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true --xla_tpu_enable_concurrent_sparse_core_offloading=true --xla_tpu_assign_all_reduce_scatter_layout=true" +# Timings only compare across runs passing the same extra flags. +export LIBTPU_INIT_ARGS="${LIBTPU_INIT_ARGS} ${EXTRA_LIBTPU:-}" + +# fixed-m on by default: faster, and covered by tests/ring_fixed_m_test.py. +if [ "${FIXEDM:-1}" = "1" ]; then + ATTENTION=ulysses_ring_custom_fixed_m + BQ=6400; BKV=2048 +else + ATTENTION=ulysses_ring_custom + BQ=9472; BKV=1024 +fi if [ "$MODEL" = "21" ]; then CONFIG=src/maxdiffusion/configs/base_wan_14b.yml @@ -57,6 +73,7 @@ fi PROMPT_ARG=() [ -n "$PROMPT" ] && PROMPT_ARG=("prompt=$PROMPT") RUN_NAME="wan${MODEL}_fast_$(date +%m%d-%H%M%S)" +echo "== ${ATTENTION} | tile ${BQ}/${BKV} | ${STEPS} steps" # libtpu's XLA:CPU AOT feature-mismatch log is cosmetic and ignores every # log-level env var; filter just that message from stderr. @@ -66,7 +83,7 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \ jax_cache_dir="$CACHE_ROOT/jax" \ aot_cache_dir="$CACHE_ROOT/aot_wan$MODEL" \ converted_weights_dir="$CACHE_ROOT/converted" \ - attention=ulysses_ring_custom \ + attention=$ATTENTION \ ulysses_shards=2 \ ici_data_parallelism=2 ici_fsdp_parallelism=1 \ ici_context_parallelism=4 ici_tensor_parallelism=1 \ @@ -78,7 +95,7 @@ python src/maxdiffusion/generate_wan.py "$CONFIG" \ text_encoder_dtype=bfloat16 compile_text_encoder="${COMPILE_TE:-false}" use_batched_text_encoder=false \ use_base2_exp=true use_experimental_scheduler=true \ fps=16 $GUIDANCE_ARGS \ - flash_block_sizes='{"block_q":9472,"block_kv":1024,"block_kv_compute":1024,"block_kv_compute_in":1024,"heads_per_tile":1,"vmem_limit_bytes":67108864,"block_q_dkv":9472,"block_kv_dkv":1024,"block_kv_dkv_compute":1024}' \ + flash_block_sizes="{\"block_q\":$BQ,\"block_kv\":$BKV,\"block_kv_compute\":$BKV,\"block_kv_compute_in\":1024,\"heads_per_tile\":1,\"vmem_limit_bytes\":67108864,\"block_q_dkv\":$BQ,\"block_kv_dkv\":$BKV,\"block_kv_dkv_compute\":$BKV}" \ "${PROMPT_ARG[@]}" \ 2> >(grep -vE --line-buffered 'cpu_aot_loader|machine type for execution' >&2) diff --git a/src/maxdiffusion/aot_cache.py b/src/maxdiffusion/aot_cache.py index ba722160a..1a905b9f3 100644 --- a/src/maxdiffusion/aot_cache.py +++ b/src/maxdiffusion/aot_cache.py @@ -425,3 +425,23 @@ def warmup_mode(): yield finally: _STATE.warmup_only = False + + +def in_warmup() -> bool: + """True inside `warmup_mode()`, so callers can add warmup-only work.""" + return _STATE.warmup_only + + +@contextlib.contextmanager +def real_execution(): + """Temporarily suspends zero-execution warmup so a call really runs. + + Warmup skips execution, leaving first-execution costs to the first real + request. Wrap a targeted priming call in this to pay them during warmup. + """ + previous = _STATE.warmup_only + _STATE.warmup_only = False + try: + yield + finally: + _STATE.warmup_only = previous diff --git a/src/maxdiffusion/kernels/custom_splash_attention.py b/src/maxdiffusion/kernels/custom_splash_attention.py index 3f4ef71be..6bd3f3493 100644 --- a/src/maxdiffusion/kernels/custom_splash_attention.py +++ b/src/maxdiffusion/kernels/custom_splash_attention.py @@ -34,7 +34,13 @@ class _BlockSizes: __slots__ = ("block_q", "block_kv", "block_kv_compute", "block_kv_compute_in") - def __init__(self, block_q: int, block_kv: int, block_kv_compute: int | None = None, block_kv_compute_in: int = 256): + def __init__( + self, + block_q: int, + block_kv: int, + block_kv_compute: int | None = None, + block_kv_compute_in: int = 256, + ): self.block_q = block_q self.block_kv = block_kv self.block_kv_compute = block_kv_compute if block_kv_compute is not None else block_kv @@ -52,6 +58,12 @@ def __init__(self, block_q: int, block_kv: int, block_kv_compute: int | None = N # exceeds the gate fall back to online softmax (the "sink" heads). _FIXED_M_RECENTER = 88.0 _FIXED_M_SAFE_BOUND = 213.0 +# Ring-path gate: the ring processes UN-smoothed K shards (no ring rank holds +# the full K to compute a mean, and a per-shard mean would shift each hop's +# logits differently, breaking the cross-shard merge). Without k-smoothing the +# per-row max logit has no >=0 guarantee, so the safe bound halves (calibrated +# for ring_size=2, matching DiffusionServing's ring gate). +_FIXED_M_RING_SAFE_BOUND = _FIXED_M_SAFE_BOUND / 2.0 def _flash_attention_kernel( @@ -76,6 +88,7 @@ def _flash_attention_kernel( use_base2_exp: bool = True, fuse_reciprocal: bool = True, use_fixed_m: bool = False, + uniform_fixed_m: bool = False, ): float32 = jnp.float32 head_dim_v_repeats, rem = divmod(head_dim_v, NUM_SUBLANES) @@ -86,24 +99,45 @@ def _flash_attention_kernel( exp = jnp.exp2 if use_base2_exp else jnp.exp sv_dims = (((0,), (0,)), ((), ())) + # `uniform_fixed_m` is the caller's compile-time promise that EVERY head + # passed the gate (the ring's accumulate merge runs under exactly that + # predicate). It buys two things at once: the per-head dispatch collapses to + # a single body per block, and the fixed bound stays PINNED through the + # ragged last KV block, so every hop reports the identical m and the hops + # combine by plain accumulation. + # + # Both must move together. Pinning without the uniform promise needs a + # SECOND body in the last block (fixed and online), and a two-body last + # block degrades the instruction schedule of the WHOLE grid -- measured 3x + # slower end to end, which is the cliff the design doc's D3 warns about. + # Keeping this one flag rather than two makes that combination unspellable. + fixed_only = use_fixed_m and uniform_fixed_m + if uniform_fixed_m and not use_fixed_m: + raise ValueError("uniform_fixed_m requires use_fixed_m.") + # Per-head dispatch: heads inside the no-flush window run fixed-m, the rest # keep online softmax. Branch once per head (body level), never per step. - is_fixed = (mk_ref[1, h] > 0.5) if use_fixed_m else False + is_fixed = (mk_ref[1, h] > 0.5) if (use_fixed_m and not fixed_only) else False + + def _write_fixed_m(): + # Per-query Cauchy-Schwarz bound m_i = ceil(||q_i|| * max_j||k_j||) - C. + qf = q_ref[...].astype(float32) + qn = jnp.sqrt((qf * qf).sum(axis=1))[None, :] # (1, bq) per-query norm + bound = qn * mk_ref[0, h] + m_fixed = jnp.ceil(bound) - _FIXED_M_RECENTER + m_scratch_ref[...] = jnp.broadcast_to(m_fixed, m_scratch_ref.shape) @pl.when(j == 0) def init(): o_scratch_ref[...] = jnp.zeros_like(o_scratch_ref) l_scratch_ref[...] = jnp.zeros_like(l_scratch_ref) - if use_fixed_m: + if fixed_only: + _write_fixed_m() + elif use_fixed_m: @pl.when(is_fixed) def _init_fixed(): - # Per-query Cauchy-Schwarz bound m_i = ceil(||q_i|| * max_j||k_j||) - C. - qf = q_ref[...].astype(float32) - qn = jnp.sqrt((qf * qf).sum(axis=1))[None, :] # (1, bq) per-query norm - bound = qn * mk_ref[0, h] - m_fixed = jnp.ceil(bound) - _FIXED_M_RECENTER - m_scratch_ref[...] = jnp.broadcast_to(m_fixed, m_scratch_ref.shape) + _write_fixed_m() @pl.when(jnp.logical_not(is_fixed)) def _init_online(): @@ -185,9 +219,28 @@ def last_compute_body_online(kv_compute_index): m_scratch_ref[...], l_scratch_ref[...] = m_prev, l_prev o_scratch_ref[:] = o_prev + def last_compute_body_fixed(kv_compute_index): + # Ragged tail for the pinned fixed-m path: exact slice (padded keys are + # never touched -- with a pinned m their exp2(0 - m_fixed) would be huge + # garbage, so slicing, not masking, is load-bearing here). + q = q_ref[...] + slice_k_len = kv_seq_len % bkv_compute + slice_k = pl.ds(kv_compute_index * bkv_compute, slice_k_len) + qk = lax.dot_general(k_ref[slice_k, :], q, NT_DIM_NUMBERS, preferred_element_type=float32) + v_chunk = v_ref[slice_k, :] + l_prev, o_prev = _fixed_inner(qk, v_chunk, m_scratch_ref[...], l_scratch_ref[...], o_scratch_ref[:]) + l_scratch_ref[...] = l_prev + o_scratch_ref[:] = o_prev + assert bkv % bkv_compute == 0 - if use_fixed_m: + if fixed_only: + + @pl.when(j != grid_width - 1) + def _body_uniform_fixed(): + lax.fori_loop(0, (bkv // bkv_compute), compute_body_fixed, None, unroll=True) + + elif use_fixed_m: @pl.when((j != grid_width - 1) & is_fixed) def _body_fixed(): @@ -203,11 +256,16 @@ def _body_online(): def body(): lax.fori_loop(0, (bkv // bkv_compute), compute_body_online, None, unroll=True) - # The final KV block always runs online for every head. Fixed-m heads arrive - # with m_scratch = ceil(bound) - C and o/l at that reference; the online - # alpha = exp2(m_prev - m_next) rescale renormalizes them transparently. - @pl.when(j == grid_width - 1) - def last_body(): + # Exactly ONE of these runs in the final KV block -- never both (see the + # note on `uniform_fixed_m` above). + # + # `_last_online` is the hybrid default: a fixed-m head arrives with + # m_scratch = ceil(bound) - C, and since that is an upper bound on every + # logit the online step's max leaves it unchanged and its rescale factor is + # exp2(0) = 1, so running the last block online is exact for fixed heads too. + # `_last_fixed` keeps the bound pinned instead, which is what lets the ring's + # accumulate merge assume every hop reports the identical m. + def _last_online(): if kv_seq_len % bkv == 0: iter_num = bkv // bkv_compute lax.fori_loop(0, iter_num, compute_body_online, None, unroll=True) @@ -220,6 +278,31 @@ def last_body(): lax.fori_loop(0, iter_num - 1, compute_body_online, None, unroll=True) last_compute_body_online(iter_num - 1) + def _last_fixed(): + if kv_seq_len % bkv == 0: + iter_num = bkv // bkv_compute + lax.fori_loop(0, iter_num, compute_body_fixed, None, unroll=True) + else: + remain_kv_seq_len = kv_seq_len % bkv + iter_num = (remain_kv_seq_len + bkv_compute - 1) // bkv_compute + if remain_kv_seq_len % bkv_compute == 0: + lax.fori_loop(0, iter_num, compute_body_fixed, None, unroll=True) + else: + lax.fori_loop(0, iter_num - 1, compute_body_fixed, None, unroll=True) + last_compute_body_fixed(iter_num - 1) + + if fixed_only: + # ONE body in the last block (the whole point of uniform mode). + @pl.when(j == grid_width - 1) + def last_body_uniform_fixed(): + _last_fixed() + + else: + + @pl.when(j == grid_width - 1) + def last_body(): + _last_online() + @pl.when(j == grid_width - 1) def end(): l = l_scratch_ref[...] @@ -494,6 +577,7 @@ def _splash_attention_forward_ring( vmem_limit_bytes: int | None = None, use_fixed_m: bool = False, mk: jax.Array | None = None, + uniform_fixed_m: bool = False, ): """Ring-specific forward path that returns pre-reciprocal fp32 accumulators. @@ -578,6 +662,7 @@ def v_index_map(h, i, j, *_): use_base2_exp=use_base2_exp, fuse_reciprocal=False, use_fixed_m=use_fixed_m, + uniform_fixed_m=uniform_fixed_m, ), grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=1, diff --git a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py index 5b1bd6187..bc49c5af7 100644 --- a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py +++ b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py @@ -123,7 +123,9 @@ def _ring_attention_forward( l_init = jnp.zeros((o_shape[0], o_shape[1]), jnp.float32) m_init = jnp.full_like(l_init, mask_value, dtype=jnp.float32) - def body(carry, i: int) -> tuple[tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array, SegmentIds | None], None]: + def body( + carry, i: int + ) -> tuple[tuple[jax.Array, jax.Array, jax.Array, jax.Array, jax.Array, SegmentIds | None], None,]: m_prev, l_prev, o_prev, k_current, v_current, segment_ids_current = carry current_kv_shard_idx = (ring_axis_idx - i) % ring_axis_size @@ -784,14 +786,26 @@ def _merge(m, l, o, mc, lc, oc, valid): m_next = jnp.maximum(m, mc) alpha = exp_fn(m - m_next) beta = exp_fn(mc - m_next) - return m_next, alpha * l + beta * lc, alpha[..., None] * o + beta[..., None] * oc + return ( + m_next, + alpha * l + beta * lc, + alpha[..., None] * o + beta[..., None] * oc, + ) # t=0: own shard (always valid). _attn returns (o, m, l). o, m, l = _attn(k, v) # Non-wrapping one-hop shifts (line ends send/receive nothing). - shift_r = partial(lax.ppermute, axis_name=ring_axis, perm=[(i, i + 1) for i in range(axis_size - 1)]) - shift_l = partial(lax.ppermute, axis_name=ring_axis, perm=[(i, i - 1) for i in range(1, axis_size)]) + shift_r = partial( + lax.ppermute, + axis_name=ring_axis, + perm=[(i, i + 1) for i in range(axis_size - 1)], + ) + shift_l = partial( + lax.ppermute, + axis_name=ring_axis, + perm=[(i, i - 1) for i in range(1, axis_size)], + ) # Prime buffers for t=1 (one hop each direction): device i -> KV_{i-1}, KV_{i+1}. kr, vr = shift_r(k), shift_r(v) @@ -844,7 +858,7 @@ def _custom_ring_attention_forward( perm: list[tuple[int, int]] | None = None, bidirectional: bool = False, use_fixed_m: bool = False, - mk: jax.Array | None = None, + fixed_m_norms: tuple[jax.Array, jax.Array] | None = None, ) -> jax.Array: """Forward-only ring attention using the custom dense splash kernel. @@ -911,15 +925,182 @@ def _custom_ring_attention_forward( num_q_heads = q.shape[0] head_dim_v = v.shape[-1] + + if use_fixed_m: + # Fixed-m ring: each hop gates PER (head, K-shard) against the halved + # un-smoothed bound, so a head can be fixed on one shard and online on + # another. A fixed hop returns m = the Cauchy-Schwarz upper bound (not the + # rowmax); the naive (m, l) merge below would then flush the other hop's + # partial (exp(m_other - m_bound) underflows once the overshoot exceeds + # the f32 window). Merge in LSE space instead: lse = m + log(l) is + # invariant to the kernel's m convention, so overshoot cancels exactly. + # The K-shard norms rotate WITH K/V (a (heads,)-sized ppermute) instead of + # being re-reduced per hop, which would stall the kernel's scalar prefetch. + if fixed_m_norms is None: + raise ValueError("use_fixed_m on the ring path requires fixed_m_norms=(qn_max, mk_h).") + log_fn = jnp.log2 if use_base2_exp else jnp.log + qn_max, mk_h_init = fixed_m_norms + tiny = jnp.finfo(jnp.float32).tiny + # Finite (not -inf) init: the first merge computes exp(init - lse_new) = 0.0 + # exactly; a -inf init meeting an empty partial would produce inf - inf = NaN. + lse_init = -1e30 + + # Every rank's K-shard norms, gathered ONCE before the scan: (R, heads). + # Rotating mk alongside K/V instead (a third per-hop ppermute feeding the + # kernel's scalar prefetch) serialized the K/V rotation against the kernel + # (trace: collective-permute-done 0.004s -> 0.467s per window); a local + # index into a pre-gathered array keeps the per-hop gate collective-free. + # A caller holding the full table already (e.g. a static weight-derived + # bound, identical on every rank) passes it as (ring_size, heads) and + # skips the gather -- an all_gather of a constant is NOT folded by XLA + # and would still occupy the async-collective machinery every call. + if mk_h_init.ndim == 2: + mk_all = mk_h_init + else: + mk_all = lax.all_gather(mk_h_init, ring_axis) # (axis_size, heads) + my_ring_index = lax.axis_index(ring_axis) + + # GLOBAL bound = max over every shard's mk. When ALL local heads pass the + # gate at this single bound, every hop's pinned m is IDENTICAL (it depends + # only on the stationary q rows and the global bound), so hop partials + # combine by PURE ACCUMULATION: o += o_hop, l += l_hop, one normalize at + # the end -- no per-hop LSE math or [H,S,D] divides. The predicate is + # device-uniform ALONG THE RING (the caller pmaxes qn over the ring axis + # and mk_all is a gathered table), so every ppermute participant takes the + # same lax.cond branch; ulysses ranks may diverge freely (no ulysses + # collective lives inside the branches). + mk_global = mk_all.max(axis=0) # (heads,) + all_fixed_global = jnp.all( + qn_max * mk_global <= custom_splash._FIXED_M_RING_SAFE_BOUND # pylint: disable=protected-access + ) + + def _accumulate_scan(_): + mk_arr = jnp.stack([mk_global, jnp.ones_like(mk_global)]) + o_sum = jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32) + l_sum = jnp.zeros((num_q_heads, orig_q_seq_len), jnp.float32) + k_current, v_current = k, v + # Python loop over the (static) ring size rather than a scan: it lets the + # LAST hop skip its rotation. A scan body must rotate unconditionally, and + # the trailing ppermute is NOT dead-code-eliminated (collectives carry + # cross-device pairing), so a scan pays ring_size rotations to consume + # ring_size - 1 shards. At 2 x 387MB/hop over ~16 GB/s of unidirectional + # ICI that wasted hop is ~24ms/layer of wire time -- enough to make the + # ring ICI-bound and hide any kernel-side win. + for hop in range(ring_size): + is_last_hop = hop == ring_size - 1 + if not is_last_hop: + # Issue the next shard's rotation before computing on this one so the + # transfer overlaps the kernel. + k_next = shift(k_current) + v_next = shift(v_current) + o_curr, _, l_curr = custom_splash._splash_attention_forward_ring( # pylint: disable=protected-access + q, + k_current, + v_current, + block_sizes, + q_seq_len=orig_q_seq_len, + kv_seq_len=orig_kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + use_fixed_m=True, + mk=mk_arr, + # This branch only runs under `all_fixed_global`, so the kernel is + # told at compile time that every head is fixed: no per-head scalar + # dispatch, and -- load-bearing -- a single body in the ragged last + # KV block instead of two (a two-body last block is what triggers + # the Mosaic scheduler cliff on the whole grid). + uniform_fixed_m=True, + ) + o_sum = o_sum + o_curr.astype(jnp.float32) + l_sum = l_sum + l_curr.astype(jnp.float32) + if not is_last_hop: + k_current, v_current = k_next, v_next + l_inv = jnp.where(l_sum == 0.0, 0.0, 1.0 / l_sum) + # Narrow INSIDE the branch: lax.cond has to move its output between the + # branch buffer and its own, and this array is [heads, seq, head_dim] -- + # 387MB in f32. Casting here halves that copy (measured 1.04ms/layer for + # the conditional, ~half of fixed-m's whole kernel win). + return (o_sum * l_inv[..., None]).astype(q.dtype) + + def fixed_body(carry, hop, is_last_hop): + o_run, lse_run, k_current, v_current = carry + # Prefetch the next shard while computing on this one. The last hop skips + # it: nothing consumes the rotated shard, and the collective would still + # occupy ICI (see the accumulate path's note). + if is_last_hop: + k_next, v_next = k_current, v_current + else: + k_next = shift(k_current) + v_next = shift(v_current) + + # perm src i -> dst i+1: after `hop` shifts this rank holds the K shard + # of ring rank (my_index - hop) mod R; its norms come from the local table. + mk_h = jax.lax.dynamic_index_in_dim(mk_all, (my_ring_index - hop) % axis_size, keepdims=False) + fixed_ok = (qn_max * mk_h <= custom_splash._FIXED_M_RING_SAFE_BOUND).astype( # pylint: disable=protected-access + jnp.float32 + ) + mk_arr = jnp.stack([mk_h, fixed_ok]) + + o_curr, m_curr, l_curr = custom_splash._splash_attention_forward_ring( # pylint: disable=protected-access + q, + k_current, + v_current, + block_sizes, + q_seq_len=orig_q_seq_len, + kv_seq_len=orig_kv_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + use_fixed_m=True, + mk=mk_arr, + ) + m_curr = m_curr.astype(jnp.float32) + l_curr = l_curr.astype(jnp.float32) + o_curr = o_curr.astype(jnp.float32) + + # Partial -> (normalized output, LSE); empty rows map to lse = -inf. + l_safe = jnp.maximum(l_curr, tiny) + lse_curr = jnp.where(l_curr > 0.0, m_curr + log_fn(l_safe), -jnp.inf) + o_norm = o_curr / l_safe[..., None] + + # LSE-space merge of two normalized partials over disjoint KV sets. + lse_new = jnp.maximum(lse_run, lse_curr) + w_run = exp_fn(lse_run - lse_new) + w_curr = exp_fn(lse_curr - lse_new) + denom = w_run + w_curr + o_new = (w_run[..., None] * o_run + w_curr[..., None] * o_norm) / denom[..., None] + return (o_new, lse_new + log_fn(denom), k_next, v_next), None + + fixed_init = ( + jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32), + jnp.full((num_q_heads, orig_q_seq_len), lse_init, jnp.float32), + k, + v, + ) + + def _lse_scan(_): + carry = fixed_init + for hop in range(ring_size): + carry, _ = fixed_body(carry, hop, hop == ring_size - 1) + return carry[0].astype(q.dtype) + + return lax.cond(all_fixed_global, _accumulate_scan, _lse_scan, None) + o_init = jnp.zeros((num_q_heads, orig_q_seq_len, head_dim_v), jnp.float32) l_init = jnp.zeros((num_q_heads, orig_q_seq_len), jnp.float32) m_init = jnp.full((num_q_heads, orig_q_seq_len), mask_value, jnp.float32) - def body(carry, i): - m_prev, l_prev, o_prev, k_current, v_current = carry - # Prefetch the next shard while we compute on the current one. - k_next = shift(k_current) - v_next = shift(v_current) + m_final, l_final, o_final = m_init, l_init, o_init + k_current, v_current = k, v + # Static python loop, not lax.scan: the last hop must not rotate (see the + # fixed-m accumulate path -- a trailing ppermute survives DCE and burns a full + # shard's worth of ICI bandwidth per layer). + for hop in range(ring_size): + if hop != ring_size - 1: + # Prefetch the next shard while we compute on the current one. + k_next = shift(k_current) + v_next = shift(v_current) o_curr, m_curr, l_curr = custom_splash._splash_attention_forward_ring( # pylint: disable=protected-access q, @@ -931,28 +1112,19 @@ def body(carry, i): use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, - use_fixed_m=use_fixed_m, - mk=mk, ) m_curr = m_curr.astype(jnp.float32) l_curr = l_curr.astype(jnp.float32) o_curr = o_curr.astype(jnp.float32) - m_next = jnp.maximum(m_prev, m_curr) - alpha = exp_fn(m_prev - m_next) + m_next = jnp.maximum(m_final, m_curr) + alpha = exp_fn(m_final - m_next) beta = exp_fn(m_curr - m_next) - l_next = alpha * l_prev + beta * l_curr - o_next = alpha[..., None] * o_prev + beta[..., None] * o_curr - return (m_next, l_next, o_next, k_next, v_next), None - - initial_carry = (m_init, l_init, o_init, k, v) - (_, l_final, o_final, _, _), _ = lax.scan( - body, - initial_carry, - xs=jnp.arange(0, ring_size), - length=ring_size, - unroll=True, - ) + l_final = alpha * l_final + beta * l_curr + o_final = alpha[..., None] * o_final + beta[..., None] * o_curr + m_final = m_next + if hop != ring_size - 1: + k_current, v_current = k_next, v_next l_inv = jnp.where(l_final == 0.0, 0.0, 1.0 / l_final) out = (o_final * l_inv[..., None]).astype(q.dtype) @@ -973,7 +1145,7 @@ def make_custom_ring_attention( perm: list[tuple[int, int]] | None = None, bidirectional: bool = False, use_fixed_m: bool = False, - mk: jax.Array | None = None, + fixed_m_norms: tuple[jax.Array, jax.Array] | None = None, ): """Builds a forward-only ring-attention callable around the custom kernel. @@ -1008,7 +1180,7 @@ def _ring(q, k, v): perm=perm, bidirectional=bidirectional, use_fixed_m=use_fixed_m, - mk=mk, + fixed_m_norms=fixed_m_norms, ) return _ring diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index c56d8c64c..c809b859d 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -15,7 +15,6 @@ import contextlib import functools import math -import os from typing import Optional, Callable, Tuple, Any, Dict import flax.linen as nn from flax import nnx @@ -410,10 +409,19 @@ def _build_padding_segment_ids( kv_mask_for_batch = attention_mask[0, :mask_len] # Tokens past the mask but within key_seq_len are assumed valid. if key_seq_len > mask_len: - kv_mask_for_batch = jnp.concatenate([kv_mask_for_batch, jnp.ones((key_seq_len - mask_len,), jnp.int32)], axis=0) + kv_mask_for_batch = jnp.concatenate( + [kv_mask_for_batch, jnp.ones((key_seq_len - mask_len,), jnp.int32)], + axis=0, + ) # Tokens past key_seq_len are padding. if kv_padded_len > key_seq_len: - kv_mask_for_batch = jnp.concatenate([kv_mask_for_batch, jnp.zeros((kv_padded_len - key_seq_len,), jnp.int32)], axis=0) + kv_mask_for_batch = jnp.concatenate( + [ + kv_mask_for_batch, + jnp.zeros((kv_padded_len - key_seq_len,), jnp.int32), + ], + axis=0, + ) kv_segment_ids = (kv_segment_ids * kv_mask_for_batch).astype(jnp.int32) return segment_ids_cls(q=q_segment_ids, kv=kv_segment_ids) @@ -538,7 +546,14 @@ def wrap_flash_attention(query, key, value): if attention_kernel == "tokamax_ring_custom": # Ring attention backed by the custom dense splash kernel. q stays local, # k/v rotate over the "context" axis (handled inside the ring kernel). - bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes = _extract_custom_block_sizes(flash_block_sizes) + ( + bq, + bkv, + bkv_compute, + bkv_compute_in, + heads_per_tile, + vmem_limit_bytes, + ) = _extract_custom_block_sizes(flash_block_sizes) if heads_per_tile > 1: raise NotImplementedError("tokamax_ring_custom currently supports heads_per_tile == 1 only.") query_local = query * LOG2E if use_base2_exp else query @@ -547,7 +562,10 @@ def wrap_flash_attention(query, key, value): value_local, _, _ = _pad_data_for_flash(value, heads, bkv) bsizes = custom_splash._BlockSizes( - block_q=bq, block_kv=bkv, block_kv_compute=bkv_compute, block_kv_compute_in=bkv_compute_in + block_q=bq, + block_kv=bkv, + block_kv_compute=bkv_compute, + block_kv_compute_in=bkv_compute_in, ) ring_kernel = tokamax_ring_attention_kernel.make_custom_ring_attention( block_sizes=bsizes, @@ -592,7 +610,12 @@ def wrap_flash_attention(query, key, value): tokamax_splash_base.SegmentIds if attention_kernel == "tokamax_ring" else splash_attention_kernel.SegmentIds ) segment_ids = _build_padding_segment_ids( - query_seq_len, query.shape[2], key_seq_len, key.shape[2], attention_mask, segment_ids_cls + query_seq_len, + query.shape[2], + key_seq_len, + key.shape[2], + attention_mask, + segment_ids_cls, ) # make_splash_mha is wrapped around shardmap and seq and head is already @@ -780,7 +803,14 @@ def wrap_ulysses_attention(query, key, value): "The custom dense splash kernel (use_custom_kernel) does not support attention_mask " "(it only handles padding via orig_seq_len); got a non-None attention_mask." ) - bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes = _extract_custom_block_sizes(flash_block_sizes) + ( + bq, + bkv, + bkv_compute, + bkv_compute_in, + heads_per_tile, + vmem_limit_bytes, + ) = _extract_custom_block_sizes(flash_block_sizes) if use_base2_exp: query = query * LOG2E @@ -808,7 +838,10 @@ def wrap_ulysses_attention(query, key, value): mk_arr = jnp.stack([mk_h, fixed_ok]) # (2, local_heads) bsizes = custom_splash._BlockSizes( - block_q=bq, block_kv=bkv, block_kv_compute=bkv_compute, block_kv_compute_in=bkv_compute_in + block_q=bq, + block_kv=bkv, + block_kv_compute=bkv_compute, + block_kv_compute_in=bkv_compute_in, ) splash_kernel = custom_splash.make_splash_mha( @@ -869,7 +902,13 @@ def wrap_ulysses_attention(query, key, value): attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype) # Restore original layout: head-sharded/full-sequence -> sequence-sharded/full-heads. - attention_output = jax.lax.all_to_all(attention_output, axis_name=axis_name, split_axis=2, concat_axis=1, tiled=True) + attention_output = jax.lax.all_to_all( + attention_output, + axis_name=axis_name, + split_axis=2, + concat_axis=1, + tiled=True, + ) return attention_output devices_in_batch_sharding = mesh.shape["data"] * (mesh.shape["fsdp"] if "fsdp" in mesh.shape else 1) @@ -887,7 +926,12 @@ def wrap_ulysses_attention(query, key, value): # a2a tensors inside the scanned layers (measured 7.0 -> expected ~3.5 # s/step at 720p 81f cp8 CFG). batch = query.shape[0] - fold_batch = batch > 1 and (batch * num_heads) % num_shards == 0 + # Only foldable when nothing else shards the batch: with data/fsdp > 1 the + # shard_map still maps axis 0 onto those mesh axes, and a folded (size-1) + # batch is not divisible by them. Those configs are already batch=1 per + # device inside the shard_map, so they never hit the T(2,128) tiling problem + # the fold exists to avoid. + fold_batch = batch > 1 and devices_in_batch_sharding == 1 and (batch * num_heads) % num_shards == 0 if fold_batch: query = query.reshape(1, batch * num_heads, *query.shape[2:]) key = key.reshape(1, batch * num_heads, *key.shape[2:]) @@ -986,7 +1030,11 @@ def _ulysses_ring_attention( @functools.partial( jax.shard_map, mesh=internal_mesh, - in_specs=(internal_q_axis_names, internal_kv_axis_names, internal_kv_axis_names), + in_specs=( + internal_q_axis_names, + internal_kv_axis_names, + internal_kv_axis_names, + ), out_specs=internal_q_axis_names, check_vma=False, ) @@ -1021,7 +1069,12 @@ def wrap_ulysses_ring_attention(query, key, value): # ring shard pads identically so every shard shares the same per-shard ids # and rotation is unneeded. segment_ids = _build_padding_segment_ids( - query_seq_len, q_padded_len, key_seq_len, kv_padded_len, attention_mask, tokamax_splash_base.SegmentIds + query_seq_len, + q_padded_len, + key_seq_len, + kv_padded_len, + attention_mask, + tokamax_splash_base.SegmentIds, ) if not mask_padding_tokens: @@ -1079,6 +1132,13 @@ def wrap_ulysses_ring_attention(query, key, value): return x +def _max_row_norm_per_head(x: jax.Array) -> jax.Array: + """Largest row L2 norm per head of a `[B, H, S, D]` activation.""" + row_sq = jnp.square(x).sum(axis=-1, dtype=jnp.float32) + # 1.01 keeps the result an upper bound despite bf16 mantissa loss. + return jnp.sqrt(row_sq.max(axis=(0, 2))) * 1.01 + + def _ulysses_ring_custom_attention( query: jax.Array, key: jax.Array, @@ -1139,7 +1199,14 @@ def _ulysses_ring_custom_attention( if num_heads % num_ulysses_shards != 0: raise ValueError(f"Ulysses+Ring requires heads divisible by U={num_ulysses_shards}, got heads={num_heads}.") - bq, bkv, bkv_compute, bkv_compute_in, heads_per_tile, vmem_limit_bytes = _extract_custom_block_sizes(flash_block_sizes) + ( + bq, + bkv, + bkv_compute, + bkv_compute_in, + heads_per_tile, + vmem_limit_bytes, + ) = _extract_custom_block_sizes(flash_block_sizes) if heads_per_tile > 1: raise NotImplementedError("ulysses_ring_custom currently supports heads_per_tile == 1 only.") @@ -1155,11 +1222,46 @@ def _ulysses_ring_custom_attention( @functools.partial( jax.shard_map, mesh=internal_mesh, - in_specs=(internal_q_axis_names, internal_kv_axis_names, internal_kv_axis_names), + in_specs=( + internal_q_axis_names, + internal_kv_axis_names, + internal_kv_axis_names, + ), out_specs=internal_q_axis_names, check_vma=False, ) def wrap_ulysses_ring_attention(query, key, value): + fixed_m_norms = None + if use_fixed_m and num_ring_shards > 1: + # Fixed-m's Cauchy-Schwarz inputs, reduced on the PRE-a2a activation so + # the reduction overlaps the all_to_all instead of stalling the first + # ring step (taking them after the a2a measured +8% end to end). + # + # The barrier is load-bearing: the norms are a second consumer of these + # activations, and without it XLA duplicates the producer chain into the + # norm fusion instead of materializing once -- worth 1.46 ms/layer, the + # difference between fixed-m breaking even and winning. + # + # Reducing them further upstream (on the flat [B, S, H*D] form, where + # head_dim is contiguous) is exact and looks cheaper, but there the array + # is still globally sharded, so the reduction becomes a per-layer + # all-reduce over the context axis: measured WORSE (+54 ms per forward). + query, key = jax.lax.optimization_barrier((query, key)) + qn_local = _max_row_norm_per_head(query) + kn_local = _max_row_norm_per_head(key) + if use_base2_exp: + qn_local = qn_local * LOG2E + # The accumulate-vs-LSE lax.cond predicate must be uniform along the RING + # axis (every ppermute participant takes the same branch). + qn_all = jax.lax.pmax(qn_local, (ring_axis, ulysses_axis)) + mk_all = jax.lax.pmax(kn_local, ulysses_axis) + heads_per_dev = qn_all.shape[0] // num_ulysses_shards + start_head = jax.lax.axis_index(ulysses_axis) * heads_per_dev + fixed_m_norms = ( + jax.lax.dynamic_slice_in_dim(qn_all, start_head, heads_per_dev), + jax.lax.dynamic_slice_in_dim(mk_all, start_head, heads_per_dev), + ) + # (1) Ulysses all-to-all over the (intra-chip) ulysses axis: heads -> sequence, # so each device holds the full ring-chunk sequence with heads/U heads. a2a = functools.partial(jax.lax.all_to_all, axis_name=ulysses_axis, tiled=True) @@ -1170,12 +1272,12 @@ def wrap_ulysses_ring_attention(query, key, value): if use_base2_exp: query = query * LOG2E - if use_fixed_m: - # K-smoothing precondition for fixed-m, computed PER SHARD (no ring pmean). - # A global mean would be a perfectly-uniform per-query logit shift, but the - # per-shard local mean differs from it by only O(1/sqrt(local_seq)), and the - # ring's outer online-softmax merge re-normalizes across shards anyway, so we - # drop the per-layer ring collective and accept the negligible shift error. + if use_fixed_m and num_ring_shards == 1: + # K-smoothing precondition for fixed-m (R=1 / pure-ulysses semantics, + # same as _ulysses_attention). The R>1 ring path deliberately does NOT + # smooth: no ring rank holds the full K to compute a mean, and a per- + # shard mean would shift each hop's logits differently, breaking the + # cross-shard merge; it gates on the un-smoothed halved bound instead. kbar = jnp.mean(key, axis=2, keepdims=True) key = key - kbar @@ -1184,26 +1286,19 @@ def wrap_ulysses_ring_attention(query, key, value): value, _, _ = _pad_data_for_flash(value, heads, bkv) mk_arr = None - if use_fixed_m: - # Per-(local-)head Cauchy-Schwarz inputs, all LOCAL to this ring shard. The - # outer ring merge does an online softmax across shards, so each shard's - # kernel may use its own local max||k|| as the fixed-m bound for its own - # local keys -- no global ring pmax is needed for correctness. This removes - # the second per-layer ring collective. + if use_fixed_m and num_ring_shards == 1: qf = query.astype(jnp.float32) kf = key.astype(jnp.float32) qn_max = jnp.sqrt((qf * qf).sum(-1)).max(axis=(0, 2)) # (local_heads,) mk_h = jnp.sqrt((kf * kf).sum(-1)).max(axis=(0, 2)) # (local_heads,) local fixed_ok = (qn_max * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(jnp.float32) - if os.environ.get("FIXED_M_FORCE_ALL", "0") == "1": - # PERF PROBE ONLY (unsafe): force every head onto the fixed-m fast path, - # bypassing the safety gate, to measure fixed-m's speed CEILING on the - # ring kernel. Output may be garbage; timing is still valid. - fixed_ok = jnp.ones_like(fixed_ok) mk_arr = jnp.stack([mk_h, fixed_ok]) # (2, local_heads) bsizes = custom_splash._BlockSizes( - block_q=bq, block_kv=bkv, block_kv_compute=bkv_compute, block_kv_compute_in=bkv_compute_in + block_q=bq, + block_kv=bkv, + block_kv_compute=bkv_compute, + block_kv_compute_in=bkv_compute_in, ) if num_ring_shards == 1: # (2a) R=1: the ring is trivial (no rotation) -> use the lighter dedicated @@ -1221,7 +1316,11 @@ def wrap_ulysses_ring_attention(query, key, value): use_fixed_m=use_fixed_m, ) if use_fixed_m: - attention_output = jnp.swapaxes(jax.vmap(splash_kernel, in_axes=(0, 0, 0, None))(query, key, value, mk_arr), 2, 3) + attention_output = jnp.swapaxes( + jax.vmap(splash_kernel, in_axes=(0, 0, 0, None))(query, key, value, mk_arr), + 2, + 3, + ) else: attention_output = jnp.swapaxes(jax.vmap(splash_kernel, in_axes=(0, 0, 0))(query, key, value), 2, 3) else: @@ -1239,7 +1338,7 @@ def wrap_ulysses_ring_attention(query, key, value): ring_size=num_ring_shards, bidirectional=bidirectional, use_fixed_m=use_fixed_m, - mk=mk_arr, + fixed_m_norms=fixed_m_norms, ) attention_output = jax.vmap(ring_kernel, in_axes=(0, 0, 0))(query, key, value) attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype) @@ -1667,7 +1766,14 @@ def _apply_attention( seq_len_idx = 2 can_use_flash_attention = True - if attention_kernel in ["flash", "tokamax_flash", "ulysses", "ulysses_custom", "ulysses_custom_fixed_m", "ulysses_ring"]: + if attention_kernel in [ + "flash", + "tokamax_flash", + "ulysses", + "ulysses_custom", + "ulysses_custom_fixed_m", + "ulysses_ring", + ]: can_use_flash_attention = ( query.shape[seq_len_idx] >= flash_min_seq_length and key.shape[seq_len_idx] >= flash_min_seq_length @@ -1977,7 +2083,13 @@ def __init__( self.mask_padding_tokens = mask_padding_tokens self.residual_checkpoint_name = residual_checkpoint_name - def apply_attention(self, query: Array, key: Array, value: Array, attention_mask: Array = None): + def apply_attention( + self, + query: Array, + key: Array, + value: Array, + attention_mask: Array = None, + ): return _apply_attention( query=query, key=key, @@ -2141,7 +2253,10 @@ def __init__( axis_names_kv = (BATCH, CROSS_ATTN_HEAD, CROSS_ATTN_KV_LENGTH, D_KV) if attention_kernel in ("tokamax_ring", "tokamax_ring_custom", "ulysses_ring") and not is_self_attention: attention_kernel = "tokamax_flash" # do not use ring attention for cross attention - if attention_kernel in ("ulysses_ring_custom", "ulysses_ring_custom_bidir") and not is_self_attention: + if ( + attention_kernel in ("ulysses_ring_custom", "ulysses_ring_custom_bidir", "ulysses_ring_custom_fixed_m") + and not is_self_attention + ): attention_kernel = "ulysses_custom" # plain ulysses (no ring) for cross attention self.added_kv_proj_dim = added_kv_proj_dim # New for I2V self.image_seq_len = image_seq_len # New for I2V diff --git a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py index 8d628b493..5aac990a7 100644 --- a/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py +++ b/src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py @@ -23,6 +23,8 @@ from ...models.wan.transformers.transformer_wan import WanModel from typing import List, Union, Optional from ...pyconfig import HyperParameters +from ... import aot_cache +import os import concurrent.futures from functools import partial import time @@ -394,12 +396,20 @@ def run_inference_2_2( # Select transformer + guidance for this phase. if step_uses_high[step]: - graphdef, state, rest = high_noise_graphdef, high_noise_state, high_noise_rest + graphdef, state, rest = ( + high_noise_graphdef, + high_noise_state, + high_noise_rest, + ) guidance_scale = guidance_scale_high kv_cache = kv_cache_high encoder_attention_mask = encoder_attention_mask_high else: - graphdef, state, rest = low_noise_graphdef, low_noise_state, low_noise_rest + graphdef, state, rest = ( + low_noise_graphdef, + low_noise_state, + low_noise_rest, + ) guidance_scale = guidance_scale_low kv_cache = kv_cache_low encoder_attention_mask = encoder_attention_mask_low @@ -485,7 +495,14 @@ def run_inference_2_2( # SenCache state ref_noise_pred = jnp.zeros( - (bsz * 2, latents.shape[1], latents.shape[2], latents.shape[3], latents.shape[4]), dtype=latents.dtype + ( + bsz * 2, + latents.shape[1], + latents.shape[2], + latents.shape[3], + latents.shape[4], + ), + dtype=latents.dtype, ) ref_latent = jnp.zeros_like(latents) ref_timestep = jnp.array(0.0, dtype=jnp.float32) @@ -788,7 +805,10 @@ def low_noise_branch(operands): ) if scan_diffusion_loop: - scheduler_state = scheduler_state.replace(last_sample=jnp.zeros_like(latents), step_index=jnp.array(0, dtype=jnp.int32)) + scheduler_state = scheduler_state.replace( + last_sample=jnp.zeros_like(latents), + step_index=jnp.array(0, dtype=jnp.int32), + ) def scan_body(carry, t): current_latents, current_scheduler_state = carry @@ -830,7 +850,50 @@ def scan_body(carry, t): final_latents, _ = final_carry return final_latents + # Warmup only runs a couple of steps and the scheduler puts both on the + # high-noise transformer, so the low-noise weights are never executed; their + # first real call then lands mid-generation and costs 26.4s against a 2.6s + # step. Touch every weight set here instead, so the first generation already + # runs at steady-state latency. + if aot_cache.in_warmup() and do_classifier_free_guidance: + with aot_cache.real_execution(): + priming_latents = jnp.concatenate([latents] * 2) + priming_timestep = jnp.broadcast_to(timesteps[0], bsz * 2) + for _gd, _st, _rest, _gs, _kv, _mask in ( + ( + high_noise_graphdef, + high_noise_state, + high_noise_rest, + guidance_scale_high, + kv_cache_high, + encoder_attention_mask_high, + ), + ( + low_noise_graphdef, + low_noise_state, + low_noise_rest, + guidance_scale_low, + kv_cache_low, + encoder_attention_mask_low, + ), + ): + jax.block_until_ready( + transformer_forward_pass_full_cfg( + _gd, + _st, + _rest, + priming_latents, + priming_timestep, + prompt_embeds_combined, + guidance_scale=_gs, + kv_cache=_kv, + rotary_emb=rotary_emb, + encoder_attention_mask=_mask, + ) + ) + profiler = None + _inflight_latents = [] for step in range(num_inference_steps): if config and max_utils.profiler_enabled(config) and step == first_profiling_step: profiler = max_utils.Profiler(config) @@ -886,6 +949,22 @@ def scan_body(carry, t): latents, scheduler_state = scheduler.step(scheduler_state, noise_pred, t, latents).to_tuple() + # Bound the async dispatch queue. The python loop otherwise queues every + # remaining step on the device; past ~30 in-flight steps the runtime + # degrades ~20% (measured: 40-step fixed-m 3.30 s/step unbounded vs 2.69 + # with a mid-run drain; 4/12/24-step runs unaffected). A periodic drain + # keeps the queue shallow and costs nothing (the device stays busy). + # Sliding-window bound on in-flight dispatched steps: block on the + # output of step (N - depth) so the device queue holds at most `depth` + # steps while the pipeline stays full (a periodic FULL drain would + # empty the pipe and cost a bubble per drain). The safe depth shrinks + # as the per-step executable's footprint grows (bkv2048 tiles need + # ~2-4; bkv1024 tolerates 8+); past it the TPU runtime degrades ~20%. + _inflight_latents.append(latents) + depth = int(os.environ.get("MAXD_QUEUE_MAX_INFLIGHT", "4")) + if depth > 0 and len(_inflight_latents) > depth: + _inflight_latents.pop(0).block_until_ready() + if config and max_utils.profiler_enabled(config) and step == last_profiling_step: if profiler: latents.block_until_ready() diff --git a/src/maxdiffusion/tests/ring_fixed_m_test.py b/src/maxdiffusion/tests/ring_fixed_m_test.py new file mode 100644 index 000000000..0f4cc6876 --- /dev/null +++ b/src/maxdiffusion/tests/ring_fixed_m_test.py @@ -0,0 +1,193 @@ +""" +Copyright 2026 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +"""Unit tests for the fixed-m path of the custom RING attention. + +The ring path gates fixed-m PER (head, K-shard) against the halved +un-smoothed bound, rotates each K shard's max row norm alongside K/V, and +merges the per-hop partials in LSE space (invariant to fixed-m's bound +overshoot). These tests check, against an f32 dense-softmax reference: + + * the untouched online ring path (regression guard), + * fixed-m with every (head, shard) eligible, + * a sink head ineligible on every shard (all-online fallback), + * a head eligible on one shard but not the other -- the mixed + fixed/online partial case that requires the LSE merge. +""" + +import functools +import math +import unittest + +import jax +import jax.numpy as jnp +import numpy as np + +from maxdiffusion.kernels import custom_splash_attention as custom_splash +from maxdiffusion.kernels.splash_attention import ring_attention_kernel + +_LOG2E = math.log2(math.e) +_RING_AXIS = "ring" +_RING_SIZE = 2 + + +class RingFixedMTest(unittest.TestCase): + """Numerical tests for the fixed-m custom ring attention.""" + + num_heads = 4 + shard_len = 2048 # per-device sequence; total = shard_len * ring_size + head_dim = 128 + + def setUp(self): + super().setUp() + if jax.default_backend() != "tpu": + self.skipTest("Only supported on TPUs.") + if len(jax.devices()) < _RING_SIZE: + self.skipTest(f"Requires {_RING_SIZE} devices.") + self.scale = 1.0 / math.sqrt(self.head_dim) + self.block_sizes = custom_splash._BlockSizes(block_q=1024, block_kv=1024, block_kv_compute=512, block_kv_compute_in=256) + devices = np.asarray(jax.devices()[:_RING_SIZE]) + self.mesh = jax.sharding.Mesh(devices, (_RING_AXIS,)) + + def _random_qkv(self, q_gain=None, k_gain=None): + """bf16 (q, k, v), [heads, total_seq, dim]; optional (head, row-slice) gains.""" + total = self.shard_len * _RING_SIZE + shape = (self.num_heads, total, self.head_dim) + q = jax.random.normal(jax.random.PRNGKey(0), shape, jnp.bfloat16) + k = jax.random.normal(jax.random.PRNGKey(1), shape, jnp.bfloat16) + v = jax.random.normal(jax.random.PRNGKey(2), shape, jnp.bfloat16) + if q_gain is not None: + head, rows, gain = q_gain + q = q.at[head, rows].multiply(gain) + if k_gain is not None: + head, rows, gain = k_gain + k = k.at[head, rows].multiply(gain) + return q, k, v + + def _scaled_inputs(self, q, k): + """The EXACT bf16 tensors the kernel sees (attention_flax's contract): + k pre-scaled by the softmax scale, q pre-scaled by LOG2E (base-2 + kernel). The reference must consume these same tensors -- comparing + against raw f32 inputs instead double-rounds k, and on an amplified + head (logits ~2^9) the bf16 rounding alone shifts softmax weights by + factors of ~2^2, drowning the kernel error being tested.""" + q_in = (q * _LOG2E).astype(q.dtype) + k_in = (k.astype(jnp.float32) * self.scale).astype(k.dtype) + return q_in, k_in + + def _reference(self, q_in, k_in, v): + """Dense f32 log2-domain softmax on the kernel's own bf16 inputs.""" + qf, kf, vf = (x.astype(jnp.float32) for x in (q_in, k_in, v)) + logits = jnp.einsum("hqd,hkd->hqk", qf, kf) # LOG2E & scale pre-folded + return jax.nn.softmax(logits * math.log(2.0), axis=-1) @ vf + + def _run_ring(self, q_in, k_in, v, use_fixed_m): + """Runs the custom ring under shard_map with per-rank fixed_m_norms + from the LOCAL q / initial K shard.""" + spec = jax.sharding.PartitionSpec(None, _RING_AXIS, None) + + @functools.partial( + jax.shard_map, + mesh=self.mesh, + in_specs=(spec, spec, spec), + out_specs=spec, + check_vma=False, + ) + def _body(ql, kl, vl): + fixed_m_norms = None + if use_fixed_m: + qf = ql.astype(jnp.float32) + kf = kl.astype(jnp.float32) + qn_max = jnp.sqrt((qf * qf).sum(-1)).max(axis=1) # (heads,) + mk_h = jnp.sqrt((kf * kf).sum(-1)).max(axis=1) # (heads,) local shard + fixed_m_norms = (qn_max, mk_h) + ring = ring_attention_kernel.make_custom_ring_attention( + block_sizes=self.block_sizes, + orig_q_seq_len=self.shard_len, + orig_kv_seq_len=self.shard_len, + use_base2_exp=True, + ring_axis=_RING_AXIS, + ring_size=_RING_SIZE, + use_fixed_m=use_fixed_m, + fixed_m_norms=fixed_m_norms, + ) + return ring(ql, kl, vl) + + return _body(q_in, k_in, v) + + def _gate_per_shard(self, q_in, k_in): + """(heads, ring_size) eligibility against the halved un-smoothed bound.""" + qf = q_in.astype(jnp.float32) + kf = k_in.astype(jnp.float32) + qn = jnp.sqrt((qf * qf).sum(-1)) # (heads, total) + kn = jnp.sqrt((kf * kf).sum(-1)) + gates = [] + for r in range(_RING_SIZE): + rows = slice(r * self.shard_len, (r + 1) * self.shard_len) + # Stationary q max is per-RANK, but for the gate check we use the global + # q max: it upper-bounds every rank's local max, so "eligible globally" + # implies eligible on every rank. + bound = qn.max(axis=1) * kn[:, rows].max(axis=1) + gates.append(bound <= custom_splash._FIXED_M_RING_SAFE_BOUND) + return jnp.stack(gates, axis=1) + + def _run_and_compare(self, q, k, v, use_fixed_m): + q_in, k_in = self._scaled_inputs(q, k) + out = self._run_ring(q_in, k_in, v, use_fixed_m=use_fixed_m).astype(jnp.float32) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + return float(jnp.max(jnp.abs(out - self._reference(q_in, k_in, v)))) + + def _gate(self, q, k): + return self._gate_per_shard(*self._scaled_inputs(q, k)) + + def test_online_ring_matches_reference(self): + q, k, v = self._random_qkv() + self.assertLess(self._run_and_compare(q, k, v, use_fixed_m=False), 2e-2) + + def test_fixed_m_all_eligible_matches_reference(self): + q, k, v = self._random_qkv() + self.assertTrue(bool(jnp.all(self._gate(q, k)))) + self.assertLess(self._run_and_compare(q, k, v, use_fixed_m=True), 2e-2) + + def test_sink_head_falls_back_everywhere(self): + total = self.shard_len * _RING_SIZE + q, k, v = self._random_qkv(q_gain=(0, slice(0, total), 40.0)) + gate = self._gate(q, k) + self.assertFalse(bool(jnp.any(gate[0]))) # head 0 online on every shard + self.assertTrue(bool(jnp.all(gate[1:]))) + self.assertLess(self._run_and_compare(q, k, v, use_fixed_m=True), 2e-2) + + def test_fixed_m_accumulate_ragged_tail(self): + # All-eligible (accumulate merge) with tiles that leave a ragged last KV + # block (2048 %% 768 = 512) and a ragged inner chunk (512 %% 384 = 128), + # covering the pinned fixed-m path's exact-slice tail handling. + self.block_sizes = custom_splash._BlockSizes(block_q=1024, block_kv=768, block_kv_compute=384, block_kv_compute_in=384) + q, k, v = self._random_qkv() + self.assertTrue(bool(jnp.all(self._gate(q, k)))) + self.assertLess(self._run_and_compare(q, k, v, use_fixed_m=True), 2e-2) + + def test_mixed_fixed_online_across_shards(self): + # Amplify head 0's keys on shard 1 only: head 0 is fixed on shard 0 but + # online on shard 1 -- the mixed-partial merge the LSE space exists for. + q, k, v = self._random_qkv(k_gain=(0, slice(self.shard_len, self.shard_len * _RING_SIZE), 40.0)) + gate = self._gate(q, k) + self.assertTrue(bool(gate[0, 0])) + self.assertFalse(bool(gate[0, 1])) + self.assertLess(self._run_and_compare(q, k, v, use_fixed_m=True), 2e-2) + + +if __name__ == "__main__": + unittest.main()