Dual MI50, part 2: from 10 to 74 tokens/s, a 200k-token context, and two attention kernels
Where we left off
Part 1 (in Arabic) covered getting two AMD MI50 16 GB cards working with vLLM. The blocker turned out to be the motherboard, not the software. With everything running, two questions remained:
- DFlash2 is the fastest speculative decoding method for Qwen3.8. Why does the gfx906 vLLM fork list it as “dead” on these cards?
- Can I get a 150–200k-token context? I run Hermes Agent and opencode, and coding agents read whole codebases.
The short answers:
- DFlash2 wasn’t dead. Two bugs drove its acceptance rate to exactly zero. With both fixed it went from 10 to 74 tokens/s.
- I now run an agent server with a 200k-token context on two eight-year-old GPUs.
- Two attention kernels and a prefix-cache fix later, that context is also fast: 48 t/s of decoding at 61k tokens (was 29), and a follow-up agent turn starts in about 5 seconds instead of 12–14.
TL;DR
| Before | After | |
|---|---|---|
| DFlash2 (k=7) | 10.2 t/s, every draft rejected | 73.7 t/s mean, 85 peak |
| MTP (k=3) | 49.9 t/s | 65.1 t/s with custom kernels |
| Max context for agents | ~96k | 200k (143k tested end to end) |
| GPU-to-GPU all-reduce | ~72 µs | 28 µs |
| Decode at 61k context | 29.1 t/s | 48.2 t/s |
| Prefill, cold 63k prompt | 90 tok/s | 160 tok/s |
| Agent turn, 16k conversation | 12–14.5 s to first token | 4.4–5.3 s |
Both cards are capped at 90 W (the driver default is 190 W).
How I worked: me and Claude Code
How the work was done is part of the story, so I’ll start there. I ran Claude Code as an agent inside the LXC under fixed rules:
- never touch the host’s ROCm, and do everything in Docker;
- ask before any apt change on the host;
- run builds in tmux with logs (the container has few cores);
- every reported number has to come from a real run.
I set the goals, owned the hardware and the decisions, and reviewed the results. The agent did the heavy lifting: reading the fork’s code, writing the kernels and fixes, and running the A/B tests. Directing AI agents on real engineering problems is what I do as an AI engineer, and this was a hard one.
Bug 1: fp16 overflows
The MI50 (gfx906) has no bf16, so everything runs in fp16. The DFlash2 drafter was trained in bf16, which has a much wider range. The largest value fp16 can hold is 65504.
By its second layer, the drafter’s residual stream goes past 65504 and becomes inf, then NaN. The drafts turn into garbage (token ids 0, 1, 2…), the target model rejects all of them, and acceptance is exactly zero. This is a known vLLM issue (#55250).
The fix. Every branch in a layer reads the residual stream through an RMSNorm, and an RMSNorm gives the same output when its input is divided by a constant. So at load time I divide the stream by 32, a power of two, which makes the division exact in floating point. The matching change goes into the o_proj and down_proj weights, once. The outputs are identical and the values stay inside fp16 range, with zero runtime cost.
Bug 2: the drafter read freed memory
With bug 1 fixed, DFlash2 reached 48 t/s, but only on short prompts. Past 2–4k tokens of context, acceptance fell back to zero.
The drafter uses a 2048-token sliding window: it should only see the most recent 2048 tokens. The fork added a non-causal attention path for the drafter, and that path set the window to zero, so the drafter attended over the whole context. Meanwhile, vLLM’s cache manager frees KV blocks that fall outside a sliding layer’s window and reuses them. The drafter was reading freed memory.
The fix. The window is now expressed with the mask the existing kernel already supports. The change is Python only, with no kernel rebuild. A regression test fills every out-of-window block with NaN, so any read of one fails the test. It fails on the original code and passes with the fix.
With only the fp16 fix, acceptance collapses after 4.6k tokens. With both fixes it holds at 3.7–3.9 accepted tokens per step up to 24k tokens.
Kernels: why the generic code is slow here
Speculative decoding verifies k+1 tokens per step, so most matrix multiplies have 8 rows or fewer. At that size they are bound by memory bandwidth, not compute, and the generic GEMM paths aren’t built for it. So I wrote three custom HIP kernels for gfx906:
| Kernel | What it does | Result |
|---|---|---|
| Skinny fp16 GEMM | M ≤ 8 matmul with coalesced 16-byte loads and activations staged in LDS | lm_head: 10.3 → 3.4 ms |
| W4 GEMM | 4-bit (AWQ) weights, each dequantized once and reused for every row, with fp32 accumulation | 1.1–1.7× faster per layer, and more accurate |
| Host all-reduce | GPU-to-GPU reduction through host memory, since these cards have no P2P | 72 → 28 µs, bit-identical |
The last one matters more than it looks. One verify step issues about 140 all-reduces. RCCL costs ~72 µs per call whatever the size, so communication alone was eating a large part of every step.
The new kernel works like this: each GPU writes its slice into shared host memory, waits for the other GPU’s flag, and adds the two slices. Over 3,000 random calls its results were bit-identical to RCCL’s.
200k context for coding agents
My real workload is Hermes Agent and opencode, and an agent needs many files in view at once. With an fp16 KV cache the ceiling was ~96k tokens. I tried three options:
| KV cache | Max context | Verdict |
|---|---|---|
| fp16 | ~96k | not enough |
| TurboQuant k8v4 | ~220k | can’t run MTP k=3; decode drops to 13 t/s |
| int8 per-token-head | 200k | the one I use |
The int8 option didn’t even boot at first. A third bug in the fork stopped it: triton_attn.py read a field that vLLM 0.30 removed. A one-line fix.
| Prompt length | Decode speed | Follow-up turn |
|---|---|---|
| short | 70 t/s | — |
| 63k | 29–32 t/s | starts in ~45 s |
| 143k | 17.6 t/s | starts in ~75 s |
The 143k test used a prompt of real source code with a fact hidden at the very top. The model recalled it correctly at the end, and nothing ran out of memory. (The config boots with a 200k limit, but 143k is the longest prompt I’ve actually run.)
Why MTP rather than DFlash2 for agents? DFlash2’s drafter takes 1.8 GB per card and needs prefix caching turned off. Without prefix caching, an agent re-reads everything on every turn.
It worked, but it was slow: a cold 143k-token prompt took 42 minutes to read, and every agent turn waited the better part of a minute. Fixing that is the next part of the story.
Making long context fast
A 200k context is only useful if the model can work in it. So I measured where the time went, piece by piece. Prefill and decode had the same culprit: attention. The int8 KV cache that makes 200k fit is read by vLLM’s Triton kernel, and on gfx906 that kernel reached only 1.1–1.3 TFLOPS. At a 143k prompt, attention was about 80 % of prefill, by my estimate.
Prefill: teach the fork’s fast kernel to read the int8 cache
Two steps:
- Bigger tiles. The Triton kernel handled 16 query rows at a time, so it reloaded and dequantized every block of keys over and over. With 64 rows it’s 1.8× faster, with the same results.
- A better kernel. The fork already ships a fast gfx906 FlashAttention kernel (ported from llama.cpp, using the card’s int8 dot-product instruction), but only for its own cache formats. I taught it to read vLLM’s int8 cache in place: any block size, the per-row scales, no copies. Comparing against a float64 reference turned up two precision issues, so the scale product and the accumulator are now fp32. Result: 4.45 TFLOPS, 3.5× where I started.
| Cold prompt | Before | After |
|---|---|---|
| 7.8k tokens | 254 tok/s | 286 tok/s |
| 31.5k tokens | 155 tok/s | 224 tok/s |
| 62.9k tokens | 90 tok/s | 160 tok/s |
Decode: a new kernel for the shape MTP actually runs
With MTP, every decode step checks 4 tokens, and 6 query heads share each KV head, so 24 query rows read the same keys and values. That makes decode attention compute-bound, and Triton handled it at ~1.1 TFLOPS. I tried 39 Triton configurations, and none was faster than the default.
So I wrote a decode kernel for exactly this shape. Each GPU thread takes one key, converts its int8 row to fp16 once, and computes all 24 scores with the card’s fp16 dot-product instruction. The work is split along the context so 60 compute units stay busy, and the partial results are merged at the end. It allocates nothing and never waits on the CPU, so it runs inside vLLM’s captured GPU graphs. It also covers the MTP draft step.
On its own it is 3.2–4.5× faster than the Triton path. End to end:
| Context | Before | After |
|---|---|---|
| 0.9k | 65.9 t/s | 66.7 t/s |
| 15.9k | 54.2 t/s | 63.8 t/s |
| 30.5k | 42.8 t/s | 57.5 t/s |
| 61.4k | 29.1 t/s | 48.2 t/s |
Agent turns: stop re-reading what’s already cached
An agent sends the whole conversation every turn, and the prefix cache should make only the new part cost anything. On this model it didn’t. Qwen3.8 mixes attention with Gated DeltaNet layers, whose recurrent state vLLM saves only at cache-block boundaries, and with int8 KV those blocks are 1,584 tokens long. With MTP, vLLM also drops the last matched block of every hit. So each turn re-read 2.4–2.9k tokens it had already seen.
The fix was two flags vLLM already had: match prefixes every 16 tokens instead of every block, and save a state checkpoint where sessions that share a prefix diverge. No code change:
On a 16k conversation, turns start in 4.4–5.3 s instead of 12–14.5 s, a repeated request in 1.4 s, and on a ~60k conversation the follow-up turn starts in 6.2 s (it was ~45 s). The replies were byte-identical with and without the change.
Did quality change?
Faster kernels are worthless if the answers drift, so both kernels were checked against a float64 reference, and every change was compared end to end with the old path on the same prompts:
- Prefill kernel: the next-token probabilities after 4k–34k-token prompts match the old path closely (KL divergence ≤ 0.0016), and the greedy answers are identical.
- Decode kernel: 192 of 192 greedy tokens are identical at 3.7k context. At 16k the two texts split at token 14, on a near tie between the top two words (0.047 apart in log-probability); before that point the tokens were identical.
- A fact hidden at the top of a 59k-token prompt was still recalled.
The repo has the kernels, the tests, the raw logs and a build script that reproduces the image I run.
Contributing back
I didn’t want these fixes to stay on my machine, so I prepared a pull request for the fork (KIntegrated/vllm-gfx906-mobydick). It contains:
- the fp16 fix and the window fix, with tests (127 passing);
- an A/B on the fork’s own code: 10.2 → 46.5 t/s with the stock kernels;
- updated docs that had marked DFlash2 as “dead”.
A second, small PR fixes the int8 crash.
- The PR texts and
git am-ready patches:upstream/in the repo. - Everything else (build, serve profiles, kernels, benchmarks, raw logs): github.com/0c33/mi50-vllm-lab
What I learned
- Exactly zero acceptance means a bug, not a weak model. A weak drafter gives low acceptance. Exactly zero means something is broken.
- Re-check “dead” verdicts. Every earlier experiment ran the drafter in broken fp16, so every conclusion built on them was wrong.
- A/B everything. Each fix was measured on its own, so I know what actually made the difference.
- Old hardware has headroom. A 2018 GPU at 90 W runs a 27B model at 74 t/s once the code fits the real workload shape.
- Measure before you optimize. Attention was the bottleneck in both prefill and decode, but for different reasons, and the fix for each came from the measurement, not from a guess.
- Old hardware needs care. These cards are passively cooled and can’t hold 90 W for long. One card hit its thermal limit during an unwatched tuning run. Since then every GPU job runs under a watchdog that stops it at 80 °C.
- An AI agent with clear rules can do serious engineering. You still have to own the goal, review the work and verify the numbers.
Wrapping up
I started with one question: why doesn’t DFlash2 work? I ended up with:
- three bugs fixed;
- five custom kernels, two of them for attention over the int8 KV cache;
- an agent server I use every day with Hermes and opencode, with a 200k-token context on two MI50s that is also fast to work in.
If you have MI50s or any gfx906 card, everything you need is in the repo. Questions are welcome.