minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Revert "[AMD][DSV4] Fix unified-KV pool sizing and SWA ring accounting" (#38163)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-09-05 17:28:46 -07:00
committed by GitHub
parent 09daea94ac
commit f5819b09bf
23 changed files with 159 additions and 791 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/ops/attention/test_c4_v2.py
+1 -90
View File
@@ -7,11 +7,7 @@ import pytest
import torch
import triton
from sglang.kernels.ops.attention.dsv4 import (
CompressorDecodePlan,
CompressorPrefillPlan,
compress_forward,
)
from sglang.kernels.ops.attention.dsv4 import compress_forward
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.kernels.deepseek_v4.common import (
@@ -126,91 +122,6 @@ def _make_inputs(
# -----------------------------------------------------------------------------
@pytest.mark.parametrize("ring_size", [8, 16])
@pytest.mark.parametrize(
("gpu_inputs", "use_cuda_graph"),
[(False, False), (True, False), (True, True)],
)
def test_unified_request_ring_plans_ignore_full_to_state(
ring_size: int, gpu_inputs: bool, use_cuda_graph: bool
) -> None:
"""C4 plans must address state by request slot, not the full-cache map."""
device = torch.device(get_device())
req_pool_indices = torch.tensor([2, 5], dtype=torch.int64, device=device)
req_to_token = torch.zeros((6, 16), dtype=torch.int32, device=device)
full_to_state = torch.zeros(1, dtype=torch.int64, device=device)
seq_lens = torch.tensor([8, 12], dtype=torch.int64)
extend_lens = torch.tensor([4, 4], dtype=torch.int64)
if gpu_inputs:
seq_lens = seq_lens.to(device)
extend_lens = extend_lens.to(device)
prefill = CompressorPrefillPlan.generate(
compress_ratio=RATIO,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
extend_lens=extend_lens,
req_to_token=req_to_token,
full_to_state=full_to_state,
swa_page_size=256,
ring_size=ring_size,
num_q_tokens=8,
use_cuda_graph=use_cuda_graph,
use_req_ring=True,
)
plan_c = prefill.plan_c.view(torch.int32).reshape(-1, 4).cpu()
plan_w = prefill.plan_w.view(torch.int32).reshape(-1, 2).cpu()
valid_c = plan_c[plan_c[:, 2] >= 0]
got_reads = {
int(row[1].item()) & 0xFFFF: (int(row[2].item()), int(row[3].item()))
for row in valid_c
}
expected_reads = {
3: (
(2 * ring_size + 3 % ring_size) // RATIO,
(2 * ring_size + 7 % ring_size) // RATIO,
),
7: (
(5 * ring_size + 7 % ring_size) // RATIO,
(5 * ring_size + 11 % ring_size) // RATIO,
),
}
assert got_reads == expected_reads
valid_w = plan_w[plan_w[:, 1] >= 0]
got_writes = {int(row[0].item()): int(row[1].item()) for row in valid_w}
expected_writes = {
**{j: 2 * ring_size + (4 + j) % ring_size for j in range(4)},
**{4 + j: 5 * ring_size + (8 + j) % ring_size for j in range(4)},
}
assert got_writes == expected_writes
assert {got_writes[j] for j in range(4)}.isdisjoint(
{got_writes[j] for j in range(4, 8)}
)
decode = CompressorDecodePlan.generate(
compress_ratio=RATIO,
req_pool_indices=req_pool_indices,
req_to_token=req_to_token,
full_to_state=full_to_state,
seq_lens=torch.tensor([8, 12], dtype=torch.int64, device=device),
swa_page_size=256,
ring_size=ring_size,
use_req_ring=True,
)
got_decode = decode.plan_d.view(torch.int32).reshape(-1, 4).cpu()
expected_decode = torch.tensor(
[
[8, 2 * ring_size + 7 % ring_size, *expected_reads[3]],
[12, 5 * ring_size + 11 % ring_size, *expected_reads[7]],
],
dtype=torch.int32,
)
assert torch.equal(got_decode, expected_decode)
assert got_decode[0, 1] != got_decode[1, 1]
@pytest.mark.parametrize("mode", ["legacy", "paged"])
@pytest.mark.parametrize("seq_len", [4, 8, 32, 256, 1024])
def test_prefill_no_context(mode: str, seq_len: int) -> None: