Optimizing TTFT for a Hybrid MoE model

Inference engines like vLLM and SGLang are well optimized out of the box to serve many models efficiently across a wide range of request sizes. But the defaults won't always work for your specific use case. Each engine exposes a large space of knobs like max_num_seqs, max_num_batched_tokens, max_model_len, gpu_memory_utilization, CUDA-graph capture etc. If the model has a hybrid architecture, there is more configuration around kernel and KV-cache layout.

As with any optimization problem, the right choice depends on the input workload and the metric you're optimizing for. This article works through one such exercise, where we optimize TTFT for Qwen3.6-35B-A3B for a specific workload defined below. This model is a hybrid attention/Mamba MoE. We'll be using vLLM on an H100 throughout.

Constraints and setup

• Input tokens up to ~8K, with the majority under 4K.
• Output tokens up to 256
• Latency over throughput.

We will use the following setup to benchmark:

• Single request (concurrency 1) : the latency a request sees with nothing queued ahead of it.
• Single GPU : no tensor parallelism, no replicas.
• Cold (prefix caching off) : every prompt is a full prefill. This gives us the worst-case number that reflects raw compute without any cache hits.

Understanding the model

We need to understand the model architecture together with the workload above to set the initial knobs on vLLM. Then we can profile and iterate until we achieve the desired result. The following are the parts of the model that matter for inference:

A few things stand out from that table:

Hybrid 3:1: Only 10 of 40 layers run full attention. This means the O(L²) cost from attention is only on 10 layers. The other 30 are Gated DeltaNet linear-attention layers which would be O(L) compute with a fixed-size recurrent state. So long context and KV cache should be cheap in both compute and memory, dominated by just 10 layers.

MoE: Only ~3B of 35B params are active per token and the compute is spread across 256 experts with an intermediate of just 512 (for comparison, a big-expert MoE like Mixtral uses 14336). This means there are many small grouped GEMMs, so kernel efficiency should be more important than FLOPs for optimized inference

KV cache: Each full attention layer stores just 2 KB/token as there are only 2 KV heads and 256 head_dim. Only 10 / 40 layers keep KV cache, so even an 8K request costs only ~160 MB. For comparison, an all-attention model like Google's Gemma 4 26B — 30 KV layers, 8 KV heads — needs ~240 KB/token = ~1.8 GB at 8K.

The first run

Before the first run, let's look at the kernels we expect vLLM to use for this model.

Here is the vLLM serve command for our first run without any optimizations. We are running with max_model_len = 8192, max_num_seq = 2, max_num_batched_tokens = 8192 and without prefix caching. This baseline reflects raw work and nothing clever.

vllm serve Qwen/Qwen3.6-35B-A3B-FP8 \
  --served-model-name qwen-3.6-a3b \
  --host 0.0.0.0 --port 8000 \
  --tensor-parallel-size 1 \
  --max-model-len 8192 \
  --max-num-seqs 2 \
  --max-num-batched-tokens 8192 \
  --gpu-memory-utilization 0.85 \
  --enable-chunked-prefill \
  --no-enable-prefix-caching \
  --trust-remote-code \
  --reasoning-parser qwen3 --tool-call-parser qwen3_coder --enable-auto-tool-choice \
  --default-chat-template-kwargs '{"enable_thinking": false}' \
  --dtype auto \
  --limit-mm-per-prompt '{"image": 0, "video": 0}' \
  --enforce-eager     

The curve sits near a ~68 ms floor, almost all kernel-launch overhead. Short prompts land around ~80 ms, and first-token latency crosses 100 ms at only ~2.3K tokens, reaching ~206 ms by 8K. That overhead is the first obvious thing to cut.

The second run with cuda graphs

A CUDA graph records a whole sequence of GPU kernel launches once and replays it as a single unit. Normally the CPU issues each kernel one at a time; for a batch of one, that per-launch overhead is most of the forward pass. Capturing the pass as a graph lets the GPU replay the entire thing from a single launch, driving that overhead toward zero.

So the second run is the same command with graphs on, because our prompts reach 8K, we widen the capture range to cover them:

vllm serve Qwen/Qwen3.6-35B-A3B-FP8 \
  --served-model-name qwen-3.6-a3b \
  --host 0.0.0.0 --port 8000 \
  --tensor-parallel-size 1 \
  --max-model-len 8192 \
  --max-num-seqs 2 \
  --max-num-batched-tokens 8192 \
  --gpu-memory-utilization 0.85 \
  --enable-chunked-prefill \
  --no-enable-prefix-caching \
  --trust-remote-code \
  --reasoning-parser qwen3 --tool-call-parser qwen3_coder --enable-auto-tool-choice \
  --default-chat-template-kwargs '{"enable_thinking": false}' \
  --dtype auto \
  --limit-mm-per-prompt '{"image": 0, "video": 0}' \
  --compilation-config '{"cudagraph_capture_sizes":
    [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128,
     136, 144, 152, 160, 168, 176, 184, 192, 200, 208, 216, 224, 232, 240, 248,
     256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480,
     496, 512, 640, 768, 896, 1024, 1056, 1152, 1280, 1408, 1536, 1664, 1792,
     1920, 2048, 2112, 3168, 4224, 5280, 6336, 7392, 8192]}'

When we start the vLLM server with above command, the engine dies in the CUDA-graph memory-profiling pass:

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 16.50 GiB.
GPU 0 has a total capacity of 79.18 GiB, 11.40 GiB free; 67.77 GiB in use.

A 35 GB model on an 80 GB H100 wants an extra 16.5 GB beyond the KV cache and model weights. This doesn't add up as we measured the KV cache at only ~160 MB (8K) and the weights are ~35 GB, and CUDA graphs are cheap in memory to begin with, a few MB per size. Something in how vLLM sizes this is miscounted.

Another reason that it's model-specific is that the same config, same H100 boots fine on a non-hybrid model (Google Gemma 4) with cuda graph pool under 1 GB. So whatever is wrong is specific to this architecture and we have to debug before going further.

Debugging the OOM issue

1. Give more room to cuda graphs. The first instinct is to change gpu_memory_utilization but it didn't help: as it freed some memory but the OOM persisted.
2. Trace it in vLLM. vLLM captures the CUDA graphs by running a dummy forward pass, and that pass needs a scratch KV cache that is allocated just for the capture and freed straight after. vLLM sizes it as min_blocks blocks, and sets min_blocks to the largest capture size (gpu_model_runner.py):

min_blocks = self.compilation_config.max_cudagraph_capture_size or 1   # 8192 = our largest capture size

# allocate a "minimal" scratch KV of min_blocks BLOCKS for the dummy capture pass
self.cache_config.num_gpu_blocks_override = min_blocks
minimal_config = get_kv_cache_config_from_groups(self.vllm_config, kv_cache_groups, available_memory=0)
self.initialize_kv_cache(minimal_config, is_profiling=True)

So the scratch is min_blocks = 8192 blocks. On a normal model a block is 16 tokens, so 8192 blocks is a trivial amount of memory. But when we traced it on this model, each block turned out to be 1056 tokens which is about 22 MB (1056 tokens × ~20 KB/token across the 10 attention layers), so those 8192 blocks come to 8192 × ~22 MB ≈ 180 GB, far past the 80 GB VRAM. There's the exact OOM. But why is one block 1056 tokens?

Why the page is 1056 tokens?

On boot, vLLM logs something you won't see for a Llama-class model:

Setting attention block size to 1056 tokens to ensure attention page size >= mamba page size.
Padding mamba page size by 0.76% to ensure mamba page size and attention page size are exactly equal.

A Llama-style model uses a KV block size of 16 tokens. But since this model is hybrid, and vLLM wants one unified block pool for both layer types, it chooses the max of both. The GDN layer's recurrent state is the big one:

GDN state ≈ num_value_heads(32) × key_head_dim(128) × value_head_dim(128) × 4 B (fp32)
          ≈ 2.0 MiB  per linear layer per sequence    (+ a small conv state)

to make ONE attention page hold the same bytes:
block_size = mamba_page_bytes / per_token_attn_KV  ≈  2.16 MB / 2 KB  ≈  1056 tokens

Which is exactly what vLLM computes (vllm/platforms/interface.py):

attn_page_size_1_token = FullAttentionSpec(block_size=1, ...).page_size_bytes   # ~2 KB
mamba_page_size        = MambaSpec(...).page_size_bytes                         # ~2.1 MB (GDN state)

# smallest block_size whose attention page holds >= one mamba state
attn_block_size = align * cdiv(mamba_page_size, align * attn_page_size_1_token) # -> 1056

Attention doesn't need 1056-token pages, it inherits them so the allocator can treat a sequence as N identical blocks regardless of layer type. The price is internal fragmentation, which a sequence pays by rounding up to whole 1056-token blocks on each of the 10 attention layers:

The fix

Since the dummy capture only ever needs a handful of blocks, let's fix by capping it:

_max_cg = self.compilation_config.max_cudagraph_capture_size or 1
min_blocks = min(_max_cg, 64)   # 64 blocks × 1056 tok ≈ 1.4 GB, not 180 GB

The Gemma row is the proof: same GPU, same capture size but it never had the problem, because a 16-token page makes the block miscount trivial. The patch (vllm-patches/patch_minimal_kv_cache.sh) is Mamba/hybrid-specific and a harmless no-op everywhere else. Now the engine boots, and we can finally measure with cuda graphs enabled.

The final run

With the engine serving at our capture sizes, here's TTFT across the workload range — single request, cold, on H100 (config/roofline_sweep.py; 30 cold reps/length, in-process):

H100 80GB, vLLM 0.23.0, Qwen3.6-35B-A3B-FP8, c=1, prefix-caching off.

Note that the with-graphs curve isn't perfectly smooth: there's a repeatable ~24 ms step up just past 2048, 4096, and 6144 tokens because we capture graphs at those lengths. For prompts which fall in the middle, vLLM pads the prompts to the nearest size and hence the steps in the graph.

Profiling

Profile the 4K prefill (profile/profile_phases.py, share of GPU device time on H100):

MoE is the single largest prefill cost on H100, ~39% (~76 ms including routing), with the fused_moe_kernel alone at ~35%, ~68 ms. That's huge for a model whose MoE is only ~13% of the FLOPs.

We also measured the same MoE on a B200 and it takes ~17 ms, a 4.5× gap. vLLM auto-selects different kernels per GPU. On H100 the MoE runs on untuned Triton (~76 ms) while attention uses a fused FA3 path (~9 ms); on B200 the MoE runs on FlashInfer-TRTLLM (~17 ms) while attention runs on an FP8 bmm path (~41 ms).

So the 4.5× MoE gap is mostly not hardware: B200 gets a fast, purpose-built kernel for this shape and H100 gets an untuned generic one.

vLLM ships an autotuner (benchmark_moe.py) that searches block sizes, warp counts, and pipeline stages for a given expert shape and writes the config the kernel loads. There's no entry for our architecture, so we ran it at 640 configurations for E=256, N=512 FP8 block-scale on H100 (config/tune_moe.sh). The result:

About 1%. The best config the search found was essentially the default.

Prefill vs decode

A single request has two phases: one big prefill forward (→ TTFT, ~101 ms at 4K) and then one decode forward per output token. We profiled a decode run too (128-token prompt, 256 outputs). Its kernel mix looks nothing like prefill:

Where prefill is compute bound (MoE dominates at ~39%), decode is memory bound: GEMM/GEMV — matrix-vector ops that read the whole weight matrix for a single token climb from ~25% to ~35% and tie MoE.

And it's why single-request decode isn't worth optimizing on its own. At batch 1 the GPU is ~5% active, so the measured ~2.9 ms/token TPOT is ~20× overhead. Batching amortizes it: throughput climbs ~12× while TPOT stays flat out to c ≈ 16:

H100, 128-token prompt / 128 output, prefix-caching off.

What's next

The open thread is the small-expert MoE kernel, essentially a grouped GEMM without padding for the N≤512 regime. That's the one lever that could move the H100 number, and it's a worklog of its own. And the other levers like prefix caching, concurrency / continuous batching, and multi-GPU will only move the numbers down from here.