Hit a wall trying to shard the mamba scan under tensor parallelism. GGML_OP_SSM_SCAN writes one fused F32 buffer: [ y : nelements(x) | K state snapshots : K·d_state·d_inner·n_seqs ] (ggml.c:5652). We shard every mamba weight and activation along channels, but per-device ownership of this single output can't be expressed as contiguous slices — the y region interleaves channels across token rows while the states region repeats whole D rows.
Right now any sharded tensor reaching the scan forces all-mirrored mamba (LLAMA_SSM_SPLIT in the fork branch is an inert no-op since 8791d5ee), which leaves ~40% of non-embedding weight bytes and the entire S_cache unsharded.
Here's the numbers: Nemotron-3.5-Lightning-30B-A3B Q4_0 (nemotron_h_moe) has n_embd 2688, d_inner 4096, head_dim 64 × n_head 64, d_state 128, n_group 8, d_conv 4. Per-layer S_cache slice alone is 524288 f32 elements. Halving it per GPU plus halving ssm_in/ssm_out reads is worth ~+8% decode and doubles recurrent-context headroom (64k → 131k) at equal tok/s. Same kernel family gates qwen3next/qwen35/falcon-h1/laguna/bailingmoe3 GDN-style hybrid TP work later.
What would unblock this:
Either of these would work:
-
Decoupled outputs (preferred): split ggml_ssm_scan into two ops or add a flag so y and the state result are separate tensors, each individually contiguous. Then a channel-split scan becomes fully expressible with existing meta primitives: sharded x/dt/B/C/states-in → sharded y (lane-owned head halves) + sharded states-out; state copy-back cpy targets lane cache slabs; closure needs exactly one AR after ssm_out per layer (+26 ARs/token, already budgeted).
-
Sharded-execution contract for the fused op: keep the fused buffer but define that when x/dt/A/D/ssm-states arrive as channel-sliced views (with matching sliced declared shapes), each device's kernel invocation processes only its heads and writes only its heads' slots into its local fused buffer ([y_half | states_half]), where lane-local layout matches the global one with head_dim·(n_head_local) width per token row and (n_head_local/n_group) group re-basing.
For option 2, the group re-basing works out because Mamba2 scans are per-head independent given (x_h, dt_h, A_h, D_h, B_g(h), C_g(h)) with g(h) = h·n_group/n_head. For widths dividing evenly, a contiguous head half maps to whole group slabs; if each lane receives B/C slices relabeled from its first owned group onward (e.g. lane 1 gets groups 4..7 exposed as local 0..3 via pointer offset), then locally computed g(h') indexes the correct global group. Our arch satisfies alignment trivially: heads 0-31 ↔ groups 0-3, heads 32-63 ↔ groups 4-7. The kernel just needs to derive group index positionally from local shape only (no absolute head ids anywhere) and not read across its own slab. Option 2 would also need CPU fallback path parity so graph verification tools agree, and snapshot/rollback (K>1) lanes follow the same slicing.
Meta-backend side we'll handle: given either contract, handle_ssm_scan in ggml/src/ggml-backend-meta.cpp derives output segmentation from src0 states + x segments (handler skeleton already flagged in-tree); static routing groundwork (ssm_in 5-segment over {x,B,C,gates}, conv 3-segment, per-head cache slabs) exists in branch tp-bailingmoe3 history up to commit 8791d5ee^.
Validation on our end: we can rebuild llama-server only, recreate the probe profile nemotron35-tp-ssm-ngram (clone nemotron35-tp-gqa-ngram with image = "modelbench/mx-llama:tp-nemo-ssm-next"), and run:
- decode ≥ 355 tok/s @2048c1 (mirrored baseline 337-348), AR count ≈ +26/token, zero amdgpu faults
- greedy accuracy core.jsonl: 33-34/35 (±1 noise)
- parity probe vs LLAMA_SSM_SPLIT=0 byte-identical greedy outputs at fixed seed on smoke-tiny
Hit a wall trying to shard the mamba scan under tensor parallelism.
GGML_OP_SSM_SCANwrites one fused F32 buffer:[ y : nelements(x) | K state snapshots : K·d_state·d_inner·n_seqs ](ggml.c:5652). We shard every mamba weight and activation along channels, but per-device ownership of this single output can't be expressed as contiguous slices — the y region interleaves channels across token rows while the states region repeats whole D rows.Right now any sharded tensor reaching the scan forces all-mirrored mamba (
LLAMA_SSM_SPLITin the fork branch is an inert no-op since 8791d5ee), which leaves ~40% of non-embedding weight bytes and the entire S_cache unsharded.Here's the numbers: Nemotron-3.5-Lightning-30B-A3B Q4_0 (nemotron_h_moe) has n_embd 2688, d_inner 4096, head_dim 64 × n_head 64, d_state 128, n_group 8, d_conv 4. Per-layer S_cache slice alone is 524288 f32 elements. Halving it per GPU plus halving ssm_in/ssm_out reads is worth ~+8% decode and doubles recurrent-context headroom (64k → 131k) at equal tok/s. Same kernel family gates qwen3next/qwen35/falcon-h1/laguna/bailingmoe3 GDN-style hybrid TP work later.
What would unblock this:
Either of these would work:
Decoupled outputs (preferred): split ggml_ssm_scan into two ops or add a flag so y and the state result are separate tensors, each individually contiguous. Then a channel-split scan becomes fully expressible with existing meta primitives: sharded x/dt/B/C/states-in → sharded y (lane-owned head halves) + sharded states-out; state copy-back cpy targets lane cache slabs; closure needs exactly one AR after ssm_out per layer (+26 ARs/token, already budgeted).
Sharded-execution contract for the fused op: keep the fused buffer but define that when x/dt/A/D/ssm-states arrive as channel-sliced views (with matching sliced declared shapes), each device's kernel invocation processes only its heads and writes only its heads' slots into its local fused buffer ([y_half | states_half]), where lane-local layout matches the global one with head_dim·(n_head_local) width per token row and (n_head_local/n_group) group re-basing.
For option 2, the group re-basing works out because Mamba2 scans are per-head independent given (x_h, dt_h, A_h, D_h, B_g(h), C_g(h)) with g(h) = h·n_group/n_head. For widths dividing evenly, a contiguous head half maps to whole group slabs; if each lane receives B/C slices relabeled from its first owned group onward (e.g. lane 1 gets groups 4..7 exposed as local 0..3 via pointer offset), then locally computed g(h') indexes the correct global group. Our arch satisfies alignment trivially: heads 0-31 ↔ groups 0-3, heads 32-63 ↔ groups 4-7. The kernel just needs to derive group index positionally from local shape only (no absolute head ids anywhere) and not read across its own slab. Option 2 would also need CPU fallback path parity so graph verification tools agree, and snapshot/rollback (K>1) lanes follow the same slicing.
Meta-backend side we'll handle: given either contract, handle_ssm_scan in ggml/src/ggml-backend-meta.cpp derives output segmentation from src0 states + x segments (handler skeleton already flagged in-tree); static routing groundwork (ssm_in 5-segment over {x,B,C,gates}, conv 3-segment, per-head cache slabs) exists in branch tp-bailingmoe3 history up to commit 8791d5ee^.
Validation on our end: we can rebuild llama-server only, recreate the probe profile nemotron35-tp-ssm-ngram (clone nemotron35-tp-gqa-ngram with image = "modelbench/mx-llama:tp-nemo-ssm-next"), and run: