32B模型首token从4.3秒到？实测通信优化降38%延迟 (English)

Generated: 2026-06-22 08:43:45

---

Large Model Inference Optimization: The Pits I Fell Into, and You Might Be Jumping Right In

To be honest, I held off writing this article for over half a year before I finally dared to put pen to paper.

The reason is simple—too many people oversimplify large model inference optimization. The articles online either throw a bunch of jargon at you that makes your scalp tingle, or they just hand you conclusions to copy. I've been in this field for ten years, and I know one thing for sure: What you read on paper is shallow; the real pitfalls only show up when you actually do it.

Today, I'm going to lay out my year's worth of hands-on records and faceplant experiences, just like chatting with an old friend.

---

The Moment That Broke Me First

In early 2024, I took on an inference service optimization project.

The model was 32B-level, running logic problems like MATH. I thought swapping in an A100 would be enough, but when I tested it, an input of 300 tokens forced an output of a 32K-token chain of thought.

Think about it: one request, and the first token took 4.3 seconds to appear.

A full 4.3 seconds! In the world of AI, that's like waiting a lifetime.

This is a classic characteristic of LRMs (Large Reasoning Models)—inputs are as short as a chat, outputs are as long as a thesis. I call it "generation asymmetry," and the computational bottleneck is entirely in the Decode phase.

Even worse, many people think it's just a memory issue.

I thought the same at first.

Then I crashed and burned.

---

Step One: Don't Be Fooled by the Word "Sparsification"

When I first saw the RTPurbo approach, my immediate reaction was: Isn't this just Top-k sparse attention? Anyone who's worked with Transformers knows that, right?

But in practice, I found that static Top-k strategies are a trap in long-context scenarios.

Here's the first pitfall I fell into: I set a fixed sparsity for a 32K context, say keeping 1024 tokens per layer.

It worked fine for short texts, but as soon as I switched to long texts, it collapsed—some queries suddenly needed a much larger context window, and the fixed sparsity caused the model to lose critical information, dropping accuracy from 85% to 67%.

See? That's not optimization; that's self-harm!

Later, after reading the RTPurbo paper, I understood the real issue: Attention's token budget is highly query-aware.

Simply put, different queries have completely different context needs.

Under a 32K context, RTPurbo's simplest "needle in a haystack" task kept an average of only 468.8 active tokens, but complex reasoning tasks automatically expanded to 2462.1 tokens—a dynamic range of 5x.

Allocate on demand, not with a one-size-fits-all approach.

I remember telling my team: This isn't an algorithm problem; it's a methodology problem. What you need to do isn't pick a better sparsification strategy, but let the model decide for itself which tokens are important.

---

Step Two: Blood and Tears of Communication Optimization

This is the part that made me want to curse the most.

Our inference cluster was four DGX A100s, each with 8 GPUs, totaling 32 GPUs. You'd think that's enough?

Naive.

The first time I ran a 32B model with tensor parallelism, I found that All-Reduce took up over 60% of the single-step computation time.

Two reasons:

1. I was using the default ring all-reduce, and the communication latency for 32 cards exploded.
2. The inter-node connection used RoCEv2, which had nearly 3x higher latency than NVLink.

I spent two months trying out solutions.

First, I switched from Ring All-Reduce to Tree All-Reduce.

For a scale of 32 cards, tree reduction has a clear latency advantage. I benchmarked with NCCL's ncclAllReduce, and tree reduction reduced single-step communication time by 38% for 24+ cards.

Second, topology-aware scheduling isn't voodoo.

I enabled NVIDIA's NVTAGS to keep communication within the NVLink domain as much as possible. Reducing cross-node communication naturally lowered latency.

But here's a trap: many people think enabling NVTAGS is enough.

It's not.

You have to tune the CARVEOUT parameter based on the actual topology—that's the allocation ratio between L1 cache and shared memory. This parameter directly affects kernel execution efficiency.

How many times did I tune it?

16 times.

I went from the default 0.5 to 0.625, and the GEMM operator performance improved by 12%.

Why? Because in communication-intensive TP scenarios, shared memory pressure is lower than in compute-intensive scenarios, so giving more resources to L1 cache actually improves data reuse.

This taught me two things:

• There's no silver bullet; if you're going to tune parameters, be prepared to tune.
• What you think of as "default configuration" is probably optimized for a different scenario.

---

Step Three: MoE's "Tail Effect" Almost Made Me Give Up

Midway through the project this year, I tried switching the model to a MoE architecture, thinking expert parallelism would surely reduce costs.

Reality slapped me hard.

The problem was the "tail effect"—when certain experts become hot spots, the overloaded gate slows down the entire pipeline.

Logs recorded an extreme case: one expert's load was 17 times that of others, and effective throughput was cut in half.

17 times! That's not optimization; that's a disaster.

My first solution was expert replication. I cloned the hot experts onto additional GPUs to share the load. Before the experiment, I thought it was simple—just a bit more memory.

It wasn't.

Without routing optimization, replicating experts actually made communication messier—different tokens from the same request could be routed to different GPU copies of the same expert, causing the complexity of All-To-All communication to grow exponentially.

Later, I switched to an adaptive routing scheme.

The core logic was simple: before each batch, dynamically compute the gate based on real-time load.

But here's a critical detail—you can't recompute every step.

I set a 100ms update window, and within that window, I reused the routing decision.

Why? Because recomputing the gate itself has a cost. Too frequent updates prevent communication and computation from overlapping, and the gain doesn't offset the loss.

Final test results: adaptive routing with a reasonable window restored throughput from 47% to 89%.

---

Step Four: KV Cache Management—A Small Problem, a Big Pitfall

I used to think KV Cache wasn't worth optimizing.

Until one day, an online service ran out of memory (OOM).

Investigation revealed that a service using continuous batching was handling 32 long-context requests at peak, and the KV Cache ate up 196GB of memory, exceeding the total 640GB of 8 A100s.

Wait, how could 640GB be exhausted?

The problem was memory fragmentation.

PyTorch's BFC allocator, when frequently allocating and releasing KV Cache blocks of different sizes, creates a lot of fragmentation. The actual usable contiguous memory blocks are far smaller than the theoretical value.

My solution had three lines:

Line one: Quantize the cache.

I tried INT8 and FP8. INT8 had negligible accuracy loss (tested on MMLU-PRO, less than 0.3% difference) and cut memory usage in half. I used HQQ, a calibration-free quantization method that doesn't need extra calibration datasets—very practical.

Line two: Slab allocator.

I replaced the default BFC with Slab. The approach was simple: monitor the size distribution of KV Cache in production, pre-allocate fixed-size objects. When a request arrives, just grab from the pool—no need to call cudaMallocAsync frequently.

Result? Fragmentation dropped from 32% to 4.7%.

Line three: Sparse KV Cache.

Honestly, I didn't use this in production because the risk was too high. But LazyLLM's idea is interesting—selectively compute KV cache for important tokens. It maintains over 90% accuracy even at 95% sparsity.

---

Step Five: Scheduling Is the Real Killer

All the optimizations above ultimately come down to scheduling.

My first attempt was Chunked Prefill—splitting a large prefill request into smaller chunks and interleaving them with decode requests. The SARATHI scheduler does exactly that.

In practice, I found a few useful details:

First, chunk size can't be fixed.

I started with a fixed 256-token chunk, but for short inputs (say 50 tokens), chunking added extra scheduling overhead. I switched to dynamic chunking based on sequence length—short inputs aren't chunked, long inputs are split into at most 4 chunks.

Second, Decode-Maximal Batching is effective but not a panacea.

This strategy prioritizes decode tasks