17 KiB

Raw Blame History

Qwen3.5-MoE Model Implementation

Pantheon implementation of the Nex-N2-mini architecture: a 40-layer hybrid Gated-Delta-Net SSM / sparse GQA attention stack with 256-expert MoE on every layer.

Source files:

src/models/qwen35moe.c — loader, config parser, GGUF weight mapping
src/models/qwen35moe_internal.h — config/block structs, dual-state types, cross-file kernel declarations
src/models/qwen35moe.h — public API, PN_QWEN35MOE_MAX_SEQ_CAP
src/models/qwen35moe_forward.c — CPU forward, dual-state lifecycle, KV ring init
src/models/qwen35moe_ssm.c — Gated-Delta-Net SSM block kernel
src/models/qwen35moe_forward_cuda.c — CUDA forward: 2-die split, GPU weight residency, decode slots, async step machine (7 348 lines)

Architecture

Dims (from `pn_qwen35moe_config`, `qwen35moe_internal.h:36–60`)

Field	Value	GGUF key
`n_layer`	40	`qwen35moe.block_count`
`full_attention_interval`	4	`qwen35moe.full_attention_interval`
`n_embd`	2048	`qwen35moe.embedding_length`
`n_head`	16	`qwen35moe.attention.head_count`
`n_head_kv`	2	`qwen35moe.attention.head_count_kv`
`head_dim`	256	`qwen35moe.attention.key_length`
`rope_dim`	64	`qwen35moe.rope.dimension_count` (partial NEOX)
`n_expert`	256	`qwen35moe.expert_count`
`n_expert_used`	8	`qwen35moe.expert_used_count`
`expert_ff`	512	`qwen35moe.expert_feed_forward_length`
`shared_ff`	512	`qwen35moe.expert_shared_feed_forward_length`
`ssm_state`	128	`qwen35moe.ssm.state_size`
`ssm_n_group`	16	`qwen35moe.ssm.group_count` (K-heads, GQA key/query)
`ssm_dt_rank`	32	`qwen35moe.ssm.time_step_rank` (V-heads)
`ssm_d_inner`	4096	`qwen35moe.ssm.inner_size`
`ssm_conv`	4	`qwen35moe.ssm.conv_kernel`
`conv_dim`	8192	derived: read from `blk.0.attn_qkv.weight` outermost dim (`qwen35moe.c:163–167`)

Layer layout

Layer type is set by qwen35moe_is_attn_layer() (qwen35moe.c:216–220):

return ((i + 1) % interval) == 0;   /* interval default 4 */

With full_attention_interval=4 and 40 layers, this yields:

Layers 3, 7, 11, 15, 19, 23, 27, 31, 35, 39 → GQA full-attention (10 layers, n_attn=10)
All other 30 layers → Gated-Delta-Net SSM (n_ssm=30)

block_type[i] and block_index[i] in pn_qwen35moe_model (qwen35moe.c:59–60) map any global layer index to its type and per-type slot. These are the only two arrays that must be consulted to dispatch a layer.

Per-layer block structure

Every layer (SSM or attention) shares the same outer shell:

attn_norm → (SSM block | GQA attn block) → +residual
         → post_attention_norm → shared-MoE FFN → +residual

(qwen35moe_forward.c:316–369)

MoE routing

Every layer carries an identical pn_qwen35moe_moe block (qwen35moe_internal.h:75–84):

Router: ffn_gate_inp [n_expert, n_embd] → logits → softmax → top-8 selection → norm_topk_prob renormalization → scaled by expert_weights_scale (default 1.0; key absent in current GGUFs, qwen35moe.c:141–145).
Routed experts: 256 experts, zero-copy sliced from merged 3-D GGUF tensors blk.N.ffn_gate_exps.weight / ffn_up_exps.weight / ffn_down_exps.weight (qwen35moe.c:246–293). Each expert is a SwiGLU: silu(gate·x) * (up·x) → down. moe_out = Σ weight_e * SwiGLU_e(x).
Shared expert (always on): separate ffn_gate_shexp / ffn_up_shexp / ffn_down_shexp + a scalar sigmoid gate ffn_gate_inp_shexp [n_embd]. shexp_out = silu(gate·x) * (up·x) → down, then scaled by sigmoid(gate_inp_shexp · x).
Combine: out = moe_out + shexp_out. Residual add is the caller's responsibility.

(qwen35moe_internal.h:278–308)

GQA attention block (`qwen35moe_internal.h:160–196`)

Op order for attention layers:

attn_q [2*head_dim*n_head, n_embd] → split: first head_dim per head is Q, second is gate z.
QK-RMSNorm on Q (per head via attn_q_norm [head_dim]) and on K (via attn_k_norm [head_dim]), pre-RoPE.
attn_k [head_dim*n_head_kv, n_embd] and attn_v [head_dim*n_head_kv, n_embd].
Partial NEOX RoPE on the first rope_dim=64 dims of every Q and K head. Multi-section RoPE is parsed (rope_sections[4], qwen35moe.c:111–130) but collapses to standard NEOX for text.
GQA: 16 Q heads broadcast across 2 KV heads (each KV head serves 8 Q heads), scale 1/sqrt(256). K/V appended to per-layer ring before attention.
Output gate: attn_out *= sigmoid(z).
attn_output [n_embd, head_dim*n_head].

SSM (Gated-Delta-Net) block (`qwen35moe_ssm.c:229–351`)

Op order for SSM layers:

Projections from attn_norm-ed input: attn_qkv → qkv [conv_dim=8192]; attn_gate → z [ssm_d_inner=4096]; ssm_beta → beta [32] → sigmoid; ssm_alpha → alpha [32].
Decay gate per V-head (qwen35moe_ssm.c:86–93): g[h] = ssm_a[h] * softplus(alpha[h] + ssm_dt[h]). ssm_a is stored as −exp(A_log) (negative), so exp(g) ∈ (0,1).
Causal depthwise conv1d (kernel size 4) over conv_dim channels + SiLU (qwen35moe_ssm.c:110–135). Rolls conv_state in place (oldest-first, channel-major layout).
Split convolved output into q [128*16], k [128*16], v [128*32]; L2-normalize q/k per K-head (qwen35moe_ssm.c:70–76); GQA-tile from 16 K-heads to 32 V-heads using interleaved tiling (vh % n_k_heads == kh, qwen35moe_ssm.c:311).
Delta-rule recurrence per V-head (qwen35moe_ssm.c:163–195): for each of 32 [128×128] state matrices (stored transposed, row j = column j of S): decay → pred → delta → rank-1 write → read-after-write y. Scale 1/sqrt(128) applied once at read.
Gated RMSNorm: RMSNorm(y, ssm_norm) * silu(z) per V-head (qwen35moe_ssm.c:206–223), then ssm_out [n_embd, ssm_d_inner].

The recurrence is parallelized across V-heads via fw_for when USE_THREADS is set and n_v_heads >= 4 (qwen35moe_ssm.c:332–336).

2-Die Split and Placement

Constants (`qwen35moe_forward_cuda.c:89–95`)

#define GPU_DIE_SPLIT_LAYER  40        /* default = disabled (all on one die) */
#define GPU_DIE_A_DEFAULT     0
#define GPU_DIE_B_DEFAULT     1

GPU_DIE_SPLIT_LAYER=40 means no split by default; the single-die path is the validated baseline. A split is engaged only when QWEN35_GPU_SPLIT_LAYER < n_layer or explicit placement provides two dies.

Boundary resolution (`gpu_split_layer`, `qwen35moe_forward_cuda.c:696–711`)

pn_load_on(2 dies)  → split = n_layer/2 = 20
pn_load_on(1 die)   → split = n_layer (single-die forced)
no placement        → GPU_DIE_SPLIT_LAYER (40 = disabled)
QWEN35_GPU_SPLIT_LAYER env → overrides boundary (unless 1-die placement)

Clamped to (0, n_layer].

Die id resolution (`gpu_die_ids`, `qwen35moe_forward_cuda.c:717–736`)

Priority: pn_load_on explicit placement > QWEN35_GPU_DIE_A / QWEN35_GPU_DIE_B env vars > GPU_DIE_A_DEFAULT=0 / GPU_DIE_B_DEFAULT=1.

Die ids are ordinals relative to CUDA_VISIBLE_DEVICES. On a K80 host, one K80 card exposes two GK110 dies; CUDA_VISIBLE_DEVICES maps them to ordinals 0 and 1.

Weight placement

Die A (ctx_a): layers [0, split) + token embedding (read host-side, row uploaded per token).
Die B (ctx_b): layers [split, n_layer) + output_norm + LM head.

Each gpu_ssm_layer.ctx / gpu_attn_layer.ctx stores a back-pointer to its owning die context (qwen35moe_forward_cuda.c:205, 285).

2-die multiplexing (`gpu_select_die`, `qwen35moe_forward_cuda.c:751–764`)

On legacy large-BAR K80 rigs (single MMIO slot), barswap time-multiplexes the two dies. The forward:

Acquires die A via bs_acquire.
Runs front layers on ctx_a.
Syncs, D2H copies the boundary hidden [n_embd float32].
Acquires die B via bs_acquire.
H2D uploads hidden to ctx_b.
Runs back layers + LM head on ctx_b.

In small-BAR mode (modern driver, real dual-device), both contexts are live simultaneously and g->bs == NULL; no barswap or per-layer gpu_select_die is needed. The decode-slot bind path and the async path only run in small-BAR mode (qwen35moe_forward_cuda.c:4655–4658).

KV Cache

Allocation (`qwen35moe_kv_init`, `qwen35moe_forward.c:129–176`)

One pn_qwen35moe_kv per attention layer, 10 total. Layout [n_head_kv][max_seq][head_dim] (head-major).

Q4 mode (QWEN35_KV_Q4, default ON, qwen35moe_forward.c:47–56):

Cache stored as ggml-Q4_0 rows: kq [n_head_kv * max_seq * pn_kvq4_row_bytes(head_dim)], same for vq.
fp16 sink shadow k_sink/v_sink [n_head_kv * PN_KVQ4_SINKMAX * head_dim] — attention sink tokens kept at fp16 to prevent argmax flips on near-ties.
fp16 window shadow k_win/v_win [n_head_kv * PN_KVQ4_WMAX * head_dim] — the most recent W tokens similarly kept at fp16.
PPL delta vs fp32 KV: 0.024% (validated, qwen35moe_forward.c:45).
Disable with QWEN35_KV_Q4=0.

fp32 mode: k/v [n_head_kv * max_seq * head_dim float32].

The reactive memory ladder rung 4a (qwen35moe_forward.c:142–152) auto-forces Q4 if the box is under memory pressure at allocation time even without the env var.

KV window size

PN_QWEN35MOE_MAX_SEQ = 4096 (compile-time default, qwen35moe_forward.c:70).

QWEN35_MAX_SEQ env var overrides, clamped to:

minimum: 16
maximum: PN_QWEN35MOE_MAX_SEQ_CAP = 262 144 (qwen35moe.h:29)

The cap is the single source of truth shared between the CPU KV ring allocator (qwen35moe_max_seq, qwen35moe_forward.c:76–91) and the GPU admission gate (gpu_admit_max_ctx, qwen35moe_forward_cuda.c:619–629). Memory scales linearly: n_head_kv * max_seq * head_dim * 8 bytes per attention layer on the host (same again on device in GPU mode).

Ring behavior

KV is a fixed-size ring; pos % max_seq gives the slot. There is no dynamic resizing. If pos >= max_seq the oldest entry is overwritten. The ring is never grown after allocation.

Deferred host mirror (B2)

After a GPU decode token, the GPU KV ring is the live source of truth. Writing back to the host mirror (kv->k/v or kq/vq) per token was too expensive. Instead, kv->kv_mirror_stale = 1 is set (qwen35moe_forward_cuda.c:3547) and the D2H copy is deferred to pn_qwen35moe_gpu_state_to_host() (qwen35moe_forward_cuda.c:6779), which is called by the prefix cache on session eviction. After pn_qwen35moe_state_reset() the flag is cleared to 0 (qwen35moe_forward.c:308).

Decode Slots / Multi-User Batching

Slot capacity

PN_QWEN35MOE_GPU_MAX_SLOTS = 16 (qwen35moe_forward_cuda.c:115).

Memory cost: ~2.1 MB per SSM layer per slot for the per-slot conv_slots / ssm_slots images. At 16 slots across 30 SSM layers ≈ 1 GB on the die pair.

The reactive ladder rung 2 (pn_qwen35moe_gpu_ladder_max_slots, qwen35moe_forward_cuda.c:4605–4613) caps live slots: OK → 16, TIGHT → 8, CRITICAL → 1.

Slot protocol

1. prefill each user's pn_state (single-token or batched)
2. immediately after each user's prefill:
   pn_qwen35moe_decode_slot_bind(model, state_u, u, K)
   — D2D snapshots the live GPU SSM state into slot u;
     host mirror is stale after GPU prefill so bind must precede
     any other user's forward
3. step all K users: pn_qwen35moe_decode_slots(model, tokens, states, K, logits)
   — ONE forward pass, weights stream once for all K users

(qwen35moe.h:126–138)

Slot-major per-user SSM images live in gpu_ssm_layer.conv_slots / ssm_slots (qwen35moe_forward_cuda.c:247–250). Bind snapshots via D2D from the shared conv_state_g / ssm_state_g right after prefill (qwen35moe_forward_cuda.c:4616–4638). A slot-count change invalidates all existing bindings.

pn_qwen35moe_decode_slot_move(src, dst, n_slots) (qwen35moe.h:154) compacts the active set via D2D copy without a host round-trip.

Async step machine (C2, `QWEN35_ASYNC_SCHED`, default ON)

Three-state enum (qwen35moe_forward_cuda.c:7062–7065):

QW_ASYNC_IDLE — no step in flight
QW_ASYNC_WAIT_A — 2-die only: front segment (die A) in flight, CUDA event pending
QW_ASYNC_WAIT_FINAL — back segment + LM head in flight (or single-die full step)

pn_qwen35moe_forward_step_begin() launches the step and returns without blocking (returns 1 = not async-eligible and runs synchronously in that case). pn_qwen35moe_forward_step_poll() reaps via fc_event_query, advances to the next phase, and on completion sets *logits_out and increments state->pos. This lets a 2-core host host keep up to 8 cards' generations pipelined.

Requires: small-BAR mode (real dual-device, g->bs == NULL); LM head must have an on-device kernel (Q3_K/Q4_K/Q5_K/Q6_K/Q8_0/F32, qwen35moe_forward_cuda.c:7083–7098). Disable with QWEN35_ASYNC_SCHED=0.

The MTP / nextn Trap

The GGUF metadata for Qwen3.5-MoE may advertise MTP (multi-token prediction) draft-head weights under blk.N.nextn.* keys. The weights do not ship with public GGUF conversions; nextn_predict_layers must always be 0.

The loader explicitly documents this and never touches any blk.N.nextn.* tensor:

/* Honor the MTP trap: nextn_predict_layers=0 — never touch blk.N.nextn.* . */

(qwen35moe.c:475)

If a future GGUF includes these weights, loading them would require a separate code path not yet implemented. Attempting to load them without that path would silently skip them (no tensor lookup is performed) — so correctness is maintained but draft-head acceleration would simply not engage.

Gotchas and Fragile Spots

KV ring is fixed at allocation — no resize path

qwen35moe_kv_init (qwen35moe_forward.c:129) allocates the ring once at max_seq. There is no resize-on-demand path. If the ring reallocates (e.g., cache_cap changes on the GPU side under a geometry change), the segment graph must be dropped and recaptured (qwen35moe_forward_cuda.c:4941). Under memory pressure the reactive ladder forces Q4 at init time but cannot shrink the ring after the fact.

Segfault risk: if pos >= max_seq is not checked before the GPU KV-append kernel writes pos % max_seq, and the GPU ring's cache_cap was allocated for a different max_seq than the host ring, the write goes out of bounds. The GPU ring capacity is tracked separately as d->cache_cap (qwen35moe_forward_cuda.c:306); a mismatch after a partial realloc failure would be silent until the OOB write.

kv_mirror_stale + session switch

After a GPU decode step, kv->kv_mirror_stale = 1. If a session is evicted without calling pn_qwen35moe_gpu_state_to_host() first (e.g., the prefix cache path is bypassed or fails partway), the host mirror remains stale. A subsequent resume re-syncs from the stale host mirror instead of the GPU ring, silently corrupting the KV cache with zeros or old content.

The pn_qwen35moe_gpu_forget_state() call (qwen35moe.h:212) must be paired with pn_state_reset() on session reuse; omitting it leaves the GPU SSM buffers binding to the old session and the next prefill reads stale GPU residue.

GQA repeat tiling — interleaved, not blocked

The SSM's GQA repeat tiles K-heads to V-heads using vh % n_k_heads == kh (interleaved), not the blocked vh / kv_groups mapping (qwen35moe_ssm.c:311–312). This matches ggml_repeat's tiling in llama.cpp. Getting this wrong produces plausible but numerically wrong y vectors — the error is hard to detect without the oracle diff.

conv_state layout: channel-major, oldest-first

conv_state[c * (ssm_conv-1) + t], t=0 oldest. The roll-forward in qwen35moe_conv1d_silu (qwen35moe_ssm.c:131–133) shifts t=0..km1-2 forward and appends the new raw input at t=km1-1. The GPU path must replicate this exactly or the conv state diverges after the first token.

ssm_state is stored transposed

ssm_state[h * S * S + j * S + i] == S_h[i][j] — row j (contiguous over i) is column j of S (qwen35moe_internal.h:214–215). This matches llama.cpp's gated_delta_net layout. Any GPU kernel that reinterprets the layout as standard row-major will compute wrong dot products.

Batched CUDA prefill and CUDA-graph capture ordering

The per-layer segment graph is only captured after the layer has run DIRECTLY (uncaptured) at least once — body_ran gates the attempt (qwen35moe_forward_cuda.c:229–232). A batched prefill that precedes the first single-token forward never sets body_ran for those layers' fused descriptors, so the first single-token decode after a batch-only session must run uncaptured to build the fused descs. Capturing mid-allocation (before the fused descs exist) would call cuMemAlloc inside the capture, invalidating the graph. The body_ran gate prevents this.

Embedding: lazy vs materialized

For untied models (a separate output.weight exists), token_embd is left NULL and each row is dequantized on demand from the mmap (pn_qwen35moe_model_embed_row, qwen35moe.c:620–629). For tied models, the full ~2 GB F32 embedding is materialized and also serves as the LM head (qwen35moe.c:437–453). The fallback for untied models where n_embd % block_size != 0 also materializes the full heap.

Numeric status: A8 not done

qwen35moe_forward.c and qwen35moe_forward_cuda.c both carry an open @todo A8: validate logits vs the llama.cpp oracle. End-to-end numeric correctness has not been verified. Several operations in the CUDA path have no dedicated GPU kernel and fall back to host-side computation with a D2H/H2D round trip (marked TODO kernel: inline in the CUDA file): joint Q+gate split, GQA head broadcast + L2-norm, decay gate / softplus, sigmoid gates, SwiGLU for non-Q3_K experts, attention softmax.

17 KiB Raw Blame History Unescape Escape