fix(kernels): drop the D2H sync from the varlen GDN/KDA prefill conv

[dejay2]

What

causal_conv1d_varlen only needs the longest request in the batch to size its
triton launch grid, but the only place that number lived was on the device, so the
triton fallback derived it with int(seq_lens.max().item()).

build_fla_metadata already computes the per-request lengths on the host, so it
now carries FLAMetadata.max_seq_len, and the three linear-attention ops that call
the conv (qwen3_5_moe/gdn.py, qwen4_exp/gdn.py, glm5_next/kda.py) pass it
down. The new kwarg is optional; with it omitted the kernel wrapper derives the
value on device exactly as before.

Why

int(seq_lens.max().item()) is a device-to-host sync on every prefill: the whole
pipeline stalls to read back a number the scheduler computed on the host in the
first place. It is also illegal inside a CUDA stream capture, so its presence alone
makes the prefill forward of every GDN/KDA model uncapturable. Passing the host
value removes both problems without touching the kernel.

How it was tested

Windows 11, RTX 5090, triton fallback path (no sgl_kernel installed).

python -m pytest -q \
  tests/kernels/test_causal_conv1d_capture.py \
  tests/models/qwen4_exp/test_gdn.py \
  tests/models/test_glm5_next_kda_snapshot.py \
  tests/models/test_glm5_next_kda_op.py \
  tests/kvcache/test_linear_state_pool_alloc.py

• before this commit: 21 passed
• after this commit: 26 passed (5 new)

The new tests/kernels/test_causal_conv1d_capture.py pins that the host-metadata
path performs no .item() at all, that the default device-derived path still
works, that both produce bit-identical output and conv-state updates, and that the
call captures into and replays from a torch.cuda.CUDAGraph.

What is NOT included

• Nothing in kernel/triton/causal_conv1d_triton.py: it already accepts an optional
  max_seq_len and falls back to the device-side max when it is None. This PR
  only supplies the value.
• The fork branch this comes from also primes the fla chunk-index cache before
  graph capture. That helper has no caller outside the fork's speculative-decoding
  graph runner, so it is left out here.

🤖 Generated with Claude Code

https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK

[MT-z]

Thanks for this one — the sync you are removing is real, and the reasoning in the description
matches what I see in the tree. I pulled the branch onto an Ada card to try it and the fix itself
is fine; what I ran into is that the three new tests all fail here, and the cause is that they
assume the triton fallback is the active path. Since sgl is an optional extra, I suspect the
tests pass on your install and fail on any freetoken[accel] one, which would make this hard to
land without someone hitting it.

Setup. RTX 4090 24 GB (sm_89), CUDA toolkit 13.3, sglang-kernel 0.4.5 installed
(freetoken[accel]). This PR on top of main @ 03c28d2 plus PR #337 @ a6bd5c0. Every number
below is from this box.

tests/kernels/test_causal_conv1d_capture.py    3 failed, 2 passed

Cause 1 — the two backends disagree about whether x is mutated (2 of the 3)

causal_conv1d_varlen dispatches on is_sgl_kernel_installed(), and the two branches differ:

backend	returns	mutates x	mutates conv_states
sgl_kernel (native)	x itself	yes	yes
triton fallback	a new tensor	no	yes

Measured by calling _call() from your own _conv_inputs() with the dispatch forced each way.

test_varlen_conv_with_host_metadata_matches_the_device_derived_result and
test_varlen_conv_replays_inside_a_cuda_graph both restore conv_states between the two calls
they compare, but not x. On the triton path that is correct, because x is untouched. With
sgl_kernel the second call convolves an already-convolved x, so the two results differ and the
torch.equal assertions fail. Nothing is wrong with the kernel or with your change here — the
comparison is just not starting from the same input twice.

Cause 2 — the D2H sync only exists on the fallback path (the third)

test_varlen_conv_still_derives_max_seq_len_on_device_by_default asserts that .item() is called
when max_seq_len is omitted. Counting torch.Tensor.item calls with the dispatch forced each
way:

sgl_kernel path     .item() called 0 times   -> assertion fails
triton fallback     .item() called 1 time    -> assertion holds

Which is exactly right: int(seq_lens.max().item()) lives in
kernel/triton/causal_conv1d_triton.py, and the native kernel does not need the value at all. So
the test is asserting a property of the fallback while running whichever backend happens to be
installed.

Suggested test fix — verified, 5 passed here

Restore x alongside conv_states, and pin the third test to the path whose behaviour it
describes:

     baseline_states = inputs["conv_states"].clone()
+    baseline_x = inputs["x"].clone()
...
     inputs["conv_states"].copy_(baseline_states)
+    inputs["x"].copy_(baseline_x)

+    # The sync this asserts on lives in the triton fallback; on an install with
+    # sgl_kernel the native kernel needs no max_seq_len and calls no .item().
+    import freetoken.kernel.backend as _backend
+    monkeypatch.setattr(_backend, "is_sgl_kernel_installed", lambda: False)
     monkeypatch.setattr(torch.Tensor, "item", counted_item)

With both applied: 5 passed on this box. The second one is the part I would argue for on its
own merits — forcing the fallback means the test proves the claim on any install, rather than
only where the fallback happens to be selected. (A skipif(is_sgl_kernel_installed()) would also
go green, but it would stop testing the thing on the machines most likely to run CI.)

I have not opened this as a PR; the patch is small enough to paste, and it is your branch. Happy
to send it if you would rather have it that way.

One separate observation, offered as a note rather than a request

The x-mutation divergence above is not caused by this PR and is harmless in the tree today: all
three call sites (qwen3_5_moe/gdn.py:122, qwen4_exp/gdn.py:131, glm5_next/kda.py:183) build
x = conv_in.transpose(0, 1).contiguous() immediately before the call and never read it again, so
nobody depends on x surviving. But the wrapper's docstring does not say which contract holds,
and a future caller that keeps x would break on one install shape and not the other — the same
way these tests just did. Might be worth a line in the wrapper's docstring while this file is
open. I have not audited beyond the three call sites above.

For what it is worth, the fix is a no-op on my own serving path for the same reason — with
sgl_kernel installed the sync never executes — so I cannot give you a before/after timing. On a
default install (no [accel]), where the fallback is the path, the change should do exactly
what you describe.

_{Written with AI assistance; every number above was measured on my hardware (RTX 4090, sm_89)
and I can reproduce it.}