메뉴
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에서 현재 제공할 수 있는 파일을 다운로드하세요.
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: