eagle3 issue

Welcome to the ATOM wiki!

ATOM Eagle3 Spec Decode 持续负载相变 Bug — 根因分析

日期：2026-05-05 复现状态：100%（7+ 次独立复现） 严重性：高 — gsm8k lm_eval 全跑 acc 跌至 ~47%（vs vLLM 69%） 根因定位：FlyDSL pa_decode_fp8.py:141 _chunk_buffer_resource_for_block 函数在 KV cache > 2GB 时跨 chunk 边界数据错位

---

1. 问题现象

ATOM 在 --method eagle3 --num-speculative-tokens 3 配 Kimi-K2.5-MXFP4 + kimi-k2.5-eagle3，跑 lm_eval gsm8k（concurrent=64, 1319 题）：

• vLLM 同模型同数据：raw accept rate 69%, toks/fwd 3.07
• ATOM 同模型同数据：raw accept rate 47%, toks/fwd 2.46

差距 -22pp。但 gsm8k 答案准确率两边都 ~93%，说明只是接受率差，最终答案对。

关键特征：服务器持续跑约 17-19 wave × 64 并发后，接受率从 80% 断崖跌到 25%，永远回不去（attractor），重启 server 才恢复。

---

2. 触发条件

条件	是否必要
Eagle3 + Kimi K2.5 模型	（未测其他组合，可能更广）
spec decode 开（mtp_k≥1）	✅ 必要
持续 64 并发负载	✅ 必要
KV cache 总大小 > 2 GB	✅ 充要
异质 batch / 同质 batch	❌ 都触发
concurrent ≥ 32	（未测更小并发）
cudagraph 开	❌ --level 0 也触发

崩盘后所有请求都掉到 ~30%（不只触发它的请求）。 单 curl post-crash 也 30%（与并发无关，纯全局状态污染）。 60s 静默后单 curl 仍 30%（污染源不会自然衰减）。

---

3. 实验汇总

3.1 基础复现（7 次）

实验	wave 触发	post-crash acc
全量 gsm8k lm_eval × 2	全跑 1319 题最终 47%	—
100-wave exp_c #1	wave 7（累计）	25%
11-wave + 续 14 wave	wave 17（累计）	22%
30-wave mtp_trace	wave 17	22%
30-wave mtp_k=1	wave 19	31%
30-wave 同质 batch（64×同 prompt）	wave 19	25%
30-wave 限 KV (util=0.4)	不崩	80% (全程)

3.2 排除项

嫌疑	实验	结论
Heterogeneous batch composition	同质 64×d934 也崩 wave 19	❌ 排除
MTP 链（step 1/2 echo）	mtp_k=1 也崩	❌ 排除
Draft KV pool / draft 内部 manager	mtp_k=1 时 draft 不写自己 KV 也崩	❌ 排除
eagle.py:337 slot_mapping off-by-one	trace 显示 slot 数值健康/崩盘期一致	❌ 排除
deferred-output remap (model_runner.py:339)	3625 trace 全部合法，无越界	❌ 排除
cudagraph 状态污染	--level 0 也崩	❌ 排除
Block manager / free-list 复用顺序	全新 fp prompt 也崩	❌ 排除
Slot 内容残留污染	同上	❌ 排除
In-flight buffer 短期污染	60s idle 后单 curl 仍 30%	❌ 排除
Target 模型出问题	trace 显示 T_hs/T_lg/aux_hs 全程稳定	❌ 排除

3.3 token-level 行为变化（崩盘后）

通过 ATOM_DUMP_EAGLE_VERIFY 抓的 verify 记录 + ATOM_DUMP_MTP_TRACE 抓的 propose 内部状态：

指标	健康	崩盘后
Draft step 0 pos0 alive accept rate	83%	49%
Draft step 1 pos1 alive accept rate	82%	18%
Draft step 2 pos2 alive accept rate	70%	21%
Draft top-1 logit 中位数	~17.5	~12.9
step 1 draft == step 0 draft (per req)	26%	63%
Reject 时 target_argmax ∈ draft top-5	72%	51%
Bucket-3 (full mtp_k=3 accept)	50%	31%
Bucket-0 (zero accept)	15%	30%

draft 模型在崩盘后输出的 token 分布从"数字 (2/3/5)"漂到"格式符 (空格/换行/' the'/' =')"，明显是信心降低的 fallback 模式。

3.4 决定性实验：buffer hash 阈值定位

加 ATOM_DUMP_FWDVARS_HASH 每 100 步抓所有 forward_vars buffer 的 amax/sum：

step	block_tables_amax	kv_indices_amax	acc 状态
1100	57,497	919,952	✅ 健康
1200	61,167	978,684	✅ 健康
1300	68,045 ← 跨 2^16	1,088,672 ← 跨 2^20	(transition)
1400	71,644	1,146,312	✅ 健康（最后一步）
1500	75,312	1,204,992	❌ 崩盘

block_id 和 slot_id 同步增长（slot = block × 16 + offset），同时跨过 2^16 / 2^20 阈值。

3.5 限 KV 验证实验

--gpu-memory-utilization 0.4 让 num_kvcache_blocks = 30,091（< 65,536）：

wave 1-30: acc 78-82% 全程稳定，无任何崩盘
block_id 在 [7k, 30k] 区间循环（用满后回收复用）

证明 block_id 跨 65,535 是充要触发条件。

---

4. 根因定位

4.1 嫌疑函数

/root/FlyDSL/kernels/pa_decode_fp8.py 第 141 行 _chunk_buffer_resource_for_block：

def _chunk_buffer_resource_for_block(memref_val, phys_block, block_stride):
    # `buffer_load` multiplies the i32 logical offset by element bytes before
    # issuing the raw buffer instruction, so each descriptor window must stay
    # within signed-i32 byte reach.
    max_chunk_bytes_i64 = fx.Int64(0x7FFF0000)
    one_i64 = fx.Int64(1)
    phys_block_i64 = _widen_nonnegative_i32_to_i64(phys_block)
    block_stride_i64 = _widen_nonnegative_i32_to_i64(block_stride)
    chunk_blocks_i64 = arith.select(
        block_stride_i64 < max_chunk_bytes_i64,
        max_chunk_bytes_i64 // block_stride_i64,    # ← 65,535 for our case
        one_i64,
    )
    chunk_span_bytes_i64 = chunk_blocks_i64 * block_stride_i64
    chunk_idx_i64 = phys_block_i64 // chunk_blocks_i64                # 0 or 1+
    local_phys_block_i64 = phys_block_i64 % chunk_blocks_i64
    chunk_byte_offset_i64 = chunk_idx_i64 * chunk_span_bytes_i64       # jumps 2.14GB
    rsrc = _make_shifted_buffer_resource(memref_val, chunk_byte_offset_i64, chunk_span_bytes_i64)
    local_phys_block_i32 = arith.ArithValue(local_phys_block_i64).trunci(T.i32)
    return rsrc, local_phys_block_i32

4.2 触发数学

draft KV cache shape：[2, 1, num_blocks, 16, 8, 128] × bf16

• per-block stride = 16 × 8 × 128 × 2 = 32,768 bytes

Chunking 触发条件：

chunk_blocks = 0x7FFF0000 // 32,768 = 65,535
KV cache total = num_blocks × 32,768 bytes

num_blocks	KV total	是否触发 chunking	是否崩盘
30,091 (util=0.4)	0.92 GB	❌ 不触发（所有 block 都在 chunk 0）	✅ 不崩
65,536	2.0 GB	⚠️ 边界	（未测）
134,041 (util=0.8)	4.4 GB	✅ 触发（block ≥ 65,535 时跨 chunk）	❌ 崩 wave 17-19

4.3 推测的实际 bug

_chunk_buffer_resource_for_block 的设计意图是：当 KV cache > 2GB 时，因为 buffer descriptor 的偏移字段是 signed i32（最大 ~2GB），单个 descriptor 装不下，所以需要按 chunk 切分，每个 chunk 用一个独立 descriptor。

跨 chunk 时：

• 新 buffer descriptor 的 base shift 到 chunk_idx × chunk_span_bytes 字节位置
• local_phys_block 重置到 chunk 内的本地索引
• _compute_block_base_dw(local_phys_block, ...) 算的是相对新 base 的偏移

最可能的 bug：

• (1) descriptor base shift 与 KV cache 的实际 GPU 内存布局不一致（如果 KV cache 不是按 chunk_span 对齐 alloc 的，跨 chunk 后会读到 padding 或邻接 tensor）
• (2) _make_shifted_buffer_resource 的 num_records_bytes 设置导致 GPU 的 OOB 检测把跨 chunk 的访问当成越界，硬件返回 0 / 未定义值
• (3) 跨 chunk 时调用方某处仍用了原 phys_block（未 chunk）算偏移，与新 descriptor 不匹配 → 地址错位

需要 FlyDSL 团队对照 K/V cache 实际 stride/shape 检查 chunking 逻辑。

---

5. 修复方案

5.1 临时绕过（已验证）

# 让 num_kvcache_blocks × block_stride < 2GB
--gpu-memory-utilization 0.4

对当前配置（MI355X × 8, Kimi-K2.5-MXFP4 + Eagle3, bf16 KV）：

• util=0.4 → 30,091 blocks → 0.92 GB ✅
• util=0.5 → ~37k blocks → ~1.2 GB ✅（推测安全）
• util=0.7 → ~52k blocks → ~1.7 GB ✅（推测安全）
• util=0.8 → 134k blocks → 4.4 GB ❌

安全阈值：num_kvcache_blocks < 65,535。可以让 ATOM 启动时打印警告，或自动 clamp。

5.2 真正修复

需要 FlyDSL 团队修 _chunk_buffer_resource_for_block。两种方向：

方向 A：避免 chunking

• 改用 64-bit 偏移的 buffer descriptor（如果 GFX 架构支持）
• 或用多个 register 存 base + 高位

方向 B：修 chunking bug

• 审查跨 chunk 时的所有偏移计算路径，确保所有依赖 phys_block 的算式都用 local_phys_block + chunk_byte_offset
• 验证 KV cache memory layout 与 chunk_span_bytes 对齐

---

6. 复现脚本

参考 /root/ATOM/case_verify_d934/exp_post_crash_single_curl/run_no_restart.sh。

最小复现：

# 启 server
ATOM_DUMP_PROMPTS=/tmp/atom_prompts.jsonl \
ATOM_DUMP_EAGLE_VERIFY=/tmp/atom_verify.jsonl \
ATOM_DUMP_MAX_RECORDS=300000 \
HSA_NO_SCRATCH_RECLAIM=1 AITER_LOG_LEVEL=WARNING \
nohup python -m atom.entrypoints.openai_server \
  --model /root/models/Kimi-K2.5-MXFP4 --kv_cache_dtype bf16 -tp 8 \
  --port 8000 --trust-remote-code --gpu-memory-utilization 0.8 \
  --method eagle3 --draft-model /root/models/kimi-k2.5-eagle3 \
  --num-speculative-tokens 1 \
  > /tmp/atom.log 2>&1 &

# 等 Uvicorn ready，跑 25 wave 64 并发同 prompt
for w in {1..25}; do
  for i in {1..64}; do
    curl -s -o /dev/null -H "Content-Type: application/json" \
      -d "$(jq -n --rawfile p /tmp/prompt.txt '{model:"/root/models/Kimi-K2.5-MXFP4",prompt:$p,max_tokens:256,temperature:0,stop:["Question:"]}')" \
      http://localhost:8000/v1/completions &
  done
  wait
done

# 看 [MTP Stats] 累积接受率，应该跌到 ~30%
grep "MTP Stats" /tmp/atom.log | tail -3

切换 --gpu-memory-utilization 0.4 → 同样 25 wave，acc 永远稳定 ~80%。

---

7. 实验产物路径

文件	内容
/root/ATOM/case_verify_d934/atom/atom_verify.jsonl	每 verify step 的 per-position dump
/root/ATOM/case_verify_d934/atom/atom_prompts.jsonl	prompt fp → token ids 映射
/tmp/mtp_trace.jsonl	EagleProposer 内部 hidden_state / draft_id / position / context_len trace（per mtp step）
/tmp/fwdvars_hash.jsonl	forward_vars 所有 buffer 每 100 step 的 amax / sum 哈希
/tmp/exp_b/wave_bounds.csv	每 wave 的 verify dump 行号边界
/root/ATOM/case_verify_d934/exp_post_crash_single_curl/	单 curl post-crash 验证实验脚本

8. 添加的诊断代码（ATOM 侧）

文件	改动
atom/utils/envs.py	+ATOM_DUMP_EAGLE_VERIFY / ATOM_DUMP_PROMPTS / ATOM_DUMP_MAX_RECORDS / ATOM_DUMP_MTP_TRACE / ATOM_DUMP_FWDVARS_HASH
atom/spec_decode/eagle.py	EagleProposer 加 verify dump（per req×step）+ MTP 内部 trace
atom/model_engine/scheduler.py	ScheduledBatch 加 prompt_token_ids 字段（dump 用）
atom/model_engine/model_runner.py	postprocess 加 flush_verify_records 调用 + propose_draft_token_ids 加 prepare_dump 调用 + forward 末尾加 fwdvars hash dump

所有改动都是只读 trace，不修改 forward 路径数值（dump-off 时 acc 一致）。

---

9. 摘要

• 现象：ATOM Eagle3 spec decode 在持续 64 并发负载下，约 17-19 wave (~5万 token) 后接受率从 80% 永久跌至 25%
• 触发：KV cache block_id 跨过 65,535 时
• 机制：/root/FlyDSL/kernels/pa_decode_fp8.py:_chunk_buffer_resource_for_block 在 KV > 2GB 触发 chunking 时，跨 chunk 边界数据错位
• 绕过：--gpu-memory-utilization 0.4（让 num_blocks < 65,535）
• 修复：FlyDSL 团队修 chunking 路径
• 验证：限 KV 的 30 wave 全程 80% acc，零崩盘

---

10. 为什么 target MLA 没事，draft MHA 出问题（详细原理）

10.1 GPU Buffer Descriptor 是什么

AMD GPU 读显存需要先创建一个 buffer descriptor（"内存通行证"）：

• base 地址：内存从哪开始
• size：最多能读多少字节

GPU 内核读数据：buffer_load(descriptor, offset)，offset 是 signed int32，最大 0x7FFFFFFF ≈ 2 GB。

关键限制：单个 descriptor 只能覆盖 2 GB 内存范围。

10.2 KV cache 超 2 GB 怎么办：chunking

_chunk_buffer_resource_for_block 的设计意图：当 KV cache > 2 GB 时，把它切成多个 chunk，每个 chunk 一个独立 descriptor，按需要切换。

┌─────── 整个 KV cache (4.4 GB) ───────┐
│ chunk 0 (block 0~65534, ~2 GB)      │ ← descriptor A
├──────────────────────────────────────┤
│ chunk 1 (block 65535~131069, ~2 GB) │ ← descriptor B
├──────────────────────────────────────┤
│ chunk 2 (block 131070~134040, 残余) │ ← descriptor C
└──────────────────────────────────────┘

每个 chunk 装 65,535 个 block（因为 65,535 × 32 KB ≈ 2 GB，正好顶到 i32 上限）。

关键数字：chunk_blocks = 0x7FFF0000 / 32768 = 65,535

10.3 跨 chunk 边界时 bug 发生

GPU 内核读 block 65,535 时：

• Block 65,534（chunk 0 内）：用 descriptor A，offset = 65,534 × 32 KB（i32 内 OK）
• Block 65,535（chunk 1 第一个）：切换到 descriptor B，base 偏移到 2 GB 处，offset = 0

理论上正确。但实际上 跨 chunk 边界那一刻数据错位 → KV 读到错误内存 → draft 输出垃圾。

具体的 bug 在 chunking 实现的某个细节（buffer descriptor 的 base shift 与实际 GPU memory layout 不一致 / num_records 越界检测 / 调用方某处仍用了原 phys_block 算偏移），需要 FlyDSL 团队对照 K/V cache 实际 GPU 内存布局精确审查。

10.4 为什么 target MLA 不触发

因素	Target MLA	Draft Eagle3 MHA
Attention 类型	MLA (Multi-Head Latent)	MHA (Llama-style)
Kernel 路径	aiter_mla.AiterMLABackend → aiter MLA 专用 kernel	aiter_attention.AiterBackend → FlyDSL pa_decode_fp8.py
是否走 _chunk_buffer_resource_for_block	❌ 不经过	✅ 经过
单 block 字节数	~18 KB（latent dim 576 × 2）	32 KB（K + V × 8 head × 128 × 2）
Bug 是否触发	❌ 安全（不走 chunking 函数）	✅ 跨 65k blocks 时炸

target 用的是 MLA 专用 kernel，根本不走 FlyDSL 的 chunking 函数。即使 target KV total = 134k × 18 KB = 2.4 GB（也超过 2 GB），它走另一条 kernel 路径，不踩这个 bug。

10.5 为什么 Eagle3 draft（1 层）特别容易触发

KV cache shape：[2, num_layers, num_blocks, block_size, kv_heads, head_dim]。per-block-stride 与 num_layers 无关（block 在 dim 2，per-block-stride = block_size × kv_heads × head_dim × 2 bytes）。

但 KV cache 是 每层独立张量。每层独立判断"该层是不是超 2 GB"。

模型	num_layers	每层 KV 总大小	是否超 2 GB	是否触发 chunking
普通 60 层 LLM (per layer)	60	134k × 32 KB / 60 ≈ 73 MB	❌ 不超	❌ 不触发
Kimi K2.5 target MLA (per layer)	61	134k × 18 KB / 61 ≈ 40 MB	❌ 不超	❌（且不走这条路径）
Eagle3 draft (per layer)	1	134k × 32 KB / 1 = 4.4 GB	✅ 超	✅ 必触发

Eagle3 draft 的特殊性：因为只有 1 层，所有 KV 容量全砸到这一层 → 单层 4.4 GB → 必触发 chunking → 必踩 bug。

普通多层模型每层 KV cache 平摊容量后通常 < 2 GB，从来不触发这条路径，所以 bug 一直没被发现。

10.6 类比帮助理解

你有一本超过 1000 页的书（4.4 GB KV），但你的书签只能标到 999 页（descriptor offset < 2 GB 上限）。所以你需要把书撕成两半（chunk），分别用两个书签。换书签时如果记错了页码，你就翻到别人的书去了。

• Target 用的是另一种"压缩本"（MLA），整本不到 1000 页 + 用另一种翻页系统，根本不需要撕开。
• 普通 60 层模型用的也是"完整本"，但分成 60 个分册，每册不到 1000 页，每册一个书签就够用。
• Draft Eagle3 用的是"完整本 + 单分册 4400 页"，必须撕开换书签 → 必踩书签错位 bug。

10.7 Bug 暴露的本质条件（精确版）

单层 KV cache 总大小 > 2 GB                                ← 触发 chunking
    AND
该层走的 attention kernel 是 FlyDSL pa_decode_fp8.py        ← 才会跨 chunk 出错

模型	条件 1 (单层 > 2GB)	条件 2 (走 FlyDSL)	触发 bug
普通 60 层 LLM	❌	✅	❌
Kimi K2.5 Target MLA	✅	❌	❌
Kimi K2.5 Eagle3 Draft	✅	✅	✅

这解释了为什么这个 bug 之前没被发现：必须满足"少数大层 + paged_attention 路径"这种少见组合。Eagle3 draft 正好踩到。

10.8 推论：什么模型也可能踩

任何满足以下条件的模型 + ATOM 部署都可能踩这个 bug：

• Draft model 使用 MHA (非 MLA) attention
• Draft model layer 数很少（1-2 层）
• KV cache 容量很大（多 GPU 大显存）

包括：

• Eagle / Eagle2 / Eagle3 各种 draft（一般 1-2 层）
• MTP head 很少层的模型
• 任何 plain Llama-style draft

—

11. 给 FlyDSL / aiter 团队的精确 issue

Title: pa_decode_fp8.py _chunk_buffer_resource_for_block: data corruption when block_id crosses chunk boundary

Repro:
  - 8x MI355X, ATOM + Kimi-K2.5-MXFP4 + Eagle3 draft
  - --gpu-memory-utilization 0.8 → num_kvcache_blocks=134041
  - Eagle3 draft single layer, K+V per-block stride = 32,768 bytes
  - chunk_blocks = 0x7FFF0000 / 32768 = 65,535
  - Run 25 wave × 64 concurrent same-prompt requests
  - Crash: when block_tables_amax crosses 65,535, draft accept rate drops 80% → 25% permanently

Workaround: --gpu-memory-utilization 0.4 (num_blocks=30k < chunk_blocks)

Hypothesis: when phys_block ≥ chunk_blocks (65535), the new buffer descriptor's
base shift to chunk_idx * chunk_span_bytes (= 0x7FFF0000) doesn't match the
GPU memory layout of the contiguous KV cache tensor, causing buffer_load to
read wrong physical addresses.

Affected scenarios: any draft model with single attention layer + KV cache > 2GB.
Target MLA path doesn't trigger (uses different aiter MLA kernel, doesn't go
through pa_decode_fp8 chunking).

Reference function: /root/FlyDSL/kernels/pa_decode_fp8.py:141 _chunk_buffer_resource_for_block