diff --git a/src/maxdiffusion/kernels/internal_ring_attention.py b/src/maxdiffusion/kernels/internal_ring_attention.py new file mode 100644 index 000000000..727fba685 --- /dev/null +++ b/src/maxdiffusion/kernels/internal_ring_attention.py @@ -0,0 +1,525 @@ +""" +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. +""" + +"""Ring attention with the KV permutation INSIDE the Pallas kernel. + +Contrast with `kernels/splash_attention/ring_attention_kernel.py` +(`_custom_ring_attention_forward`), which is the design MaxDiffusion ships +today: there the ring lives in XLA. Per ring hop it issues a `lax.ppermute`, +launches a *separate* pallas_call over the whole local shard, and merges the +per-hop `(m, l, o)` partials with an online-softmax merge written in JAX. R hops +therefore cost R kernel launches and a full `[H, S, D]` fp32 numerator round trip +through HBM: R written by the kernels (a pallas_call is an opaque custom call, so +XLA cannot elide its outputs), R read back by the merge. The merge chain itself +IS fused -- XLA collapses all R hops' merges plus the final normalize into one +kLoop fusion -- so the cost is the 2R numerator traffic, not R-1 separate passes. + +Here the hop is a `pltpu.make_async_remote_copy` issued from inside the kernel +body, so a single pallas_call sees every ring shard. The loop nest is the one +the PCP design note specifies -- the ring is the INNERMOST loop, under the kv +block loop: + + grid = (b, h, i, j, r) + for each q block i: + for each kv block j: + for each hop r: <-- one bkv-sized KV block arrives over ICI + accumulate online softmax in VMEM + +Two consequences follow from putting `r` innermost. + +1. The online-softmax accumulator for q block `i` never leaves VMEM: `(m, l, o)` + are carried across every hop and every kv block, normalized once, and written + to HBM exactly once. There is no cross-hop merge and no per-hop numerator + round trip. + +2. Only ONE kv block per tensor is ever in flight, so the ring buffer is + `2 x bkv x head_dim` and not `2 x kv_seq x head_dim`. That is what makes the + design fit at all: a whole-shard rotation needs ~40 MB of the 64 MB VMEM + budget at the 2D-ring shapes (20 heads x 37800 tokens) and OOMs. + +The price is ICI volume. `r` under `i` means the whole KV shard is rotated once +per q block, so the wire traffic is `num_q_blocks x (R-1) x |KV shard|` against +the external design's `(R-1) x |KV shard|`. That trade is why this design wins +for LLM inference -- GQA shrinks the KV shard by the group ratio G, so rotating +it `num_q_blocks` times is cheap -- and is a much closer call for diffusion +self-attention, where G = 1 and KV is exactly as large as Q. + +Flow control. Slot `(t+1) % 2` is the slot the receiver was computing out of at +hop t-1, so an unsynchronised push can land on a lagging neighbour's live block. +A one-credit protocol closes that: after finishing a slot a rank signals its +ring-UPSTREAM neighbour, and a rank waits for one credit before pushing. One +credit is exactly right -- with two slots a sender may run at most one hop ahead +of its receiver. +""" + +import functools + +import jax +import jax.numpy as jnp +from jax import lax +from jax.experimental import pallas as pl +from jax.experimental.pallas import tpu as pltpu + +from maxdiffusion.kernels import custom_splash_attention as custom_splash + +DEFAULT_MASK_VALUE = custom_splash.DEFAULT_MASK_VALUE +NUM_SUBLANES = custom_splash.NUM_SUBLANES +NT_DIM_NUMBERS = custom_splash.NT_DIM_NUMBERS + +# collective_id namespace for this kernel's barrier semaphore. Must not collide +# with another Pallas collective live in the same program; nothing else in +# MaxDiffusion uses in-kernel collectives today. +_COLLECTIVE_ID = 7 + + +def _neighbor(axis_names, ring_axis, offset): + """Mesh-index tuple of the rank `offset` steps along `ring_axis`.""" + ring_size = lax.axis_size(ring_axis) + idx = lax.axis_index(ring_axis) + nxt = lax.rem(idx + offset + ring_size, ring_size) + return tuple(nxt if a == ring_axis else lax.axis_index(a) for a in axis_names) + + +def _internal_ring_kernel( + # scalar prefetch + mk_ref, # (2, num_q_heads) f32: [0,h]=global max_j||k_j|| over the WHOLE ring, + # [1,h]=eligibility (unused when uniform) + # inputs + q_ref, # VMEM (bq, head_dim_qk) -- BlockSpec pipelined + k_hbm, # ANY (batch * num_kv_heads, kv_pad, head_dim_qk) + v_hbm, # ANY (batch * num_kv_heads, kv_pad, head_dim_v) + # outputs + o_ref, # VMEM (head_dim_v, bq) + # scratch + k_buf, # VMEM (2, bkv, head_dim_qk) + v_buf, # VMEM (2, bkv, head_dim_v) + m_scratch_ref, # VMEM (NUM_SUBLANES, bq) f32 + l_scratch_ref, # VMEM (NUM_SUBLANES, bq) f32 + o_scratch_ref, # VMEM (head_dim_v, bq) f32 + local_sem, # DMA (2,) [k, v] + send_sem, # DMA (2,) [k, v] + recv_sem, # DMA (2,) [k, v] + credit_sem, # REGULAR + *, + mask_value: float, + grid_width: int, + ring_size: int, + bkv: int, + bkv_compute: int, + bkv_compute_in: int, + head_dim_v: int, + kv_seq_len: int, + q_heads_per_kv_head: int, + num_kv_heads: int, + use_base2_exp: bool, + use_fixed_m: bool, + axis_names: tuple[str, ...], + ring_axis: str, +): + float32 = jnp.float32 + head_dim_v_repeats, rem = divmod(head_dim_v, NUM_SUBLANES) + if rem != 0: + raise NotImplementedError(f"{head_dim_v=} should be a multiple of {NUM_SUBLANES}") + + b, h, i, j, r = (pl.program_id(n) for n in range(5)) + exp = jnp.exp2 if use_base2_exp else jnp.exp + sv_dims = (((0,), (0,)), ((), ())) + + # The KV HBM refs arrive flattened to (batch * num_kv_heads, kv_pad, d) so one + # `pl.ds` picks a shard: `.at[]` on an ANY-space ref does not squeeze, so the + # source and destination ranks have to line up by construction. + hk = b * num_kv_heads + h // q_heads_per_kv_head + + # Hop counter within one (b, h, i). Slot parity follows it so the block a hop + # computes from and the block the next hop streams into are always different. + t = j * ring_size + r + slot = lax.rem(t, 2) + nslot = 1 - slot + + is_first_hop = (j == 0) & (r == 0) + is_last_hop = (j == grid_width - 1) & (r == ring_size - 1) + + upstream = _neighbor(axis_names, ring_axis, -1) + downstream = _neighbor(axis_names, ring_axis, +1) + + def _local_load(dst_slot, block): + """Own KV block, HBM -> VMEM: what the stock BlockSpec pipeline would fetch. + The ring never re-reads a neighbour's block from HBM, only over ICI.""" + src = (pl.ds(hk, 1), pl.ds(block * bkv, bkv)) + dst = pl.ds(dst_slot, 1) + return ( + pltpu.make_async_copy(k_hbm.at[src], k_buf.at[dst], local_sem.at[0]), + pltpu.make_async_copy(v_hbm.at[src], v_buf.at[dst], local_sem.at[1]), + ) + + def _remote_push(src_slot, dst_slot): + """One ring hop: the KV block I hold now -> the downstream rank's next slot.""" + src, dst = pl.ds(src_slot, 1), pl.ds(dst_slot, 1) + return ( + pltpu.make_async_remote_copy(k_buf.at[src], k_buf.at[dst], send_sem.at[0], recv_sem.at[0], device_id=downstream), + pltpu.make_async_remote_copy(v_buf.at[src], v_buf.at[dst], send_sem.at[1], recv_sem.at[1], device_id=downstream), + ) + + # ------------------------------------------------------------------- DMA -- + # (0) Once per launch: rendezvous with both ring neighbours before any remote + # write or remote semaphore signal can be issued (the Pallas all_gather + # example's `main_barrier`). At R == 2 upstream and downstream are the same + # rank, which the 2-signal / 2-wait form still handles. + @pl.when((b == 0) & (h == 0) & (i == 0) & is_first_hop) + def _barrier(): + sem = pltpu.get_barrier_semaphore() + pl.semaphore_signal(sem, 1, device_id=upstream) + pl.semaphore_signal(sem, 1, device_id=downstream) + pl.semaphore_wait(sem, 2) + + # (1) First hop of a q block has no prefetch behind it: load and block. Not + # prefetching across the (b, h, i) boundary keeps the head/block index out of + # the DMA schedule at the cost of one exposed 2 x bkv HBM read per q block. + @pl.when(is_first_hop) + def _prime(): + for dma in _local_load(0, 0): + dma.start() + for dma in _local_load(0, 0): + dma.wait() + + # (2) Wait for the block this hop computes from. + @pl.when(jnp.logical_not(is_first_hop)) + def _await_current(): + @pl.when(r == 0) + def _await_local(): + # My own block j, prefetched from HBM by hop (j-1, R-1). + for dma in _local_load(slot, j): + dma.wait() + + @pl.when(r > 0) + def _await_remote(): + # Pushed by the upstream rank one hop ago: its source slot was `nslot`, + # my destination slot is `slot`. Same shapes => same semaphore credit. + for dma in _remote_push(nslot, slot): + dma.wait_recv() + + # (3) Retire my own previous send before its source buffer is reused as this + # hop's DMA destination (hop t-1's source slot == hop t's destination slot). + # Hop t-1 pushed iff its r was < R-1: always true when r > 0, never when + # r == 0 (the r == R-1 hop reloads from HBM instead of pushing). + @pl.when(r > 0) + def _retire_send(): + for dma in _remote_push(nslot, slot): + dma.wait_send() + + # (3b) Release `nslot` -- the block I finished computing at hop t-1 -- to the + # upstream rank, which is about to push into it. + # + # THIS MUST LIVE IN THE PROLOGUE OF HOP t, NOT THE EPILOGUE OF HOP t-1. + # Emitted from the epilogue it sits in the same grid step as the compute that + # reads `k_buf[slot]`, and nothing orders a `semaphore_signal` after those + # vector loads -- Mosaic may hoist it, letting the upstream overwrite a buffer + # that is still being read. That is a genuine race: it reproduced at R=8 with + # a ragged tail, non-deterministically (rel err 0.30 / 0.41 / 0.52 on repeat + # runs of one config), while R<=4 happened to schedule safely. Grid iteration + # order IS a real ordering guarantee, so releasing here puts the signal + # provably after hop t-1's compute. + # + # Condition: release iff a push targets `nslot` at THIS hop, i.e. r < R-1. + # At the last hop of a q block (r == R-1) the slot is refilled from HBM + # instead, so no credit is owed -- which also closes the ledger at zero + # (Pallas checks semaphores are drained at kernel exit) with no seed needed: + # at hop 0 `nslot` has never been written, so releasing it is correct and it + # is exactly the credit the downstream rank's first push consumes. + @pl.when((r < ring_size - 1) & (ring_size > 1)) + def _release(): + pl.semaphore_signal(credit_sem, 1, device_id=upstream) + + # (4) Issue the next hop's transfer. + @pl.when(jnp.logical_not(is_last_hop)) + def _prefetch_next(): + @pl.when(r < ring_size - 1) + def _push(): + # One credit == "the downstream rank has retired the slot I am about to + # write". Without it a push lands on a lagging neighbour's live block. + pl.semaphore_wait(credit_sem, 1) + for dma in _remote_push(slot, nslot): + dma.start() + + @pl.when(r == ring_size - 1) + def _reload_own(): + # End of a ring cycle: the next kv block starts from my own shard again. + for dma in _local_load(nslot, j + 1): + dma.start() + + # ------------------------------------------------------------- accumulate -- + @pl.when(is_first_hop) + def _init(): + o_scratch_ref[...] = jnp.zeros_like(o_scratch_ref) + l_scratch_ref[...] = jnp.zeros_like(l_scratch_ref) + if use_fixed_m: + # Cauchy-Schwarz bound m_i = ceil(||q_i|| * max_j||k_j||) - C, pinned for + # the WHOLE ring. Unlike the external ring this needs no per-hop gating and + # no LSE merge: one kernel sees every shard and the accumulator never + # leaves VMEM, so every hop is already subtracting the identical m. + # `mk_ref[0, h]` is max||k|| reduced over the ring AND ulysses axes by the + # caller, so the bound covers keys this rank never holds. + qf = q_ref[...].astype(float32) + qn = jnp.sqrt((qf * qf).sum(axis=1))[None, :] + m_fixed = jnp.ceil(qn * mk_ref[0, h]) - custom_splash._FIXED_M_RECENTER # pylint: disable=protected-access + m_scratch_ref[...] = jnp.broadcast_to(m_fixed, m_scratch_ref.shape) + else: + m_scratch_ref[...] = jnp.full_like(m_scratch_ref, mask_value) + + def _online_inner(qk, v_chunk, m_prev, l_prev, o_prev): + step = bkv_compute_in + for c in range(0, qk.shape[0], step): + qk_slice = qk[c : c + step] + m_curr = qk_slice.max(axis=0)[None, :] + m_next = jnp.maximum(m_prev, m_curr) + s_curr = exp(qk_slice - m_next[0:1]) + l_curr = s_curr.sum(axis=0, keepdims=True) + alpha = exp(m_prev - m_next) + l_next = l_curr + alpha * l_prev + o_curr = lax.dot_general( + v_chunk[c : c + step], + s_curr.astype(q_ref.dtype), + sv_dims, + preferred_element_type=float32, + ) + o_prev = alpha[0:1, ...] * o_prev + o_curr + m_prev, l_prev = m_next, l_next + return m_prev, l_prev, o_prev + + def _fixed_inner(qk, v_chunk, m_fix, l_prev, o_prev): + # m is constant: no reduce-max over the block and no alpha rescale of o. + step = bkv_compute_in + for c in range(0, qk.shape[0], step): + s_curr = exp(qk[c : c + step] - m_fix[0:1]) + l_prev = l_prev + s_curr.sum(axis=0, keepdims=True) + o_prev = o_prev + lax.dot_general( + v_chunk[c : c + step], + s_curr.astype(q_ref.dtype), + sv_dims, + preferred_element_type=float32, + ) + return l_prev, o_prev + + def _step(offset, length): + q = q_ref[...] + sl = pl.ds(offset, length) + qk = lax.dot_general(k_buf[slot, sl, :], q, NT_DIM_NUMBERS, preferred_element_type=float32) + v_chunk = v_buf[slot, sl, :] + if use_fixed_m: + l_prev, o_prev = _fixed_inner(qk, v_chunk, m_scratch_ref[...], l_scratch_ref[...], o_scratch_ref[:]) + l_scratch_ref[...] = l_prev + else: + m_prev, l_prev, o_prev = _online_inner(qk, v_chunk, m_scratch_ref[...], l_scratch_ref[...], o_scratch_ref[:]) + m_scratch_ref[...], l_scratch_ref[...] = m_prev, l_prev + o_scratch_ref[:] = o_prev + + def compute_body(kv_compute_index, _): + _step(kv_compute_index * bkv_compute, bkv_compute) + + assert bkv % bkv_compute == 0 + + @pl.when(j != grid_width - 1) + def _body(): + lax.fori_loop(0, bkv // bkv_compute, compute_body, None, unroll=True) + + @pl.when(j == grid_width - 1) + def _last_body(): + # Ragged tail. `kv_seq_len` is the un-padded shard length and every ring rank + # pads identically, so the same tail applies on every hop. + if kv_seq_len % bkv == 0: + lax.fori_loop(0, bkv // bkv_compute, compute_body, None, unroll=True) + else: + remain = kv_seq_len % bkv + iter_num = (remain + bkv_compute - 1) // bkv_compute + if remain % bkv_compute == 0: + lax.fori_loop(0, iter_num, compute_body, None, unroll=True) + else: + lax.fori_loop(0, iter_num - 1, compute_body, None, unroll=True) + _step((iter_num - 1) * bkv_compute, remain % bkv_compute) + + # -------------------------------------------------------------- epilogue -- + # Nothing to release here: the credit for this hop's slot is emitted in the + # NEXT hop's prologue (see 3b), which is the only placement that is provably + # ordered after this hop's compute. + @pl.when(is_last_hop) + def _write_out(): + l = l_scratch_ref[...] + l_inv = jnp.tile(1.0 / l, (head_dim_v_repeats, 1)) + o_ref[...] = (o_scratch_ref[...] * l_inv).astype(o_ref.dtype) + + +def internal_ring_attention_forward( + q: jax.Array, + k: jax.Array, + v: jax.Array, + block_sizes: "custom_splash._BlockSizes", + *, + q_seq_len: int, + kv_seq_len: int, + ring_axis: str, + ring_size: int, + axis_names: tuple[str, ...], + use_base2_exp: bool = True, + use_fixed_m: bool = False, + mk: jax.Array | None = None, + use_experimental_scheduler: bool = False, + vmem_limit_bytes: int | None = None, + mask_value: float = DEFAULT_MASK_VALUE, +) -> jax.Array: + """Single-launch ring attention; the hop is a remote DMA inside the kernel. + + Args: + q: `(batch, num_q_heads, q_seq_padded, head_dim_qk)`, already LOG2E-scaled + by the caller when `use_base2_exp`. + k, v: `(batch, num_kv_heads, kv_seq_padded, head_dim)` local ring shard. + q_seq_len / kv_seq_len: un-padded lengths (grid bounds / ragged tail). + ring_axis / ring_size: mesh axis the KV rotates over and its (static) size. + axis_names: `mesh.axis_names` of the enclosing shard_map, needed to spell a + neighbour as a full mesh-index tuple. + + Returns: + `(batch, num_q_heads, q_seq_len, head_dim_v)`, softmax-normalized. + """ + batch, num_q_heads, _, head_dim_qk = q.shape + num_kv_heads = k.shape[1] + kv_pad = k.shape[2] + head_dim_v = v.shape[-1] + q_heads_per_kv_head = num_q_heads // num_kv_heads + + bq, bkv = block_sizes.block_q, block_sizes.block_kv + bkv_compute = block_sizes.block_kv_compute + bkv_compute_in = block_sizes.block_kv_compute_in + + # Scalar-prefetch operand: mk[0,h] = max_j||k_j|| over EVERY ring shard, + # mk[1,h] = per-head eligibility. A dummy keeps the signature uniform when the + # online path is compiled. + if mk is None: + mk = jnp.zeros((2, num_q_heads), jnp.float32) + + grid_width = (kv_seq_len + bkv - 1) // bkv + grid_height = (q_seq_len + bq - 1) // bq + grid = (batch, num_q_heads, grid_height, grid_width, ring_size) + + # `*_` absorbs the scalar-prefetch operand, which Pallas appends to every + # index_map's argument list once num_scalar_prefetch > 0. + def q_index_map(b, h, i, j, r, *_): + return (b, h, i, 0) + + def out_index_map(b, h, i, j, r, *_): + return (b, h, 0, i) + + in_specs = [ + pl.BlockSpec((None, None, bq, head_dim_qk), q_index_map), + pl.BlockSpec(memory_space=pl.ANY), + pl.BlockSpec(memory_space=pl.ANY), + ] + out_specs = pl.BlockSpec((None, None, head_dim_v, bq), out_index_map) + out_shape = jax.ShapeDtypeStruct((batch, num_q_heads, head_dim_v, q_seq_len), q.dtype) + + scratch_shapes = [ + pltpu.VMEM((2, bkv, head_dim_qk), k.dtype), + pltpu.VMEM((2, bkv, head_dim_v), v.dtype), + pltpu.VMEM((NUM_SUBLANES, bq), jnp.float32), + pltpu.VMEM((NUM_SUBLANES, bq), jnp.float32), + pltpu.VMEM((head_dim_v, bq), jnp.float32), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.DMA((2,)), + pltpu.SemaphoreType.REGULAR, + ] + + out = pl.pallas_call( + functools.partial( + _internal_ring_kernel, + mask_value=mask_value, + grid_width=grid_width, + ring_size=ring_size, + bkv=bkv, + bkv_compute=bkv_compute, + bkv_compute_in=bkv_compute_in, + head_dim_v=head_dim_v, + kv_seq_len=kv_seq_len, + q_heads_per_kv_head=q_heads_per_kv_head, + num_kv_heads=num_kv_heads, + use_base2_exp=use_base2_exp, + use_fixed_m=use_fixed_m, + axis_names=tuple(axis_names), + ring_axis=ring_axis, + ), + grid_spec=pltpu.PrefetchScalarGridSpec( + num_scalar_prefetch=1, + in_specs=in_specs, + out_specs=out_specs, + grid=grid, + scratch_shapes=scratch_shapes, + ), + compiler_params=pltpu.CompilerParams( + # Every dimension is "arbitrary": the in-kernel ring makes grid order + # semantically load-bearing (hop r must follow hop r-1, and the credit + # protocol assumes every rank walks the grid in lockstep), so no + # dimension may be reordered or split across cores. In particular `h` + # cannot be "parallel" the way the stock kernel has it. + dimension_semantics=("arbitrary",) * 5, + flags={"XLA_TPU_FORCE_LP_LLO_SCHEDULER": use_experimental_scheduler}, + disable_bounds_checks=True, + vmem_limit_bytes=vmem_limit_bytes, + collective_id=_COLLECTIVE_ID, + has_side_effects=True, + ), + out_shape=out_shape, + )(mk, q, k.reshape(batch * num_kv_heads, kv_pad, head_dim_qk), v.reshape(batch * num_kv_heads, kv_pad, head_dim_v)) + return jnp.swapaxes(out, 2, 3) + + +def make_internal_ring_attention( + *, + block_sizes: "custom_splash._BlockSizes", + orig_q_seq_len: int, + orig_kv_seq_len: int, + ring_axis: str, + ring_size: int, + axis_names: tuple[str, ...], + use_base2_exp: bool = True, + use_experimental_scheduler: bool = False, + vmem_limit_bytes: int | None = None, + mask_value: float = DEFAULT_MASK_VALUE, + use_fixed_m: bool = False, +): + """Batched `(b, h, s, d) -> (b, h, s, d)` callable. Deliberately NOT vmapped: + the batch axis is a grid dimension, because vmapping a pallas_call that owns + semaphores and a collective_id would duplicate the collective per batch + element.""" + + def _ring(q, k, v, mk=None): + return internal_ring_attention_forward( + q, + k, + v, + block_sizes, + q_seq_len=orig_q_seq_len, + kv_seq_len=orig_kv_seq_len, + ring_axis=ring_axis, + ring_size=ring_size, + axis_names=axis_names, + use_base2_exp=use_base2_exp, + use_fixed_m=use_fixed_m, + mk=mk, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + mask_value=mask_value, + ) + + return _ring diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 7b2ba0df7..0236830b6 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -27,6 +27,7 @@ from maxdiffusion.kernels.splash_attention import splash_attention_mask as tokamax_splash_attention_mask from maxdiffusion.kernels.splash_attention import splash_attention_kernel as tokamax_splash_attention_kernel from maxdiffusion.kernels.splash_attention import ring_attention_kernel as tokamax_ring_attention_kernel +from maxdiffusion.kernels import internal_ring_attention as internal_ring_kernel_mod from maxdiffusion.kernels.splash_attention import base as tokamax_splash_base from einops import rearrange from .. import common_types, max_logging @@ -1321,6 +1322,8 @@ def _ulysses_ring_custom_attention( use_experimental_scheduler: bool = False, bidirectional: bool = False, use_fixed_m: bool = False, + fixed_m_uncond: bool = False, + internal_perm: bool = False, ulysses_attention_chunks: int = 1, ) -> jax.Array: """Hybrid Ulysses + Ring (USP) with the CUSTOM splash kernel on main's mesh. @@ -1355,6 +1358,8 @@ def _ulysses_ring_custom_attention( f"got context_shards={num_context_shards} and ulysses_shards={num_ulysses_shards}." ) num_ring_shards = num_context_shards // num_ulysses_shards + if internal_perm and bidirectional: + raise NotImplementedError("internal-permutation ring does not implement the bidirectional schedule.") query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_context_shards) key, _ = _reshape_data_for_flash(key, heads, num_context_shards) @@ -1418,7 +1423,11 @@ def wrap_ulysses_ring_attention(query, key, value): # 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) + # The internal kernel pins ONE m for the whole ring, so its bound must + # cover keys held by every ring rank -- reduce over the ring axis too. + # (The external path deliberately keeps this per-shard and gathers it + # inside `_custom_ring_attention_forward`, which gates hop by hop.) + mk_all = jax.lax.pmax(kn_local, (ring_axis, ulysses_axis) if internal_perm else 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 = ( @@ -1491,20 +1500,77 @@ def wrap_ulysses_ring_attention(query, key, value): # (2b) Ring (full ppermute over the cross-chip ring axis) with the custom kernel. # bidirectional=True -> wrap-free schedule (streams K/V both directions one hop # at a time), for a non-wrapping ring axis. Selected by attention=ulysses_ring_custom_bidir. - ring_kernel = tokamax_ring_attention_kernel.make_custom_ring_attention( - block_sizes=bsizes, - orig_q_seq_len=query_seq_len, - orig_kv_seq_len=key_seq_len, - use_base2_exp=use_base2_exp, - use_experimental_scheduler=use_experimental_scheduler, - vmem_limit_bytes=vmem_limit_bytes, - ring_axis=ring_axis, - ring_size=num_ring_shards, - bidirectional=bidirectional, - use_fixed_m=use_fixed_m, - fixed_m_norms=fixed_m_norms, - ) - attention_output = jax.vmap(ring_kernel, in_axes=(0, 0, 0))(query, key, value) + if internal_perm and use_fixed_m: + # (2c-fm) Internal ring + fixed m. The gate is device-uniform along the + # ring (qn is pmaxed over it and mk is the global bound), so every rank + # takes the same lax.cond branch and the two pallas_calls stay in step. + # Uniform (all-heads) gating rather than per-head dispatch: a two-body + # ragged last block is what trips the Mosaic scheduler cliff. + qn_max, mk_h = fixed_m_norms + all_fixed = jnp.all(qn_max * mk_h <= custom_splash._FIXED_M_RING_SAFE_BOUND) + mk_arr = jnp.stack([mk_h, jnp.ones_like(mk_h)]) + + def _mk_internal(fixed): + return internal_ring_kernel_mod.make_internal_ring_attention( + block_sizes=bsizes, + orig_q_seq_len=query_seq_len, + orig_kv_seq_len=key_seq_len, + ring_axis=ring_axis, + ring_size=num_ring_shards, + axis_names=internal_mesh.axis_names, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + use_fixed_m=fixed, + ) + + if fixed_m_uncond: + # MEASUREMENT VARIANT. Skips the lax.cond and always takes the fixed + # branch, to isolate the conditional's cost from fixed-m's own. Safe + # only while the gate actually holds -- `all_fixed` is computed above + # and its value is asserted post hoc by comparing this variant's output + # against the guarded one (if the gate ever failed, the output would be + # visibly wrong, not subtly so). Not for production. + del all_fixed + attention_output = _mk_internal(True)(query, key, value, mk_arr) + else: + attention_output = jax.lax.cond( + all_fixed, + lambda: _mk_internal(True)(query, key, value, mk_arr), + lambda: _mk_internal(False)(query, key, value, None), + ) + elif internal_perm: + # (2c) Permutation INSIDE the kernel: one pallas_call for the whole ring, + # the hop is a remote DMA between neighbours' VMEM. No ppermute, no + # per-hop (m, l, o) merge in XLA, no vmap (the batch axis is a grid dim + # because the kernel owns semaphores and a collective_id). + ring_kernel = internal_ring_kernel_mod.make_internal_ring_attention( + block_sizes=bsizes, + orig_q_seq_len=query_seq_len, + orig_kv_seq_len=key_seq_len, + ring_axis=ring_axis, + ring_size=num_ring_shards, + axis_names=internal_mesh.axis_names, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + ) + attention_output = ring_kernel(query, key, value) + else: + ring_kernel = tokamax_ring_attention_kernel.make_custom_ring_attention( + block_sizes=bsizes, + orig_q_seq_len=query_seq_len, + orig_kv_seq_len=key_seq_len, + use_base2_exp=use_base2_exp, + use_experimental_scheduler=use_experimental_scheduler, + vmem_limit_bytes=vmem_limit_bytes, + ring_axis=ring_axis, + ring_size=num_ring_shards, + bidirectional=bidirectional, + use_fixed_m=use_fixed_m, + 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) # (3) Ulysses all-to-all back: sequence -> heads, restoring the layout. @@ -1701,6 +1767,90 @@ def ulysses_ring_custom_kernel(q, k, v, context): ) +@register_kernel("ulysses_ring_custom_iperm") +def ulysses_ring_custom_iperm_kernel(q, k, v, context): + """INTERNAL-permutation variant of ulysses_ring_custom: the ring hop is a + remote DMA issued from inside the Pallas kernel (one launch for the whole + ring, accumulator never leaves VMEM) instead of an XLA `lax.ppermute` plus a + per-hop online-softmax merge. Same USP split as ulysses_ring_custom; U=1 + gives a pure 1D internal ring.""" + return _ulysses_ring_custom_attention( + q, + k * context["scale"], + v, + context["heads"], + context["mesh"], + context["axis_names_q"], + context["axis_names_kv"], + context["flash_block_sizes"], + context["dtype"], + mask_padding_tokens=context["mask_padding_tokens"], + residual_checkpoint_name=context["residual_checkpoint_name"], + attention_mask=context["attention_mask"], + ulysses_shards=context["ulysses_shards"], + use_base2_exp=context.get("use_base2_exp", True), + use_experimental_scheduler=context.get("use_experimental_scheduler", False), + internal_perm=True, + ulysses_attention_chunks=context["ulysses_attention_chunks"], + ) + + +@register_kernel("ulysses_ring_custom_iperm_fixed_m") +def ulysses_ring_custom_iperm_fixed_m_kernel(q, k, v, context): + """Internal-permutation ring + Cauchy-Schwarz fixed m. Simpler than the + external ring's fixed-m: one kernel spans every shard and the accumulator + stays in VMEM, so a single pinned bound is exact for the whole ring -- no + per-hop gate, no LSE merge, no rotating norms.""" + return _ulysses_ring_custom_attention( + q, + k * context["scale"], + v, + context["heads"], + context["mesh"], + context["axis_names_q"], + context["axis_names_kv"], + context["flash_block_sizes"], + context["dtype"], + mask_padding_tokens=context["mask_padding_tokens"], + residual_checkpoint_name=context["residual_checkpoint_name"], + attention_mask=context["attention_mask"], + ulysses_shards=context["ulysses_shards"], + use_base2_exp=context.get("use_base2_exp", True), + use_experimental_scheduler=context.get("use_experimental_scheduler", False), + internal_perm=True, + use_fixed_m=True, + ulysses_attention_chunks=context["ulysses_attention_chunks"], + ) + + +@register_kernel("ulysses_ring_custom_iperm_fixed_m_nocond") +def ulysses_ring_custom_iperm_fixed_m_nocond_kernel(q, k, v, context): + """Measurement-only twin of `ulysses_ring_custom_iperm_fixed_m` with the + eligibility `lax.cond` removed, to separate the conditional's cost (a + [H, S, D] copy between branch buffers) from fixed-m's own kernel win.""" + return _ulysses_ring_custom_attention( + q, + k * context["scale"], + v, + context["heads"], + context["mesh"], + context["axis_names_q"], + context["axis_names_kv"], + context["flash_block_sizes"], + context["dtype"], + mask_padding_tokens=context["mask_padding_tokens"], + residual_checkpoint_name=context["residual_checkpoint_name"], + attention_mask=context["attention_mask"], + ulysses_shards=context["ulysses_shards"], + use_base2_exp=context.get("use_base2_exp", True), + use_experimental_scheduler=context.get("use_experimental_scheduler", False), + internal_perm=True, + use_fixed_m=True, + fixed_m_uncond=True, + ulysses_attention_chunks=context["ulysses_attention_chunks"], + ) + + @register_kernel("ulysses_ring_custom_fixed_m") def ulysses_ring_custom_fixed_m_kernel(q, k, v, context): """fixed-m variant of ulysses_ring_custom: the per-shard custom splash kernel @@ -2470,6 +2620,9 @@ def __init__( "ulysses_ring_custom", "ulysses_ring_custom_fixed_m", "ulysses_ring_custom_bidir", + "ulysses_ring_custom_iperm", + "ulysses_ring_custom_iperm_fixed_m", + "ulysses_ring_custom_iperm_fixed_m_nocond", "ulysses_custom", "ulysses_custom_fixed_m", ) @@ -2492,7 +2645,18 @@ def __init__( elif 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 elif ( - attention_kernel in ("ulysses_ring_custom", "ulysses_ring_custom_bidir", "ulysses_ring_custom_fixed_m") + attention_kernel + in ( + "ulysses_ring_custom", + "ulysses_ring_custom_bidir", + "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_iperm", + "ulysses_ring_custom_iperm_fixed_m", + "ulysses_ring_custom_iperm_fixed_m_nocond", + "ulysses_ring_custom_iperm_fixed_m_nocond", + "ulysses_ring_custom_iperm_fixed_m", + "ulysses_ring_custom_iperm_fixed_m_nocond", + ) and not is_self_attention ): attention_kernel = "ulysses_custom" # plain ulysses (no ring) for cross attention diff --git a/src/maxdiffusion/pyconfig.py b/src/maxdiffusion/pyconfig.py index d1121ca3f..2f5770e0f 100644 --- a/src/maxdiffusion/pyconfig.py +++ b/src/maxdiffusion/pyconfig.py @@ -230,6 +230,9 @@ def user_init(raw_keys): "ulysses_ring_custom", "ulysses_ring_custom_fixed_m", "ulysses_ring_custom_bidir", + "ulysses_ring_custom_iperm", + "ulysses_ring_custom_iperm_fixed_m", + "ulysses_ring_custom_iperm_fixed_m_nocond", } if attention in ulysses_ring_attentions and raw_keys.get("ulysses_shards", -1) <= 0: raise ValueError(f"{attention} requires ulysses_shards to be set from config or command line.") diff --git a/src/maxdiffusion/tests/tile_size_grid_search_test.py b/src/maxdiffusion/tests/tile_size_grid_search_test.py index 3419c142f..2f93662e8 100644 --- a/src/maxdiffusion/tests/tile_size_grid_search_test.py +++ b/src/maxdiffusion/tests/tile_size_grid_search_test.py @@ -86,10 +86,21 @@ def test_bq_fewest_tile_ladder(self): self.assertEqual(bqs[0], 9472) self.assertTrue(all(b % VPU_LANE == 0 for b in bqs)) + def test_bkv_ceiling_is_per_family_and_measured(self): + # Measured at 64 MB, bq=9472: external tops out at 1152, internal at 1408. + # The old single-fraction model capped BOTH at 1024, hiding the internal + # kernel's optimum (9472/1408) -- worth 3-12% e2e. See report/internal_perm. + self.assertEqual(vmem_bkv_ceiling(9472, vmem_bytes=VMEM_64MB, family="external"), 1152) + self.assertEqual(vmem_bkv_ceiling(9472, vmem_bytes=VMEM_64MB, family="internal"), 1408) + self.assertGreater( + vmem_bkv_ceiling(9472, vmem_bytes=VMEM_64MB, family="internal"), + vmem_bkv_ceiling(9472, vmem_bytes=VMEM_64MB, family="external"), + ) + def test_bkv_largest_fits_includes_winner(self): ceil = vmem_bkv_ceiling(9472, vmem_bytes=VMEM_64MB) - bkvs = bkv_candidates(RING_SEQ, k=3, max_block=ceil) - self.assertIn(1024, bkvs) # the measured optimum, largest 256-mult that fits at bq=9472 + bkvs = bkv_candidates(RING_SEQ, k=4, max_block=ceil) + self.assertIn(1024, bkvs) # still offered: 1024 beats 1152 in all 13 measured configs def test_smart_grid_pairs_winner(self): pairs = smart_grid(RING_SEQ, RING_SEQ, vmem_bytes=VMEM_64MB, dtype_bytes=4) @@ -153,9 +164,15 @@ def test_process_measurements_reject_candidate_if_any_host_fails(self): self.assertIsNone(mean_ms) def test_smart_search_picks_measured_winner(self): + # Expectation updated when the VMEM ceiling was corrected. The mock's cost + # model rewards fewer KV blocks (0.9 * n_kv) and caps the compute bonus at + # min(cmp, 1024), so 1280 genuinely beats 1024 under its own physics -- the + # old assertion only held because the 0.65-fraction ceiling never OFFERED + # anything above 1024. It was testing the cap, not the search. res = grid_search(_MockRingBench(), mode="smart", iters=10, log=lambda *a, **k: None) self.assertIsNotNone(res.best) - self.assertEqual((res.best.bq, res.best.bkv), (9472, 1024)) + self.assertEqual(res.best.bq, 9472) + self.assertGreaterEqual(res.best.bkv, 1024) def test_oom_configs_pruned_not_raised(self): # a tiny VMEM budget OOMs the big pairs; search must still return (or None), never raise. diff --git a/src/maxdiffusion/utils/tile_size_grid_search.py b/src/maxdiffusion/utils/tile_size_grid_search.py index 3b7afb90b..9a62aa5a7 100644 --- a/src/maxdiffusion/utils/tile_size_grid_search.py +++ b/src/maxdiffusion/utils/tile_size_grid_search.py @@ -47,6 +47,9 @@ "ulysses_ring_custom", "ulysses_ring_custom_fixed_m", "ulysses_ring_custom_bidir", + "ulysses_ring_custom_iperm", + "ulysses_ring_custom_iperm_fixed_m", + "ulysses_ring_custom_iperm_fixed_m_nocond", }) @@ -132,21 +135,6 @@ def full_axis_candidates( return list(range(min_block, max_block + 1, step)) -def vmem_bq_ceiling( - bkv: int, - *, - vmem_bytes: int, - dtype_bytes: int = 4, - score_fraction: float = 0.65, - align: int = VPU_LANE, -) -> int: - """Approx largest bq whose score tile [bq, bkv]*dtype_bytes fits VMEM (mirror of - `vmem_bkv_ceiling`). Use a SMALL reference bkv so the fewest-tile BQ ladder isn't - over-constrained; big-bq x big-bkv corners are OOM-pruned by the orchestrator.""" - budget = int(vmem_bytes * score_fraction) - return max(align, _floor_to(budget // (bkv * dtype_bytes), align)) - - def bq_candidates( seq_len: int, *, @@ -184,21 +172,183 @@ def bq_candidates( return sorted(set(out), reverse=True) +# --------------------------------------------------------------------------- +# VMEM ceiling model. +# +# Replaces a single `score_tile <= vmem * score_fraction` fraction, which cannot +# be right: VMEM also holds terms that scale with bq INDEPENDENTLY of bkv (the q +# block double-buffered, the fp32 `o` accumulator, the output block). Proof -- +# two tiles with an identical score-tile product: 9472x1024 FITS, 18944x512 OOMs. +# A one-term model gets tuned to the worst corner and under-caps everywhere else. +# +# Fitting `used/bq = a*(4*bkv) + b` to the ok/OOM boundary of ~3,300 measured +# sweep points (45 bracketed bq rungs; brackets the true ceiling on 36 of 45): +# +# external (ulysses_ring_custom*) a=0.973 b=1989 26 rungs +# internal (ulysses_ring_custom_iperm*) a=0.982 b=1295 19 rungs +# +# The families differ because the external ring keeps fp32 online-softmax +# residual windows while the internal-permutation kernel carries (m, l, o) in one +# VMEM scratch. Measured: external OOMs at 9472/1280 (3/3 reps) where internal +# runs it. Ceilings at bq=9472: external ~1152, internal 1408. +# +# Impact measured end to end: the old 0.65 fraction capped bkv at 1024 for +# bq=9472, under-tiling every internal cell by 3-12%; external ONLINE was already +# at its optimum (dense sweep confirmed the old pick 3/3). +# +# CAVEAT -- external + fixed-m. There the 1-block bench disagrees with e2e: dense +# sweeps proposed tiles that measured WORSE end to end in 3/3 cases (1.570 vs +# 1.547, 1.610 vs 1.528, 1.900 vs 1.898 s/step). A correct ceiling widens the +# search space, and for that combination a wider space can surface a tile the +# bench over-rates. Verify any fixed-m winner end to end before trusting it. +# See report/internal_perm/ (sections 4.2b, 4.2c) for the data. +_VMEM_FIT = { # family -> (score scale a, per-bq bytes b) + "external": (0.973, 1989), + "internal": (0.982, 1295), +} + +INTERNAL_PERM_KERNELS = frozenset({ + "ulysses_ring_custom_iperm", + "ulysses_ring_custom_iperm_fixed_m", + "ulysses_ring_custom_iperm_fixed_m_nocond", +}) + + +# Measured ceilings from the sweeps (~3,300 points), MIN across configs so the +# search never proposes a tile that OOMs for some shape. Preferred over the +# fit wherever the exact bq was swept -- a 2-term model provably cannot bracket +# both ends (at bq=9472 internal needs b<=1226 to reach its true 1408 ceiling; at +# bq=18944 it needs b>1367 to avoid over-capping -- contradictory), and the rung +# that matters most, 9472, is exactly where the fit under-caps. +# +# NOTE: these are for a 64 MiB VMEM budget. They are scaled linearly for other +# budgets, which is only approximate -- the per-bq term does not scale with the +# score tile. For a different VMEM size the fit is used instead. +_MEASURED_BKV_CEILING = { # (family, bq) -> min ceiling across configs @64MB + "external": { + 1024: 4096, + 1280: 4096, + 1536: 4096, + 1792: 4096, + 2048: 4096, + 2304: 4096, + 2560: 4096, + 2816: 4096, + 3072: 4096, + 3328: 4096, + 3584: 3968, + 3840: 3840, + 4096: 3328, + 4352: 3200, + 4608: 2944, + 4864: 2816, + 5376: 2560, + 5632: 2304, + 6144: 2048, + 6400: 2048, + 7168: 1664, + 7680: 1536, + 7936: 1536, + 8192: 1408, + 8448: 1408, + 8704: 1280, + 8960: 1280, + 9472: 1152, + 10240: 1024, + 11008: 896, + 11264: 896, + 12032: 768, + 12800: 768, + 14080: 640, + 14336: 512, + 15360: 512, + 15872: 512, + 17152: 384, + 18944: 256, + 22272: 256, + 25344: 256, + }, + "internal": { + 1024: 4096, + 1280: 4096, + 1536: 4096, + 1792: 4096, + 2048: 4096, + 2304: 4096, + 2560: 4096, + 2816: 4096, + 3072: 4096, + 3328: 4096, + 3584: 4096, + 3840: 3968, + 4096: 3456, + 4352: 3328, + 4608: 3072, + 4864: 2944, + 5376: 2688, + 5632: 2432, + 6144: 2304, + 6400: 2176, + 7168: 1920, + 7680: 1664, + 7936: 1664, + 8704: 1536, + 9472: 1408, + 11264: 1024, + 12800: 896, + 14336: 768, + 15872: 640, + 18944: 384, + }, +} + +# The sweeps ran at the kernel's vmem_limit_bytes, which is 64 MiB (67,108,864), +# NOT 64e6. Getting this wrong silently disables the table and falls back to the +# fit -- caught by tile_size_grid_search_test.test_bkv_ceiling_is_per_family. +_MEASURED_VMEM_BYTES = 64 * 1024 * 1024 + + +def vmem_family(attention: str) -> str: + """Which calibrated VMEM fit applies to this attention kernel.""" + return "internal" if attention in INTERNAL_PERM_KERNELS else "external" + + def vmem_bkv_ceiling( bq: int, *, vmem_bytes: int, dtype_bytes: int = 4, - score_fraction: float = 0.65, - align: int = MXU_TILE, + align: int = VPU_LANE, + family: str = "external", ) -> int: - """Approx largest bkv_compute whose score tile [bq, bkv]*dtype_bytes fits the fraction - of VMEM left after the kernel's other resident tiles (ring fp32 residual windows, K/V, - Q, accumulators). APPROXIMATE — OOM-pruning in the orchestrator is the real guard. - Default 0.65 fits the ulysses_ring_custom kernel (measured: bkv 1152 fits, 1280 OOMs - at bq=9472 / 64 MB).""" - budget = int(vmem_bytes * score_fraction) - return max(align, _floor_to(budget // (bq * dtype_bytes), align)) + """Largest bkv that fits at this bq. Returns 0 when bq alone exhausts VMEM. + + align defaults to VPU_LANE (128), not MXU_TILE: the measured optima include + 1152, 1280 and 1408, none of which a 256-aligned ladder can express. + """ + # Prefer the measured ceiling when this exact bq was swept at this VMEM size. + if abs(vmem_bytes - _MEASURED_VMEM_BYTES) < 512 * 1024: + hit = _MEASURED_BKV_CEILING.get(family, {}).get(bq) + if hit is not None: + return _floor_to(hit, align) + a, b = _VMEM_FIT[family] + per_row = vmem_bytes / max(bq, 1) - b + if per_row <= 0: + return 0 + return max(0, _floor_to(int(per_row / (a * dtype_bytes)), align)) + + +def vmem_bq_ceiling( + bkv: int, + *, + vmem_bytes: int, + dtype_bytes: int = 4, + align: int = VPU_LANE, + family: str = "external", +) -> int: + """Largest bq that fits at this bkv (inverse of `vmem_bkv_ceiling`).""" + a, b = _VMEM_FIT[family] + return max(align, _floor_to(int(vmem_bytes / (a * bkv * dtype_bytes + b)), align)) def bkv_candidates(seq_len: int, *, k: int = 3, align: int = VPU_LANE, max_block: int) -> list[int]: @@ -292,7 +442,7 @@ def vmem_bytes(self) -> int: raise NotImplementedError def dtype_bytes(self) -> int: - return 2 # bf16 q/k/v; score tile is f32 (4) -> handled by score_fraction tuning + return 2 # bf16 q/k/v; the score tile is f32 (4) -> see _MEASURED_BKV_CEILING def run( self, @@ -399,10 +549,10 @@ def smart_grid( vmem_bytes: int, dtype_bytes: int = 4, k_bq: int = 3, - k_bkv: int = 3, + k_bkv: int = 4, spread_bq: int = 2, - score_fraction: float = 0.65, min_bkv_ref: int = 1024, + family: str = "external", ) -> list[tuple[int, int]]: """Nested candidate pairs: BQ = VMEM-capped fewest-tile ladder + spread; then for EACH bq, BKV = largest-that-fits at that bq (so bkv is VMEM-correct for its partner, not globally). @@ -413,21 +563,11 @@ def smart_grid( so the fewest-tile ladder starts at the single-tile end (which OOMs for a large per-shard seq) and the feasible moderate-BQ optimum (e.g. bq=9472 at seq 37800) falls in the ladder's gap. """ - bq_cap = vmem_bq_ceiling( - min_bkv_ref, - vmem_bytes=vmem_bytes, - dtype_bytes=dtype_bytes, - score_fraction=score_fraction, - ) + bq_cap = vmem_bq_ceiling(min_bkv_ref, vmem_bytes=vmem_bytes, dtype_bytes=dtype_bytes, family=family) bqs = bq_candidates(q_seq, k=k_bq, spread=spread_bq, max_block=bq_cap) pairs: list[tuple[int, int]] = [] for bq in bqs: - bkv_cap = vmem_bkv_ceiling( - bq, - vmem_bytes=vmem_bytes, - dtype_bytes=dtype_bytes, - score_fraction=score_fraction, - ) + bkv_cap = vmem_bkv_ceiling(bq, vmem_bytes=vmem_bytes, dtype_bytes=dtype_bytes, family=family) for bkv in bkv_candidates(kv_seq, k=k_bkv, max_block=bkv_cap): pairs.append((bq, bkv)) return pairs @@ -461,8 +601,11 @@ def grid_search( winner (lowest mean_ms among status=='ok'). `mode`: 'smart' (candidate ladders) | 'full'. """ q_seq, kv_seq = bench.tiled_seq_lens() + # The VMEM fit is per kernel family (see `_VMEM_FIT`); read the attention name + # off the bench when it exposes one, else assume the external ring. + family = vmem_family(getattr(bench, "_attention", "") or "") if mode == "smart": - pairs = smart_grid(q_seq, kv_seq, vmem_bytes=bench.vmem_bytes(), dtype_bytes=4, k_bq=k, k_bkv=k) + pairs = smart_grid(q_seq, kv_seq, vmem_bytes=bench.vmem_bytes(), dtype_bytes=4, k_bq=k, k_bkv=max(k, 4), family=family) elif mode == "full": log( "Warning: tile_search mode is 'full', not 'smart' -- this is an exhaustive O(N^2) 2D BQ x BKV" @@ -473,7 +616,10 @@ def grid_search( else: raise ValueError(f"mode must be 'smart' or 'full', got {mode!r}") - log(f"[tile-search] {bench.label}: q_seq={q_seq} kv_seq={kv_seq} mode={mode} " f"-> {len(pairs)} configs (iters={iters})") + log( + f"[tile-search] {bench.label}: q_seq={q_seq} kv_seq={kv_seq} mode={mode} " + f"family={family} -> {len(pairs)} configs (iters={iters})" + ) results: list[BenchResult] = [] for i, (bq, bkv) in enumerate(pairs, 1): if jax.process_count() > 1: diff --git a/src/maxdiffusion/utils/wan_block_benchmark.py b/src/maxdiffusion/utils/wan_block_benchmark.py index 40667f88a..1beb246dd 100644 --- a/src/maxdiffusion/utils/wan_block_benchmark.py +++ b/src/maxdiffusion/utils/wan_block_benchmark.py @@ -76,6 +76,9 @@ _RING_VARIANTS = { "ulysses_ring_custom", "ulysses_ring_custom_bidir", + "ulysses_ring_custom_iperm", + "ulysses_ring_custom_iperm_fixed_m", + "ulysses_ring_custom_iperm_fixed_m_nocond", "tokamax_ring", "ring", }