RFC-0034: SGL-JAX GDN 路径统一对比与移植决策

概述

本提案定义一套分层、可复现且防止“假快”的 GDN benchmark，用统一输入和正确性契约比较：

SGL-JAX 当前实现；
SGL-JAX PR #1466 的 GDN chunkwise 路径；
primatrix/pallas-kernel 的 KDA/GDN 相关实现；
vllm-project/tpu-inference 的 GDN 实现。

比较将从 recurrence/KDA math core、fused GDN core 和完整单层逐级推进。Prefill 与 decode 可以分别选择不同赢家；只有在完整 GDN layer 和实际模型性能上仍占优的路径，才进入隔离的 SGL-JAX 移植分支。第一阶段使用 Qwen3.5-35B-A3B 和 TPU v6e-16，最终再扩展到 Qwen3.5-397B-A17B 的 BF16 与 FP8 版本。

背景与动机

SGL-JAX 的既有 GDN 优化将 prefill recurrence 适配到 chunk_kda，已经证明 chunkwise recurrence 能显著降低逐 token recurrence 的成本。然而，后续分析也表明：

同一个 kernel 在独立 Stage、完整 KDA、完整 GDN layer 和端到端 serving 中的收益可能不同；
为解决长序列容量问题引入的布局转换或 HBM materialization，可能在较短序列造成回归；
primatrix/pallas-kernel 与 vllm-project/tpu-inference 的实现边界、gate 语义、融合程度和 benchmark 口径并不相同；
只比较某个 Pallas kernel 的延迟，无法回答“哪条路径最适合 SGL-JAX 用户”。

因此，本提案不预设某个仓库一定胜出，也不把历史实验结果直接当作新比较的排名证据。所有候选都必须在固定 commit、固定环境和统一 harness 下重新验证；仓库自带 benchmark 只作为 native calibration 和接口核验。

目标

建立可复用的 GDN 分层 benchmark harness，覆盖 math core、fused GDN core、完整单层、完整模型和 serving。
在同一硬件、shape、dtype、输入、state 和计时规则下比较四条候选路径。
分别为 prefill 与 decode 选择最佳路径，不强制两阶段共用同一实现。
以完整 GDN layer 和实际模型性能作为最终选择依据，而不是以 isolated kernel 的单点结果决定移植。
在隔离实验分支中移植赢家并保留 fallback；是否进入正式 SGL-JAX PR，由 benchmark 结果和 mentor review 决定。
形成可复核的 raw JSON、manifest、XProf、对比表、迁移风险和决策报告。

成功标准

四个候选的代码 revision、依赖环境、输入、输出和产物可追溯。
所有进入性能排名的 case 均通过统一正确性、有限值和 anti-fallback 门禁。
Prefill 与 decode 都有 kernel/core、完整 layer 和模型层证据。
被选路径在主要完整 GDN core/layer workload 上相对基准改善至少 10%，其他重要 workload 回归不超过 5%。
移植到 SGL-JAX 后，模型层主要指标改善至少 5%，端到端 serving 不回归；差异小于 5% 时优先选择更简单、维护成本更低的路径。
报告明确区分统一 benchmark 结果、各仓库 native benchmark 结果和历史背景数据。

非目标

第一阶段不要求完成 Qwen3.5-397B-A17B 的最终 BF16/FP8 测试。
第一轮比较不修改候选仓库的生产 kernel；只允许为统一接口新增 thin adapter。
不把 pallas-kernel 的未合并实验 PR 当作正式候选；首轮只使用实验开始时固定的 main。
不在 benchmark 阶段直接把候选实现合入 SGL-JAX main。
不用旧的 #324/#325/#1465/#1466 数字替代本提案的新 A/B 数据。
不为了获得更好排名而更换 workload、删除样本、放宽容差或让失败候选静默 fallback。

方案

1. 候选定义与版本冻结

实验 preflight 必须把候选解析为不可变 commit，并记录源文件 SHA：

ID	候选	固定规则	角色
S0	`sgl-project/sglang-jax`	实验开始时固定 `main` commit	SGL-JAX 当前基准
S1	SGL-JAX PR #1466	实验开始时固定 PR head commit	现有 GDN chunkwise 候选
P	`primatrix/pallas-kernel`	实验开始时固定 `main` commit	Primatrix KDA/GDN 候选
V	`vllm-project/tpu-inference`	实验开始时固定 `main` commit	TPU Inference GDN 候选

Preflight 必须完成只读代码审计，逐一回答：

实现是否直接使用 KDA recurrence，还是使用 GDN 专用 kernel。
Conv1D、gate activation、recurrence、state gather/scatter 的融合边界。
Prefill、decode、fresh state、continuing state 和 varlen/ragged 的支持情况。
输入布局、sharding、state contract、dtype 和 accumulation precision。
是否存在 native benchmark、正确性测试、XProf 脚本或已知 fallback。
适配 SGL-JAX 需要的 layout conversion 和潜在额外 HBM 流量。

审计结论属于 Phase 0 交付物。未完成审计前不得实现 adapter 或开始性能运行。

2. 分层 benchmark

统一 benchmark 分为五层：

层级	名称	计时范围	决策作用
L1	Recurrence / KDA math core	已预格式化 Q/K/V/gate/beta/state；只计 recurrence math core	定位 kernel 上限和算法成本
L2	Fused GDN core	Conv1D、activation、recurrence、state I/O，以及进入候选 kernel 所需 layout conversion	判断真实 GDN core 成本
L3	Full GDN layer	L2 加 input projection、normalization、output projection，使用真实 35B 单层权重	第一阶段最终候选选择
L4	Full-model offline	完整模型、固定 token IDs；不含 HTTP、tokenizer 和 scheduler	验证移植后的模型实际性能
L5	Serving end-to-end	SGL-JAX server 与现有 serving benchmark	验证用户可感知吞吐和延迟

L1–L3 用统一 harness 比较全部候选。L4–L5 只比较经过 L1–L3 选择并移植到隔离 SGL-JAX 分支的赢家与 S0。

Prefill 与 decode 在 L3 后分别决策：

text

Prefill winner = best eligible candidate on prefill L2/L3
Decode winner  = best eligible candidate on decode L2/L3

二者可以来自不同仓库。若采用混合方案，L4/L5 必须验证 prefill→decode state continuity 和 dispatch 正确性。

3. 统一 workload

第一阶段固定：

模型：Qwen3.5-35B-A3B；
硬件：TPU v6e-16，4×4；
dtype：BF16；
SGL-JAX 并行配置：TP16 / DP4 / EP16；
同一物理 slice 串行执行候选；
固定 random seed；
每个候选使用独立、初始为空的 compilation cache。

预期的 35B per-shard GDN shape 为：

text

n_kq = 1
n_v = 2
d_k = d_v = 128
mixed_qkv = [T, 512]

该 shape 只是设计输入，实验 preflight 必须从固定模型 revision 的 config 和实际 projection 切分重新推导，并写入 manifest。推导不一致时停止，不得沿用上述预期值。

3.1 性能矩阵

Prefill：

ID	T	State	目的
P4K-F	4096	fresh	常用短上下文
P4K-C	4096	continuing	state restore/update
P16K-F	16384	fresh	中长上下文
P16K-C	16384	continuing	中长上下文 state 路径
P32K-F	32768	fresh	容量和长上下文
P32K-C	32768	continuing	长上下文 state 路径

Decode：

ID	Batch / concurrency	State	目的
D1	1	continuing	单请求单 token 延迟
D16	16	continuing	持续生成吞吐

3.2 正确性边界矩阵

正确性单独覆盖：

T = 1, 63, 64, 65, 127, 128, 129；
zero-length request；
ragged / packed requests；
partial chunk；
fresh、continuing 和混合 state；
prefill→decode continuation；
dummy slot、unused slot 和 slot isolation。

正确性边界 case 不自动进入性能统计。任何候选若在边界 case 失败，必须先判定该失败是否命中目标使用 contract；命中时候选失去排名资格。

4. 计时与运行规则

Warmup：每个独立 case 3 次。
L1/L2：每轮 30 个 measured samples。
L3：每轮 10 个 measured samples。
每个 case 至少 2 个独立 round。
四个 rank 分别保存原始样本，不先聚合再丢弃。
每次计时必须显式 block 到 device 完成，并消费 output 与 final state。
统计报告同时给出 per-rank raw samples、per-rank median、pooled median、min/max 和离散度。
不得静默删除 outlier；若测量无效，整轮标记失败并保留证据。

为减小顺序和温度偏差，候选顺序固定为：

text

Round 1: S0 → S1 → P → V
Round 2: V → P → S1 → S0

每个 variant 的进程、venv、cache、日志、XProf 和 output 路径相互隔离。前一个候选残留进程、端口或 cache 未清理完成时，不得启动下一个候选。

5. Native benchmark 与统一 benchmark 的关系

统一 harness 是候选排名的唯一直接依据。各仓库自带 benchmark 仍需执行，但用途不同：

验证调用方式和环境与仓库维护者预期一致；
复现候选仓库公开或自带的 baseline；
发现统一 adapter 是否遗漏重要 fast path；
记录 native result 与 unified result 的口径差异。

Native benchmark 结果不得与统一 harness 数值直接拼表排名。报告必须分成：

Unified L1/L2/L3；
Repository-native benchmark；
Ported SGL-JAX L4/L5。

SGL-JAX 的现有 .claude/skills/model-speed-benchmark 用于 L5 serving 验证，不代替 L1–L3；L4 full-model offline 是否已有可复用 runner，必须在 Phase 0 审计中确认。若没有，只实现完成决策所需的最小 runner。

6. 备选方案

只运行各仓库 native benchmark：执行成本最低，但输入、计时边界、依赖和统计方法不同，不能支持公平排名。Native benchmark 因此保留为 calibration，不作为唯一决策证据。
只比较 recurrence/KDA kernel：能快速判断理论上限，但会漏掉 layout conversion、Conv1D、gate activation、state I/O 和 projection，无法代表完整 GDN layer，故 L1 只用于定位。
直接把 tpu-inference 或 pallas-kernel 整套实现复制进 SGL-JAX：可能携带训练、backward、CP 或其他不需要的依赖，也难以隔离真实收益。主方案先比较，再做最小移植。
第一轮强制统一 JAX/libtpu：有利于控制 runtime 变量，但可能使某个候选离开其维护者支持的依赖组合并产生假失败。主方案先使用各仓库 native pinned environment，并把统一兼容 runtime 作为可选第二轮。
直接使用 397B 做首轮选择：最贴近最终目标，但资源、编译和调试成本过高。主方案先用 35B 完成接口与决策闭环，再分别验证 397B BF16/FP8。
Prefill 与 decode 强制选择同一路径：实现较简单，但可能牺牲一个阶段的真实性能。主方案允许分别选择，并在 L4/L5 验证跨阶段 state continuity。

7. 选择与移植

7.1 候选资格

候选必须同时满足：

统一 reference correctness 通过；
output、final recurrent state 和 final Conv1D state 为有限值；
requested implementation 等于 effective implementation；
无 silent fallback；
XProf 中出现预期候选 kernel/op；
artifact manifest 和 SHA readback 完整。

不满足时，候选可以保留诊断报告，但不能进入性能排名。

7.2 选择规则

对 prefill 和 decode 分别应用：

L1 只判断 math core 上限，不直接决定赢家。
L2 判断 integration conversion、Conv1D 和 state I/O 是否抵消 kernel 收益。
L3 作为第一阶段主要选择依据。
主要 L2/L3 workload 改善至少 10%，其他重要 workload 回归不超过 5%。
两个候选的 L3 差异小于 5% 时，优先选择接口简单、SGL-JAX 改动少、fallback 清楚且维护成本低的候选。

7.3 移植规则

赢家先进入独立分支：

text

experiment/gdn-path-comparison

移植要求：

保留 S0 reference fallback；
selector 在初始化时冻结，不在 hot path 读取环境变量；
requested/effective/fallback reason 进入日志和 manifest；
只移植赢家所需的最小生产代码；
不直接复制候选仓库不需要的训练、backward、CP 或其他模型能力；
不更改 scheduler、state-pool schema 或无关 backend；
未经 mentor 确认，不创建正式 SGL-JAX PR。

8. 最终模型计划

第一阶段完成 35B 选择和移植闭环。之后再执行：

Qwen/Qwen3.5-397B-A17B，BF16；
Qwen/Qwen3.5-397B-A17B-FP8。

两种模型分别验收，不允许用 BF16 结果推断 FP8，也不允许用 FP8 结果替代 BF16。397B 阶段需要重新确认：

模型 config 和 per-shard GDN shape；
权重路径与 revision；
TPU topology、TP/DP/EP；
内存容量；
FP8 quantization contract；
compilation cache 和 warmup；
L4/L5 workload。

397B 不属于本周第一阶段的完成门槛。

设计细节

1. 统一接口

所有 adapter 暴露逻辑等价的接口；精确 Python 类型在 Phase 0 源码审计后冻结：

python

class GDNInputs:
    mixed_qkv
    raw_gate
    beta
    conv_state
    recurrent_state
    cu_seqlens
    state_indices
    track_indices
    metadata

class GDNOutputs:
    output
    final_conv_state
    final_recurrent_state
    aux

Adapter 只负责：

layout / axis conversion；
candidate API 参数转换；
state gather/scatter contract 对齐；
输出恢复为统一 layout；
provenance 和 effective-path 记录。

Adapter 不得：

重写候选 kernel；
提前计算候选本应融合的核心 math；
隐藏 fallback；
把 layout conversion 排除在 L2/L3 计时之外。

2. Gate 与 state 语义

每个候选都必须记录：

gate 输入是 raw parameter 还是 activated decay；
softplus、exp、sigmoid 的唯一执行位置；
A_log、dt_bias 的 shape 与 broadcast；
Q/K normalization 和 scale；
recurrence accumulation dtype；
zero-length 的 final-state 行为；
continuing state 的读写与 slot isolation。

禁止 adapter 与 candidate kernel 同时激活 gate。若两个实现采用不同数学形式，必须先证明与统一 reference 等价，再进行性能比较。

3. 独立 reference

统一 reference 使用 pure JAX 实现，关键 reduction 和 recurrence 尽可能采用 FP32 accumulation。它不调用四个候选的 production kernel。

所有候选比较：

output；
final recurrent state；
final Conv1D state；
fresh/continuing；
prefill→decode；
input immutability；
dummy/unused/track slot；
finite values。

BF16 默认容差：

模式	rtol	atol
单设备 / 非 DP	`2e-2`	`1e-2`
DP/TP sharded	`2e-2`	`5e-2`

同时报告 max absolute error、max relative error、mismatch count 和 NaN/Inf count。容差不是自动豁免；token divergence 或 state contract 异常仍需停止。

4. 防止假一致与假加速

每个 rank 的 artifact 至少记录：

repo URL、commit SHA、dirty state；
candidate source file SHA；
adapter SHA；
Python/JAX/libtpu/image 版本；
TPU type/topology/device IDs；
requested/effective implementation 和 fallback reason；
input/output/final-state SHA；
logical/aligned/allocated length；
fixture seed；
cache path 和 cache 初始状态；
raw samples；
profiler run/host/step；
completion marker。

运行时门禁：

使用非零、非恒定、固定 seed 的输入；
output 和 final state 必须被 device-block 并消费；
continuing case 必须证明 state 发生预期变化；
reference 与 candidate 的输入 SHA 一致；
XProf 必须出现预期 candidate kernel；
发现 fallback、NaN/Inf、OOM、rank mapping 异常或 artifact SHA 不一致时立即暂停。

5. XProf

对通过正确性和初步 timing 的候选采集代表性 XProf：

Prefill：4K、16K、32K 中至少包含主要决策 shape；
Decode：D1 和 D16 至少各一项；
L2 与 L3 分开采集或通过明确 trace annotation 隔离；
四个 host 的 trace 都必须非空并可回读；
比较 operator self-time、HBM traffic、Pallas program grid、collective、peak memory、host spread 和 idle。

XProf 用于解释收益和回归，不用于替代端到端 timing。若 XProf coverage 不同，必须披露 normalization 方法。

6. 目录与 artifact

计划中的 tracked harness：

text

benchmark/kernels/gdn_compare/
├── README.md
├── reference.py
├── fixtures.py
├── schemas.py
├── runner.py
├── adapters/
│   ├── sglang_main.py
│   ├── sglang_1466.py
│   ├── primatrix_kda.py
│   └── tpu_inference.py
├── configs/
│   └── qwen35_v6e16.yaml
└── tests/

外部仓库和依赖环境：

text

/tmp/gdn-path-comparison/
├── repos/{sglang-main,sglang-1466,pallas-kernel,tpu-inference}
├── venvs/{s0,s1,p,v}
├── caches/{s0,s1,p,v}
└── staging/

共享 fixture：

text

/data/user/dongbohao/gdn-path-comparison/fixtures/

正式证据：

text

/data/user/dongbohao/gdn-path-comparison/evidence/{run_id}/
├── manifest/
├── raw-json/
├── logs/
├── xprof/
├── comparisons/
├── reports/
└── .complete

只有通过 schema、finite、size、SHA readback 和 completion marker 的 run 才能进入报告。

7. 资源与进程边界

优先复用 perf-16，但执行前必须实时只读确认：

使用授权仍有效；
Job 和四个 Pod 为 Ready；
四个 TPU 节点健康；
没有其他任务占用 server、benchmark、profiler 端口或 TPU；
旧的 #324/#325 evidence 不会被覆盖。

所有新路径使用唯一 run ID。只允许停止本次实验明确标记的进程；不得按宽泛进程名 kill。若 perf-16 不可用，需要通过 Falcon 或新资源执行时，必须单独请求批准。

测试计划

Verification Group A：Harness 单元测试

schema 和 manifest；
fixed-seed fixture 可复现；
candidate adapter 输入/输出 layout；
gate activation 只执行一次；
output/final-state 被消费；
requested/effective/fallback 校验；
artifact local staging、exclusive publish、SHA readback 和 marker；
非 identity PJRT rank permutation；
sample 不得静默丢弃。

Verification Group B：正确性

L1、L2、L3 对 independent reference；
性能矩阵和边界矩阵；
single-device 与 DP/TP sharded；
fresh/continuing、zero-length、ragged、partial chunk；
output、final recurrent state、final Conv1D state；
prefill→decode continuity；
finite 和 state-slot isolation。

Verification Group C：统一性能 A/B

按固定顺序执行两轮；
L1/L2 每轮 30 samples，L3 每轮 10 samples；
四 rank raw samples 和完整 provenance；
先通过 correctness 才生成 performance comparison。

Verification Group D：Native benchmark

使用各仓库固定 commit 的自带 runner；
保存原始命令、默认参数和版本；
只做 calibration，不与统一结果混为一张排名表。

Verification Group E：移植后模型与 serving

L4：完整 Qwen3.5-35B-A3B，固定 token IDs，隔离 model compute；
L5：SGL-JAX server，使用 .claude/skills/model-speed-benchmark 的标准健康检查、并发 sweep 和报告契约；
same-revision reference/optimized A/B；
greedy token IDs 一致，必要时运行 accuracy regression；
primary model metric ≥5% 改善；
serving 不回归，理想改善 ≥5%。

影响范围

第一阶段 tracked 变更只包含 benchmark harness、adapter、测试、配置和文档；候选生产 kernel 保持不变。选出赢家后，隔离分支可能修改 SGL-JAX 的 GDN kernel/backend 和最小依赖 KDA 代码，但范围必须由新的移植计划逐文件列出。

可能受影响的团队：

SGL-JAX inference / GDN maintainers；
pallas-kernel maintainers；
TPU inference kernel maintainers；
TPU profiling / infra 使用者。

本 RFC 不要求其他仓库接受 SGL-JAX 的 adapter，也不把实验分支变成长期 fork。

实施计划

每阶段结束后暂停并提交检查点；未经批准不自动进入下一阶段。

Phase	Depends on	内容	产物	进入下一阶段门禁
0	无	四仓库只读审计、固定 commits、接口和 native benchmark 盘点	discovery report、commit/source SHA、接口表	用户批准候选与实际接口
1	0	独立 worktree、目录、schema、fixture、reference	harness RED→GREEN、manifest contract	本地单元测试和 `git diff --check` 通过
2	1	S0/S1/P/V thin adapters	adapter source、adapter contract tests	无 production kernel 修改；正确性 smoke 通过
3	2	L1 unified correctness + performance	raw JSON、comparison、XProf pilot	候选资格门禁通过
4	3	L2 fused GDN core	core A/B、conversion breakdown、XProf	full core 改善/回归可解释
5	4	L3 full GDN layer	layer A/B、prefill/decode decision table	分别确定 prefill/decode winner
6	2	Native benchmark calibration	四仓库 native reports	与 unified 差异有书面解释
7	5, 6	赢家移植到隔离 SGL-JAX 分支	fallback、测试、migration report	mentor 批准进入模型验证
8	7	L4 35B full-model offline	model A/B、token/state correctness	primary metric ≥5%，无重要回归
9	8	L5 35B serving	model-speed-benchmark report	serving 不回归，token IDs/accuracy 通过
10	9	397B BF16/FP8 后续验证	两套独立报告	另行批准资源与矩阵
11	9；需要最终模型结论时再依赖 10	正式代码决策	mentor one-pager、PR proposal 或停止结论	mentor 明确决定是否提交正式 PR

停止条件

出现以下任一情况立即暂停，保留日志，不做猜测性修复：

固定 commit 或模型 shape 无法复现；
需要修改候选 production kernel 才能完成首轮比较；
independent reference 不一致；
requested/effective implementation 不一致或 silent fallback；
NaN/Inf、OOM、rank mapping、state contract 或 token divergence；
artifact publish/SHA/marker 失败；
性能离散度异常且无法由预设规则解释；
perf-16 有资源冲突；
需要创建 Falcon/GKE/SkyPilot 新资源；
需要本 RFC 已批准发布之外的 GitHub 写操作、代码 PR、merge 或 Issue 状态变更。

风险与缓解

风险	缓解
Kernel 快但 integration 慢	L2 把 layout conversion、Conv1D 和 state I/O 纳入计时；L3 再做完整层决策
不同仓库依赖版本导致不可比	首轮 native pinned env；记录全部版本；可选第二轮统一兼容 JAX/libtpu
候选走 fallback 形成假一致	requested/effective/fallback、source SHA、XProf op 三重核验
Gate 语义不同导致假正确	独立 reference、raw/activated gate contract、单次 activation 测试
顺序、温度和 cache 偏差	正反运行顺序、独立空 cache、至少两轮、四 rank raw samples
只在单点 shape 胜出	覆盖 4K/16K/32K、fresh/continuing、decode 1/16 和边界 lengths
过早移植导致范围扩张	Phase 0–6 禁止修改候选 production kernel；Phase 7 另写逐文件计划
历史结果污染新排名	历史数据只放背景；统一实验使用新 run ID 和完整 manifest
397B 资源和内存不确定	作为后续独立阶段，重新批准 topology、权重和 BF16/FP8 矩阵
赢家差异很小却引入复杂代码	L3 差异小于 5% 时优先可维护方案

Review 问题

L1/L2/L3 的边界是否与四个候选的真实接口兼容？
Prefill 与 decode 分别选择是否符合 SGL-JAX 的 dispatch 和 state continuity 约束？
10% layer、5% model 和 5% regression 门槛是否合适？
第一阶段是否还需要加入其他正式候选 commit？
L4 是否已有应优先复用的 SGL-JAX full-model offline runner？
397B BF16/FP8 的资源、模型 revision 和最终 workload 应由谁确认？

RFC-0034 Phase 0 Discovery Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: 在不修改四个候选仓库生产代码、不启动 TPU 的前提下，冻结 S0、S1、P、V 与 Qwen3.5-35B-A3B 的不可变版本，核验真实 GDN/KDA 调用链、接口、native benchmark 和 35B per-shard shape，并形成 Phase 1 adapter 计划所需的完整输入。

Architecture: Phase 0 只在独立 SGL-JAX experiment worktree 中新增 discovery 脚本、机器可读 manifest 和审计报告。固定版本通过 Git remote refs 解析，不依赖本地陈旧 branch；源码索引基于固定 checkout 和 SHA-256；需要判断力的调用链审计由人工阅读完成，再由 schema validator 阻止空字段、未引用源码或未解释的能力结论进入门禁。

Tech Stack: Git、Python 3.12、标准库 json/hashlib/pathlib/subprocess、 pytest、Hugging Face Hub config、Markdown。

Global Constraints

本计划只执行 RFC Phase 0；完成后必须暂停，不得自动进入 Phase 1。
不修改 S0、S1、P、V 的 production kernel，不创建 adapter，不运行性能测试或 XProf。
不启动或停止 perf-16、Falcon、GKE、SkyPilot 或任何 TPU 进程。
S0、P、V 固定实验执行时各自 main 的远端 SHA；S1 固定 refs/pull/1466/head 的远端 SHA。
模型固定为 Qwen/Qwen3.5-35B-A3B 的不可变 Hugging Face revision；只下载 config/tokenizer metadata，不下载权重。
所有外部仓库 checkout 必须位于 /tmp/gdn-path-comparison/repos/，且 checkout 后 git status --porcelain 为空；绝对路径只写入 git-ignored runtime manifest，不得进入 tracked artifact。
所有 tracked Phase 0 文件必须位于 benchmark/kernels/gdn_compare/discovery/ 或 benchmark/kernels/gdn_compare/tests/。
不用 #324/#325/#1465/#1466 的历史性能数值作为新候选排名证据。
任何远端 ref、模型 revision、真实调用链或 shape 无法核验时，大声失败并暂停。

Task 1: 创建隔离 worktree 并冻结候选与模型 revision

Files:

Create: benchmark/kernels/gdn_compare/discovery/freeze_candidates.py
Create: benchmark/kernels/gdn_compare/discovery/candidates.json
Create: benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py

Interfaces:

Consumes: Git remote refs、Hugging Face model metadata。
Produces: tracked candidates.json，只包含四个候选的 repo/ref/commit/checkout key 和模型 repo/revision；git-ignored /tmp/gdn-path-comparison/runtime-checkouts.json 保存本机绝对路径。后续所有 Phase 0 任务只读取这两个冻结文件，不重新解析 floating ref。
[ ] Step 1: 先冻结 S0 SHA，再从该 SHA 创建隔离 worktree

在本地 sglang-jax 仓库执行：

bash

S0_SHA=$(
  git ls-remote https://github.com/sgl-project/sglang-jax.git \
    refs/heads/main | awk '{print $1}'
)
test "${#S0_SHA}" -eq 40
git fetch origin "$S0_SHA"
git worktree add \
  .claude/worktrees/experiment/gdn-path-comparison \
  -b experiment/gdn-path-comparison \
  "$S0_SHA"
cd .claude/worktrees/experiment/gdn-path-comparison
test "$(git rev-parse HEAD)" = "$S0_SHA"
printf '%s\n' "$S0_SHA" > /tmp/gdn-path-comparison-s0-sha
git status --short --branch

Expected:

text

## experiment/gdn-path-comparison

若目标 branch/worktree 已存在，先只读执行 git worktree list --porcelain 和 git status --short --branch；仅在 HEAD 与本步骤冻结的 S0_SHA 一致且 worktree clean 时复用，不删除或重置既有 worktree。若不一致则暂停，不重置既有分支。

[ ] Step 2: 写候选 manifest 的失败测试

test_discovery_manifest.py 必须先包含：

python

import json
import re
from pathlib import Path


MANIFEST = (
    Path(__file__).parents[1] / "discovery" / "candidates.json"
)
SHA_RE = re.compile(r"^[0-9a-f]{40}$")


def test_candidate_manifest_is_frozen_and_complete():
    data = json.loads(MANIFEST.read_text())
    assert data["schema_version"] == 1
    assert set(data["candidates"]) == {"S0", "S1", "P", "V"}
    for key, candidate in data["candidates"].items():
        assert candidate["repo_url"].startswith("https://github.com/")
        assert SHA_RE.fullmatch(candidate["commit"])
        assert candidate["checkout_key"] == key.lower()
        assert "checkout" not in candidate
        assert not any(value.startswith("/") for value in candidate.values())
    assert data["model"]["repo_id"] == "Qwen/Qwen3.5-35B-A3B"
    assert SHA_RE.fullmatch(data["model"]["revision"])


def test_candidate_refs_match_the_approved_contract():
    data = json.loads(MANIFEST.read_text())
    assert data["candidates"]["S0"]["ref"] == "refs/heads/main"
    assert data["candidates"]["S1"]["ref"] == "refs/pull/1466/head"
    assert data["candidates"]["P"]["ref"] == "refs/heads/main"
    assert data["candidates"]["V"]["ref"] == "refs/heads/main"

[ ] Step 3: 运行测试并确认 RED

Run:

bash

uv run --with pytest python -m pytest \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py -q

Expected: FAIL，因为 candidates.json 尚不存在。

[ ] Step 4: 实现不可变 ref 解析和 checkout

freeze_candidates.py 使用以下完整候选定义：

python

from __future__ import annotations

import argparse
import json
import subprocess
from datetime import datetime, timezone
from pathlib import Path

from huggingface_hub import HfApi


CANDIDATES = {
    "S0": (
        "https://github.com/sgl-project/sglang-jax.git",
        "refs/heads/main",
        "sglang-main",
    ),
    "S1": (
        "https://github.com/sgl-project/sglang-jax.git",
        "refs/pull/1466/head",
        "sglang-1466",
    ),
    "P": (
        "https://github.com/primatrix/pallas-kernel.git",
        "refs/heads/main",
        "pallas-kernel",
    ),
    "V": (
        "https://github.com/vllm-project/tpu-inference.git",
        "refs/heads/main",
        "tpu-inference",
    ),
}
MODEL_ID = "Qwen/Qwen3.5-35B-A3B"


def run(*args: str, cwd: Path | None = None) -> str:
    return subprocess.run(
        args,
        cwd=cwd,
        check=True,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    ).stdout.strip()


def resolve_ref(repo_url: str, ref: str) -> str:
    rows = run("git", "ls-remote", repo_url, ref).splitlines()
    if len(rows) != 1:
        raise RuntimeError(
            f"expected one remote row for {repo_url} {ref}, got {rows}"
        )
    commit, resolved_ref = rows[0].split()
    if resolved_ref != ref or len(commit) != 40:
        raise RuntimeError(f"invalid ref resolution: {rows[0]}")
    return commit


def checkout(repo_url: str, commit: str, path: Path) -> None:
    if not path.exists():
        run("git", "clone", "--filter=blob:none", "--no-checkout", repo_url, str(path))
    run("git", "fetch", "--depth=1", "origin", commit, cwd=path)
    run("git", "checkout", "--detach", commit, cwd=path)
    if run("git", "status", "--porcelain", cwd=path):
        raise RuntimeError(f"dirty checkout: {path}")
    if run("git", "rev-parse", "HEAD", cwd=path) != commit:
        raise RuntimeError(f"HEAD mismatch: {path}")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--runtime-output", type=Path, required=True)
    parser.add_argument("--s0-commit", required=True)
    args = parser.parse_args()
    if len(args.s0_commit) != 40:
        raise RuntimeError(f"invalid frozen S0 commit: {args.s0_commit}")
    repos_root = args.root / "repos"
    repos_root.mkdir(parents=True, exist_ok=True)

    frozen = {}
    runtime = {"schema_version": 1, "checkouts": {}}
    for key, (repo_url, ref, dirname) in CANDIDATES.items():
        commit = args.s0_commit if key == "S0" else resolve_ref(repo_url, ref)
        path = repos_root / dirname
        checkout(repo_url, commit, path)
        frozen[key] = {
            "repo_url": repo_url,
            "ref": ref,
            "commit": commit,
            "checkout_key": key.lower(),
        }
        runtime["checkouts"][key.lower()] = {
            "path": str(path.resolve()),
            "commit": commit,
            "dirty": False,
        }

    model_revision = HfApi().model_info(MODEL_ID).sha
    if model_revision is None or len(model_revision) != 40:
        raise RuntimeError(f"invalid model revision: {model_revision}")

    result = {
        "schema_version": 1,
        "captured_at_utc": datetime.now(timezone.utc).isoformat(),
        "candidates": frozen,
        "model": {"repo_id": MODEL_ID, "revision": model_revision},
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    args.runtime_output.parent.mkdir(parents=True, exist_ok=True)
    args.runtime_output.write_text(
        json.dumps(runtime, indent=2, sort_keys=True) + "\n"
    )


if __name__ == "__main__":
    main()

[ ] Step 5: 生成 manifest 并验证 GREEN

Run:

bash

uv run --with huggingface_hub \
  python benchmark/kernels/gdn_compare/discovery/freeze_candidates.py \
  --root /tmp/gdn-path-comparison \
  --output benchmark/kernels/gdn_compare/discovery/candidates.json \
  --runtime-output /tmp/gdn-path-comparison/runtime-checkouts.json \
  --s0-commit "$(cat /tmp/gdn-path-comparison-s0-sha)"
uv run --with pytest python -m pytest \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py -q

Expected: 2 passed，tracked manifest 不含绝对路径；runtime manifest 中四个 checkout 均为 detached、clean、HEAD 与 tracked manifest commit 一致。

[ ] Step 6: 提交 revision freeze

bash

git add \
  benchmark/kernels/gdn_compare/discovery/freeze_candidates.py \
  benchmark/kernels/gdn_compare/discovery/candidates.json \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py
git commit -m "test(gdn): freeze comparison candidates"

Task 2: 生成固定源码索引和 SHA 证据

Files:

Create: benchmark/kernels/gdn_compare/discovery/index_sources.py
Create: benchmark/kernels/gdn_compare/discovery/source-index.json
Modify: benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py

Interfaces:

Consumes: Task 1 的 tracked candidates.json、git-ignored runtime-checkouts.json 和四个 frozen checkout。
Produces: 按候选分组的相关源码、测试、benchmark、skill/docs 路径及 SHA-256，供调用链审计逐条引用。
[ ] Step 1: 为源码索引增加失败测试

追加：

python



SOURCE_INDEX = (
    Path(__file__).parents[1] / "discovery" / "source-index.json"
)


def test_source_index_has_verifiable_files_for_every_candidate():
    data = json.loads(SOURCE_INDEX.read_text())
    assert set(data["candidates"]) == {"S0", "S1", "P", "V"}
    for key, files in data["candidates"].items():
        assert files, f"{key} has no relevant source hits"
        for item in files:
            assert item["checkout_key"] == key.lower()
            path = Path(item["relative_path"])
            assert not path.is_absolute()
            assert ".." not in path.parts
            assert "absolute_path" not in item
            assert len(item["sha256"]) == 64
            assert item["category"] in {
                "gdn",
                "kda",
                "model",
                "test",
                "benchmark",
                "documentation",
            }

[ ] Step 2: 确认测试 RED

Run:

bash

uv run --with pytest python -m pytest \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py \
  -q -k source_index

Expected: FAIL，因为 source-index.json 尚不存在。

[ ] Step 3: 实现确定性索引器

索引器只读取以下文本后缀：

python

TEXT_SUFFIXES = {".py", ".md", ".rst", ".toml", ".yaml", ".yml", ".json"}
KEYWORDS = {
    "gdn": ("gdn", "gated_delta", "gated delta"),
    "kda": ("kda", "chunk_kda", "delta_rule"),
    "model": ("qwen3_5", "qwen3.5", "qwen35"),
}

分类优先级固定为路径中的 test、benchmark、docs，再按内容关键字：

python

def classify(relative: str, text: str) -> str | None:
    lowered_path = relative.lower()
    lowered_text = text.lower()
    if "test" in lowered_path:
        return "test"
    if "bench" in lowered_path:
        return "benchmark"
    if "doc" in lowered_path or relative.endswith((".md", ".rst")):
        if any(word in lowered_text for words in KEYWORDS.values() for word in words):
            return "documentation"
    for category in ("gdn", "kda", "model"):
        if any(
            word in lowered_path or word in lowered_text
            for word in KEYWORDS[category]
        ):
            return category
    return None

完整遍历必须：

跳过 .git、venv、cache、build 和生成目录；
从 runtime-checkouts.json 解析本机根目录，但对每个 tracked 命中只记录 candidate、checkout key、relative path、category、size 和 SHA-256；
以 relative path 排序；
输出 tracked manifest 中的 candidate commit，并在生成时断言 runtime checkout HEAD 没有漂移。
[ ] Step 4: 生成索引并验证 GREEN

Run:

bash

python benchmark/kernels/gdn_compare/discovery/index_sources.py \
  --candidates benchmark/kernels/gdn_compare/discovery/candidates.json \
  --runtime-checkouts /tmp/gdn-path-comparison/runtime-checkouts.json \
  --output benchmark/kernels/gdn_compare/discovery/source-index.json
uv run --with pytest python -m pytest \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py -q

Expected: 全部测试 PASS，四个候选均至少有一个源码命中，所有 SHA 回读一致。

[ ] Step 5: 提交源码索引

bash

git add benchmark/kernels/gdn_compare/discovery/index_sources.py \
  benchmark/kernels/gdn_compare/discovery/source-index.json \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py
git commit -m "docs(gdn): index candidate implementations"

Task 3: 审计真实调用链、数学语义和 adapter 边界

Files:

Create: benchmark/kernels/gdn_compare/discovery/interface-audit.md
Create: benchmark/kernels/gdn_compare/discovery/interface-audit.json
Create: benchmark/kernels/gdn_compare/discovery/validate_audit.py
Modify: benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py

Interfaces:

Consumes: frozen checkout、source-index.json。
Produces: S0/S1/P/V 的 prefill/decode 调用链、KDA 使用情况、融合边界、gate/state contract 与最小 adapter boundary；这是 Phase 1–2 计划的权威输入。
[ ] Step 1: 为审计 schema 写失败测试

每个候选 JSON 必须含有以下非空字段：

python

REQUIRED_AUDIT_FIELDS = {
    "uses_kda",
    "prefill_entrypoint",
    "decode_entrypoint",
    "call_chain",
    "conv1d_boundary",
    "gate_contract",
    "recurrence_contract",
    "state_io_contract",
    "supported_modes",
    "input_layout",
    "output_layout",
    "accumulation_dtype",
    "fallback_behavior",
    "adapter_requirements",
    "source_evidence",
}

测试还必须断言：

uses_kda 是 boolean；
每个 source_evidence 路径与 SHA 存在于 source-index.json；
call chain 至少包含入口和 recurrence 两个 symbol；
prefill 与 decode 明确标记 supported 或 unsupported，不能留空；
gate_contract 明确 raw/activated 和 softplus/exp/sigmoid 的执行位置；
state_io_contract 明确 fresh/continuing/zero-length/final-state。
[ ] Step 2: 运行 schema test 并确认 RED

Run:

bash

uv run --with pytest python -m pytest \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py \
  -q -k audit

Expected: FAIL，因为审计产物尚不存在。

[ ] Step 3: 按候选逐个做源码审计

只使用固定 checkout，优先采用仓库 CodeGraph（存在 .codegraph/ 时）：

bash

codegraph explore \
  "GDN prefill decode KDA call chain gate activation state update benchmark"

没有 .codegraph/ 时使用 source index 给出的文件，并对 symbol 执行：

bash

uv run python -c \
  'import json; d=json.load(open("/tmp/gdn-path-comparison/runtime-checkouts.json")); print(*[v["path"] for v in d["checkouts"].values()], sep="\n")' \
| while IFS= read -r checkout; do
    rg -n \
      "GDN|gated_delta|chunk_kda|delta_rule|prefill|decode|conv|state|softplus|sigmoid" \
      "$checkout"
  done

每条事实必须在 interface-audit.json 的 source_evidence 中引用固定 checkout 的 relative path、symbol/line 和 Task 2 SHA。interface-audit.md 将机器数据展开成以下统一表：

text

Candidate
├── Prefill entry → gate/conv → recurrence → state write
├── Decode entry  → gate/conv → recurrence → state write
├── KDA usage: yes/no and exact symbol
├── Fused boundary
├── Layout and sharding
├── Supported/unsupported modes
├── Native fallback
└── Minimal SGL-JAX adapter operations and timed conversions

[ ] Step 4: 实现 validator 并验证 GREEN

validate_audit.py 必须读取三个文件：

text

candidates.json
source-index.json
interface-audit.json

并拒绝：

缺字段、空字符串、unknown、TBD、TODO；
不在索引中的 evidence；
candidate commit 与 frozen manifest 不一致；
只有结论没有 symbol/path/SHA；
adapter 把 candidate core math 提前算完；
silent fallback 或未说明 effective path。

Run:

bash

python benchmark/kernels/gdn_compare/discovery/validate_audit.py
uv run --with pytest python -m pytest \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py -q

Expected: validator exit 0，全部测试 PASS。

[ ] Step 5: 提交接口审计

bash

git add benchmark/kernels/gdn_compare/discovery/interface-audit.md \
  benchmark/kernels/gdn_compare/discovery/interface-audit.json \
  benchmark/kernels/gdn_compare/discovery/validate_audit.py \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py
git commit -m "docs(gdn): audit candidate call paths"

Task 4: 盘点 native benchmark、测试和 profiling 能力

Files:

Create: benchmark/kernels/gdn_compare/discovery/native-benchmark-audit.md
Create: benchmark/kernels/gdn_compare/discovery/native-benchmark-audit.json
Modify: benchmark/kernels/gdn_compare/discovery/validate_audit.py
Modify: benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py

Interfaces:

Consumes: frozen source index 与接口审计。
Produces: 每个候选可执行的 native correctness/benchmark/profile 命令、环境入口、计时边界和与统一 L1/L2/L3 的映射；Phase 6 只允许使用这里批准的命令。
[ ] Step 1: 扩展失败测试

每个候选必须记录：

python

{
    "correctness_commands": [],
    "benchmark_commands": [],
    "profiling_commands": [],
    "environment_files": [],
    "timing_boundary": "",
    "reported_metrics": [],
    "maps_to_unified_layers": [],
    "limitations": [],
    "source_evidence": [],
}

允许某类命令为空，但必须在 limitations 中用源码证据解释仓库为什么没有该能力。测试必须拒绝把 native benchmark 标记为统一排名数据。

[ ] Step 2: 只读核验仓库原生入口

逐仓库检查：

bash

uv run python -c \
  'import json; d=json.load(open("/tmp/gdn-path-comparison/runtime-checkouts.json")); print(*[v["path"] for v in d["checkouts"].values()], sep="\n")' \
| while IFS= read -r checkout; do
    find "$checkout" -maxdepth 4 \
      \( -iname '*bench*' -o -iname '*profile*' -o -iname '*test*gdn*' \
         -o -iname '*test*kda*' -o -iname 'pyproject.toml' \
         -o -iname 'requirements*.txt' -o -iname 'Dockerfile*' \) -print
  done

对发现的 CLI 只读其 parser/入口源码，推导精确 --help 和运行命令；Phase 0 不 import 候选 package，不执行 runner。这样避免 import 触发 JAX/TPU 初始化，也不会把本地环境缺依赖误写成候选功能失败。

[ ] Step 3: 写 native audit 并通过 validator

native-benchmark-audit.md 必须逐候选区分：

repository-native correctness；
repository-native performance；
XProf/profiler；
环境/镜像；
输入 shape、dtype、warmup、samples、统计；
L1/L2/L3 中可映射的层级；
不能公平比较的原因。

Run:

bash

python benchmark/kernels/gdn_compare/discovery/validate_audit.py
uv run --with pytest python -m pytest \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py -q

Expected: exit 0，全部测试 PASS。

[ ] Step 4: 提交 native benchmark 盘点

bash

git add benchmark/kernels/gdn_compare/discovery/native-benchmark-audit.md \
  benchmark/kernels/gdn_compare/discovery/native-benchmark-audit.json \
  benchmark/kernels/gdn_compare/discovery/validate_audit.py \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py
git commit -m "docs(gdn): audit native benchmark coverage"

Task 5: 从固定模型与实际 projection 推导 35B per-shard shape

Files:

Create: benchmark/kernels/gdn_compare/discovery/derive_qwen35_shape.py
Create: benchmark/kernels/gdn_compare/discovery/qwen35-shape.json
Create: benchmark/kernels/gdn_compare/discovery/qwen35-shape.md
Modify: benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py

Interfaces:

Consumes: frozen model revision、S0 模型 config/implementation、TP16/DP4/EP16。
Produces: L1–L3 fixture 使用的权威 global/per-rank heads、dims、projection widths、 gate/state shapes 和推导证据。
[ ] Step 1: 写 shape 失败测试

测试必须断言：

python

def test_qwen35_shape_matches_frozen_model_and_tp16():
    data = json.loads(
        (MANIFEST.parent / "qwen35-shape.json").read_text()
    )
    candidates = json.loads(MANIFEST.read_text())
    assert data["model_revision"] == candidates["model"]["revision"]
    assert data["parallelism"] == {"tp": 16, "dp": 4, "ep": 16}
    per_rank = data["per_rank_gdn"]
    assert per_rank["n_kq"] == 1
    assert per_rank["n_v"] == 2
    assert per_rank["d_k"] == 128
    assert per_rank["d_v"] == 128
    assert per_rank["mixed_qkv_width"] == 512
    assert data["evidence"]["config_sha256"]
    assert data["evidence"]["projection_symbol"]
    assert data["evidence"]["projection_source_sha256"]

该断言来自 RFC 的预期设计输入；若固定 revision 的真实 config/代码推导不一致，测试保持失败并停止 Phase 0，不为通过测试而修改真实推导。

[ ] Step 2: 只下载固定 revision 的 config metadata

Run:

bash

MODEL_SHA=$(
  python -c \
  'import json; print(json.load(open("benchmark/kernels/gdn_compare/discovery/candidates.json"))["model"]["revision"])'
)
hf download Qwen/Qwen3.5-35B-A3B config.json \
  --revision "$MODEL_SHA" \
  --local-dir /tmp/gdn-path-comparison/model-config
sha256sum /tmp/gdn-path-comparison/model-config/config.json

不得下载 safetensors 权重。

[ ] Step 3: 实现 shape 推导

derive_qwen35_shape.py 必须：

读取固定 config.json；
读取 frozen S0 的 Qwen3.5 config/model projection 定义；
从 global head/dim/projection 字段计算 TP16 per-rank shape；
断言涉及的 global head 数可以被 TP16 合法切分，或记录源码中实际复制/扩展规则；
输出 config SHA、projection source SHA、symbol 和公式；
输出 gate、Conv1D state、recurrent state、fresh/continuing state shape；
绝不把 RFC 预期值直接写成计算结果。

qwen35-shape.md 必须展示：

text

config field → global shape → TP/replication rule → per-rank shape

[ ] Step 4: 运行 shape 和全套 discovery 测试

Run:

bash

python benchmark/kernels/gdn_compare/discovery/derive_qwen35_shape.py \
  --candidates benchmark/kernels/gdn_compare/discovery/candidates.json \
  --model-config /tmp/gdn-path-comparison/model-config/config.json \
  --output benchmark/kernels/gdn_compare/discovery/qwen35-shape.json \
  --report benchmark/kernels/gdn_compare/discovery/qwen35-shape.md
uv run --with pytest python -m pytest \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py -q

Expected: 全部 PASS，并确认或明确否定 RFC 中的 n_kq=1, n_v=2, d_k=d_v=128, mixed_qkv=[T,512] 预期。

[ ] Step 5: 提交 shape 推导

bash

git add benchmark/kernels/gdn_compare/discovery/derive_qwen35_shape.py \
  benchmark/kernels/gdn_compare/discovery/qwen35-shape.json \
  benchmark/kernels/gdn_compare/discovery/qwen35-shape.md \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py
git commit -m "docs(gdn): derive Qwen3.5 comparison shapes"

Task 6: Phase 0 总审、报告与停止门禁

Files:

Create: benchmark/kernels/gdn_compare/discovery/phase0-report.md
Create: benchmark/kernels/gdn_compare/discovery/phase0-evidence.json
Modify: benchmark/kernels/gdn_compare/discovery/validate_audit.py

Interfaces:

Consumes: Tasks 1–5 的全部 manifest、索引和报告。
Produces: Phase 0 结论、candidate/interface/native benchmark/shape 总表，以及 Phase 1–2 详细实施计划的批准输入。
[ ] Step 1: 生成 evidence manifest

phase0-evidence.json 必须记录 Tasks 1–5 所有 tracked artifact 的 relative path、 size、SHA-256 和生成 commit。validator 必须拒绝：

文件缺失或 SHA 不符；
worktree dirty；
frozen checkout HEAD 漂移；
审计结论缺少 source evidence；
Qwen shape 未确认；
任一候选 prefill/decode 能力未明确；
Phase 0 中出现 timing、XProf 或 TPU 产物。
[ ] Step 2: 写 Phase 0 报告

phase0-report.md 必须包含：

frozen S0/S1/P/V/model SHA；
四条 prefill/decode 调用链图；
“是否使用 KDA”的逐候选结论；
gate、Conv1D、recurrence、state I/O 融合边界；
layout/sharding/dtype/accumulation 对照表；
native correctness/benchmark/profile 命令表；
35B global/per-rank shape 推导；
L1/L2/L3 adapter 所需转换和必须纳入计时的 conversion；
明确 unsupported/fallback/风险；
Phase 1–2 需要创建/修改的精确文件与 API 建议；
发现的 blocker 与是否建议继续；
“未运行 TPU、未形成性能排名、未修改 production kernel”的范围声明。

[ ] Step 3: 执行最终验证

Run:

bash

python benchmark/kernels/gdn_compare/discovery/validate_audit.py
uv run --with pytest python -m pytest \
  benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py -q
git diff --check origin/main...HEAD
git status --short

Expected:

validator exit 0；
所有 discovery tests PASS；
git diff --check 无输出；
git status --short 无输出。
[ ] Step 4: 独立自审

逐条把 RFC Phase 0 的六个审计问题映射到 phase0-report.md 的 section 和 evidence SHA。检查：

没有把 keyword hit 当作真实调用链；
没有把 native benchmark 数字当作统一排名；
没有把 S1 误当作 S0；
没有把 P 或 V 的 training/backward/CP 能力自动纳入 SGL-JAX 移植范围；
没有用历史 #324/#325 数字选择赢家；
没有未解释的 unknown、TBD、TODO 或 silent skip。
[ ] Step 5: 提交 Phase 0 报告并暂停

bash

git add benchmark/kernels/gdn_compare/discovery
git commit -m "docs(gdn): complete candidate discovery"
git status --short --branch

完成后向用户报告：

frozen commits 和 model revision；
V 是否使用 KDA；
四条真实调用链与融合边界；
可复用的 native benchmark/skill；
确认后的 35B per-rank shape；
blocker 和推荐的 Phase 1–2 文件/API；
commit 与验证结果。

然后暂停。未经用户批准，不写 adapter、不运行 benchmark、不使用 TPU，也不追加 Phase 1–6 实施计划。

Phase 6 运行时兼容补充：JAX 0.10.2 下显式对齐 FusedEPMoE v1 输入分片

状态：设计已批准，生产实现尚未开始。日期：2026-07-27。

1. 触发条件与已验证现象

Phase 6 的 common-runtime gate 将四条候选统一到 Python 3.12、JAX/JAXLIB 0.10.2 和 libtpu 0.0.43。固定环境为 Qwen3.5-35B-A3B、TPU v6e-16、TP16/DP1/EP16、四个 rank。候选 revision 为：

S0：ffe9548f3fcd45dcc91c1d459b1c60c4bdc071a7；
S1：eb66c49f40a651dff2f3b151a6ef2bfd82c75665；
V：5d888d6279e7afd7eb20327ac77225bc1188ab2a。

S0 在进入目标 GDN 测量前，于 Qwen3.5 的 FusedEPMoE v1 路径稳定失败；四个 rank 的错误一致：

text

ValueError: in_specs passed to shard_map:
P(('data', 'tensor'), None, None)
does not match input sharding
P('data', None, None)
for bfloat16[64@data,2,1024].
JAX 0.10.2 no longer auto-reshards.

调用链为：

text

bench_one_batch
→ ModelRunner.forward
→ Qwen3_5
→ FusedEPMoE.__call__
→ fused_ep_moe
→ python/sgl_jax/srt/kernels/fused_moe/v1/kernel.py 的 jax.shard_map

因此当前 runtime_gate=blocked、ranking_eligible=false。S1 和 V 尚未形成排名，不能把此次失败解释成任何候选的性能输赢。

2. 根因

python/sgl_jax/srt/layers/fused_moe.py 中的 FusedEPMoE.__call__ 将 hidden_states、topk_weights 和 topk_ids 原样传给 v1 fused_ep_moe。模型侧实际输入只沿 data 轴分片，而 v1 kernel 内部的 jax.shard_map 明确要求 token 维同时沿 data 与 tensor 轴分片。

JAX 0.8.1 会在该边界隐式完成 reshard；JAX 0.10.2 不再自动 reshard，于是运行时在执行 GDN 之前拒绝该输入。

同文件的 FusedEPMoEV2.__call__ 已采用明确、经过现有代码使用的兼容模式：在进入 kernel 前构造 NamedSharding(mesh, P(("data", "tensor"), None))，并显式 reshard 三个输入。该既有实现是本补充设计的直接代码依据。

3. 批准的最小设计

只在 FusedEPMoE.__call__ 的 v1 layer-to-kernel 边界增加：

python

kernel_sharding = jax.sharding.NamedSharding(
    self.mesh, P(("data", "tensor"), None)
)
hidden_states = jax.sharding.reshard(hidden_states, kernel_sharding)
topk_weights = jax.sharding.reshard(topk_weights, kernel_sharding)
topk_ids = jax.sharding.reshard(topk_ids, kernel_sharding)

随后继续调用现有 fused_ep_moe。现有 output reshard、dtype、shape、routing、collective、weight sharding 和 kernel 参数保持不变。

该修改采用与 V2 相同的 layer-boundary 约定，不在 kernel 内隐藏额外转换，也不改变统一 GDN benchmark 的计时边界。

4. 明确不做

不修改 python/sgl_jax/srt/kernels/fused_moe/v1/kernel.py 的 shard_map 规格；
不把 v1 kernel 的输入规格放宽为仅 P("data", ...)；
不改变 Qwen3.5 或其他模型产生 hidden state 的上游分片；
不修改 GDN、KDA、scheduler、ModelRunner、候选 adapter 或 V 路径；
不借此修复、重构或调优任何相邻 fused-MoE 逻辑；
不在 S0 gate 重新通过前恢复 S1/V 排名。

放宽 kernel in_specs 会改变 EP/TP 语义；修改模型上游分片会扩大到全部消费者。两者都比 layer boundary 显式 reshard 风险更高，因此拒绝。

5. TDD 与验证门禁

实施必须遵守以下顺序，且避免在每个微小编辑后重复完整回归：

新增 focused layer contract test，拦截 v1 kernel 调用并验证三个输入实际采用 P(("data", "tensor"), None)；测试必须证明真实 dispatch 消费了 reshard 后的对象，而不是只验证构造了一个未使用的 sharding。
在 f83e9adf80b06075803d94c3b4647b0648b4e2d6 上先运行 focused test，预期 RED；失败原因必须是 v1 尚未显式 reshard，不能是 import、fixture、设备或依赖错误。
只加入第 3 节的最小生产修改，重跑 focused test，预期 GREEN。
修改完成后集中运行一次本地回归：focused test、现有 Qwen3.5 CPU component suite，以及可在当前环境执行的 fused-MoE contract tests。不能执行的 TPU test 必须明确列为远端门禁，不能静默跳过。
在原 Phase 6 四-rank TPU 环境、相同模型 revision、相同 parallelism、相同 runtime 和独立空 compilation cache 下重跑 S0 common-runtime gate。
只有四个 rank 都越过 fused-MoE 初始化并到达预期 GDN measurement boundary，S0 才恢复 ranking_eligible=true；随后才允许恢复 S1/V。

远端 gate 必须继续记录 requested rank 与 PJRT process index 的完整排列，不要求逐 Pod identity；四 rank 输入、revision、runtime、模型、shape 和 artifact SHA 必须回读一致。

6. 验收标准与停止条件

验收全部满足时才可关闭本兼容 blocker：

focused RED/GREEN 证据完整；
生产 diff 仅限 FusedEPMoE.__call__ 的显式 input reshard；
本地集中回归通过，或对不能本地执行的测试给出明确远端结果；
JAX 0.10.2 四-rank S0 gate 不再出现 input-sharding mismatch；
S0 到达 Phase 6 设计中的 GDN 测量边界；
worktree、commit、runtime、cache 和远端证据可追溯。

出现以下任一情况立即暂停，不扩张修复：

explicit reshard 后出现新的 fused-MoE 数值或 collective 问题；
需要修改 v1 kernel 内部 shard_map、Qwen 上游分片或其他生产文件；
四 rank revision/runtime/输入不一致；
远端失败发生在新的非目标组件，且无法由本补充设计解释。

7. 证据边界

实验分支：experiment/gdn-path-comparison@f83e9adf80b06075803d94c3b4647b0648b4e2d6；
Gate 记录：benchmark/kernels/gdn_compare/discovery/phase6-runtime-gate.md；
远端证据根：/tmp/gdn-path-comparison-phase6-dfbfc538/runtime-gate/s0-jax0102-dfbfc538-aligned16；
source archive SHA-256：029ac9192956e518492c81cb6a84604246122631f9b2128152f1b3056ec1eb30；
deployed bench_one_batch.py SHA-256：024dbbbdad2172794f765a4b7d05ffc0cf61b664ff4ac41de4bb16e24eee5114。

本补充只定义兼容 blocker 的修复与门禁，不改变 RFC-0034 的候选、workload、正确性、性能选择或最终移植规则。

Phase 6 JAX 0.10.2 Compatibility Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development or superpowers:executing-plans to implement this plan task-by-task. Use superpowers:test-driven-development for Task 1 and exec-remote for Task 2. Steps use checkbox (- [ ]) syntax for tracking.

Goal: 在不改变 fused-MoE kernel 规格、Qwen3.5 上游分片或 GDN 候选逻辑的前提下，让 FusedEPMoE v1 在 JAX 0.10.2 中显式满足 kernel 输入分片契约，并重新打开 S0 common-runtime gate。

Architecture: 修改严格位于 layer-to-kernel boundary。FusedEPMoE v1 在真实调用中，把 hidden states、top-k weights 和 top-k ids 显式 reshard 到 NamedSharding(mesh, P(("data", "tensor"), None))，随后原样调用现有 fused_ep_moe；输出继续沿用现有 reshard。先用真实 CPU JAX array、真实 NamedSharding 和真实 FusedEPMoE.__call__ 构造 focused contract RED/GREEN，只替换 TPU Pallas kernel boundary；完成后集中运行一次本地回归。远端只恢复 S0 四-rank JAX 0.10.2 gate，不运行 S1/V 排名。

Tech Stack: Python 3.12、JAX/JAXLIB 0.8.1（本地 contract）、JAX/JAXLIB 0.10.2 + libtpu 0.0.43（远端 gate）、pytest、GKE/kubectl、TPU v6e-16。

Global Constraints

基线 worktree：/Users/feibo/projects/intership/sglang-jax/.claude/worktrees/experiment/gdn-path-comparison。
基线 branch/HEAD：experiment/gdn-path-comparison@f83e9adf80b06075803d94c3b4647b0648b4e2d6。
生产修改只允许出现在 python/sgl_jax/srt/layers/fused_moe.py 的 FusedEPMoE.__call__。
新测试只允许新增 python/sgl_jax/test/layers/test_fused_moe_v1_sharding.py。
不修改 v1 kernel 的 shard_map、FusedEPMoEV2、Qwen3.5 模型、GDN/KDA、scheduler、ModelRunner、候选 adapter 或 V checkout。
本地完整回归只在最小实现完成后集中运行一次；不在每个三行编辑后重复整套测试。
远端使用原有 perf-16，不创建、删除或扩容资源；只停止本次 run 明确记录的 PID。
任一新失败若不再是已批准的 input-sharding mismatch，立即暂停，不做第二项生产修复。
S0 gate 通过只表示 S0 恢复排名资格；不表示 S1/V 已运行、已排名或已有赢家。

Task 1: 用 TDD 修复 FusedEPMoE v1 layer-to-kernel 分片边界

Files:

Create: python/sgl_jax/test/layers/test_fused_moe_v1_sharding.py
Modify: python/sgl_jax/srt/layers/fused_moe.py
Verify only: python/sgl_jax/test/models/test_qwen3_5.py
Remote-only existing suite: python/sgl_jax/test/kernels/fused_moe_v1_test.py

Interfaces:

Consumes: arrays carrying model-side P("data", None) sharding.
Produces: the same arrays explicitly resharded to P(("data", "tensor"), None) before the v1 kernel call.
Preserves: output shape、dtype、routing、collective、weights and existing output reshard.
Public API/config changes: none.
[ ] Step 1: 确认基线和文件边界

Run:

bash

cd /Users/feibo/projects/intership/sglang-jax/.claude/worktrees/experiment/gdn-path-comparison
test "$(git branch --show-current)" = "experiment/gdn-path-comparison"
test "$(git rev-parse HEAD)" = "f83e9adf80b06075803d94c3b4647b0648b4e2d6"
git status --short
git diff --check

Expected:

branch 与 HEAD 精确匹配；
git status --short 无输出；
git diff --check 无输出。

若不匹配，停止并报告，不 reset、不覆盖用户改动。

[ ] Step 2: 新增 focused contract test

Create python/sgl_jax/test/layers/test_fused_moe_v1_sharding.py with:

python

from types import SimpleNamespace

import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, NamedSharding
from jax.sharding import PartitionSpec as P

from sgl_jax.srt.layers import fused_moe as fused_moe_module


def _param(value):
    return SimpleNamespace(value=value)


def _layer_stub(mesh):
    return SimpleNamespace(
        mesh=mesh,
        w1=_param(jnp.zeros((2, 8, 8), dtype=jnp.bfloat16)),
        w2=_param(jnp.zeros((2, 8, 8), dtype=jnp.bfloat16)),
        w3=_param(jnp.zeros((2, 8, 8), dtype=jnp.bfloat16)),
        w1_shared=None,
        w2_shared=None,
        w3_shared=None,
        w1_scale=None,
        w2_scale=None,
        w3_scale=None,
        w1_shared_scale=None,
        w2_shared_scale=None,
        w3_shared_scale=None,
        quant_block_k=None,
        num_experts_per_tok=1,
        use_grouped_topk=False,
        num_groups=1,
        top_k_groups=1,
        renormalize_topk_logits=False,
        routed_scaling_factor=None,
        activation="silu",
        disable_a2a=False,
        disable_dynamic_ffn1=False,
        disable_dynamic_ffn2=False,
        disable_weight_load=False,
        disable_a2a_s_tile_read=False,
        disable_a2a_s_acc_tile_write=False,
        disable_shared_expert=False,
        disable_all_reduce_metadata=False,
        disable_sync_barrier=False,
        use_jax_allreduce_metadata=True,
    )


def test_v1_layer_reshards_kernel_inputs_over_ep_axes(monkeypatch):
    mesh = Mesh(
        np.asarray(jax.devices()[:1], dtype=object).reshape(1, 1),
        ("data", "tensor"),
    )
    captured = {}

    def fake_fused_ep_moe(**kwargs):
        for name in ("tokens", "topk_weights", "topk_ids"):
            captured[name] = kwargs[name].sharding.spec
        return kwargs["tokens"]

    monkeypatch.setattr(fused_moe_module, "fused_ep_moe", fake_fused_ep_moe)

    data_sharding = NamedSharding(mesh, P("data", None))
    hidden_states = jax.device_put(
        jnp.zeros((2, 8), dtype=jnp.bfloat16), data_sharding
    )
    topk_weights = jax.device_put(
        jnp.ones((2, 1), dtype=jnp.float32), data_sharding
    )
    topk_ids = jax.device_put(
        jnp.zeros((2, 1), dtype=jnp.int32), data_sharding
    )

    output = fused_moe_module.FusedEPMoE.__call__(
        _layer_stub(mesh), hidden_states, topk_weights, topk_ids
    )

    expected = P(("data", "tensor"), None)
    assert captured == {
        "tokens": expected,
        "topk_weights": expected,
        "topk_ids": expected,
    }
    assert output.shape == hidden_states.shape

该测试不实例化 FusedEPMoE.__init__，避免 CPU Auto mesh 初始化噪声；它直接执行真实 production __call__。只把 TPU kernel boundary 替换为 capture spy，因此测试会在 dispatch 未消费显式 reshard 时失败。

[ ] Step 3: 运行 focused RED

Run:

bash

JAX_PLATFORMS=cpu PYTHONWARNINGS=ignore PYTHONPATH=python \
/private/tmp/gdn-path-comparison/venvs/s0-baseline/bin/python -m pytest \
  python/sgl_jax/test/layers/test_fused_moe_v1_sharding.py \
  -q --tb=short \
  -o cache_dir=/private/tmp/gdn-path-comparison/pytest-cache/phase6-fused-moe-red

Expected: 1 failed。失败必须显示 captured spec 仍为 P("data", None)，而 expected 为 P(("data", "tensor"), None)。

以下失败都不是有效 RED，必须先停止并分类：import error、fixture error、没有 CPU device、依赖缺失、mock 未命中真实 dispatch。

[ ] Step 4: 实现最小 GREEN

In FusedEPMoE.__call__, immediately before the existing v1 fused_ep_moe call, add exactly:

python

kernel_sharding = jax.sharding.NamedSharding(
    self.mesh, P(("data", "tensor"), None)
)
hidden_states = jax.sharding.reshard(hidden_states, kernel_sharding)
topk_weights = jax.sharding.reshard(topk_weights, kernel_sharding)
topk_ids = jax.sharding.reshard(topk_ids, kernel_sharding)

Do not alter the call arguments except that the three existing variables now refer to the resharded arrays. Do not move, delete, or replace the existing output reshard.

[ ] Step 5: 运行 focused GREEN

Run the same focused command with a new cache:

bash

JAX_PLATFORMS=cpu PYTHONWARNINGS=ignore PYTHONPATH=python \
/private/tmp/gdn-path-comparison/venvs/s0-baseline/bin/python -m pytest \
  python/sgl_jax/test/layers/test_fused_moe_v1_sharding.py \
  -q --tb=short \
  -o cache_dir=/private/tmp/gdn-path-comparison/pytest-cache/phase6-fused-moe-green

Expected: 1 passed。

[ ] Step 6: 一次性集中本地回归

Run once after the production edit is final:

bash

JAX_PLATFORMS=cpu PYTHONWARNINGS=ignore PYTHONPATH=python \
/private/tmp/gdn-path-comparison/venvs/s0-baseline/bin/python -m pytest \
  python/sgl_jax/test/layers/test_fused_moe_v1_sharding.py \
  python/sgl_jax/test/models/test_qwen3_5.py \
  -q --tb=short \
  -o cache_dir=/private/tmp/gdn-path-comparison/pytest-cache/phase6-fused-moe-regression

PYTHONPYCACHEPREFIX=/private/tmp/gdn-path-comparison/pycache/phase6 \
/private/tmp/gdn-path-comparison/venvs/s0-baseline/bin/python -m py_compile \
  python/sgl_jax/srt/layers/fused_moe.py \
  python/sgl_jax/test/layers/test_fused_moe_v1_sharding.py

git diff --check
git diff -- \
  python/sgl_jax/srt/layers/fused_moe.py \
  python/sgl_jax/test/layers/test_fused_moe_v1_sharding.py

Expected:

pytest: 19 passed；
py_compile exit 0；
git diff --check 无输出；
生产 diff 只有三个 input reshard 和必要的 local variable；
test diff 只有一个 focused contract file。

python/sgl_jax/test/kernels/fused_moe_v1_test.py 是 TPU-only suite。本地不得声称执行；它由 Task 2 的真实四-host gate 覆盖。

[ ] Step 7: 提交本地修复

Run:

bash

git add \
  python/sgl_jax/srt/layers/fused_moe.py \
  python/sgl_jax/test/layers/test_fused_moe_v1_sharding.py
git commit -m "fix(moe): reshard v1 inputs at kernel boundary"
git status --short --branch

Expected: commit 成功，worktree clean。记录新 HEAD，Task 2 全程固定使用该 SHA。

Task 2: 在四-host JAX 0.10.2 环境关闭 S0 runtime gate

Files:

Modify after successful gate only: benchmark/kernels/gdn_compare/discovery/phase6-runtime-gate.md
Create transient controller wrapper: /private/tmp/gdn-path-comparison/phase6_s0_gate_wrapper.py
Remote source root: /tmp/gdn-path-comparison-phase6-compat-$SHORT_SHA/source
Remote run root: /tmp/gdn-path-comparison-phase6-compat-$SHORT_SHA/runtime-gate/s0-jax0102
Remote interpreter: /tmp/gdn-path-comparison-phase4-v-jax0102/bin/python
Model: /data/user/dongbohao/beaver-324/models/Qwen3.5-35B-A3B-59d61f3

Interfaces:

Consumes: the exact clean Task 1 commit, frozen JAX 0.10.2 interpreter, existing perf-16 four-host slice, and the fixed 16-line prompt fixture.
Produces: four-rank S0 gate evidence plus an updated phase6-runtime-gate.md.
Gate state on success: runtime_gate=passed, ranking_eligible=true.
Explicitly does not produce: S1/V timing, ranking, winner, production PR, or new TPU resources.

The frozen V checkout is deliberately absent from PYTHONPATH: this gate runs only S0, and the approved runtime-gate evidence already established that V was present as external source but was not imported. The shared runtime contract is the pinned interpreter and package set, not a V source import.

Fixed infrastructure:

text

context=gke_tpu-service-473302_asia-northeast1-b_niu-v6e-4x4-asia
namespace=default
container=jax-tpu
rank 0=perf-16-0-jgb5c
rank 1=perf-16-1-zhn6d
rank 2=perf-16-2-hs7kc
rank 3=perf-16-3-vqw2x
coordinator=perf-16-0.perf-16-headless-svc:5000

[ ] Step 1: 只读远端 preflight

Run:

bash

set -euo pipefail
CTX=gke_tpu-service-473302_asia-northeast1-b_niu-v6e-4x4-asia
NS=default
PODS=(
  perf-16-0-jgb5c
  perf-16-1-zhn6d
  perf-16-2-hs7kc
  perf-16-3-vqw2x
)
kubectl --context "$CTX" -n "$NS" get job perf-16 -o wide
kubectl --context "$CTX" -n "$NS" get pods -l job-name=perf-16 -o wide
kubectl --context "$CTX" -n "$NS" get pods \
  perf-16-0-jgb5c perf-16-1-zhn6d perf-16-2-hs7kc perf-16-3-vqw2x \
  -o jsonpath='{range .items[*]}{.metadata.name}{"\t"}{.status.phase}{"\t"}{.status.containerStatuses[?(@.name=="jax-tpu")].ready}{"\t"}{.status.containerStatuses[?(@.name=="jax-tpu")].restartCount}{"\n"}{end}'

Expected: Job active，四个 exact Pod 均 Running、jax-tpu ready=true、restartCount 未增加。

随后对四个 Pod 只读检查：

远端 interpreter 存在；
模型目录存在；
没有本次 run root；
5000 协调端口和本次明确使用的 PID 不冲突；
不停止其他人的进程。

任一不满足时暂停并报告。

[ ] Step 2: 打包固定 HEAD 并部署到四个 Pod

Run locally from the clean Task 1 HEAD:

bash

HEAD_SHA=$(git rev-parse HEAD)
SHORT_SHA=$(git rev-parse --short=12 HEAD)
ARCHIVE=/private/tmp/gdn-path-comparison/phase6-compat-$SHORT_SHA.tar
REMOTE_ROOT=/tmp/gdn-path-comparison-phase6-compat-$SHORT_SHA
RUN_ROOT=$REMOTE_ROOT/runtime-gate/s0-jax0102
git archive --format=tar --prefix=source/ "$HEAD_SHA" -o "$ARCHIVE"
shasum -a 256 "$ARCHIVE"

Use the fixed PODS array from Step 1. For each exact Pod, assert the new root does not exist, create it, and stream the same archive:

bash

for POD in "${PODS[@]}"; do
  if kubectl --context "$CTX" -n "$NS" exec -c jax-tpu "$POD" -- \
    test -e "$REMOTE_ROOT"; then
    echo "refusing to reuse existing root on $POD: $REMOTE_ROOT" >&2
    exit 1
  fi
  kubectl --context "$CTX" -n "$NS" exec -c jax-tpu "$POD" -- \
    mkdir -p "$REMOTE_ROOT"
  kubectl --context "$CTX" -n "$NS" exec -i -c jax-tpu "$POD" -- \
    tar -xf - -C "$REMOTE_ROOT" < "$ARCHIVE"
done

After deployment, read back on all four Pods:

bash

for POD in "${PODS[@]}"; do
  kubectl --context "$CTX" -n "$NS" exec -c jax-tpu "$POD" -- \
    env REMOTE_ROOT="$REMOTE_ROOT" bash -lc '
      set -euo pipefail
      cd "$REMOTE_ROOT/source"
      sha256sum python/sgl_jax/bench_one_batch.py
      sha256sum python/sgl_jax/srt/layers/fused_moe.py
      /tmp/gdn-path-comparison-phase4-v-jax0102/bin/python -c \
        "import importlib.metadata as m; print(m.version(\"jax\"), m.version(\"jaxlib\"), m.version(\"libtpu\"))"
    '
done

Expected: four hosts report the same source archive SHA、bench SHA、fused_moe.py SHA and 0.10.2 / 0.10.2 / 0.0.43。

[ ] Step 3: 创建相同 prompt 和独立空 cache

First create the transient controller wrapper at /private/tmp/gdn-path-comparison/phase6_s0_gate_wrapper.py:

python

import json
import os
from pathlib import Path
import runpy
import traceback

import jax


requested_rank = int(os.environ["REQUESTED_RANK"])
run_root = Path(os.environ["RUN_ROOT"])
identity_path = run_root / "identity" / f"rank-{requested_rank}.json"
status = 0

try:
    runpy.run_module("sgl_jax.bench_one_batch", run_name="__main__")
except SystemExit as exc:
    status = int(exc.code or 0)
except BaseException:
    traceback.print_exc()
    status = 1
finally:
    identity = {"requested_rank": requested_rank}
    try:
        identity.update(
            {
                "pjrt_process_index": int(jax.process_index()),
                "process_count": int(jax.process_count()),
                "local_device_count": int(jax.local_device_count()),
                "local_devices": [str(device) for device in jax.local_devices()],
            }
        )
    except BaseException as exc:
        identity["identity_error"] = repr(exc)
    temporary = identity_path.with_suffix(".json.tmp")
    temporary.write_text(json.dumps(identity, indent=2, sort_keys=True) + "\n")
    os.replace(temporary, identity_path)

raise SystemExit(status)

Syntax-check it with the frozen local baseline interpreter:

bash

PYTHONPYCACHEPREFIX=/private/tmp/gdn-path-comparison/pycache/phase6-wrapper \
/private/tmp/gdn-path-comparison/venvs/s0-baseline/bin/python -m py_compile \
  /private/tmp/gdn-path-comparison/phase6_s0_gate_wrapper.py

Copy that same wrapper to each Pod, create the fixed 16-line prompt, and create all four initially empty per-rank cache directories:

bash

for POD in "${PODS[@]}"; do
  kubectl --context "$CTX" -n "$NS" exec -c jax-tpu "$POD" -- \
    env RUN_ROOT="$RUN_ROOT" bash -lc '
      set -euo pipefail
      mkdir -p "$RUN_ROOT"/{cache,identity,logs,sha}
      for rank in 0 1 2 3; do
        mkdir "$RUN_ROOT/cache/rank-$rank"
      done
      /tmp/gdn-path-comparison-phase4-v-jax0102/bin/python -c \
        "from pathlib import Path; Path(\"$RUN_ROOT/prompts.txt\").write_text(\"The capital of France is Paris.\\n\" * 16)"
      test "$(wc -l < "$RUN_ROOT/prompts.txt")" -eq 16
      find "$RUN_ROOT/cache" -type f -print -quit | \
        grep -q . && { echo "cache is not empty" >&2; exit 1; } || true
      sha256sum "$RUN_ROOT/prompts.txt" > "$RUN_ROOT/sha/prompt.sha256"
    '
  kubectl --context "$CTX" -n "$NS" exec -i -c jax-tpu "$POD" -- \
    env RUN_ROOT="$RUN_ROOT" bash -lc \
      'cat > "$RUN_ROOT/phase6_s0_gate_wrapper.py"' \
    < /private/tmp/gdn-path-comparison/phase6_s0_gate_wrapper.py
done

Read back all four prompt and wrapper SHA values and require exact equality:

bash

for POD in "${PODS[@]}"; do
  kubectl --context "$CTX" -n "$NS" exec -c jax-tpu "$POD" -- \
    env RUN_ROOT="$RUN_ROOT" bash -lc '
      set -euo pipefail
      cat "$RUN_ROOT/sha/prompt.sha256"
      sha256sum "$RUN_ROOT/phase6_s0_gate_wrapper.py"
      for rank in 0 1 2 3; do
        test -d "$RUN_ROOT/cache/rank-$rank"
        test -z "$(find "$RUN_ROOT/cache/rank-$rank" -type f -print -quit)"
      done
    '
done

[ ] Step 4: 四个 Pod 并发运行 S0 gate

Launch all four ranks concurrently from one Bash controlling shell. Each rank writes only its own log, exit file and identity JSON under the remote run root:

bash

CONTROLLER_LOG_ROOT=/private/tmp/gdn-path-comparison/phase6-controller-$SHORT_SHA
mkdir -p "$CONTROLLER_LOG_ROOT"
declare -a CONTROLLER_PIDS=()

for RANK in 0 1 2 3; do
  POD=${PODS[$RANK]}
  kubectl --context "$CTX" -n "$NS" exec -c jax-tpu "$POD" -- \
    env \
      REQUESTED_RANK="$RANK" \
      HEAD_SHA="$HEAD_SHA" \
      REMOTE_ROOT="$REMOTE_ROOT" \
      RUN_ROOT="$RUN_ROOT" \
    bash -lc '
      set -o pipefail
      export SGLANG_JAX_GDN_PREFILL_IMPL=reference
      export JAX_COMPILATION_CACHE_DIR="$RUN_ROOT/cache/rank-$REQUESTED_RANK"
      export PYTHONPATH="$REMOTE_ROOT/source/python"
      {
        printf "requested_rank=%s\n" "$REQUESTED_RANK"
        printf "head_sha=%s\n" "$HEAD_SHA"
        sha256sum "$RUN_ROOT/prompts.txt"
        sha256sum "$RUN_ROOT/phase6_s0_gate_wrapper.py"
        sha256sum "$REMOTE_ROOT/source/python/sgl_jax/bench_one_batch.py"
        sha256sum "$REMOTE_ROOT/source/python/sgl_jax/srt/layers/fused_moe.py"
      } > "$RUN_ROOT/sha/rank-$REQUESTED_RANK.txt"
      /tmp/gdn-path-comparison-phase4-v-jax0102/bin/python \
        "$RUN_ROOT/phase6_s0_gate_wrapper.py" \
        --model-path /data/user/dongbohao/beaver-324/models/Qwen3.5-35B-A3B-59d61f3 \
        --trust-remote-code \
        --device tpu \
        --dtype bfloat16 \
        --tp-size 16 \
        --data-parallel-size 1 \
        --ep-size 16 \
        --nnodes 4 \
        --node-rank "$REQUESTED_RANK" \
        --dist-init-addr perf-16-0.perf-16-headless-svc:5000 \
        --max-total-num-tokens 2048 \
        --max-prefill-tokens 512 \
        --chunked-prefill-size 512 \
        --page-size 64 \
        --max-running-requests 64 \
        --disable-radix-cache \
        --disable-overlap-schedule \
        --correctness-test \
        --prompt-filename "$RUN_ROOT/prompts.txt" \
        --output-len 2 \
        --cut-len 4 \
        2>&1 | tee "$RUN_ROOT/logs/rank-$REQUESTED_RANK.log"
      status=${PIPESTATUS[0]}
      printf "%s\n" "$status" > \
        "$RUN_ROOT/logs/rank-$REQUESTED_RANK.exit"
      exit "$status"
    ' > "$CONTROLLER_LOG_ROOT/rank-$RANK.log" 2>&1 &
  CONTROLLER_PIDS[$RANK]=$!
done

overall=0
for RANK in 0 1 2 3; do
  if ! wait "${CONTROLLER_PIDS[$RANK]}"; then
    overall=1
  fi
done
test "$overall" -eq 0

The wrapper must persist, per rank:

text

logs/rank-N.log
logs/rank-N.exit
identity/rank-N.json
sha/rank-N.txt

Identity JSON must record requested rank, PJRT process index, process count, local device count and device IDs. Validation accepts a complete unique PJRT permutation such as [1,0,2,3]; it does not require requested rank to equal PJRT index per Pod.

[ ] Step 5: 判定 gate

PASS requires all of the following:

four exit files equal 0；
no log contains the previous in_specs passed to shard_map mismatch；
all four ranks complete distributed initialization and load the same 14 weight shards；
requested/effective GDN implementation are both reference and fallback reason is empty；
all four ranks reach and complete the expected GDN measurement/correctness boundary；
revision、runtime、model、prompt、controller wrapper、bench file and fused_moe.py SHA agree across ranks；
requested ranks and PJRT indices are each complete unique sets of four values；
no NaN/Inf、OOM、collective failure or artifact SHA mismatch。

If explicit reshard reveals a new fused-MoE numeric/collective error, or any failure occurs in another non-target component, stop immediately. Do not modify a second production file and do not run S1/V.

[ ] Step 6: 更新 gate 文档并集中验证

Only after PASS, append a dated follow-up section to benchmark/kernels/gdn_compare/discovery/phase6-runtime-gate.md containing:

compatibility commit SHA；
focused RED and GREEN result；
local consolidated regression result；
exact remote runtime、model、topology and parallelism；
four requested-to-PJRT mappings；
archive、prompt、bench、fused_moe.py and per-rank evidence SHA；
the statement that the old input-sharding mismatch is absent；
the statement that S0 reached the GDN boundary；
runtime_gate=passed and ranking_eligible=true；
an explicit statement that S1/V ranking has not yet run。

Run:

bash

PYTHONPYCACHEPREFIX=/private/tmp/gdn-path-comparison/pycache/phase6-final \
/private/tmp/gdn-path-comparison/venvs/s0-baseline/bin/python -m py_compile \
  python/sgl_jax/srt/layers/fused_moe.py \
  python/sgl_jax/test/layers/test_fused_moe_v1_sharding.py
git diff --check
git status --short

Expected: only the gate document is uncommitted at this point; production and test files remain exactly at the Task 1 commit.

[ ] Step 7: 提交 gate closure 并暂停

Run:

bash

git add benchmark/kernels/gdn_compare/discovery/phase6-runtime-gate.md
git commit -m "docs(gdn): close JAX 0.10.2 S0 runtime gate"
git status --short --branch

Expected: worktree clean，two new commits after f83e9adf：

text

fix(moe): reshard v1 inputs at kernel boundary
docs(gdn): close JAX 0.10.2 S0 runtime gate

Report the focused RED/GREEN、local regression、four-rank gate result、evidence paths and final HEAD, then pause. Do not start S1/V until the user explicitly approves resuming Phase 6 ranking.

Phase 8 修订：Qwen3.5-397B BF16 / TPU v7 四层最小 A/B

修订状态与覆盖关系

本节由用户于 2026-07-28 批准，覆盖本 RFC 中“397B 同时执行 BF16/FP8”以及本地旧 Phase 8 计划中的 v6e 候选、FP8 compatibility track 和完整 five-cell sweep。既有 35B 实验、Phase 6/7 结论与证据保持历史有效，但不得替代本节的新 397B 测量。

本阶段只回答一个问题：

在 Qwen/Qwen3.5-397B-A17B、BF16 和同一 TPU v7 slice 上，当前 SGL-JAX reference 与实验分支 hybrid（chunkwise prefill + reference decode）在 kernel、完整 GDN layer、完整模型和 serving 四个边界分别表现如何？

本阶段不执行 FP8，不创建正式代码 PR，不合入生产分支，不跑完整 five-cell，不默认采集 XProf。只有当四层结果相互矛盾、出现无法解释的回归或无法确认 effective path 时，才另行批准 XProf。

固定对比对象

Variant	Prefill	Decode
`reference`	`SGLANG_JAX_GDN_PREFILL_IMPL=reference`	现有 reference decode
`hybrid`	`SGLANG_JAX_GDN_PREFILL_IMPL=chunkwise`	与 reference 相同的 decode

两侧必须使用同一个实验 HEAD、同一模型 revision、同一 BF16 权重、同一 v7 slice、同一并行配置、同一 workload、同一 seed 和分别独立且初始为空的 compilation cache。唯一允许不同的是 selector 及其隔离的 cache/output 路径。

资源与拓扑

模型：Qwen/Qwen3.5-397B-A17B，BF16 only。
资源优先级：Falcon 中可用的 TPU v7x 64-chip 4x4x4 slice。
首选 Falcon cluster：当前 active 的 tpu-training-antgroup；若创建前状态变化，可使用另一条已授权且实时可用的 v7x 资源，但必须在 manifest 中记录 cluster ID、reservation、topology 和 experiment ID。
初始 mesh candidate：data=4, tensor=16，对应 TP16 / DP4。EP 必须在静态模型配置和真实 Fused-MoE sharding gate 后冻结，并由 reference/hybrid 共用；不得为了某一 variant 调整。
资源创建前必须证明 GDN heads、full-attention heads/KV heads、专家数、Fused-MoE 分片和模型 HBM budget 可构造。若 gate 不通过，停止而不是猜测新拓扑。

四层最小性能矩阵

层级	边界	最小 workload	计时
Kernel	直接 recurrence：reference 与 chunkwise	Prefill `T=4096`、`T=32768`；使用 397B 实际 per-rank GDN shape	3 warmup + 10 measured
Full GDN layer	Conv1D、projection、normalization、recurrence、state I/O、output projection	Prefill `T=4096` fresh；Decode `batch=16` continuing	3 warmup + 5 measured
Model	仓库原生 `python -m sgl_jax.bench_one_batch`	Prefill `batch=16, ISL=4096, OSL=1`；Decode `batch=16, ISL=1024, OSL=100`	shape warmup 后 5 独立 rounds
Serving	仓库 `.claude/skills/model-speed-benchmark`	Random dataset：`4096→1, concurrency=16`；`1024→100, concurrency=16`	每 cell 5 独立 rounds

Kernel 层用于解释 recurrence 上限，不能单独决定最终结论。最终选择以完整 GDN layer、model 和 serving 三层的共同结果为准；若 kernel 很快但上层没有收益，报告必须明确指出 integration 成本。

正确性与 anti-fallback 门禁

每个 variant 在进入计时前必须满足：

模型和权重 revision 完整、四 rank/host 一致；
requested/effective selector 一致，fallback reason 为空；
health、model discovery、固定 seed prefill/extend/decode smoke 通过；
output、logits、recurrent state 和 Conv state 均为有限值；
variant 内重复执行结果稳定；
reference/hybrid 的 numerical tolerance 和 exact token identity 分别报告；
无 OOM、collective、XLA、sharding、请求或 artifact 错误。

此前 35B 在第 32 个 token 出现 exact-token divergence。该事实必须保留在 397B 报告中。若 397B 再次出现 token mismatch，性能数据仍可作为实验结果保存，但不得声称 production parity 或 production-ready；必须单独给出首次 divergence 位置和 numerical metrics。

统计与产物

两侧串行运行，不允许资源重叠。
每层都保留每个 measured sample，不静默删除 outlier。
报告 median、min/max、per-host spread 和相对变化。
serving 报告 input/output throughput、TTFT、TPOT；model 报告 prefill throughput/latency 和 decode throughput/latency。
四层产物必须包含 source/model/runtime/topology provenance、完整命令、selector、cache 初始状态、raw samples、SHA-256 readback 和 completion marker。
正式证据目录：

text

/data/user/dongbohao/gdn-path-comparison/evidence/phase8-397b-bf16-v7/{run_id}/
├── preflight/
├── kernel/
├── layer/
├── model/
├── serving/
├── comparisons/
├── reports/
└── .complete

Phase 8 Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:executing-plans. 按 Task 检查点执行；在 Task 1–2 完成前不得创建 Falcon 实验，在 Task 3–6 中任一门禁失败时停止。

Goal: 在一个固定的 64-chip TPU v7x slice 上完成 397B BF16 reference/hybrid 四层最小 A/B，并形成可回读的性能报告。

Architecture: 复用实验分支现有统一 GDN harness、bench_one_batch 和 model-speed-benchmark。只新增 397B 配置/shape 证据和 transient Falcon 执行包装；不修改 GDN/KDA production path。四层逐级执行，上一层的 correctness、effective-path 和 artifact 门禁是下一层的前置条件。

Tech Stack: SGL-JAX、JAX/libtpu、Falcon、TPU v7x、BF16、现有 GDN comparison harness。

Task 1：冻结 397B 配置、代码和 topology

Files:

Create: benchmark/kernels/gdn_compare/configs/qwen397_v7x64.yaml
Create: benchmark/kernels/gdn_compare/discovery/qwen397-shape.md
Create: benchmark/kernels/gdn_compare/discovery/qwen397-shape.json
Modify only if needed for config loading: benchmark/kernels/gdn_compare/discovery/derive_qwen35_shape.py
Test: benchmark/kernels/gdn_compare/tests/test_discovery_manifest.py

Produces: immutable experiment HEAD、model revision、per-rank GDN shape、mesh/TP/DP/EP、HBM budget 和 source/config SHA。

[ ] 回读实验 branch/HEAD、dirty state、Falcon auth、cluster status 和 v7x 资源。
[ ] 解析官方 397B config.json 和 weight index，不加载 TPU workload。
[ ] 推导 64-chip candidate 的 GDN/full-attention/MoE sharding；冻结一个 reference/hybrid 共用配置。
[ ] 新增 focused config/shape test，再一次性运行相关 discovery tests、py_compile 和 git diff --check。
[ ] 若 heads、experts、mesh 或 HBM 任一不成立，写 blocker 并停止。

Task 2：创建 Falcon v7 实验并通过 load/runtime smoke

Transient files:

/tmp/gdn-phase8-397b/falcon-v7x64.yaml
/tmp/gdn-phase8-397b/preflight-manifest.json

Produces: Falcon experiment ID、64-chip v7x runtime、可读模型路径、四 host/rank identity 和 reference/hybrid smoke 结果。

[ ] 使用 falcon cluster list 和 falcon cluster resources get <cluster_id> --output json 重新确认 tpu-training-antgroup 或另一已授权 v7x cluster。
[ ] 渲染 BF16-only 4x4x4 Falcon YAML；挂载固定模型存储，只使用一个 single-writer 负责缺失权重。
[ ] 执行 falcon exp create -f /tmp/gdn-phase8-397b/falcon-v7x64.yaml --output json，记录返回的 exp_id。
[ ] reference 先完成 load、health、model discovery、prefill/extend/decode smoke；停止并清空其进程后再运行 hybrid。
[ ] 校验 requested/effective selector、finite、token identity/numerical metrics、rank permutation、端口和空 cache；失败即停止。

Task 3：Kernel recurrence A/B

Uses: benchmark.kernels.gdn_compare.runner --level L1

Produces: 4K/32K reference/hybrid raw samples 和 comparison。

[ ] 使用 Task 1 冻结的 397B config，分别运行 --mode prefill --isl 4096 与 --isl 32768。
[ ] 两侧均使用 --warmup 3 --iterations 10 --dtype bfloat16，每份 artifact 先写 Pod-local /tmp，校验后发布到正式 evidence 目录。
[ ] 要求 output/final recurrent state parity、finite、source SHA、effective path 和 readback 全部通过。
[ ] 生成 kernel comparison；不在此处决定最终赢家。

Task 4：完整 GDN layer A/B

Uses: benchmark.kernels.gdn_compare.runner --level L3

Produces: Prefill 4K fresh 与 Decode batch 16 continuing 的完整 layer timing。

[ ] Prefill 使用 397B 真实 hidden size 和 Task 1 推导的 per-rank projection shape。
[ ] Decode fixture 固定为 16 个 continuing requests，计入 Conv/state gather-scatter 和完整 layer shell。
[ ] 每侧每 cell 3 warmup + 5 measured；校验 final hidden、Conv state、recurrent state、finite 和 anti-fallback。
[ ] 若 L1 收益在 L3 消失或反转，记录 integration delta；默认不启动 XProf。

Task 5：完整模型 offline A/B

Uses: python -m sgl_jax.bench_one_batch

Produces: 两个 model cells 各 5 rounds 的 raw JSON 和 median comparison。

[ ] Reference 使用 batch=16, input_len=4096, output_len=1，完成后完全停止进程。
[ ] Hybrid 使用相同配置和独立空 cache。
[ ] 重新按相同顺序执行 batch=16, input_len=1024, output_len=100。
[ ] 每个 cell 先做 shape-matched warmup，再做 5 个独立 rounds；报告 prefill throughput/latency 和 median decode throughput/latency。
[ ] 保留 exact-token 与 tolerance 两套判定；token mismatch 阻断 production-ready 结论，但不删除性能样本。

Task 6：Serving A/B

Uses: .claude/skills/model-speed-benchmark/scripts/speed_benchmark.sh

Produces: 两个 serving cells 的 validated summary.csv、result.jsonl 和 report。

[ ] 分别启动 reference 与 hybrid server，固定模型、serve args、禁用 radix cache、端口、runtime 和 topology。
[ ] 每侧运行 random 4096→1, concurrency=16 和 1024→100, concurrency=16，各 5 独立 rounds。
[ ] 每轮检查 health、模型 ID、PID、cache hit rate、completed/error counts、finite 和 output artifact。
[ ] 发布后回读 SHA，并停止本次明确 PID；不得 kill 其他实验进程。

Task 7：收口与决策

Files:

Create: benchmark/kernels/gdn_compare/discovery/phase8-397b-bf16-v7-report.md

Produces: 四层一张对比表、可追溯证据和 mentor 决策输入。

[ ] 汇总 kernel、layer、model、serving 的 reference/hybrid median 与变化。
[ ] 以 layer/model/serving 共同结果作结论，kernel 只解释来源。
[ ] 明确披露 35B 历史 token divergence、397B exact-token 结果和 production-readiness 边界。
[ ] 仅在层级结果矛盾或 effective path 不清楚时提出新的 XProf 计划；本阶段不自动执行。
[ ] 保持实验分支本地、无生产 PR、无 merge；等待 mentor 决定是否提取最小 GDN 代码到最新 main。

完成条件

Phase 8 只有在以下条件全部成立时才可称为完成：

四层 reference/hybrid 均在同一 397B BF16/v7 配置上完成；
所有 raw samples、manifest、SHA readback 和 .complete 存在；
no fallback、no nonfinite、no OOM/collective/XLA/artifact failure；
exact token 与 numerical tolerance 分开报告；
完整 layer、model、serving 的结论一致或差异已明确解释；
报告没有把 35B、v6e 或 FP8 数据混入新的 397B 排名。

若只完成部分层级，状态必须写为 partial，并逐项列出未运行内容。

Phase 9–10 收口记录（2026-07-29）

本节记录 Phase 8 修订之后的实际执行结果，并作为原 Phase 9/10 表格的 authoritative closeout。它不授权 production merge 或默认 selector 变更。

Phase 9：397B BF16 serving 路线选择

Phase 9 已完成 serving-only 路线比较。固定环境保持为 Qwen/Qwen3.5-397B-A17B、BF16、TPU v7x 4x4x4、JAX/jaxlib 0.10.2、libtpu 0.0.43、TP128/DP8/EP128。

四条形成完整、finite serving summary 的路线为：

路线	Prefill	Decode
Reference	SGL-JAX reference	SGL-JAX reference
Hybrid	SGL-JAX chunkwise KDA/Pallas	SGL-JAX reference
TPU-Inference full	TPU-Inference fused	TPU-Inference fused
Selected composition	TPU-Inference fused	SGL-JAX reference

正式 concurrency-16 结果：

路线	4096→1 input throughput	4096→1 TTFT	1024→100 output throughput	1024→100 TPOT
Reference	10,428.18 tok/s	5,733.42 ms	294.67 tok/s	35.797 ms
Hybrid	30,378.93 tok/s	1,965.01 ms	352.23 tok/s	35.816 ms
TPU-Inference full	67,411.91 tok/s	914.98 ms	314.87 tok/s	47.280 ms
Selected composition	67,588.443 tok/s	913.453 ms	421.975 tok/s	34.952 ms

最终 serving 赢家是 TPU-Inference prefill 与 SGL-JAX reference decode 的组合。相对 Reference，其 prefill throughput 提升 548.133%，decode output throughput 提升 43.203%；相对 TPU-Inference full，其 prefill 基本持平，同时 decode output throughput 提升 34.016%、TPOT 降低 26.074%。

primatrix_latest serving attempt 在 Falcon 配置中记录为 disqualified-nonfinite，没有形成 ranking-eligible formal summary，因此不进入上表。该结论只约束本次 serving attempt，不外推为对 Primatrix upstream KDA 的通用性能判断。

最终 Falcon experiment 为 exp-3emc9xmuuq，artifact 为 art-6abvn78e1z。外层 experiment 因手工恢复和 orchestration cleanup 结束为 failed，但 formal benchmark summary 通过；两种状态必须分开披露。Lease npl-x5kxku2n3b 已释放，16 个 TPU v7 nodes 已归还。

Phase 10：397B BF16-only 矩阵

Phase 10 在 2026-07-28 批准的修订范围下完成。该修订明确覆盖原 BF16/FP8 双矩阵要求：

执行 Qwen3.5-397B-A17B BF16；
使用 TPU v7；
验证 kernel、完整 GDN layer、完整模型和 serving 四层；
FP8 不在本次范围内。

因此 FP8 未运行是批准后的 scope decision，不是缺失验收项。Phase 8 完成 Reference/Hybrid 四层 A/B 和 64-token exact-match smoke；Phase 9 完成额外 serving 路线选择。Phase 10 的技术状态为 complete。

正式报告位于实验分支：

benchmark/kernels/gdn_compare/discovery/phase8-397b-bf16-v7-report.md
benchmark/kernels/gdn_compare/discovery/phase9-serving-four-route-report.md
benchmark/kernels/gdn_compare/discovery/phase10-397b-bf16-closeout.md

下一门禁

Phase 11 尚未开始。下一步只能在 mentor 明确决策后选择以下之一：

把 selected composition 的最小实现移植到最新 SGL-JAX main 的新 production branch；
追加 correctness、maintainability 或 upstream-sync 证据；
保留实验结论并停止，不创建 production PR。

Phase 11：Selected composition 生产化（mentor 已批准）

决策与目标

Mentor 已批准把 Phase 9 选出的 serving 赢家生产化：

text

Prefill: vllm-project/tpu-inference v3 fused Conv1D + GDN
Decode:  SGL-JAX existing reference decode

本阶段在最新 SGL-JAX main 上创建新的 production branch 和新的 Draft PR。它不修改或复用 PR #1466 的分支历史；#1466 继续表示旧的 “chunkwise KDA prefill + reference decode”方案，只有在新 PR 已建立且 mentor 确认 supersede 后才单独处理其状态。

成功标准是把实验中的 selected composition 变为 SGL-JAX 仓库内可独立构建、测试和运行的实现，并用生产分支重新验证 Phase 9 的 correctness 与 serving 结论。实验分支已有数据是设计依据，不替代生产端复验。

方案选择

评估过三种方式：

新建自包含 production PR（采用）：从最新 main 建立独立分支，把冻结的 TPU-Inference v3 prefill 内核移植到 SGL-JAX，保留 reference decode。依赖边界清晰，review 可以独立于 #1466 进行。
重写 PR #1466（拒绝）：会把旧 KDA 路线、旧 35B/v6e 证据与新的 TPU-Inference/v7 路线混在同一历史中，难以 review 和回滚。
保留外部 checkout bridge（拒绝）：实验用 PYTHONPATH/动态 import 可验证路线，但不满足生产代码自包含、可打包和 CI 可复现要求。

固定上游与代码边界

TPU-Inference 源码冻结为：

Repository: https://github.com/vllm-project/tpu-inference
Commit: a9072c881843622226efc101de1a62c731ab572f
Source: tpu_inference/kernels/gdn/v3/
License: Apache License 2.0

该目录由 8 个紧耦合 Python 模块组成，共 2153 行：

text

__init__.py
compute_conv1d.py
compute_gdn.py
config.py
memory_ref.py
metadata.py
vmem_ldst.py
wrapper.py

这些模块在冻结 revision 中只依赖 JAX/Pallas 和同目录模块。生产移植必须：

保留各文件的 Google copyright 与 Apache-2.0 license header；
把内部 import 改为 SGL-JAX 包内路径；
在模块或 provenance 文档中记录冻结 repository、commit 和 source path；
不引入运行时 external checkout、PYTHONPATH 或动态 import；
不复制 TPU-Inference 中与 v3 fused GDN prefill 无关的训练、服务或其他 kernel 代码。

如果实现过程中发现 v3 目录外的隐藏运行时依赖，必须暂停并修订 RFC，不能静默扩大 vendoring scope。

SGL-JAX 集成设计

从实施开始时实时读取的最新 origin/main 创建隔离 worktree。当前审计到的 main GDN backend 只包含 reference extend/decode；实现不得基于本地过期 main 或实验分支直接提 PR。

初始化时新增一个明确、冻结的 GDN prefill selector：

text

reference        -> existing SGL-JAX reference prefill
tpu_inference_v3 -> vendored fused Conv1D + GDN prefill

首个 production PR 的默认值保持 reference。只有显式 opt-in 才使用 tpu_inference_v3；未经 mentor 单独批准，不把新路径设为默认。初始化必须记录：

requested_prefill_impl
effective_prefill_impl
fallback_reason

非法 selector 立即 ValueError。请求 tpu_inference_v3 但 shape、 platform、dtype 或 capability 不满足时，默认不得 silent fallback；应在启动阶段以明确错误停止。若后续确需 fallback，必须另行设计和批准。

数据流固定为：

沿用 SGL-JAX 现有 TP-striped mixed_qkv、state pool 和 forward metadata；
prefill adapter 只做 TPU-Inference v3 API 所需的 layout、metadata、weight 和 state contract 转换；
forward_extend 调用 vendored fused_conv1d_gdn；
返回并消费 output、final Conv state 与 final recurrent state；
forward_decode 始终继续调用现有 SGL-JAX decode_gated_delta_rule_ref 和现有 Conv1D update；
不修改 scheduler、state-pool schema、DP/TP/EP 语义、其他 attention backend 或模型结构。

Selector 在 backend 初始化时冻结，hot path 不读取环境变量。PR 中不包含 primatrix_latest、tpu_inference_latest、统一 benchmark 的其他候选 adapter 或实验 orchestration。

正确性与行为契约

先以 TDD 新增失败测试，再写生产代码。至少验证：

默认 reference、显式 tpu_inference_v3、非法值和 capability failure；
requested/effective/fallback 日志字段完整，selector 初始化后不可变化；
explicit path 实际调用 vendored fused kernel，不能别名 reference 或丢弃 optimized output/state；
decode 在两种 prefill selector 下都调用同一个 existing reference path；
fused prefill 与 independent reference 对比：
- output；
- final Conv state；
- final recurrent state；
- finite values；
fresh/continuing state、packed/ragged requests、partial tile、zero-length、 dummy/unused slot、slot isolation 和 track/checkpoint 语义；
prefill→decode continuity；
single-device/interpret correctness 与 TPU sharded correctness；
固定 seed 的 model smoke 同时报告 tolerance 与 exact token identity。

BF16 容差沿用本 RFC 的既有契约：

模式	rtol	atol
single-device / non-DP	`2e-2`	`1e-2`
DP/TP sharded	`2e-2`	`5e-2`

如果 frozen upstream v3 对某个 SGL-JAX 目标语义不支持，必须在测试中暴露并停止；不能通过跳过目标 case、放宽容差或 fallback 获得假 GREEN。

性能复验与验收

生产移植完成后，在同一生产 commit 上重新执行最小 397B BF16 / TPU v7 A/B。不得直接把 Phase 9 实验分支结果当作 production 结果。

固定环境：

Model: Qwen/Qwen3.5-397B-A17B
Model revision: 8472618112abcbd45acbcdc58436aff4233c23f7
Runtime target: JAX/jaxlib 0.10.2、libtpu 0.0.43
TPU: v7x 4x4x4
Parallelism: TP128 / DP8 / EP128
dtype: BF16

最小复验矩阵：

层级	Reference	Production selected composition
focused kernel/layer correctness	existing reference	vendored fused prefill + reference decode
model smoke	fixed seed reference	same prompts, same serving args
serving prefill	4096→1, concurrency 16	same
serving decode	1024→100, concurrency 16	same

每个 serving cell 在 shape-matched warmup 后运行 5 个独立 rounds，报告 raw samples、median、min/max、input/output throughput、TTFT 和 TPOT。两侧使用同一模型、runtime、slice、serve args、dataset、seed、并行配置，以及分别独立且初始为空的 compilation cache。

验收要求：

no fallback、no OOM、no NaN/Inf、no collective/XLA/artifact failure；
output/state correctness 和 fixed-seed smoke 通过；
prefill throughput 相对 reference 保持显著正收益；
decode 使用相同 reference 路径，TPOT 回归不得超过 5%；
production selected composition 必须在 XProf/启动日志中证明 effective prefill 是 vendored tpu_inference_v3；
raw JSON、manifest、source SHA、runtime/topology、SHA readback 和 completion marker 完整。

若 production 数字与 Phase 9 差异明显，优先诊断移植、runtime、shape、 cache 或 serving 参数差异，不得选择性沿用更好的一组旧数据。

PR 边界与停止条件

新 Draft PR 只包含：

vendored v3 GDN prefill 内核与 attribution/provenance；
最小 SGL-JAX prefill adapter/selector；
focused correctness 与 dispatch tests；
必要的文档和 benchmark command/result summary。

明确不包含：

raw logs、XProf trace、模型权重或实验 evidence 目录；
Phase 0–10 的统一候选 harness；
Primatrix KDA 路线；
scheduler/state-pool/其他 backend 重构；
默认 selector 切换；
#1466 的历史重写。

任一情况发生时立即暂停：

需要 v3 目录外未批准的 source dependency；
output、Conv state、recurrent state 或 token identity 违反批准契约；
requested/effective 不一致或发生 fallback；
nonfinite、OOM、collective、XLA 或 sharding failure；
decode TPOT 回归超过 5%；
为通过测试需要修改 scheduler、state-pool schema 或无关 backend。

完成本地、TPU correctness、397B serving 复验和独立 whole-branch review 后，才创建新的 Draft PR。新 PR 建立后，由 mentor 决定是否将 #1466 标记为 superseded/关闭；本阶段不得提前关闭 #1466。

Phase 11 完成条件

只有以下条件全部满足，Phase 11 才可标记完成：

production branch 基于实施时最新 origin/main；
vendored kernel 自包含且 license/provenance 完整；
reference decode 保持不变；
focused、sharded、model smoke 和 serving A/B 全部通过；
production evidence 可回读；
新 Draft PR 创建且 diff 不包含实验垃圾或无关改动；
#1466 状态处理由 mentor 单独确认。

Phase 11 Selected Composition Productionization Implementation Plan

For agentic workers: REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (- [ ]) syntax for tracking.

Goal: 在实施时最新的 SGL-JAX main 上，自包含移植冻结的 TPU-Inference v3 fused Conv1D+GDN prefill，继续使用现有 SGL-JAX reference decode，并在创建新 Draft PR 前完成本地、TPU sharded、397B model smoke 和 serving A/B 验证。

Architecture: 冻结的八个 v3 模块作为带 attribution 的包内 vendor 代码进入 sgl_jax.srt.kernels.gdn.tpu_inference_v3；薄 adapter 负责 SGL-JAX layout、metadata 和 state-pool/track scatter；GDNAttnBackend 在初始化时把 reference 或 tpu_inference_v3 冻结为 prefill callable。 forward_decode 保持现有 reference 实现，不经过 vendor 代码。默认 selector 保持 reference，请求不受支持的新路径时明确失败，不 silent fallback。

Tech Stack: Python 3.12、JAX/Pallas、SGL-JAX、pytest、Falcon、TPU v7x 4x4x4、Qwen3.5-397B-A17B BF16、GitHub Draft PR。

Global Constraints

实施开始先执行 git fetch origin main。当前已观察到的远端基线是 a7353325e8c00d287294c2cd679a77173f1a4594；若实时 origin/main 已变化，先审计从该 SHA 到新 HEAD 是否触及 GDN、Qwen3.5、 recurrent state 或测试注册，再记录实际 BASE_SHA。不得直接使用本地落后的 main checkout。
新分支固定命名为 feat/gdn-tpu-inference-v3-prefill，独立 worktree 为 .claude/worktrees/feat/gdn-tpu-inference-v3-prefill。不得继承 experiment/gdn-path-comparison 或 PR #1466 的提交历史。
Vendor source 只能来自 vllm-project/tpu-inference@a9072c881843622226efc101de1a62c731ab572f 的 tpu_inference/kernels/gdn/v3/ 八个模块。若发现目录外运行时依赖，立即暂停并修订 RFC。
环境变量固定为 SGLANG_JAX_GDN_PREFILL_IMPL，仅接受 reference、tpu_inference_v3；默认 reference。
新路径只改变 prefill。Decode、scheduler、state-pool schema、DP/TP/EP 语义、其他 backend 和模型结构均不得改变。
RED 必须因目标能力缺失而失败；import、fixture、依赖、路径或设备初始化失败不算有效 RED。
每个 Task 完成后提交一个外科手术式 commit、重跑该 Task 的 focused verification 并暂停检查。不要在每个小编辑后重跑完整回归；集中回归只在 Task 4、Task 5 和 Task 7 执行。
Track/checkpoint、zero-length、state-slot 与 prefill→decode 是硬契约。无法实现时暂停，不能 skip、放宽容差、丢弃 state 或 fallback。
Raw logs、模型权重、XProf、Falcon artifacts 和 evidence 目录不得进入 Git。正式 evidence 使用 /data/user/dongbohao/gdn-path-comparison/evidence/phase11-production/{run_id}/。
在 Task 1–7 全部通过、whole-branch review 无阻断 finding 前，不 push、不创建 Draft PR、不修改 #1466 状态。

Task 1：冻结 production 基线并 vendor v3 import closure

Files:

Create: python/sgl_jax/srt/kernels/gdn/tpu_inference_v3/__init__.py
Create: python/sgl_jax/srt/kernels/gdn/tpu_inference_v3/compute_conv1d.py
Create: python/sgl_jax/srt/kernels/gdn/tpu_inference_v3/compute_gdn.py
Create: python/sgl_jax/srt/kernels/gdn/tpu_inference_v3/config.py
Create: python/sgl_jax/srt/kernels/gdn/tpu_inference_v3/memory_ref.py
Create: python/sgl_jax/srt/kernels/gdn/tpu_inference_v3/metadata.py
Create: python/sgl_jax/srt/kernels/gdn/tpu_inference_v3/vmem_ldst.py
Create: python/sgl_jax/srt/kernels/gdn/tpu_inference_v3/wrapper.py
Create: python/sgl_jax/srt/kernels/gdn/tpu_inference_v3/PROVENANCE.md
Create: python/sgl_jax/test/kernels/test_gdn_tpu_inference_v3_vendor.py

Interface: vendor package must export the frozen upstream fused_conv1d_gdn entrypoint without importing outside its package except JAX/Pallas or Python stdlib.

[ ] Step 1.1 — Create the isolated worktree from live origin/main.

Run:

bash

cd /Users/feibo/projects/intership/sglang-jax
git fetch origin main
git rev-parse origin/main
git log --oneline a7353325e..origin/main -- \
  python/sgl_jax/srt/kernels/gdn \
  python/sgl_jax/srt/layers/attention/linear/gdn_backend.py \
  python/sgl_jax/srt/models/qwen3_5.py \
  python/sgl_jax/srt/configs/qwen3_5.py
git worktree add \
  .claude/worktrees/feat/gdn-tpu-inference-v3-prefill \
  -b feat/gdn-tpu-inference-v3-prefill origin/main

Expected: one actual BASE_SHA is recorded; the new worktree is clean and its HEAD equals origin/main. If the scoped log is non-empty, review those commits before proceeding.

[ ] Step 1.2 — Write the vendor provenance/import-closure RED test.
Test assertions:
- exactly eight .py modules exist;
- each upstream-derived module retains the Google/Apache-2.0 header;
- package imports do not reference tpu_inference.* or filesystem paths;
- PROVENANCE.md records repository, commit, source path and license;
- package exports fused_conv1d_gdn.
Run:
bash
```
uv run --project python --extra cpu --with pytest --frozen \
  python -m pytest \
  python/sgl_jax/test/kernels/test_gdn_tpu_inference_v3_vendor.py \
  -q --tb=short
```
Expected RED: collection succeeds, then assertions fail because the vendor package/provenance do not exist. A dependency or import-collection error is not an accepted RED.
[ ] Step 1.3 — Copy only the approved frozen modules and rewrite internal imports.
Fetch the upstream repository into /tmp/gdn-phase11-production/tpu-inference, checkout exact commit a9072c881843622226efc101de1a62c731ab572f, verify clean state, copy the eight approved files, then replace only package-internal imports with sgl_jax.srt.kernels.gdn.tpu_inference_v3.*. Preserve source bodies and license headers.

[ ] Step 1.4 — Add exact provenance and verify GREEN.

PROVENANCE.md must record:

text

repository: https://github.com/vllm-project/tpu-inference
commit: a9072c881843622226efc101de1a62c731ab572f
source: tpu_inference/kernels/gdn/v3
license: Apache-2.0

Run:

bash

uv run --project python --extra cpu --with pytest --frozen \
  python -m pytest \
  python/sgl_jax/test/kernels/test_gdn_tpu_inference_v3_vendor.py \
  -q --tb=short
PYTHONPYCACHEPREFIX=/tmp/gdn-phase11-production/pycache \
  uv run --project python --extra cpu --frozen \
  python -m py_compile \
  python/sgl_jax/srt/kernels/gdn/tpu_inference_v3/*.py
git diff --check

Expected GREEN: pytest exits 0; all eight modules compile; no whitespace errors.

[ ] Step 1.5 — Commit the frozen vendor snapshot.

bash

git add python/sgl_jax/srt/kernels/gdn/tpu_inference_v3 \
  python/sgl_jax/test/kernels/test_gdn_tpu_inference_v3_vendor.py
git commit -m "feat(gdn): vendor TPU-Inference v3 prefill kernel"

Task 2：实现 selector、capability gate 与冻结 dispatch

Files:

Create: python/sgl_jax/srt/kernels/gdn/tpu_inference_adapter.py
Modify: python/sgl_jax/srt/kernels/gdn/__init__.py
Modify: python/sgl_jax/srt/layers/attention/linear/gdn_backend.py
Create: python/sgl_jax/test/test_gdn_tpu_inference_prefill_dispatch.py

Interface:

python

requested_prefill_impl: str
effective_prefill_impl: str
fallback_reason: str | None
_prefill_callable: Callable[..., tuple[Array, Array, Array]]

The accepted selector values are exactly reference and tpu_inference_v3.

[ ] Step 2.1 — Add dispatch RED tests.
Cover:
- unset env selects reference;
- explicit reference selects existing prefill;
- explicit tpu_inference_v3 selects a distinct vendor adapter callable;
- invalid value raises ValueError during backend initialization;
- requested unsupported platform/dtype/shape raises a typed startup error;
- initialized callable remains frozen after env changes;
- startup log contains requested/effective/fallback fields once;
- both selector values leave decode bound to decode_gated_delta_rule_ref.
Run:
bash
```
uv run --project python --extra cpu --with pytest --frozen \
  python -m pytest \
  python/sgl_jax/test/test_gdn_tpu_inference_prefill_dispatch.py \
  -q --tb=short
```
Expected RED: tests fail only because selector fields, validation and callables do not exist.
[ ] Step 2.2 — Implement initialization-only selector and capability gate.
Read SGLANG_JAX_GDN_PREFILL_IMPL once in GDNAttnBackend.__init__. Keep default reference. For tpu_inference_v3, validate from the supplied mesh/config only:
- all mesh devices are TPU, or Pallas interpret is explicitly enabled for local correctness;
- dtype is BF16;
- head_k_dim, head_v_dim, head counts and kernel size satisfy vendor static constraints.
Do not call global jax.devices()/jax.default_backend() in the gate. Do not set an automatic fallback reason; unsupported explicit requests must raise before serving.
[ ] Step 2.3 — Freeze prefill and decode callables and verify GREEN.
forward_extend must call _prefill_callable; forward_decode must keep the existing reference decode and Conv1D update without branching on the prefill selector.
Run the focused test from Step 2.1. Expected: exit 0 and log assertions pass.

[ ] Step 2.4 — Commit selector/dispatch.

bash

git add \
  python/sgl_jax/srt/kernels/gdn/__init__.py \
  python/sgl_jax/srt/kernels/gdn/tpu_inference_adapter.py \
  python/sgl_jax/srt/layers/attention/linear/gdn_backend.py \
  python/sgl_jax/test/test_gdn_tpu_inference_prefill_dispatch.py
git commit -m "feat(gdn): select TPU-Inference v3 prefill at startup"

Task 3：对齐 metadata、state-pool、track/checkpoint 契约

Files:

Modify: python/sgl_jax/srt/kernels/gdn/tpu_inference_adapter.py
Create: python/sgl_jax/test/test_gdn_tpu_inference_state_contract.py

Interface:

python

def fused_conv1d_gdn_prefill(
    mixed_qkv,
    b,
    a,
    conv_state,
    recurrent_state,
    conv_weight,
    a_log,
    dt_bias,
    cu_seqlens,
    state_indices,
    track_indices,
    has_initial_state,
    seq_lens,
    *,
    n_kq,
    n_v,
    d_k,
    d_v,
    kernel_size,
) -> tuple[output, new_conv_state_pool, new_recurrent_state_pool]:
    ...

[ ] Step 3.1 — Add state-contract RED tests using a deterministic fake vendor kernel.
Cover:
- TP-striped mixed_qkv and Conv weight/state layout conversion;
- query_start_loc, distribution, state_indices, seq_lens;
- fresh and continuing state gather;
- scatter to running state slot;
- optional scatter to distinct track/checkpoint slot;
- dummy slot 0 remains unchanged;
- unused slots remain unchanged;
- zero-length request preserves gathered initial state;
- duplicate/invalid track indices fail explicitly rather than corrupt pool;
- output, Conv state and recurrent state from the vendor callable are all consumed.
Run:
bash
```
uv run --project python --extra cpu --with pytest --frozen \
  python -m pytest \
  python/sgl_jax/test/test_gdn_tpu_inference_state_contract.py \
  -q --tb=short
```
Expected RED: failures stop at missing metadata/scatter behavior, not at fixtures or imports.
[ ] Step 3.2 — Implement the thin adapter without duplicating GDN math.
The adapter may perform layout conversion, metadata construction, gather/mask/scatter and dtype restoration. It must not precompute fused Conv1D/GDN math, activate gate twice or call reference prefill when the requested path is tpu_inference_v3.
[ ] Step 3.3 — Enforce the track/checkpoint hard gate.
Run the focused state-contract suite. Expected GREEN: all pool/track assertions pass.
STOP: if the frozen vendor return values cannot implement exact running and track state semantics without changing files outside the approved v3 package/adapter/backend boundary, write a blocker report and stop before Task 4. Do not reject all non-null track_indices as a production shortcut.

[ ] Step 3.4 — Commit the contract adapter.

bash

git add \
  python/sgl_jax/srt/kernels/gdn/tpu_inference_adapter.py \
  python/sgl_jax/test/test_gdn_tpu_inference_state_contract.py
git commit -m "feat(gdn): adapt v3 prefill state contracts"

Task 4：完成真实 numerical TDD 与本地集中回归

Files:

Create: python/sgl_jax/test/test_gdn_tpu_inference_prefill.py
Create: python/sgl_jax/test/test_gdn_tpu_inference_prefill_dp.py
Modify only when registering tests: test/srt/run_suite.py
[ ] Step 4.1 — Add actual vendor-vs-reference RED cases.
Use non-zero, non-constant, fixed-seed BF16 inputs. Compare actual vendor output, final Conv state and final recurrent state with the independent existing GDN reference for:
- lengths 1, 63, 64, 65, 127, 128, 129;
- packed/ragged and partial tile;
- fresh, continuing and mixed state;
- zero-length request;
- dummy/unused/track slot isolation;
- prefill→decode continuity;
- finite values and input immutability.
Single-device tolerance: rtol=2e-2, atol=1e-2. Sharded tolerance: rtol=2e-2, atol=5e-2.
Run:
bash
```
PALLAS_INTERPRET=true \
uv run --project python --extra cpu --with pytest --frozen \
  python -m pytest \
  python/sgl_jax/test/test_gdn_tpu_inference_prefill.py \
  -q --tb=short
```
Expected initial RED: the real vendor path is invoked and numerical/state assertions expose adapter or kernel-contract gaps. A test that accidentally calls reference for both sides is invalid.
[ ] Step 4.2 — Make the smallest approved production corrections.
Corrections may touch only:
- tpu_inference_v3/;
- tpu_inference_adapter.py;
- gdn_backend.py;
- gdn/__init__.py.
Preserve the frozen upstream math unless a clearly documented SGL-JAX integration conversion is required. Any new dependency outside v3 is a stop condition.

[ ] Step 4.3 — Run focused single-device and four-device GREEN.

bash

PALLAS_INTERPRET=true \
uv run --project python --extra cpu --with pytest --frozen \
  python -m pytest \
  python/sgl_jax/test/test_gdn_tpu_inference_prefill.py \
  -q --tb=short

XLA_FLAGS=--xla_force_host_platform_device_count=4 \
PALLAS_INTERPRET=true \
uv run --project python --extra cpu --with pytest --frozen \
  python -m pytest \
  python/sgl_jax/test/test_gdn_tpu_inference_prefill_dp.py \
  -q --tb=short

Expected GREEN: both commands exit 0; no skip for target contract cases; all outputs/states are finite.

[ ] Step 4.4 — Run one concentrated local regression batch.

bash

PALLAS_INTERPRET=true \
uv run --project python --extra cpu --with pytest --frozen \
  python -m pytest \
  python/sgl_jax/test/kernels/test_gdn_tpu_inference_v3_vendor.py \
  python/sgl_jax/test/test_gdn_tpu_inference_prefill_dispatch.py \
  python/sgl_jax/test/test_gdn_tpu_inference_state_contract.py \
  python/sgl_jax/test/test_gdn_tpu_inference_prefill.py \
  python/sgl_jax/test/test_gdn_attention.py \
  python/sgl_jax/test/models/test_qwen3_5.py \
  -q --tb=short

XLA_FLAGS=--xla_force_host_platform_device_count=4 \
PALLAS_INTERPRET=true \
uv run --project python --extra cpu --with pytest --frozen \
  python -m pytest \
  python/sgl_jax/test/test_gdn_tpu_inference_prefill_dp.py \
  python/sgl_jax/test/test_gdn_attention_dp.py \
  -q --tb=short

PYTHONPYCACHEPREFIX=/tmp/gdn-phase11-production/pycache \
  uv run --project python --extra cpu --frozen \
  python -m py_compile \
  python/sgl_jax/srt/kernels/gdn/tpu_inference_v3/*.py \
  python/sgl_jax/srt/kernels/gdn/tpu_inference_adapter.py \
  python/sgl_jax/srt/layers/attention/linear/gdn_backend.py
git diff --check

Expected: all pytest commands exit 0; py_compile and diff check exit 0.

[ ] Step 4.5 — Register tests and commit numerical GREEN.

Add the single-device suite to the TPU-v6e-1 lane and the sharded suite to the TPU-v6e-4 lane, matching adjacent GDN registrations.

bash

git add \
  python/sgl_jax/test/test_gdn_tpu_inference_prefill.py \
  python/sgl_jax/test/test_gdn_tpu_inference_prefill_dp.py \
  test/srt/run_suite.py \
  python/sgl_jax/srt/kernels/gdn \
  python/sgl_jax/srt/layers/attention/linear/gdn_backend.py
git commit -m "test(gdn): validate TPU-Inference v3 prefill"

Task 5：在 TPU v7x 完成 sharded correctness 与 397B model smoke

Transient files:

/tmp/gdn-phase11-production/falcon-v7x64.yaml
/tmp/gdn-phase11-production/runtime-manifest.json
/tmp/gdn-phase11-production/model-smoke/{reference,tpu_inference_v3}/

Tracked files: none unless a verified result summary is later added in Task 7.

[ ] Step 5.1 — Reconfirm runtime/resource preflight.
Run falcon auth status, falcon cluster list, and the cluster resource query used by Phase 9. Select one authorized v7x 4x4x4 slice; record cluster ID, reservation, experiment ID, image, JAX/jaxlib, libtpu and topology. Require:
text
```
model revision = 8472618112abcbd45acbcdc58436aff4233c23f7
JAX/jaxlib = 0.10.2
libtpu = 0.0.43
TP/DP/EP = 128/8/128
dtype = bfloat16
```
If no authorized slice is immediately available, pause without creating a sleeping holder.
[ ] Step 5.2 — Render a finite Falcon workload.
Reuse the proven Phase 9 mount/runtime pattern, but the command must:
1. checkout this production branch commit;
2. run focused sharded GDN correctness;
3. run fixed-seed 397B prefill/extend/decode smoke for reference;
4. stop its processes;
5. repeat with tpu_inference_v3;
6. publish evidence and exit.
Do not end with sleep infinity.
[ ] Step 5.3 — Verify effective path and state correctness on all hosts.
Require every rank/host to record:
- source/commit SHA and clean state;
- requested/effective selector and empty fallback;
- rank→PJRT permutation;
- output/logit, Conv state and recurrent-state finite checks;
- tolerance metrics;
- 64-token exact identity and first divergence if unequal;
- non-empty vendor Pallas op evidence for the optimized prefill.
Expected: four-host exit code 0, no OOM/collective/XLA/sharding failure.
[ ] Step 5.4 — Publish/read back evidence and release resources.
Publish through Pod-local staging, verify JSON schema, size and SHA-256, create .complete, then confirm no server/benchmark process or occupied port remains and release the Falcon lease.
STOP: any token/state contract violation, fallback or nonfinite result blocks Task 6 and Draft PR creation.

Task 6：执行 production 397B serving A/B

Uses:

.claude/skills/model-speed-benchmark/scripts/speed_benchmark.sh
/data/user/dongbohao/gdn-path-comparison/evidence/phase11-production/{run_id}/

Variants:

Variant	Prefill	Decode
`reference`	existing reference	existing reference
`tpu_inference_v3`	vendored fused v3	existing reference

[ ] Step 6.1 — Freeze one serve config and two empty caches.
Both variants use the Task 5 model/runtime/topology, identical serve args, random dataset, seed and workload. The only differences are SGLANG_JAX_GDN_PREFILL_IMPL and isolated cache/output paths. Record cache emptiness before launch.
[ ] Step 6.2 — Run reference serially.
After shape-matched warmup, run:
- 4096→1, concurrency 16, five independent rounds;
- 1024→100, concurrency 16, five independent rounds.
Require health/model discovery, completed/error counts, finite metrics, raw JSON and completion marker. Stop only the exact recorded reference PIDs before proceeding.
[ ] Step 6.3 — Run tpu_inference_v3 serially.
Repeat the identical matrix with a fresh empty cache. Require startup log requested=effective=tpu_inference_v3, no fallback, and trace/startup evidence that the vendored prefill callable ran.
[ ] Step 6.4 — Compute the locked comparison.
Report all five raw samples plus median/min/max for:
- input throughput and TTFT on prefill;
- output throughput and TPOT on decode.
Acceptance:
- prefill retains a clear positive improvement over reference;
- decode TPOT regression is at most 5%;
- no sample is silently removed;
- exact-token and numerical results are reported separately.
STOP: if decode TPOT regresses by more than 5%, do not create a Draft PR. Diagnose shared runtime/scheduling differences before changing decode code, because both variants use the same decode implementation.
[ ] Step 6.5 — Publish and release.
Verify manifests and SHA readback, write .complete, stop exact experiment processes, release the v7 lease and record final resource state.

Task 7：文档、whole-branch review 与新 Draft PR

Files:

Modify: docs/features/server_arguments.md
Create: docs/developer_guide/gdn_tpu_inference_v3.md
Modify only for validated fixes from review: files already in this plan.
[ ] Step 7.1 — Document opt-in behavior and provenance.
Document:
- SGLANG_JAX_GDN_PREFILL_IMPL;
- default reference;
- explicit tpu_inference_v3;
- TPU/BF16/capability requirements;
- no silent fallback;
- reference decode unchanged;
- frozen upstream repository/commit/license;
- exact local and TPU verification commands;
- production serving A/B summary with hardware/runtime/model/workload.
[ ] Step 7.2 — Run final concentrated verification.
Re-run Task 4’s two regression commands, formatter/lint versions used by repository CI, py_compile, git diff --check, and the production evidence validator. Verify:
bash
```
git status --short
git diff --name-only origin/main...HEAD
git log --oneline origin/main..HEAD
```
Expected: worktree clean after committing; diff contains only approved vendor, adapter/backend, focused tests, test registration and docs. No experiment harness, model, logs or evidence.
[ ] Step 7.3 — Request independent whole-branch review.
Review specifically for:
- source/license attribution;
- hidden external dependencies;
- fake optimized/reference aliasing;
- lost output/Conv/recurrent state;
- track/checkpoint and zero-length correctness;
- selector hot-path reads or silent fallback;
- decode changes;
- scope creep;
- claims not backed by production evidence.
Resolve all Critical/Important findings and rerun only the affected focused tests plus the final concentrated verification once.

[ ] Step 7.4 — Commit documentation and final fixes.

bash

git add docs/features/server_arguments.md \
  docs/developer_guide/gdn_tpu_inference_v3.md
git commit -m "docs(gdn): document TPU-Inference v3 prefill"

[ ] Step 7.5 — Push and create a new Draft PR.
Only after every prior gate passes:
bash
```
git push -u origin feat/gdn-tpu-inference-v3-prefill
```
Draft PR title:
text
```
feat(gdn): add TPU-Inference v3 prefill path
```
PR body must include:
- Phase 9 decision and why this is a new PR rather than #1466;
- frozen upstream commit/license/provenance;
- file/scope boundary;
- focused, sharded and model-smoke results;
- full production serving A/B raw-sample summary;
- decode TPOT gate;
- evidence URI and manifest SHA;
- default remains reference;
- #1466 remains unchanged pending mentor decision.
Keep the PR Draft. Do not request review, merge, close #1466 or change its status until mentor reviews this new PR.

Final Definition of Done

Phase 11 is implementation-complete only when:

Tasks 1–4 are GREEN on the isolated production branch;
Task 3 track/checkpoint hard gate passes without scope expansion;
Task 5 four-host TPU correctness/model smoke passes;
Task 6 production serving A/B satisfies the decode TPOT gate;
all evidence is complete and read back by SHA;
independent whole-branch review has no unresolved Critical/Important issue;
a self-contained new Draft PR exists and remains Draft;
main, PR #1466, scheduler/state-pool schema and unrelated backends were not modified.

Post-#1505 production alignment amendment (2026-08-17)

PR #1498 was rebased onto SGL-JAX main after #1505 landed. This changes the production route names and startup selection contract, but does not rewrite the fixed-revision comparison or Phase 11 execution history above.

The native production route is now chunked_jax prefill plus the existing base-JAX decode implementation.
The optimized production route is fused_chunk_parallel prefill plus the same base-JAX decode implementation.
--gdn-prefill-impl is the only production selector. Its accepted values are chunked_jax and fused_chunk_parallel; the default is fused_chunk_parallel.
The selector is frozen when the backend is initialized. Unsupported fused configurations fail at startup and never silently fall back.
token_scan, reference, and tpu_inference_v3 remain valid labels only for the historical fixed-revision experiments that used those routes. They are not current production selector values.
Where the Phase 11 implementation plan conflicts with this amendment, this section is authoritative for naming, selector transport, default behavior, and the native route after #1505.

Phase 11 execution result (2026-07-30)

Tracking

Task: primatrix/projects#324
Current design RFC: primatrix/wiki#204
Historical chunkwise RFC: primatrix/wiki#201
Earlier implementation candidate: sgl-project/sglang-jax#1466
Selected production candidate: sgl-project/sglang-jax#1498

RFC-0033 / wiki #201 and code PR #1466 record the earlier chunkwise-KDA candidate. This RFC and PR #1498 record the route comparison and the selected production path. PR #1498 is rebased on SGL-JAX main after #1505 and therefore retains the native chunked_jax path. It does not include PR #1466's older chunkwise-KDA adapter.

Selected implementation

The final production interface exposes two prefill choices:

`--gdn-prefill-impl`	Prefill	Decode
`chunked_jax`	Native JAX chunk-parallel recurrence from #1505	Existing base-JAX decode
`fused_chunk_parallel`	Pallas fused Conv1D+GDN chunk-parallel kernel	Existing base-JAX decode

fused_chunk_parallel is the default. The choice is frozen during backend initialization. Unsupported fused configurations fail at startup rather than silently falling back. The pre-#1505 token_scan baseline and the historical tpu_inference_v3 experiment label are retained only in fixed-revision evidence.

Qwen3.5-397B-A17B BF16 serving A/B

Configuration:

model revision: 8472618112abcbd45acbcdc58436aff4233c23f7
TPU: v7x64, 4x4x4, 16 hosts
runtime: JAX/jaxlib 0.10.2, libtpu 0.0.43
parallelism: TP/DP/EP 128/8/128
five measured rounds per cell, 80 requests per round
exact 64-token output identity and zero request errors

Workload	Metric	Reference	TPU-Inference v3	Change
4096→1, c16	input throughput	9,485.999241 tok/s	40,821.238522 tok/s	+330.331455%
4096→1, c16	mean TTFT	6,515.492437 ms	1,513.621697 ms	-76.768883%
1024→100, c16	output throughput	275.949103 tok/s	398.225380 tok/s	+44.311170%
1024→100, c16	mean TPOT	36.643958 ms	36.005768 ms	-1.741596%

Verification and current gate

Numerical, recurrent-state, Conv-state, zero-length, track/checkpoint, slot-isolation, and prefill→decode contracts passed.
Qwen3.5-397B-A17B reference/optimized smoke produced identical 64-token outputs with finite values and zero errors.
PR #1498 Ready CI completed with 21 successful checks, 3 expected skips, and 0 failures.
Technical implementation and validation are complete. PR #1498 merged into main on 2026-08-18 as merge commit f1f5a513279b0eb5a1d6319239f5ec216107652c; #324 is closed as completed.
PR #1466 remains the historical earlier chunkwise-KDA candidate.

RFC-0034: SGL-JAX GDN 路径统一对比与移植决策 ​

概述 ​

背景与动机 ​

目标 ​

成功标准 ​

非目标 ​

方案 ​

1. 候选定义与版本冻结 ​

2. 分层 benchmark ​

3. 统一 workload ​

3.1 性能矩阵 ​

3.2 正确性边界矩阵 ​

4. 计时与运行规则 ​

5. Native benchmark 与统一 benchmark 的关系 ​

6. 备选方案 ​

7. 选择与移植 ​

7.1 候选资格 ​

7.2 选择规则 ​

7.3 移植规则 ​

8. 最终模型计划 ​

设计细节 ​

1. 统一接口 ​

2. Gate 与 state 语义 ​

3. 独立 reference ​

4. 防止假一致与假加速 ​

5. XProf ​

6. 目录与 artifact ​

7. 资源与进程边界 ​

测试计划 ​

Verification Group A：Harness 单元测试 ​

Verification Group B：正确性 ​

Verification Group C：统一性能 A/B ​

Verification Group D：Native benchmark ​

Verification Group E：移植后模型与 serving ​

影响范围 ​

实施计划 ​

停止条件 ​

风险与缓解 ​

Review 问题 ​

RFC-0034 Phase 0 Discovery Implementation Plan ​

Global Constraints ​

Task 1: 创建隔离 worktree 并冻结候选与模型 revision ​

Task 2: 生成固定源码索引和 SHA 证据 ​

Task 3: 审计真实调用链、数学语义和 adapter 边界 ​

Task 4: 盘点 native benchmark、测试和 profiling 能力 ​

Task 5: 从固定模型与实际 projection 推导 35B per-shard shape ​

Task 6: Phase 0 总审、报告与停止门禁 ​

Phase 6 运行时兼容补充：JAX 0.10.2 下显式对齐 FusedEPMoE v1 输入分片 ​

1. 触发条件与已验证现象 ​

2. 根因 ​

3. 批准的最小设计 ​

4. 明确不做 ​

5. TDD 与验证门禁 ​

6. 验收标准与停止条件 ​

7. 证据边界 ​

Phase 6 JAX 0.10.2 Compatibility Implementation Plan ​

Global Constraints ​

Task 1: 用 TDD 修复 FusedEPMoE v1 layer-to-kernel 分片边界 ​

Task 2: 在四-host JAX 0.10.2 环境关闭 S0 runtime gate ​

Phase 8 修订：Qwen3.5-397B BF16 / TPU v7 四层最小 A/B ​

修订状态与覆盖关系 ​

固定对比对象 ​

资源与拓扑 ​

四层最小性能矩阵 ​

正确性与 anti-fallback 门禁 ​

统计与产物 ​

Phase 8 Implementation Plan ​

Task 1：冻结 397B 配置、代码和 topology ​

Task 2：创建 Falcon v7 实验并通过 load/runtime smoke ​

Task 3：Kernel recurrence A/B ​

Task 4：完整 GDN layer A/B ​

Task 5：完整模型 offline A/B ​

Task 6：Serving A/B ​

Task 7：收口与决策 ​

完成条件 ​

Phase 9–10 收口记录（2026-07-29） ​

Phase 9：397B BF16 serving 路线选择 ​

Phase 10：397B BF16-only 矩阵 ​

下一门禁 ​

Phase 11：Selected composition 生产化（mentor 已批准） ​

RFC-0034: SGL-JAX GDN 路径统一对比与移植决策

概述

背景与动机

目标

成功标准

非目标

方案

1. 候选定义与版本冻结

2. 分层 benchmark

3. 统一 workload

3.1 性能矩阵

3.2 正确性边界矩阵

4. 计时与运行规则

5. Native benchmark 与统一 benchmark 的关系

6. 备选方案

7. 选择与移植

7.1 候选资格

7.2 选择规则

7.3 移植规则

8. 最终模型计划

设计细节

1. 统一接口

2. Gate 与 state 语义

3. 独立 reference

4. 防止假一致与假加速

5. XProf

6. 目录与 artifact

7. 资源与进程边界

测试计划

Verification Group A：Harness 单元测试

Verification Group B：正确性

Verification Group C：统一性能 A/B

Verification Group D：Native benchmark

Verification Group E：移植后模型与 serving

影响范围

实施计划

停止条件

风险与缓解

Review 问题

RFC-0034 Phase 0 Discovery Implementation Plan

Global Constraints

Task 1: 创建隔离 worktree 并冻结候选与模型 revision

Task 2: 生成固定源码索引和 SHA 证据

Task 3: 审计真实调用链、数学语义和 adapter 边界

Task 4: 盘点 native benchmark、测试和 profiling 能力

Task 5: 从固定模型与实际 projection 推导 35B per-shard shape

Task 6: Phase 0 总审、报告与停止门禁

Phase 6 运行时兼容补充：JAX 0.10.2 下显式对齐 FusedEPMoE v1 输入分片

1. 触发条件与已验证现象

2. 根因

3. 批准的最小设计

4. 明确不做

5. TDD 与验证门禁

6. 验收标准与停止条件

7. 证据边界

Phase 6 JAX 0.10.2 Compatibility Implementation Plan

Global Constraints

Task 1: 用 TDD 修复 FusedEPMoE v1 layer-to-kernel 分片边界

Task 2: 在四-host JAX 0.10.2 环境关闭 S0 runtime gate

Phase 8 修订：Qwen3.5-397B BF16 / TPU v7 四层最小 A/B

修订状态与覆盖关系

固定对比对象

资源与拓扑

四层最小性能矩阵

正确性与 anti-fallback 门禁

统计与产物

Phase 8 Implementation Plan

Task 1：冻结 397B 配置、代码和 topology

Task 2：创建 Falcon v7 实验并通过 load/runtime smoke

Task 3：Kernel recurrence A/B

Task 4：完整 GDN layer A/B

Task 5：完整模型 offline A/B

Task 6：Serving A/B

Task 7：收口与决策

完成条件

Phase 9–10 收口记录（2026-07-29）

Phase 9：397B BF16 serving 路线选择

Phase 10：397B BF16-only 矩阵

下一门禁

Phase 11：Selected composition 生产化（mentor 已批准）