feat(qwen3next): GDN spec-decode verify path + linear-att cache split

Gated-delta-net (linear attention) speculative-decode verify path for qwen3next: a per-sequence spec causal_conv1d kernel; a widened conv working slot split from the committed (narrow) persisted slot; MTP draft full-attn KV-slot accounting across the linear-att cache config, mem operator and req manager; and removal of the dead gen_b_req_mtp_start_loc kernel.

Author: sufubao

lightllm/common/basemodel/triton_kernel/linear_att_copy.py

@@ -5,12 +5,13 @@
 
 @triton.jit
 def _copy_linear_att_state_to_kv_buffer(
-    gpu_conv_ptr,  # [linear_layer_num, size_num, xdim]
+    gpu_conv_ptr,  # [linear_layer_num, size_num, conv_dim * gpu_widened_width] (uint8 tail)
     gpu_ssm_ptr,  # [linear_layer_num, size_num, xxdim]
-    cpu_kv_conv_ptr,  # [size, linear_layer_num, xdim]
+    cpu_kv_conv_ptr,  # [size, linear_layer_num, conv_dim * width_narrow] (uint8 tail)
     cpu_kv_ssm_ptr,  # [size, linear_layer_num, xxdim]
     b_req_idx,  # [batch_size,]
     big_page_buffer_ids,  # [batch_size,]
+    num_accepted_tokens_ptr,  # [batch_size,]
     gpu_conv_stride_l,
     gpu_conv_stride_s,
     gpu_conv_stride_d,
@@ -24,7 +25,9 @@ def _copy_linear_att_state_to_kv_buffer(
     cpu_kv_ssm_stride_l,
     cpu_kv_ssm_stride_d,
     mtp_step,
-    gpu_conv_tail_dim,
+    conv_dim,  # number of conv rows (the d dimension)
+    gpu_conv_row_bytes,  # widened per-row byte length: gpu_widened_width * itemsize
+    conv_narrow_row_bytes,  # narrow per-row byte length: width_narrow * itemsize
     gpu_ssm_tail_dim,
     BLOCK: tl.constexpr,
 ):
@@ -40,28 +43,26 @@ def _copy_linear_att_state_to_kv_buffer(
         return
 
     cur_req_idx = tl.load(b_req_idx + cur_batch).to(tl.int64)
-    cur_state_req_idx = (cur_req_idx * (mtp_step + 1)).to(tl.int64)
+    accept_len = tl.load(num_accepted_tokens_ptr + cur_batch).to(tl.int64)
+    canonical_off = accept_len - 1
 
-    for i in range(tl.cdiv(gpu_conv_tail_dim, BLOCK)):
-        gpu_start_off = i * BLOCK + tl.arange(0, BLOCK)
-        mask = gpu_start_off < gpu_conv_tail_dim
-        conv_data = tl.load(
-            gpu_conv_ptr + cur_layer * gpu_conv_stride_l + cur_state_req_idx * gpu_conv_stride_s + gpu_start_off,
-            mask=mask,
-        )
-        dest_conv_ptr = (
-            cpu_kv_conv_ptr
-            + big_page_buffer_idx * cpu_kv_conv_stride_s
-            + cur_layer * cpu_kv_conv_stride_l
-            + gpu_start_off
-        )
-        tl.store(dest_conv_ptr, conv_data, mask=mask)
+    conv_src_slot = cur_req_idx
+    conv_off_bytes = canonical_off * gpu_conv_stride_d
+    gpu_conv_base = gpu_conv_ptr + cur_layer * gpu_conv_stride_l + conv_src_slot * gpu_conv_stride_s + conv_off_bytes
+    cpu_conv_base = cpu_kv_conv_ptr + big_page_buffer_idx * cpu_kv_conv_stride_s + cur_layer * cpu_kv_conv_stride_l
+    for d in range(conv_dim):
+        for i in range(tl.cdiv(conv_narrow_row_bytes, BLOCK)):
+            off = i * BLOCK + tl.arange(0, BLOCK)
+            mask = off < conv_narrow_row_bytes
+            conv_data = tl.load(gpu_conv_base + d * gpu_conv_row_bytes + off, mask=mask)
+            tl.store(cpu_conv_base + d * cpu_kv_conv_stride_d + off, conv_data, mask=mask)
 
+    ssm_src_slot = (cur_req_idx * (mtp_step + 1) + canonical_off).to(tl.int64)
     for i in range(tl.cdiv(gpu_ssm_tail_dim, BLOCK)):
         gpu_start_off = i * BLOCK + tl.arange(0, BLOCK)
         mask = gpu_start_off < gpu_ssm_tail_dim
         ssm_data = tl.load(
-            gpu_ssm_ptr + cur_layer * gpu_ssm_stride_l + cur_state_req_idx * gpu_ssm_stride_s + gpu_start_off,
+            gpu_ssm_ptr + cur_layer * gpu_ssm_stride_l + ssm_src_slot * gpu_ssm_stride_s + gpu_start_off,
             mask=mask,
         )
         dest_ssm_ptr = (
@@ -75,32 +76,51 @@ def _copy_linear_att_state_to_kv_buffer(
 def copy_linear_att_state_to_kv_buffer(
     b_req_idx: torch.Tensor,
     big_page_buffer_ids: torch.Tensor,
-    gpu_conv_state: torch.Tensor,  # [linear_layer_num, s, ...]
-    gpu_ssm_state: torch.Tensor,  # [linear_layer_num, s, ...]
-    cpu_kv_conv_state: torch.Tensor,  # [s, linear_layer_num, ...]
-    cpu_kv_ssm_state: torch.Tensor,  # [s, linear_layer_num, ...]
+    gpu_conv_state: torch.Tensor,  # [linear_layer_num, s_widened, conv_dim, gpu_widened_width]
+    gpu_ssm_state: torch.Tensor,  # [linear_layer_num, s_block, ...]
+    cpu_kv_conv_state: torch.Tensor,  # [size, linear_layer_num, conv_dim, width_narrow]
+    cpu_kv_ssm_state: torch.Tensor,  # [size, linear_layer_num, ...]
     mtp_step: int,
+    b_num_accepted_tokens: torch.Tensor,  # [batch_size,] per-req post-accept count (>=1)
 ):
     assert len(b_req_idx) == big_page_buffer_ids.shape[0]
+    assert len(b_req_idx) == b_num_accepted_tokens.shape[0]
     BLOCK = 4096
-    gpu_conv_state = gpu_conv_state.view(gpu_conv_state.shape[0], gpu_conv_state.shape[1], -1).view(dtype=torch.uint8)
-    gpu_ssm_state = gpu_ssm_state.view(gpu_ssm_state.shape[0], gpu_ssm_state.shape[1], -1).view(dtype=torch.uint8)
-    cpu_kv_conv_state = cpu_kv_conv_state.view(cpu_kv_conv_state.shape[0], cpu_kv_conv_state.shape[1], -1).view(
-        dtype=torch.uint8
+
+    assert gpu_conv_state.dim() >= 4, "gpu_conv_state must be [layer, s, conv_dim, widened_width]"
+    assert cpu_kv_conv_state.dim() >= 4, "cpu_kv_conv_state must be [size, layer, conv_dim, width_narrow]"
+    # #6: the byte snapshot hardcodes gpu_conv_stride_d=conv_itemsize, which is only valid when the
+    # widened-width axis is element-contiguous (stride 1). Fail fast instead of snapshotting wrong bytes.
+    assert gpu_conv_state.stride(3) == 1, (
+        "gpu_conv_state widened-width axis must be element-contiguous (stride 1); "
+        "gpu_conv_stride_d=conv_itemsize assumes it"
+    )
+    # #18: canonical_off = accept_len - 1 indexes into the widened slot; bound it to [0, mtp_step]
+    # (accept_len in [1, mtp_step+1]) so a stale/oversized accept-count can't slice past the slot.
+    assert int(b_num_accepted_tokens.min()) >= 1 and int(b_num_accepted_tokens.max()) <= mtp_step + 1, (
+        f"b_num_accepted_tokens out of range [1, {mtp_step + 1}]: "
+        f"min={int(b_num_accepted_tokens.min())} max={int(b_num_accepted_tokens.max())}"
     )
+    conv_itemsize = gpu_conv_state.element_size()
+    gpu_conv_state = gpu_conv_state.view(
+        gpu_conv_state.shape[0], gpu_conv_state.shape[1], gpu_conv_state.shape[2], -1
+    ).view(dtype=torch.uint8)
+    cpu_kv_conv_state = cpu_kv_conv_state.view(
+        cpu_kv_conv_state.shape[0], cpu_kv_conv_state.shape[1], cpu_kv_conv_state.shape[2], -1
+    ).view(dtype=torch.uint8)
+
+    gpu_ssm_state = gpu_ssm_state.view(gpu_ssm_state.shape[0], gpu_ssm_state.shape[1], -1).view(dtype=torch.uint8)
     cpu_kv_ssm_state = cpu_kv_ssm_state.view(cpu_kv_ssm_state.shape[0], cpu_kv_ssm_state.shape[1], -1).view(
         dtype=torch.uint8
     )
-    assert gpu_conv_state.shape[-1] == cpu_kv_conv_state.shape[-1]
+
+    assert gpu_conv_state.shape[2] == cpu_kv_conv_state.shape[2], "conv_dim mismatch between gpu and cpu conv buffers"
     assert gpu_ssm_state.shape[-1] == cpu_kv_ssm_state.shape[-1]
-    assert (
-        gpu_conv_state.stride(-1)
-        == gpu_ssm_state.stride(-1)
-        == cpu_kv_conv_state.stride(-1)
-        == cpu_kv_ssm_state.stride(-1)
-    )
 
-    gpu_conv_tail_dim = gpu_conv_state.shape[-1]
+    conv_dim = gpu_conv_state.shape[2]
+    gpu_conv_row_bytes = gpu_conv_state.shape[-1]  # widened per-row byte length
+    conv_narrow_row_bytes = cpu_kv_conv_state.shape[-1]  # narrow per-row byte length
+    assert conv_narrow_row_bytes <= gpu_conv_row_bytes
     gpu_ssm_tail_dim = gpu_ssm_state.shape[-1]
 
     layer_num = gpu_conv_state.shape[0]
@@ -114,9 +134,10 @@ def copy_linear_att_state_to_kv_buffer(
         cpu_kv_ssm_ptr=cpu_kv_ssm_state,
         b_req_idx=b_req_idx,
         big_page_buffer_ids=big_page_buffer_ids,
+        num_accepted_tokens_ptr=b_num_accepted_tokens,
         gpu_conv_stride_l=gpu_conv_state.stride(0),
         gpu_conv_stride_s=gpu_conv_state.stride(1),
-        gpu_conv_stride_d=gpu_conv_state.stride(2),
+        gpu_conv_stride_d=conv_itemsize,
         gpu_ssm_stride_l=gpu_ssm_state.stride(0),
         gpu_ssm_stride_s=gpu_ssm_state.stride(1),
         gpu_ssm_stride_d=gpu_ssm_state.stride(2),
@@ -127,7 +148,9 @@ def copy_linear_att_state_to_kv_buffer(
         cpu_kv_ssm_stride_l=cpu_kv_ssm_state.stride(1),
         cpu_kv_ssm_stride_d=cpu_kv_ssm_state.stride(2),
         mtp_step=mtp_step,
-        gpu_conv_tail_dim=gpu_conv_tail_dim,
+        conv_dim=conv_dim,
+        gpu_conv_row_bytes=gpu_conv_row_bytes,
+        conv_narrow_row_bytes=conv_narrow_row_bytes,
         gpu_ssm_tail_dim=gpu_ssm_tail_dim,
         BLOCK=BLOCK,
     )

lightllm/common/basemodel/triton_kernel/linear_att_cpu_cache_copy.py

@@ -193,11 +193,7 @@ def copy_kv_buffer_to_cpu_cache(
     cpu_kv_ssm_tail_dim = cpu_kv_ssm_state.shape[-1]
     full_att_layer_num = gpu_kv_full_att_state.shape[-2]
 
-    assert (
-        full_att_layer_num
-        == (linear_config.all_layer_num // linear_config.full_attention_interval)
-        == (linear_config.all_layer_num - linear_config.linear_layer_num)
-    )
+    assert full_att_layer_num == linear_config.get_persisted_full_att_layer_num()
     assert gpu_full_att_tail_dim == cpu_cache_full_att.shape[-1]
     assert cpu_cache_conv.shape[-1] == cpu_kv_conv_state.shape[-1]
     assert cpu_cache_ssm.shape[-1] == cpu_kv_ssm_state.shape[-1]
@@ -428,6 +424,7 @@ def copy_cpu_cache_to_kv_buffer(
     cpu_kv_ssm_tail_dim = cpu_kv_ssm_state.shape[-1]
     full_att_layer_num = gpu_full_att_kv_state.shape[-2]
 
+    assert full_att_layer_num == linear_config.get_persisted_full_att_layer_num()
     assert gpu_full_att_tail_dim == cpu_cache_full_att.shape[-1]
     assert cpu_cache_conv.shape[-1] == cpu_kv_conv_state.shape[-1]
     assert cpu_cache_ssm.shape[-1] == cpu_kv_ssm_state.shape[-1]

lightllm/common/basemodel/triton_kernel/mtp_utils.py

@@ -148,38 +148,6 @@ def mtp_scatter_next_token_ids(
     )
 
 
-@triton.jit
-def _fwd_kernel_gen_b_req_mtp_start_loc(
-    b_mtp_index,
-    b_req_mtp_start_loc,
-    num_reqs: tl.constexpr,
-    batch_size: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    offset = tl.arange(0, BLOCK_SIZE)
-    cur_mtp_index = tl.load(b_mtp_index + offset, mask=offset < batch_size, other=-1)
-    non_zero_mask = tl.where(cur_mtp_index == 0, 1, 0)  # 1 0 1 0 0
-    output_offset = tl.cumsum(non_zero_mask) - 1
-    tl.store(b_req_mtp_start_loc + output_offset, offset, mask=non_zero_mask == 1)
-    return
-
-
-def gen_b_req_mtp_start_loc(b_mtp_index: torch.Tensor, num_reqs: int):
-    b_req_mtp_start_loc = torch.empty((num_reqs,), dtype=torch.int32, device=b_mtp_index.device)
-    BLOCK_SIZE = triton.next_power_of_2(b_mtp_index.shape[0])
-    batch_size = b_mtp_index.shape[0]
-    grid = (1,)
-    _fwd_kernel_gen_b_req_mtp_start_loc[grid](
-        b_mtp_index=b_mtp_index,
-        b_req_mtp_start_loc=b_req_mtp_start_loc,
-        num_reqs=num_reqs,
-        batch_size=batch_size,
-        BLOCK_SIZE=BLOCK_SIZE,
-        num_warps=8,
-    )
-    return b_req_mtp_start_loc
-
-
 def test_mtp_verify():
     req_to_next_token_ids = torch.tensor(
         [[1, 2, -2, -1, -1], [1, 2, 0, -1, -1], [1, 3, 4, 4, 5]], dtype=torch.int32, device="cuda"
@@ -201,13 +169,5 @@ def test_mtp_verify():
     print(accepted_index)
 
 
-def test_gen_b_req_mtp_start_loc():
-    b_mtp_index = torch.tensor([0, 1, 0, 1, 2], dtype=torch.int32, device="cuda")
-    gt_output = torch.where(b_mtp_index == 0)[0]
-    b_req_mtp_start_loc = gen_b_req_mtp_start_loc(b_mtp_index, 2)
-    print(b_req_mtp_start_loc, gt_output)
-
-
 if __name__ == "__main__":
     test_mtp_verify()
-    # test_gen_b_req_mtp_start_loc()

lightllm/common/kv_cache_mem_manager/operator/linear_att.py

@@ -24,6 +24,16 @@ def __init__(self, mem_manager):
         super().__init__(mem_manager)
         self.linear_config = LinearAttCacheConfig.load_from_args()
 
+    @staticmethod
+    def _get_persisted_full_att_layer_num(mem_manager) -> int:
+        persisted_full_att = getattr(mem_manager, "persisted_full_att_layer_num", None)
+        if persisted_full_att is None:
+            main_full_att = getattr(mem_manager, "main_full_att_layer_num", mem_manager.kv_buffer.shape[0])
+            draft_full_att = getattr(mem_manager, "draft_full_att_layers", 0)
+            persisted_full_att = main_full_att + draft_full_att
+        assert 0 < persisted_full_att <= mem_manager.kv_buffer.shape[0]
+        return int(persisted_full_att)
+
     def load_cpu_cache_to_gpu(
         self,
         mem_indexes: torch.Tensor,
@@ -76,11 +86,14 @@ def load_cpu_cache_to_gpu(
             copy_cpu_cache_to_kv_buffer,
         )
 
+        # Restore the persisted full-attn slice: main slots followed by MTP draft slots.
+        persisted_full_att = self._get_persisted_full_att_layer_num(mem_manager)
+
         copy_cpu_cache_to_kv_buffer(
             mem_indexes=mem_indexes,
             big_page_buffer_ids=big_page_buffer_ids_gpu,
             page_indexes=page_indexes,
-            gpu_full_att_kv_state=mem_manager.kv_buffer,
+            gpu_full_att_kv_state=mem_manager.kv_buffer[:persisted_full_att],
             cpu_kv_conv_state=mem_manager.linear_att_big_page_buffers.conv_state_cache.buffer,
             cpu_kv_ssm_state=mem_manager.linear_att_big_page_buffers.ssm_state_cache.buffer,
             cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor,
@@ -169,12 +182,15 @@ def offload_gpu_kv_to_cpu_cache(
             copy_kv_buffer_to_cpu_cache,
         )
 
+        # Persist the full-attn slice used for prefix reuse: main slots followed by MTP draft slots.
+        persisted_full_att = self._get_persisted_full_att_layer_num(mem_manager)
+
         copy_kv_buffer_to_cpu_cache(
             mem_indexes=mem_indexes,
             page_indexes=page_indexes,
             page_readies=page_readies,
             big_page_buffer_ids=big_page_buffer_ids_gpu,
-            gpu_kv_full_att_state=mem_manager.kv_buffer,
+            gpu_kv_full_att_state=mem_manager.kv_buffer[:persisted_full_att],
             cpu_kv_conv_state=mem_manager.linear_att_big_page_buffers.conv_state_cache.buffer,
             cpu_kv_ssm_state=mem_manager.linear_att_big_page_buffers.ssm_state_cache.buffer,
             cpu_cache_tensor=cpu_cache_client.cpu_kv_cache_tensor,

lightllm/common/linear_att_cache_manager/config_objs.py

@@ -1,13 +1,18 @@
 import torch
 import dataclasses
 import triton
-from lightllm.utils.envs_utils import get_env_start_args
+from lightllm.utils.envs_utils import get_env_start_args, _mtp_added_layer_num
 from lightllm.utils.log_utils import init_logger
 from lightllm.utils.torch_dtype_utils import get_torch_dtype
 
 logger = init_logger(__name__)
 
 
+def get_mtp_draft_full_att_layer_num(args) -> int:
+    # Delegates to the single source of truth in envs_utils (#9).
+    return _mtp_added_layer_num(getattr(args, "mtp_mode", None), getattr(args, "mtp_step", 0))
+
+
 @dataclasses.dataclass
 class LinearAttCacheConfig:
     tp_world_size: int
@@ -30,6 +35,7 @@ class LinearAttCacheConfig:
     ssm_state_dtype: torch.dtype
     full_attention_interval: int
     all_layer_num: int  # 包括 linear att 和 full att 的层加起来的层数
+    draft_full_att_layer_num: int = 0
 
     def get_conv_dim(self):
         # 第一项对应q的参数，第二项对应k的参数，第三项对应v的参数
@@ -41,9 +47,25 @@ def get_conv_dim(self):
             + self.head_linear_v_dim * self.num_linear_v_heads
         )
 
-    def get_conv_state_shape(self):
+    def get_main_full_att_layer_num(self):
+        main_full_att_layer_num = self.all_layer_num - self.linear_layer_num
+        assert main_full_att_layer_num == self.all_layer_num // self.full_attention_interval
+        return main_full_att_layer_num
+
+    def get_persisted_full_att_layer_num(self):
+        return self.get_main_full_att_layer_num() + self.draft_full_att_layer_num
+
+    def get_persisted_conv_state_shape(self):
+        # NARROW shape used for the CPU/disk persisted page and ALL byte math.
+        # Persisted state is always the committed (narrow) sliding window.
         return (self.get_conv_dim(), self.conv_kernel_size - 1)
 
+    def get_gpu_conv_state_shape(self, mtp_step: int):
+        # WIDENED working shape for the GPU buffer: holds the tentatively
+        # rolled-in S speculative tokens before acceptance. width-1 + S, where
+        # S = mtp_step (a verify step has seqlen=S+1 -> width-1+(seqlen-1)).
+        return (self.get_conv_dim(), (self.conv_kernel_size - 1) + mtp_step)
+
     def get_ssm_state_shape(self):
         return (self.num_linear_v_heads, self.head_linear_k_dim, self.head_linear_v_dim)
 
@@ -66,7 +88,7 @@ def get_cpu_cache_full_att_bytes(self):
         )
         assert big_page_token_num == get_env_start_args().cpu_cache_token_page_size
         full_att_bytes = 2 * self.full_att_all_num_kv_heads * self.full_att_head_dim * self.full_att_dtype.itemsize
-        a = full_att_bytes * (self.all_layer_num - self.linear_layer_num) * big_page_token_num
+        a = full_att_bytes * self.get_persisted_full_att_layer_num() * big_page_token_num
         return a
 
     def get_cpu_cache_conv_bytes(self):
@@ -113,4 +135,5 @@ def load_from_args() -> "LinearAttCacheConfig":
             ssm_state_dtype=get_torch_dtype(args.linear_att_ssm_data_type),
             full_attention_interval=llm_config["full_attention_interval"],
             all_layer_num=n_layer,
+            draft_full_att_layer_num=get_mtp_draft_full_att_layer_num(args),
         )

lightllm/common/linear_att_cache_manager/linear_att_buffer_manager.py

@@ -24,7 +24,7 @@ def __init__(
         self.conv_state_cache = LayerCache(
             size=self.size,
             dtype=self.linear_config.conv_state_dtype,
-            shape=self.linear_config.get_conv_state_shape(),
+            shape=self.linear_config.get_persisted_conv_state_shape(),
             layer_num=self.linear_config.linear_layer_num,
             device="cpu",
             size_first=True,