ssm_conv: channels-major input mode to drop the delta-net transpose#52
Draft
roberteg16 wants to merge 2 commits into
Draft
ssm_conv: channels-major input mode to drop the delta-net transpose#52roberteg16 wants to merge 2 commits into
roberteg16 wants to merge 2 commits into
Conversation
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).
f0659c2 to
18abc43
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Qwen3.5/3.6 hybrid models (Gated DeltaNet + MoE) fed
ggml_ssm_conva time-majorconv input
[time, channels, seqs], while the delta-net projection produceschannels-major
[channels, time, seqs]. The required transpose showed up inprefill as a
concat_non_contkernel scaling super-linearly with the ubatch, andafter forcing it contiguous it remained as a ~221 ms
CONT(transpose) kernel atub=4096.This PR removes it in two steps:
delta-net: force contiguous conv-state concat input— materialize thetransposed input so the conv-state concat takes the coalesced
concat_contpathinstead of
concat_non_cont(which scaled super-linearly with the ubatch).ssm_conv: add channels-major input mode; drop delta-net transpose— add anopt-in channels-major layout to
ggml_ssm_conv(
ggml_ssm_conv_channels_major, selected viaop_params[0]=1): input is[d_inner, d_conv-1+n_t, n_s](channels contiguous), output layout unchanged.build_conv_statekeepsqkv_mixedchannels-major, prepends the recurrent convstate along the time axis (concat dim 1), and calls the new entry point —
eliminating the transpose entirely.
supports_oprejectop_params[0]!=0so the scheduler falls back toCPU; existing (time-major) callers are bit-identical.
qwen35,qwen35moe, andqwen3nextgraphs.Measured impact (Qwen3.6-35B-A3B, prefill pp4096, gfx1151)
Max ubatch (
-ub 4096 -b 4096→ M=4096) — per-commit breakdownconcat_non_contblowup (1166 → 34 ms) but relocatesthe transpose cost into a 221 ms
CONTkernel: −899 ms (−23.5%).CONT221 → 10 ms) with no regression elsewhere:a further −208 ms.
Default batch (
-ub 4096 -b 2048→ M=2048; the realistic config, sincen_batchdefaults to 2048 andn_ubatchis clamped to it)At the smaller chunk the
concat_non_contcost 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.ggml-backend-meta.cpp(tensor-parallel split) still assumes time-major — untouched, out of scope.