diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py index a30c4d633..7cfba422d 100644 --- a/src/maxdiffusion/generate_wan.py +++ b/src/maxdiffusion/generate_wan.py @@ -353,7 +353,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): s0 = time.perf_counter() # Restore original profiler setting for the profiling run config.get_keys()["enable_profiler"] = original_enable_profiler - if max_utils.profiler_enabled(config): + if original_enable_profiler: # Injecting user requested XLA tracing flags xla_flags = os.environ.get("XLA_FLAGS", "") new_flags = "--xla_enable_mxu_trace=true --xla_jf_dump_llo_html=true --xla_tpu_enable_llo_profiling=true" diff --git a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel_test.py b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel_test.py index 5c0b7d189..78d168a6c 100644 --- a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel_test.py +++ b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel_test.py @@ -160,5 +160,87 @@ def ring_attn(ring_kernel, q, k, v, segment_ids): self._assert_allclose(dv, dv_ref, rtol=1e-2, atol=1e-2) +class RingAttentionHeadsPerTileTest(test_utils.SplashAttentionTestCase): + """`heads_per_tile` (multi-head-per-tile) invariance for ring attention. + + heads_per_tile is a pure tiling/scheduling choice for the forward kernel: for a + fixed input, running with heads_per_tile=N must produce the same output as + heads_per_tile=1. This guards against the block-index class of bug (a wrong + head-tile mapping compiles and runs but silently returns garbage). + """ + + def setUp(self): + if jax.default_backend() != "tpu": + self.skipTest("Multi-head-per-tile ring attention runs on TPU.") + super().setUp() + + @parameterized.product( + heads_per_tile=[2, 4], + head_dim=[128], + dtype=[jnp.bfloat16], + ) + def test_heads_per_tile_matches_single_head(self, heads_per_tile, head_dim, dtype): + ring_size = 2 + num_heads = 4 # MHA (num_q_heads == num_kv_heads); divisible by heads_per_tile. + if len(jax.devices()) < ring_size: + self.skipTest(f"This test requires {ring_size} devices, but has only {len(jax.devices())}.") + + ring_axis = "ring" + devices = np.asarray(jax.devices()[:ring_size]).reshape(1, ring_size) + mesh = jax.sharding.Mesh(devices, ("heads", ring_axis)) + seq_len = 1024 * ring_size + + k1, k2, k3 = random.split(random.key(0), 3) + scale = head_dim**-0.5 + q = random.normal(k1, (num_heads, seq_len, head_dim), dtype=dtype) * scale + k = random.normal(k2, (num_heads, seq_len, head_dim), dtype=dtype) * scale + v = random.normal(k3, (num_heads, seq_len, head_dim), dtype=dtype) * scale + + # The mhpt fast path supports full MHA + static FullMask + HEAD_DIM_MINOR only. + mask = mask_lib.FullMask(_shape=(seq_len, seq_len)) + q_spec = P(None, ring_axis, None) + kv_spec = q_spec + + def run(hpt): + config = splash.SplashConfig.get_default() + config = dataclasses.replace( + config, + use_base2_exp=False, + fuse_reciprocal=True, + heads_per_tile=hpt, + ) + ring_kernel = ring_attention_kernel.make_ring_attention( + mask, + is_mqa=False, + ring_axis=ring_axis, + config=config, + save_residuals=False, + q_seq_shards=ring_size, + kv_seq_shards=ring_size, + ) + kernel_spec = ring_kernel.manual_sharding_spec() + + @partial( + jax.shard_map, + mesh=mesh, + in_specs=(kernel_spec, q_spec, kv_spec, kv_spec, None), + out_specs=q_spec, + check_vma=False, + ) + def ring_attn(ring_kernel, q, k, v, segment_ids): + return ring_kernel(q, k, v, segment_ids) + + return ring_attn(ring_kernel, q, k, v, None) + + out_ref = run(1) # baseline: single head per tile (flash_attention_kernel) + out_mhpt = run(heads_per_tile) # multi-head-per-tile (flash_attention_kernel_mhpt) + + # Pure tiling => numerically equivalent to the single-head-per-tile baseline. + # The ring mesh is jax.devices()[:ring_size] (all on process 0), so the + # outputs are only addressable there; use the multi-controller-safe compare + # so this passes on every host, not just the owner. + self.assert_allclose_mcjax(out_mhpt, out_ref, rtol=5e-3, atol=5e-3) + + if __name__ == "__main__": absltest.main() diff --git a/src/maxdiffusion/kernels/splash_attention/splash_attention_kernel.py b/src/maxdiffusion/kernels/splash_attention/splash_attention_kernel.py index 77b248477..40a98fc8a 100644 --- a/src/maxdiffusion/kernels/splash_attention/splash_attention_kernel.py +++ b/src/maxdiffusion/kernels/splash_attention/splash_attention_kernel.py @@ -155,6 +155,11 @@ class SplashConfig: dq_reduction_steps: int | None = None # An experimental scheduler that sometimes produces better softmax overlap. use_experimental_scheduler: bool = False + heads_per_tile: int = 1 + # Overrides Mosaic's scoped-VMEM budget for the kernel (None = compiler default). + # heads_per_tile > 1 scales the m/l/o scratch linearly, so larger tiles may need + # this raised above the default ceiling to compile. + vmem_limit_bytes: int | None = None def __post_init__(self): if self.block_kv_compute is None: @@ -162,6 +167,8 @@ def __post_init__(self): if self.block_kv_dkv_compute is None: object.__setattr__(self, "block_kv_dkv_compute", self.block_kv_dkv) + if self.heads_per_tile < 1: + raise ValueError(f"Invalid heads_per_tile: {self.heads_per_tile}, expected >= 1.") if self.dq_reduction_steps is not None and self.dq_reduction_steps != 3: raise ValueError(f"Invalid dq_reduction_steps: {self.dq_reduction_steps}, only 3 or" " None are supported.") if not self.use_fused_bwd_kernel: @@ -482,6 +489,131 @@ def end(): max_logits_ref[...] = m.astype(max_logits_ref.dtype) +def flash_attention_kernel_mhpt( + # Prefetched inputs + active_rows_ref, + active_cols_ref, + mask_next_ref, + bounds_start_ref, + bounds_end_ref, + block_mask_ref, + # Inputs + q_ref, + k_ref, + v_ref, + q_segment_ids_ref, + kv_segment_ids_ref, + sinks_ref, + mask_ref, + q_sequence_ref, + max_logit_value_ref, + # Outputs + o_ref, + logsumexp_ref, + l_linear_ref, + max_logits_ref, + # Scratch + m_scratch_ref, + l_scratch_ref, + o_scratch_ref, + *, + mask_value: float, + kv_steps: int, + bq: int, + bkv: int, + bkv_compute: int, + head_dim_v: int, + heads_per_tile: int, + config: SplashConfig, +): + """Multi-head-per-tile (mhpt) forward kernel for ring splash attention. + + Processes `heads_per_tile` heads per grid program (a Python loop over `h_local`) + instead of one, amortizing per-tile launch/prefetch/pipelining overhead. Each head + keeps its own online-softmax scratch slice `(heads_per_tile, bq, ...)` and, on the + last kv step, writes the raw residuals (`o`, `l_linear`, `max_logits`) that the ring + wrapper merges across shards. + + This is a lean, dedicated kernel rather than a branch inside `flash_attention_kernel` + on purpose: mhpt supports only a narrow regime (full MHA, static FullMask, + HEAD_DIM_MINOR, no sinks / max-logit / soft-cap), so keeping it separate avoids adding + a head loop crossed with all of `flash_attention_kernel`'s branches to the hot path. + Mirrors `custom_splash_attention._flash_attention_kernel_mhpt` (the non-ring mhpt path). + Supported regime is enforced by guards in `_splash_attention_forward_ring_raw`. + """ + del active_rows_ref, active_cols_ref, mask_next_ref + del bounds_start_ref, bounds_end_ref, block_mask_ref + del sinks_ref, mask_ref, q_sequence_ref, max_logit_value_ref + + grid_idx = pl.program_id(1) + j = grid_idx % kv_steps + should_initialize = j == 0 + should_write = j == kv_steps - 1 + head_dim_v_repeats = pl.cdiv(head_dim_v, NUM_LANES) + exp = jnp.exp2 if config.use_base2_exp else jnp.exp + + @pl.when(should_initialize) + def init(): + o_scratch_ref[...] = jnp.zeros_like(o_scratch_ref) + m_scratch_ref[...] = jnp.full_like(m_scratch_ref, mask_value) + l_scratch_ref[...] = jnp.zeros_like(l_scratch_ref) + + def body(kv_compute_index, _): + slice_k = pl.ds(kv_compute_index * bkv_compute, bkv_compute) + bkv_repeats, rem = divmod(bkv_compute, NUM_LANES) + if rem != 0: + raise NotImplementedError(f"{bkv_compute=} should be a multiple of {NUM_LANES}") + + for h_local in range(heads_per_tile): + m_prev = m_scratch_ref[h_local] + l_prev = l_scratch_ref[h_local] + q = q_ref[h_local] + if config.use_base2_exp: + q *= LOG2E + + qk = lax.dot_general(q, k_ref[h_local, slice_k, :], NT_DIM_NUMBERS, preferred_element_type=jnp.float32) + qk = _apply_mask_and_soft_cap( + qk, + mask_value, + None, + None, + q_segment_ids_ref, + kv_segment_ids_ref, + attn_logits_soft_cap=None, + k_slice=slice_k, + k_offset=j * bkv + kv_compute_index * bkv_compute, + bq=bq, + mask_function=None, + has_partial_mask=False, + ) + + m_curr = qk.max(axis=-1)[:, None] # pytype: disable=attribute-error + m_next = jnp.maximum(m_prev, m_curr) + s_curr = exp(qk - jnp.tile(m_next, (1, bkv_repeats))) + l_curr = jax.lax.broadcast_in_dim(s_curr.sum(axis=-1), l_prev.shape, (0,)) + alpha = exp(m_prev - m_next) + m_scratch_ref[h_local] = m_next + l_scratch_ref[h_local] = l_curr + alpha * l_prev + + o_curr = lax.dot_general(s_curr, v_ref[h_local, slice_k, :], NN_DIM_NUMBERS) + alpha_o = jnp.tile(alpha, (1, head_dim_v_repeats)) + alpha_o = alpha_o[..., : o_scratch_ref.shape[-1]] + o_scratch_ref[h_local] = alpha_o * o_scratch_ref[h_local] + o_curr + + assert bkv % bkv_compute == 0 + lax.fori_loop(0, k_ref.shape[1] // bkv_compute, body, None, unroll=True) + + @pl.when(should_write) + def end(): + for h_local in range(heads_per_tile): + l = l_scratch_ref[h_local] + m = m_scratch_ref[h_local] + o_ref[h_local] = o_scratch_ref[h_local].astype(o_ref.dtype) + if l_linear_ref is not None: + l_linear_ref[h_local] = l.astype(l_linear_ref.dtype) + if max_logits_ref is not None: + max_logits_ref[h_local] = m.astype(max_logits_ref.dtype) + def _div(dividend: int, divisor: int): if divisor == 1: return dividend @@ -793,6 +925,7 @@ def _fwd_cost_estimate( compiler_params=pltpu.CompilerParams( dimension_semantics=("parallel", "arbitrary"), flags={"XLA_TPU_FORCE_LP_LLO_SCHEDULER": (config.use_experimental_scheduler)}, + vmem_limit_bytes=config.vmem_limit_bytes, ), out_shape=out_shapes, name=kernel_name, @@ -918,6 +1051,9 @@ def _splash_attention_forward_ring_raw( kv_steps = kv_seq_len // bkv q_heads_per_kv_head = num_q_heads // num_kv_heads dynamic_grid = mask_info.active_rows is not None + heads_per_tile = config.heads_per_tile + if num_q_heads % heads_per_tile != 0: + raise ValueError(f"num_q_heads {num_q_heads} must be divisible by heads_per_tile {heads_per_tile}.") if segment_ids is not None: assert isinstance(segment_ids.q, jax.Array) @@ -940,6 +1076,37 @@ def _splash_attention_forward_ring_raw( q_layout = config.q_layout k_layout = config.k_layout v_layout = config.v_layout + use_heads_per_tile = heads_per_tile > 1 + if use_heads_per_tile: + if dynamic_grid: + raise NotImplementedError("heads_per_tile > 1 is only implemented for static ring attention grids.") + if is_mqa or q_heads_per_kv_head != 1: + # The mhpt head-dim BlockSpec tiles Q and K/V with the same head_block + # size, so the K/V block index must equal the Q head-tile index. That + # only holds for full MHA (num_q_heads == num_kv_heads); GQA/MQA would + # need a distinct K/V tiling. Matches custom_splash_attention.py's assert. + raise NotImplementedError("heads_per_tile > 1 is only implemented for MHA ring attention (num_q_heads == num_kv_heads).") + if ( + mask_info.block_mask is not None + or mask_info.partial_mask_blocks is not None + or mask_info.q_sequence is not None + or mask_function is not None + ): + raise NotImplementedError("heads_per_tile > 1 is only implemented for static FullMask ring attention.") + if sinks is not None: + raise NotImplementedError("heads_per_tile > 1 does not support attention sinks.") + if max_logit_value is not None or config.max_logit_const is not None: + raise NotImplementedError("heads_per_tile > 1 does not support max logit bounds.") + if config.attn_logits_soft_cap is not None: + raise NotImplementedError("heads_per_tile > 1 does not support attention logit soft cap.") + if ( + q_layout != QKVLayout.HEAD_DIM_MINOR + or k_layout != QKVLayout.HEAD_DIM_MINOR + or v_layout != QKVLayout.HEAD_DIM_MINOR + ): + raise NotImplementedError("heads_per_tile > 1 only supports HEAD_DIM_MINOR Q/K/V layouts.") + head_programs = num_q_heads // heads_per_tile + head_block = heads_per_tile if use_heads_per_tile else None def unravel(f): def index_map(h, grid_idx, rows_ref, cols_ref, *_): @@ -953,6 +1120,12 @@ def index_map(h, grid_idx, rows_ref, cols_ref, *_): return index_map + # `h` is program_id(0), which ranges over head *tiles* (grid dim 0 is + # num_q_heads // heads_per_tile). The head-dim BlockSpec block size is + # `head_block` (= heads_per_tile when tiling), so the block index is simply + # `h`: Pallas multiplies it by the block size to get the element offset + # (h * heads_per_tile). This matches the reference mhpt kernel in + # custom_splash_attention.py. def create_kv_index_map(layout): def index_map(h, i, j): del i @@ -975,13 +1148,13 @@ def mask_index_map(h, grid_idx, rows_ref, cols_ref, mask_next_ref=None, *_): kv_segment_ids_index_map = unravel(lambda h, i, j: (0, j)) in_specs = [ - pl.BlockSpec(from_head_minor((None, bq, head_dim_qk), q_layout), q_index_map), + pl.BlockSpec(from_head_minor((head_block, bq, head_dim_qk), q_layout), q_index_map), pl.BlockSpec( - from_head_minor((bkv, head_dim_qk) if is_mqa else (None, bkv, head_dim_qk), k_layout), + from_head_minor((bkv, head_dim_qk) if is_mqa else (head_block, bkv, head_dim_qk), k_layout), k_index_map, ), pl.BlockSpec( - from_head_minor((bkv, head_dim_v) if is_mqa else (None, bkv, head_dim_v), v_layout), + from_head_minor((bkv, head_dim_v) if is_mqa else (head_block, bkv, head_dim_v), v_layout), v_index_map, ), ] @@ -1046,15 +1219,17 @@ def mask_index_map(h, grid_idx, rows_ref, cols_ref, mask_next_ref=None, *_): jax.ShapeDtypeStruct((num_q_heads, q_seq_len, NUM_LANES), jnp.float32), ] out_specs = [ - pl.BlockSpec((None, bq, head_dim_v), out_index_map), + pl.BlockSpec((head_block, bq, head_dim_v), out_index_map), None, - pl.BlockSpec((None, bq, NUM_LANES), logsumexp_index_map), - pl.BlockSpec((None, bq, NUM_LANES), logsumexp_index_map), + pl.BlockSpec((head_block, bq, NUM_LANES), logsumexp_index_map), + pl.BlockSpec((head_block, bq, NUM_LANES), logsumexp_index_map), ] kernel_name = ( f"{get_kernel_name(is_mqa=is_mqa, save_residuals=True, is_segmented=segment_ids is not None, phase='fwd')}_ring_raw" ) + if use_heads_per_tile: + kernel_name = f"{kernel_name}_hpt{heads_per_tile}" metadata = {"xprof_metadata": json.dumps(dataclasses.asdict(config))} vmem_inputs = [q, k, v, q_segment_ids, kv_segment_ids, mask_info.partial_mask_blocks] @@ -1087,40 +1262,58 @@ def _fwd_cost_estimate( if dynamic_grid: num_active_blocks = mask_info.num_active_blocks[0] - grid = (num_q_heads, num_active_blocks) + grid = (head_programs, num_active_blocks) is_empty_attention_block = num_active_blocks == 0 else: - grid = (num_q_heads, kv_steps * (q_seq_len // bq)) + grid = (head_programs, kv_steps * (q_seq_len // bq)) is_empty_attention_block = False + # Dispatch to the dedicated mhpt kernel when tiling >1 head per program. The mhpt + # path is a guarded branch *inside* this forward (rather than a separate forward like + # custom_splash_attention._splash_attention_forward_mhpt) so it reuses all the shared + # setup below/above -- index maps, in/out specs, mask & segment plumbing -- and only + # the kernel fn, scratch shapes, and head-dim block size differ. + kernel = flash_attention_kernel_mhpt if use_heads_per_tile else flash_attention_kernel + kernel_kwargs = { + "mask_value": mask_value, + "kv_steps": kv_steps, + "bq": bq, + "bkv": bkv, + "bkv_compute": bkv_compute, + "head_dim_v": head_dim_v, + "config": config, + } + if use_heads_per_tile: + kernel_kwargs["heads_per_tile"] = heads_per_tile + else: + kernel_kwargs["fuse_reciprocal"] = False + kernel_kwargs["mask_function"] = mask_function + scratch_shapes = [ + pltpu.VMEM((bq, NUM_LANES), jnp.float32), + pltpu.VMEM((bq, NUM_LANES), jnp.float32), + pltpu.VMEM((bq, head_dim_v), jnp.float32), + ] + if use_heads_per_tile: + scratch_shapes = [ + pltpu.VMEM((heads_per_tile, bq, NUM_LANES), jnp.float32), + pltpu.VMEM((heads_per_tile, bq, NUM_LANES), jnp.float32), + pltpu.VMEM((heads_per_tile, bq, head_dim_v), jnp.float32), + ] + with jax.named_scope(kernel_name): all_out = pl.pallas_call( - partial( - flash_attention_kernel, - mask_value=mask_value, - kv_steps=kv_steps, - bq=bq, - bkv=bkv, - bkv_compute=bkv_compute, - head_dim_v=head_dim_v, - fuse_reciprocal=False, - config=config, - mask_function=mask_function, - ), + partial(kernel, **kernel_kwargs), grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=6, in_specs=in_specs, out_specs=out_specs, grid=grid, - scratch_shapes=[ - pltpu.VMEM((bq, NUM_LANES), jnp.float32), - pltpu.VMEM((bq, NUM_LANES), jnp.float32), - pltpu.VMEM((bq, head_dim_v), jnp.float32), - ], + scratch_shapes=scratch_shapes, ), compiler_params=pltpu.CompilerParams( dimension_semantics=("parallel", "arbitrary"), flags={"XLA_TPU_FORCE_LP_LLO_SCHEDULER": (config.use_experimental_scheduler)}, + vmem_limit_bytes=config.vmem_limit_bytes, ), out_shape=out_shapes, name=kernel_name, diff --git a/src/maxdiffusion/kernels/splash_attention/splash_attention_test_utils.py b/src/maxdiffusion/kernels/splash_attention/splash_attention_test_utils.py index 6622008e3..f10902785 100644 --- a/src/maxdiffusion/kernels/splash_attention/splash_attention_test_utils.py +++ b/src/maxdiffusion/kernels/splash_attention/splash_attention_test_utils.py @@ -16,6 +16,7 @@ import unittest from absl.testing import parameterized import jax +from jax.experimental import multihost_utils import jax.numpy as jnp import numpy as np @@ -76,6 +77,43 @@ def _assert_allclose(self, x, y, **kwargs): self.assertTupleEqual(x.shape, y.shape) np.testing.assert_allclose(x, y, **kwargs) + def assert_allclose_mcjax(self, x, y, *, rtol, atol): + """`allclose` that is safe under multi-controller (multi-host) JAX. + + Some tests build their device mesh from a subset of the global devices -- + e.g. `jax.devices()[:ring_size]`, which all live on process 0 -- so the + result `jax.Array`s are only addressable on that one process. Pulling them + to host with `np.testing.assert_allclose` (as `_assert_allclose` does) + raises `RuntimeError: ... spans non-addressable devices` on every *other* + process, failing the test on all but one host. + + Instead, evaluate the comparison on-device (works on all processes, no host + fetch), read the scalar verdict only on the owning process (process 0, which + holds `jax.devices()[:ring_size]`), and broadcast it to every process with a + single collective. Every process runs the same two `broadcast_one_to_all` + calls in the same order, so there is no collective-participation mismatch / + deadlock. The result is identical on all hosts and genuinely reflects the + owner's computation. + + Only use this for tests whose mesh is a subset of one process's devices; for + fully-sharded (every-process-owns-a-shard) arrays use `_assert_allclose`. + """ + if x.dtype == np.dtype(jnp.bfloat16): + x = x.astype(jnp.float32) + if y.dtype == np.dtype(jnp.bfloat16): + y = y.astype(jnp.float32) + self.assertTupleEqual(x.shape, y.shape) + ok = jnp.all(jnp.abs(x - y) <= atol + rtol * jnp.abs(y)) + max_err = jnp.max(jnp.abs(x - y)) + is_owner = jax.process_index() == 0 + local_ok = np.asarray(ok) if is_owner else np.array(True) + local_err = np.asarray(max_err) if is_owner else np.array(0.0, np.float32) + global_err = float(multihost_utils.broadcast_one_to_all(local_err)) + self.assertTrue( + bool(multihost_utils.broadcast_one_to_all(local_ok)), + f"arrays differ: max abs err {global_err:.3e} exceeds rtol={rtol} atol={atol}", + ) + def create_segment_ids(seq_len: int, num_breaks: int = 2) -> base.SegmentIds: break_indices = np.random.choice(range(1, seq_len), num_breaks, replace=False) diff --git a/src/maxdiffusion/max_utils.py b/src/maxdiffusion/max_utils.py index 1ab5407ac..b365ea965 100644 --- a/src/maxdiffusion/max_utils.py +++ b/src/maxdiffusion/max_utils.py @@ -633,11 +633,29 @@ class CustomFlashBlockSizes: vmem_limit_bytes: int | None = None +@dataclasses.dataclass(frozen=True) +class TokamaxRingFlashBlockSizes: + """Hashable carrier for tokamax ring block sizes plus heads_per_tile.""" + + block_q: int | None = None + block_kv: int | None = None + block_kv_compute: int | None = None + block_q_dkv: int | None = None + block_kv_dkv: int | None = None + block_kv_dkv_compute: int | None = None + block_q_dq: int | None = None + block_kv_dq: int | None = None + use_fused_bwd_kernel: bool | None = None + heads_per_tile: int | None = None + vmem_limit_bytes: int | None = None + + def get_flash_block_sizes(config): """Create custom flash attention BlockSizes.""" flash_block_sizes = None if len(config.flash_block_sizes.keys()) > 0: attention_is_tokamax = "tokamax" in config.attention or config.attention == "ulysses_ring" + attention_uses_tokamax_ring = config.attention in ("tokamax_ring", "ulysses_ring") user_block_sizes: Dict[str, int] = config.flash_block_sizes # The custom splash kernel reads flash_block_sizes via getattr and needs # fields the JAX BlockSizes dataclass cannot hold. Return a frozen, hashable @@ -651,6 +669,20 @@ def get_flash_block_sizes(config): heads_per_tile=user_block_sizes.get("heads_per_tile"), vmem_limit_bytes=user_block_sizes.get("vmem_limit_bytes"), ) + if attention_uses_tokamax_ring and "heads_per_tile" in user_block_sizes: + return TokamaxRingFlashBlockSizes( + block_q=user_block_sizes.get("block_q_dkv", user_block_sizes["block_kv"]), + block_kv=user_block_sizes["block_kv"], + block_kv_compute=user_block_sizes["block_kv_compute"], + block_q_dkv=user_block_sizes["block_q_dkv"], + block_kv_dkv=user_block_sizes["block_kv_dkv"], + block_kv_dkv_compute=user_block_sizes["block_kv_dkv_compute"], + block_q_dq=None, + block_kv_dq=None, + use_fused_bwd_kernel=True, + heads_per_tile=user_block_sizes.get("heads_per_tile"), + vmem_limit_bytes=user_block_sizes.get("vmem_limit_bytes"), + ) if attention_is_tokamax: max_logging.log( "Tokamax kernel specified, Note: Tokamax only supports fused backward kernel." diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index edc9f4f7b..13d08ec80 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -344,6 +344,8 @@ def convert_to_tokamax_splash_config( max_logit_const=max_logit_const, interpret=interpret, dq_reduction_steps=dq_reduction_steps, + heads_per_tile=getattr(block_sizes, "heads_per_tile", None) or 1, + vmem_limit_bytes=getattr(block_sizes, "vmem_limit_bytes", None), )