move dead sglang.test files to test/manual (#25316)
This commit is contained in:
Liangsheng Yin
@@ -0,0 +1,559 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.configs.model_config import AttentionArch
|
||||
from sglang.srt.layers.attention.flashattention_backend import (
|
||||
FlashAttentionBackend,
|
||||
draft_decode_set_expand_metadata,
|
||||
)
|
||||
from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
|
||||
from sglang.srt.layers.radix_attention import RadixAttention
|
||||
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class MockModelRunner:
|
||||
def __init__(
|
||||
self,
|
||||
page_size=1,
|
||||
num_heads=2,
|
||||
head_dim=8,
|
||||
):
|
||||
self.device = "cuda"
|
||||
self.dtype = torch.float16
|
||||
self.kv_cache_dtype = torch.float16
|
||||
self.is_hybrid_swa = False
|
||||
self.attention_chunk_size = None
|
||||
attention_arch = AttentionArch.MHA
|
||||
# Max batch size for the test.
|
||||
max_batch_size = 160
|
||||
# Total tokens(prefix + extend + decode) in the test should not exceed this length.
|
||||
max_context_len = 2048
|
||||
self.model_config = type(
|
||||
"ModelConfig",
|
||||
(),
|
||||
{
|
||||
"context_len": max_context_len,
|
||||
"is_multimodal": False,
|
||||
"attention_arch": attention_arch,
|
||||
"is_encoder_decoder": False,
|
||||
"is_local_attention_model": False,
|
||||
},
|
||||
)()
|
||||
self.sliding_window_size = None
|
||||
self.kv_cache_dtype = (
|
||||
self.dtype
|
||||
) # torch dtype, required by FlashAttentionBackend
|
||||
|
||||
# server_args is still needed for string-based config (kv_cache_dtype_str)
|
||||
self.server_args = type(
|
||||
"ServerArgs",
|
||||
(),
|
||||
{
|
||||
"kv_cache_dtype": "auto", # string version for kv_cache_dtype_str
|
||||
"speculative_eagle_topk": None,
|
||||
"speculative_num_draft_tokens": 0,
|
||||
"enable_deterministic_inference": False,
|
||||
},
|
||||
)
|
||||
self.attn_cp_size = 1
|
||||
# Create a large enough req_to_token_pool to fit the test usage.
|
||||
self.req_to_token_pool = type(
|
||||
"TokenPool",
|
||||
(),
|
||||
{
|
||||
# A typical max_bs * max_context_len for cuda graph decode
|
||||
"size": max_batch_size,
|
||||
# Add req_to_token attribute
|
||||
"req_to_token": torch.zeros(
|
||||
max_batch_size,
|
||||
max_context_len,
|
||||
dtype=torch.int32,
|
||||
device=self.device,
|
||||
),
|
||||
},
|
||||
)()
|
||||
self.page_size = page_size
|
||||
max_total_num_tokens = max_batch_size * max_context_len
|
||||
self.token_to_kv_pool = MHATokenToKVPool(
|
||||
size=max_total_num_tokens,
|
||||
page_size=page_size,
|
||||
dtype=self.dtype,
|
||||
head_num=num_heads,
|
||||
head_dim=head_dim,
|
||||
layer_num=1, # only consider layer=1 for unit test
|
||||
device=self.device,
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
|
||||
class TestFlashAttentionBackend(CustomTestCase):
|
||||
def setUp(self):
|
||||
# Test parameters
|
||||
self.batch_size = 2
|
||||
self.seq_len = 256
|
||||
self.num_heads = 2
|
||||
self.head_dim = 8
|
||||
self.device = "cuda"
|
||||
self.dtype = torch.float16
|
||||
|
||||
def _init_model_runner(self, page_size=1):
|
||||
self.model_runner = MockModelRunner(
|
||||
page_size=page_size,
|
||||
num_heads=self.num_heads,
|
||||
head_dim=self.head_dim,
|
||||
)
|
||||
self.backend = FlashAttentionBackend(self.model_runner)
|
||||
self.ref_backend = TorchNativeAttnBackend(self.model_runner)
|
||||
self.model_runner.model_config.num_attention_heads = self.num_heads
|
||||
|
||||
def _mock_write_to_req_to_token_pool(self, batch_size, seq_len, page_size):
|
||||
# if page_size > 1, the token pool stores the index to the page.
|
||||
# so we need to multiply the index by page_size.
|
||||
self.req_to_token = (
|
||||
torch.arange(0, batch_size, dtype=torch.int32, device=self.device)[:, None]
|
||||
* seq_len
|
||||
+ torch.arange(0, seq_len, dtype=torch.int32, device=self.device)[None, :]
|
||||
+ page_size
|
||||
)
|
||||
self.model_runner.req_to_token_pool.req_to_token[:batch_size, :seq_len] = (
|
||||
self.req_to_token
|
||||
)
|
||||
|
||||
def _create_attention_layer(self):
|
||||
"""Create attention layer for testing."""
|
||||
return RadixAttention(
|
||||
num_heads=self.num_heads,
|
||||
head_dim=self.head_dim,
|
||||
scaling=1.0,
|
||||
num_kv_heads=self.num_heads,
|
||||
layer_id=0,
|
||||
)
|
||||
|
||||
def _create_qkv_tensors(self, tokens_len):
|
||||
"""Create q, k, v tensors for testing."""
|
||||
shape = (tokens_len, self.num_heads, self.head_dim)
|
||||
return (
|
||||
torch.randn(shape, dtype=self.dtype, device=self.device),
|
||||
torch.randn(shape, dtype=self.dtype, device=self.device),
|
||||
torch.randn(shape, dtype=self.dtype, device=self.device),
|
||||
)
|
||||
|
||||
def _run_reference_forward(
|
||||
self, mode, q, k, v, layer, forward_batch, expected_shape
|
||||
):
|
||||
"""Run reference forward pass using native backend."""
|
||||
if mode == ForwardMode.EXTEND:
|
||||
output = self.ref_backend.forward_extend(q, k, v, layer, forward_batch)
|
||||
else: # ForwardMode.DECODE
|
||||
output = self.ref_backend.forward_decode(q, k, v, layer, forward_batch)
|
||||
return output.view(expected_shape)
|
||||
|
||||
def _verify_output(self, output, expected_shape, output_ref=None):
|
||||
"""Verify output tensor shape, dtype, and values."""
|
||||
self.assertEqual(
|
||||
output.shape,
|
||||
expected_shape,
|
||||
f"Expected shape {expected_shape}, got {output.shape}",
|
||||
)
|
||||
self.assertEqual(output.dtype, self.dtype)
|
||||
self.assertEqual(output.device.type, "cuda")
|
||||
self.assertEqual(
|
||||
torch.isnan(output).sum().item(), 0, "Output contains NaN values"
|
||||
)
|
||||
|
||||
if output_ref is not None:
|
||||
if not torch.allclose(output, output_ref, atol=1e-1, rtol=0.0):
|
||||
# Check where the values differ beyond the given tolerances
|
||||
diff_mask = ~torch.isclose(output, output_ref, atol=1e-1, rtol=0.0)
|
||||
|
||||
# Find the first index where the difference occurs
|
||||
if diff_mask.any():
|
||||
first_mismatch_idx = diff_mask.nonzero()[0]
|
||||
print(
|
||||
"First mismatch at index:", tuple(first_mismatch_idx.tolist())
|
||||
)
|
||||
print("output:", output[tuple(first_mismatch_idx.tolist())])
|
||||
print("output_ref:", output_ref[tuple(first_mismatch_idx.tolist())])
|
||||
raise AssertionError(
|
||||
"Attention output is not close to the torch native backend output"
|
||||
)
|
||||
|
||||
def _create_forward_batch(
|
||||
self, mode, q_len=None, prefix_len=0, page_size=1, attn_cp_size=1
|
||||
):
|
||||
"""Create a forward batch for testing based on mode and lengths."""
|
||||
self._init_model_runner(page_size=page_size)
|
||||
|
||||
# Default to self.seq_len if not specified
|
||||
q_len = q_len or self.seq_len
|
||||
|
||||
if mode == ForwardMode.EXTEND:
|
||||
total_len = prefix_len + q_len
|
||||
out_cache_start = prefix_len * self.batch_size
|
||||
out_cache_end = total_len * self.batch_size
|
||||
|
||||
forward_batch = ForwardBatch(
|
||||
batch_size=self.batch_size,
|
||||
input_ids=torch.randint(
|
||||
0, 100, (self.batch_size, q_len), device=self.device
|
||||
),
|
||||
out_cache_loc=torch.arange(
|
||||
out_cache_start, out_cache_end, device=self.device
|
||||
),
|
||||
seq_lens_sum=self.batch_size * total_len,
|
||||
forward_mode=mode,
|
||||
req_pool_indices=torch.arange(self.batch_size, device=self.device),
|
||||
seq_lens=torch.tensor(
|
||||
[total_len] * self.batch_size, device=self.device
|
||||
),
|
||||
seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
|
||||
extend_prefix_lens=torch.tensor(
|
||||
[prefix_len] * self.batch_size, device=self.device
|
||||
),
|
||||
extend_prefix_lens_cpu=torch.tensor(
|
||||
[prefix_len] * self.batch_size, device="cpu"
|
||||
),
|
||||
extend_seq_lens=torch.tensor(
|
||||
[q_len] * self.batch_size, device=self.device
|
||||
),
|
||||
extend_seq_lens_cpu=torch.tensor(
|
||||
[q_len] * self.batch_size, device="cpu"
|
||||
),
|
||||
attn_backend=self.backend,
|
||||
)
|
||||
if attn_cp_size > 1:
|
||||
forward_batch.attn_cp_metadata = type(
|
||||
"AttnCPMetadata",
|
||||
(),
|
||||
{
|
||||
"kv_len_prev_tensor": torch.tensor(
|
||||
[q_len // 2] * self.batch_size,
|
||||
dtype=torch.int32,
|
||||
device=self.device,
|
||||
),
|
||||
"kv_len_next_tensor": torch.tensor(
|
||||
[q_len] * self.batch_size,
|
||||
dtype=torch.int32,
|
||||
device=self.device,
|
||||
),
|
||||
"actual_seq_q_prev": q_len // 2,
|
||||
"actual_seq_q_next": q_len // 2,
|
||||
},
|
||||
)
|
||||
else: # ForwardMode.DECODE
|
||||
decode_len = q_len # Assuming 1 for decode testing
|
||||
total_len = self.seq_len + decode_len
|
||||
if mode == ForwardMode.DECODE and page_size > 1:
|
||||
# Get next page_size multiple of self.seq_len
|
||||
out_cache_start = (
|
||||
self.batch_size * self.seq_len // page_size + 1
|
||||
) * page_size
|
||||
# out_cache_end is the start of the next block
|
||||
out_cache_end = out_cache_start + decode_len * page_size
|
||||
else:
|
||||
out_cache_start = self.batch_size * self.seq_len
|
||||
out_cache_end = self.batch_size * total_len
|
||||
|
||||
forward_batch = ForwardBatch(
|
||||
batch_size=self.batch_size,
|
||||
input_ids=torch.randint(
|
||||
0, 100, (self.batch_size, decode_len), device=self.device
|
||||
),
|
||||
out_cache_loc=torch.tensor(
|
||||
[out_cache_start, out_cache_end], device=self.device
|
||||
),
|
||||
seq_lens_sum=self.batch_size * total_len,
|
||||
forward_mode=mode,
|
||||
req_pool_indices=torch.arange(self.batch_size, device=self.device),
|
||||
seq_lens=torch.tensor(
|
||||
[total_len] * self.batch_size, device=self.device
|
||||
),
|
||||
seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
|
||||
attn_backend=self.backend,
|
||||
)
|
||||
|
||||
# Add token pool
|
||||
forward_batch.req_to_token_pool = self.model_runner.req_to_token_pool
|
||||
|
||||
# Write current batch's req_to_token to req_to_token_pool
|
||||
self._mock_write_to_req_to_token_pool(self.batch_size, total_len, page_size)
|
||||
# Add kv pool for this forward batch
|
||||
forward_batch.token_to_kv_pool = self.model_runner.token_to_kv_pool
|
||||
|
||||
return forward_batch
|
||||
|
||||
def _setup_kv_cache(self, forward_batch, layer, cache_len):
|
||||
# Create constant values for the prefix cache for easy debugging
|
||||
cache_k = torch.ones(
|
||||
self.batch_size * cache_len,
|
||||
self.num_heads,
|
||||
self.head_dim,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
cache_v = (
|
||||
torch.ones(
|
||||
self.batch_size * cache_len,
|
||||
self.num_heads,
|
||||
self.head_dim,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
* 2
|
||||
)
|
||||
|
||||
# Set the prefix KV cache
|
||||
forward_batch.token_to_kv_pool.set_kv_buffer(
|
||||
layer,
|
||||
torch.arange(self.batch_size * cache_len, device=self.device),
|
||||
cache_k,
|
||||
cache_v,
|
||||
layer.k_scale,
|
||||
layer.v_scale,
|
||||
)
|
||||
|
||||
def _run_attention_test(self, mode, q_len, prefix_len=0, page_size=1):
|
||||
"""
|
||||
Run an attention test with the specified parameters.
|
||||
Args:
|
||||
mode: ForwardMode.EXTEND or ForwardMode.DECODE
|
||||
q_len: Length of the query sequence. For decode mode, q_len is 1.
|
||||
prefix_len: Length of the prefix sequence for extend mode
|
||||
page_size: Page size for the KV cache
|
||||
"""
|
||||
layer = self._create_attention_layer()
|
||||
|
||||
# Create forward batch and set up
|
||||
forward_batch = self._create_forward_batch(mode, q_len, prefix_len, page_size)
|
||||
|
||||
# Create QKV tensors for the input
|
||||
q, k, v = self._create_qkv_tensors(self.batch_size * q_len)
|
||||
|
||||
# KV cache for prefixed extend is prefix_len
|
||||
# KV cache for decode is same as seq_len
|
||||
# No KV cache for extend without prefix
|
||||
if mode == ForwardMode.EXTEND:
|
||||
if prefix_len > 0:
|
||||
self._setup_kv_cache(forward_batch, layer, prefix_len)
|
||||
else:
|
||||
self._setup_kv_cache(forward_batch, layer, self.seq_len)
|
||||
|
||||
self.backend.init_forward_metadata(forward_batch)
|
||||
|
||||
if mode == ForwardMode.EXTEND:
|
||||
expected_shape = (
|
||||
self.batch_size * q_len,
|
||||
self.num_heads * self.head_dim,
|
||||
)
|
||||
output = self.backend.forward_extend(q, k, v, layer, forward_batch)
|
||||
else:
|
||||
expected_shape = (self.batch_size, self.num_heads * self.head_dim)
|
||||
output = self.backend.forward_decode(q, k, v, layer, forward_batch)
|
||||
|
||||
output_ref = self._run_reference_forward(
|
||||
mode, q, k, v, layer, forward_batch, expected_shape
|
||||
)
|
||||
|
||||
self._verify_output(output, expected_shape, output_ref)
|
||||
|
||||
return output
|
||||
|
||||
def _run_attention_cp_test(self, mode, q_len, prefix_len=0, page_size=1):
|
||||
layer = self._create_attention_layer()
|
||||
|
||||
# Create forward batch and set up
|
||||
forward_batch = self._create_forward_batch(
|
||||
mode, q_len, prefix_len, page_size, attn_cp_size=2
|
||||
)
|
||||
self.backend.attn_cp_size = 2
|
||||
|
||||
# Create QKV tensors for the input
|
||||
q, k, v = self._create_qkv_tensors(self.batch_size * q_len)
|
||||
|
||||
# KV cache for prefixed extend is prefix_len
|
||||
# KV cache for decode is same as seq_len
|
||||
# No KV cache for extend without prefix
|
||||
# Setup KV cache for CP testing - need KV cache to have actual values
|
||||
# For extend with CP, we need KV cache populated so attention has something to attend to
|
||||
self._setup_kv_cache(forward_batch, layer, q_len)
|
||||
|
||||
self.backend.init_forward_metadata(forward_batch)
|
||||
|
||||
# if mode == ForwardMode.EXTEND:
|
||||
expected_shape = (
|
||||
self.batch_size * q_len,
|
||||
self.num_heads * self.head_dim,
|
||||
)
|
||||
|
||||
output = self.backend.forward_extend(q, k, v, layer, forward_batch)
|
||||
# else:
|
||||
# expected_shape = (self.batch_size, self.num_heads * self.head_dim)
|
||||
# output = self.backend.forward_decode(q, k, v, layer, forward_batch)
|
||||
|
||||
output_ref = self._run_reference_forward(
|
||||
mode, q, k, v, layer, forward_batch, expected_shape
|
||||
)
|
||||
|
||||
self._verify_output(output, expected_shape, output_ref)
|
||||
|
||||
return output
|
||||
|
||||
def test_forward_extend_cp(self):
|
||||
"""Test the standard extend operation with context parallel."""
|
||||
self._run_attention_cp_test(ForwardMode.EXTEND, q_len=self.seq_len)
|
||||
|
||||
# def test_forward_extend_cp_with_prefix(self):
|
||||
# """Test the standard extend operation with context parallel and prefix."""
|
||||
# prefix_len = self.seq_len // 2
|
||||
# extend_len = self.seq_len - prefix_len
|
||||
# self._run_attention_cp_test(ForwardMode.EXTEND, q_len=extend_len, prefix_len=prefix_len)
|
||||
|
||||
# def test_forward_extend(self):
|
||||
# """Test the standard extend operation."""
|
||||
# self._run_attention_test(ForwardMode.EXTEND, q_len=self.seq_len)
|
||||
|
||||
# def test_forward_decode(self):
|
||||
# """Test the decode operation with cached tokens."""
|
||||
# self._run_attention_test(ForwardMode.DECODE, q_len=1)
|
||||
|
||||
# def test_forward_extend_with_prefix(self):
|
||||
# """Test extending from cached prefix tokens."""
|
||||
# prefix_len = self.seq_len // 2
|
||||
# extend_len = self.seq_len - prefix_len
|
||||
# self._run_attention_test(
|
||||
# ForwardMode.EXTEND, q_len=extend_len, prefix_len=prefix_len
|
||||
# )
|
||||
|
||||
# def test_forward_extend_with_page_size_greater_than_1(self):
|
||||
# """Test extending from cached prefix tokens with page size greater than 1."""
|
||||
# self._run_attention_test(ForwardMode.EXTEND, q_len=self.seq_len, page_size=64)
|
||||
|
||||
# def test_forward_decode_with_page_size_greater_than_1(self):
|
||||
# """Test decode operation with page size greater than 1."""
|
||||
# self._run_attention_test(ForwardMode.DECODE, q_len=1, page_size=64)
|
||||
|
||||
|
||||
class TestUpdateDraftDecodeSetExpandMetadata(CustomTestCase):
|
||||
"""
|
||||
All the test cases examples have 1 additional cache location than the decode length.
|
||||
This is to align with the current allocation logic. It does not affect the correctness.
|
||||
"""
|
||||
|
||||
def test_draft_decode_set_expand_metadata(self):
|
||||
bs, topk, page_size = 1, 2, 4
|
||||
|
||||
cases = [
|
||||
(
|
||||
torch.tensor(
|
||||
[
|
||||
[23, 24],
|
||||
[31, 32],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
),
|
||||
torch.tensor(
|
||||
[
|
||||
[5, 6],
|
||||
[7, 8],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
),
|
||||
1,
|
||||
),
|
||||
# Decode span multiple pages:
|
||||
# duplicated kv cache: 24, 25, 26
|
||||
# decode locations: 27, 28, 29, 30, 31, 32
|
||||
# We need 3 pages in total.
|
||||
(
|
||||
torch.tensor(
|
||||
[
|
||||
[27, 28, 29, 30, 31, 32],
|
||||
[35, 36, 37, 38, 39, 40],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
),
|
||||
torch.tensor(
|
||||
[
|
||||
[6, 7, 8, 0, 0, 0],
|
||||
[8, 9, 10, 0, 0, 0],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
),
|
||||
5,
|
||||
),
|
||||
]
|
||||
|
||||
last_page_lens = torch.tensor([3], dtype=torch.int32)
|
||||
for cache_loc, expected_page_table, decode_length in cases:
|
||||
cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32)
|
||||
page_table = torch.zeros_like(cache_loc, dtype=torch.int32)
|
||||
draft_decode_set_expand_metadata(
|
||||
cache_seqlens_int32=cache_seqlens_int32,
|
||||
page_table=page_table,
|
||||
last_page_lens=last_page_lens,
|
||||
decode_length=decode_length,
|
||||
cache_loc=cache_loc,
|
||||
topk=topk,
|
||||
page_size=page_size,
|
||||
)
|
||||
|
||||
expected_cache_seqlens = torch.tensor(
|
||||
[decode_length + 3, decode_length + 3], dtype=torch.int32
|
||||
)
|
||||
self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens))
|
||||
self.assertTrue(torch.equal(page_table, expected_page_table))
|
||||
|
||||
def test_update_draft_decode_set_expand_metadata_multi_batch(self):
|
||||
"""
|
||||
Ensure expand metadata works when batch size > 1 and last pages differ.
|
||||
"""
|
||||
bs, topk, decode_length, page_size = 3, 2, 3, 4
|
||||
cache_loc = torch.tensor(
|
||||
[
|
||||
# First batch: last page duplicate is 1, consecutive pages
|
||||
[1, 2, 3, 4],
|
||||
[6, 7, 8, 9],
|
||||
# Second batch: last page duplicate is 3, non-consecutive pages
|
||||
[3, 8, 9, 10],
|
||||
[14, 15, 16, 17],
|
||||
# Third batch: last page duplicate is 0, consecutive pages
|
||||
[0, 1, 2, 3],
|
||||
[4, 5, 6, 7],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
)
|
||||
cache_seqlens_int32 = torch.zeros(bs * topk, dtype=torch.int32)
|
||||
last_page_lens = torch.tensor([1, 3, 0], dtype=torch.int32)
|
||||
page_table = torch.zeros_like(cache_loc, dtype=torch.int32)
|
||||
draft_decode_set_expand_metadata(
|
||||
cache_seqlens_int32=cache_seqlens_int32,
|
||||
page_table=page_table,
|
||||
last_page_lens=last_page_lens,
|
||||
decode_length=decode_length,
|
||||
cache_loc=cache_loc,
|
||||
topk=topk,
|
||||
page_size=page_size,
|
||||
)
|
||||
|
||||
expected_cache_seqlens = torch.tensor([4, 4, 6, 6, 3, 3], dtype=torch.int32)
|
||||
expected_page_table = torch.tensor(
|
||||
[
|
||||
[0, 1, 0, 0],
|
||||
[1, 2, 0, 0],
|
||||
[0, 2, 0, 0],
|
||||
[3, 4, 0, 0],
|
||||
[0, 0, 0, 0],
|
||||
[1, 0, 0, 0],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
)
|
||||
self.assertTrue(torch.equal(cache_seqlens_int32, expected_cache_seqlens))
|
||||
self.assertTrue(torch.equal(page_table, expected_page_table))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,331 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.configs.model_config import AttentionArch
|
||||
from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
|
||||
from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
|
||||
from sglang.srt.layers.radix_attention import RadixAttention
|
||||
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class MockModelRunner:
|
||||
def __init__(
|
||||
self,
|
||||
kv_lora_rank,
|
||||
qk_rope_head_dim,
|
||||
):
|
||||
attention_arch = AttentionArch.MLA
|
||||
self.device = "cuda"
|
||||
self.dtype = torch.float16
|
||||
self.is_hybrid_swa = False
|
||||
context_len = 2048
|
||||
self.model_config = type(
|
||||
"ModelConfig",
|
||||
(),
|
||||
{
|
||||
"context_len": context_len,
|
||||
"attention_arch": attention_arch,
|
||||
"is_encoder_decoder": False,
|
||||
"is_local_attention_model": False,
|
||||
},
|
||||
)
|
||||
self.sliding_window_size = None
|
||||
# Add server_args attribute
|
||||
self.server_args = type(
|
||||
"ServerArgs",
|
||||
(),
|
||||
{
|
||||
"kv_cache_dtype": torch.float16,
|
||||
"speculative_eagle_topk": None,
|
||||
"speculative_num_draft_tokens": 0,
|
||||
"enable_deterministic_inference": False,
|
||||
},
|
||||
)
|
||||
self.kv_cache_dtype = self.server_args.kv_cache_dtype
|
||||
|
||||
batch_size = 160
|
||||
# Create a proper req_to_token_pool with the req_to_token attribute
|
||||
self.req_to_token_pool = type(
|
||||
"TokenPool",
|
||||
(),
|
||||
{
|
||||
# A typical max_bs * max_context_len for cuda graph decode
|
||||
"size": batch_size,
|
||||
# Add req_to_token attribute
|
||||
"req_to_token": torch.zeros(
|
||||
batch_size, context_len, dtype=torch.int32, device=self.device
|
||||
),
|
||||
},
|
||||
)
|
||||
self.page_size = 1
|
||||
max_total_num_tokens = batch_size * context_len
|
||||
self.token_to_kv_pool = MLATokenToKVPool(
|
||||
size=max_total_num_tokens,
|
||||
page_size=self.page_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
kv_lora_rank=kv_lora_rank,
|
||||
qk_rope_head_dim=qk_rope_head_dim,
|
||||
layer_num=1, # only consider layer=1 for unit test
|
||||
device=self.device,
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
|
||||
|
||||
class MockReqToTokenPool:
|
||||
def __init__(self, batch_size, seq_len, device):
|
||||
self.req_to_token = (
|
||||
torch.arange(batch_size * seq_len, device=device)
|
||||
.reshape(batch_size, seq_len)
|
||||
.to(torch.int32)
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
|
||||
class TestFlashAttentionMLABackend(CustomTestCase):
|
||||
def setUp(self):
|
||||
# MLA with different V headdim requires Hopper architecture (compute capability >= 9.0)
|
||||
if torch.cuda.is_available():
|
||||
compute_capability = torch.cuda.get_device_capability()
|
||||
if compute_capability[0] < 9:
|
||||
self.skipTest(
|
||||
f"MLA requires Hopper GPU (compute capability >= 9.0), "
|
||||
f"but found compute capability {compute_capability[0]}.{compute_capability[1]}"
|
||||
)
|
||||
|
||||
# Test parameters
|
||||
self.batch_size = 2
|
||||
self.seq_len = 360
|
||||
self.num_heads = 2
|
||||
self.device = "cuda"
|
||||
self.dtype = torch.float16
|
||||
self.kv_lora_rank = 512
|
||||
self.q_lora_rank = 128
|
||||
self.qk_rope_head_dim = 64
|
||||
self.qk_head_dim = self.qk_rope_head_dim + self.kv_lora_rank
|
||||
# Assume no rope scaling
|
||||
self.scaling = self.qk_head_dim**-0.5
|
||||
# Initialize model runner and backend
|
||||
self._init_model_runner()
|
||||
self.backend = FlashAttentionBackend(self.model_runner)
|
||||
self.ref_backend = TorchNativeAttnBackend(self.model_runner)
|
||||
self.num_local_heads = 2
|
||||
|
||||
def _init_model_runner(self):
|
||||
self.model_runner = MockModelRunner(
|
||||
kv_lora_rank=self.kv_lora_rank,
|
||||
qk_rope_head_dim=self.qk_rope_head_dim,
|
||||
)
|
||||
|
||||
def _create_attention_layer(self):
|
||||
"""Create attention layer for testing."""
|
||||
self.attn_mqa = RadixAttention(
|
||||
num_heads=self.num_local_heads,
|
||||
head_dim=self.kv_lora_rank + self.qk_rope_head_dim,
|
||||
scaling=self.scaling,
|
||||
num_kv_heads=1,
|
||||
layer_id=0,
|
||||
v_head_dim=self.kv_lora_rank,
|
||||
prefix="attn_mqa",
|
||||
)
|
||||
return self.attn_mqa
|
||||
|
||||
def _run_reference_forward(
|
||||
self, mode, q, k, v, layer, forward_batch, expected_shape
|
||||
):
|
||||
"""Run reference forward pass using native backend."""
|
||||
if mode == ForwardMode.EXTEND:
|
||||
output = self.ref_backend.forward_extend(q, k, v, layer, forward_batch)
|
||||
else: # ForwardMode.DECODE
|
||||
output = self.ref_backend.forward_decode(q, k, v, layer, forward_batch)
|
||||
return output.view(expected_shape)
|
||||
|
||||
def _verify_output(self, output, expected_shape):
|
||||
"""Verify output tensor shape, dtype, and values."""
|
||||
self.assertEqual(
|
||||
output.shape,
|
||||
expected_shape,
|
||||
f"Expected shape {expected_shape}, got {output.shape}",
|
||||
)
|
||||
self.assertEqual(output.dtype, self.dtype)
|
||||
self.assertEqual(output.device.type, "cuda")
|
||||
self.assertEqual(
|
||||
torch.isnan(output).sum().item(), 0, "Output contains NaN values"
|
||||
)
|
||||
|
||||
def _create_forward_batch(self, mode, q_len=None, prefix_len=0):
|
||||
"""Create a forward batch for testing based on mode and lengths."""
|
||||
# Default to self.seq_len if not specified
|
||||
q_len = q_len or self.seq_len
|
||||
|
||||
if mode == ForwardMode.EXTEND:
|
||||
total_len = prefix_len + q_len
|
||||
out_cache_start = prefix_len * self.batch_size
|
||||
out_cache_end = total_len * self.batch_size
|
||||
|
||||
forward_batch = ForwardBatch(
|
||||
batch_size=self.batch_size,
|
||||
input_ids=torch.randint(
|
||||
0, 100, (self.batch_size, q_len), device=self.device
|
||||
),
|
||||
out_cache_loc=torch.arange(
|
||||
out_cache_start, out_cache_end, device=self.device
|
||||
),
|
||||
seq_lens_sum=self.batch_size * total_len,
|
||||
forward_mode=mode,
|
||||
req_pool_indices=torch.arange(self.batch_size, device=self.device),
|
||||
seq_lens=torch.tensor(
|
||||
[total_len] * self.batch_size, device=self.device
|
||||
),
|
||||
seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
|
||||
extend_prefix_lens=torch.tensor(
|
||||
[prefix_len] * self.batch_size, device=self.device
|
||||
),
|
||||
extend_prefix_lens_cpu=torch.tensor(
|
||||
[prefix_len] * self.batch_size, device="cpu"
|
||||
),
|
||||
extend_seq_lens=torch.tensor(
|
||||
[q_len] * self.batch_size, device=self.device
|
||||
),
|
||||
extend_seq_lens_cpu=torch.tensor(
|
||||
[q_len] * self.batch_size, device="cpu"
|
||||
),
|
||||
attn_backend=self.backend,
|
||||
)
|
||||
|
||||
else: # ForwardMode.DECODE
|
||||
decode_len = q_len # typically 1 for decode mode
|
||||
total_len = self.seq_len + decode_len
|
||||
out_cache_start = self.batch_size * self.seq_len
|
||||
out_cache_end = self.batch_size * total_len
|
||||
|
||||
forward_batch = ForwardBatch(
|
||||
batch_size=self.batch_size,
|
||||
input_ids=torch.randint(
|
||||
0, 100, (self.batch_size, decode_len), device=self.device
|
||||
),
|
||||
out_cache_loc=torch.arange(
|
||||
out_cache_start, out_cache_end, device=self.device
|
||||
),
|
||||
seq_lens_sum=self.batch_size * total_len,
|
||||
forward_mode=mode,
|
||||
req_pool_indices=torch.arange(self.batch_size, device=self.device),
|
||||
seq_lens=torch.tensor(
|
||||
[total_len] * self.batch_size, device=self.device
|
||||
),
|
||||
seq_lens_cpu=torch.tensor([total_len] * self.batch_size, device="cpu"),
|
||||
attn_backend=self.backend,
|
||||
)
|
||||
|
||||
# Add token pool from model runner to forward batch
|
||||
forward_batch.req_to_token_pool = self.model_runner.req_to_token_pool
|
||||
|
||||
# Add KV cache from model runner to forward batch
|
||||
forward_batch.token_to_kv_pool = self.model_runner.token_to_kv_pool
|
||||
|
||||
return forward_batch
|
||||
|
||||
def _setup_kv_cache(self, forward_batch, layer, cache_len):
|
||||
"""Set up KV cache with prefix tokens."""
|
||||
if cache_len <= 0:
|
||||
return
|
||||
|
||||
# For MLA, create separate nope and rope caches
|
||||
cache_k_nope = torch.ones(
|
||||
self.batch_size * cache_len,
|
||||
1, # latent cache has only one head in MQA
|
||||
self.kv_lora_rank,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
cache_k_rope = torch.ones(
|
||||
self.batch_size * cache_len,
|
||||
1, # latent cache has only one head in MQA
|
||||
self.qk_rope_head_dim,
|
||||
dtype=self.dtype,
|
||||
device=self.device,
|
||||
)
|
||||
|
||||
# Set the prefix KV cache using MLA-specific method
|
||||
forward_batch.token_to_kv_pool.set_mla_kv_buffer(
|
||||
layer,
|
||||
torch.arange(self.batch_size * cache_len, device=self.device),
|
||||
cache_k_nope,
|
||||
cache_k_rope,
|
||||
)
|
||||
|
||||
def _run_attention_test(self, mode, q_len, prefix_len=0):
|
||||
"""
|
||||
Run an attention test with the specified parameters.
|
||||
Args:
|
||||
mode: ForwardMode.EXTEND or ForwardMode.DECODE
|
||||
q_len: Length of the query sequence. For decode mode, q_len is 1.
|
||||
prefix_len: Length of the prefix sequence for extend mode
|
||||
"""
|
||||
layer = self._create_attention_layer()
|
||||
|
||||
# Create forward batch and set up
|
||||
forward_batch = self._create_forward_batch(mode, q_len, prefix_len)
|
||||
|
||||
# Create q, kv_compressed for testing
|
||||
q_shape = (self.batch_size * q_len, self.num_heads, self.qk_head_dim)
|
||||
kv_shape = (self.batch_size * q_len, self.qk_head_dim)
|
||||
q = torch.randn(q_shape, dtype=self.dtype, device=self.device)
|
||||
kv_compressed = torch.randn(kv_shape, dtype=self.dtype, device=self.device)
|
||||
|
||||
# For MLA, split kv_compressed into k_nope and k_rope
|
||||
# k_nope has dimension kv_lora_rank, k_rope has dimension qk_rope_head_dim
|
||||
k_nope = kv_compressed[:, : self.kv_lora_rank]
|
||||
k_rope = kv_compressed[:, self.kv_lora_rank :]
|
||||
|
||||
# k_nope needs to be unsqueezed for the num_heads dimension
|
||||
k = k_nope.unsqueeze(1)
|
||||
# k_rope also needs to be unsqueezed
|
||||
k_rope = k_rope.unsqueeze(1)
|
||||
|
||||
# v is not used for mqa
|
||||
v = torch.randn((1), dtype=self.dtype, device=self.device)
|
||||
|
||||
self._setup_kv_cache(forward_batch, layer, prefix_len)
|
||||
|
||||
self.backend.init_forward_metadata(forward_batch)
|
||||
|
||||
expected_shape = (
|
||||
self.batch_size * q_len,
|
||||
self.num_heads * self.kv_lora_rank,
|
||||
)
|
||||
|
||||
if mode == ForwardMode.EXTEND:
|
||||
output = self.backend.forward_extend(
|
||||
q, k, v, layer, forward_batch, k_rope=k_rope
|
||||
)
|
||||
else:
|
||||
output = self.backend.forward_decode(
|
||||
q, k, v, layer, forward_batch, k_rope=k_rope
|
||||
)
|
||||
|
||||
self._verify_output(output, expected_shape)
|
||||
return output
|
||||
|
||||
def test_forward_extend(self):
|
||||
"""Test the standard extend operation."""
|
||||
self._run_attention_test(ForwardMode.EXTEND, q_len=self.seq_len)
|
||||
|
||||
def test_forward_decode(self):
|
||||
"""Test the decode operation with cached tokens."""
|
||||
self._run_attention_test(ForwardMode.DECODE, q_len=1)
|
||||
|
||||
def test_forward_extend_with_prefix(self):
|
||||
"""Test extending from cached prefix tokens."""
|
||||
prefix_len = self.seq_len // 2
|
||||
extend_len = self.seq_len - prefix_len
|
||||
self._run_attention_test(
|
||||
ForwardMode.EXTEND, q_len=extend_len, prefix_len=prefix_len
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,228 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
||||
from sglang.srt.utils.common import get_device
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
TEST_CASES = [
|
||||
# Sequence with same prefix lens
|
||||
{
|
||||
"batch_size": 3,
|
||||
"prefix_lens": [64, 64, 64],
|
||||
"max_chunk_capacity": 48,
|
||||
"prefix_chunk_len": 16,
|
||||
"num_prefix_chunks": 4,
|
||||
"prefix_chunk_starts": torch.tensor(
|
||||
[
|
||||
[0, 0, 0],
|
||||
[16, 16, 16],
|
||||
[32, 32, 32],
|
||||
[48, 48, 48],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
),
|
||||
"prefix_chunk_seq_lens": torch.tensor(
|
||||
[
|
||||
[16, 16, 16],
|
||||
[16, 16, 16],
|
||||
[16, 16, 16],
|
||||
[16, 16, 16],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
),
|
||||
},
|
||||
# Sequence with different prefix lens
|
||||
{
|
||||
"batch_size": 4,
|
||||
"prefix_lens": [16, 32, 48, 64],
|
||||
"max_chunk_capacity": 64,
|
||||
"prefix_chunk_len": 16,
|
||||
"num_prefix_chunks": 4,
|
||||
"prefix_chunk_starts": torch.tensor(
|
||||
[
|
||||
[0, 0, 0, 0],
|
||||
[16, 16, 16, 16],
|
||||
[32, 32, 32, 32],
|
||||
[48, 48, 48, 48],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
),
|
||||
"prefix_chunk_seq_lens": torch.tensor(
|
||||
[
|
||||
[16, 16, 16, 16],
|
||||
[0, 16, 16, 16],
|
||||
[0, 0, 16, 16],
|
||||
[0, 0, 0, 16],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
),
|
||||
},
|
||||
# Sequence with irregular shapes
|
||||
{
|
||||
"batch_size": 2,
|
||||
"prefix_lens": [1, 64],
|
||||
"max_chunk_capacity": 31,
|
||||
"prefix_chunk_len": 15,
|
||||
"num_prefix_chunks": 5,
|
||||
"prefix_chunk_starts": torch.tensor(
|
||||
[
|
||||
[0, 0],
|
||||
[15, 15],
|
||||
[30, 30],
|
||||
[45, 45],
|
||||
[60, 60],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
),
|
||||
"prefix_chunk_seq_lens": torch.tensor(
|
||||
[
|
||||
[1, 15],
|
||||
[0, 15],
|
||||
[0, 15],
|
||||
[0, 15],
|
||||
[0, 4],
|
||||
],
|
||||
dtype=torch.int32,
|
||||
),
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
class MockForwardBatch(ForwardBatch):
|
||||
def __init__(self, max_chunk_capacity: int, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.max_chunk_capacity = max_chunk_capacity
|
||||
|
||||
def get_max_chunk_capacity(self):
|
||||
return self.max_chunk_capacity
|
||||
|
||||
|
||||
class MockReqToTokenPool:
|
||||
def __init__(self, batch_size, seq_len, device):
|
||||
self.req_to_token = (
|
||||
torch.arange(batch_size * seq_len, device=device)
|
||||
.reshape(batch_size, seq_len)
|
||||
.to(torch.int32)
|
||||
)
|
||||
|
||||
|
||||
# Test correctness of triton kernel for computing kv indices
|
||||
def check_kv_indices(forward_batch):
|
||||
for i in range(forward_batch.num_prefix_chunks):
|
||||
computed_kv_indices = forward_batch.prefix_chunk_kv_indices[i]
|
||||
req_to_token = forward_batch.req_to_token_pool.req_to_token[
|
||||
: forward_batch.batch_size, :
|
||||
]
|
||||
ref_kv_indices = torch.empty(
|
||||
forward_batch.prefix_chunk_num_tokens[i],
|
||||
dtype=torch.int32,
|
||||
device=computed_kv_indices.device,
|
||||
)
|
||||
running_ptr = 0
|
||||
for j in range(forward_batch.batch_size):
|
||||
seq_start = forward_batch.prefix_chunk_starts[i, j].item()
|
||||
seq_len = forward_batch.prefix_chunk_seq_lens[i, j].item()
|
||||
ref_kv_indices[running_ptr : running_ptr + seq_len].copy_(
|
||||
req_to_token[j, seq_start : seq_start + seq_len]
|
||||
)
|
||||
running_ptr += seq_len
|
||||
assert torch.allclose(computed_kv_indices, ref_kv_indices)
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
not (torch.cuda.is_available() or torch.xpu.is_available()),
|
||||
"Test requires CUDA or XPU",
|
||||
)
|
||||
class TestPrefixChunkInfo(CustomTestCase):
|
||||
def setUp(self):
|
||||
# Common test parameters
|
||||
self.num_local_heads = 128
|
||||
self.kv_lora_rank = 512
|
||||
self.qk_rope_head_dim = 64
|
||||
self.device = get_device()
|
||||
self.dtype = torch.bfloat16
|
||||
self.extend_len = 64
|
||||
self.max_bs = 4
|
||||
self.max_seq_len = 128
|
||||
|
||||
# req_to_token_pool
|
||||
self.req_to_token_pool = MockReqToTokenPool(
|
||||
self.max_bs,
|
||||
self.max_seq_len,
|
||||
self.device,
|
||||
)
|
||||
|
||||
# token_to_kv_pool
|
||||
self.token_to_kv_pool = MLATokenToKVPool(
|
||||
size=self.max_bs * self.max_seq_len,
|
||||
page_size=1, # only consider page=1 for unit test
|
||||
dtype=self.dtype,
|
||||
kv_lora_rank=self.kv_lora_rank,
|
||||
qk_rope_head_dim=self.qk_rope_head_dim,
|
||||
layer_num=1, # only consider layer=1 for unit test
|
||||
device=self.device,
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
|
||||
def test_prefix_chunk_info(self):
|
||||
"""Test the standard extend operation."""
|
||||
|
||||
for test_case in TEST_CASES:
|
||||
print(
|
||||
f"Test case with batch_size={test_case['batch_size']}, prefix_lens={test_case['prefix_lens']}, max_chunk_capacity={test_case['max_chunk_capacity']}"
|
||||
)
|
||||
batch_size = test_case["batch_size"]
|
||||
prefix_lens_cpu = test_case["prefix_lens"]
|
||||
assert len(prefix_lens_cpu) == batch_size
|
||||
prefix_lens = torch.tensor(prefix_lens_cpu, device=self.device)
|
||||
max_chunk_capacity = test_case["max_chunk_capacity"]
|
||||
seq_lens_cpu = [
|
||||
self.extend_len + prefix_lens_cpu[i] for i in range(batch_size)
|
||||
]
|
||||
seq_lens = torch.tensor(seq_lens_cpu, device=self.device)
|
||||
|
||||
# Create forward batch
|
||||
# input_ids and out_cache_loc are dummy tensors in this test
|
||||
forward_batch = MockForwardBatch(
|
||||
max_chunk_capacity=max_chunk_capacity,
|
||||
batch_size=batch_size,
|
||||
input_ids=torch.randint(
|
||||
0, 100, (batch_size, self.extend_len), device=self.device
|
||||
),
|
||||
out_cache_loc=torch.arange(
|
||||
self.max_bs * self.max_seq_len - batch_size * self.extend_len,
|
||||
self.max_bs * self.max_seq_len,
|
||||
device=self.device,
|
||||
),
|
||||
seq_lens_sum=sum(seq_lens_cpu),
|
||||
forward_mode=ForwardMode.EXTEND,
|
||||
req_pool_indices=torch.arange(batch_size, device=self.device),
|
||||
seq_lens=seq_lens,
|
||||
seq_lens_cpu=seq_lens_cpu,
|
||||
extend_prefix_lens=prefix_lens,
|
||||
extend_prefix_lens_cpu=prefix_lens_cpu,
|
||||
)
|
||||
forward_batch.req_to_token_pool = self.req_to_token_pool
|
||||
forward_batch.token_to_kv_pool = self.token_to_kv_pool
|
||||
|
||||
forward_batch.prepare_chunked_prefix_cache_info(self.device)
|
||||
assert forward_batch.get_max_chunk_capacity() == max_chunk_capacity
|
||||
assert forward_batch.prefix_chunk_len == test_case["prefix_chunk_len"]
|
||||
assert forward_batch.num_prefix_chunks == test_case["num_prefix_chunks"]
|
||||
assert torch.allclose(
|
||||
forward_batch.prefix_chunk_starts,
|
||||
test_case["prefix_chunk_starts"].to(self.device),
|
||||
)
|
||||
assert torch.allclose(
|
||||
forward_batch.prefix_chunk_seq_lens,
|
||||
test_case["prefix_chunk_seq_lens"].to(self.device),
|
||||
)
|
||||
|
||||
check_kv_indices(forward_batch)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+1436
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user