[FullCG] Preserve attention LSE through the custom-op boundary (#31050)

python/sglang/srt/layers/radix_attention.py

@@ -147,6 +147,7 @@ class RadixAttention(nn.Module):
         v,
         forward_batch: ForwardBatch,
         save_kv_cache: bool = True,
+        key_value_num_tokens: Optional[int] = None,
         **kwargs,
     ):
         if k is not None:
@@ -158,9 +159,10 @@ class RadixAttention(nn.Module):
             else:
                 k = k.view(-1, self.tp_k_head_num, self.v_head_dim)
 
+        context = get_tc_piecewise_forward_context()
         if (
             forward_batch.forward_mode.is_extend()
-            and get_tc_piecewise_forward_context() is not None
+            and context is not None
             # ``_force_eager_attn`` is only set inside Inkling's eager
             # norm+attn+sconv region, never during tc-piecewise capture. Reading
             # the ContextVar under the fullgraph torch.compile trace is
@@ -232,14 +234,39 @@ class RadixAttention(nn.Module):
                         q, k, v, output, save_kv_cache, self.layer_id, kwargs
                     )
                 return output
+            # Chunked-prefix MHA needs LSE to merge independently normalized
+            # suffix and cached-prefix attention states.
+            return_lse = bool(forward_batch.mha_return_lse)
+            mha_companion_layers = context.mha_companion_layers
+            use_mha_companion = (
+                mha_companion_layers is not None
+                and mha_companion_layers[self.layer_id] is self
+            )
             if is_in_breakable_cuda_graph():
-                breakable_unified_attention_with_output(
-                    q, k, v, output, save_kv_cache, self.layer_id, **kwargs
+                op = (
+                    breakable_unified_attention_with_output_and_lse
+                    if return_lse
+                    else breakable_unified_attention_with_output
                 )
             else:
-                unified_attention_with_output(
-                    q, k, v, output, save_kv_cache, self.layer_id, **kwargs
+                op = (
+                    unified_attention_with_output_and_lse
+                    if return_lse
+                    else unified_attention_with_output
                 )
+            lse = op(
+                q,
+                k,
+                v,
+                output,
+                save_kv_cache,
+                self.layer_id,
+                use_mha_companion=use_mha_companion,
+                key_value_num_tokens=key_value_num_tokens,
+                **kwargs,
+            )
+            if return_lse:
+                return output.view(-1, self.tp_q_head_num, self.v_head_dim), lse
             return output
         else:
             return get_attn_backend().forward(
@@ -253,16 +280,17 @@ class RadixAttention(nn.Module):
             )
 
 
-@register_custom_op(mutates_args=["output"])
-@register_split_op()
-def unified_attention_with_output(
+def _unified_attention_with_output_impl(
     query: torch.Tensor,
     key: Optional[torch.Tensor],
     value: Optional[torch.Tensor],
     output: torch.Tensor,
     save_kv_cache: bool,
     layer_id: int,
+    use_mha_companion: bool,
+    return_lse: bool,
     *,
+    key_value_num_tokens: Optional[int] = None,
     q_rope: Optional[torch.Tensor] = None,
     k_rope: Optional[torch.Tensor] = None,
     sinks: Optional[torch.Tensor] = None,
@@ -272,29 +300,38 @@ def unified_attention_with_output(
     is_neox: Optional[bool] = None,
     llama_4_scaling: Optional[torch.Tensor] = None,
     topk_indices: Optional[torch.Tensor] = None,
-) -> None:
+) -> Optional[torch.Tensor]:
     context = get_tc_piecewise_forward_context()
     forward_batch = context.forward_batch
     attention_layers = context.attention_layers
     attention_layer = attention_layers[layer_id]
-    real_num_tokens = forward_batch.num_token_non_padded_cpu
+    real_query_num_tokens = forward_batch.num_token_non_padded_cpu
+    # Ordinary PCG attention pads Q/K/V to the same token bucket. Prefix MHA
+    # instead supplies a fixed-capacity K/V chunk whose extent is independent
+    # of the suffix queries, so its caller must preserve that separate extent.
+    if key_value_num_tokens is None:
+        key_value_num_tokens = real_query_num_tokens
 
-    query = query[:real_num_tokens]
+    query = query[:real_query_num_tokens]
     if key is not None:
-        key = key[:real_num_tokens]
+        key = key[:key_value_num_tokens]
     if value is not None:
-        value = value[:real_num_tokens]
+        value = value[:key_value_num_tokens]
 
-    if not save_kv_cache and context.mha_companion_layers is not None:
-        mha_companion_layer = context.mha_companion_layers[layer_id]
-        if mha_companion_layer is not None:
-            attention_layer = mha_companion_layer
+    # DeepSeek MLA has two RadixAttention instances per layer (attn_mqa and
+    # attn_mha) that share the same layer_id. Preserve the calling instance's
+    # identity through the custom-op boundary; save_kv_cache is not an identity
+    # signal because absorbed MLA can also disable a redundant cache store.
+    if use_mha_companion:
+        assert context.mha_companion_layers is not None
+        attention_layer = context.mha_companion_layers[layer_id]
+        assert attention_layer is not None
 
     kwargs = {}
     if q_rope is not None:
-        kwargs["q_rope"] = q_rope[:real_num_tokens]
+        kwargs["q_rope"] = q_rope[:real_query_num_tokens]
     if k_rope is not None:
-        kwargs["k_rope"] = k_rope[:real_num_tokens]
+        kwargs["k_rope"] = k_rope[:key_value_num_tokens]
     if sinks is not None:
         kwargs["sinks"] = sinks
     if cos_sin_cache is not None:
@@ -304,17 +341,17 @@ def unified_attention_with_output(
     if llama_4_scaling is not None:
         kwargs["llama_4_scaling"] = llama_4_scaling
     if topk_indices is not None:
-        kwargs["topk_indices"] = topk_indices[:real_num_tokens]
+        kwargs["topk_indices"] = topk_indices[:real_query_num_tokens]
 
     original_out_cache_loc = forward_batch.out_cache_loc
     # Keep the original ForwardBatch object and only narrow cache locations for
     # this backend call so model/backend state is still written to the same batch.
-    forward_batch.out_cache_loc = original_out_cache_loc[:real_num_tokens]
+    forward_batch.out_cache_loc = original_out_cache_loc[:real_query_num_tokens]
 
     # Store pre-allocated output for FA backend to write directly into.
-    # Must slice to real_num_tokens to match the narrowed query shape —
+    # Must slice to real_query_num_tokens to match the narrowed query shape —
     # the FA kernel validates out.size(0) == q.size(0).
-    forward_batch._attn_output = output[:real_num_tokens]
+    forward_batch._attn_output = output[:real_query_num_tokens]
 
     ret = get_attn_backend().forward(
         query,
@@ -327,18 +364,119 @@ def unified_attention_with_output(
     )
     forward_batch.out_cache_loc = original_out_cache_loc
 
+    lse = None
+    if return_lse:
+        assert isinstance(ret, tuple)
+        ret, lse, *_ = ret
+    else:
+        assert isinstance(ret, torch.Tensor)
+
     if ret.data_ptr() != output.data_ptr():
-        output[:real_num_tokens].view(ret.shape).copy_(ret)
+        output[:real_query_num_tokens].view(ret.shape).copy_(ret)
 
     # During PCG replay the attention backend writes only the narrowed
-    # real-token slice (output[:real_num_tokens]) and leaves padded positions
+    # real-token slice (output[:real_query_num_tokens]) and leaves padded positions
     # as uninitialized torch.empty garbage. Zero them so garbage (NaN/Inf) does
     # not propagate through residual connections, MoE routing, and allreduce.
     # This affects every backend that varlen-writes under PCG, not just ROCm.
     # Use context.raw_num_tokens (pre-padding count from PCG runner) instead of
     # forward_batch.extend_num_tokens, which is None for TARGET_VERIFY batches.
     _zero_padded_pcg_tail(output, context)
-    return
+    if lse is not None and lse.shape[0] != output.shape[0]:
+        padded_lse = lse.new_zeros((output.shape[0], *lse.shape[1:]))
+        padded_lse[:real_query_num_tokens].copy_(lse)
+        lse = padded_lse
+    return lse
+
+
+@register_custom_op(mutates_args=["output"])
+@register_split_op()
+def unified_attention_with_output(
+    query: torch.Tensor,
+    key: Optional[torch.Tensor],
+    value: Optional[torch.Tensor],
+    output: torch.Tensor,
+    save_kv_cache: bool,
+    layer_id: int,
+    *,
+    use_mha_companion: bool = False,
+    key_value_num_tokens: Optional[int] = None,
+    q_rope: Optional[torch.Tensor] = None,
+    k_rope: Optional[torch.Tensor] = None,
+    sinks: Optional[torch.Tensor] = None,
+    cos_sin_cache: Optional[torch.Tensor] = None,
+    is_neox: Optional[bool] = None,
+    llama_4_scaling: Optional[torch.Tensor] = None,
+    topk_indices: Optional[torch.Tensor] = None,
+) -> None:
+    _unified_attention_with_output_impl(
+        query,
+        key,
+        value,
+        output,
+        save_kv_cache,
+        layer_id,
+        use_mha_companion,
+        False,
+        key_value_num_tokens=key_value_num_tokens,
+        q_rope=q_rope,
+        k_rope=k_rope,
+        sinks=sinks,
+        cos_sin_cache=cos_sin_cache,
+        is_neox=is_neox,
+        llama_4_scaling=llama_4_scaling,
+        topk_indices=topk_indices,
+    )
+
+
+def _unified_attention_with_output_and_lse_fake(
+    query: torch.Tensor, *args, **kwargs
+) -> torch.Tensor:
+    return query.new_empty((query.shape[0], query.shape[1]), dtype=torch.float32)
+
+
+@register_custom_op(
+    mutates_args=["output"], fake_impl=_unified_attention_with_output_and_lse_fake
+)
+@register_split_op()
+def unified_attention_with_output_and_lse(
+    query: torch.Tensor,
+    key: Optional[torch.Tensor],
+    value: Optional[torch.Tensor],
+    output: torch.Tensor,
+    save_kv_cache: bool,
+    layer_id: int,
+    *,
+    use_mha_companion: bool = False,
+    key_value_num_tokens: Optional[int] = None,
+    q_rope: Optional[torch.Tensor] = None,
+    k_rope: Optional[torch.Tensor] = None,
+    sinks: Optional[torch.Tensor] = None,
+    cos_sin_cache: Optional[torch.Tensor] = None,
+    is_neox: Optional[bool] = None,
+    llama_4_scaling: Optional[torch.Tensor] = None,
+    topk_indices: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    lse = _unified_attention_with_output_impl(
+        query,
+        key,
+        value,
+        output,
+        save_kv_cache,
+        layer_id,
+        use_mha_companion,
+        True,
+        key_value_num_tokens=key_value_num_tokens,
+        q_rope=q_rope,
+        k_rope=k_rope,
+        sinks=sinks,
+        cos_sin_cache=cos_sin_cache,
+        is_neox=is_neox,
+        llama_4_scaling=llama_4_scaling,
+        topk_indices=topk_indices,
+    )
+    assert lse is not None
+    return lse
 
 
 @register_custom_op(mutates_args=["attn_out", "idx_out"])
@@ -401,6 +539,9 @@ def unified_sparse_attention_with_output(
 breakable_unified_attention_with_output = eager_on_graph(True)(
     unified_attention_with_output
 )
+breakable_unified_attention_with_output_and_lse = eager_on_graph(True)(
+    unified_attention_with_output_and_lse
+)
 
 
 def attention_with_output_extra_kwargs(

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py

@@ -141,7 +141,6 @@ def _forward_dsa_indexer_for_mha(
 
 
 class DeepseekMHAForwardMixin:
-
     def init_mha_forward(self: DeepseekV2AttentionMLA):
         self.disable_chunked_prefix_cache = (
             get_server_args().disable_chunked_prefix_cache
@@ -224,7 +223,6 @@ class DeepseekMHAForwardMixin:
                 )
                 q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
             elif _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.float8_e4m3fn:
-
                 q, _, _, _ = fused_rms_fp8_group_quant(
                     q,
                     self.q_a_layernorm.weight,
@@ -256,7 +254,6 @@ class DeepseekMHAForwardMixin:
         latent_cache = latent_cache.unsqueeze(1)
 
         if _use_aiter_gfx95 and self.kv_b_proj.weight.dtype == torch.float8_e4m3fn:
-
             kv_a_quanted, kv_a, _, _ = fused_rms_fp8_group_quant(
                 kv_a,
                 self.kv_a_layernorm.weight,
@@ -464,7 +461,6 @@ class DeepseekMHAForwardMixin:
         accum_lse: torch.Tensor,
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
-
         # kv_b_proj needs BF16 input, but legacy q.dtype was BF16 by accident.
         backend = _resolve_attn_backend(forward_batch)
         pack_fn = getattr(backend, "pack_prefix_chunk_kv", None)
@@ -507,7 +503,17 @@ class DeepseekMHAForwardMixin:
                 k[..., : self.qk_nope_head_dim] = k_nope
                 k[..., self.qk_nope_head_dim :] = k_pe
 
-            output, lse = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False)
+            output, lse = self.attn_mha(
+                q,
+                k,
+                v,
+                forward_batch,
+                save_kv_cache=False,
+                # Prefix K/V is independent of the suffix query length. Under
+                # FullCG this is the fixed captured chunk extent; per-request
+                # active lengths remain encoded in the backend metadata.
+                key_value_num_tokens=k.shape[0],
+            )
             tmp_output = torch.empty_like(accum_output)
             tmp_lse = torch.empty_like(accum_lse)
             merge_state_v2(output, lse, accum_output, accum_lse, tmp_output, tmp_lse)

test/registered/unit/layers/test_radix_attention.py

@@ -0,0 +1,321 @@
+"""CPU unit tests for the graph-safe ``RadixAttention`` interface."""
+
+import unittest
+from contextlib import ExitStack
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+import sglang.srt.layers.radix_attention as radix_attention_module
+from sglang.srt.layers.radix_attention import RadixAttention
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+class _RecordingAttentionBackend:
+    def __init__(self, *, return_lse=True):
+        self.calls = []
+        self.return_lse = return_lse
+
+    def forward(
+        self,
+        query,
+        key,
+        value,
+        attention_layer,
+        forward_batch,
+        save_kv_cache,
+        **kwargs,
+    ):
+        self.calls.append(
+            SimpleNamespace(
+                query=query,
+                key=key,
+                value=value,
+                attention_layer=attention_layer,
+                output=forward_batch._attn_output,
+                out_cache_loc=forward_batch.out_cache_loc.clone(),
+                save_kv_cache=save_kv_cache,
+                kwargs=kwargs,
+            )
+        )
+        output = torch.full_like(query, 3)
+        lse = torch.full((query.shape[0], query.shape[1]), 7, dtype=torch.float32)
+        return (output, lse) if self.return_lse else output
+
+
+class TestRadixAttentionGraphInterface(CustomTestCase):
+    @staticmethod
+    def _new_layer() -> RadixAttention:
+        layer = RadixAttention(
+            num_heads=2,
+            head_dim=3,
+            scaling=1.0,
+            num_kv_heads=2,
+            layer_id=0,
+        )
+        return layer
+
+    @staticmethod
+    def _new_impl_context(
+        attention_layers,
+        *,
+        mha_companion_layers=None,
+        num_tokens=4,
+        real_num_tokens=2,
+    ):
+        forward_batch = SimpleNamespace(
+            num_token_non_padded_cpu=real_num_tokens,
+            out_cache_loc=torch.arange(num_tokens, dtype=torch.int64),
+            _attn_output=None,
+        )
+        return SimpleNamespace(
+            forward_batch=forward_batch,
+            attention_layers=attention_layers,
+            mha_companion_layers=mha_companion_layers,
+            num_tokens=None,
+            raw_num_tokens=None,
+        )
+
+    def test_forward_dispatches_all_graph_and_lse_variants(self):
+        layer = self._new_layer()
+        query = torch.zeros((4, 2, 3))
+        key = torch.zeros_like(query)
+        value = torch.zeros_like(query)
+
+        op_names = {
+            (False, False): "unified_attention_with_output",
+            (False, True): "unified_attention_with_output_and_lse",
+            (True, False): "breakable_unified_attention_with_output",
+            (True, True): "breakable_unified_attention_with_output_and_lse",
+        }
+
+        for breakable in (False, True):
+            for return_lse in (False, True):
+                with self.subTest(breakable=breakable, return_lse=return_lse):
+                    forward_batch = SimpleNamespace(
+                        forward_mode=ForwardMode.EXTEND,
+                        mha_return_lse=return_lse,
+                    )
+                    calls = []
+
+                    def output_only(*args, **kwargs):
+                        args[3].fill_(5)
+                        calls.append(kwargs)
+
+                    def output_and_lse(*args, **kwargs):
+                        args[3].fill_(5)
+                        calls.append(kwargs)
+                        return torch.full((4, 2), 11, dtype=torch.float32)
+
+                    with ExitStack() as stack:
+                        stack.enter_context(
+                            patch.object(
+                                radix_attention_module,
+                                "get_tc_piecewise_forward_context",
+                                return_value=SimpleNamespace(
+                                    mha_companion_layers=[layer]
+                                ),
+                            )
+                        )
+                        stack.enter_context(
+                            patch.object(
+                                radix_attention_module,
+                                "is_in_breakable_cuda_graph",
+                                return_value=breakable,
+                            )
+                        )
+                        mocks = {
+                            name: stack.enter_context(
+                                patch.object(
+                                    radix_attention_module,
+                                    name,
+                                    side_effect=(
+                                        output_and_lse
+                                        if name.endswith("and_lse")
+                                        else output_only
+                                    ),
+                                )
+                            )
+                            for name in op_names.values()
+                        }
+                        result = layer(
+                            query,
+                            key,
+                            value,
+                            forward_batch,
+                            key_value_num_tokens=3,
+                        )
+
+                    selected_name = op_names[(breakable, return_lse)]
+                    for name, mock in mocks.items():
+                        self.assertEqual(mock.call_count, int(name == selected_name))
+
+                    self.assertEqual(
+                        calls,
+                        [
+                            {
+                                "use_mha_companion": True,
+                                "key_value_num_tokens": 3,
+                            }
+                        ],
+                    )
+                    if return_lse:
+                        output, lse = result
+                        self.assertEqual(lse.shape, (4, 2))
+                        self.assertTrue(torch.all(lse == 11))
+                    else:
+                        output = result
+                    self.assertEqual(output.shape, query.shape)
+                    self.assertTrue(torch.all(output == 5))
+
+    def test_impl_preserves_attention_identity_and_lse(self):
+        mqa = SimpleNamespace()
+        mha = SimpleNamespace()
+        context = self._new_impl_context([mqa], mha_companion_layers=[mha])
+        forward_batch = context.forward_batch
+        original_out_cache_loc = forward_batch.out_cache_loc
+        backend = _RecordingAttentionBackend()
+        query = torch.zeros((4, 2, 3))
+
+        with (
+            patch.object(
+                radix_attention_module,
+                "get_tc_piecewise_forward_context",
+                return_value=context,
+            ),
+            patch.object(
+                radix_attention_module, "get_attn_backend", return_value=backend
+            ),
+        ):
+            for use_mha_companion, expected_layer in ((False, mqa), (True, mha)):
+                with self.subTest(use_mha_companion=use_mha_companion):
+                    output = torch.empty_like(query)
+                    lse = radix_attention_module._unified_attention_with_output_impl(
+                        query,
+                        query,
+                        query,
+                        output,
+                        False,
+                        0,
+                        use_mha_companion,
+                        True,
+                    )
+
+                    call_record = backend.calls[-1]
+                    self.assertIs(call_record.attention_layer, expected_layer)
+                    self.assertEqual(call_record.query.shape, (2, 2, 3))
+                    self.assertEqual(call_record.key.shape, (2, 2, 3))
+                    self.assertEqual(call_record.value.shape, (2, 2, 3))
+                    self.assertEqual(call_record.output.shape, (2, 2, 3))
+                    self.assertEqual(call_record.out_cache_loc.tolist(), [0, 1])
+                    self.assertFalse(call_record.save_kv_cache)
+                    self.assertTrue(torch.all(output[:2] == 3))
+                    self.assertEqual(lse.shape, (4, 2))
+                    self.assertTrue(torch.all(lse[:2] == 7))
+                    self.assertTrue(torch.all(lse[2:] == 0))
+                    self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc)
+
+    def test_impl_uses_independent_query_and_key_value_extents(self):
+        attention_layer = SimpleNamespace()
+        context = self._new_impl_context([attention_layer])
+        forward_batch = context.forward_batch
+        original_out_cache_loc = forward_batch.out_cache_loc
+        backend = _RecordingAttentionBackend()
+        query = torch.zeros((4, 2, 3))
+        key = torch.zeros((6, 2, 3))
+        value = torch.zeros((6, 2, 3))
+        k_rope = torch.zeros((6, 2, 1))
+        output = torch.empty_like(query)
+
+        with (
+            patch.object(
+                radix_attention_module,
+                "get_tc_piecewise_forward_context",
+                return_value=context,
+            ),
+            patch.object(
+                radix_attention_module, "get_attn_backend", return_value=backend
+            ),
+        ):
+            lse = radix_attention_module._unified_attention_with_output_impl(
+                query,
+                key,
+                value,
+                output,
+                False,
+                0,
+                False,
+                True,
+                key_value_num_tokens=5,
+                k_rope=k_rope,
+            )
+
+        call_record = backend.calls[-1]
+        self.assertEqual(call_record.query.shape, (2, 2, 3))
+        self.assertEqual(call_record.key.shape, (5, 2, 3))
+        self.assertEqual(call_record.value.shape, (5, 2, 3))
+        self.assertEqual(call_record.kwargs["k_rope"].shape, (5, 2, 1))
+        self.assertEqual(call_record.output.shape, (2, 2, 3))
+        self.assertEqual(lse.shape, (4, 2))
+        self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc)
+
+    def test_impl_preserves_output_only_contract(self):
+        attention_layer = SimpleNamespace()
+        context = self._new_impl_context([attention_layer])
+        forward_batch = context.forward_batch
+        original_out_cache_loc = forward_batch.out_cache_loc
+        backend = _RecordingAttentionBackend(return_lse=False)
+        query = torch.zeros((4, 2, 3))
+        output = torch.empty_like(query)
+
+        with (
+            patch.object(
+                radix_attention_module,
+                "get_tc_piecewise_forward_context",
+                return_value=context,
+            ),
+            patch.object(
+                radix_attention_module, "get_attn_backend", return_value=backend
+            ),
+        ):
+            lse = radix_attention_module._unified_attention_with_output_impl(
+                query,
+                query,
+                query,
+                output,
+                False,
+                0,
+                False,
+                False,
+            )
+
+        self.assertIsNone(lse)
+        self.assertIs(backend.calls[-1].attention_layer, attention_layer)
+        self.assertTrue(torch.all(output[:2] == 3))
+        self.assertIs(forward_batch.out_cache_loc, original_out_cache_loc)
+
+    def test_lse_fake_impl_declares_shape_and_dtype(self):
+        query = torch.empty((5, 3, 7), dtype=torch.float16)
+        output = torch.empty_like(query)
+
+        lse = radix_attention_module._unified_attention_with_output_and_lse_fake(
+            query,
+            None,
+            None,
+            output,
+            False,
+            0,
+        )
+
+        self.assertEqual(lse.shape, (5, 3))
+        self.assertEqual(lse.dtype, torch.float32)
+        self.assertEqual(lse.device, query.device)
+
+
+if __name__ == "__main__":
+    unittest.main()