feat(mem_cache): page-major (layer-major within a page) KV/state layout (#29533)
Co-authored-by: lch1475369 <lch1475369@gmail.com>
This commit is contained in:
Cheng Wan
@@ -0,0 +1,69 @@
|
||||
"""
|
||||
End-to-end accuracy test for the page-major KV layout on a hybrid-SWA MoE model.
|
||||
|
||||
Launches gpt-oss-20b with ``--enable-page-major-kv-layout`` on the Triton
|
||||
attention backend and checks that GSM8K accuracy holds. This exercises the
|
||||
SWA + full-attention KV pools under the page-granularity envelope layout
|
||||
(SWAKVPool routes both sub-pools through PageMajorMHATokenToKVPool).
|
||||
|
||||
Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
|
||||
|
||||
Usage:
|
||||
python3 -m unittest test_page_major_gpt_oss
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
|
||||
from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE
|
||||
|
||||
register_cuda_ci(est_time=420, stage="extra-a", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
class TestPageMajorGptOss(DefaultServerBase):
|
||||
"""Page-major KV layout on gpt-oss-20b (hybrid-SWA MoE), Triton backend."""
|
||||
|
||||
model = DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE
|
||||
|
||||
gsm8k_threshold = 0.45
|
||||
num_gsm8k_questions = 200
|
||||
num_shots = 5
|
||||
parallel = 32
|
||||
|
||||
other_args = [
|
||||
"--enable-page-major-kv-layout",
|
||||
# The envelope's strided 4-D K/V views are only read by the Triton
|
||||
# attention kernels (the layout's validator enforces this).
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--mem-fraction-static",
|
||||
"0.70",
|
||||
"--cuda-graph-backend-prefill=disabled",
|
||||
]
|
||||
|
||||
def test_gsm8k(self):
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_few_shot_gsm8k
|
||||
|
||||
url = urlparse(self.base_url)
|
||||
args = SimpleNamespace(
|
||||
num_shots=self.num_shots,
|
||||
data_path=None,
|
||||
num_questions=self.num_gsm8k_questions,
|
||||
max_new_tokens=512,
|
||||
parallel=self.parallel,
|
||||
host=f"http://{url.hostname}",
|
||||
port=int(url.port),
|
||||
)
|
||||
metrics = run_few_shot_gsm8k(args)
|
||||
print(
|
||||
f"[{self.__class__.__name__}] GSM8K accuracy: {metrics['accuracy']:.3f} "
|
||||
f"(threshold: {self.gsm8k_threshold})"
|
||||
)
|
||||
self.assertGreaterEqual(metrics["accuracy"], self.gsm8k_threshold)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,79 @@
|
||||
"""
|
||||
End-to-end accuracy test for the page-major KV layout on a GDN-hybrid model.
|
||||
|
||||
Launches Qwen3.5-4B (a gated-delta-net / linear-attention hybrid) with
|
||||
``--enable-page-major-kv-layout`` on the Triton attention + linear-attn + Mamba
|
||||
backends and checks that GSM8K accuracy holds. This exercises the page-major
|
||||
path most prone to subtle bugs: the Mamba conv/SSM state stored as a strided
|
||||
envelope view, plus the full-attention KV pool, both read/written by the GDN
|
||||
prefill and decode kernels.
|
||||
|
||||
Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
|
||||
|
||||
Usage:
|
||||
python3 -m unittest test_page_major_qwen_hybrid
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
|
||||
from sglang.test.test_utils import DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
register_cuda_ci(est_time=300, stage="extra-a", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
class TestPageMajorQwenHybrid(DefaultServerBase):
|
||||
"""Page-major KV layout on Qwen3.5-4B (GDN-hybrid), Triton backends."""
|
||||
|
||||
model = DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
|
||||
|
||||
# Measured in this harness: baseline (no page-major) and page-major both
|
||||
# ~0.86; the 0.80 threshold leaves margin for run-to-run noise while still
|
||||
# catching the prefill-state corruption that page-major hit before the
|
||||
# gather/scatter fix in gdn_backend.forward_extend (which dropped it to ~0.61).
|
||||
gsm8k_threshold = 0.80
|
||||
num_gsm8k_questions = 200
|
||||
num_shots = 5
|
||||
parallel = 32
|
||||
|
||||
other_args = [
|
||||
"--trust-remote-code",
|
||||
"--mem-fraction-static",
|
||||
"0.85",
|
||||
"--enable-page-major-kv-layout",
|
||||
# Only the Triton attention / linear-attn / Mamba kernels read the
|
||||
# strided envelope K/V and conv/SSM state (enforced by the validator).
|
||||
"--attention-backend",
|
||||
"triton",
|
||||
"--linear-attn-backend",
|
||||
"triton",
|
||||
"--mamba-backend",
|
||||
"triton",
|
||||
]
|
||||
|
||||
def test_gsm8k(self):
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_few_shot_gsm8k
|
||||
|
||||
url = urlparse(self.base_url)
|
||||
args = SimpleNamespace(
|
||||
num_shots=self.num_shots,
|
||||
data_path=None,
|
||||
num_questions=self.num_gsm8k_questions,
|
||||
max_new_tokens=512,
|
||||
parallel=self.parallel,
|
||||
host=f"http://{url.hostname}",
|
||||
port=int(url.port),
|
||||
)
|
||||
metrics = run_few_shot_gsm8k(args)
|
||||
print(
|
||||
f"[{self.__class__.__name__}] GSM8K accuracy: {metrics['accuracy']:.3f} "
|
||||
f"(threshold: {self.gsm8k_threshold})"
|
||||
)
|
||||
self.assertGreaterEqual(metrics["accuracy"], self.gsm8k_threshold)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -60,7 +60,6 @@ class TestMamba(unittest.TestCase):
|
||||
head_num=head_num,
|
||||
head_dim=head_dim,
|
||||
full_attention_layer_ids=full_attention_layer_ids,
|
||||
enable_kvcache_transpose=False,
|
||||
device=device,
|
||||
enable_memory_saver=False,
|
||||
mamba_pool=None,
|
||||
@@ -475,7 +474,6 @@ class TestMamba(unittest.TestCase):
|
||||
head_num=head_num,
|
||||
head_dim=head_dim,
|
||||
full_attention_layer_ids=full_attention_layer_ids,
|
||||
enable_kvcache_transpose=False,
|
||||
device=device,
|
||||
enable_memory_saver=False,
|
||||
mamba_pool=req_to_token_pool.mamba_pool,
|
||||
|
||||
@@ -0,0 +1,163 @@
|
||||
"""CPU correctness tests for the page-major layer-major envelope layout.
|
||||
|
||||
Covers the standalone view builders (no allocator / shared pool):
|
||||
|
||||
- ``build_page_major_mha_views``: 4-D K/V views with correct addressing at
|
||||
page_size 1 (token-granularity envelope) and > 1 (layer-major within a page),
|
||||
and no aliasing across layers / slots.
|
||||
- ``build_page_major_mamba_views``: conv / temporal state views.
|
||||
- ``move_kv_cache_native`` 4-D branch: relocating token rows preserves data.
|
||||
|
||||
Runs on CPU — pure-torch advanced indexing, no Triton.
|
||||
|
||||
python -m pytest test/registered/unit/mem_cache/test_page_major_layout.py -v
|
||||
"""
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=6, suite="base-a-test-cpu")
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.mem_cache.layout.page_major import (
|
||||
build_page_major_mamba_views,
|
||||
build_page_major_mha_views,
|
||||
mamba_entry_bytes,
|
||||
mha_entry_bytes,
|
||||
)
|
||||
from sglang.srt.mem_cache.memory_pool import move_kv_cache_native
|
||||
|
||||
_DEV = "cpu"
|
||||
_DT = torch.float32
|
||||
|
||||
|
||||
def _make_mha_views(layer_num, head_num, head_dim, v_head_dim, page_size, num_pages):
|
||||
entry = mha_entry_bytes(
|
||||
layer_num=layer_num,
|
||||
head_num=head_num,
|
||||
head_dim=head_dim,
|
||||
v_head_dim=v_head_dim,
|
||||
itemsize=_DT.itemsize,
|
||||
)
|
||||
raw = torch.zeros(num_pages * page_size * entry, dtype=torch.uint8, device=_DEV)
|
||||
k, v = build_page_major_mha_views(
|
||||
raw,
|
||||
layer_num=layer_num,
|
||||
head_num=head_num,
|
||||
head_dim=head_dim,
|
||||
v_head_dim=v_head_dim,
|
||||
store_dtype=_DT,
|
||||
page_size=page_size,
|
||||
num_pages=num_pages,
|
||||
)
|
||||
return raw, k, v
|
||||
|
||||
|
||||
class TestPageMajorMHAViews(unittest.TestCase):
|
||||
def test_view_shapes(self):
|
||||
_, k, v = _make_mha_views(3, 2, 4, 4, page_size=2, num_pages=4)
|
||||
self.assertEqual(len(k), 3)
|
||||
for t in k:
|
||||
self.assertEqual(tuple(t.shape), (4, 2, 2, 4))
|
||||
for t in v:
|
||||
self.assertEqual(tuple(t.shape), (4, 2, 2, 4))
|
||||
|
||||
def test_no_aliasing_ps1(self):
|
||||
# Every (layer, slot) cell must be independently addressable.
|
||||
layer_num, slots = 3, 5
|
||||
_, k, v = _make_mha_views(layer_num, 2, 4, 4, page_size=1, num_pages=slots)
|
||||
for L in range(layer_num):
|
||||
for s in range(slots):
|
||||
k[L][s, 0] = float(100 + L * 10 + s)
|
||||
v[L][s, 0] = float(200 + L * 10 + s)
|
||||
for L in range(layer_num):
|
||||
for s in range(slots):
|
||||
self.assertTrue(torch.all(k[L][s, 0] == float(100 + L * 10 + s)))
|
||||
self.assertTrue(torch.all(v[L][s, 0] == float(200 + L * 10 + s)))
|
||||
|
||||
def test_page_slot_addressing_ps_gt1(self):
|
||||
# token id t -> page t // ps, slot t % ps; no aliasing across tokens.
|
||||
ps, pages = 2, 4
|
||||
total = ps * pages
|
||||
_, k, _ = _make_mha_views(2, 1, 2, 2, page_size=ps, num_pages=pages)
|
||||
for L in range(2):
|
||||
for t in range(total):
|
||||
k[L][t // ps, t % ps, 0] = float(1000 + L * 100 + t)
|
||||
for L in range(2):
|
||||
for t in range(total):
|
||||
self.assertEqual(
|
||||
float(k[L][t // ps, t % ps, 0, 0].item()), 1000 + L * 100 + t
|
||||
)
|
||||
|
||||
def test_asymmetric_v_head_dim(self):
|
||||
_, k, v = _make_mha_views(2, 2, 6, 4, page_size=1, num_pages=3)
|
||||
self.assertEqual(tuple(k[0].shape), (3, 1, 2, 6))
|
||||
self.assertEqual(tuple(v[0].shape), (3, 1, 2, 4))
|
||||
|
||||
|
||||
class TestPageMajorMove(unittest.TestCase):
|
||||
def test_move_ps1(self):
|
||||
slots = 6
|
||||
_, k, v = _make_mha_views(2, 1, 4, 4, page_size=1, num_pages=slots)
|
||||
for L in range(2):
|
||||
for s in range(slots):
|
||||
k[L][s, 0] = float(s + 1)
|
||||
v[L][s, 0] = float(-(s + 1))
|
||||
tgt = torch.tensor([0, 1], dtype=torch.int64)
|
||||
src = torch.tensor([4, 5], dtype=torch.int64)
|
||||
move_kv_cache_native(k, v, tgt, src, page_size=1)
|
||||
for L in range(2):
|
||||
self.assertTrue(torch.all(k[L][0, 0] == 5.0))
|
||||
self.assertTrue(torch.all(k[L][1, 0] == 6.0))
|
||||
self.assertTrue(torch.all(v[L][0, 0] == -5.0))
|
||||
|
||||
def test_move_ps_gt1(self):
|
||||
ps, pages = 2, 4
|
||||
total = ps * pages
|
||||
_, k, v = _make_mha_views(1, 1, 2, 2, page_size=ps, num_pages=pages)
|
||||
for t in range(total):
|
||||
k[0][t // ps, t % ps, 0] = float(t + 1)
|
||||
tgt = torch.tensor([0, 3], dtype=torch.int64) # page0 slot0, page1 slot1
|
||||
src = torch.tensor([6, 7], dtype=torch.int64) # page3 slot0, page3 slot1
|
||||
move_kv_cache_native(k, v, tgt, src, page_size=ps)
|
||||
self.assertEqual(float(k[0][0, 0, 0, 0].item()), 7.0)
|
||||
self.assertEqual(float(k[0][1, 1, 0, 0].item()), 8.0)
|
||||
|
||||
|
||||
class TestMambaEnvelopeViews(unittest.TestCase):
|
||||
def test_conv_temporal_shapes_no_alias(self):
|
||||
layers, slots = 2, 4
|
||||
conv_shapes = [(2, 3)]
|
||||
temp_shape = (2, 2)
|
||||
conv_dt, temp_dt = torch.bfloat16, torch.float32
|
||||
entry = mamba_entry_bytes(
|
||||
layer_num=layers,
|
||||
conv_state_shapes=conv_shapes,
|
||||
conv_dtype=conv_dt,
|
||||
temporal_state_shape=temp_shape,
|
||||
temporal_dtype=temp_dt,
|
||||
)
|
||||
raw = torch.zeros(slots * entry, dtype=torch.uint8, device=_DEV)
|
||||
conv_views, temporal = build_page_major_mamba_views(
|
||||
raw,
|
||||
layer_num=layers,
|
||||
conv_state_shapes=conv_shapes,
|
||||
conv_dtype=conv_dt,
|
||||
temporal_state_shape=temp_shape,
|
||||
temporal_dtype=temp_dt,
|
||||
max_slots=slots,
|
||||
)
|
||||
self.assertEqual(tuple(conv_views[0].shape), (layers, slots, 2, 3))
|
||||
self.assertEqual(tuple(temporal.shape), (layers, slots, 2, 2))
|
||||
for L in range(layers):
|
||||
for s in range(slots):
|
||||
temporal[L, s] = float(s + L * 10 + 1)
|
||||
for L in range(layers):
|
||||
for s in range(slots):
|
||||
self.assertTrue(torch.all(temporal[L, s] == float(s + L * 10 + 1)))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,529 @@
|
||||
"""Parity tests for the `store_cache_4d` Triton kernel.
|
||||
|
||||
The kernel writes K/V into the 4-D page-major envelope view. These tests prove
|
||||
it produces byte-identical output to the legacy advanced-indexing path on
|
||||
representative fixtures:
|
||||
|
||||
- ``page_size = 1`` (envelope-degenerate, the critical compatibility case)
|
||||
- ``page_size > 1`` (layer-major within page)
|
||||
- both int32 and int64 ``loc`` dtypes
|
||||
- bf16 and fp8_e5m2 view dtypes
|
||||
- asymmetric ``head_dim != v_head_dim``
|
||||
- empty ``loc`` (no-op)
|
||||
|
||||
Skipped on CPU — Triton requires a GPU.
|
||||
|
||||
python -m pytest test/registered/unit/mem_cache/test_store_cache_4d.py -v
|
||||
"""
|
||||
|
||||
import importlib.util
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
_HAS_CUDA = torch.cuda.is_available()
|
||||
# The set_kv_buffer integration test needs SharedMHATokenToKVPool, which only
|
||||
# exists once the shared-memory-pool feature lands; skip it where absent.
|
||||
_HAS_SHARED_POOL = (
|
||||
importlib.util.find_spec("sglang.srt.mem_cache.shared_memory_pool") is not None
|
||||
)
|
||||
|
||||
register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
def _legacy_advanced_indexing_write(
|
||||
k_view: torch.Tensor,
|
||||
v_view: torch.Tensor,
|
||||
cache_k: torch.Tensor,
|
||||
cache_v: torch.Tensor,
|
||||
loc: torch.Tensor,
|
||||
page_size: int,
|
||||
) -> None:
|
||||
"""Reference implementation: the legacy bypass-super() advanced-indexing
|
||||
path that the Triton kernel replaces. Used as the byte-identity oracle
|
||||
for the parity tests below.
|
||||
"""
|
||||
if page_size == 1:
|
||||
k_view[loc, 0] = cache_k
|
||||
v_view[loc, 0] = cache_v
|
||||
else:
|
||||
page_id = loc // page_size
|
||||
tok_in_p = loc % page_size
|
||||
k_view[page_id, tok_in_p] = cache_k
|
||||
v_view[page_id, tok_in_p] = cache_v
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_CUDA, "Triton kernels require CUDA")
|
||||
class TestStoreCache4D(unittest.TestCase):
|
||||
"""Byte-identity parity vs the legacy advanced-indexing write path."""
|
||||
|
||||
def _make_view_and_cache(
|
||||
self,
|
||||
num_pages: int,
|
||||
page_size: int,
|
||||
head_num: int,
|
||||
head_dim: int,
|
||||
v_head_dim: int,
|
||||
N: int,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
loc_dtype: torch.dtype = torch.int64,
|
||||
seed: int = 0xC0FFEE,
|
||||
):
|
||||
torch.manual_seed(seed)
|
||||
# The shared pool's views are 4-D `(num_pages, page_size, head_num,
|
||||
# head_dim)` with the trailing two dims contiguous. We allocate two
|
||||
# independent contiguous buffers (one for the kernel-under-test,
|
||||
# one as the legacy-path target) so we can compare them.
|
||||
k_view = torch.zeros(
|
||||
(num_pages, page_size, head_num, head_dim),
|
||||
dtype=dtype,
|
||||
device="cuda",
|
||||
)
|
||||
v_view = torch.zeros(
|
||||
(num_pages, page_size, head_num, v_head_dim),
|
||||
dtype=dtype,
|
||||
device="cuda",
|
||||
)
|
||||
cache_k = torch.randn(
|
||||
(N, head_num, head_dim), dtype=torch.float32, device="cuda"
|
||||
).to(dtype)
|
||||
cache_v = torch.randn(
|
||||
(N, head_num, v_head_dim), dtype=torch.float32, device="cuda"
|
||||
).to(dtype)
|
||||
# Valid loc values in [0, num_pages * page_size); generate without
|
||||
# duplicates so the comparison is unambiguous (advanced-indexing
|
||||
# with duplicates is order-undefined for both paths).
|
||||
total_slots = num_pages * page_size
|
||||
assert N <= total_slots
|
||||
loc = torch.randperm(total_slots, device="cuda")[:N].to(loc_dtype)
|
||||
return k_view, v_view, cache_k, cache_v, loc
|
||||
|
||||
def _check_parity(
|
||||
self,
|
||||
num_pages: int,
|
||||
page_size: int,
|
||||
head_num: int,
|
||||
head_dim: int,
|
||||
v_head_dim: int,
|
||||
N: int,
|
||||
dtype: torch.dtype = torch.bfloat16,
|
||||
loc_dtype: torch.dtype = torch.int64,
|
||||
):
|
||||
from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d
|
||||
|
||||
# Two independent target buffers — one for the kernel, one for the
|
||||
# legacy reference path.
|
||||
k_kernel, v_kernel, cache_k, cache_v, loc = self._make_view_and_cache(
|
||||
num_pages,
|
||||
page_size,
|
||||
head_num,
|
||||
head_dim,
|
||||
v_head_dim,
|
||||
N,
|
||||
dtype=dtype,
|
||||
loc_dtype=loc_dtype,
|
||||
)
|
||||
k_legacy = k_kernel.clone()
|
||||
v_legacy = v_kernel.clone()
|
||||
|
||||
# Kernel-under-test
|
||||
store_cache_4d(k_kernel, v_kernel, cache_k, cache_v, loc, page_size)
|
||||
# Legacy reference
|
||||
_legacy_advanced_indexing_write(
|
||||
k_legacy, v_legacy, cache_k, cache_v, loc, page_size
|
||||
)
|
||||
|
||||
# Byte-identical comparison — the kernel must reproduce the
|
||||
# advanced-indexing path bit-for-bit, NOT just numerically close.
|
||||
# For fp8 dtypes, torch.equal works on the integer bit pattern.
|
||||
self.assertTrue(
|
||||
torch.equal(k_kernel, k_legacy),
|
||||
f"K view mismatch: ps={page_size}, dtype={dtype}, "
|
||||
f"loc_dtype={loc_dtype}, N={N}",
|
||||
)
|
||||
self.assertTrue(
|
||||
torch.equal(v_kernel, v_legacy),
|
||||
f"V view mismatch: ps={page_size}, dtype={dtype}, "
|
||||
f"loc_dtype={loc_dtype}, N={N}",
|
||||
)
|
||||
|
||||
# ---- Test 1: ps=1 envelope-degenerate (the critical compat case) ----
|
||||
|
||||
def test_store_cache_4d_ps1_byte_identical(self):
|
||||
"""At page_size=1 the kernel constexpr-folds to the slot-major
|
||||
envelope view. Output must be byte-identical to advanced indexing.
|
||||
This protects the Stage 1/2/3 green eval matrix from regression."""
|
||||
self._check_parity(
|
||||
num_pages=64,
|
||||
page_size=1,
|
||||
head_num=4,
|
||||
head_dim=128,
|
||||
v_head_dim=128,
|
||||
N=16,
|
||||
)
|
||||
|
||||
# ---- Test 2: ps>1 layer-major within page ----
|
||||
|
||||
def test_store_cache_4d_ps_gt1_byte_identical(self):
|
||||
"""At page_size > 1 the kernel splits loc into (page_id, tok_in_p)
|
||||
and writes via the 4-D stride. Output must match the equivalent
|
||||
advanced-indexing write."""
|
||||
self._check_parity(
|
||||
num_pages=8,
|
||||
page_size=64,
|
||||
head_num=4,
|
||||
head_dim=128,
|
||||
v_head_dim=128,
|
||||
N=128,
|
||||
)
|
||||
|
||||
# ---- Test 3: int32 loc dtype ----
|
||||
|
||||
def test_store_cache_4d_int32_loc(self):
|
||||
"""The SWA-side path passes int32 loc (matches the SWA Triton
|
||||
kernel contract). PyTorch advanced indexing tolerates either
|
||||
int32 or int64; the kernel must too."""
|
||||
self._check_parity(
|
||||
num_pages=32,
|
||||
page_size=1,
|
||||
head_num=4,
|
||||
head_dim=64,
|
||||
v_head_dim=64,
|
||||
N=10,
|
||||
loc_dtype=torch.int32,
|
||||
)
|
||||
|
||||
# ---- Test 4: int64 loc dtype (already exercised, explicit) ----
|
||||
|
||||
def test_store_cache_4d_int64_loc(self):
|
||||
"""The full-side path passes int64 loc (matches the v2p table
|
||||
dtype)."""
|
||||
self._check_parity(
|
||||
num_pages=32,
|
||||
page_size=1,
|
||||
head_num=4,
|
||||
head_dim=64,
|
||||
v_head_dim=64,
|
||||
N=10,
|
||||
loc_dtype=torch.int64,
|
||||
)
|
||||
|
||||
# ---- Test 5: bf16 dtype (the production case) ----
|
||||
|
||||
def test_store_cache_4d_dtype_bf16(self):
|
||||
"""bf16 is the production K/V dtype for gpt-oss-20b, Falcon-H1."""
|
||||
self._check_parity(
|
||||
num_pages=16,
|
||||
page_size=64,
|
||||
head_num=4,
|
||||
head_dim=128,
|
||||
v_head_dim=128,
|
||||
N=64,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
# ---- Test 6: fp8_e5m2 dtype ----
|
||||
|
||||
def test_store_cache_4d_dtype_fp8_e5m2(self):
|
||||
"""fp8_e5m2 is used for KV-cache quantization. Caller is responsible
|
||||
for the cast (Phase 1); the kernel sees same-dtype source and
|
||||
destination."""
|
||||
self._check_parity(
|
||||
num_pages=16,
|
||||
page_size=64,
|
||||
head_num=4,
|
||||
head_dim=128,
|
||||
v_head_dim=128,
|
||||
N=64,
|
||||
dtype=torch.float8_e5m2,
|
||||
)
|
||||
|
||||
# ---- Test 7: empty loc (no-op) ----
|
||||
|
||||
def test_store_cache_4d_empty_loc(self):
|
||||
"""N=0 must be a no-op: no kernel launch, no exception, no buffer
|
||||
mutation."""
|
||||
from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d
|
||||
|
||||
k_view = torch.zeros((8, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
v_view = torch.zeros((8, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
k_before = k_view.clone()
|
||||
v_before = v_view.clone()
|
||||
cache_k = torch.empty((0, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
cache_v = torch.empty((0, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
loc = torch.empty((0,), dtype=torch.int64, device="cuda")
|
||||
|
||||
store_cache_4d(k_view, v_view, cache_k, cache_v, loc, page_size=4)
|
||||
|
||||
# Buffers must be unchanged.
|
||||
self.assertTrue(torch.equal(k_view, k_before))
|
||||
self.assertTrue(torch.equal(v_view, v_before))
|
||||
|
||||
# ---- Test 8: head_dim != v_head_dim (asymmetric, e.g. MLA-style) ----
|
||||
|
||||
def test_store_cache_4d_v_head_dim_differs(self):
|
||||
"""When v_head_dim != head_dim, the kernel's K and V branches use
|
||||
different per-token strides. Exercises the stride_k_tok ≠
|
||||
stride_v_tok branch."""
|
||||
self._check_parity(
|
||||
num_pages=8,
|
||||
page_size=16,
|
||||
head_num=2,
|
||||
head_dim=128,
|
||||
v_head_dim=64,
|
||||
N=16,
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_CUDA, "Triton kernels require CUDA")
|
||||
class TestStoreCache4DAssertions(unittest.TestCase):
|
||||
"""The wrapper's contract assertions must fire on bad inputs."""
|
||||
|
||||
def test_rejects_non_contiguous_view_trailing_dim(self):
|
||||
"""Wrapper requires `stride[-1] == 1` and `stride[-2] == head_dim`
|
||||
(the trailing two dims must be contiguous). A permutation that
|
||||
breaks this should trigger AssertionError."""
|
||||
from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d
|
||||
|
||||
# Build a 4-D view, then permute the last two dims → trailing
|
||||
# contiguity violated.
|
||||
k_view = torch.zeros(
|
||||
(4, 4, 4, 64), dtype=torch.bfloat16, device="cuda"
|
||||
).permute(
|
||||
0, 1, 3, 2
|
||||
) # now shape (4, 4, 64, 4); strides broken
|
||||
v_view = torch.zeros((4, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
cache_k = torch.zeros((2, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
cache_v = torch.zeros((2, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
loc = torch.arange(2, dtype=torch.int64, device="cuda")
|
||||
with self.assertRaises(AssertionError):
|
||||
store_cache_4d(k_view, v_view, cache_k, cache_v, loc, page_size=4)
|
||||
|
||||
def test_rejects_dtype_mismatch(self):
|
||||
"""All four tensors must share a dtype; the caller is responsible
|
||||
for any cast before the call."""
|
||||
from sglang.srt.mem_cache.triton_ops.cache_move import store_cache_4d
|
||||
|
||||
k_view = torch.zeros((4, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
v_view = torch.zeros((4, 4, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
cache_k = torch.zeros((2, 4, 64), dtype=torch.float16, device="cuda")
|
||||
cache_v = torch.zeros((2, 4, 64), dtype=torch.bfloat16, device="cuda")
|
||||
loc = torch.arange(2, dtype=torch.int64, device="cuda")
|
||||
with self.assertRaises(AssertionError):
|
||||
store_cache_4d(k_view, v_view, cache_k, cache_v, loc, page_size=4)
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
_HAS_CUDA and _HAS_SHARED_POOL,
|
||||
"Triton kernels require CUDA; SharedMHATokenToKVPool required",
|
||||
)
|
||||
class TestStoreCache4DThroughSetKVBuffer(unittest.TestCase):
|
||||
"""Integration parity test — exercises the kernel through the FULL
|
||||
``SharedMHATokenToKVPool.set_kv_buffer`` path, including the
|
||||
``_external_allocator`` v2p translation and the dtype cast. Confirms the
|
||||
production code path produces bit-identical output to a PyTorch
|
||||
advanced-indexing reference write.
|
||||
"""
|
||||
|
||||
def _build_pool_and_stub_alloc(self, page_size: int, v2p=None):
|
||||
"""Build a small SharedMHATokenToKVPool wired to a stub allocator.
|
||||
|
||||
By default `virtual_to_physical` is identity (the kernel-vs-legacy
|
||||
parity tests don't exercise virtual-id semantics). Pass an explicit
|
||||
`v2p` tensor (sized `max_slots + 1`) to exercise a NON-identity
|
||||
translation — used by the `set_full_loc` fast-path parity test, which
|
||||
needs virtual != physical so the precomputed-physical fast path is
|
||||
meaningfully different from the per-call gather."""
|
||||
import torch as _t
|
||||
|
||||
from sglang.srt.mem_cache.shared_memory_pool import (
|
||||
MHASubPoolSpec,
|
||||
SharedMemoryPool,
|
||||
SharedMHATokenToKVPool,
|
||||
)
|
||||
|
||||
spec = MHASubPoolSpec(
|
||||
name="full",
|
||||
layer_num=2,
|
||||
head_num=4,
|
||||
head_dim=64,
|
||||
store_dtype=_t.bfloat16,
|
||||
grow_direction="up",
|
||||
)
|
||||
total = spec.entry_bytes() * 64
|
||||
# Use a peer to satisfy the two-sub-pool contract.
|
||||
peer = MHASubPoolSpec(
|
||||
name="swa",
|
||||
layer_num=1,
|
||||
head_num=4,
|
||||
head_dim=64,
|
||||
store_dtype=_t.bfloat16,
|
||||
grow_direction="down",
|
||||
)
|
||||
pool = SharedMemoryPool(
|
||||
total_bytes=total + peer.entry_bytes() * 16,
|
||||
sub_pool_specs=[spec, peer],
|
||||
device="cuda",
|
||||
enable_memory_saver=False,
|
||||
page_size=page_size,
|
||||
)
|
||||
kv_pool = SharedMHATokenToKVPool(
|
||||
shared_buffer=pool,
|
||||
sub_pool_name="full",
|
||||
page_size=page_size,
|
||||
start_layer=0,
|
||||
end_layer=2,
|
||||
enable_alt_stream=False,
|
||||
)
|
||||
|
||||
# Stub allocator with an identity (default) or caller-supplied v2p.
|
||||
max_slots = pool.max_slots("full")
|
||||
if v2p is None:
|
||||
v2p = _t.arange(max_slots + 1, dtype=_t.int64, device="cuda")
|
||||
|
||||
class _StubAllocator:
|
||||
virtual_to_physical = v2p
|
||||
|
||||
kv_pool.attach_allocator(_StubAllocator())
|
||||
return kv_pool
|
||||
|
||||
def _run_set_kv_buffer_and_compare(self, page_size: int):
|
||||
import torch as _t
|
||||
|
||||
kv_pool = self._build_pool_and_stub_alloc(page_size)
|
||||
|
||||
# A fake `layer` object with the minimum interface
|
||||
# `set_kv_buffer` reads: `.layer_id`.
|
||||
class _FakeLayer:
|
||||
layer_id = 0
|
||||
|
||||
layer = _FakeLayer()
|
||||
head_num, head_dim = 4, 64
|
||||
N = 16
|
||||
# Generate valid loc in range [0, num_pages * page_size).
|
||||
num_pages = kv_pool.k_buffer[0].shape[0]
|
||||
total = num_pages * page_size
|
||||
assert N <= total
|
||||
loc = _t.randperm(total, device="cuda")[:N].to(_t.int64)
|
||||
cache_k = _t.randn((N, head_num, head_dim), dtype=_t.bfloat16, device="cuda")
|
||||
cache_v = _t.randn((N, head_num, head_dim), dtype=_t.bfloat16, device="cuda")
|
||||
|
||||
# Production path: the Triton `store_cache_4d` kernel via set_kv_buffer.
|
||||
kv_pool.set_kv_buffer(layer, loc, cache_k.clone(), cache_v.clone())
|
||||
k_kernel = kv_pool.k_buffer[0].clone()
|
||||
v_kernel = kv_pool.v_buffer[0].clone()
|
||||
|
||||
# Reference: PyTorch advanced-indexing into a fresh view. The stub
|
||||
# allocator's v2p is identity, so physical loc == virtual loc and no
|
||||
# dtype cast happens (store_dtype == dtype), making this the exact
|
||||
# write the kernel performs.
|
||||
kv_pool.k_buffer[0].zero_()
|
||||
kv_pool.v_buffer[0].zero_()
|
||||
k_view = kv_pool.k_buffer[0]
|
||||
v_view = kv_pool.v_buffer[0]
|
||||
if page_size == 1:
|
||||
k_view[loc, 0] = cache_k
|
||||
v_view[loc, 0] = cache_v
|
||||
else:
|
||||
page_id = loc // page_size
|
||||
tok_in_p = loc % page_size
|
||||
k_view[page_id, tok_in_p] = cache_k
|
||||
v_view[page_id, tok_in_p] = cache_v
|
||||
k_ref = kv_pool.k_buffer[0].clone()
|
||||
v_ref = kv_pool.v_buffer[0].clone()
|
||||
|
||||
self.assertTrue(
|
||||
_t.equal(k_kernel, k_ref),
|
||||
f"K view mismatch through set_kv_buffer at ps={page_size}",
|
||||
)
|
||||
self.assertTrue(
|
||||
_t.equal(v_kernel, v_ref),
|
||||
f"V view mismatch through set_kv_buffer at ps={page_size}",
|
||||
)
|
||||
|
||||
def test_integration_ps1(self):
|
||||
self._run_set_kv_buffer_and_compare(page_size=1)
|
||||
|
||||
def test_integration_ps64(self):
|
||||
self._run_set_kv_buffer_and_compare(page_size=64)
|
||||
|
||||
def _run_full_loc_fast_path_parity(self, page_size: int):
|
||||
"""Stage 3.5 fast-path byte-identity: writing through the precomputed
|
||||
full-physical loc (`set_loc` fast path) must produce a byte-identical
|
||||
KV buffer to writing the virtual loc and letting `set_kv_buffer`
|
||||
translate per call. Uses a NON-identity v2p so the two paths are
|
||||
genuinely different code (fast path skips the gather)."""
|
||||
import torch as _t
|
||||
|
||||
# Non-identity v2p: reverse-map the physical slot space so virtual i
|
||||
# lands on a different physical slot. Keep slot 0 -> 0 (padding sink).
|
||||
# Build the pool once to learn max_slots, then rebuild with the v2p.
|
||||
probe = self._build_pool_and_stub_alloc(page_size)
|
||||
max_slots = probe.k_buffer[0].shape[0] * page_size
|
||||
v2p = _t.arange(max_slots + 1, dtype=_t.int64, device="cuda")
|
||||
# Shuffle the interior [1, max_slots) so virtual != physical, leave
|
||||
# 0 (sink) and the trailing sentinel (max_slots -> itself) alone.
|
||||
interior = _t.randperm(max_slots - 1, device="cuda") + 1
|
||||
v2p[1:max_slots] = interior
|
||||
|
||||
kv_pool = self._build_pool_and_stub_alloc(page_size, v2p=v2p)
|
||||
|
||||
class _FakeLayer:
|
||||
layer_id = 0
|
||||
|
||||
layer = _FakeLayer()
|
||||
head_num, head_dim = 4, 64
|
||||
N = 16
|
||||
num_pages = kv_pool.k_buffer[0].shape[0]
|
||||
total = num_pages * page_size
|
||||
# Draw virtual ids from [1, total) (avoid the padding sink at 0).
|
||||
loc = (_t.randperm(total - 1, device="cuda")[:N] + 1).to(_t.int64)
|
||||
cache_k = _t.randn((N, head_num, head_dim), dtype=_t.bfloat16, device="cuda")
|
||||
cache_v = _t.randn((N, head_num, head_dim), dtype=_t.bfloat16, device="cuda")
|
||||
|
||||
# SLOW path: no precompute pinned -> per-call v2p gather inside
|
||||
# set_kv_buffer translates virtual -> physical.
|
||||
kv_pool.set_loc(None)
|
||||
kv_pool.set_kv_buffer(layer, loc, cache_k.clone(), cache_v.clone())
|
||||
k_slow = kv_pool.k_buffer[0].clone()
|
||||
v_slow = kv_pool.v_buffer[0].clone()
|
||||
|
||||
# FAST path: precompute the full-physical loc exactly as
|
||||
# `set_kv_buffer`'s page math would, pin it via set_loc, and pass
|
||||
# it as `loc` so the data-ptr fast path fires (no gather).
|
||||
if page_size == 1:
|
||||
phys = _t.clamp_min(v2p[loc], 0)
|
||||
else:
|
||||
virt_pages = loc // page_size
|
||||
offsets = loc % page_size
|
||||
phys = _t.clamp_min(v2p[virt_pages] * page_size + offsets, 0)
|
||||
kv_pool.k_buffer[0].zero_()
|
||||
kv_pool.v_buffer[0].zero_()
|
||||
kv_pool.set_loc(phys)
|
||||
try:
|
||||
kv_pool.set_kv_buffer(layer, phys, cache_k.clone(), cache_v.clone())
|
||||
k_fast = kv_pool.k_buffer[0].clone()
|
||||
v_fast = kv_pool.v_buffer[0].clone()
|
||||
finally:
|
||||
kv_pool.set_loc(None)
|
||||
|
||||
self.assertTrue(
|
||||
_t.equal(k_fast, k_slow),
|
||||
f"K mismatch: full_loc fast path != per-call translate at ps={page_size}",
|
||||
)
|
||||
self.assertTrue(
|
||||
_t.equal(v_fast, v_slow),
|
||||
f"V mismatch: full_loc fast path != per-call translate at ps={page_size}",
|
||||
)
|
||||
|
||||
def test_full_loc_fast_path_parity_ps1(self):
|
||||
self._run_full_loc_fast_path_parity(page_size=1)
|
||||
|
||||
def test_full_loc_fast_path_parity_ps64(self):
|
||||
self._run_full_loc_fast_path_parity(page_size=64)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -75,7 +75,6 @@ def _build_swa_tree(page_size, sliding_window_size, kv_size=1024, kv_size_swa=51
|
||||
head_dim=head_dim,
|
||||
swa_attention_layer_ids=swa_ids,
|
||||
full_attention_layer_ids=full_ids,
|
||||
enable_kvcache_transpose=False,
|
||||
device=device,
|
||||
)
|
||||
allocator = SWATokenToKVPoolAllocator(
|
||||
|
||||
@@ -61,7 +61,6 @@ def _build_tree(
|
||||
head_dim=head_dim,
|
||||
swa_attention_layer_ids=swa_ids,
|
||||
full_attention_layer_ids=full_ids,
|
||||
enable_kvcache_transpose=False,
|
||||
device=device,
|
||||
)
|
||||
allocator = SWATokenToKVPoolAllocator(
|
||||
|
||||
@@ -77,7 +77,6 @@ def _build_swa_tree(
|
||||
head_dim=head_dim,
|
||||
swa_attention_layer_ids=swa_attention_layer_ids,
|
||||
full_attention_layer_ids=full_attention_layer_ids,
|
||||
enable_kvcache_transpose=False,
|
||||
device=device,
|
||||
)
|
||||
allocator = SWATokenToKVPoolAllocator(
|
||||
@@ -226,7 +225,6 @@ class TestSWA(unittest.TestCase):
|
||||
head_dim=head_dim,
|
||||
swa_attention_layer_ids=swa_attention_layer_ids,
|
||||
full_attention_layer_ids=full_attention_layer_ids,
|
||||
enable_kvcache_transpose=False,
|
||||
device=device,
|
||||
)
|
||||
alloc = SWATokenToKVPoolAllocator(
|
||||
@@ -310,7 +308,6 @@ class TestSWA(unittest.TestCase):
|
||||
head_dim=head_dim,
|
||||
swa_attention_layer_ids=swa_attention_layer_ids,
|
||||
full_attention_layer_ids=full_attention_layer_ids,
|
||||
enable_kvcache_transpose=False,
|
||||
device=device,
|
||||
)
|
||||
# setup token to kv pool allocator
|
||||
@@ -468,7 +465,6 @@ class TestSWA(unittest.TestCase):
|
||||
head_dim=head_dim,
|
||||
swa_attention_layer_ids=swa_attention_layer_ids,
|
||||
full_attention_layer_ids=full_attention_layer_ids,
|
||||
enable_kvcache_transpose=False,
|
||||
device=device,
|
||||
)
|
||||
# setup token to kv pool allocator
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
"""Triton-kernel parity test for the page-aware decode / extend kernels.
|
||||
|
||||
Verifies that the modified decode / extend Triton kernels produce
|
||||
bit-identical output when called against:
|
||||
|
||||
(a) the legacy 3-D ``[N, head, dim]`` KV view (PAGE_SIZE=1 default),
|
||||
(b) the new 4-D ``[num_pages, page_size, head, dim]`` view with
|
||||
``page_size=1`` (degenerate envelope — same physical bytes as (a)),
|
||||
(c) the new 4-D view with ``page_size>1`` (layer-major), using the
|
||||
same logical KV data but routed via page-aware address math.
|
||||
|
||||
Output for (a) vs (b) must be bit-identical at PAGE_SIZE=1 (the kernel
|
||||
specializes to the legacy branch). Output for (c) must match a hand-
|
||||
computed reference SDPA result (same logical attention; different byte
|
||||
layout).
|
||||
|
||||
Skipped on CPU — Triton requires a GPU.
|
||||
|
||||
python -m pytest test/registered/unit/mem_cache/test_triton_kernel_layout.py -v
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
_HAS_CUDA = torch.cuda.is_available()
|
||||
|
||||
register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_CUDA, "Triton kernels require CUDA")
|
||||
class TestTritonKernelLayoutParity(unittest.TestCase):
|
||||
"""Decode + extend kernel parity across (3-D, 4-D ps=1, 4-D ps>1)."""
|
||||
|
||||
def _setup_decode_inputs(
|
||||
self, bs=2, head_num=2, head_dim=8, num_slots=64, dtype=torch.float16
|
||||
):
|
||||
torch.manual_seed(0xC0FFEE)
|
||||
# Logical KV: shape [num_slots, head_num, head_dim]
|
||||
logical_kv_k = torch.randn(
|
||||
num_slots, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
logical_kv_v = torch.randn(
|
||||
num_slots, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
q = torch.randn(bs, head_num, head_dim, dtype=dtype, device="cuda")
|
||||
# All requests use the first `seq_len` slots.
|
||||
seq_len = 16
|
||||
kv_indices_per_req = torch.arange(seq_len, dtype=torch.int64, device="cuda")
|
||||
kv_indices = kv_indices_per_req.repeat(bs) # [bs * seq_len]
|
||||
kv_indptr = torch.tensor(
|
||||
[i * seq_len for i in range(bs + 1)], dtype=torch.int32, device="cuda"
|
||||
)
|
||||
return q, logical_kv_k, logical_kv_v, kv_indptr, kv_indices, seq_len
|
||||
|
||||
def _run_decode(self, q, k_buf, v_buf, kv_indptr, kv_indices, page_size):
|
||||
from sglang.srt.layers.attention.triton_ops.decode_attention import (
|
||||
decode_attention_fwd,
|
||||
)
|
||||
|
||||
bs, head_num, head_dim = q.shape
|
||||
max_kv_splits = 4
|
||||
attn_logits = torch.empty(
|
||||
(bs, head_num, max_kv_splits, head_dim),
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
)
|
||||
attn_lse = torch.empty(
|
||||
(bs, head_num, max_kv_splits),
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
)
|
||||
o = torch.empty_like(q)
|
||||
num_kv_splits = torch.full(
|
||||
(bs,), max_kv_splits, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
decode_attention_fwd(
|
||||
q,
|
||||
k_buf,
|
||||
v_buf,
|
||||
o,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
attn_logits,
|
||||
attn_lse,
|
||||
num_kv_splits,
|
||||
max_kv_splits,
|
||||
sm_scale=1.0 / (head_dim**0.5),
|
||||
k_scale=1.0,
|
||||
v_scale=1.0,
|
||||
logit_cap=0.0,
|
||||
page_size=page_size,
|
||||
)
|
||||
return o
|
||||
|
||||
def test_decode_3d_vs_4d_ps1_byte_identical(self):
|
||||
"""(a) vs (b): same physical bytes, different view shape.
|
||||
Triton specializes PAGE_SIZE=1 to the legacy branch; output must
|
||||
be bit-identical (modulo non-deterministic FP add ordering, which
|
||||
we sidestep here since the kernels use deterministic reductions
|
||||
for fixed input + grid)."""
|
||||
q, k, v, kv_indptr, kv_indices, seq_len = self._setup_decode_inputs()
|
||||
# (a) legacy 3-D view
|
||||
o_3d = self._run_decode(q, k, v, kv_indptr, kv_indices, page_size=1)
|
||||
# (b) 4-D view: reshape SAME physical bytes to (num_pages=N, 1, head, dim)
|
||||
num_slots = k.shape[0]
|
||||
k_4d = k.view(num_slots, 1, *k.shape[1:])
|
||||
v_4d = v.view(num_slots, 1, *v.shape[1:])
|
||||
o_4d_ps1 = self._run_decode(q, k_4d, v_4d, kv_indptr, kv_indices, page_size=1)
|
||||
# bit-identical (same byte layout, same PAGE_SIZE specialization)
|
||||
self.assertTrue(torch.equal(o_3d, o_4d_ps1))
|
||||
|
||||
def test_extend_3d_vs_4d_ps1_byte_identical(self):
|
||||
"""Same parity check for extend kernel."""
|
||||
from sglang.srt.layers.attention.triton_ops.extend_attention import (
|
||||
extend_attention_fwd,
|
||||
)
|
||||
|
||||
torch.manual_seed(0xDEADBEEF)
|
||||
# head_dim must be >= 16: the extend kernel's QK^T tl.dot requires the
|
||||
# contraction dim K (= head_dim) >= 16 on modern GPU archs (Hopper+).
|
||||
head_num, head_dim = 2, 32
|
||||
num_slots = 32
|
||||
dtype = torch.float16
|
||||
bs = 2
|
||||
prefix_len = 8
|
||||
extend_len = 4
|
||||
|
||||
k_buffer = torch.randn(
|
||||
num_slots, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
v_buffer = torch.randn(
|
||||
num_slots, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
q_extend = torch.randn(
|
||||
bs * extend_len, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
k_extend = torch.randn(
|
||||
bs * extend_len, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
v_extend = torch.randn(
|
||||
bs * extend_len, head_num, head_dim, dtype=dtype, device="cuda"
|
||||
)
|
||||
o = torch.empty_like(q_extend)
|
||||
|
||||
qo_indptr = torch.tensor(
|
||||
[i * extend_len for i in range(bs + 1)], dtype=torch.int32, device="cuda"
|
||||
)
|
||||
kv_indptr = torch.tensor(
|
||||
[i * prefix_len for i in range(bs + 1)], dtype=torch.int32, device="cuda"
|
||||
)
|
||||
kv_indices = torch.arange(prefix_len, dtype=torch.int64, device="cuda").repeat(
|
||||
bs
|
||||
)
|
||||
|
||||
def run(k_buf, v_buf, page_size):
|
||||
o_out = torch.empty_like(q_extend)
|
||||
extend_attention_fwd(
|
||||
q_extend,
|
||||
k_extend,
|
||||
v_extend,
|
||||
o_out,
|
||||
k_buf,
|
||||
v_buf,
|
||||
qo_indptr,
|
||||
kv_indptr,
|
||||
kv_indices,
|
||||
custom_mask=None,
|
||||
is_causal=True,
|
||||
mask_indptr=None,
|
||||
max_len_extend=extend_len,
|
||||
k_scale=1.0,
|
||||
v_scale=1.0,
|
||||
sm_scale=1.0 / (head_dim**0.5),
|
||||
page_size=page_size,
|
||||
)
|
||||
return o_out
|
||||
|
||||
o_3d = run(k_buffer, v_buffer, page_size=1)
|
||||
k_4d = k_buffer.view(num_slots, 1, *k_buffer.shape[1:])
|
||||
v_4d = v_buffer.view(num_slots, 1, *v_buffer.shape[1:])
|
||||
o_4d_ps1 = run(k_4d, v_4d, page_size=1)
|
||||
self.assertTrue(torch.equal(o_3d, o_4d_ps1))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -191,7 +191,6 @@ def create_bench_cache(
|
||||
head_dim=_HEAD_DIM,
|
||||
swa_attention_layer_ids=_non_full_layer_ids(),
|
||||
full_attention_layer_ids=_full_attention_layer_ids(),
|
||||
enable_kvcache_transpose=False,
|
||||
device=device,
|
||||
)
|
||||
allocator = SWATokenToKVPoolAllocator(
|
||||
@@ -211,7 +210,6 @@ def create_bench_cache(
|
||||
head_num=_HEAD_NUM,
|
||||
head_dim=_HEAD_DIM,
|
||||
full_attention_layer_ids=_full_attention_layer_ids(),
|
||||
enable_kvcache_transpose=False,
|
||||
device=device,
|
||||
enable_memory_saver=False,
|
||||
mamba_pool=req_to_token_pool.mamba_pool if has_mamba else None,
|
||||
|
||||
@@ -268,7 +268,6 @@ def build_fixture(cfg: CacheConfig, *, enable_kv_cache_events: bool = False):
|
||||
head_dim=cfg.head_dim,
|
||||
swa_attention_layer_ids=cfg.non_full_layer_ids,
|
||||
full_attention_layer_ids=cfg.full_attention_layer_ids,
|
||||
enable_kvcache_transpose=False,
|
||||
device=device,
|
||||
)
|
||||
allocator = SWATokenToKVPoolAllocator(
|
||||
@@ -288,7 +287,6 @@ def build_fixture(cfg: CacheConfig, *, enable_kv_cache_events: bool = False):
|
||||
head_num=cfg.head_num,
|
||||
head_dim=cfg.head_dim,
|
||||
full_attention_layer_ids=cfg.full_attention_layer_ids,
|
||||
enable_kvcache_transpose=False,
|
||||
device=device,
|
||||
enable_memory_saver=False,
|
||||
mamba_pool=req_to_token_pool.mamba_pool,
|
||||
|
||||
Reference in New Issue
Block a user