diff --git a/python/sglang/srt/layers/radix_attention.py b/python/sglang/srt/layers/radix_attention.py index 2f69e6387..4bcfa541b 100644 --- a/python/sglang/srt/layers/radix_attention.py +++ b/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( diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index 214244a1a..31f7e73d0 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/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) diff --git a/test/registered/unit/layers/test_radix_attention.py b/test/registered/unit/layers/test_radix_attention.py new file mode 100644 index 000000000..b30bc2f85 --- /dev/null +++ b/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()