菜单
Fuse categorical sampling directly into LM-head matrix multiplication to eliminate logits materialization. Use Gumbel noise during computation and hierarchical reduction to achieve 19% token-level latency reduction.
默认使用会先检查来源的 Prompt；你也可以切换为直接命令，或下载本地副本。
决定是否安装前，请先阅读 SKILL.md，以及 SkillsMP 当前展示的配套文件。
用 Codex 或 Claude 帮你安装 复制这段 Prompt，粘贴到 Codex、Claude 或其他助手里，让它检查 Skill 页面并帮你完成安装。
直接命令不会经过审查 Prompt；运行前请先检查来源。
npx skills add https://github.com/ADu2021/skillXiv --skill flash-sampling-efficient-decoding命令会保持在同一行。复制前请横向滚动并检查完整内容。
想先保存到本地？可下载 SkillsMP 当前能够提供的文件。
Convert arXiv papers into ready-to-use agent skills using category-aware extraction. First classifies the paper into one or more of 11 research categories, then applies a specialized extraction pipeline for each category — because different types of papers produce different types of usable knowledge. A single paper can yield multiple skills if it spans categories. Use this skill whenever the user wants to turn a paper into a skill, extract practical techniques from research, build a skill library from papers, convert arXiv papers into reusable agent instructions, or batch-process multiple papers into skills. Also trigger when someone asks about extracting actionable knowledge from papers, making research practical for LLM agents, or systematically converting academic contributions into structured agent capabilities.
Establish regret bounds for behavior cloning with discretized actions combining statistical error and quantization error terms. Prove smoothness requirements for safe quantizer design, show that learning-based quantizers fail these requirements, and propose model-based augmentation to reduce error dependence from H² to H.
Dynamically allocate LoRA ranks per-layer during fine-tuning instead of using fixed uniform ranks. Learn optimal rank for each layer and subject via variational framework with discretized exponential distribution, reducing memory footprint while maintaining fidelity and text-alignment.
基于 SOC 职业分类
正在显示 SKILL.md
| name | flash-sampling-efficient-decoding |
| title | FlashSampling: Fast and Memory-Efficient Exact Sampling |
| version | 0.0.2 |
| engine | skillxiv-v0.0.2-claude-opus-4.6 |
| license | MIT |
| url | https://arxiv.org/abs/2603.15854 |
| keywords | ["Sampling Efficiency","Exact Sampling","Gumbel Trick","Bandwidth Optimization","Token Generation"] |
| description | Fuse categorical sampling directly into LM-head matrix multiplication to eliminate logits materialization. Use Gumbel noise during computation and hierarchical reduction to achieve 19% token-level latency reduction. |
Sampling from the output logits is a hidden performance bottleneck in language model inference. Current approaches materialize large logits tensors in high-bandwidth memory (HBM), consuming significant memory bandwidth and requiring multiple GPU kernels after LM-head computation. FlashSampling solves this by fusing sampling directly into the LM-head matrix multiplication: logits are computed tile-by-tile on-chip, Gumbel noise is applied during computation, and only the maximum per-tile is tracked—enabling exact categorical sampling without materializing full logits. This achieves up to 19% reduction in per-token latency across modern GPUs.
The approach remains mathematically exact because argmax decomposes over partitions (max of maximums equals global maximum).
FlashSampling exploits the mathematical property that sampling from a categorical distribution can be reformulated as:
Standard Sampling:
logits = LM_head(hidden) # [vocab_size]
probs = softmax(logits)
token = categorical_sample(probs)
FlashSampling:
# Compute logits tile-by-tile, maintaining only running maximum
max_logit = -inf
max_idx = -1
for tile_idx in range(num_tiles):
tile_logits = LM_head_partial(hidden, tile_idx)
tile_logits += Gumbel_noise
tile_max_idx = argmax(tile_logits)
tile_max = tile_logits[tile_max_idx]
if tile_max > max_logit:
max_logit = tile_max
max_idx = tile_idx * tile_size + tile_max_idx
# Token with highest Gumbel-perturbed logit is sampled token
token = max_idx
Because argmax commutes with Gumbel perturbation, this is mathematically equivalent to exact sampling.
Start by implementing the basic tile-based sampling logic.
import torch
import torch.nn.functional F
numpy np
:
():
.vocab_size = vocab_size
.hidden_dim = hidden_dim
.tile_size = tile_size
.num_tiles = (vocab_size + tile_size - ) // tile_size
():
u = torch.rand_like(logits)
gumbel_noise = -torch.log(-torch.log(u + ) + )
perturbed = (logits + gumbel_noise) / temperature
torch.argmax(perturbed, dim=-)
() -> torch.Tensor:
batch_size = hidden.size()
best_indices = torch.zeros(batch_size, dtype=torch.long)
best_scores = torch.full((batch_size,), ())
tile_idx (.num_tiles):
start_vocab = tile_idx * .tile_size
end_vocab = (start_vocab + .tile_size, .vocab_size)
tile_weights = lm_head_weight[start_vocab:end_vocab, :]
tile_logits = torch.matmul(hidden, tile_weights.t())
u = torch.rand_like(tile_logits)
gumbel = -torch.log(-torch.log(u + ) + )
tile_scores = (tile_logits + gumbel) / temperature
tile_best_scores, tile_best_local_idx = torch.(tile_scores, dim=)
improved = tile_best_scores > best_scores
best_scores[improved] = tile_best_scores[improved]
best_indices[improved] = (
start_vocab + tile_best_local_idx[improved]
)
best_indices
() -> torch.Tensor:
group_vocab_size = .vocab_size // group_size
best_indices = torch.zeros(hidden.size(), dtype=torch.long)
best_scores = torch.full((hidden.size(),), ())
group_idx (group_size):
group_start = group_idx * group_vocab_size
group_end = (group_idx + ) * group_vocab_size
group_weights = lm_head_weight[group_start:group_end, :]
tile_in_group ((group_vocab_size + .tile_size - )
// .tile_size):
tile_start = group_start + tile_in_group * .tile_size
tile_end = (tile_start + .tile_size, group_end)
tile_weights = lm_head_weight[tile_start:tile_end, :]
tile_logits = torch.matmul(hidden, tile_weights.t())
u = torch.rand_like(tile_logits)
gumbel = -torch.log(-torch.log(u + ) + )
tile_scores = (tile_logits + gumbel) / temperature
tile_best_scores, tile_best_idx = torch.(tile_scores, dim=)
improved = tile_best_scores > best_scores
best_scores[improved] = tile_best_scores[improved]
best_indices[improved] = tile_start + tile_best_idx[improved]
best_indices
Now implement the kernel-fused version that operates at the hardware level.
class FusedFlashSamplingKernel:
"""Fused kernel combining LM-head and sampling."""
def __init__(self, vocab_size, hidden_dim, tile_size=256, device='cuda'):
self.vocab_size = vocab_size
self.hidden_dim = hidden_dim
self.tile_size = tile_size
self.device = device
# Pre-allocate buffers for tile computation
self.tile_buffer = torch.zeros((tile_size, hidden_dim),
device=device)
self.scratch_space = torch.zeros((tile_size,), device=device)
def fused_lm_head_sample(self, hidden: torch.Tensor,
lm_head_weight: torch.Tensor,
lm_head_bias: torch.Tensor = None,
temperature: float = 1.0) -> torch.Tensor:
"""
Single fused kernel: LM-head matrix multiplication + sampling.
This represents what would be implemented in CUDA/Triton.
"""
batch_size = hidden.size(0)
best_indices = torch.zeros(batch_size, dtype=torch.long,
device=self.device)
best_scores = torch.full((batch_size,), float('-inf'),
device=self.device)
# Simulate on-chip tile processing
for tile_idx in range((self.vocab_size + self.tile_size - 1)
// self.tile_size):
start_vocab = tile_idx * self.tile_size
end_vocab = min(start_vocab + self.tile_size, self.vocab_size)
tile_vocab_size = end_vocab - start_vocab
# In actual kernel: all computation happens on-chip
# Simulating: fetch weight tile, compute tile
tile_weights = lm_head_weight[start_vocab:end_vocab, :]
# Matrix multiply
tile_logits = torch.matmul(hidden, tile_weights.t())
if lm_head_bias is not None:
tile_bias = lm_head_bias[start_vocab:end_vocab]
tile_logits = tile_logits + tile_bias
# Generate Gumbel noise
u = torch.rand_like(tile_logits)
gumbel = -torch.log(-torch.log(u + 1e-20) + 1e-20)
# Perturb and scale
tile_scores = (tile_logits + gumbel) / temperature
# Find tile maximum
tile_max_scores, tile_max_indices = torch.max(tile_scores, dim=1)
# Update global maximum
improved = tile_max_scores > best_scores
best_scores[improved] = tile_max_scores[improved]
best_indices[improved] = start_vocab + tile_max_indices[improved]
return best_indices
def benchmark_vs_standard(self, batch_size=32, hidden_dim=4096,
vocab_size=128000, num_iterations=100):
"""Compare latency with standard sampling."""
import time
hidden = torch.randn(batch_size, hidden_dim, device=self.device)
lm_head_weight = torch.randn(vocab_size, hidden_dim,
device=self.device)
# Warm up
for _ in range(5):
_ = self.fused_lm_head_sample(hidden, lm_head_weight)
# FlashSampling timing
torch.cuda.synchronize()
start = time.time()
for _ in range(num_iterations):
_ = self.fused_lm_head_sample(hidden, lm_head_weight)
torch.cuda.synchronize()
flash_time = time.time() - start
# Standard sampling timing
torch.cuda.synchronize()
start = time.time()
for _ in range(num_iterations):
logits = torch.matmul(hidden, lm_head_weight.t())
u = torch.rand_like(logits)
gumbel = -torch.log(-torch.log(u + 1e-20) + 1e-20)
_ = torch.argmax(logits + gumbel, dim=-1)
torch.cuda.synchronize()
standard_time = time.time() - start
speedup = standard_time / flash_time
reduction_percent = (1 - 1/speedup) * 100
print(f"FlashSampling: {flash_time:.3f}s")
print(f"Standard: {standard_time:.3f}s")
print(f"Speedup: {speedup:.2f}x ({reduction_percent:.1f}% reduction)")
return speedup
Finally, integrate into inference pipeline and demonstrate usage.
class VLLMWithFlashSampling:
"""vLLM-compatible interface with FlashSampling."""
def __init__(self, model_name, use_flash_sampling=True):
self.model = load_model(model_name)
self.tokenizer = load_tokenizer(model_name)
self.use_flash_sampling = use_flash_sampling
if use_flash_sampling:
vocab_size = len(self.tokenizer)
hidden_dim = self.model.config.hidden_size
self.sampler = FusedFlashSamplingKernel(vocab_size, hidden_dim)
def generate(self, prompt: str, max_tokens=100, temperature=0.7):
"""Generate with optional FlashSampling."""
input_ids = self.tokenizer.encode(prompt)
generated = input_ids.copy()
for _ in range(max_tokens):
# Forward pass
with torch.no_grad():
outputs = self.model(torch.tensor([generated]))
hidden = outputs.hidden_states[-1][:, -1, :] # Last token
lm_head_weight = self.model.lm_head.weight
# Sample next token
if self.use_flash_sampling:
next_token = self.sampler.fused_lm_head_sample(
hidden, lm_head_weight, temperature=temperature
)
else:
# Standard sampling
logits = torch.matmul(hidden, lm_head_weight.t())
u = torch.rand_like(logits)
gumbel = -torch.log(-torch.log(u + 1e-20) + 1e-20)
next_token = torch.argmax(logits + gumbel, dim=-1)
generated.append(next_token.item())
if next_token.item() == self.tokenizer.eos_token_id:
break
return self.tokenizer.decode(generated)
Hyperparameters and When to Use:
When NOT to use:
Common Pitfalls: