diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py index b6e11b8bea..fd6d83835b 100644 --- a/tests/jax/test_distributed_fused_attn.py +++ b/tests/jax/test_distributed_fused_attn.py @@ -399,6 +399,12 @@ def test_softcap_score_mod_with_aux_params_backward( ), ] +DISTRIBUTED_CONTEXT_SELF_ATTN_MAX_LOGIT_CP_MODES = [ + pytest.param(CPStrategy.ALL_GATHER, False, id="AG"), + pytest.param(CPStrategy.RING, False, id="RING-NO_SCAN"), + pytest.param(CPStrategy.RING, True, id="RING-SCAN"), +] + class TestDistributedContextParallelSelfAttn: # TODO(KshitijLakhani): parametrize num_segments_per_seq for all CP tests @@ -419,6 +425,8 @@ def impl_test_context_parallel_attn( window_size=None, stripe_size=None, num_segments_per_seq=None, + return_max_logit=False, + check_forward_output=True, ): if qkv_layout.is_thd(): if not load_balanced and ( @@ -513,9 +521,89 @@ def check_has_backend_for_mask(mask_type): if num_head % kv_groups != 0 or (num_head // kv_groups) % tp_size != 0: pytest.skip(f"Skipping {kv_groups=} not multiple of {data_shape=} or {tp_size=}") - runner.test_backward() + if return_max_logit: + runner.test_forward( + return_max_logit=True, + check_output=check_forward_output, + ) + else: + runner.test_backward() del os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"] + @pytest_parametrize_wrapper( + "device_count,mesh_shape,mesh_axes,mesh_resource", + generate_context_parallel_configs_for_attn(), + ) + @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1]) + @pytest.mark.parametrize("kv_groups", [1, 8]) + @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")]) + @pytest.mark.parametrize( + "qkv_layout, attn_mask_type", + DISTRIBUTED_CONTEXT_SELF_ATTN_LAYOUTS_MASKS, + ) + @pytest.mark.parametrize( + "cp_strategy, use_scan_ring", + DISTRIBUTED_CONTEXT_SELF_ATTN_MAX_LOGIT_CP_MODES, + ) + @pytest.mark.parametrize( + "window_size", + [ + pytest.param((-1, -1), id="NO_SWA"), + pytest.param((20, 0), id="SWA"), + ], + ) + def test_context_parallel_return_max_logit( + self, + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + kv_groups, + dtype, + qkv_layout, + attn_mask_type, + cp_strategy, + window_size, + use_scan_ring, + ): + """Check CP fused attention returns global per-head max_logit.""" + is_thd = qkv_layout.is_thd() + supports_swa = is_thd and ( + cp_strategy == CPStrategy.ALL_GATHER + or (cp_strategy == CPStrategy.RING and not use_scan_ring) + ) + if window_size != (-1, -1) and not supports_swa: + pytest.skip("CP SWA requires THD All-Gather or unrolled THD Ring.") + # TODO: Evaluate cuDNN Max mismatches observed for striped multi-segment THD Ring GQA. + if is_thd and cp_strategy == CPStrategy.RING and kv_groups > 1: + pytest.skip("THD Ring GQA Max mismatches require further evaluation.") + + stripe_size = 64 if is_thd and cp_strategy == CPStrategy.ALL_GATHER else None + if is_thd and cp_strategy == CPStrategy.RING: + stripe_size = 1 + num_segments_per_seq = 5 if is_thd else None + check_forward_output = not (is_thd and cp_strategy == CPStrategy.RING) + self.impl_test_context_parallel_attn( + device_count, + mesh_shape, + mesh_axes, + mesh_resource, + data_shape, + kv_groups, + attn_mask_type, + dtype, + qkv_layout, + True, + cp_strategy, + use_scan_ring=use_scan_ring, + window_size=window_size, + stripe_size=stripe_size, + num_segments_per_seq=num_segments_per_seq, + return_max_logit=True, + check_forward_output=check_forward_output, + ) + @pytest_parametrize_wrapper( "device_count,mesh_shape,mesh_axes,mesh_resource", generate_context_parallel_configs_for_attn(), diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 352ab64a0d..d68e409331 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -64,7 +64,7 @@ def init(): yield -@partial(jax.jit, static_argnums=(6, 7, 8, 9, 11, 12)) +@partial(jax.jit, static_argnums=(6, 7, 8, 9, 11, 12, 13)) def general_dot_product_attention( query: ArrayLike, key: ArrayLike, @@ -79,27 +79,23 @@ def general_dot_product_attention( dropout_rng: ArrayLike, dtype: DTypeLike, score_mod_reference: Optional[Callable[[Array], Array]] = None, + is_max_logit_enabled: bool = False, ) -> Array: """ Similar to flax.linen.dot_product_attention but with GQA support """ query, key, value, bias = promote_dtype(query, key, value, bias, dtype=dtype) dtype = query.dtype - b, s_q, h_q, d = query.shape _, s_kv, h_kv, _ = key.shape assert (h_q % h_kv == 0) and (h_q >= h_kv) num_groups = h_q // h_kv grouped_query = jnp.reshape(query, (b, s_q, h_kv, num_groups, d)) - # logits with shape (b, h_kv, num_groups, s_q, s_kv) logits = scale_factor * jnp.einsum("...qhgd,...khd->...hgqk", grouped_query, key) if bias is not None: - # reshape logits without groups logits = logits.reshape((b, h_kv * num_groups, s_q, s_kv)) - # apply post-scale bias logits = logits + bias - # reshape logits back to original logits = logits.reshape((b, h_kv, num_groups, s_q, s_kv)) if mask is not None: @@ -110,6 +106,8 @@ def general_dot_product_attention( if score_mod_reference is not None: # Kernel tests use NO_MASK; fused_attn rejects mask+score_mod before this reference path. logits = score_mod_reference(logits.astype(jnp.float32)) + if is_max_logit_enabled: + return jnp.max(logits.reshape((b, h_q, s_q, s_kv)), axis=(0, 2, 3)).astype(dtype) match softmax_type: case AttnSoftmaxType.VANILLA_SOFTMAX: @@ -268,7 +266,17 @@ def _split_valid_and_invalid(primitive, reference, pad): return primitive_valid, primitive_invalid, reference_valid, reference_invalid -def jax_dpa(query, key, value, bias, softmax_offset, mask, dropout_rng, **kwargs): +def jax_dpa( + query, + key, + value, + bias, + softmax_offset, + mask, + dropout_rng, + is_max_logit_enabled=False, + **kwargs, +): """ JAX native dot product attention implementation """ @@ -308,6 +316,7 @@ def jax_dpa(query, key, value, bias, softmax_offset, mask, dropout_rng, **kwargs dropout_rng=dropout_rng, dtype=jnp.float32, score_mod_reference=score_mod_reference, + is_max_logit_enabled=is_max_logit_enabled, ) return output.astype(query.dtype) @@ -339,9 +348,13 @@ def customcall_fused_dpa( qkv_args = (query, key, value) case _: raise ValueError(f"Unsupported {qkv_layout=}") - return fused_attn( + result = fused_attn( qkv_args, bias, sequence_descriptor, dropout_rng, softmax_offset=softmax_offset, **kwargs - ).astype(query.dtype) + ) + if isinstance(result, tuple): + output, max_logit = result + return output.astype(query.dtype), max_logit + return result.astype(query.dtype) def test_fused_attn_score_mod_rejects_masks_before_cudnn_frontend(): @@ -936,7 +949,7 @@ def to_dp_shardings(x): self.seq_length_offset_pspec = PartitionSpec(self.mesh_resource.dp_resource, None) self.seq_length_offset_sharding = NamedSharding(self.mesh, self.seq_length_offset_pspec) - def test_forward(self): + def test_forward(self, return_max_logit=False, check_output=True): """ Test forward with JITted primitive and unJITted reference """ @@ -984,6 +997,7 @@ def test_forward(self): "score_mod_bprop": self.score_mod_bprop, "score_mod_tensors": self.score_mod_tensors, "score_mod_bprop_tensors": self.score_mod_bprop_tensors, + "return_max_logit": return_max_logit, } reference_kwargs = {**kwargs, "score_mod_reference": self.score_mod_reference} @@ -1003,31 +1017,37 @@ def test_forward(self): with self.mesh, autocast(mesh_resource=self.mesh_resource): primitive_out = customcall_fused_dpa_jit(*customcall_args) + if return_max_logit: + primitive_out, primitive_max_logit = primitive_out primitive_out = self.cp_inverse_reorder_fn(primitive_out) - reference_out = jax_dpa(*args, **reference_kwargs) + if return_max_logit: + reference_max_logit = jax_dpa(*args, is_max_logit_enabled=True, **reference_kwargs) - if self.is_training and self.dropout_prob > 0.0: - return + if check_output and not (self.is_training and self.dropout_prob > 0.0): + reference_out = jax_dpa(*args, **reference_kwargs) - primitive_valid, primitive_invalid, reference_valid, reference_invalid = ( - _split_valid_and_invalid(primitive_out, reference_out, self.pad_q) - ) + primitive_valid, primitive_invalid, reference_valid, _ = _split_valid_and_invalid( + primitive_out, reference_out, self.pad_q + ) - assert_allclose( - primitive_invalid, - jnp.zeros_like(primitive_invalid), - rtol=self.rtol, - atol=self.atol, - dtype=self.dtype, - ) - assert_allclose( - primitive_valid, - reference_valid, - rtol=self.rtol, - atol=self.atol, - dtype=self.dtype, - ) + assert_allclose( + primitive_invalid, + jnp.zeros_like(primitive_invalid), + rtol=self.rtol, + atol=self.atol, + dtype=self.dtype, + ) + assert_allclose( + primitive_valid, + reference_valid, + rtol=self.rtol, + atol=self.atol, + dtype=self.dtype, + ) + + if return_max_logit: + assert_allclose(primitive_max_logit, reference_max_logit, dtype=self.dtype) if self.coll_count_ref is not None: with self.mesh, autocast(mesh_resource=self.mesh_resource): @@ -1036,7 +1056,7 @@ def test_forward(self): ) assert_equal_collectives(target_hlo, self.coll_count_ref) - def test_backward(self): + def test_backward(self, return_max_logit=False): """ Test value_and_grad with JIT, which includes both forward and backward. @@ -1064,6 +1084,8 @@ def grad_func( if self.attn_mask_type.is_causal(): gradient_multiplier /= 10 output = func(q, k, v, bias, softmax_offset, sequence_descriptor, dropout_rng, **kwargs) + if isinstance(output, tuple): + output, _ = output if cp_reverse_out: output = self.cp_inverse_reorder_fn(output) # Keep only valid result for the gradient @@ -1129,6 +1151,7 @@ def grad_func( "score_mod_bprop": self.score_mod_bprop, "score_mod_tensors": self.score_mod_tensors, "score_mod_bprop_tensors": self.score_mod_bprop_tensors, + "return_max_logit": return_max_logit, } reference_kwargs = {**kwargs, "score_mod_reference": self.score_mod_reference} @@ -1279,6 +1302,123 @@ def check_dqkv(primitive, reference, pad, idx): assert_equal_collectives(target_hlo, self.coll_count_ref) +FUSED_ATTN_MAX_LOGIT_QKV_LAYOUTS = [ + pytest.param( + QKVLayout.BSHD_BSHD_BSHD, + 8, + 8, + id="BSHD_SEPARATE", + ), + pytest.param( + QKVLayout.BS3HD, + 8, + 8, + id="BS3HD", + ), + pytest.param( + QKVLayout.BSHD_BS2HD, + 8, + 4, + id="BSHD_KV_PACKED-GQA", + ), + pytest.param( + QKVLayout.T3HD, + 8, + 8, + id="THD_QKV_PACKED", + ), + pytest.param( + QKVLayout.THD_THD_THD, + 8, + 8, + id="THD_SEPARATE", + ), +] + + +class TestFusedAttnMaxLogit: + """Targeted non-CP max_logit coverage.""" + + @staticmethod + @pytest.mark.parametrize( + "qkv_layout, num_heads_q, num_heads_kv", + FUSED_ATTN_MAX_LOGIT_QKV_LAYOUTS, + ) + @pytest.mark.parametrize( + "attn_bias_type, bias_shape", + [ + pytest.param(AttnBiasType.NO_BIAS, None, id="NO_BIAS"), + pytest.param( + AttnBiasType.POST_SCALE_BIAS, + BiasShape._1HSS, + id="POST_SCALE_BIAS-1HSS", + ), + ], + ) + @pytest.mark.parametrize( + "attn_mask_type", + [ + pytest.param(AttnMaskType.NO_MASK, id="NO_MASK"), + pytest.param(AttnMaskType.PADDING_MASK, id="PADDING_MASK"), + pytest.param(AttnMaskType.CAUSAL_MASK, id="CAUSAL_MASK"), + pytest.param(AttnMaskType.PADDING_CAUSAL_MASK, id="PADDING_CAUSAL_MASK"), + ], + ) + def test_forward( + qkv_layout, + num_heads_q, + num_heads_kv, + attn_bias_type, + bias_shape, + attn_mask_type, + ): + """Check non-CP JAX fused attention can expose framework-compatible max_logit.""" + runner = FusedAttnRunner( + batch_size=2, + max_seqlen_q=128, + max_seqlen_kv=128, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + head_dim_qk=64, + head_dim_v=64, + attn_bias_type=attn_bias_type, + attn_mask_type=attn_mask_type, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_prob=0.0, + dtype=jnp.bfloat16, + is_training=True, + qkv_layout=qkv_layout, + bias_shape=bias_shape, + window_size=None, + seq_desc_format=SeqDescFormat.Seqlens, + ) + runner.test_forward(return_max_logit=True) + + @staticmethod + def test_backward(): + """Ensure aux-return cotangents do not break the fused attention backward path.""" + runner = FusedAttnRunner( + batch_size=2, + max_seqlen_q=128, + max_seqlen_kv=128, + num_heads_q=8, + num_heads_kv=8, + head_dim_qk=64, + head_dim_v=64, + attn_bias_type=AttnBiasType.NO_BIAS, + attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK, + softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX, + dropout_prob=0.0, + dtype=jnp.bfloat16, + is_training=True, + qkv_layout=QKVLayout.BSHD_BSHD_BSHD, + bias_shape=None, + window_size=None, + seq_desc_format=SeqDescFormat.Seqlens, + ) + runner.test_backward(return_max_logit=True) + + def _get_swa_window_size_for_test(s_kv: int, attn_mask_type: AttnMaskType) -> Tuple[int, int]: """Pick a sliding-window size for SWA tests, gated on cuDNN version. diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py index b1f165f491..7d1137b21f 100644 --- a/tests/jax/test_fused_attn_score_mod.py +++ b/tests/jax/test_fused_attn_score_mod.py @@ -426,6 +426,7 @@ def fake_fused_attn( score_mod_bprop=None, score_mod_tensors=None, score_mod_bprop_tensors=None, + return_max_logit=False, ): captured.update( qkv=qkv, diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index ecca4a3871..6d7c823e12 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -339,6 +339,7 @@ def is_fused_attn_kernel_available( head_dim_qk, head_dim_v, window_size: Optional[Tuple[int, int]] = None, + return_max_logit: bool = False, ): """ To check whether the fused attention kernel is supported @@ -362,6 +363,7 @@ def make_helper(attn_mask_type): head_dim_qk, head_dim_v, window_size_tuple, + return_max_logit, ) return make_helper(attn_mask_type).is_fused_attn_kernel_available() @@ -1053,6 +1055,7 @@ def _legacy_fused_attn( context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, + return_max_logit: bool = False, ): """ Perform non-THD (non-packed) cuDNN fused attention. @@ -1084,8 +1087,15 @@ def _legacy_fused_attn( context_parallel_causal_load_balanced (bool): Indicates the sequences are ordered for causal mask load balancing when running context parallelism. context_parallel_axis (str): The name of the context parallel axis. + softmax_offset (Optional[jnp.ndarray]): An optional learnable softmax offset tensor with shape + [1, num_heads, 1, 1]. Used when softmax_type is AttnSoftmaxType.LEARNABLE_SOFTMAX. + return_max_logit (bool): If True, also return per-head maximum attention logits + with shape ``[h]``. Returns: - (jnp.ndarray): The output tensor from the fused attention. + jnp.ndarray: + Attention output when ``return_max_logit`` is False. + tuple[jnp.ndarray, jnp.ndarray]: + ``(output, max_logit)`` when ``return_max_logit`` is True. """ assert ( not qkv_layout.is_thd() @@ -1139,6 +1149,7 @@ def _legacy_fused_attn( context_parallel_strategy=context_parallel_strategy, context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, + return_max_logit=return_max_logit, ) return output @@ -1164,6 +1175,7 @@ def fused_attn_thd( context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", softmax_offset: Optional[jnp.ndarray] = None, + return_max_logit: bool = False, ): """ Deprecated THD fused attn, please use fusd_attn with SequenceDescriptor @@ -1218,12 +1230,16 @@ def fused_attn_thd( context_parallel_strategy=context_parallel_strategy, context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, + return_max_logit=return_max_logit, ) return output -@partial(jax.custom_vjp, nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18)) +@partial( + jax.custom_vjp, + nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19), +) def _fused_attn( qkv: Tuple[jnp.ndarray, ...], bias: Optional[jnp.ndarray], @@ -1244,6 +1260,7 @@ def _fused_attn( context_parallel_axis: str, context_checkpoint_name: str = "context", stripe_size: int | None = None, + return_max_logit: bool = False, ): output, _ = _fused_attn_fwd_rule( qkv, @@ -1265,6 +1282,7 @@ def _fused_attn( context_parallel_axis, context_checkpoint_name=context_checkpoint_name, stripe_size=stripe_size, + return_max_logit=return_max_logit, ) return output @@ -1289,8 +1307,9 @@ def _fused_attn_fwd_rule( context_parallel_axis, context_checkpoint_name, stripe_size, + return_max_logit, ): - output, softmax_aux, rng_state = tex.fused_attn_fwd( + output, softmax_aux, rng_state, max_logit = tex.fused_attn_fwd( qkv, bias, softmax_offset, @@ -1309,11 +1328,14 @@ def _fused_attn_fwd_rule( context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, stripe_size=stripe_size, + return_max_logit=return_max_logit, ) output = checkpoint_name(output, context_checkpoint_name) softmax_aux = checkpoint_name(softmax_aux, context_checkpoint_name) rng_state = checkpoint_name(rng_state, context_checkpoint_name) - return output, ( + max_logit = checkpoint_name(max_logit, context_checkpoint_name) + attn_output = (output, max_logit) if return_max_logit else output + return attn_output, ( qkv, bias, sequence_descriptor, @@ -1339,10 +1361,13 @@ def _fused_attn_bwd_rule( context_parallel_axis, context_checkpoint_name, stripe_size, + return_max_logit, ctx, dz, ): del context_checkpoint_name + if return_max_logit: + dz, _ = dz ( qkv, bias, @@ -1468,6 +1493,7 @@ def fused_attn( score_mod_bprop: Optional[Callable] = None, score_mod_tensors: Optional[Mapping[str, Any]] = None, score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None, + return_max_logit: bool = False, ): """ Perform cuDNN fused attention. @@ -1524,8 +1550,13 @@ def fused_attn( non-differentiable auxiliary inputs. score_mod_bprop_tensors (Optional[Mapping[str, Any]]): Additional tensors or Python/NumPy scalars made available to `score_mod_bprop`. + return_max_logit (bool): If True, also return per-head maximum attention logits + with shape ``[h]``. Returns: - (jnp.ndarray): The output tensor from the fused attention. + jnp.ndarray: + Attention output when ``return_max_logit`` is False. + tuple[jnp.ndarray, jnp.ndarray]: + ``(output, max_logit)`` when ``return_max_logit`` is True. Examples (non-THD, also known as non-packed): >>> # q_segment_ids = [[1, 1, 1, 0], [1, 1, 0, 0]], 0 means padded tokens @@ -1569,6 +1600,8 @@ def fused_attn( if score_mod_only_args: raise ValueError(f"{', '.join(score_mod_only_args)} require score_mod to be provided.") else: + if return_max_logit: + raise ValueError("return_max_logit is not supported with score_mod fused_attn.") tex.validate_fused_attn_score_mod( qkv, bias, @@ -1628,6 +1661,7 @@ def fused_attn( context_parallel_causal_load_balanced=context_parallel_causal_load_balanced, context_parallel_axis=context_parallel_axis, softmax_offset=softmax_offset, + return_max_logit=return_max_logit, ) if max_segments_per_seq > 1 and not qkv_layout.is_thd(): warnings.warn( @@ -1658,5 +1692,6 @@ def fused_attn( context_parallel_axis=context_parallel_axis, context_checkpoint_name=context_checkpoint_name, stripe_size=stripe_size, + return_max_logit=return_max_logit, ) return output diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 489bfde997..6ea54195ab 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -75,6 +75,7 @@ "cp_axis", "cp_striped_window_size", "stripe_size", + "return_max_logit", ], ) @dataclass(frozen=True) @@ -99,6 +100,7 @@ class _FusedAttnConfig: stripe_size: ( int | None ) # Only for CP + Striped. For Ring P2P, stripe_size=1 only.For AG, stripe_size>=1. + return_max_logit: bool = False @dataclass(frozen=True) @@ -122,6 +124,7 @@ class FusedAttnHelper: head_dim_qk: int head_dim_v: int window_size: Tuple[int, int] + return_max_logit: bool = False def is_fused_attn_kernel_available(self): """Check if there is available fused attention kernel""" @@ -146,6 +149,7 @@ def get_fused_attn_backend(self): self.head_dim_v, self.window_size[0], self.window_size[1], + self.return_max_logit, not self.is_non_deterministic_allowed(), ) @@ -351,6 +355,7 @@ def abstract( q_head_dim, v_head_dim, config.window_size, + config.return_max_logit, ).get_fused_attn_backend() if backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: @@ -371,6 +376,17 @@ def abstract( else: raise ValueError(f"Unsupported {backend=}") softmax_aux_aval = q_aval.update(shape=softmax_shape, dtype=softmax_dtype) + if config.return_max_logit: + # cuDNN Max is row-wise over S_kv. Dense and SM120 THD use + # [..., H, S_q, 1]; cuDNN >= 9.6 non-SM120 THD uses [..., S_q, H, 1]. + # Both raw layouts are reduced to the public per-head [H] result below. + if FusedAttnFwdPrimitive._uses_thd_ragged_max_tensor(config): + max_tensor_shape = (*batch_shape, q_max_seqlen, attn_heads, 1) + else: + max_tensor_shape = (*batch_shape, attn_heads, q_max_seqlen, 1) + else: + max_tensor_shape = (0,) + max_tensor_aval = q_aval.update(shape=max_tensor_shape, dtype=softmax_dtype) # JAX does not enable 64-bit int by default so we get XLA to allocate x8 memory with # 32-bit unsigned int to get the buffer size we need in the C++ kernel @@ -417,6 +433,7 @@ def abstract( config.max_segments_per_seq, config.window_size[0], config.window_size[1], + config.return_max_logit, bottom_right_diagonal, ) wkspace_aval = q_aval.update( @@ -437,17 +454,19 @@ def abstract( f" {softmax_offset_aval.shape}" ) - return out_aval, softmax_aux_aval, rng_state_aval, wkspace_aval + return out_aval, softmax_aux_aval, max_tensor_aval, rng_state_aval, wkspace_aval @staticmethod def outer_abstract(*args, **kwargs): """ Fused attention fwd outer primitive abstract """ - out_aval, softmax_aux_aval, rng_state_aval, _ = FusedAttnFwdPrimitive.abstract( + out_aval, softmax_aux_aval, _, rng_state_aval, _ = FusedAttnFwdPrimitive.abstract( *args, **kwargs ) - return out_aval, softmax_aux_aval, rng_state_aval + max_logit_shape = (out_aval.shape[-2],) if kwargs["config"].return_max_logit else (0,) + max_logit_aval = out_aval.update(shape=max_logit_shape, dtype=out_aval.dtype) + return out_aval, softmax_aux_aval, rng_state_aval, max_logit_aval @staticmethod def lowering( @@ -531,6 +550,7 @@ def lowering( mask_type=int(config.attn_mask_type.value), qkv_layout=int(config.qkv_layout.value), is_training=config.is_training, + return_max_logit=config.return_max_logit, deterministic=not FusedAttnHelper.is_non_deterministic_allowed(), window_size_left=window_size_left, window_size_right=window_size_right, @@ -574,6 +594,9 @@ def impl( config.max_segments_per_seq, ) ) + raw_q_seqlen = q_seqlen + raw_q_seq_offsets = q_seq_offsets + if config.qkv_layout.is_thd(): def _fix_len_take(x, condition, fill_value=-1): @@ -630,7 +653,7 @@ def convert_to_2d(offsets, batch, max_seqlen): q_cu_seqlen = generate_cu_seqlen(q_seqlen.flatten()) kv_cu_seqlen = generate_cu_seqlen(kv_seqlen.flatten()) - output, softmax_aux, rng_state, _ = FusedAttnFwdPrimitive.inner_primitive.bind( + output, softmax_aux, max_tensor, rng_state, _ = FusedAttnFwdPrimitive.inner_primitive.bind( q, k, v, @@ -647,7 +670,91 @@ def convert_to_2d(offsets, batch, max_seqlen): _kv_segment_pos, config=config, ) - return output, softmax_aux, rng_state + # Reduce cuDNN's raw Max tensor to TE's public per-head [H] max_logit. + max_logit = FusedAttnFwdPrimitive._reduce_max_logit( + max_tensor, output, raw_q_seqlen, raw_q_seq_offsets, config + ) + return output, softmax_aux, rng_state, max_logit + + @staticmethod + def _reduce_max_logit(max_tensor, output, q_seqlen, q_seq_offsets, config): + """Reduce cuDNN's row-wise Max tensor to the public per-head max_logit. + + Dense and SM120 THD use ``[..., H, S_q, 1]``; cuDNN >= 9.6 non-SM120 + THD uses ``[..., S_q, H, 1]``. A rank-3 THD result is ``[T_q, H, 1]``. + All layouts reduce to ``[H]``. Static THD buffers can contain invalid query + rows, so those rows are masked before reduction. + """ + if not config.return_max_logit: + return jnp.zeros((0,), dtype=output.dtype) + + uses_thd_ragged_max_tensor = FusedAttnFwdPrimitive._uses_thd_ragged_max_tensor(config) + if config.qkv_layout.is_thd() and max_tensor.ndim == 4: + # Dense BSHD Max rows are expected to be masked by cuDNN before TE reduces them. + # THD Max can include static holes/unwritten rows, so mask valid query rows here. + q_seqlen = jnp.where(q_seqlen > 0, q_seqlen, 0) + q_seq_offsets = jnp.where(q_seq_offsets >= 0, q_seq_offsets, -1) + num_segments = min(q_seqlen.shape[-1], q_seq_offsets.shape[-1]) + q_seqlen = q_seqlen[..., :num_segments] + q_seq_offsets = q_seq_offsets[..., :num_segments] + token_idx = jnp.arange(output.shape[-3], dtype=q_seq_offsets.dtype) + valid = jnp.any( + (q_seq_offsets[..., None] >= 0) + & (token_idx >= q_seq_offsets[..., None]) + & (token_idx < (q_seq_offsets[..., None] + q_seqlen[..., None])), + axis=-2, + ) + if uses_thd_ragged_max_tensor: + max_tensor = jnp.where(valid[:, :, None, None], max_tensor, -jnp.inf) + else: + max_tensor = jnp.where(valid[:, None, :, None], max_tensor, -jnp.inf) + + if max_tensor.ndim == 3: + amax_dims = (0, 2) + elif uses_thd_ragged_max_tensor: + amax_dims = (0, 1, 3) + else: + amax_dims = (0, 2, 3) + return jnp.max(max_tensor, axis=amax_dims).astype(output.dtype) + + @staticmethod + def _uses_thd_ragged_max_tensor(config): + """Return whether cuDNN writes THD Max with BSH-like ragged-stats layout.""" + return ( + config.qkv_layout.is_thd() + and get_cudnn_version() >= (9, 6, 0) + and 120 not in get_all_device_compute_capability() + ) + + @staticmethod + def _empty_or_neg_inf_max_logit(head, dtype, config): + """Return the neutral value for per-head max_logit accumulation.""" + if config.return_max_logit: + return jnp.full((head,), -jnp.inf, dtype=dtype) + return jnp.zeros((0,), dtype=dtype) + + @staticmethod + def _max_logit_reduce_axes(mesh, max_logit_sharding): + """Return mesh axes to reduce while preserving max_logit's head sharding.""" + # max_logit is [H], so axes that shard H (typically TP) are preserved. + # Axes for collapsed dimensions such as batch/sequence (DP/CP) must pmax. + head_axes = set() + for axis in max_logit_sharding.spec: + if axis is None: + continue + if isinstance(axis, tuple): + head_axes.update(axis) + else: + head_axes.add(axis) + return tuple(axis for axis in mesh.axis_names if axis not in head_axes) + + @staticmethod + def _reduce_max_logit_across_mesh(max_logit, mesh, reduce_axes, config): + """Reduce max_logit across mesh axes absent from the [H] result.""" + if config.return_max_logit: + for axis in reduce_axes: + max_logit = lax_paral_op(max_logit, lax.pmax, axis, mesh=mesh) + return max_logit @staticmethod def batcher(batched_args, batch_dims, *, config): @@ -659,7 +766,8 @@ def batcher(batched_args, batch_dims, *, config): q_bdim, _, _, _, _, seed_bdim, *_ = batch_dims # Pass through; segment_ids/segment_pos may have different batch dims (e.g. vmapped ids, # replicated pos). get_seqlens_and_offsets() in attention.py handles conversion without expanding. - out_bdims = q_bdim, q_bdim, seed_bdim + max_logit_bdim = q_bdim if config.return_max_logit else None + out_bdims = q_bdim, q_bdim, seed_bdim, max_logit_bdim return ( FusedAttnFwdPrimitive.outer_primitive.bind(*batched_args, config=config), out_bdims, @@ -713,12 +821,16 @@ def infer_sharding_from_operands(config, mesh, arg_infos, result_infos): raise ValueError(f"Unsupported {config.qkv_layout=}") rng_state_sharding = NamedSharding(mesh, PartitionSpec(get_all_mesh_axes(), None)) - return (out_sharding, softmax_aux_sharding, rng_state_sharding) + max_logit_sharding = NamedSharding( + mesh, PartitionSpec(q_spec[-2] if config.return_max_logit else None) + ) + return (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding) @staticmethod def partition(config, mesh, arg_infos, result_infos): out_sharding = result_infos[0].sharding softmax_aux_sharding = result_infos[1].sharding + max_logit_sharding = result_infos[3].sharding rng_state_sharding = seed_sharding = NamedSharding( mesh, PartitionSpec(get_all_mesh_axes(), None) ) @@ -727,8 +839,23 @@ def partition(config, mesh, arg_infos, result_infos): arg_shardings[-1] = arg_shardings[-3] arg_shardings[-2] = arg_shardings[-4] arg_shardings = tuple(arg_shardings) - out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) - impl = partial(FusedAttnFwdPrimitive.impl, config=config) + out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding) + max_logit_reduce_axes = ( + FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding) + if config.return_max_logit + else () + ) + + def impl(*args): + output, softmax_aux, rng_state, max_logit = FusedAttnFwdPrimitive.impl( + *args, config=config + ) + # Globalize the rank-local [H] max across DP/CP while preserving TP head sharding. + max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh( + max_logit, mesh, max_logit_reduce_axes, config + ) + return output, softmax_aux, rng_state, max_logit + return mesh, impl, out_shardings, arg_shardings @staticmethod @@ -756,8 +883,10 @@ def shardy_sharding_rule(config, mesh, value_types, result_types): else: softmax_aux_sharding = ("…0", "head", "seqlen", "i") + max_logit_sharding = ("head",) if config.return_max_logit else ("max_logit",) return SdyShardingRule( - tuple(input_spec), (out_sharding, softmax_aux_sharding, rng_sharding) + tuple(input_spec), + (out_sharding, softmax_aux_sharding, rng_sharding, max_logit_sharding), ) @@ -1426,6 +1555,7 @@ def get_step_config(self) -> _FusedAttnConfig: cp_axis=self.config.cp_axis, cp_striped_window_size=None, stripe_size=self.config.stripe_size, + return_max_logit=self.config.return_max_logit, ) def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig: @@ -1446,6 +1576,7 @@ def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig: cp_axis=self.config.cp_axis, cp_striped_window_size=None, stripe_size=self.config.stripe_size, + return_max_logit=self.config.return_max_logit, ) def all_gather_kv(self, k, v): @@ -1816,13 +1947,19 @@ def partition(config, mesh, arg_infos, result_infos): out_sharding = result_infos[0].sharding softmax_aux_sharding = result_infos[1].sharding + max_logit_sharding = result_infos[3].sharding rng_state_sharding = seed_sharding = NamedSharding( mesh, PartitionSpec(get_all_mesh_axes(), None) ) arg_shardings = [arg_i.sharding for arg_i in arg_infos] arg_shardings[5] = seed_sharding arg_shardings = tuple(arg_shardings) - out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) + out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding) + max_logit_reduce_axes = ( + FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding) + if config.return_max_logit + else () + ) def impl( q, @@ -1870,7 +2007,8 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed): q_seqlen_for_step = q_seqlen / (cp_size * 2) num_kv_chunks = kv_max_seqlen // kv_seqlens_for_rank[sub_idx] kv_seqlen_for_step = (kv_seqlen / (cp_size * 2)) * num_kv_chunks - output, softmax_aux, rng_state = FusedAttnFwdPrimitive.impl( + # max_logit returned here is already reduced to shape [H] + output, softmax_aux, rng_state, max_logit = FusedAttnFwdPrimitive.impl( q_split[sub_idx], k_unmasked, v_unmasked, @@ -1887,13 +2025,15 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed): _kv_segment_pos, config=helper.get_step_config(), ) - results.append((output, softmax_aux, rng_state)) + results.append((output, softmax_aux, rng_state, max_logit)) output = jnp.concatenate((results[0][0], results[1][0]), axis=1) softmax_aux = jnp.concatenate((results[0][1], results[1][1]), axis=2) rng_state = results[1][2] # Use the final RNG state + # Rank-local [H] max across both local dual-chunk query pieces. + max_logit = jnp.maximum(results[0][3], results[1][3]) - return output, softmax_aux, rng_state + return output, softmax_aux, rng_state, max_logit k_ag, v_ag = helper.all_gather_kv(k, v) @@ -1904,7 +2044,12 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed): for idx in range(cp_size) ] - return lax.switch(cp_rank, functions) + output, softmax_aux, rng_state, max_logit = lax.switch(cp_rank, functions) + # Globalize the rank-local [H] max across DP/CP while preserving TP head sharding. + max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh( + max_logit, mesh, max_logit_reduce_axes, config + ) + return output, softmax_aux, rng_state, max_logit return mesh, impl, out_shardings, arg_shardings @@ -2109,13 +2254,19 @@ def partition(config, mesh, arg_infos, result_infos): out_sharding = result_infos[0].sharding softmax_aux_sharding = result_infos[1].sharding + max_logit_sharding = result_infos[3].sharding rng_state_sharding = seed_sharding = NamedSharding( mesh, PartitionSpec(get_all_mesh_axes(), None) ) arg_shardings = [arg_i.sharding for arg_i in arg_infos] arg_shardings[5] = seed_sharding arg_shardings = tuple(arg_shardings) - out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) + out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding) + max_logit_reduce_axes = ( + FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding) + if config.return_max_logit + else () + ) def impl( q, @@ -2179,7 +2330,7 @@ def _cross_attn( max_segments_per_seq=adjusted_max_segments_per_seq, ) - output, softmax_aux, rng_state = FusedAttnFwdPrimitive.impl( + output, softmax_aux, rng_state, max_logit = FusedAttnFwdPrimitive.impl( q, # sharded for rank k, # ag v, # ag @@ -2198,7 +2349,7 @@ def _cross_attn( max_seqlen=kv_max_seqlen, cp_size=cp_size ), ) - return output, softmax_aux, rng_state + return output, softmax_aux, rng_state, max_logit # AG the k, v, kv_segment_ids and kv_segment_pos k_ag, v_ag = helper.all_gather_kv(k, v) @@ -2219,7 +2370,12 @@ def _cross_attn( ) for _ in range(cp_size) ] - return lax.switch(cp_rank, functions) + output, softmax_aux, rng_state, max_logit = lax.switch(cp_rank, functions) + # Globalize the rank-local [H] max across DP/CP while preserving TP head sharding. + max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh( + max_logit, mesh, max_logit_reduce_axes, config + ) + return output, softmax_aux, rng_state, max_logit return mesh, impl, out_shardings, arg_shardings @@ -2492,6 +2648,7 @@ def get_step_config(self, attn_mask_type) -> _FusedAttnConfig: cp_axis=self.config.cp_axis, cp_striped_window_size=None, stripe_size=self.config.stripe_size, + return_max_logit=self.config.return_max_logit, ) def stack_kv(self, k, v): @@ -2559,6 +2716,7 @@ def partition(config, mesh, arg_infos, result_infos): out_sharding = result_infos[0].sharding softmax_aux_sharding = result_infos[1].sharding + max_logit_sharding = result_infos[3].sharding rng_state_sharding = seed_sharding = NamedSharding( mesh, PartitionSpec(get_all_mesh_axes(), None) ) @@ -2568,7 +2726,12 @@ def partition(config, mesh, arg_infos, result_infos): arg_shardings[-1] = arg_shardings[-3] arg_shardings[-2] = arg_shardings[-4] arg_shardings = tuple(arg_shardings) - out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) + out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding) + max_logit_reduce_axes = ( + FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding) + if config.return_max_logit + else () + ) def ring_attn_fwd_impl( q, @@ -2606,9 +2769,10 @@ def ring_attn_fwd_impl( # support dropout currently. rng_state_shape = (seed.shape[0], *result_infos[2].shape[1:]) rng_state = jnp.zeros(rng_state_shape).astype(result_infos[2].dtype) + max_logit = FusedAttnFwdPrimitive._empty_or_neg_inf_max_logit(head, q.dtype, config) def scan_kv_block(idx, carry): - kv, output, softmax_aux = carry + kv, output, softmax_aux, max_logit = carry # Send KV block to next step so we can overlap compute. kv_next = helper.permute_kv(kv, cp_perm) @@ -2616,24 +2780,26 @@ def scan_kv_block(idx, carry): def mask_compute(attn_mask_type): q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx) kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx) - output_per_step, softmax_aux_per_step, _ = FusedAttnFwdPrimitive.impl( - q, - kv, - _not_used, - bias, - _softmax_offset, - seed, - q_seqlen_per_step, - kv_seqlen_per_step, - q_seq_offsets, - k_seq_offsets, - _q_segment_ids, - _kv_segment_ids, - _q_segment_pos, - _kv_segment_pos, - config=helper.get_step_config(attn_mask_type), + output_per_step, softmax_aux_per_step, _, max_logit_per_step = ( + FusedAttnFwdPrimitive.impl( + q, + kv, + _not_used, + bias, + _softmax_offset, + seed, + q_seqlen_per_step, + kv_seqlen_per_step, + q_seq_offsets, + k_seq_offsets, + _q_segment_ids, + _kv_segment_ids, + _q_segment_pos, + _kv_segment_pos, + config=helper.get_step_config(attn_mask_type), + ) ) - return output_per_step, softmax_aux_per_step + return output_per_step, softmax_aux_per_step, max_logit_per_step causal_mask_compute = partial(mask_compute, AttnMaskType.CAUSAL_MASK) no_mask_compute = partial(mask_compute, AttnMaskType.NO_MASK) @@ -2642,45 +2808,49 @@ def half_kv_no_mask_compute(): q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx) kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx) // 2 kv_part = lax.slice_in_dim(kv, 0, kv.shape[1] // 2, axis=1) - output_per_step, softmax_aux_per_step, _ = FusedAttnFwdPrimitive.impl( - q, - kv_part, - _not_used, - bias, - _softmax_offset, - seed, - q_seqlen_per_step, - kv_seqlen_per_step, - q_seq_offsets, - k_seq_offsets, - _q_segment_ids, - _kv_segment_ids, - _q_segment_pos, - _kv_segment_pos, - config=helper.get_step_config(AttnMaskType.NO_MASK), + output_per_step, softmax_aux_per_step, _, max_logit_per_step = ( + FusedAttnFwdPrimitive.impl( + q, + kv_part, + _not_used, + bias, + _softmax_offset, + seed, + q_seqlen_per_step, + kv_seqlen_per_step, + q_seq_offsets, + k_seq_offsets, + _q_segment_ids, + _kv_segment_ids, + _q_segment_pos, + _kv_segment_pos, + config=helper.get_step_config(AttnMaskType.NO_MASK), + ) ) - return output_per_step, softmax_aux_per_step + return output_per_step, softmax_aux_per_step, max_logit_per_step def half_q_no_mask_compute(): q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx) // 2 kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx) q_part = lax.slice_in_dim(q, q_max_seqlen // 2, q_max_seqlen, axis=1) - output_per_step, softmax_aux_per_step, _ = FusedAttnFwdPrimitive.impl( - q_part, - kv, - _not_used, - bias, - _softmax_offset, - seed, - q_seqlen_per_step, - kv_seqlen_per_step, - q_seq_offsets, - k_seq_offsets, - _q_segment_ids, - _kv_segment_ids, - _q_segment_pos, - _kv_segment_pos, - config=helper.get_step_config(AttnMaskType.NO_MASK), + output_per_step, softmax_aux_per_step, _, max_logit_per_step = ( + FusedAttnFwdPrimitive.impl( + q_part, + kv, + _not_used, + bias, + _softmax_offset, + seed, + q_seqlen_per_step, + kv_seqlen_per_step, + q_seq_offsets, + k_seq_offsets, + _q_segment_ids, + _kv_segment_ids, + _q_segment_pos, + _kv_segment_pos, + config=helper.get_step_config(AttnMaskType.NO_MASK), + ) ) output_per_step = jnp.concat([jnp.zeros_like(q_part), output_per_step], axis=1) softmax_aux_per_step = jnp.concat( @@ -2690,14 +2860,17 @@ def half_q_no_mask_compute(): ], axis=2, ) - return output_per_step, softmax_aux_per_step + return output_per_step, softmax_aux_per_step, max_logit_per_step def skip_compute(): output_per_step = jnp.zeros_like(q) softmax_aux_per_step = jnp.full( (batch, head, q.shape[1], 1), -jnp.inf, dtype=jnp.float32 ) - return output_per_step, softmax_aux_per_step + max_logit_per_step = FusedAttnFwdPrimitive._empty_or_neg_inf_max_logit( + head, q.dtype, config + ) + return output_per_step, softmax_aux_per_step, max_logit_per_step if config.attn_mask_type == AttnMaskType.CAUSAL_MASK: # This is for nested jax.lax.cond @@ -2708,11 +2881,11 @@ def jax_cond_wrap(): ) return lax.cond((idx <= cp_rank), no_mask_compute, skip_compute) - output_per_step, softmax_aux_per_step = lax.cond( + output_per_step, softmax_aux_per_step, max_logit_per_step = lax.cond( idx == 0, causal_mask_compute, jax_cond_wrap ) else: - output_per_step, softmax_aux_per_step = no_mask_compute() + output_per_step, softmax_aux_per_step, max_logit_per_step = no_mask_compute() def skip_correction(output, softmax_aux, output_per_step, softmax_aux_per_step): # No correction done here but we cast outputs to float32 and perform reduction @@ -2735,19 +2908,25 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step): output_per_step, softmax_aux_per_step, ) + # Running per-head max over all ring steps for this rank. + max_logit = jnp.maximum(max_logit, max_logit_per_step) - return (kv_next, output, softmax_aux) + return (kv_next, output, softmax_aux, max_logit) - carry = (kv, output, softmax_aux) + carry = (kv, output, softmax_aux, max_logit) if helper.use_scanloop(): carry = lax.fori_loop(0, cp_size, scan_kv_block, carry) else: for i in range(0, cp_size): carry = scan_kv_block(i, carry) - (kv, output, softmax_aux) = carry + (kv, output, softmax_aux, max_logit) = carry output = output.astype(q.dtype) - return output, softmax_aux, rng_state + # Globalize the rank-local running [H] max across DP/CP. + max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh( + max_logit, mesh, max_logit_reduce_axes, config + ) + return output, softmax_aux, rng_state, max_logit return mesh, ring_attn_fwd_impl, out_shardings, arg_shardings @@ -3065,6 +3244,7 @@ def partition(config, mesh, arg_infos, result_infos): out_sharding = result_infos[0].sharding softmax_aux_sharding = result_infos[1].sharding + max_logit_sharding = result_infos[3].sharding rng_state_sharding = seed_sharding = NamedSharding( mesh, PartitionSpec(get_all_mesh_axes(), None) ) @@ -3074,7 +3254,12 @@ def partition(config, mesh, arg_infos, result_infos): arg_shardings[-1] = arg_shardings[-3] arg_shardings[-2] = arg_shardings[-4] arg_shardings = tuple(arg_shardings) - out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding) + out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding) + max_logit_reduce_axes = ( + FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding) + if config.return_max_logit + else () + ) def fwd_impl( q, @@ -3117,9 +3302,10 @@ def fwd_impl( # support dropout currently. rng_state_shape = (seed.shape[0], *result_infos[2].shape[1:]) rng_state = jnp.zeros(rng_state_shape).astype(result_infos[2].dtype) + max_logit = FusedAttnFwdPrimitive._empty_or_neg_inf_max_logit(head, q.dtype, config) def scan_kv_block(idx, carry): - kv, kv_segment_ids, kv_segment_pos, output, softmax_aux = carry + kv, kv_segment_ids, kv_segment_pos, output, softmax_aux, max_logit = carry # TODO(rewang): To check whether we need special handle for the last idx # Send KV block to next step so we can overlap compute. @@ -3157,7 +3343,9 @@ def compute(config): ) else: current_config = subblock_config - output_per_step, softmax_aux_per_step, _ = compute(current_config) + output_per_step, softmax_aux_per_step, _, max_logit_per_step = compute( + current_config + ) softmax_aux_per_step = softmax_aux_per_step.reshape((batch, q_max_seqlen, head, 1)) @@ -3183,18 +3371,32 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step): output_per_step, softmax_aux_per_step, ) + # Running per-head max over all ring steps for this rank. + max_logit = jnp.maximum(max_logit, max_logit_per_step) - return (kv_next, kv_segment_ids_next, kv_segment_pos_next, output, softmax_aux) + return ( + kv_next, + kv_segment_ids_next, + kv_segment_pos_next, + output, + softmax_aux, + max_logit, + ) - carry = (kv, kv_segment_ids, kv_segment_pos, output, softmax_aux) + carry = (kv, kv_segment_ids, kv_segment_pos, output, softmax_aux, max_logit) if helper.use_scanloop(): carry = lax.fori_loop(0, cp_size, scan_kv_block, carry) else: for i in range(0, cp_size): carry = scan_kv_block(i, carry) - (_, _, _, output, softmax_aux) = carry + (_, _, _, output, softmax_aux, max_logit) = carry - return output.astype(q.dtype), softmax_aux, rng_state + output = output.astype(q.dtype) + # Globalize the rank-local running [H] max across DP/CP. + max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh( + max_logit, mesh, max_logit_reduce_axes, config + ) + return output, softmax_aux, rng_state, max_logit return mesh, fwd_impl, out_shardings, arg_shardings @@ -3375,6 +3577,7 @@ def fused_attn_fwd( context_parallel_causal_load_balanced: bool = False, context_parallel_axis: str = "", stripe_size: int | None = None, + return_max_logit: bool = False, ) -> jnp.ndarray: """ Perform the forward pass of with cuDNN fused attention implementations. @@ -3414,6 +3617,7 @@ def fused_attn_fwd( Indicates the sequences are ordered for causal mask load balancing when running context parallelism. context_parallel_axis (str): The name of the context parallel axis. stripe_size (int | None): Indicates the striping height to be used for ReorderStrategy.Striped Load Balancing + return_max_logit (bool): Whether to return the per-head maximum attention logit. Returns: (jnp.ndarray): The output tensor from the fused attention. """ @@ -3489,6 +3693,7 @@ def fused_attn_fwd( cp_axis=_maybe_context_parallel_axis(context_parallel_axis), cp_striped_window_size=None, stripe_size=stripe_size, + return_max_logit=return_max_logit, ) primitive = None @@ -3506,7 +3711,7 @@ def fused_attn_fwd( primitive = FusedRingAttnFwdPrimitive.outer_primitive seq_desc_flatten, _ = jax.tree.flatten(sequence_descriptor) - output, softmax_aux, rng_state = primitive.bind( + output, softmax_aux, rng_state, max_logit = primitive.bind( *qkv_for_primitive, bias, softmax_offset, @@ -3515,7 +3720,7 @@ def fused_attn_fwd( config=fused_config, ) rng_state = with_sharding_constraint(rng_state, PartitionSpec(get_all_mesh_axes(), None)) - return (output, softmax_aux, rng_state) + return (output, softmax_aux, rng_state, max_logit) def fused_attn_bwd( diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 9bd0940c4b..580219baf2 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -156,7 +156,7 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend( NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen, size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left, - int64_t window_size_right, bool deterministic); + int64_t window_size_right, bool return_max_logit, bool deterministic); pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, @@ -164,7 +164,7 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, bool is_training, size_t max_segments_per_seq, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal); + int64_t window_size_right, bool return_max_logit, bool bottom_right_diagonal); pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 3fd6780d6d..41f87ae71b 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -29,12 +29,12 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend( NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen, size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left, - int64_t window_size_right, bool deterministic) { + int64_t window_size_right, bool return_max_logit, bool deterministic) { auto backend = nvte_get_fused_attn_backend( is_training, static_cast(q_dtype), static_cast(kv_dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, q_attn_heads, kv_attn_heads, q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false, false, deterministic); + return_max_logit, false, deterministic); return backend; } @@ -48,8 +48,8 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t const size_t bias_heads, const size_t q_max_seqlen, const size_t kv_max_seqlen, DType dtype, NVTE_Bias_Type bias_type, NVTE_Fused_Attn_Backend backend, - void *softmax_buf, void *rng_state_buf = nullptr, - void *bias_buf = nullptr, + void *softmax_buf, void *max_logits_buf = nullptr, + void *rng_state_buf = nullptr, void *bias_buf = nullptr, void *softmax_offset_buf = nullptr) { // all backends need softmax but expect different shapes/dtypes tensor_pack->size = 1; @@ -65,8 +65,28 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t // arbitrary sequence length backend needs the RNG state and a different shape/dtype softmax if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { - tensor_pack->size = 2; - NVTETensor &rng_state_aux = tensor_pack->tensors[1]; + int size = 1; // Start after softmax. + auto next_aux_tensor = [&]() -> NVTETensor & { + NVTE_CHECK(size < NVTETensorPack::MAX_SIZE, + "Fused attention auxiliary tensor pack capacity exceeded."); + return tensor_pack->tensors[size++]; + }; + + if (max_logits_buf != nullptr) { + NVTETensor &max_aux = next_aux_tensor(); + NVTEBasicTensor max_aux_data; + max_aux_data.data_ptr = max_logits_buf; + max_aux_data.shape = {}; + max_aux_data.shape.ndim = 4; + max_aux_data.shape.data[0] = input_batch; + max_aux_data.shape.data[1] = attn_heads; + max_aux_data.shape.data[2] = q_max_seqlen; + max_aux_data.shape.data[3] = 1; + max_aux_data.dtype = static_cast(DType::kFloat32); + nvte_set_tensor_param(&max_aux, kNVTERowwiseData, &max_aux_data); + } + + NVTETensor &rng_state_aux = next_aux_tensor(); NVTEBasicTensor rng_state_aux_data; rng_state_aux_data.data_ptr = rng_state_buf; rng_state_aux_data.shape = {}; @@ -77,12 +97,9 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t softmax_aux_data.shape.data[3] = 1; // {B,H,Qs,Ks} -> {B,H,Qs,1} softmax_aux_data.dtype = static_cast(DType::kFloat32); - int size = 2; // Start at 2 (we have softmax and rng_state at indices 0, 1) - // include bias if enabled if (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS && bias_type != NVTE_Bias_Type::NVTE_ALIBI) { - NVTETensor &bias_aux = tensor_pack->tensors[size]; - size++; + NVTETensor &bias_aux = next_aux_tensor(); NVTEBasicTensor bias_aux_data; bias_aux_data.data_ptr = bias_buf; bias_aux_data.shape.ndim = 4; @@ -96,8 +113,7 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t // include softmax_offset if provided if (softmax_offset_buf != nullptr) { - NVTETensor &softmax_offset_aux = tensor_pack->tensors[size]; - size++; + NVTETensor &softmax_offset_aux = next_aux_tensor(); NVTEBasicTensor softmax_offset_aux_data; softmax_offset_aux_data.data_ptr = softmax_offset_buf; softmax_offset_aux_data.shape.ndim = 4; @@ -136,7 +152,7 @@ void PrepareFusedAttnBackwardAuxTensors(NVTETensorPack *tensor_pack, const size_ auto dummy_backend = NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; PrepareFusedAttnForwardAuxTensors(tensor_pack, input_batch, bias_batch, attn_heads, bias_heads, q_max_seqlen, kv_max_seqlen, dtype, dummy_bias_type, - dummy_backend, softmax_buf, rng_state_buf, bias_buf, + dummy_backend, softmax_buf, nullptr, rng_state_buf, bias_buf, softmax_offset_buf); } @@ -146,7 +162,7 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, bool is_training, size_t max_segments_per_seq, int64_t window_size_left, - int64_t window_size_right, bool bottom_right_diagonal) { + int64_t window_size_right, bool return_max_logit, bool bottom_right_diagonal) { auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; auto q_shape = is_ragged ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} @@ -200,8 +216,8 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), ragged_offset_tensor.data(), ragged_offset_tensor.data(), dummy_page_table_tensor.data(), dummy_page_table_tensor.data(), - dummy_rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, false, false, - scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), + dummy_rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, return_max_logit, + false, scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, query_workspace_tensor.data(), nullptr); } @@ -242,13 +258,14 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( static void FusedAttnForwardImpl( cudaStream_t stream, void *q, void *k, void *v, void *bias, void *softmax_offset, void *seed, void *q_cu_seqlens, void *kv_cu_seqlens, void *q_seq_offsets, void *k_seq_offsets, void *output, - void *softmax_aux, void *rng_state, void *workspace, size_t input_batch, size_t bias_batch, - size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, size_t num_gqa_groups, - size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, size_t max_segments_per_seq, - size_t wkspace_size, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, - DType dtype, DType wkspace_dtype, bool is_training, bool deterministic, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal) { + void *softmax_aux, void *max_tensor, void *rng_state, void *workspace, size_t input_batch, + size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, + size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, + size_t max_segments_per_seq, size_t wkspace_size, float scaling_factor, + float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, DType wkspace_dtype, + bool is_training, bool return_max_logit, bool deterministic, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal) { FUSED_ATTN_IMPL_COMMON_BLOCK; /* Input tensors */ @@ -263,6 +280,9 @@ static void FusedAttnForwardImpl( // Memset to 0xF0 for filling large negative numbers auto softmax_aux_size = input_batch * q_max_seqlen * attn_heads; cudaMemsetAsync(softmax_aux, 0xF0, softmax_aux_size * sizeof(float), stream); + if (return_max_logit) { + cudaMemsetAsync(max_tensor, 0xF0, softmax_aux_size * sizeof(float), stream); + } } /* Output tensors */ @@ -278,7 +298,7 @@ static void FusedAttnForwardImpl( is_training, static_cast(dtype), static_cast(dtype), qkv_layout, bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups, q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right, - false, false, deterministic); + return_max_logit, false, deterministic); nvte_populate_rng_state_async(rng_state, seed, q_max_seqlen, kv_max_seqlen, backend, stream); /* Auxiliary tensors (to be propagated to the backward pass later) */ @@ -286,7 +306,8 @@ static void FusedAttnForwardImpl( nvte_tensor_pack_create(&aux_output_tensors); PrepareFusedAttnForwardAuxTensors(&aux_output_tensors, input_batch, bias_batch, attn_heads, bias_heads, q_max_seqlen, kv_max_seqlen, dtype, bias_type, - backend, softmax_aux, softmax_offset); + backend, softmax_aux, return_max_logit ? max_tensor : nullptr, + rng_state, bias, softmax_offset); /* Call the underlying NVTE API */ auto dummy_page_table_tensor = TensorWrapper(nullptr, std::vector{1}, DType::kInt32); @@ -344,7 +365,7 @@ static void FusedAttnForwardImpl( softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors, q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), q_seq_offsets_tensor.data(), k_seq_offsets_tensor.data(), dummy_page_table_tensor.data(), dummy_page_table_tensor.data(), - rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, false, false, + rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, return_max_logit, false, scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, workspace_tensor.data(), stream); @@ -378,6 +399,7 @@ static void FusedAttnForwardImpl( NVTE_QKV_Layout qkv_layout = \ static_cast(get_attr_value(attrs, "qkv_layout")); \ bool is_training = get_attr_value(attrs, "is_training"); \ + bool return_max_logit = get_attr_value_or_default(attrs, "return_max_logit", false); \ bool deterministic = get_attr_value(attrs, "deterministic"); \ auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; \ size_t wkspace_size = product(workspace_buf->dimensions()); \ @@ -390,8 +412,9 @@ Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Ty Buffer_Type q_cu_seqlens_buf, Buffer_Type kv_cu_seqlens_buf, Buffer_Type q_seq_offsets_buf, Buffer_Type k_seq_offsets_buf, Variadic_Buffer_Type _unused_args, Result_Type output_buf, - Result_Type softmax_aux_buf, Result_Type rng_state_buf, - Result_Type workspace_buf, Dictionary attrs) { + Result_Type softmax_aux_buf, Result_Type max_tensor_buf, + Result_Type rng_state_buf, Result_Type workspace_buf, + Dictionary attrs) { FUSED_ATTN_FFI_GET_ATTRS; FusedAttnForwardImpl( @@ -400,11 +423,12 @@ Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Ty q_cu_seqlens_buf.untyped_data(), kv_cu_seqlens_buf.untyped_data(), is_ragged ? q_seq_offsets_buf.untyped_data() : nullptr, is_ragged ? k_seq_offsets_buf.untyped_data() : nullptr, output_buf->untyped_data(), - softmax_aux_buf->untyped_data(), rng_state_buf->untyped_data(), workspace_buf->untyped_data(), - input_batch, bias_batch, q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, - qk_head_dim, v_head_dim, max_segments_per_seq, wkspace_size, scaling_factor, - dropout_probability, bias_type, mask_type, softmax_type, qkv_layout, dtype, wkspace_dtype, - is_training, deterministic, window_size_left, window_size_right, bottom_right_diagonal); + softmax_aux_buf->untyped_data(), max_tensor_buf->untyped_data(), + rng_state_buf->untyped_data(), workspace_buf->untyped_data(), input_batch, bias_batch, + q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim, + max_segments_per_seq, wkspace_size, scaling_factor, dropout_probability, bias_type, mask_type, + softmax_type, qkv_layout, dtype, wkspace_dtype, is_training, return_max_logit, deterministic, + window_size_left, window_size_right, bottom_right_diagonal); return ffi_with_cuda_error_check(); } @@ -424,6 +448,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnForwardHandler, FusedAttnForwardFFI, .RemainingArgs() // _cp_aux_args unused .Ret() // output .Ret() // softmax_aux + .Ret() // max_tensor .Ret() // rng_state .Ret() // workspace .Attrs(), diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 76922d2b55..4b497826cc 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -308,6 +308,7 @@ class _FusedDotProductAttention(nn.Module): # pylint: disable=too-few-public-me score_mod: Optional[Callable] = None score_mod_bprop: Optional[Callable] = None score_mod_requested: bool = False + return_max_logit: bool = False @nn.compact def __call__( @@ -363,6 +364,7 @@ def __call__( "score_mod_bprop": self.score_mod_bprop, "score_mod_tensors": score_mod_tensors, "score_mod_bprop_tensors": score_mod_bprop_tensors, + "return_max_logit": self.return_max_logit, } if self.qkv_layout.is_qkvpacked(): @@ -434,12 +436,17 @@ def __call__( else: raise ValueError(f"Unsupported {self.qkv_layout=}.") + if self.return_max_logit: + x, max_logit = x + if self.transpose_batch_sequence: x = x.transpose([1, 0, 2, 3]) assert ( x.dtype == query.dtype ), f"output dtype {x.dtype} does not match query dtype {query.dtype}" + if self.return_max_logit: + return x, max_logit return x @@ -619,6 +626,9 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods argument to keep tensor operands as normal JAX inputs. score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None Additional tensors or pass-by-value scalars for ``score_mod_bprop``. + return_max_logit: bool, default = False + If True, return ``(output, max_logit)`` where ``max_logit`` contains the per-head + maximum attention logits with shape ``[h]``. This path requires fused attention. Optimization parameters ----------------------- @@ -647,6 +657,7 @@ class DotProductAttention(nn.Module): # pylint: disable=too-few-public-methods softmax_type: str = "vanilla" score_mod: Optional[Callable] = None score_mod_bprop: Optional[Callable] = None + return_max_logit: bool = False def __post_init__(self): # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12 @@ -717,8 +728,8 @@ def __call__( Returns ------- - outputs: jax.numpy.ndarray - Output tensors. + outputs: jax.numpy.ndarray or tuple[jax.numpy.ndarray, jax.numpy.ndarray] + Output tensor, or ``(output, max_logit)`` when ``return_max_logit`` is enabled. """ input_dtype = query.dtype @@ -777,6 +788,8 @@ def __call__( # Use fused attn (if kernel check below passes) by default enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "1")) + if self.return_max_logit and not enable_fused_attn: + raise ValueError("return_max_logit requires fused attention, but NVTE_FUSED_ATTN=0.") sequence_dim = 0 if self.transpose_batch_sequence else 1 seqlen_q = query.shape[sequence_dim] @@ -815,11 +828,17 @@ def __call__( head_dim_qk, head_dim_v, self.window_size, + return_max_logit=self.return_max_logit, ) if score_mod_requested and not has_fused_attn_kernel: raise ValueError( "score_mod requires fused attention, but no fused attention kernel is available." ) + if self.return_max_logit and not has_fused_attn_kernel: + raise ValueError( + "return_max_logit requires fused attention, but no fused attention kernel is " + "available." + ) use_fused_attn = enable_fused_attn and has_fused_attn_kernel @@ -916,6 +935,7 @@ def __call__( score_mod=self.score_mod, score_mod_bprop=self.score_mod_bprop, score_mod_requested=score_mod_requested, + return_max_logit=self.return_max_logit, )( query, key, @@ -927,7 +947,10 @@ def __call__( score_mod_tensors=score_mod_tensors, score_mod_bprop_tensors=score_mod_bprop_tensors, ) - assert x.dtype == input_dtype, f"output_dtype={x.dtype}, input_dtype={input_dtype}" + output = x[0] if self.return_max_logit else x + assert ( + output.dtype == input_dtype + ), f"output_dtype={output.dtype}, input_dtype={input_dtype}" return x