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

[DSA] Drop the redundant 512 from the top-k transform entry-point names (#36831)

Browse Source 
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
DarkSharpness
2026-09-01 22:37:03 +08:00
committed by GitHub
co-authored by Claude Opus 5
parent 00689c0c94
commit b6c06e1efb
7 changed files with 31 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/kernels/benchmark/attention/bench_topk.py
+4 -4
View File
@@ -3,8 +3,8 @@ import torch
from sglang.kernels.jit.benchmark import marker
from sglang.kernels.ops.attention.dsv4.topk import (
plan_topk_v2,
topk_transform_512,
topk_transform_512_v2,
topk_transform_paged,
topk_transform_paged_v2,
topk_transform_ragged_v2,
)
from sglang.test.ci.ci_register import register_cuda_ci
@@ -42,10 +42,10 @@ def _build_paged_fn(
def fn(scores, seq_lens, page_table):
if provider == "jit_v1":
topk_transform_512(scores, seq_lens, page_table, out, N)
topk_transform_paged(scores, seq_lens, page_table, out, N)
return out
elif provider == "jit_v2":
topk_transform_512_v2(scores, seq_lens, page_table, out, N, metadata)
topk_transform_paged_v2(scores, seq_lens, page_table, out, N, metadata)
return out
elif provider == "flashinfer":
from flashinfer import top_k_page_table_transform
test/registered/kernels/ops/attention/test_topk_v2.py
+3 -3
View File
@@ -31,7 +31,7 @@ import torch
from sglang.kernels.ops.attention.dsv4.topk import (
plan_topk_v2,
topk_transform_512_v2,
topk_transform_paged_v2,
topk_transform_ragged_v2,
)
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
@@ -165,7 +165,7 @@ def _run(scores, seq_lens, page_table, inv_cpu, k):
batch = scores.shape[0]
metadata = _plan(seq_lens)
out = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device)
topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata)
topk_transform_paged_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata)
torch.cuda.synchronize()
out_cpu = out.cpu().tolist()
return [_invert(out_cpu[i], inv_cpu[i]) for i in range(batch)]
@@ -177,7 +177,7 @@ def _run_raw(scores, seq_lens, k):
batch = scores.shape[0]
metadata = _plan(seq_lens)
out = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device)
topk_transform_512_v2(scores, seq_lens, None, out, PAGE_SIZE, metadata)
topk_transform_paged_v2(scores, seq_lens, None, out, PAGE_SIZE, metadata)
torch.cuda.synchronize()
out_cpu = out.cpu().tolist()
return [[v for v in out_cpu[i] if v != -1] for i in range(batch)]