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()