lm-head-topk

The vocabulary projection fused with top-k selection, so the full logit vector is never written, loadable through kernels. The reference baseline is torch.topk on the materialized logits, matched index for index (torch.equal). Feeds logits-processor and spec-decode-ops; binary weights use the packed layout of binary-gemm.

The last operation of every decoded token turns a hidden vector into one score per vocabulary entry. At vocabulary 248,320 and hidden 5,120 that layer holds 1.27 billion weights, more than the rest of a 1-bit model's body, and it emits a quarter-million-wide vector of which sampling keeps a few dozen entries. This kernel computes the projection from packed binary weights and keeps only the exact top-k as it goes, so the score vector never exists and the model's largest layer runs from a fourteenth of the memory.

One pixel per vocabulary entry: the chunk scan sweeps 248,320 logits, each 2,048-entry chunk offering its own candidates, and the exact top-50 collapse to a sorted row, torch.equal to torch.topk on the materialized logits; 3.56 ms and 2,426 MB held for the bf16 head against 1.36 ms and 171 MB fused.

Usage

import torch
from kernels import get_kernel

lt = get_kernel("phanerozoic/lm-head-topk", version=1, trust_remote_code=True)

# binary head, the path this kernel is for
values, indices = lt.lm_head_topk_binary(act_q, act_scale, wq, wscale,
                                         bias=None, group_size=128, k=50)

# dense head, correctness reference
values, indices = lt.lm_head_topk_dense(hidden, lm_head.weight, k=50)

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Returns (values [M, k] fp32, indices [M, k] int64), sorted descending.

Method

The vocabulary is cut into 2,048-entry chunks. A block computes its chunk's scores into shared memory with one warp per vocabulary entry, so the 32 lanes stride a single weight row and the loads coalesce, then extracts that chunk's own top k by repeated block-wide argmax and writes only those k. A second pass merges the per-chunk candidates. Each chunk must offer k candidates rather than fewer, because the global top k can lie entirely inside one chunk; that is what makes the merge exact rather than a heuristic, and it is tested by planting every winner in one chunk. Ties resolve to the lower vocabulary index, matching torch.topk.

Measured

Vocabulary 248,320, hidden 5,120, k = 50. The comparison a deployer faces is a bf16 head with materialized logits against a binary head with fused selection:

The dense bf16 entry point exists as the correctness reference, not for speed: a dense head is 2.4 GB against 1 MB of logits, so there is nothing to save by not writing them, and cuBLAS runs the projection on tensor cores that a scalar GEMV cannot follow:

Correctness

• Indices are exactly torch's, torch.equal, across rows 1 to 8, hidden 512 to 2,048, vocabularies 8,192 to 128,256, and k 1 to 100, on both paths, with and without bias.
• Values agree to 1e-5 (dense) and 1e-2 (binary, against a dequantized reference); k = 1 equals torch.argmax; values return sorted descending.
• A ragged final chunk returns only valid indices and still matches torch; with every winner planted inside a single chunk, so a per-chunk heuristic would fail, the result still equals torch.topk.
• Deterministic: repeated calls are bitwise identical.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+.
• k at most 2,048 and at most the vocabulary size.
• Single-token kernel: every row re-reads the whole head, so the advantage is gone by four rows; route prefill elsewhere.
• Forward only; the binary path expects weights packed by binary-gemm's pack_weights.

References

Fused selection as it appears in speculative-decoding and sampling stacks; torch.topk tie and ordering semantics; the packed binary layout of phanerozoic/binary-gemm.

License

Apache-2.0.