Remove dead TorchAX wrapper and tpu_custom_attention

Since maxdiffusion internally uses attention_flax.py (which calls the core splash kernels directly via make_splash_mha), the high-level tpu_custom_attention orchestration and make_custom_splash_sdpa TorchAX wrapper are effectively dead code in this repository. Removing them to keep custom_splash_attention.py strictly focused on Pallas kernel definitions.

Author: Perseus14

src/maxdiffusion/kernels/custom_splash_attention.py

@@ -17,27 +17,19 @@
 """Custom Pallas flash attention kernel for TPU."""
 
 import functools
-import math
 
 import jax
 import jax.numpy as jnp
 import numpy as np
 from jax import lax
 from jax.experimental import pallas as pl
 from jax.experimental.pallas import tpu as pltpu
-from jax.experimental.shard_map import shard_map
-from jax.sharding import PartitionSpec as P
 
 DEFAULT_MASK_VALUE = -0.7 * float(np.finfo(np.dtype("float32")).max)
 NUM_LANES = 128
 NUM_SUBLANES = 8
 NT_DIM_NUMBERS = (((1,), (1,)), ((), ()))
 
-# Default block sizes (tuned for 720p Wan2.1 on v6e/v7x)
-DEFAULT_BQSIZE = 3328
-DEFAULT_BKVSIZE = 2816
-# Cranked up to 1024 for massive MXU throughput
-DEFAULT_BKVCOMPUTESIZE = 1024
 # Kept at 256 to protect VPU registers (V1 Optimization)
 DEFAULT_BKVCOMPUTEINSIZE = 256
 
@@ -569,227 +561,3 @@ def _splash_attention(q, k, v):
   return _splash_attention
 
 
-# ---------------------------------------------------------------------------
-# High-level attention function with shard_map
-# TODO: Move `tpu_custom_attention` and `make_custom_splash_sdpa` 
-# out of this file and into `attention_flax.py` where other orchestration 
-# functions live. This file should be strictly Pallas kernel definitions.
-# ---------------------------------------------------------------------------
-
-def tpu_custom_attention(
-    query,
-    key,
-    value,
-    mesh,
-    *,
-    scale=None,
-    block_q=None,
-    block_kv=None,
-    block_kv_compute=None,
-    block_kv_compute_in=None,
-    heads_per_tile=None,
-    use_base2_exp=True,
-    use_experimental_scheduler=False,
-    vmem_limit_bytes=None,
-    flash_block_sizes=None,
-):
-  _LOG2_E = 1.44269504
-  num_heads = query.shape[1]
-
-  if flash_block_sizes is not None:
-    block_q = flash_block_sizes.get("block_q", block_q)
-    block_kv = flash_block_sizes.get("block_kv", block_kv)
-    block_kv_compute = flash_block_sizes.get("block_kv_compute", block_kv_compute)
-    block_kv_compute_in = flash_block_sizes.get("block_kv_compute_in", block_kv_compute_in)
-    heads_per_tile = flash_block_sizes.get("heads_per_tile", heads_per_tile)
-    vmem_limit_bytes = flash_block_sizes.get("vmem_limit_bytes", vmem_limit_bytes)
-
-  block_q = block_q if block_q is not None else DEFAULT_BQSIZE
-  block_kv = block_kv if block_kv is not None else DEFAULT_BKVSIZE
-  block_kv_compute = block_kv_compute if block_kv_compute is not None else DEFAULT_BKVCOMPUTESIZE
-  block_kv_compute_in = block_kv_compute_in if block_kv_compute_in is not None else DEFAULT_BKVCOMPUTEINSIZE
-  heads_per_tile = heads_per_tile if heads_per_tile is not None else 1
-
-  def _attention_on_slices(q, k, v):
-    scale_factor = 1.0 / math.sqrt(q.shape[-1]) if scale is None else scale
-    if use_base2_exp:
-      q = q * scale_factor * _LOG2_E
-    else:
-      q = q * scale_factor
-
-    def _pad_to_multiple(x, multiple, axis):
-      seq_len = x.shape[axis]
-      pad_len = (multiple - seq_len % multiple) % multiple
-      if pad_len == 0:
-        return x, seq_len
-      pad_width = [(0, 0)] * x.ndim
-      pad_width[axis] = (0, pad_len)
-      return jnp.pad(x, pad_width), seq_len
-
-    def _kernel_3d(q_3d, k_3d, v_3d):
-      q_orig_len = q_3d.shape[1]
-      kv_orig_len = k_3d.shape[1]
-
-      q_3d_padded, _ = _pad_to_multiple(q_3d, block_q, axis=1)
-      k_3d_padded, _ = _pad_to_multiple(k_3d, block_kv, axis=1)
-      v_3d_padded, _ = _pad_to_multiple(v_3d, block_kv, axis=1)
-
-      padded_q_seq_len = q_3d_padded.shape[1]
-      padded_kv_seq_len = k_3d_padded.shape[1]
-
-      bsizes = _BlockSizes(
-          block_q=min(block_q, padded_q_seq_len),
-          block_kv=min(block_kv, padded_kv_seq_len),
-          block_kv_compute=min(block_kv_compute, padded_kv_seq_len),
-      )
-      splash_kernel = make_splash_mha(
-          block_sizes=bsizes,
-          bkv_compute_in=block_kv_compute_in,
-          orig_q_seq_len=q_orig_len,
-          orig_kv_seq_len=kv_orig_len,
-          heads_per_tile=heads_per_tile,
-          use_base2_exp=use_base2_exp,
-          use_experimental_scheduler=use_experimental_scheduler,
-          vmem_limit_bytes=vmem_limit_bytes,
-      )
-      out = splash_kernel(
-          q_3d_padded.astype(jnp.bfloat16),
-          k_3d_padded,
-          v_3d_padded,
-      )
-      out = jnp.swapaxes(out, 1, 2)
-      return out[:, :q_orig_len, ...]
-
-    return jax.vmap(_kernel_3d, in_axes=(0, 0, 0), out_axes=0)(q, k, v)
-
-  batch_size = query.shape[0]
-  if num_heads < mesh.size:
-    q_partition_spec = P()
-    kv_partition_spec = P()
-    out_constraint = P()
-  else:
-    # Name-based axis resolution for the standard 4D mesh ("data", "fsdp", "context", "tensor").
-    # Axes are grouped by semantic role:
-    #   - Batch-parallel: "data", "fsdp"  (shard the batch dimension)
-    #   - Head-parallel:  "context", "tensor" (shard the head dimension)
-    # Only axes with size > 1 participate in sharding to avoid shard_map errors.
-    # Q and KV are sharded identically (symmetric head-parallel), matching the
-    # strategy that the legacy 2D (dp, tp) path used.
-    batch_axes = tuple(a for a in ("data", "fsdp") if a in mesh.shape and mesh.shape[a] > 1)
-    head_axes = tuple(a for a in ("context", "tensor") if a in mesh.shape and mesh.shape[a] > 1)
-
-    batch_parallel_size = 1
-    for a in batch_axes:
-      batch_parallel_size *= mesh.shape[a]
-    head_parallel_size = 1
-    for a in head_axes:
-      head_parallel_size *= mesh.shape[a]
-
-    # Flatten singleton tuples for clean PartitionSpecs.
-    batch_spec = batch_axes[0] if len(batch_axes) == 1 else (batch_axes if batch_axes else None)
-    head_spec = head_axes[0] if len(head_axes) == 1 else (head_axes if head_axes else None)
-
-    if (
-        batch_parallel_size > 1
-        and batch_size >= batch_parallel_size
-        and num_heads % head_parallel_size == 0
-    ):
-      # Head-parallel: batch on data/fsdp, heads on context/tensor.
-      q_partition_spec = P(batch_spec, head_spec, None, None)
-      kv_partition_spec = P(batch_spec, head_spec, None, None)
-      out_constraint = P(batch_spec, None, head_spec, None)
-    elif head_parallel_size > 1 or batch_parallel_size > 1:
-      # Batch too small for batch-parallel; shard all parallelism onto heads.
-      all_axes = batch_axes + head_axes
-      all_parallel_size = batch_parallel_size * head_parallel_size
-      all_spec = all_axes[0] if len(all_axes) == 1 else (all_axes if all_axes else None)
-      if num_heads % all_parallel_size == 0:
-        q_partition_spec = P(None, all_spec, None, None)
-        kv_partition_spec = P(None, all_spec, None, None)
-        out_constraint = P(None, None, all_spec, None)
-      else:
-        # Heads not divisible by available parallelism; fall back to replicated.
-        q_partition_spec = P()
-        kv_partition_spec = P()
-        out_constraint = P()
-    else:
-      # All axes are size-1; fully replicated (e.g. single-device).
-      q_partition_spec = P()
-      kv_partition_spec = P()
-      out_constraint = P()
-
-  sharded_fn = shard_map(
-      _attention_on_slices,
-      mesh=mesh,
-      in_specs=(q_partition_spec, kv_partition_spec, kv_partition_spec),
-      out_specs=q_partition_spec,
-      check_rep=False,
-  )
-  out = sharded_fn(query, key, value)
-  out = jax.lax.with_sharding_constraint(out, out_constraint)
-  return out
-
-
-# ---------------------------------------------------------------------------
-# TorchAX SDPA wrapper
-# ---------------------------------------------------------------------------
-
-
-def make_custom_splash_sdpa(mesh, env, **kwargs):
-  flash_block_sizes = kwargs.get("flash_block_sizes", None)
-  bq = kwargs.get("block_q", DEFAULT_BQSIZE)
-  bkv = kwargs.get("block_kv", DEFAULT_BKVSIZE)
-  bkv_compute = kwargs.get("block_kv_compute", DEFAULT_BKVCOMPUTESIZE)
-  bkv_compute_in = kwargs.get("block_kv_compute_in", DEFAULT_BKVCOMPUTEINSIZE)
-  hpt = kwargs.get("heads_per_tile", 1)
-  use_k_smooth = kwargs.get("use_k_smooth", True)
-  use_base2_exp = kwargs.get("use_base2_exp", True)
-  use_experimental_scheduler = kwargs.get("use_experimental_scheduler", False)
-  vmem_limit_bytes = kwargs.get("vmem_limit_bytes", None)
-
-  def _simple_attention(q, k, v, scale=None):
-    s = scale if scale is not None else 1.0 / math.sqrt(q.shape[-1])
-    attn = jnp.einsum("bhsd,bhtd->bhst", q * s, k)
-    attn = jax.nn.softmax(attn.astype(jnp.float32), axis=-1).astype(q.dtype)
-    return jnp.einsum("bhst,bhtd->bhsd", attn, v)
-
-  def _sdpa(
-      query,
-      key,
-      value,
-      attn_mask=None,
-      dropout_p=0.0,
-      is_causal=False,
-      scale=None,
-      enable_gqa=False,
-  ):
-    jquery, jkey, jvalue = env.t2j_iso((query, key, value))
-    num_heads = jquery.shape[1]
-
-    if num_heads <= 8:
-      result = _simple_attention(jquery, jkey, jvalue, scale=scale)
-      return env.j2t_iso(result)
-
-    if use_k_smooth:
-      key_mean = jnp.mean(jkey, axis=2, keepdims=True)
-      jkey = jkey - key_mean
-
-    result = tpu_custom_attention(
-        jquery,
-        jkey,
-        jvalue,
-        mesh,
-        scale=scale,
-        block_q=bq,
-        block_kv=bkv,
-        block_kv_compute=bkv_compute,
-        block_kv_compute_in=bkv_compute_in,
-        heads_per_tile=hpt,
-        use_base2_exp=use_base2_exp,
-        use_experimental_scheduler=use_experimental_scheduler,
-        vmem_limit_bytes=vmem_limit_bytes,
-        flash_block_sizes=flash_block_sizes,
-    )
-    return env.j2t_iso(result)
-
-  return _sdpa