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

python/freetoken/attention/linear.py

@@ -26,12 +26,16 @@ class FLAMetadata:
       fresh_state_indices prefill only: the state-pool slots whose sequence is fresh
                           (cached_len == 0) and must be zeroed before the chunk kernel
                           reads them in place. None if there are none / for decode.
+      max_seq_len         prefill only: the longest extend_len, as a host int. It sizes the
+                          varlen conv's launch grid; carrying it here keeps that launch free
+                          of the D2H sync a device-side max would need (illegal under capture).
     """
 
     cu_seqlens: torch.Tensor
     cache_indices: torch.Tensor
     has_initial_state: torch.Tensor | None = None
     fresh_state_indices: torch.Tensor | None = None
+    max_seq_len: int | None = None
 
     # --- hybrid-radix track-checkpoint (extra_buffer) fields; all None when not caching ---
     # For each request crossing a chunk-aligned (×CHUNK) boundary this forward, snapshot its
@@ -87,6 +91,7 @@ def gdn_slot(r):
         fresh_state_indices=(
             fresh_host.to(device, non_blocking=True) if fresh_host is not None else None
         ),
+        max_seq_len=max(lens),
         **track,
     )
 

python/freetoken/kernel/causal_conv1d.py

@@ -22,9 +22,14 @@ def causal_conv1d_varlen(
     cu_seqlens: torch.Tensor,   # [batch+1] int32 prefix sums of per-request lengths
     cache_indices: torch.Tensor,    # [batch] int32 slot id per request
     has_initial_state: torch.Tensor,  # [batch] bool (carry conv state across chunks)
+    max_seq_len: int | None = None,   # host-known longest extend_len; None derives it on device
 ) -> torch.Tensor:
     """Varlen (prefill) depthwise causal conv with silu; writes silu(conv) into ``x``
-    in place and refreshes ``conv_states[cache_indices]`` with each request's tail."""
+    in place and refreshes ``conv_states[cache_indices]`` with each request's tail.
+
+    ``max_seq_len`` only sizes the triton launch grid. Deriving it from ``cu_seqlens`` costs a
+    D2H sync per prefill, which is illegal under CUDA graph capture, so the scheduler passes the
+    host value it already has."""
     from freetoken.kernel.backend import is_sgl_kernel_installed
 
     if not is_sgl_kernel_installed():
@@ -33,7 +38,8 @@ def causal_conv1d_varlen(
         )
 
         return triton_causal_conv1d_varlen(
-            x, weight, conv_states, cu_seqlens, cache_indices, has_initial_state
+            x, weight, conv_states, cu_seqlens, cache_indices, has_initial_state,
+            max_seq_len=max_seq_len,
         )
 
     from sgl_kernel import causal_conv1d_fwd

python/freetoken/models/glm5_next/kda.py

@@ -183,6 +183,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             mixed = causal_conv1d_varlen(
                 x, self._conv_weight(), pool.conv_states[li],
                 fla.cu_seqlens, fla.cache_indices, fla.has_initial_state,
+                max_seq_len=fla.max_seq_len,  # host-known: no D2H sync sizing the launch
             ).transpose(0, 1)
             q, k, v = (
                 t.reshape(1, total, h, d).to(dtype)

python/freetoken/models/qwen3_5_moe/gdn.py

@@ -113,14 +113,17 @@ def _gate_params(self, a: torch.Tensor, b: torch.Tensor):
     def _conv_weight(self) -> torch.Tensor:
         return self.conv1d.weight.squeeze(1)  # [conv_dim, kernel] for the fused kernel
 
-    def _conv_prefill(self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state) -> torch.Tensor:
+    def _conv_prefill(self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state,
+                      max_seq_len=None) -> torch.Tensor:
         """Varlen causal conv (fused sgl_kernel) with silu; reads/updates each request's
         conv state in place by ``cache_indices`` slot. ``conv_in`` [total, conv_dim].
-        ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` come from FLAMetadata."""
+        ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` / ``max_seq_len`` come from
+        FLAMetadata; the last is host-known so the launch needs no D2H sync."""
         li = pool.local_index(self.layer_id)
         x = conv_in.transpose(0, 1).contiguous()  # [conv_dim, total]
         out = causal_conv1d_varlen(x, self._conv_weight(), pool.conv_states[li],
-                                   cu_seqlens, cache_indices, has_initial_state)
+                                   cu_seqlens, cache_indices, has_initial_state,
+                                   max_seq_len=max_seq_len)
         return out.transpose(0, 1)  # [total, conv_dim]
 
     def _conv_decode(self, conv_in: torch.Tensor, table_idx: torch.Tensor, pool) -> torch.Tensor:
@@ -189,7 +192,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             )
         else:
             mixed = self._conv_prefill(
-                conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state)
+                conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state,
+                fla.max_seq_len)
             # fla chunk handles GQA in-kernel: q/k stay at num_k_heads, v at num_v_heads.
             qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
             q = qf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype)

python/freetoken/models/qwen4_exp/gdn.py

@@ -122,14 +122,17 @@ def _gate_params(self, a: torch.Tensor, b: torch.Tensor):
     def _conv_weight(self) -> torch.Tensor:
         return self.conv1d.weight.squeeze(1)  # [conv_dim, kernel] for the fused kernel
 
-    def _conv_prefill(self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state) -> torch.Tensor:
+    def _conv_prefill(self, conv_in, pool, cu_seqlens, cache_indices, has_initial_state,
+                      max_seq_len=None) -> torch.Tensor:
         """Varlen causal conv (fused sgl_kernel) with silu; reads/updates each request's
         conv state in place by ``cache_indices`` slot. ``conv_in`` [total, conv_dim].
-        ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` come from FLAMetadata."""
+        ``cu_seqlens`` / ``cache_indices`` / ``has_initial_state`` / ``max_seq_len`` come from
+        FLAMetadata; the last is host-known so the launch needs no D2H sync (graph capture)."""
         li = pool.local_index(self.layer_id)
         x = conv_in.transpose(0, 1).contiguous()  # [conv_dim, total]
         out = causal_conv1d_varlen(x, self._conv_weight(), pool.conv_states[li],
-                                   cu_seqlens, cache_indices, has_initial_state)
+                                   cu_seqlens, cache_indices, has_initial_state,
+                                   max_seq_len=max_seq_len)
         return out.transpose(0, 1)  # [total, conv_dim]
 
     def _conv_decode(self, conv_in: torch.Tensor, table_idx: torch.Tensor, pool) -> torch.Tensor:
@@ -198,7 +201,8 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
             )
         else:
             mixed = self._conv_prefill(
-                conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state)
+                conv_in, pool, fla.cu_seqlens, fla.cache_indices, fla.has_initial_state,
+                fla.max_seq_len)
             # fla chunk handles GQA in-kernel: q/k stay at num_k_heads, v at num_v_heads.
             qf, kf, vf = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1)
             q = qf.reshape(1, total, self.num_k_heads, self.head_k_dim).to(dtype)

tests/kernels/test_causal_conv1d_capture.py

@@ -0,0 +1,147 @@
+"""The varlen GDN conv must be CUDA-graph capturable: its ``max_seq_len`` is host-known
+metadata everywhere the scheduler calls it, and deriving it on device costs a D2H sync that
+is illegal inside a capture region."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from freetoken.kernel.causal_conv1d import causal_conv1d_varlen
+
+requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
+
+
+def _conv_inputs(device, *, lens=(4, 2), conv_dim=8, kernel=4):
+    total = sum(lens)
+    cu = torch.tensor([0, *lens], dtype=torch.int32).cumsum(0).to(torch.int32).to(device)
+    return dict(
+        x=torch.randn(conv_dim, total, device=device, dtype=torch.bfloat16),
+        weight=torch.randn(conv_dim, kernel, device=device, dtype=torch.bfloat16),
+        conv_states=torch.randn(
+            len(lens) + 1, conv_dim, kernel - 1, device=device, dtype=torch.bfloat16
+        ),
+        cu_seqlens=cu,
+        cache_indices=torch.arange(1, len(lens) + 1, dtype=torch.int32, device=device),
+        has_initial_state=torch.ones(len(lens), dtype=torch.bool, device=device),
+    )
+
+
+def _call(inputs, **extra):
+    return causal_conv1d_varlen(
+        inputs["x"],
+        inputs["weight"],
+        inputs["conv_states"],
+        inputs["cu_seqlens"],
+        inputs["cache_indices"],
+        inputs["has_initial_state"],
+        **extra,
+    )
+
+
+@requires_cuda
+def test_varlen_conv_skips_the_device_to_host_sync_when_max_seq_len_is_given(monkeypatch):
+    device = torch.device("cuda")
+    inputs = _conv_inputs(device)
+    original_item = torch.Tensor.item
+    calls = []
+
+    def counted_item(self):
+        calls.append(tuple(self.shape))
+        return original_item(self)
+
+    monkeypatch.setattr(torch.Tensor, "item", counted_item)
+
+    _call(inputs, max_seq_len=4)
+    torch.cuda.synchronize()
+
+    assert calls == []
+
+
+@requires_cuda
+def test_varlen_conv_still_derives_max_seq_len_on_device_by_default(monkeypatch):
+    device = torch.device("cuda")
+    inputs = _conv_inputs(device)
+    original_item = torch.Tensor.item
+    calls = []
+
+    def counted_item(self):
+        calls.append(tuple(self.shape))
+        return original_item(self)
+
+    monkeypatch.setattr(torch.Tensor, "item", counted_item)
+
+    _call(inputs)
+    torch.cuda.synchronize()
+
+    assert calls, "the default path must still work without host-side metadata"
+
+
+@requires_cuda
+def test_varlen_conv_with_host_metadata_matches_the_device_derived_result():
+    device = torch.device("cuda")
+    inputs = _conv_inputs(device)
+    baseline_states = inputs["conv_states"].clone()
+
+    device_derived = _call(inputs).clone()
+    device_states = inputs["conv_states"].clone()
+
+    inputs["conv_states"].copy_(baseline_states)
+    host_known = _call(inputs, max_seq_len=4).clone()
+
+    assert torch.equal(host_known, device_derived)
+    assert torch.equal(inputs["conv_states"], device_states)
+
+
+@requires_cuda
+def test_varlen_conv_replays_inside_a_cuda_graph():
+    device = torch.device("cuda")
+    inputs = _conv_inputs(device)
+    baseline_states = inputs["conv_states"].clone()
+    expected = _call(inputs, max_seq_len=4).clone()
+    expected_states = inputs["conv_states"].clone()
+
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            inputs["conv_states"].copy_(baseline_states)
+            _call(inputs, max_seq_len=4)
+    torch.cuda.current_stream().wait_stream(stream)
+    torch.cuda.synchronize()
+
+    graph = torch.cuda.CUDAGraph()
+    inputs["conv_states"].copy_(baseline_states)
+    with torch.cuda.graph(graph, stream=stream, capture_error_mode="thread_local"):
+        captured = _call(inputs, max_seq_len=4)
+    inputs["conv_states"].copy_(baseline_states)
+    graph.replay()
+    torch.cuda.synchronize()
+
+    assert torch.equal(captured, expected)
+    assert torch.equal(inputs["conv_states"], expected_states)
+
+
+def test_prefill_fla_metadata_carries_the_host_known_longest_extend_len():
+    from freetoken.attention.linear import build_fla_metadata
+    from freetoken.core import Batch, Req, SamplingParams
+
+    reqs = []
+    for index, length in enumerate((5, 2, 9)):
+        req = Req(
+            input_ids=torch.arange(length, dtype=torch.int32),
+            table_idx=index,
+            cached_len=0,
+            output_len=1,
+            uid=index,
+            sampling_params=SamplingParams(),
+            cache_handle=None,
+        )
+        req.linear_slot_idx = index + 1
+        reqs.append(req)
+    batch = Batch(reqs=reqs, phase="prefill")
+    batch.padded_reqs = batch.reqs
+
+    fla = build_fla_metadata(batch, torch.device("cpu"))
+
+    assert fla.max_seq_len == 9