Skip to content

feat(qwen4): share the Qwen3.5 GDN decode wrappers and prove their reuse - #1159

Merged
FeathBow merged 6 commits into
pegainfer-project:mainfrom
sysuTobo:feat/qwen4-gdn-reuse
Oct 20, 2026
Merged

FeathBow merged 6 commits into
pegainfer-project:mainfrom
sysuTobo:feat/qwen4-gdn-reuse

Conversation

@sysuTobo

@sysuTobo sysuTobo commented Oct 14, 2026 •

Copy link
Copy Markdown
Contributor

Summary

The Qwen3.8-Flash-Next (qwen4_exp) line has 36 Gated DeltaNet layers whose geometry
(16 key / 48 value heads at head dim 128, 4-tap depthwise conv) is identical to Qwen3.5's,
so its decode path reuses Qwen3.5's GDN and conv1d kernels instead of growing a second
implementation.

The FFI was already shared; the host-side wrapper — argument marshalling, shape
assertions, state-pointer handling — lived only inside pegainfer-qwen35, which is
entirely feature-gated. This PR moves that wrapper and the three kernel limits it asserts
into the kernels crate's existing Qwen3.5 ops module, beside the HD256 wrappers that
already live there, so both lines reach one owner through pegainfer_kernels::ops. It adds
no feature to pegainfer-core, no new crate and no second copy, and leaves the shipped
line's call sites and compiled kernel set unchanged.

Commits

  • cea0b149 — promote the wrappers, repoint Qwen3.5 at them, and add a real-weight
    component test behind a qwen4 device feature.
  • 1397a754 — review response: move the wrappers into pegainfer-kernels/src/ops/qwen35.rs
    (under the existing kernels qwen35 feature), drop the component test and the qwen4
    device feature with it, correct the docs, and check the now-public wrapper's qkv width
    and dt_bias/A_log lengths against the head counts.
  • 47cd8514 — add a GPU unit test for conv1d_decode_batch_into at the production channel
    width, assert the GDN wrapper's head-group invariant, and correct the qwen4 wording about
    the CPU-only constant mirror.

The final diff is 9 files, +339/−112.

File Change
pegainfer-kernels/src/ops/qwen35.rs the two decode wrappers and the three kernel limits; checks for qkv width, dt_bias/A_log lengths, and the key/value head-group invariant
pegainfer-kernels/src/ops.rs re-export them under the existing qwen35 feature
pegainfer-qwen35/src/recurrent.rs delete the two wrapper bodies; add a conv1d decode unit test; prefill wrappers and the vec path untouched
pegainfer-qwen35/src/ops.rs consume the wrappers from pegainfer_kernels::ops
pegainfer-qwen35/src/config/{mod,model}.rs drop three dead re-exports; re-export the constants from the kernels owner
pegainfer-qwen4/src/{lib,config}.rs point the docs at the kernels owner and at the CPU-only mirror
docs/models/qwen4/bring-up.md record what is pinned and what slice B still proves

Why the reuse is sound

  • State layout: [num_value_heads, key_dim, val_dim] f32 with the value dim contiguous
    (FLA [H, K, V]), addressed through a device array of per-request pointers; both lines
    allocate f32 state (this line's config asserts mamba_ssm_dtype: float32).
  • Geometry: 16 key / 48 value heads at 128/128 and a 4-tap conv is the frozen checkpoint's
    geometry and the kernels' baked-in limits (GDR_KEY_DIM/GDR_VAL_DIM,
    LINEAR_CONV_MAX_KERNEL_DIM). Head counts are runtime arguments, so 48 value heads needs
    no rebuild.
  • Tensor names: byte-identical down to the model.language_model.layers.N. prefix.

Why this owner

  • The wrappers mount the same csrc/qwen35/{gated_delta_rule,conv1d}.cu behind the same
    qwen35 feature as pegainfer-kernels/src/ops/qwen35.rs's HD256 wrappers.
  • pegainfer-qwen35 already imports that module directly; every kernel-owning model crate
    does the same (gemma4, dsv2-lite, glm52, k3, kimi-k2, qwen3), and feat(qwen4): score QSA token-indexer blocks in native CUDA #1160 puts qwen4's
    device ops in pegainfer-kernels/src/ops/qwen4.rs for the same reason.
  • pegainfer-core has no model feature today; promoting the wrapper there would have added
    the first one and a second owner. Moving into the existing module needs no change to
    pegainfer-kernels/{Cargo.toml,build.rs}: the feature exists and the csrc/qwen35 path
    filter already covers the sources.

What is proven, and what is not

Pinned without a GPU:

  • the tensor contract derives every GDN shape from the frozen config and validates all
    1658 headers (the_contract_is_the_frozen_checkpoint);
  • the config pins the two head dims and the conv width to CPU-only mirrors of the kernels'
    constants;
  • the frozen tensor names match Qwen3.5's.

The mirror is the one deliberate exception to the single-owner rule: qwen4's config must
run with no device, so it cannot import the kernels' constants. Nothing binds the mirror to
the kernels until slice B wires the device path, and this line does not serve until then.

Not proven here: the reuse on this line's activations. The component test from cea0b149
loaded layer 0's real tensors and drove the kernels, but it compared them with a
transcription of the same kernels from a zero state (a dropped feedback term would still
pass), and its relative-error reduction degraded to an absolute bound below 1 and skipped
NaN. It was removed rather than presented as a reuse proof. The first correct prefill and
decode through this line's own graph is #1103's M0 exit (slice B, #1127), where the GDN
block gets a reference to compare against (transformers' recurrence, not a kernel
transcription).

Verification

All commands run from the PR worktree; the CUDA ones inside docker/dev.sh
(pegainfer-dev:cu132, CUDA 12.9.1, Triton 3.7.1). GPU is a shared RTX 5060 Laptop
(sm_120, 8 GB) — neither support-matrix row.

  • cargo fmt --all -- --check — clean.
  • cargo test --release --locked -p pegainfer-qwen4 — 48 passed (default, CUDA-free).
  • cargo clippy --locked -p pegainfer-qwen35 -p pegainfer-kernels --features qwen35 --all-targets -- -D warnings — clean.
  • cargo test --release --locked -p pegainfer-qwen35 --features qwen35 --lib recurrent::tests::conv1d_decode_batch_matches_reference_steps — ok (10240 channels, 4 steps, 3 slots with distinct states; output within one bf16 ulp, state bit-exact).
  • cargo test --release --locked -p pegainfer-qwen35 --features qwen35 --lib recurrent::tests::gdr_decode_batch_matches_single_slot_reference — ok (16K/48V@128, batch 3).
  • CI on 47cd8514: all jobs pass, including Qwen3.5 CUDA Clippy (sm_80), CPU Clippy, CPU unit tests, DCO and formatting.

Not runnable here: qwen35's hf_golden_gate / e2e_scheduler need a Qwen3.5 checkpoint;
this host has the Qwen3.8-Flash-Next partial shards and no Qwen3.5 weights. The moved
paths are covered by the two unit tests above.

Refs #1103

The Qwen3.8-Flash-Next line's 36 Gated DeltaNet layers have Qwen3.5's
decode geometry (16 key / 48 value heads at head dim 128, 4-tap conv), so
they must drive Qwen3.5's GDN and conv1d kernels rather than a second
implementation. The raw FFI was already shared in pegainfer-kernels, but
the host-side wrapper — marshalling, shape assertions and state-pointer
handling — lived only inside the fully feature-gated pegainfer-qwen35,
which the new line cannot depend on.

Move that wrapper (and the three kernel limits it asserts) into
pegainfer-core::ops::gdn, behind a new qwen35 feature that forwards to the
kernels feature the shipped line already enabled, so the kernel build is
unchanged; Qwen3.5 consumes the promoted functions through its ops facade
with its call sites untouched.

Prove the reuse on layer 0 of the real checkpoint: a test behind the
crate's qwen4 feature loads the nine GDN tensors from the real shards,
holds their headers against the frozen contract, derives the geometry the
shapes imply, and drives both kernels on a GPU — the conv against the
reference formula and an exact state handoff, the GDN step against a host
transcription of the documented recurrence, each request's own state slot,
and the real A_log/dt_bias decaying a state by exactly the host factor.

No reference oracle for this checkpoint exists locally, so the test states
plainly what it is: a real-weight invariant check, not a parity check.

Signed-off-by: bohaotu <bohaotu@pony.ai>
@github-actions github-actions Bot added qwen35 Qwen3.5 model crate (pegainfer-qwen35) qwen4 Qwen4-architecture model crate (pegainfer-qwen4) enhancement New feature or request labels Oct 14, 2026
@sysuTobo

Copy link
Copy Markdown
Contributor Author

Code review

审计对象:本 PR 改动集(da39fbc8..cea0b149,14 文件,含两个新文件);基线 da39fbc8。
正向路径结论:共享包装器与原线行为等价、新线真权重复用已跑通,正向路径全部符合预期;发现 8 条,其中 4 条已修(2 条 M、2 条 L),4 条为已确认无缺陷的判断项/已知边界(明确留给谁)。

ID 级别 一句话 证据(文件:行)
R1 M 状态衰减断言每头只读 1 个元素,j 切片只衰减一半这类 bug 逃得过 pegainfer-qwen4/tests/gdn_decode_reuse.rs(原 device_state[head * key_dim * val_dim],单元素)→ 现 :814-824(整头逐元素,断言 :821)
R2 M 几何推导用整除截断:qkv 宽度不自洽时会静默取整,而下游断言恰好仍成立 pegainfer-qwen4/tests/gdn_decode_reuse.rs:99-121,断言在 :115-121(原 (qkv[0] - value_heads * head_dim) / (2 * head_dim))
R3 L 文档 TL;DR 的 "All of it is CPU-only" 在本 PR 后不再为真 docs/models/qwen4/bring-up.md:3
R4 L 新模块 doc 写"两条线都 drive them",qwen4 目前只有测试在调(无执行器) pegainfer-core/src/ops/gdn.rs:5-7
R5 L 同一份 csrc/qwen35 源集的包装器现有两个家(attention 在 kernels::ops::qwen35,GDN 解码在 core::ops::gdn) pegainfer-kernels/src/ops/qwen35.rs:1-2 vs pegainfer-core/src/ops/gdn.rs:1-11
R6 L 缺权重时测试打印 SKIP 后仍报 ok(仓库既有约定;CI 不跑 GPU) pegainfer-qwen4/tests/gdn_decode_reuse.rs:250-285;同形先例 pegainfer-qwen35/tests/common/model_fixture.rs:85-88
R7 L 局部下载无 config.json,测试无法核对 revision,只能按 name+dtype+shape 比 pegainfer-qwen4/tests/gdn_decode_reuse.rs:255-285(只查 index.json 与张量名,无 config.json 可核)
R8 L --lib recurrent 在 gdn_chunkwise_* 上 SIGSEGV;stash 后基线同样复现 ⇒ 与本改动无关,本机 sm_120 chunk-wise GDN prefill 路径 pegainfer-qwen35/src/recurrent.rs:57-92(该路径非本 PR 移动的两个包装器)

覆盖清单(14/14 已认领):pegainfer-core/src/ops/gdn.rs(新,两个包装器+3 常量)、pegainfer-core/src/ops.rs+Cargo.toml(re-export、qwen35 feature)、pegainfer-qwen35/src/recurrent.rs(删两个函数体,prefill 包装器与 #[cfg(test)] vec 路径未动)、pegainfer-qwen35/src/ops.rs(2 行 use 改指向)、pegainfer-qwen35/src/config/{mod,model}.rs(3 个死 re-export 删除、3 常量改 re-export)、pegainfer-qwen35/Cargo.toml(feature 追加转发)、pegainfer-qwen4/{Cargo.toml,src/config.rs,src/lib.rs}(qwen4 feature、镜像常量注释、crate doc)、pegainfer-qwen4/tests/gdn_decode_reuse.rs(新,2 个测试)、docs/models/qwen4/bring-up.md(决策+边界)、Cargo.lock(qwen4 依赖列表 4 行)。

三条数据流逐条模拟(全部 ✓ 符合预期):

  1. 真权重 → 复用内核:shard model-00001-of-00131 → load_real_tensor(tests/gdn_decode_reuse.rs:301-338)→ host 投影(真实 in_proj_{qkv,a,b},:549-558)→ conv1d_decode_batch_into(tests/…:639 → pegainfer-core/src/ops/gdn.rs:114)→ gated_delta_rule_decode_batch_into(tests/…:704 → pegainfer-core/src/ops/gdn.rs:51)→ 输出/状态 D2H。代入真实规模:qkv 10240 = 2×16×128 + 48×128,state 3 slot × 48×128×128 f32 = 9.4 MB;四步 conv 逐步与 host 参考逐元素比,GDN 与 host 递推逐头逐元素比。
  2. 生产解码路径:pegainfer-qwen35/src/batch_decode.rs:832,840 → pegainfer-qwen35/src/ops.rs → pegainfer_core::ops::gdn → pegainfer_kernels::ffi::…。函数体逐字未改(断言顺序、FFI 实参、stream 均同),唯一差异是跨 crate 调用;qwen35 = ["pegainfer-kernels/qwen35", "pegainfer-core/qwen35"] 只追加转发,编译的 .cu 集合/nvcc 参数/Triton AOT 调用不变 ⇒ 数值路径逐位不变。
  3. 常量单一真相源:pegainfer-core/src/ops/gdn.rs:28-36 ←re-export— pegainfer-qwen35/src/config/model.rs:61-67;qwen4 侧 src/config.rs 因 CPU-only 保留镜像,由 MODEL_LINE.probe(frozen_config) + 形状推导等式钉住(tests/gdn_decode_reuse.rs:145-203),镜像漂移会成为测试失败而非静默误读。

改动五问:①层——本体落在 FFI 与同一 feature 门所在的层,pegainfer-core::ops 作为 model crate 门面(与本仓库 pegainfer-qwen35/src/ops.rs:2-18 既有用法一致);②路径——qwen35 feature 是 qwen35 线唯一编译路径,qwen4 侧由测试 target 的 required-features 覆盖,生产 server 的 qwen4 feature 未接(有意保持 CUDA-free);③作用域——pegainfer-core 追加 feature 为纯增量,唯一可能被影响的调用方是 qwen35 线,已逐一列出并跑测;④公共形状——无带默认值的新字段;删掉的 3 个 re-export 经全 target 编译确认无其他调用点;⑤测试价值——删掉包装器里的状态指针串表会让 R1/行隔离断言失败,删掉 b/a 传参会让参考比对失败(非空转);the_frozen_geometry_is_the_shared_kernel_geometry 无权重无设备也总是运行。

修改措施

  • R1 已修:衰减断言改为整头逐元素(tests/gdn_decode_reuse.rs:808-825,worst 取该头所有元素的最大相对偏差),仍报 1e-5 线;重跑测试通过。
  • R2 已修:先断言 key_width % (2 * head_dim) == 0 再整除(tests/gdn_decode_reuse.rs:115-121),形状不自洽即响亮失败。
  • R3 已修:TL;DR 改为 "The library is CPU-only; the one device path so far is the GDN decode reuse test…"(docs/models/qwen4/bring-up.md:3),并把 "Nothing runs" 收敛为 "Nothing serves"。
  • R4 已修:模块 doc 改为 "The Qwen3.5 line calls them from its batched decode, and the Qwen3.8-Flash-Next line calls them from its reuse test today and its executor next"(pegainfer-core/src/ops/gdn.rs:5-8)。
  • R5 明确留给装配切片:本 PR 的决策就是"上提 pegainfer-core::ops"(PR 正文有备选方案对比);把 pegainfer-kernels::ops::qwen35 的 attention 包装也上提是一次独立重构,届时两处应合并为一家。当前不一致是有意识的、不制造重复实现。
  • R6 明确留给 CI 决策:缺权重 skip 是仓库既有约定(model_fixture.rs:85-88 同形),且 CI 不跑 GPU;PR 正文写明运行命令与前置。若将来 CI 真跑 GPU,应改为"缺权重即失败"。
  • R7 明确留给下载配方:局部下载无 config.json,只能按 name+dtype+shape 比对;已在 PR 正文"未验证"一节写明(revision 未被本测试核对)。
  • R8 已就地更正归属:先怀疑本改动,git stash -u 后在同基线复跑仍 SIGSEGV ⇒ 判定为基线既有问题(本机 sm_120 chunk-wise GDN prefill),写进 PR 正文"Regression status",不在本 PR 修复。该路径不在本 PR 移动的两个包装器上;本 PR 实际触及的两个包装器由 recurrent::tests::gdr_decode_batch_matches_single_slot_reference 覆盖并通过。

@sysuTobo

Copy link
Copy Markdown
Contributor Author

Review receipt (independent audit + reproduction)

Audited the full diff (da39fbc8..cea0b149): the promotion is a verbatim body move (same asserts in the same order, same FFI calls and stream), the qwen35 feature chain (pegainfer-core/qwen35 → pegainfer-kernels/qwen35) compiles exactly the kernel set Qwen3.5 already built, the dead re-exports are removed, and the qwen4 mirror constants are bound to the shared ones by the CPU test (the_frozen_geometry_is_the_shared_kernel_geometry reads the frozen shapes, derives 16K/48V@128/128 with a 4-tap conv, and asserts equality with the shared constants and the frozen config). The blast-radius section on the shipped Qwen3.5 line checks out against the diff.

No changes required. Reproduced on this box (RTX 5060, sm_120, shared):

  • cargo test --release --locked -p pegainfer-qwen4 (default, CUDA-free) — ok
  • gdn_decode_reuse with the real shards: 2/2 pass — 9 real tensors against the frozen contract, gates decay in (8.408e-44, 0.999992), conv1d matches the reference formula with an exact bf16 state handoff, GDN matches the host recurrence, state slots isolated (bit-identical slot 0 across runs, distinct slots stay distinct)
  • promoted-wrapper regression on the shipped line: recurrent::tests::gdr_decode_batch_matches_single_slot_reference and conv1d_prefill_handoff_matches_single_prefill both pass
  • cargo fmt --all --check clean

The documented non-goals (no HF oracle, synthetic activations, block composition and prefill path open, sm_89 row untested) are correctly scoped in the PR and the bring-up doc.

@sysuTobo
sysuTobo marked this pull request as ready for review October 17, 2026 10:03

@FeathBow FeathBow left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for this. The move itself is clean. The two bodies are byte for byte the code that left recurrent.rs, and the tensor names do match Qwen3.5's. On one sm_89 GPU, fmt and the Qwen3.5 clippy pass, pegainfer-qwen4 passes its 48 default tests, and the CPU half of the reuse test passes. Requesting changes on where the wrappers land and on the new test, which the description presents as a proof of reuse on the real weights but which mostly re-proves what the tree already pins.

The Qwen3.5 wrappers that need the kernels' qwen35 feature already live in pegainfer-kernels/src/ops/qwen35.rs, behind that same feature and over the same csrc/qwen35 sources, and pegainfer-qwen35 imports them from pegainfer_kernels::ops directly. pegainfer-core has no model feature today and re-exports kernels items one by one, so this adds a second owner, which the self review's R5 already notes. Moving the two functions into the kernels module needs no change to the kernels Cargo.toml or build.rs, so the reason given for deferring the kernels side does not apply.

Check in gdn_decode_reuse.rs Status
GDN tensor shapes and dtypes match the frozen contract already held by the_contract_is_the_frozen_checkpoint
Frozen config passes the geometry rules already held by the validator's head dim and conv checks
Gate in (0, 1) on the real A_log and dt_bias an identity, computed on the host
conv1d and GDN against a host transcription copies the kernel and starts from a zero state (inline)
Per request state slots already held by gdr_decode_batch_matches_single_slot_reference, which runs this exact 16/48/128 geometry at batch 3

What is new is the check of the mirror against the shared constants, and the file only builds with the qwen4 feature, so nothing runs it unless someone builds that feature by hand. Under that feature cargo clippy -p pegainfer-qwen4 --features qwen4 --all-targets -- -D warnings fails with three errors in the test (chunks_exact_to_as_chunks, needless_range_loop, float_cmp), while the description reports clippy for pegainfer-qwen35 only.

Before merge we would ask for three things.

  1. Move the two wrappers and the three constants into pegainfer-kernels/src/ops/qwen35.rs and drop the new pegainfer-core feature (inline).
  2. Drop gdn_decode_reuse.rs, and with it the new qwen4 feature, which otherwise pulls the whole Qwen3.5 kernel set (HD256 attention and the Triton AOT prefill included) into this crate for one test. The roadmap in #1103 covers the reused GDN through the model graph fixtures and makes the first correct prefill and decode its M0 exit, so that is where the reuse gets proven. If a component check should land now, run several decode steps from a live state against the Transformers recurrence (inline).
  3. Correct the description and the bring-up note where they call this a proof of reuse on the real weights (inline).

Smaller points.

  • The two wrappers go from crate private to public here, so the qkv width and the dt_bias and a_log lengths are worth checking against the head counts like the other dimensions already are.
  • If the test stays, an explicitly set but unusable PEGAINFER_TEST_MODEL_PATH should fail instead of skipping, and in_proj_z, norm.weight and out_proj (about 60 MiB) are loaded but never used.
  • The two audit comments on this thread come from the author's account, so we read their conclusions as claims and checked them here.

Comment thread pegainfer-core/Cargo.toml Outdated
Comment on lines +33 to +36
# Gated DeltaNet decode ops. Their kernels are Qwen3.5's (`csrc/qwen35/…`), which
# both the Qwen3.5 and the Qwen3.8-Flash-Next lines build; the flag names the
# kernel set, not the model line.
qwen35 = ["pegainfer-kernels/qwen35"]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would be the first model feature on pegainfer-core. The qwen35 gated wrappers otherwise sit in pegainfer-kernels/src/ops/qwen35.rs behind this same kernels feature, so the two functions can go there with no edit to the kernels Cargo.toml or build.rs, and this feature and its forwarding in pegainfer-qwen35/Cargo.toml go away.

Comment on lines +727 to +758
for row in 0..batch {
let mut reference_state = vec![0.0f32; head_state_len];
let expected = reference_gdn_step(
&gdn_input[row * qkv_dim..(row + 1) * qkv_dim],
&b_rows[row],
&a_rows[row],
&dt_bias,
&a_log,
&mut reference_state,
key_heads,
value_heads,
key_dim,
val_dim,
);
let got = &out_host[row * out_dim..(row + 1) * out_dim];
let diff = max_rel_diff(
&got.iter().map(|v| v.to_f32()).collect::<Vec<_>>(),
&expected.iter().map(|v| v.to_f32()).collect::<Vec<_>>(),
);
ensure!(
diff <= 8e-3,
"gdn decode row {row}: the reused kernel differs from the reference recurrence by a \
relative {diff}"
);
let device_state: Vec<f32> = ctx.stream.clone_dtoh(&gdn_slots[row])?;
ctx.sync()?;
let state_diff = max_rel_diff(&device_state, &reference_state);
ensure!(
state_diff <= 1e-4,
"gdn decode row {row}: the updated state differs from the reference by a relative \
{state_diff}"
);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Both the device state and reference_state start at zero here, so k · S is zero and a kernel that dropped the feedback subtraction would still match. The state comparison also has no finite check, and fold(0.0, f32::max) skips NaN, so an all NaN state reports zero error. A component check would need several steps from a nonzero state, with both arrays checked finite before the reduction.

Comment on lines +395 to +433
fn reference_gdn_step(
qkv: &[bf16],
b_proj: &[bf16],
a_proj: &[bf16],
dt_bias: &[f32],
a_log: &[f32],
state: &mut [f32],
num_key_heads: usize,
num_value_heads: usize,
key_dim: usize,
val_dim: usize,
) -> Vec<bf16> {
let q_dim = key_dim * num_key_heads;
let out_dim = val_dim * num_value_heads;
let mut out = vec![bf16::ZERO; out_dim];
let head_stride = key_dim * val_dim;
for v_head in 0..num_value_heads {
let k_head = v_head * num_key_heads / num_value_heads;
let q: Vec<f32> = (0..key_dim)
.map(|i| qkv[k_head * key_dim + i].to_f32())
.collect();
let k: Vec<f32> = (0..key_dim)
.map(|i| qkv[q_dim + k_head * key_dim + i].to_f32())
.collect();
let v: Vec<f32> = (0..val_dim)
.map(|i| qkv[2 * q_dim + v_head * val_dim + i].to_f32())
.collect();

let q_sq: f32 = q.iter().map(|value| value * value).sum();
let k_sq: f32 = k.iter().map(|value| value * value).sum();
let q_scale = 1.0 / (q_sq + 1e-12).sqrt() / (key_dim as f32).sqrt();
let k_scale = 1.0 / (k_sq + 1e-12).sqrt();
let q: Vec<f32> = q.iter().map(|value| value * q_scale).collect();
let k: Vec<f32> = k.iter().map(|value| value * k_scale).collect();

let x = a_proj[v_head].to_f32() + dt_bias[v_head];
let softplus = if x > 20.0 { x } else { (1.0 + x.exp()).ln() };
let decay = (-a_log[v_head].exp() * softplus).exp();
let beta = 1.0 / (1.0 + (-b_proj[v_head].to_f32()).exp());

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This transcribes the kernel, including its 1e-12 L2 epsilon, f32 normalization and f32 sigmoid, while the Transformers 5.16.1 reference uses l2norm(..., eps=1e-6) on the bf16 q and k before promoting them and takes b.sigmoid() in bf16. Agreement with it shows the kernel matches itself, which the recurrent.rs tests already cover at this geometry.

Comment on lines +572 to +592
// 3. The gate path on the real scalars: the reused kernel computes
// `-exp(A_log) * softplus(a + dt_bias)`; on this checkpoint that must be a
// decay in (0, 1) for every value head, or the state would grow.
let mut decays = Vec::with_capacity(value_heads);
for head in 0..value_heads {
for row in 0..batch {
let x = a_rows[row][head].to_f32() + dt_bias[head];
let softplus = if x > 20.0 { x } else { (1.0 + x.exp()).ln() };
decays.push((-a_log[head].exp() * softplus).exp());
}
}
ensure!(
decays.iter().all(|d| d.is_finite() && *d > 0.0 && *d < 1.0),
"the real A_log/dt_bias must put every gate in (0, 1), got {:?}",
decays.iter().fold(f32::INFINITY, |a, b| a.min(*b))
);
println!(
"{TEST}: real gates decay in ({:.3e}, {:.6})",
decays.iter().fold(f32::INFINITY, |a, b| a.min(*b)),
decays.iter().fold(0.0f32, |a, b| a.max(*b))
);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This runs on the host, and exp(-exp(A_log) * softplus(x)) lies in (0, 1] for any finite input because both factors in the exponent are positive. It can only fail at the ends through underflow or rounding, so it says nothing about the kernel or about the regime the description mentions.

Comment on lines +469 to +478
/// The largest deviation relative to the reference's own magnitude, so a check is
/// as tight at 1e-3 as at 10. One bf16 rounding is 2^-9 relative; the reference and
/// the kernel differ only by summation order and by `exp`/`ln` implementations.
fn max_rel_diff(a: &[f32], b: &[f32]) -> f32 {
assert_eq!(a.len(), b.len());
a.iter()
.zip(b)
.map(|(x, y)| (x - y).abs() / y.abs().max(1.0))
.fold(0.0f32, f32::max)
}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

.max(1.0) makes this an absolute bound for anything below magnitude 1, so it is not as tight at 1e-3 as at 10. An expected 0.001 against an output of 0 reports 0.001 and passes both the conv 5e-3 and the GDN 8e-3 bounds.

Comment thread pegainfer-qwen4/src/config.rs Outdated
Comment on lines +20 to +27
/// The GDN and full-attention kernel limits: the first three are owned by the
/// shared Gated DeltaNet decode wrappers in `pegainfer_core::ops::gdn` (which this
/// line drives with the same numbers as Qwen3.5), the last by the HD256
/// full-attention kernels. This crate's config layer runs with no device — a probe
/// must not need CUDA — so the values are mirrored here rather than imported.
/// `tests/gdn_decode_reuse.rs` binds the frozen checkpoint's geometry to the
/// shared constants, so a drift between this mirror and the kernels is a test
/// failure rather than a silent mis-read.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The test that binds this mirror only builds with the qwen4 feature, so a drift stays silent until someone builds that feature by hand, which is weaker than the test failure this comment promises.

Comment thread docs/models/qwen4/bring-up.md Outdated
Comment on lines +137 to +143
The reuse is proven **on the real weights**. `pegainfer-qwen4/tests/gdn_decode_reuse.rs`
(behind this crate's `qwen4` feature, pointed at a checkpoint with
`PEGAINFER_TEST_MODEL_PATH`) loads layer 0's nine GDN tensors from the real shards, holds
their headers against `tests/frozen_tensors.json`, and drives both kernels on a GPU: the
conv against the reference formula plus an exact state handoff, the GDN step against a
host transcription of the same recurrence, each request's own state slot, and the real
`A_log`/`dt_bias` decaying a state by exactly the host-computed factor.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The kernels already run at 16 key and 48 value heads in gdr_decode_batch_matches_single_slot_reference, and the checks here compare them with a transcription of themselves, so this is not a proof of reuse on the real weights. Please describe it as what it is, and leave the exit criteria row open until slice B's GDN layer has a reference to compare against.

The two decode wrappers and the three kernel limits belong beside the
other Qwen3.5 wrappers in pegainfer-kernels/src/ops/qwen35.rs, behind
the kernels' qwen35 feature. That removes the first model feature from
pegainfer-core and keeps Qwen3.5 on the owner it already imports for
HD256; qwen4 will import the same owner when slice B wires its executor.

Drop the real-weight reuse test and the qwen4 device feature with it.
The tensor contract and config already pin the geometry, names, dtypes
and 1-D lengths, while the component test compared the kernels with a
transcription of themselves from a zero state. The reuse gets a
reference check in slice B's model-graph fixture, which is pegainfer-project#1103's M0
exit.

Also check the now-public wrapper's qkv width and dt_bias/A_log lengths
against the head counts.

Refs pegainfer-project#1103

Signed-off-by: bohaotu <bohaotu@pony.ai>
@sysuTobo

Copy link
Copy Markdown
Contributor Author

改动内容:

  1. 归属移到 kernels:两个 wrapper 和三个常量并入 pegainfer-kernels/src/ops/qwen35.rs;pegainfer-qwen35 改为从 pegainfer_kernels::ops 导入。删除了 pegainfer-core 的 qwen35 feature、ops/gdn.rs 及相关转发。
  2. 删除组件测试与 qwen4 device feature。
  3. 新增:gated_delta_rule_decode_batch_into 校验 qkv 宽度和 dt_bias/a_log 长度与 head counts 一致。
  4. 文档改为"wrapper 已共享、geometry 由契约钉住、语义复用留给 slice B 图级 fixture";退出条件改回 open;lib.rs/config.rs 的旧说法一并清理。

验证结果:

  • cargo fmt --all --check 通过
  • cargo test --release -p pegainfer-qwen4:48 passed
  • 容器内 cargo clippy -p pegainfer-qwen35 -p pegainfer-kernels --features qwen35 --all-targets -- -D warnings:通过
  • GPU 测试 recurrent::tests::gdr_decode_batch_matches_single_slot_reference:通过
  • PR 全部 CI job:通过

@sysuTobo
sysuTobo marked this pull request as draft October 17, 2026 17:07
@sysuTobo
sysuTobo marked this pull request as ready for review October 17, 2026 17:18
@sysuTobo
sysuTobo requested a review from FeathBow October 17, 2026 17:18
Add a GPU unit test for `conv1d_decode_batch_into` at the production
channel width (10240), four decode steps, and three slots with distinct
initial states. It compares the output with the kernel's documented
recurrence within one bf16 ulp and the state hand-off bit for bit, so
the shared wrapper keeps a direct guard after the qwen4 component test
was removed.

Also assert the GDN wrapper's head-group invariant (positive key heads,
value heads a whole multiple of them) before the kernel's integer head
mapping runs, and correct the qwen4 notes: the wrappers will run once
slice B wires the text graph, and this line's config pins CPU-only
mirrors rather than sharing the kernels' definition.

Refs pegainfer-project#1103

Signed-off-by: bohaotu <bohaotu@pony.ai>
@FeathBow

Copy link
Copy Markdown
Collaborator

@codex review pelase.

@chatgpt-codex-connector

chatgpt-codex-connector Bot commented Oct 17, 2026 •

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ✅ Completed 2026-10-17T21:15:33.552514Z 47cd851 Manual request
ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 47cd85144c

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment on lines +19 to +20
pub(crate) use pegainfer_kernels::ops::conv1d_decode_batch_into;
pub(crate) use pegainfer_kernels::ops::gated_delta_rule_decode_batch_into;

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Run the production Qwen3.5 decode gates

Before redirecting the shipped Qwen3.5 decode path through these relocated wrappers, run the production hf_golden_gate/e2e_scheduler: the supplied verification explicitly leaves both unrun, while the added GPU unit tests exercise isolated conv/GDR calls rather than real weights, scheduler integration, or captured decode. The repository requires production E2E evidence for feature and fix changes, so this relocation is not yet verified to preserve serving behavior.

AGENTS.md reference: AGENTS.md:L126-L126

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks — Codex is right that the relocation had not been exercised on a real checkpoint. I ran both production gates on the L20 (sm_89) at head 47cd8514, with Qwen3.5-0.8B :

image

Environment: L20 46 GB, CUDA 12.9.86, Triton 3.7.1, FlashInfer 4d75a33f / CCCL 16bd510c, PEGAINFER_CUDA_SM=89, CUDA_VISIBLE_DEVICES=0.

The relocated bodies are unchanged apart from the new shape and head-group asserts, so this is the production regression loop on top of the byte-level diff. The 48-value-head geometry remains covered by gdr_decode_batch_matches_single_slot_reference (16K/48V@128, batch 3); the 0.8B E2E exercises the same wrapper, scheduler and graph path at 16 value heads.

Commands (run inside the L20 pod):

export PEGAINFER_TEST_MODEL_PATH=/home/pony-trainer/models/Qwen3.5-0.8B
export PEGAINFER_TEST_MODEL_REVISION=2fc06364715b967f1860aea9cf38778875588b17
export PEGAINFER_CUDA_SM=89 PEGAINFER_TRITON_PYTHON=/usr/bin/python3.10
export CUDA_VISIBLE_DEVICES=0

cargo test --release --locked -p pegainfer-qwen35 --features qwen35 \                                           
  --test hf_golden_gate -- --exact --nocapture --test-threads 1 \
  pega_logprobs_match_hf_golden_within_qwen35_tolerance

cargo test --release --locked -p pegainfer-qwen35 --features qwen35 \
  --test hf_golden_gate -- --exact --nocapture --test-threads 1 \
  pega_logprobs_match_hf_long_golden_within_qwen35_tolerance

cargo test --release --locked -p pegainfer-qwen35 --features qwen35 \
  --test e2e_scheduler -- --exact --nocapture --test-threads 1 \
  test_e2e_qwen35_scheduler

@github-actions github-actions Bot added the needs-rebase Conflicts with main, needs a rebase label Oct 19, 2026
@github-actions

Copy link
Copy Markdown

This pull request has merge conflicts with main. Please rebase it so it can be reviewed and merged.

Signed-off-by: Feathbow <feathbow@gmail.com>
… it to one bf16 ulp

Signed-off-by: Feathbow <feathbow@gmail.com>
…s as future work

Signed-off-by: Feathbow <feathbow@gmail.com>
@github-actions

Copy link
Copy Markdown

Commit attribution notice

This pull request contains commits attributed to an account other than the PR author @sysuTobo without a verified signature from that account:

If you are mentioned above and did not take part in this pull request, please say so here. Attributing a commit to someone else normally requires that person to be involved; maintainers should treat unacknowledged attribution as a review blocker.

This check never fails the build.

@FeathBow

Copy link
Copy Markdown
Collaborator

We pushed three commits to this branch to close the review points, a merge of main, a test fix and a docs fix.

  • Merged main. perf(qwen35): fuse the linear-attention decode projections into two GEMMs #1073 changed the two GDN decode wrappers this PR moves. They now read a band of the fused qkvz and ba projections and pass the slot stride of the whole tensor to the kernels. The merge carries that into the moved copies in pegainfer-kernels/src/ops/qwen35.rs and keeps your qkv, dt_bias, A_log and head group checks. The old copies are gone from recurrent.rs, and the band tests from perf(qwen35): fuse the linear-attention decode projections into two GEMMs #1073 now call the shared wrappers.
  • conv1d_decode_batch_matches_reference_steps now fails on any output that is not finite and allows at most one bf16 ulp. The loop runs six decode steps (kernel_size + 2).
  • config/model.rs and the config paragraph of bring-up.md now say the line will drive the shared wrappers once slice B gives it a device path. Until then nothing in the Qwen3.8-Flash-Next line runs them.
  • The GDN section of bring-up.md no longer retells the removed test. It now says where the shared wrappers live, that src/config.rs keeps a CPU-only mirror of the three limits, and what slice B still has to verify. The file now keeps each paragraph on one line, so most of its diff is line joins with the wording unchanged.

We checked the test by breaking the conv decode kernel on purpose.

Kernel change Test at the previous head Test now
none passes passes, output equal to the reference
output written as NaN passes fails at step 0, output not finite
2^-7 added to the output passes fails at step 0, output 30463 bf16 ulps off

Verification at the new head on two sm_90 GPUs with CUDA 13.1.

Check Result
clippy with -D warnings for the Qwen3.5 row, the Qwen3 row, pegainfer-kernels --features qwen35 and pegainfer-qwen4 clean
pegainfer-qwen4 tests 48 passed
pegainfer-qwen35 lib tests 112 passed
Qwen3.5-4B hf_golden_gate, short and long fixtures passed
Qwen3.5-4B e2e_scheduler passed
Qwen3.5-4B TP2 hf_golden_gate (eager, graph, long) 3 passed
Qwen3.5-4B TP2 e2e_scheduler (eager, graph) and serving_tp2 3 passed

@github-actions github-actions Bot removed the needs-rebase Conflicts with main, needs a rebase label Oct 20, 2026
@FeathBow
FeathBow merged commit cbe68f5 into pegainfer-project:main Oct 20, 2026
14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request qwen4 Qwen4-architecture model crate (pegainfer-qwen4) qwen35 Qwen3.5 model crate (pegainfer-qwen35)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants