Max logits fix (#239)

* refactor max_logits code and add max_logits feature for sdpa backend and test_pipeline_sdpa

* add max_logits_print in test_case of test_pipeline_sdpa

* enable sdpa backend when test max_logits in test_dist_attn

* fix unit test bugs

* refactor code

* update test_dist_attn

Author: lijinnn

magi_attention/common/forward_meta.py

@@ -19,5 +19,16 @@
 
 @dataclass
 class AttnForwardMeta:
+    """Attention forward metadata.
+
+    Attributes:
+        lse: Log-sum-exp of the attention weights. In a distributed setting, this is a
+            local tensor where each device holds the LSE computed from its local query
+            shards.
+        max_logits: Maximum logits per query head. In a distributed setting,
+            this is a replicated tensor where each device holds the global maximum
+            computed across the entire sequence, ensuring consistency across all devices.
+    """
+
     lse: torch.Tensor | None
     max_logits: torch.Tensor | None

magi_attention/dist_attn_runtime_mgr.py

@@ -150,6 +150,7 @@ def calc_attn(
         sink: torch.Tensor | None = None,
         softmax_scale: float | None = None,
         softcap: float = 0.0,
+        return_max_logits: bool = False,
     ) -> tuple[torch.Tensor, AttnForwardMeta]:
         return dist_attn_func(
             q=q,
@@ -159,6 +160,7 @@ def calc_attn(
             sink=sink,
             softmax_scale=softmax_scale,
             softcap=softcap,
+            return_max_logits=return_max_logits,
         )
 
     def get_xattn_args(

magi_attention/functional/dist_attn.py

@@ -230,7 +230,13 @@ def apply_fwd_partial_attn(
                     q=q,
                     sink=sink,
                 )
-                return partial_out, AttnForwardMeta(lse=partial_lse, max_logits=None)
+                partial_max_logits = self._init_max_logits_skipped_host_stage(
+                    q=q,
+                    return_max_logits=return_max_logits,
+                )
+                return partial_out, AttnForwardMeta(
+                    lse=partial_lse, max_logits=partial_max_logits
+                )
             return None, None
 
         # attention forward pass
@@ -1056,17 +1062,17 @@ def _launch_attn_fwd_kernel(
         return_max_logits: bool = False,
     ) -> tuple[torch.Tensor, AttnForwardMeta]:
         if return_max_logits:
-            assert not (
-                self.use_sdpa_backend or self.use_fa4_backend
-            ), "SDPA and FA4 backend do not support return max logits"
+            assert (
+                not self.use_fa4_backend
+            ), "FA4 backend does not support return max logits"
         with nvtx.add_nvtx_event(
             f"attn-fwd: "
             f"{attn_arg.total_area=} | "
             f"{attn_arg.q_ranges=} | "
             f"{attn_arg.k_ranges=}"
         ):
             if self.use_sdpa_backend:
-                partial_out, partial_lse = sdpa_fwd(
+                partial_out, meta = sdpa_fwd(
                     q=q,
                     k=k,
                     v=v,
@@ -1077,8 +1083,12 @@ def _launch_attn_fwd_kernel(
                     softmax_scale=softmax_scale,
                     softcap=softcap,
                     sink_layout="sh",
+                    return_max_logits=return_max_logits,
                 )
-                meta = AttnForwardMeta(lse=partial_lse, max_logits=None)
+                if return_max_logits and max_logits_acc is not None:
+                    assert meta.max_logits is not None
+                    torch.maximum(max_logits_acc, meta.max_logits, out=max_logits_acc)
+                    meta.max_logits = max_logits_acc
             elif self.use_fa4_backend:
                 partial_out, partial_lse = fa4_fwd(
                     q=q,
@@ -1920,6 +1930,20 @@ def _init_out_lse_skipped_host_stage(
 
         return out, lse
 
+    def _init_max_logits_skipped_host_stage(
+        self,
+        q: torch.Tensor,
+        return_max_logits: bool,
+    ) -> torch.Tensor | None:
+        if return_max_logits:
+            return torch.full(
+                (q.size(1),),  # [nhq]
+                fill_value=float("-inf"),
+                dtype=q.dtype,
+                device=q.device,
+            )
+        return None
+
     def _init_dq_dkv_dsink_skipped_host_stage(
         self,
         qo_do: FusedOrTupleTensor,

magi_attention/functional/flex_flash_attn.py

@@ -397,7 +397,7 @@ def _flex_flash_attn_forward(
         assert q.size(1) <= 128, (
             f"num_qheads ({q.size(1)}) must be <= 128 because the epilogue shmem "
             "for max_logits reduction is fixed at 128 in C++ code. You can increase "
-            "the shmem size by increasing the `smem_max_logitss` in `epilogue_fwd.hpp`."
+            "the shmem size by increasing the `smem_max_logits` in `epilogue_fwd.hpp`."
         )
 
     if ref_block_size is not None:

magi_attention/functional/sdpa.py

@@ -17,6 +17,7 @@
 from einops import reduce
 
 from magi_attention.common.enum import AttnSinkLayout
+from magi_attention.common.forward_meta import AttnForwardMeta
 from magi_attention.meta.collection.calc_meta import AttnArg
 from magi_attention.utils import make_attn_mask_from_ffa_args, to_higher_fp_dtype
 
@@ -94,22 +95,30 @@ def sdpa_fwd_calc(
     v: torch.Tensor,
     attn_bias: torch.Tensor,
     softmax_scale: float,
-) -> tuple[torch.Tensor, torch.Tensor]:
+    return_max_logits: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
     attn_weight = to_higher_fp_dtype(
         q @ k.transpose(-2, -1) * softmax_scale,
         lowest_precision=torch.float32,
     )
     attn_weight += attn_bias
 
     lse = attn_weight.logsumexp(dim=-1, keepdim=True)
+    if return_max_logits:
+        # compute per-head max logits over score matrix
+        # attn_weight shape: [batch_size, num_heads, num_tokens_q, num_tokens_k]
+        bsz, nhq = attn_weight.shape[:2]
+        max_logits = attn_weight.view(bsz, nhq, -1).max(dim=-1).values.contiguous()
+    else:
+        max_logits = None
 
     # NOTE: pytorch softmax has many limitations and bugs
     # thus we use our own safe_softmax with lse involved
     attn_weight = safe_softmax(attn_weight, lse).to(v.dtype)
 
     out = attn_weight @ v
 
-    return out, lse.squeeze(-1)
+    return out, lse.squeeze(-1), max_logits
 
 
 def _sdpa_fwd(
@@ -119,14 +128,17 @@ def _sdpa_fwd(
     attn_mask: torch.Tensor | None = None,
     is_causal: bool = False,
     softmax_scale: float | None = None,
-) -> tuple[torch.Tensor, torch.Tensor]:
+    return_max_logits: bool = False,
+) -> tuple[torch.Tensor, AttnForwardMeta]:
     q, k, v, attn_bias, softmax_scale, _ = sdpa_fwd_preprocess(
         q, k, v, attn_mask, is_causal, softmax_scale
     )
 
-    out, lse = sdpa_fwd_calc(q, k, v, attn_bias, softmax_scale)
+    out, lse, max_logits = sdpa_fwd_calc(
+        q, k, v, attn_bias, softmax_scale, return_max_logits
+    )
 
-    return out, lse
+    return out, AttnForwardMeta(lse=lse, max_logits=max_logits)
 
 
 @torch.no_grad()
@@ -139,7 +151,8 @@ def sdpa_fwd(
     softmax_scale: float | None = None,
     softcap: float = 0.0,
     sink_layout: AttnSinkLayout = "sh",
-) -> tuple[torch.Tensor, torch.Tensor]:
+    return_max_logits: bool = False,
+) -> tuple[torch.Tensor, AttnForwardMeta]:
     """SDPA forward function
 
     Args:
@@ -163,12 +176,19 @@ def sdpa_fwd(
 
         sink_layout (AttnSinkLayout, optional): sink layout. Defaults to "sh".
 
+        return_max_logits (bool, optional): whether to return max logits.
+            Defaults to ``False``.
+
     Returns:
         torch.Tensor: out with shape [num_tokens_q, num_heads_q, head_dim]
             or [batch_size, num_heads_q, num_tokens_q, head_dim]
 
-        torch.Tensor: lse with shape [num_tokens_q, num_heads_q]
-            or [batch_size, num_heads_q, num_tokens_q]
+        AttnForwardMeta: metadata for attention forward, including lse and max_logits.
+            - lse (torch.Tensor): [num_tokens_q, num_heads_q]
+                or [batch_size, num_heads_q, num_tokens_q]
+            - max_logits (torch.Tensor or None): [num_heads_q]
+                or [batch_size, num_heads_q]
+                or None if return_max_logits is False
     """
     assert softcap == 0.0, "non-zero softcap is not supported by now"
 
@@ -187,17 +207,21 @@ def sdpa_fwd(
         device=torch.cuda.current_device(),
     )
 
-    out, lse = _sdpa_fwd(
+    out, meta = _sdpa_fwd(
         q,
         k,
         v,
         attn_mask=attn_mask,
         is_causal=False,
         softmax_scale=softmax_scale,
+        return_max_logits=return_max_logits,
     )
+    lse, max_logits = meta.lse, meta.max_logits
 
     if rearrange:
         out, lse = sdpa_fwd_out_lse_rearrange(out, lse)
+        if max_logits is not None:
+            max_logits = max_logits.squeeze(0)
 
     if sink is not None:
         assert rearrange
@@ -209,7 +233,7 @@ def sdpa_fwd(
             inplace=True,
         )
 
-    return out, lse
+    return out, AttnForwardMeta(lse=lse, max_logits=max_logits)
 
 
 # ------------------        sdpa bwd       ------------------ #

magi_attention/testing/ref_attn.py

@@ -700,6 +700,7 @@ def ref_attn_func(
     # maybe cast input to high precision
     org_dtype = q.dtype
     lse_dtype = max_fp_dtype(org_dtype, torch.float32)
+    max_logits_dtype = max_fp_dtype(org_dtype, torch.float32)
     if high_precision:  # use fp64 as ground-truth
         q = q.to(torch.float64)
         k = k.to(torch.float64)
@@ -743,6 +744,6 @@ def ref_attn_func(
     if return_max_logits:
         assert meta is not None  # mypy
         assert meta.max_logits is not None  # mypy
-        meta.max_logits = meta.max_logits.to(torch.float32)
+        meta.max_logits = meta.max_logits.to(max_logits_dtype)
 
     return out, meta

tests/test_api/test_interface.py

@@ -123,7 +123,7 @@ def world_size(self) -> int:
 
     @property
     def timeout(self) -> int:
-        return 600
+        return 1200
 
     @property
     def seed(self) -> int:

tests/test_attn/test_dist_attn.py

@@ -118,7 +118,7 @@ def world_size(self) -> int:
 
     @property
     def timeout(self) -> int:
-        return 1200
+        return 1800
 
     @property
     def seed(self) -> int:
@@ -163,10 +163,6 @@ def test_full_attn(
             if use_native_grpcoll:
                 return
 
-        # sdpa backend do not support return max logits
-        if return_max_logits and use_sdpa_backend:
-            return
-
         # switch the env flags
         switch_back = switch_envvars(
             envvar_name_list=[
@@ -341,7 +337,7 @@ def test_full_attn(
                 local_max_logits,
                 total_max_logits_ref,
                 atol=EPSILON,
-                rtol=1e-3,
+                rtol=1e-2 if use_sdpa_backend else 1e-3,
                 mismatch_threshold=0.01,
                 test_case="max_logits",
             )
@@ -373,7 +369,7 @@ def test_full_attn(
             assert_close(
                 total_dsink,
                 total_dsink_ref,
-                atol=1e-3,
+                atol=5e-3,
                 rtol=0.1,
                 mismatch_threshold=max(1 / (seqlen_sink * nhq), 5e-2),
                 test_case="dsink",