137 lines
4.8 KiB
Python
137 lines
4.8 KiB
Python
"""Test for the kpool top-k transform JIT kernel.
|
|
|
|
Ported from the former AOT sgl-kernel test (sgl-kernel/tests/test_topk.py).
|
|
The kernel selects pool groups at pool granularity, expands each selected group
|
|
to ``pool_size`` token indices, and optionally transforms those token indices
|
|
through a page table or a ragged offset.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from typing import Optional
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from sglang.jit_kernel.kpool_topk_transform import fast_kpool_topk_transform_fused
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
|
|
register_cuda_ci(est_time=60, suite="base-b-kernel-unit-1-gpu-large")
|
|
|
|
|
|
def _ref_torch_kpool_transform_impl(
|
|
score: torch.Tensor,
|
|
lengths: torch.Tensor,
|
|
pool_size: int,
|
|
topk: int,
|
|
page_table: Optional[torch.Tensor] = None,
|
|
topk_indices_offset: Optional[torch.Tensor] = None,
|
|
seq_lens: Optional[torch.Tensor] = None,
|
|
) -> torch.Tensor:
|
|
rows = score.shape[0]
|
|
group_topk = topk // pool_size
|
|
offsets = torch.arange(pool_size, dtype=torch.int32, device=score.device)
|
|
out_cols = topk + (pool_size - 1 if seq_lens is not None else 0)
|
|
out = torch.full((rows, out_cols), -1, dtype=torch.int32, device=score.device)
|
|
for i in range(rows):
|
|
length = int(lengths[i].item())
|
|
valid_count = min(length, group_topk)
|
|
write_pos = 0
|
|
if valid_count == 0:
|
|
token_ids = torch.empty((0,), dtype=torch.int32, device=score.device)
|
|
elif length <= group_topk:
|
|
selected = torch.arange(length, dtype=torch.int32, device=score.device)
|
|
token_ids = (selected.unsqueeze(1) * pool_size + offsets).reshape(-1)
|
|
else:
|
|
selected = torch.topk(
|
|
score[i, :length], group_topk, dim=-1, sorted=False
|
|
).indices.to(torch.int32)
|
|
token_ids = (selected.unsqueeze(1) * pool_size + offsets).reshape(-1)
|
|
if token_ids.numel() > 0:
|
|
if page_table is not None:
|
|
token_ids = page_table[i, token_ids.long()].to(torch.int32)
|
|
elif topk_indices_offset is not None:
|
|
token_ids = token_ids + topk_indices_offset[i].to(torch.int32)
|
|
write_pos = valid_count * pool_size
|
|
out[i, :write_pos] = token_ids[:write_pos]
|
|
if seq_lens is not None:
|
|
tail_count = int(seq_lens[i].item()) % pool_size
|
|
if tail_count > 0:
|
|
raw_tail = length * pool_size + torch.arange(
|
|
tail_count, dtype=torch.int32, device=score.device
|
|
)
|
|
if page_table is not None:
|
|
tail = page_table[i, raw_tail.long()].to(torch.int32)
|
|
elif topk_indices_offset is not None:
|
|
tail = raw_tail + topk_indices_offset[i].to(torch.int32)
|
|
else:
|
|
tail = raw_tail
|
|
out[i, write_pos : write_pos + tail_count] = tail
|
|
return out
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"pool_size,group_topk",
|
|
[(16, 128), (16, 160), (16, 192), (16, 224), (8, 256), (4, 512)],
|
|
)
|
|
@pytest.mark.parametrize("mode", ["raw", "paged", "ragged"])
|
|
@pytest.mark.parametrize("append_tail", [False, True])
|
|
@torch.inference_mode()
|
|
def test_kpool_topk_transform_kernel(
|
|
pool_size: int, group_topk: int, mode: str, append_tail: bool
|
|
) -> None:
|
|
torch.manual_seed(42)
|
|
bs = 17
|
|
topk = pool_size * group_topk
|
|
num_groups = 4096
|
|
score = torch.randn(bs, num_groups, dtype=torch.float32, device="cuda")
|
|
lengths = torch.randint(
|
|
group_topk + 1, num_groups + 1, (bs,), dtype=torch.int32, device="cuda"
|
|
)
|
|
|
|
page_table = None
|
|
topk_indices_offset = None
|
|
seq_lens = None
|
|
tail_counts = torch.randint(0, pool_size, (bs,), dtype=torch.int32, device="cuda")
|
|
if append_tail:
|
|
seq_lens = lengths * pool_size + tail_counts
|
|
if mode == "paged":
|
|
page_table = torch.arange(
|
|
bs * (num_groups * pool_size + pool_size),
|
|
dtype=torch.int32,
|
|
device="cuda",
|
|
).view(bs, num_groups * pool_size + pool_size)
|
|
elif mode == "ragged":
|
|
topk_indices_offset = torch.randint(
|
|
0, 2048, (bs,), dtype=torch.int32, device="cuda"
|
|
)
|
|
|
|
out_ref = _ref_torch_kpool_transform_impl(
|
|
score,
|
|
lengths,
|
|
pool_size,
|
|
topk,
|
|
page_table=page_table,
|
|
topk_indices_offset=topk_indices_offset,
|
|
seq_lens=seq_lens,
|
|
)
|
|
out_our = fast_kpool_topk_transform_fused(
|
|
score,
|
|
lengths,
|
|
pool_size,
|
|
topk,
|
|
page_table=page_table,
|
|
topk_indices_offset=topk_indices_offset,
|
|
seq_lens=seq_lens,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
out_ref = torch.sort(out_ref, dim=-1).values
|
|
out_our = torch.sort(out_our, dim=-1).values
|
|
torch.testing.assert_close(out_our, out_ref, atol=0, rtol=0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main([__file__, "-v", "-s"]))
|