From 66c8eb6904f4ce95b08651de1c330e789845f064 Mon Sep 17 00:00:00 2001 From: Rishabh Manoj Date: Sun, 13 Sep 2026 05:57:23 +0000 Subject: [PATCH] feat(attention,wan): fixed-m splash and 2D-ring attention with Wan2.2 TPU optimizations Implements exact fixed-m splash attention with Global Virtual K-Centering, 2D Ulysses+Ring attention, and Wan2.2 fast serving optimizations for Cloud TPUs. Key Features & Safety: 1. Fixed-m Splash Attention: - Dynamic C(N) headroom constants guaranteeing FP32 accumulator safety. - Fixed-m dtype safety checks (rejecting FP16/FP8 exponent overflow). - Global V-magnitude bounds (|V| <= 256) with safe online softmax fallback. 2. 2D Ulysses + Ring Attention: - Global Virtual K-Centering via cross-ring pmean(k). - Cross-ring v_ok reduction via jax.lax.pmin preventing unsafe hop-level fixed-m. - Per-Q-block fixed-m accumulation with tuned tile size (block_q=9472). - Full GQA support across arbitrary Ulysses and Ring splits. 3. Wan2.2 Fast Serving: - Persistent AOT compilation caching and zero-exec warmup. - Tuned TPU inference recipes in run_wan_fast_inference.sh. Hardware Verification (Cloud TPU v6e-8): - 44 tests passed in ring_fixed_m_test and custom_splash_fixed_m_test. - 25-prompt VBench evaluation: zero quality regression vs main. - Step latency: 3.445 s/step (tile 9472) vs 3.526 s/step (main), delivering a +2.32% denoise speedup. --- end_to_end/tpu/run_wan_fast_inference.sh | 83 +- src/maxdiffusion/aot_cache.py | 41 +- src/maxdiffusion/configs/base_wan_14b.yml | 2 +- src/maxdiffusion/configs/base_wan_27b.yml | 2 +- src/maxdiffusion/configs/base_wan_animate.yml | 2 +- src/maxdiffusion/configs/base_wan_i2v_14b.yml | 2 +- src/maxdiffusion/configs/base_wan_i2v_27b.yml | 2 +- src/maxdiffusion/generate_wan.py | 127 ++- .../kernels/custom_splash_attention.py | 257 +++++- src/maxdiffusion/kernels/fused_producers.py | 102 ++ .../splash_attention/ring_attention_kernel.py | 229 +++-- src/maxdiffusion/max_utils.py | 10 +- src/maxdiffusion/models/attention_flax.py | 868 +++++++++++++----- .../wan/transformers/transformer_wan.py | 34 +- src/maxdiffusion/pyconfig.py | 1 + .../tests/custom_splash_fixed_m_test.py | 499 +++++++++- src/maxdiffusion/tests/ring_fixed_m_test.py | 716 +++++++++++++-- .../tests/wan/wan_transformer_test.py | 309 ++++++- src/maxdiffusion/utils/export_utils.py | 15 +- 19 files changed, 2781 insertions(+), 520 deletions(-) create mode 100644 src/maxdiffusion/kernels/fused_producers.py diff --git a/end_to_end/tpu/run_wan_fast_inference.sh b/end_to_end/tpu/run_wan_fast_inference.sh index 89680760b..926b5651b 100755 --- a/end_to_end/tpu/run_wan_fast_inference.sh +++ b/end_to_end/tpu/run_wan_fast_inference.sh @@ -35,6 +35,7 @@ set -u MODEL=${1:-22} STEPS=${2:-40} PROMPT=${3:-""} +shift $(($# > 3 ? 3 : $#)) PROJECT_ROOT="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &> /dev/null && pwd)" cd "$PROJECT_ROOT" || exit 1 @@ -42,38 +43,71 @@ export PYTHONPATH="$PROJECT_ROOT/src:${PYTHONPATH:-}" export HF_HUB_ENABLE_HF_TRANSFER=1 export JAX_DEFAULT_MATMUL_PRECISION=bfloat16 export TORCHINDUCTOR_FX_GRAPH_CACHE=1 +export XLA_PYTHON_CLIENT_MEM_FRACTION=0.95 CACHE_ROOT=${WAN_CACHE_ROOT:-$HOME/.cache/maxdiffusion_wan} -OUTPUT_DIR=${OUTPUT_DIR:-/tmp/wan_out} -mkdir -p "$CACHE_ROOT/jax" "$CACHE_ROOT/aot_wan$MODEL" "$CACHE_ROOT/converted" "$OUTPUT_DIR" +OUTPUT_DIR=${OUTPUT_DIR:-$HOME/maxdiffusion_wan_output} +export TMPDIR=${TMPDIR:-$CACHE_ROOT/tmp} +mkdir -p "$CACHE_ROOT/jax" "$CACHE_ROOT/aot_wan$MODEL" "$CACHE_ROOT/converted" "$OUTPUT_DIR" "$TMPDIR" -# 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:-}" +# Tuned collective/scheduler flag set for TPU. +# One line: libtpu stops parsing at a literal backslash. +DEFAULT_LIBTPU="--xla_tpu_spmd_rng_bit_generator_unsafe=true --xla_tpu_enable_dot_strength_reduction=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_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" +export LIBTPU_INIT_ARGS="${DEFAULT_LIBTPU} ${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 +# Attention selection +# +# Default is the 2D-ring per-q-block fixed-m kernel. Measured on v6e-8, 720p +# 81f, 40 steps, 25 prompts, matched config (U=4, tile 6400/1024, kv cache on): +# +# ulysses_ring_custom_fixed_m_per_q_block 3.4919 s/step <- default +# ulysses_custom_fixed_m 3.5338 s/step +# ulysses_custom_fixed_m_per_q_block 3.5328 s/step +# +# For reference, `main`'s kernel on its best recipe is 3.5274 s/step, i.e. the +# per-head variant that used to be the default here is slightly SLOWER than +# main; the win is specifically in the ring per-q-block path. +if [ -z "${ATTENTION:-}" ]; then + if [ "${FIXEDM:-1}" = "1" ]; then + ATTENTION="ulysses_ring_custom_fixed_m_per_q_block" + else + ATTENTION="ulysses_custom" + fi +fi + +if [ "$ATTENTION" = "ulysses_custom_fixed_m" ] || + [ "$ATTENTION" = "ulysses_custom_fixed_m_per_q_block" ] || + [ "$ATTENTION" = "ulysses_ring_custom_fixed_m_per_q_block" ] || + [ "$ATTENTION" = "ulysses_custom" ]; then + DEFAULT_U=4 + DEFAULT_BQ=9472 + DEFAULT_BKV=1024 +elif [ "${FIXEDM:-1}" = "1" ]; then + DEFAULT_U=2 + DEFAULT_BQ=6400 + DEFAULT_BKV=2048 else - ATTENTION=ulysses_ring_custom - BQ=9472; BKV=1024 + DEFAULT_U=2 + DEFAULT_BQ=9472 + DEFAULT_BKV=1024 fi +ULYSSES_SHARDS=${ULYSSES_SHARDS:-$DEFAULT_U} +BQ=${BQ:-$DEFAULT_BQ} +BKV=${BKV:-$DEFAULT_BKV} + if [ "$MODEL" = "21" ]; then CONFIG=src/maxdiffusion/configs/base_wan_14b.yml - GUIDANCE_ARGS="" + GUIDANCE_ARGS=() else CONFIG=src/maxdiffusion/configs/base_wan_27b.yml - GUIDANCE_ARGS="guidance_scale_low=3.0 guidance_scale_high=4.0" + GUIDANCE_ARGS=(guidance_scale_low=3.0 guidance_scale_high=4.0) 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" +echo "== ${ATTENTION} | U=${ULYSSES_SHARDS} | 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. @@ -83,25 +117,26 @@ 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=$ATTENTION \ - ulysses_shards=2 \ + attention="$ATTENTION" \ + ulysses_shards="$ULYSSES_SHARDS" \ ici_data_parallelism=2 ici_fsdp_parallelism=1 \ ici_context_parallelism=4 ici_tensor_parallelism=1 \ per_device_batch_size=0.125 \ num_inference_steps="$STEPS" num_frames=81 width=1280 height=720 \ weights_dtype=bfloat16 activations_dtype=bfloat16 \ - vae_spatial=4 vae_decode_chunk=-1 \ + vae_spatial=8 vae_decode_chunk="${VAE_DECODE_CHUNK:-1}" \ vae_weights_dtype=bfloat16 vae_dtype=bfloat16 \ 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 \ + use_kv_cache=true use_base2_exp=true use_experimental_scheduler=true \ + fps=16 "${GUIDANCE_ARGS[@]}" \ + seed="${SEED:-12345}" \ 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) -mp4=$(ls -t wan_output_*.mp4 2>/dev/null | head -1) +mp4=$(ls -t "$OUTPUT_DIR"/${RUN_NAME}*.mp4 "$OUTPUT_DIR"/wan_output_*.mp4 2>/dev/null | head -1) if [ -n "$mp4" ]; then - mv "$mp4" "$OUTPUT_DIR/${RUN_NAME}.mp4" echo "" - echo "=== video saved: $OUTPUT_DIR/${RUN_NAME}.mp4 ===" + echo "=== video saved: $mp4 ===" fi diff --git a/src/maxdiffusion/aot_cache.py b/src/maxdiffusion/aot_cache.py index c07af3e24..4143dc2a4 100644 --- a/src/maxdiffusion/aot_cache.py +++ b/src/maxdiffusion/aot_cache.py @@ -96,7 +96,8 @@ def _dynamic_signature(args: tuple, kwargs: dict) -> str: parts = [] for path, leaf in leaves_with_paths: if hasattr(leaf, "shape") and hasattr(leaf, "dtype"): - desc = f"{tuple(leaf.shape)}:{leaf.dtype}" + weak = getattr(leaf, "weak_type", False) + desc = f"{tuple(leaf.shape)}:{leaf.dtype}:weak={weak}" else: desc = re.sub(r"0x[0-9a-fA-F]+", "@", repr(leaf)) parts.append(f"{jax.tree_util.keystr(path)}={desc}") @@ -197,7 +198,23 @@ def __call__(self, *args, **kwargs): # Under an outer trace a deserialized executable cannot be applied # and tracers must not be recorded -- inline like a nested jit. return self.jitted(**dynamic, **static) - signature = _dynamic_signature((), {**dynamic, **static}) + + # Fast-path signature cache: avoid tree_flatten_with_path + SHA256 string hashing on repeated steps + shapes_dtypes = tuple( + (leaf.shape, leaf.dtype, getattr(leaf, "weak_type", False)) + if hasattr(leaf, "shape") and hasattr(leaf, "dtype") + else (type(leaf), re.sub(r"0x[0-9a-fA-F]+", "@", repr(leaf))) + for leaf in leaves + ) + cache_key = (treedef, shapes_dtypes, tuple(sorted(static.items()))) + signature = getattr(self, "_sig_cache", {}).get(cache_key) + if signature is None: + signature = _dynamic_signature((), {**dynamic, **static}) + if not hasattr(self, "_sig_cache"): + self._sig_cache = {} + if len(self._sig_cache) < 64: + self._sig_cache[cache_key] = signature + if _STATE.warmup_only: # Compilation only needs avals; skip the (possibly seconds-long) # real execution and hand back correctly-shaped/sharded zeros so @@ -212,17 +229,14 @@ def __call__(self, *args, **kwargs): return zeros compiled = self._compiled.get(signature) if compiled is not None: - flat = self._align_inputs(compiled, leaves) - if flat is not None: - return compiled(flat) - # Fewer expected shardings than leaves: XLA pruned unused inputs - # (e.g. encoder params in a decode-only executable). Compiled keeps - # the full in_tree and prunes internally, so hand it the raw leaves; - # sharding/structure problems surface as catchable Python errors. try: + flat = self._align_inputs(compiled, leaves) + if flat is not None: + return compiled(flat) return compiled(leaves) except Exception as e: # noqa: BLE001 - any failure means "use jit" - max_logging.log(f"[aot] {self.name}: compiled call failed ({e}); using jit") + max_logging.log(f"[aot] fast-path execution failed for {self.name}/{signature} ({e}); falling back to JIT") + return self.jitted(**dynamic, **static) with self._lock: if signature not in self._pending and signature not in self._compiled: self._pending[signature] = (leaves, treedef, static) @@ -243,12 +257,17 @@ def _align_inputs(self, compiled: Any, leaves: list): # Fewer expected shardings than leaves = XLA pruned unused inputs; # the caller retries via Compiled's own pruning path. Not an error. return None + + # Fast path: check if all leaves already match expected sharding exactly + if all(getattr(leaf, "sharding", None) is expected for leaf, expected in zip(leaves, flat_expected)): + return leaves + aligned = [] for leaf, expected in zip(leaves, flat_expected): if not hasattr(leaf, "shape"): # python scalar traced as weak array leaf = jnp.asarray(leaf) sharding = getattr(leaf, "sharding", None) - if sharding is not None and sharding.is_equivalent_to(expected, leaf.ndim): + if sharding is not None and (sharding is expected or sharding.is_equivalent_to(expected, leaf.ndim)): aligned.append(leaf) else: aligned.append(jax.device_put(leaf, expected)) diff --git a/src/maxdiffusion/configs/base_wan_14b.yml b/src/maxdiffusion/configs/base_wan_14b.yml index 837bbe98b..a0166712d 100644 --- a/src/maxdiffusion/configs/base_wan_14b.yml +++ b/src/maxdiffusion/configs/base_wan_14b.yml @@ -83,7 +83,7 @@ jit_initializers: True # Set true to load weights from pytorch from_pt: True split_head_dim: True -attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block use_base2_exp: True use_experimental_scheduler: True # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml index bf8e1c740..2b22717a0 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/src/maxdiffusion/configs/base_wan_27b.yml @@ -83,7 +83,7 @@ jit_initializers: True # Set true to load weights from pytorch from_pt: True split_head_dim: True -attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, tokamax_ring_custom, ulysses, ulysses_custom, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_bidir +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, tokamax_ring_custom, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block, ulysses_ring_custom_bidir # # Best 2D-ring / USP (Ulysses x ring) configs for WAN2.2-T2V-A14B (720x1280, 81 frames) # Set attention=ulysses_ring_custom and ulysses_shards=U (ring degree R=CP/U): diff --git a/src/maxdiffusion/configs/base_wan_animate.yml b/src/maxdiffusion/configs/base_wan_animate.yml index 5e9df7d0d..48b3747f6 100644 --- a/src/maxdiffusion/configs/base_wan_animate.yml +++ b/src/maxdiffusion/configs/base_wan_animate.yml @@ -81,7 +81,7 @@ jit_initializers: True # Set true to load weights from pytorch from_pt: True split_head_dim: True -attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block use_base2_exp: True use_experimental_scheduler: True # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. diff --git a/src/maxdiffusion/configs/base_wan_i2v_14b.yml b/src/maxdiffusion/configs/base_wan_i2v_14b.yml index a129ff66c..dfdb691af 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_14b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_14b.yml @@ -83,7 +83,7 @@ jit_initializers: True # Set true to load weights from pytorch from_pt: True split_head_dim: True -attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block use_base2_exp: True use_experimental_scheduler: True # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. diff --git a/src/maxdiffusion/configs/base_wan_i2v_27b.yml b/src/maxdiffusion/configs/base_wan_i2v_27b.yml index 6a28986fc..ba4ca546e 100644 --- a/src/maxdiffusion/configs/base_wan_i2v_27b.yml +++ b/src/maxdiffusion/configs/base_wan_i2v_27b.yml @@ -83,7 +83,7 @@ jit_initializers: True # Set true to load weights from pytorch from_pt: True split_head_dim: True -attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_custom_fixed_m, ulysses_custom_fixed_m_per_q_block, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_fixed_m, ulysses_ring_custom_fixed_m_per_q_block use_base2_exp: True use_experimental_scheduler: True # For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. diff --git a/src/maxdiffusion/generate_wan.py b/src/maxdiffusion/generate_wan.py index 662069ffd..5fe619dff 100644 --- a/src/maxdiffusion/generate_wan.py +++ b/src/maxdiffusion/generate_wan.py @@ -16,6 +16,7 @@ import jax import time import os +import uuid from maxdiffusion.checkpointing.wan_checkpointer_2_1 import WanCheckpointer2_1 from maxdiffusion.checkpointing.wan_checkpointer_2_2 import WanCheckpointer2_2 from maxdiffusion.checkpointing.wan_checkpointer_i2v_2p1 import WanCheckpointerI2V_2_1 @@ -37,6 +38,28 @@ jax.config.update("jax_use_shardy_partitioner", True) +def _non_reusable_aot_revision(): + """Returns a unique identity so unversioned/dirty development source can never hit old HLO.""" + return f"unversioned:{uuid.uuid4().hex}" + + +def _resolve_wan_aot_source_revision(config, commit_hash=None): + """Prefers an explicit Git revision, then a packaged-build revision.""" + for revision in (commit_hash, getattr(config, "aot_build_revision", None)): + if revision is not None and str(revision).strip(): + return str(revision).strip() + return None + + +def _is_reusable_aot_revision(source_revision) -> bool: + if source_revision is None or not str(source_revision).strip(): + return False + s = str(source_revision).strip() + if s.startswith(("dirty:", "unversioned:")) or s.endswith("-dirty"): + return False + return True + + def call_pipeline(config, pipeline, prompt, negative_prompt, num_inference_steps=None): model_key = config.model_name model_type = config.model_type @@ -298,23 +321,45 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): # Per-shape AOT executable cache: deserialization starts on background # threads now and overlaps the remaining setup; unknown shapes silently # fall back to jit and are serialized by save_pending() after warmup. + detected_revision = commit_hash if commit_hash is not None else max_utils.get_git_commit_hash(check_dirty=True) + source_revision = _resolve_wan_aot_source_revision(config, detected_revision) + aot_cache_dir = getattr(config, "aot_cache_dir", "") + if aot_cache_dir and not _is_reusable_aot_revision(source_revision): + max_logging.log( + "[aot] No clean Git commit or aot_build_revision was supplied; " + "persistent Wan AOT caching is disabled for this development run." + ) + aot_cache_dir = "" + + aot_metadata = { + "model": config.pretrained_model_name_or_path, + "attention": config.attention, + # Kernel block sizes change the lowered graph, not the input + # shapes — they must key the executable or a re-tuned config + # would silently hit stale binaries. + "flash_block_sizes": str(config.flash_block_sizes), + "mesh_shape": str(pipeline.mesh.shape), + "vae_spatial": str(getattr(config, "vae_spatial", 8)), + "vae_decode_chunk": str(getattr(config, "vae_decode_chunk", 1)), + "weights_dtype": str(config.weights_dtype), + "activations_dtype": str(config.activations_dtype), + "scan_layers": str(config.scan_layers), + "ulysses_shards": str(getattr(config, "ulysses_shards", 1)), + "ulysses_attention_chunks": str(getattr(config, "ulysses_attention_chunks", 1)), + "flash_min_seq_length": str(getattr(config, "flash_min_seq_length", 4096)), + "mask_padding_tokens": str(getattr(config, "mask_padding_tokens", True)), + "precision": str(getattr(config, "precision", "default")), + "logical_axis_rules": str(getattr(config, "logical_axis_rules", ())), + "device_kind": str(jax.devices()[0].device_kind if jax.devices() else "unknown"), + "process_count": str(jax.process_count()), + "use_base2_exp": str(getattr(config, "use_base2_exp", True)), + "use_experimental_scheduler": str(getattr(config, "use_experimental_scheduler", False)), + "jax": jax.__version__, + "source_revision": source_revision if source_revision else _non_reusable_aot_revision(), + } aot_cache.install( - getattr(config, "aot_cache_dir", ""), - meta={ - "model": config.pretrained_model_name_or_path, - "attention": config.attention, - # Kernel block sizes change the lowered graph, not the input - # shapes — they must key the executable or a re-tuned config - # would silently hit stale binaries. - "flash_block_sizes": str(config.flash_block_sizes), - "mesh_shape": str(pipeline.mesh.shape), - "vae_spatial": str(config.vae_spatial), - "vae_decode_chunk": str(config.vae_decode_chunk), - "weights_dtype": str(config.weights_dtype), - "activations_dtype": str(config.activations_dtype), - "scan_layers": str(config.scan_layers), - "jax": jax.__version__, - }, + aot_cache_dir, + meta=aot_metadata, mesh=pipeline.mesh, ) # Deserialization is seconds and warmup must see the loaded executables @@ -391,12 +436,23 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): else: videos = outputs trace = {} - for i in range(len(videos)): - video_path = f"{filename_prefix}wan_output_{config.seed}_{i}.mp4" - export_to_video(videos[i], video_path, fps=config.fps) - saved_video_path.append(video_path) - if gcs_output_path: - max_utils.upload_file_to_gcs(gcs_output_path, video_path, subdir="videos") + if jax.process_index() == 0: + import numpy as np + + for i in range(len(videos)): + if getattr(config, "output_dir", "") and not config.output_dir.startswith("gs://"): + os.makedirs(config.output_dir, exist_ok=True) + video_path = os.path.join(config.output_dir, f"{config.run_name}_{config.seed}_{i}.mp4") + else: + video_path = f"{filename_prefix}wan_output_{config.seed}_{i}.mp4" + frames_np = np.asarray(videos[i]) + export_to_video(frames_np, video_path, fps=config.fps) + saved_video_path.append(video_path) + max_logging.log(f"Saved video to {video_path}") + if gcs_output_path: + max_utils.upload_file_to_gcs(gcs_output_path, video_path, subdir="videos") + elif getattr(config, "output_dir", "").startswith("gs://"): + max_utils.upload_file_to_gcs(os.path.join(config.output_dir, config.run_name), video_path, subdir="videos") else: trace = {} for i, padded_chunk, actual_chunk_len in max_utils.chunk_and_pad(prompts, batch_size): @@ -407,13 +463,24 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): videos, trace = outputs else: videos = outputs - for j in range(actual_chunk_len): - prompt_idx = i + j - video_path = f"{filename_prefix}wan_output_{config.seed}_{prompt_idx}.mp4" - export_to_video(videos[j], video_path, fps=config.fps) - saved_video_path.append(video_path) - if gcs_output_path: - max_utils.upload_file_to_gcs(gcs_output_path, video_path, subdir="videos") + if jax.process_index() == 0: + import numpy as np + + for j in range(actual_chunk_len): + prompt_idx = i + j + if getattr(config, "output_dir", "") and not config.output_dir.startswith("gs://"): + os.makedirs(config.output_dir, exist_ok=True) + video_path = os.path.join(config.output_dir, f"{config.run_name}_{config.seed}_{prompt_idx}.mp4") + else: + video_path = f"{filename_prefix}wan_output_{config.seed}_{prompt_idx}.mp4" + frames_np = np.asarray(videos[j]) + export_to_video(frames_np, video_path, fps=config.fps) + saved_video_path.append(video_path) + max_logging.log(f"Saved video to {video_path}") + if gcs_output_path: + max_utils.upload_file_to_gcs(gcs_output_path, video_path, subdir="videos") + elif getattr(config, "output_dir", "").startswith("gs://"): + max_utils.upload_file_to_gcs(os.path.join(config.output_dir, config.run_name), video_path, subdir="videos") generation_time = time.perf_counter() - s0 max_logging.log(f"generation_time: {generation_time}") @@ -482,7 +549,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None): def main(argv: Sequence[str]) -> None: - commit_hash = max_utils.get_git_commit_hash() + commit_hash = max_utils.get_git_commit_hash(check_dirty=True) pyconfig.initialize(argv) try: flax.config.update("flax_always_shard_variable", False) diff --git a/src/maxdiffusion/kernels/custom_splash_attention.py b/src/maxdiffusion/kernels/custom_splash_attention.py index 6bd3f3493..a23abb310 100644 --- a/src/maxdiffusion/kernels/custom_splash_attention.py +++ b/src/maxdiffusion/kernels/custom_splash_attention.py @@ -17,6 +17,7 @@ """Custom Pallas flash attention kernel for TPU.""" import functools +import math import jax import jax.numpy as jnp @@ -47,23 +48,89 @@ def __init__( self.block_kv_compute_in = block_kv_compute_in -# Fixed-m softmax-bound constants. Instead of tracking the online-softmax -# running max per KV block, eligible heads subtract a precomputed per-query -# upper bound on the logits (Cauchy-Schwarz: max_j q_i.k_j <= ||q_i|| * -# max_j||k_j||). _FIXED_M_RECENTER (C) shifts the exp2 exponents up so the -# largest surviving term stays above the f32 subnormal-flush floor 2^-126: -# with k-smoothing the per-row max is >= 0, so the max term has exponent -# >= -ceil(bound) + C, which stays > -126 while ceil(bound) <= -# _FIXED_M_SAFE_BOUND (= C + 126 - 1 of margin). Heads whose worst-case bound -# 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 +FP32_OUTPUT_HEADROOM_BITS = 8.0 # Assumes default activation |V| <= 2**FP32_OUTPUT_HEADROOM_BITS = 256.0 +DEFAULT_MAX_V_BOUND = 256.0 + + +def fixed_m_dtype_is_safe(dtype, recenter: float) -> bool: + """Whether `dtype` can hold the fixed-m softmax weights without overflowing. + + Fixed-m deliberately parks the un-normalized weights at up to `2**recenter`, + a range derived against FP32's exponent (see `get_fixed_m_constants`). The + kernel then narrows them to the activation dtype for the S@V matmul + (`s_curr.astype(q_ref.dtype)`), so a dtype with a *smaller exponent range* + silently overflows to inf even though the FP32 bound analysis passed. + + bfloat16 and float32 both have 8-bit exponents (maxexp 128) and are safe for + every C(N) this module produces. float16 has a 5-bit exponent (maxexp 16) and + is not: at N=4096 with |V| <= 256, C(N) = 107 and 2**107 is far beyond + float16's 65504 ceiling. The fp8 formats fail for the same reason. + + This is deliberately expressed in terms of the exponent range rather than an + allowlist so narrower formats are rejected automatically. + + Args: + dtype: Activation dtype the kernel will narrow the weights to. + recenter: The fixed-m constant C(N) from `get_fixed_m_constants`. + + Returns: + True if `2**recenter` is representable in `dtype`. + """ + return float(jnp.finfo(jnp.dtype(dtype)).maxexp) > float(recenter) + + +def get_fixed_m_constants( + kv_seq_len: int, + is_ring: bool = False, + v_max_bound: float = DEFAULT_MAX_V_BOUND, +) -> tuple[float, float]: + """Computes dynamic fixed-m constants C(N) and safe bounds based on KV sequence length. + + Mathematical Derivations: + 1. Overflow Ceiling: + For a given upper bound on value activation magnitude |V| <= V_max: + output_headroom_bits = ceil(log2(V_max)). + The ceiling constant C(N) = 127.0 - ceil(log2(N)) - output_headroom_bits guarantees that: + - Denominator accumulator: l = sum_j 2^{z_j - m} <= N * 2^C(N) <= 2^{127 - headroom} < 2^{128} + - Numerator accumulator: |o_d| = |sum_j V_{j,d} 2^{z_j - m}| <= V_max * N * 2^C(N) <= 2^{127} < 2^{128} + preventing IEEE-754 FP32 overflow for all activations |V| <= V_max. + + 2. Subnormal Underflow Floor (Cauchy-Schwarz Proof): + Let U_i = max_i ||q_i|| * max_j ||k_j|| be the Cauchy-Schwarz bound on query-key inner products. + By Cauchy-Schwarz inequality, for all tokens j: + z_j = Q_i . K_j >= -||Q_i|| * ||K_j|| >= -U_i. + With the fixed-m base shift defined as m_i = ceil(U_i) - C(N): + z_j - m_i >= -U_i - (ceil(U_i) - C(N)) = C(N) - (U_i + ceil(U_i)). + To guarantee that no term underflows into the subnormal range (requiring minimal shifted exponent >= -125.0, + providing 1 bit of margin above IEEE-754 normal floor -126.0): + - Ulysses (Centered, M >= 0): logit centering guarantees row max >= 0, so ceil(U) <= C(N) + 125.0 = W(N). + - Ring (Uncentered, M >= -U): U + ceil(U) <= W(N) => U <= floor(W(N) / 2). + If U_i <= floor(W(N) / 2), it is mathematically impossible to underflow below -125.0, even without K-centering! + """ + if kv_seq_len is None or kv_seq_len <= 0: + raise ValueError(f"kv_seq_len must be a positive integer to compute dynamic fixed-m constants, got {kv_seq_len=}") + if v_max_bound <= 0.0: + raise ValueError(f"v_max_bound must be a positive float, got {v_max_bound=}") + + fp32_max_exp = 128.0 + fp32_min_normal_exp = -126.0 + output_headroom_bits = float(max(0, math.ceil(math.log2(float(v_max_bound))))) + + max_accumulation_bits = float(math.ceil(math.log2(float(kv_seq_len)))) + + # C(N) = 127.0 - max_accumulation_bits - output_headroom_bits + recenter = fp32_max_exp - max_accumulation_bits - output_headroom_bits - 1.0 + + # Safe window W(N) = C(N) - (-126.0) - 1.0 = C(N) + 125.0 + safe_window = recenter - fp32_min_normal_exp - 1.0 + + if is_ring: + # For integer threshold K = floor(W / 2), U <= K guarantees U + ceil(U) <= 2K <= W + safe_bound = float(int(safe_window // 2)) + else: + safe_bound = safe_window + + return recenter, safe_bound def _flash_attention_kernel( @@ -71,6 +138,7 @@ def _flash_attention_kernel( q_ref, k_ref, v_ref, + k_mean_ref, m_scratch_ref, l_scratch_ref, o_scratch_ref, @@ -89,13 +157,23 @@ def _flash_attention_kernel( fuse_reciprocal: bool = True, use_fixed_m: bool = False, uniform_fixed_m: bool = False, + fixed_m_recenter: float | None = None, + q_heads_per_kv_head: int = 1, ): + """Pallas Mosaic TPU flash attention kernel with fixed-m support. + + Scalar Prefetch Multiplexing: + `mk_ref` is a multiplexed scalar prefetch buffer of shape `(2, num_heads, num_q_blocks)` + passing both the precomputed block fixed-m base shift and discrete predicate in a single scalar memory slot: + - `mk_ref[0, h, i]`: Precomputed block shift m_B = ceil(max_i ||q_i|| * max_j ||k_j||) - C. + - `mk_ref[1, h, i]`: Gating eligibility predicate (1.0 for fixed-m, 0.0 for online). + """ 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}") - h, _, j = pl.program_id(0), pl.program_id(1), pl.program_id(2) + h, i, j = pl.program_id(0), pl.program_id(1), pl.program_id(2) exp = jnp.exp2 if use_base2_exp else jnp.exp sv_dims = (((0,), (0,)), ((), ())) @@ -105,26 +183,31 @@ def _flash_attention_kernel( # 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 and not fixed_only) else False + if use_fixed_m and fixed_m_recenter is None: + raise ValueError("fixed_m_recenter must be specified when use_fixed_m=True.") + + # Per-(head, Q-block) dispatch: heads / Q-blocks inside the no-flush window run + # fixed-m, the rest keep online softmax. + if use_fixed_m and not fixed_only: + is_fixed = mk_ref[1, h, i] > 0.5 + else: + is_fixed = 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 + # Precomputed block bound m_B = ceil(max_i ||q_i|| * max_j ||k_j||) - C. + # Virtual K-centering applies the row-specific projection: m_i = m_B + q_i^T \bar{k}. + m_base = mk_ref[0, h, i] + if k_mean_ref is not None: + qf = q_ref[...].astype(float32) + km = k_mean_ref[h // q_heads_per_kv_head, :].astype(float32) + mu = (qf * km[None, :]).sum(axis=1)[None, :] + m_fixed = m_base + mu + else: + m_fixed = m_base m_scratch_ref[...] = jnp.broadcast_to(m_fixed, m_scratch_ref.shape) @pl.when(j == 0) @@ -232,7 +315,8 @@ def last_compute_body_fixed(kv_compute_index): l_scratch_ref[...] = l_prev o_scratch_ref[:] = o_prev - assert bkv % bkv_compute == 0 + if bkv % bkv_compute != 0: + raise ValueError(f"block_kv ({bkv}) must be divisible by block_kv_compute ({bkv_compute})") if fixed_only: @@ -443,7 +527,8 @@ def last_compute_body(kv_compute_index): l_scratch_ref[h_local] = l_prev o_scratch_ref[h_local] = o_prev - assert bkv % bkv_compute == 0 + if bkv % bkv_compute != 0: + raise ValueError(f"block_kv ({bkv}) must be divisible by block_kv_compute ({bkv_compute})") @pl.when(j != grid_width - 1) def body(): @@ -483,14 +568,12 @@ def _splash_attention_forward( vmem_limit_bytes: int | None = None, use_fixed_m: bool = False, mk: jax.Array | None = None, + fixed_m_recenter: float | None = None, + uniform_fixed_m: bool = False, + k_mean: jax.Array | None = None, ): num_q_heads, padded_q_seq_len, head_dim_qk = q.shape head_dim_v = v.shape[-1] - # Scalar-prefetch operand carrying per-head fixed-m data: - # mk[0, h] = max_j||k_j|| (Cauchy-Schwarz factor), mk[1, h] = eligibility. - # A dummy is supplied for online callers; the kernel ignores it. - if mk is None: - mk = jnp.zeros((2, num_q_heads), jnp.float32) 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 @@ -499,7 +582,36 @@ def _splash_attention_forward( actual_q_seq_len = q_seq_len if q_seq_len is not None else padded_q_seq_len actual_kv_seq_len = kv_seq_len if kv_seq_len is not None else padded_kv_seq_len + if num_q_heads % num_kv_heads != 0: + raise ValueError(f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA.") q_heads_per_kv_head = num_q_heads // num_kv_heads + grid_width = (actual_kv_seq_len + bkv - 1) // bkv + grid_height = (actual_q_seq_len + bq - 1) // bq + grid = (num_q_heads, grid_height, grid_width) + + if use_fixed_m and fixed_m_recenter is None: + raise ValueError("`fixed_m_recenter` must be explicitly specified when `use_fixed_m=True`.") + + # Scalar-prefetch operand carrying per-head / per-Q-block fixed-m data: + # mk[0, h, i] = m_B (precomputed block fixed-m base shift), mk[1, h, i] = eligibility. + # A dummy is supplied for online callers; the kernel ignores it. + if use_fixed_m and mk is None: + raise ValueError("`mk` metadata array is required when `use_fixed_m=True`.") + if mk is None: + mk = jnp.zeros((2, num_q_heads, grid_height), jnp.float32) + elif mk.ndim == 2: + mk = jnp.broadcast_to(mk[:, :, None], (2, num_q_heads, grid_height)) + + if mk.shape[0] != 2 or mk.shape[1] != num_q_heads or mk.shape[2] != grid_height: + raise ValueError(f"mk must have shape (2, {num_q_heads}, {grid_height}), got {mk.shape}") + + if k_mean is None: + k_mean = jnp.zeros((num_kv_heads, head_dim_qk), dtype=jnp.float32) + elif k_mean.shape[0] == num_q_heads and num_q_heads != num_kv_heads: + k_mean = k_mean[::q_heads_per_kv_head] + + if k_mean.shape[0] != num_kv_heads or k_mean.shape[1] != head_dim_qk: + raise ValueError(f"k_mean must have shape ({num_kv_heads}, {head_dim_qk}), got {k_mean.shape}") def q_index_map(h, i, j, *_): return (h, i, 0) @@ -517,6 +629,7 @@ def v_index_map(h, i, j, *_): pl.BlockSpec((None, bq, head_dim_qk), q_index_map), pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), pl.BlockSpec((None, bkv, head_dim_v), v_index_map), + pl.BlockSpec((k_mean.shape[0], head_dim_qk), lambda *_: (0, 0)), ] out_shapes = [ jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), @@ -530,9 +643,6 @@ def v_index_map(h, i, j, *_): pl.BlockSpec((head_dim_v, bq), lambda *_: (0, 0)), pl.BlockSpec((None, head_dim_v, bq), out_index_map), ] - grid_width = (actual_kv_seq_len + bkv - 1) // bkv - grid_height = (actual_q_seq_len + bq - 1) // bq - grid = (num_q_heads, grid_height, grid_width) all_out = pl.pallas_call( functools.partial( @@ -546,6 +656,9 @@ def v_index_map(h, i, j, *_): kv_seq_len=actual_kv_seq_len, use_base2_exp=use_base2_exp, use_fixed_m=use_fixed_m, + uniform_fixed_m=uniform_fixed_m, + fixed_m_recenter=fixed_m_recenter, + q_heads_per_kv_head=q_heads_per_kv_head, ), grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=1, @@ -561,7 +674,7 @@ def v_index_map(h, i, j, *_): vmem_limit_bytes=vmem_limit_bytes, ), out_shape=out_shapes, - )(mk, q, k, v) + )(mk, q, k, v, k_mean) return all_out[-1] @@ -578,6 +691,8 @@ def _splash_attention_forward_ring( use_fixed_m: bool = False, mk: jax.Array | None = None, uniform_fixed_m: bool = False, + fixed_m_recenter: float | None = None, + k_mean: jax.Array | None = None, ): """Ring-specific forward path that returns pre-reciprocal fp32 accumulators. @@ -603,8 +718,24 @@ def _splash_attention_forward_ring( actual_q_seq_len = q_seq_len if q_seq_len is not None else padded_q_seq_len actual_kv_seq_len = kv_seq_len if kv_seq_len is not None else padded_kv_seq_len + if num_q_heads % num_kv_heads != 0: + raise ValueError(f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA.") q_heads_per_kv_head = num_q_heads // num_kv_heads + if use_fixed_m and fixed_m_recenter is None: + raise ValueError("`fixed_m_recenter` must be explicitly specified when `use_fixed_m=True`.") + + if use_fixed_m and mk is None: + raise ValueError("`mk` metadata array is required when `use_fixed_m=True`.") + + if k_mean is None: + k_mean = jnp.zeros((num_kv_heads, head_dim_qk), dtype=jnp.float32) + elif k_mean.shape[0] == num_q_heads and num_q_heads != num_kv_heads: + k_mean = k_mean[::q_heads_per_kv_head] + + if k_mean.shape[0] != num_kv_heads or k_mean.shape[1] != head_dim_qk: + raise ValueError(f"k_mean must have shape ({num_kv_heads}, {head_dim_qk}), got {k_mean.shape}") + def q_index_map(h, i, j, *_): return (h, i, 0) @@ -621,6 +752,7 @@ def v_index_map(h, i, j, *_): pl.BlockSpec((None, bq, head_dim_qk), q_index_map), pl.BlockSpec((None, bkv, head_dim_qk), k_index_map), pl.BlockSpec((None, bkv, head_dim_v), v_index_map), + pl.BlockSpec((k_mean.shape[0], head_dim_qk), lambda *_: (0, 0)), ] out_shapes = [ jax.ShapeDtypeStruct((NUM_SUBLANES, bq), jnp.float32), @@ -642,12 +774,16 @@ def v_index_map(h, i, j, *_): grid_height = (actual_q_seq_len + bq - 1) // bq grid = (num_q_heads, grid_height, grid_width) - # Scalar-prefetch operand carrying per-head fixed-m data (same convention as - # `_splash_attention_forward`): mk[0, h] = max_j||k_j|| over ALL ring shards - # (the caller all-reduces this over the ring axis), mk[1, h] = eligibility. + # Scalar-prefetch operand carrying per-head / per-Q-block fixed-m data: + # mk[0, h, i] = max_j||k_j|| over ALL ring shards + # (the caller all-reduces this over the ring axis), mk[1, h, i] = eligibility. # A dummy is supplied for online callers; the kernel ignores it. + if use_fixed_m and mk is None: + raise ValueError("`mk` metadata array is required when `use_fixed_m=True`.") if mk is None: - mk = jnp.zeros((2, num_q_heads), jnp.float32) + mk = jnp.zeros((2, num_q_heads, grid_height), jnp.float32) + elif mk.ndim == 2: + mk = jnp.broadcast_to(mk[:, :, None], (2, num_q_heads, grid_height)) all_out = pl.pallas_call( functools.partial( @@ -663,6 +799,8 @@ def v_index_map(h, i, j, *_): fuse_reciprocal=False, use_fixed_m=use_fixed_m, uniform_fixed_m=uniform_fixed_m, + fixed_m_recenter=fixed_m_recenter, + q_heads_per_kv_head=q_heads_per_kv_head, ), grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=1, @@ -678,7 +816,7 @@ def v_index_map(h, i, j, *_): vmem_limit_bytes=vmem_limit_bytes, ), out_shape=out_shapes, - )(mk, q, k, v) + )(mk, q, k, v, k_mean) out = jnp.swapaxes(all_out[3], 1, 2) # (h, head_dim_v, s) -> (h, s, head_dim_v) l = all_out[4][:, 0, :] # (h, s) m = all_out[5][:, 0, :] # (h, s) @@ -707,8 +845,10 @@ def _splash_attention_forward_mhpt( actual_kv_seq_len = kv_seq_len if kv_seq_len is not None else k.shape[1] hpt = heads_per_tile - assert num_q_heads % hpt == 0, f"num_heads {num_q_heads} must be divisible by heads_per_tile {hpt}" - assert num_q_heads == num_kv_heads, "MHPT currently requires num_q_heads == num_kv_heads (no GQA)" + if num_q_heads % hpt != 0: + raise ValueError(f"num_heads {num_q_heads} must be divisible by heads_per_tile {hpt}") + if num_q_heads != num_kv_heads: + raise ValueError(f"MHPT currently requires num_q_heads == num_kv_heads (no GQA), got {num_q_heads=} vs {num_kv_heads=}") def q_index_map(h, i, j, *_): return (h, i, 0) @@ -783,8 +923,20 @@ def make_splash_mha( use_experimental_scheduler: bool = False, vmem_limit_bytes: int | None = None, use_fixed_m: bool = False, + uniform_fixed_m: bool = False, ): - def _splash_attention(q, k, v, mk=None): + if use_fixed_m: + if not use_base2_exp: + raise NotImplementedError( + "fixed-m softmax bounds are derived strictly for base-2 exponents. Please set use_base2_exp=True." + ) + recenter, _ = get_fixed_m_constants(orig_kv_seq_len, is_ring=False) + else: + recenter = None + + def _splash_attention(q, k, v, mk=None, k_mean=None): + if use_fixed_m and mk is None: + raise ValueError("`mk` metadata array is required when `use_fixed_m=True`.") if heads_per_tile > 1: if use_fixed_m: raise NotImplementedError("fixed-m is not supported with heads_per_tile > 1") @@ -812,6 +964,9 @@ def _splash_attention(q, k, v, mk=None): vmem_limit_bytes=vmem_limit_bytes, use_fixed_m=use_fixed_m, mk=mk, + uniform_fixed_m=uniform_fixed_m, + fixed_m_recenter=recenter, + k_mean=k_mean, ) return _splash_attention diff --git a/src/maxdiffusion/kernels/fused_producers.py b/src/maxdiffusion/kernels/fused_producers.py new file mode 100644 index 000000000..708be811b --- /dev/null +++ b/src/maxdiffusion/kernels/fused_producers.py @@ -0,0 +1,102 @@ +"""Optimized fused producers for Wan Transformer Block and Attention.""" + +from typing import Tuple +import jax +import jax.numpy as jnp + + +def fused_ln_adaln( + x: jax.Array, + scale_msa: jax.Array, + shift_msa: jax.Array, + eps: float = 1e-6, +) -> jax.Array: + """Fusion-friendly FP32 LayerNorm + AdaLN scale/shift modulation producer.""" + x_fp32 = x.astype(jnp.float32) + mean = jnp.mean(x_fp32, axis=-1, keepdims=True) + diff = x_fp32 - mean + var = jnp.mean(jnp.square(diff), axis=-1, keepdims=True) + x_ln = diff * jax.lax.rsqrt(var + eps) + scale_term = 1.0 + scale_msa.astype(jnp.float32) + x_mod = (x_ln * scale_term + shift_msa.astype(jnp.float32)).astype(x.dtype) + return x_mod + + +def fused_rmsnorm_rope( + raw_q: jax.Array, + raw_k: jax.Array, + q_norm_scale: jax.Array, + k_norm_scale: jax.Array, + freqs_cis: jax.Array, + q_heads: int = 40, + kv_heads: int | None = None, + dim_head: int = 128, + eps: float = 1e-6, + heads: int | None = None, +) -> Tuple[jax.Array, jax.Array]: + """Fusion-friendly FP32 RMSNorm + BF16 RoPE + Head Transposition producer. + + Performs FP32 RMSNorm normalization for maximum stability, casts to input dtype + (e.g. BF16), and applies RoPE rotation and head transposition in BF16 precision, + avoiding excess FP32 VPU/VMEM cycles on long sequence lengths. Fully supports GQA + where q_heads != kv_heads. + + Args: + raw_q: Raw query projection of shape [B, Sq, Dq] (where Dq = q_heads * dim_head). + raw_k: Raw key projection of shape [B, Sk, Dk] (where Dk = kv_heads * dim_head). + q_norm_scale: RMSNorm scale parameter for query of shape [Dq]. + k_norm_scale: RMSNorm scale parameter for key of shape [Dk]. + freqs_cis: Complex rotary embedding tensor of shape [1, 1, S, dim_head // 2]. + q_heads: Number of query attention heads. + kv_heads: Number of key/value attention heads (defaults to q_heads for MHA). + dim_head: Dimension of each attention head. + eps: Epsilon for RMSNorm numerical stability. + heads: Deprecated alias for q_heads. + + Returns: + Transposed and RoPE-rotated (q_out, k_out) of shapes [B, q_heads, Sq, dim_head] + and [B, kv_heads, Sk, dim_head]. + """ + if heads is not None: + q_heads = heads + kv_heads = q_heads if kv_heads is None else kv_heads + B, Sq, Dq = raw_q.shape + _, Sk, Dk = raw_k.shape + + if Dq != q_heads * dim_head: + raise ValueError(f"raw_q feature dim ({Dq}) must equal q_heads ({q_heads}) * dim_head ({dim_head})") + if Dk != kv_heads * dim_head: + raise ValueError(f"raw_k feature dim ({Dk}) must equal kv_heads ({kv_heads}) * dim_head ({dim_head})") + + # 1. FP32 RMSNorm for stability, then cast directly to target activation dtype + q_fp32 = raw_q.astype(jnp.float32) + q_rms = jax.lax.rsqrt(jnp.mean(jnp.square(q_fp32), axis=-1, keepdims=True) + eps) + q_norm = (q_fp32 * q_rms * q_norm_scale.astype(jnp.float32)).astype(raw_q.dtype) + + k_fp32 = raw_k.astype(jnp.float32) + k_rms = jax.lax.rsqrt(jnp.mean(jnp.square(k_fp32), axis=-1, keepdims=True) + eps) + k_norm = (k_fp32 * k_rms * k_norm_scale.astype(jnp.float32)).astype(raw_k.dtype) + + # 2. Reshape and transpose to [B, heads, S, dim_head] + q_h = q_norm.reshape(B, Sq, q_heads, dim_head).transpose(0, 2, 1, 3) + k_h = k_norm.reshape(B, Sk, kv_heads, dim_head).transpose(0, 2, 1, 3) + + # 3. Direct RoPE with freqs_cis [1, 1, S, dim_head // 2] in input dtype + cos = jnp.real(freqs_cis).astype(raw_q.dtype) + sin = jnp.imag(freqs_cis).astype(raw_q.dtype) + cos_q, sin_q = cos[:, :, :Sq, :], sin[:, :, :Sq, :] + cos_k, sin_k = cos[:, :, :Sk, :], sin[:, :, :Sk, :] + + q_pairs = q_h.reshape(B, q_heads, Sq, -1, 2) + q_0, q_1 = q_pairs[..., 0], q_pairs[..., 1] + q_out_0 = q_0 * cos_q - q_1 * sin_q + q_out_1 = q_0 * sin_q + q_1 * cos_q + q_out = jnp.concatenate([q_out_0[..., None], q_out_1[..., None]], axis=-1).reshape(B, q_heads, Sq, dim_head) + + k_pairs = k_h.reshape(B, kv_heads, Sk, -1, 2) + k_0, k_1 = k_pairs[..., 0], k_pairs[..., 1] + k_out_0 = k_0 * cos_k - k_1 * sin_k + k_out_1 = k_0 * sin_k + k_1 * cos_k + k_out = jnp.concatenate([k_out_0[..., None], k_out_1[..., None]], axis=-1).reshape(B, kv_heads, Sk, dim_head) + + return q_out, k_out diff --git a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py index bc49c5af7..11611ebca 100644 --- a/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py +++ b/src/maxdiffusion/kernels/splash_attention/ring_attention_kernel.py @@ -168,7 +168,8 @@ def body( unroll=True, ) # type: ignore[arg-type] # Final normalization - assert l_final.dtype == jnp.float32 + if l_final.dtype != jnp.float32: + raise TypeError(f"l_final must have dtype float32, got {l_final.dtype}") l_inv = jnp.where(l_final == 0.0, 0.0, 1.0 / l_final) out = (o_final * l_inv[..., None]).astype(q.dtype) # Final logsumexp for residuals @@ -697,7 +698,8 @@ def make_ring_attention( is_dkv=True, return_dynamic_grid=config.dq_reduction_steps == 3, ) - assert (mask_function_fwd is None) == (mask_function_dkv is None) + if (mask_function_fwd is None) != (mask_function_dkv is None): + raise ValueError("mask_function_fwd and mask_function_dkv must both be None or both be provided") dkv_mask_sparsity = float(np.mean(dkv_mask_info.block_mask != 0)) dkv_mask_info = tree_util.tree_map(jnp.array, dkv_mask_info) @@ -760,6 +762,8 @@ def _custom_bidirectional_ring_forward( axis (no sub-group perm). """ axis_size = lax.axis_size(ring_axis) + effective_kv_seq_len = orig_kv_seq_len * axis_size + recenter, ring_safe_bound = custom_splash.get_fixed_m_constants(effective_kv_seq_len, is_ring=True) idx = lax.axis_index(ring_axis) exp_fn = jnp.exp2 if use_base2_exp else jnp.exp @@ -774,6 +778,7 @@ def _attn(kc, vc): use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, + fixed_m_recenter=recenter, ) return o.astype(jnp.float32), m.astype(jnp.float32), l.astype(jnp.float32) @@ -859,6 +864,11 @@ def _custom_ring_attention_forward( bidirectional: bool = False, use_fixed_m: bool = False, fixed_m_norms: tuple[jax.Array, jax.Array] | None = None, + per_q_block: bool = False, + pregathered_mk: bool = False, + k_mean: jax.Array | None = None, + uniform_fixed_m: bool | None = None, + v_ok: jax.Array | bool | None = None, ) -> jax.Array: """Forward-only ring attention using the custom dense splash kernel. @@ -882,17 +892,53 @@ def _custom_ring_attention_forward( mask_value: Initial running-max value for the online softmax. ring_axis: Name of the mesh axis to rotate K/V over (e.g. "context"). ring_size: Number of ring steps to scan over. Defaults to the full size of - `ring_axis`. For a hybrid Ulysses+Ring (USP) split this is the ring - sub-group size R (< full axis size), so each device only rotates within its - ring sub-group. + `ring_axis`. For fixed-m, ring_size must equal the size of ring_axis (2D + Ulysses+Ring should use a dedicated ring mesh axis). For online ring on a + flattened axis, this is the ring sub-group size R (< full axis size). perm: Explicit `ppermute` permutation. Defaults to a full-axis +1 rotation. - For the hybrid split, pass a perm that rotates K/V *within each ring - sub-group only* (built by the caller from the U x R factorization). + For fixed-m, the canonical ring permutation is required. For the online + hybrid split on a flattened axis, pass a perm that rotates K/V within each + ring sub-group only. + k_mean: Per-head global key mean vector (num_kv_heads, head_dim_qk) for + Global Virtual K-Centering. When provided or automatically reduced across + the ring axis, logit centering guarantees row max >= 0 across all hops. Returns: Normalized attention output, shape `(num_q_heads, q_seq_len, head_dim_v)`. """ axis_size = lax.axis_size(ring_axis) + effective_ring_size = ring_size if ring_size is not None else axis_size + effective_kv_seq_len = orig_kv_seq_len * effective_ring_size + + num_q_heads = q.shape[0] + num_kv_heads = k.shape[0] + head_dim_v = v.shape[-1] + + if use_fixed_m and not use_base2_exp: + raise NotImplementedError( + "fixed-m softmax bounds are derived strictly for base-2 exponents. Please set use_base2_exp=True." + ) + + if use_fixed_m and k_mean is None: + # Compute global key mean across ring ranks automatically on real (unpadded) tokens if not provided + actual_len = orig_kv_seq_len if orig_kv_seq_len is not None else k.shape[1] + k_mean_local = jnp.mean(k[:, :actual_len, :].astype(jnp.float32), axis=1) + k_mean = lax.pmean(k_mean_local, ring_axis) + + if use_fixed_m and num_q_heads != num_kv_heads: + if num_q_heads % num_kv_heads != 0: + raise ValueError( + f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA ring fixed-m." + ) + q_heads_per_kv_head = num_q_heads // num_kv_heads + if k_mean is not None and k_mean.shape[0] == num_kv_heads: + k_mean = jnp.repeat(k_mean, q_heads_per_kv_head, axis=0) + + if use_fixed_m and k_mean is not None and k_mean.shape[-1] < q.shape[-1]: + k_mean = jnp.pad(k_mean, ((0, 0), (0, q.shape[-1] - k_mean.shape[-1]))) + + global_recenter, global_centered_bound = custom_splash.get_fixed_m_constants(effective_kv_seq_len, is_ring=False) + local_recenter, per_shard_bound = custom_splash.get_fixed_m_constants(orig_kv_seq_len, is_ring=True) if bidirectional: if perm is not None or (ring_size is not None and ring_size != axis_size): raise ValueError( @@ -914,68 +960,91 @@ def _custom_ring_attention_forward( mask_value=mask_value, ring_axis=ring_axis, ) + if use_fixed_m and ring_size is not None and ring_size != axis_size: + raise NotImplementedError( + f"fixed-m ring attention requires ring_size == ring axis size (got ring_size={ring_size}, axis_size={axis_size}); " + "use a dedicated ring mesh axis for 2D Ulysses+Ring." + ) if ring_size is None: ring_size = axis_size + canonical_perm = [(i, (i + 1) % axis_size) for i in range(axis_size)] + if use_fixed_m and perm is not None and perm != canonical_perm: + raise NotImplementedError( + "fixed-m ring attention currently requires the canonical ring permutation " + f"[(i, (i + 1) % axis_size)], got perm={perm}." + ) if perm is None: - perm = [(i, (i + 1) % axis_size) for i in range(axis_size)] + perm = canonical_perm shift = partial(lax.ppermute, axis_name=ring_axis, perm=perm) exp_fn = jnp.exp2 if use_base2_exp else jnp.exp - 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. + # Fixed-m ring: with Global Virtual K-Centering, the keys are centered against + # the global mean across all ring ranks (k_mean = pmean(mean(k), ring_axis)). + # Because logit centering guarantees max_j (q^T (k_j - k_mean)) >= 0 globally, + # the centered Cauchy-Schwarz bound applies across the entire distributed sequence. + # We gather each rank's squared K-shard norms once before the scan: mk_all_sq (R, heads), + # and form mk_global_sq = mk_all_sq.max(axis=0). When all_fixed_global holds, every hop + # evaluates identical m_fixed, enabling direct FP32 (o_sum, l_sum) accumulation. + # In the hybrid fallback branch, individual hops do NOT have a zero-mean guarantee + # in isolation, so per-hop eligibility uses the conservative two-sided per_shard_bound. + # All Cauchy-Schwarz gating is computed in exact squared-norm space (|q|^2 * R_k^2 <= W^2) + # eliminating square roots from all query block and hop gating checks. if fixed_m_norms is None: - raise ValueError("use_fixed_m on the ring path requires fixed_m_norms=(qn_max, mk_h).") + raise ValueError("use_fixed_m on the ring path requires fixed_m_norms=(qn_max_sq, mk_h_sq).") log_fn = jnp.log2 if use_base2_exp else jnp.log - qn_max, mk_h_init = fixed_m_norms + qn_max_sq, mk_h_init_sq = fixed_m_norms + if num_q_heads != num_kv_heads and mk_h_init_sq.shape[-1] == num_kv_heads: + q_heads_per_kv_head = num_q_heads // num_kv_heads + mk_h_init_sq = jnp.repeat(mk_h_init_sq, q_heads_per_kv_head, axis=-1) 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 + # Every rank's squared K-shard norms, gathered ONCE before the scan: (R, heads). + # A pre-gathered array keeps the per-hop gate collective-free and avoids + # serializing a third ppermute alongside K/V transfers. + if pregathered_mk or (mk_h_init_sq.ndim == 2 and mk_h_init_sq.shape[0] == axis_size): + mk_all_sq = mk_h_init_sq else: - mk_all = lax.all_gather(mk_h_init, ring_axis) # (axis_size, heads) + mk_all_sq = lax.all_gather(mk_h_init_sq, 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 - ) + num_q_blocks = (orig_q_seq_len + block_sizes.block_q - 1) // block_sizes.block_q + mk_global_sq = mk_all_sq.max(axis=0) # (heads,) + + global_centered_bound_sq = global_centered_bound**2 + per_shard_bound_sq = per_shard_bound**2 + + # Global V-magnitude / dtype safety verdict. Unlike the Cauchy-Schwarz + # norm bounds this is NOT re-derivable from a single hop's Q/K, so it has + # to be carried in and applied to every eligibility decision below -- + # including the per-hop ones in `fixed_body`. + v_gate = True if v_ok is None else v_ok + + if not per_q_block: + bound_sq_1d = qn_max_sq * mk_global_sq + all_fixed_local = jnp.all(bound_sq_1d <= global_centered_bound_sq) & v_gate + all_fixed_global = lax.pmin(all_fixed_local, ring_axis) + m_base_1d = jnp.ceil(jnp.sqrt(bound_sq_1d)) - global_recenter + m_base_expanded = jnp.broadcast_to(m_base_1d[:, None], (num_q_heads, num_q_blocks)) + fixed_ok_expanded = jnp.ones_like(m_base_expanded) + mk_arr = jnp.stack([m_base_expanded, fixed_ok_expanded], axis=0) + qn_blocks_sq = jnp.broadcast_to(qn_max_sq[:, None], (num_q_heads, num_q_blocks)) + else: + qn_blocks_sq = qn_max_sq + bound_blocks_sq = qn_blocks_sq * mk_global_sq[:, None] + fixed_ok_local = bound_blocks_sq <= global_centered_bound_sq # pylint: disable=protected-access + all_fixed_local = jnp.all(fixed_ok_local) & v_gate + all_fixed_global = lax.pmin(all_fixed_local, ring_axis) + m_base = jnp.ceil(jnp.sqrt(bound_blocks_sq)) - global_recenter + fixed_ok_expanded = jnp.ones_like(m_base) + mk_arr = jnp.stack([m_base, fixed_ok_expanded], axis=0) # (2, heads, num_q_blocks) 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 @@ -1003,8 +1072,10 @@ def _accumulate_scan(_): use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, + fixed_m_recenter=global_recenter, use_fixed_m=True, mk=mk_arr, + k_mean=k_mean, # 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 @@ -1036,11 +1107,15 @@ def fixed_body(carry, hop, is_last_hop): # 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]) + mk_h_sq = jax.lax.dynamic_index_in_dim(mk_all_sq, (my_ring_index - hop) % axis_size, keepdims=False) + bound_hop_sq = qn_blocks_sq * mk_h_sq[:, None] + # `v_gate` is load-bearing here. The Cauchy-Schwarz term is per-hop, but + # V-magnitude and dtype safety are global; recomputing eligibility from + # Q/K norms alone would re-enable fixed-m on this hop even when the + # caller's global V check already rejected it, overflowing to inf. + fixed_ok = ((bound_hop_sq <= per_shard_bound_sq) & v_gate).astype(jnp.float32) # pylint: disable=protected-access + m_base_hop = jnp.ceil(jnp.sqrt(bound_hop_sq)) - local_recenter + mk_arr = jnp.stack([m_base_hop, fixed_ok], axis=0) o_curr, m_curr, l_curr = custom_splash._splash_attention_forward_ring( # pylint: disable=protected-access q, @@ -1052,8 +1127,10 @@ def fixed_body(carry, hop, is_last_hop): use_base2_exp=use_base2_exp, use_experimental_scheduler=use_experimental_scheduler, vmem_limit_bytes=vmem_limit_bytes, + fixed_m_recenter=local_recenter, use_fixed_m=True, mk=mk_arr, + k_mean=k_mean, ) m_curr = m_curr.astype(jnp.float32) l_curr = l_curr.astype(jnp.float32) @@ -1085,7 +1162,12 @@ def _lse_scan(_): 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) + if uniform_fixed_m is True: + return _accumulate_scan(None) + elif uniform_fixed_m is False: + return _lse_scan(None) + else: + 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) @@ -1132,7 +1214,6 @@ def _lse_scan(_): def make_custom_ring_attention( - *, block_sizes: "custom_splash._BlockSizes", orig_q_seq_len: int, orig_kv_seq_len: int, @@ -1146,24 +1227,33 @@ def make_custom_ring_attention( bidirectional: bool = False, use_fixed_m: bool = False, fixed_m_norms: tuple[jax.Array, jax.Array] | None = None, + per_q_block: bool = True, + pregathered_mk: bool = False, + k_mean: jax.Array | None = None, + uniform_fixed_m: bool | None = None, + v_ok: jax.Array | bool | None = None, ): """Builds a forward-only ring-attention callable around the custom kernel. The returned function takes a single (un-batched) `(q, k, v)` triple of shape - `(num_heads, seq, head_dim)` and is meant to be `jax.vmap`-ped over the batch - axis inside the attention `shard_map` (the `ppermute` rotates over `ring_axis`, - which is a mesh axis and independent of the vmap batch axis). + `(num_heads, seq, head_dim)` and optional per-batch `fixed_m_norms=(qn_max_sq, mk_h_sq)` + and `k_mean` to be `jax.vmap`-ped over the batch axis inside the attention `shard_map`. - `ring_size` / `perm` let a caller restrict the rotation to a ring sub-group of - the axis (for the hybrid Ulysses+Ring / USP split); when omitted the rotation - covers the whole `ring_axis`. - - `bidirectional=True` selects the wrap-free schedule (streams K/V both directions - one hop at a time) for a NON-wrapping ring axis, avoiding the diameter-length - wrap hop. Requires `perm=None` and the full real ring axis (no sub-group). + `v_ok` is a global (already cross-ring-reduced) scalar predicate asserting that + the value magnitudes and activation dtype are safe for fixed-m. It is closed + over rather than passed per call, since it is invariant across the batch. + Leaving it None preserves the previous behaviour of trusting the caller. """ + if use_fixed_m and not use_base2_exp: + raise NotImplementedError( + "fixed-m softmax bounds are derived strictly for base-2 exponents. Please set use_base2_exp=True." + ) + default_fixed_m_norms = fixed_m_norms + default_k_mean = k_mean - def _ring(q, k, v): + def _ring(q, k, v, fixed_m_norms=None, k_mean=None): + norms = fixed_m_norms if fixed_m_norms is not None else default_fixed_m_norms + km = k_mean if k_mean is not None else default_k_mean return _custom_ring_attention_forward( q, k, @@ -1180,7 +1270,12 @@ def _ring(q, k, v): perm=perm, bidirectional=bidirectional, use_fixed_m=use_fixed_m, - fixed_m_norms=fixed_m_norms, + fixed_m_norms=norms, + per_q_block=per_q_block, + pregathered_mk=pregathered_mk, + k_mean=km, + uniform_fixed_m=uniform_fixed_m, + v_ok=v_ok, ) return _ring diff --git a/src/maxdiffusion/max_utils.py b/src/maxdiffusion/max_utils.py index 4471e3bd1..ca0fa7b0d 100644 --- a/src/maxdiffusion/max_utils.py +++ b/src/maxdiffusion/max_utils.py @@ -468,10 +468,16 @@ def delete_file(file_path: str): max_logging.log(f"The file '{file_path}' does not exist.") -def get_git_commit_hash(): +def get_git_commit_hash(check_dirty: bool = True): """Tries to get the current Git commit hash, for run provenance.""" try: - return subprocess.check_output(["git", "rev-parse", "HEAD"]).strip().decode("utf-8") + repo_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + commit = subprocess.check_output(["git", "-C", repo_dir, "rev-parse", "HEAD"]).strip().decode("utf-8") + if check_dirty: + status = subprocess.check_output(["git", "-C", repo_dir, "status", "--porcelain"]).strip().decode("utf-8") + if status: + return f"{commit}-dirty" + return commit except subprocess.CalledProcessError: max_logging.log("Warning: 'git rev-parse HEAD' failed. Not running in a git repo?") return None diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py index 7b2ba0df7..a75040671 100644 --- a/src/maxdiffusion/models/attention_flax.py +++ b/src/maxdiffusion/models/attention_flax.py @@ -68,6 +68,18 @@ def _coerce_tokamax_block_sizes(block_sizes): + if isinstance(block_sizes, dict): + return splash_attention_kernel.BlockSizes( + block_q=block_sizes.get("block_q", 512), + block_kv=block_sizes.get("block_kv", 512), + block_kv_compute=block_sizes.get("block_kv_compute", 512), + block_q_dkv=block_sizes.get("block_q_dkv", 512), + block_kv_dkv=block_sizes.get("block_kv_dkv", 512), + block_kv_dkv_compute=block_sizes.get("block_kv_dkv_compute", 512), + block_q_dq=block_sizes.get("block_q_dq", None), + block_kv_dq=block_sizes.get("block_kv_dq", None), + use_fused_bwd_kernel=block_sizes.get("use_fused_bwd_kernel", False), + ) # Tokamax requires fused bwd; convert if needed. if getattr(block_sizes, "use_fused_bwd_kernel", False): return block_sizes @@ -523,14 +535,14 @@ def _ulysses_head_chunk_ranges(num_heads: int, ulysses_shards: int, num_chunks: head_groups = num_heads // ulysses_shards num_chunks = min(num_chunks, head_groups) - regular_groups_per_chunk = max(1, head_groups // num_chunks) + base_groups = head_groups // num_chunks + remainder_groups = head_groups % num_chunks ranges = [] start_group = 0 for chunk_idx in range(num_chunks): - end_group = head_groups if chunk_idx == num_chunks - 1 else min(start_group + regular_groups_per_chunk, head_groups) - if start_group >= end_group: - break + groups_in_chunk = base_groups + (1 if chunk_idx < remainder_groups else 0) + end_group = start_group + groups_in_chunk ranges.append((start_group * ulysses_shards, end_group * ulysses_shards)) start_group = end_group return ranges @@ -563,6 +575,11 @@ def _run_chunked_ulysses_attention( Returns: The concatenated attention output tensor. """ + if query.shape[1] != key.shape[1] and ulysses_attention_chunks > 1: + raise NotImplementedError( + f"GQA (query heads {query.shape[1]} != key heads {key.shape[1]}) with " + f"ulysses_attention_chunks={ulysses_attention_chunks} > 1 is not supported." + ) head_chunk_ranges = _ulysses_head_chunk_ranges(num_heads, ulysses_shards, ulysses_attention_chunks) if len(head_chunk_ranges) > 1: chunk_outputs = [ @@ -859,6 +876,8 @@ def _ulysses_attention( use_fixed_m: bool = False, ulysses_attention_chunks: int = 1, preserve_asymmetric_block_sizes: bool = False, + per_q_block: bool = True, + kv_heads: Optional[int] = None, ) -> jax.Array: """Ulysses sequence-parallel attention. @@ -869,24 +888,33 @@ def _ulysses_attention( """ axis_name = CONTEXT num_shards = mesh.shape[axis_name] + if kv_heads is None: + kv_heads = heads query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_shards) - key, _ = _reshape_data_for_flash(key, heads, num_shards) - value, _ = _reshape_data_for_flash(value, heads, num_shards) + key, orig_kv_seq_len = _reshape_data_for_flash(key, kv_heads, num_shards) + value, _ = _reshape_data_for_flash(value, kv_heads, num_shards) attention_mask = _prepare_attention_mask_for_shard_map(attention_mask, query.shape[0], key.shape[2]) if attention_mask is not None and use_custom_kernel: raise NotImplementedError( "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." ) - num_heads = query.shape[1] + num_q_heads = query.shape[1] + num_kv_heads = key.shape[1] # Ulysses only redistributes existing heads across the context mesh; unlike # the earlier draft, we fail fast instead of padding synthetic heads. - if num_heads % num_shards != 0: + if num_q_heads % num_shards != 0: + raise ValueError( + "Ulysses attention requires the number of query heads to be divisible by the context shard count, " + f"got q_heads={num_q_heads} and context_shards={num_shards}." + ) + if num_kv_heads % num_shards != 0: raise ValueError( - "Ulysses attention requires the number of heads to be divisible by the context shard count, " - f"got heads={num_heads} and context_shards={num_shards}." + "Ulysses attention requires the number of KV heads to be divisible by the context shard count, " + f"got kv_heads={num_kv_heads} and context_shards={num_shards}." ) + num_heads = num_q_heads if not use_custom_kernel: block_sizes = _select_flash_block_sizes( @@ -923,27 +951,42 @@ def wrap_ulysses_attention(query, key, value, attention_mask): if use_base2_exp: query = query * LOG2E + raw_key = key + raw_query = query + raw_value = value + context_q_seq_len = raw_query.shape[2] + actual_kv_seq_len = orig_kv_seq_len + + real_key = raw_key[:, :, :actual_kv_seq_len, :] + + recenter, safe_bound = custom_splash.get_fixed_m_constants(actual_kv_seq_len, is_ring=False) + + k_mean = None if use_fixed_m: - # k-smoothing (output-invariant): subtracting the per-row key mean - # forces every logit row to have mean 0, hence row-max >= 0 — the - # precondition that keeps the fixed-m Cauchy-Schwarz bound flush-free. - key = key - jnp.mean(key, axis=2, keepdims=True) + # Virtual k-centering (output-invariant): project q^T \bar{k} inside the + # kernel registers without writing back / materializing (K - \bar{k}) in HBM. + # Computed strictly on real (unpadded) tokens. + k_mean = jnp.mean(real_key.astype(jnp.float32), axis=2) + if k_mean.shape[-1] < 128: + k_mean = jnp.pad(k_mean, ((0, 0), (0, 0), (0, 128 - k_mean.shape[-1]))) - query, kv_size, query_seq_len = _pad_data_for_flash(query, heads, bq) - key, _, key_seq_len = _pad_data_for_flash(key, heads, bkv) - value, _, _ = _pad_data_for_flash(value, heads, bkv) + query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) + key, _, key_seq_len = _pad_data_for_flash(raw_key, heads, bkv) + value, _, _ = _pad_data_for_flash(raw_value, heads, bkv) mk_arr = None + all_fixed = None if use_fixed_m: - # Per-(local-)head Cauchy-Schwarz inputs over the (batch, seq) slice; - # padded rows have zero norm and never raise the max. mk[0] feeds the - # in-kernel per-query bound, mk[1] flags heads within the no-flush gate. - 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,) - fixed_ok = (qn_max * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(jnp.float32) - mk_arr = jnp.stack([mk_h, fixed_ok]) # (2, local_heads) + mk_arr, all_fixed = _compute_fixed_m_metadata( + query, + real_key, + block_q=bq, + safe_bound=safe_bound, + recenter=recenter, + per_q_block=per_q_block, + k_mean=k_mean, + value=value, + ) bsizes = custom_splash._BlockSizes( block_q=bq, @@ -952,25 +995,61 @@ def wrap_ulysses_attention(query, key, value, attention_mask): block_kv_compute_in=bkv_compute_in, ) - splash_kernel = custom_splash.make_splash_mha( - block_sizes=bsizes, - orig_q_seq_len=query_seq_len, - orig_kv_seq_len=key_seq_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, - use_fixed_m=use_fixed_m, - ) - if use_fixed_m: - vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0, None)) - attention_output = vmapped_splash(query, key, value, mk_arr) + splash_kernel_uniform = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=context_q_seq_len, + orig_kv_seq_len=actual_kv_seq_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, + use_fixed_m=True, + uniform_fixed_m=True, + ) + splash_kernel_hybrid = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=context_q_seq_len, + orig_kv_seq_len=actual_kv_seq_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, + use_fixed_m=True, + uniform_fixed_m=False, + ) + + def _run_uniform(q, k, v, m, km): + return jax.vmap(splash_kernel_uniform, in_axes=(0, 0, 0, 0, 0))(q, k, v, m, km) + + def _run_hybrid(q, k, v, m, km): + return jax.vmap(splash_kernel_hybrid, in_axes=(0, 0, 0, 0, 0))(q, k, v, m, km) + + attention_output = jax.lax.cond(all_fixed, _run_uniform, _run_hybrid, query, key, value, mk_arr, k_mean) else: + splash_kernel = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=context_q_seq_len, + orig_kv_seq_len=actual_kv_seq_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, + use_fixed_m=False, + ) vmapped_splash = jax.vmap(splash_kernel, in_axes=(0, 0, 0)) attention_output = vmapped_splash(query, key, value) - attention_output = jnp.swapaxes(attention_output, 2, 3) - attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype) + attention_output = attention_output[:, :, :kv_size, :context_q_seq_len].astype(query.dtype) + # Restore original layout: head-sharded/full-sequence -> sequence-sharded/full-heads. + # Sequence axis is at index 3, heads axis is at index 1. + attention_output = jax.lax.all_to_all( + attention_output, + axis_name=axis_name, + split_axis=3, + concat_axis=1, + tiled=True, + ) + return attention_output else: # Run the same local splash kernel as standard TPU flash attention, but now # on full-sequence / fewer-heads tensors produced by the all-to-all above. @@ -1010,15 +1089,15 @@ def wrap_ulysses_attention(query, key, value, attention_mask): attention_output = vmapped_splash(query, key, value, segment_ids) 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, - ) - return attention_output + # 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, + ) + return attention_output devices_in_batch_sharding = mesh.shape["data"] * (mesh.shape["fsdp"] if "fsdp" in mesh.shape else 1) if not (query.shape[0] / devices_in_batch_sharding).is_integer(): @@ -1043,7 +1122,11 @@ def wrap_ulysses_attention(query, key, value, attention_mask): # Folding batch into heads destroys the one-mask-per-example association. # Keep the optimization for the common unmasked path only. fold_batch = ( - attention_mask is None and batch > 1 and devices_in_batch_sharding == 1 and (batch * num_heads) % num_shards == 0 + attention_mask is None + and batch > 1 + and devices_in_batch_sharding == 1 + and num_q_heads == num_kv_heads + and (batch * num_heads) % num_shards == 0 ) if fold_batch: query = query.reshape(1, batch * num_heads, *query.shape[2:]) @@ -1053,12 +1136,18 @@ def wrap_ulysses_attention(query, key, value, attention_mask): else: effective_num_heads = num_heads + out_q_axis_names = ( + jax.sharding.PartitionSpec(q_axis_names[0], q_axis_names[1], q_axis_names[3], q_axis_names[2]) + if use_custom_kernel + else q_axis_names + ) + if attention_mask is None: sharded_ulysses_attention = jax.shard_map( lambda q, k, v: wrap_ulysses_attention(q, k, v, None), mesh=mesh, in_specs=(q_axis_names, kv_axis_names, kv_axis_names), - out_specs=q_axis_names, + out_specs=out_q_axis_names, check_vma=False, ) @@ -1070,7 +1159,7 @@ def run_ulysses_attention(q, k, v): wrap_ulysses_attention, mesh=mesh, in_specs=(q_axis_names, kv_axis_names, kv_axis_names, mask_axis_names), - out_specs=q_axis_names, + out_specs=out_q_axis_names, check_vma=False, ) @@ -1087,10 +1176,19 @@ def run_ulysses_attention(q, k, v): run_ulysses_attention, ) - if fold_batch: - x = x.reshape(batch, num_heads, *x.shape[2:]) - x = x[:, :, :orig_q_seq_len, :] - x = _reshape_heads_to_head_dim(x) + if use_custom_kernel: + if fold_batch: + x = x.reshape(batch, num_heads, *x.shape[2:]) + x = x[:, :, :, :orig_q_seq_len] + b, h, d, s = x.shape + x = jnp.transpose(x, (0, 3, 1, 2)).reshape(b, -1, h * d) + axis_names = nn.logical_to_mesh_axes((BATCH, LENGTH, HEAD)) + x = jax.lax.with_sharding_constraint(x, axis_names) + else: + if fold_batch: + x = x.reshape(batch, num_heads, *x.shape[2:]) + x = x[:, :, :orig_q_seq_len, :] + x = _reshape_heads_to_head_dim(x) return x @@ -1115,6 +1213,7 @@ def _ulysses_ring_attention( ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, preserve_asymmetric_block_sizes: bool = False, + kv_heads: int | None = None, ) -> jax.Array: """2D context-parallel attention using a private Ulysses x ring mesh. @@ -1123,6 +1222,8 @@ def _ulysses_ring_attention( Ulysses all-to-all over the hidden Ulysses axis, and rotates K/V over the hidden ring axis. """ + if kv_heads is None: + kv_heads = heads context_axis = CONTEXT if context_axis not in mesh.shape: @@ -1139,9 +1240,14 @@ def _ulysses_ring_attention( ) if heads % num_ulysses_shards != 0: raise ValueError( - "Ulysses ring attention requires the number of heads to be divisible by the requested Ulysses shard count, " + "Ulysses ring attention requires the number of query heads to be divisible by the requested Ulysses shard count, " f"got heads={heads} and ulysses_shards={num_ulysses_shards}." ) + if kv_heads % num_ulysses_shards != 0: + raise ValueError( + "Ulysses ring attention requires the number of KV heads to be divisible by the requested Ulysses shard count, " + f"got kv_heads={kv_heads} and ulysses_shards={num_ulysses_shards}." + ) num_ring_shards = num_context_shards // num_ulysses_shards internal_mesh = _create_internal_ulysses_ring_mesh( mesh, @@ -1154,8 +1260,8 @@ def _ulysses_ring_attention( num_sequence_shards = num_context_shards query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_sequence_shards) - key, _ = _reshape_data_for_flash(key, heads, num_sequence_shards) - value, _ = _reshape_data_for_flash(value, heads, num_sequence_shards) + key, _ = _reshape_data_for_flash(key, kv_heads, num_sequence_shards) + value, _ = _reshape_data_for_flash(value, kv_heads, num_sequence_shards) attention_mask = _prepare_attention_mask_for_shard_map(attention_mask, query.shape[0], key.shape[2]) num_heads = query.shape[1] @@ -1198,8 +1304,8 @@ def wrap_ulysses_ring_attention(query, key, value, attention_mask): block_q = max(*block_q_sizes) query, kv_size, query_seq_len = _pad_data_for_flash(query, heads, block_q) block_kv = max(*block_kv_sizes) - key, _, key_seq_len = _pad_data_for_flash(key, heads, block_kv) - value, _, _ = _pad_data_for_flash(value, heads, block_kv) + key, _, key_seq_len = _pad_data_for_flash(key, kv_heads, block_kv) + value, _, _ = _pad_data_for_flash(value, kv_heads, block_kv) q_padded_len = query.shape[2] kv_padded_len = key.shape[2] @@ -1297,10 +1403,98 @@ def run_ulysses_ring_attention(q, k, v): 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 + """Largest row L2 norm estimate per head of a `[..., H, S, D]` activation in FP32 with conservative exponent margin, preserving batch dims.""" + row_sq = (x.astype(jnp.float32) ** 2).sum(axis=-1) + return jnp.sqrt(row_sq.max(axis=-1)) + + +def _compute_fixed_m_metadata( + query: jax.Array, + key: jax.Array, + block_q: int, + safe_bound: float | None = None, + recenter: float | None = None, + per_q_block: bool = True, + k_mean: jax.Array | None = None, + value: jax.Array | None = None, + v_max_bound: float = 256.0, +) -> tuple[jax.Array, jax.Array]: + """Computes Cauchy-Schwarz norm bounds and per-Q-block (or per-head) fixed-m metadata. + + Args: + query: Padded query activation, shape `(batch, local_heads, padded_q_len, head_dim)`. + key: Key activation (raw unpadded or padded), shape `(batch, local_heads, kv_len, head_dim)`. + block_q: Query tile block size. + safe_bound: Maximum safe norm product threshold before falling back to online softmax. + recenter: Fixed-m dynamic recenter constant C(N). + per_q_block: If True, evaluates gating independently per query tile. If False, + evaluates monolithic gating per head. + k_mean: Optional mean key vector for Virtual K-centering, shape `(batch, local_heads, head_dim)`. + value: Optional value activation, shape `(batch, local_heads, kv_len, head_dim_v)`, used to + verify that |V| <= v_max_bound to guarantee against FP32 overflow. + v_max_bound: Maximum safe value magnitude (default 256.0). + + Returns: + mk_arr: Gating metadata array of shape `(batch, 2, local_heads, num_q_blocks)` + multiplexing precomputed block base shifts and binary fixed-m gating predicates into a single + Pallas scalar prefetch memory slot: + - `mk_arr[:, 0, h, i]`: Precomputed block base shift m_B = ceil(max_i ||q_i|| * max_j ||k_j||) - C. + - `mk_arr[:, 1, h, i]`: Discrete eligibility predicate (1.0 for fixed-m, 0.0 for online). + all_fixed: Boolean scalar indicating if all elements are eligible for uniform fixed-m. + """ + batch_size, num_q_heads, q_len, _ = query.shape + num_kv_heads = key.shape[1] + if safe_bound is None or recenter is None: + rec, bnd = custom_splash.get_fixed_m_constants(key.shape[2], is_ring=False, v_max_bound=v_max_bound) + if safe_bound is None: + safe_bound = bnd + if recenter is None: + recenter = rec + safe_bound_sq = safe_bound**2 + if k_mean is not None: + centered_k = key.astype(jnp.float32) - k_mean[:, :, None, : key.shape[-1]] + mk_h_sq = (centered_k**2).sum(axis=-1).max(axis=-1) + else: + mk_h_sq = (key.astype(jnp.float32) ** 2).sum(axis=-1).max(axis=-1) # (batch, num_kv_heads) + + if num_q_heads != num_kv_heads: + if num_q_heads % num_kv_heads != 0: + raise ValueError(f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA fixed-m.") + q_heads_per_kv_head = num_q_heads // num_kv_heads + mk_h_sq = jnp.repeat(mk_h_sq, q_heads_per_kv_head, axis=1) # (batch, num_q_heads) + + # Fixed-m weights reach 2**recenter before being narrowed to the activation + # dtype for the S@V matmul. If that dtype's exponent range cannot hold them + # (fp16, fp8), the FP32 bound analysis is irrelevant -- the narrowing itself + # overflows to inf -- so disqualify every head up front. + dtype_safe = custom_splash.fixed_m_dtype_is_safe(query.dtype, recenter) + v_ok = 1.0 if dtype_safe else 0.0 + if dtype_safe and value is not None: + v_max_sq = (value.astype(jnp.float32) ** 2).max() + v_ok = (v_max_sq <= (v_max_bound**2)).astype(jnp.float32) + + num_q_blocks = q_len // block_q + if per_q_block: + norm_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1) # (batch, num_q_heads, q_len) + qn_max_sq = norm_sq.reshape(batch_size, num_q_heads, num_q_blocks, block_q).max( + axis=-1 + ) # (batch, num_q_heads, num_q_blocks) + bound_sq = qn_max_sq * mk_h_sq[:, :, None] + fixed_ok = (bound_sq <= safe_bound_sq).astype(jnp.float32) * v_ok + m_base = jnp.ceil(jnp.sqrt(bound_sq)) - recenter + mk_arr = jnp.stack([m_base, fixed_ok], axis=1) # (batch, 2, num_q_heads, num_q_blocks) + all_fixed = jnp.all(fixed_ok > 0.5) + else: + qn_max_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1).max(axis=-1) # (batch, num_q_heads) + bound_sq_1d = qn_max_sq * mk_h_sq + fixed_ok_1d = (bound_sq_1d <= safe_bound_sq).astype(jnp.float32) * v_ok + m_base_1d = jnp.ceil(jnp.sqrt(bound_sq_1d)) - recenter + m_base_expanded = jnp.broadcast_to(m_base_1d[:, :, None], (batch_size, num_q_heads, num_q_blocks)) + fixed_ok_expanded = jnp.broadcast_to(fixed_ok_1d[:, :, None], (batch_size, num_q_heads, num_q_blocks)) + mk_arr = jnp.stack([m_base_expanded, fixed_ok_expanded], axis=1) # (batch, 2, num_q_heads, num_q_blocks) + all_fixed = jnp.all(fixed_ok_1d > 0.5) + + return mk_arr, all_fixed def _ulysses_ring_custom_attention( @@ -1322,28 +1516,15 @@ def _ulysses_ring_custom_attention( bidirectional: bool = False, use_fixed_m: bool = False, ulysses_attention_chunks: int = 1, + per_q_block: bool = True, + kv_heads: int | None = None, ) -> jax.Array: - """Hybrid Ulysses + Ring (USP) with the CUSTOM splash kernel on main's mesh. + """2D USP attention (Ulysses + Ring) using custom splash kernel with exact Fixed-m support.""" + if kv_heads is None: + kv_heads = heads - Uses origin/main's explicit internal `(ring, ulysses)` mesh - (`_create_internal_ulysses_ring_mesh`, commit c104db51) instead of single-axis - collective sub-groups: the public `context` axis is reshaped with the Ulysses - axis innermost, so the Ulysses all-to-all stays INTRA-chip and the ring rotates - ACROSS chips. The per-shard attention is our custom splash kernel - (`make_custom_ring_attention`), not the tokamax_ring kernel main uses. - - 1. all-to-all over the (intra-chip) Ulysses axis: trade sequence for heads; - 2. ring (full ppermute) over the (cross-chip) ring axis, online-softmax merge; - 3. all-to-all back to restore the sequence-sharded / full-heads layout. - - U = ulysses_shards (from config); R = context // U. U=context -> pure - Ulysses, U=1 -> pure Ring (all on the same custom kernel). - """ if attention_mask is not None: - raise NotImplementedError( - "ulysses_ring_custom does not support attention_mask (the custom splash kernels only " - "handle padding via orig_seq_len); got a non-None attention_mask." - ) + raise NotImplementedError("ulysses_ring_custom does not support attention_mask.") axis_name = "context" num_context_shards = mesh.shape[axis_name] num_ulysses_shards = ulysses_shards @@ -1355,29 +1536,24 @@ 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 - query, orig_q_seq_len = _reshape_data_for_flash(query, heads, num_context_shards) - key, _ = _reshape_data_for_flash(key, heads, num_context_shards) - value, _ = _reshape_data_for_flash(value, heads, num_context_shards) + key, orig_kv_seq_len = _reshape_data_for_flash(key, kv_heads, num_context_shards) + value, _ = _reshape_data_for_flash(value, kv_heads, num_context_shards) num_heads = query.shape[1] 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) - if heads_per_tile > 1: - raise NotImplementedError("ulysses_ring_custom currently supports heads_per_tile == 1 only.") - + raise ValueError(f"Ulysses+Ring requires query heads divisible by U={num_ulysses_shards}, got heads={num_heads}.") + if kv_heads % num_ulysses_shards != 0: + raise ValueError(f"Ulysses+Ring requires KV heads divisible by U={num_ulysses_shards}, got kv_heads={kv_heads}.") + if num_ring_shards > 1 and orig_kv_seq_len % num_context_shards != 0: + raise ValueError( + f"2D Ulysses+Ring attention requires sequence length to be divisible by context_shards={num_context_shards}, " + f"got orig_kv_seq_len={orig_kv_seq_len}." + ) + 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 and num_ring_shards > 1: + raise NotImplementedError("heads_per_tile > 1 is not supported for multi-shard ring attention.") internal_mesh = _create_internal_ulysses_ring_mesh(mesh, num_ring_shards, num_ulysses_shards) - ring_axis = INTERNAL_RING_AXIS - ulysses_axis = INTERNAL_ULYSSES_AXIS - + ring_axis, ulysses_axis = INTERNAL_RING_AXIS, INTERNAL_ULYSSES_AXIS q_axis_names = nn.logical_to_mesh_axes(axis_names_q) kv_axis_names = nn.logical_to_mesh_axes(axis_names_kv) internal_q_axis_names = _replace_mesh_axis_names(q_axis_names, axis_name, (ring_axis, ulysses_axis)) @@ -1386,48 +1562,12 @@ 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. + # (1) Ulysses All-to-All: heads -> sequence a2a = functools.partial(jax.lax.all_to_all, axis_name=ulysses_axis, tiled=True) query = a2a(query, split_axis=1, concat_axis=2) key = a2a(key, split_axis=1, concat_axis=2) @@ -1436,86 +1576,220 @@ def wrap_ulysses_ring_attention(query, key, value): if use_base2_exp: query = query * LOG2E - 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 - - query, kv_size, query_seq_len = _pad_data_for_flash(query, heads, bq) - key, _, key_seq_len = _pad_data_for_flash(key, heads, bkv) - value, _, _ = _pad_data_for_flash(value, heads, bkv) - - mk_arr = None - 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) - 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, - ) + raw_key = key + raw_query = query + raw_value = value + context_q_seq_len = raw_query.shape[2] + actual_kv_seq_len = orig_kv_seq_len if num_ring_shards == 1 else raw_key.shape[2] + + real_key = raw_key[:, :, :actual_kv_seq_len, :] + + k_mean = None + if use_fixed_m: + # Global Virtual K-Centering: + # Compute local key mean after Ulysses A2A on real unpadded tokens, then all-reduce across ring axis if R > 1 + k_mean_local = jnp.mean(real_key.astype(jnp.float32), axis=2) + if num_ring_shards > 1: + k_mean = jax.lax.pmean(k_mean_local, axis_name=ring_axis) + else: + k_mean = k_mean_local + if k_mean.shape[-1] < 128: + k_mean = jnp.pad(k_mean, ((0, 0), (0, 0), (0, 128 - k_mean.shape[-1]))) + + query, kv_size, query_seq_len = _pad_data_for_flash(raw_query, heads, bq) + key, _, key_seq_len = _pad_data_for_flash(raw_key, kv_heads, bkv) + value, _, _ = _pad_data_for_flash(raw_value, kv_heads, bkv) + + mk_arr, all_fixed, qn_dev, mk_dev = None, None, None, None + if use_fixed_m: + if num_ring_shards == 1: + recenter, safe_bound = custom_splash.get_fixed_m_constants(actual_kv_seq_len, is_ring=False) + mk_arr, all_fixed = _compute_fixed_m_metadata( + query, + real_key, + bq, + safe_bound=safe_bound, + recenter=recenter, + per_q_block=per_q_block, + k_mean=k_mean, + value=raw_value, + ) + else: + batch_size, num_q_heads, q_seq, _ = query.shape + num_kv_heads = raw_key.shape[1] + norm_sq = (query.astype(jnp.float32) ** 2).sum(axis=-1) + if per_q_block: + num_q_blocks = q_seq // bq + qn_dev = norm_sq.reshape(batch_size, num_q_heads, num_q_blocks, bq).max(axis=-1) + else: + qn_dev = norm_sq.max(axis=-1) + kf_centered = raw_key.astype(jnp.float32) - k_mean[:, :, None, : raw_key.shape[-1]] + mk_dev = jnp.sum(kf_centered**2, axis=-1).max(axis=-1) + if num_q_heads != num_kv_heads: + if num_q_heads % num_kv_heads != 0: + raise ValueError( + f"num_q_heads ({num_q_heads}) must be divisible by num_kv_heads ({num_kv_heads}) for GQA ring fixed-m." + ) + q_heads_per_kv_head = num_q_heads // num_kv_heads + mk_dev = jnp.repeat(mk_dev, q_heads_per_kv_head, axis=1) + k_mean = jnp.repeat(k_mean, q_heads_per_kv_head, axis=1) + + bsizes = custom_splash._BlockSizes(bq, bkv, bkv_compute, bkv_compute_in) + + # (2a) R=1: Dedicated single-device splash kernel with fixed-m or online softmax if num_ring_shards == 1: - # (2a) R=1: the ring is trivial (no rotation) -> use the lighter dedicated - # splash kernel (fuse_reciprocal, no fp32 online-softmax residual windows). - # Same math as the 1-step ring, and it fits BQ=8448 where the ring kernel - # OOMs (its 3x residual windows). make_splash_mha returns [H, D, S]. - splash_kernel = custom_splash.make_splash_mha( - block_sizes=bsizes, - orig_q_seq_len=query_seq_len, - orig_kv_seq_len=key_seq_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, - 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, + splash_kernel_uniform = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=context_q_seq_len, + orig_kv_seq_len=actual_kv_seq_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, + use_fixed_m=True, + uniform_fixed_m=True, ) + splash_kernel_hybrid = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=context_q_seq_len, + orig_kv_seq_len=actual_kv_seq_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, + use_fixed_m=True, + uniform_fixed_m=False, + ) + + def _run_uniform(q, k, v, m, km): + return jax.vmap(splash_kernel_uniform, in_axes=(0, 0, 0, 0, 0))(q, k, v, m, km) + + def _run_hybrid(q, k, v, m, km): + return jax.vmap(splash_kernel_hybrid, in_axes=(0, 0, 0, 0, 0))(q, k, v, m, km) + + raw_out = jax.lax.cond(all_fixed, _run_uniform, _run_hybrid, query, key, value, mk_arr, k_mean) else: - attention_output = jnp.swapaxes(jax.vmap(splash_kernel, in_axes=(0, 0, 0))(query, key, value), 2, 3) + splash_kernel = custom_splash.make_splash_mha( + block_sizes=bsizes, + orig_q_seq_len=context_q_seq_len, + orig_kv_seq_len=actual_kv_seq_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, + use_fixed_m=False, + ) + raw_out = jax.vmap(splash_kernel, in_axes=(0, 0, 0))(query, key, value) + attention_output = jnp.swapaxes(raw_out, 2, 3) + + # (2b) Ring: Cross-chip ppermute schedule with custom ring kernel else: - # (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) - attention_output = attention_output[:, :, :query_seq_len, :kv_size].astype(query.dtype) + if use_fixed_m: + effective_kv_seq_len = key_seq_len * num_ring_shards + global_recenter, global_centered_bound = custom_splash.get_fixed_m_constants(effective_kv_seq_len, is_ring=False) + global_centered_bound_sq = global_centered_bound**2 + + # Pre-gather mk_dev across the ring to evaluate global eligibility outside vmap + mk_all_gathered = jax.lax.all_gather(mk_dev, ring_axis) # (ring_size, batch, heads) + mk_all_sq = jnp.swapaxes(mk_all_gathered, 0, 1) # (batch, ring_size, heads) + mk_global_sq = mk_all_sq.max(axis=1) # (batch, heads) + + # V-magnitude and dtype safety are properties of the *whole* distributed + # problem, not of any single hop. Both are reduced to one global scalar + # here and -- critically -- also handed to the kernels themselves: the + # LSE fallback re-derives per-hop eligibility from Q/K norms alone, so + # without this it would happily re-enable fixed-m on a hop after the + # global V check had already rejected it. + dtype_safe = custom_splash.fixed_m_dtype_is_safe(query.dtype, global_recenter) + v_max_sq = (raw_value.astype(jnp.float32) ** 2).max() + v_ok_local = (v_max_sq <= (custom_splash.DEFAULT_MAX_V_BOUND**2)) & dtype_safe + v_ok = jax.lax.pmin(v_ok_local, ring_axis) + + if not per_q_block: + bound_sq_1d = qn_dev * mk_global_sq + fixed_ok = (bound_sq_1d <= global_centered_bound_sq) & v_ok + else: + bound_blocks_sq = qn_dev * mk_global_sq[:, :, None] + fixed_ok = (bound_blocks_sq <= global_centered_bound_sq) & v_ok - # (3) Ulysses all-to-all back: sequence -> heads, restoring the layout. - attention_output = a2a(attention_output, split_axis=2, concat_axis=1) - return attention_output + all_fixed_local = jnp.all(fixed_ok) + all_fixed_global = jax.lax.pmin(all_fixed_local, ring_axis) + + ring_kernel_accumulate = 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=True, + per_q_block=per_q_block, + pregathered_mk=True, + uniform_fixed_m=True, + v_ok=v_ok, + ) + ring_kernel_lse = 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=True, + per_q_block=per_q_block, + pregathered_mk=True, + uniform_fixed_m=False, + v_ok=v_ok, + ) + + def _run_ring_accumulate(q, k, v, norms, km): + return jax.vmap(ring_kernel_accumulate, in_axes=(0, 0, 0, (0, 0), 0))(q, k, v, norms, km) + + def _run_ring_lse(q, k, v, norms, km): + return jax.vmap(ring_kernel_lse, in_axes=(0, 0, 0, (0, 0), 0))(q, k, v, norms, km) + + attention_output = jax.lax.cond( + all_fixed_global, + _run_ring_accumulate, + _run_ring_lse, + query, + key, + value, + (qn_dev, mk_all_sq), + k_mean, + ) + 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=False, + ) + attention_output = jax.vmap(ring_kernel, in_axes=(0, 0, 0))(query, key, value) + + attention_output = attention_output[:, :, :context_q_seq_len, :kv_size].astype(query.dtype) + + # (3) Ulysses All-to-All back: sequence -> heads + return a2a(attention_output, split_axis=2, concat_axis=1) x = _run_chunked_ulysses_attention( query, key, value, - num_heads, + heads, num_ulysses_shards, ulysses_attention_chunks, wrap_ulysses_ring_attention, @@ -1538,17 +1812,33 @@ def _apply_attention_dot( float32_qk_product: bool, use_memory_efficient_attention: bool, attention_mask: Array = None, + kv_heads: int | None = None, ): """Apply Attention.""" + effective_kv_heads = kv_heads if kv_heads is not None else heads if split_head_dim: b = key.shape[0] query_states = jnp.reshape(query, (b, -1, heads, dim_head)) - key_states = jnp.reshape(key, (b, -1, heads, dim_head)) - value_states = jnp.reshape(value, (b, -1, heads, dim_head)) + key_states = jnp.reshape(key, (b, -1, effective_kv_heads, dim_head)) + value_states = jnp.reshape(value, (b, -1, effective_kv_heads, dim_head)) + if heads != effective_kv_heads: + num_repeats = heads // effective_kv_heads + key_states = jnp.repeat(key_states, num_repeats, axis=2) + value_states = jnp.repeat(value_states, num_repeats, axis=2) else: query_states = _reshape_heads_to_batch_dim(query, heads) - key_states = _reshape_heads_to_batch_dim(key, heads) - value_states = _reshape_heads_to_batch_dim(value, heads) + key_states = _reshape_heads_to_batch_dim(key, effective_kv_heads) + value_states = _reshape_heads_to_batch_dim(value, effective_kv_heads) + if heads != effective_kv_heads: + num_repeats = heads // effective_kv_heads + b = query.shape[0] + s_k = key_states.shape[1] + key_states = jnp.repeat(key_states.reshape(b, effective_kv_heads, s_k, -1), num_repeats, axis=1).reshape( + b * heads, s_k, -1 + ) + value_states = jnp.repeat(value_states.reshape(b, effective_kv_heads, s_k, -1), num_repeats, axis=1).reshape( + b * heads, s_k, -1 + ) if float32_qk_product: query_states = query_states.astype(jnp.float32) @@ -1654,6 +1944,7 @@ def dot_product_kernel(q, k, v, context): context["float32_qk_product"], context["use_memory_efficient_attention"], context["attention_mask"], + kv_heads=context.get("kv_heads", None), ) @@ -1676,6 +1967,7 @@ def ulysses_custom_kernel(q, k, v, context): use_base2_exp=context.get("use_base2_exp", True), use_experimental_scheduler=context.get("use_experimental_scheduler", False), ulysses_attention_chunks=context["ulysses_attention_chunks"], + kv_heads=context.get("kv_heads", None), ) @@ -1698,16 +1990,13 @@ def ulysses_ring_custom_kernel(q, k, v, context): use_base2_exp=context.get("use_base2_exp", True), use_experimental_scheduler=context.get("use_experimental_scheduler", False), ulysses_attention_chunks=context["ulysses_attention_chunks"], + kv_heads=context.get("kv_heads", None), ) @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 - uses the Cauchy-Schwarz fixed-m softmax bound (no in-kernel running-max - rescale). max||k|| and the K-smoothing mean are taken LOCALLY per ring shard - (no per-layer ring collective); the outer ring online-softmax merge still - re-normalizes across shards, so per-shard bounds stay correct.""" + """fixed-m variant of ulysses_ring_custom with monolithic per-head gating.""" return _ulysses_ring_custom_attention( q, k * context["scale"], @@ -1725,7 +2014,35 @@ def ulysses_ring_custom_fixed_m_kernel(q, k, v, context): use_base2_exp=context.get("use_base2_exp", True), use_experimental_scheduler=context.get("use_experimental_scheduler", False), use_fixed_m=True, + per_q_block=False, ulysses_attention_chunks=context.get("ulysses_attention_chunks", 1), + kv_heads=context.get("kv_heads", None), + ) + + +@register_kernel("ulysses_ring_custom_fixed_m_per_q_block") +def ulysses_ring_custom_fixed_m_per_q_block_kernel(q, k, v, context): + """fixed-m variant of ulysses_ring_custom with per-Q-block gating.""" + 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), + use_fixed_m=True, + per_q_block=True, + ulysses_attention_chunks=context.get("ulysses_attention_chunks", 1), + kv_heads=context.get("kv_heads", None), ) @@ -1752,6 +2069,7 @@ def ulysses_ring_custom_bidir_kernel(q, k, v, context): use_experimental_scheduler=context.get("use_experimental_scheduler", False), bidirectional=True, ulysses_attention_chunks=context["ulysses_attention_chunks"], + kv_heads=context.get("kv_heads", None), ) @@ -1774,6 +2092,34 @@ def ulysses_custom_fixed_m_kernel(q, k, v, context): use_base2_exp=context.get("use_base2_exp", True), use_experimental_scheduler=context.get("use_experimental_scheduler", False), use_fixed_m=True, + per_q_block=False, + ulysses_attention_chunks=context["ulysses_attention_chunks"], + kv_heads=context.get("kv_heads", None), + ) + + +@register_kernel("ulysses_custom_fixed_m_per_q_block") +def ulysses_custom_fixed_m_per_q_block_kernel(q, k, v, context): + return _ulysses_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"], + use_custom_kernel=True, + use_base2_exp=context.get("use_base2_exp", True), + use_experimental_scheduler=context.get("use_experimental_scheduler", False), + use_fixed_m=True, + per_q_block=True, + ulysses_attention_chunks=context["ulysses_attention_chunks"], + kv_heads=context.get("kv_heads", None), ) @@ -1794,6 +2140,7 @@ def ulysses_kernel(q, k, v, context): attention_mask=context["attention_mask"], ulysses_attention_chunks=context["ulysses_attention_chunks"], preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False), + kv_heads=context.get("kv_heads", None), ) @@ -1817,6 +2164,7 @@ def ulysses_ring_kernel(q, k, v, context): ulysses_shards=context["ulysses_shards"], ulysses_attention_chunks=context["ulysses_attention_chunks"], preserve_asymmetric_block_sizes=context.get("preserve_asymmetric_block_sizes", False), + kv_heads=context.get("kv_heads", None), ) @@ -1942,6 +2290,7 @@ def _apply_attention( ulysses_attention_chunks: int = 1, is_causal: bool = False, preserve_asymmetric_block_sizes: bool = False, + kv_heads: Optional[int] = None, ): """Routes to different attention kernels using a module-level registry.""" @@ -1957,7 +2306,12 @@ def _apply_attention( "ulysses", "ulysses_custom", "ulysses_custom_fixed_m", + "ulysses_custom_fixed_m_per_q_block", "ulysses_ring", + "ulysses_ring_custom", + "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_fixed_m_per_q_block", + "ulysses_ring_custom_bidir", ]: can_use_flash_attention = ( query.shape[seq_len_idx] >= flash_min_seq_length @@ -1989,6 +2343,7 @@ def _apply_attention( context = { "heads": heads, + "kv_heads": kv_heads, "mesh": mesh, "axis_names_q": axis_names_q, "axis_names_kv": axis_names_kv, @@ -2245,6 +2600,7 @@ def __init__( use_experimental_scheduler: bool = False, ulysses_shards: int = -1, ulysses_attention_chunks: int = 1, + kv_heads: Optional[int] = None, ): self.dpa_layer = None self.use_base2_exp = use_base2_exp @@ -2275,6 +2631,7 @@ def __init__( self.mesh = mesh self.scale = scale self.heads = heads + self.kv_heads = kv_heads self.dim_head = dim_head self.attention_kernel = attention_kernel self.use_memory_efficient_attention = use_memory_efficient_attention @@ -2323,6 +2680,7 @@ def apply_attention( ulysses_shards=(self.ulysses_shards if hasattr(self, "ulysses_shards") else -1), ulysses_attention_chunks=(self.ulysses_attention_chunks if hasattr(self, "ulysses_attention_chunks") else 1), preserve_asymmetric_block_sizes=preserve_asymmetric_block_sizes, + kv_heads=self.kv_heads, ) @@ -2346,6 +2704,7 @@ class AttentionOp(nn.Module): ulysses_shards: int = -1 ulysses_attention_chunks: int = 1 is_causal: bool = False + kv_heads: Optional[int] = None def setup(self): self.dpa_layer = None @@ -2404,6 +2763,7 @@ def apply_attention( ulysses_attention_chunks=self.ulysses_attention_chunks, is_causal=self.is_causal, preserve_asymmetric_block_sizes=preserve_asymmetric_block_sizes, + kv_heads=self.kv_heads, ) @@ -2462,6 +2822,8 @@ def __init__( self.value_axis_names = value_axis_names self.out_axis_names = out_axis_names self.enable_jax_named_scopes = enable_jax_named_scopes + self.is_self_attention = is_self_attention + self.eps = eps cross_attention_remapped_to_flash = not is_self_attention and attention_kernel in ( "tokamax_ring", @@ -2469,9 +2831,11 @@ def __init__( "ulysses_ring", "ulysses_ring_custom", "ulysses_ring_custom_fixed_m", + "ulysses_ring_custom_fixed_m_per_q_block", "ulysses_ring_custom_bidir", "ulysses_custom", "ulysses_custom_fixed_m", + "ulysses_custom_fixed_m_per_q_block", ) cross_attention_uses_local_kv = not is_self_attention and ( cross_attention_remapped_to_flash or attention_kernel in ("flash", "tokamax_flash", "cudnn_flash_te") @@ -2500,6 +2864,7 @@ def __init__( self.image_seq_len = image_seq_len # New for I2V tpu_type = get_tpu_type() self.alignment = 256 if tpu_type in [TpuType.TPU_V6_LITE, TpuType.TPU_7X] else 128 + self.precision = precision self.attention_op = NNXAttentionOp( mesh=mesh, @@ -2674,9 +3039,9 @@ def _apply_rope(self, xq: jax.Array, xk: jax.Array, freqs_cis: jax.Array) -> Tup xk_out_0 = xk_0 * cos - xk_1 * sin xk_out_1 = xk_0 * sin + xk_1 * cos - # 5. Stack and reshape back to original - xq_out = jnp.stack([xq_out_0, xq_out_1], axis=-1).reshape(xq.shape) - xk_out = jnp.stack([xk_out_0, xk_out_1], axis=-1).reshape(xk.shape) + # 5. Concatenate along last axis instead of stack + reshape to prevent layout fragmentation + xq_out = jnp.concatenate([xq_out_0[..., None], xq_out_1[..., None]], axis=-1).reshape(xq.shape) + xk_out = jnp.concatenate([xk_out_0[..., None], xk_out_1[..., None]], axis=-1).reshape(xk.shape) return xq_out, xk_out @@ -2698,7 +3063,10 @@ def __call__( if encoder_hidden_states is not None: encoder_hidden_states = nn.with_logical_constraint(encoder_hidden_states, (BATCH, LENGTH, HEAD)) dtype = hidden_states.dtype - is_self_attention = encoder_hidden_states is None + if encoder_hidden_states is not None: + is_self_attention = False + else: + is_self_attention = getattr(self, "is_self_attention", True) if encoder_hidden_states is None: encoder_hidden_states = hidden_states @@ -2713,11 +3081,12 @@ def __call__( with jax.named_scope("query_proj"): query_proj = self.query(hidden_states) - if self.qk_norm: - with self.conditional_named_scope("attn_q_norm"): - query_proj = self.norm_q(query_proj) - - if not is_self_attention and cached_kv is not None and "text" in cached_kv: + if is_self_attention: + with jax.named_scope("key_proj"): + key_proj = self.key(hidden_states) + with jax.named_scope("value_proj"): + value_proj = self.value(hidden_states) + elif cached_kv is not None and "text" in cached_kv: key_proj, value_proj = cached_kv["text"] else: with jax.named_scope("key_proj"): @@ -2725,17 +3094,38 @@ def __call__( with jax.named_scope("value_proj"): value_proj = self.value(encoder_hidden_states) - if self.qk_norm: - with self.conditional_named_scope("attn_k_norm"): - key_proj = self.norm_k(key_proj) - - if rotary_emb is not None: - with self.conditional_named_scope("attn_rope"): - query_proj = _unflatten_heads(query_proj, self.heads) - key_proj = _unflatten_heads(key_proj, self.heads) + if rotary_emb is not None and self.qk_norm and is_self_attention: + from maxdiffusion.kernels.fused_producers import fused_rmsnorm_rope + + with self.conditional_named_scope("fused_rmsnorm_rope"): + q_scale = self.norm_q.scale[...] + k_scale = self.norm_k.scale[...] + query_proj, key_proj = fused_rmsnorm_rope( + query_proj, + key_proj, + q_scale, + k_scale, + rotary_emb, + q_heads=self.heads, + dim_head=self.dim_head, + eps=self.eps, + ) value_proj = _unflatten_heads(value_proj, self.heads) - # output of _unflatten_heads Batch, heads, seq_len, head_dim - query_proj, key_proj = self._apply_rope(query_proj, key_proj, rotary_emb) + else: + if self.qk_norm: + with self.conditional_named_scope("attn_q_norm"): + query_proj = self.norm_q(query_proj) + if not (not is_self_attention and cached_kv is not None and "text" in cached_kv): + with self.conditional_named_scope("attn_k_norm"): + key_proj = self.norm_k(key_proj) + + if rotary_emb is not None: + with self.conditional_named_scope("attn_rope"): + query_proj = _unflatten_heads(query_proj, self.heads) + key_proj = _unflatten_heads(key_proj, self.heads) + value_proj = _unflatten_heads(value_proj, self.heads) + # output of _unflatten_heads Batch, heads, seq_len, head_dim + query_proj, key_proj = self._apply_rope(query_proj, key_proj, rotary_emb) query_proj = checkpoint_name(query_proj, "query_proj") key_proj = checkpoint_name(key_proj, "key_proj") diff --git a/src/maxdiffusion/models/wan/transformers/transformer_wan.py b/src/maxdiffusion/models/wan/transformers/transformer_wan.py index 4cdfd0ca1..2b829369b 100644 --- a/src/maxdiffusion/models/wan/transformers/transformer_wan.py +++ b/src/maxdiffusion/models/wan/transformers/transformer_wan.py @@ -233,6 +233,7 @@ def __init__( weights_dtype: jnp.dtype = jnp.float32, precision: jax.lax.Precision = None, ): + self.precision = precision self.proj = nnx.Linear( rngs=rngs, in_features=dim_in, @@ -253,8 +254,7 @@ def __init__( def __call__(self, x: jax.Array) -> jax.Array: with jax.named_scope("gelu"): - x = self.proj(x) - return nnx.gelu(x) + return nnx.gelu(self.proj(x)) class WanFeedForward(nnx.Module): @@ -280,6 +280,7 @@ def __init__( dim_out = dim_out if dim_out is not None else dim self.enable_jax_named_scopes = enable_jax_named_scopes + self.precision = precision self.act_fn = nnx.data(None) if activation_fn == "gelu-approximate": self.act_fn = ApproximateGELU( @@ -322,12 +323,13 @@ def __call__( deterministic: bool = True, rngs: nnx.Rngs = None, ) -> jax.Array: - hidden_states = self.act_fn(hidden_states) # Output is (4, 75600, 13824) + hidden_states = self.act_fn(hidden_states) hidden_states = checkpoint_name(hidden_states, "ffn_activation") if self.drop_out.rate > 0: hidden_states = self.drop_out(hidden_states, deterministic=deterministic, rngs=rngs) with jax.named_scope("proj_out"): - return self.proj_out(hidden_states) # output is (4, 75600, 5120) + hidden_states = self.proj_out(hidden_states) + return hidden_states class WanTransformerBlock(nnx.Module): @@ -484,9 +486,9 @@ def __call__( # 1. Self-attention with self.conditional_named_scope("self_attn"): with self.conditional_named_scope("self_attn_norm"): - norm_hidden_states = (self.norm1(hidden_states.astype(jnp.float32)) * (1 + scale_msa) + shift_msa).astype( - hidden_states.dtype - ) + from maxdiffusion.kernels.fused_producers import fused_ln_adaln + + norm_hidden_states = fused_ln_adaln(hidden_states, scale_msa, shift_msa, eps=self.norm1.layer_norm.epsilon) with self.conditional_named_scope("self_attn_attn"): attn_output = self.attn1( hidden_states=norm_hidden_states, @@ -496,12 +498,14 @@ def __call__( rngs=rngs, ) with self.conditional_named_scope("self_attn_residual"): - hidden_states = (hidden_states.astype(jnp.float32) + attn_output * gate_msa).astype(hidden_states.dtype) + hidden_states = (hidden_states.astype(jnp.float32) + attn_output.astype(jnp.float32) * gate_msa).astype( + hidden_states.dtype + ) # 2. Cross-attention with self.conditional_named_scope("cross_attn"): with self.conditional_named_scope("cross_attn_norm"): - norm_hidden_states = self.norm2(hidden_states.astype(jnp.float32)).astype(hidden_states.dtype) + norm_hidden_states = self.norm2(hidden_states) with self.conditional_named_scope("cross_attn_attn"): attn_output = self.attn2( hidden_states=norm_hidden_states, @@ -512,14 +516,14 @@ def __call__( cached_kv=cached_kv, ) with self.conditional_named_scope("cross_attn_residual"): - hidden_states = hidden_states + attn_output + hidden_states = (hidden_states.astype(jnp.float32) + attn_output.astype(jnp.float32)).astype(hidden_states.dtype) # 3. Feed-forward with self.conditional_named_scope("mlp"): with self.conditional_named_scope("mlp_norm"): - norm_hidden_states = (self.norm3(hidden_states.astype(jnp.float32)) * (1 + c_scale_msa) + c_shift_msa).astype( - hidden_states.dtype - ) + from maxdiffusion.kernels.fused_producers import fused_ln_adaln + + norm_hidden_states = fused_ln_adaln(hidden_states, c_scale_msa, c_shift_msa, eps=self.norm3.layer_norm.epsilon) with self.conditional_named_scope("mlp_ffn"): ff_output = self.ffn(norm_hidden_states, deterministic=deterministic, rngs=rngs) with self.conditional_named_scope("mlp_residual"): @@ -934,7 +938,9 @@ def layer_forward(hidden_states, l_kv): scale = scale.squeeze(2) # [B, sl, dim] else: shift, scale = jnp.split(self.scale_shift_table + jnp.expand_dims(temb, axis=1), 2, axis=1) - hidden_states = (self.norm_out(hidden_states.astype(jnp.float32)) * (1 + scale) + shift).astype(hidden_states.dtype) + from maxdiffusion.kernels.fused_producers import fused_ln_adaln + + hidden_states = fused_ln_adaln(hidden_states, scale, shift, eps=self.norm_out.layer_norm.epsilon) with jax.named_scope("proj_out"): hidden_states = self.proj_out(hidden_states) diff --git a/src/maxdiffusion/pyconfig.py b/src/maxdiffusion/pyconfig.py index ddec32c78..78dbba86a 100644 --- a/src/maxdiffusion/pyconfig.py +++ b/src/maxdiffusion/pyconfig.py @@ -234,6 +234,7 @@ def user_init(raw_keys): "ulysses_ring_custom", "ulysses_ring_custom_fixed_m", "ulysses_ring_custom_bidir", + "ulysses_ring_custom_fixed_m_per_q_block", } 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/custom_splash_fixed_m_test.py b/src/maxdiffusion/tests/custom_splash_fixed_m_test.py index 0f80a8306..80cd58c2f 100644 --- a/src/maxdiffusion/tests/custom_splash_fixed_m_test.py +++ b/src/maxdiffusion/tests/custom_splash_fixed_m_test.py @@ -87,8 +87,11 @@ def _run_kernel(self, q: jax.Array, k: jax.Array, v: jax.Array, use_fixed_m: boo k_in = k_in - jnp.mean(k_in, axis=1, keepdims=True) qn = jnp.sqrt((q_in.astype(jnp.float32) ** 2).sum(-1)).max(axis=1) mk_h = jnp.sqrt((k_in.astype(jnp.float32) ** 2).sum(-1)).max(axis=1) - eligible = (qn * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(jnp.float32) - mk = jnp.stack([mk_h, eligible]) + recenter, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len, is_ring=False) + bound = qn * mk_h + eligible = (bound <= safe_bound).astype(jnp.float32) + m_base = jnp.ceil(bound) - recenter + mk = jnp.stack([m_base, eligible]) kernel = custom_splash.make_splash_mha( block_sizes=self.block_sizes, orig_q_seq_len=self.seq_len, @@ -121,6 +124,39 @@ def test_fixed_m_matches_reference(self): fixed, _ = self._run_kernel(q, k, v, use_fixed_m=True) self.assertLess(float(jnp.max(jnp.abs(fixed - self._reference(q, k, v)))), 2e-2) + def _run_kernel_per_q_block( + self, q: jax.Array, k: jax.Array, v: jax.Array, uniform_fixed_m: bool = False + ) -> tuple[jax.Array, jax.Array]: + """Runs the custom kernel with 3D per-Q-block mk inputs.""" + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in = k * self.scale + k_in = k_in - jnp.mean(k_in, axis=1, keepdims=True) + + bq = self.block_sizes.block_q + num_q_blocks = self.seq_len // bq + qf = q_in.astype(jnp.float32) + kf = k_in.astype(jnp.float32) + qf_blocks = qf.reshape(self.num_heads, num_q_blocks, bq, self.head_dim) + qn_max = jnp.sqrt((qf_blocks * qf_blocks).sum(-1)).max(axis=-1) # (heads, num_q_blocks) + mk_h = jnp.sqrt((kf * kf).sum(-1)).max(axis=1) # (heads,) + recenter, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len, is_ring=False) + bound = qn_max * mk_h[:, None] + fixed_ok = (bound <= safe_bound).astype(jnp.float32) + m_base = jnp.ceil(bound) - recenter + mk = jnp.stack([m_base, fixed_ok], axis=0) # (2, heads, num_q_blocks) + + kernel = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + uniform_fixed_m=uniform_fixed_m, + ) + out = kernel(q_in, k_in, v, mk) + out = jnp.swapaxes(out, 1, 2) + return out.astype(jnp.float32), mk + def test_sink_head_falls_back_to_online(self): """An out-of-gate head is flagged ineligible and stays finite (no flush).""" q, k, v = self._random_qkv(q_gain=6.0, k_gain=6.0) @@ -129,6 +165,465 @@ def test_sink_head_falls_back_to_online(self): self.assertTrue(bool(jnp.all(mk[1][1:] > 0.5))) # the rest stay eligible self.assertTrue(bool(jnp.all(jnp.isfinite(fixed)))) + def test_per_q_block_sink_fallback(self): + """Per-Q-block eligibility keeps normal Q-blocks fixed while sinking outlier blocks.""" + q, k, v = self._random_qkv(k_gain=2.0) + # Amplify only Q-block 1 of Head 0 (bq = 2048, so indices 2048:4096) + q = q.at[0, 2048:].multiply(10.0) + + fixed, mk = self._run_kernel_per_q_block(q, k, v) + # Head 0, Block 0 should be eligible (1.0) + self.assertEqual(float(mk[1, 0, 0]), 1.0) + # Head 0, Block 1 should be ineligible (0.0) due to amplified Q outlier + self.assertEqual(float(mk[1, 0, 1]), 0.0) + # All other heads should be eligible across both blocks + self.assertTrue(bool(jnp.all(mk[1, 1:, :] > 0.5))) + self.assertTrue(bool(jnp.all(jnp.isfinite(fixed)))) + # Check numerical agreement against online kernel running on the same centered inputs + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in_centered = (k * self.scale) - jnp.mean(k * self.scale, axis=1, keepdims=True) + kernel_online = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=False, + ) + online_centered = jnp.swapaxes(kernel_online(q_in, k_in_centered, v), 1, 2).astype(jnp.float32) + self.assertLess(float(jnp.max(jnp.abs(fixed - online_centered))), 1e-2) + + def test_batched_fixed_m_isolation(self): + """Batch-isolated gating ensures outliers in one sample do not contaminate other samples.""" + q0, k0, v0 = self._random_qkv(k_gain=2.0) + # Sample 0 has an outlier in Q-block 1 of Head 0 + q0 = q0.at[0, 2048:].multiply(10.0) + + # Sample 1 is completely clean + q1, k1, v1 = self._random_qkv(k_gain=1.0) + + q = jnp.stack([q0, q1], axis=0) # (2, heads, seq, dim) + k = jnp.stack([k0, k1], axis=0) + v = jnp.stack([v0, v1], axis=0) + + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in = k * self.scale + k_in = k_in - jnp.mean(k_in, axis=2, keepdims=True) + + bq = self.block_sizes.block_q + num_q_blocks = self.seq_len // bq + qf = q_in.astype(jnp.float32) + kf = k_in.astype(jnp.float32) + qf_blocks = qf.reshape(2, self.num_heads, num_q_blocks, bq, self.head_dim) + qn_max_sq = (qf_blocks * qf_blocks).sum(-1).max(axis=-1) # (2, heads, num_q_blocks) + mk_h_sq = (kf * kf).sum(-1).max(axis=-1) # (2, heads) + recenter, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len, is_ring=False) + bound_sq = qn_max_sq * mk_h_sq[:, :, None] + fixed_ok = (bound_sq <= (safe_bound**2)).astype(jnp.float32) + m_base = jnp.ceil(jnp.sqrt(bound_sq)) - recenter + mk_arr = jnp.stack([m_base, fixed_ok], axis=1) # (2, 2, heads, num_q_blocks) + + # Verify Sample 0 has Head 0 Block 1 disqualified (0.0) + self.assertEqual(float(mk_arr[0, 1, 0, 1]), 0.0) + # Verify Sample 1 has ALL heads and ALL blocks eligible (1.0) - zero contamination! + self.assertTrue(bool(jnp.all(mk_arr[1, 1] > 0.5))) + + kernel = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + uniform_fixed_m=False, + ) + vmapped_kernel = jax.vmap(kernel, in_axes=(0, 0, 0, 0)) + out = vmapped_kernel(q_in, k_in, v, mk_arr) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + + def test_uniform_fixed_matches_hybrid(self): + """Uniform-fixed kernel matches hybrid kernel and f32 reference when all eligible.""" + q, k, v = self._random_qkv() + hybrid_out, mk = self._run_kernel_per_q_block(q, k, v, uniform_fixed_m=False) + uniform_out, _ = self._run_kernel_per_q_block(q, k, v, uniform_fixed_m=True) + ref = self._reference(q, k, v) + + self.assertTrue(bool(jnp.all(mk[1] > 0.5))) + self.assertLess(float(jnp.max(jnp.abs(uniform_out - hybrid_out))), 5e-3) + self.assertLess(float(jnp.max(jnp.abs(uniform_out - ref))), 2e-2) + + def test_missing_mk_raises_value_error(self): + """When use_fixed_m=True, omitting mk raises an immediate ValueError.""" + q, k, v = self._random_qkv() + kernel = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + ) + with self.assertRaises(ValueError): + kernel(q, k, v, mk=None) + + def test_phase_transition_boundary_continuity(self): + """Verifies seamless output continuity between fixed-m and online mode across the dynamic safe bound threshold.""" + q_base, k_base, v = self._random_qkv() + q_normed = q_base / jnp.sqrt((q_base.astype(jnp.float32) ** 2).sum(-1, keepdims=True)) + k_normed = k_base / jnp.sqrt((k_base.astype(jnp.float32) ** 2).sum(-1, keepdims=True)) + + _, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len, is_ring=False) + test_bounds = [ + safe_bound - 2.0, + safe_bound - 0.5, + safe_bound - 0.01, + safe_bound, + safe_bound + 0.01, + safe_bound + 0.5, + safe_bound + 2.0, + ] + for target_bound in test_bounds: + factor = math.sqrt(target_bound / _LOG2E / self.scale) + q = (q_normed * factor).astype(jnp.bfloat16) + k = (k_normed * factor).astype(jnp.bfloat16) + + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in = k * self.scale + k_in = k_in - jnp.mean(k_in, axis=1, keepdims=True) + + out_gated, _ = self._run_kernel_per_q_block(q, k, v) + kernel_online = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=False, + ) + out_online = jnp.swapaxes(kernel_online(q_in, k_in, v), 1, 2).astype(jnp.float32) + + self.assertTrue(bool(jnp.all(jnp.isfinite(out_gated)))) + diff = float(jnp.max(jnp.abs(out_gated - out_online))) + self.assertLess(diff, 2e-2, f"Discontinuity at bound={target_bound}, diff={diff}") + + def test_cpu_proof_invariant_bounds(self): + """Verifies that mathematical underflow and overflow invariants hold across sequence lengths.""" + # Test sequence lengths across short, medium, and production Wan2.2 dimensions + test_lengths = [1, 2, 512, 1024, 16384, 75600, 151200] + + for n in test_lengths: + # 1. Pure Ulysses (Centered): M >= 0 + recenter, safe_bound = custom_splash.get_fixed_m_constants(n, is_ring=False) + # Minimal shifted exponent at the boundary U = safe_bound + exponent_centered = 0.0 - (math.ceil(safe_bound) - recenter) + self.assertGreaterEqual( + exponent_centered, + -125.0, + f"Underflow violation on Ulysses: {exponent_centered=} for N={n}", + ) + # Non-overflow check with explicit 8-bit FP32 output headroom: log2(N) + C(N) + FP32_OUTPUT_HEADROOM_BITS <= 127 + max_accum_bits = math.ceil(math.log2(n)) + self.assertLessEqual( + max_accum_bits + recenter + custom_splash.FP32_OUTPUT_HEADROOM_BITS, + 127.0, + f"Overflow violation on Ulysses: max bits={max_accum_bits + recenter + custom_splash.FP32_OUTPUT_HEADROOM_BITS} for N={n}", + ) + + # 2. Ring Attention (Uncentered across R hops): M >= -U + for ring_size in [2, 4, 8]: + n_total = n * ring_size + ring_recenter, ring_safe_bound = custom_splash.get_fixed_m_constants(n_total, is_ring=True) + # Minimal shifted exponent at worst-case extremum M = -U + exponent_ring = -ring_safe_bound - (math.ceil(ring_safe_bound) - ring_recenter) + self.assertGreaterEqual( + exponent_ring, + -125.0, + f"Underflow violation on Ring (R={ring_size}): {exponent_ring=} for N={n}", + ) + # Ring direct accumulation non-overflow check: log2(N_total) + C_ring + FP32_OUTPUT_HEADROOM_BITS <= 127 + ring_max_bits = math.ceil(math.log2(n_total)) + self.assertLessEqual( + ring_max_bits + ring_recenter + custom_splash.FP32_OUTPUT_HEADROOM_BITS, + 127.0, + f"Overflow violation on Ring (R={ring_size}): max bits={ring_max_bits + ring_recenter + custom_splash.FP32_OUTPUT_HEADROOM_BITS} for N={n}", + ) + + def test_non_divisible_sequence_context_padding_fixed_m(self): + """Verifies that non-divisible sequences (e.g. S=1001 padded for 8 shards) are correctly masked without zero-padding pollution.""" + seq_len = 1001 + context_shards = 8 + rem = seq_len % context_shards + padded_seq_len = seq_len + (context_shards - rem) # 1008 + heads = 4 + dim = 64 + bq = 512 + + q_raw = jax.random.normal(jax.random.PRNGKey(101), (heads, seq_len, dim), jnp.bfloat16) + k_raw = jax.random.normal(jax.random.PRNGKey(102), (heads, seq_len, dim), jnp.bfloat16) + v_raw = jax.random.normal(jax.random.PRNGKey(103), (heads, seq_len, dim), jnp.bfloat16) + + # Reference dense attention on true unpadded inputs + ref_out = self._reference(q_raw, k_raw, v_raw) + + # Pad inputs as _reshape_data_for_flash would for context sharding + q_pad = jnp.pad(q_raw, ((0, 0), (0, padded_seq_len - seq_len), (0, 0))) + k_pad = jnp.pad(k_raw, ((0, 0), (0, padded_seq_len - seq_len), (0, 0))) + v_pad = jnp.pad(v_raw, ((0, 0), (0, padded_seq_len - seq_len), (0, 0))) + + # Compute unpadded K centering and metadata + k_mean = jnp.mean(k_raw.astype(jnp.float32) * self.scale, axis=1) # (heads, dim) + recenter, safe_bound = custom_splash.get_fixed_m_constants(seq_len, is_ring=False) + + q_in = (q_pad * _LOG2E).astype(jnp.bfloat16) + k_in = (k_pad * self.scale).astype(jnp.bfloat16) + + num_q_blocks = math.ceil(padded_seq_len / bq) + # Pad to systolic block_q boundary + pad_bq = num_q_blocks * bq + q_in_padded = jnp.pad(q_in, ((0, 0), (0, pad_bq - padded_seq_len), (0, 0))) + k_in_padded = jnp.pad(k_in, ((0, 0), (0, pad_bq - padded_seq_len), (0, 0))) + v_in_padded = jnp.pad(v_pad, ((0, 0), (0, pad_bq - padded_seq_len), (0, 0))) + + # Metadata computed on real keys + k_centered = (k_raw.astype(jnp.float32) * self.scale) - k_mean[:, None, :] + mk_h = jnp.sqrt((k_centered**2).sum(-1)).max(axis=-1) + qf_blocks = q_in_padded.astype(jnp.float32).reshape(heads, num_q_blocks, bq, dim) + qn_max = jnp.sqrt((qf_blocks * qf_blocks).sum(-1)).max(axis=-1) + bound = qn_max * mk_h[:, None] + fixed_ok = (bound <= safe_bound).astype(jnp.float32) + m_base = jnp.ceil(bound) - recenter + mk = jnp.stack([m_base, fixed_ok], axis=0) + + block_sizes = custom_splash._BlockSizes(block_q=bq, block_kv=bq, block_kv_compute=bq, block_kv_compute_in=bq) + kernel = custom_splash.make_splash_mha( + block_sizes=block_sizes, + orig_q_seq_len=padded_seq_len, + orig_kv_seq_len=seq_len, + use_base2_exp=True, + use_fixed_m=True, + uniform_fixed_m=True, + ) + out = jnp.swapaxes(kernel(q_in_padded, k_in_padded, v_in_padded, mk, k_mean), 1, 2).astype(jnp.float32) + out_sliced = out[:, :seq_len, :] + + diff = float(jnp.max(jnp.abs(out_sliced - ref_out))) + self.assertTrue(bool(jnp.all(jnp.isfinite(out_sliced)))) + self.assertLess(diff, 2e-2, f"Non-divisible sequence output diverged from reference: {diff=}") + + def test_pathological_keys_extreme_negative_logits(self): + """Verifies stability when logits are heavily negative and close to underflow.""" + q, k, v = self._random_qkv() + # Shift keys far into negative space so dot products are mostly negative + k_pathological = k - 30.0 + out_fixed, _ = self._run_kernel_per_q_block(q, k_pathological, v) + out_online, _ = self._run_kernel(q, k_pathological, v, use_fixed_m=False) + + self.assertTrue(bool(jnp.all(jnp.isfinite(out_fixed)))) + diff = float(jnp.max(jnp.abs(out_fixed - out_online))) + self.assertLess(diff, 2e-2) + + def test_extreme_dynamic_range_inputs(self): + """Verifies that norm computation and gating remain robust with wide dynamic ranges.""" + shape = (self.num_heads, self.seq_len, self.head_dim) + scales = jnp.array([1e-3, 0.1, 0.5, 1.0, 1.5])[:, None, None] + q = (jax.random.normal(jax.random.PRNGKey(42), shape, jnp.bfloat16) * scales).astype(jnp.bfloat16) + k = (jax.random.normal(jax.random.PRNGKey(43), shape, jnp.bfloat16) * scales).astype(jnp.bfloat16) + v = jax.random.normal(jax.random.PRNGKey(44), shape, jnp.bfloat16) + + out, mk = self._run_kernel_per_q_block(q, k, v) + ref = self._reference(q, k, v) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + self.assertLess(float(jnp.max(jnp.abs(out - ref))), 3e-2) + + def test_virtual_k_centering_matches_explicit(self): + """Virtual K-centering with raw keys matches explicit K-centering numerically.""" + q, k, v = self._random_qkv() + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in_raw = (k * self.scale).astype(jnp.bfloat16) + k_mean = jnp.mean(k_in_raw.astype(jnp.float32), axis=1) + + k_in_centered = k_in_raw.astype(jnp.float32) - k_mean[:, None, :] + mk_h_sq = (k_in_centered**2).sum(axis=-1).max(axis=1) + mk_h = jnp.sqrt(mk_h_sq) + + bq = self.block_sizes.block_q + num_q_blocks = self.seq_len // bq + qf = q_in.astype(jnp.float32) + qf_blocks = qf.reshape(self.num_heads, num_q_blocks, bq, self.head_dim) + qn_max = jnp.sqrt((qf_blocks * qf_blocks).sum(-1)).max(axis=-1) + recenter, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len, is_ring=False) + bound = qn_max * mk_h[:, None] + fixed_ok = (bound <= safe_bound).astype(jnp.float32) + m_base = jnp.ceil(bound) - recenter + mk = jnp.stack([m_base, fixed_ok], axis=0) + + # Virtual K-centering with raw uncentered keys + kernel_virtual = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + uniform_fixed_m=True, + ) + out_virtual = jnp.swapaxes(kernel_virtual(q_in, k_in_raw, v, mk, k_mean), 1, 2).astype(jnp.float32) + + # Explicit centering with centered keys + k_centered_bf16 = k_in_centered.astype(jnp.bfloat16) + kernel_explicit = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + uniform_fixed_m=True, + ) + out_explicit = jnp.swapaxes(kernel_explicit(q_in, k_centered_bf16, v, mk), 1, 2).astype(jnp.float32) + + ref = self._reference(q, k, v) + diff_virtual_explicit = float(jnp.max(jnp.abs(out_virtual - out_explicit))) + diff_virtual_ref = float(jnp.max(jnp.abs(out_virtual - ref))) + + self.assertLess(diff_virtual_explicit, 2e-3) + self.assertLess(diff_virtual_ref, 2e-2) + self.assertTrue(bool(jnp.all(jnp.isfinite(out_virtual)))) + + def test_virtual_k_centering_per_q_block_hybrid_fallback(self): + """Exercises Virtual K-Centering + Per-Q-Block Hybrid dispatch with mixed fixed/online tiles.""" + q, k, v = self._random_qkv() + bq = self.block_sizes.block_q + num_q_blocks = self.seq_len // bq + q_in = (q * _LOG2E).astype(jnp.bfloat16) + k_in_raw = (k * self.scale).astype(jnp.bfloat16) + k_mean = jnp.mean(k_in_raw.astype(jnp.float32), axis=1) + + k_in_centered = k_in_raw.astype(jnp.float32) - k_mean[:, None, :] + mk_h_sq = (k_in_centered**2).sum(axis=-1).max(axis=1) + mk_h = jnp.sqrt(mk_h_sq) + + # Test hybrid dispatch where Head 0 Block 0 is Fixed-M and Block 1 is Online Fallback + recenter, safe_bound = custom_splash.get_fixed_m_constants(self.seq_len, is_ring=False) + qf_blocks = q_in.astype(jnp.float32).reshape(self.num_heads, num_q_blocks, bq, self.head_dim) + qn_max = jnp.sqrt((qf_blocks * qf_blocks).sum(-1)).max(axis=-1) + bound = qn_max * mk_h[:, None] + m_base = jnp.ceil(bound) - recenter + fixed_ok = jnp.ones((self.num_heads, num_q_blocks), dtype=jnp.float32).at[0, 1].set(0.0) + mk = jnp.stack([m_base, fixed_ok], axis=0) + + # Verify Block 0 is fixed (1.0), Block 1 is online fallback (0.0) on Head 0 + self.assertEqual(float(mk[1, 0, 0]), 1.0) + self.assertEqual(float(mk[1, 0, 1]), 0.0) + + # Hybrid kernel with raw uncentered keys + k_mean + kernel_hybrid = custom_splash.make_splash_mha( + block_sizes=self.block_sizes, + orig_q_seq_len=self.seq_len, + orig_kv_seq_len=self.seq_len, + use_base2_exp=True, + use_fixed_m=True, + uniform_fixed_m=False, + ) + out_hybrid = jnp.swapaxes(kernel_hybrid(q_in, k_in_raw, v, mk, k_mean), 1, 2).astype(jnp.float32) + + # Dense f32 reference + ref = self._reference(q, k, v) + diff = float(jnp.max(jnp.abs(out_hybrid - ref))) + + self.assertTrue(bool(jnp.all(jnp.isfinite(out_hybrid)))) + self.assertLess(diff, 2e-2, f"Hybrid virtual K output diverged from reference: diff={diff}") + + def test_gqa_fixed_m_metadata_broadcast(self): + """Verifies that _compute_fixed_m_metadata correctly handles GQA (num_q_heads != num_kv_heads).""" + from maxdiffusion.models.attention_flax import _compute_fixed_m_metadata + + batch = 2 + num_q_heads = 8 + num_kv_heads = 2 + seq_len = 2048 + dim = 64 + bq = 512 + q = jax.random.normal(jax.random.PRNGKey(10), (batch, num_q_heads, seq_len, dim), jnp.bfloat16) + k = jax.random.normal(jax.random.PRNGKey(11), (batch, num_kv_heads, seq_len, dim), jnp.bfloat16) + mk_arr, all_fixed = _compute_fixed_m_metadata(q, k, block_q=bq) + expected_blocks = seq_len // bq + self.assertEqual(mk_arr.shape, (batch, 2, num_q_heads, expected_blocks)) + self.assertTrue(bool(jnp.all(jnp.isfinite(mk_arr)))) + + def test_adversarial_v_magnitude_safely_disqualifies_fixed_m(self): + """Verifies that |V| > v_max_bound (e.g. V=512) disqualifies fixed-m gating to prevent FP32 overflow.""" + from maxdiffusion.models.attention_flax import _compute_fixed_m_metadata + + batch = 1 + num_heads = 4 + seq_len = 4096 + dim = 64 + bq = 512 + + q = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.bfloat16) + k = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.bfloat16) + v_overflow = jnp.full((batch, num_heads, seq_len, dim), 512.0, dtype=jnp.bfloat16) + + # With adversarial V=512 (> 256 default bound), fixed_ok must be 0.0, safely falling back to online + mk_arr, all_fixed = _compute_fixed_m_metadata(q, k, block_q=bq, value=v_overflow) + self.assertFalse(bool(all_fixed)) + self.assertTrue(bool(jnp.all(mk_arr[:, 1] == 0.0))) + + # With normal V <= 256, fixed_ok should remain 1.0 (all eligible) + v_normal = jnp.full((batch, num_heads, seq_len, dim), 1.0, dtype=jnp.bfloat16) + mk_arr_normal, all_fixed_normal = _compute_fixed_m_metadata(q, k, block_q=bq, value=v_normal) + self.assertTrue(bool(all_fixed_normal)) + self.assertTrue(bool(jnp.all(mk_arr_normal[:, 1] == 1.0))) + + +class FixedMDtypeSafetyTest(unittest.TestCase): + """P2 regression: dtypes that cannot represent 2**C(N) must not use fixed-m. + + Fixed-m parks the un-normalized softmax weights at up to 2**C(N), a range + derived against FP32's exponent. The kernel narrows them to the activation + dtype for the S@V matmul, so a dtype with a smaller exponent range overflows + to inf even when the FP32 bound analysis passes. + + Backend-agnostic on purpose: this gate is pure Python/jnp, so it should be + enforced in CI even where no TPU is attached. + """ + + def test_float16_is_rejected(self): + recenter, _ = custom_splash.get_fixed_m_constants(4096, is_ring=False) + # C(4096) with |V| <= 256 is 107; float16 tops out at 2**16. + self.assertGreater(recenter, 16.0) + self.assertFalse(custom_splash.fixed_m_dtype_is_safe(jnp.float16, recenter)) + + def test_bfloat16_and_float32_are_accepted(self): + recenter, _ = custom_splash.get_fixed_m_constants(4096, is_ring=False) + self.assertTrue(custom_splash.fixed_m_dtype_is_safe(jnp.bfloat16, recenter)) + self.assertTrue(custom_splash.fixed_m_dtype_is_safe(jnp.float32, recenter)) + + def test_gate_tracks_recenter_not_a_hardcoded_allowlist(self): + """A small enough C(N) is representable even in float16.""" + self.assertTrue(custom_splash.fixed_m_dtype_is_safe(jnp.float16, 4.0)) + self.assertFalse(custom_splash.fixed_m_dtype_is_safe(jnp.float16, 200.0)) + + def test_float16_query_disqualifies_fixed_m_metadata(self): + """The reviewer's case: fp16, N=4096, Q=K=0, |V|=1 previously reported all_fixed=True.""" + from maxdiffusion.models.attention_flax import _compute_fixed_m_metadata + + batch, num_heads, seq_len, dim, bq = 1, 2, 4096, 128, 512 + q = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.float16) + k = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.float16) + v = jnp.full((batch, num_heads, seq_len, dim), 1.0, dtype=jnp.float16) + + mk_arr, all_fixed = _compute_fixed_m_metadata(q, k, block_q=bq, value=v) + self.assertFalse(bool(all_fixed), "fp16 must not be eligible for fixed-m") + self.assertTrue(bool(jnp.all(mk_arr[:, 1] == 0.0))) + + def test_bfloat16_same_case_remains_eligible(self): + """Control: the identical case in bf16 must still take the fast path.""" + from maxdiffusion.models.attention_flax import _compute_fixed_m_metadata + + batch, num_heads, seq_len, dim, bq = 1, 2, 4096, 128, 512 + q = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.bfloat16) + k = jnp.zeros((batch, num_heads, seq_len, dim), dtype=jnp.bfloat16) + v = jnp.full((batch, num_heads, seq_len, dim), 1.0, dtype=jnp.bfloat16) + + _, all_fixed = _compute_fixed_m_metadata(q, k, block_q=bq, value=v) + self.assertTrue(bool(all_fixed)) + if __name__ == "__main__": unittest.main() diff --git a/src/maxdiffusion/tests/ring_fixed_m_test.py b/src/maxdiffusion/tests/ring_fixed_m_test.py index 0f4cc6876..8f6555428 100644 --- a/src/maxdiffusion/tests/ring_fixed_m_test.py +++ b/src/maxdiffusion/tests/ring_fixed_m_test.py @@ -14,18 +14,13 @@ 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. +"""Unit tests for the fixed-m path of the custom RING attention with Global Virtual K-Centering. + +The ring path computes a global Key mean across the ring axis (k_mean = pmean(mean(k), ring_axis)), +which mathematically guarantees max_j (q^T (k_j - k_mean)) >= 0 across the entire distributed sequence. +Key norms are gathered across ranks once before the scan (mk_global = mk_all.max(axis=0)) to evaluate +identical fixed m bounds across all ring hops, enabling direct FP32 accumulation and full centered +safe bounds (W(N) = 127 - ceil(log2 N) + 125). """ import functools @@ -38,6 +33,8 @@ from maxdiffusion.kernels import custom_splash_attention as custom_splash from maxdiffusion.kernels.splash_attention import ring_attention_kernel +from maxdiffusion.models import attention_flax +from flax.linen import partitioning as nn_partitioning _LOG2E = math.log2(math.e) _RING_AXIS = "ring" @@ -45,7 +42,7 @@ class RingFixedMTest(unittest.TestCase): - """Numerical tests for the fixed-m custom ring attention.""" + """Numerical tests for the fixed-m custom ring attention across topologies.""" num_heads = 4 shard_len = 2048 # per-device sequence; total = shard_len * ring_size @@ -55,16 +52,18 @@ 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): + def _mesh_for_size(self, ring_size: int): + if len(jax.devices()) < ring_size: + self.skipTest(f"Requires {ring_size} devices, but only {len(jax.devices())} available.") + devices = np.asarray(jax.devices()[:ring_size]) + return jax.sharding.Mesh(devices, (_RING_AXIS,)) + + def _random_qkv(self, ring_size: int = 2, 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 + 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) @@ -78,12 +77,7 @@ def _random_qkv(self, q_gain=None, k_gain=None): 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.""" + """The EXACT bf16 tensors the kernel sees (attention_flax's contract).""" q_in = (q * _LOG2E).astype(q.dtype) k_in = (k.astype(jnp.float32) * self.scale).astype(k.dtype) return q_in, k_in @@ -94,99 +88,691 @@ def _reference(self, 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.""" + def _run_ring(self, q_in, k_in, v, ring_size: int = 2, use_fixed_m: bool = True, v_ok=None): + """Runs the custom ring under shard_map with per-rank fixed_m_norms.""" + mesh = self._mesh_for_size(ring_size) spec = jax.sharding.PartitionSpec(None, _RING_AXIS, None) @functools.partial( jax.shard_map, - mesh=self.mesh, + mesh=mesh, in_specs=(spec, spec, spec), out_specs=spec, check_vma=False, ) def _body(ql, kl, vl): fixed_m_norms = None + k_mean = 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) + k_mean_local = jnp.mean(kf, axis=1) # (heads, dim) + k_mean = jax.lax.pmean(k_mean_local, axis_name=_RING_AXIS) + bq = self.block_sizes.block_q + num_q_blocks = qf.shape[1] // bq + qf_blocks = qf.reshape(qf.shape[0], num_q_blocks, bq, qf.shape[-1]) + qn_blocks_sq = (qf_blocks * qf_blocks).sum(-1).max(axis=-1) # (heads, num_q_blocks) + kf_centered = kf - k_mean[:, None, :] + mk_h_sq = (kf_centered * kf_centered).sum(-1).max(axis=1) # (heads,) local shard + fixed_m_norms = (qn_blocks_sq, mk_h_sq) + v_ok_effective = v_ok + if v_ok is None and use_fixed_m: + v_max_sq = (vl.astype(jnp.float32) ** 2).max() + v_ok_local = v_max_sq <= (custom_splash.DEFAULT_MAX_V_BOUND**2) + v_ok_effective = jax.lax.pmin(v_ok_local, axis_name=_RING_AXIS) 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, + ring_size=ring_size, use_fixed_m=use_fixed_m, fixed_m_norms=fixed_m_norms, + k_mean=k_mean, + v_ok=v_ok_effective, ) 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.""" + def _global_v_ok(self, v, ring_size: int = 2): + """The V-safety verdict attention_flax computes, reduced over the whole ring.""" + v_max_sq = (v.astype(jnp.float32) ** 2).max() + return bool(v_max_sq <= custom_splash.DEFAULT_MAX_V_BOUND**2) + + def _gate_per_shard(self, q_in, k_in, ring_size: int = 2): + """(heads, ring_size) eligibility against the dynamic centered bound with Global Virtual K-Centering.""" qf = q_in.astype(jnp.float32) kf = k_in.astype(jnp.float32) + k_mean_global = jnp.mean(kf, axis=1, keepdims=True) + kf_centered = kf - k_mean_global qn = jnp.sqrt((qf * qf).sum(-1)) # (heads, total) - kn = jnp.sqrt((kf * kf).sum(-1)) + kn = jnp.sqrt((kf_centered * kf_centered).sum(-1)) + _, safe_bound = custom_splash.get_fixed_m_constants(self.shard_len * ring_size, is_ring=False) gates = [] - for r in range(_RING_SIZE): + 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) + gates.append(bound <= safe_bound) return jnp.stack(gates, axis=1) - def _run_and_compare(self, q, k, v, use_fixed_m): + def _run_and_compare(self, q, k, v, ring_size: int = 2, use_fixed_m: bool = True): 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) + out = self._run_ring(q_in, k_in, v, ring_size=ring_size, 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 _gate(self, q, k, ring_size: int = 2): + return self._gate_per_shard(*self._scaled_inputs(q, k), ring_size=ring_size) 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) + q, k, v = self._random_qkv(ring_size=2) + self.assertLess(self._run_and_compare(q, k, v, ring_size=2, 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) + q, k, v = self._random_qkv(ring_size=2) + self.assertTrue(bool(jnp.all(self._gate(q, k, ring_size=2)))) + self.assertLess(self._run_and_compare(q, k, v, ring_size=2, 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) + total = self.shard_len * 2 + q, k, v = self._random_qkv(ring_size=2, q_gain=(0, slice(0, total), 40.0)) + gate = self._gate(q, k, ring_size=2) 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) + self.assertLess(self._run_and_compare(q, k, v, ring_size=2, 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) + q, k, v = self._random_qkv(ring_size=2) + self.assertTrue(bool(jnp.all(self._gate(q, k, ring_size=2)))) + self.assertLess(self._run_and_compare(q, k, v, ring_size=2, 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) + q, k, v = self._random_qkv(ring_size=2, k_gain=(0, slice(self.shard_len, self.shard_len * 2), 40.0)) + gate = self._gate(q, k, ring_size=2) 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) + self.assertLess(self._run_and_compare(q, k, v, ring_size=2, use_fixed_m=True), 2e-2) + + def test_per_q_block_sink_tile_ring(self): + q, k, v = self._random_qkv(ring_size=2, q_gain=(0, slice(0, self.block_sizes.block_q), 40.0)) + self.assertLess(self._run_and_compare(q, k, v, ring_size=2, use_fixed_m=True), 2e-2) + + def test_fixed_m_multi_hop_ring_size_4(self): + """Verifies wrap-around collective correctness and LSE accumulation for R=4.""" + q, k, v = self._random_qkv(ring_size=4) + self.assertTrue(bool(jnp.all(self._gate(q, k, ring_size=4)))) + self.assertLess(self._run_and_compare(q, k, v, ring_size=4, use_fixed_m=True), 2e-2) + + def test_mixed_fixed_online_ring_size_4(self): + """Verifies 4-hop ring with mixed fixed/online shards on separate ranks.""" + # Shard 2 is amplified: head 0 will be fixed on shards 0, 1, 3 and online on shard 2 + q, k, v = self._random_qkv(ring_size=4, k_gain=(0, slice(self.shard_len * 2, self.shard_len * 3), 40.0)) + gate = self._gate(q, k, ring_size=4) + self.assertTrue(bool(gate[0, 0])) + self.assertTrue(bool(gate[0, 1])) + self.assertFalse(bool(gate[0, 2])) + self.assertTrue(bool(gate[0, 3])) + self.assertLess(self._run_and_compare(q, k, v, ring_size=4, use_fixed_m=True), 2e-2) + + def test_fixed_m_multi_hop_ring_size_8(self): + """Verifies 8-device full torus ring rotation and fixed-m numerical equivalence.""" + q, k, v = self._random_qkv(ring_size=8) + self.assertTrue(bool(jnp.all(self._gate(q, k, ring_size=8)))) + self.assertLess(self._run_and_compare(q, k, v, ring_size=8, use_fixed_m=True), 2e-2) + + def test_batched_cfg_isolation_ring(self): + """Verifies that CFG batch items (batch=2) are strictly isolated with zero cross-contamination.""" + # Batch 0: normal bounded activations (all fixed-m) + # Batch 1: massive sink outlier token (forces online fallback) + ring_size = 2 + q0, k0, v0 = self._random_qkv(ring_size=ring_size) + q1, k1, v1 = self._random_qkv(ring_size=ring_size, q_gain=(0, slice(0, self.block_sizes.block_q), 50.0)) + q_batch = jnp.stack([q0, q1], axis=0) # (2, heads, total_seq, dim) + k_batch = jnp.stack([k0, k1], axis=0) + v_batch = jnp.stack([v0, v1], axis=0) + q_in, k_in = self._scaled_inputs(q_batch, k_batch) + + # Run batch through shard_map with vmap over batch + mesh = self._mesh_for_size(ring_size) + spec = jax.sharding.PartitionSpec(None, None, _RING_AXIS, None) + + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=(spec, spec, spec), + out_specs=spec, + check_vma=False, + ) + def _body_batched(qb, kb, vb): + batch_size, num_h, q_seq, _ = qb.shape + bq = self.block_sizes.block_q + num_q_blocks = q_seq // bq + qfb = qb.astype(jnp.float32) + kfb = kb.astype(jnp.float32) + k_mean_local = jnp.mean(kfb, axis=2) # (batch, heads, dim) + k_mean = jax.lax.pmean(k_mean_local, axis_name=_RING_AXIS) + norm_sq = (qfb * qfb).sum(axis=-1) + qn_dev = norm_sq.reshape(batch_size, num_h, num_q_blocks, bq).max(axis=-1) # (batch, heads, num_q_blocks) + kfb_centered = kfb - k_mean[:, :, None, :] + mk_dev = (kfb_centered * kfb_centered).sum(axis=-1).max(axis=-1) # (batch, heads) + ring_kernel = 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=True, + ) + return jax.vmap(ring_kernel, in_axes=(0, 0, 0, (0, 0), 0))(qb, kb, vb, (qn_dev, mk_dev), k_mean) + + out_batched = _body_batched(q_in, k_in, v_batch).astype(jnp.float32) + + # Verify batch item 0 (clean prompt) matches exact unbatched fixed-m reference + ref0 = self._reference(q_in[0], k_in[0], v0) + diff0 = float(jnp.max(jnp.abs(out_batched[0] - ref0))) + self.assertLess(diff0, 2e-2) + + # Verify batch item 1 (sink outlier prompt) matches exact unbatched reference + ref1 = self._reference(q_in[1], k_in[1], v1) + diff1 = float(jnp.max(jnp.abs(out_batched[1] - ref1))) + self.assertLess(diff1, 2e-2) + + def test_ring_phase_transition_boundary_continuity(self): + """Verifies seamless continuity between fixed-m and online mode across the dynamic centered Ring safe bound threshold.""" + ring_size = 2 + q_base, k_base, v = self._random_qkv(ring_size=ring_size) + q_normed = q_base / jnp.sqrt((q_base.astype(jnp.float32) ** 2).sum(-1, keepdims=True)) + k_normed = k_base / jnp.sqrt((k_base.astype(jnp.float32) ** 2).sum(-1, keepdims=True)) + + _, safe_bound = custom_splash.get_fixed_m_constants(self.shard_len * ring_size, is_ring=False) + test_bounds = [ + safe_bound - 2.0, + safe_bound - 0.5, + safe_bound - 0.01, + safe_bound, + safe_bound + 0.01, + safe_bound + 0.5, + safe_bound + 2.0, + ] + for target_bound in test_bounds: + factor = math.sqrt(target_bound / _LOG2E / self.scale) + q = (q_normed * factor).astype(jnp.bfloat16) + k = (k_normed * factor).astype(jnp.bfloat16) + + q_in, k_in = self._scaled_inputs(q, k) + out_fixed = self._run_ring(q_in, k_in, v, ring_size=ring_size, use_fixed_m=True).astype(jnp.float32) + out_online = self._run_ring(q_in, k_in, v, ring_size=ring_size, use_fixed_m=False).astype(jnp.float32) + + self.assertTrue(bool(jnp.all(jnp.isfinite(out_fixed)))) + diff = float(jnp.max(jnp.abs(out_fixed - out_online))) + self.assertLess(diff, 2e-2, f"Ring discontinuity at bound={target_bound}, diff={diff}") + + def test_per_head_single_rank_outlier_sync_ring(self): + """Verifies that in per_q_block=False mode, an outlier on a single ring rank synchronizes via pmin across all ranks.""" + ring_size = 2 + # Create Q with outlier only on Rank 1 (rows self.shard_len to 2 * self.shard_len) + q, k, v = self._random_qkv(ring_size=ring_size, q_gain=(0, slice(self.shard_len, self.shard_len * 2), 40.0)) + q_in, k_in = self._scaled_inputs(q, k) + mesh = self._mesh_for_size(ring_size) + spec = jax.sharding.PartitionSpec(None, _RING_AXIS, None) + + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=(spec, spec, spec), + out_specs=spec, + check_vma=False, + ) + def _body_per_head(ql, kl, vl): + qf = ql.astype(jnp.float32) + kf = kl.astype(jnp.float32) + k_mean_local = jnp.mean(kf, axis=1) + k_mean = jax.lax.pmean(k_mean_local, axis_name=_RING_AXIS) + qn_max_sq = (qf * qf).sum(-1).max(axis=1) # (heads,) per-head 1D squared norm + kf_centered = kf - k_mean[:, None, :] + mk_h_sq = (kf_centered * kf_centered).sum(-1).max(axis=1) + 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=True, + per_q_block=False, + fixed_m_norms=(qn_max_sq, mk_h_sq), + k_mean=k_mean, + ) + return ring(ql, kl, vl) + + out = _body_per_head(q_in, k_in, v).astype(jnp.float32) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + ref = self._reference(q_in, k_in, v) + diff = float(jnp.max(jnp.abs(out - ref))) + self.assertLess(diff, 2e-2) + + def test_adversarial_unsmoothed_negative_keys_ring(self): + """Verifies that un-smoothed keys with large negative bias merge safely without underflow NaNs.""" + ring_size = 2 + q, k, v = self._random_qkv(ring_size=ring_size) + k_negative = k - 25.0 # Force strong negative bias across both ring shards + q_in, k_in = self._scaled_inputs(q, k_negative) + + out_fixed = self._run_ring(q_in, k_in, v, ring_size=ring_size, use_fixed_m=True).astype(jnp.float32) + out_online = self._run_ring(q_in, k_in, v, ring_size=ring_size, use_fixed_m=False).astype(jnp.float32) + + self.assertTrue(bool(jnp.all(jnp.isfinite(out_fixed)))) + diff = float(jnp.max(jnp.abs(out_fixed - out_online))) + self.assertLess(diff, 2e-2) + + def test_adversarial_hybrid_ring_negative_shard(self): + """Verifies hybrid LSE fallback when one shard is heavily negative and another positive (global mean 0).""" + ring_size = 2 + q, k, v = self._random_qkv(ring_size=ring_size) + # Shard 0: heavily negative, Shard 1: heavily positive + k = k.at[:, : self.shard_len, :].add(-20.0) + k = k.at[:, self.shard_len :, :].add(20.0) + # Force rank 0 to exceed the global bound so execution enters the hybrid fallback branch + q = q.at[0, : self.block_sizes.block_q, :].multiply(40.0) + q_in, k_in = self._scaled_inputs(q, k) + + out_hybrid = self._run_ring(q_in, k_in, v, ring_size=ring_size, use_fixed_m=True).astype(jnp.float32) + ref = self._reference(q_in, k_in, v) + + self.assertTrue(bool(jnp.all(jnp.isfinite(out_hybrid)))) + diff = float(jnp.max(jnp.abs(out_hybrid - ref))) + self.assertLess(diff, 2e-2, f"Adversarial hybrid ring output diverged from reference: diff={diff}") + + # --- P1 regression: the ring LSE fallback must honour the global V check --- + # + # The fallback re-derives per-hop eligibility from Q/K norms alone. Those are + # the only quantities recoverable from a single hop, so a failed V check used + # to be silently discarded and individual hops re-enabled fixed-m -- parking + # weights at 2**C with an out-of-contract |V| and overflowing to inf. These + # tests execute the fallback and assert on the *output*, not on metadata. + + def test_oversized_v_ring_fallback_is_finite(self): + """|V| beyond the contract must fall back to online arithmetic, not overflow.""" + ring_size = 2 + q, k, v = self._random_qkv(ring_size=ring_size) + v_big = (v.astype(jnp.float32) * 1024.0).astype(v.dtype) + self.assertFalse(self._global_v_ok(v_big), "test setup: V should violate the bound") + q_in, k_in = self._scaled_inputs(q, k) + + out = self._run_ring(q_in, k_in, v_big, ring_size=ring_size, use_fixed_m=True, v_ok=False).astype(jnp.float32) + ref = self._reference(q_in, k_in, v_big) + + self.assertTrue(bool(jnp.all(jnp.isfinite(out))), "oversized |V| overflowed the ring fixed-m fallback") + # Tolerance scales with |V|: the 2e-2 used elsewhere is for unit-scale V. + self.assertLess(float(jnp.max(jnp.abs(out - ref))), 2e-2 * 1024.0) + + def test_oversized_v_on_single_rank_disqualifies_whole_ring(self): + """Only rank 1 has oversized V; the globally reduced verdict must protect every hop.""" + ring_size = 2 + q, k, v = self._random_qkv(ring_size=ring_size) + # Rank 1 owns the second shard_len slice of the sequence. + v_mixed = v.at[:, self.shard_len :, :].set((v[:, self.shard_len :, :].astype(jnp.float32) * 1024.0).astype(v.dtype)) + self.assertFalse(self._global_v_ok(v_mixed), "test setup: global V should violate the bound") + q_in, k_in = self._scaled_inputs(q, k) + + # v_ok is a *global* reduction, so a single offending rank disqualifies all. + # We do NOT pass v_ok=False: the automatic cross-rank pmin reduction must compute it. + out = self._run_ring(q_in, k_in, v_mixed, ring_size=ring_size, use_fixed_m=True).astype(jnp.float32) + ref = self._reference(q_in, k_in, v_mixed) + + self.assertTrue(bool(jnp.all(jnp.isfinite(out))), "single-rank oversized |V| overflowed the ring fallback") + self.assertLess(float(jnp.max(jnp.abs(out - ref))), 2e-2 * 1024.0) + + def test_v_gate_is_what_prevents_the_overflow(self): + """Witness: without the gate the same inputs are unsafe. + + Guards against the gate being quietly dropped again. If a future change makes + the fallback safe by construction this test should be deleted, not muted -- + but it must never be allowed to pass by accident. + """ + ring_size = 2 + q, k, v = self._random_qkv(ring_size=ring_size) + v_big = (v.astype(jnp.float32) * 1024.0).astype(v.dtype) + q_in, k_in = self._scaled_inputs(q, k) + + gated = self._run_ring(q_in, k_in, v_big, ring_size=ring_size, use_fixed_m=True, v_ok=False).astype(jnp.float32) + self.assertTrue(bool(jnp.all(jnp.isfinite(gated)))) + + ungated = self._run_ring(q_in, k_in, v_big, ring_size=ring_size, use_fixed_m=True, v_ok=True).astype(jnp.float32) + if bool(jnp.all(jnp.isfinite(ungated))): + self.skipTest("ungated path happens to stay finite for these inputs; gate still required in general") + + def test_gqa_ring_fixed_m_shard_map(self): + """Verifies that GQA (4 Q heads, 2 KV heads) works seamlessly across ring ranks.""" + ring_size = 2 + num_q_heads = 4 + num_kv_heads = 2 + total_seq = self.shard_len * ring_size + q = jax.random.normal(jax.random.PRNGKey(101), (num_q_heads, total_seq, self.head_dim), jnp.bfloat16) + k = jax.random.normal(jax.random.PRNGKey(102), (num_kv_heads, total_seq, self.head_dim), jnp.bfloat16) + v = jax.random.normal(jax.random.PRNGKey(103), (num_kv_heads, total_seq, self.head_dim), jnp.bfloat16) + q_in, k_in = self._scaled_inputs(q, k) + + mesh = self._mesh_for_size(ring_size) + spec_q = jax.sharding.PartitionSpec(None, _RING_AXIS, None) + spec_kv = jax.sharding.PartitionSpec(None, _RING_AXIS, None) + + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=(spec_q, spec_kv, spec_kv), + out_specs=spec_q, + check_vma=False, + ) + def _body_gqa(ql, kl, vl): + qf = ql.astype(jnp.float32) + kf = kl.astype(jnp.float32) + k_mean_local = jnp.mean(kf, axis=1) # (kv_heads, dim) + k_mean = jax.lax.pmean(k_mean_local, axis_name=_RING_AXIS) + bq = self.block_sizes.block_q + num_q_blocks = qf.shape[1] // bq + qf_blocks = qf.reshape(num_q_heads, num_q_blocks, bq, self.head_dim) + qn_blocks_sq = (qf_blocks * qf_blocks).sum(-1).max(axis=-1) # (q_heads, num_q_blocks) + kf_centered = kf - k_mean[:, None, :] + mk_h_sq = (kf_centered * kf_centered).sum(-1).max(axis=1) # (kv_heads,) + 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=True, + fixed_m_norms=(qn_blocks_sq, mk_h_sq), + k_mean=k_mean, + ) + return ring(ql, kl, vl) + + out = _body_gqa(q_in, k_in, v).astype(jnp.float32) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + # GQA Reference + q_rep = q_in + k_rep = jnp.repeat(k_in, num_q_heads // num_kv_heads, axis=0) + v_rep = jnp.repeat(v, num_q_heads // num_kv_heads, axis=0) + ref = self._reference(q_rep, k_rep, v_rep) + diff = float(jnp.max(jnp.abs(out - ref))) + self.assertLess(diff, 2e-2, f"GQA ring output diverged from reference: diff={diff}") + + def test_fixed_m_mismatched_ring_size_raises(self): + """Verifies that ring_size != axis_size raises NotImplementedError when use_fixed_m=True.""" + q, k, v = self._random_qkv(ring_size=2) + q_in, k_in = self._scaled_inputs(q, k) + mesh = self._mesh_for_size(2) + spec = jax.sharding.PartitionSpec(None, _RING_AXIS, None) + + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=(spec, spec, spec), + out_specs=spec, + check_vma=False, + ) + def _body(ql, kl, vl): + qf = ql.astype(jnp.float32) + kf = kl.astype(jnp.float32) + k_mean_local = jnp.mean(kf, axis=1) + k_mean = jax.lax.pmean(k_mean_local, axis_name=_RING_AXIS) + qn_max_sq = (qf * qf).sum(-1).max(axis=1) + kf_centered = kf - k_mean[:, None, :] + mk_h_sq = (kf_centered * kf_centered).sum(-1).max(axis=1) + 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=1, # Mismatched: ring_size=1 != axis_size=2 + use_fixed_m=True, + fixed_m_norms=(qn_max_sq, mk_h_sq), + k_mean=k_mean, + ) + return ring(ql, kl, vl) + + with self.assertRaises(NotImplementedError): + _body(q_in, k_in, v) + + def test_fixed_m_non_canonical_perm_raises(self): + """Verifies that non-canonical perm raises NotImplementedError when use_fixed_m=True.""" + q, k, v = self._random_qkv(ring_size=2) + q_in, k_in = self._scaled_inputs(q, k) + mesh = self._mesh_for_size(2) + spec = jax.sharding.PartitionSpec(None, _RING_AXIS, None) + + @functools.partial( + jax.shard_map, + mesh=mesh, + in_specs=(spec, spec, spec), + out_specs=spec, + check_vma=False, + ) + def _body(ql, kl, vl): + qf = ql.astype(jnp.float32) + kf = kl.astype(jnp.float32) + k_mean_local = jnp.mean(kf, axis=1) + k_mean = jax.lax.pmean(k_mean_local, axis_name=_RING_AXIS) + qn_max_sq = (qf * qf).sum(-1).max(axis=1) + kf_centered = kf - k_mean[:, None, :] + mk_h_sq = (kf_centered * kf_centered).sum(-1).max(axis=1) + 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=2, + perm=[(0, 0), (1, 1)], # Non-canonical identity permutation + use_fixed_m=True, + fixed_m_norms=(qn_max_sq, mk_h_sq), + k_mean=k_mean, + ) + return ring(ql, kl, vl) + + with self.assertRaises(NotImplementedError): + _body(q_in, k_in, v) + + def test_gqa_with_chunked_ulysses_raises(self): + """Verifies that GQA (Hq != Hkv) with ulysses_attention_chunks > 1 raises NotImplementedError.""" + q = jnp.zeros((1, 8, 128, 64), dtype=jnp.float32) + k = jnp.zeros((1, 2, 128, 64), dtype=jnp.float32) + v = jnp.zeros((1, 2, 128, 64), dtype=jnp.float32) + + with self.assertRaises(NotImplementedError): + attention_flax._run_chunked_ulysses_attention( + q, + k, + v, + num_heads=8, + ulysses_shards=2, + ulysses_attention_chunks=2, + attention_fn=lambda q, k, v: q, + ) + + def test_2d_gqa_ulysses_ring_attention(self): + """Verifies that 2D Ulysses+Ring attention correctly executes GQA (Hq=8, Hkv=2) with chunks=1.""" + if len(jax.devices()) < 4: + self.skipTest("Requires 4 devices for 2D Ulysses+Ring test.") + + devices = np.array(jax.devices()[:4]).reshape(1, 1, 4, 1) + mesh = jax.sharding.Mesh(devices, ("data", "fsdp", "context", "tensor")) + axis_rules = ( + (attention_flax.BATCH, "data"), + (attention_flax.LENGTH, "context"), + (attention_flax.HEAD, None), + (attention_flax.SELF_ATTN_HEAD, None), + (attention_flax.SELF_ATTN_Q_LENGTH, "context"), + (attention_flax.SELF_ATTN_KV_LENGTH, "context"), + (attention_flax.D_KV, None), + ) + + batch = 1 + length = 2048 + q_heads = 8 + kv_heads = 2 + head_dim = 128 + + q = jax.random.normal(jax.random.PRNGKey(10), (batch, length, q_heads * head_dim), dtype=jnp.bfloat16) + k = jax.random.normal(jax.random.PRNGKey(11), (batch, length, kv_heads * head_dim), dtype=jnp.bfloat16) + v = jax.random.normal(jax.random.PRNGKey(12), (batch, length, kv_heads * head_dim), dtype=jnp.bfloat16) + + flash_block_sizes = { + "block_q": 1024, + "block_kv": 1024, + "block_kv_compute": 512, + "block_kv_compute_in": 256, + "heads_per_tile": 1, + "vmem_limit_bytes": 67108864, + } + + with mesh, nn_partitioning.axis_rules(axis_rules): + out = attention_flax._ulysses_ring_custom_attention( + q, + k * (1.0 / math.sqrt(head_dim)), + v, + heads=q_heads, + mesh=mesh, + axis_names_q=( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_Q_LENGTH, + attention_flax.D_KV, + ), + axis_names_kv=( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_KV_LENGTH, + attention_flax.D_KV, + ), + flash_block_sizes=flash_block_sizes, + dtype=jnp.bfloat16, + ulysses_shards=2, + use_base2_exp=True, + use_fixed_m=True, + per_q_block=True, + ulysses_attention_chunks=1, + kv_heads=kv_heads, + ) + self.assertEqual(out.shape, (batch, length, q_heads * head_dim)) + self.assertTrue(bool(jnp.all(jnp.isfinite(out)))) + + # Reference calculation: repeat KV heads to match Q heads (Hq=8, Hkv=2 => repeat factor 4) + q_unflat = q.reshape((batch, length, q_heads, head_dim)).swapaxes(1, 2) # [B, Hq, L, D] + k_unflat = k.reshape((batch, length, kv_heads, head_dim)).swapaxes(1, 2) # [B, Hkv, L, D] + v_unflat = v.reshape((batch, length, kv_heads, head_dim)).swapaxes(1, 2) # [B, Hkv, L, D] + + k_repeated = jnp.repeat(k_unflat, q_heads // kv_heads, axis=1) # [B, Hq, L, D] + v_repeated = jnp.repeat(v_unflat, q_heads // kv_heads, axis=1) # [B, Hq, L, D] + + # Reference scaled dot product attention in FP32 + scores = jnp.einsum( + "bhqd,bhkd->bhqk", + q_unflat.astype(jnp.float32) * (1.0 / math.sqrt(head_dim)), + k_repeated.astype(jnp.float32), + ) + attn_weights = jax.nn.softmax(scores, axis=-1) + ref_out = jnp.einsum("bhqk,bhkd->bhqd", attn_weights, v_repeated.astype(jnp.float32)) + ref_out = ref_out.swapaxes(1, 2).reshape((batch, length, q_heads * head_dim)) + + np.testing.assert_allclose( + np.array(out, dtype=np.float32), + np.array(ref_out, dtype=np.float32), + rtol=2e-2, + atol=2e-2, + ) + + def test_oversized_v_single_rank_production_dispatch(self): + """Verifies that attention_flax._ulysses_ring_custom_attention reduces v_ok across ranks. + + Only rank 1 has oversized V. If v_ok were local, rank 0 would retain v_ok=True and + overflow to inf on the second hop. The cross-ring reduction ensures both ranks fall back. + """ + if len(jax.devices()) < 4: + self.skipTest("Requires 4 devices for 2D Ulysses+Ring test.") + + devices = np.array(jax.devices()[:4]).reshape(1, 1, 4, 1) + mesh = jax.sharding.Mesh(devices, ("data", "fsdp", "context", "tensor")) + axis_rules = ( + (attention_flax.BATCH, "data"), + (attention_flax.LENGTH, "context"), + (attention_flax.HEAD, None), + (attention_flax.SELF_ATTN_HEAD, None), + (attention_flax.SELF_ATTN_Q_LENGTH, "context"), + (attention_flax.SELF_ATTN_KV_LENGTH, "context"), + (attention_flax.D_KV, None), + ) + + batch = 1 + length = 8192 + q_heads = 4 + kv_heads = 4 + head_dim = 128 + + # Q=K=0 and constant V=1 on rank 0, V=1024 on rank 1. + # With 4 context devices and ulysses=2, ring_size=2. Each ring rank gets 4096 tokens (4 blocks of 1024). + # If v_ok is not reduced across ranks, rank 0 retains v_ok=True and attempts fixed-m accumulation + # on rank 1's 4 blocks of V=1024, overflowing FP32 (4096 * 1024 * 2^107 = 2^129 > 2^128). + q = jnp.zeros((batch, length, q_heads * head_dim), dtype=jnp.bfloat16) + k = jnp.zeros((batch, length, kv_heads * head_dim), dtype=jnp.bfloat16) + half = length // 2 + v_mixed = jnp.ones((batch, length, kv_heads * head_dim), dtype=jnp.bfloat16) + v_mixed = v_mixed.at[:, half:, :].set(1024.0) + + flash_block_sizes = { + "block_q": 1024, + "block_kv": 1024, + "block_kv_compute": 512, + "block_kv_compute_in": 256, + "heads_per_tile": 1, + "vmem_limit_bytes": 67108864, + } + + with mesh, nn_partitioning.axis_rules(axis_rules): + out = attention_flax._ulysses_ring_custom_attention( + q, + k, + v_mixed, + heads=q_heads, + mesh=mesh, + axis_names_q=( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_Q_LENGTH, + attention_flax.D_KV, + ), + axis_names_kv=( + attention_flax.BATCH, + attention_flax.SELF_ATTN_HEAD, + attention_flax.SELF_ATTN_KV_LENGTH, + attention_flax.D_KV, + ), + flash_block_sizes=flash_block_sizes, + dtype=jnp.bfloat16, + ulysses_shards=2, + use_base2_exp=True, + use_fixed_m=True, + per_q_block=True, + ulysses_attention_chunks=1, + kv_heads=kv_heads, + ) + self.assertTrue(bool(jnp.all(jnp.isfinite(out))), "Production dispatch overflowed on mixed-rank oversized V") + + # With Q=K=0, attention weights are uniform 1/N. Expected output is exactly (1.0 + 1024.0)/2 = 512.5 everywhere. + expected_val = (1.0 + 1024.0) / 2.0 + out_f32 = np.array(out, dtype=np.float32) + np.testing.assert_allclose( + out_f32, + np.full_like(out_f32, expected_val), + rtol=2e-2, + atol=2.0, + ) if __name__ == "__main__": diff --git a/src/maxdiffusion/tests/wan/wan_transformer_test.py b/src/maxdiffusion/tests/wan/wan_transformer_test.py index 69bed9a6a..1685ec6f2 100644 --- a/src/maxdiffusion/tests/wan/wan_transformer_test.py +++ b/src/maxdiffusion/tests/wan/wan_transformer_test.py @@ -34,9 +34,11 @@ ) from maxdiffusion.models.embeddings_flax import NNXTimestepEmbedding, NNXPixArtAlphaTextProjection from maxdiffusion.models.normalization_flax import FP32LayerNorm -from maxdiffusion.models.attention_flax import FlaxWanAttention +from maxdiffusion.models.attention_flax import FlaxWanAttention, _unflatten_heads +from maxdiffusion.kernels.fused_producers import fused_ln_adaln, fused_rmsnorm_rope from maxdiffusion.pyconfig import HyperParameters from maxdiffusion.pipelines.wan.wan_pipeline import WanPipeline +import numpy as np import qwix import flax @@ -250,6 +252,311 @@ def test_wan_attention(self): except NotImplementedError: pass + def test_fused_ln_adaln_parity(self): + """Verifies numerical parity of fused FP32 LayerNorm + AdaLN modulation.""" + key = jax.random.PRNGKey(101) + k1, k2, k3 = jax.random.split(key, 3) + B, S, D = 1, 1024, 5120 + x = jax.random.normal(k1, (B, S, D), dtype=jnp.bfloat16) + scale = jax.random.normal(k2, (B, 1, D), dtype=jnp.bfloat16) + shift = jax.random.normal(k3, (B, 1, D), dtype=jnp.bfloat16) + + # Reference implementation + x_fp32 = x.astype(jnp.float32) + mean = jnp.mean(x_fp32, axis=-1, keepdims=True) + var = jnp.var(x_fp32, axis=-1, keepdims=True) + x_ln = (x_fp32 - mean) * jax.lax.rsqrt(var + 1e-6) + ref_out = (x_ln * (1.0 + scale.astype(jnp.float32)) + shift.astype(jnp.float32)).astype(x.dtype) + + # Fused producer + fused_out = fused_ln_adaln(x, scale, shift, eps=1e-6) + + np.testing.assert_allclose( + np.array(ref_out, dtype=np.float32), + np.array(fused_out, dtype=np.float32), + atol=0.01, + rtol=1e-2, + ) + + def test_fused_rmsnorm_rope_parity(self): + """Verifies numerical parity of fused RMSNorm + RoPE against unfused reference with complex freqs_cis.""" + key = jax.random.PRNGKey(202) + k1, k2, k3, k4, k5, k6 = jax.random.split(key, 6) + B, S, D, H, DH = 1, 1024, 5120, 40, 128 + eps = 1e-6 + + raw_q = jax.random.normal(k1, (B, S, D), dtype=jnp.bfloat16) + raw_k = jax.random.normal(k2, (B, S, D), dtype=jnp.bfloat16) + q_scale = jax.random.normal(k3, (D,), dtype=jnp.bfloat16) + k_scale = jax.random.normal(k4, (D,), dtype=jnp.bfloat16) + + freqs_real = jax.random.normal(k5, (1, 1, S, DH // 2), dtype=jnp.float32) + freqs_imag = jax.random.normal(k6, (1, 1, S, DH // 2), dtype=jnp.float32) + freqs_cis = jax.lax.complex(freqs_real, freqs_imag) + + # Unfused reference: FP32 RMSNorm -> unflatten -> RoPE + def ref_norm(x, scale): + var = jnp.mean(jnp.square(x.astype(jnp.float32)), axis=-1, keepdims=True) + return (x.astype(jnp.float32) * jax.lax.rsqrt(var + eps) * scale.astype(jnp.float32)).astype(x.dtype) + + def ref_unflatten(x, heads): + b, s, d = x.shape + return x.reshape(b, s, heads, d // heads).transpose(0, 2, 1, 3) + + def ref_apply_rope(xq, xk, freqs_cis): + cos = jnp.real(freqs_cis).astype(xq.dtype) + sin = jnp.imag(freqs_cis).astype(xq.dtype) + xq_reshaped = xq.reshape(*xq.shape[:-1], -1, 2) + xk_reshaped = xk.reshape(*xk.shape[:-1], -1, 2) + xq_0, xq_1 = xq_reshaped[..., 0], xq_reshaped[..., 1] + xk_0, xk_1 = xk_reshaped[..., 0], xk_reshaped[..., 1] + xq_out_0 = xq_0 * cos - xq_1 * sin + xq_out_1 = xq_0 * sin + xq_1 * cos + xk_out_0 = xk_0 * cos - xk_1 * sin + xk_out_1 = xk_0 * sin + xk_1 * cos + xq_out = jnp.concatenate([xq_out_0[..., None], xq_out_1[..., None]], axis=-1).reshape(xq.shape) + xk_out = jnp.concatenate([xk_out_0[..., None], xk_out_1[..., None]], axis=-1).reshape(xk.shape) + return xq_out, xk_out + + q_ref, k_ref = ref_apply_rope( + ref_unflatten(ref_norm(raw_q, q_scale), H), + ref_unflatten(ref_norm(raw_k, k_scale), H), + freqs_cis, + ) + + # Fused producer + q_fused, k_fused = fused_rmsnorm_rope( + raw_q, + raw_k, + q_scale, + k_scale, + freqs_cis, + q_heads=H, + dim_head=DH, + eps=eps, + ) + + np.testing.assert_allclose( + np.array(q_ref, dtype=np.float32), + np.array(q_fused, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + np.testing.assert_allclose( + np.array(k_ref, dtype=np.float32), + np.array(k_fused, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + + def test_fused_rmsnorm_rope_with_wan_rotary_embed(self): + """Verifies numerical parity of fused RMSNorm + RoPE against unfused reference with real WanRotaryPosEmbed frequencies.""" + key = jax.random.PRNGKey(404) + k1, k2, k3, k4 = jax.random.split(key, 4) + B, S, D, H, DH = 1, 1024, 5120, 40, 128 + eps = 1e-6 + + # Generate real RoPE frequencies on the complex unit circle using WanRotaryPosEmbed + wan_rot_embed = WanRotaryPosEmbed(attention_head_dim=DH, patch_size=[1, 2, 2], max_seq_len=1024) + dummy_video = jnp.ones((B, 1, 64, 64, 16)) + freqs_cis = wan_rot_embed(dummy_video) # (1, 1, 1024, 64) + + raw_q = jax.random.normal(k1, (B, S, D), dtype=jnp.bfloat16) + raw_k = jax.random.normal(k2, (B, S, D), dtype=jnp.bfloat16) + q_scale = jax.random.normal(k3, (D,), dtype=jnp.bfloat16) + k_scale = jax.random.normal(k4, (D,), dtype=jnp.bfloat16) + + def ref_norm(x, scale): + var = jnp.mean(jnp.square(x.astype(jnp.float32)), axis=-1, keepdims=True) + return (x.astype(jnp.float32) * jax.lax.rsqrt(var + eps) * scale.astype(jnp.float32)).astype(x.dtype) + + def ref_unflatten(x, heads): + b, s, d = x.shape + return x.reshape(b, s, heads, d // heads).transpose(0, 2, 1, 3) + + def ref_apply_rope(xq, xk, freqs_cis): + cos = jnp.real(freqs_cis).astype(xq.dtype) + sin = jnp.imag(freqs_cis).astype(xq.dtype) + xq_reshaped = xq.reshape(*xq.shape[:-1], -1, 2) + xk_reshaped = xk.reshape(*xk.shape[:-1], -1, 2) + xq_0, xq_1 = xq_reshaped[..., 0], xq_reshaped[..., 1] + xk_0, xk_1 = xk_reshaped[..., 0], xk_reshaped[..., 1] + xq_out_0 = xq_0 * cos - xq_1 * sin + xq_out_1 = xq_0 * sin + xq_1 * cos + xk_out_0 = xk_0 * cos - xk_1 * sin + xk_out_1 = xk_0 * sin + xk_1 * cos + xq_out = jnp.concatenate([xq_out_0[..., None], xq_out_1[..., None]], axis=-1).reshape(xq.shape) + xk_out = jnp.concatenate([xk_out_0[..., None], xk_out_1[..., None]], axis=-1).reshape(xk.shape) + return xq_out, xk_out + + q_ref, k_ref = ref_apply_rope( + ref_unflatten(ref_norm(raw_q, q_scale), H), + ref_unflatten(ref_norm(raw_k, k_scale), H), + freqs_cis, + ) + + q_fused, k_fused = fused_rmsnorm_rope( + raw_q, + raw_k, + q_scale, + k_scale, + freqs_cis, + q_heads=H, + dim_head=DH, + eps=eps, + ) + + np.testing.assert_allclose( + np.array(q_ref, dtype=np.float32), + np.array(q_fused, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + np.testing.assert_allclose( + np.array(k_ref, dtype=np.float32), + np.array(k_fused, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + + def test_fused_rmsnorm_rope_gqa_parity(self): + """Verifies that fused_rmsnorm_rope correctly handles asymmetric GQA shapes (e.g. q_heads=8, kv_heads=2).""" + key = jax.random.PRNGKey(505) + k1, k2, k3, k4, k5 = jax.random.split(key, 5) + + B = 2 + S = 64 + Q_H = 8 + KV_H = 2 + DH = 128 + D_q = Q_H * DH + D_kv = KV_H * DH + eps = 1e-6 + + raw_q = jax.random.normal(k1, (B, S, D_q), dtype=jnp.bfloat16) + raw_k = jax.random.normal(k2, (B, S, D_kv), dtype=jnp.bfloat16) + q_scale = jax.random.normal(k3, (D_q,), dtype=jnp.float32) + k_scale = jax.random.normal(k4, (D_kv,), dtype=jnp.float32) + freqs_cis = jax.random.normal(k5, (1, 1, S, DH // 2), dtype=jnp.float32) + 1j * jax.random.normal( + key, (1, 1, S, DH // 2), dtype=jnp.float32 + ) + + def ref_norm(x, scale): + x_fp32 = x.astype(jnp.float32) + rms = jax.lax.rsqrt(jnp.mean(jnp.square(x_fp32), axis=-1, keepdims=True) + eps) + return (x_fp32 * rms * scale).astype(x.dtype) + + def ref_unflatten(x, heads): + return x.reshape(B, S, heads, DH).transpose(0, 2, 1, 3) + + def ref_apply_rope(xq, xk, freqs): + cos = jnp.real(freqs).astype(xq.dtype) + sin = jnp.imag(freqs).astype(xq.dtype) + xq_reshaped = xq.reshape(*xq.shape[:-1], -1, 2) + xk_reshaped = xk.reshape(*xk.shape[:-1], -1, 2) + xq_0, xq_1 = xq_reshaped[..., 0], xq_reshaped[..., 1] + xk_0, xk_1 = xk_reshaped[..., 0], xk_reshaped[..., 1] + xq_out_0 = xq_0 * cos - xq_1 * sin + xq_out_1 = xq_0 * sin + xq_1 * cos + xk_out_0 = xk_0 * cos - xk_1 * sin + xk_out_1 = xk_0 * sin + xk_1 * cos + xq_out = jnp.concatenate([xq_out_0[..., None], xq_out_1[..., None]], axis=-1).reshape(xq.shape) + xk_out = jnp.concatenate([xk_out_0[..., None], xk_out_1[..., None]], axis=-1).reshape(xk.shape) + return xq_out, xk_out + + q_ref, k_ref = ref_apply_rope( + ref_unflatten(ref_norm(raw_q, q_scale), Q_H), + ref_unflatten(ref_norm(raw_k, k_scale), KV_H), + freqs_cis, + ) + + q_fused, k_fused = fused_rmsnorm_rope( + raw_q, + raw_k, + q_scale, + k_scale, + freqs_cis, + q_heads=Q_H, + kv_heads=KV_H, + dim_head=DH, + eps=eps, + ) + + self.assertEqual(q_fused.shape, (B, Q_H, S, DH)) + self.assertEqual(k_fused.shape, (B, KV_H, S, DH)) + + np.testing.assert_allclose( + np.array(q_ref, dtype=np.float32), + np.array(q_fused, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + np.testing.assert_allclose( + np.array(k_ref, dtype=np.float32), + np.array(k_fused, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + + def test_wan_self_attention_is_self_attention_dispatch(self): + """Verifies that FlaxWanAttention with is_self_attention=True correctly dispatches fused RMSNorm+RoPE and matches reference.""" + key = jax.random.PRNGKey(303) + k1, k2, k3, k4 = jax.random.split(key, 4) + rngs = nnx.Rngs(k1) + + batch_size = 1 + seq_len = 1024 + query_dim = 5120 + heads = 40 + dim_head = 128 + + flash_block_sizes = get_flash_block_sizes(self.config) + with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + attn = FlaxWanAttention( + rngs=rngs, + query_dim=query_dim, + heads=heads, + dim_head=dim_head, + attention_kernel="flash", + mesh=self.mesh, + flash_block_sizes=flash_block_sizes, + is_self_attention=True, + ) + self.assertTrue(attn.is_self_attention) + + hidden_states = jax.random.normal(k2, (batch_size, seq_len, query_dim), dtype=jnp.bfloat16) + freqs_real = jax.random.normal(k3, (1, 1, seq_len, dim_head // 2), dtype=jnp.float32) + freqs_imag = jax.random.normal(k4, (1, 1, seq_len, dim_head // 2), dtype=jnp.float32) + rotary_emb = jax.lax.complex(freqs_real, freqs_imag) + + # Wan self-attention call passes encoder_hidden_states=hidden_states + out = attn( + hidden_states=hidden_states, + encoder_hidden_states=hidden_states, + rotary_emb=rotary_emb, + ) + self.assertEqual(out.shape, (batch_size, seq_len, query_dim)) + + # Reference unfused path execution with identical weights + raw_q = attn.query(hidden_states) + raw_k = attn.key(hidden_states) + raw_v = attn.value(hidden_states) + q_norm = attn.norm_q(raw_q) + k_norm = attn.norm_k(raw_k) + q_h = _unflatten_heads(q_norm, heads) + k_h = _unflatten_heads(k_norm, heads) + v_h = _unflatten_heads(raw_v, heads) + q_rope, k_rope = attn._apply_rope(q_h, k_h, rotary_emb) + ref_attn_out = attn.attention_op.apply_attention(q_rope, k_rope, v_h, attention_mask=None) + ref_out = attn.proj_attn(ref_attn_out) + + np.testing.assert_allclose( + np.array(out, dtype=np.float32), + np.array(ref_out, dtype=np.float32), + atol=0.08, + rtol=1e-2, + ) + @pytest.mark.skipif(IN_GITHUB_ACTIONS, reason="Don't run smoke tests on Github Actions") def test_wan_model(self): pyconfig.initialize( diff --git a/src/maxdiffusion/utils/export_utils.py b/src/maxdiffusion/utils/export_utils.py index 279ad1e90..46c51086f 100644 --- a/src/maxdiffusion/utils/export_utils.py +++ b/src/maxdiffusion/utils/export_utils.py @@ -212,21 +212,18 @@ def export_to_video( if output_video_path is None: output_video_path = tempfile.NamedTemporaryFile(suffix=".mp4").name - if isinstance(video_frames, np.ndarray): - if video_frames.dtype != np.uint8: - video_frames = (video_frames * 255).astype(np.uint8) - elif isinstance(video_frames[0], np.ndarray): - video_frames = np.stack(video_frames) - if video_frames.dtype != np.uint8: - video_frames = (video_frames * 255).astype(np.uint8) - elif isinstance(video_frames[0], PIL.Image.Image): + if isinstance(video_frames, list) and len(video_frames) > 0 and isinstance(video_frames[0], PIL.Image.Image): video_frames = np.stack([np.asarray(frame) for frame in video_frames]) + else: + video_frames = np.asarray(video_frames) + if video_frames.dtype != np.uint8: + video_frames = (video_frames * 255).clip(0, 255).astype(np.uint8) with imageio.get_writer( output_video_path, fps=fps, quality=quality, bitrate=bitrate, macro_block_size=macro_block_size ) as writer: for frame in video_frames: - writer.append_data(frame) + writer.append_data(np.asarray(frame)) return output_video_path