Skip to content

ssm_conv: channels-major input mode to drop the delta-net transpose#52

Draft
roberteg16 wants to merge 2 commits into
gfx11from
rogarcia.qwen3.6-moe-optimization
Draft

ssm_conv: channels-major input mode to drop the delta-net transpose#52
roberteg16 wants to merge 2 commits into
gfx11from
rogarcia.qwen3.6-moe-optimization

Conversation

@roberteg16

@roberteg16 roberteg16 commented Jul 15, 2026

Copy link
Copy Markdown

Summary

Qwen3.5/3.6 hybrid models (Gated DeltaNet + MoE) fed ggml_ssm_conv a time-major
conv input [time, channels, seqs], while the delta-net projection produces
channels-major [channels, time, seqs]. The required transpose showed up in
prefill as a concat_non_cont kernel scaling super-linearly with the ubatch, and
after forcing it contiguous it remained as a ~221 ms CONT (transpose) kernel at
ub=4096.

This PR removes it in two steps:

  1. delta-net: force contiguous conv-state concat input — materialize the
    transposed input so the conv-state concat takes the coalesced concat_cont path
    instead of concat_non_cont (which scaled super-linearly with the ubatch).
  2. ssm_conv: add channels-major input mode; drop delta-net transpose — add an
    opt-in channels-major layout to ggml_ssm_conv
    (ggml_ssm_conv_channels_major, selected via op_params[0]=1): input is
    [d_inner, d_conv-1+n_t, n_s] (channels contiguous), output layout unchanged.
    build_conv_state keeps qkv_mixed channels-major, prepends the recurrent conv
    state along the time axis (concat dim 1), and calls the new entry point —
    eliminating the transpose entirely.
  • Implemented on CPU and CUDA (short- and long-token kernels; coalesced smem load).
  • Other backends' supports_op reject op_params[0]!=0 so the scheduler falls back to
    CPU; existing (time-major) callers are bit-identical.
  • Shared by the qwen35, qwen35moe, and qwen3next graphs.

Measured impact (Qwen3.6-35B-A3B, prefill pp4096, gfx1151)

Max ubatch (-ub 4096 -b 4096 → M=4096) — per-commit breakdown

op baseline after commit 1 (contiguous concat) after commit 2 (channels-major)
CONCAT 1166 ms 34 ms 34 ms
CONT (transpose) 10 ms 221 ms 10 ms
SSM_CONV 47 ms 47 ms 43 ms
total prefill 3825 ms 2925 ms 2717 ms
  • Commit 1 removes the quadratic concat_non_cont blowup (1166 → 34 ms) but relocates
    the transpose cost into a 221 ms CONT kernel: −899 ms (−23.5%).
  • Commit 2 eliminates that transpose (CONT 221 → 10 ms) with no regression elsewhere:
    a further −208 ms.
  • Combined: 3825 → 2717 ms, −1108 ms (−29%).

Default batch (-ub 4096 -b 2048 → M=2048; the realistic config, since n_batch defaults to 2048 and n_ubatch is clamped to it)

op baseline this PR (channels-major)
CONCAT 283 ms 35 ms
CONT (transpose) 11 ms 11 ms
SSM_CONV 51 ms 43 ms
total prefill 2786 ms 2590 ms

At the smaller chunk the concat_non_cont cost is lower (the blowup is super-linear),
so the win is smaller but still meaningful: −196 ms (−7.0%), driven by CONCAT
283 → 35 ms with no added transpose.

Test plan

  • test-backend-ops -o SSM_CONV — new channels-major cases pass (CUDA vs CPU reference), short and long-token paths.
  • Perplexity healthy across the three archs sharing the helper: qwen3next (Qwen3.6-35B-A3B) 3.40, qwen35moe (Qwen3.5-35B-A3B) 2.18, qwen35 (Qwen3.5-0.8B) 16.28.
  • Non-CUDA backends only exercised via the CPU fallback guard; not run on Metal/Vulkan/SYCL/etc.
  • ggml-backend-meta.cpp (tensor-parallel split) still assumes time-major — untouched, out of scope.

Materialize the transposed qkv_mixed before ggml_concat so the conv-state
prepend takes the coalesced concat_cont path instead of concat_non_cont,
which scaled super-linearly with chunk size (CONCAT ~1166ms -> ~34ms at
ub=4096, -899ms total prefill).
Delta-net fed ggml_ssm_conv a transposed (time-major) conv input, forcing a
physical transpose. Option A removed the quadratic concat but left a 221ms
CONT (transpose) kernel at ub=4096.

Add an opt-in channels-major layout to ggml_ssm_conv (op_params[0]=1):
input is [d_inner, d_conv-1+n_t, n_s] (channels contiguous); output layout
unchanged. Implemented on CPU and CUDA (both short and long-token kernels,
coalesced smem load); other backends' supports_op reject mode!=0 so the
scheduler falls back to CPU. build_conv_state now keeps qkv_mixed channels-
major, prepends the recurrent conv state along the time axis (concat dim 1),
and calls ggml_ssm_conv_channels_major (qwen35/qwen35moe/qwen3next).

Prefill pp4096 ub=4096: CONT 221->10ms, total 2925->2717ms (-1108ms vs the
pre-optimization baseline). Validated: test-backend-ops SSM_CONV channels-
major cases pass (CUDA vs CPU); Qwen3.6 perplexity 3.40 (healthy).
@roberteg16 roberteg16 force-pushed the rogarcia.qwen3.6-moe-optimization branch from f0659c2 to 18abc43 Compare July 15, 2026 10:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant