Overview
Python wrapper functions for top-k selection operations used in attention mechanisms, specifically optimized for DeepSeek V3.2 style models requiring topk=2048.
Description
This module provides the top-k selection operations API for the SGLang kernel library. It exposes four functions that handle efficient top-k index selection across different KV cache layouts (paged and ragged).
fast_topk is a lightweight dispatcher: for topk=1 it uses torch.max along the specified dimension, and for larger k values it delegates to torch.topk. This function does not require custom CUDA kernels.
The remaining three functions are thin wrappers around custom CUDA kernels registered via torch.ops.sgl_kernel and are currently optimized exclusively for topk=2048 (the DeepSeek V3.2 sparse attention configuration):
- fast_topk_v2 selects the top-k indices from a 2D score tensor, supporting both ragged and paged key layouts via the optional row_starts parameter.
- fast_topk_transform_fused combines top-k selection with page table index transformation (page_size=1), mapping selected indices through a page table for paged KV cache access.
- fast_topk_transform_ragged_fused combines top-k selection with ragged KV index offset transformation, used specifically during extend operations (not draft extend).
All three CUDA-backed functions assert that topk == 2048 and operate on 2D score tensors of shape (B, L).
Usage
Use fast_topk as a general-purpose top-k selection function for any value of k. Use the specialized functions (fast_topk_v2, fast_topk_transform_fused, fast_topk_transform_ragged_fused) when performing sparse attention with topk=2048 in DeepSeek V3.2 style models, where the fused index transformation eliminates extra memory round-trips during decode or extend phases.
Code Reference
Source Location
Signature
def fast_topk(values, topk, dim):
...
def fast_topk_v2(
score: torch.Tensor,
lengths: torch.Tensor,
topk: int,
row_starts: Optional[torch.Tensor] = None,
) -> torch.Tensor:
...
def fast_topk_transform_fused(
score: torch.Tensor,
lengths: torch.Tensor,
page_table_size_1: torch.Tensor,
cu_seqlens_q: torch.Tensor,
topk: int,
row_starts: Optional[torch.Tensor] = None,
) -> torch.Tensor:
...
def fast_topk_transform_ragged_fused(
score: torch.Tensor,
lengths: torch.Tensor,
topk_indices_offset: torch.Tensor,
topk: int,
row_starts: Optional[torch.Tensor] = None,
) -> torch.Tensor:
...
Import
from sgl_kernel import (
fast_topk,
fast_topk_v2,
fast_topk_transform_fused,
fast_topk_transform_ragged_fused,
)
I/O Contract
Inputs
fast_topk
| Name |
Type |
Required |
Description
|
| values |
torch.Tensor |
Yes |
Input tensor from which to select top-k values
|
| topk |
int |
Yes |
Number of top elements to select
|
| dim |
int |
Yes |
Dimension along which to perform the top-k selection
|
fast_topk_v2
| Name |
Type |
Required |
Description
|
| score |
torch.Tensor |
Yes |
Score tensor of shape (B, L) containing logits between query and key
|
| lengths |
torch.Tensor |
Yes |
Lengths tensor of shape (B) indicating valid length per row
|
| topk |
int |
Yes |
Number of top-k indices to select; must be 2048
|
| row_starts |
Optional[torch.Tensor] |
No |
Start index of each row in the score tensor of shape (B); required when key layout is ragged
|
fast_topk_transform_fused
| Name |
Type |
Required |
Description
|
| score |
torch.Tensor |
Yes |
Score tensor of shape (B, L) containing logits between query and key
|
| lengths |
torch.Tensor |
Yes |
Lengths tensor of shape (B) indicating valid length per row
|
| page_table_size_1 |
torch.Tensor |
Yes |
Page table tensor of shape (Batch, topk) with page_size=1
|
| cu_seqlens_q |
torch.Tensor |
Yes |
Cumulative sequence lengths tensor of shape (Batch + 1)
|
| topk |
int |
Yes |
Number of top-k indices to select; must be 2048
|
| row_starts |
Optional[torch.Tensor] |
No |
Start index of each row in the score tensor of shape (B); required when key layout is ragged
|
fast_topk_transform_ragged_fused
| Name |
Type |
Required |
Description
|
| score |
torch.Tensor |
Yes |
Score tensor of shape (B, L) containing logits between query and key
|
| lengths |
torch.Tensor |
Yes |
Lengths tensor of shape (B) indicating valid length per row
|
| topk_indices_offset |
torch.Tensor |
Yes |
Offset of topk indices in ragged KV of shape (B)
|
| topk |
int |
Yes |
Number of top-k indices to select; must be 2048
|
| row_starts |
Optional[torch.Tensor] |
No |
Start index of each row in the score tensor of shape (B); can be None if all lengths <= topk
|
Outputs
fast_topk
| Name |
Type |
Description
|
| (values, indices) |
tuple[torch.Tensor, torch.Tensor] |
When topk=1, returns (max_values, max_indices) with keepdim=True; otherwise returns standard torch.topk output
|
fast_topk_v2
| Name |
Type |
Description
|
| topk_indices |
torch.Tensor |
Top-k indices tensor of shape (B, topk) with dtype int32
|
fast_topk_transform_fused
| Name |
Type |
Description
|
| dst_page_table |
torch.Tensor |
Transformed page table tensor of shape (B, topk) with dtype int32, mapping top-k indices through the source page table
|
fast_topk_transform_ragged_fused
| Name |
Type |
Description
|
| topk_indices_ragged |
torch.Tensor |
Top-k indices tensor of shape (B, topk) with dtype int32, offset-adjusted for ragged KV layout
|
Usage Examples
import torch
from sgl_kernel import fast_topk, fast_topk_v2
# --- fast_topk: General-purpose top-k selection ---
logits = torch.randn(4, 32000, device="cuda")
# Get top-1 (uses torch.max internally)
top1_values, top1_indices = fast_topk(logits, topk=1, dim=-1)
# Get top-10 (uses torch.topk internally)
top10_values, top10_indices = fast_topk(logits, topk=10, dim=-1)
# --- fast_topk_v2: Optimized CUDA kernel for topk=2048 ---
batch_size = 8
max_len = 8192
score = torch.randn(batch_size, max_len, device="cuda", dtype=torch.float16)
lengths = torch.full((batch_size,), max_len, dtype=torch.int32, device="cuda")
topk_indices = fast_topk_v2(score, lengths, topk=2048)
# topk_indices shape: (8, 2048), dtype: int32
Related Pages
