# Copyright 2024 The JAX Authors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # https://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from functools import partial from absl.testing import absltest import numpy as np import jax import jax.numpy as jnp from jax.sharding import Mesh from jax.sharding import PartitionSpec, NamedSharding from jax._src import config from jax._src import test_util as jtu from jax._src.cudnn.fused_attention_stablehlo import ( dot_product_attention, check_cudnn_version, MaskType, ) from typing import Callable, Tuple config.parse_flags_with_absl() Array = jnp.ndarray fp8_meta_names = [ "amax_dQ", "amax_dK", "amax_dV", "amax_dP", "descale_q", "descale_k", "descale_v", "descale_s", "scale_s", "scale_o", "descale_o", "descale_dO", "descale_dP", "scale_dQ", "scale_dK", "scale_dV", "scale_dP", ] def quantize_to_fp8(x, q_dtype, compute_dtype, scale): # Explicitly cast the max values to the compute dtype to avoid unnecessary # casting to FP32 during the subsequent math operations." assert q_dtype in ( jnp.float8_e4m3fn, jnp.float8_e5m2, jnp.float8_e4m3fnuz, jnp.float8_e5m2fnuz, ) dtype_max = jnp.finfo(q_dtype).max.astype(compute_dtype) scaled_x = x / jnp.broadcast_to( jnp.asarray(scale, dtype=compute_dtype), x.shape ) clipped_x = jnp.clip(scaled_x, -dtype_max, dtype_max) return clipped_x.astype(q_dtype) def quantize_dequantize_fp8(x, q_dtype, scale, compute_dtype): qx = quantize_to_fp8(x, q_dtype, compute_dtype, scale) out = qx.astype(x.dtype) * jnp.broadcast_to( jnp.asarray(scale, dtype=x.dtype), qx.shape ) return out cast_to_representable = partial( quantize_dequantize_fp8, scale=1, compute_dtype=jnp.bfloat16 ) quantize = partial(quantize_to_fp8, scale=1) def get_large_negative_number(dtype): return 0.7 * jnp.finfo(dtype).min def sdpa_train(query: Array, key: Array, value: Array, grad: Array, bias: Array | None = None, mask: Array | None = None, q_seqlen: Array | None = None, kv_seqlen: Array | None = None, q_offsets: Array | None = None, kv_offsets: Array | None = None, score_mod_args: Tuple[Array, ...] = (), scale: float = 0.5, mask_type: MaskType = MaskType.NO_MASK, is_bnth: bool = False, dropout_rate: float = 0.1, sliding_window_length: int | None = None, score_mod: Callable[[Array], Array] | None = None) -> Array: if mask_type == MaskType.PADDING: if is_bnth: B, _, S, _ = query.shape else: B, S, _, _ = query.shape q_seqlen = kv_seqlen = jnp.full((B,), S // 2, jnp.int32) out, sdpa_vjp = jax.vjp( partial(dot_product_attention, scale=scale, mask_type=mask_type, dropout_rate=dropout_rate, qkv_layout="BNTH" if is_bnth else "BTNH", sliding_window_length=sliding_window_length, score_mod=score_mod), query, key, value, bias, mask, q_seqlen, kv_seqlen, q_offsets, kv_offsets, None, score_mod_args) query_grad, key_grad, value_grad, bias_grad = sdpa_vjp(grad)[:4] if bias is not None: # has dbias return out, (query_grad, key_grad, value_grad, bias_grad) return out, (query_grad, key_grad, value_grad) def sdpa_ref(query: Array, key: Array, value: Array, bias: Array | None = None, mask: Array | None = None, score_mod_args: Tuple[Array, ...] = (), scale: float = 0.5, mask_type: MaskType = MaskType.NO_MASK, is_bnth: bool = False, dropout_rate: float = 0.1, sliding_window_length: int | None = None, score_mod: Callable[[Array], Array] | None = None) -> Array: def get_causal_mask(logits): large_negative_number = get_large_negative_number(logits.dtype) t = logits.shape[-2] col_idx = jax.lax.broadcasted_iota(np.int32, (t, t), 1) row_idx = jax.lax.broadcasted_iota(np.int32, (t, t), 0) mask = (row_idx < col_idx).astype(logits.dtype) * large_negative_number return mask[(*([jnp.newaxis]*(len(logits.shape) - 2)), ...)] def get_padding_mask(logits): S, T = logits.shape[-2:] large_negative_number = get_large_negative_number(logits.dtype) q_padding = (jax.lax.iota(np.int32, S) >= S // 2).reshape((S, 1)) kv_padding = (jax.lax.iota(np.int32, T) >= T // 2).reshape((1, T)) combined_padding = \ (q_padding + kv_padding).astype(logits.dtype) * large_negative_number return jax.lax.broadcast(combined_padding, logits.shape[:-2]) def get_encoded_padding_mask(encoded, is_bnth): dim = 2 if is_bnth else 1 T = encoded.shape[dim] encoded_padding = (jax.lax.iota(np.int32, T) < T // 2).astype(encoded.dtype) return jax.lax.broadcast_in_dim( encoded_padding, encoded.shape, broadcast_dimensions=[dim]) def get_sliding_window_mask(logits, window_length): large_negative_number = get_large_negative_number(logits.dtype) T = logits.shape[-2] col_idx = jax.lax.broadcasted_iota(np.int32, (T, T), 1) row_idx = jax.lax.broadcasted_iota(np.int32, (T, T), 0) mask = jnp.logical_or( row_idx < col_idx, col_idx <= row_idx - window_length).astype(logits.dtype) * large_negative_number return mask[(*([jnp.newaxis]*(len(logits.shape) - 2)), ...)] if is_bnth: B, qN, T, H = query.shape _, kN, _, _ = key.shape logits = jnp.einsum("bhqd,bhkd->bhqk", query, key, preferred_element_type=jnp.float32) else: B, T, qN, H = query.shape _, _, kN, _ = key.shape logits = jnp.einsum("bqhd,bkhd->bhqk", query, key, preferred_element_type=jnp.float32) if scale != 1.0: logits = logits * scale if mask_type == MaskType.CAUSAL: bias = get_causal_mask(logits) elif mask_type == MaskType.PADDING: bias = get_padding_mask(logits) elif sliding_window_length is not None: if sliding_window_length <= 0: raise ValueError( f"Expect sliding_window_length > 0, got {sliding_window_length}.") bias = get_sliding_window_mask(logits, sliding_window_length) if mask is not None: large_negative_number = get_large_negative_number(logits.dtype) mask = jnp.where(mask, 0, large_negative_number) # combine bias and mask if bias is None: bias = mask elif mask is not None: bias = bias.astype(logits.dtype) bias += mask # apply bias to logits if bias is not None: if bias.shape != logits.shape: bias = jnp.broadcast_to(bias, logits.shape) logits = logits + bias.astype(logits.dtype) if score_mod is not None: logits = score_mod(logits, *score_mod_args).astype(jnp.float32) probs = jax.nn.softmax(logits, axis=-1).astype(query.dtype) if dropout_rate > 0.: keep_prob = 1.0 - dropout_rate dropout_rng = jax.random.key(0) keep = jax.random.bernoulli(dropout_rng, keep_prob, probs.shape) probs = jax.lax.select(keep, probs / keep_prob, jnp.zeros_like(probs)) if is_bnth: encoded = jnp.einsum("bhqk,bhkd->bhqd", probs, value, preferred_element_type=jnp.float32) else: encoded = jnp.einsum("bhqk,bkhd->bqhd", probs, value, preferred_element_type=jnp.float32) if mask_type == MaskType.PADDING: # cuDNN padding mask generation will mask out output accordingly # make sure the behavior is the same encoded_mask = get_encoded_padding_mask(encoded, is_bnth) encoded = encoded * encoded_mask return encoded.astype(query.dtype) def sdpa_train_ref(query: Array, key: Array, value: Array, grad: Array, bias: Array | None = None, mask: Array | None = None, scale: float = 0.5, mask_type: MaskType = MaskType.NO_MASK, is_bnth: bool = False, dropout_rate: float = 0.1, sliding_window_length: int | None = None, score_mod: Callable[[Array], Array] | None = None, score_mod_args: Tuple[Array, ...] = ()) -> Array: out_ref, sdpa_vjp_ref = jax.vjp( partial( sdpa_ref, scale=scale, mask_type=mask_type, dropout_rate=dropout_rate, sliding_window_length=sliding_window_length, is_bnth=is_bnth, score_mod=score_mod), query, key, value, bias, mask, score_mod_args) query_grad_ref, key_grad_ref, value_grad_ref, bias_grad_ref = sdpa_vjp_ref(grad)[:4] if bias is not None: return out_ref, (query_grad_ref, key_grad_ref, value_grad_ref, bias_grad_ref) return out_ref, (query_grad_ref, key_grad_ref, value_grad_ref) def sdpa_train_fp8( query: Array, key: Array, value: Array, grad: Array, fp8_metas: dict[Array], scale: float = 0.5, mask_type: MaskType = MaskType.NO_MASK, ): def dot_product_attention_fp8(query, key, value, fp8_metas): f_p = partial( dot_product_attention, scale=scale, mask_type=mask_type, use_fp8=True ) return f_p(query, key, value, fp8_params=fp8_metas) out, sdpa_vjp = jax.vjp( dot_product_attention_fp8, query, key, value, fp8_metas ) grad_amax_s = jnp.ones((1, 1, 1, 1), dtype=jnp.float32) grad_amax_o = jnp.ones((1, 1, 1, 1), dtype=jnp.float32) query_grad, key_grad, value_grad, *_ = sdpa_vjp( (grad, grad_amax_s, grad_amax_o) ) return out[0], (query_grad, key_grad, value_grad) class DotProductAttentionTest(jtu.JaxTestCase): def setUp(self): super().setUp() if not jtu.is_cuda_compute_capability_at_least("8.0"): self.skipTest("Requires at least Ampere arch") if jtu.is_cuda_version_at_least(13, 0): self.skipTest("cuDNN creates no execution plans on CUDA 13.0.") self.enter_context(jtu.ignore_warning( category=DeprecationWarning, message='`with mesh:` context manager')) @jtu.sample_product( batch_size=[4], seq_len=[1024], num_heads=[8], head_dim=[64, 128], use_mask=[False, True], use_bias=[False, True], mask_type=[MaskType.NO_MASK], is_bnth=[False, True], dropout_rate=[0], scale=[0.5], dtype=[jnp.float16, jnp.bfloat16] ) @jtu.run_on_devices("cuda") def test_sdpa(self, batch_size: int, seq_len: int, num_heads: int, head_dim: int, use_mask: bool, use_bias: bool, mask_type: MaskType, is_bnth: bool, dropout_rate: float, scale: float, dtype: jnp.dtype): if len(jax.local_devices()) < 4: self.skipTest("Require at least 4 devices to run sharding tests.") if use_mask and mask_type != MaskType.NO_MASK: self.skipTest("Either pass in mask or generate mask directly in cuDNN.") k1, k2, k3, k4, k5, k6 = jax.random.split(jax.random.key(0), 6) if is_bnth: qkv_shape = (batch_size, num_heads, seq_len, head_dim) else: qkv_shape = (batch_size, seq_len, num_heads, head_dim) query = jax.random.normal(k1, qkv_shape, dtype=dtype) key = jax.random.normal(k2, qkv_shape, dtype=dtype) value = jax.random.normal(k3, qkv_shape, dtype=dtype) grad = jax.random.normal(k4, qkv_shape, dtype=dtype) if use_bias: bias = jax.random.normal( k5, (batch_size, num_heads, seq_len, seq_len), dtype=dtype) else: bias = None if use_mask: mask = jax.random.bernoulli( k6, 0.5, (batch_size, num_heads, seq_len, seq_len)) else: mask = None devices = np.array(jax.local_devices()[:4]) devices = devices.reshape((2, 2)) with Mesh(devices, ("dp", "tp")) as mesh: if is_bnth: qkv_spec = PartitionSpec("dp", "tp", None, None) else: qkv_spec = PartitionSpec("dp", None, "tp", None) qkv_sharding = NamedSharding(mesh, qkv_spec) if bias is not None: bias_spec = PartitionSpec("dp", "tp", None, None) else: bias_spec = PartitionSpec() if mask is not None: mask_spec = PartitionSpec("dp", "tp", None, None) else: mask_spec = PartitionSpec() bias_sharding = NamedSharding(mesh, bias_spec) mask_sharding = NamedSharding(mesh, mask_spec) query = jax.device_put(query, qkv_sharding) key = jax.device_put(key, qkv_sharding) value = jax.device_put(value, qkv_sharding) if bias is not None: bias = jax.device_put(bias, bias_sharding) if mask is not None: mask = jax.device_put(mask, mask_sharding) grad = jax.device_put(grad, qkv_sharding) in_shardings = (qkv_sharding, qkv_sharding, qkv_sharding, qkv_sharding, bias_sharding, mask_sharding) if use_bias: out_shardings = (qkv_sharding, (qkv_sharding, qkv_sharding, qkv_sharding, bias_sharding)) else: out_shardings = (qkv_sharding, (qkv_sharding, qkv_sharding, qkv_sharding)) jitted_sdpa_train = jax.jit( partial( sdpa_train, scale=scale, mask_type=mask_type, dropout_rate=dropout_rate, is_bnth=is_bnth), in_shardings=in_shardings, out_shardings=out_shardings ) jitted_sdpa_train_ref = jax.jit( partial( sdpa_train_ref, scale=scale, mask_type=mask_type, dropout_rate=dropout_rate, is_bnth=is_bnth), in_shardings=in_shardings, out_shardings=out_shardings ) if use_bias: out, (query_grad, key_grad, value_grad, _) = \ jitted_sdpa_train(query, key, value, grad, bias, mask) out_ref, (query_grad_ref, key_grad_ref, value_grad_ref, _) = \ jitted_sdpa_train_ref(query, key, value, grad, bias, mask) else: out, (query_grad, key_grad, value_grad) = \ jitted_sdpa_train(query, key, value, grad, bias, mask) out_ref, (query_grad_ref, key_grad_ref, value_grad_ref) = \ jitted_sdpa_train_ref(query, key, value, grad, bias, mask) self.assertArraysAllClose(out_ref, out, rtol=2e-2, atol=2e-2) self.assertArraysAllClose( query_grad_ref, query_grad, rtol=2e-1, atol=2e-1) self.assertArraysAllClose( key_grad_ref, key_grad, rtol=2e-1, atol=2e-1) self.assertArraysAllClose( value_grad_ref, value_grad, rtol=2e-1, atol=2e-1) @jtu.run_on_devices("cuda") def test_sdpa_inference(self): if jax.device_count() < 4: self.skipTest("Requires more than 4 devices.") k1, k2, k3 = jax.random.split(jax.random.key(0), 3) query = jax.random.normal( k1, (4, 1024, 4, 64), dtype=jnp.bfloat16) key = jax.random.normal( k2, (4, 1024, 4, 64), dtype=jnp.bfloat16) value = jax.random.normal( k3, (4, 1024, 4, 64), dtype=jnp.bfloat16) devices = np.array(jax.local_devices()[:4]) devices = devices.reshape((2, 2)) with Mesh(devices, ("dp", "tp")) as mesh: qkv_spec = PartitionSpec("dp", None, "tp", None) qkv_sharding = NamedSharding(mesh, qkv_spec) in_shardings = ( qkv_sharding, qkv_sharding, qkv_sharding) out_shardings = qkv_sharding query = jax.device_put(query, qkv_sharding) key = jax.device_put(key, qkv_sharding) value = jax.device_put(value, qkv_sharding) jitted_sdpa_inference = jax.jit( partial( dot_product_attention, scale=1.0, mask_type=MaskType.NO_MASK, dropout_rate=0), in_shardings=in_shardings, out_shardings=out_shardings ) jitted_sdpa_inference_ref = jax.jit( partial( sdpa_ref, scale=1.0, mask_type=MaskType.NO_MASK, dropout_rate=0), in_shardings=in_shardings, out_shardings=out_shardings ) out = jitted_sdpa_inference(query, key, value) out_ref = jitted_sdpa_inference_ref(query, key, value) self.assertArraysAllClose(out_ref, out, rtol=2e-2, atol=2e-2) @jtu.run_on_devices("cuda") def test_sdpa_flex_attention_train(self): if len(jax.local_devices()) < 4: self.skipTest("Require at least 4 devices to run sharding tests.") if not jtu.is_cuda_compute_capability_at_least("9.0"): self.skipTest("Requires at least Hopper arch") try: cudnn_version = check_cudnn_version() except RuntimeError as e: self.skipTest(str(e)) if cudnn_version < 91300: self.skipTest("Requires >= cuDNN 9.13.0") k1, k2, k3, k4, k5 = jax.random.split(jax.random.key(0), 5) batch_size, seq_len, num_heads, head_dim = 4, 1024, 4, 64 qkv_shape = (batch_size, seq_len, num_heads, head_dim) query = jax.random.normal(k1, qkv_shape, dtype=jnp.bfloat16) key = jax.random.normal(k2, qkv_shape, dtype=jnp.bfloat16) value = jax.random.normal(k3, qkv_shape, dtype=jnp.bfloat16) grad = jax.random.normal(k4, qkv_shape, dtype=jnp.bfloat16) soft_cap_scalar = jax.random.normal( k5, (batch_size, num_heads, seq_len, seq_len), dtype=jnp.float32) def soft_cap(attn_score, soft_cap_scalar): return soft_cap_scalar * jax.lax.tanh(attn_score / soft_cap_scalar) devices = np.array(jax.local_devices()[:4]) devices = devices.reshape((2, 2)) with Mesh(devices, ("dp", "tp")) as mesh: qkv_spec = PartitionSpec("dp", None, "tp", None) score_mod_spec = PartitionSpec("dp", "tp", None, None) qkv_sharding = NamedSharding(mesh, qkv_spec) score_mod_sharding = NamedSharding(mesh, score_mod_spec) query = jax.device_put(query, qkv_sharding) key = jax.device_put(key, qkv_sharding) value = jax.device_put(value, qkv_sharding) grad = jax.device_put(grad, qkv_sharding) soft_cap_scalar = jax.device_put(soft_cap_scalar, score_mod_sharding) in_shardings = (qkv_sharding, qkv_sharding, qkv_sharding, qkv_sharding, (score_mod_sharding,)) out_shardings = (qkv_sharding, (qkv_sharding, qkv_sharding, qkv_sharding)) jitted_sdpa = jax.jit( lambda q, k, v, g, sm_args: sdpa_train( q, k, v, g, score_mod_args=sm_args, scale=1.0, mask_type=MaskType.NO_MASK, dropout_rate=0, score_mod=soft_cap), in_shardings=in_shardings, out_shardings=out_shardings, ) jitted_sdpa_ref = jax.jit( lambda q, k, v, g, sm_args: sdpa_train_ref( q, k, v, g, score_mod_args=sm_args, scale=1.0, mask_type=MaskType.NO_MASK, dropout_rate=0, score_mod=soft_cap), in_shardings=in_shardings, out_shardings=out_shardings, ) out, (query_grad, key_grad, value_grad) = \ jitted_sdpa(query, key, value, grad, (soft_cap_scalar,)) out_ref, (query_grad_ref, key_grad_ref, value_grad_ref) = \ jitted_sdpa_ref(query, key, value, grad, (soft_cap_scalar,)) self.assertArraysAllClose(out_ref, out, rtol=1e-2, atol=1e-2) self.assertArraysAllClose(query_grad_ref, query_grad, rtol=1e-2, atol=1e-2) self.assertArraysAllClose(key_grad_ref, key_grad, rtol=1e-2, atol=1e-2) self.assertArraysAllClose(value_grad_ref, value_grad, rtol=1e-2, atol=1e-2) @jtu.run_on_devices("cuda") def test_sdpa_var_seq(self): if jax.device_count() < 4: self.skipTest("Requires more than 4 devices.") k1, k2, k3, k4 = jax.random.split(jax.random.key(0), 4) query = jax.random.normal( k1, (4, 1024, 4, 64), dtype=jnp.bfloat16) key = jax.random.normal( k2, (4, 1024, 4, 64), dtype=jnp.bfloat16) value = jax.random.normal( k3, (4, 1024, 4, 64), dtype=jnp.bfloat16) grad = jax.random.normal( k4, (4, 1024, 4, 64), dtype=jnp.bfloat16) jitted_sdpa_train = jax.jit( partial( sdpa_train, scale=1.0, mask_type=MaskType.PADDING, dropout_rate=0), ) jitted_sdpa_train_ref = jax.jit( partial( sdpa_train_ref, scale=1.0, mask_type=MaskType.PADDING, dropout_rate=0), ) out, (query_grad, key_grad, value_grad) = \ jitted_sdpa_train(query, key, value, grad) out_ref, (query_grad_ref, key_grad_ref, value_grad_ref) = \ jitted_sdpa_train_ref(query, key, value, grad) self.assertArraysAllClose(out_ref, out, rtol=2e-2, atol=2e-2) self.assertArraysAllClose(query_grad_ref, query_grad, rtol=2e-1, atol=2e-1) self.assertArraysAllClose(key_grad_ref, key_grad, rtol=2e-1, atol=2e-1) self.assertArraysAllClose(value_grad_ref, value_grad, rtol=2e-1, atol=2e-1) @jtu.sample_product( broadcast_dims=[(), (0,), (1,), (0,1)], ) @jtu.run_on_devices("cuda") def test_sdpa_broadcast_bias_and_dbias(self, broadcast_dims): if jax.device_count() < 4: self.skipTest("Requires more than 4 devices.") if not jtu.is_cuda_compute_capability_at_least("9.0"): self.skipTest("Requires at least Hopper arch") k1, k2, k3, k4, k5 = jax.random.split(jax.random.key(0), 5) query = jax.random.normal( k1, (4, 1024, 4, 64), dtype=jnp.bfloat16) key = jax.random.normal( k2, (4, 1024, 4, 64), dtype=jnp.bfloat16) value = jax.random.normal( k3, (4, 1024, 4, 64), dtype=jnp.bfloat16) grad = jax.random.normal( k4, (4, 1024, 4, 64), dtype=jnp.bfloat16) if broadcast_dims == (): bias = jax.random.normal( k5, (4, 4, 1024, 1024), dtype=jnp.bfloat16) bias_spec = PartitionSpec("dp", "tp", None, None) elif broadcast_dims == (0,): bias = jax.random.normal( k5, (1, 4, 1024, 1024), dtype=jnp.bfloat16) bias_spec = PartitionSpec(None, "tp", None, None) elif broadcast_dims == (1,): bias = jax.random.normal( k5, (4, 1, 1024, 1024), dtype=jnp.bfloat16) bias_spec = PartitionSpec("dp", None, None, None) else: bias = jax.random.normal( k5, (1, 1, 1024, 1024), dtype=jnp.bfloat16) bias_spec = PartitionSpec(None, None, None, None) devices = np.array(jax.local_devices()[:4]) devices = devices.reshape((2, 2)) with Mesh(devices, ("dp", "tp")) as mesh: qkv_spec = PartitionSpec("dp", None, "tp", None) qkv_sharding = NamedSharding(mesh, qkv_spec) bias_sharding = NamedSharding(mesh, bias_spec) in_shardings = (qkv_sharding, qkv_sharding, qkv_sharding, qkv_sharding, bias_sharding) out_shardings = (qkv_sharding, (qkv_sharding, qkv_sharding, qkv_sharding, bias_sharding)) query = jax.device_put(query, qkv_sharding) key = jax.device_put(key, qkv_sharding) value = jax.device_put(value, qkv_sharding) grad = jax.device_put(grad, qkv_sharding) bias = jax.device_put(bias, bias_sharding) jitted_sdpa_train = jax.jit( partial( sdpa_train, scale=1.0, mask_type=MaskType.NO_MASK, dropout_rate=0), in_shardings=in_shardings, out_shardings=out_shardings ) jitted_sdpa_train_ref = jax.jit( partial( sdpa_train_ref, scale=1.0, mask_type=MaskType.NO_MASK, dropout_rate=0), in_shardings=in_shardings, out_shardings=out_shardings ) out, (query_grad, key_grad, value_grad, bias_grad) = \ jitted_sdpa_train(query, key, value, grad, bias) out_ref, (query_grad_ref, key_grad_ref, value_grad_ref, bias_grad_ref) = \ jitted_sdpa_train_ref(query, key, value, grad, bias) self.assertArraysAllClose(out_ref, out, rtol=2e-2, atol=2e-2) self.assertArraysAllClose(query_grad_ref, query_grad, rtol=2e-1, atol=2e-1) self.assertArraysAllClose(key_grad_ref, key_grad, rtol=2e-1, atol=2e-1) self.assertArraysAllClose(value_grad_ref, value_grad, rtol=2e-1, atol=2e-1) self.assertArraysAllClose(bias_grad_ref, bias_grad, rtol=2e-1, atol=2e-1) @jtu.sample_product( batch_size=[1, 16], ) @jtu.run_on_devices("cuda") def test_sdpa_dbias(self, batch_size: int): if jax.device_count() < 4: self.skipTest("Requires more than 4 devices.") # cuDNN only supports dbias when batch size is 1. If the batch size is # greater, dbias is silently set to all zeros. This test verifies this # behavior for both vmap and regular use cases. # TODO: Remove this test once cuDNN adds broader dbias support. dtype = jnp.bfloat16 x_shape = (batch_size, 512, 16, 48) bias_shape = (batch_size, 16, 512, 512) mask_shape = (1, 1, 512) keys = jax.random.split(jax.random.key(0), 2) x = jax.random.normal(keys[0], x_shape, dtype=dtype) bias = jax.random.normal(keys[1], bias_shape, dtype=dtype) mask = jnp.ones(mask_shape, dtype=jnp.bool_) def attn(x, bias, mask): return dot_product_attention(x, x, x, bias, mask) def attn_vjp(x, bias, mask, target_fn): _, f_vjp = jax.vjp(target_fn, x, bias, mask) return f_vjp(x) attn_vmap = jax.vmap(attn, in_axes=(0, 0, None)) attn_ref = jax.jit(partial(attn_vjp, target_fn=attn)) attn_ans = jax.jit(partial(attn_vjp, target_fn=attn_vmap)) _, dbias_ref, _ = attn_ref(x, bias, mask) x = jnp.expand_dims(x, axis=1) bias = jnp.expand_dims(bias, axis=1) _, dbias_ans, _ = attn_ans(x, bias, mask) dbias_ans = jnp.squeeze(dbias_ans, axis=1) self.assertArraysAllClose(dbias_ans, dbias_ref) @jtu.run_on_devices("cuda") def test_sdpa_sliding_window_length(self): if jax.device_count() < 4: self.skipTest("Requires more than 4 devices.") k1, k2, k3, k4 = jax.random.split(jax.random.key(0), 4) query = jax.random.normal( k1, (4, 1024, 4, 64), dtype=jnp.bfloat16) key = jax.random.normal( k2, (4, 1024, 4, 64), dtype=jnp.bfloat16) value = jax.random.normal( k3, (4, 1024, 4, 64), dtype=jnp.bfloat16) grad = jax.random.normal( k4, (4, 1024, 4, 64), dtype=jnp.bfloat16) jitted_sdpa_train = jax.jit( partial( sdpa_train, scale=1.0, mask_type=MaskType.CAUSAL, dropout_rate=0, sliding_window_length=64), ) # for reference implementation # sliding_window_length option itself will setup correct mask jitted_sdpa_train_ref = jax.jit( partial( sdpa_train_ref, scale=1.0, mask_type=MaskType.NO_MASK, dropout_rate=0, sliding_window_length=64), ) out, (query_grad, key_grad, value_grad) = \ jitted_sdpa_train(query, key, value, grad) out_ref, (query_grad_ref, key_grad_ref, value_grad_ref) = \ jitted_sdpa_train_ref(query, key, value, grad) self.assertArraysAllClose(out_ref, out, rtol=2e-2, atol=2e-2) self.assertArraysAllClose(query_grad_ref, query_grad, rtol=2e-1, atol=2e-1) self.assertArraysAllClose(key_grad_ref, key_grad, rtol=2e-1, atol=2e-1) self.assertArraysAllClose(value_grad_ref, value_grad, rtol=2e-1, atol=2e-1) @jtu.sample_product( num_kv_heads=[4, 2], mask_type=[MaskType.NO_MASK, MaskType.PADDING_CAUSAL], ) @jtu.run_on_devices("cuda") def test_sdpa_packed_layout(self, num_kv_heads, mask_type): if jax.device_count() < 4: self.skipTest("Requires more than 4 devices.") if not jtu.is_cuda_compute_capability_at_least("9.0"): self.skipTest("Requires at least Hopper arch") num_heads = 4 head_dim = 64 q_per_kv = num_heads // num_kv_heads k1, k2, k3, k4 = jax.random.split(jax.random.key(0), 4) query = jax.random.normal( k1, (4, 512, num_heads, head_dim), dtype=jnp.bfloat16) key = jax.random.normal( k2, (4, 512, num_kv_heads, head_dim), dtype=jnp.bfloat16) value = jax.random.normal( k3, (4, 512, num_kv_heads, head_dim), dtype=jnp.bfloat16) grad = jax.random.normal( k4, (4, 512, num_heads, head_dim), dtype=jnp.bfloat16) def generate_padding_mask(segment_ids, padding_id, shape, dtype): # segment_ids [B, T] encoded_padding = jnp.where(segment_ids >= padding_id, 0, 1).astype(dtype) return jax.lax.broadcast_in_dim( encoded_padding, shape, broadcast_dimensions=[0, 1]) def generate_segment_mask(segment_ids, dtype): segment_ids_1 = jnp.expand_dims(segment_ids, axis=-1) # segment_ids_1 = jnp.where(segment_ids_1 == 3, 4, segment_ids_1) segment_ids_2 = jnp.expand_dims(segment_ids, axis=1) mask = jnp.not_equal(segment_ids_1, segment_ids_2).astype(dtype) # broadcast to [B, N, T, T] mask = jnp.expand_dims(mask, 1) mask *= get_large_negative_number(dtype) return mask def generate_causal_mask(seq_len, dtype): col_idx = jax.lax.broadcasted_iota(np.int32, (seq_len, seq_len), 1) row_idx = jax.lax.broadcasted_iota(np.int32, (seq_len, seq_len), 0) mask = (row_idx < col_idx).astype(dtype) * get_large_negative_number(dtype) return mask[jnp.newaxis, jnp.newaxis] # starting pos of each segment q_offsets = jnp.asarray([ [0, 170, 340, -1], # 3 segments [0, 150, 340, -1], # 3 segments [0, 190, -1, -1], # 2 segments [0, -1, -1, -1] # 1 segment ], dtype=np.int32) # actual seqlen of each segment without padding q_seqlen = jnp.asarray([ [170, 170, 172], # No padding inside each segment [150, 187, 172], # 3 padding tokens inside second segment [190, 190, -1], # 132 padding tokens inside last segment [400, -1, -1], # 112 padding tokens inside last segment ], dtype=np.int32) # maximum number of segments is id for padding token segment_ids = jnp.asarray([ [0]*170 + [1]*170 + [2]*172, [0]*150 + [1]*187 + [3]*3 + [2]*172, [0]*190 + [1]*190 + [3]*132, [0]*400 + [3]*112, ], dtype=np.int32) kv_offsets = q_offsets.copy() kv_seqlen = q_seqlen.copy() mask = generate_padding_mask(segment_ids, q_seqlen.shape[1], query.shape, query.dtype) kv_mask = generate_padding_mask(segment_ids, q_seqlen.shape[1], key.shape, key.dtype) # The reference always runs with NO_MASK, so the variants differ only here. # Causal within a segment is global causal restricted to that segment. bias = generate_segment_mask(segment_ids, jnp.float32) if mask_type == MaskType.PADDING_CAUSAL: # `minimum`, not `+`: the large negative value is scaled to leave headroom, # so summing two of them overflows to -inf and gives NaN gradients. bias = jnp.minimum(bias, generate_causal_mask(query.shape[1], jnp.float32)) devices = np.array(jax.local_devices()[:4]) devices = devices.reshape((2, 2)) with Mesh(devices, ("dp", "tp")) as mesh: qkv_spec = PartitionSpec("dp", None, "tp", None) qkv_sharding = NamedSharding(mesh, qkv_spec) bias_spec = PartitionSpec("dp", None, None, None) bias_sharding = NamedSharding(mesh, bias_spec) offsets_specs = PartitionSpec("dp", None) offsets_sharding = NamedSharding(mesh, offsets_specs) query = jax.device_put(query, qkv_sharding) key = jax.device_put(key, qkv_sharding) value = jax.device_put(value, qkv_sharding) grad = jax.device_put(grad, qkv_sharding) bias = jax.device_put(bias, bias_sharding) q_offsets = jax.device_put(q_offsets, offsets_sharding) kv_offsets = jax.device_put(kv_offsets, offsets_sharding) q_seqlen = jax.device_put(q_seqlen, offsets_sharding) kv_seqlen = jax.device_put(kv_seqlen, offsets_sharding) jitted_sdpa_train = jax.jit( partial( sdpa_train, scale=0.1, mask_type=mask_type, dropout_rate=0), in_shardings=(qkv_sharding, qkv_sharding, qkv_sharding, qkv_sharding, None, None, offsets_sharding, offsets_sharding, offsets_sharding, offsets_sharding), out_shardings=(qkv_sharding, (qkv_sharding, qkv_sharding, qkv_sharding)) ) jitted_sdpa_train_ref = jax.jit( partial( sdpa_train_ref, scale=0.1, mask_type=MaskType.NO_MASK, dropout_rate=0), in_shardings=(qkv_sharding, qkv_sharding, qkv_sharding, qkv_sharding, bias_sharding), out_shardings=(qkv_sharding, (qkv_sharding, qkv_sharding, qkv_sharding, bias_sharding)) ) query = query * mask key = key * kv_mask value = value * kv_mask grad = grad * mask out, (query_grad, key_grad, value_grad) = \ jitted_sdpa_train(query, key, value, grad, None, None, q_seqlen, kv_seqlen, q_offsets, kv_offsets) # No grouped-query reference, so expand kv: query head i attends to kv # head i // q_per_kv, which is what repeat produces. key_ref = jnp.repeat(key, q_per_kv, axis=2) value_ref = jnp.repeat(value, q_per_kv, axis=2) out_ref, (query_grad_ref, key_grad_ref, value_grad_ref, _) = \ jitted_sdpa_train_ref(query, key_ref, value_ref, grad, bias) def fold_heads(x): b, s, _, h = x.shape return x.reshape(b, s, num_kv_heads, q_per_kv, h).sum(axis=3) key_grad_ref = fold_heads(key_grad_ref) value_grad_ref = fold_heads(value_grad_ref) out = out * mask out_ref = out_ref * mask query_grad = query_grad * mask query_grad_ref = query_grad_ref * mask key_grad = key_grad * kv_mask key_grad_ref = key_grad_ref * kv_mask value_grad = value_grad * kv_mask value_grad_ref = value_grad_ref * kv_mask self.assertArraysAllClose(out_ref, out, rtol=1e-2, atol=1e-2) # Causal masking leaves fewer terms in each sum, so the gradients disagree # more than the output does: the largest gap measured here needs an atol of # 1.1e-2, against 1.5e-3 for the non-causal cases. self.assertArraysAllClose(query_grad_ref, query_grad, rtol=1e-2, atol=2e-2) self.assertArraysAllClose(key_grad_ref, key_grad, rtol=1e-2, atol=2e-2) self.assertArraysAllClose(value_grad_ref, value_grad, rtol=1e-2, atol=2e-2) @jtu.run_on_devices("cuda") @jtu.ignore_warning(category=DeprecationWarning, message='`with mesh:` context manager') def test_sdpa_residual(self): k1, k2, k3, k4, k5 = jax.random.split(jax.random.key(0), 5) query = jax.random.normal( k1, (4, 1024, 4, 64), dtype=jnp.bfloat16) key = jax.random.normal( k2, (4, 1024, 4, 64), dtype=jnp.bfloat16) value = jax.random.normal( k3, (4, 1024, 4, 64), dtype=jnp.bfloat16) grad = jax.random.normal( k4, (4, 1024, 4, 64), dtype=jnp.bfloat16) grad_stat = jax.random.normal( k5, (4, 4, 1024), dtype=jnp.float32) devices = np.array(jax.local_devices()[:2]) with Mesh(devices, ("dp")) as mesh: qkv_spec = PartitionSpec("dp", None, None, None) stat_spec = PartitionSpec("dp", None, None) qkv_sharding = NamedSharding(mesh, qkv_spec) stat_sharding = NamedSharding(mesh, stat_spec) query = jax.device_put(query, qkv_sharding) key = jax.device_put(key, qkv_sharding) value = jax.device_put(value, qkv_sharding) grad = jax.device_put(grad, qkv_sharding) grad_stat = jax.device_put(grad_stat, stat_sharding) jitted_sdpa_inference = jax.jit( partial( dot_product_attention, scale=1.0, mask_type=MaskType.NO_MASK, dropout_rate=0, return_residual=True), in_shardings=(qkv_sharding, qkv_sharding, qkv_sharding), out_shardings=(qkv_sharding, stat_sharding) ) outs = jitted_sdpa_inference(query, key, value) assert len(outs) == 2 def train(query, key, value, grads): outs, grad_fn = jax.vjp(partial( dot_product_attention, scale=1.0, mask_type=MaskType.NO_MASK, dropout_rate=0, return_residual=True), query, key, value) return outs, grad_fn(grads) jitted_sdpa_train = jax.jit(train, in_shardings=(qkv_sharding, qkv_sharding, qkv_sharding, (qkv_sharding, stat_sharding)), out_shardings=((qkv_sharding, stat_sharding), (qkv_sharding, qkv_sharding, qkv_sharding))) outs = jitted_sdpa_train(query, key, value, (grad, grad_stat)) assert len(outs) == 2 @jtu.run_on_devices("cuda") def test_sdpa_mla(self): if jax.device_count() < 4: self.skipTest("Requires more than 4 devices.") try: cudnn_version = check_cudnn_version() except RuntimeError as e: self.skipTest(str(e)) if cudnn_version < 91000: self.skipTest("Requires >= cuDNN 9.10.0") if not jtu.is_cuda_compute_capability_at_least("9.0"): self.skipTest("Requires at least Hopper arch") k1, k2, k3 = jax.random.split(jax.random.key(0), 3) query = jax.random.normal( k1, (4, 1024, 4, 128), dtype=jnp.bfloat16) key = jax.random.normal( k2, (4, 1024, 4, 128), dtype=jnp.bfloat16) value = jax.random.normal( k3, (4, 1024, 4, 64), dtype=jnp.bfloat16) devices = np.array(jax.local_devices()[:4]) devices = devices.reshape((2, 2)) with Mesh(devices, ("dp", "tp")) as mesh: qkv_spec = PartitionSpec("dp", None, "tp", None) qkv_sharding = NamedSharding(mesh, qkv_spec) in_shardings = ( qkv_sharding, qkv_sharding, qkv_sharding) out_shardings = qkv_sharding query = jax.device_put(query, qkv_sharding) key = jax.device_put(key, qkv_sharding) value = jax.device_put(value, qkv_sharding) jitted_sdpa_inference = jax.jit( partial( dot_product_attention, scale=1.0, mask_type=MaskType.NO_MASK, dropout_rate=0), in_shardings=in_shardings, out_shardings=out_shardings ) jitted_sdpa_inference_ref = jax.jit( partial( sdpa_ref, scale=1.0, mask_type=MaskType.NO_MASK, dropout_rate=0), in_shardings=in_shardings, out_shardings=out_shardings ) out = jitted_sdpa_inference(query, key, value) out_ref = jitted_sdpa_inference_ref(query, key, value) self.assertArraysAllClose(out_ref, out, rtol=2e-2, atol=2e-2) @jtu.run_on_devices("cuda") def test_sdpa_grad_sharded_regression_25986(self): """Regression test for https://github.com/jax-ml/jax/issues/25986. Taking the gradient of jnp.sum(dot_product_attention(...)) with batch-sharded inputs causes grad_output to be inferred as replicated (from the sum's backward broadcast), while the query is batch-sharded. Without the fix in the bwd partition function, this triggers an SPMD partitioning error because the cuDNN backward custom-call is lowered with the batch dimension taken from the partitioned query. """ if len(jax.local_devices()) < 2: self.skipTest("Requires at least 2 devices.") k1, k2, k3 = jax.random.split(jax.random.key(0), 3) query = jax.random.normal(k1, (2, 128, 2, 64), dtype=jnp.bfloat16) key = jax.random.normal(k2, (2, 128, 2, 64), dtype=jnp.bfloat16) value = jax.random.normal(k3, (2, 128, 2, 64), dtype=jnp.bfloat16) devices = np.array(jax.local_devices()[:2]) with Mesh(devices, ("dp",)) as mesh: qkv_spec = PartitionSpec("dp", None, None, None) qkv_sharding = NamedSharding(mesh, qkv_spec) query = jax.device_put(query, qkv_sharding) key = jax.device_put(key, qkv_sharding) value = jax.device_put(value, qkv_sharding) # This pattern -- jax.grad of jnp.sum(dot_product_attention(...)) -- # causes the SPMD partitioner to infer grad_output as replicated. @jax.jit def compute_grad(q, k, v): def loss_fn(q, k, v): out = dot_product_attention( q, k, v, scale=0.5, mask_type=MaskType.NO_MASK, dropout_rate=0) return jnp.sum(out) return jax.grad(loss_fn, argnums=(0, 1, 2))(q, k, v) # Should not raise an SPMD partitioning error. dq, dk, dv = compute_grad(query, key, value) # Sanity-check that the output is batch-sharded. self.assertTrue(dq.sharding.is_equivalent_to(qkv_sharding, ndim=4)) # Compute reference gradients without sharding using the pure-JAX # reference implementation. grad = jnp.ones_like(query) _, (dq_ref, dk_ref, dv_ref) = sdpa_train_ref( query, key, value, grad, scale=0.5, mask_type=MaskType.NO_MASK, dropout_rate=0) self.assertArraysAllClose(dq_ref, dq, rtol=2e-1, atol=2e-1) self.assertArraysAllClose(dk_ref, dk, rtol=2e-1, atol=2e-1) self.assertArraysAllClose(dv_ref, dv, rtol=2e-1, atol=2e-1) @jtu.run_on_devices("cuda") def test_sdpa_shard_map_grad_regression_36008(self): """Regression test for https://github.com/jax-ml/jax/issues/36008. Taking the gradient of a cuDNN attention inside a `shard_map` with manual (dp, tp) axes must preserve the varying manual axes (vma) of the inputs in the backward pass. Without propagating vma in the backward primitive's abstract eval, the custom_vjp backward rule produces cotangents whose manual-axis type does not match the primal inputs, raising a ValueError. """ if len(jax.local_devices()) < 4: self.skipTest("Requires at least 4 devices.") devices = np.array(jax.local_devices()[:4]).reshape(2, 2) mesh = Mesh(devices, axis_names=("dp", "tp")) qkv_spec = PartitionSpec("dp", None, "tp", None) @partial( jax.shard_map, mesh=mesh, in_specs=qkv_spec, out_specs=qkv_spec, ) def attn_fn(qkv): q, k, v = jnp.split(qkv, 3, axis=-1) return dot_product_attention( q, k, v, scale=1.0, mask_type=MaskType.CAUSAL, dropout_rate=0) def loss_fn(qkv): return attn_fn(qkv).sum() # (batch=4, seq=128, num_heads=4, 3*head_dim=192) -> head_dim=64. qkv = jax.random.normal( jax.random.key(0), (4, 128, 4, 192), dtype=jnp.bfloat16) qkv = jax.device_put(qkv, NamedSharding(mesh, qkv_spec)) # Should not raise a custom_vjp varying-manual-axes mismatch error. grads = jax.grad(loss_fn)(qkv) self.assertEqual(grads.shape, qkv.shape) if __name__ == "__main__": absltest.main(testLoader=jtu.JaxTestLoader())