83 lines
2.9 KiB
Python
83 lines
2.9 KiB
Python
|
"""Python launcher for the indexer score+topk kernel.
|
||
|
|
|
||
|
|
Provides run_indexer_score_topk() which takes FP32 query tensors
|
||
|
|
and an IndexerView from the cache, runs the fused score + ReLU +
|
||
|
|
weighted sum + top-k kernel, and returns [T, top_k] compressed
|
||
|
|
entry indices.
|
||
|
|
|
||
|
|
Phase 1: FP32 dot products. Correct, testable.
|
||
|
|
Phase 2: FP4 tcgen05 MMA swap (optimization on known-correct base).
|
||
|
|
"""
|
||
|
|
from __future__ import annotations
|
||
|
|
import os
|
||
|
|
import torch
|
||
|
|
from typing import TYPE_CHECKING
|
||
|
|
|
||
|
|
if TYPE_CHECKING:
|
||
|
|
from dsv4.cache.handle import IndexerView
|
||
|
|
|
||
|
|
_kernel_module = None
|
||
|
|
|
||
|
|
|
||
|
|
def _get_kernel_module():
|
||
|
|
global _kernel_module
|
||
|
|
if _kernel_module is not None:
|
||
|
|
return _kernel_module
|
||
|
|
kernel_dir = os.path.join(os.path.dirname(__file__), "..", "cuda")
|
||
|
|
_kernel_module = torch.utils.cpp_extension.load(
|
||
|
|
name="indexer_score_topk",
|
||
|
|
sources=[os.path.join(kernel_dir, "indexer_score_topk.cu")],
|
||
|
|
extra_cuda_cflags=["-O3", "--generate-code=arch=compute_100a,code=[sm_100a]"],
|
||
|
|
verbose=False,
|
||
|
|
)
|
||
|
|
return _kernel_module
|
||
|
|
|
||
|
|
|
||
|
|
def run_indexer_score_topk(
|
||
|
|
q_I: torch.Tensor, # [T, n_heads * head_dim] BF16 — indexer queries
|
||
|
|
w_h: torch.Tensor, # [T, n_heads] FP32 — per-head weights
|
||
|
|
indexer_view: "IndexerView",
|
||
|
|
num_heads: int,
|
||
|
|
head_dim: int,
|
||
|
|
top_k: int,
|
||
|
|
entries_per_block: int,
|
||
|
|
) -> torch.Tensor:
|
||
|
|
"""Returns [T, top_k] int32 of selected compressed entry indices.
|
||
|
|
|
||
|
|
The kernel computes:
|
||
|
|
I[t,s] = Σ_h w_h[t,h] · ReLU(q_I[t,h] · K^IComp[s,h])
|
||
|
|
topk_indices = argtopk(I[t,:], k=top_k)
|
||
|
|
|
||
|
|
q_I is passed as BF16 and dequantized to FP32 before the kernel.
|
||
|
|
The indexer keys are stored FP4 in the cache and dequantized
|
||
|
|
inside the kernel.
|
||
|
|
"""
|
||
|
|
mod = _get_kernel_module()
|
||
|
|
T = q_I.shape[0]
|
||
|
|
|
||
|
|
# Dequantize q_I from BF16 to FP32 and reshape to [T, n_heads, head_dim]
|
||
|
|
q_I_f32 = q_I.float().reshape(T, num_heads, head_dim).contiguous()
|
||
|
|
|
||
|
|
# Compute valid lens from block_lens
|
||
|
|
valid_lens = indexer_view.block_lens * entries_per_block # [B] int32
|
||
|
|
# We need per-query valid lens. block_lens is [B] where B = batch.
|
||
|
|
# For a single request, this is just the one value.
|
||
|
|
# For batched, repeat across tokens belonging to the same request.
|
||
|
|
# Simplification: assume T == B for now (one token per request in decode).
|
||
|
|
if valid_lens.shape[0] != T:
|
||
|
|
# Prefill: T > B. We need to map tokens to requests.
|
||
|
|
# For now, broadcast the first request's valid_lens.
|
||
|
|
# TODO: proper per-token valid_lens from request_ids mapping.
|
||
|
|
valid_lens = valid_lens[:1].expand(T).contiguous()
|
||
|
|
|
||
|
|
out = torch.full((T, top_k), -1, dtype=torch.int32, device=q_I.device)
|
||
|
|
|
||
|
|
mod.indexer_score_topk_fp32(
|
||
|
|
q_I_f32, w_h,
|
||
|
|
indexer_view.keys_fp4, indexer_view.scale, indexer_view.global_scale,
|
||
|
|
indexer_view.block_table, valid_lens,
|
||
|
|
out,
|
||
|
|
num_heads, head_dim, top_k, entries_per_block,
|
||
|
|
)
|
||
|
|
return out
|