Sparse Activation and Multi-Token Prediction: LLM Inference Acceleration in Practice

1. Background: Parameter Scale vs. Inference Cost

The capability of large language models usually grows with parameter scale, but inference cost rises as well. More parameters mean more computation per generated token and higher memory usage. In real projects, developers commonly face three problems: response latency exceeds the limit for interactive experience, the number of requests served per unit time is insufficient, and long documents or large codebases cannot fit into the context at once.

Solutions generally fall into two categories. The first reduces computation per inference step so that the model keeps its knowledge capacity while fewer parameters participate in each step. The second improves hardware utilization so that each decoding step produces more tokens. Most current mainstream approaches combine engineering optimizations in both directions.

Understanding these approaches requires distinguishing two concepts. Total parameters determine how much knowledge the model can store and how much memory the weights occupy. The number of parameters actually involved in a single inference step determines the computational cost per token. The two can differ, and this is the starting point for sparse architectures.

2. Sparse Activation: Mixture of Experts

2.1 Basic Structure

A Mixture of Experts (MoE) architecture replaces the feed-forward network layer with multiple parallel “expert” sub-networks and adds a router. For each input token, the router computes scores for the experts and selects only the few experts with the highest scores to participate in computation, leaving the rest idle.

As a result, the total number of model parameters can be large, while only a small fraction is activated for each token. From a compute perspective, the forward cost of one token is close to that of a small model. From a storage perspective, all expert weights must reside in GPU memory or system memory.

2.2 Routing Example

The following code demonstrates the core logic of top-k routing using only the standard library:

import math
def softmax(xs):
    m = max(xs)
    exps = [math.exp(x - m) for x in xs]
    s = sum(exps)
    return [e / s for e in exps]
def top_k_route(logits, k):
    probs = softmax(logits)
    ranked = sorted(range(len(probs)), key=lambda i: probs[i], reverse=True)[:k]
    total = sum(probs[i] for i in ranked)
    # Keep only top-k experts and renormalize their weights
    return [(i, probs[i] / total) for i in ranked]
# Assume 8 experts in total and only 2 are activated per token
router_logits = [0.1, 2.3, -0.5, 1.7, 0.0, -1.2, 0.9, 0.4]
print(top_k_route(router_logits, k=2))

The output will show experts 1 and 3 with renormalized weights. The other six experts incur no computation for this token.

2.3 Load Balancing

If the router always favors a small number of experts, two problems arise. Popular experts become bottlenecks, and less-used experts do not receive sufficient training. During training, an auxiliary loss is usually added to encourage expert invocation frequencies to approach a uniform distribution. During inference, imbalanced load affects batching efficiency, so deployment should monitor statistics on expert distribution.

2.4 Use Cases and Trade-offs

Sparse architectures are suitable for the following situations:

• A large knowledge capacity is needed, but single-response latency is sensitive;
• The server processes requests in batches, which can amortize expert switching and communication overhead;
• Sufficient GPU memory or unified memory is available to hold all expert weights.

Its limitations are also clear. Total parameters determine weight storage requirements, so even if only a small portion is computed per step, the deployment threshold still depends on total parameters. Communication between experts introduces additional latency in multi-GPU scenarios and must be controlled through parallelism strategies and communication optimization.

3. Multi-Token Prediction for Faster Decoding

3.1 Bottleneck of Autoregressive Decoding

Standard autoregressive generation outputs only one token at a time, and each step must read all relevant weights. When the batch size is small, the decoding phase is usually limited by memory bandwidth, and compute units are idle most of the time. Therefore, even if the theoretical compute capacity of the hardware is high, actual generation speed may be slow.

3.2 Basic Idea of Multi-Token Prediction

Multi-Token Prediction (MTP) adds several prediction heads on top of the model backbone. Each head predicts the token at future position 1, 2, …, k. During inference, the model can propose multiple candidate tokens in one forward pass, then use a verification step to confirm which candidates match the standard autoregressive output and accept the consecutive correct portion.

This idea is related to speculative decoding. The difference is that the MTP draft model shares most parameters with the backbone, so the additional overhead is smaller.

3.3 Estimating Speedup

The speedup depends on the acceptance rate of candidate tokens. Assume each draft token is independently accepted with probability p, and at most k draft tokens are proposed per verification step. The average number of tokens produced per verification step is:

E = (1 - p^(k+1)) / (1 - p)

The following code calculates the expected output under different acceptance rates:

def expected_tokens_per_step(p, k):
    if p == 1:
        return k + 1
    return (1 - p ** (k + 1)) / (1 - p)
for p in [0.5, 0.7, 0.8, 0.9]:
    print(p, round(expected_tokens_per_step(p, k=3), 2))

This formula is based on a simplified independent and identically distributed assumption. In real scenarios, acceptance rates fluctuate with context. Code, structured text, and highly repetitive content usually have higher acceptance rates, while open-ended creation usually has lower acceptance rates. Therefore, actual speedup must be measured with real workloads rather than relying only on theoretical values.

3.4 Practical Considerations

• With larger batches, decoding may become compute-bound, and the benefit of MTP decreases;
• Extra computation in the verification phase must be included in total latency;
• For scenarios with high requirements on output determinism, confirm that the verification mechanism is consistent with sampling settings.

4. Long Context: Hybrid Sliding-Window and Global Attention

4.1 Cost of Full Attention

The computation of standard full attention grows quadratically with sequence length. When sequence length increases from 4K to 256K, attention computation grows by about 4096 times. This makes very long contexts expensive in practical deployment.

4.2 Sliding Window Attention

Sliding Window Attention (SWA) limits each position to attending only to the most recent w tokens. The computational complexity of a single layer becomes linear with sequence length, and memory is easier to control.

The drawback is that information outside the window cannot be passed directly, and the model must obtain long-distance information indirectly through multiple stacked layers.

4.3 Hybrid Structure

A common approach is to interleave the two types of layers in a ratio, for example, one global attention layer after every three sliding-window attention layers. Local layers handle nearby context efficiently, while global layers aggregate information across distances. This structure retains the ability to model long-range dependencies while controlling cost.

The following code generates a causal sliding-window mask to help illustrate the difference between the two attention types:

def build_attention_mask(seq_len, window, full_attention=False):
    mask = []
    for i in range(seq_len):
        row = []
        for j in range(seq_len):
            if j > i:
                row.append(0)  # Causal constraint: cannot attend to future positions
            elif full_attention or i - j < window:
                row.append(1)
            else:
                row.append(0)
        mask.append(row)
    return mask
local = build_attention_mask(6, window=3)
for row in local:
    print(row)

4.4 KV Cache Memory Estimation

The main memory overhead of long-context inference comes from the KV cache. Its size can be estimated as follows:

2 × layers × KV heads × head dimension × sequence length × bytes per element × batch size

For layers using sliding windows, the cache length only needs to retain the window size, so a hybrid structure can significantly reduce cache growth.

def kv_cache_gb(layers, kv_heads, head_dim, seq_len, bytes_per_elem=2, batch=1):
    total_bytes = 2 * layers * kv_heads * head_dim * seq_len * bytes_per_elem * batch
    return total_bytes / 1e9
# Example: 40 layers, 8 KV heads, 128 dimensions per head, 256K context, FP16 storage
print(round(kv_cache_gb(40, 8, 128, 256_000), 2), "GB")

This number is only an estimation example. Actual cache usage is also affected by framework implementation, paging strategy, and alignment.

5. Evaluation Methods for Coding and Agent Tasks

5.1 Limitations of a Single Metric

A single public benchmark can only reflect performance under a specific task distribution. A model may have seen similar problems in training data, making test results higher. Sampling temperature, prompt format, and number of runs also affect final scores.

5.2 Building an Internal Evaluation Set

A more reliable approach is to build an evaluation set based on your own business, including the following task types:

• Locate and fix defects in real code repositories;
• Generate functions that pass unit tests according to requirement descriptions;
• Complete multi-step operations in a restricted terminal environment;
• Perform security reviews or suggest refactoring for existing code.

The following is a simplified evaluation framework:

def evaluate(model_fn, tasks, runs=3):
    results = []
    for task in tasks:
        passes = 0
        for _ in range(runs):
            output = model_fn(task["prompt"])
            if task["check"](output):
                passes += 1
        results.append({"id": task["id"], "pass_rate": passes / runs})
    return results

Multiple runs can expose fluctuations caused by randomness and avoid treating accidental success as stable capability.

5.3 Record Speed and Quality Together

Evaluation should record first-token latency, overall throughput, pass rate, and cost at the same time. Looking only at throughput while ignoring correctness can lead to biased selection conclusions. It is recommended to compare speed, quality, and cost in the same table and note the test environment.

6. Local and Consumer Hardware Deployment

6.1 Memory and GPU Memory Estimation

The storage requirement for model weights can be estimated from parameter count and quantization bit width:

def weight_memory_gb(params_billion, bits):
    return params_billion * bits / 8
# Example: 100B parameters, 4-bit quantization
print(weight_memory_gb(100, 4), "GB")  # About 50 GB

For mixture-of-experts models, all expert weights must reside in memory, so storage requirements are calculated based on total parameters. Computation speed is closer to the scale of activated parameters.

6.2 Choosing Hardware Form Factors

• Discrete GPUs: high memory bandwidth, but limited per-card capacity; large models require quantization or multi-card partitioning;
• Desktop workstations with unified memory: larger capacity, suitable for holding large model weights, but bandwidth is usually lower than high-end GPUs, affecting decoding speed;
• Server-grade accelerators: sufficient capacity and bandwidth, suitable for multi-user concurrent serving.

When choosing, first determine the target metric: single-user response speed, multi-user total throughput, or whether very long contexts can be accommodated.

6.3 Quantization and Precision

Quantization can reduce memory usage for weights and caches, but it introduces some precision loss. It is recommended to compare quality differences across bit widths on target tasks, especially for error-sensitive scenarios such as code generation and numerical computation.

7. Selection and Deployment Recommendations

1. Identify the business bottleneck: which matters most among latency, throughput, context length, or cost;
2. Estimate resources first: GPU memory required for weights, KV cache, and batching, and confirm whether the hardware can meet it;
3. Test with an internal evaluation set: record quality and speed together, and avoid relying on a single leaderboard;
4. Measure rather than estimate: test acceleration ratios and other claims on real workloads.