Skip to content

Feature: split-aware GGML_OP_SSM_SCAN execution #6

Description

@assistmeister

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:

  1. 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).

  2. 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

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions