tps-maxxing on amd gpu

Nvidia is estimated to generate 74% of inference tokens, while AMD is less than 20%. This gap is surprising, looking at the hardware specs both offer:

At Hopper, we are big users of H100s for voice inference. So we started optimizing inference on an AMD MI300X to see how good we could get it. Let's first setup the system and the goals we are targeting.

• Model: Qwen3.6-35B-A3B-FP8
• GPU: AMD MI300X
• TP1, no batching : we optimize single cold prompt performance
• Achieve highest possible TTFT and throughput across 1–8K prompts, maintaining quality (>97% in τ²-bench)

Baseline

We started with a plain vLLM serving containing full hip-graph (think cuda graphs for AMD) capture with a ROCM_ATTN attention backend. Here is the launch command:

vllm serve "$MODEL_DIR" \
  --served-model-name qwen-3.6-a3b \
  --tensor-parallel-size 1 \
  --max-model-len 262144 \
  --max-num-seqs 1 \
  --max-num-batched-tokens 12288 \
  --gpu-memory-utilization 0.9 \
  --enable-chunked-prefill \
  --default-chat-template-kwargs '{"enable_thinking": false}' \
  --speculative-config null \
  --attention-backend auto \
  --dtype auto \
  --limit-mm-per-prompt '{"image": 0, "video": 0}' \
  --compilation-config '{
    "cudagraph_capture_sizes": [
      1,2,3,4,5,6,7,8,12,16,24,32,40,48,56,64,72,80,88,96,
      104,112,120,128,256,512,768,1024,1280,1536,1792,2048,
      2304,2560,2816,3072,3328,3584,3840,4096,4352,4608,4864,
      5120,5376,5632,5888,6144,6400,6656,6912,7168,7424,7680,
      7936,8192
    ],
    "max_cudagraph_capture_size": 8192,
    "cudagraph_mode": "FULL_AND_PIECEWISE"
  }'

The following graph shows the TTFT sweep on this server covering 91 exact prompt lengths from 1 to 8,192 :

The ridges in the middle are because the input sizes match the graph capture boundaries. This will be important later when we are tuning the kernel selection.

Following are the decode numbers on same server:

Now, lets see where the baseline spends its GPU time. The table below shows the costliest kernels at 4k prompt length.

Few observations :

• The largest prefill kernel is the routed MoE at 40.8 ms. This is expected with the model's architecture : Qwen3.6 contains 40 layers with a 256-expert, top-8 MoE in every layer.
• Full attention was still substantial at 22.0 ms, but it did not dominate prefill as it would in an all-attention model because only a quarter of Qwen3.6’s layers use full attention.
• The logs also shows that a 4096 prompt is executed as 2 model forwards, 3168 first and then 928 in another. This is because vLLM selected a page size of 1,056 tokens.
• During decode, the costliest kernel is the paged attention at 5.157 ms because every new token has to query KV cache which starts at 4K tokens and grows throughout generation.

Optimizing kernels for MI300X

The next step is to optimize performance for each individual kernel, by specializing on shapes needed by our model. We adapted AMD's AITER for this. AITER provides optimized Triton, Composable Kernel, and assembly implementations for attention, GEMM, MoE, quantization, and other common inference operators. A kernel that performs well for M=4096 prefill may be poor for M=1 decode. So, we benchmark and find compatible implementations for each (M,N,K) or MoE shape.

The exact-4K tuning invocation looks like:

MODEL=qwen3.6-35b-a3b-fp8 \
bash profile/run.sh dev tune_mi300x_aiter.py \
  --component all \
  --dense-m-values 4096 \
  --dense-shape all \
  --dense-libtype both \
  --moe-tokens 4096 \
  --warmup 5 \
  --iters 30 \
  --out results/mi300x_aiter_tuning

gfx942 is AMD’s compiler target (think AMD's sm_90) for the CDNA3 architecture used by MI300X. It tells AITER which hardware-specific kernels and tuning configurations are valid for the GPU. AITER maintains tuned dispatch tables where each row effectively says which kernel to choose given a GPU, data_type and (M,N,K) shape.

For example, here is an entry from the block-scaled A8W8 GEMM tuning table used in our MI300X run:

"gfx": "gfx942",
"cu_num": 304,
"M": 2048,
"N": 7168,
"K": 2048,
"libtype": "ck",
"kernelId": 0,
"splitK": 0,
"us": 78.7607

This tells AITER that, on a gfx942 GPU, an A8W8 block-scaled mat mul with shape M=2048, N=7168, and K=2048 should use Composable Kernel configuration 0 without splitting the K dimension and this configuration took 78.76 microseconds in tuning.

During the tuning sweep, a 2,048-token request took 108.95 ms TTFT, while the 2,049-token request fell to 72.64 ms. We found that this was because the 2048-request selected a slower kernel. The dispatch table contained optimized rows for M=1 and M=4096, but not for the M=2048 bucket. The 2,048-token request therefore used fallback kernels, while the next bucket was padded into an already-optimized shape (M=4096). So, this tuning process adds the most perfomant kernels at every shape.

From multiple such passes across 1-8k tokens, the following prefill and decode kernels gave the best results:

At exact 4K, TTFT fell from 166.666 ms to 95.076 ms (43.0% reduction). This improvement persisted across the longer prompts. The following chart shows this improvement compared to baseline :

With decode attention reduced by 95%, throughput nearly doubled from 77.878 tok/s to 152.891 tok/s. Here are the details for a 4K prompt:

Speculative decoding: native MTP vs. DFlash

With the individual kernels optimized, to improve throughput further we tested two speculative decoding branches : Qwen's native MTP vs DFlash.

Native MTP reuses the prediction layer already embedded in the Qwen checkpoint. Because this model has one MTP layer, vLLM invokes it sequentially three times to propose three draft tokens.

DFlash instead uses a separate six-layer BF16 draft model to predict fifteen draft positions in parallel.

Speculator models take the tokens generated so far and propose the next K candidate tokens (K=3 for MTP, K=15 for DFlash). The target model scores those candidates in one verification pass, and the system accepts the valid prefix until the first rejection. The fraction of proposed draft tokens that survive is the acceptance rate.

When acceptance rate is high, one verification pass can commit several tokens. On the controlled workload with high acceptance rate >90%, MTP reached 329–331 tok/s, while DFlash reached 1,063 tok/s.

Here is how the three configurations behaved across lower- and high-acceptance workloads:

On varied MI300X prompts, MTP's three draft positions were accepted at approximately 77.7%, 53.5%, and 36.8%. DFlash began at 71.7%, then fell to 44.5%, 27.7%, 18.2%, 11.7%, and 7.9% over its first six positions. DFlash still committed slightly more text per verification round, but its wider draft and verification work erased that advantage.

The DFlash's wider proposer also costs more to start. At 4K input tokens, TTFT was 95.1 ms for the tuned target, 102.1 ms with MTP K3, and 114.3 ms with DFlash B16. MTP added approximately 7.4% to MI300X's one-token TTFT, while DFlash added 20.2%.

While MTP K3 performs well on general workloads, DFlash would be very beneficial if we could train the speculator model on predictable, domain-specific outputs. We are seeing >800 tok/s with narrow, domain-specific speculators on healthcare and insurance datasets.

Results and comparison with H100

First, here is the full progression on MI300X. The TTFT chart uses the full 1–8K one-token sweep, while the decode table reports median native throughput after the first token at 4K.

MI300X progression

H100 comparison

MI300X delivered lower TTFT in every matched configuration: 95.1 vs. 99.7 ms without spec-decoding, 102.1 vs. 104.8 ms with native MTP, and 114.3 vs. 124.3 ms with DFlash. That advantage reversed at longer contexts, with H100 achieving roughly 21% lower TTFT at 8K. While H100 remains substantially stronger for general decode and longer prompts, MI300X is highly competitive at shorter prompts (4k) prefill.