#029

topk_filter

fp32 vllm · decode · Sampler.forward · importance 3.5%

Reference Implementation

reference.py

import torch
import torch.nn as nn

class Model(nn.Module):

    def __init__(self, top_k: int=50) -> None:
        super().__init__()
        self.top_k = top_k

    def forward(self, logits: torch.Tensor) -> torch.Tensor:
        top_k = min(self.top_k, logits.shape[-1])
        threshold = torch.topk(logits, top_k, dim=-1).values[..., -1:]
        return logits.masked_fill(logits < threshold, float('-inf'))

Shapes

T_SOL hardware:

#	seq_len	hidden_size	T_SOL(XPU-A)	T_Prod	S
0	1	151936	0.23 us	299.70 us	0.1%
1	1	152064	0.23 us	328.60 us	0.1%
2	922	2048	2.85 us	187.40 us	1.5%
3	3	4002	0.02 us	31.30 us	0.1%
4	4098	4096	25.34 us	805.10 us	3.1%
5	8179	4096	50.57 us	2.01 ms	2.5%
6	15381	4096	95.10 us	2.93 ms	3.2%
7	3	4215	0.02 us	30.70 us	0.1%

Input Generation

input.py

import torch


def _make_inputs(seq_len: int, hidden_size: int) -> dict[str, torch.Tensor]:
    logits = torch.randn(seq_len, hidden_size, dtype=torch.float32, device='cuda')
    return {'logits': logits}