ADu2021/flash-sampling-efficient-decoding

FlashSampling: Fusing Sampling into Matrix Multiplication

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).

Core Concept

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.

Architecture Overview

• LM-Head Tiling — Decompose matrix multiplication into independent tiles
• On-Chip Computation — Keep tiles in fast cache, avoid HBM spills
• Gumbel Noise Injection — Apply noise during tile computation
• Running Maximization — Track (value, index) of best tile element
• Grouped Variant — Hierarchical reduction for tensor parallelism
• Kernel Fusion — Single GPU kernel replaces compute + memory stages
• Batch Processing — Maintain efficiency across batch dimensions

Implementation Steps

Start by implementing the basic tile-based sampling logic.

import torch
import torch.nn.functional as F
import numpy as np

class FlashSampler:
    """Efficient sampling via tile-based computation."""

    def __init__(self, vocab_size, hidden_dim, tile_size=256):
        self.vocab_size = vocab_size
        self.hidden_dim = hidden_dim
        self.tile_size = tile_size
        self.num_tiles = (vocab_size + tile_size - 1) // tile_size

    @staticmethod
    def gumbel_sample(logits: torch.Tensor, temperature: float = 1.0):
        """Sample using Gumbel-max trick with temperature scaling."""
        # Generate Gumbel noise
        u = torch.rand_like(logits)
        gumbel_noise = -torch.log(-torch.log(u + 1e-20) + 1e-20)

        # Apply noise and temperature
        perturbed = (logits + gumbel_noise) / temperature

        # Argmax gives sample
        return torch.argmax(perturbed, dim=-1)

    def sample_tiled(self, hidden: torch.Tensor, lm_head_weight: torch.Tensor,
                    temperature: float = 1.0) -> torch.Tensor:
        """
        Sample from categorical distribution using tile-based computation.

        Args:
            hidden: [batch_size, hidden_dim]
            lm_head_weight: [vocab_size, hidden_dim]
            temperature: sampling temperature
        """
        batch_size = hidden.size(0)
        best_indices = torch.zeros(batch_size, dtype=torch.long)
        best_scores = torch.full((batch_size,), float('-inf'))

        # Iterate over vocabulary tiles
        for tile_idx in range(self.num_tiles):
            start_vocab = tile_idx * self.tile_size
            end_vocab = min(start_vocab + self.tile_size, self.vocab_size)

            # Get tile of LM-head weights
            tile_weights = lm_head_weight[start_vocab:end_vocab, :]
            # [tile_size, hidden_dim]

            # Compute logits for this tile
            tile_logits = torch.matmul(hidden, tile_weights.t())
            # [batch_size, tile_size]

            # Generate Gumbel noise
            u = torch.rand_like(tile_logits)
            gumbel = -torch.log(-torch.log(u + 1e-20) + 1e-20)

            # Apply noise and temperature
            tile_scores = (tile_logits + gumbel) / temperature
            # [batch_size, tile_size]

            # Find best in this tile
            tile_best_scores, tile_best_local_idx = torch.max(tile_scores, dim=1)

            # Update global best
            improved = tile_best_scores > best_scores
            best_scores[improved] = tile_best_scores[improved]
            best_indices[improved] = (
                start_vocab + tile_best_local_idx[improved]
            )

        return best_indices

    def sample_grouped(self, hidden: torch.Tensor,
                      lm_head_weight: torch.Tensor,
                      temperature: float = 1.0,
                      group_size: int = 4) -> torch.Tensor:
        """
        Grouped variant for tensor-parallel settings.
        Each device handles vocab_size // group_size tokens.
        """
        # Hierarchical reduction across groups
        group_vocab_size = self.vocab_size // group_size

        best_indices = torch.zeros(hidden.size(0), dtype=torch.long)
        best_scores = torch.full((hidden.size(0),), float('-inf'))

        for group_idx in range(group_size):
            # Get weights for this group
            group_start = group_idx * group_vocab_size
            group_end = (group_idx + 1) * group_vocab_size
            group_weights = lm_head_weight[group_start:group_end, :]

            # Tile within group
            for tile_in_group in range((group_vocab_size + self.tile_size - 1)
                                      // self.tile_size):
                tile_start = group_start + tile_in_group * self.tile_size
                tile_end = min(tile_start + self.tile_size, group_end)

                tile_weights = lm_head_weight[tile_start:tile_end, :]
                tile_logits = torch.matmul(hidden, tile_weights.t())

                # Gumbel noise
                u = torch.rand_like(tile_logits)
                gumbel = -torch.log(-torch.log(u + 1e-20) + 1e-20)
                tile_scores = (tile_logits + gumbel) / temperature

                # Update best
                tile_best_scores, tile_best_idx = torch.max(tile_scores, dim=1)
                improved = tile_best_scores > best_scores
                best_scores[improved] = tile_best_scores[improved]
                best_indices[improved] = tile_start + tile_best_idx[improved]

        return 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)

Practical Guidance

Hyperparameters and When to Use:

• Tile size 256-512 works well; smaller tiles improve cache locality, larger reduce kernel launch overhead
• Temperature 0.7-1.0 for typical generation; lower temperatures sharpen distribution
• Use when performing high-throughput inference on modern GPUs (H100, B200, etc.)
• Particularly effective for large vocabulary models (100K+ tokens)
• Benefit amplifies with larger batch sizes (less kernel launch overhead amortization)

When NOT to use:

• For CPU inference (no CUDA optimization benefit)
• For very small vocabulary models (< 10K) where logits materialization is already fast
• When using older GPU architectures without good on-chip memory

Common Pitfalls:

• Gumbel noise generation becoming bottleneck; batch random number generation across tiles
• Tile boundaries causing alignment issues; use aligned tile sizes
• Numerical instability in log-probability computation; use log-sum-exp tricks
• Not accounting for biases; apply bias during tile computation

Reference

Paper: FlashSampling: Fast and Memory-Efficient Exact Sampling