The GDN prefill convolution writes q/k/v itself: 32 GB of copies gone per 32K prompt, 128 MiB of workspace freed

[MichaelDementii]

The chunked GDN prefill convolved into one packed [C,T] buffer and then pulled q, k and v out of
it with three cudaMemcpy2DAsync per layer per chunk. The convolution already knows which channel
each thread owns, so it can address the destination directly. causal_conv1d_silu_split does that.

Rebased onto 6e8b2e2a and squashed to one commit. This description replaces the original one. The
four review points are answered below; the second one turned up a route-selection problem in the
shared part of this Op that the same measurement had to resolve.

Scope

This PR predates the issue-first rule in CONTRIBUTING.md. Its problem and direction were
established in your review rather than in an Issue; if you would rather it be re-filed as an Issue
first, say so and I will close this and open one.

One decision: the convolution writes its three destinations instead of writing one plane that is
then copied three ways. Everything in the diff belongs to that decision, including the route
selection, which is the route table of this Op.

1. Compile-time geometry

CausalConvSplitOutput3 now carries the partition in the type, following w8_rowsplit_output.cuh:

template <class Element, std::int32_t Rows0, std::int32_t Rows1, std::int32_t Rows2>
struct CausalConvSplitOutput3 {
    Element* __restrict__ out0;
    Element* __restrict__ out1;
    Element* __restrict__ out2;
    __device__ __forceinline__ CausalConvOutputColumn<Element> column(std::int64_t row) const;
};

The struct holds only pointers; the boundaries and the three leading dimensions are constants inside
column(). The pointers carry __restrict__ to record what the split entry proves - though I should
say that it buys nothing today: column() hands back a plain pointer, and the SASS is identical
instruction for instruction with and without the qualifier. CausalConvContiguousOutput keeps a
plain pointer and a runtime leading dimension: one destination admits any row count, and the packed
entries do not prove their destination disjoint from x.

The wrapper resolves (x.ne[0], out0.ne[0], out1.ne[0], out2.ne[0]) to one of the two registered
profiles and throws otherwise, so dispatch stays where op-development.md section 4.1 puts it. Each
launcher takes the resolved tag, asserts that the geometry covers the channel extent it was handed,
and instantiates one kernel.

Two narrowings follow, both stated in include/ninfer/ops/causal_conv1d_silu.h. The partition
parity the paired route used to test at runtime is a static_assert. And the entry requires
four-byte aligned operands, which deletes the scalar split instantiation; the requirement is stated
for the entry rather than per route so a caller need not know which route its column count selects.
That alignment narrowing goes past what the review asked for - it buys one instantiation and one
host branch, and it is unobservable to the caller, whose tensors come off a 256-byte-aligned arena.

2. Route selection, measured

Both registered geometries, T = 1..72, cold L2, one process, median of 40:

./build/bench/ninfer_causal_conv1d_silu_bench --split --cache cold --channels 8192  --tokens $(seq -s, 1 72)
./build/bench/ninfer_causal_conv1d_silu_bench --split --cache cold --channels 10240 --tokens $(seq -s, 1 72)

On T == 1 the answer is no. On this commit, cold: the split entry measures 4.13 us, the
dedicated decode kernel 4.10, and a same-byte device copy of the same bytes 4.10. All three sit at
the launch floor of this harness, so there is nothing for a split decode route to recover, and none
was added.

The interval above the small-T bound was the finding. The sequence route, which both entries
selected for 17 <= T <= 64, measured 12.29 us at T = 17 and 28.67 us at T = 64, while the
prefill route it handed over to at T = 65 measured 8.19 us for strictly more work. Separately,
the small-T kernel had never been measured above its bound of 16, although it launches to 32
(kCausalConvChannelTile x T threads, capped at 1024). It does step there - its block is 32 x T, so
it holds two CTAs per SM to T = 24 and one above - but even at its ceiling it stays under the
prefill route.

So the route table is now small-T to 32 and prefill above, and the sequence kernel, its launcher,
its declaration and kCausalConvSequenceMaxTokens are removed rather than bypassed - its only call
sites were the two entries above, and nothing else in the tree used them.

T	before	after
17	12.29 us	6.11	x2.01
32	16.38	6.11	x2.68
48	22.53	8.22	x2.74
64	28.67	8.19	x3.50

C = 10240 is the same shape. Cold medians in this region land on a roughly 2 us grid, so read the
small-T numbers as ~4.1 and ~6.1 and the prefill ones as ~8.2; the route finding is three to ten grid
steps and survives that, but single-step differences within one route are not signal. The remaining
seam is 32 -> 33, ~6.1 to ~8.2 us, and it is the small-T kernel's structural limit rather than a
tuning choice: at T = 33 its block would need 32 x 33 = 1056 threads.

I took the removal as part of this decision rather than a separate PR because it is the route table
of the same Op and the measurement is the one point 2 asked for. Two consequences reach entries this
PR was not otherwise touching, and neither changes output. The packed entries' dispatch changes the
same way. And moving the small-T bound also moves the snapshot entry's, which reads the same
constant: T in 17..32 now takes causal_conv1d_snapshot_smallt_kernel instead of
causal_conv1d_sequence_snapshot_kernel. Measured on that entry, cold, C = 8192: 6.11 us at
T = 8, 16, 17 and 24, and 8.16 at 32 - flat across the bound that moved. If you would rather have
the removal separately, I will split it out.

3. The stage, not half of it

--legacy-stage times what the stage was - one packed convolution into a [C,T] plane, then three
extract_bf16_columns - against the split entry, in one process under one set of timing conditions:

./build/bench/ninfer_causal_conv1d_silu_bench --legacy-stage --split --cache cold --channels 8192 \
  --tokens 1,2,4,8,16,17,32,33,48,64,65,128,256,512,1024,2048,4096,8192

Speedup of the split entry over that stage, cold, by column count:

C	1-16	17-64	128	1024	8192
8192	x2.49-2.50	x1.68-2.01	x1.60	x1.73	x1.74
10240	x1.68-2.50	x1.74-2.01	x1.79	x1.78	x1.75

The 10240 band is not homogeneous: its worst point is T = 16 at x1.68, because that geometry is
already at ~6.1 us from T = 15 where 8192 is still at ~4.1.

Absolute, C = 8192, cold: 10.27 -> 4.10 us at T = 16, 53.22 -> 30.75 at 1024, 360.48 -> 206.85 at
8192. Across the whole ladder on both geometries the range is x1.56 to x2.50.

So "the split convolution pays 3.2-5.0%" was measured against the packed convolution alone; against
the stage it replaces it is faster everywhere. The --legacy-stage arm is the decision benchmark you
allowed; it is here so you can reproduce the comparison, and I will drop it with the scatter.h
include on request.

4. Documents, names, state rule

The two model documents are reverted to their master text: the mathematics did not change, and
they should not carry a function name or a materialization decision.

The Op parameters are out0/out1/out2. For the record, no public header in include/ninfer/ops/
uses ordinal names today - attn_input_proj.h and gdn_input_proj.h name destinations by role - but
those Ops are themselves about q/k/v and this one is a convolution, so I read your point as applying
here specifically.

The state rule is enforced, on both distinct-state entries rather than only the new one:

void require_state_alias_rule(const Tensor& conv_state_in, const Tensor& conv_state_out) {
    if (conv_state_in.data == conv_state_out.data) { return; }
    if (overlaps(conv_state_in, conv_state_out)) { throw std::invalid_argument(...); }
}

The packed entry had no overlap validation at all, and causal_conv1d_decode_launch selects
causal_conv1d_decode_distinct_kernel whenever the two pointers differ - a kernel that declares both
state pointers __restrict__. A partially overlapping pair was therefore undefined behaviour there.
The packed entry now also enforces the non-overlap the family contract states for out and
conv_state_out, which the route change made load-bearing: the prefill route publishes state in a
second launch that reads x, so an out-of-contract out == x call that the removed sequence kernel
happened to tolerate would otherwise corrupt state silently.

Resources

Workspace peak, read from the engine's own report, 35B-A3B, same prompt on both arms:

chunk	master	this branch
8192	963.06 MiB	835.06 MiB
1024	120.38 MiB	104.38 MiB

Exactly convolution_dim * chunk * 2 in each case - the buffer that no longer exists. No graph node
appears or disappears; prefill is not captured (src/core/decode_graph.cpp is the only capture
site). Per-kernel resources, from cuobjdump --dump-resource-usage on the two binaries: the split
small-T instantiations match the packed one at 38 registers and 1472 bytes of shared memory, and the
split prefill-pairs instantiations at 38 registers and no shared memory. Templating the packed
prefill kernel on its output map cost it two registers, 28 to 30.

After this change nothing under src/ calls either packed causal_conv1d_silu overload or
extract_bf16_columns; their remaining uses are the tests, bench/ops/gdn_layer_bench.cu, which
composes the unfused snapshot alternative, and the --legacy-stage arm this PR adds. I have not
deleted them - that is a separate decision and I would rather you make it.

Correctness

cd build && ctest -j1

92 pass, 1 skipped (27b_load_plan, no matching artifact on this box), 1 fails.
ninfer_qwen3_6_27b_prefix_real_test fails with Host checkpoint restore changed greedy output: restored=64,1248, baseline=64,56127, restored_spec=0/0/0/1 baseline_spec=0/0/0/1 reused=305 transfers=1/1/3/3 - identical on three consecutive runs of this commit. Unmodified master rebuilt
here fails with the same line, character for character; that comparison was run on 9dbc0740 and
6e8b2e2a during earlier work rather than on this commit, which changes nothing it touches. Both
arms of the test's own comparison run in one process on one artifact, so it is not an artifact or an
environment effect. I will report it separately.

The Op suite covers both registered geometries at T = 1, 2, 7, 15, 16, 17, 32, 63, 64, 65, 257, 1024, the exact-alias state form on both at T = 1, 2, 15, 16, 17, 32, 33, 64, 65, 257,
destinations offset by one pair
(four-byte aligned, which the contract admits, but not the 256-byte alignment the arena gives), and
thirteen rejection classes: a null x, an FP32 weight, a rank-3 destination, an unregistered row
profile, the other geometry's profile on these channels, an FP32 destination, a null destination, a
short column count, two destinations in one buffer, a destination overlapping x, a destination
overlapping the input state, a destination that is not four-byte aligned, and a state pair that
overlaps without being the same storage.

End to end on 35B-A3B, three rounds, arms alternating inside each round, greedy: prefill +1.34%
at 8,515 tokens, +0.93% at 33,031, +0.76% at 33,031 with chunk 1024. All nine generations
byte-identical to master.

Checks not run

• No ncu counters: RmProfilingAdminOnly is set on this host. Kernel-level claims come from the Op
  benchmark and cuobjdump, not hardware counters.
• Qwen3.8-27B NVFP4 was not re-measured end to end on this rework. It reaches the 10240 geometry
  through the same entry and the Op benchmark covers that geometry, but the end-to-end rows are 35B
  only.
• The route removal is qualified by the Op suite and the end-to-end gate. I did not run the serve
  corpus.

RTX 5090, sm_120a, driver 580.105.08, CUDA 13.1.115, Release, -DCMAKE_CUDA_ARCHITECTURES=120a.
Base is 6e8b2e2a.

[Neroued]

Note: This comment is being posted by an agent on behalf of the maintainer. The maintainer has reviewed and approved the core technical conclusions below.

Thank you for splitting this change out of #96. This PR is much easier to review. The overall direction is sound: having the convolution write q/k/v directly, removing the intermediate buffer and the three extraction operations, and reducing workspace are all valuable. The numerical qualification and route-boundary coverage are also thorough.

There are still a few points that should be tightened before merging:

1. CausalConvSplitOutput3 currently keeps row counts, partition selection, and leading dimensions as runtime values in order to support arbitrary partitions. The production layouts are only 2048 + 2048 + 4096 and 2048 + 2048 + 6144, and your own measurements show that the split convolution itself pays a 3.2%-5.0% cost. Please use compile-time specializations for the registered geometries, with wrapper dispatch based on geometry, so the production route does not pay for unused generality.

2. The new entry sends T == 1 through the small-T kernel instead of the existing dedicated decode-shaped kernel, but there is no short-interval performance evidence. Please measure both registered geometries at T = 1/2/7/15/16/17/64/65 and select the route from those results, including a dedicated split decode route if the data supports it.

3. The current microbenchmark compares the packed convolution alone with the split entry. The old complete stage also includes the three extraction operations. Please provide a direct, same-condition comparison of:

   packed convolution + 3 x extraction

   versus

   split convolution

   This may be a temporary decision benchmark; there is no need to retain the legacy stage after the implementation decision is closed.

4. The model documents should not record the concrete function name or whether a private intermediate tensor is materialized; the model mathematics did not change. Please also keep the generic Op parameter names independent of the first caller (out0/out1/out2 rather than out_q/out_k/out_v) and enforce the documented state rule that input/output state storage must be either disjoint or an exact alias.

With these changes, I expect this PR to be a strong merge candidate.

Once the requested changes are ready, please squash the PR into a single commit and rebase it onto the latest master HEAD before final review.

[MichaelDementii]

Reworked against all four points, rebased onto 6e8b2e2a, squashed to one commit. The description
above is rewritten; the short version:

1. The output map is now CausalConvSplitOutput3<Element, Rows0, Rows1, Rows2> following
   w8_rowsplit_output.cuh - pointers only, boundaries and leading dimensions constant. The wrapper
   resolves the row profile to one of the two registered geometries and refuses the rest, which also
   let the scalar split instantiation go.
2. Measured. At T == 1 the split entry, the decode kernel and a same-byte copy all sit at the
   4.1 us launch floor, so no split decode route was added. The sweep over T = 1..72 did turn up
   something else: the sequence route both entries selected for 17 <= T <= 64 cost 12.29 us at
   T=17 and 28.67 at T=64 against 8.19 for the prefill route above it, and the small-T kernel had
   never been measured above its bound of 16 although it launches to 32. The route table is now
   small-T to 32 and prefill above, and the sequence kernel is removed rather than bypassed.
3. --legacy-stage times the packed convolution plus its three extractions against the split entry
   in one process. The stage it replaces is x1.56 to x2.50 slower across both geometries; the
   earlier 3.2-5.0% was measured against the convolution alone.
4. Model documents reverted, parameters out0/out1/out2, and the state rule enforced on both
   distinct-state entries - the packed one had no overlap validation at all, and its decode route
   declares both state pointers __restrict__.

Two things are yours to decide rather than mine. The route removal changes the packed and snapshot
entries' dispatch, which this PR was not otherwise touching; output is unchanged and I measured the
snapshot entry across the moved bound, but if you would rather have it as a separate change I will
split it out. And this PR predates the issue-first rule - say the word and I will close it and open
an Issue instead.

Unrelated to this branch: ninfer_qwen3_6_27b_prefix_real_test fails here on a clean tree, three
runs, byte-identical message. I have opened it separately rather than leaving it as a footnote.