From adc22458bf4e929da2e8fb2c4c27907005ec8743 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Sun, 30 Aug 2026 19:38:17 +0000 Subject: [PATCH 1/4] feat: add indexed attention training kernels --- README.md | 8 +- prime_kernels/__init__.py | 9 +- prime_kernels/indexed_attention/__init__.py | 3 + prime_kernels/indexed_attention/backward.py | 357 ++++++++++++++++++++ prime_kernels/indexed_attention/forward.py | 304 +++++++++++++++++ prime_kernels/kernels.toml | 6 + 6 files changed, 685 insertions(+), 2 deletions(-) create mode 100644 prime_kernels/indexed_attention/__init__.py create mode 100644 prime_kernels/indexed_attention/backward.py create mode 100644 prime_kernels/indexed_attention/forward.py diff --git a/README.md b/README.md index 1420a28..3560bf8 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ CUDA kernels for Prime Intellect training stacks, shipped as one wheel, `prime-k │ ├── __init__.py # Python surface: op wrappers, fake tensors │ ├── mxfp8.py │ └── csrc/ # the C++/CUDA sources compiled into prime_kernels.flash_moe._C + ├── indexed_attention/ # Python-only TileLang indexed GQA forward + backward ├── mxfp8_moe/ # Python-only MXFP8 MoE runtime kernels └── rmsnorm/ ├── __init__.py @@ -55,6 +56,9 @@ neither built nor shipped in the wheel, and the registry does not list it. `flash_moe` is currently dormant in prime-rl. `mxfp8_moe` provides differentiable MXFP8 grouped GEMM and MXFP8 expert-parallel transport. It is registered as Python-only because it orchestrates PyTorch and torchao kernels rather than compiling a `_C` extension here. +`indexed_attention` provides differentiable grouped-query attention over an explicit token +selection for each query. Its TileLang kernels accept different query and KV lengths so the +caller can gather KV for context parallelism without gathering queries. ## Installing @@ -99,7 +103,9 @@ cxx-std = 20 For a Python-only kernel, set `python-only = true`, omit `ops` and `sources`, and expose the differentiable Python surface from `__init__.py`. Optional import requirements belong -in the manifest's `requires` list so `is_available()` fails during setup. +in the manifest's `requires` list so `is_available()` fails during setup. Python-only ops +may use `torch.library.custom_op`; register fake and autograd implementations so they remain +visible to `torch.compile` and training. Whatever the kernel requires of its inputs — block sizes, alignments, layouts — belongs here, not in the caller: `TORCH_CHECK` it in the binding, and export the constants diff --git a/prime_kernels/__init__.py b/prime_kernels/__init__.py index b48d3d4..ca7ba99 100644 --- a/prime_kernels/__init__.py +++ b/prime_kernels/__init__.py @@ -21,6 +21,7 @@ ] _SPECS: dict[str, KernelSpec] = _load_manifest(Path(__file__).parent) +_LOADED_MODULES: dict[str, ModuleType] = {} KERNELS: tuple[str, ...] = tuple(_SPECS) @@ -61,10 +62,16 @@ def is_available(name: str, device: int | None = None) -> bool: def load(name: str, device: int | None = None) -> ModuleType: + try: + return _LOADED_MODULES[name] + except KeyError: + pass reason = unavailable_reason(name, device) if reason is not None: raise RuntimeError(reason) - return importlib.import_module(spec(name).module) + module = importlib.import_module(spec(name).module) + _LOADED_MODULES[name] = module + return module def status(device: int | None = None) -> dict[str, str]: diff --git a/prime_kernels/indexed_attention/__init__.py b/prime_kernels/indexed_attention/__init__.py new file mode 100644 index 0000000..6363734 --- /dev/null +++ b/prime_kernels/indexed_attention/__init__.py @@ -0,0 +1,3 @@ +from prime_kernels.indexed_attention.forward import indexed_attention, unsupported_shape_reason + +__all__ = ["indexed_attention", "unsupported_shape_reason"] diff --git a/prime_kernels/indexed_attention/backward.py b/prime_kernels/indexed_attention/backward.py new file mode 100644 index 0000000..3587b68 --- /dev/null +++ b/prime_kernels/indexed_attention/backward.py @@ -0,0 +1,357 @@ +# Vendored from tile-ai/tilelang (Apache 2.0), modified for indexed GQA. + +import tilelang +import torch +from tilelang import language as T + + +@tilelang.jit(out_idx=[-1]) +def attention_delta( + num_heads: int, + head_dim: int, + block_tokens: int = 32, + num_stages: int = 5, +): + batch_size = T.dynamic("batch_size") + query_length = T.dynamic("query_length") + shape = [batch_size, query_length, num_heads, head_dim] + + @T.prim_func + def kernel( + output: T.Tensor(shape, T.bfloat16), + grad_output: T.Tensor(shape, T.bfloat16), + delta: T.Tensor([batch_size, query_length, num_heads], T.float32), + ): + with T.Kernel(num_heads, T.ceildiv(query_length, block_tokens), batch_size) as (head, token_block, batch): + output_fragment = T.alloc_fragment([block_tokens, block_tokens], T.float32) + grad_output_fragment = T.alloc_fragment([block_tokens, block_tokens], T.float32) + products = T.alloc_fragment([block_tokens, block_tokens], T.float32) + row_sum = T.alloc_fragment([block_tokens], T.float32) + + T.clear(products) + for dim_block in T.Pipelined(T.ceildiv(head_dim, block_tokens), num_stages=num_stages): + T.copy( + output[ + batch, + token_block * block_tokens : (token_block + 1) * block_tokens, + head, + dim_block * block_tokens : (dim_block + 1) * block_tokens, + ], + output_fragment, + ) + T.copy( + grad_output[ + batch, + token_block * block_tokens : (token_block + 1) * block_tokens, + head, + dim_block * block_tokens : (dim_block + 1) * block_tokens, + ], + grad_output_fragment, + ) + for token, dim in T.Parallel(block_tokens, block_tokens): + products[token, dim] += output_fragment[token, dim] * grad_output_fragment[token, dim] + + T.reduce_sum(products, row_sum, dim=1) + T.copy( + row_sum, + delta[ + batch, + token_block * block_tokens : (token_block + 1) * block_tokens, + head, + ], + ) + + return kernel + + +@tilelang.jit( + out_idx=[-3], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + # TileLang can otherwise merge buffers used on opposite sides of an atomic update. + tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: False, + }, +) +def indexed_attention_backward_kernel( + num_query_heads: int, + num_kv_heads: int, + head_dim: int, + selection_width: int, + scale: float, + block_tokens: int = 32, + threads: int = 128, +): + batch_size = T.dynamic("batch_size") + query_length = T.dynamic("query_length") + kv_length = T.dynamic("kv_length") + + query_heads_per_kv_head = num_query_heads // num_kv_heads + head_tile = max(tilelang.math.next_power_of_2(query_heads_per_kv_head), 16) + num_token_blocks = tilelang.cdiv(selection_width, block_tokens) + log2_scale = scale * 1.44269504 + + query_shape = [batch_size, query_length, num_query_heads, head_dim] + kv_shape = [batch_size, kv_length, num_kv_heads, head_dim] + indices_shape = [batch_size, query_length, selection_width] + statistics_shape = [batch_size, query_length, num_query_heads] + + @T.prim_func + def kernel( + query: T.Tensor(query_shape, T.bfloat16), + key: T.Tensor(kv_shape, T.bfloat16), + value: T.Tensor(kv_shape, T.bfloat16), + grad_output: T.Tensor(query_shape, T.bfloat16), + indices: T.Tensor(indices_shape, T.int32), + logsumexp: T.Tensor(statistics_shape, T.float32), + delta: T.Tensor(statistics_shape, T.float32), + grad_query: T.Tensor(query_shape, T.bfloat16), + grad_key: T.Tensor(kv_shape, T.float32), + grad_value: T.Tensor(kv_shape, T.float32), + ): + with T.Kernel(query_length, batch_size, num_kv_heads, threads=threads) as (query_token, batch, kv_head): + query_shared = T.alloc_shared([head_tile, head_dim], T.bfloat16) + key_shared = T.alloc_shared([block_tokens, head_dim], T.bfloat16) + value_shared = T.alloc_shared([block_tokens, head_dim], T.bfloat16) + grad_output_shared = T.alloc_shared([head_tile, head_dim], T.bfloat16) + probability_shared = T.alloc_shared([head_tile, block_tokens], T.bfloat16) + grad_score_shared = T.alloc_shared([head_tile, block_tokens], T.bfloat16) + grad_query_shared = T.alloc_shared([head_tile, head_dim], T.bfloat16) + logsumexp_shared = T.alloc_shared([head_tile], T.float32) + delta_shared = T.alloc_shared([head_tile], T.float32) + atomic_store = T.alloc_shared([block_tokens, head_dim], T.float32) + + valid_index = T.alloc_fragment([block_tokens], "bool") + probability = T.alloc_fragment([head_tile, block_tokens], T.float32) + grad_probability = T.alloc_fragment([head_tile, block_tokens], T.float32) + grad_query_accumulator = T.alloc_fragment([head_tile, head_dim], T.float32) + grad_kv_accumulator = T.alloc_fragment([block_tokens, head_dim], T.float32) + + first_query_head = kv_head * query_heads_per_kv_head + last_valid_kv = kv_length - 2 + + T.copy( + query[ + batch, + query_token, + first_query_head : first_query_head + query_heads_per_kv_head, + :, + ], + query_shared[:query_heads_per_kv_head, :], + ) + T.copy( + grad_output[ + batch, + query_token, + first_query_head : first_query_head + query_heads_per_kv_head, + :, + ], + grad_output_shared[:query_heads_per_kv_head, :], + ) + T.copy( + logsumexp[ + batch, + query_token, + first_query_head : first_query_head + query_heads_per_kv_head, + ], + logsumexp_shared[:query_heads_per_kv_head], + ) + T.copy( + delta[ + batch, + query_token, + first_query_head : first_query_head + query_heads_per_kv_head, + ], + delta_shared[:query_heads_per_kv_head], + ) + for head, dim in T.Parallel(head_tile - query_heads_per_kv_head, head_dim): + query_shared[query_heads_per_kv_head + head, dim] = 0 + grad_output_shared[query_heads_per_kv_head + head, dim] = 0 + for head in T.Parallel(head_tile - query_heads_per_kv_head): + logsumexp_shared[query_heads_per_kv_head + head] = 0 + delta_shared[query_heads_per_kv_head + head] = 0 + + T.clear(grad_query_accumulator) + for selection_block in T.Pipelined(num_token_blocks, num_stages=0): + for selected_token in T.Parallel(block_tokens): + valid_index[selected_token] = ( + indices[batch, query_token, selection_block * block_tokens + selected_token] <= last_valid_kv + ) + + for head, selected_token in T.Parallel(head_tile, block_tokens): + probability[head, selected_token] = T.if_then_else( + valid_index[selected_token], + 0, + -T.infinity(T.float32), + ) + for selected_token, dim in T.Parallel(block_tokens, head_dim): + key_shared[selected_token, dim] = key[ + batch, + indices[batch, query_token, selection_block * block_tokens + selected_token], + kv_head, + dim, + ] + T.gemm( + query_shared, + key_shared, + probability, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + + for head, selected_token in T.Parallel(head_tile, block_tokens): + if head < query_heads_per_kv_head: + probability[head, selected_token] = T.exp2( + probability[head, selected_token] * log2_scale - logsumexp_shared[head] + ) + else: + probability[head, selected_token] = 0 + T.copy(probability, probability_shared) + + for selected_token, dim in T.Parallel(block_tokens, head_dim): + value_shared[selected_token, dim] = value[ + batch, + indices[batch, query_token, selection_block * block_tokens + selected_token], + kv_head, + dim, + ] + T.gemm( + grad_output_shared, + value_shared, + grad_probability, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + clear_accum=True, + ) + + for head, selected_token in T.Parallel(head_tile, block_tokens): + if head < query_heads_per_kv_head: + grad_probability[head, selected_token] = ( + probability[head, selected_token] + * (grad_probability[head, selected_token] - delta_shared[head]) + * scale + ) + else: + grad_probability[head, selected_token] = 0 + T.copy(grad_probability, grad_score_shared) + + T.gemm( + grad_score_shared, + key_shared, + grad_query_accumulator, + policy=T.GemmWarpPolicy.FullRow, + ) + + T.gemm( + grad_score_shared, + query_shared, + grad_kv_accumulator, + transpose_A=True, + policy=T.GemmWarpPolicy.FullRow, + clear_accum=True, + ) + T.copy(grad_kv_accumulator, atomic_store) + for selected_token, dim in T.Parallel(block_tokens, head_dim // 4): + T.atomic_addx4( + grad_key[ + batch, + indices[ + batch, + query_token, + selection_block * block_tokens + selected_token, + ], + kv_head, + dim * 4, + ], + atomic_store[selected_token, dim * 4], + ) + + T.gemm( + probability_shared, + grad_output_shared, + grad_kv_accumulator, + transpose_A=True, + policy=T.GemmWarpPolicy.FullRow, + clear_accum=True, + ) + T.copy(grad_kv_accumulator, atomic_store) + for selected_token, dim in T.Parallel(block_tokens, head_dim // 4): + T.atomic_addx4( + grad_value[ + batch, + indices[ + batch, + query_token, + selection_block * block_tokens + selected_token, + ], + kv_head, + dim * 4, + ], + atomic_store[selected_token, dim * 4], + ) + + T.copy(grad_query_accumulator, grad_query_shared) + T.copy( + grad_query_shared[:query_heads_per_kv_head, :], + grad_query[ + batch, + query_token, + first_query_head : first_query_head + query_heads_per_kv_head, + :, + ], + ) + + return kernel + + +@torch.library.custom_op("prime_kernels::indexed_attention_backward", mutates_args=()) +def indexed_attention_backward( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + output: torch.Tensor, + grad_output: torch.Tensor, + indices: torch.Tensor, + logsumexp: torch.Tensor, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + grad_output = grad_output.contiguous() + _, _, num_query_heads, head_dim = query.shape + _, _, num_kv_heads, _ = key.shape + selection_width = indices.shape[-1] + + delta = attention_delta(num_query_heads, head_dim)(output, grad_output) + grad_key = torch.zeros_like(key, dtype=torch.float32) + grad_value = torch.zeros_like(value, dtype=torch.float32) + grad_query = indexed_attention_backward_kernel( + num_query_heads, + num_kv_heads, + head_dim, + selection_width, + scale, + )( + query, + key, + value, + grad_output, + indices, + logsumexp, + delta, + grad_key, + grad_value, + ) + return grad_query, grad_key.to(query.dtype), grad_value.to(query.dtype) + + +@indexed_attention_backward.register_fake +def indexed_attention_backward_fake( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + output: torch.Tensor, + grad_output: torch.Tensor, + indices: torch.Tensor, + logsumexp: torch.Tensor, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return torch.empty_like(query), torch.empty_like(key), torch.empty_like(value) diff --git a/prime_kernels/indexed_attention/forward.py b/prime_kernels/indexed_attention/forward.py new file mode 100644 index 0000000..93e2a30 --- /dev/null +++ b/prime_kernels/indexed_attention/forward.py @@ -0,0 +1,304 @@ +# Vendored from tile-ai/tilelang (Apache 2.0), modified for indexed GQA. + +import tilelang +import torch +import torch.nn.functional as F +from tilelang import language as T + +from prime_kernels.indexed_attention.backward import indexed_attention_backward + +FORWARD_BLOCK_TOKENS = 64 +MAX_QUERY_HEADS_PER_KV_HEAD = 16 + + +@tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, +) +def indexed_attention_forward_kernel( + num_query_heads: int, + num_kv_heads: int, + head_dim: int, + selection_width: int, + scale: float, + block_tokens: int = FORWARD_BLOCK_TOKENS, + num_stages: int = 2, + threads: int = 256, +): + batch_size = T.dynamic("batch_size") + query_length = T.dynamic("query_length") + kv_length = T.dynamic("kv_length") + + query_heads_per_kv_head = num_query_heads // num_kv_heads + head_tile = max(tilelang.math.next_power_of_2(query_heads_per_kv_head), 16) + num_token_blocks = tilelang.cdiv(selection_width, block_tokens) + log2_scale = scale * 1.44269504 + + query_shape = [batch_size, query_length, num_query_heads, head_dim] + kv_shape = [batch_size, kv_length, num_kv_heads, head_dim] + indices_shape = [batch_size, query_length, selection_width] + output_shape = query_shape + logsumexp_shape = [batch_size, query_length, num_query_heads] + + @T.prim_func + def kernel( + query: T.Tensor(query_shape, T.bfloat16), + key: T.Tensor(kv_shape, T.bfloat16), + value: T.Tensor(kv_shape, T.bfloat16), + indices: T.Tensor(indices_shape, T.int32), + output: T.Tensor(output_shape, T.bfloat16), + logsumexp: T.Tensor(logsumexp_shape, T.float32), + ): + with T.Kernel(query_length, batch_size, num_kv_heads, threads=threads) as (query_token, batch, kv_head): + query_shared = T.alloc_shared([head_tile, head_dim], T.bfloat16) + key_shared = T.alloc_shared([block_tokens, head_dim], T.bfloat16) + value_shared = T.alloc_shared([block_tokens, head_dim], T.bfloat16) + probability_shared = T.alloc_shared([head_tile, block_tokens], T.bfloat16) + output_shared = T.alloc_shared([head_tile, head_dim], T.bfloat16) + logsumexp_shared = T.alloc_shared([head_tile], T.float32) + + valid_index = T.alloc_fragment([block_tokens], "bool") + scores = T.alloc_fragment([head_tile, block_tokens], T.float32) + output_accumulator = T.alloc_fragment([head_tile, head_dim], T.float32) + row_sum = T.alloc_fragment([head_tile], T.float32) + block_row_sum = T.alloc_fragment([head_tile], T.float32) + rescale = T.alloc_fragment([head_tile], T.float32) + row_max = T.alloc_fragment([head_tile], T.float32) + previous_row_max = T.alloc_fragment([head_tile], T.float32) + + first_query_head = kv_head * query_heads_per_kv_head + last_valid_kv = kv_length - 2 + + T.copy( + query[ + batch, + query_token, + first_query_head : first_query_head + query_heads_per_kv_head, + :, + ], + query_shared[:query_heads_per_kv_head, :], + ) + for head, dim in T.Parallel(head_tile - query_heads_per_kv_head, head_dim): + query_shared[query_heads_per_kv_head + head, dim] = 0 + + T.clear(output_accumulator) + T.clear(row_sum) + T.fill(row_max, -(2**30)) + + for selection_block in T.Pipelined(num_token_blocks, num_stages=num_stages): + for selected_token in T.Parallel(block_tokens): + valid_index[selected_token] = ( + indices[batch, query_token, selection_block * block_tokens + selected_token] <= last_valid_kv + ) + for head, selected_token in T.Parallel(head_tile, block_tokens): + scores[head, selected_token] = T.if_then_else( + valid_index[selected_token], + 0, + -T.infinity(T.float32), + ) + for selected_token, dim in T.Parallel(block_tokens, head_dim): + key_shared[selected_token, dim] = key[ + batch, + indices[batch, query_token, selection_block * block_tokens + selected_token], + kv_head, + dim, + ] + T.gemm( + query_shared, + key_shared, + scores, + transpose_B=True, + policy=T.GemmWarpPolicy.FullRow, + ) + + T.copy(row_max, previous_row_max) + T.reduce_max(scores, row_max, dim=1, clear=False) + for head in T.Parallel(head_tile): + row_max[head] = T.max(row_max[head], previous_row_max[head]) + rescale[head] = T.exp2((previous_row_max[head] - row_max[head]) * log2_scale) + for head, selected_token in T.Parallel(head_tile, block_tokens): + scores[head, selected_token] = T.exp2( + scores[head, selected_token] * log2_scale - row_max[head] * log2_scale + ) + T.reduce_sum(scores, block_row_sum, dim=1) + for head in T.Parallel(head_tile): + row_sum[head] = row_sum[head] * rescale[head] + block_row_sum[head] + for head, dim in T.Parallel(head_tile, head_dim): + output_accumulator[head, dim] *= rescale[head] + + T.copy(scores, probability_shared) + for selected_token, dim in T.Parallel(block_tokens, head_dim): + value_shared[selected_token, dim] = value[ + batch, + indices[batch, query_token, selection_block * block_tokens + selected_token], + kv_head, + dim, + ] + T.gemm( + probability_shared, + value_shared, + output_accumulator, + policy=T.GemmWarpPolicy.FullRow, + ) + + for head, dim in T.Parallel(head_tile, head_dim): + output_accumulator[head, dim] /= row_sum[head] + for head in T.Parallel(head_tile): + row_sum[head] = T.log2(row_sum[head]) + row_max[head] * log2_scale + + T.copy(output_accumulator, output_shared) + T.copy(row_sum, logsumexp_shared) + T.copy( + output_shared[:query_heads_per_kv_head, :], + output[ + batch, + query_token, + first_query_head : first_query_head + query_heads_per_kv_head, + :, + ], + ) + T.copy( + logsumexp_shared[:query_heads_per_kv_head], + logsumexp[ + batch, + query_token, + first_query_head : first_query_head + query_heads_per_kv_head, + ], + ) + + return kernel + + +def unsupported_shape_reason( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + indices: torch.Tensor, +) -> str | None: + if query.ndim != 4 or key.ndim != 4 or value.ndim != 4: + return "query, key, and value must have shape [batch, tokens, heads, head_dim]" + if indices.ndim != 3: + return "indices must have shape [batch, query_tokens, selected_tokens]" + if query.dtype != torch.bfloat16 or key.dtype != torch.bfloat16 or value.dtype != torch.bfloat16: + return "query, key, and value must use bfloat16" + if indices.dtype != torch.int32: + return "indices must use int32" + if not query.is_cuda or not key.is_cuda or not value.is_cuda or not indices.is_cuda: + return "all inputs must be CUDA tensors" + if not (query.device == key.device == value.device == indices.device): + return "all inputs must be on the same CUDA device" + batch_size, query_length, num_query_heads, head_dim = query.shape + key_batch, kv_length, num_kv_heads, key_head_dim = key.shape + if value.shape != key.shape: + return "key and value must have the same shape" + if key_batch != batch_size or key_head_dim != head_dim: + return "query, key, and value batch and head dimensions must match" + if indices.shape[:2] != (batch_size, query_length): + return "indices batch and query dimensions must match query" + if query_length == 0 or kv_length == 0 or indices.shape[-1] == 0: + return "query, key, value, and selection dimensions must be non-empty" + if num_kv_heads == 0 or num_query_heads % num_kv_heads: + return "query heads must be divisible by KV heads" + if num_query_heads // num_kv_heads > MAX_QUERY_HEADS_PER_KV_HEAD: + return f"at most {MAX_QUERY_HEADS_PER_KV_HEAD} query heads per KV head are supported" + if head_dim < 16 or head_dim & (head_dim - 1): + return "head_dim must be a power of two and at least 16" + return None + + +@torch.library.custom_op("prime_kernels::indexed_attention_forward", mutates_args=()) +def indexed_attention_forward( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + indices: torch.Tensor, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor]: + _, _, num_query_heads, head_dim = query.shape + _, _, num_kv_heads, _ = key.shape + selection_width = indices.shape[-1] + return indexed_attention_forward_kernel( + num_query_heads, + num_kv_heads, + head_dim, + selection_width, + scale, + )(query, key, value, indices) + + +@indexed_attention_forward.register_fake +def indexed_attention_forward_fake( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + indices: torch.Tensor, + scale: float, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.empty_like(query), query.new_empty(query.shape[:-1], dtype=torch.float32) + + +def indexed_attention_setup_context(ctx, inputs, output) -> None: + query, key, value, indices, scale = inputs + attention_output, logsumexp = output + ctx.save_for_backward(query, key, value, attention_output, indices, logsumexp) + ctx.scale = scale + ctx.mark_non_differentiable(logsumexp) + + +def indexed_attention_autograd_backward(ctx, grad_output: torch.Tensor, _grad_logsumexp: torch.Tensor | None): + query, key, value, output, indices, logsumexp = ctx.saved_tensors + grad_query, grad_key, grad_value = indexed_attention_backward( + query.detach(), + key.detach(), + value.detach(), + output.detach(), + grad_output, + indices, + logsumexp.detach(), + ctx.scale, + ) + return grad_query, grad_key, grad_value, None, None + + +indexed_attention_forward.register_autograd( + indexed_attention_autograd_backward, + setup_context=indexed_attention_setup_context, +) + + +def indexed_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + indices: torch.Tensor, + scale: float | None = None, +) -> torch.Tensor: + """Apply GQA over each query's selected tokens. + + Indices are in the key sequence's coordinate space. The value ``key.shape[1]`` is a + sentinel for unused entries, and every query must select at least one real token. + """ + reason = unsupported_shape_reason(query, key, value, indices) + if reason is not None: + raise ValueError(reason) + + scale = query.shape[-1] ** -0.5 if scale is None else scale + sentinel = key.shape[1] + key = torch.cat((key, key.new_zeros((key.shape[0], 1, key.shape[2], key.shape[3]))), dim=1) + value = torch.cat((value, value.new_zeros((value.shape[0], 1, value.shape[2], value.shape[3]))), dim=1) + + padding = (-indices.shape[-1]) % FORWARD_BLOCK_TOKENS + if padding: + indices = F.pad(indices, (0, padding), value=sentinel) + + output, _ = indexed_attention_forward( + query.contiguous(), + key.contiguous(), + value.contiguous(), + indices.contiguous(), + scale, + ) + return output diff --git a/prime_kernels/kernels.toml b/prime_kernels/kernels.toml index 2c6f5ec..3383b7c 100644 --- a/prime_kernels/kernels.toml +++ b/prime_kernels/kernels.toml @@ -30,6 +30,12 @@ python-only = true requires = ["torchao"] arch = ["10.0"] +[indexed_attention] +description = "Training forward and backward for token-indexed grouped-query attention" +python-only = true +requires = ["tilelang"] +arch = ["8.0", "9.0", "10.0"] + # rmsnorm is not built yet: only its sources are committed. Uncomment the table below to # put it back into the build (and to make the registry report on it). # From a4b79e655fdb0af57898790d2153234b93d5a4d3 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Mon, 31 Aug 2026 13:33:07 +0000 Subject: [PATCH 2/4] feat: add exact sparse index selection kernel --- prime_kernels/indexed_attention/__init__.py | 3 +- prime_kernels/indexed_attention/selection.py | 148 +++++++++++++++++++ 2 files changed, 150 insertions(+), 1 deletion(-) create mode 100644 prime_kernels/indexed_attention/selection.py diff --git a/prime_kernels/indexed_attention/__init__.py b/prime_kernels/indexed_attention/__init__.py index 6363734..4cfe0c3 100644 --- a/prime_kernels/indexed_attention/__init__.py +++ b/prime_kernels/indexed_attention/__init__.py @@ -1,3 +1,4 @@ from prime_kernels.indexed_attention.forward import indexed_attention, unsupported_shape_reason +from prime_kernels.indexed_attention.selection import select_indexed_blocks -__all__ = ["indexed_attention", "unsupported_shape_reason"] +__all__ = ["indexed_attention", "select_indexed_blocks", "unsupported_shape_reason"] diff --git a/prime_kernels/indexed_attention/selection.py b/prime_kernels/indexed_attention/selection.py new file mode 100644 index 0000000..0a78e00 --- /dev/null +++ b/prime_kernels/indexed_attention/selection.py @@ -0,0 +1,148 @@ +import tilelang +import torch +import torch.nn.functional as F +from tilelang import language as T + +# A 32K-token sequence produces a 1 GiB score matrix and runs in one pass. +SCORE_WORKSPACE_BYTES = 1024**3 + + +@tilelang.jit(out_idx=[-1]) +def indexed_selection_scores_kernel( + num_query_heads: int, + head_dim: int, + block_queries: int = 64, + block_keys: int = 128, + threads: int = 256, +): + query_tokens = T.dynamic("query_tokens") + key_blocks = T.dynamic("key_blocks") + + query_shape = [query_tokens * num_query_heads, head_dim] + key_shape = [key_blocks, head_dim] + bounds_shape = [query_tokens] + scores_shape = [query_tokens, key_blocks] + + @T.prim_func + def kernel( + query: T.Tensor(query_shape, T.bfloat16), + key: T.Tensor(key_shape, T.bfloat16), + starts: T.Tensor(bounds_shape, T.int32), + ends: T.Tensor(bounds_shape, T.int32), + scores: T.Tensor(scores_shape, T.float32), + ): + with T.Kernel( + T.ceildiv(key_blocks, block_keys), + T.ceildiv(query_tokens, block_queries), + threads=threads, + ) as (key_block, query_block): + query_shared = T.alloc_shared([block_queries, head_dim], T.float32) + key_shared = T.alloc_shared([block_keys, head_dim], T.float32) + head_scores = T.alloc_fragment([block_queries, block_keys], T.float32) + combined_scores = T.alloc_fragment([block_queries, block_keys], T.float32) + + for row, dim in T.Parallel(block_keys, head_dim): + key_index = key_block * block_keys + row + key_shared[row, dim] = T.if_then_else( + key_index < key_blocks, + T.cast(key[key_index, dim], T.float32), + 0, + ) + T.clear(combined_scores) + for head in T.serial(num_query_heads): + for row, dim in T.Parallel(block_queries, head_dim): + query_index = query_block * block_queries + row + query_shared[row, dim] = T.if_then_else( + query_index < query_tokens, + T.cast(query[query_index * num_query_heads + head, dim], T.float32), + 0, + ) + T.gemm( + query_shared, + key_shared, + head_scores, + transpose_B=True, + clear_accum=True, + policy=T.GemmWarpPolicy.FullRow, + ) + for row, column in T.Parallel(block_queries, block_keys): + combined_scores[row, column] += T.max(head_scores[row, column], 0) + + for row, column in T.Parallel(block_queries, block_keys): + query_index = query_block * block_queries + row + key_index = key_block * block_keys + column + if query_index < query_tokens and key_index < key_blocks: + if key_index < starts[query_index] or key_index >= ends[query_index]: + combined_scores[row, column] = -T.infinity(T.float32) + + T.copy( + combined_scores, + scores[query_block * block_queries, key_block * block_keys], + ) + + return kernel + + +def _select_blocks( + query: torch.Tensor, + key: torch.Tensor, + starts: torch.Tensor, + ends: torch.Tensor, + topk: int, +) -> torch.Tensor: + num_blocks = key.shape[0] + if num_blocks == 0: + return torch.zeros(query.shape[0], topk, dtype=torch.int32, device=query.device) + + chunk_size = max(1, SCORE_WORKSPACE_BYTES // (num_blocks * torch.float32.itemsize)) + selected_chunks = [] + selected_count = min(topk, num_blocks) + key = key.contiguous() + for query_chunk, start_chunk, end_chunk in zip( + query.split(chunk_size), + starts.split(chunk_size), + ends.split(chunk_size), + strict=True, + ): + scores = indexed_selection_scores_kernel( + query.shape[1], + query.shape[2], + )( + query_chunk.flatten(0, 1).contiguous(), + key, + start_chunk.contiguous(), + end_chunk.contiguous(), + ) + selected = scores.topk(selected_count, dim=-1).indices + if selected_count < topk: + selected = F.pad(selected, (0, topk - selected_count), value=num_blocks) + selected.masked_fill_( + (selected < start_chunk[:, None]) | (selected >= end_chunk[:, None]), + num_blocks, + ) + selected_chunks.append(selected.to(torch.int32)) + if len(selected_chunks) == 1: + return selected_chunks[0] + return torch.cat(selected_chunks) + + +@torch.library.custom_op("prime_kernels::select_indexed_blocks", mutates_args=()) +def select_indexed_blocks( + query: torch.Tensor, + key: torch.Tensor, + starts: torch.Tensor, + ends: torch.Tensor, + topk: int, +) -> torch.Tensor: + return _select_blocks(query, key, starts, ends, topk) + + +@select_indexed_blocks.register_fake +def select_indexed_blocks_fake( + query: torch.Tensor, + key: torch.Tensor, + starts: torch.Tensor, + ends: torch.Tensor, + topk: int, +) -> torch.Tensor: + return query.new_empty((query.shape[0], topk), dtype=torch.int32) From e309f53044f245b613693d7b623fd3a627920130 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Mon, 31 Aug 2026 14:15:43 +0000 Subject: [PATCH 3/4] perf: vendor vLLM QSA selection --- prime_kernels/indexed_attention/LICENSE.vllm | 201 +++ prime_kernels/indexed_attention/__init__.py | 17 + .../csrc/persistent_topk.cuh | 1363 +++++++++++++++++ prime_kernels/indexed_attention/csrc/topk.cu | 287 ++++ .../csrc/topk_histogram_4096.cuh | 563 +++++++ prime_kernels/indexed_attention/selection.py | 267 ++-- prime_kernels/kernels.toml | 7 +- setup.py | 2 + 8 files changed, 2593 insertions(+), 114 deletions(-) create mode 100644 prime_kernels/indexed_attention/LICENSE.vllm create mode 100644 prime_kernels/indexed_attention/csrc/persistent_topk.cuh create mode 100644 prime_kernels/indexed_attention/csrc/topk.cu create mode 100644 prime_kernels/indexed_attention/csrc/topk_histogram_4096.cuh diff --git a/prime_kernels/indexed_attention/LICENSE.vllm b/prime_kernels/indexed_attention/LICENSE.vllm new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/prime_kernels/indexed_attention/LICENSE.vllm @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/prime_kernels/indexed_attention/__init__.py b/prime_kernels/indexed_attention/__init__.py index 4cfe0c3..569b33b 100644 --- a/prime_kernels/indexed_attention/__init__.py +++ b/prime_kernels/indexed_attention/__init__.py @@ -1,4 +1,21 @@ +from __future__ import annotations + +import torch + +from . import _C # noqa: F401 from prime_kernels.indexed_attention.forward import indexed_attention, unsupported_shape_reason from prime_kernels.indexed_attention.selection import select_indexed_blocks __all__ = ["indexed_attention", "select_indexed_blocks", "unsupported_shape_reason"] + + +@torch.library.register_fake("prime_indexed_attention::persistent_topk") +def _persistent_topk_fake( + logits: torch.Tensor, + lengths: torch.Tensor, + output: torch.Tensor, + workspace: torch.Tensor, + k: int, + max_seq_len: int, +) -> None: + return None diff --git a/prime_kernels/indexed_attention/csrc/persistent_topk.cuh b/prime_kernels/indexed_attention/csrc/persistent_topk.cuh new file mode 100644 index 0000000..9ae0fd4 --- /dev/null +++ b/prime_kernels/indexed_attention/csrc/persistent_topk.cuh @@ -0,0 +1,1363 @@ +/* + * Persistent TopK Scheduler for DSA Indexer + */ + +#ifndef PERSISTENT_TOPK_CUH_ +#define PERSISTENT_TOPK_CUH_ + +#include +#include +#include +#include +#include + +#include "topk_histogram_4096.cuh" + +namespace vllm { +namespace persistent { + +// ============================================================================ +// Constants +// ============================================================================ + +constexpr int kThreadsPerBlock = 1024; +constexpr int RADIX = 256; + +// Medium path: all shared state in dynamic smem (no static __shared__, +// which would inflate the kernel's smem footprint and kill occupancy +// for the decode/trivial paths). +constexpr size_t kMediumHistBytes = 2 * (RADIX + 128) * sizeof(int); // 3072 +constexpr size_t kMediumScalarsBytes = 5 * sizeof(int); // 20 +constexpr size_t kMediumHeaderSize = + (kMediumHistBytes + kMediumScalarsBytes + 127) & ~size_t(127); // 3200 +constexpr int MAX_BUFFERED_ITEMS = 4096; +constexpr size_t kSmemMedium = + kMediumHeaderSize + 2 * MAX_BUFFERED_ITEMS * sizeof(int); // 35968 +constexpr uint32_t RADIX_THRESHOLD = 32768; + +// Decode path constants +constexpr int kDecodeBins = 2048; +constexpr uint32_t HIST2048_THRESHOLD = 8192; + +// Large path: fixed shared memory for histograms + scalars +constexpr size_t kFixedSmemLarge = + ((RADIX + RADIX + 5) * sizeof(uint32_t) + 15) & ~size_t(15); + +// ============================================================================ +// Common helpers +// ============================================================================ + +__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { + uint32_t bits = __float_as_uint(x); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t { + __half h = __float2half_rn(x); + uint16_t bits = __half_as_ushort(h); + uint16_t key = (bits & 0x8000) ? static_cast(~bits) + : static_cast(bits | 0x8000); + return static_cast(key >> 8); +} + +// ============================================================================ +// Vectorized load helpers +// ============================================================================ + +// Unconditional float4 load with cache hint (.cg = cache at global level only). +__device__ __forceinline__ void load_float4(const float* ptr, float& v0, + float& v1, float& v2, float& v3) { + uint32_t r0, r1, r2, r3; + asm volatile("ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%4];\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "l"(ptr)); + v0 = __uint_as_float(r0); + v1 = __uint_as_float(r1); + v2 = __uint_as_float(r2); + v3 = __uint_as_float(r3); +} + +// Per-element predicated scalar loads with -inf default. +__device__ __forceinline__ void load_float4_predicated(const float* ptr, + int base, int seq_len, + float& v0, float& v1, + float& v2, float& v3) { + uint32_t r0, r1, r2, r3; + int p0 = (base < seq_len); + int p1 = (base + 1 < seq_len); + int p2 = (base + 2 < seq_len); + int p3 = (base + 3 < seq_len); + asm volatile( + "{\n" + " .reg .pred pr0, pr1, pr2, pr3;\n" + " setp.ne.u32 pr0, %4, 0;\n" + " setp.ne.u32 pr1, %5, 0;\n" + " setp.ne.u32 pr2, %6, 0;\n" + " setp.ne.u32 pr3, %7, 0;\n" + " mov.u32 %0, 0xFF800000;\n" + " mov.u32 %1, 0xFF800000;\n" + " mov.u32 %2, 0xFF800000;\n" + " mov.u32 %3, 0xFF800000;\n" + " @pr0 ld.global.cg.u32 %0, [%8];\n" + " @pr1 ld.global.cg.u32 %1, [%8+4];\n" + " @pr2 ld.global.cg.u32 %2, [%8+8];\n" + " @pr3 ld.global.cg.u32 %3, [%8+12];\n" + "}\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr)); + v0 = __uint_as_float(r0); + v1 = __uint_as_float(r1); + v2 = __uint_as_float(r2); + v3 = __uint_as_float(r3); +} + +// ============================================================================ +// Large path: inter-CTA coordination state (one per group) +// ============================================================================ + +struct RadixRowState { + uint32_t histogram[3][256]; // Triple-buffered histograms + uint32_t remaining_k; + uint32_t prefix; + int arrival_counter; + int output_counter; +}; + +// ============================================================================ +// Kernel parameters +// ============================================================================ + +struct PersistentTopKParams { + const float* __restrict__ input; // [num_rows, stride] + int32_t* __restrict__ output; // [num_rows, top_k] + const int32_t* __restrict__ lengths; // [num_rows] + RadixRowState* row_states; // large path: per-group state + uint32_t num_rows; + uint32_t stride; + uint32_t top_k; // actual k value for output stride + uint32_t chunk_size; // large path: elements per CTA + uint32_t ctas_per_group; // 1=medium, >1=large + uint32_t max_seq_len; // max seq_len across all rows (for early CTA exit) +}; + +// ============================================================================ +// Decode path: 2048-bin histogram for short sequences (seq_len <= 8192) +// Uses 11-bit half-precision bins for fine granularity. +// One histogram pass typically suffices since 8192/2048 = 4 elements/bin avg. +// ============================================================================ + +// 11-bit bin from half-precision representation (ascending: high values -> high +// bins) +__device__ __forceinline__ uint32_t decode_bin(float x) { + __half hx = __float2half(x); + uint16_t bits = __half_as_ushort(hx); + uint16_t key = (bits & 0x8000) ? static_cast(~bits) + : static_cast(bits | 0x8000); + return key >> 5; +} + +template +__device__ __noinline__ void histogram_2048_topk( + const float* __restrict__ logits, int32_t* __restrict__ output_indices, + int32_t seq_len) { + extern __shared__ int decode_smem[]; + const int tx = threadIdx.x; + const int lane = tx & 31; + + // ---- Layout constants ---- + constexpr int SBASE = 8192 - 8; // 8184 + constexpr int RHIST = RADIX + 128; // 384 + constexpr int BOFF = 2 * RHIST; // 768 + constexpr int DBUF = (SBASE - BOFF) / 2; // 3708 + constexpr int MAX_ITEMS_PER_THREAD = + (HIST2048_THRESHOLD + kThreadsPerBlock - 1) / kThreadsPerBlock; + + enum : int { sTHR = 0, sOUT = 1, sREF = 2, sFIN = 3, sBUF0 = 4, sBUF1 = 5 }; + + // ---- Initialize scalars (prevents stale data from prior rows) ---- + if (tx < 8) { + decode_smem[SBASE + tx] = 0; + } + + // ---- Phase 1: Build 2048-bin histogram with float4 vectorized loads ---- + int* histo = decode_smem; + uint16_t reg_bins[MAX_ITEMS_PER_THREAD]; + int nitems = 0; + + for (int i = tx; i < kDecodeBins; i += kThreadsPerBlock) { + histo[i] = 0; + } + __syncthreads(); + + const int n_vec = (seq_len + 3) >> 2; + const bool row_aligned = ((reinterpret_cast(logits) & 15) == 0); + + for (int i = tx; i < n_vec; i += kThreadsPerBlock) { + const int base = i << 2; + float v0, v1, v2, v3; + + if (row_aligned && base + 3 < seq_len) { + load_float4(logits + base, v0, v1, v2, v3); + } else { + load_float4_predicated(logits + base, base, seq_len, v0, v1, v2, v3); + } + + const uint16_t b0 = static_cast(decode_bin(v0)); + const uint16_t b1 = static_cast(decode_bin(v1)); + const uint16_t b2 = static_cast(decode_bin(v2)); + const uint16_t b3 = static_cast(decode_bin(v3)); + reg_bins[nitems++] = b0; + reg_bins[nitems++] = b1; + reg_bins[nitems++] = b2; + reg_bins[nitems++] = b3; + atomicAdd(&histo[b0], 1); + atomicAdd(&histo[b1], 1); + atomicAdd(&histo[b2], 1); + atomicAdd(&histo[b3], 1); + } + __syncthreads(); + + // ---- CUB suffix sum ---- + using BlockScanT = cub::BlockScan; + const int h0 = histo[2 * tx]; + const int pair_sum = h0 + histo[2 * tx + 1]; + + auto& scan_storage = *reinterpret_cast( + decode_smem + kDecodeBins); + + int pair_prefix, total; + BlockScanT(scan_storage).ExclusiveSum(pair_sum, pair_prefix, total); + + // Find threshold bin purely from registers + const int pair_suffix = total - pair_prefix; + + if (pair_suffix >= TopK && (pair_suffix - h0) < TopK) { + decode_smem[SBASE + sTHR] = 2 * tx; + } + { + const int right_suf = pair_suffix - h0; + const int next_suf = pair_suffix - pair_sum; + if (right_suf >= TopK && next_suf < TopK) { + decode_smem[SBASE + sTHR] = 2 * tx + 1; + } + } + __syncthreads(); + + const int threshold = decode_smem[SBASE + sTHR]; + + // ---- Phase 2: Collection with warp-aggregated atomicAdds ---- + int* bufs[2] = {decode_smem + BOFF, decode_smem + BOFF + DBUF}; + const int sOUT_abs = SBASE + sOUT; + const int sBUF0_abs = SBASE + sBUF0; + + { + const uint32_t uthr = static_cast(threshold); + int item = 0; + const int n_vec_iters = (n_vec + kThreadsPerBlock - 1) / kThreadsPerBlock; + + for (int iter = 0; iter < n_vec_iters; iter++) { + const int i = tx + iter * kThreadsPerBlock; + const bool vec_valid = (i < n_vec); + const int base_idx = i << 2; + +#pragma unroll 4 + for (int sub = 0; sub < 4; sub++) { + const int elem_idx = base_idx + sub; + uint32_t bin = 0; + if (vec_valid) bin = reg_bins[item++]; + const bool is_above = vec_valid && (bin > uthr); + const bool is_equal = vec_valid && (bin == uthr); + + const uint32_t above_mask = __ballot_sync(0xffffffff, is_above); + if (above_mask) { + const int above_count = __popc(above_mask); + const int above_rank = __popc(above_mask & ((1u << lane) - 1)); + int above_base; + if (lane == 0) { + above_base = atomicAdd(&decode_smem[sOUT_abs], above_count); + } + above_base = __shfl_sync(0xffffffff, above_base, 0); + if (is_above) { + output_indices[above_base + above_rank] = elem_idx; + } + } + + const uint32_t equal_mask = __ballot_sync(0xffffffff, is_equal); + if (equal_mask) { + const int equal_count = __popc(equal_mask); + const int equal_rank = __popc(equal_mask & ((1u << lane) - 1)); + int equal_base; + if (lane == 0) { + equal_base = atomicAdd(&decode_smem[sBUF0_abs], equal_count); + } + equal_base = __shfl_sync(0xffffffff, equal_base, 0); + if (is_equal && __builtin_expect(equal_base + equal_rank < DBUF, 1)) { + bufs[0][equal_base + equal_rank] = elem_idx; + } + } + } + } + } + __syncthreads(); + + int remaining_k = TopK - decode_smem[SBASE + sOUT]; + if (remaining_k <= 0) return; + + // If all buffered elements fit, output them all (common for short seqs) + const int raw_buf0 = decode_smem[SBASE + sBUF0]; + if (raw_buf0 <= remaining_k) { + const int nb = (raw_buf0 < DBUF) ? raw_buf0 : DBUF; + const int base = decode_smem[SBASE + sOUT]; + for (int i = tx; i < nb; i += kThreadsPerBlock) { + output_indices[base + i] = bufs[0][i]; + } + __syncthreads(); + return; + } + + // ---- Phase 3: Deferred refinement (rare path) ---- + int* refine[2] = {decode_smem, decode_smem + RHIST}; + const int num_buf0 = (raw_buf0 < DBUF) ? raw_buf0 : DBUF; + + for (int i = tx; i < RHIST; i += kThreadsPerBlock) { + refine[0][i] = 0; + } + __syncthreads(); + + for (int i = tx; i < num_buf0; i += kThreadsPerBlock) { + const uint32_t fp32 = convert_to_uint32_v2(logits[bufs[0][i]]); + atomicAdd(&refine[0][(fp32 >> 24) & 0xFF], 1); + } + __syncthreads(); + + auto compute_suffix_sum = [&]() { +#pragma unroll 8 + for (int i = 0; i < 8; ++i) { + if (tx < RADIX) { + const int stride = 1 << i; + const int s = i & 1; + const int d = s ^ 1; + int value = refine[s][tx]; + if (tx < RADIX - stride) value += refine[s][tx + stride]; + refine[d][tx] = value; + } + __syncthreads(); + } + }; + +#pragma unroll 4 + for (int pass = 0; pass < 4; ++pass) { + const int src = pass & 1; + const int dst = src ^ 1; + + const int raw_buf = decode_smem[SBASE + sBUF0 + src]; + const int num_buffered = (raw_buf < DBUF) ? raw_buf : DBUF; + + compute_suffix_sum(); + + if (tx < RADIX && refine[0][tx] > remaining_k && + refine[0][tx + 1] <= remaining_k) { + decode_smem[SBASE + sREF] = tx; + decode_smem[SBASE + sBUF0 + dst] = 0; + decode_smem[SBASE + sFIN] = remaining_k - refine[0][tx + 1]; + } + __syncthreads(); + + const int ref_thr = decode_smem[SBASE + sREF]; + remaining_k -= refine[0][ref_thr + 1]; + const int bit_offset = 24 - pass * 8; + + if (remaining_k == 0) { + for (int i = tx; i < num_buffered; i += kThreadsPerBlock) { + const int idx = bufs[src][i]; + const uint32_t fp32 = convert_to_uint32_v2(logits[idx]); + if (((fp32 >> bit_offset) & 0xFF) > static_cast(ref_thr)) { + const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1); + output_indices[pos] = idx; + } + } + __syncthreads(); + break; + } + + __syncthreads(); + if (tx < RADIX + 1) refine[0][tx] = 0; + __syncthreads(); + + for (int i = tx; i < num_buffered; i += kThreadsPerBlock) { + const int idx = bufs[src][i]; + const float logit_val = logits[idx]; + const uint32_t fp32 = convert_to_uint32_v2(logit_val); + const int bin = (fp32 >> bit_offset) & 0xFF; + + if (bin > ref_thr) { + const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1); + output_indices[pos] = idx; + } else if (bin == ref_thr) { + if (pass == 3) { + const int slot = atomicAdd(&decode_smem[SBASE + sFIN], -1); + if (slot > 0) output_indices[TopK - slot] = idx; + } else { + const int bp = atomicAdd(&decode_smem[SBASE + sBUF0 + dst], 1); + if (__builtin_expect(bp < DBUF, 1)) { + bufs[dst][bp] = idx; + const int nbo = bit_offset - 8; + atomicAdd(&refine[0][(fp32 >> nbo) & 0xFF], 1); + } + } + } + } + __syncthreads(); + } +} + +// ============================================================================ +// Medium path: coarse FP16 histogram + 4-pass FP32 radix refinement +// For sequences 8K < seq_len <= 64K. +// ============================================================================ + +// Adapted from: +// https://github.com/sgl-project/sglang/blob/v0.5.8/sgl-kernel/csrc/elementwise/topk.cu#L87 +// by: DarkSharpness +// which at the same time is an optimized topk kernel copied from tilelang +// kernel +template +__device__ __noinline__ void histogram_256_topk( + const float* __restrict__ logits, int* __restrict__ output_indices, + int logits_offset, int seq_len) { + // All shared state lives in dynamic shared memory to avoid static + extern __shared__ char medium_smem[]; + + int (*shared_histogram)[RADIX + 128] = + reinterpret_cast(medium_smem); + int* medium_scalars = reinterpret_cast(medium_smem + kMediumHistBytes); + int& shared_output_count = medium_scalars[0]; + int& shared_threshold_bin = medium_scalars[1]; + int* shared_buffered_count = &medium_scalars[2]; + int& shared_final_k = medium_scalars[4]; + int (*buffered_indices)[MAX_BUFFERED_ITEMS] = + reinterpret_cast(medium_smem + + kMediumHeaderSize); + + const int thread_id = threadIdx.x; + int remaining_k = TopK; + + if (thread_id < RADIX + 1) { + shared_histogram[0][thread_id] = 0; + } + __syncthreads(); + + for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { + const auto bin = convert_to_uint8(logits[idx + logits_offset]); + atomicAdd(&shared_histogram[0][bin], 1); + } + __syncthreads(); + + auto compute_cumulative_sum = [&]() { +#pragma unroll 8 + for (int i = 0; i < 8; ++i) { + if (__builtin_expect(thread_id < RADIX, 1)) { + const int stride = 1 << i; + const int src_buffer = i & 1; + const int dst_buffer = src_buffer ^ 1; + int value = shared_histogram[src_buffer][thread_id]; + if (thread_id < RADIX - stride) { + value += shared_histogram[src_buffer][thread_id + stride]; + } + shared_histogram[dst_buffer][thread_id] = value; + } + __syncthreads(); + } + }; + + compute_cumulative_sum(); + + if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && + shared_histogram[0][thread_id + 1] <= remaining_k) { + shared_threshold_bin = thread_id; + shared_buffered_count[0] = 0; + shared_output_count = 0; + } + __syncthreads(); + + const int threshold_bin = shared_threshold_bin; + remaining_k -= shared_histogram[0][threshold_bin + 1]; + + if (remaining_k == 0) { + for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { + const int bin = convert_to_uint8(logits[idx + logits_offset]); + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + output_indices[output_pos] = idx; + } + } + __syncthreads(); + return; + } + + __syncthreads(); + if (thread_id < RADIX + 1) { + shared_histogram[0][thread_id] = 0; + } + __syncthreads(); + + for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { + const float logit_value = logits[idx + logits_offset]; + const int bin = convert_to_uint8(logit_value); + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + output_indices[output_pos] = idx; + } else if (bin == threshold_bin) { + const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1); + if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { + buffered_indices[0][buffer_pos] = idx; + const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); + const int next_bin = (fp32_bits >> 24) & 0xFF; + atomicAdd(&shared_histogram[0][next_bin], 1); + } + } + } + __syncthreads(); + +#pragma unroll 4 + for (int pass = 0; pass < 4; ++pass) { + const int src_buffer = pass % 2; + const int dst_buffer = src_buffer ^ 1; + const int raw_buffered = shared_buffered_count[src_buffer]; + const int num_buffered = + (raw_buffered < MAX_BUFFERED_ITEMS) ? raw_buffered : MAX_BUFFERED_ITEMS; + + compute_cumulative_sum(); + + if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && + shared_histogram[0][thread_id + 1] <= remaining_k) { + shared_threshold_bin = thread_id; + shared_buffered_count[dst_buffer] = 0; + shared_final_k = remaining_k - shared_histogram[0][thread_id + 1]; + } + __syncthreads(); + + const int threshold_bin = shared_threshold_bin; + remaining_k -= shared_histogram[0][threshold_bin + 1]; + const int bit_offset = 24 - pass * 8; + + if (remaining_k == 0) { + for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { + const int idx = buffered_indices[src_buffer][i]; + const uint32_t fp32_bits = + convert_to_uint32_v2(logits[idx + logits_offset]); + const int bin = (fp32_bits >> bit_offset) & 0xFF; + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + output_indices[output_pos] = idx; + } + } + __syncthreads(); + break; + } + + __syncthreads(); + if (thread_id < RADIX + 1) { + shared_histogram[0][thread_id] = 0; + } + __syncthreads(); + + for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { + const int idx = buffered_indices[src_buffer][i]; + const float logit_value = logits[idx + logits_offset]; + const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); + const int bin = (fp32_bits >> bit_offset) & 0xFF; + if (bin > threshold_bin) { + const int output_pos = atomicAdd(&shared_output_count, 1); + output_indices[output_pos] = idx; + } else if (bin == threshold_bin) { + if (pass == 3) { + const int slot = atomicAdd(&shared_final_k, -1); + if (slot > 0) { + output_indices[TopK - slot] = idx; + } + } else { + const int buffer_pos = + atomicAdd(&shared_buffered_count[dst_buffer], 1); + if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { + buffered_indices[dst_buffer][buffer_pos] = idx; + const int next_bit_offset = bit_offset - 8; + const int next_bin = (fp32_bits >> next_bit_offset) & 0xFF; + atomicAdd(&shared_histogram[0][next_bin], 1); + } + } + } + } + __syncthreads(); + } +} + +// ============================================================================ +// Inter-CTA sync primitives +// ============================================================================ + +__device__ __forceinline__ int ld_acquire(int* ptr) { + int state = 0; +#if (__CUDA_ARCH__ >= 700) + asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" + : "=r"(state) + : "l"(ptr)); +#else + asm volatile("ld.cg.global.b32 %0, [%1];\n" : "=r"(state) : "l"(ptr)); +#endif + return state; +} + +__device__ __forceinline__ void red_release(int* ptr, int val) { +#if (__CUDA_ARCH__ >= 700) + asm volatile("fence.acq_rel.gpu;\n"); + asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" + : + : "l"(ptr), "r"(val)); +#else + __threadfence(); + atomicAdd(ptr, val); +#endif +} + +__device__ __forceinline__ void st_release(int* ptr, int val) { +#if (__CUDA_ARCH__ >= 700) + asm volatile("fence.acq_rel.gpu;\n"); + asm volatile("st.release.gpu.global.b32 [%0], %1;\n" : : "l"(ptr), "r"(val)); +#else + __threadfence(); + atomicExch(ptr, val); +#endif +} + +__device__ __forceinline__ void wait_ge(int* ptr, int target_val, + int thread_idx) { + if (thread_idx == 0) { +#pragma unroll 1 + while (ld_acquire(ptr) < target_val) { + } + } + __syncthreads(); +} + +// ============================================================================ +// Large path: multi-CTA radix select for sequences > 64K +// +// Each row is processed by a group of CTAs. Each CTA loads its chunk into +// shared memory as ordered uint32, then participates in 4 rounds of +// coordinated radix select via global-memory histograms and barriers. +// ============================================================================ + +// ============================================================================ +// Multi-CTA cooperative RadixTopK for a single large row. +// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 +// ============================================================================ + +template +__device__ void radix_topk(const float* __restrict__ row_input, + int32_t* __restrict__ row_output, uint32_t seq_len, + uint32_t my_chunk_start, uint32_t chunk_size, + uint32_t* local_histogram, uint32_t* suffix_sum, + uint32_t* shared_scalars, uint32_t* shared_ordered, + RadixRowState* state, uint32_t cta_in_group, + uint32_t ctas_per_group, int& barrier_phase, + uint32_t radix_iter, uint32_t tx) { + const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len) + ? my_chunk_start + chunk_size + : seq_len; + const uint32_t actual_chunk_size = + (my_chunk_start < seq_len) ? (my_chunk_end - my_chunk_start) : 0; + + // -- Stage 1: Load chunk to shared memory as ordered uint32 -- + { + const uint32_t aligned_size = (actual_chunk_size / VEC_SIZE) * VEC_SIZE; + + for (uint32_t i = tx * VEC_SIZE; i < aligned_size; + i += kThreadsPerBlock * VEC_SIZE) { + const float* src = row_input + my_chunk_start + i; + if constexpr (VEC_SIZE == 4) { + float4 v = *reinterpret_cast(src); + shared_ordered[i] = convert_to_uint32_v2(v.x); + shared_ordered[i + 1] = convert_to_uint32_v2(v.y); + shared_ordered[i + 2] = convert_to_uint32_v2(v.z); + shared_ordered[i + 3] = convert_to_uint32_v2(v.w); + } else if constexpr (VEC_SIZE == 2) { + float2 v = *reinterpret_cast(src); + shared_ordered[i] = convert_to_uint32_v2(v.x); + shared_ordered[i + 1] = convert_to_uint32_v2(v.y); + } else { + shared_ordered[i] = convert_to_uint32_v2(*src); + } + } + for (uint32_t i = aligned_size + tx; i < actual_chunk_size; + i += kThreadsPerBlock) { + shared_ordered[i] = convert_to_uint32_v2(row_input[my_chunk_start + i]); + } + } + __syncthreads(); + + // -- Init radix select state -- + if (tx == 0) { + shared_scalars[0] = 0; // prefix + shared_scalars[1] = TopK; // remaining_k + } + __syncthreads(); + + // -- Initial barrier -- + if (tx == 0) { + red_release(&state->arrival_counter, 1); + } + wait_ge(&state->arrival_counter, + (barrier_phase + 1) * static_cast(ctas_per_group), tx); + barrier_phase++; + __syncthreads(); + + if (cta_in_group == 0 && tx == 0) { + st_release(&state->output_counter, 0); + } + + // -- Stage 2: 4 rounds of radix select -- + for (uint32_t round = 0; round < 4; round++) { + const uint32_t global_round = radix_iter * 4 + round; + const uint32_t shift = 24 - round * 8; + const uint32_t prefix = shared_scalars[0]; + const uint32_t remaining_k = shared_scalars[1]; + + uint32_t* current_hist = state->histogram[global_round % 3]; + uint32_t* next_hist = state->histogram[(global_round + 1) % 3]; + + for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { + local_histogram[i] = 0; + } + __syncthreads(); + + for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { + uint32_t ordered = shared_ordered[i]; + uint32_t mask = (round == 0) ? 0u : (~0u << (32 - round * 8)); + if ((ordered & mask) == prefix) { + uint32_t bucket = (ordered >> shift) & 0xFF; + atomicAdd(&local_histogram[bucket], 1); + } + } + __syncthreads(); + + for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { + if (local_histogram[i] > 0) { + atomicAdd(¤t_hist[i], local_histogram[i]); + } + } + + if (cta_in_group == 0) { + for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { + next_hist[i] = 0; + } + } + + if (tx == 0) { + red_release(&state->arrival_counter, 1); + } + wait_ge(&state->arrival_counter, + (barrier_phase + 1) * static_cast(ctas_per_group), tx); + barrier_phase++; + __syncthreads(); + + for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { + suffix_sum[i] = current_hist[i]; + } + __syncthreads(); + + for (uint32_t stride = 1; stride < RADIX; stride *= 2) { + uint32_t val = 0; + if (tx < RADIX) { + val = suffix_sum[tx]; + if (tx + stride < RADIX) val += suffix_sum[tx + stride]; + } + __syncthreads(); + if (tx < RADIX) suffix_sum[tx] = val; + __syncthreads(); + } + + if (tx == 0) { + shared_scalars[2] = 0; + shared_scalars[3] = remaining_k; + } + __syncthreads(); + + if (tx < RADIX) { + uint32_t count_ge = suffix_sum[tx]; + uint32_t count_gt = (tx + 1 < RADIX) ? suffix_sum[tx + 1] : 0; + if (count_ge >= remaining_k && count_gt < remaining_k) { + shared_scalars[2] = tx; + shared_scalars[3] = remaining_k - count_gt; + } + } + __syncthreads(); + + if (tx == 0) { + shared_scalars[0] = prefix | (shared_scalars[2] << shift); + shared_scalars[1] = shared_scalars[3]; + } + __syncthreads(); + } // end 4 radix rounds + + // -- Count local > pivot elements -- + const uint32_t ordered_pivot = shared_scalars[0]; + + if (tx == 0) suffix_sum[0] = 0; + __syncthreads(); + + uint32_t my_gt_count = 0; + for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { + if (shared_ordered[i] > ordered_pivot) my_gt_count++; + } + for (int offset = 16; offset > 0; offset /= 2) { + my_gt_count += __shfl_down_sync(0xffffffff, my_gt_count, offset); + } + if (tx % 32 == 0 && my_gt_count > 0) { + atomicAdd(&suffix_sum[0], my_gt_count); + } + __syncthreads(); + const uint32_t local_gt_count = suffix_sum[0]; + + // -- Stage 3: Collect top-k indices -- + if (tx == 0) { + local_histogram[0] = 0; + if (local_gt_count > 0) { + local_histogram[1] = + atomicAdd(&state->output_counter, static_cast(local_gt_count)); + } + } + __syncthreads(); + + for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { + if (shared_ordered[i] > ordered_pivot) { + uint32_t local_pos = atomicAdd(&local_histogram[0], 1); + int pos = static_cast(local_histogram[1]) + local_pos; + row_output[pos] = static_cast(my_chunk_start + i); + } + } + + if (tx == 0) { + red_release(&state->arrival_counter, 1); + } + wait_ge(&state->arrival_counter, + (barrier_phase + 1) * static_cast(ctas_per_group), tx); + barrier_phase++; + __syncthreads(); + + for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { + if (shared_ordered[i] == ordered_pivot) { + int pos = atomicAdd(&state->output_counter, 1); + if (pos < TopK) { + row_output[pos] = static_cast(my_chunk_start + i); + } + } + } +} + +// ============================================================================ +// Persistent kernel — BS≤32, decode/medium/large paths with RadixTopK +// BS>32 uses standalone histogram_256_buffered_topk (separate kernel, +// see filtered_topk.cuh) +// ============================================================================ + +template +__global__ void __launch_bounds__(kThreadsPerBlock, 2) + persistent_topk_kernel(PersistentTopKParams params) { + const uint32_t tx = threadIdx.x; + extern __shared__ uint8_t smem_raw[]; + + // ======================================================================== + // Group mode: multi-CTA groups with static round-robin row assignment. + // Non-large rows: CTA-0 handles trivial/decode/medium. + // Large rows: all CTAs in the group cooperate via RadixTopK. + // ======================================================================== + const uint32_t ctas_per_group = params.ctas_per_group; + const uint32_t group_id = blockIdx.x / ctas_per_group; + const uint32_t cta_in_group = blockIdx.x % ctas_per_group; + const uint32_t num_groups = gridDim.x / ctas_per_group; + const uint32_t chunk_size = params.chunk_size; + + if (blockIdx.x >= num_groups * ctas_per_group) return; + + // Early exit: non-CTA-0 threads are never needed if no large rows exist + if (cta_in_group != 0 && params.max_seq_len <= RADIX_THRESHOLD) return; + + uint32_t* local_histogram = reinterpret_cast(smem_raw); + uint32_t* suffix_sum = local_histogram + RADIX; + uint32_t* shared_scalars = suffix_sum + RADIX; + uint32_t* shared_ordered = + reinterpret_cast(smem_raw + kFixedSmemLarge); + + // RadixRowState for multi-CTA cooperative radix. + // Zero-initialization is done host-side via cudaMemsetAsync in topk.cu + // before launch — that gives a stream-ordered happens-before edge for all + // CTAs, which the previous in-kernel init (CTA-0 only + intra-CTA + // __syncthreads) did not provide and which manifested as a race against + // CTA-1+'s first red_release on arrival_counter. + RadixRowState* state = ¶ms.row_states[group_id]; + + int barrier_phase = 0; + uint32_t radix_iter = 0; + const uint32_t total_iters = (params.num_rows + num_groups - 1) / num_groups; + + for (uint32_t iter = 0; iter < total_iters; iter++) { + // Static round-robin: all CTAs in the group implicitly agree on the row + uint32_t row_idx = group_id + iter * num_groups; + if (row_idx >= params.num_rows) break; + + // Clamp the row length before any decision is made on it. + // + // `lengths` is int32 and is consumed here as uint32, so a negative value + // (e.g. a padded decode slot whose per-token context length underflowed) + // would reinterpret as ~4e9 and sail past every threshold below. Any + // value beyond the row width would also read into the next row. + // + // Clamping to max_seq_len additionally keeps this per-row decision + // consistent with the `cta_in_group != 0` early exit above, which is + // taken from the host-side scalar: when max_seq_len <= RADIX_THRESHOLD + // the non-leader CTAs return immediately, so a leader that reached the + // cooperative radix path would wait on the inter-CTA barrier for peers + // that no longer exist and spin until the kernel is killed. + const int32_t raw_len = params.lengths[row_idx]; + const uint32_t row_bound = + params.stride < params.max_seq_len ? params.stride : params.max_seq_len; + const uint32_t non_negative_len = + raw_len > 0 ? static_cast(raw_len) : 0u; + const uint32_t seq_len = + non_negative_len < row_bound ? non_negative_len : row_bound; + int32_t* row_output = params.output + row_idx * params.top_k; + const float* row_input = params.input + row_idx * params.stride; + + if (seq_len <= RADIX_THRESHOLD) { + if (cta_in_group == 0) { + if (seq_len <= static_cast(TopK)) { + // Trivial case: seq_len <= TopK + for (uint32_t i = tx; i < static_cast(TopK); + i += kThreadsPerBlock) { + row_output[i] = (i < seq_len) ? static_cast(i) : -1; + } + } else if (seq_len <= static_cast(HIST2048_THRESHOLD)) { + histogram_2048_topk(row_input, row_output, seq_len); + } else { + histogram_256_topk(row_input, row_output, 0, seq_len); + } + } + continue; + } + + const uint32_t my_chunk_start = cta_in_group * chunk_size; + radix_topk( + row_input, row_output, seq_len, my_chunk_start, chunk_size, + local_histogram, suffix_sum, shared_scalars, shared_ordered, state, + cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx); + radix_iter++; + } +} + +} // namespace persistent + +// ============================================================================ +// ============================================================================ +// Optimized FilteredTopK — single CTA per row for bs > 32. +// Kept with persistent_topk so the portable fallback owns the non-cluster path. +// ============================================================================ +namespace filtered_topk { + +namespace hist4096 = topk_histogram_4096; + +// ============================================================================ +// FilteredTopK — single CTA per row for bs > 32 +// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 +// ============================================================================ + +#define FLASHINFER_CUDA_CALL(func, ...) \ + { \ + cudaError_t e = (func); \ + if (e != cudaSuccess) { \ + return e; \ + } \ + } + +#define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__ + +template +struct vec_t { + T data[N]; + + FLASHINFER_INLINE T& operator[](size_t i) { return data[i]; } + FLASHINFER_INLINE const T& operator[](size_t i) const { return data[i]; } + + FLASHINFER_INLINE void cast_load(const T* ptr) { +#pragma unroll + for (size_t i = 0; i < N; ++i) { + data[i] = ptr[i]; + } + } +}; +#undef FLASHINFER_INLINE + +// FilteredTopK traits for different data types +template +struct FilteredTopKTraits; + +// Specialization for float (32-bit): coarse histogram uses FP16 high 8 bits, 4 +// refinement rounds +template <> +struct FilteredTopKTraits { + using OrderedType = uint32_t; + static constexpr int NUM_REFINE_ROUNDS = 4; + static constexpr int FIRST_REFINE_SHIFT = 24; + + __device__ __forceinline__ static uint8_t ToCoarseKey(float x) { + // Convert to FP16 representation and extract high 8 bits + __half h = __float2half_rn(x); + uint16_t bits = __half_as_ushort(h); + uint16_t key = (bits & 0x8000) ? static_cast(~bits) + : static_cast(bits | 0x8000); + return static_cast(key >> 8); + } + + __device__ __forceinline__ static OrderedType ToOrdered(float x) { + uint32_t bits = __float_as_uint(x); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); + } +}; + +constexpr uint32_t FILTERED_TOPK_BLOCK_THREADS = 1024; +constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE = + 16 * 1024; // 16K indices per buffer +constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC = + sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB + +/*! + * \brief Filtered Top-K kernel for ragged sequences. + * + * \tparam DType Data type (float, half, nv_bfloat16) + * \tparam IdType Index type (int32_t) + * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) + */ +template +__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) + FilteredTopKUnifiedKernel(const DType* __restrict__ input, + IdType* __restrict__ output, + const IdType* __restrict__ lengths, + uint32_t num_rows, uint32_t top_k, + uint32_t max_len) { + constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS; + constexpr int RADIX = 256; + constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE; + + const uint32_t bid = blockIdx.x; + const int tx = threadIdx.x; + + if (bid >= num_rows) return; + + const int length = + (lengths != nullptr) ? lengths[bid] : static_cast(max_len); + const DType* score = input + bid * max_len; + IdType* dst = output + bid * top_k; + + // Trivial case: length <= top_k + if (length <= static_cast(top_k)) { + for (int i = tx; i < static_cast(top_k); i += BLOCK_SIZE) { + dst[i] = (i < length) ? static_cast(i) : static_cast(-1); + } + return; + } + + // Short path + if (length <= 32768) { + extern __shared__ uint8_t _smem_reg[]; + if constexpr (UsePredicatedShortLoads) { + hist4096::histogram_4096_topk_predicated(score, dst, length, + _smem_reg); + } else { + hist4096::histogram_4096_topk(score, dst, length, + _smem_reg); + } + return; + } + + // Static shared memory + alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128]; + alignas(128) __shared__ int s_counter; + alignas(128) __shared__ int s_threshold_bin_id; + alignas(128) __shared__ int s_num_input[2]; + alignas(128) __shared__ int s_indices[MAX_K]; + + auto& s_histogram = s_histogram_buf[0]; + + // Dynamic shared memory for input double buffer + extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE]; + + using Traits = FilteredTopKTraits; + int topk = top_k; + + // Stage 1: 8-bit coarse histogram with vectorized loads + if (tx < RADIX + 1) s_histogram[tx] = 0; + __syncthreads(); + + vec_t score_vec; + + const int aligned_length = (length / VEC_SIZE) * VEC_SIZE; +#pragma unroll 2 + for (int base = tx * VEC_SIZE; base < aligned_length; + base += BLOCK_SIZE * VEC_SIZE) { + score_vec.cast_load(&score[base]); +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + const auto bin = Traits::ToCoarseKey(score_vec[j]); + atomicAdd(&s_histogram[bin], 1); + } + } + // Handle tail + for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { + const auto bin = Traits::ToCoarseKey(score[i]); + atomicAdd(&s_histogram[bin], 1); + } + __syncthreads(); + + // Suffix sum + const auto run_cumsum = [&]() { +#pragma unroll 8 + for (int i = 0; i < 8; ++i) { + if (tx < RADIX) { + const auto j = 1 << i; + const auto k = i & 1; + auto value = s_histogram_buf[k][tx]; + if (tx < RADIX - j) { + value += s_histogram_buf[k][tx + j]; + } + s_histogram_buf[k ^ 1][tx] = value; + } + __syncthreads(); + } + }; + + run_cumsum(); + if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { + s_threshold_bin_id = tx; + s_num_input[0] = 0; + s_counter = 0; + } + __syncthreads(); + + const auto threshold_bin = s_threshold_bin_id; + topk -= s_histogram[threshold_bin + 1]; + + constexpr int NUM_ROUNDS = Traits::NUM_REFINE_ROUNDS; + constexpr int FIRST_SHIFT = Traits::FIRST_REFINE_SHIFT; + + if (topk == 0) { + // Collect indices where bin > threshold +#pragma unroll 2 + for (int base = tx * VEC_SIZE; base < aligned_length; + base += BLOCK_SIZE * VEC_SIZE) { + score_vec.cast_load(&score[base]); +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + const auto bin = static_cast(Traits::ToCoarseKey(score_vec[j])); + if (bin > threshold_bin) { + const auto pos = atomicAdd(&s_counter, 1); + s_indices[pos] = base + j; + } + } + } + // Handle tail + for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { + const auto bin = static_cast(Traits::ToCoarseKey(score[i])); + if (bin > threshold_bin) { + const auto pos = atomicAdd(&s_counter, 1); + s_indices[pos] = i; + } + } + __syncthreads(); + } else { + __syncthreads(); + if (tx < RADIX + 1) s_histogram[tx] = 0; + __syncthreads(); + + // Filter + histogram for refinement + auto filter_and_add_to_histogram = [&](auto raw_input, int index) { + const auto bin = static_cast(Traits::ToCoarseKey(raw_input)); + if (bin > threshold_bin) { + const auto pos = atomicAdd(&s_counter, 1); + s_indices[pos] = index; + } else if (bin == threshold_bin) { + const auto pos = atomicAdd(&s_num_input[0], 1); + if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) { + s_input_idx[0][pos] = index; + const auto ordered = Traits::ToOrdered(raw_input); + const auto sub_bin = (ordered >> FIRST_SHIFT) & 0xFF; + atomicAdd(&s_histogram[sub_bin], 1); + } + } + }; +#pragma unroll 2 + for (int base = tx * VEC_SIZE; base < aligned_length; + base += BLOCK_SIZE * VEC_SIZE) { + score_vec.cast_load(&score[base]); +#pragma unroll + for (int j = 0; j < VEC_SIZE; ++j) { + filter_and_add_to_histogram(score_vec[j], base + j); + } + } + // Handle tail + for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { + filter_and_add_to_histogram(score[i], i); + } + __syncthreads(); + + // Stage 2: refine with 8bit radix passes +#pragma unroll + for (int round = 0; round < NUM_ROUNDS; ++round) { + __shared__ int s_last_remain; + const auto r_idx = round % 2; + + const auto _raw_num_input = s_num_input[r_idx]; + const auto num_input = + (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE; + + run_cumsum(); + if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { + s_threshold_bin_id = tx; + s_num_input[r_idx ^ 1] = 0; + s_last_remain = topk - s_histogram[tx + 1]; + } + __syncthreads(); + + const auto threshold = s_threshold_bin_id; + topk -= s_histogram[threshold + 1]; + + const int offset = FIRST_SHIFT - round * 8; + const bool is_last_round = (round == NUM_ROUNDS - 1); + + if (topk == 0) { + for (int i = tx; i < num_input; i += BLOCK_SIZE) { + const auto idx = s_input_idx[r_idx][i]; + const auto bin = (Traits::ToOrdered(score[idx]) >> offset) & 0xFF; + if (static_cast(bin) > threshold) { + const auto pos = atomicAdd(&s_counter, 1); + s_indices[pos] = idx; + } + } + __syncthreads(); + break; + } else { + __syncthreads(); + if (tx < RADIX + 1) s_histogram[tx] = 0; + __syncthreads(); + for (int i = tx; i < num_input; i += BLOCK_SIZE) { + const auto idx = s_input_idx[r_idx][i]; + const auto raw_input = score[idx]; + const auto bin = (Traits::ToOrdered(raw_input) >> offset) & 0xFF; + if (static_cast(bin) > threshold) { + const auto pos = atomicAdd(&s_counter, 1); + s_indices[pos] = idx; + } else if (static_cast(bin) == threshold) { + if (is_last_round) { + const auto pos = atomicAdd(&s_last_remain, -1); + if (pos > 0) { + s_indices[top_k - pos] = idx; + } + } else { + const auto pos = atomicAdd(&s_num_input[r_idx ^ 1], 1); + if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) { + s_input_idx[r_idx ^ 1][pos] = idx; + const auto bin32 = Traits::ToOrdered(raw_input); + const auto sub_bin = (bin32 >> (offset - 8)) & 0xFF; + atomicAdd(&s_histogram[sub_bin], 1); + } + } + } + } + __syncthreads(); + } + } + } + + // Output phase - mode-specific +#pragma unroll 2 + for (int base = tx; base < static_cast(top_k); base += BLOCK_SIZE) { + const int idx = s_indices[base]; + dst[base] = static_cast(idx); + } +} + +// Helper to compute GCD for VEC_SIZE selection +constexpr uint32_t gcd(uint32_t a, uint32_t b) { + while (b != 0) { + uint32_t t = b; + b = a % b; + a = t; + } + return a; +} + +// Compute optimal VEC_SIZE based on max_len and dtype +// Returns 1, 2, 4, or 8 +template +constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) { + constexpr int MAX_VEC = 16 / sizeof(DType); // 4 for float32, 8 for fp16/bf16 + // Use GCD to find largest power-of-2 divisor + const uint32_t g = gcd(max_len, static_cast(MAX_VEC)); + return static_cast(g); +} + +template +cudaError_t FilteredTopKRaggedTransform(const DType* input, + IdType* output_indices, + const IdType* lengths, + uint32_t num_rows, uint32_t top_k_val, + uint32_t max_len, + cudaStream_t stream = 0) { + constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC; + constexpr int MAX_VEC = 16 / sizeof(DType); + + dim3 grid(num_rows); + dim3 block(FILTERED_TOPK_BLOCK_THREADS); + void* args[] = {&input, &output_indices, &lengths, + &num_rows, &top_k_val, &max_len}; + + const int vec_size = ComputeFilteredTopKVecSize(max_len); + +#define DISPATCH_VEC_SIZE(VS) \ + if (vec_size == VS) { \ + auto kernel = \ + FilteredTopKUnifiedKernel; \ + FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \ + FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \ + smem_size, stream)); \ + return cudaSuccess; \ + } + + DISPATCH_VEC_SIZE(1) + DISPATCH_VEC_SIZE(2) + DISPATCH_VEC_SIZE(4) + if constexpr (MAX_VEC >= 8) { + DISPATCH_VEC_SIZE(8) + } +#undef DISPATCH_VEC_SIZE + + return cudaSuccess; +} + +} // namespace filtered_topk + +template +cudaError_t FilteredTopKRaggedTransform(const DType* input, + IdType* output_indices, + const IdType* lengths, + uint32_t num_rows, uint32_t top_k_val, + uint32_t max_len, + cudaStream_t stream = 0) { + return filtered_topk::FilteredTopKRaggedTransform( + input, output_indices, lengths, num_rows, top_k_val, max_len, stream); +} + +} // namespace vllm + +#endif // PERSISTENT_TOPK_CUH_ diff --git a/prime_kernels/indexed_attention/csrc/topk.cu b/prime_kernels/indexed_attention/csrc/topk.cu new file mode 100644 index 0000000..eadcc11 --- /dev/null +++ b/prime_kernels/indexed_attention/csrc/topk.cu @@ -0,0 +1,287 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// Adapted from vLLM's libtorch-stable persistent top-k interface for the +// prime_indexed_attention dispatcher namespace. + +#include +#include +#include +#include +#include +#include +#include + +#include "persistent_topk.cuh" + +namespace { + +template +void launch_persistent_topk(const torch::Tensor& logits, + const torch::Tensor& lengths, + torch::Tensor& output, + torch::Tensor& workspace, + int64_t max_seq_len) { + namespace P = vllm::persistent; + + const at::cuda::OptionalCUDAGuard device_guard{device_of(logits)}; + const int64_t num_rows = logits.size(0); + const int64_t stride = logits.stride(0); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + static int num_sms = 0; + static int max_smem_per_block = 0; + if (num_sms == 0) { + const cudaDeviceProp* device_prop = at::cuda::getDeviceProperties(logits.get_device()); + num_sms = device_prop->multiProcessorCount; + max_smem_per_block = device_prop->sharedMemPerBlockOptin; + } + + if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { + cudaError_t status = + vllm::FilteredTopKRaggedTransform( + logits.const_data_ptr(), output.data_ptr(), + lengths.const_data_ptr(), static_cast(num_rows), + static_cast(TopK), static_cast(stride), stream); + TORCH_CHECK(status == cudaSuccess, + "FilteredTopK failed: ", cudaGetErrorString(status)); + } else { + TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor"); + TORCH_CHECK( + workspace.scalar_type() == torch::kUInt8, + "workspace must be uint8"); + + int effective_max_smem; + if (num_rows <= 4) { + effective_max_smem = + std::min(max_smem_per_block, static_cast(P::kSmemMedium)); + } else if (num_rows <= 8) { + constexpr int kSmemCapMedium = 48 * 1024; + effective_max_smem = std::min(max_smem_per_block, kSmemCapMedium); + } else { + effective_max_smem = max_smem_per_block; + } + + size_t available_for_ordered = + static_cast(effective_max_smem) - P::kFixedSmemLarge; + uint32_t max_chunk_elements = + static_cast(available_for_ordered / sizeof(uint32_t)); + + uint32_t vec_size = 1; + if (stride % 4 == 0) + vec_size = 4; + else if (stride % 2 == 0) + vec_size = 2; + + max_chunk_elements = (max_chunk_elements / vec_size) * vec_size; + uint32_t min_chunk = vec_size * P::kThreadsPerBlock; + if (max_chunk_elements < min_chunk) max_chunk_elements = min_chunk; + + uint32_t ctas_per_group = + (static_cast(stride) + max_chunk_elements - 1) / + max_chunk_elements; + uint32_t chunk_size = + (static_cast(stride) + ctas_per_group - 1) / ctas_per_group; + chunk_size = ((chunk_size + vec_size - 1) / vec_size) * vec_size; + if (chunk_size > max_chunk_elements) chunk_size = max_chunk_elements; + + size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t); + if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium; + + // Query occupancy for the instantiation that will actually launch; + // overestimating it deadlocks the cooperative barrier. + int occupancy = 1; + cudaError_t occ_err = cudaSuccess; + if (vec_size == 4) { + occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &occupancy, P::persistent_topk_kernel, P::kThreadsPerBlock, + smem_size); + } else if (vec_size == 2) { + occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &occupancy, P::persistent_topk_kernel, P::kThreadsPerBlock, + smem_size); + } else { + occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &occupancy, P::persistent_topk_kernel, P::kThreadsPerBlock, + smem_size); + } + TORCH_CHECK(occ_err == cudaSuccess, + "persistent_topk occupancy query failed: ", + cudaGetErrorString(occ_err)); + if (occupancy < 1) occupancy = 1; + + // The cooperative spin-wait barrier only runs when at least one row hits + // the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs + // early-exit, so oversubscription can't deadlock and headroom is wasted. + const bool needs_cooperative = + static_cast(max_seq_len) > P::RADIX_THRESHOLD; + + const uint32_t hw_resident_cap = + static_cast(num_sms) * static_cast(occupancy); + uint32_t max_resident_ctas = hw_resident_cap; + if (needs_cooperative) { + // Reserve one CTA per SM when occupancy allows; fall back to a single + // CTA when occupancy == 1 (the most deadlock-prone case — any straggler + // kernel that takes the only slot on one SM hangs the barrier). Never + // drop below one full group's worth. + uint32_t headroom = (occupancy > 1) ? static_cast(num_sms) : 1u; + if (max_resident_ctas >= headroom + ctas_per_group) { + max_resident_ctas -= headroom; + } + } + uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group, + static_cast(num_rows)); + if (num_groups == 0) num_groups = 1; + uint32_t total_ctas = num_groups * ctas_per_group; + + // If the cooperative launch wouldn't fit, fall back to FilteredTopK + // instead of deadlocking. Only relevant when needs_cooperative. + if (needs_cooperative && total_ctas > hw_resident_cap) { + TORCH_CHECK( + max_smem_per_block >= 128 * 1024, + "persistent_topk would oversubscribe and the FilteredTopK " + "fallback requires >=128KB smem per block (have ", + max_smem_per_block, "). total_ctas=", total_ctas, + " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK, + ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group, + ", smem=", smem_size, ")."); + cudaError_t status = + vllm::FilteredTopKRaggedTransform( + logits.const_data_ptr(), + output.data_ptr(), + lengths.const_data_ptr(), + static_cast(num_rows), static_cast(TopK), + static_cast(stride), stream); + TORCH_CHECK(status == cudaSuccess, "FilteredTopK fallback failed: ", + cudaGetErrorString(status)); + return; + } + + size_t state_bytes = num_groups * sizeof(P::RadixRowState); + TORCH_CHECK(workspace.size(0) >= static_cast(state_bytes), + "workspace too small, need ", state_bytes, " bytes"); + + // Zero the per-group RadixRowState region before launch. + // + // Issued UNCONDITIONALLY so the memset is captured as its own node in + // the cudagraph (a separate cudaMemsetAsync node, sequenced before the + // persistent_topk_kernel launch on the same stream). The previous + // host-side guard `if (needs_cooperative)` was evaluated at capture time; + // when capture-time max_seq_len <= RADIX_THRESHOLD (always true under + // FULL_DECODE_ONLY with max_model_len < 32 K) the memset would NOT be + // captured, leaving the workspace state to accumulate across replays. + // That's a latent correctness bug if the runtime data ever takes the + // radix path, and removes one variable while debugging hangs in the + // decode/medium paths. + // + // Cost is sub-microsecond: state_bytes = num_groups * sizeof(RadixRowState) + // is ~3 KB per group, ~100 KB for the largest grids on this hardware. + // + // Why the memset is required (regardless of which path the kernel takes): + // 1. arrival_counter accumulates within a launch and is never reset, + // so a prior call leaves it at a large positive value. Without this + // reset, the very first wait_ge in the next call sees counter >> + // target and returns instantly, breaking the barrier. + // 2. The previous in-kernel init only ran in CTA-0 with intra-CTA + // __syncthreads(), so it had no happens-before edge to CTA-1+'s + // first red_release. cudaMemsetAsync is stream-ordered: the zero + // is globally visible before any CTA runs. + { + cudaError_t mz_err = cudaMemsetAsync( + workspace.data_ptr(), 0, state_bytes, stream); + TORCH_CHECK(mz_err == cudaSuccess, + "row_states memset failed: ", cudaGetErrorString(mz_err)); + } + + P::PersistentTopKParams params; + params.input = logits.const_data_ptr(); + params.output = output.data_ptr(); + params.lengths = lengths.const_data_ptr(); + params.num_rows = static_cast(num_rows); + params.stride = static_cast(stride); + params.top_k = static_cast(TopK); + params.chunk_size = chunk_size; + params.row_states = reinterpret_cast( + workspace.data_ptr()); + params.ctas_per_group = ctas_per_group; + params.max_seq_len = static_cast(max_seq_len); + + #define LAUNCH_PERSISTENT(TOPK_VAL, VS) \ + do { \ + auto kernel = &P::persistent_topk_kernel; \ + cudaError_t err = cudaFuncSetAttribute( \ + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \ + TORCH_CHECK(err == cudaSuccess, \ + "Failed to set smem: ", cudaGetErrorString(err)); \ + kernel<<>>(params); \ + } while (0) + + if (vec_size == 4) { + LAUNCH_PERSISTENT(TopK, 4); + } else if (vec_size == 2) { + LAUNCH_PERSISTENT(TopK, 2); + } else { + LAUNCH_PERSISTENT(TopK, 1); + } + #undef LAUNCH_PERSISTENT + } + + cudaError_t err = cudaGetLastError(); + TORCH_CHECK(err == cudaSuccess, + "persistent_topk failed: ", cudaGetErrorString(err)); +} + +} // anonymous namespace + +void persistent_topk(const torch::Tensor& logits, + const torch::Tensor& lengths, + torch::Tensor& output, + torch::Tensor& workspace, + int64_t k, + int64_t max_seq_len) { + TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor"); + TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor"); + TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor"); + TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor"); + TORCH_CHECK(logits.scalar_type() == torch::kFloat32, + "logits must be float32"); + TORCH_CHECK(lengths.scalar_type() == torch::kInt32, + "lengths must be int32"); + TORCH_CHECK(output.scalar_type() == torch::kInt32, + "output must be int32"); + TORCH_CHECK(workspace.scalar_type() == torch::kUInt8, + "workspace must be uint8"); + TORCH_CHECK(logits.dim() == 2, "logits must be 2D"); + TORCH_CHECK(lengths.dim() == 1, "lengths must be 1D"); + TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous"); + TORCH_CHECK(output.dim() == 2, "output must be 2D"); + TORCH_CHECK(logits.is_contiguous(), "logits must be contiguous"); + TORCH_CHECK(output.is_contiguous(), "output must be contiguous"); + + const int64_t num_rows = logits.size(0); + TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch"); + TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k, + "output size mismatch"); + TORCH_CHECK(k == 512 || k == 1024 || k == 2048, + "persistent_topk supports k=512, k=1024, or k=2048, got ", k); + + if (k == 512) { + launch_persistent_topk<512>(logits, lengths, output, workspace, + max_seq_len); + } else if (k == 1024) { + launch_persistent_topk<1024>(logits, lengths, output, workspace, + max_seq_len); + } else { + launch_persistent_topk<2048>(logits, lengths, output, workspace, + max_seq_len); + } +} + +TORCH_LIBRARY_FRAGMENT(prime_indexed_attention, m) { + m.def( + "persistent_topk(Tensor logits, Tensor lengths, " + "Tensor(a!) output, Tensor(b!) workspace, int k, int max_seq_len) -> ()"); + m.impl("persistent_topk", torch::kCUDA, &persistent_topk); +} + +PYBIND11_MODULE(_C, m) {} diff --git a/prime_kernels/indexed_attention/csrc/topk_histogram_4096.cuh b/prime_kernels/indexed_attention/csrc/topk_histogram_4096.cuh new file mode 100644 index 0000000..5f9f823 --- /dev/null +++ b/prime_kernels/indexed_attention/csrc/topk_histogram_4096.cuh @@ -0,0 +1,563 @@ +/* + * Shared 4096-bin single-CTA TopK helpers. + */ + +#ifndef TOPK_HISTOGRAM_4096_CUH_ +#define TOPK_HISTOGRAM_4096_CUH_ + +#include +#include +#include + +namespace vllm { +namespace topk_histogram_4096 { + +constexpr uint32_t kBlockSize = 1024; +constexpr uint32_t RADIX = 256; +constexpr uint32_t kMaxTies = 1024; +static_assert(kMaxTies <= kBlockSize, + "tie_handle requires kMaxTies <= kBlockSize"); +constexpr uint32_t kWarpSize = 32; +constexpr uint32_t kNumWarps = kBlockSize / kWarpSize; + +// Register path +constexpr uint32_t kHist4096VecsPerThread = 4; +constexpr uint32_t kHist4096MaxLen = + kHist4096VecsPerThread * 4 * kBlockSize; // 16384 + +struct alignas(16) MatchBin { + uint32_t bin, above_count, equal_count; +}; +struct alignas(8) Tie { + uint32_t idx; + float score; +}; + +__device__ __forceinline__ void load_float4_predicated(const float* ptr, + int base, int seq_len, + float& v0, float& v1, + float& v2, float& v3) { + uint32_t r0, r1, r2, r3; + const int p0 = (base < seq_len); + const int p1 = (base + 1 < seq_len); + const int p2 = (base + 2 < seq_len); + const int p3 = (base + 3 < seq_len); + asm volatile( + "{\n" + " .reg .pred pr0, pr1, pr2, pr3;\n" + " setp.ne.u32 pr0, %4, 0;\n" + " setp.ne.u32 pr1, %5, 0;\n" + " setp.ne.u32 pr2, %6, 0;\n" + " setp.ne.u32 pr3, %7, 0;\n" + " mov.u32 %0, 0xFF800000;\n" + " mov.u32 %1, 0xFF800000;\n" + " mov.u32 %2, 0xFF800000;\n" + " mov.u32 %3, 0xFF800000;\n" + " @pr0 ld.global.cg.u32 %0, [%8];\n" + " @pr1 ld.global.cg.u32 %1, [%8+4];\n" + " @pr2 ld.global.cg.u32 %2, [%8+8];\n" + " @pr3 ld.global.cg.u32 %3, [%8+12];\n" + "}\n" + : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) + : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr)); + v0 = __uint_as_float(r0); + v1 = __uint_as_float(r1); + v2 = __uint_as_float(r2); + v3 = __uint_as_float(r3); +} + +// converts the float32 score to a 32-bit ordered unsigned integer — the full +// precision key for radix sorting +__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { + uint32_t bits = __float_as_uint(x); + return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); +} + +// Converts each score to a 12-bit bin (FP16 sign-magnitude -> top 12 bits -> +// bin 0-4095) +template +__device__ __forceinline__ uint32_t extract_coarse_bin_N(float x) { + __half h = __float2half_rn(x); + uint16_t bits = __half_as_ushort(h); + uint16_t key = (bits & 0x8000) ? static_cast(~bits) + : static_cast(bits | 0x8000); + return key >> (16 - kBits); +} + +// running sum within each warp — thread 0 gets its own value, thread 1 gets +// thread 0 + thread 1, thread 2 gets threads 0+1+2, etc. +__device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane, + uint32_t v) { +#pragma unroll + for (uint32_t o = 1; o < 32; o *= 2) { + uint32_t n = __shfl_up_sync(0xFFFFFFFF, v, o); + if (lane >= o) v += n; + } + return v; +} + +// Returns the sum of a value across all 32 threads in the warp, and every +// thread gets the same result. SM80+ uses redux.sync.add.u32, a single PTX +// instruction for hardware warp-wide reduction. Older targets use the +// __shfl_xor_sync butterfly tree, like warp::reduce_sum() (5 shuffles for 32 +// lanes). +__device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + uint32_t r; + asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v)); + return r; +#else + #pragma unroll + for (uint32_t mask = kWarpSize >> 1; mask > 0; mask >>= 1) { + v += __shfl_xor_sync(0xFFFFFFFF, v, mask); + } + return v; +#endif +} + +// ============================================================================ +// Tie refinement (single CTA): 4-round radix-256 topK on the full FP32 ordered +// key Each round narrows by 8 bits until ties are fully resolved +// ============================================================================ + +template +__device__ void tie_handle(const Tie* ties, uint32_t num_ties, + uint32_t num_above, int32_t* output, void* _smem) { + struct TS { + alignas(128) uint32_t counter; + alignas(128) MatchBin match; + uint32_t histogram[RADIX]; + uint32_t warp_sum[kNumWarps]; + }; + auto* s = static_cast(_smem); + const auto tx = threadIdx.x; + const auto li = tx % kWarpSize, wi = tx / kWarpSize; + + // Each thread loads one tie element. + const bool has = tx < num_ties; + const auto tie = has ? ties[tx] : Tie{0, 0.0f}; + const uint32_t key = convert_to_uint32_v2(tie.score); + + bool active = has; // tracks whether this thread's tie is still a candidate. + uint32_t remain = + TopK - num_above; // decreases each round as ties are resolved. + uint32_t wpos = TopK; // wpos will hold the final output position. + s->counter = 0; + __syncthreads(); + + // The 4-round radix loop - each round narrows by 8 bits until ties are fully + // resolved +#pragma unroll + for (int r = 0; r < 4; r++) { + uint32_t sh = 24 - r * 8; // round 0: bits 31-24, round 1: 23-16, etc. + uint32_t bin = (key >> sh) & 0xFF; // this tie's 8-bit bin for this round + + // Step 1: Build 256-bin histogram. + if (tx < RADIX) s->histogram[tx] = 0; + __syncthreads(); + if (active) atomicAdd(&s->histogram[bin], 1); + __syncthreads(); + + // Step 2: Prefix scan to find threshold + uint32_t hv = 0, wi2 = 0; + if (tx < RADIX) { + hv = s->histogram[tx]; + wi2 = warp_inclusive_sum(li, hv); + if (li == kWarpSize - 1) s->warp_sum[wi] = wi2; + } + __syncthreads(); + + if (tx < RADIX) { + auto tmp = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; + auto tot = warp_reduce_sum_full(tmp); + auto inter = warp_reduce_sum_full(li < wi ? tmp : 0); + auto above = tot - (inter + wi2); + if (above < remain && above + hv >= remain) { + s->match = {tx, above, remain - above}; + } + } + __syncthreads(); + + // Step 3: Scatter + auto [thr, na, _] = s->match; // threshold bin, num above, unused + if (active) { + if (bin > thr) { + wpos = num_above + + atomicAdd(&s->counter, 1); // above -> place in output directly + active = false; + } else if (bin < thr) + active = false; // below -> discard + else if (r == 3) + wpos = TopK - atomicAdd(&s->match.equal_count, + -1u); // last round: place remaining + } + remain -= na; + if (!remain) break; // all ties resolved early + } + // Final write + if (wpos < TopK) output[wpos] = tie.idx; +} + +// Extended tie_handle for TopK > kBlockSize (e.g. TopK=2048). +// tie_handle assumes 1 tie per thread (max 1024). +// This version handles 2 ties per thread via kPerThread=2 +template +__device__ void tie_handle_large(const Tie* ties, uint32_t num_ties, + uint32_t num_above, int32_t* output, + void* _smem) { + static_assert(TopK > kBlockSize); + struct TS { + alignas(128) uint32_t counter; + alignas(128) MatchBin match; + uint32_t histogram[RADIX]; + uint32_t warp_sum[kNumWarps]; + }; + auto* s = static_cast(_smem); + const auto tx = threadIdx.x; + const auto li = tx % kWarpSize; + const auto wi = tx / kWarpSize; + + constexpr uint32_t kPerThread = (TopK + kBlockSize - 1) / kBlockSize; + Tie my_ties[kPerThread]; + uint32_t keys[kPerThread]; + bool active[kPerThread]; + + for (uint32_t e = 0; e < kPerThread; e++) { + uint32_t idx = e * kBlockSize + tx; + if (idx < num_ties) { + my_ties[e] = ties[idx]; + keys[e] = convert_to_uint32_v2(ties[idx].score); + active[e] = true; + } else { + my_ties[e] = {0, 0.0f}; + keys[e] = 0; + active[e] = false; + } + } + + uint32_t remain = TopK - num_above; + s->counter = 0; + __syncthreads(); + + for (int r = 0; r < 4; r++) { + uint32_t sh = 24 - r * 8; + if (tx < RADIX) { + s->histogram[tx] = 0; + } + __syncthreads(); + + for (uint32_t e = 0; e < kPerThread; e++) { + if (active[e]) { + atomicAdd(&s->histogram[(keys[e] >> sh) & 0xFF], 1); + } + } + __syncthreads(); + + uint32_t hv = 0; + if (tx < RADIX) { + hv = s->histogram[tx]; + auto wi2 = warp_inclusive_sum(li, hv); + if (li == kWarpSize - 1) { + s->warp_sum[wi] = wi2; + } + } + __syncthreads(); + if (tx < RADIX) { + auto tmp2 = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; + auto total = warp_reduce_sum_full(tmp2); + auto inter = warp_reduce_sum_full(li < wi ? tmp2 : 0); + auto wi2 = warp_inclusive_sum(li, hv); + auto above = total - (inter + wi2); + if (above < remain && above + hv >= remain) { + s->match = { + .bin = tx, .above_count = above, .equal_count = remain - above}; + } + } + __syncthreads(); + + auto thr = s->match.bin; + auto na = s->match.above_count; + + for (uint32_t e = 0; e < kPerThread; e++) { + if (!active[e]) { + continue; + } + uint32_t bin = (keys[e] >> sh) & 0xFF; + if (bin > thr) { + uint32_t wpos = num_above + atomicAdd(&s->counter, 1); + if (wpos < TopK) { + output[wpos] = my_ties[e].idx; + } + active[e] = false; + } else if (bin < thr) { + active[e] = false; + } else if (r == 3) { + uint32_t wpos = TopK - atomicAdd(&s->match.equal_count, -1u); + if (wpos < TopK) { + output[wpos] = my_ties[e].idx; + } + } + } + + num_above += na; + remain -= na; + __syncthreads(); + s->counter = 0; + __syncthreads(); + } +} + +// ============================================================================ +// Register-based single-CTA fast path for seq_len <= 16384 +// 4 float4 per thread × 1024 threads = 16384 elements max +// Uses 4096-bin (12-bit) histogram for better precision +// ============================================================================ + +template +struct Histogram4096Smem { + static constexpr uint32_t HIST_BINS = 1 << HIST_BITS; + static constexpr uint32_t TIE_CAPACITY = TopK > kMaxTies ? TopK : kMaxTies; + alignas(128) uint32_t counter_gt; + alignas(128) uint32_t counter_eq; + MatchBin match; + uint32_t warp_sum[kNumWarps]; + union { + uint32_t histogram[HIST_BINS]; + Tie tie_buffer[TIE_CAPACITY]; + }; +}; + +template +__device__ void histogram_4096_topk(const float* __restrict__ scores, + int32_t* __restrict__ output, + uint32_t length, void* _smem) { + constexpr uint32_t HIST_BINS = 1 << HIST_BITS; + constexpr uint32_t ITEMS_PER_THREAD = HIST_BINS / kBlockSize; + static_assert(HIST_BINS >= kBlockSize, + "HIST_BITS must give >= kBlockSize bins"); + + using Smem = Histogram4096Smem; + auto* smem = static_cast(_smem); + const auto tx = threadIdx.x; + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + + // Phase 1: Load all data into RF + build histogram + float4 + vecs[VECS_PER_THREAD]; // 4 vectors x 4 floats = 16 elements per thread + if constexpr (ITEMS_PER_THREAD >= 4) { + // Zero the histogram (SMEM writes) + for (uint32_t i = 0; i < ITEMS_PER_THREAD / 4; i++) + reinterpret_cast( + smem->histogram)[tx * (ITEMS_PER_THREAD / 4) + i] = + make_uint4(0, 0, 0, 0); + } else { + if (tx < HIST_BINS) smem->histogram[tx] = 0; + } + if (tx == 0) { + smem->counter_gt = 0; + smem->counter_eq = 0; + } + if constexpr (UsePredicatedLoads) { + const bool row_aligned = (reinterpret_cast(scores) & 0xFu) == 0; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { + const uint32_t base = (tx + v * kBlockSize) * 4; + if (base < length) { + if (row_aligned && base + 3 < length) { + vecs[v] = *reinterpret_cast(scores + base); + } else { + load_float4_predicated(scores + base, static_cast(base), + static_cast(length), vecs[v].x, vecs[v].y, + vecs[v].z, vecs[v].w); + } + } + } + } else { +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { + const uint32_t base = (tx + v * kBlockSize) * 4; + if (base < length) { + vecs[v] = *reinterpret_cast(scores + base); + } + } + } + __syncthreads(); + + // Build histogram from RF via atomic adds into the shared histogram + bool done = false; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + atomicAdd(&smem->histogram[extract_coarse_bin_N(elems[e])], + 1); + } + } + } + __syncthreads(); + + // Phase 2: Prefix scan to find threshold bin + // Multi-element scan (4096 bins: 4 per thread) + uint32_t orig[ITEMS_PER_THREAD]; + uint32_t local_sum = 0; + + // Step 1: Each thread sums its 4 bins +#pragma unroll + for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { + orig[i] = smem->histogram[tx * ITEMS_PER_THREAD + i]; + local_sum += orig[i]; + } + + // Step 2: Warp-level inclusive prefix sum on local_sum + const auto warp_inc = warp_inclusive_sum(lane_id, local_sum); + if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc; + __syncthreads(); + + // Step 3: Inter-warp prefix across warp sums. + const auto tmp = smem->warp_sum[lane_id]; + uint32_t prefix = warp_reduce_sum_full( + lane_id < warp_id ? tmp : 0); // sum of all prior warps + prefix += + warp_inc - local_sum; // exclusive prefix within this thread's position + + // Step 4: Find threshold - scan 4 bins, accumulate prefix +#pragma unroll + for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { + prefix += orig[i]; + const auto above = length - prefix; // elements in bins ABOVE this one + if (above < TopK && above + orig[i] >= TopK) { + smem->match = {.bin = tx * ITEMS_PER_THREAD + i, + .above_count = above, + .equal_count = orig[i]}; + } + } + + __syncthreads(); + + // Phase 3: Scatter from registers + const auto [thr_bin, num_above, num_equal] = smem->match; + const bool need_tie = (num_equal + num_above > TopK); + + done = false; +#pragma unroll + for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { + const float* elems = reinterpret_cast(&vecs[v]); +#pragma unroll + for (uint32_t e = 0; e < 4 && !done; e++) { + const uint32_t idx = (tx + v * kBlockSize) * 4 + e; + if (idx >= length) { + done = true; + } else { + const uint32_t bin = extract_coarse_bin_N(elems[e]); + if (bin > thr_bin) { + output[atomicAdd(&smem->counter_gt, 1)] = + idx; // above -> output directly + } else if (bin == thr_bin) { + const auto pos = atomicAdd(&smem->counter_eq, 1); + if (!need_tie) { + if (pos + num_above < TopK) { + output[pos + num_above] = idx; // all fit + } + } else { + if (pos < TopK) { + smem->tie_buffer[pos] = {idx, elems[e]}; // store for refirement + } + } + } + // else: bin < thr_bin - discard (not in top-k) + } + } + } + + // Phase 4: Tie-breaking + if (!need_tie) return; + __syncthreads(); + + // Fast warp-ballot tie-breaking for small tie counts + const uint32_t num_ties = min(num_equal, static_cast(TopK)); + const uint32_t topk_remain = + TopK - num_above; // pick exactly remaining elements to fill topK + + auto is_greater = [](const Tie& a, const Tie& b) { + return (a.score > b.score) || (a.score == b.score && a.idx < b.idx); + }; + + if (num_ties <= kWarpSize) { + // <=32 ties - Use warp ballot + // All-to-all comparison in one __ballot_sync. 32 ties x 32 warps = 1024 + // comparisons in one instruction per warp. O(1) work. + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + if (lane_id >= num_ties || warp_id >= num_ties) return; + const uint32_t mask = (1ull << num_ties) - 1u; + const auto tie = smem->tie_buffer[lane_id]; // each lane holds one tie + const auto target = + smem->tie_buffer[warp_id]; // each warp evaluates one candidate + const bool pred = + is_greater(tie, target); // compare all ties against target + const auto rank = static_cast( + __popc(__ballot_sync(mask, pred))); // count how many are greater + if (lane_id == 0 && rank < topk_remain) { + output[num_above + rank] = target.idx; // place at correct position + } + } else if (num_ties <= + kWarpSize * + 2) { // TODO (roberto): try to refactor this with <=32 case + // Same idea but each thread handles 2 tie elements + const auto lane_id = tx % kWarpSize; + const auto warp_id = tx / kWarpSize; + const auto lane1 = lane_id + kWarpSize; + const auto warp1 = warp_id + kWarpSize; + const auto invalid = Tie{0xFFFFFFFF, -__FLT_MAX__}; + const auto tie0 = smem->tie_buffer[lane_id]; + const auto tie1 = lane1 < num_ties ? smem->tie_buffer[lane1] : invalid; + if (warp_id < num_ties) { + const auto target = smem->tie_buffer[warp_id]; + const auto r0 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); + const auto r1 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); + if (lane_id == 0 && r0 + r1 < topk_remain) + output[num_above + r0 + r1] = target.idx; + } + if (warp1 < num_ties) { + const auto target = smem->tie_buffer[warp1]; + const auto r0 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); + const auto r1 = + __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); + if (lane_id == 0 && r0 + r1 < topk_remain) + output[num_above + r0 + r1] = target.idx; + } + } else { + // Large tie count: fall back to 4-round radix-256 sort + if constexpr (TopK <= kBlockSize) { + tie_handle(smem->tie_buffer, num_ties, num_above, output, smem); + } else { + tie_handle_large(smem->tie_buffer, num_ties, num_above, output, + smem); + } + } +} + +template +__device__ __noinline__ void histogram_4096_topk_predicated( + const float* __restrict__ scores, int32_t* __restrict__ output, + uint32_t length, void* _smem) { + histogram_4096_topk(scores, output, + length, _smem); +} + +} // namespace topk_histogram_4096 +} // namespace vllm + +#endif // TOPK_HISTOGRAM_4096_CUH_ diff --git a/prime_kernels/indexed_attention/selection.py b/prime_kernels/indexed_attention/selection.py index 0a78e00..b214844 100644 --- a/prime_kernels/indexed_attention/selection.py +++ b/prime_kernels/indexed_attention/selection.py @@ -1,89 +1,136 @@ -import tilelang +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import math + import torch -import torch.nn.functional as F -from tilelang import language as T +import triton +import triton.language as tl -# A 32K-token sequence produces a 1 GiB score matrix and runs in one pass. +# A 32K-token packed batch produces at most a 1 GiB score workspace. SCORE_WORKSPACE_BYTES = 1024**3 +TOPK_WORKSPACE_BYTES = 1024**2 + + +@triton.jit +def _selection_scores_kernel( + query_ptr, + key_ptr, + starts_ptr, + ends_ptr, + visible_blocks_ptr, + scores_ptr, + stride_query_row, + stride_query_head, + stride_query_dim, + stride_key_row, + stride_key_dim, + stride_scores_row, + rows, + columns, + key_blocks, + score_divisor, + NUM_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_D: tl.constexpr, + TILES_PER_PROGRAM: tl.constexpr, + STAGES: tl.constexpr, + MAX_N: tl.constexpr, +) -> None: + row = tl.program_id(0) + dimensions = tl.arange(0, BLOCK_D) + heads = tl.arange(0, MAX_N) + start = tl.load(starts_ptr + row) + end = tl.load(ends_ptr + row) + visible = end - start + if tl.program_id(1) == 0: + tl.store(visible_blocks_ptr + row, visible) + + tile_start = tl.program_id(1) * TILES_PER_PROGRAM + if tile_start * BLOCK_N >= visible: + return + tile_end = tl.minimum(tile_start + TILES_PER_PROGRAM, tl.cdiv(visible, BLOCK_N)) + tile_end = tl.minimum(tile_end, tl.cdiv(columns, BLOCK_N)) + + query = tl.load( + query_ptr + + row * stride_query_row + + heads[None, :] * stride_query_head + + dimensions[:, None] * stride_query_dim, + mask=(heads[None, :] < NUM_HEADS) & (dimensions[:, None] < HEAD_DIM), + other=0.0, + ) + column_offsets = tl.arange(0, BLOCK_N) + for tile in tl.range(tile_start, tile_end, num_stages=STAGES): + columns_in_tile = tile * BLOCK_N + column_offsets + key_rows = start + columns_in_tile + live = (columns_in_tile < visible) & (key_rows < key_blocks) + keys = tl.load( + key_ptr + key_rows[:, None].to(tl.int64) * stride_key_row + dimensions[None, :] * stride_key_dim, + mask=live[:, None] & (dimensions[None, :] < HEAD_DIM), + other=0.0, + eviction_policy="evict_first", + ) + head_scores = tl.dot(keys, query, out_dtype=tl.float32) + head_scores = tl.where(heads[None, :] < NUM_HEADS, tl.maximum(head_scores, 0.0), 0.0) + scores = tl.sum(head_scores, axis=1) / score_divisor + tl.store( + scores_ptr + row * stride_scores_row + columns_in_tile, + tl.where(live, scores, -float("inf")), + mask=columns_in_tile < columns, + ) -@tilelang.jit(out_idx=[-1]) -def indexed_selection_scores_kernel( - num_query_heads: int, - head_dim: int, - block_queries: int = 64, - block_keys: int = 128, - threads: int = 256, -): - query_tokens = T.dynamic("query_tokens") - key_blocks = T.dynamic("key_blocks") - - query_shape = [query_tokens * num_query_heads, head_dim] - key_shape = [key_blocks, head_dim] - bounds_shape = [query_tokens] - scores_shape = [query_tokens, key_blocks] - - @T.prim_func - def kernel( - query: T.Tensor(query_shape, T.bfloat16), - key: T.Tensor(key_shape, T.bfloat16), - starts: T.Tensor(bounds_shape, T.int32), - ends: T.Tensor(bounds_shape, T.int32), - scores: T.Tensor(scores_shape, T.float32), - ): - with T.Kernel( - T.ceildiv(key_blocks, block_keys), - T.ceildiv(query_tokens, block_queries), - threads=threads, - ) as (key_block, query_block): - query_shared = T.alloc_shared([block_queries, head_dim], T.float32) - key_shared = T.alloc_shared([block_keys, head_dim], T.float32) - head_scores = T.alloc_fragment([block_queries, block_keys], T.float32) - combined_scores = T.alloc_fragment([block_queries, block_keys], T.float32) - - for row, dim in T.Parallel(block_keys, head_dim): - key_index = key_block * block_keys + row - key_shared[row, dim] = T.if_then_else( - key_index < key_blocks, - T.cast(key[key_index, dim], T.float32), - 0, - ) - T.clear(combined_scores) - for head in T.serial(num_query_heads): - for row, dim in T.Parallel(block_queries, head_dim): - query_index = query_block * block_queries + row - query_shared[row, dim] = T.if_then_else( - query_index < query_tokens, - T.cast(query[query_index * num_query_heads + head, dim], T.float32), - 0, - ) - T.gemm( - query_shared, - key_shared, - head_scores, - transpose_B=True, - clear_accum=True, - policy=T.GemmWarpPolicy.FullRow, - ) - for row, column in T.Parallel(block_queries, block_keys): - combined_scores[row, column] += T.max(head_scores[row, column], 0) - - for row, column in T.Parallel(block_queries, block_keys): - query_index = query_block * block_queries + row - key_index = key_block * block_keys + column - if query_index < query_tokens and key_index < key_blocks: - if key_index < starts[query_index] or key_index >= ends[query_index]: - combined_scores[row, column] = -T.infinity(T.float32) - - T.copy( - combined_scores, - scores[query_block * block_queries, key_block * block_keys], - ) - - return kernel - - -def _select_blocks( +def selection_scores( + query: torch.Tensor, + key: torch.Tensor, + starts: torch.Tensor, + ends: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + columns = key.shape[0] + scores = torch.empty((query.shape[0], columns), dtype=torch.float32, device=query.device) + visible_blocks = torch.empty(query.shape[0], dtype=torch.int32, device=query.device) + if not query.shape[0] or not columns: + return scores, visible_blocks + + block_n = 64 + block_d = max(16, triton.next_power_of_2(query.shape[2])) + max_n = max(16, triton.next_power_of_2(query.shape[1])) + tiles_per_program = 1 if query.shape[0] <= 32 else 8 + _selection_scores_kernel[(query.shape[0], triton.cdiv(columns, block_n * tiles_per_program))]( + query, + key, + starts, + ends, + visible_blocks, + scores, + query.stride(0), + query.stride(1), + query.stride(2), + key.stride(0), + key.stride(1), + scores.stride(0), + query.shape[0], + columns, + key.shape[0], + math.sqrt(query.shape[2]), + NUM_HEADS=query.shape[1], + HEAD_DIM=query.shape[2], + BLOCK_N=block_n, + BLOCK_D=block_d, + TILES_PER_PROGRAM=tiles_per_program, + STAGES=2, + MAX_N=max_n, + num_warps=2, + ) + return scores, visible_blocks + + +@torch.library.custom_op("prime_kernels::select_indexed_blocks", mutates_args=()) +def select_indexed_blocks( query: torch.Tensor, key: torch.Tensor, starts: torch.Tensor, @@ -91,52 +138,48 @@ def _select_blocks( topk: int, ) -> torch.Tensor: num_blocks = key.shape[0] - if num_blocks == 0: + if not num_blocks: return torch.zeros(query.shape[0], topk, dtype=torch.int32, device=query.device) - chunk_size = max(1, SCORE_WORKSPACE_BYTES // (num_blocks * torch.float32.itemsize)) + rows_per_chunk = max(1, SCORE_WORKSPACE_BYTES // (num_blocks * torch.float32.itemsize)) selected_chunks = [] - selected_count = min(topk, num_blocks) - key = key.contiguous() + workspace = torch.empty(TOPK_WORKSPACE_BYTES, dtype=torch.uint8, device=query.device) for query_chunk, start_chunk, end_chunk in zip( - query.split(chunk_size), - starts.split(chunk_size), - ends.split(chunk_size), + query.split(rows_per_chunk), + starts.split(rows_per_chunk), + ends.split(rows_per_chunk), strict=True, ): - scores = indexed_selection_scores_kernel( - query.shape[1], - query.shape[2], - )( - query_chunk.flatten(0, 1).contiguous(), + scores, visible_blocks = selection_scores( + query_chunk, key, - start_chunk.contiguous(), - end_chunk.contiguous(), + start_chunk, + end_chunk, + ) + selected = torch.full( + (query_chunk.shape[0], topk), + -1, + dtype=torch.int32, + device=query.device, ) - selected = scores.topk(selected_count, dim=-1).indices - if selected_count < topk: - selected = F.pad(selected, (0, topk - selected_count), value=num_blocks) - selected.masked_fill_( - (selected < start_chunk[:, None]) | (selected >= end_chunk[:, None]), + torch.ops.prime_indexed_attention.persistent_topk( + scores, + visible_blocks, + selected, + workspace, + topk, num_blocks, ) - selected_chunks.append(selected.to(torch.int32)) + valid = (selected >= 0) & (selected < visible_blocks[:, None]) + selected.add_(start_chunk[:, None]) + selected.masked_fill_(~valid, num_blocks) + selected_chunks.append(selected) + if len(selected_chunks) == 1: return selected_chunks[0] return torch.cat(selected_chunks) -@torch.library.custom_op("prime_kernels::select_indexed_blocks", mutates_args=()) -def select_indexed_blocks( - query: torch.Tensor, - key: torch.Tensor, - starts: torch.Tensor, - ends: torch.Tensor, - topk: int, -) -> torch.Tensor: - return _select_blocks(query, key, starts, ends, topk) - - @select_indexed_blocks.register_fake def select_indexed_blocks_fake( query: torch.Tensor, diff --git a/prime_kernels/kernels.toml b/prime_kernels/kernels.toml index 3383b7c..557c55e 100644 --- a/prime_kernels/kernels.toml +++ b/prime_kernels/kernels.toml @@ -32,9 +32,12 @@ arch = ["10.0"] [indexed_attention] description = "Training forward and backward for token-indexed grouped-query attention" -python-only = true -requires = ["tilelang"] +ops = "prime_indexed_attention" +sources = ["csrc/topk.cu"] +include-dirs = ["csrc"] +requires = ["tilelang", "triton"] arch = ["8.0", "9.0", "10.0"] +cxx-std = 20 # rmsnorm is not built yet: only its sources are committed. Uncomment the table below to # put it back into the build (and to make the registry report on it). diff --git a/setup.py b/setup.py index 06c81a4..471910d 100644 --- a/setup.py +++ b/setup.py @@ -98,8 +98,10 @@ def _extension(kernel) -> CUDAExtension: # Listed explicitly: the kernel folders carry C++/CUDA sources next to their Python, and # only the Python surface plus the compiled extension belongs in the wheel. packages=["prime_kernels", *(f"prime_kernels.{name}" for name in kernels)], + include_package_data=False, package_data={ "prime_kernels": ["kernels.toml"], + "prime_kernels.indexed_attention": ["LICENSE.vllm"], "prime_kernels.mxfp8_moe": ["LICENSE.torchao"], }, ext_modules=extensions, From 51e6ed6a04895c5cdb3cd9b0b89de6a41af4bf4f Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Mon, 31 Aug 2026 16:02:55 +0000 Subject: [PATCH 4/4] refactor: implement indexed selection in TileLang --- README.md | 5 +- prime_kernels/indexed_attention/LICENSE.vllm | 201 --- prime_kernels/indexed_attention/__init__.py | 17 - .../csrc/persistent_topk.cuh | 1363 ----------------- prime_kernels/indexed_attention/csrc/topk.cu | 287 ---- .../csrc/topk_histogram_4096.cuh | 563 ------- prime_kernels/indexed_attention/selection.py | 389 +++-- prime_kernels/kernels.toml | 7 +- setup.py | 2 - 9 files changed, 245 insertions(+), 2589 deletions(-) delete mode 100644 prime_kernels/indexed_attention/LICENSE.vllm delete mode 100644 prime_kernels/indexed_attention/csrc/persistent_topk.cuh delete mode 100644 prime_kernels/indexed_attention/csrc/topk.cu delete mode 100644 prime_kernels/indexed_attention/csrc/topk_histogram_4096.cuh diff --git a/README.md b/README.md index 3560bf8..999a791 100644 --- a/README.md +++ b/README.md @@ -57,8 +57,9 @@ neither built nor shipped in the wheel, and the registry does not list it. grouped GEMM and MXFP8 expert-parallel transport. It is registered as Python-only because it orchestrates PyTorch and torchao kernels rather than compiling a `_C` extension here. `indexed_attention` provides differentiable grouped-query attention over an explicit token -selection for each query. Its TileLang kernels accept different query and KV lengths so the -caller can gather KV for context parallelism without gathering queries. +selection for each query. Its TileLang kernels compute selection scores and radix selection +as well as attention, and accept different query and KV lengths so the caller can gather KV +for context parallelism without gathering queries. ## Installing diff --git a/prime_kernels/indexed_attention/LICENSE.vllm b/prime_kernels/indexed_attention/LICENSE.vllm deleted file mode 100644 index 261eeb9..0000000 --- a/prime_kernels/indexed_attention/LICENSE.vllm +++ /dev/null @@ -1,201 +0,0 @@ - Apache License - Version 2.0, January 2004 - http://www.apache.org/licenses/ - - TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION - - 1. Definitions. - - "License" shall mean the terms and conditions for use, reproduction, - and distribution as defined by Sections 1 through 9 of this document. - - "Licensor" shall mean the copyright owner or entity authorized by - the copyright owner that is granting the License. - - "Legal Entity" shall mean the union of the acting entity and all - other entities that control, are controlled by, or are under common - control with that entity. For the purposes of this definition, - "control" means (i) the power, direct or indirect, to cause the - direction or management of such entity, whether by contract or - otherwise, or (ii) ownership of fifty percent (50%) or more of the - outstanding shares, or (iii) beneficial ownership of such entity. - - "You" (or "Your") shall mean an individual or Legal Entity - exercising permissions granted by this License. - - "Source" form shall mean the preferred form for making modifications, - including but not limited to software source code, documentation - source, and configuration files. - - "Object" form shall mean any form resulting from mechanical - transformation or translation of a Source form, including but - not limited to compiled object code, generated documentation, - and conversions to other media types. - - "Work" shall mean the work of authorship, whether in Source or - Object form, made available under the License, as indicated by a - copyright notice that is included in or attached to the work - (an example is provided in the Appendix below). - - "Derivative Works" shall mean any work, whether in Source or Object - form, that is based on (or derived from) the Work and for which the - editorial revisions, annotations, elaborations, or other modifications - represent, as a whole, an original work of authorship. For the purposes - of this License, Derivative Works shall not include works that remain - separable from, or merely link (or bind by name) to the interfaces of, - the Work and Derivative Works thereof. - - "Contribution" shall mean any work of authorship, including - the original version of the Work and any modifications or additions - to that Work or Derivative Works thereof, that is intentionally - submitted to Licensor for inclusion in the Work by the copyright owner - or by an individual or Legal Entity authorized to submit on behalf of - the copyright owner. For the purposes of this definition, "submitted" - means any form of electronic, verbal, or written communication sent - to the Licensor or its representatives, including but not limited to - communication on electronic mailing lists, source code control systems, - and issue tracking systems that are managed by, or on behalf of, the - Licensor for the purpose of discussing and improving the Work, but - excluding communication that is conspicuously marked or otherwise - designated in writing by the copyright owner as "Not a Contribution." - - "Contributor" shall mean Licensor and any individual or Legal Entity - on behalf of whom a Contribution has been received by Licensor and - subsequently incorporated within the Work. - - 2. Grant of Copyright License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - copyright license to reproduce, prepare Derivative Works of, - publicly display, publicly perform, sublicense, and distribute the - Work and such Derivative Works in Source or Object form. - - 3. Grant of Patent License. Subject to the terms and conditions of - this License, each Contributor hereby grants to You a perpetual, - worldwide, non-exclusive, no-charge, royalty-free, irrevocable - (except as stated in this section) patent license to make, have made, - use, offer to sell, sell, import, and otherwise transfer the Work, - where such license applies only to those patent claims licensable - by such Contributor that are necessarily infringed by their - Contribution(s) alone or by combination of their Contribution(s) - with the Work to which such Contribution(s) was submitted. If You - institute patent litigation against any entity (including a - cross-claim or counterclaim in a lawsuit) alleging that the Work - or a Contribution incorporated within the Work constitutes direct - or contributory patent infringement, then any patent licenses - granted to You under this License for that Work shall terminate - as of the date such litigation is filed. - - 4. Redistribution. You may reproduce and distribute copies of the - Work or Derivative Works thereof in any medium, with or without - modifications, and in Source or Object form, provided that You - meet the following conditions: - - (a) You must give any other recipients of the Work or - Derivative Works a copy of this License; and - - (b) You must cause any modified files to carry prominent notices - stating that You changed the files; and - - (c) You must retain, in the Source form of any Derivative Works - that You distribute, all copyright, patent, trademark, and - attribution notices from the Source form of the Work, - excluding those notices that do not pertain to any part of - the Derivative Works; and - - (d) If the Work includes a "NOTICE" text file as part of its - distribution, then any Derivative Works that You distribute must - include a readable copy of the attribution notices contained - within such NOTICE file, excluding those notices that do not - pertain to any part of the Derivative Works, in at least one - of the following places: within a NOTICE text file distributed - as part of the Derivative Works; within the Source form or - documentation, if provided along with the Derivative Works; or, - within a display generated by the Derivative Works, if and - wherever such third-party notices normally appear. The contents - of the NOTICE file are for informational purposes only and - do not modify the License. You may add Your own attribution - notices within Derivative Works that You distribute, alongside - or as an addendum to the NOTICE text from the Work, provided - that such additional attribution notices cannot be construed - as modifying the License. - - You may add Your own copyright statement to Your modifications and - may provide additional or different license terms and conditions - for use, reproduction, or distribution of Your modifications, or - for any such Derivative Works as a whole, provided Your use, - reproduction, and distribution of the Work otherwise complies with - the conditions stated in this License. - - 5. Submission of Contributions. Unless You explicitly state otherwise, - any Contribution intentionally submitted for inclusion in the Work - by You to the Licensor shall be under the terms and conditions of - this License, without any additional terms or conditions. - Notwithstanding the above, nothing herein shall supersede or modify - the terms of any separate license agreement you may have executed - with Licensor regarding such Contributions. - - 6. Trademarks. This License does not grant permission to use the trade - names, trademarks, service marks, or product names of the Licensor, - except as required for reasonable and customary use in describing the - origin of the Work and reproducing the content of the NOTICE file. - - 7. Disclaimer of Warranty. Unless required by applicable law or - agreed to in writing, Licensor provides the Work (and each - Contributor provides its Contributions) on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or - implied, including, without limitation, any warranties or conditions - of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A - PARTICULAR PURPOSE. You are solely responsible for determining the - appropriateness of using or redistributing the Work and assume any - risks associated with Your exercise of permissions under this License. - - 8. Limitation of Liability. In no event and under no legal theory, - whether in tort (including negligence), contract, or otherwise, - unless required by applicable law (such as deliberate and grossly - negligent acts) or agreed to in writing, shall any Contributor be - liable to You for damages, including any direct, indirect, special, - incidental, or consequential damages of any character arising as a - result of this License or out of the use or inability to use the - Work (including but not limited to damages for loss of goodwill, - work stoppage, computer failure or malfunction, or any and all - other commercial damages or losses), even if such Contributor - has been advised of the possibility of such damages. - - 9. Accepting Warranty or Additional Liability. While redistributing - the Work or Derivative Works thereof, You may choose to offer, - and charge a fee for, acceptance of support, warranty, indemnity, - or other liability obligations and/or rights consistent with this - License. However, in accepting such obligations, You may act only - on Your own behalf and on Your sole responsibility, not on behalf - of any other Contributor, and only if You agree to indemnify, - defend, and hold each Contributor harmless for any liability - incurred by, or claims asserted against, such Contributor by reason - of your accepting any such warranty or additional liability. - - END OF TERMS AND CONDITIONS - - APPENDIX: How to apply the Apache License to your work. - - To apply the Apache License to your work, attach the following - boilerplate notice, with the fields enclosed by brackets "[]" - replaced with your own identifying information. (Don't include - the brackets!) The text should be enclosed in the appropriate - comment syntax for the file format. We also recommend that a - file or class name and description of purpose be included on the - same "printed page" as the copyright notice for easier - identification within third-party archives. - - Copyright [yyyy] [name of copyright owner] - - Licensed under the Apache License, Version 2.0 (the "License"); - you may not use this file except in compliance with the License. - You may obtain a copy of the License at - - http://www.apache.org/licenses/LICENSE-2.0 - - Unless required by applicable law or agreed to in writing, software - distributed under the License is distributed on an "AS IS" BASIS, - WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - See the License for the specific language governing permissions and - limitations under the License. diff --git a/prime_kernels/indexed_attention/__init__.py b/prime_kernels/indexed_attention/__init__.py index 569b33b..4cfe0c3 100644 --- a/prime_kernels/indexed_attention/__init__.py +++ b/prime_kernels/indexed_attention/__init__.py @@ -1,21 +1,4 @@ -from __future__ import annotations - -import torch - -from . import _C # noqa: F401 from prime_kernels.indexed_attention.forward import indexed_attention, unsupported_shape_reason from prime_kernels.indexed_attention.selection import select_indexed_blocks __all__ = ["indexed_attention", "select_indexed_blocks", "unsupported_shape_reason"] - - -@torch.library.register_fake("prime_indexed_attention::persistent_topk") -def _persistent_topk_fake( - logits: torch.Tensor, - lengths: torch.Tensor, - output: torch.Tensor, - workspace: torch.Tensor, - k: int, - max_seq_len: int, -) -> None: - return None diff --git a/prime_kernels/indexed_attention/csrc/persistent_topk.cuh b/prime_kernels/indexed_attention/csrc/persistent_topk.cuh deleted file mode 100644 index 9ae0fd4..0000000 --- a/prime_kernels/indexed_attention/csrc/persistent_topk.cuh +++ /dev/null @@ -1,1363 +0,0 @@ -/* - * Persistent TopK Scheduler for DSA Indexer - */ - -#ifndef PERSISTENT_TOPK_CUH_ -#define PERSISTENT_TOPK_CUH_ - -#include -#include -#include -#include -#include - -#include "topk_histogram_4096.cuh" - -namespace vllm { -namespace persistent { - -// ============================================================================ -// Constants -// ============================================================================ - -constexpr int kThreadsPerBlock = 1024; -constexpr int RADIX = 256; - -// Medium path: all shared state in dynamic smem (no static __shared__, -// which would inflate the kernel's smem footprint and kill occupancy -// for the decode/trivial paths). -constexpr size_t kMediumHistBytes = 2 * (RADIX + 128) * sizeof(int); // 3072 -constexpr size_t kMediumScalarsBytes = 5 * sizeof(int); // 20 -constexpr size_t kMediumHeaderSize = - (kMediumHistBytes + kMediumScalarsBytes + 127) & ~size_t(127); // 3200 -constexpr int MAX_BUFFERED_ITEMS = 4096; -constexpr size_t kSmemMedium = - kMediumHeaderSize + 2 * MAX_BUFFERED_ITEMS * sizeof(int); // 35968 -constexpr uint32_t RADIX_THRESHOLD = 32768; - -// Decode path constants -constexpr int kDecodeBins = 2048; -constexpr uint32_t HIST2048_THRESHOLD = 8192; - -// Large path: fixed shared memory for histograms + scalars -constexpr size_t kFixedSmemLarge = - ((RADIX + RADIX + 5) * sizeof(uint32_t) + 15) & ~size_t(15); - -// ============================================================================ -// Common helpers -// ============================================================================ - -__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { - uint32_t bits = __float_as_uint(x); - return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); -} - -__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t { - __half h = __float2half_rn(x); - uint16_t bits = __half_as_ushort(h); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return static_cast(key >> 8); -} - -// ============================================================================ -// Vectorized load helpers -// ============================================================================ - -// Unconditional float4 load with cache hint (.cg = cache at global level only). -__device__ __forceinline__ void load_float4(const float* ptr, float& v0, - float& v1, float& v2, float& v3) { - uint32_t r0, r1, r2, r3; - asm volatile("ld.global.cg.v4.u32 {%0,%1,%2,%3}, [%4];\n" - : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) - : "l"(ptr)); - v0 = __uint_as_float(r0); - v1 = __uint_as_float(r1); - v2 = __uint_as_float(r2); - v3 = __uint_as_float(r3); -} - -// Per-element predicated scalar loads with -inf default. -__device__ __forceinline__ void load_float4_predicated(const float* ptr, - int base, int seq_len, - float& v0, float& v1, - float& v2, float& v3) { - uint32_t r0, r1, r2, r3; - int p0 = (base < seq_len); - int p1 = (base + 1 < seq_len); - int p2 = (base + 2 < seq_len); - int p3 = (base + 3 < seq_len); - asm volatile( - "{\n" - " .reg .pred pr0, pr1, pr2, pr3;\n" - " setp.ne.u32 pr0, %4, 0;\n" - " setp.ne.u32 pr1, %5, 0;\n" - " setp.ne.u32 pr2, %6, 0;\n" - " setp.ne.u32 pr3, %7, 0;\n" - " mov.u32 %0, 0xFF800000;\n" - " mov.u32 %1, 0xFF800000;\n" - " mov.u32 %2, 0xFF800000;\n" - " mov.u32 %3, 0xFF800000;\n" - " @pr0 ld.global.cg.u32 %0, [%8];\n" - " @pr1 ld.global.cg.u32 %1, [%8+4];\n" - " @pr2 ld.global.cg.u32 %2, [%8+8];\n" - " @pr3 ld.global.cg.u32 %3, [%8+12];\n" - "}\n" - : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) - : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr)); - v0 = __uint_as_float(r0); - v1 = __uint_as_float(r1); - v2 = __uint_as_float(r2); - v3 = __uint_as_float(r3); -} - -// ============================================================================ -// Large path: inter-CTA coordination state (one per group) -// ============================================================================ - -struct RadixRowState { - uint32_t histogram[3][256]; // Triple-buffered histograms - uint32_t remaining_k; - uint32_t prefix; - int arrival_counter; - int output_counter; -}; - -// ============================================================================ -// Kernel parameters -// ============================================================================ - -struct PersistentTopKParams { - const float* __restrict__ input; // [num_rows, stride] - int32_t* __restrict__ output; // [num_rows, top_k] - const int32_t* __restrict__ lengths; // [num_rows] - RadixRowState* row_states; // large path: per-group state - uint32_t num_rows; - uint32_t stride; - uint32_t top_k; // actual k value for output stride - uint32_t chunk_size; // large path: elements per CTA - uint32_t ctas_per_group; // 1=medium, >1=large - uint32_t max_seq_len; // max seq_len across all rows (for early CTA exit) -}; - -// ============================================================================ -// Decode path: 2048-bin histogram for short sequences (seq_len <= 8192) -// Uses 11-bit half-precision bins for fine granularity. -// One histogram pass typically suffices since 8192/2048 = 4 elements/bin avg. -// ============================================================================ - -// 11-bit bin from half-precision representation (ascending: high values -> high -// bins) -__device__ __forceinline__ uint32_t decode_bin(float x) { - __half hx = __float2half(x); - uint16_t bits = __half_as_ushort(hx); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return key >> 5; -} - -template -__device__ __noinline__ void histogram_2048_topk( - const float* __restrict__ logits, int32_t* __restrict__ output_indices, - int32_t seq_len) { - extern __shared__ int decode_smem[]; - const int tx = threadIdx.x; - const int lane = tx & 31; - - // ---- Layout constants ---- - constexpr int SBASE = 8192 - 8; // 8184 - constexpr int RHIST = RADIX + 128; // 384 - constexpr int BOFF = 2 * RHIST; // 768 - constexpr int DBUF = (SBASE - BOFF) / 2; // 3708 - constexpr int MAX_ITEMS_PER_THREAD = - (HIST2048_THRESHOLD + kThreadsPerBlock - 1) / kThreadsPerBlock; - - enum : int { sTHR = 0, sOUT = 1, sREF = 2, sFIN = 3, sBUF0 = 4, sBUF1 = 5 }; - - // ---- Initialize scalars (prevents stale data from prior rows) ---- - if (tx < 8) { - decode_smem[SBASE + tx] = 0; - } - - // ---- Phase 1: Build 2048-bin histogram with float4 vectorized loads ---- - int* histo = decode_smem; - uint16_t reg_bins[MAX_ITEMS_PER_THREAD]; - int nitems = 0; - - for (int i = tx; i < kDecodeBins; i += kThreadsPerBlock) { - histo[i] = 0; - } - __syncthreads(); - - const int n_vec = (seq_len + 3) >> 2; - const bool row_aligned = ((reinterpret_cast(logits) & 15) == 0); - - for (int i = tx; i < n_vec; i += kThreadsPerBlock) { - const int base = i << 2; - float v0, v1, v2, v3; - - if (row_aligned && base + 3 < seq_len) { - load_float4(logits + base, v0, v1, v2, v3); - } else { - load_float4_predicated(logits + base, base, seq_len, v0, v1, v2, v3); - } - - const uint16_t b0 = static_cast(decode_bin(v0)); - const uint16_t b1 = static_cast(decode_bin(v1)); - const uint16_t b2 = static_cast(decode_bin(v2)); - const uint16_t b3 = static_cast(decode_bin(v3)); - reg_bins[nitems++] = b0; - reg_bins[nitems++] = b1; - reg_bins[nitems++] = b2; - reg_bins[nitems++] = b3; - atomicAdd(&histo[b0], 1); - atomicAdd(&histo[b1], 1); - atomicAdd(&histo[b2], 1); - atomicAdd(&histo[b3], 1); - } - __syncthreads(); - - // ---- CUB suffix sum ---- - using BlockScanT = cub::BlockScan; - const int h0 = histo[2 * tx]; - const int pair_sum = h0 + histo[2 * tx + 1]; - - auto& scan_storage = *reinterpret_cast( - decode_smem + kDecodeBins); - - int pair_prefix, total; - BlockScanT(scan_storage).ExclusiveSum(pair_sum, pair_prefix, total); - - // Find threshold bin purely from registers - const int pair_suffix = total - pair_prefix; - - if (pair_suffix >= TopK && (pair_suffix - h0) < TopK) { - decode_smem[SBASE + sTHR] = 2 * tx; - } - { - const int right_suf = pair_suffix - h0; - const int next_suf = pair_suffix - pair_sum; - if (right_suf >= TopK && next_suf < TopK) { - decode_smem[SBASE + sTHR] = 2 * tx + 1; - } - } - __syncthreads(); - - const int threshold = decode_smem[SBASE + sTHR]; - - // ---- Phase 2: Collection with warp-aggregated atomicAdds ---- - int* bufs[2] = {decode_smem + BOFF, decode_smem + BOFF + DBUF}; - const int sOUT_abs = SBASE + sOUT; - const int sBUF0_abs = SBASE + sBUF0; - - { - const uint32_t uthr = static_cast(threshold); - int item = 0; - const int n_vec_iters = (n_vec + kThreadsPerBlock - 1) / kThreadsPerBlock; - - for (int iter = 0; iter < n_vec_iters; iter++) { - const int i = tx + iter * kThreadsPerBlock; - const bool vec_valid = (i < n_vec); - const int base_idx = i << 2; - -#pragma unroll 4 - for (int sub = 0; sub < 4; sub++) { - const int elem_idx = base_idx + sub; - uint32_t bin = 0; - if (vec_valid) bin = reg_bins[item++]; - const bool is_above = vec_valid && (bin > uthr); - const bool is_equal = vec_valid && (bin == uthr); - - const uint32_t above_mask = __ballot_sync(0xffffffff, is_above); - if (above_mask) { - const int above_count = __popc(above_mask); - const int above_rank = __popc(above_mask & ((1u << lane) - 1)); - int above_base; - if (lane == 0) { - above_base = atomicAdd(&decode_smem[sOUT_abs], above_count); - } - above_base = __shfl_sync(0xffffffff, above_base, 0); - if (is_above) { - output_indices[above_base + above_rank] = elem_idx; - } - } - - const uint32_t equal_mask = __ballot_sync(0xffffffff, is_equal); - if (equal_mask) { - const int equal_count = __popc(equal_mask); - const int equal_rank = __popc(equal_mask & ((1u << lane) - 1)); - int equal_base; - if (lane == 0) { - equal_base = atomicAdd(&decode_smem[sBUF0_abs], equal_count); - } - equal_base = __shfl_sync(0xffffffff, equal_base, 0); - if (is_equal && __builtin_expect(equal_base + equal_rank < DBUF, 1)) { - bufs[0][equal_base + equal_rank] = elem_idx; - } - } - } - } - } - __syncthreads(); - - int remaining_k = TopK - decode_smem[SBASE + sOUT]; - if (remaining_k <= 0) return; - - // If all buffered elements fit, output them all (common for short seqs) - const int raw_buf0 = decode_smem[SBASE + sBUF0]; - if (raw_buf0 <= remaining_k) { - const int nb = (raw_buf0 < DBUF) ? raw_buf0 : DBUF; - const int base = decode_smem[SBASE + sOUT]; - for (int i = tx; i < nb; i += kThreadsPerBlock) { - output_indices[base + i] = bufs[0][i]; - } - __syncthreads(); - return; - } - - // ---- Phase 3: Deferred refinement (rare path) ---- - int* refine[2] = {decode_smem, decode_smem + RHIST}; - const int num_buf0 = (raw_buf0 < DBUF) ? raw_buf0 : DBUF; - - for (int i = tx; i < RHIST; i += kThreadsPerBlock) { - refine[0][i] = 0; - } - __syncthreads(); - - for (int i = tx; i < num_buf0; i += kThreadsPerBlock) { - const uint32_t fp32 = convert_to_uint32_v2(logits[bufs[0][i]]); - atomicAdd(&refine[0][(fp32 >> 24) & 0xFF], 1); - } - __syncthreads(); - - auto compute_suffix_sum = [&]() { -#pragma unroll 8 - for (int i = 0; i < 8; ++i) { - if (tx < RADIX) { - const int stride = 1 << i; - const int s = i & 1; - const int d = s ^ 1; - int value = refine[s][tx]; - if (tx < RADIX - stride) value += refine[s][tx + stride]; - refine[d][tx] = value; - } - __syncthreads(); - } - }; - -#pragma unroll 4 - for (int pass = 0; pass < 4; ++pass) { - const int src = pass & 1; - const int dst = src ^ 1; - - const int raw_buf = decode_smem[SBASE + sBUF0 + src]; - const int num_buffered = (raw_buf < DBUF) ? raw_buf : DBUF; - - compute_suffix_sum(); - - if (tx < RADIX && refine[0][tx] > remaining_k && - refine[0][tx + 1] <= remaining_k) { - decode_smem[SBASE + sREF] = tx; - decode_smem[SBASE + sBUF0 + dst] = 0; - decode_smem[SBASE + sFIN] = remaining_k - refine[0][tx + 1]; - } - __syncthreads(); - - const int ref_thr = decode_smem[SBASE + sREF]; - remaining_k -= refine[0][ref_thr + 1]; - const int bit_offset = 24 - pass * 8; - - if (remaining_k == 0) { - for (int i = tx; i < num_buffered; i += kThreadsPerBlock) { - const int idx = bufs[src][i]; - const uint32_t fp32 = convert_to_uint32_v2(logits[idx]); - if (((fp32 >> bit_offset) & 0xFF) > static_cast(ref_thr)) { - const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1); - output_indices[pos] = idx; - } - } - __syncthreads(); - break; - } - - __syncthreads(); - if (tx < RADIX + 1) refine[0][tx] = 0; - __syncthreads(); - - for (int i = tx; i < num_buffered; i += kThreadsPerBlock) { - const int idx = bufs[src][i]; - const float logit_val = logits[idx]; - const uint32_t fp32 = convert_to_uint32_v2(logit_val); - const int bin = (fp32 >> bit_offset) & 0xFF; - - if (bin > ref_thr) { - const int pos = atomicAdd(&decode_smem[SBASE + sOUT], 1); - output_indices[pos] = idx; - } else if (bin == ref_thr) { - if (pass == 3) { - const int slot = atomicAdd(&decode_smem[SBASE + sFIN], -1); - if (slot > 0) output_indices[TopK - slot] = idx; - } else { - const int bp = atomicAdd(&decode_smem[SBASE + sBUF0 + dst], 1); - if (__builtin_expect(bp < DBUF, 1)) { - bufs[dst][bp] = idx; - const int nbo = bit_offset - 8; - atomicAdd(&refine[0][(fp32 >> nbo) & 0xFF], 1); - } - } - } - } - __syncthreads(); - } -} - -// ============================================================================ -// Medium path: coarse FP16 histogram + 4-pass FP32 radix refinement -// For sequences 8K < seq_len <= 64K. -// ============================================================================ - -// Adapted from: -// https://github.com/sgl-project/sglang/blob/v0.5.8/sgl-kernel/csrc/elementwise/topk.cu#L87 -// by: DarkSharpness -// which at the same time is an optimized topk kernel copied from tilelang -// kernel -template -__device__ __noinline__ void histogram_256_topk( - const float* __restrict__ logits, int* __restrict__ output_indices, - int logits_offset, int seq_len) { - // All shared state lives in dynamic shared memory to avoid static - extern __shared__ char medium_smem[]; - - int (*shared_histogram)[RADIX + 128] = - reinterpret_cast(medium_smem); - int* medium_scalars = reinterpret_cast(medium_smem + kMediumHistBytes); - int& shared_output_count = medium_scalars[0]; - int& shared_threshold_bin = medium_scalars[1]; - int* shared_buffered_count = &medium_scalars[2]; - int& shared_final_k = medium_scalars[4]; - int (*buffered_indices)[MAX_BUFFERED_ITEMS] = - reinterpret_cast(medium_smem + - kMediumHeaderSize); - - const int thread_id = threadIdx.x; - int remaining_k = TopK; - - if (thread_id < RADIX + 1) { - shared_histogram[0][thread_id] = 0; - } - __syncthreads(); - - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const auto bin = convert_to_uint8(logits[idx + logits_offset]); - atomicAdd(&shared_histogram[0][bin], 1); - } - __syncthreads(); - - auto compute_cumulative_sum = [&]() { -#pragma unroll 8 - for (int i = 0; i < 8; ++i) { - if (__builtin_expect(thread_id < RADIX, 1)) { - const int stride = 1 << i; - const int src_buffer = i & 1; - const int dst_buffer = src_buffer ^ 1; - int value = shared_histogram[src_buffer][thread_id]; - if (thread_id < RADIX - stride) { - value += shared_histogram[src_buffer][thread_id + stride]; - } - shared_histogram[dst_buffer][thread_id] = value; - } - __syncthreads(); - } - }; - - compute_cumulative_sum(); - - if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && - shared_histogram[0][thread_id + 1] <= remaining_k) { - shared_threshold_bin = thread_id; - shared_buffered_count[0] = 0; - shared_output_count = 0; - } - __syncthreads(); - - const int threshold_bin = shared_threshold_bin; - remaining_k -= shared_histogram[0][threshold_bin + 1]; - - if (remaining_k == 0) { - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const int bin = convert_to_uint8(logits[idx + logits_offset]); - if (bin > threshold_bin) { - const int output_pos = atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } - } - __syncthreads(); - return; - } - - __syncthreads(); - if (thread_id < RADIX + 1) { - shared_histogram[0][thread_id] = 0; - } - __syncthreads(); - - for (int idx = thread_id; idx < seq_len; idx += kThreadsPerBlock) { - const float logit_value = logits[idx + logits_offset]; - const int bin = convert_to_uint8(logit_value); - if (bin > threshold_bin) { - const int output_pos = atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } else if (bin == threshold_bin) { - const int buffer_pos = atomicAdd(&shared_buffered_count[0], 1); - if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { - buffered_indices[0][buffer_pos] = idx; - const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); - const int next_bin = (fp32_bits >> 24) & 0xFF; - atomicAdd(&shared_histogram[0][next_bin], 1); - } - } - } - __syncthreads(); - -#pragma unroll 4 - for (int pass = 0; pass < 4; ++pass) { - const int src_buffer = pass % 2; - const int dst_buffer = src_buffer ^ 1; - const int raw_buffered = shared_buffered_count[src_buffer]; - const int num_buffered = - (raw_buffered < MAX_BUFFERED_ITEMS) ? raw_buffered : MAX_BUFFERED_ITEMS; - - compute_cumulative_sum(); - - if (thread_id < RADIX && shared_histogram[0][thread_id] > remaining_k && - shared_histogram[0][thread_id + 1] <= remaining_k) { - shared_threshold_bin = thread_id; - shared_buffered_count[dst_buffer] = 0; - shared_final_k = remaining_k - shared_histogram[0][thread_id + 1]; - } - __syncthreads(); - - const int threshold_bin = shared_threshold_bin; - remaining_k -= shared_histogram[0][threshold_bin + 1]; - const int bit_offset = 24 - pass * 8; - - if (remaining_k == 0) { - for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { - const int idx = buffered_indices[src_buffer][i]; - const uint32_t fp32_bits = - convert_to_uint32_v2(logits[idx + logits_offset]); - const int bin = (fp32_bits >> bit_offset) & 0xFF; - if (bin > threshold_bin) { - const int output_pos = atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } - } - __syncthreads(); - break; - } - - __syncthreads(); - if (thread_id < RADIX + 1) { - shared_histogram[0][thread_id] = 0; - } - __syncthreads(); - - for (int i = thread_id; i < num_buffered; i += kThreadsPerBlock) { - const int idx = buffered_indices[src_buffer][i]; - const float logit_value = logits[idx + logits_offset]; - const uint32_t fp32_bits = convert_to_uint32_v2(logit_value); - const int bin = (fp32_bits >> bit_offset) & 0xFF; - if (bin > threshold_bin) { - const int output_pos = atomicAdd(&shared_output_count, 1); - output_indices[output_pos] = idx; - } else if (bin == threshold_bin) { - if (pass == 3) { - const int slot = atomicAdd(&shared_final_k, -1); - if (slot > 0) { - output_indices[TopK - slot] = idx; - } - } else { - const int buffer_pos = - atomicAdd(&shared_buffered_count[dst_buffer], 1); - if (__builtin_expect(buffer_pos < MAX_BUFFERED_ITEMS, 1)) { - buffered_indices[dst_buffer][buffer_pos] = idx; - const int next_bit_offset = bit_offset - 8; - const int next_bin = (fp32_bits >> next_bit_offset) & 0xFF; - atomicAdd(&shared_histogram[0][next_bin], 1); - } - } - } - } - __syncthreads(); - } -} - -// ============================================================================ -// Inter-CTA sync primitives -// ============================================================================ - -__device__ __forceinline__ int ld_acquire(int* ptr) { - int state = 0; -#if (__CUDA_ARCH__ >= 700) - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" - : "=r"(state) - : "l"(ptr)); -#else - asm volatile("ld.cg.global.b32 %0, [%1];\n" : "=r"(state) : "l"(ptr)); -#endif - return state; -} - -__device__ __forceinline__ void red_release(int* ptr, int val) { -#if (__CUDA_ARCH__ >= 700) - asm volatile("fence.acq_rel.gpu;\n"); - asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" - : - : "l"(ptr), "r"(val)); -#else - __threadfence(); - atomicAdd(ptr, val); -#endif -} - -__device__ __forceinline__ void st_release(int* ptr, int val) { -#if (__CUDA_ARCH__ >= 700) - asm volatile("fence.acq_rel.gpu;\n"); - asm volatile("st.release.gpu.global.b32 [%0], %1;\n" : : "l"(ptr), "r"(val)); -#else - __threadfence(); - atomicExch(ptr, val); -#endif -} - -__device__ __forceinline__ void wait_ge(int* ptr, int target_val, - int thread_idx) { - if (thread_idx == 0) { -#pragma unroll 1 - while (ld_acquire(ptr) < target_val) { - } - } - __syncthreads(); -} - -// ============================================================================ -// Large path: multi-CTA radix select for sequences > 64K -// -// Each row is processed by a group of CTAs. Each CTA loads its chunk into -// shared memory as ordered uint32, then participates in 4 rounds of -// coordinated radix select via global-memory histograms and barriers. -// ============================================================================ - -// ============================================================================ -// Multi-CTA cooperative RadixTopK for a single large row. -// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 -// ============================================================================ - -template -__device__ void radix_topk(const float* __restrict__ row_input, - int32_t* __restrict__ row_output, uint32_t seq_len, - uint32_t my_chunk_start, uint32_t chunk_size, - uint32_t* local_histogram, uint32_t* suffix_sum, - uint32_t* shared_scalars, uint32_t* shared_ordered, - RadixRowState* state, uint32_t cta_in_group, - uint32_t ctas_per_group, int& barrier_phase, - uint32_t radix_iter, uint32_t tx) { - const uint32_t my_chunk_end = (my_chunk_start + chunk_size < seq_len) - ? my_chunk_start + chunk_size - : seq_len; - const uint32_t actual_chunk_size = - (my_chunk_start < seq_len) ? (my_chunk_end - my_chunk_start) : 0; - - // -- Stage 1: Load chunk to shared memory as ordered uint32 -- - { - const uint32_t aligned_size = (actual_chunk_size / VEC_SIZE) * VEC_SIZE; - - for (uint32_t i = tx * VEC_SIZE; i < aligned_size; - i += kThreadsPerBlock * VEC_SIZE) { - const float* src = row_input + my_chunk_start + i; - if constexpr (VEC_SIZE == 4) { - float4 v = *reinterpret_cast(src); - shared_ordered[i] = convert_to_uint32_v2(v.x); - shared_ordered[i + 1] = convert_to_uint32_v2(v.y); - shared_ordered[i + 2] = convert_to_uint32_v2(v.z); - shared_ordered[i + 3] = convert_to_uint32_v2(v.w); - } else if constexpr (VEC_SIZE == 2) { - float2 v = *reinterpret_cast(src); - shared_ordered[i] = convert_to_uint32_v2(v.x); - shared_ordered[i + 1] = convert_to_uint32_v2(v.y); - } else { - shared_ordered[i] = convert_to_uint32_v2(*src); - } - } - for (uint32_t i = aligned_size + tx; i < actual_chunk_size; - i += kThreadsPerBlock) { - shared_ordered[i] = convert_to_uint32_v2(row_input[my_chunk_start + i]); - } - } - __syncthreads(); - - // -- Init radix select state -- - if (tx == 0) { - shared_scalars[0] = 0; // prefix - shared_scalars[1] = TopK; // remaining_k - } - __syncthreads(); - - // -- Initial barrier -- - if (tx == 0) { - red_release(&state->arrival_counter, 1); - } - wait_ge(&state->arrival_counter, - (barrier_phase + 1) * static_cast(ctas_per_group), tx); - barrier_phase++; - __syncthreads(); - - if (cta_in_group == 0 && tx == 0) { - st_release(&state->output_counter, 0); - } - - // -- Stage 2: 4 rounds of radix select -- - for (uint32_t round = 0; round < 4; round++) { - const uint32_t global_round = radix_iter * 4 + round; - const uint32_t shift = 24 - round * 8; - const uint32_t prefix = shared_scalars[0]; - const uint32_t remaining_k = shared_scalars[1]; - - uint32_t* current_hist = state->histogram[global_round % 3]; - uint32_t* next_hist = state->histogram[(global_round + 1) % 3]; - - for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { - local_histogram[i] = 0; - } - __syncthreads(); - - for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { - uint32_t ordered = shared_ordered[i]; - uint32_t mask = (round == 0) ? 0u : (~0u << (32 - round * 8)); - if ((ordered & mask) == prefix) { - uint32_t bucket = (ordered >> shift) & 0xFF; - atomicAdd(&local_histogram[bucket], 1); - } - } - __syncthreads(); - - for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { - if (local_histogram[i] > 0) { - atomicAdd(¤t_hist[i], local_histogram[i]); - } - } - - if (cta_in_group == 0) { - for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { - next_hist[i] = 0; - } - } - - if (tx == 0) { - red_release(&state->arrival_counter, 1); - } - wait_ge(&state->arrival_counter, - (barrier_phase + 1) * static_cast(ctas_per_group), tx); - barrier_phase++; - __syncthreads(); - - for (uint32_t i = tx; i < RADIX; i += kThreadsPerBlock) { - suffix_sum[i] = current_hist[i]; - } - __syncthreads(); - - for (uint32_t stride = 1; stride < RADIX; stride *= 2) { - uint32_t val = 0; - if (tx < RADIX) { - val = suffix_sum[tx]; - if (tx + stride < RADIX) val += suffix_sum[tx + stride]; - } - __syncthreads(); - if (tx < RADIX) suffix_sum[tx] = val; - __syncthreads(); - } - - if (tx == 0) { - shared_scalars[2] = 0; - shared_scalars[3] = remaining_k; - } - __syncthreads(); - - if (tx < RADIX) { - uint32_t count_ge = suffix_sum[tx]; - uint32_t count_gt = (tx + 1 < RADIX) ? suffix_sum[tx + 1] : 0; - if (count_ge >= remaining_k && count_gt < remaining_k) { - shared_scalars[2] = tx; - shared_scalars[3] = remaining_k - count_gt; - } - } - __syncthreads(); - - if (tx == 0) { - shared_scalars[0] = prefix | (shared_scalars[2] << shift); - shared_scalars[1] = shared_scalars[3]; - } - __syncthreads(); - } // end 4 radix rounds - - // -- Count local > pivot elements -- - const uint32_t ordered_pivot = shared_scalars[0]; - - if (tx == 0) suffix_sum[0] = 0; - __syncthreads(); - - uint32_t my_gt_count = 0; - for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { - if (shared_ordered[i] > ordered_pivot) my_gt_count++; - } - for (int offset = 16; offset > 0; offset /= 2) { - my_gt_count += __shfl_down_sync(0xffffffff, my_gt_count, offset); - } - if (tx % 32 == 0 && my_gt_count > 0) { - atomicAdd(&suffix_sum[0], my_gt_count); - } - __syncthreads(); - const uint32_t local_gt_count = suffix_sum[0]; - - // -- Stage 3: Collect top-k indices -- - if (tx == 0) { - local_histogram[0] = 0; - if (local_gt_count > 0) { - local_histogram[1] = - atomicAdd(&state->output_counter, static_cast(local_gt_count)); - } - } - __syncthreads(); - - for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { - if (shared_ordered[i] > ordered_pivot) { - uint32_t local_pos = atomicAdd(&local_histogram[0], 1); - int pos = static_cast(local_histogram[1]) + local_pos; - row_output[pos] = static_cast(my_chunk_start + i); - } - } - - if (tx == 0) { - red_release(&state->arrival_counter, 1); - } - wait_ge(&state->arrival_counter, - (barrier_phase + 1) * static_cast(ctas_per_group), tx); - barrier_phase++; - __syncthreads(); - - for (uint32_t i = tx; i < actual_chunk_size; i += kThreadsPerBlock) { - if (shared_ordered[i] == ordered_pivot) { - int pos = atomicAdd(&state->output_counter, 1); - if (pos < TopK) { - row_output[pos] = static_cast(my_chunk_start + i); - } - } - } -} - -// ============================================================================ -// Persistent kernel — BS≤32, decode/medium/large paths with RadixTopK -// BS>32 uses standalone histogram_256_buffered_topk (separate kernel, -// see filtered_topk.cuh) -// ============================================================================ - -template -__global__ void __launch_bounds__(kThreadsPerBlock, 2) - persistent_topk_kernel(PersistentTopKParams params) { - const uint32_t tx = threadIdx.x; - extern __shared__ uint8_t smem_raw[]; - - // ======================================================================== - // Group mode: multi-CTA groups with static round-robin row assignment. - // Non-large rows: CTA-0 handles trivial/decode/medium. - // Large rows: all CTAs in the group cooperate via RadixTopK. - // ======================================================================== - const uint32_t ctas_per_group = params.ctas_per_group; - const uint32_t group_id = blockIdx.x / ctas_per_group; - const uint32_t cta_in_group = blockIdx.x % ctas_per_group; - const uint32_t num_groups = gridDim.x / ctas_per_group; - const uint32_t chunk_size = params.chunk_size; - - if (blockIdx.x >= num_groups * ctas_per_group) return; - - // Early exit: non-CTA-0 threads are never needed if no large rows exist - if (cta_in_group != 0 && params.max_seq_len <= RADIX_THRESHOLD) return; - - uint32_t* local_histogram = reinterpret_cast(smem_raw); - uint32_t* suffix_sum = local_histogram + RADIX; - uint32_t* shared_scalars = suffix_sum + RADIX; - uint32_t* shared_ordered = - reinterpret_cast(smem_raw + kFixedSmemLarge); - - // RadixRowState for multi-CTA cooperative radix. - // Zero-initialization is done host-side via cudaMemsetAsync in topk.cu - // before launch — that gives a stream-ordered happens-before edge for all - // CTAs, which the previous in-kernel init (CTA-0 only + intra-CTA - // __syncthreads) did not provide and which manifested as a race against - // CTA-1+'s first red_release on arrival_counter. - RadixRowState* state = ¶ms.row_states[group_id]; - - int barrier_phase = 0; - uint32_t radix_iter = 0; - const uint32_t total_iters = (params.num_rows + num_groups - 1) / num_groups; - - for (uint32_t iter = 0; iter < total_iters; iter++) { - // Static round-robin: all CTAs in the group implicitly agree on the row - uint32_t row_idx = group_id + iter * num_groups; - if (row_idx >= params.num_rows) break; - - // Clamp the row length before any decision is made on it. - // - // `lengths` is int32 and is consumed here as uint32, so a negative value - // (e.g. a padded decode slot whose per-token context length underflowed) - // would reinterpret as ~4e9 and sail past every threshold below. Any - // value beyond the row width would also read into the next row. - // - // Clamping to max_seq_len additionally keeps this per-row decision - // consistent with the `cta_in_group != 0` early exit above, which is - // taken from the host-side scalar: when max_seq_len <= RADIX_THRESHOLD - // the non-leader CTAs return immediately, so a leader that reached the - // cooperative radix path would wait on the inter-CTA barrier for peers - // that no longer exist and spin until the kernel is killed. - const int32_t raw_len = params.lengths[row_idx]; - const uint32_t row_bound = - params.stride < params.max_seq_len ? params.stride : params.max_seq_len; - const uint32_t non_negative_len = - raw_len > 0 ? static_cast(raw_len) : 0u; - const uint32_t seq_len = - non_negative_len < row_bound ? non_negative_len : row_bound; - int32_t* row_output = params.output + row_idx * params.top_k; - const float* row_input = params.input + row_idx * params.stride; - - if (seq_len <= RADIX_THRESHOLD) { - if (cta_in_group == 0) { - if (seq_len <= static_cast(TopK)) { - // Trivial case: seq_len <= TopK - for (uint32_t i = tx; i < static_cast(TopK); - i += kThreadsPerBlock) { - row_output[i] = (i < seq_len) ? static_cast(i) : -1; - } - } else if (seq_len <= static_cast(HIST2048_THRESHOLD)) { - histogram_2048_topk(row_input, row_output, seq_len); - } else { - histogram_256_topk(row_input, row_output, 0, seq_len); - } - } - continue; - } - - const uint32_t my_chunk_start = cta_in_group * chunk_size; - radix_topk( - row_input, row_output, seq_len, my_chunk_start, chunk_size, - local_histogram, suffix_sum, shared_scalars, shared_ordered, state, - cta_in_group, ctas_per_group, barrier_phase, radix_iter, tx); - radix_iter++; - } -} - -} // namespace persistent - -// ============================================================================ -// ============================================================================ -// Optimized FilteredTopK — single CTA per row for bs > 32. -// Kept with persistent_topk so the portable fallback owns the non-cluster path. -// ============================================================================ -namespace filtered_topk { - -namespace hist4096 = topk_histogram_4096; - -// ============================================================================ -// FilteredTopK — single CTA per row for bs > 32 -// Adapted from https://github.com/flashinfer-ai/flashinfer/pull/2215 -// ============================================================================ - -#define FLASHINFER_CUDA_CALL(func, ...) \ - { \ - cudaError_t e = (func); \ - if (e != cudaSuccess) { \ - return e; \ - } \ - } - -#define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__ - -template -struct vec_t { - T data[N]; - - FLASHINFER_INLINE T& operator[](size_t i) { return data[i]; } - FLASHINFER_INLINE const T& operator[](size_t i) const { return data[i]; } - - FLASHINFER_INLINE void cast_load(const T* ptr) { -#pragma unroll - for (size_t i = 0; i < N; ++i) { - data[i] = ptr[i]; - } - } -}; -#undef FLASHINFER_INLINE - -// FilteredTopK traits for different data types -template -struct FilteredTopKTraits; - -// Specialization for float (32-bit): coarse histogram uses FP16 high 8 bits, 4 -// refinement rounds -template <> -struct FilteredTopKTraits { - using OrderedType = uint32_t; - static constexpr int NUM_REFINE_ROUNDS = 4; - static constexpr int FIRST_REFINE_SHIFT = 24; - - __device__ __forceinline__ static uint8_t ToCoarseKey(float x) { - // Convert to FP16 representation and extract high 8 bits - __half h = __float2half_rn(x); - uint16_t bits = __half_as_ushort(h); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return static_cast(key >> 8); - } - - __device__ __forceinline__ static OrderedType ToOrdered(float x) { - uint32_t bits = __float_as_uint(x); - return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); - } -}; - -constexpr uint32_t FILTERED_TOPK_BLOCK_THREADS = 1024; -constexpr uint32_t FILTERED_TOPK_SMEM_INPUT_SIZE = - 16 * 1024; // 16K indices per buffer -constexpr size_t FILTERED_TOPK_SMEM_DYNAMIC = - sizeof(int) * 2 * FILTERED_TOPK_SMEM_INPUT_SIZE; // 128KB - -/*! - * \brief Filtered Top-K kernel for ragged sequences. - * - * \tparam DType Data type (float, half, nv_bfloat16) - * \tparam IdType Index type (int32_t) - * \tparam VEC_SIZE Vector size for input loads (1, 2, 4, or 8) - */ -template -__global__ void __launch_bounds__(FILTERED_TOPK_BLOCK_THREADS) - FilteredTopKUnifiedKernel(const DType* __restrict__ input, - IdType* __restrict__ output, - const IdType* __restrict__ lengths, - uint32_t num_rows, uint32_t top_k, - uint32_t max_len) { - constexpr uint32_t BLOCK_SIZE = FILTERED_TOPK_BLOCK_THREADS; - constexpr int RADIX = 256; - constexpr int SMEM_INPUT_SIZE = FILTERED_TOPK_SMEM_INPUT_SIZE; - - const uint32_t bid = blockIdx.x; - const int tx = threadIdx.x; - - if (bid >= num_rows) return; - - const int length = - (lengths != nullptr) ? lengths[bid] : static_cast(max_len); - const DType* score = input + bid * max_len; - IdType* dst = output + bid * top_k; - - // Trivial case: length <= top_k - if (length <= static_cast(top_k)) { - for (int i = tx; i < static_cast(top_k); i += BLOCK_SIZE) { - dst[i] = (i < length) ? static_cast(i) : static_cast(-1); - } - return; - } - - // Short path - if (length <= 32768) { - extern __shared__ uint8_t _smem_reg[]; - if constexpr (UsePredicatedShortLoads) { - hist4096::histogram_4096_topk_predicated(score, dst, length, - _smem_reg); - } else { - hist4096::histogram_4096_topk(score, dst, length, - _smem_reg); - } - return; - } - - // Static shared memory - alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128]; - alignas(128) __shared__ int s_counter; - alignas(128) __shared__ int s_threshold_bin_id; - alignas(128) __shared__ int s_num_input[2]; - alignas(128) __shared__ int s_indices[MAX_K]; - - auto& s_histogram = s_histogram_buf[0]; - - // Dynamic shared memory for input double buffer - extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE]; - - using Traits = FilteredTopKTraits; - int topk = top_k; - - // Stage 1: 8-bit coarse histogram with vectorized loads - if (tx < RADIX + 1) s_histogram[tx] = 0; - __syncthreads(); - - vec_t score_vec; - - const int aligned_length = (length / VEC_SIZE) * VEC_SIZE; -#pragma unroll 2 - for (int base = tx * VEC_SIZE; base < aligned_length; - base += BLOCK_SIZE * VEC_SIZE) { - score_vec.cast_load(&score[base]); -#pragma unroll - for (int j = 0; j < VEC_SIZE; ++j) { - const auto bin = Traits::ToCoarseKey(score_vec[j]); - atomicAdd(&s_histogram[bin], 1); - } - } - // Handle tail - for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { - const auto bin = Traits::ToCoarseKey(score[i]); - atomicAdd(&s_histogram[bin], 1); - } - __syncthreads(); - - // Suffix sum - const auto run_cumsum = [&]() { -#pragma unroll 8 - for (int i = 0; i < 8; ++i) { - if (tx < RADIX) { - const auto j = 1 << i; - const auto k = i & 1; - auto value = s_histogram_buf[k][tx]; - if (tx < RADIX - j) { - value += s_histogram_buf[k][tx + j]; - } - s_histogram_buf[k ^ 1][tx] = value; - } - __syncthreads(); - } - }; - - run_cumsum(); - if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { - s_threshold_bin_id = tx; - s_num_input[0] = 0; - s_counter = 0; - } - __syncthreads(); - - const auto threshold_bin = s_threshold_bin_id; - topk -= s_histogram[threshold_bin + 1]; - - constexpr int NUM_ROUNDS = Traits::NUM_REFINE_ROUNDS; - constexpr int FIRST_SHIFT = Traits::FIRST_REFINE_SHIFT; - - if (topk == 0) { - // Collect indices where bin > threshold -#pragma unroll 2 - for (int base = tx * VEC_SIZE; base < aligned_length; - base += BLOCK_SIZE * VEC_SIZE) { - score_vec.cast_load(&score[base]); -#pragma unroll - for (int j = 0; j < VEC_SIZE; ++j) { - const auto bin = static_cast(Traits::ToCoarseKey(score_vec[j])); - if (bin > threshold_bin) { - const auto pos = atomicAdd(&s_counter, 1); - s_indices[pos] = base + j; - } - } - } - // Handle tail - for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { - const auto bin = static_cast(Traits::ToCoarseKey(score[i])); - if (bin > threshold_bin) { - const auto pos = atomicAdd(&s_counter, 1); - s_indices[pos] = i; - } - } - __syncthreads(); - } else { - __syncthreads(); - if (tx < RADIX + 1) s_histogram[tx] = 0; - __syncthreads(); - - // Filter + histogram for refinement - auto filter_and_add_to_histogram = [&](auto raw_input, int index) { - const auto bin = static_cast(Traits::ToCoarseKey(raw_input)); - if (bin > threshold_bin) { - const auto pos = atomicAdd(&s_counter, 1); - s_indices[pos] = index; - } else if (bin == threshold_bin) { - const auto pos = atomicAdd(&s_num_input[0], 1); - if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) { - s_input_idx[0][pos] = index; - const auto ordered = Traits::ToOrdered(raw_input); - const auto sub_bin = (ordered >> FIRST_SHIFT) & 0xFF; - atomicAdd(&s_histogram[sub_bin], 1); - } - } - }; -#pragma unroll 2 - for (int base = tx * VEC_SIZE; base < aligned_length; - base += BLOCK_SIZE * VEC_SIZE) { - score_vec.cast_load(&score[base]); -#pragma unroll - for (int j = 0; j < VEC_SIZE; ++j) { - filter_and_add_to_histogram(score_vec[j], base + j); - } - } - // Handle tail - for (int i = aligned_length + tx; i < length; i += BLOCK_SIZE) { - filter_and_add_to_histogram(score[i], i); - } - __syncthreads(); - - // Stage 2: refine with 8bit radix passes -#pragma unroll - for (int round = 0; round < NUM_ROUNDS; ++round) { - __shared__ int s_last_remain; - const auto r_idx = round % 2; - - const auto _raw_num_input = s_num_input[r_idx]; - const auto num_input = - (_raw_num_input < SMEM_INPUT_SIZE) ? _raw_num_input : SMEM_INPUT_SIZE; - - run_cumsum(); - if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) { - s_threshold_bin_id = tx; - s_num_input[r_idx ^ 1] = 0; - s_last_remain = topk - s_histogram[tx + 1]; - } - __syncthreads(); - - const auto threshold = s_threshold_bin_id; - topk -= s_histogram[threshold + 1]; - - const int offset = FIRST_SHIFT - round * 8; - const bool is_last_round = (round == NUM_ROUNDS - 1); - - if (topk == 0) { - for (int i = tx; i < num_input; i += BLOCK_SIZE) { - const auto idx = s_input_idx[r_idx][i]; - const auto bin = (Traits::ToOrdered(score[idx]) >> offset) & 0xFF; - if (static_cast(bin) > threshold) { - const auto pos = atomicAdd(&s_counter, 1); - s_indices[pos] = idx; - } - } - __syncthreads(); - break; - } else { - __syncthreads(); - if (tx < RADIX + 1) s_histogram[tx] = 0; - __syncthreads(); - for (int i = tx; i < num_input; i += BLOCK_SIZE) { - const auto idx = s_input_idx[r_idx][i]; - const auto raw_input = score[idx]; - const auto bin = (Traits::ToOrdered(raw_input) >> offset) & 0xFF; - if (static_cast(bin) > threshold) { - const auto pos = atomicAdd(&s_counter, 1); - s_indices[pos] = idx; - } else if (static_cast(bin) == threshold) { - if (is_last_round) { - const auto pos = atomicAdd(&s_last_remain, -1); - if (pos > 0) { - s_indices[top_k - pos] = idx; - } - } else { - const auto pos = atomicAdd(&s_num_input[r_idx ^ 1], 1); - if (__builtin_expect(pos < SMEM_INPUT_SIZE, 1)) { - s_input_idx[r_idx ^ 1][pos] = idx; - const auto bin32 = Traits::ToOrdered(raw_input); - const auto sub_bin = (bin32 >> (offset - 8)) & 0xFF; - atomicAdd(&s_histogram[sub_bin], 1); - } - } - } - } - __syncthreads(); - } - } - } - - // Output phase - mode-specific -#pragma unroll 2 - for (int base = tx; base < static_cast(top_k); base += BLOCK_SIZE) { - const int idx = s_indices[base]; - dst[base] = static_cast(idx); - } -} - -// Helper to compute GCD for VEC_SIZE selection -constexpr uint32_t gcd(uint32_t a, uint32_t b) { - while (b != 0) { - uint32_t t = b; - b = a % b; - a = t; - } - return a; -} - -// Compute optimal VEC_SIZE based on max_len and dtype -// Returns 1, 2, 4, or 8 -template -constexpr int ComputeFilteredTopKVecSize(uint32_t max_len) { - constexpr int MAX_VEC = 16 / sizeof(DType); // 4 for float32, 8 for fp16/bf16 - // Use GCD to find largest power-of-2 divisor - const uint32_t g = gcd(max_len, static_cast(MAX_VEC)); - return static_cast(g); -} - -template -cudaError_t FilteredTopKRaggedTransform(const DType* input, - IdType* output_indices, - const IdType* lengths, - uint32_t num_rows, uint32_t top_k_val, - uint32_t max_len, - cudaStream_t stream = 0) { - constexpr size_t smem_size = FILTERED_TOPK_SMEM_DYNAMIC; - constexpr int MAX_VEC = 16 / sizeof(DType); - - dim3 grid(num_rows); - dim3 block(FILTERED_TOPK_BLOCK_THREADS); - void* args[] = {&input, &output_indices, &lengths, - &num_rows, &top_k_val, &max_len}; - - const int vec_size = ComputeFilteredTopKVecSize(max_len); - -#define DISPATCH_VEC_SIZE(VS) \ - if (vec_size == VS) { \ - auto kernel = \ - FilteredTopKUnifiedKernel; \ - FLASHINFER_CUDA_CALL(cudaFuncSetAttribute( \ - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); \ - FLASHINFER_CUDA_CALL(cudaLaunchKernel((void*)kernel, grid, block, args, \ - smem_size, stream)); \ - return cudaSuccess; \ - } - - DISPATCH_VEC_SIZE(1) - DISPATCH_VEC_SIZE(2) - DISPATCH_VEC_SIZE(4) - if constexpr (MAX_VEC >= 8) { - DISPATCH_VEC_SIZE(8) - } -#undef DISPATCH_VEC_SIZE - - return cudaSuccess; -} - -} // namespace filtered_topk - -template -cudaError_t FilteredTopKRaggedTransform(const DType* input, - IdType* output_indices, - const IdType* lengths, - uint32_t num_rows, uint32_t top_k_val, - uint32_t max_len, - cudaStream_t stream = 0) { - return filtered_topk::FilteredTopKRaggedTransform( - input, output_indices, lengths, num_rows, top_k_val, max_len, stream); -} - -} // namespace vllm - -#endif // PERSISTENT_TOPK_CUH_ diff --git a/prime_kernels/indexed_attention/csrc/topk.cu b/prime_kernels/indexed_attention/csrc/topk.cu deleted file mode 100644 index eadcc11..0000000 --- a/prime_kernels/indexed_attention/csrc/topk.cu +++ /dev/null @@ -1,287 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project -// -// Adapted from vLLM's libtorch-stable persistent top-k interface for the -// prime_indexed_attention dispatcher namespace. - -#include -#include -#include -#include -#include -#include -#include - -#include "persistent_topk.cuh" - -namespace { - -template -void launch_persistent_topk(const torch::Tensor& logits, - const torch::Tensor& lengths, - torch::Tensor& output, - torch::Tensor& workspace, - int64_t max_seq_len) { - namespace P = vllm::persistent; - - const at::cuda::OptionalCUDAGuard device_guard{device_of(logits)}; - const int64_t num_rows = logits.size(0); - const int64_t stride = logits.stride(0); - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - static int num_sms = 0; - static int max_smem_per_block = 0; - if (num_sms == 0) { - const cudaDeviceProp* device_prop = at::cuda::getDeviceProperties(logits.get_device()); - num_sms = device_prop->multiProcessorCount; - max_smem_per_block = device_prop->sharedMemPerBlockOptin; - } - - if (num_rows > 32 && max_smem_per_block >= 128 * 1024) { - cudaError_t status = - vllm::FilteredTopKRaggedTransform( - logits.const_data_ptr(), output.data_ptr(), - lengths.const_data_ptr(), static_cast(num_rows), - static_cast(TopK), static_cast(stride), stream); - TORCH_CHECK(status == cudaSuccess, - "FilteredTopK failed: ", cudaGetErrorString(status)); - } else { - TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor"); - TORCH_CHECK( - workspace.scalar_type() == torch::kUInt8, - "workspace must be uint8"); - - int effective_max_smem; - if (num_rows <= 4) { - effective_max_smem = - std::min(max_smem_per_block, static_cast(P::kSmemMedium)); - } else if (num_rows <= 8) { - constexpr int kSmemCapMedium = 48 * 1024; - effective_max_smem = std::min(max_smem_per_block, kSmemCapMedium); - } else { - effective_max_smem = max_smem_per_block; - } - - size_t available_for_ordered = - static_cast(effective_max_smem) - P::kFixedSmemLarge; - uint32_t max_chunk_elements = - static_cast(available_for_ordered / sizeof(uint32_t)); - - uint32_t vec_size = 1; - if (stride % 4 == 0) - vec_size = 4; - else if (stride % 2 == 0) - vec_size = 2; - - max_chunk_elements = (max_chunk_elements / vec_size) * vec_size; - uint32_t min_chunk = vec_size * P::kThreadsPerBlock; - if (max_chunk_elements < min_chunk) max_chunk_elements = min_chunk; - - uint32_t ctas_per_group = - (static_cast(stride) + max_chunk_elements - 1) / - max_chunk_elements; - uint32_t chunk_size = - (static_cast(stride) + ctas_per_group - 1) / ctas_per_group; - chunk_size = ((chunk_size + vec_size - 1) / vec_size) * vec_size; - if (chunk_size > max_chunk_elements) chunk_size = max_chunk_elements; - - size_t smem_size = P::kFixedSmemLarge + chunk_size * sizeof(uint32_t); - if (smem_size < P::kSmemMedium) smem_size = P::kSmemMedium; - - // Query occupancy for the instantiation that will actually launch; - // overestimating it deadlocks the cooperative barrier. - int occupancy = 1; - cudaError_t occ_err = cudaSuccess; - if (vec_size == 4) { - occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &occupancy, P::persistent_topk_kernel, P::kThreadsPerBlock, - smem_size); - } else if (vec_size == 2) { - occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &occupancy, P::persistent_topk_kernel, P::kThreadsPerBlock, - smem_size); - } else { - occ_err = cudaOccupancyMaxActiveBlocksPerMultiprocessor( - &occupancy, P::persistent_topk_kernel, P::kThreadsPerBlock, - smem_size); - } - TORCH_CHECK(occ_err == cudaSuccess, - "persistent_topk occupancy query failed: ", - cudaGetErrorString(occ_err)); - if (occupancy < 1) occupancy = 1; - - // The cooperative spin-wait barrier only runs when at least one row hits - // the radix path (seq_len > RADIX_THRESHOLD). Below that, non-CTA-0 CTAs - // early-exit, so oversubscription can't deadlock and headroom is wasted. - const bool needs_cooperative = - static_cast(max_seq_len) > P::RADIX_THRESHOLD; - - const uint32_t hw_resident_cap = - static_cast(num_sms) * static_cast(occupancy); - uint32_t max_resident_ctas = hw_resident_cap; - if (needs_cooperative) { - // Reserve one CTA per SM when occupancy allows; fall back to a single - // CTA when occupancy == 1 (the most deadlock-prone case — any straggler - // kernel that takes the only slot on one SM hangs the barrier). Never - // drop below one full group's worth. - uint32_t headroom = (occupancy > 1) ? static_cast(num_sms) : 1u; - if (max_resident_ctas >= headroom + ctas_per_group) { - max_resident_ctas -= headroom; - } - } - uint32_t num_groups = std::min(max_resident_ctas / ctas_per_group, - static_cast(num_rows)); - if (num_groups == 0) num_groups = 1; - uint32_t total_ctas = num_groups * ctas_per_group; - - // If the cooperative launch wouldn't fit, fall back to FilteredTopK - // instead of deadlocking. Only relevant when needs_cooperative. - if (needs_cooperative && total_ctas > hw_resident_cap) { - TORCH_CHECK( - max_smem_per_block >= 128 * 1024, - "persistent_topk would oversubscribe and the FilteredTopK " - "fallback requires >=128KB smem per block (have ", - max_smem_per_block, "). total_ctas=", total_ctas, - " > num_sms*occupancy=", hw_resident_cap, " (TopK=", TopK, - ", vec_size=", vec_size, ", ctas_per_group=", ctas_per_group, - ", smem=", smem_size, ")."); - cudaError_t status = - vllm::FilteredTopKRaggedTransform( - logits.const_data_ptr(), - output.data_ptr(), - lengths.const_data_ptr(), - static_cast(num_rows), static_cast(TopK), - static_cast(stride), stream); - TORCH_CHECK(status == cudaSuccess, "FilteredTopK fallback failed: ", - cudaGetErrorString(status)); - return; - } - - size_t state_bytes = num_groups * sizeof(P::RadixRowState); - TORCH_CHECK(workspace.size(0) >= static_cast(state_bytes), - "workspace too small, need ", state_bytes, " bytes"); - - // Zero the per-group RadixRowState region before launch. - // - // Issued UNCONDITIONALLY so the memset is captured as its own node in - // the cudagraph (a separate cudaMemsetAsync node, sequenced before the - // persistent_topk_kernel launch on the same stream). The previous - // host-side guard `if (needs_cooperative)` was evaluated at capture time; - // when capture-time max_seq_len <= RADIX_THRESHOLD (always true under - // FULL_DECODE_ONLY with max_model_len < 32 K) the memset would NOT be - // captured, leaving the workspace state to accumulate across replays. - // That's a latent correctness bug if the runtime data ever takes the - // radix path, and removes one variable while debugging hangs in the - // decode/medium paths. - // - // Cost is sub-microsecond: state_bytes = num_groups * sizeof(RadixRowState) - // is ~3 KB per group, ~100 KB for the largest grids on this hardware. - // - // Why the memset is required (regardless of which path the kernel takes): - // 1. arrival_counter accumulates within a launch and is never reset, - // so a prior call leaves it at a large positive value. Without this - // reset, the very first wait_ge in the next call sees counter >> - // target and returns instantly, breaking the barrier. - // 2. The previous in-kernel init only ran in CTA-0 with intra-CTA - // __syncthreads(), so it had no happens-before edge to CTA-1+'s - // first red_release. cudaMemsetAsync is stream-ordered: the zero - // is globally visible before any CTA runs. - { - cudaError_t mz_err = cudaMemsetAsync( - workspace.data_ptr(), 0, state_bytes, stream); - TORCH_CHECK(mz_err == cudaSuccess, - "row_states memset failed: ", cudaGetErrorString(mz_err)); - } - - P::PersistentTopKParams params; - params.input = logits.const_data_ptr(); - params.output = output.data_ptr(); - params.lengths = lengths.const_data_ptr(); - params.num_rows = static_cast(num_rows); - params.stride = static_cast(stride); - params.top_k = static_cast(TopK); - params.chunk_size = chunk_size; - params.row_states = reinterpret_cast( - workspace.data_ptr()); - params.ctas_per_group = ctas_per_group; - params.max_seq_len = static_cast(max_seq_len); - - #define LAUNCH_PERSISTENT(TOPK_VAL, VS) \ - do { \ - auto kernel = &P::persistent_topk_kernel; \ - cudaError_t err = cudaFuncSetAttribute( \ - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); \ - TORCH_CHECK(err == cudaSuccess, \ - "Failed to set smem: ", cudaGetErrorString(err)); \ - kernel<<>>(params); \ - } while (0) - - if (vec_size == 4) { - LAUNCH_PERSISTENT(TopK, 4); - } else if (vec_size == 2) { - LAUNCH_PERSISTENT(TopK, 2); - } else { - LAUNCH_PERSISTENT(TopK, 1); - } - #undef LAUNCH_PERSISTENT - } - - cudaError_t err = cudaGetLastError(); - TORCH_CHECK(err == cudaSuccess, - "persistent_topk failed: ", cudaGetErrorString(err)); -} - -} // anonymous namespace - -void persistent_topk(const torch::Tensor& logits, - const torch::Tensor& lengths, - torch::Tensor& output, - torch::Tensor& workspace, - int64_t k, - int64_t max_seq_len) { - TORCH_CHECK(logits.is_cuda(), "logits must be CUDA tensor"); - TORCH_CHECK(lengths.is_cuda(), "lengths must be CUDA tensor"); - TORCH_CHECK(output.is_cuda(), "output must be CUDA tensor"); - TORCH_CHECK(workspace.is_cuda(), "workspace must be CUDA tensor"); - TORCH_CHECK(logits.scalar_type() == torch::kFloat32, - "logits must be float32"); - TORCH_CHECK(lengths.scalar_type() == torch::kInt32, - "lengths must be int32"); - TORCH_CHECK(output.scalar_type() == torch::kInt32, - "output must be int32"); - TORCH_CHECK(workspace.scalar_type() == torch::kUInt8, - "workspace must be uint8"); - TORCH_CHECK(logits.dim() == 2, "logits must be 2D"); - TORCH_CHECK(lengths.dim() == 1, "lengths must be 1D"); - TORCH_CHECK(lengths.is_contiguous(), "lengths must be contiguous"); - TORCH_CHECK(output.dim() == 2, "output must be 2D"); - TORCH_CHECK(logits.is_contiguous(), "logits must be contiguous"); - TORCH_CHECK(output.is_contiguous(), "output must be contiguous"); - - const int64_t num_rows = logits.size(0); - TORCH_CHECK(lengths.numel() == num_rows, "lengths size mismatch"); - TORCH_CHECK(output.size(0) == num_rows && output.size(1) == k, - "output size mismatch"); - TORCH_CHECK(k == 512 || k == 1024 || k == 2048, - "persistent_topk supports k=512, k=1024, or k=2048, got ", k); - - if (k == 512) { - launch_persistent_topk<512>(logits, lengths, output, workspace, - max_seq_len); - } else if (k == 1024) { - launch_persistent_topk<1024>(logits, lengths, output, workspace, - max_seq_len); - } else { - launch_persistent_topk<2048>(logits, lengths, output, workspace, - max_seq_len); - } -} - -TORCH_LIBRARY_FRAGMENT(prime_indexed_attention, m) { - m.def( - "persistent_topk(Tensor logits, Tensor lengths, " - "Tensor(a!) output, Tensor(b!) workspace, int k, int max_seq_len) -> ()"); - m.impl("persistent_topk", torch::kCUDA, &persistent_topk); -} - -PYBIND11_MODULE(_C, m) {} diff --git a/prime_kernels/indexed_attention/csrc/topk_histogram_4096.cuh b/prime_kernels/indexed_attention/csrc/topk_histogram_4096.cuh deleted file mode 100644 index 5f9f823..0000000 --- a/prime_kernels/indexed_attention/csrc/topk_histogram_4096.cuh +++ /dev/null @@ -1,563 +0,0 @@ -/* - * Shared 4096-bin single-CTA TopK helpers. - */ - -#ifndef TOPK_HISTOGRAM_4096_CUH_ -#define TOPK_HISTOGRAM_4096_CUH_ - -#include -#include -#include - -namespace vllm { -namespace topk_histogram_4096 { - -constexpr uint32_t kBlockSize = 1024; -constexpr uint32_t RADIX = 256; -constexpr uint32_t kMaxTies = 1024; -static_assert(kMaxTies <= kBlockSize, - "tie_handle requires kMaxTies <= kBlockSize"); -constexpr uint32_t kWarpSize = 32; -constexpr uint32_t kNumWarps = kBlockSize / kWarpSize; - -// Register path -constexpr uint32_t kHist4096VecsPerThread = 4; -constexpr uint32_t kHist4096MaxLen = - kHist4096VecsPerThread * 4 * kBlockSize; // 16384 - -struct alignas(16) MatchBin { - uint32_t bin, above_count, equal_count; -}; -struct alignas(8) Tie { - uint32_t idx; - float score; -}; - -__device__ __forceinline__ void load_float4_predicated(const float* ptr, - int base, int seq_len, - float& v0, float& v1, - float& v2, float& v3) { - uint32_t r0, r1, r2, r3; - const int p0 = (base < seq_len); - const int p1 = (base + 1 < seq_len); - const int p2 = (base + 2 < seq_len); - const int p3 = (base + 3 < seq_len); - asm volatile( - "{\n" - " .reg .pred pr0, pr1, pr2, pr3;\n" - " setp.ne.u32 pr0, %4, 0;\n" - " setp.ne.u32 pr1, %5, 0;\n" - " setp.ne.u32 pr2, %6, 0;\n" - " setp.ne.u32 pr3, %7, 0;\n" - " mov.u32 %0, 0xFF800000;\n" - " mov.u32 %1, 0xFF800000;\n" - " mov.u32 %2, 0xFF800000;\n" - " mov.u32 %3, 0xFF800000;\n" - " @pr0 ld.global.cg.u32 %0, [%8];\n" - " @pr1 ld.global.cg.u32 %1, [%8+4];\n" - " @pr2 ld.global.cg.u32 %2, [%8+8];\n" - " @pr3 ld.global.cg.u32 %3, [%8+12];\n" - "}\n" - : "=r"(r0), "=r"(r1), "=r"(r2), "=r"(r3) - : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "l"(ptr)); - v0 = __uint_as_float(r0); - v1 = __uint_as_float(r1); - v2 = __uint_as_float(r2); - v3 = __uint_as_float(r3); -} - -// converts the float32 score to a 32-bit ordered unsigned integer — the full -// precision key for radix sorting -__device__ __forceinline__ auto convert_to_uint32_v2(float x) -> uint32_t { - uint32_t bits = __float_as_uint(x); - return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u); -} - -// Converts each score to a 12-bit bin (FP16 sign-magnitude -> top 12 bits -> -// bin 0-4095) -template -__device__ __forceinline__ uint32_t extract_coarse_bin_N(float x) { - __half h = __float2half_rn(x); - uint16_t bits = __half_as_ushort(h); - uint16_t key = (bits & 0x8000) ? static_cast(~bits) - : static_cast(bits | 0x8000); - return key >> (16 - kBits); -} - -// running sum within each warp — thread 0 gets its own value, thread 1 gets -// thread 0 + thread 1, thread 2 gets threads 0+1+2, etc. -__device__ __forceinline__ uint32_t warp_inclusive_sum(uint32_t lane, - uint32_t v) { -#pragma unroll - for (uint32_t o = 1; o < 32; o *= 2) { - uint32_t n = __shfl_up_sync(0xFFFFFFFF, v, o); - if (lane >= o) v += n; - } - return v; -} - -// Returns the sum of a value across all 32 threads in the warp, and every -// thread gets the same result. SM80+ uses redux.sync.add.u32, a single PTX -// instruction for hardware warp-wide reduction. Older targets use the -// __shfl_xor_sync butterfly tree, like warp::reduce_sum() (5 shuffles for 32 -// lanes). -__device__ __forceinline__ uint32_t warp_reduce_sum_full(uint32_t v) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) - uint32_t r; - asm("redux.sync.add.u32 %0, %1, 0xFFFFFFFF;" : "=r"(r) : "r"(v)); - return r; -#else - #pragma unroll - for (uint32_t mask = kWarpSize >> 1; mask > 0; mask >>= 1) { - v += __shfl_xor_sync(0xFFFFFFFF, v, mask); - } - return v; -#endif -} - -// ============================================================================ -// Tie refinement (single CTA): 4-round radix-256 topK on the full FP32 ordered -// key Each round narrows by 8 bits until ties are fully resolved -// ============================================================================ - -template -__device__ void tie_handle(const Tie* ties, uint32_t num_ties, - uint32_t num_above, int32_t* output, void* _smem) { - struct TS { - alignas(128) uint32_t counter; - alignas(128) MatchBin match; - uint32_t histogram[RADIX]; - uint32_t warp_sum[kNumWarps]; - }; - auto* s = static_cast(_smem); - const auto tx = threadIdx.x; - const auto li = tx % kWarpSize, wi = tx / kWarpSize; - - // Each thread loads one tie element. - const bool has = tx < num_ties; - const auto tie = has ? ties[tx] : Tie{0, 0.0f}; - const uint32_t key = convert_to_uint32_v2(tie.score); - - bool active = has; // tracks whether this thread's tie is still a candidate. - uint32_t remain = - TopK - num_above; // decreases each round as ties are resolved. - uint32_t wpos = TopK; // wpos will hold the final output position. - s->counter = 0; - __syncthreads(); - - // The 4-round radix loop - each round narrows by 8 bits until ties are fully - // resolved -#pragma unroll - for (int r = 0; r < 4; r++) { - uint32_t sh = 24 - r * 8; // round 0: bits 31-24, round 1: 23-16, etc. - uint32_t bin = (key >> sh) & 0xFF; // this tie's 8-bit bin for this round - - // Step 1: Build 256-bin histogram. - if (tx < RADIX) s->histogram[tx] = 0; - __syncthreads(); - if (active) atomicAdd(&s->histogram[bin], 1); - __syncthreads(); - - // Step 2: Prefix scan to find threshold - uint32_t hv = 0, wi2 = 0; - if (tx < RADIX) { - hv = s->histogram[tx]; - wi2 = warp_inclusive_sum(li, hv); - if (li == kWarpSize - 1) s->warp_sum[wi] = wi2; - } - __syncthreads(); - - if (tx < RADIX) { - auto tmp = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; - auto tot = warp_reduce_sum_full(tmp); - auto inter = warp_reduce_sum_full(li < wi ? tmp : 0); - auto above = tot - (inter + wi2); - if (above < remain && above + hv >= remain) { - s->match = {tx, above, remain - above}; - } - } - __syncthreads(); - - // Step 3: Scatter - auto [thr, na, _] = s->match; // threshold bin, num above, unused - if (active) { - if (bin > thr) { - wpos = num_above + - atomicAdd(&s->counter, 1); // above -> place in output directly - active = false; - } else if (bin < thr) - active = false; // below -> discard - else if (r == 3) - wpos = TopK - atomicAdd(&s->match.equal_count, - -1u); // last round: place remaining - } - remain -= na; - if (!remain) break; // all ties resolved early - } - // Final write - if (wpos < TopK) output[wpos] = tie.idx; -} - -// Extended tie_handle for TopK > kBlockSize (e.g. TopK=2048). -// tie_handle assumes 1 tie per thread (max 1024). -// This version handles 2 ties per thread via kPerThread=2 -template -__device__ void tie_handle_large(const Tie* ties, uint32_t num_ties, - uint32_t num_above, int32_t* output, - void* _smem) { - static_assert(TopK > kBlockSize); - struct TS { - alignas(128) uint32_t counter; - alignas(128) MatchBin match; - uint32_t histogram[RADIX]; - uint32_t warp_sum[kNumWarps]; - }; - auto* s = static_cast(_smem); - const auto tx = threadIdx.x; - const auto li = tx % kWarpSize; - const auto wi = tx / kWarpSize; - - constexpr uint32_t kPerThread = (TopK + kBlockSize - 1) / kBlockSize; - Tie my_ties[kPerThread]; - uint32_t keys[kPerThread]; - bool active[kPerThread]; - - for (uint32_t e = 0; e < kPerThread; e++) { - uint32_t idx = e * kBlockSize + tx; - if (idx < num_ties) { - my_ties[e] = ties[idx]; - keys[e] = convert_to_uint32_v2(ties[idx].score); - active[e] = true; - } else { - my_ties[e] = {0, 0.0f}; - keys[e] = 0; - active[e] = false; - } - } - - uint32_t remain = TopK - num_above; - s->counter = 0; - __syncthreads(); - - for (int r = 0; r < 4; r++) { - uint32_t sh = 24 - r * 8; - if (tx < RADIX) { - s->histogram[tx] = 0; - } - __syncthreads(); - - for (uint32_t e = 0; e < kPerThread; e++) { - if (active[e]) { - atomicAdd(&s->histogram[(keys[e] >> sh) & 0xFF], 1); - } - } - __syncthreads(); - - uint32_t hv = 0; - if (tx < RADIX) { - hv = s->histogram[tx]; - auto wi2 = warp_inclusive_sum(li, hv); - if (li == kWarpSize - 1) { - s->warp_sum[wi] = wi2; - } - } - __syncthreads(); - if (tx < RADIX) { - auto tmp2 = (li < RADIX / kWarpSize) ? s->warp_sum[li] : 0; - auto total = warp_reduce_sum_full(tmp2); - auto inter = warp_reduce_sum_full(li < wi ? tmp2 : 0); - auto wi2 = warp_inclusive_sum(li, hv); - auto above = total - (inter + wi2); - if (above < remain && above + hv >= remain) { - s->match = { - .bin = tx, .above_count = above, .equal_count = remain - above}; - } - } - __syncthreads(); - - auto thr = s->match.bin; - auto na = s->match.above_count; - - for (uint32_t e = 0; e < kPerThread; e++) { - if (!active[e]) { - continue; - } - uint32_t bin = (keys[e] >> sh) & 0xFF; - if (bin > thr) { - uint32_t wpos = num_above + atomicAdd(&s->counter, 1); - if (wpos < TopK) { - output[wpos] = my_ties[e].idx; - } - active[e] = false; - } else if (bin < thr) { - active[e] = false; - } else if (r == 3) { - uint32_t wpos = TopK - atomicAdd(&s->match.equal_count, -1u); - if (wpos < TopK) { - output[wpos] = my_ties[e].idx; - } - } - } - - num_above += na; - remain -= na; - __syncthreads(); - s->counter = 0; - __syncthreads(); - } -} - -// ============================================================================ -// Register-based single-CTA fast path for seq_len <= 16384 -// 4 float4 per thread × 1024 threads = 16384 elements max -// Uses 4096-bin (12-bit) histogram for better precision -// ============================================================================ - -template -struct Histogram4096Smem { - static constexpr uint32_t HIST_BINS = 1 << HIST_BITS; - static constexpr uint32_t TIE_CAPACITY = TopK > kMaxTies ? TopK : kMaxTies; - alignas(128) uint32_t counter_gt; - alignas(128) uint32_t counter_eq; - MatchBin match; - uint32_t warp_sum[kNumWarps]; - union { - uint32_t histogram[HIST_BINS]; - Tie tie_buffer[TIE_CAPACITY]; - }; -}; - -template -__device__ void histogram_4096_topk(const float* __restrict__ scores, - int32_t* __restrict__ output, - uint32_t length, void* _smem) { - constexpr uint32_t HIST_BINS = 1 << HIST_BITS; - constexpr uint32_t ITEMS_PER_THREAD = HIST_BINS / kBlockSize; - static_assert(HIST_BINS >= kBlockSize, - "HIST_BITS must give >= kBlockSize bins"); - - using Smem = Histogram4096Smem; - auto* smem = static_cast(_smem); - const auto tx = threadIdx.x; - const auto lane_id = tx % kWarpSize; - const auto warp_id = tx / kWarpSize; - - // Phase 1: Load all data into RF + build histogram - float4 - vecs[VECS_PER_THREAD]; // 4 vectors x 4 floats = 16 elements per thread - if constexpr (ITEMS_PER_THREAD >= 4) { - // Zero the histogram (SMEM writes) - for (uint32_t i = 0; i < ITEMS_PER_THREAD / 4; i++) - reinterpret_cast( - smem->histogram)[tx * (ITEMS_PER_THREAD / 4) + i] = - make_uint4(0, 0, 0, 0); - } else { - if (tx < HIST_BINS) smem->histogram[tx] = 0; - } - if (tx == 0) { - smem->counter_gt = 0; - smem->counter_eq = 0; - } - if constexpr (UsePredicatedLoads) { - const bool row_aligned = (reinterpret_cast(scores) & 0xFu) == 0; -#pragma unroll - for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { - const uint32_t base = (tx + v * kBlockSize) * 4; - if (base < length) { - if (row_aligned && base + 3 < length) { - vecs[v] = *reinterpret_cast(scores + base); - } else { - load_float4_predicated(scores + base, static_cast(base), - static_cast(length), vecs[v].x, vecs[v].y, - vecs[v].z, vecs[v].w); - } - } - } - } else { -#pragma unroll - for (uint32_t v = 0; v < VECS_PER_THREAD; v++) { - const uint32_t base = (tx + v * kBlockSize) * 4; - if (base < length) { - vecs[v] = *reinterpret_cast(scores + base); - } - } - } - __syncthreads(); - - // Build histogram from RF via atomic adds into the shared histogram - bool done = false; -#pragma unroll - for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { - const float* elems = reinterpret_cast(&vecs[v]); -#pragma unroll - for (uint32_t e = 0; e < 4 && !done; e++) { - const uint32_t idx = (tx + v * kBlockSize) * 4 + e; - if (idx >= length) { - done = true; - } else { - atomicAdd(&smem->histogram[extract_coarse_bin_N(elems[e])], - 1); - } - } - } - __syncthreads(); - - // Phase 2: Prefix scan to find threshold bin - // Multi-element scan (4096 bins: 4 per thread) - uint32_t orig[ITEMS_PER_THREAD]; - uint32_t local_sum = 0; - - // Step 1: Each thread sums its 4 bins -#pragma unroll - for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { - orig[i] = smem->histogram[tx * ITEMS_PER_THREAD + i]; - local_sum += orig[i]; - } - - // Step 2: Warp-level inclusive prefix sum on local_sum - const auto warp_inc = warp_inclusive_sum(lane_id, local_sum); - if (lane_id == kWarpSize - 1) smem->warp_sum[warp_id] = warp_inc; - __syncthreads(); - - // Step 3: Inter-warp prefix across warp sums. - const auto tmp = smem->warp_sum[lane_id]; - uint32_t prefix = warp_reduce_sum_full( - lane_id < warp_id ? tmp : 0); // sum of all prior warps - prefix += - warp_inc - local_sum; // exclusive prefix within this thread's position - - // Step 4: Find threshold - scan 4 bins, accumulate prefix -#pragma unroll - for (uint32_t i = 0; i < ITEMS_PER_THREAD; i++) { - prefix += orig[i]; - const auto above = length - prefix; // elements in bins ABOVE this one - if (above < TopK && above + orig[i] >= TopK) { - smem->match = {.bin = tx * ITEMS_PER_THREAD + i, - .above_count = above, - .equal_count = orig[i]}; - } - } - - __syncthreads(); - - // Phase 3: Scatter from registers - const auto [thr_bin, num_above, num_equal] = smem->match; - const bool need_tie = (num_equal + num_above > TopK); - - done = false; -#pragma unroll - for (uint32_t v = 0; v < VECS_PER_THREAD && !done; v++) { - const float* elems = reinterpret_cast(&vecs[v]); -#pragma unroll - for (uint32_t e = 0; e < 4 && !done; e++) { - const uint32_t idx = (tx + v * kBlockSize) * 4 + e; - if (idx >= length) { - done = true; - } else { - const uint32_t bin = extract_coarse_bin_N(elems[e]); - if (bin > thr_bin) { - output[atomicAdd(&smem->counter_gt, 1)] = - idx; // above -> output directly - } else if (bin == thr_bin) { - const auto pos = atomicAdd(&smem->counter_eq, 1); - if (!need_tie) { - if (pos + num_above < TopK) { - output[pos + num_above] = idx; // all fit - } - } else { - if (pos < TopK) { - smem->tie_buffer[pos] = {idx, elems[e]}; // store for refirement - } - } - } - // else: bin < thr_bin - discard (not in top-k) - } - } - } - - // Phase 4: Tie-breaking - if (!need_tie) return; - __syncthreads(); - - // Fast warp-ballot tie-breaking for small tie counts - const uint32_t num_ties = min(num_equal, static_cast(TopK)); - const uint32_t topk_remain = - TopK - num_above; // pick exactly remaining elements to fill topK - - auto is_greater = [](const Tie& a, const Tie& b) { - return (a.score > b.score) || (a.score == b.score && a.idx < b.idx); - }; - - if (num_ties <= kWarpSize) { - // <=32 ties - Use warp ballot - // All-to-all comparison in one __ballot_sync. 32 ties x 32 warps = 1024 - // comparisons in one instruction per warp. O(1) work. - const auto lane_id = tx % kWarpSize; - const auto warp_id = tx / kWarpSize; - if (lane_id >= num_ties || warp_id >= num_ties) return; - const uint32_t mask = (1ull << num_ties) - 1u; - const auto tie = smem->tie_buffer[lane_id]; // each lane holds one tie - const auto target = - smem->tie_buffer[warp_id]; // each warp evaluates one candidate - const bool pred = - is_greater(tie, target); // compare all ties against target - const auto rank = static_cast( - __popc(__ballot_sync(mask, pred))); // count how many are greater - if (lane_id == 0 && rank < topk_remain) { - output[num_above + rank] = target.idx; // place at correct position - } - } else if (num_ties <= - kWarpSize * - 2) { // TODO (roberto): try to refactor this with <=32 case - // Same idea but each thread handles 2 tie elements - const auto lane_id = tx % kWarpSize; - const auto warp_id = tx / kWarpSize; - const auto lane1 = lane_id + kWarpSize; - const auto warp1 = warp_id + kWarpSize; - const auto invalid = Tie{0xFFFFFFFF, -__FLT_MAX__}; - const auto tie0 = smem->tie_buffer[lane_id]; - const auto tie1 = lane1 < num_ties ? smem->tie_buffer[lane1] : invalid; - if (warp_id < num_ties) { - const auto target = smem->tie_buffer[warp_id]; - const auto r0 = - __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); - const auto r1 = - __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); - if (lane_id == 0 && r0 + r1 < topk_remain) - output[num_above + r0 + r1] = target.idx; - } - if (warp1 < num_ties) { - const auto target = smem->tie_buffer[warp1]; - const auto r0 = - __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie0, target))); - const auto r1 = - __popc(__ballot_sync(0xFFFFFFFF, is_greater(tie1, target))); - if (lane_id == 0 && r0 + r1 < topk_remain) - output[num_above + r0 + r1] = target.idx; - } - } else { - // Large tie count: fall back to 4-round radix-256 sort - if constexpr (TopK <= kBlockSize) { - tie_handle(smem->tie_buffer, num_ties, num_above, output, smem); - } else { - tie_handle_large(smem->tie_buffer, num_ties, num_above, output, - smem); - } - } -} - -template -__device__ __noinline__ void histogram_4096_topk_predicated( - const float* __restrict__ scores, int32_t* __restrict__ output, - uint32_t length, void* _smem) { - histogram_4096_topk(scores, output, - length, _smem); -} - -} // namespace topk_histogram_4096 -} // namespace vllm - -#endif // TOPK_HISTOGRAM_4096_CUH_ diff --git a/prime_kernels/indexed_attention/selection.py b/prime_kernels/indexed_attention/selection.py index b214844..ea02410 100644 --- a/prime_kernels/indexed_attention/selection.py +++ b/prime_kernels/indexed_attention/selection.py @@ -1,179 +1,270 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import tilelang +import torch +from tilelang import language as T -from __future__ import annotations +SCORE_WORKSPACE_BYTES = 1024**3 +RADIX_BITS = 8 +RADIX_SIZE = 1 << RADIX_BITS -import math -import torch -import triton -import triton.language as tl +@tilelang.jit( + out_idx=[-2, -1], + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + }, +) +def selection_score_kernel( + num_query_heads: int, + head_dim: int, + block_keys: int = 64, + key_tiles_per_program: int = 8, + threads: int = 64, +): + query_tokens = T.dynamic("query_tokens") + key_blocks = T.dynamic("key_blocks") + head_tile = max(tilelang.math.next_power_of_2(num_query_heads), 16) -# A 32K-token packed batch produces at most a 1 GiB score workspace. -SCORE_WORKSPACE_BYTES = 1024**3 -TOPK_WORKSPACE_BYTES = 1024**2 - - -@triton.jit -def _selection_scores_kernel( - query_ptr, - key_ptr, - starts_ptr, - ends_ptr, - visible_blocks_ptr, - scores_ptr, - stride_query_row, - stride_query_head, - stride_query_dim, - stride_key_row, - stride_key_dim, - stride_scores_row, - rows, - columns, - key_blocks, - score_divisor, - NUM_HEADS: tl.constexpr, - HEAD_DIM: tl.constexpr, - BLOCK_N: tl.constexpr, - BLOCK_D: tl.constexpr, - TILES_PER_PROGRAM: tl.constexpr, - STAGES: tl.constexpr, - MAX_N: tl.constexpr, -) -> None: - row = tl.program_id(0) - dimensions = tl.arange(0, BLOCK_D) - heads = tl.arange(0, MAX_N) - start = tl.load(starts_ptr + row) - end = tl.load(ends_ptr + row) - visible = end - start - if tl.program_id(1) == 0: - tl.store(visible_blocks_ptr + row, visible) - - tile_start = tl.program_id(1) * TILES_PER_PROGRAM - if tile_start * BLOCK_N >= visible: - return - tile_end = tl.minimum(tile_start + TILES_PER_PROGRAM, tl.cdiv(visible, BLOCK_N)) - tile_end = tl.minimum(tile_end, tl.cdiv(columns, BLOCK_N)) - - query = tl.load( - query_ptr - + row * stride_query_row - + heads[None, :] * stride_query_head - + dimensions[:, None] * stride_query_dim, - mask=(heads[None, :] < NUM_HEADS) & (dimensions[:, None] < HEAD_DIM), - other=0.0, - ) - column_offsets = tl.arange(0, BLOCK_N) - for tile in tl.range(tile_start, tile_end, num_stages=STAGES): - columns_in_tile = tile * BLOCK_N + column_offsets - key_rows = start + columns_in_tile - live = (columns_in_tile < visible) & (key_rows < key_blocks) - keys = tl.load( - key_ptr + key_rows[:, None].to(tl.int64) * stride_key_row + dimensions[None, :] * stride_key_dim, - mask=live[:, None] & (dimensions[None, :] < HEAD_DIM), - other=0.0, - eviction_policy="evict_first", - ) - head_scores = tl.dot(keys, query, out_dtype=tl.float32) - head_scores = tl.where(heads[None, :] < NUM_HEADS, tl.maximum(head_scores, 0.0), 0.0) - scores = tl.sum(head_scores, axis=1) / score_divisor - tl.store( - scores_ptr + row * stride_scores_row + columns_in_tile, - tl.where(live, scores, -float("inf")), - mask=columns_in_tile < columns, - ) + query_shape = [query_tokens, num_query_heads, head_dim] + key_shape = [key_blocks, head_dim] + bounds_shape = [query_tokens] + scores_shape = [query_tokens, key_blocks] + + @T.prim_func + def kernel( + query: T.Tensor(query_shape, T.bfloat16), + key: T.Tensor(key_shape, T.bfloat16), + block_starts: T.Tensor(bounds_shape, T.int32), + block_ends: T.Tensor(bounds_shape, T.int32), + scores: T.Tensor(scores_shape, T.float32), + visible_block_counts: T.Tensor(bounds_shape, T.int32), + ): + with T.Kernel( + T.ceildiv(key_blocks, block_keys * key_tiles_per_program), + query_tokens, + threads=threads, + ) as (key_group, query_token): + query_shared = T.alloc_shared([head_tile, head_dim], T.bfloat16) + key_shared = T.alloc_shared([block_keys, head_dim], T.bfloat16) + head_scores = T.alloc_fragment([block_keys, head_tile], T.float32) + combined_scores = T.alloc_fragment([block_keys], T.float32) + + start = T.alloc_var(T.int32) + visible_count = T.alloc_var(T.int32) + start = block_starts[query_token] + visible_count = block_ends[query_token] - start + + if key_group == 0: + visible_block_counts[query_token] = visible_count + + for head, dim in T.Parallel(head_tile, head_dim): + query_shared[head, dim] = T.if_then_else( + head < num_query_heads, + query[query_token, head, dim], + 0, + ) + + for key_tile in T.serial(key_tiles_per_program): + first_local_key = (key_group * key_tiles_per_program + key_tile) * block_keys + if first_local_key < visible_count: + for local_key, dim in T.Parallel(block_keys, head_dim): + key_index = start + first_local_key + local_key + key_shared[local_key, dim] = T.if_then_else( + first_local_key + local_key < visible_count and key_index < key_blocks, + key[key_index, dim], + 0, + ) + + T.gemm( + key_shared, + query_shared, + head_scores, + transpose_B=True, + clear_accum=True, + policy=T.GemmWarpPolicy.FullRow, + ) + for local_key, head in T.Parallel(block_keys, head_tile): + head_scores[local_key, head] = T.if_then_else( + head < num_query_heads, + T.max(head_scores[local_key, head], 0), + 0, + ) + T.reduce_sum(head_scores, combined_scores, dim=1) + for local_key in T.Parallel(block_keys): + if first_local_key + local_key < visible_count: + scores[query_token, first_local_key + local_key] = combined_scores[local_key] * ( + head_dim**-0.5 + ) + return kernel -def selection_scores( + +@tilelang.jit(out_idx=[-1]) +def radix_select_kernel(topk: int, threads: int = 256): + rows = T.dynamic("rows") + columns = T.dynamic("columns") + + @T.prim_func + def kernel( + scores: T.Tensor([rows, columns], T.float32), + visible_block_counts: T.Tensor([rows], T.int32), + block_starts: T.Tensor([rows], T.int32), + selected_blocks: T.Tensor([rows, topk], T.int32), + ): + with T.Kernel(rows, threads=threads) as row: + thread = T.get_thread_binding() + histogram = T.alloc_shared([RADIX_SIZE], T.int32) + threshold_prefix = T.alloc_shared([1], T.uint32) + threshold_prefix_mask = T.alloc_shared([1], T.uint32) + remaining_count = T.alloc_shared([1], T.int32) + threshold_digit = T.alloc_shared([1], T.int32) + threshold_count = T.alloc_shared([1], T.int32) + greater_digit_count = T.alloc_shared([1], T.int32) + output_count = T.alloc_shared([1], T.int32) + + score_bits = T.alloc_var(T.uint32) + digit = T.alloc_var(T.int32) + running_count = T.alloc_var(T.int32) + bin_count = T.alloc_var(T.int32) + output_position = T.alloc_var(T.int32) + visible_count = T.alloc_var(T.int32) + column = T.alloc_var(T.int32) + + visible_count = T.min(visible_block_counts[row], columns) + for output_position in T.Parallel(topk): + selected_blocks[row, output_position] = columns + if thread == 0: + threshold_prefix[0] = 0 + threshold_prefix_mask[0] = 0 + remaining_count[0] = T.min(topk, visible_count) + threshold_count[0] = 0 + T.sync_threads() + + for radix_pass in T.serial(32 // RADIX_BITS): + if remaining_count[0] > 0: + T.fill(histogram, 0) + T.sync_threads() + for column_group in T.serial(T.ceildiv(columns, threads)): + column = column_group * threads + thread + if column < visible_count: + score_bits = T.reinterpret(scores[row, column], T.uint32) + if (score_bits & threshold_prefix_mask[0]) == threshold_prefix[0]: + digit = T.cast( + (score_bits >> (32 - RADIX_BITS * (radix_pass + 1))) & (RADIX_SIZE - 1), + T.int32, + ) + T.atomic_add(histogram[digit], 1) + T.sync_threads() + + if thread == 0: + running_count = 0 + threshold_digit[0] = 0 + threshold_count[0] = 0 + greater_digit_count[0] = 0 + for digit_offset in T.serial(RADIX_SIZE): + digit = RADIX_SIZE - 1 - digit_offset + bin_count = histogram[digit] + if running_count < remaining_count[0] and running_count + bin_count >= remaining_count[0]: + threshold_digit[0] = digit + threshold_count[0] = bin_count + greater_digit_count[0] = running_count + running_count += bin_count + remaining_count[0] -= greater_digit_count[0] + threshold_prefix[0] |= T.cast(threshold_digit[0], T.uint32) << ( + 32 - RADIX_BITS * (radix_pass + 1) + ) + threshold_prefix_mask[0] |= T.cast(RADIX_SIZE - 1, T.uint32) << ( + 32 - RADIX_BITS * (radix_pass + 1) + ) + T.sync_threads() + + if thread == 0: + output_count[0] = 0 + T.sync_threads() + for column_group in T.serial(T.ceildiv(columns, threads)): + column = column_group * threads + thread + if column < visible_count: + score_bits = T.reinterpret(scores[row, column], T.uint32) + if score_bits > threshold_prefix[0]: + output_position = T.atomic_add(output_count[0], 1, return_prev=True) + selected_blocks[row, output_position] = block_starts[row] + column + T.sync_threads() + if threshold_count[0] == remaining_count[0]: + for column_group in T.serial(T.ceildiv(columns, threads)): + column = column_group * threads + thread + if column < visible_count: + score_bits = T.reinterpret(scores[row, column], T.uint32) + if score_bits == threshold_prefix[0]: + output_position = T.atomic_add(output_count[0], 1, return_prev=True) + if output_position < topk: + selected_blocks[row, output_position] = block_starts[row] + column + elif thread == 0: + output_position = output_count[0] + for tied_column in T.serial(columns): + if tied_column < visible_count and output_position < topk: + score_bits = T.reinterpret(scores[row, tied_column], T.uint32) + if score_bits == threshold_prefix[0]: + selected_blocks[row, output_position] = block_starts[row] + tied_column + output_position += 1 + + return kernel + + +def compute_selection_scores( query: torch.Tensor, key: torch.Tensor, - starts: torch.Tensor, - ends: torch.Tensor, + block_starts: torch.Tensor, + block_ends: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: - columns = key.shape[0] - scores = torch.empty((query.shape[0], columns), dtype=torch.float32, device=query.device) - visible_blocks = torch.empty(query.shape[0], dtype=torch.int32, device=query.device) - if not query.shape[0] or not columns: - return scores, visible_blocks - - block_n = 64 - block_d = max(16, triton.next_power_of_2(query.shape[2])) - max_n = max(16, triton.next_power_of_2(query.shape[1])) - tiles_per_program = 1 if query.shape[0] <= 32 else 8 - _selection_scores_kernel[(query.shape[0], triton.cdiv(columns, block_n * tiles_per_program))]( - query, - key, - starts, - ends, - visible_blocks, - scores, - query.stride(0), - query.stride(1), - query.stride(2), - key.stride(0), - key.stride(1), - scores.stride(0), - query.shape[0], - columns, - key.shape[0], - math.sqrt(query.shape[2]), - NUM_HEADS=query.shape[1], - HEAD_DIM=query.shape[2], - BLOCK_N=block_n, - BLOCK_D=block_d, - TILES_PER_PROGRAM=tiles_per_program, - STAGES=2, - MAX_N=max_n, - num_warps=2, + if not query.shape[0] or not key.shape[0]: + return ( + torch.empty((query.shape[0], key.shape[0]), dtype=torch.float32, device=query.device), + torch.empty(query.shape[0], dtype=torch.int32, device=query.device), + ) + return selection_score_kernel(query.shape[1], query.shape[2])( + query.contiguous(), + key.contiguous(), + block_starts.contiguous(), + block_ends.contiguous(), ) - return scores, visible_blocks @torch.library.custom_op("prime_kernels::select_indexed_blocks", mutates_args=()) def select_indexed_blocks( query: torch.Tensor, key: torch.Tensor, - starts: torch.Tensor, - ends: torch.Tensor, + block_starts: torch.Tensor, + block_ends: torch.Tensor, topk: int, ) -> torch.Tensor: num_blocks = key.shape[0] + if not query.shape[0] or not topk: + return torch.empty((query.shape[0], topk), dtype=torch.int32, device=query.device) if not num_blocks: - return torch.zeros(query.shape[0], topk, dtype=torch.int32, device=query.device) + return torch.zeros((query.shape[0], topk), dtype=torch.int32, device=query.device) rows_per_chunk = max(1, SCORE_WORKSPACE_BYTES // (num_blocks * torch.float32.itemsize)) selected_chunks = [] - workspace = torch.empty(TOPK_WORKSPACE_BYTES, dtype=torch.uint8, device=query.device) + contiguous_key = key.contiguous() for query_chunk, start_chunk, end_chunk in zip( query.split(rows_per_chunk), - starts.split(rows_per_chunk), - ends.split(rows_per_chunk), + block_starts.split(rows_per_chunk), + block_ends.split(rows_per_chunk), strict=True, ): - scores, visible_blocks = selection_scores( - query_chunk, - key, - start_chunk, - end_chunk, - ) - selected = torch.full( - (query_chunk.shape[0], topk), - -1, - dtype=torch.int32, - device=query.device, + scores, visible_block_counts = selection_score_kernel(query.shape[1], query.shape[2])( + query_chunk.contiguous(), + contiguous_key, + start_chunk.contiguous(), + end_chunk.contiguous(), ) - torch.ops.prime_indexed_attention.persistent_topk( - scores, - visible_blocks, - selected, - workspace, - topk, - num_blocks, + selected_chunks.append( + radix_select_kernel(topk)( + scores, + visible_block_counts, + start_chunk.contiguous(), + ) ) - valid = (selected >= 0) & (selected < visible_blocks[:, None]) - selected.add_(start_chunk[:, None]) - selected.masked_fill_(~valid, num_blocks) - selected_chunks.append(selected) if len(selected_chunks) == 1: return selected_chunks[0] @@ -184,8 +275,8 @@ def select_indexed_blocks( def select_indexed_blocks_fake( query: torch.Tensor, key: torch.Tensor, - starts: torch.Tensor, - ends: torch.Tensor, + block_starts: torch.Tensor, + block_ends: torch.Tensor, topk: int, ) -> torch.Tensor: return query.new_empty((query.shape[0], topk), dtype=torch.int32) diff --git a/prime_kernels/kernels.toml b/prime_kernels/kernels.toml index 557c55e..3383b7c 100644 --- a/prime_kernels/kernels.toml +++ b/prime_kernels/kernels.toml @@ -32,12 +32,9 @@ arch = ["10.0"] [indexed_attention] description = "Training forward and backward for token-indexed grouped-query attention" -ops = "prime_indexed_attention" -sources = ["csrc/topk.cu"] -include-dirs = ["csrc"] -requires = ["tilelang", "triton"] +python-only = true +requires = ["tilelang"] arch = ["8.0", "9.0", "10.0"] -cxx-std = 20 # rmsnorm is not built yet: only its sources are committed. Uncomment the table below to # put it back into the build (and to make the registry report on it). diff --git a/setup.py b/setup.py index 471910d..06c81a4 100644 --- a/setup.py +++ b/setup.py @@ -98,10 +98,8 @@ def _extension(kernel) -> CUDAExtension: # Listed explicitly: the kernel folders carry C++/CUDA sources next to their Python, and # only the Python surface plus the compiled extension belongs in the wheel. packages=["prime_kernels", *(f"prime_kernels.{name}" for name in kernels)], - include_package_data=False, package_data={ "prime_kernels": ["kernels.toml"], - "prime_kernels.indexed_attention": ["LICENSE.vllm"], "prime_kernels.mxfp8_moe": ["LICENSE.torchao"], }, ext_modules=extensions,