[FullCG] Preserve attention LSE through the custom-op boundary (#31050)
This commit is contained in:
paulzhang-tm
@@ -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()
|
||||
Reference in New Issue
Block a user