feat(mtp): add windowed draft KV cache for DFlash-style models

lightllm/common/basemodel/attention/fa3/windowed_mtp.py

@@ -0,0 +1,49 @@
+import weakref
+
+import torch
+
+from ..base_att import AttControl, BaseAttBackend, BaseDecodeAttState
+from lightllm.utils.sgl_utils import flash_attn_with_kvcache
+
+
+class WindowedMTPAttBackend(BaseAttBackend):
+    """Non-causal block attention over retained draft KV and temporary noise KV."""
+
+    def create_att_decode_state(self, infer_state):
+        return WindowedMTPDecodeAttState(backend=self, _infer_state_weak_ref=weakref.ref(infer_state))
+
+
+class WindowedMTPDecodeAttState(BaseDecodeAttState):
+    def init_state(self):
+        model = self.backend.model
+        self.block_size = model.block_size
+        self.b_req_idx = self.infer_state.b_req_idx[:: self.block_size].contiguous()
+        # The base attention-state copy refreshes these tensors on graph replay.
+        store = self.infer_state.mem_manager.windowed_draft_kv
+        self.cache_seqlens = (
+            store.req_to_seq_len.index_select(0, self.b_req_idx).clamp(max=store.capacity).to(torch.int32)
+            + self.block_size
+        )
+
+    def prepare_kv(self, layer_index, noise_kv, alloc_func=torch.empty):
+        """Pack retained history and the current block in the backend's KV layout."""
+        scratch = self.infer_state.mem_manager.windowed_draft_kv.pack(
+            layer_index - self.backend.model.draft_layer_start,
+            self.b_req_idx,
+            noise_kv,
+            self.block_size,
+            alloc_func=alloc_func,
+        )
+        head_num = noise_kv.shape[1] // 2
+        return scratch[:, :, :head_num], scratch[:, :, head_num:]
+
+    def decode_att(self, q, k, v, att_control: AttControl = AttControl(), alloc_func=torch.empty):
+        return flash_attn_with_kvcache(
+            q.view(-1, self.block_size, q.shape[-2], q.shape[-1]),
+            k,
+            v,
+            cache_seqlens=self.cache_seqlens,
+            causal=False,
+            window_size=att_control.sliding_window if att_control.use_sliding_window else (-1, -1),
+            softmax_scale=q.shape[-1] ** -0.5,
+        ).view(q.shape)

lightllm/common/basemodel/basemodel.py

@@ -4,7 +4,6 @@
 import gc
 import copy
 import math
-import json
 import torch
 import torch.nn.functional as F
 import triton
@@ -15,8 +14,7 @@
 from lightllm.common.basemodel.infer_struct import InferStateInfo
 from lightllm.common.kv_cache_mem_manager import MemoryManager
 from lightllm.common.kv_cache_mem_manager.mem_utils import select_mem_manager_class
-from lightllm.common.req_manager import ReqManager
-from lightllm.common.build_utils import repair_config
+from lightllm.common.req_manager import ReqManager, ReqManagerForWindowedMTP
 from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import (
     select_kv_index_from_req,
     select_kv_index_from_req_prefill,
@@ -27,6 +25,7 @@
 from lightllm.common.quantization import Quantcfg
 from lightllm.common.basemodel.triton_kernel.gather_token_id import gather_token, gather_token_prefill_decode_mixed
 from lightllm.utils.log_utils import init_logger
+from lightllm.utils.config_utils import load_config
 from lightllm.utils.dist_utils import get_dp_world_size
 from lightllm.utils.profile_max_tokens import profile_mtp_weight_memory
 from lightllm.utils.envs_utils import (
@@ -146,12 +145,7 @@ def __init__(self, kvargs):
         return
 
     def _init_config(self):
-        with open(os.path.join(self.weight_dir_, "config.json"), "r") as json_file:
-            self.config = json.load(json_file)
-        # rename keys
-        repair_config(self.config, same_names=["num_attention_heads", "n_head"])
-        repair_config(self.config, same_names=["hidden_size", "n_embd", "n_embed"])
-        repair_config(self.config, same_names=["num_hidden_layers", "n_layer"])
+        self.config = load_config(self.weight_dir_)
         if self.finetune_config:
             self.config["vocab_size"] = self.finetune_config.vocab_size
         return
@@ -237,7 +231,8 @@ def _init_req_manager(self):
         if self.max_seq_length is not None:
             create_max_seq_len = max(create_max_seq_len, self.max_seq_length)
 
-        self.req_manager = ReqManager(self.max_req_num, create_max_seq_len, None)
+        req_manager_class = ReqManagerForWindowedMTP if self.args.mtp_windowed else ReqManager
+        self.req_manager = req_manager_class(self.max_req_num, create_max_seq_len, None)
         return
 
     def _init_infer_layer(self, start_layer_index=0):

lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py

@@ -29,8 +29,9 @@ def _fwd_kernel_destindex_copy_kv(
     k_ptrs = K + cur_index * stride_k_bs + stride_k_h * offs_h[:, None] + stride_k_d * offs_d[None, :]
     o_ptrs = Out + dest_index * stride_o_bs + stride_o_h * offs_h[:, None] + stride_o_d * offs_d[None, :]
 
-    k = tl.load(k_ptrs, mask=offs_h[:, None] < head_num, other=0.0)
-    tl.store(o_ptrs, k, mask=offs_h[:, None] < head_num)
+    mask = (dest_index >= 0) & (offs_h[:, None] < head_num)
+    k = tl.load(k_ptrs, mask=mask, other=0.0)
+    tl.store(o_ptrs, k, mask=mask)
     return
 
 

lightllm/common/basemodel/triton_kernel/hybrid_cpu_cache_copy.py

@@ -0,0 +1,54 @@
+"""Stream-ordered copies for pinned checkpoint pages."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _copy_checkpoint_pages(
+    src_state,
+    dst_state,
+    src_mem_index,
+    dst_mem_index,
+    page_readies,
+    src_stride_s,
+    dst_stride_s,
+    state_dim,
+    page_num,
+    SKIP_READY: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    for page in range(page_num):
+        src_id = tl.load(src_mem_index + page).to(tl.int64)
+        dst_id = tl.load(dst_mem_index + page).to(tl.int64)
+        valid = (src_id >= 0) & (dst_id >= 0)
+        if SKIP_READY:
+            valid = valid & (tl.load(page_readies + page) == 0)
+        if valid:
+            for block in range(tl.program_id(0), tl.cdiv(state_dim, BLOCK), tl.num_programs(0)):
+                offsets = block * BLOCK + tl.arange(0, BLOCK)
+                data = tl.load(src_state + src_id * src_stride_s + offsets, offsets < state_dim, other=0)
+                tl.store(dst_state + dst_id * dst_stride_s + offsets, data, offsets < state_dim, cache_modifier=".wt")
+
+
+def copy_checkpoint_pages(sources, destinations, source_ids, destination_ids, page_readies=None):
+    """Copy slot-indexed, contiguous rows; preserve each tensor's byte representation."""
+    assert len(source_ids) == len(destination_ids)
+    for source, destination in zip(sources, destinations):
+        source = source.reshape(source.shape[0], -1).view(torch.uint8)
+        destination = destination.reshape(destination.shape[0], -1).view(torch.uint8)
+        assert source.shape[1] == destination.shape[1]
+        _copy_checkpoint_pages[(12,)](
+            source,
+            destination,
+            source_ids,
+            destination_ids,
+            page_readies,
+            source.stride(0),
+            destination.stride(0),
+            source.shape[1],
+            len(source_ids),
+            SKIP_READY=page_readies is not None,
+            BLOCK=4096,
+        )

lightllm/common/basemodel/triton_kernel/kv_cache_offload.py

@@ -61,6 +61,7 @@ def _offload_gpu_kv_to_cpu(
 
         token_range = block_index * TOKEN_BLOCK + tl.arange(0, TOKEN_BLOCK)
         token_indexes = tl.load(token_indexes_ptr + token_range).to(tl.int64)
+        token_mask = token_indexes >= 0
         head_dim_range = tl.arange(0, BLOCK_HEAD_DIM)
         head_dim_mask = head_dim_range < head_dim
         scale_head_dim_mask = head_dim_range < scale_head_dim
@@ -77,7 +78,7 @@ def _offload_gpu_kv_to_cpu(
                     + gpu_k_head_index.to(tl.int64) * gpu_stride2
                     + head_dim_range[None, :]
                 )
-                gpu_data = tl.load(gpu_ptr, mask=head_dim_mask[None, :], other=0.0)
+                gpu_data = tl.load(gpu_ptr, mask=token_mask[:, None] & head_dim_mask[None, :], other=0.0)
                 cpu_ptr = (
                     cpu_kv_cache_ptr
                     + cpu_page_index * cpu_stride0
@@ -89,7 +90,7 @@ def _offload_gpu_kv_to_cpu(
                 tl.store(
                     cpu_ptr,
                     gpu_data,
-                    mask=head_dim_mask[None, :],
+                    mask=token_mask[:, None] & head_dim_mask[None, :],
                     cache_modifier=".wt",
                 )
                 if HAS_SCALE:
@@ -127,7 +128,7 @@ def _offload_gpu_kv_to_cpu(
                     + gpu_v_head_index.to(tl.int64) * gpu_stride2
                     + head_dim_range[None, :]
                 )
-                gpu_data = tl.load(gpu_ptr, mask=head_dim_mask[None, :], other=0.0)
+                gpu_data = tl.load(gpu_ptr, mask=token_mask[:, None] & head_dim_mask[None, :], other=0.0)
                 cpu_ptr = (
                     cpu_kv_cache_ptr
                     + cpu_page_index * cpu_stride0
@@ -139,7 +140,7 @@ def _offload_gpu_kv_to_cpu(
                 tl.store(
                     cpu_ptr,
                     gpu_data,
-                    mask=head_dim_mask[None, :],
+                    mask=token_mask[:, None] & head_dim_mask[None, :],
                     cache_modifier=".wt",
                 )
                 if HAS_SCALE:
@@ -407,6 +408,7 @@ def _load_cpu_cache_to_gpu(
         gpu_mem_indexes = tl.load(gpu_mem_indexes_ptr + token_range, mask=token_mask).to(tl.int64)
         cpu_mem_indexes = tl.load(cpu_mem_indexes_ptr + token_range, mask=token_mask).to(tl.int64)
         cpu_page_indexes = tl.load(cpu_page_indexes_ptr + token_range, mask=token_mask).to(tl.int64)
+        token_mask = token_mask & (gpu_mem_indexes >= 0) & (cpu_page_indexes >= 0)
 
         head_dim_range = tl.arange(0, BLOCK_HEAD_DIM)
         head_dim_mask = head_dim_range < head_dim

lightllm/common/basemodel/triton_kernel/windowed_mtp.py

@@ -0,0 +1,151 @@
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _prepare_window_write_indices(
+    b_req_idx,
+    b_start_loc,
+    b_ready_cache_len,
+    b_commit_len,
+    out_mem_index,
+    req_to_seq_len,
+    window_size: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_req_idx = tl.load(b_req_idx + cur_batch).to(tl.int64)
+    cur_ready_cache_len = tl.load(b_ready_cache_len + cur_batch)
+    cur_commit_len = tl.load(b_commit_len + cur_batch)
+    cur_seq_len = cur_ready_cache_len + cur_commit_len
+    cur_start_loc = tl.load(b_start_loc + cur_batch)
+
+    # Only the retained tail may write: older rows would collide in the ring.
+    tail_start = tl.maximum(0, cur_commit_len - window_size)
+    for start in range(tail_start, cur_commit_len, BLOCK):
+        row = start + tl.arange(0, BLOCK)
+        position = cur_ready_cache_len + row
+        mem_index = cur_req_idx * window_size + position % window_size
+        tl.store(out_mem_index + cur_start_loc + row, mem_index, mask=row < cur_commit_len)
+    # Publish the new length now; the caller must order all layer KV writes before any state reads.
+    tl.store(req_to_seq_len + cur_req_idx, cur_seq_len)
+
+
+def prepare_window_write_indices(
+    b_req_idx,
+    b_start_loc,
+    b_ready_cache_len,
+    b_commit_len,
+    req_to_seq_len,
+    window_size,
+    token_num,
+    alloc_func=torch.empty,
+):
+    """Return one destination per source token; rejected and expired rows stay -1."""
+    mem_index = alloc_func((token_num,), dtype=torch.int32, device=b_req_idx.device)
+    mem_index.fill_(-1)
+    if b_req_idx.numel():
+        _prepare_window_write_indices[(b_req_idx.numel(),)](
+            b_req_idx,
+            b_start_loc,
+            b_ready_cache_len,
+            b_commit_len,
+            mem_index,
+            req_to_seq_len,
+            window_size=window_size,
+            BLOCK=256,
+        )
+    return mem_index
+
+
+@triton.jit
+def _pack_window_kv(
+    kv_buffer,
+    noise_kv,
+    b_req_idx,
+    req_to_seq_len,
+    out_kv,
+    stride_kv_bs,
+    stride_kv_h,
+    stride_kv_d,
+    stride_noise_bs,
+    stride_noise_h,
+    stride_noise_d,
+    stride_out_b,
+    stride_out_s,
+    stride_out_h,
+    stride_out_d,
+    window_size: tl.constexpr,
+    block_size: tl.constexpr,
+    head_num: tl.constexpr,
+    head_dim: tl.constexpr,
+    BLOCK_HEAD: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+):
+    cur_token = tl.program_id(0)
+    cur_batch = tl.program_id(1)
+    cur_req_idx = tl.load(b_req_idx + cur_batch).to(tl.int64)
+    cur_seq_len = tl.load(req_to_seq_len + cur_req_idx)
+    cur_cache_len = tl.minimum(cur_seq_len, window_size)
+    cur_position = tl.maximum(0, cur_seq_len - window_size) + cur_token
+    cur_mem_index = cur_req_idx * window_size + cur_position % window_size
+    cur_noise_index = cur_batch * block_size + cur_token - cur_cache_len
+    offs_h = tl.arange(0, BLOCK_HEAD)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    head_mask = (offs_h[:, None] < head_num) & (offs_d[None, :] < head_dim)
+
+    kv = tl.load(
+        kv_buffer + cur_mem_index * stride_kv_bs + offs_h[:, None] * stride_kv_h + offs_d[None, :] * stride_kv_d,
+        mask=(cur_token < cur_cache_len) & head_mask,
+        other=0,
+    )
+    noise = tl.load(
+        noise_kv
+        + cur_noise_index.to(tl.int64) * stride_noise_bs
+        + offs_h[:, None] * stride_noise_h
+        + offs_d[None, :] * stride_noise_d,
+        mask=(cur_token >= cur_cache_len) & (cur_token < cur_cache_len + block_size) & head_mask,
+        other=0,
+    )
+    out_ptr = (
+        out_kv
+        + cur_batch * stride_out_b
+        + cur_token * stride_out_s
+        + offs_h[:, None] * stride_out_h
+        + offs_d[None, :] * stride_out_d
+    )
+    tl.store(out_ptr, tl.where(cur_token < cur_cache_len, kv, noise), mask=head_mask)
+
+
+def pack_window_kv(kv_buffer, noise_kv, b_req_idx, req_to_seq_len, window_size, block_size, alloc_func=torch.empty):
+    head_num, head_dim = kv_buffer.shape[-2:]
+    out_kv = alloc_func(
+        (b_req_idx.numel(), window_size + block_size, head_num, head_dim), dtype=noise_kv.dtype, device=noise_kv.device
+    )
+    if b_req_idx.numel():
+        _pack_window_kv[(window_size + block_size, b_req_idx.numel())](
+            kv_buffer,
+            noise_kv,
+            b_req_idx,
+            req_to_seq_len,
+            out_kv,
+            stride_kv_bs=kv_buffer.stride(0),
+            stride_kv_h=kv_buffer.stride(1),
+            stride_kv_d=kv_buffer.stride(2),
+            stride_noise_bs=noise_kv.stride(0),
+            stride_noise_h=noise_kv.stride(1),
+            stride_noise_d=noise_kv.stride(2),
+            stride_out_b=out_kv.stride(0),
+            stride_out_s=out_kv.stride(1),
+            stride_out_h=out_kv.stride(2),
+            stride_out_d=out_kv.stride(3),
+            window_size=window_size,
+            block_size=block_size,
+            head_num=head_num,
+            head_dim=head_dim,
+            BLOCK_HEAD=triton.next_power_of_2(head_num),
+            BLOCK_DMODEL=triton.next_power_of_2(head_dim),
+            num_warps=4,
+        )
+    return out_kv