From 261e0796d2cca687fc4082049eab0f8a659cf5c3 Mon Sep 17 00:00:00 2001 From: dejay2 <218806300+dejay2@users.noreply.github.com> Date: Wed, 2 Sep 2026 05:02:10 +0100 Subject: [PATCH] fix(kernels): drop the D2H sync from the varlen GDN/KDA prefill conv causal_conv1d_varlen sized its triton launch grid from the longest request in the batch, and the only place that number existed was on the device: the triton fallback fell back to int(seq_lens.max().item()), a D2H sync. Every prefill therefore paid a full pipeline stall to read back a number the scheduler already knew, and a sync is illegal inside a stream capture, so the prefill forward of every GDN/KDA model was uncapturable. build_fla_metadata computes the per-request lengths on the host, so carry the max there (FLAMetadata.max_seq_len) and thread it down through the three linear-attention ops (qwen3_5_moe, qwen4_exp, glm5_next) into the kernel wrapper. The kwarg is optional and the device-derived path is unchanged when it is omitted, so no other caller has to change. Tested on an RTX 5090 (triton fallback path, no sgl_kernel): 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 26 passed (21 before this change, 5 new). Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01RG8BXfsSZi1nh4wMZnhJQK --- python/freetoken/attention/linear.py | 5 + python/freetoken/kernel/causal_conv1d.py | 10 +- python/freetoken/models/glm5_next/kda.py | 1 + python/freetoken/models/qwen3_5_moe/gdn.py | 12 +- python/freetoken/models/qwen4_exp/gdn.py | 12 +- tests/kernels/test_causal_conv1d_capture.py | 147 ++++++++++++++++++++ 6 files changed, 177 insertions(+), 10 deletions(-) create mode 100644 tests/kernels/test_causal_conv1d_capture.py diff --git a/python/freetoken/attention/linear.py b/python/freetoken/attention/linear.py index f3717e58e..c174c976a 100644 --- a/python/freetoken/attention/linear.py +++ b/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, ) diff --git a/python/freetoken/kernel/causal_conv1d.py b/python/freetoken/kernel/causal_conv1d.py index 81f1fab7b..ea5a78c5c 100644 --- a/python/freetoken/kernel/causal_conv1d.py +++ b/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 diff --git a/python/freetoken/models/glm5_next/kda.py b/python/freetoken/models/glm5_next/kda.py index 1d44e3a44..d2aba7c50 100644 --- a/python/freetoken/models/glm5_next/kda.py +++ b/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) diff --git a/python/freetoken/models/qwen3_5_moe/gdn.py b/python/freetoken/models/qwen3_5_moe/gdn.py index 2e7320051..361db4289 100644 --- a/python/freetoken/models/qwen3_5_moe/gdn.py +++ b/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) diff --git a/python/freetoken/models/qwen4_exp/gdn.py b/python/freetoken/models/qwen4_exp/gdn.py index 69838153f..c238de095 100644 --- a/python/freetoken/models/qwen4_exp/gdn.py +++ b/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) diff --git a/tests/kernels/test_causal_conv1d_capture.py b/tests/kernels/test_causal_conv1d_capture.py new file mode 100644 index 000000000..025bd46cc --- /dev/null +++ b/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