[CI] Migrate Attention Backend tests to test/registered/attention/ (#15563)
This commit is contained in:
Alison Shao
@@ -54,12 +54,10 @@ suites = {
|
||||
TestFile("dllm/test_llada2_mini.py", 520),
|
||||
TestFile("test_abort.py", 131),
|
||||
TestFile("test_chunked_prefill.py", 312),
|
||||
TestFile("test_create_kvindices.py", 7),
|
||||
TestFile("test_deterministic.py", 228),
|
||||
TestFile("test_constrained_decoding.py", 111),
|
||||
TestFile("test_eval_fp8_accuracy.py", 250),
|
||||
TestFile("test_external_models.py", 30),
|
||||
TestFile("test_fa3.py", 420),
|
||||
TestFile("test_flashmla.py", 230),
|
||||
TestFile("test_fp8_utils.py", 9),
|
||||
TestFile("rotary_embedding/test_mrope.py", 10),
|
||||
@@ -67,7 +65,6 @@ suites = {
|
||||
TestFile("test_gpt_oss_1gpu.py", 402),
|
||||
TestFile("test_harmony_parser.py", 6),
|
||||
TestFile("test_hidden_states.py", 55),
|
||||
TestFile("test_hybrid_attn_backend.py", 379),
|
||||
TestFile("test_input_embeddings.py", 38),
|
||||
TestFile("test_io_struct.py", 8),
|
||||
TestFile("test_jinja_template_utils.py", 7),
|
||||
@@ -91,8 +88,6 @@ suites = {
|
||||
TestFile("test_prefill_adder.py", 1),
|
||||
TestFile("test_priority_scheduling.py", 130),
|
||||
TestFile("test_pytorch_sampling_backend.py", 66),
|
||||
TestFile("test_radix_attention.py", 105),
|
||||
TestFile("test_radix_cache_unit.py", 8),
|
||||
TestFile("test_reasoning_parser.py", 5),
|
||||
TestFile("test_request_queue_validation.py", 47),
|
||||
TestFile("test_retract_decode.py", 259),
|
||||
@@ -110,13 +105,8 @@ suites = {
|
||||
TestFile("test_torch_compile.py", 190),
|
||||
TestFile("test_torch_compile_moe.py", 210),
|
||||
TestFile("test_triton_fused_moe.py", 12),
|
||||
TestFile("test_torch_native_attention_backend.py", 221),
|
||||
TestFile("test_torchao.py", 103),
|
||||
TestFile("test_triton_attention_kernels.py", 4),
|
||||
TestFile("test_triton_attention_backend.py", 203),
|
||||
TestFile("test_triton_attention_kernels.py", 4),
|
||||
TestFile("test_triton_moe_channel_fp8_kernel.py", 16),
|
||||
TestFile("test_triton_sliding_window.py", 84),
|
||||
TestFile("test_utils_update_weights.py", 29),
|
||||
TestFile("test_video_utils.py", 5),
|
||||
TestFile("test_vision_chunked_prefill.py", 150),
|
||||
@@ -144,7 +134,6 @@ suites = {
|
||||
"per-commit-4-gpu": [
|
||||
TestFile("models/test_qwen3_next_models.py", 650),
|
||||
TestFile("test_gpt_oss_4gpu.py", 300),
|
||||
TestFile("test_local_attn.py", 411),
|
||||
TestFile("test_multi_instance_release_memory_occupation.py", 64),
|
||||
TestFile("test_pp_single_node.py", 500),
|
||||
TestFile("test_epd_disaggregation.py", 150),
|
||||
@@ -166,7 +155,6 @@ suites = {
|
||||
],
|
||||
"per-commit-4-gpu-b200": [
|
||||
TestFile("test_deepseek_v3_fp4_4gpu.py", 1500),
|
||||
TestFile("test_flash_attention_4.py", 90),
|
||||
TestFile("test_fp8_blockwise_gemm.py", 280),
|
||||
TestFile("test_gpt_oss_4gpu.py", 700),
|
||||
TestFile("test_llama31_fp4.py", 90),
|
||||
@@ -253,7 +241,6 @@ suite_amd = {
|
||||
TestFile("test_abort.py", 51),
|
||||
TestFile("test_bench_typebaseddispatcher.py", 10),
|
||||
TestFile("test_chunked_prefill.py", 312),
|
||||
TestFile("test_create_kvindices.py", 2),
|
||||
TestFile("test_eval_fp8_accuracy.py", 303),
|
||||
TestFile("test_fused_moe.py", 30),
|
||||
TestFile("test_harmony_parser.py", 20),
|
||||
@@ -268,7 +255,6 @@ suite_amd = {
|
||||
TestFile("test_page_size.py", 60),
|
||||
TestFile("test_penalty.py", 180),
|
||||
TestFile("test_pytorch_sampling_backend.py", 66),
|
||||
TestFile("test_radix_attention.py", 105),
|
||||
TestFile("test_reasoning_parser.py", 5),
|
||||
TestFile("test_constrained_decoding.py", 120),
|
||||
TestFile("test_retract_decode.py", 450),
|
||||
@@ -279,12 +265,7 @@ suite_amd = {
|
||||
TestFile("test_srt_engine.py", 261),
|
||||
TestFile("test_torch_compile.py", 169),
|
||||
# TestFile("test_torch_compile_moe.py", 210), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/13107
|
||||
TestFile("test_torch_native_attention_backend.py", 123),
|
||||
# TestFile("test_triton_attention_kernels.py", 4),
|
||||
TestFile("test_triton_attention_backend.py", 150),
|
||||
TestFile("test_triton_sliding_window.py", 250),
|
||||
TestFile("test_type_based_dispatcher.py", 10),
|
||||
TestFile("test_wave_attention_kernels.py", 2),
|
||||
# Disabled temporarily
|
||||
# TestFile("test_vlm_input_format.py", 300),
|
||||
# TestFile("models/test_embedding_models.py", 73), # Disabled temporarily, see https://github.com/sgl-project/sglang/issues/11127
|
||||
|
||||
@@ -1,74 +0,0 @@
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestCreateKvIndices(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
if not torch.cuda.is_available():
|
||||
raise unittest.SkipTest("CUDA is not available")
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
def _run_test(self, batch, max_batch, max_context_len):
|
||||
req_to_token = torch.arange(
|
||||
max_batch * max_context_len, dtype=torch.int32, device="cuda"
|
||||
).reshape((max_batch, max_context_len))
|
||||
req_pool_indices = torch.tensor(
|
||||
torch.from_numpy(
|
||||
np.random.choice(range(max_batch), size=batch, replace=False)
|
||||
),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
paged_kernel_lens = torch.tensor(
|
||||
torch.from_numpy(
|
||||
np.random.choice(range(max_context_len), size=batch, replace=False)
|
||||
),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
kv_indptr = torch.zeros((batch + 1,), dtype=torch.int32, device="cuda")
|
||||
kv_indptr[1:] = torch.cumsum(paged_kernel_lens, dim=0)
|
||||
|
||||
# ref
|
||||
req_pool_indices_cpu = req_pool_indices.cpu().numpy()
|
||||
paged_kernel_lens_cpu = paged_kernel_lens.cpu().numpy()
|
||||
kv_indices_ref = torch.cat(
|
||||
[
|
||||
req_to_token[req_pool_indices_cpu[i], : paged_kernel_lens_cpu[i]]
|
||||
for i in range(batch)
|
||||
],
|
||||
dim=0,
|
||||
).contiguous()
|
||||
|
||||
# triton
|
||||
kv_indices_triton = torch.empty(kv_indptr[-1], dtype=torch.int32, device="cuda")
|
||||
create_flashinfer_kv_indices_triton[(batch,)](
|
||||
req_to_token,
|
||||
req_pool_indices,
|
||||
paged_kernel_lens,
|
||||
kv_indptr,
|
||||
None,
|
||||
kv_indices_triton,
|
||||
req_to_token.size(1),
|
||||
)
|
||||
|
||||
# Check
|
||||
self.assertTrue(torch.equal(kv_indices_ref, kv_indices_triton))
|
||||
|
||||
def test_create_kvindices(self):
|
||||
BATCH = [1, 37, 1786]
|
||||
MAX_BATCH = 4096
|
||||
MAX_CONTEXT_LEN = 4096
|
||||
for batch in BATCH:
|
||||
self._run_test(batch, MAX_BATCH, MAX_CONTEXT_LEN)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,263 +0,0 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import get_device_sm, kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_EAGLE3,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
GSM_DATASET_PATH = None
|
||||
|
||||
# In case of some machine lack internet connection, we can set OFFLINE_MODE to True.
|
||||
OFFLINE_MODE = False
|
||||
|
||||
# Change the path below when OFFLINE_MODE is True.
|
||||
OFFLINE_PATH_DICT = {
|
||||
DEFAULT_MODEL_NAME_FOR_TEST: "/shared/public/elr-models/meta-llama/Meta-Llama-3.1-8B-Instruct",
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_EAGLE3: "/shared/public/elr-models/jamesliu1/sglang-EAGLE3-Llama-3.1-Instruct-8B",
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA: "/shared/public/sharing/deepseek/dsv3-test/snapshots/",
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN: "/shared/public/sharing/deepseek/dsv3-test-NextN/snapshots/",
|
||||
GSM_DATASET_PATH: "/shared/public/data/gsm8k/test.jsonl",
|
||||
}
|
||||
|
||||
|
||||
if OFFLINE_MODE:
|
||||
DEFAULT_MODEL_NAME_FOR_TEST = OFFLINE_PATH_DICT[DEFAULT_MODEL_NAME_FOR_TEST]
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_EAGLE3 = OFFLINE_PATH_DICT[
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_EAGLE3
|
||||
]
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA = OFFLINE_PATH_DICT[DEFAULT_MODEL_NAME_FOR_TEST_MLA]
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN = OFFLINE_PATH_DICT[
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN
|
||||
]
|
||||
GSM_DATASET_PATH = OFFLINE_PATH_DICT[GSM_DATASET_PATH]
|
||||
|
||||
|
||||
# Default server arguments shared across all tests
|
||||
DEFAULT_SERVER_ARGS = [
|
||||
"--trust-remote-code",
|
||||
"--cuda-graph-max-bs",
|
||||
"8",
|
||||
"--attention-backend",
|
||||
"fa3",
|
||||
]
|
||||
|
||||
"""
|
||||
Integration test for python/sglang/srt/layers/attention/flashattention_backend.py
|
||||
"""
|
||||
|
||||
|
||||
@unittest.skipIf(get_device_sm() < 90, "Test requires CUDA SM 90 or higher")
|
||||
class BaseFlashAttentionTest(CustomTestCase):
|
||||
"""Base class for testing FlashAttention3."""
|
||||
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
accuracy_threshold = 0.65 # derived tests need to override this
|
||||
speculative_decode = False
|
||||
spec_decode_threshold = 1.0 # derived spec decoding tests need to override this
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
"""Return the arguments for the server launch. Override in subclasses."""
|
||||
return DEFAULT_SERVER_ARGS
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# disable deep gemm precompile to make launch server faster
|
||||
# please don't do this if you want to make your inference workload faster
|
||||
with (
|
||||
envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.override(False),
|
||||
envs.SGLANG_ENABLE_JIT_DEEPGEMM.override(False),
|
||||
):
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=cls.get_server_args(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
requests.get(self.base_url + "/flush_cache")
|
||||
|
||||
args = SimpleNamespace(
|
||||
num_shots=4,
|
||||
num_questions=100,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
data_path=GSM_DATASET_PATH,
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
# Use the appropriate metric key based on the test class
|
||||
metric_key = "accuracy"
|
||||
self.assertGreater(metrics[metric_key], self.accuracy_threshold)
|
||||
|
||||
if self.speculative_decode:
|
||||
server_info = requests.get(self.base_url + "/get_server_info")
|
||||
avg_spec_accept_length = server_info.json()["internal_states"][0][
|
||||
"avg_spec_accept_length"
|
||||
]
|
||||
print(f"{avg_spec_accept_length=}")
|
||||
self.assertGreater(avg_spec_accept_length, self.spec_decode_threshold)
|
||||
|
||||
|
||||
class TestFlashAttention3MLA(BaseFlashAttentionTest):
|
||||
"""Test FlashAttention3 with MLA, e.g. deepseek v3 test model"""
|
||||
|
||||
accuracy_threshold = 0.60
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
return DEFAULT_SERVER_ARGS
|
||||
|
||||
|
||||
class TestFlashAttention3SpeculativeDecode(BaseFlashAttentionTest):
|
||||
"""Test FlashAttention3 with speculative decode enabled with Llama 3.1 8B and its eagle3 model"""
|
||||
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
accuracy_threshold = 0.65
|
||||
speculative_decode = True
|
||||
spec_decode_threshold = 1.5
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
args = DEFAULT_SERVER_ARGS
|
||||
args.extend(
|
||||
[
|
||||
"--cuda-graph-max-bs",
|
||||
"4",
|
||||
"--speculative-algorithm",
|
||||
"EAGLE3",
|
||||
"--speculative-draft-model-path",
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_EAGLE3,
|
||||
"--speculative-num-steps",
|
||||
"3",
|
||||
"--speculative-eagle-topk",
|
||||
"1",
|
||||
"--speculative-num-draft-tokens",
|
||||
"4",
|
||||
"--dtype",
|
||||
"float16",
|
||||
]
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
class TestFlashAttention3SpeculativeDecodeTopk(BaseFlashAttentionTest):
|
||||
"""Tests FlashAttention3 with enhanced speculative decoding using Llama 3.1 8B and EAGLE3.
|
||||
This test will be using top-k value > 1 which would verify the other branches of the FA3 code
|
||||
"""
|
||||
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
accuracy_threshold = 0.65
|
||||
speculative_decode = True
|
||||
spec_decode_threshold = 1.6
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
args = DEFAULT_SERVER_ARGS
|
||||
args.extend(
|
||||
[
|
||||
"--cuda-graph-max-bs",
|
||||
"4",
|
||||
"--speculative-algorithm",
|
||||
"EAGLE3",
|
||||
"--speculative-draft-model-path",
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_EAGLE3,
|
||||
"--speculative-num-steps",
|
||||
"5",
|
||||
"--speculative-eagle-topk",
|
||||
"4",
|
||||
"--speculative-num-draft-tokens",
|
||||
"8",
|
||||
"--dtype",
|
||||
"float16",
|
||||
]
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
class TestFlashAttention3MLASpeculativeDecode(BaseFlashAttentionTest):
|
||||
"""Test FlashAttention3 with speculative decode enabled with deepseek v3 test model and its nextN model"""
|
||||
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
|
||||
accuracy_threshold = 0.60
|
||||
speculative_decode = True
|
||||
spec_decode_threshold = 2.5
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
args = DEFAULT_SERVER_ARGS
|
||||
args.extend(
|
||||
[
|
||||
"--cuda-graph-max-bs",
|
||||
"4",
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-draft-model-path",
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
|
||||
"--speculative-num-steps",
|
||||
"3",
|
||||
"--speculative-eagle-topk",
|
||||
"1",
|
||||
"--speculative-num-draft-tokens",
|
||||
"4",
|
||||
]
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
class TestFlashAttention3MLASpeculativeDecodeTopk(BaseFlashAttentionTest):
|
||||
"""Test FlashAttention3 with speculative decode enabled with deepseek v3 test model and its nextN model
|
||||
This test will be using top-k value > 1 which would verify the other branches of the FA3 code
|
||||
"""
|
||||
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
|
||||
accuracy_threshold = 0.60
|
||||
speculative_decode = True
|
||||
spec_decode_threshold = 2.95
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
args = DEFAULT_SERVER_ARGS
|
||||
args.extend(
|
||||
[
|
||||
"--cuda-graph-max-bs",
|
||||
"4",
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-draft-model-path",
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA_NEXTN,
|
||||
"--speculative-num-steps",
|
||||
"5",
|
||||
"--speculative-eagle-topk",
|
||||
"4",
|
||||
"--speculative-num-draft-tokens",
|
||||
"8",
|
||||
]
|
||||
)
|
||||
return args
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,56 +0,0 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from sglang.srt.utils import get_device_sm, kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipIf(get_device_sm() < 100, "Test requires CUDA SM 100 or higher")
|
||||
class TestFlashAttention4(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = "Qwen/Qwen3-8B"
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--prefill-attention-backend",
|
||||
"fa4",
|
||||
"--decode-attention-backend",
|
||||
"flashinfer",
|
||||
]
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=other_args,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
parsed_url = urlparse(self.base_url)
|
||||
args = SimpleNamespace(
|
||||
num_shots=5,
|
||||
data_path=None,
|
||||
num_questions=1319,
|
||||
max_new_tokens=512,
|
||||
parallel=200,
|
||||
host=f"{parsed_url.scheme}://{parsed_url.hostname}",
|
||||
port=parsed_url.port,
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(metrics)
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.89)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,163 +0,0 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import get_device_sm, kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
|
||||
DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_MLA,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
GSM_DATASET_PATH = None
|
||||
|
||||
# Default server arguments shared across all tests
|
||||
DEFAULT_SERVER_ARGS = [
|
||||
"--trust-remote-code",
|
||||
"--cuda-graph-max-bs",
|
||||
"8",
|
||||
"--prefill-attention-backend",
|
||||
"fa3",
|
||||
"--decode-attention-backend",
|
||||
"flashinfer",
|
||||
]
|
||||
|
||||
|
||||
@unittest.skipIf(get_device_sm() < 90, "Test requires CUDA SM 90 or higher")
|
||||
class TestHybridAttnBackendBase(CustomTestCase):
|
||||
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
accuracy_threshold = 0.65 # derived tests need to override this
|
||||
speculative_decode = False
|
||||
spec_decode_threshold = 2.2 # derived spec decoding tests need to override this
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
"""Return the arguments for the server launch. Override in subclasses."""
|
||||
return DEFAULT_SERVER_ARGS
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# disable deep gemm precompile to make launch server faster
|
||||
# please don't do this if you want to make your inference workload faster
|
||||
with (
|
||||
envs.SGLANG_JIT_DEEPGEMM_PRECOMPILE.override(False),
|
||||
envs.SGLANG_ENABLE_JIT_DEEPGEMM.override(False),
|
||||
):
|
||||
if cls.speculative_decode:
|
||||
model = DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST
|
||||
else:
|
||||
model = cls.model
|
||||
cls.process = popen_launch_server(
|
||||
model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=cls.get_server_args(),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
requests.get(self.base_url + "/flush_cache")
|
||||
|
||||
args = SimpleNamespace(
|
||||
num_shots=4,
|
||||
num_questions=100,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
data_path=GSM_DATASET_PATH,
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
# Use the appropriate metric key based on the test class
|
||||
metric_key = "accuracy"
|
||||
self.assertGreater(metrics[metric_key], self.accuracy_threshold)
|
||||
|
||||
if self.speculative_decode:
|
||||
server_info = requests.get(self.base_url + "/get_server_info")
|
||||
avg_spec_accept_length = server_info.json()["internal_states"][0][
|
||||
"avg_spec_accept_length"
|
||||
]
|
||||
print(f"{avg_spec_accept_length=}")
|
||||
self.assertGreater(avg_spec_accept_length, self.spec_decode_threshold)
|
||||
|
||||
|
||||
class TestHybridAttnBackendMLA(TestHybridAttnBackendBase):
|
||||
accuracy_threshold = 0.60
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST_MLA
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
return DEFAULT_SERVER_ARGS
|
||||
|
||||
|
||||
class TestHybridAttnBackendTorchCompile(TestHybridAttnBackendBase):
|
||||
accuracy_threshold = 0.65
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
return DEFAULT_SERVER_ARGS + ["--enable-torch-compile"]
|
||||
|
||||
|
||||
class TestHybridAttnBackendSpeculativeDecodingPrefillBackend(TestHybridAttnBackendBase):
|
||||
speculative_decode = True
|
||||
# This eagle test uses a very small model, so the accuracy is low.
|
||||
accuracy_threshold = 0.2
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
return DEFAULT_SERVER_ARGS + [
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-draft-model-path",
|
||||
DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
|
||||
"--speculative-num-steps",
|
||||
"3",
|
||||
"--speculative-eagle-topk",
|
||||
"2",
|
||||
"--speculative-num-draft-tokens",
|
||||
"4",
|
||||
"--speculative-attention-mode",
|
||||
"prefill",
|
||||
]
|
||||
|
||||
|
||||
class TestHybridAttnBackendSpeculativeDecodingDecodeBackend(TestHybridAttnBackendBase):
|
||||
speculative_decode = True
|
||||
# This eagle test uses a very small model, so the accuracy is low.
|
||||
accuracy_threshold = 0.2
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
return DEFAULT_SERVER_ARGS + [
|
||||
"--speculative-algorithm",
|
||||
"EAGLE",
|
||||
"--speculative-draft-model-path",
|
||||
DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
|
||||
"--speculative-num-steps",
|
||||
"3",
|
||||
"--speculative-eagle-topk",
|
||||
"2",
|
||||
"--speculative-num-draft-tokens",
|
||||
"4",
|
||||
"--speculative-attention-mode",
|
||||
"decode",
|
||||
]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,72 +0,0 @@
|
||||
import os
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import get_device_sm, kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST_LOCAL_ATTENTION,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipIf(get_device_sm() < 90, "Test requires CUDA SM 90 or higher")
|
||||
class TestFlashAttention3LocalAttn(CustomTestCase):
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST_LOCAL_ATTENTION
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
accuracy_threshold = 0.90
|
||||
|
||||
@classmethod
|
||||
def get_server_args(cls):
|
||||
return [
|
||||
"--cuda-graph-max-bs",
|
||||
"2",
|
||||
"--attention-backend",
|
||||
"fa3",
|
||||
"--tp",
|
||||
"4",
|
||||
"--context-length",
|
||||
"1000000",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=cls.get_server_args(),
|
||||
env=os.environ,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
requests.get(self.base_url + "/flush_cache")
|
||||
|
||||
args = SimpleNamespace(
|
||||
num_shots=4,
|
||||
num_questions=100,
|
||||
max_new_tokens=512,
|
||||
parallel=128,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
data_path=None,
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"{metrics=}")
|
||||
|
||||
# Use the appropriate metric key based on the test class
|
||||
metric_key = "accuracy"
|
||||
self.assertGreater(metrics[metric_key], self.accuracy_threshold)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,88 +0,0 @@
|
||||
import unittest
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.test.kits.radix_cache_server_kit import run_radix_attention_test
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
kill_process_tree,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestRadixCacheFCFS(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--chunked-prefill-size",
|
||||
"128",
|
||||
"--max-total-tokens",
|
||||
"20000",
|
||||
"--schedule-policy",
|
||||
"fcfs",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_radix_attention(self):
|
||||
run_radix_attention_test(self.base_url)
|
||||
|
||||
|
||||
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
|
||||
class TestRadixCacheLPM(TestRadixCacheFCFS):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--chunked-prefill-size",
|
||||
"128",
|
||||
"--max-total-tokens",
|
||||
"20000",
|
||||
"--schedule-policy",
|
||||
"lpm",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class TestRadixCacheNonOverlapLPM(TestRadixCacheFCFS):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--disable-overlap-schedule",
|
||||
"--chunked-prefill-size",
|
||||
"128",
|
||||
"--max-total-tokens",
|
||||
"20000",
|
||||
"--schedule-policy",
|
||||
"lpm",
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
envs.SGLANG_TEST_RETRACT.set(True)
|
||||
envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.set(1)
|
||||
unittest.main()
|
||||
@@ -1,646 +0,0 @@
|
||||
"""
|
||||
Unit tests for the RadixCache implementation.
|
||||
|
||||
This module tests the core functionality of RadixCache, RadixKey, and TreeNode
|
||||
following SGLang testing patterns.
|
||||
|
||||
Test Coverage:
|
||||
- RadixKey: token ID management, slicing, iteration, representation
|
||||
- TreeNode: node properties, reference counting, hash values
|
||||
- RadixCache: insert/match operations, eviction, page alignment, error handling
|
||||
- Cache events and request handling
|
||||
- Boundary conditions with parameterized testing
|
||||
|
||||
Usage:
|
||||
python test_radix_cache_unit.py
|
||||
python -m pytest test_radix_cache_unit.py -v
|
||||
python -m pytest test_radix_cache_unit.py::TestRadixCache::test_insert_basic
|
||||
"""
|
||||
|
||||
import time
|
||||
import unittest
|
||||
import unittest.mock
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored
|
||||
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
|
||||
|
||||
# Test constants
|
||||
DEFAULT_PAGE_SIZE = 4
|
||||
|
||||
|
||||
class TestRadixKey(unittest.TestCase):
|
||||
"""Test cases for RadixKey class."""
|
||||
|
||||
def test_init_basic(self):
|
||||
"""Test basic initialization of RadixKey."""
|
||||
token_ids = [1, 2, 3, 4]
|
||||
key = RadixKey(token_ids)
|
||||
self.assertEqual(key.token_ids, token_ids)
|
||||
self.assertIsNone(key.extra_key)
|
||||
|
||||
def test_init_with_extra_key(self):
|
||||
"""Test initialization with extra_key."""
|
||||
token_ids = [1, 2, 3]
|
||||
extra_key = "test_key"
|
||||
key = RadixKey(token_ids, extra_key)
|
||||
self.assertEqual(key.token_ids, token_ids)
|
||||
self.assertEqual(key.extra_key, extra_key)
|
||||
|
||||
def test_len(self):
|
||||
"""Test __len__ method."""
|
||||
key = RadixKey([1, 2, 3])
|
||||
self.assertEqual(len(key), 3)
|
||||
|
||||
empty_key = RadixKey([])
|
||||
self.assertEqual(len(empty_key), 0)
|
||||
|
||||
def test_iter(self):
|
||||
"""Test __iter__ method."""
|
||||
token_ids = [1, 2, 3, 4]
|
||||
key = RadixKey(token_ids)
|
||||
self.assertEqual(list(key), token_ids)
|
||||
|
||||
def test_len_and_iter(self):
|
||||
"""Test __len__ and __iter__ methods."""
|
||||
test_cases = [
|
||||
([1, 2, 3], 3),
|
||||
([], 0),
|
||||
([42], 1),
|
||||
]
|
||||
|
||||
for tokens, expected in test_cases:
|
||||
with self.subTest(tokens=tokens):
|
||||
key = RadixKey(tokens)
|
||||
self.assertEqual(len(key), expected)
|
||||
self.assertEqual(list(key), tokens)
|
||||
|
||||
def test_getitem_int(self):
|
||||
"""Test __getitem__ with int index."""
|
||||
test_cases = [
|
||||
([10, 20, 30], 0, [10]),
|
||||
([10, 20, 30], -1, [30]),
|
||||
([10, 20, 30], 2, [30]),
|
||||
]
|
||||
|
||||
for tokens, index, expected in test_cases:
|
||||
with self.subTest(tokens=tokens, index=index):
|
||||
key = RadixKey(tokens)
|
||||
result = key[index]
|
||||
self.assertIsInstance(result, RadixKey)
|
||||
self.assertEqual(result.token_ids, expected)
|
||||
|
||||
def test_getitem_slice(self):
|
||||
"""Test __getitem__ with slice and edge cases."""
|
||||
key = RadixKey([1, 2, 3, 4, 5], "extra")
|
||||
|
||||
# Basic slice
|
||||
sliced = key[1:4]
|
||||
self.assertIsInstance(sliced, RadixKey)
|
||||
self.assertEqual(sliced.token_ids, [2, 3, 4])
|
||||
self.assertEqual(sliced.extra_key, "extra")
|
||||
|
||||
# Edge cases
|
||||
self.assertEqual(key[2:2].token_ids, []) # Empty slice
|
||||
self.assertEqual(key[:].token_ids, [1, 2, 3, 4, 5]) # Full slice
|
||||
|
||||
def test_getitem_invalid_index(self):
|
||||
"""Test __getitem__ with invalid indices."""
|
||||
key = RadixKey([1, 2, 3])
|
||||
with self.assertRaises(IndexError):
|
||||
_ = key[10] # Out of bounds
|
||||
|
||||
def test_repr(self):
|
||||
"""Test __repr__ method."""
|
||||
key = RadixKey([1, 2, 3], "test")
|
||||
repr_str = repr(key)
|
||||
self.assertIn("RadixKey", repr_str)
|
||||
self.assertIn("extra_key='test'", repr_str)
|
||||
self.assertIn("[1, 2, 3]", repr_str)
|
||||
|
||||
def test_repr_long_token_ids(self):
|
||||
"""Test __repr__ with long token_ids."""
|
||||
long_tokens = list(range(15))
|
||||
key = RadixKey(long_tokens)
|
||||
repr_str = repr(key)
|
||||
self.assertIn("...", repr_str) # Should be truncated
|
||||
|
||||
|
||||
class TestTreeNode(unittest.TestCase):
|
||||
"""Test cases for TreeNode class."""
|
||||
|
||||
def setUp(self):
|
||||
"""Reset the counter before each test."""
|
||||
TreeNode.counter = 0
|
||||
|
||||
def test_init_basic(self):
|
||||
"""Test basic initialization of TreeNode."""
|
||||
node = TreeNode()
|
||||
self.assertEqual(node.id, 0)
|
||||
self.assertEqual(len(node.children), 0)
|
||||
self.assertIsNone(node.parent)
|
||||
self.assertIsNone(node.key)
|
||||
self.assertIsNone(node.value)
|
||||
self.assertEqual(node.lock_ref, 0)
|
||||
self.assertEqual(node.hit_count, 0)
|
||||
self.assertEqual(node.host_ref_counter, 0)
|
||||
self.assertIsNone(node.host_value)
|
||||
self.assertIsNone(node.hash_value)
|
||||
|
||||
def test_init_with_id(self):
|
||||
"""Test initialization with custom ID."""
|
||||
node = TreeNode(id=42)
|
||||
self.assertEqual(node.id, 42)
|
||||
node2 = TreeNode()
|
||||
self.assertEqual(node2.id, 1) # Counter was incremented
|
||||
|
||||
def test_counter_increment(self):
|
||||
"""Test that counter increments properly."""
|
||||
node1 = TreeNode()
|
||||
node2 = TreeNode()
|
||||
self.assertEqual(node1.id, 0)
|
||||
self.assertEqual(node2.id, 1)
|
||||
|
||||
def test_evicted_backuped_properties(self):
|
||||
"""Test evicted and backuped properties."""
|
||||
test_cases = [
|
||||
(False, False, True, False),
|
||||
(True, False, False, False),
|
||||
(True, True, False, True),
|
||||
(False, True, True, True),
|
||||
]
|
||||
|
||||
for (
|
||||
has_value,
|
||||
has_host_value,
|
||||
expected_evicted,
|
||||
expected_backuped,
|
||||
) in test_cases:
|
||||
with self.subTest(has_value=has_value, has_host_value=has_host_value):
|
||||
node = TreeNode()
|
||||
|
||||
if has_value:
|
||||
node.value = torch.tensor([1, 2, 3])
|
||||
if has_host_value:
|
||||
node.host_value = torch.tensor([4, 5, 6])
|
||||
|
||||
self.assertEqual(node.evicted, expected_evicted)
|
||||
self.assertEqual(node.backuped, expected_backuped)
|
||||
|
||||
def test_protect_release_host(self):
|
||||
"""Test protect_host and release_host methods."""
|
||||
node = TreeNode()
|
||||
self.assertEqual(node.host_ref_counter, 0)
|
||||
|
||||
node.protect_host()
|
||||
self.assertEqual(node.host_ref_counter, 1)
|
||||
|
||||
node.release_host()
|
||||
self.assertEqual(node.host_ref_counter, 0)
|
||||
|
||||
# Test error case
|
||||
with self.assertRaises(RuntimeError):
|
||||
node.release_host()
|
||||
|
||||
def test_get_last_hash_value(self):
|
||||
"""Test get_last_hash_value method."""
|
||||
node = TreeNode()
|
||||
self.assertIsNone(node.get_last_hash_value())
|
||||
|
||||
node.hash_value = ["hash1", "hash2", "hash3"]
|
||||
self.assertEqual(node.get_last_hash_value(), "hash3")
|
||||
|
||||
def test_lt_comparison(self):
|
||||
"""Test less than comparison based on last_access_time."""
|
||||
node1 = TreeNode()
|
||||
time.sleep(0.001) # Small delay to ensure different timestamps
|
||||
node2 = TreeNode()
|
||||
|
||||
self.assertTrue(node1 < node2)
|
||||
self.assertFalse(node2 < node1)
|
||||
|
||||
|
||||
class TestRadixCache(unittest.TestCase):
|
||||
"""Test cases for RadixCache class."""
|
||||
|
||||
def setUp(self):
|
||||
"""Set up test fixtures."""
|
||||
TreeNode.counter = 0
|
||||
|
||||
def test_init_variations(self):
|
||||
"""Test cache initialization with different parameters."""
|
||||
test_cases = [
|
||||
(1, False, False),
|
||||
(4, False, True),
|
||||
(1, True, False),
|
||||
]
|
||||
|
||||
for page_size, disable, enable_events in test_cases:
|
||||
with self.subTest(
|
||||
page_size=page_size, disable=disable, enable_events=enable_events
|
||||
):
|
||||
cache = RadixCache.create_simulated(
|
||||
disable=disable,
|
||||
page_size=page_size,
|
||||
enable_kv_cache_events=enable_events,
|
||||
)
|
||||
|
||||
self.assertEqual(cache.page_size, page_size)
|
||||
self.assertEqual(cache.disable, disable)
|
||||
self.assertEqual(cache.enable_kv_cache_events, enable_events)
|
||||
self.assertEqual(cache.device, torch.device("cpu"))
|
||||
self.assertIsNotNone(cache.root_node)
|
||||
self.assertEqual(len(cache.root_node.key), 0)
|
||||
|
||||
def test_reset(self):
|
||||
"""Test reset method."""
|
||||
cache = RadixCache.create_simulated()
|
||||
|
||||
# Insert some data
|
||||
cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64))
|
||||
self.assertGreater(cache.total_size(), 0)
|
||||
|
||||
# Reset
|
||||
cache.reset()
|
||||
self.assertEqual(cache.total_size(), 0)
|
||||
self.assertEqual(cache.evictable_size(), 0)
|
||||
self.assertEqual(cache.protected_size(), 0)
|
||||
|
||||
def test_insert_and_match_basic(self):
|
||||
"""Test basic insert and match operations."""
|
||||
for disable_cache in [False, True]:
|
||||
with self.subTest(disable_cache=disable_cache):
|
||||
cache = RadixCache.create_simulated(disable=disable_cache)
|
||||
|
||||
key = RadixKey([1, 2, 3])
|
||||
value = torch.tensor([10, 20, 30], dtype=torch.int64)
|
||||
prefix_len = cache.insert(key, value)
|
||||
|
||||
if disable_cache:
|
||||
self.assertEqual(prefix_len, 0)
|
||||
self.assertEqual(cache.total_size(), 0)
|
||||
continue
|
||||
|
||||
self.assertEqual(prefix_len, 0) # No existing prefix
|
||||
self.assertEqual(cache.total_size(), 3)
|
||||
self.assertEqual(cache.evictable_size(), 3)
|
||||
|
||||
# Test match_prefix
|
||||
result = cache.match_prefix(RadixKey([1, 2, 3]))
|
||||
self.assertEqual(len(result.device_indices), 3)
|
||||
torch.testing.assert_close(result.device_indices, value)
|
||||
|
||||
# Test partial match
|
||||
result = cache.match_prefix(RadixKey([1, 2]))
|
||||
self.assertEqual(len(result.device_indices), 2)
|
||||
torch.testing.assert_close(
|
||||
result.device_indices, torch.tensor([10, 20], dtype=torch.int64)
|
||||
)
|
||||
|
||||
def test_insert_with_none_value(self):
|
||||
"""Test insert with None value (should use token_ids as list)."""
|
||||
cache = RadixCache.create_simulated()
|
||||
|
||||
key = RadixKey([1, 2, 3])
|
||||
prefix_len = cache.insert(key, None)
|
||||
|
||||
# When None is passed, it should create value from token_ids
|
||||
self.assertEqual(prefix_len, 0)
|
||||
self.assertEqual(cache.total_size(), 3)
|
||||
|
||||
def test_total_size(self):
|
||||
"""Test total_size calculation."""
|
||||
cache = RadixCache.create_simulated()
|
||||
|
||||
self.assertEqual(cache.total_size(), 0)
|
||||
|
||||
cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64))
|
||||
self.assertEqual(cache.total_size(), 3)
|
||||
|
||||
cache.insert(RadixKey([4, 5]), torch.tensor([40, 50], dtype=torch.int64))
|
||||
self.assertEqual(cache.total_size(), 5)
|
||||
|
||||
def test_kv_cache_events(self):
|
||||
"""Test KV cache events functionality."""
|
||||
test_cases = [
|
||||
(1, True),
|
||||
(2, True),
|
||||
(1, False),
|
||||
]
|
||||
|
||||
for page_size, enable_events in test_cases:
|
||||
with self.subTest(page_size=page_size, enable_events=enable_events):
|
||||
cache = RadixCache.create_simulated(
|
||||
page_size=page_size, enable_kv_cache_events=enable_events
|
||||
)
|
||||
|
||||
# Insert data
|
||||
cache.insert(RadixKey([1, 2, 3, 4, 5]), None)
|
||||
|
||||
# Take events
|
||||
events = cache.take_events()
|
||||
|
||||
if enable_events:
|
||||
self.assertGreater(len(events), 0)
|
||||
# Verify events include BlockStored events (there might be other event types)
|
||||
block_stored_events = [
|
||||
e for e in events if isinstance(e, BlockStored)
|
||||
]
|
||||
self.assertGreater(len(block_stored_events), 0)
|
||||
for event in block_stored_events:
|
||||
self.assertLessEqual(len(event.token_ids), page_size)
|
||||
else:
|
||||
self.assertEqual(len(events), 0)
|
||||
|
||||
def test_kv_cache_events_with_eviction(self):
|
||||
"""Test KV cache events include removal events."""
|
||||
mock_allocator = unittest.mock.Mock()
|
||||
mock_allocator.device = torch.device("cpu")
|
||||
|
||||
cache = RadixCache.create_simulated(
|
||||
mock_allocator=mock_allocator, enable_kv_cache_events=True
|
||||
)
|
||||
|
||||
# Insert and then evict data
|
||||
cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64))
|
||||
cache.evict(3)
|
||||
|
||||
# Take events - should include both store and remove events
|
||||
events = cache.take_events()
|
||||
self.assertGreater(len(events), 0)
|
||||
|
||||
# Check event types
|
||||
event_types = [type(event).__name__ for event in events]
|
||||
self.assertIn("BlockStored", event_types)
|
||||
|
||||
# Verify BlockRemoved event content
|
||||
remove_events = [e for e in events if isinstance(e, BlockRemoved)]
|
||||
for event in remove_events:
|
||||
self.assertGreater(len(event.block_hashes), 0)
|
||||
|
||||
def test_extra_key_isolation(self):
|
||||
"""Test that keys with different extra_key values are isolated."""
|
||||
cache = RadixCache.create_simulated()
|
||||
|
||||
# Insert same token sequence with different extra keys
|
||||
cache.insert(
|
||||
RadixKey([1, 2, 3], "key1"), torch.tensor([10, 20, 30], dtype=torch.int64)
|
||||
)
|
||||
cache.insert(
|
||||
RadixKey([1, 2, 3], "key2"), torch.tensor([40, 50, 60], dtype=torch.int64)
|
||||
)
|
||||
cache.insert(
|
||||
RadixKey([1, 2, 3], None), torch.tensor([70, 80, 90], dtype=torch.int64)
|
||||
)
|
||||
|
||||
# Keys with different extra_key should not match each other
|
||||
result1 = cache.match_prefix(RadixKey([1, 2, 3], "key1"))
|
||||
result2 = cache.match_prefix(RadixKey([1, 2, 3], "key2"))
|
||||
result3 = cache.match_prefix(RadixKey([1, 2, 3], None))
|
||||
result4 = cache.match_prefix(RadixKey([1, 2, 3], "nonexistent"))
|
||||
|
||||
# Each should match only its own data
|
||||
self.assertEqual(len(result1.device_indices), 3)
|
||||
torch.testing.assert_close(
|
||||
result1.device_indices, torch.tensor([10, 20, 30], dtype=torch.int64)
|
||||
)
|
||||
|
||||
self.assertEqual(len(result2.device_indices), 3)
|
||||
torch.testing.assert_close(
|
||||
result2.device_indices, torch.tensor([40, 50, 60], dtype=torch.int64)
|
||||
)
|
||||
|
||||
self.assertEqual(len(result3.device_indices), 3)
|
||||
torch.testing.assert_close(
|
||||
result3.device_indices, torch.tensor([70, 80, 90], dtype=torch.int64)
|
||||
)
|
||||
|
||||
# Non-existent extra_key should not match
|
||||
self.assertEqual(len(result4.device_indices), 0)
|
||||
|
||||
def test_lock_ref_operations(self):
|
||||
"""Test lock reference counting operations."""
|
||||
cache = RadixCache.create_simulated()
|
||||
|
||||
# Insert sequence
|
||||
cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64))
|
||||
|
||||
# Get node
|
||||
result = cache.match_prefix(RadixKey([1, 2, 3]))
|
||||
node = result.last_device_node
|
||||
|
||||
initial_evictable = cache.evictable_size()
|
||||
initial_protected = cache.protected_size()
|
||||
|
||||
# Lock the node
|
||||
cache.inc_lock_ref(node)
|
||||
self.assertEqual(cache.protected_size(), initial_protected + 3)
|
||||
self.assertEqual(cache.evictable_size(), initial_evictable - 3)
|
||||
|
||||
# Unlock the node
|
||||
cache.dec_lock_ref(node)
|
||||
self.assertEqual(cache.protected_size(), initial_protected)
|
||||
self.assertEqual(cache.evictable_size(), initial_evictable)
|
||||
|
||||
def test_evict_functionality(self):
|
||||
"""Test eviction functionality."""
|
||||
mock_allocator = unittest.mock.Mock()
|
||||
mock_allocator.device = torch.device("cpu")
|
||||
|
||||
cache = RadixCache.create_simulated(mock_allocator=mock_allocator)
|
||||
|
||||
# Insert sequences
|
||||
cache.insert(RadixKey([1, 2]), torch.tensor([10, 20], dtype=torch.int64))
|
||||
cache.insert(RadixKey([3, 4]), torch.tensor([30, 40], dtype=torch.int64))
|
||||
|
||||
initial_size = cache.total_size()
|
||||
|
||||
# Evict some tokens
|
||||
cache.evict(2)
|
||||
|
||||
# Should have called free and reduced size
|
||||
mock_allocator.free.assert_called()
|
||||
self.assertLess(cache.total_size(), initial_size)
|
||||
|
||||
def test_page_alignment_boundary(self):
|
||||
"""Test page alignment with different sizes."""
|
||||
test_cases = [
|
||||
(1, 5),
|
||||
(2, 5),
|
||||
(4, 6),
|
||||
]
|
||||
|
||||
for page_size, sequence_length in test_cases:
|
||||
with self.subTest(page_size=page_size, sequence_length=sequence_length):
|
||||
cache = RadixCache.create_simulated(page_size=page_size)
|
||||
|
||||
tokens = list(range(sequence_length))
|
||||
cache.insert(RadixKey(tokens), torch.tensor(tokens, dtype=torch.int64))
|
||||
|
||||
result = cache.match_prefix(RadixKey(tokens))
|
||||
self.assertGreater(len(result.device_indices), 0)
|
||||
|
||||
# Match length should be page-aligned
|
||||
match_len = len(result.device_indices)
|
||||
self.assertEqual(match_len % page_size, 0)
|
||||
|
||||
def test_pretty_print_basic(self):
|
||||
"""Test pretty_print produces output."""
|
||||
cache = RadixCache.create_simulated()
|
||||
|
||||
cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64))
|
||||
|
||||
# Just test that it doesn't crash
|
||||
try:
|
||||
cache.pretty_print()
|
||||
except Exception as e:
|
||||
self.fail(f"pretty_print raised an exception: {e}")
|
||||
|
||||
def test_all_values_flatten(self):
|
||||
"""Test all_values_flatten method."""
|
||||
cache = RadixCache.create_simulated()
|
||||
|
||||
cache.insert(RadixKey([1, 2]), torch.tensor([10, 20], dtype=torch.int64))
|
||||
cache.insert(RadixKey([3, 4]), torch.tensor([30, 40], dtype=torch.int64))
|
||||
|
||||
all_values = cache.all_values_flatten()
|
||||
self.assertEqual(len(all_values), 4)
|
||||
# Values should contain all inserted values (order may vary)
|
||||
values_set = set(all_values.tolist())
|
||||
self.assertEqual(values_set, {10, 20, 30, 40})
|
||||
|
||||
def test_advanced_prefix_match_with_node_splits(self):
|
||||
"""Advanced prefix matching: splits inside nodes and across pages."""
|
||||
for page_size in [1, 2]:
|
||||
with self.subTest(page_size=page_size):
|
||||
cache = RadixCache.create_simulated(page_size=page_size)
|
||||
|
||||
# Insert a long sequence that will be split later.
|
||||
seq1 = [1, 2, 3, 4, 5, 6, 7, 8]
|
||||
val1 = torch.tensor([x * 10 for x in seq1], dtype=torch.int64)
|
||||
cache.insert(RadixKey(seq1), val1)
|
||||
|
||||
# Insert a diverging branch to create an internal node on the path.
|
||||
seq2 = [1, 2, 9, 10]
|
||||
val2 = torch.tensor([x * 10 for x in seq2], dtype=torch.int64)
|
||||
cache.insert(RadixKey(seq2), val2)
|
||||
print(cache.pretty_print())
|
||||
|
||||
baseline_total = cache.total_size()
|
||||
expected_total = 10 # 8 + 2
|
||||
self.assertEqual(baseline_total, expected_total)
|
||||
|
||||
# Match that causes a split inside an existing node:
|
||||
# take first 4 tokens of seq1, then diverge.
|
||||
query1 = [1, 2, 3, 4, 999, 1000]
|
||||
result1 = cache.match_prefix(RadixKey(query1))
|
||||
torch.testing.assert_close(result1.device_indices, val1[:4])
|
||||
# No data change after structural split during matching.
|
||||
self.assertEqual(cache.total_size(), baseline_total)
|
||||
|
||||
# Full match of the long sequence still returns the full indices.
|
||||
result_full = cache.match_prefix(RadixKey(seq1))
|
||||
torch.testing.assert_close(result_full.device_indices, val1)
|
||||
|
||||
# Another split deeper on the path (after matching 6 tokens, then diverge).
|
||||
query2 = [1, 2, 3, 4, 5, 6, 777, 888]
|
||||
result2 = cache.match_prefix(RadixKey(query2))
|
||||
torch.testing.assert_close(result2.device_indices, val1[:6])
|
||||
self.assertEqual(cache.total_size(), baseline_total)
|
||||
|
||||
# Matching the short diverging branch should return exactly its indices.
|
||||
result_branch = cache.match_prefix(RadixKey(seq2))
|
||||
torch.testing.assert_close(result_branch.device_indices, val2)
|
||||
|
||||
def test_hash_value_storage(self):
|
||||
"""Test that hash_value is stored correctly after insert operations."""
|
||||
cache = RadixCache.create_simulated(
|
||||
page_size=4,
|
||||
enable_kv_cache_events=True,
|
||||
)
|
||||
|
||||
# Insert a sequence
|
||||
cache.insert(RadixKey([1, 2, 3, 4, 5, 6, 7, 8]), None)
|
||||
|
||||
# Trigger event emission to compute hash_value lazily
|
||||
cache.take_events()
|
||||
|
||||
# Find the inserted node (traverse from root)
|
||||
node = cache.root_node
|
||||
for i in range(0, 8, 4): # page_size=4, so 2 pages
|
||||
child_key = tuple([1, 2, 3, 4][:4]) if i == 0 else tuple([5, 6, 7, 8][:4])
|
||||
if child_key in node.children:
|
||||
node = node.children[child_key]
|
||||
break
|
||||
|
||||
# Verify hash_value is set (computed lazily during event emission)
|
||||
self.assertIsNotNone(node.hash_value)
|
||||
# Should have 2 pages (8 tokens / 4 page_size)
|
||||
self.assertEqual(len(node.hash_value), 2)
|
||||
|
||||
def test_hash_value_repeating_tokens(self):
|
||||
"""Test that repeating token patterns get different hash values."""
|
||||
cache = RadixCache.create_simulated(
|
||||
page_size=4,
|
||||
enable_kv_cache_events=True,
|
||||
)
|
||||
|
||||
# Insert a sequence with repeating token pattern: [1,2,3,4, 1,2,3,4]
|
||||
cache.insert(RadixKey([1, 2, 3, 4, 1, 2, 3, 4]), None)
|
||||
|
||||
events = cache.take_events()
|
||||
block_stored_events = [e for e in events if isinstance(e, BlockStored)]
|
||||
|
||||
# Should have 2 blocks (2 pages of size 4)
|
||||
self.assertEqual(len(block_stored_events), 2)
|
||||
|
||||
# Extract block hashes
|
||||
block_hash_1 = block_stored_events[0].block_hashes[0]
|
||||
block_hash_2 = block_stored_events[1].block_hashes[0]
|
||||
|
||||
# The two blocks should have DIFFERENT hashes despite same content
|
||||
# because they are at different positions (sequence-aware hashing)
|
||||
self.assertNotEqual(
|
||||
block_hash_1,
|
||||
block_hash_2,
|
||||
"Repeating token patterns should get different sequence-aware hashes",
|
||||
)
|
||||
|
||||
# First block should have no parent
|
||||
self.assertIsNone(block_stored_events[0].parent_block_hash)
|
||||
|
||||
# Second block's parent should be the first block's hash
|
||||
self.assertEqual(block_stored_events[1].parent_block_hash, block_hash_1)
|
||||
|
||||
def test_hash_value_split(self):
|
||||
"""Test that hash_value is split correctly when nodes are split."""
|
||||
cache = RadixCache.create_simulated(
|
||||
page_size=2,
|
||||
enable_kv_cache_events=True,
|
||||
)
|
||||
|
||||
# Insert a sequence that will cause a split
|
||||
cache.insert(RadixKey([1, 2, 3, 4]), None)
|
||||
cache.take_events() # Clear events and compute hash_value for first node
|
||||
|
||||
# Insert a diverging sequence that will cause a split at page boundary
|
||||
cache.insert(RadixKey([1, 2, 5, 6]), None)
|
||||
cache.take_events() # Trigger event emission to compute hash_value
|
||||
|
||||
# Find the split node
|
||||
node = cache.root_node
|
||||
child_key = tuple([1, 2])
|
||||
if child_key in node.children:
|
||||
node = node.children[child_key]
|
||||
# After split and event emission, hash_value should be computed
|
||||
# Note: If hash_value wasn't set before split, it will be computed lazily
|
||||
# during event emission. If it was set, it will be split.
|
||||
# Either way, after events are emitted, it should be set.
|
||||
self.assertIsNotNone(node.hash_value)
|
||||
# Should have 1 page (split at page_size=2)
|
||||
self.assertEqual(len(node.hash_value), 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,47 +0,0 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest test_triton_attention_backend.TestTritonAttnBackend.test_mmlu
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestTorchNativeAttnBackend(CustomTestCase):
|
||||
def test_mmlu(self):
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
process = popen_launch_server(
|
||||
model,
|
||||
base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=["--attention-backend", "torch_native"],
|
||||
)
|
||||
|
||||
try:
|
||||
args = SimpleNamespace(
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreaterEqual(metrics["score"], 0.65)
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,66 +0,0 @@
|
||||
"""
|
||||
Usage:
|
||||
python3 -m unittest test_triton_attention_backend.TestTritonAttnBackend.test_mmlu
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
run_bench_offline_throughput,
|
||||
)
|
||||
|
||||
|
||||
class TestTritonAttnBackend(CustomTestCase):
|
||||
def test_latency(self):
|
||||
output_throughput = run_bench_offline_throughput(
|
||||
DEFAULT_MODEL_NAME_FOR_TEST,
|
||||
[
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--enable-torch-compile",
|
||||
"--cuda-graph-max-bs",
|
||||
4,
|
||||
],
|
||||
)
|
||||
|
||||
print(f"{output_throughput=}")
|
||||
|
||||
if is_in_ci():
|
||||
self.assertGreater(output_throughput, 153)
|
||||
|
||||
def test_mmlu(self):
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST
|
||||
base_url = DEFAULT_URL_FOR_TEST
|
||||
process = popen_launch_server(
|
||||
model,
|
||||
base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=["--attention-backend", "triton"],
|
||||
)
|
||||
|
||||
try:
|
||||
args = SimpleNamespace(
|
||||
base_url=base_url,
|
||||
model=model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreaterEqual(metrics["score"], 0.65)
|
||||
finally:
|
||||
kill_process_tree(process.pid)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,783 +0,0 @@
|
||||
import random
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.srt.layers.attention.triton_ops.decode_attention import (
|
||||
decode_attention_fwd,
|
||||
decode_attention_fwd_grouped,
|
||||
decode_attention_fwd_normal,
|
||||
)
|
||||
from sglang.srt.layers.attention.triton_ops.extend_attention import (
|
||||
build_unified_kv_indices,
|
||||
extend_attention_fwd,
|
||||
extend_attention_fwd_unified,
|
||||
redundant_attention,
|
||||
)
|
||||
from sglang.srt.layers.attention.triton_ops.prefill_attention import (
|
||||
context_attention_fwd,
|
||||
)
|
||||
from sglang.srt.utils import get_device
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
def extend_attention_fwd_torch(
|
||||
q: torch.Tensor, # [extend_tokens, H_Q, D]
|
||||
k: torch.Tensor, # [extend_tokens, H_KV, D]
|
||||
v: torch.Tensor, # [extend_tokens, H_KV, D]
|
||||
o: torch.Tensor, # [extend_tokens, H_Q, D]
|
||||
k_cache: torch.Tensor, # [total_tokens, H_KV, D]
|
||||
v_cache: torch.Tensor, # [total_tokens, H_KV, D]
|
||||
qo_indptr: torch.Tensor, # [B+1]
|
||||
kv_indptr: torch.Tensor, # [B+1]
|
||||
kv_indices: torch.Tensor, # [prefix_tokens]
|
||||
sliding_window_size: int,
|
||||
):
|
||||
B = qo_indptr.size(0) - 1
|
||||
_, H_Q, D = q.shape
|
||||
_, H_KV, _ = k.shape
|
||||
|
||||
group_size = H_Q // H_KV
|
||||
scale = 1.0 / D**0.5
|
||||
|
||||
for i in range(B):
|
||||
q_start = int(qo_indptr[i].item())
|
||||
q_end = int(qo_indptr[i + 1].item())
|
||||
kv_start = int(kv_indptr[i].item())
|
||||
kv_end = int(kv_indptr[i + 1].item())
|
||||
|
||||
prefix_indices = kv_indices[kv_start:kv_end]
|
||||
k_prefix = k_cache[prefix_indices] # [prefix_len, H_KV, D]
|
||||
v_prefix = v_cache[prefix_indices] # [prefix_len, H_KV, D]
|
||||
|
||||
k_extend = k[q_start:q_end] # [extend_len, H_KV, D]
|
||||
v_extend = v[q_start:q_end] # [extend_len, H_KV, D]
|
||||
q_extend = q[q_start:q_end] # [extend_len, H_Q, D]
|
||||
|
||||
k_full = torch.cat([k_prefix, k_extend], dim=0) # [total_len, H_KV, D]
|
||||
v_full = torch.cat([v_prefix, v_extend], dim=0) # [total_len, H_KV, D]
|
||||
|
||||
if group_size != 1:
|
||||
k_full_hq = k_full.repeat_interleave(
|
||||
group_size, dim=1
|
||||
) # [total_len, H_Q, D]
|
||||
v_full_hq = v_full.repeat_interleave(
|
||||
group_size, dim=1
|
||||
) # [total_len, H_Q, D]
|
||||
else:
|
||||
k_full_hq = k_full
|
||||
v_full_hq = v_full
|
||||
|
||||
prefix_len = k_prefix.size(0)
|
||||
extend_len = k_extend.size(0)
|
||||
total_len = prefix_len + extend_len
|
||||
|
||||
# causal
|
||||
pos_keys = torch.arange(total_len, device=q.device)
|
||||
t = prefix_len + torch.arange(extend_len, device=q.device) # [extend_len]
|
||||
causal_mask = pos_keys.unsqueeze(0) <= t.unsqueeze(1)
|
||||
|
||||
# sliding window
|
||||
if sliding_window_size is not None and sliding_window_size > 0:
|
||||
start = (t - (sliding_window_size)).clamp_min(0) # [extend_len]
|
||||
else:
|
||||
start = torch.zeros_like(t)
|
||||
window_mask = pos_keys.unsqueeze(0) >= start.unsqueeze(1)
|
||||
|
||||
final_mask = causal_mask & window_mask
|
||||
|
||||
attn_scores = (
|
||||
torch.einsum("qhd,khd->qhk", q_extend, k_full_hq) * scale
|
||||
) # [extend_len, H_Q, total_len]
|
||||
attn_scores = attn_scores.masked_fill(~final_mask.unsqueeze(1), float("-inf"))
|
||||
|
||||
attn_weights = F.softmax(attn_scores, dim=-1)
|
||||
o[q_start:q_end] = torch.einsum("qhk,khd->qhd", attn_weights, v_full_hq)
|
||||
|
||||
|
||||
class TestTritonAttention(CustomTestCase):
|
||||
|
||||
def _set_all_seeds(self, seed):
|
||||
"""Set all random seeds for reproducibility."""
|
||||
random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
def setUp(self):
|
||||
# Set seeds before each test method
|
||||
self._set_all_seeds(42)
|
||||
|
||||
def _test_extend_attention_once(self, B, N_CTX, H_Q, H_KV, D):
|
||||
dtype = torch.bfloat16
|
||||
device = get_device()
|
||||
|
||||
b_seq_len_prefix = torch.randint(
|
||||
1, N_CTX // 2, (B,), dtype=torch.int32, device=device
|
||||
)
|
||||
b_seq_len_extend = torch.randint(
|
||||
1, N_CTX // 2, (B,), dtype=torch.int32, device=device
|
||||
)
|
||||
b_seq_len = b_seq_len_prefix + b_seq_len_extend
|
||||
max_len_in_batch = torch.max(b_seq_len, 0)[0].item()
|
||||
|
||||
b_req_idx = torch.arange(B, dtype=torch.int32, device=device)
|
||||
b_start_loc = torch.zeros((B,), dtype=torch.int32, device=device)
|
||||
b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
|
||||
b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device=device)
|
||||
b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
|
||||
|
||||
kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
|
||||
kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0)
|
||||
kv_indices = torch.zeros(
|
||||
(b_seq_len_prefix.sum().item(),), dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
for i in range(B):
|
||||
kv_indices[kv_indptr[i] : kv_indptr[i + 1]] = torch.arange(
|
||||
b_start_loc[i], b_start_loc[i] + b_seq_len_prefix[i]
|
||||
)
|
||||
|
||||
total_token_num = torch.sum(b_seq_len).item()
|
||||
extend_token_num = torch.sum(b_seq_len_extend).item()
|
||||
k_buffer = torch.empty(
|
||||
(total_token_num, H_KV, D), dtype=dtype, device=device
|
||||
).normal_(mean=0.1, std=0.2)
|
||||
v_buffer = torch.empty(
|
||||
(total_token_num, H_KV, D), dtype=dtype, device=device
|
||||
).normal_(mean=0.1, std=0.2)
|
||||
|
||||
k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device)
|
||||
v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device)
|
||||
q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device)
|
||||
for i in range(B):
|
||||
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
|
||||
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
|
||||
extend_start = b_start_loc_extend[i]
|
||||
extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
|
||||
k_extend[extend_start:extend_end] = k_buffer[
|
||||
extend_start_in_buffer:extend_end_in_buffer
|
||||
]
|
||||
v_extend[extend_start:extend_end] = v_buffer[
|
||||
extend_start_in_buffer:extend_end_in_buffer
|
||||
]
|
||||
q_extend[extend_start:extend_end] = torch.empty(
|
||||
(b_seq_len_extend[i], H_Q, D), dtype=dtype, device=device
|
||||
).normal_(mean=0.1, std=0.2)
|
||||
|
||||
o_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device)
|
||||
o_extend_mask = torch.empty(
|
||||
(extend_token_num, H_Q, D), dtype=dtype, device=device
|
||||
)
|
||||
o_redundant = torch.empty(
|
||||
(extend_token_num, H_Q, D), dtype=dtype, device=device
|
||||
)
|
||||
|
||||
b_seq_len_extend = b_seq_len - b_seq_len_prefix
|
||||
max_len_extend = torch.max(b_seq_len_extend, 0)[0].item()
|
||||
qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
|
||||
qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0)
|
||||
|
||||
custom_mask = None
|
||||
mask_indptr = None
|
||||
|
||||
extend_attention_fwd(
|
||||
q_extend,
|
||||
k_extend,
|
||||
v_extend,
|
||||
o_extend,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
custom_mask,
|
||||
True,
|
||||
mask_indptr,
|
||||
max_len_extend,
|
||||
)
|
||||
|
||||
b_seq_mask_len = b_seq_len_extend * b_seq_len
|
||||
custom_mask = torch.ones(
|
||||
(b_seq_mask_len.sum().item(),), dtype=torch.bool, device=device
|
||||
)
|
||||
mask_indptr = torch.zeros((B + 1,), dtype=torch.int64, device=device)
|
||||
mask_indptr[1 : B + 1] = torch.cumsum(b_seq_mask_len[:B], dim=0)
|
||||
for i in range(B):
|
||||
causal_mask = (
|
||||
torch.tril(
|
||||
torch.ones(b_seq_len_extend[i], b_seq_len_extend[i]), diagonal=0
|
||||
)
|
||||
== 1
|
||||
)
|
||||
prefix_mask = torch.ones(
|
||||
b_seq_len_extend[i], b_seq_len_prefix[i], dtype=torch.bool
|
||||
)
|
||||
mask_flatten = torch.cat([prefix_mask, causal_mask], dim=1).flatten()
|
||||
custom_mask[mask_indptr[i] : mask_indptr[i + 1]] = mask_flatten
|
||||
|
||||
extend_attention_fwd(
|
||||
q_extend,
|
||||
k_extend,
|
||||
v_extend,
|
||||
o_extend_mask,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
custom_mask,
|
||||
True,
|
||||
mask_indptr,
|
||||
max_len_extend,
|
||||
)
|
||||
|
||||
redundant_attention(
|
||||
q_extend,
|
||||
o_redundant,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
b_req_idx,
|
||||
b_start_loc,
|
||||
b_seq_len,
|
||||
b_seq_len_prefix,
|
||||
max_len_in_batch,
|
||||
)
|
||||
|
||||
self.assertTrue(torch.allclose(o_extend, o_redundant, rtol=1e-2))
|
||||
self.assertTrue(torch.allclose(o_extend_mask, o_redundant, rtol=1e-2))
|
||||
|
||||
def test_extend_attention(self):
|
||||
|
||||
# Define the varying parameter values
|
||||
attention_values = [128, 96, 80, 13]
|
||||
|
||||
# Loop through the values and call the method
|
||||
for value in attention_values:
|
||||
self._test_extend_attention_once(19, 12331, 12, 4, value)
|
||||
|
||||
def _test_extend_attention_sliding_window_once(
|
||||
self, B, N_CTX, H_Q, H_KV, D, WINDOW_SIZE
|
||||
):
|
||||
dtype = torch.bfloat16
|
||||
device = get_device()
|
||||
|
||||
b_seq_len_prefix = torch.randint(
|
||||
1, N_CTX // 2, (B,), dtype=torch.int32, device=device
|
||||
)
|
||||
b_seq_len_extend = torch.randint(
|
||||
1, N_CTX // 2, (B,), dtype=torch.int32, device=device
|
||||
)
|
||||
b_seq_len = b_seq_len_prefix + b_seq_len_extend
|
||||
|
||||
b_start_loc = torch.zeros((B,), dtype=torch.int32, device=device)
|
||||
b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
|
||||
b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device=device)
|
||||
b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
|
||||
|
||||
kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
|
||||
kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0)
|
||||
kv_indices = torch.zeros(
|
||||
(b_seq_len_prefix.sum().item(),), dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
for i in range(B):
|
||||
kv_indices[kv_indptr[i] : kv_indptr[i + 1]] = torch.arange(
|
||||
b_start_loc[i], b_start_loc[i] + b_seq_len_prefix[i]
|
||||
)
|
||||
|
||||
total_token_num = torch.sum(b_seq_len).item()
|
||||
extend_token_num = torch.sum(b_seq_len_extend).item()
|
||||
k_buffer = torch.empty(
|
||||
(total_token_num, H_KV, D), dtype=dtype, device=device
|
||||
).normal_(mean=0.1, std=0.2)
|
||||
v_buffer = torch.empty(
|
||||
(total_token_num, H_KV, D), dtype=dtype, device=device
|
||||
).normal_(mean=0.1, std=0.2)
|
||||
|
||||
k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device)
|
||||
v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device)
|
||||
q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device)
|
||||
for i in range(B):
|
||||
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
|
||||
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
|
||||
extend_start = b_start_loc_extend[i]
|
||||
extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
|
||||
k_extend[extend_start:extend_end] = k_buffer[
|
||||
extend_start_in_buffer:extend_end_in_buffer
|
||||
]
|
||||
v_extend[extend_start:extend_end] = v_buffer[
|
||||
extend_start_in_buffer:extend_end_in_buffer
|
||||
]
|
||||
q_extend[extend_start:extend_end] = torch.empty(
|
||||
(b_seq_len_extend[i], H_Q, D), dtype=dtype, device=device
|
||||
).normal_(mean=0.1, std=0.2)
|
||||
|
||||
o_extend_triton = torch.empty(
|
||||
(extend_token_num, H_Q, D), dtype=dtype, device=device
|
||||
)
|
||||
o_extend_torch = torch.empty(
|
||||
(extend_token_num, H_Q, D), dtype=dtype, device=device
|
||||
)
|
||||
|
||||
b_seq_len_extend = b_seq_len - b_seq_len_prefix
|
||||
max_len_extend = torch.max(b_seq_len_extend, 0)[0].item()
|
||||
qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
|
||||
qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0)
|
||||
|
||||
extend_attention_fwd(
|
||||
q_extend,
|
||||
k_extend,
|
||||
v_extend,
|
||||
o_extend_triton,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
custom_mask=None,
|
||||
is_causal=True,
|
||||
mask_indptr=None,
|
||||
max_len_extend=max_len_extend,
|
||||
sliding_window_size=WINDOW_SIZE,
|
||||
)
|
||||
|
||||
extend_attention_fwd_torch(
|
||||
q_extend,
|
||||
k_extend,
|
||||
v_extend,
|
||||
o_extend_torch,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
WINDOW_SIZE,
|
||||
)
|
||||
|
||||
self.assertTrue(
|
||||
torch.allclose(o_extend_triton, o_extend_torch, rtol=1e-3, atol=1e-3)
|
||||
)
|
||||
|
||||
def test_extend_attention_sliding_window(self):
|
||||
window_sizes = [-1, 127]
|
||||
for window_size in window_sizes:
|
||||
self._test_extend_attention_sliding_window_once(
|
||||
19, 12331, 64, 8, 128, window_size
|
||||
)
|
||||
|
||||
def _test_context_attention_once(self, head_dim, is_causal):
|
||||
# Set up a simple test case
|
||||
device = get_device()
|
||||
num_heads = 4
|
||||
seq_lens = [8, 12]
|
||||
max_seq_len = max(seq_lens)
|
||||
|
||||
# Create random input tensors
|
||||
q = torch.randn(sum(seq_lens), num_heads, head_dim, device=device)
|
||||
k = torch.randn(sum(seq_lens), num_heads, head_dim, device=device)
|
||||
v = torch.randn(sum(seq_lens), num_heads, head_dim, device=device)
|
||||
o = torch.zeros(sum(seq_lens), num_heads, head_dim, device=device)
|
||||
|
||||
# Create b_start_loc and b_seq_len tensors
|
||||
b_start_loc = torch.tensor([0, seq_lens[0]], device=device)
|
||||
b_seq_len = torch.tensor(seq_lens, device=device)
|
||||
|
||||
context_attention_fwd(
|
||||
q, k, v, o, b_start_loc, b_seq_len, max_seq_len, is_causal=is_causal
|
||||
)
|
||||
|
||||
cu_seq_lens = [0] * (len(seq_lens) + 1)
|
||||
for i, seq_len in enumerate(seq_lens):
|
||||
cu_seq_lens[i + 1] = cu_seq_lens[i] + seq_len
|
||||
|
||||
for i in range(len(seq_lens)):
|
||||
start, end = cu_seq_lens[i], cu_seq_lens[i + 1]
|
||||
o_torch = torch.nn.functional.scaled_dot_product_attention(
|
||||
q[start:end].permute(1, 0, 2),
|
||||
k[start:end].permute(1, 0, 2),
|
||||
v[start:end].permute(1, 0, 2),
|
||||
is_causal=is_causal,
|
||||
).permute(1, 0, 2)
|
||||
|
||||
cos_sim = torch.nn.functional.cosine_similarity(
|
||||
o[start:end].flatten(), o_torch.flatten(), dim=0
|
||||
)
|
||||
self.assertTrue(cos_sim.item() > 1 - (1e-5))
|
||||
self.assertTrue(torch.allclose(o[start:end], o_torch, atol=1e-2))
|
||||
|
||||
def test_context_attention(self):
|
||||
head_dim = [128, 96, 80, 13]
|
||||
|
||||
for dim in head_dim:
|
||||
for is_causal in [True, False]:
|
||||
self._test_context_attention_once(dim, is_causal)
|
||||
|
||||
def _test_decode_attention_once(self, B, H_Q, H_KV, D):
|
||||
device = get_device()
|
||||
dtype = torch.bfloat16
|
||||
seq_len = 10 # This represents the number of tokens already in the sequence
|
||||
total_tokens = B * seq_len
|
||||
sm_scale = 1.0 / (D**0.5)
|
||||
max_kv_splits = 8
|
||||
num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device=device)
|
||||
|
||||
# q represents the new token being generated, one per batch
|
||||
q = torch.randn(B, H_Q, D, dtype=dtype, device=device)
|
||||
|
||||
# k_buffer and v_buffer represent all previous tokens
|
||||
k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device)
|
||||
v_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device)
|
||||
|
||||
# o will have the same shape as q
|
||||
o = torch.zeros(B, H_Q, D, dtype=dtype, device=device)
|
||||
|
||||
b_seq_len = torch.full((B,), seq_len, device=device)
|
||||
|
||||
kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
|
||||
kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len[:B], dim=0)
|
||||
kv_indices = torch.arange(total_tokens, device=device)
|
||||
|
||||
attn_logits = torch.empty(
|
||||
(B, H_Q, max_kv_splits, D),
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
)
|
||||
attn_lse = torch.empty(
|
||||
(B, H_Q, max_kv_splits),
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
decode_attention_fwd(
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
o,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
attn_logits,
|
||||
attn_lse,
|
||||
num_kv_splits,
|
||||
max_kv_splits,
|
||||
sm_scale,
|
||||
)
|
||||
|
||||
def test_decode_attention(self):
|
||||
# Here we just to ensure there is no error
|
||||
# TODO: correctnesss test
|
||||
|
||||
# Test configurations
|
||||
configs = [
|
||||
(2, 4, 4, 64), # MHA
|
||||
(2, 4, 2, 64), # GQA
|
||||
(2, 4, 4, 80), # Non-standard head dim
|
||||
(2, 4, 4, 13), # Prime number head dim
|
||||
]
|
||||
|
||||
for B, H_Q, H_KV, D in configs:
|
||||
self._test_decode_attention_once(B, H_Q, H_KV, D)
|
||||
|
||||
def _test_grouped_decode_attention_once(self, B, S, H_Q, H_KV, D, D_V):
|
||||
dtype = torch.bfloat16
|
||||
device = get_device()
|
||||
seq_len = S # This represents the number of tokens already in the sequence
|
||||
total_tokens = B * seq_len
|
||||
sm_scale = 1.0 / (D**0.5)
|
||||
max_kv_splits = 8
|
||||
num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device=device)
|
||||
|
||||
# q represents the new token being generated, one per batch
|
||||
q = torch.randn(B, H_Q, D, dtype=dtype, device=device)
|
||||
|
||||
# k_buffer and v_buffer represent all previous tokens
|
||||
k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device)
|
||||
v_buffer = torch.randn(total_tokens, H_KV, D_V, dtype=dtype, device=device)
|
||||
|
||||
# o will have the same shape as q
|
||||
o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=device)
|
||||
o_grouped = torch.zeros(B, H_Q, D_V, dtype=dtype, device=device)
|
||||
|
||||
b_seq_len = torch.full((B,), seq_len, device=device)
|
||||
|
||||
kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
|
||||
kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len[:B], dim=0)
|
||||
kv_indices = torch.arange(total_tokens, device=device)
|
||||
|
||||
attn_logits = torch.empty(
|
||||
(B, H_Q, max_kv_splits, D_V),
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
)
|
||||
attn_lse = torch.empty(
|
||||
(B, H_Q, max_kv_splits),
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
decode_attention_fwd_normal(
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
o,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
attn_logits,
|
||||
attn_lse,
|
||||
num_kv_splits,
|
||||
max_kv_splits,
|
||||
sm_scale,
|
||||
)
|
||||
|
||||
attn_logits1 = torch.empty(
|
||||
(B, H_Q, max_kv_splits, D_V),
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
)
|
||||
attn_lse1 = torch.empty(
|
||||
(B, H_Q, max_kv_splits, D_V),
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
decode_attention_fwd_grouped(
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
o_grouped,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
attn_logits1,
|
||||
attn_lse1,
|
||||
num_kv_splits,
|
||||
max_kv_splits,
|
||||
sm_scale,
|
||||
)
|
||||
|
||||
cos_sim = torch.nn.functional.cosine_similarity(
|
||||
o.flatten(), o_grouped.flatten(), dim=0
|
||||
)
|
||||
print(cos_sim.item())
|
||||
self.assertTrue(cos_sim.item() > 0.99)
|
||||
self.assertTrue(torch.allclose(o, o_grouped, atol=3e-2))
|
||||
|
||||
def test_grouped_decode_attention(self):
|
||||
seq_lens = [5, 100, 128, 500]
|
||||
configs = [
|
||||
(2, 16, 16, 64, 64),
|
||||
(2, 16, 1, 64, 64),
|
||||
(2, 64, 1, 13, 13),
|
||||
(2, 128, 1, 80, 80),
|
||||
(2, 128, 2, 512, 512),
|
||||
(2, 128, 1, 576, 512),
|
||||
]
|
||||
|
||||
for S in seq_lens:
|
||||
for B, H_Q, H_KV, D, D_V in configs:
|
||||
self._test_grouped_decode_attention_once(B, S, H_Q, H_KV, D, D_V)
|
||||
|
||||
def _test_extend_attention_unified_vs_regular_once(self, B, N_CTX, H_Q, H_KV, D):
|
||||
"""Test that unified kernel produces same results as 2-stage kernel."""
|
||||
dtype = torch.bfloat16
|
||||
device = get_device()
|
||||
|
||||
b_seq_len_prefix = torch.randint(
|
||||
1, N_CTX // 2, (B,), dtype=torch.int32, device=device
|
||||
)
|
||||
b_seq_len_extend = torch.randint(
|
||||
1, N_CTX // 2, (B,), dtype=torch.int32, device=device
|
||||
)
|
||||
b_seq_len = b_seq_len_prefix + b_seq_len_extend
|
||||
|
||||
b_start_loc = torch.zeros((B,), dtype=torch.int32, device=device)
|
||||
b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
|
||||
b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device=device)
|
||||
b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
|
||||
|
||||
# Setup prefix KV indices
|
||||
kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
|
||||
kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0)
|
||||
kv_indices = torch.zeros(
|
||||
(b_seq_len_prefix.sum().item(),), dtype=torch.int64, device=device
|
||||
)
|
||||
|
||||
for i in range(B):
|
||||
kv_indices[kv_indptr[i] : kv_indptr[i + 1]] = torch.arange(
|
||||
b_start_loc[i], b_start_loc[i] + b_seq_len_prefix[i]
|
||||
)
|
||||
|
||||
total_token_num = torch.sum(b_seq_len).item()
|
||||
extend_token_num = torch.sum(b_seq_len_extend).item()
|
||||
k_buffer = torch.empty(
|
||||
(total_token_num, H_KV, D), dtype=dtype, device=device
|
||||
).normal_(mean=0.1, std=0.2)
|
||||
v_buffer = torch.empty(
|
||||
(total_token_num, H_KV, D), dtype=dtype, device=device
|
||||
).normal_(mean=0.1, std=0.2)
|
||||
|
||||
k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device)
|
||||
v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device=device)
|
||||
q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device)
|
||||
|
||||
for i in range(B):
|
||||
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
|
||||
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
|
||||
extend_start = b_start_loc_extend[i]
|
||||
extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
|
||||
k_extend[extend_start:extend_end] = k_buffer[
|
||||
extend_start_in_buffer:extend_end_in_buffer
|
||||
]
|
||||
v_extend[extend_start:extend_end] = v_buffer[
|
||||
extend_start_in_buffer:extend_end_in_buffer
|
||||
]
|
||||
q_extend[extend_start:extend_end] = torch.empty(
|
||||
(b_seq_len_extend[i], H_Q, D), dtype=dtype, device=device
|
||||
).normal_(mean=0.1, std=0.2)
|
||||
|
||||
# Setup for extend attention
|
||||
max_len_extend = torch.max(b_seq_len_extend, 0)[0].item()
|
||||
qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
|
||||
qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0)
|
||||
|
||||
# Run 2-stage kernel
|
||||
o_regular = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device)
|
||||
extend_attention_fwd(
|
||||
q_extend,
|
||||
k_extend,
|
||||
v_extend,
|
||||
o_regular,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
custom_mask=None,
|
||||
is_causal=True,
|
||||
mask_indptr=None,
|
||||
max_len_extend=max_len_extend,
|
||||
)
|
||||
|
||||
# Build unified KV indices
|
||||
extend_kv_indices = torch.arange(
|
||||
total_token_num - extend_token_num,
|
||||
total_token_num,
|
||||
dtype=torch.int64,
|
||||
device=device,
|
||||
)
|
||||
extend_start_loc = torch.zeros((B,), dtype=torch.int32, device=device)
|
||||
extend_start_loc[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
|
||||
|
||||
unified_kv_indptr, unified_kv_indices, prefix_lens = build_unified_kv_indices(
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
extend_start_loc,
|
||||
b_seq_len_extend,
|
||||
extend_kv_indices,
|
||||
B,
|
||||
)
|
||||
|
||||
# Run unified kernel
|
||||
o_unified = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device=device)
|
||||
extend_attention_fwd_unified(
|
||||
q_extend,
|
||||
o_unified,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
qo_indptr,
|
||||
unified_kv_indptr,
|
||||
unified_kv_indices,
|
||||
prefix_lens,
|
||||
max_len_extend=max_len_extend,
|
||||
custom_mask=None,
|
||||
mask_indptr=None,
|
||||
sm_scale=None,
|
||||
logit_cap=0.0,
|
||||
is_causal=True,
|
||||
)
|
||||
|
||||
# Compare results
|
||||
self.assertTrue(
|
||||
torch.allclose(o_regular, o_unified, rtol=0.15, atol=0.15),
|
||||
f"Unified kernel output differs from 2-stage kernel. "
|
||||
f"Max diff: {(o_regular - o_unified).abs().max()}",
|
||||
)
|
||||
|
||||
def test_extend_attention_unified_vs_regular(self):
|
||||
"""Test unified kernel matches 2-stage kernel across different configs."""
|
||||
configs = [
|
||||
(4, 512, 32, 8, 128), # Standard config
|
||||
(2, 2048, 32, 8, 128), # Long sequence (test 2048 specifically)
|
||||
(8, 256, 64, 8, 80), # Non-standard head dim
|
||||
]
|
||||
|
||||
for B, N_CTX, H_Q, H_KV, D in configs:
|
||||
with self.subTest(B=B, N_CTX=N_CTX, H_Q=H_Q, H_KV=H_KV, D=D):
|
||||
self._test_extend_attention_unified_vs_regular_once(
|
||||
B, N_CTX, H_Q, H_KV, D
|
||||
)
|
||||
|
||||
def test_build_unified_kv_indices(self):
|
||||
"""Test build_unified_kv_indices correctness."""
|
||||
B = 4
|
||||
dtype = torch.int64
|
||||
device = get_device()
|
||||
|
||||
# Setup test data
|
||||
prefix_lens = torch.tensor([10, 20, 15, 25], dtype=torch.int32, device=device)
|
||||
extend_lens = torch.tensor([5, 3, 7, 4], dtype=torch.int32, device=device)
|
||||
|
||||
# Build prefix indices
|
||||
prefix_kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
|
||||
prefix_kv_indptr[1:] = torch.cumsum(prefix_lens, dim=0)
|
||||
prefix_kv_indices = torch.arange(
|
||||
prefix_lens.sum().item(), dtype=dtype, device=device
|
||||
)
|
||||
|
||||
# Build extend indices
|
||||
extend_start_loc = torch.zeros((B,), dtype=torch.int32, device=device)
|
||||
extend_start_loc[1:] = torch.cumsum(extend_lens[:-1], dim=0)
|
||||
extend_kv_indices = torch.arange(
|
||||
prefix_lens.sum().item(),
|
||||
prefix_lens.sum().item() + extend_lens.sum().item(),
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
# Build unified indices
|
||||
unified_kv_indptr, unified_kv_indices, returned_prefix_lens = (
|
||||
build_unified_kv_indices(
|
||||
prefix_kv_indptr,
|
||||
prefix_kv_indices,
|
||||
extend_start_loc,
|
||||
extend_lens,
|
||||
extend_kv_indices,
|
||||
B,
|
||||
)
|
||||
)
|
||||
|
||||
# Verify unified_kv_indptr
|
||||
expected_lens = prefix_lens + extend_lens
|
||||
expected_indptr = torch.zeros((B + 1,), dtype=torch.int32, device=device)
|
||||
expected_indptr[1:] = torch.cumsum(expected_lens, dim=0)
|
||||
self.assertTrue(torch.equal(unified_kv_indptr, expected_indptr))
|
||||
|
||||
# Verify prefix_lens
|
||||
self.assertTrue(torch.equal(returned_prefix_lens, prefix_lens))
|
||||
|
||||
# Verify unified_kv_indices structure
|
||||
for i in range(B):
|
||||
start_idx = int(unified_kv_indptr[i])
|
||||
end_idx = int(unified_kv_indptr[i + 1])
|
||||
prefix_len = int(prefix_lens[i])
|
||||
extend_len = int(extend_lens[i])
|
||||
|
||||
# Check that prefix and extend are concatenated correctly
|
||||
unified_seq = unified_kv_indices[start_idx:end_idx]
|
||||
self.assertEqual(len(unified_seq), prefix_len + extend_len)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,129 +0,0 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
import requests
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
is_in_ci,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestSlidingWindowAttentionTriton(CustomTestCase):
|
||||
"""Test sliding window attention functionality with triton backend."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
"""Set up the test server with Gemma3 model and triton backend."""
|
||||
# Gemma3 model supports sliding window attention
|
||||
cls.model = "google/gemma-3-4b-it"
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
|
||||
cls.common_args = [
|
||||
"--trust-remote-code",
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--context-length",
|
||||
"8192",
|
||||
"--random-seed",
|
||||
"42",
|
||||
]
|
||||
|
||||
cls.short_context_prompt = "The capital of France is"
|
||||
|
||||
# Test prompt longer than window size
|
||||
cls.long_context_prompt = (
|
||||
"""
|
||||
Once upon a time, there was a mountain. In the mountain, there was a temple. In the temple, there was an old monk telling a story. The story was:
|
||||
"""
|
||||
* 100
|
||||
)
|
||||
cls.long_context_prompt += "\nNow, summarize the story in one sentence:"
|
||||
|
||||
def _test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=200,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
print(f"MMLU metrics with sliding window: {metrics}")
|
||||
|
||||
self.assertGreaterEqual(metrics["score"], 0.60)
|
||||
|
||||
def _test_short_context_generation(self):
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": self.short_context_prompt,
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 256,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
result = response.json()
|
||||
self.assertIn("paris", result["text"].lower())
|
||||
print(f"Short context generation result: {result['text']}")
|
||||
|
||||
def _test_long_context_generation(self):
|
||||
response = requests.post(
|
||||
self.base_url + "/generate",
|
||||
json={
|
||||
"text": self.long_context_prompt,
|
||||
"sampling_params": {
|
||||
"temperature": 0,
|
||||
"max_new_tokens": 256,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(response.status_code, 200)
|
||||
result = response.json()
|
||||
self.assertGreater(len(result["text"].strip()), 0)
|
||||
print(f"Long context generation result: {result['text'][:100]}...")
|
||||
|
||||
@unittest.skipIf(is_in_ci(), "To reduce the CI execution time.")
|
||||
def test_no_cuda_graph(self):
|
||||
self.no_cuda_graph_process = popen_launch_server(
|
||||
self.model,
|
||||
self.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=self.common_args + ["--disable-cuda-graph"],
|
||||
)
|
||||
|
||||
try:
|
||||
self._test_short_context_generation()
|
||||
self._test_long_context_generation()
|
||||
self._test_mmlu()
|
||||
finally:
|
||||
kill_process_tree(self.no_cuda_graph_process.pid)
|
||||
|
||||
def test_cuda_graph(self):
|
||||
self.cuda_graph_process = popen_launch_server(
|
||||
self.model,
|
||||
self.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=self.common_args,
|
||||
)
|
||||
|
||||
try:
|
||||
self._test_short_context_generation()
|
||||
self._test_long_context_generation()
|
||||
self._test_mmlu()
|
||||
finally:
|
||||
kill_process_tree(self.cuda_graph_process.pid)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -1,322 +0,0 @@
|
||||
import random
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.triton_ops.decode_attention import (
|
||||
decode_attention_fwd_grouped as triton_decode_attention_fwd_grouped,
|
||||
)
|
||||
from sglang.srt.layers.attention.triton_ops.extend_attention import (
|
||||
extend_attention_fwd,
|
||||
redundant_attention,
|
||||
)
|
||||
from sglang.srt.layers.attention.triton_ops.prefill_attention import (
|
||||
context_attention_fwd,
|
||||
)
|
||||
from sglang.srt.layers.attention.wave_ops.decode_attention import (
|
||||
decode_attention_intermediate_arrays_shapes,
|
||||
decode_attention_wave,
|
||||
)
|
||||
from sglang.srt.layers.attention.wave_ops.extend_attention import extend_attention_wave
|
||||
from sglang.srt.layers.attention.wave_ops.prefill_attention import (
|
||||
prefill_attention_wave,
|
||||
)
|
||||
|
||||
|
||||
class TestWaveAttention(unittest.TestCase):
|
||||
|
||||
def _set_all_seeds(self, seed):
|
||||
"""Set all random seeds for reproducibility."""
|
||||
random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
torch.backends.cudnn.benchmark = False
|
||||
|
||||
def setUp(self):
|
||||
# Set seeds before each test method
|
||||
self._set_all_seeds(42)
|
||||
|
||||
def _test_extend_attention_once(self, B, N_CTX, H_Q, H_KV, D):
|
||||
dtype = torch.float16
|
||||
extend_seq_len = 1024
|
||||
|
||||
b_seq_len_prefix = torch.full(
|
||||
(B,), N_CTX // B, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
b_seq_len_extend = torch.full(
|
||||
(B,), extend_seq_len, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
b_seq_len = b_seq_len_prefix + b_seq_len_extend
|
||||
max_len_in_batch = torch.max(b_seq_len, 0)[0].item()
|
||||
|
||||
b_req_idx = torch.arange(B, dtype=torch.int32, device="cuda")
|
||||
b_start_loc = torch.zeros((B,), dtype=torch.int32, device="cuda")
|
||||
b_start_loc[1:] = torch.cumsum(b_seq_len[:-1], 0)
|
||||
b_start_loc_extend = torch.zeros((B,), dtype=torch.int32, device="cuda")
|
||||
b_start_loc_extend[1:] = torch.cumsum(b_seq_len_extend[:-1], 0)
|
||||
|
||||
kv_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda")
|
||||
kv_indptr[1 : B + 1] = torch.cumsum(b_seq_len_prefix[:B], dim=0)
|
||||
kv_indices = torch.zeros(
|
||||
(b_seq_len_prefix.sum().item(),), dtype=torch.int32, device="cuda"
|
||||
)
|
||||
|
||||
for i in range(B):
|
||||
kv_indices[kv_indptr[i] : kv_indptr[i + 1]] = torch.arange(
|
||||
b_start_loc[i], b_start_loc[i] + b_seq_len_prefix[i]
|
||||
)
|
||||
|
||||
total_token_num = torch.sum(b_seq_len).item()
|
||||
extend_token_num = torch.sum(b_seq_len_extend).item()
|
||||
k_buffer = torch.empty(
|
||||
(total_token_num, H_KV, D), dtype=dtype, device="cuda"
|
||||
).normal_(mean=0.1, std=0.2)
|
||||
v_buffer = torch.empty(
|
||||
(total_token_num, H_KV, D), dtype=dtype, device="cuda"
|
||||
).normal_(mean=0.1, std=0.2)
|
||||
|
||||
k_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device="cuda")
|
||||
v_extend = torch.empty((extend_token_num, H_KV, D), dtype=dtype, device="cuda")
|
||||
q_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda")
|
||||
for i in range(B):
|
||||
extend_start_in_buffer = b_start_loc[i] + b_seq_len_prefix[i]
|
||||
extend_end_in_buffer = b_start_loc[i] + b_seq_len[i]
|
||||
extend_start = b_start_loc_extend[i]
|
||||
extend_end = b_start_loc_extend[i] + b_seq_len_extend[i]
|
||||
k_extend[extend_start:extend_end] = k_buffer[
|
||||
extend_start_in_buffer:extend_end_in_buffer
|
||||
]
|
||||
v_extend[extend_start:extend_end] = v_buffer[
|
||||
extend_start_in_buffer:extend_end_in_buffer
|
||||
]
|
||||
q_extend[extend_start:extend_end] = torch.empty(
|
||||
(b_seq_len_extend[i], H_Q, D), dtype=dtype, device="cuda"
|
||||
).normal_(mean=0.1, std=0.2)
|
||||
|
||||
o_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda")
|
||||
o_extend_mask = torch.empty(
|
||||
(extend_token_num, H_Q, D), dtype=dtype, device="cuda"
|
||||
)
|
||||
o_redundant = torch.empty(
|
||||
(extend_token_num, H_Q, D), dtype=dtype, device="cuda"
|
||||
)
|
||||
|
||||
b_seq_len_extend = b_seq_len - b_seq_len_prefix
|
||||
max_len_extend = torch.max(b_seq_len_extend, 0)[0].item()
|
||||
qo_indptr = torch.zeros((B + 1,), dtype=torch.int32, device="cuda")
|
||||
qo_indptr[1 : B + 1] = torch.cumsum(b_seq_len_extend[:B], dim=0)
|
||||
|
||||
custom_mask = None
|
||||
mask_indptr = None
|
||||
|
||||
redundant_attention(
|
||||
q_extend,
|
||||
o_redundant,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
b_req_idx,
|
||||
b_start_loc,
|
||||
b_seq_len,
|
||||
b_seq_len_prefix,
|
||||
max_len_in_batch,
|
||||
)
|
||||
|
||||
is_causal = True
|
||||
|
||||
o_extend = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda")
|
||||
extend_attention_fwd(
|
||||
q_extend,
|
||||
k_extend,
|
||||
v_extend,
|
||||
o_extend,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
custom_mask,
|
||||
is_causal,
|
||||
mask_indptr,
|
||||
max_len_extend,
|
||||
)
|
||||
|
||||
o_wave = torch.empty((extend_token_num, H_Q, D), dtype=dtype, device="cuda")
|
||||
extend_attention_wave(
|
||||
q_extend,
|
||||
k_extend,
|
||||
v_extend,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
custom_mask,
|
||||
mask_indptr,
|
||||
max_len_extend,
|
||||
o_wave,
|
||||
is_causal=is_causal,
|
||||
)
|
||||
|
||||
self.assertTrue(torch.allclose(o_extend, o_redundant, rtol=1e-2))
|
||||
self.assertTrue(torch.allclose(o_wave, o_redundant, rtol=1e-2))
|
||||
|
||||
def test_extend_attention(self):
|
||||
|
||||
# Define the varying parameter values
|
||||
attention_values = [128]
|
||||
|
||||
# Loop through the values and call the method
|
||||
for value in attention_values:
|
||||
self._test_extend_attention_once(32, 16384, 6, 1, value)
|
||||
|
||||
def _test_grouped_decode_attention_once(self, B, S, H_Q, H_KV, D, D_V):
|
||||
dtype = torch.float16
|
||||
seq_len = S # This represents the number of tokens already in the sequence
|
||||
total_tokens = B * seq_len
|
||||
sm_scale = 1.0 / (D**0.5)
|
||||
max_kv_splits = 8
|
||||
num_kv_splits = torch.full((B,), 4, dtype=torch.int32, device="cuda")
|
||||
|
||||
# q represents the new token being generated, one per batch
|
||||
q = torch.randn(B, H_Q, D, dtype=dtype, device="cuda")
|
||||
|
||||
# k_buffer and v_buffer represent all previous tokens
|
||||
k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device="cuda")
|
||||
v_buffer = torch.randn(total_tokens, H_KV, D_V, dtype=dtype, device="cuda")
|
||||
|
||||
# o will have the same shape as q
|
||||
o_triton = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
|
||||
o = torch.zeros(B, H_Q, D_V, dtype=dtype, device="cuda")
|
||||
|
||||
req_to_token = torch.arange(total_tokens, device="cuda", dtype=torch.int32)
|
||||
b_req_idx = torch.zeros(B + 1, device="cuda", dtype=torch.int32)
|
||||
b_seq_len = torch.full((B,), seq_len, device="cuda", dtype=torch.int32)
|
||||
b_req_idx[1 : B + 1] = torch.cumsum(b_seq_len, dim=0)
|
||||
|
||||
attn_logits = torch.empty(
|
||||
(B, H_Q, max_kv_splits, D_V + 1),
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
)
|
||||
attn_lse = torch.empty(
|
||||
(B, H_Q, max_kv_splits),
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
logit_cap = 0.0
|
||||
triton_decode_attention_fwd_grouped(
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
o_triton,
|
||||
b_req_idx,
|
||||
req_to_token,
|
||||
attn_logits,
|
||||
attn_lse,
|
||||
num_kv_splits,
|
||||
max_kv_splits,
|
||||
sm_scale,
|
||||
logit_cap,
|
||||
)
|
||||
|
||||
attn_logits_shape, attn_logits_max_shape = (
|
||||
decode_attention_intermediate_arrays_shapes(B, D_V, H_Q, max_kv_splits)
|
||||
)
|
||||
|
||||
attn_logits = torch.empty(
|
||||
attn_logits_shape,
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
attn_logits_max = torch.empty(
|
||||
attn_logits_max_shape,
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
decode_attention_wave(
|
||||
q,
|
||||
k_buffer,
|
||||
v_buffer,
|
||||
o,
|
||||
b_req_idx,
|
||||
req_to_token,
|
||||
attn_logits,
|
||||
attn_logits_max,
|
||||
num_kv_splits,
|
||||
max_kv_splits,
|
||||
sm_scale,
|
||||
logit_cap,
|
||||
)
|
||||
|
||||
cos_sim = torch.nn.functional.cosine_similarity(
|
||||
o.flatten(), o_triton.flatten(), dim=0
|
||||
)
|
||||
print(cos_sim.item())
|
||||
self.assertTrue(cos_sim.item() > 0.99)
|
||||
self.assertTrue(torch.allclose(o, o_triton, atol=3e-2))
|
||||
|
||||
def test_grouped_decode_attention(self):
|
||||
seq_lens = [5, 100, 128, 500]
|
||||
configs = [
|
||||
(2, 16, 16, 64, 64),
|
||||
(2, 16, 1, 64, 64),
|
||||
(2, 128, 1, 80, 80),
|
||||
(32, 128, 2, 512, 512),
|
||||
(2, 128, 2, 512, 512),
|
||||
(2, 128, 1, 576, 512),
|
||||
]
|
||||
|
||||
for S in seq_lens:
|
||||
for B, H_Q, H_KV, D, D_V in configs:
|
||||
self._test_grouped_decode_attention_once(B, S, H_Q, H_KV, D, D_V)
|
||||
|
||||
def _test_context_attention_once(self, head_dim, is_causal):
|
||||
# Set up a simple test case
|
||||
dtype = torch.float16
|
||||
num_heads = 4
|
||||
kv_heads = 1
|
||||
seq_lens = [128, 256]
|
||||
max_seq_len = max(seq_lens)
|
||||
|
||||
# Create random input tensors
|
||||
q = torch.randn(sum(seq_lens), num_heads, head_dim, dtype=dtype, device="cuda")
|
||||
k = torch.randn(sum(seq_lens), kv_heads, head_dim, dtype=dtype, device="cuda")
|
||||
v = torch.randn(sum(seq_lens), kv_heads, head_dim, dtype=dtype, device="cuda")
|
||||
o_triton = torch.zeros(
|
||||
sum(seq_lens), num_heads, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
o = torch.zeros(sum(seq_lens), num_heads, head_dim, dtype=dtype, device="cuda")
|
||||
|
||||
# Create b_start_loc and b_seq_len tensors
|
||||
b_start_loc = torch.tensor([0, seq_lens[0]], device="cuda")
|
||||
b_seq_len = torch.tensor(seq_lens, device="cuda")
|
||||
|
||||
context_attention_fwd(
|
||||
q, k, v, o_triton, b_start_loc, b_seq_len, max_seq_len, is_causal=is_causal
|
||||
)
|
||||
prefill_attention_wave(
|
||||
q, k, v, o, b_start_loc, b_seq_len, max_seq_len, is_causal=is_causal
|
||||
)
|
||||
cos_sim = torch.nn.functional.cosine_similarity(
|
||||
o.flatten(), o_triton.flatten(), dim=0
|
||||
)
|
||||
|
||||
print(cos_sim.item())
|
||||
self.assertTrue(torch.allclose(o, o_triton, atol=3e-2))
|
||||
self.assertTrue(cos_sim.item() > 1 - (1e-5))
|
||||
|
||||
def test_context_attention(self):
|
||||
head_dim = [128, 96]
|
||||
|
||||
for dim in head_dim:
|
||||
for is_causal in [False]:
|
||||
self._test_context_attention_once(dim, is_causal)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user