From f2c897ebf168f17716e2fac132fc14e2aff5f082 Mon Sep 17 00:00:00 2001 From: Zhao Xintong Date: Fri, 11 Sep 2026 13:14:07 +0000 Subject: [PATCH 1/9] feat(mtp): add sliding-window draft KV cache for dflash style models --- docs/CN/source/tutorial/api_server_args.rst | 22 ++ docs/EN/source/tutorial/api_server_args.rst | 25 +++ lightllm/models/windowed_mtp/__init__.py | 0 lightllm/models/windowed_mtp/kv_store.py | 191 ++++++++++++++++++ lightllm/models/windowed_mtp/layer_infer.py | 17 ++ lightllm/models/windowed_mtp/model.py | 148 ++++++++++++++ lightllm/server/api_cli.py | 8 + lightllm/server/api_start.py | 5 + lightllm/server/core/objs/start_args_type.py | 3 + .../server/router/model_infer/infer_batch.py | 2 +- .../model_infer/mode_backend/base_backend.py | 4 + .../mtp_speculative/proposers/dflash.py | 47 +++-- .../mtp_speculative/proposers/dspark.py | 47 +++-- lightllm/utils/envs_utils.py | 4 + lightllm/utils/profile_max_tokens.py | 9 +- lightllm/utils/windowed_mtp.py | 53 +++++ 16 files changed, 557 insertions(+), 28 deletions(-) create mode 100644 lightllm/models/windowed_mtp/__init__.py create mode 100644 lightllm/models/windowed_mtp/kv_store.py create mode 100644 lightllm/models/windowed_mtp/layer_infer.py create mode 100644 lightllm/models/windowed_mtp/model.py create mode 100644 lightllm/utils/windowed_mtp.py diff --git a/docs/CN/source/tutorial/api_server_args.rst b/docs/CN/source/tutorial/api_server_args.rst index e4fe2a92a7..a060681988 100644 --- a/docs/CN/source/tutorial/api_server_args.rst +++ b/docs/CN/source/tutorial/api_server_args.rst @@ -678,6 +678,28 @@ MTP 多预测参数 * ``eagle_no_att`` * ``None``: 不启用 mtp（默认） +.. option:: --mtp_draft_cache_mode + + 草稿 KV 缓存策略，默认 ``full``。dflash style 草稿模型支持 ``windowed``， + 仅保留开头和最近窗口内的 KV。target 仍使用完整上下文。 + + 当前要求 ``--run_mode normal``、``--dp 1``、未量化 KV，且设置 + ``--disable_dynamic_prompt_cache``。暂不支持 CPU/prefix 缓存、动态验证、 + diverse、混合 prefill/decode 或 overlap 模式。 + +.. option:: --mtp_draft_window + + ``windowed`` 策略保留的最近位置数，默认 ``512``，必须大于 0。 + +.. option:: --mtp_draft_sinks + + ``windowed`` 策略保留的开头位置数，默认 ``1``；设为 ``0`` 可关闭。 + 每个请求的历史容量上限为 ``mtp_draft_window + mtp_draft_sinks``。 + + 示例：在已有 dflash style 草稿模型启动命令中添加:: + + --mtp_draft_cache_mode windowed --mtp_draft_window 512 --mtp_draft_sinks 1 --disable_dynamic_prompt_cache + .. option:: --mtp_draft_model_dir MTP 多预测功能的草稿模型路径 diff --git a/docs/EN/source/tutorial/api_server_args.rst b/docs/EN/source/tutorial/api_server_args.rst index 1d4cc12851..1714937619 100644 --- a/docs/EN/source/tutorial/api_server_args.rst +++ b/docs/EN/source/tutorial/api_server_args.rst @@ -690,6 +690,31 @@ MTP Multi-Prediction Parameters * ``eagle_no_att`` * ``None``: Do not enable mtp (default) +.. option:: --mtp_draft_cache_mode + + Draft KV cache policy; defaults to ``full``. dflash style draft models support + ``windowed``, which retains only initial and recent-window KV. The target still + uses its full context. + + Currently requires ``--run_mode normal``, ``--dp 1``, unquantized KV, and + ``--disable_dynamic_prompt_cache``. CPU/prefix caching, dynamic verification, + diverse, mixed prefill/decode, and overlap modes are not supported. + +.. option:: --mtp_draft_window + + Number of recent positions retained by ``windowed``; defaults to ``512`` and + must be positive. + +.. option:: --mtp_draft_sinks + + Number of initial positions retained by ``windowed``; defaults to ``1``. + Set to ``0`` to disable sinks. Per-request history capacity is bounded by + ``mtp_draft_window + mtp_draft_sinks``. + + Add these options to an existing launch command for a dflash style draft model:: + + --mtp_draft_cache_mode windowed --mtp_draft_window 512 --mtp_draft_sinks 1 --disable_dynamic_prompt_cache + .. option:: --mtp_draft_model_dir Path to the draft model for MTP multi-prediction functionality diff --git a/lightllm/models/windowed_mtp/__init__.py b/lightllm/models/windowed_mtp/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/models/windowed_mtp/kv_store.py b/lightllm/models/windowed_mtp/kv_store.py new file mode 100644 index 0000000000..bedccb1401 --- /dev/null +++ b/lightllm/models/windowed_mtp/kv_store.py @@ -0,0 +1,191 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _gather_new( + FEATURES, + REQS, + STARTS, + FIRST, + LENGTHS, + OUT, + POS, + ENDS, + COUNTS, + N: tl.constexpr, + H: tl.constexpr, + C: tl.constexpr, + W: tl.constexpr, + S: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + batch, offset = row // N, row % N + req = tl.load(REQS + batch).to(tl.int64) + first = tl.load(FIRST + batch).to(tl.int64) + length = tl.load(LENGTHS + batch) + end = first + length + if W > 0: + sinks = tl.minimum(tl.maximum(S - first, 0), length) + recent = tl.maximum(tl.maximum(first, S), end - W) + count = sinks + tl.maximum(end - recent, 0) + position = tl.where(offset < sinks, first + offset, recent + offset - sinks) + else: + count = length + position = first + offset + valid = offset < count + dim = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + source = tl.load(STARTS + batch).to(tl.int64) + position - first + value = tl.load(FEATURES + source * H + dim, valid & (dim < H), other=0) + tl.store(OUT + row * H + dim, value, dim < H) + if tl.program_id(1) == 0: + tl.store(POS + row, tl.where(valid, position, -1)) + if offset == 0: + tl.store(ENDS + req, end) + tl.store(COUNTS + req, tl.minimum(end, C)) + + +@triton.jit +def _write_ring( + POOL, + NEW, + REQS, + POS, + N: tl.constexpr, + C: tl.constexpr, + W: tl.constexpr, + S: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + req = tl.load(REQS + row // N).to(tl.int64) + position = tl.load(POS + row) + if W > 0: + slot = tl.where(position < S, position, S + (position - S) % W) + else: + slot = position + dim = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + valid = (position >= 0) & (dim < WIDTH) + value = tl.load(NEW + row * WIDTH + dim, valid, other=0) + tl.store(POOL + (req * C + slot) * WIDTH + dim, value, valid) + + +@triton.jit +def _pack_ring( + POOL, + NOISE, + REQS, + ENDS, + OUT, + C: tl.constexpr, + B: tl.constexpr, + W: tl.constexpr, + S: tl.constexpr, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + batch, slot = row // (C + B), row % (C + B) + req = tl.load(REQS + batch).to(tl.int64) + end = tl.load(ENDS + req) + count = tl.minimum(end, C) + if W > 0: + sinks = tl.minimum(end, S) + recent = tl.maximum(S, end - W) + position = tl.where(slot < sinks, slot, recent + slot - sinks) + ring_slot = tl.where(position < S, position, S + (position - S) % W) + else: + ring_slot = slot + dim = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + context = tl.load(POOL + (req * C + ring_slot) * WIDTH + dim, (slot < count) & (dim < WIDTH), other=0) + noise = tl.load( + NOISE + (batch * B + slot - count) * WIDTH + dim, (slot >= count) & (slot < count + B) & (dim < WIDTH), other=0 + ) + tl.store(OUT + row * WIDTH + dim, context + noise, dim < WIDTH) + + +class WindowKVStore: + """Per-layer sink + recent-window K/V, indexed independently of target KV. + + Only newly accepted positions are projected and written. Each window slot + has at most one writer, including prefill chunks larger than the window. + On request reuse, the new end/count hides stale slots; all visible positions + are overwritten by the first prefill. Draft noise K/V never enters the ring. + """ + + def __init__(self, requests, capacity, layers, kv_heads, head_dim, dtype, device, window, sinks): + self.capacity, self.window, self.sinks = capacity, window, sinks + self.kv = torch.zeros((layers, requests, capacity, 2 * kv_heads, head_dim), dtype=dtype, device=device) + self.ends = torch.zeros(requests, dtype=torch.int64, device=device) + self.counts = torch.zeros(requests, dtype=torch.int32, device=device) + + def prepare(self, reqs, features, starts, first, lengths, max_new): + n = min(max_new, self.capacity) + packed = features.new_empty((reqs.numel() * n, features.shape[-1])) + positions = torch.empty(reqs.numel() * n, dtype=torch.int64, device=features.device) + _gather_new[(reqs.numel() * n, triton.cdiv(features.shape[-1], 512))]( + features, + reqs, + starts, + first, + lengths, + packed, + positions, + self.ends, + self.counts, + n, + features.shape[-1], + self.capacity, + self.window, + self.sinks, + 512, + ) + return packed, positions + + def write(self, layer, reqs, positions, new_kv): + width = new_kv.shape[-2] * new_kv.shape[-1] + _write_ring[(positions.numel(), triton.cdiv(width, 512))]( + self.kv[layer], + new_kv, + reqs, + positions, + positions.numel() // reqs.numel(), + self.capacity, + self.window, + self.sinks, + width, + 512, + ) + + def pack(self, layer, reqs, noise_kv, block): + heads, dim = self.kv.shape[-2:] + output = noise_kv.new_empty((reqs.numel(), self.capacity + block, heads, dim)) + _pack_ring[(reqs.numel() * (self.capacity + block), triton.cdiv(heads * dim, 512))]( + self.kv[layer], + noise_kv, + reqs, + self.ends, + output, + self.capacity, + block, + self.window, + self.sinks, + heads * dim, + 512, + ) + return output + + def retained_positions(self, reqs): + """Diagnostic only: logical positions, in attention order.""" + end = self.ends.index_select(0, reqs)[:, None] + slots = torch.arange(self.capacity, device=reqs.device)[None, :] + if self.window: + sinks = end.clamp_max(self.sinks) + recent = (end - self.window).clamp_min(self.sinks) + positions = torch.where(slots < sinks, slots, recent + slots - sinks) + else: + positions = slots.expand(reqs.numel(), -1) + return positions.masked_fill(slots >= end.clamp_max(self.capacity), -1) diff --git a/lightllm/models/windowed_mtp/layer_infer.py b/lightllm/models/windowed_mtp/layer_infer.py new file mode 100644 index 0000000000..735d9c021b --- /dev/null +++ b/lightllm/models/windowed_mtp/layer_infer.py @@ -0,0 +1,17 @@ +from lightllm.utils.sgl_utils import flash_attn_with_kvcache + + +class WindowedAttentionMixin: + def token_attention_forward(self, input_embdings, infer_state, layer_weight): + q, noise_kv = self._get_qkv(input_embdings, infer_state, layer_weight) + store, reqs, counts, block = infer_state.windowed_context + scratch = store.pack(self.windowed_layer_index, reqs, noise_kv, block) + output = flash_attn_with_kvcache( + q.view(-1, block, self.tp_q_head_num_, self.head_dim_), + scratch[:, :, : self.tp_k_head_num_], + scratch[:, :, self.tp_k_head_num_ :], + cache_seqlens=counts + block, + causal=False, + softmax_scale=self.head_dim_ ** -0.5, + ) + return self._get_o(output.reshape(-1, self.tp_q_head_num_ * self.head_dim_), infer_state, layer_weight) diff --git a/lightllm/models/windowed_mtp/model.py b/lightllm/models/windowed_mtp/model.py new file mode 100644 index 0000000000..ca721539e8 --- /dev/null +++ b/lightllm/models/windowed_mtp/model.py @@ -0,0 +1,148 @@ +import torch + +from lightllm.common.basemodel.triton_kernel.norm.qk_norm import qk_rmsnorm_forward +from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd +from lightllm.models.qwen3_dflash.model import Qwen3DFlashModel +from lightllm.models.windowed_mtp.kv_store import WindowKVStore +from lightllm.models.windowed_mtp.layer_infer import WindowedAttentionMixin +from lightllm.utils.windowed_mtp import window_capacity, validate_windowed_mtp +from lightllm.common.basemodel.attention.base_att import BaseAttBackend, BaseDecodeAttState +from lightllm.utils.log_utils import init_logger +from lightllm.utils.torch_memory_saver_utils import MemoryTag + + +logger = init_logger(__name__) + + +class WindowedDecodeState(BaseDecodeAttState): + def init_state(self): + pass + + def decode_att(self, *args, **kwargs): + raise RuntimeError("Windowed draft uses bounded FA3 attention directly") + + +class WindowedAttBackend(BaseAttBackend): + def create_att_decode_state(self, infer_state): + return WindowedDecodeState(backend=self, infer_state=infer_state) + + +class WindowedDraftMixin: + uses_windowed_draft_kv = True + + def _verify_params(self): + super()._verify_params() + validate_windowed_mtp(self.args) + + def _init_custom(self): + super()._init_custom() + # _init_custom runs after the base model's KV allocation region. + with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): + self.kv_store = WindowKVStore( + self.args.running_max_req_size + 1, + window_capacity(self.args), + self.config["n_layer"], + self.tp_k_head_num_, + self.head_dim_, + self.data_type, + self._cos_cached.device, + self.args.mtp_draft_window, + self.args.mtp_draft_sinks, + ) + for i, layer in enumerate(self.layers_infer): + layer.windowed_layer_index = i + store = self.kv_store + pool_bytes = sum(t.numel() * t.element_size() for t in (store.kv, store.ends, store.counts)) + logger.info( + f"{self.args.mtp_mode} windowed: full-context draft KV layers=0, " + f"window KV layers={self.config['n_layer']}, " + f"capacity={store.capacity}, KV pool bytes={pool_bytes}, target KV layers={len(self.mem_manager.kv_buffer)}" + ) + + @torch.no_grad() + def commit_features(self, model_input, hidden, starts=None, accept_len=None): + # The fused feature is transient; only its per-layer K/V is retained. + projected = self.pre_post_weight.fc_weight_.mm(hidden, use_custom_tensor_mananger=False) + projected = self.pre_post_weight.hidden_norm_weight_(projected, eps=self.pre_infer.eps_, alloc_func=torch.empty) + if model_input.is_prefill: + reqs = model_input.b_req_idx.long() + starts = model_input.b_prefill_start_loc + first = model_input.b_ready_cache_len + lengths = model_input.b_seq_len - first + max_new = model_input.max_q_seq_len + else: + reqs = model_input.b_req_idx.index_select(0, starts.long()).long() + first = model_input.b_seq_len.index_select(0, starts.long()) - 1 + lengths = accept_len + max_new = self.args.mtp_step + 1 + features, positions = self.kv_store.prepare(reqs, projected, starts, first, lengths, max_new) + cos = self._cos_cached.index_select(0, positions.clamp_min(0)) + sin = self._sin_cached.index_select(0, positions.clamp_min(0)) + for i, (layer, weight) in enumerate(zip(self.layers_infer, self.trans_layers_weight)): + kv = weight.kv_proj.mm(features, use_custom_tensor_mananger=False) + qk_rmsnorm_forward( + kv[:, : self.tp_k_head_num_ * self.head_dim_], weight.qk_norm_weight_.k_weight, layer.eps_ + ) + kv = kv.view(-1, 2 * self.tp_k_head_num_, self.head_dim_) + rotary_emb_fwd( + kv[:, : self.tp_k_head_num_], None, cos, sin, partial_rotary_factor=layer.partial_rotary_factor + ) + self.kv_store.write(i, reqs, positions, kv) + + def _token_forward(self, infer_state): + # These reads run inside CUDA Graph, using the current request IDs/state. + reqs = infer_state.b_req_idx[:: self.block_size].long() + infer_state.windowed_context = ( + self.kv_store, + reqs, + self.kv_store.counts.index_select(0, reqs), + self.block_size, + ) + try: + return super()._token_forward(infer_state) + finally: + del infer_state.windowed_context + + def _init_att_backend(self): + self.prefill_att_backend = self.decode_att_backend = WindowedAttBackend(model=self) + + def _decode(self, model_input): + # Window KV is committed by the proposer. Draft block rows never own + # target cache slots and must not overwrite its request-to-token table. + origin_batch_size = model_input.batch_size + assert origin_batch_size > 0 and origin_batch_size % self.block_size == 0 + infer_batch_size = origin_batch_size + use_cuda_graph = self.graph is not None and self.graph.can_run( + batch_size=infer_batch_size, + max_len_in_batch=max(2, model_input.max_kv_seq_len), + ) + need_capture = False + if use_cuda_graph: + infer_batch_size = self.graph.find_closest_graph_batch_size(batch_size=infer_batch_size) + need_capture = self.graph.need_capture(infer_batch_size) + model_input = self._create_padded_decode_model_input(model_input, infer_batch_size) + infer_state = self._create_inferstate(model_input) + infer_state.is_cuda_graph = need_capture + infer_state.init_some_extra_state(self) + infer_state.init_att_state() + if use_cuda_graph: + if need_capture: + output = self.graph.capture_decode(self._token_forward, infer_state) + else: + output = self.graph.replay(infer_state) + else: + output = self._token_forward(infer_state) + return self._create_unpad_decode_model_output(output, origin_batch_size=origin_batch_size) + + +def windowed_model_class(original): + if not issubclass(original, Qwen3DFlashModel): + raise ValueError("Windowed draft KV currently supports the Qwen3 parallel block backbone") + layer_class = type( + "Windowed" + original.transformer_layer_infer_class.__name__, + (WindowedAttentionMixin, original.transformer_layer_infer_class), + {}, + ) + return type( + "Windowed" + original.__name__, (WindowedDraftMixin, original), {"transformer_layer_infer_class": layer_class} + ) diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index e6c077dbe4..0927cf472d 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -767,6 +767,14 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: action="store_true", help="""Whether to enable fused shared experts for supported MoE models. It is auto-enabled when supported.""", ) + parser.add_argument( + "--mtp_draft_cache_mode", + choices=["full", "windowed"], + default="full", + help="Draft KV retention policy for parallel block MTP", + ) + parser.add_argument("--mtp_draft_window", type=int, default=512) + parser.add_argument("--mtp_draft_sinks", type=int, default=1) parser.add_argument( "--mtp_mode", choices=[ diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index e0997f971f..5b2db7eefd 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -163,6 +163,11 @@ def _launch_subprocesses(args: StartArgs): f"{sorted(allowed_ep_decode_att_backends)}; flashinfer is not supported." ) + if args.mtp_draft_cache_mode == "windowed": + from lightllm.utils.windowed_mtp import validate_windowed_mtp + + validate_windowed_mtp(args) + # mtp params check if args.mtp_mode is not None: if args.mtp_draft_model_dir is None: diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 9c89975de7..9268045808 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -199,6 +199,9 @@ class StartArgs: }, ) mtp_draft_model_dir: Optional[List[str]] = field(default=None) + mtp_draft_cache_mode: str = field(default="full", metadata={"choices": ["full", "windowed"]}) + mtp_draft_window: int = field(default=512) + mtp_draft_sinks: int = field(default=1) mtp_step: int = field(default=0) mtp_dynamic_verify: bool = field(default=False) kv_quant_calibration_config_path: Optional[str] = field(default=None) diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 18a8f6c04f..dd37dace9e 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -392,7 +392,7 @@ def get_can_alloc_token_num(self): def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferReq"]): """Snapshot request-level attention state at big/small-page boundaries.""" - if not self.is_hybrid_att_model: + if not self.is_hybrid_att_model or self.radix_cache is None: return # Request-state snapshot at a big-page boundary. diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index fe77e5f37b..a3ac168060 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -342,6 +342,10 @@ def init_mtp_draft_model(self, main_kvargs: dict): model_cfg=draft_model_cfg, spec_mode=spec_mode, ) + if self.args.mtp_draft_cache_mode == "windowed": + from lightllm.models.windowed_mtp.model import windowed_model_class + + draft_model_class = windowed_model_class(draft_model_class) self.draft_models.append(draft_model_class(draft_model_kvargs)) self.logger.info(f"loaded speculative draft model class {self.draft_models[i].__class__}") diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py index 23f81f546a..cb95c4695f 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py @@ -43,6 +43,9 @@ def fill_draft_model_kv_state( # DFlash prefill 直接复用 target prompt 的 token 布局，并注入 target # hidden 初始化唯一 draft model 的 KV。使用浅副本避免在 target 输入上 # 保留 draft 专用状态。 + if getattr(self.backend.draft_models[0], "uses_windowed_draft_kv", False) is True: + self.backend.draft_models[0].commit_features(target_model_input, target_hidden) + return draft_input = copy.copy(target_model_input) draft_input.mtp_draft_input_hiddens = target_hidden self.backend.draft_models[0].forward(draft_input) @@ -82,17 +85,31 @@ def propose_next( accepted_tail_rows = (b_req_mtp_start_loc + accept_len - 1).long() - # target verify 的行布局和 mem_indexes 对应本轮所有被验证 token。 - # 附加 target hidden 后执行一次 draft forward，将这些行提交到 DFlash - # KV cache；浅副本保证 target_model_input 本身保持不变。 - verify_draft_input = copy.copy(target_model_input) - verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden - draft_model.forward(verify_draft_input) - - # 每个请求始终展开完整 block，未被本轮 proposal 返回的 block 尾部仍会 - # 参与 parallel forward。所有临时 KV slot 在 verify 后通过 proposal - # 统一释放。 - extra_mem_indexes_cpu = mtp_utils.alloc_mem_indexes(req_num * block_size) + # full 模式通过 draft forward 写入 verify 行对应的 KV；windowed 模式 + # 只将接受的 target hidden 投影并写入独立窗口，不保留被拒绝的行。 + windowed = getattr(draft_model, "uses_windowed_draft_kv", False) is True + if windowed: + draft_model.commit_features( + target_model_input, target_model_output.mtp_collector.spec_hidden, b_req_mtp_start_loc, accept_len + ) + else: + verify_draft_input = copy.copy(target_model_input) + verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden + draft_model.forward(verify_draft_input) + + # 每个请求始终展开完整 block。full 模式分配的临时 KV slot 在 verify + # 后通过 proposal 释放；windowed 模式不占用 target KV slot。 + if windowed: + extra_mem_indexes_cpu = None + scratch_indexes = torch.full( + (req_num * block_size,), + draft_model.mem_manager.HOLD_TOKEN_MEMINDEX, + dtype=torch.int32, + device=target_next_token_ids.device, + ) + else: + extra_mem_indexes_cpu = mtp_utils.alloc_mem_indexes(req_num * block_size) + scratch_indexes = extra_mem_indexes_cpu.cuda(non_blocking=True) block_input_ids = target_next_token_ids.new_full( (req_num * block_size,), fill_value=draft_model.mask_token_id, @@ -140,7 +157,7 @@ def propose_next( .repeat_interleave(block_size) .contiguous() ) - draft_input.mem_indexes = extra_mem_indexes_cpu.cuda(non_blocking=True) + draft_input.mem_indexes = scratch_indexes draft_input.mem_indexes_cpu = None draft_input.multimodal_params = [{"images": [], "audios": []} for _ in range(draft_input.batch_size)] draft_output = draft_model.forward(draft_input) @@ -160,6 +177,10 @@ def propose_next( schedule_scores = block_draft_token_probs[:, :draft_step].float().contiguous() return DFlashSpecProposal( token_ids=proposal_token_ids, - extra_mem_indexes_cpu=[MtpMemIndexesToFree(mem_indexes_cpu=extra_mem_indexes_cpu)], + extra_mem_indexes_cpu=( + [MtpMemIndexesToFree(mem_indexes_cpu=extra_mem_indexes_cpu)] + if extra_mem_indexes_cpu is not None + else [] + ), schedule_scores=schedule_scores, ) diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py index 5e3ea5694f..2eb8af3312 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py @@ -44,6 +44,9 @@ def fill_draft_model_kv_state( # DSpark prefill 直接使用 target prompt 的 token 布局和 hidden，将 prompt # KV 写入唯一的 parallel-block draft model。使用浅副本，避免把 draft # 专用 hidden 挂到后续流程仍可能读取的 target ModelInput 上。 + if getattr(self.backend.draft_models[0], "uses_windowed_draft_kv", False) is True: + self.backend.draft_models[0].commit_features(target_model_input, target_hidden) + return draft_input = copy.copy(target_model_input) draft_input.mtp_draft_input_hiddens = target_hidden self.backend.draft_models[0].forward(draft_input) @@ -85,17 +88,31 @@ def propose_next( accepted_tail_rows = (b_req_mtp_start_loc + accept_len - 1).long() - # target verify 的行布局和 mem_indexes 已经对应本轮所有被验证 token。 - # 仅附加 target hidden 后执行一次 draft forward，即可把这些行提交到 - # DSpark KV cache；浅副本保证 target_model_input 本身不被修改。 - verify_draft_input = copy.copy(target_model_input) - verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden - draft_model.forward(verify_draft_input) - - # DSpark 每个请求固定展开一个完整 block，临时 KV 在 target verify 完成 - # 后通过 proposal 统一释放。block 第一行是 accepted-tail anchor，其余行 - # 使用 mask token，由 parallel backbone 一次并行计算。 - extra_mem_indexes_cpu = mtp_utils.alloc_mem_indexes(req_num * block_size) + # full 模式通过 draft forward 写入 verify 行对应的 KV；windowed 模式 + # 只将接受的 target hidden 投影并写入独立窗口，不保留被拒绝的行。 + windowed = getattr(draft_model, "uses_windowed_draft_kv", False) is True + if windowed: + draft_model.commit_features( + target_model_input, target_model_output.mtp_collector.spec_hidden, b_req_mtp_start_loc, accept_len + ) + else: + verify_draft_input = copy.copy(target_model_input) + verify_draft_input.mtp_draft_input_hiddens = target_model_output.mtp_collector.spec_hidden + draft_model.forward(verify_draft_input) + + # 每个请求固定展开完整 block：首行是 accepted-tail anchor，其余为 mask。 + # full 模式的临时 KV 在 verify 后释放；windowed 模式不占用 target KV slot。 + if windowed: + extra_mem_indexes_cpu = None + scratch_indexes = torch.full( + (req_num * block_size,), + draft_model.mem_manager.HOLD_TOKEN_MEMINDEX, + dtype=torch.int32, + device=target_next_token_ids.device, + ) + else: + extra_mem_indexes_cpu = mtp_utils.alloc_mem_indexes(req_num * block_size) + scratch_indexes = extra_mem_indexes_cpu.cuda(non_blocking=True) block_input_ids = target_next_token_ids.new_full( (req_num * block_size,), fill_value=draft_model.mask_token_id, @@ -143,7 +160,7 @@ def propose_next( .repeat_interleave(block_size) .contiguous() ) - draft_input.mem_indexes = extra_mem_indexes_cpu.cuda(non_blocking=True) + draft_input.mem_indexes = scratch_indexes draft_input.mem_indexes_cpu = None draft_input.multimodal_params = [{"images": [], "audios": []} for _ in range(draft_input.batch_size)] draft_output = draft_model.forward(draft_input) @@ -184,7 +201,11 @@ def propose_next( return DSparkSpecProposal( token_ids=proposal_token_ids, - extra_mem_indexes_cpu=[MtpMemIndexesToFree(mem_indexes_cpu=extra_mem_indexes_cpu)], + extra_mem_indexes_cpu=( + [MtpMemIndexesToFree(mem_indexes_cpu=extra_mem_indexes_cpu)] + if extra_mem_indexes_cpu is not None + else [] + ), schedule_scores=schedule_scores, schedule_scores_cpu=schedule_scores_cpu, ) diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index acd2ac6711..e1f93de96b 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -255,6 +255,8 @@ def enable_cpu_cache_numa_interleave() -> bool: @lru_cache(maxsize=None) def get_added_mtp_kv_layer_num() -> int: args = get_env_start_args() + if getattr(args, "mtp_draft_cache_mode", "full") == "windowed": + return 0 mtp_mode = args.mtp_mode if mtp_mode is None: @@ -288,6 +290,8 @@ def get_mtp_weight_layer_num() -> int: return args.mtp_step if mtp_mode == "eagle_no_att": return 1 + if getattr(args, "mtp_draft_cache_mode", "full") == "windowed": + return _get_mtp_draft_backbone_layer_num(args.mtp_draft_model_dir[0]) return get_added_mtp_kv_layer_num() diff --git a/lightllm/utils/profile_max_tokens.py b/lightllm/utils/profile_max_tokens.py index 5ec2a9145e..109107b9c3 100644 --- a/lightllm/utils/profile_max_tokens.py +++ b/lightllm/utils/profile_max_tokens.py @@ -38,10 +38,11 @@ def get_mtp_adjusted_mem_fraction( target_weight_bytes: int, target_layer_num: int, mtp_layer_num: int, + draft_cache_bytes: int = 0, ) -> float: mtp_weight_bytes = target_weight_bytes * mtp_layer_num / target_layer_num total_gpu_bytes = torch.cuda.get_device_properties(get_current_device_id()).total_memory - adjusted_mem_fraction = mem_fraction - mtp_weight_bytes / total_gpu_bytes + adjusted_mem_fraction = mem_fraction - (mtp_weight_bytes + draft_cache_bytes) / total_gpu_bytes # 不同 rank 的权重分片大小和 GPU 总显存可能不同。取全局最小值，保证所有 # rank 使用相同且能够安全预留 MTP 权重显存的 KV cache 比例。 @@ -74,11 +75,17 @@ def profile_mtp_weight_memory(model): weight_memory_before = torch.cuda.memory_allocated() yield target_weight_bytes = torch.cuda.memory_allocated() - weight_memory_before + draft_cache_bytes = 0 + if getattr(model.args, "mtp_draft_cache_mode", "full") == "windowed": + from lightllm.utils.windowed_mtp import window_kv_pool_bytes + + draft_cache_bytes = window_kv_pool_bytes(model.args, torch.tensor([], dtype=model.data_type).element_size()) model.mem_fraction = get_mtp_adjusted_mem_fraction( mem_fraction=model.mem_fraction, target_weight_bytes=target_weight_bytes, target_layer_num=model.config["n_layer"], mtp_layer_num=get_mtp_weight_layer_num(), + draft_cache_bytes=draft_cache_bytes, ) diff --git a/lightllm/utils/windowed_mtp.py b/lightllm/utils/windowed_mtp.py new file mode 100644 index 0000000000..656f601d04 --- /dev/null +++ b/lightllm/utils/windowed_mtp.py @@ -0,0 +1,53 @@ +"""Opt-in bounded draft KV policy, independent of the proposal algorithm.""" +import json +from pathlib import Path + +from lightllm.common.build_utils import repair_config + + +def window_capacity(args): + if args.mtp_draft_window <= 0 or args.mtp_draft_sinks < 0: + raise ValueError("mtp_draft_window must be positive and mtp_draft_sinks nonnegative") + return args.mtp_draft_window + args.mtp_draft_sinks + + +def window_kv_pool_bytes(args, element_size): + config = json.loads((Path(args.mtp_draft_model_dir[0]) / "config.json").read_text()) + # Match the existing draft initialization order: base aliases, Llama KV + # defaults, then the overrides applied by Qwen3.5 adapters. + # Do not normalize again after merging: runtime also retains those aliases. + merge_nested = config.get("model_type") in ("qwen3_5", "qwen3_5_text") + repair_config(config, same_names=["num_attention_heads", "n_head"]) + repair_config(config, same_names=["hidden_size", "n_embd", "n_embed"]) + repair_config(config, same_names=["num_hidden_layers", "n_layer"]) + config.setdefault("num_key_value_heads", config["num_attention_heads"]) + if merge_nested: + config.update(config.get("dflash_config", {})) + layers = config["n_layer"] + heads = config["num_key_value_heads"] // args.tp + dim = config.get("head_dim", config["n_embed"] // config["num_attention_heads"]) + return (args.running_max_req_size + 1) * (layers * window_capacity(args) * 2 * heads * dim * element_size + 12) + + +def validate_windowed_mtp(args): + window_capacity(args) + if args.mtp_mode not in ("dflash", "dspark"): + raise ValueError("windowed draft KV requires dflash or dspark") + if args.run_mode != "normal" or args.dp != 1: + raise ValueError("windowed draft KV currently requires normal deployment and dp=1") + if not args.disable_dynamic_prompt_cache: + raise ValueError("windowed draft KV requires --disable_dynamic_prompt_cache") + if args.llm_kv_type not in (None, "None"): + raise ValueError("windowed draft KV currently requires unquantized KV") + for option in ( + "enable_cpu_cache", + "enable_dp_prompt_cache_fetch", + "diverse_mode", + "enable_tpsp_mix_mode", + "enable_prefill_decode_mixed", + "enable_decode_microbatch_overlap", + "enable_prefill_microbatch_overlap", + "mtp_dynamic_verify", + ): + if getattr(args, option, False): + raise ValueError(f"windowed draft KV does not yet support --{option}") From cbccaf9c59f6ae3452a3cb797a889377eebe52e3 Mon Sep 17 00:00:00 2001 From: Zhao Xintong Date: Mon, 14 Sep 2026 08:04:16 +0000 Subject: [PATCH 2/9] refactor --- .../basemodel/attention/fa3/windowed_mtp.py | 31 ++++ lightllm/common/basemodel/basemodel.py | 45 ++---- .../kv_cache_mem_manager/mem_manager.py | 13 ++ .../kv_cache_mem_manager/windowed_mtp.py} | 7 + .../layer_infer/transformer_layer_infer.py | 27 +++- lightllm/models/qwen3_dflash/model.py | 56 ++++++- lightllm/models/windowed_mtp/__init__.py | 0 lightllm/models/windowed_mtp/layer_infer.py | 17 -- lightllm/models/windowed_mtp/model.py | 148 ------------------ .../model_infer/mode_backend/base_backend.py | 4 - 10 files changed, 145 insertions(+), 203 deletions(-) create mode 100644 lightllm/common/basemodel/attention/fa3/windowed_mtp.py rename lightllm/{models/windowed_mtp/kv_store.py => common/kv_cache_mem_manager/windowed_mtp.py} (96%) delete mode 100644 lightllm/models/windowed_mtp/__init__.py delete mode 100644 lightllm/models/windowed_mtp/layer_infer.py delete mode 100644 lightllm/models/windowed_mtp/model.py diff --git a/lightllm/common/basemodel/attention/fa3/windowed_mtp.py b/lightllm/common/basemodel/attention/fa3/windowed_mtp.py new file mode 100644 index 0000000000..ad613b30de --- /dev/null +++ b/lightllm/common/basemodel/attention/fa3/windowed_mtp.py @@ -0,0 +1,31 @@ +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=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].long() + # The base attention-state copy refreshes these tensors on graph replay. + store = self.infer_state.mem_manager.windowed_draft_kv + self.cache_seqlens = store.counts.index_select(0, self.b_req_idx) + self.block_size + + 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, + softmax_scale=q.shape[-1] ** -0.5, + ).view(q.shape) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index f1247e0ef4..baf3fd2fc2 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -555,6 +555,15 @@ def _prefill( model_output.prefill_mem_indexes_ready_event = prefill_mem_indexes_ready_event return model_output + def _update_decode_req_to_token(self, infer_state: InferStateInfo): + """Record decode KV slots before initializing attention state.""" + copy_kv_index_to_req( + self.req_manager.req_to_token_indexs, + infer_state.b_req_idx, + infer_state.b_seq_len, + infer_state.mem_index, + ) + def _decode( self, model_input: ModelInput, @@ -599,12 +608,7 @@ def _decode( # attention backend 会根据该标记准备 CUDA Graph capture 专用状态， # 因此必须在 init_att_state 之前完成赋值。 infer_state.is_cuda_graph = need_capture - copy_kv_index_to_req( - self.req_manager.req_to_token_indexs, - infer_state.b_req_idx, - infer_state.b_seq_len, - infer_state.mem_index, - ) + self._update_decode_req_to_token(infer_state) infer_state.init_some_extra_state(self) infer_state.init_att_state() @@ -620,7 +624,6 @@ def _decode( @final def _context_forward(self, infer_state: InferStateInfo): - input_embs = self.pre_infer.context_forward(infer_state.input_ids, infer_state, self.pre_post_weight) if self.args.enable_dp_prefill_balance: assert not self.args.enable_prefill_cudagraph, "not support now" @@ -846,23 +849,13 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 padded_model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size) infer_state0 = self._create_inferstate(padded_model_input0, 0) infer_state0.is_cuda_graph = need_capture - copy_kv_index_to_req( - self.req_manager.req_to_token_indexs, - infer_state0.b_req_idx, - infer_state0.b_seq_len, - infer_state0.mem_index, - ) + self._update_decode_req_to_token(infer_state0) infer_state0.init_some_extra_state(self) infer_state0.init_att_state() infer_state1 = self._create_inferstate(padded_model_input1, 1) infer_state1.is_cuda_graph = need_capture - copy_kv_index_to_req( - self.req_manager.req_to_token_indexs, - infer_state1.b_req_idx, - infer_state1.b_seq_len, - infer_state1.mem_index, - ) + self._update_decode_req_to_token(infer_state1) infer_state1.init_some_extra_state(self) infer_state1.init_att_state() @@ -884,22 +877,12 @@ def _microbatch_overlap_decode_cuda(self, model_input0: ModelInput, model_input1 model_input0 = self._create_padded_decode_model_input(model_input0, infer_batch_size) model_input1 = self._create_padded_decode_model_input(model_input1, infer_batch_size) infer_state0 = self._create_inferstate(model_input0, 0) - copy_kv_index_to_req( - self.req_manager.req_to_token_indexs, - infer_state0.b_req_idx, - infer_state0.b_seq_len, - infer_state0.mem_index, - ) + self._update_decode_req_to_token(infer_state0) infer_state0.init_some_extra_state(self) infer_state0.init_att_state() infer_state1 = self._create_inferstate(model_input1, 1) - copy_kv_index_to_req( - self.req_manager.req_to_token_indexs, - infer_state1.b_req_idx, - infer_state1.b_seq_len, - infer_state1.mem_index, - ) + self._update_decode_req_to_token(infer_state1) infer_state1.init_some_extra_state(self) infer_state1.init_att_state() diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 9171566599..b179857d2a 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -8,6 +8,7 @@ from lightllm.utils.log_utils import init_logger from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt from .allocator import KvCacheAllocator +from .windowed_mtp import WindowKVStore from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory from lightllm.utils.dist_utils import ( get_current_device_id, @@ -38,6 +39,8 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False self.layer_num = layer_num self.always_copy = always_copy self.dtype = dtype + self.windowed_draft_kv = None + self._windowed_draft_kv_params = None # profile the max total token num if the size is None self.profile_size(mem_fraction) @@ -112,6 +115,13 @@ def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): # 等模式下 padding 一些请求，使推理过程可以正常运行采用的，其索引值为size，存储在HOLD_TOKEN_MEMINDEX # 成员变量中，其与 req_manager 中的HOLD_REQUEST_ID具有类似的作用和意义。 self.kv_buffer = torch.empty((layer_num, size + 1, 2 * head_num, head_dim), dtype=dtype, device="cuda") + if self._windowed_draft_kv_params is not None: + self.init_windowed_draft_kv(**self._windowed_draft_kv_params) + + def init_windowed_draft_kv(self, **cache_params): + """Create draft window storage from dimensions supplied by the draft model.""" + self._windowed_draft_kv_params = cache_params + self.windowed_draft_kv = WindowKVStore(**cache_params, dtype=self.dtype, device=self.kv_buffer.device) def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: self.kv_move_buffer = torch.empty( @@ -188,6 +198,7 @@ def read_page_kv_move_buffer_to_mem( def _free_buffers(self): self.kv_buffer = None + self.windowed_draft_kv = None def alloc(self, need_size) -> torch.Tensor: return self.allocator.alloc(need_size) @@ -197,6 +208,8 @@ def free(self, free_index: Union[torch.Tensor, List[int]]) -> None: def free_all(self): self.allocator.free_all() + if self.windowed_draft_kv is not None: + self.windowed_draft_kv.reset() def resize_mem(self, new_size): """ diff --git a/lightllm/models/windowed_mtp/kv_store.py b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py similarity index 96% rename from lightllm/models/windowed_mtp/kv_store.py rename to lightllm/common/kv_cache_mem_manager/windowed_mtp.py index bedccb1401..88296b85c4 100644 --- a/lightllm/models/windowed_mtp/kv_store.py +++ b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py @@ -1,3 +1,5 @@ +"""Bounded KV storage for parallel-block speculative decoding.""" + import torch import triton import triton.language as tl @@ -122,6 +124,11 @@ def __init__(self, requests, capacity, layers, kv_heads, head_dim, dtype, device self.ends = torch.zeros(requests, dtype=torch.int64, device=device) self.counts = torch.zeros(requests, dtype=torch.int32, device=device) + def reset(self): + """Invalidate all request windows without changing captured buffer addresses.""" + self.ends.zero_() + self.counts.zero_() + def prepare(self, reqs, features, starts, first, lengths, max_new): n = min(max_new, self.capacity) packed = features.new_empty((reqs.numel() * n, features.shape[-1])) diff --git a/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py index 84ee0b4059..4573a43944 100644 --- a/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py @@ -1,5 +1,7 @@ import torch +from lightllm.common.basemodel.attention.fa3.windowed_mtp import WindowedMTPDecodeAttState + from lightllm.common.basemodel.triton_kernel.norm.qk_norm import qk_rmsnorm_forward from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd @@ -26,6 +28,11 @@ def context_forward( infer_state: Qwen3DFlashInferStateInfo, layer_weight: Qwen3DFlashTransformerLayerWeight, ) -> torch.Tensor: + cache_kv = self._get_context_kv(input_embdings, infer_state, layer_weight) + self._post_cache_kv(cache_kv, infer_state, layer_weight) + return input_embdings + + def _get_context_kv(self, input_embdings, infer_state, layer_weight): token_num, _ = input_embdings.shape cache_kv = layer_weight.kv_proj.mm(input_embdings, use_custom_tensor_mananger=False) qk_rmsnorm_forward( @@ -41,8 +48,24 @@ def context_forward( infer_state.position_sin, partial_rotary_factor=self.partial_rotary_factor, ) - self._post_cache_kv(cache_kv, infer_state, layer_weight) - return input_embdings + return cache_kv + + def token_attention_forward(self, input_embdings, infer_state, layer_weight): + att_state = infer_state.decode_att_state + if not isinstance(att_state, WindowedMTPDecodeAttState): + return super().token_attention_forward(input_embdings, infer_state, layer_weight) + + q, noise_kv = self._get_qkv(input_embdings, infer_state, layer_weight) + model = att_state.backend.model + scratch = infer_state.mem_manager.windowed_draft_kv.pack( + self.layer_num_ - model.draft_layer_start, att_state.b_req_idx, noise_kv, att_state.block_size + ) + output = att_state.decode_att( + q.view(-1, self.tp_q_head_num_, self.head_dim_), + scratch[:, :, : self.tp_k_head_num_], + scratch[:, :, self.tp_k_head_num_ :], + ) + return self._get_o(output.reshape(-1, self.tp_q_head_num_ * self.head_dim_), infer_state, layer_weight) def _get_qkv(self, input, infer_state: Qwen3DFlashInferStateInfo, layer_weight: Qwen3DFlashTransformerLayerWeight): q = layer_weight.q_proj.mm(input, use_custom_tensor_mananger=False) diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index ab0c76c149..b021d9b8f4 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -5,6 +5,8 @@ Fp8Fa3AttBackend, ) from lightllm.common.basemodel.basemodel import TpPartBaseModel +from lightllm.common.basemodel.attention.fa3.windowed_mtp import WindowedMTPAttBackend +from lightllm.utils.windowed_mtp import validate_windowed_mtp, window_capacity from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput from lightllm.models.llama.model import LlamaTpPartModel from lightllm.models.draft_registry import DraftModelRegistry @@ -21,6 +23,7 @@ class Qwen3DFlashModel(LlamaTpPartModel): """Qwen3 DFlash draft model.""" is_mtp_draft_model = True + uses_windowed_draft_kv = False pre_and_post_weight_class = Qwen3DFlashPreAndPostLayerWeight transformer_weight_class = Qwen3DFlashTransformerLayerWeight @@ -43,6 +46,9 @@ def _verify_params(self): assert self.args.mtp_step <= self.config["block_size"] self.config["block_size"] = self.args.mtp_step + self.uses_windowed_draft_kv = getattr(self.args, "mtp_draft_cache_mode", "full") == "windowed" + if self.uses_windowed_draft_kv: + validate_windowed_mtp(self.args) def _init_custom(self): self._cos_cached = self.main_model._cos_cached @@ -54,10 +60,23 @@ def _init_req_manager(self): self.req_manager = self.main_model.req_manager def _init_mem_manager(self): - # Draft KV uses target-owned cache slots. + # The shared manager owns target KV and, when enabled, draft window KV. self.mem_manager = self.main_model.mem_manager + if self.uses_windowed_draft_kv: + self.mem_manager.init_windowed_draft_kv( + requests=self.args.running_max_req_size + 1, + capacity=window_capacity(self.args), + layers=self.config["n_layer"], + kv_heads=self.config["num_key_value_heads"] // self.tp_world_size_, + head_dim=self.config.get("head_dim", self.config["n_embed"] // self.config["num_attention_heads"]), + window=self.args.mtp_draft_window, + sinks=self.args.mtp_draft_sinks, + ) def _init_att_backend(self): + if self.uses_windowed_draft_kv: + self.prefill_att_backend = self.decode_att_backend = WindowedMTPAttBackend(model=self) + return super()._init_att_backend() # FA3 is currently the only backend that supports non-causal block decode. # TODO: Remove this restriction after Triton and FlashInfer support non-causal block decode. @@ -87,7 +106,42 @@ def _init_weights(self, start_layer_index=None): for i in range(self.config["n_layer"]) ] + @torch.no_grad() + def commit_features(self, model_input, hidden, starts=None, accept_len=None): + """Commit accepted target features to the bounded draft KV store.""" + # The fused feature is transient; only its per-layer K/V is retained. + projected = self.pre_post_weight.fc_weight_.mm(hidden, use_custom_tensor_mananger=False) + projected = self.pre_post_weight.hidden_norm_weight_(projected, eps=self.pre_infer.eps_, alloc_func=torch.empty) + if model_input.is_prefill: + reqs = model_input.b_req_idx.long() + starts = model_input.b_prefill_start_loc + first = model_input.b_ready_cache_len + lengths = model_input.b_seq_len - first + max_new = model_input.max_q_seq_len + else: + reqs = model_input.b_req_idx.index_select(0, starts.long()).long() + first = model_input.b_seq_len.index_select(0, starts.long()) - 1 + lengths = accept_len + max_new = self.args.mtp_step + 1 + store = self.mem_manager.windowed_draft_kv + features, positions = store.prepare(reqs, projected, starts, first, lengths, max_new) + infer_state = self.infer_state_class() + infer_state.position_cos = self._cos_cached.index_select(0, positions.clamp_min(0)) + infer_state.position_sin = self._sin_cached.index_select(0, positions.clamp_min(0)) + for i, (layer, weight) in enumerate(zip(self.layers_infer, self.trans_layers_weight)): + kv = layer._get_context_kv(features, infer_state, weight) + store.write(i, reqs, positions, kv) + + def _update_decode_req_to_token(self, infer_state: Qwen3DFlashInferStateInfo): + # Windowed draft blocks have no target-owned KV slots to record. + if self.uses_windowed_draft_kv: + return + return super()._update_decode_req_to_token(infer_state) + def _decode(self, model_input: ModelInput) -> ModelOutput: + if self.uses_windowed_draft_kv: + assert model_input.batch_size > 0 and model_input.batch_size % self.block_size == 0 + return super()._decode(model_input) if model_input.mtp_draft_input_hiddens is None: return super()._decode(model_input) diff --git a/lightllm/models/windowed_mtp/__init__.py b/lightllm/models/windowed_mtp/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/lightllm/models/windowed_mtp/layer_infer.py b/lightllm/models/windowed_mtp/layer_infer.py deleted file mode 100644 index 735d9c021b..0000000000 --- a/lightllm/models/windowed_mtp/layer_infer.py +++ /dev/null @@ -1,17 +0,0 @@ -from lightllm.utils.sgl_utils import flash_attn_with_kvcache - - -class WindowedAttentionMixin: - def token_attention_forward(self, input_embdings, infer_state, layer_weight): - q, noise_kv = self._get_qkv(input_embdings, infer_state, layer_weight) - store, reqs, counts, block = infer_state.windowed_context - scratch = store.pack(self.windowed_layer_index, reqs, noise_kv, block) - output = flash_attn_with_kvcache( - q.view(-1, block, self.tp_q_head_num_, self.head_dim_), - scratch[:, :, : self.tp_k_head_num_], - scratch[:, :, self.tp_k_head_num_ :], - cache_seqlens=counts + block, - causal=False, - softmax_scale=self.head_dim_ ** -0.5, - ) - return self._get_o(output.reshape(-1, self.tp_q_head_num_ * self.head_dim_), infer_state, layer_weight) diff --git a/lightllm/models/windowed_mtp/model.py b/lightllm/models/windowed_mtp/model.py deleted file mode 100644 index ca721539e8..0000000000 --- a/lightllm/models/windowed_mtp/model.py +++ /dev/null @@ -1,148 +0,0 @@ -import torch - -from lightllm.common.basemodel.triton_kernel.norm.qk_norm import qk_rmsnorm_forward -from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd -from lightllm.models.qwen3_dflash.model import Qwen3DFlashModel -from lightllm.models.windowed_mtp.kv_store import WindowKVStore -from lightllm.models.windowed_mtp.layer_infer import WindowedAttentionMixin -from lightllm.utils.windowed_mtp import window_capacity, validate_windowed_mtp -from lightllm.common.basemodel.attention.base_att import BaseAttBackend, BaseDecodeAttState -from lightllm.utils.log_utils import init_logger -from lightllm.utils.torch_memory_saver_utils import MemoryTag - - -logger = init_logger(__name__) - - -class WindowedDecodeState(BaseDecodeAttState): - def init_state(self): - pass - - def decode_att(self, *args, **kwargs): - raise RuntimeError("Windowed draft uses bounded FA3 attention directly") - - -class WindowedAttBackend(BaseAttBackend): - def create_att_decode_state(self, infer_state): - return WindowedDecodeState(backend=self, infer_state=infer_state) - - -class WindowedDraftMixin: - uses_windowed_draft_kv = True - - def _verify_params(self): - super()._verify_params() - validate_windowed_mtp(self.args) - - def _init_custom(self): - super()._init_custom() - # _init_custom runs after the base model's KV allocation region. - with self.torch_memory_saver.region(tag=MemoryTag.KV_CACHE): - self.kv_store = WindowKVStore( - self.args.running_max_req_size + 1, - window_capacity(self.args), - self.config["n_layer"], - self.tp_k_head_num_, - self.head_dim_, - self.data_type, - self._cos_cached.device, - self.args.mtp_draft_window, - self.args.mtp_draft_sinks, - ) - for i, layer in enumerate(self.layers_infer): - layer.windowed_layer_index = i - store = self.kv_store - pool_bytes = sum(t.numel() * t.element_size() for t in (store.kv, store.ends, store.counts)) - logger.info( - f"{self.args.mtp_mode} windowed: full-context draft KV layers=0, " - f"window KV layers={self.config['n_layer']}, " - f"capacity={store.capacity}, KV pool bytes={pool_bytes}, target KV layers={len(self.mem_manager.kv_buffer)}" - ) - - @torch.no_grad() - def commit_features(self, model_input, hidden, starts=None, accept_len=None): - # The fused feature is transient; only its per-layer K/V is retained. - projected = self.pre_post_weight.fc_weight_.mm(hidden, use_custom_tensor_mananger=False) - projected = self.pre_post_weight.hidden_norm_weight_(projected, eps=self.pre_infer.eps_, alloc_func=torch.empty) - if model_input.is_prefill: - reqs = model_input.b_req_idx.long() - starts = model_input.b_prefill_start_loc - first = model_input.b_ready_cache_len - lengths = model_input.b_seq_len - first - max_new = model_input.max_q_seq_len - else: - reqs = model_input.b_req_idx.index_select(0, starts.long()).long() - first = model_input.b_seq_len.index_select(0, starts.long()) - 1 - lengths = accept_len - max_new = self.args.mtp_step + 1 - features, positions = self.kv_store.prepare(reqs, projected, starts, first, lengths, max_new) - cos = self._cos_cached.index_select(0, positions.clamp_min(0)) - sin = self._sin_cached.index_select(0, positions.clamp_min(0)) - for i, (layer, weight) in enumerate(zip(self.layers_infer, self.trans_layers_weight)): - kv = weight.kv_proj.mm(features, use_custom_tensor_mananger=False) - qk_rmsnorm_forward( - kv[:, : self.tp_k_head_num_ * self.head_dim_], weight.qk_norm_weight_.k_weight, layer.eps_ - ) - kv = kv.view(-1, 2 * self.tp_k_head_num_, self.head_dim_) - rotary_emb_fwd( - kv[:, : self.tp_k_head_num_], None, cos, sin, partial_rotary_factor=layer.partial_rotary_factor - ) - self.kv_store.write(i, reqs, positions, kv) - - def _token_forward(self, infer_state): - # These reads run inside CUDA Graph, using the current request IDs/state. - reqs = infer_state.b_req_idx[:: self.block_size].long() - infer_state.windowed_context = ( - self.kv_store, - reqs, - self.kv_store.counts.index_select(0, reqs), - self.block_size, - ) - try: - return super()._token_forward(infer_state) - finally: - del infer_state.windowed_context - - def _init_att_backend(self): - self.prefill_att_backend = self.decode_att_backend = WindowedAttBackend(model=self) - - def _decode(self, model_input): - # Window KV is committed by the proposer. Draft block rows never own - # target cache slots and must not overwrite its request-to-token table. - origin_batch_size = model_input.batch_size - assert origin_batch_size > 0 and origin_batch_size % self.block_size == 0 - infer_batch_size = origin_batch_size - use_cuda_graph = self.graph is not None and self.graph.can_run( - batch_size=infer_batch_size, - max_len_in_batch=max(2, model_input.max_kv_seq_len), - ) - need_capture = False - if use_cuda_graph: - infer_batch_size = self.graph.find_closest_graph_batch_size(batch_size=infer_batch_size) - need_capture = self.graph.need_capture(infer_batch_size) - model_input = self._create_padded_decode_model_input(model_input, infer_batch_size) - infer_state = self._create_inferstate(model_input) - infer_state.is_cuda_graph = need_capture - infer_state.init_some_extra_state(self) - infer_state.init_att_state() - if use_cuda_graph: - if need_capture: - output = self.graph.capture_decode(self._token_forward, infer_state) - else: - output = self.graph.replay(infer_state) - else: - output = self._token_forward(infer_state) - return self._create_unpad_decode_model_output(output, origin_batch_size=origin_batch_size) - - -def windowed_model_class(original): - if not issubclass(original, Qwen3DFlashModel): - raise ValueError("Windowed draft KV currently supports the Qwen3 parallel block backbone") - layer_class = type( - "Windowed" + original.transformer_layer_infer_class.__name__, - (WindowedAttentionMixin, original.transformer_layer_infer_class), - {}, - ) - return type( - "Windowed" + original.__name__, (WindowedDraftMixin, original), {"transformer_layer_infer_class": layer_class} - ) diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index a3ac168060..fe77e5f37b 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -342,10 +342,6 @@ def init_mtp_draft_model(self, main_kvargs: dict): model_cfg=draft_model_cfg, spec_mode=spec_mode, ) - if self.args.mtp_draft_cache_mode == "windowed": - from lightllm.models.windowed_mtp.model import windowed_model_class - - draft_model_class = windowed_model_class(draft_model_class) self.draft_models.append(draft_model_class(draft_model_kvargs)) self.logger.info(f"loaded speculative draft model class {self.draft_models[i].__class__}") From f61b3c8d3481afa8bcf18a80a243ce67a45611b7 Mon Sep 17 00:00:00 2001 From: Zhao Xintong Date: Wed, 16 Sep 2026 08:03:18 +0000 Subject: [PATCH 3/9] remove sink --- docs/CN/source/tutorial/api_server_args.rst | 17 +++--- docs/EN/source/tutorial/api_server_args.rst | 15 ++---- .../kv_cache_mem_manager/windowed_mtp.py | 52 ++++--------------- lightllm/models/qwen3_dflash/model.py | 6 +-- lightllm/server/api_cli.py | 7 ++- lightllm/server/api_start.py | 2 +- lightllm/server/core/objs/start_args_type.py | 5 +- lightllm/utils/envs_utils.py | 4 +- lightllm/utils/profile_max_tokens.py | 2 +- lightllm/utils/windowed_mtp.py | 6 +-- 10 files changed, 37 insertions(+), 79 deletions(-) diff --git a/docs/CN/source/tutorial/api_server_args.rst b/docs/CN/source/tutorial/api_server_args.rst index a060681988..69fbfc2733 100644 --- a/docs/CN/source/tutorial/api_server_args.rst +++ b/docs/CN/source/tutorial/api_server_args.rst @@ -678,27 +678,24 @@ MTP 多预测参数 * ``eagle_no_att`` * ``None``: 不启用 mtp（默认） -.. option:: --mtp_draft_cache_mode +.. option:: --mtp_draft_kv_mode - 草稿 KV 缓存策略，默认 ``full``。dflash style 草稿模型支持 ``windowed``， - 仅保留开头和最近窗口内的 KV。target 仍使用完整上下文。 + 草稿 KV 缓存策略，默认 ``full``。dflash style 草稿模型支持 ``window``， + 仅保留最近窗口内的 KV。target 仍使用完整上下文。 当前要求 ``--run_mode normal``、``--dp 1``、未量化 KV，且设置 ``--disable_dynamic_prompt_cache``。暂不支持 CPU/prefix 缓存、动态验证、 diverse、混合 prefill/decode 或 overlap 模式。 -.. option:: --mtp_draft_window +.. option:: --mtp_draft_window_size - ``windowed`` 策略保留的最近位置数，默认 ``512``，必须大于 0。 + ``window`` 策略保留的最近位置数，默认 ``512``，必须大于 0。 -.. option:: --mtp_draft_sinks - - ``windowed`` 策略保留的开头位置数，默认 ``1``；设为 ``0`` 可关闭。 - 每个请求的历史容量上限为 ``mtp_draft_window + mtp_draft_sinks``。 + 每个请求的历史容量上限为 ``mtp_draft_window_size``。 示例：在已有 dflash style 草稿模型启动命令中添加:: - --mtp_draft_cache_mode windowed --mtp_draft_window 512 --mtp_draft_sinks 1 --disable_dynamic_prompt_cache + --mtp_draft_kv_mode window --mtp_draft_window_size 512 --disable_dynamic_prompt_cache .. option:: --mtp_draft_model_dir diff --git a/docs/EN/source/tutorial/api_server_args.rst b/docs/EN/source/tutorial/api_server_args.rst index 1714937619..105d10ec1f 100644 --- a/docs/EN/source/tutorial/api_server_args.rst +++ b/docs/EN/source/tutorial/api_server_args.rst @@ -690,30 +690,25 @@ MTP Multi-Prediction Parameters * ``eagle_no_att`` * ``None``: Do not enable mtp (default) -.. option:: --mtp_draft_cache_mode +.. option:: --mtp_draft_kv_mode Draft KV cache policy; defaults to ``full``. dflash style draft models support - ``windowed``, which retains only initial and recent-window KV. The target still + ``window``, which retains only recent-window KV. The target still uses its full context. Currently requires ``--run_mode normal``, ``--dp 1``, unquantized KV, and ``--disable_dynamic_prompt_cache``. CPU/prefix caching, dynamic verification, diverse, mixed prefill/decode, and overlap modes are not supported. -.. option:: --mtp_draft_window +.. option:: --mtp_draft_window_size - Number of recent positions retained by ``windowed``; defaults to ``512`` and + Number of recent positions retained by ``window``; defaults to ``512`` and must be positive. -.. option:: --mtp_draft_sinks - - Number of initial positions retained by ``windowed``; defaults to ``1``. - Set to ``0`` to disable sinks. Per-request history capacity is bounded by - ``mtp_draft_window + mtp_draft_sinks``. Add these options to an existing launch command for a dflash style draft model:: - --mtp_draft_cache_mode windowed --mtp_draft_window 512 --mtp_draft_sinks 1 --disable_dynamic_prompt_cache + --mtp_draft_kv_mode window --mtp_draft_window_size 512 --disable_dynamic_prompt_cache .. option:: --mtp_draft_model_dir diff --git a/lightllm/common/kv_cache_mem_manager/windowed_mtp.py b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py index 88296b85c4..7ef7eb89cc 100644 --- a/lightllm/common/kv_cache_mem_manager/windowed_mtp.py +++ b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py @@ -19,8 +19,6 @@ def _gather_new( N: tl.constexpr, H: tl.constexpr, C: tl.constexpr, - W: tl.constexpr, - S: tl.constexpr, BLOCK: tl.constexpr, ): row = tl.program_id(0) @@ -29,14 +27,9 @@ def _gather_new( first = tl.load(FIRST + batch).to(tl.int64) length = tl.load(LENGTHS + batch) end = first + length - if W > 0: - sinks = tl.minimum(tl.maximum(S - first, 0), length) - recent = tl.maximum(tl.maximum(first, S), end - W) - count = sinks + tl.maximum(end - recent, 0) - position = tl.where(offset < sinks, first + offset, recent + offset - sinks) - else: - count = length - position = first + offset + recent = tl.maximum(first, end - C) + count = end - recent + position = recent + offset valid = offset < count dim = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) source = tl.load(STARTS + batch).to(tl.int64) + position - first @@ -57,18 +50,13 @@ def _write_ring( POS, N: tl.constexpr, C: tl.constexpr, - W: tl.constexpr, - S: tl.constexpr, WIDTH: tl.constexpr, BLOCK: tl.constexpr, ): row = tl.program_id(0) req = tl.load(REQS + row // N).to(tl.int64) position = tl.load(POS + row) - if W > 0: - slot = tl.where(position < S, position, S + (position - S) % W) - else: - slot = position + slot = position % C dim = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) valid = (position >= 0) & (dim < WIDTH) value = tl.load(NEW + row * WIDTH + dim, valid, other=0) @@ -84,8 +72,6 @@ def _pack_ring( OUT, C: tl.constexpr, B: tl.constexpr, - W: tl.constexpr, - S: tl.constexpr, WIDTH: tl.constexpr, BLOCK: tl.constexpr, ): @@ -94,13 +80,8 @@ def _pack_ring( req = tl.load(REQS + batch).to(tl.int64) end = tl.load(ENDS + req) count = tl.minimum(end, C) - if W > 0: - sinks = tl.minimum(end, S) - recent = tl.maximum(S, end - W) - position = tl.where(slot < sinks, slot, recent + slot - sinks) - ring_slot = tl.where(position < S, position, S + (position - S) % W) - else: - ring_slot = slot + position = tl.maximum(0, end - C) + slot + ring_slot = position % C dim = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) context = tl.load(POOL + (req * C + ring_slot) * WIDTH + dim, (slot < count) & (dim < WIDTH), other=0) noise = tl.load( @@ -110,7 +91,7 @@ def _pack_ring( class WindowKVStore: - """Per-layer sink + recent-window K/V, indexed independently of target KV. + """Per-layer recent-window K/V, indexed independently of target KV. Only newly accepted positions are projected and written. Each window slot has at most one writer, including prefill chunks larger than the window. @@ -118,9 +99,9 @@ class WindowKVStore: are overwritten by the first prefill. Draft noise K/V never enters the ring. """ - def __init__(self, requests, capacity, layers, kv_heads, head_dim, dtype, device, window, sinks): - self.capacity, self.window, self.sinks = capacity, window, sinks - self.kv = torch.zeros((layers, requests, capacity, 2 * kv_heads, head_dim), dtype=dtype, device=device) + def __init__(self, requests, layers, kv_heads, head_dim, dtype, device, window): + self.capacity = window + self.kv = torch.zeros((layers, requests, window, 2 * kv_heads, head_dim), dtype=dtype, device=device) self.ends = torch.zeros(requests, dtype=torch.int64, device=device) self.counts = torch.zeros(requests, dtype=torch.int32, device=device) @@ -146,8 +127,6 @@ def prepare(self, reqs, features, starts, first, lengths, max_new): n, features.shape[-1], self.capacity, - self.window, - self.sinks, 512, ) return packed, positions @@ -161,8 +140,6 @@ def write(self, layer, reqs, positions, new_kv): positions, positions.numel() // reqs.numel(), self.capacity, - self.window, - self.sinks, width, 512, ) @@ -178,8 +155,6 @@ def pack(self, layer, reqs, noise_kv, block): output, self.capacity, block, - self.window, - self.sinks, heads * dim, 512, ) @@ -189,10 +164,5 @@ def retained_positions(self, reqs): """Diagnostic only: logical positions, in attention order.""" end = self.ends.index_select(0, reqs)[:, None] slots = torch.arange(self.capacity, device=reqs.device)[None, :] - if self.window: - sinks = end.clamp_max(self.sinks) - recent = (end - self.window).clamp_min(self.sinks) - positions = torch.where(slots < sinks, slots, recent + slots - sinks) - else: - positions = slots.expand(reqs.numel(), -1) + positions = (end - self.capacity).clamp_min(0) + slots return positions.masked_fill(slots >= end.clamp_max(self.capacity), -1) diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index b021d9b8f4..b63a6beadc 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -46,7 +46,7 @@ def _verify_params(self): assert self.args.mtp_step <= self.config["block_size"] self.config["block_size"] = self.args.mtp_step - self.uses_windowed_draft_kv = getattr(self.args, "mtp_draft_cache_mode", "full") == "windowed" + self.uses_windowed_draft_kv = getattr(self.args, "mtp_draft_kv_mode", "full") == "window" if self.uses_windowed_draft_kv: validate_windowed_mtp(self.args) @@ -65,12 +65,10 @@ def _init_mem_manager(self): if self.uses_windowed_draft_kv: self.mem_manager.init_windowed_draft_kv( requests=self.args.running_max_req_size + 1, - capacity=window_capacity(self.args), layers=self.config["n_layer"], kv_heads=self.config["num_key_value_heads"] // self.tp_world_size_, head_dim=self.config.get("head_dim", self.config["n_embed"] // self.config["num_attention_heads"]), - window=self.args.mtp_draft_window, - sinks=self.args.mtp_draft_sinks, + window=window_capacity(self.args), ) def _init_att_backend(self): diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 0927cf472d..c425236360 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -768,13 +768,12 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="""Whether to enable fused shared experts for supported MoE models. It is auto-enabled when supported.""", ) parser.add_argument( - "--mtp_draft_cache_mode", - choices=["full", "windowed"], + "--mtp_draft_kv_mode", + choices=["full", "window"], default="full", help="Draft KV retention policy for parallel block MTP", ) - parser.add_argument("--mtp_draft_window", type=int, default=512) - parser.add_argument("--mtp_draft_sinks", type=int, default=1) + parser.add_argument("--mtp_draft_window_size", type=int, default=512) parser.add_argument( "--mtp_mode", choices=[ diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 5b2db7eefd..2734ac3b5b 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -163,7 +163,7 @@ def _launch_subprocesses(args: StartArgs): f"{sorted(allowed_ep_decode_att_backends)}; flashinfer is not supported." ) - if args.mtp_draft_cache_mode == "windowed": + if args.mtp_draft_kv_mode == "window": from lightllm.utils.windowed_mtp import validate_windowed_mtp validate_windowed_mtp(args) diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 9268045808..f0badead47 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -199,9 +199,8 @@ class StartArgs: }, ) mtp_draft_model_dir: Optional[List[str]] = field(default=None) - mtp_draft_cache_mode: str = field(default="full", metadata={"choices": ["full", "windowed"]}) - mtp_draft_window: int = field(default=512) - mtp_draft_sinks: int = field(default=1) + mtp_draft_kv_mode: str = field(default="full", metadata={"choices": ["full", "window"]}) + mtp_draft_window_size: int = field(default=512) mtp_step: int = field(default=0) mtp_dynamic_verify: bool = field(default=False) kv_quant_calibration_config_path: Optional[str] = field(default=None) diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index e1f93de96b..1b5fe746c4 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -255,7 +255,7 @@ def enable_cpu_cache_numa_interleave() -> bool: @lru_cache(maxsize=None) def get_added_mtp_kv_layer_num() -> int: args = get_env_start_args() - if getattr(args, "mtp_draft_cache_mode", "full") == "windowed": + if getattr(args, "mtp_draft_kv_mode", "full") == "window": return 0 mtp_mode = args.mtp_mode @@ -290,7 +290,7 @@ def get_mtp_weight_layer_num() -> int: return args.mtp_step if mtp_mode == "eagle_no_att": return 1 - if getattr(args, "mtp_draft_cache_mode", "full") == "windowed": + if getattr(args, "mtp_draft_kv_mode", "full") == "window": return _get_mtp_draft_backbone_layer_num(args.mtp_draft_model_dir[0]) return get_added_mtp_kv_layer_num() diff --git a/lightllm/utils/profile_max_tokens.py b/lightllm/utils/profile_max_tokens.py index 109107b9c3..a1778e1769 100644 --- a/lightllm/utils/profile_max_tokens.py +++ b/lightllm/utils/profile_max_tokens.py @@ -76,7 +76,7 @@ def profile_mtp_weight_memory(model): yield target_weight_bytes = torch.cuda.memory_allocated() - weight_memory_before draft_cache_bytes = 0 - if getattr(model.args, "mtp_draft_cache_mode", "full") == "windowed": + if getattr(model.args, "mtp_draft_kv_mode", "full") == "window": from lightllm.utils.windowed_mtp import window_kv_pool_bytes draft_cache_bytes = window_kv_pool_bytes(model.args, torch.tensor([], dtype=model.data_type).element_size()) diff --git a/lightllm/utils/windowed_mtp.py b/lightllm/utils/windowed_mtp.py index 656f601d04..d7eb568cc6 100644 --- a/lightllm/utils/windowed_mtp.py +++ b/lightllm/utils/windowed_mtp.py @@ -6,9 +6,9 @@ def window_capacity(args): - if args.mtp_draft_window <= 0 or args.mtp_draft_sinks < 0: - raise ValueError("mtp_draft_window must be positive and mtp_draft_sinks nonnegative") - return args.mtp_draft_window + args.mtp_draft_sinks + if args.mtp_draft_window_size <= 0: + raise ValueError("mtp_draft_window_size must be positive") + return args.mtp_draft_window_size def window_kv_pool_bytes(args, element_size): From 55104ca4dc320282e37518a54820824f8595668e Mon Sep 17 00:00:00 2001 From: Zhao Xintong Date: Tue, 22 Sep 2026 14:36:16 +0000 Subject: [PATCH 4/9] support cache --- docs/CN/source/tutorial/api_server_args.rst | 8 - docs/EN/source/tutorial/api_server_args.rst | 9 - lightllm/common/basemodel/basemodel.py | 5 +- .../triton_kernel/hybrid_cpu_cache_copy.py | 156 +++++++++++++++ .../kv_cache_mem_manager/mem_manager.py | 57 +++++- .../kv_cache_mem_manager/operator/hybrid.py | 179 ++++++++++++++++++ .../operator/linear_att.py | 166 ++-------------- .../operator/windowed_mtp.py | 32 ++++ .../qwen3next_mem_manager.py | 51 ++--- .../kv_cache_mem_manager/windowed_mtp.py | 87 +++++++++ lightllm/common/req_manager/__init__.py | 9 +- lightllm/common/req_manager/linear_att.py | 19 +- lightllm/common/req_manager/windowed_mtp.py | 21 ++ .../common/state_cache_manager/__init__.py | 13 +- .../common/state_cache_manager/linear_att.py | 7 + .../state_cache_manager/windowed_mtp.py | 171 +++++++++++++++++ lightllm/models/qwen3_dflash/model.py | 11 +- lightllm/server/api_start.py | 13 +- lightllm/server/core/objs/req.py | 2 +- .../model_infer/mode_backend/base_backend.py | 17 +- lightllm/utils/config_utils.py | 6 +- lightllm/utils/kv_cache_utils.py | 2 +- lightllm/utils/profile_max_tokens.py | 6 +- lightllm/utils/windowed_mtp.py | 53 ------ unit_tests/common/test_windowed_mtp.py | 134 +++++++++++++ 25 files changed, 949 insertions(+), 285 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/hybrid_cpu_cache_copy.py create mode 100644 lightllm/common/kv_cache_mem_manager/operator/hybrid.py create mode 100644 lightllm/common/kv_cache_mem_manager/operator/windowed_mtp.py create mode 100644 lightllm/common/req_manager/windowed_mtp.py create mode 100644 lightllm/common/state_cache_manager/windowed_mtp.py delete mode 100644 lightllm/utils/windowed_mtp.py create mode 100644 unit_tests/common/test_windowed_mtp.py diff --git a/docs/CN/source/tutorial/api_server_args.rst b/docs/CN/source/tutorial/api_server_args.rst index 69fbfc2733..36f02af083 100644 --- a/docs/CN/source/tutorial/api_server_args.rst +++ b/docs/CN/source/tutorial/api_server_args.rst @@ -683,20 +683,12 @@ MTP 多预测参数 草稿 KV 缓存策略，默认 ``full``。dflash style 草稿模型支持 ``window``， 仅保留最近窗口内的 KV。target 仍使用完整上下文。 - 当前要求 ``--run_mode normal``、``--dp 1``、未量化 KV，且设置 - ``--disable_dynamic_prompt_cache``。暂不支持 CPU/prefix 缓存、动态验证、 - diverse、混合 prefill/decode 或 overlap 模式。 - .. option:: --mtp_draft_window_size ``window`` 策略保留的最近位置数，默认 ``512``，必须大于 0。 每个请求的历史容量上限为 ``mtp_draft_window_size``。 - 示例：在已有 dflash style 草稿模型启动命令中添加:: - - --mtp_draft_kv_mode window --mtp_draft_window_size 512 --disable_dynamic_prompt_cache - .. option:: --mtp_draft_model_dir MTP 多预测功能的草稿模型路径 diff --git a/docs/EN/source/tutorial/api_server_args.rst b/docs/EN/source/tutorial/api_server_args.rst index 105d10ec1f..e562491eda 100644 --- a/docs/EN/source/tutorial/api_server_args.rst +++ b/docs/EN/source/tutorial/api_server_args.rst @@ -696,20 +696,11 @@ MTP Multi-Prediction Parameters ``window``, which retains only recent-window KV. The target still uses its full context. - Currently requires ``--run_mode normal``, ``--dp 1``, unquantized KV, and - ``--disable_dynamic_prompt_cache``. CPU/prefix caching, dynamic verification, - diverse, mixed prefill/decode, and overlap modes are not supported. - .. option:: --mtp_draft_window_size Number of recent positions retained by ``window``; defaults to ``512`` and must be positive. - - Add these options to an existing launch command for a dflash style draft model:: - - --mtp_draft_kv_mode window --mtp_draft_window_size 512 --disable_dynamic_prompt_cache - .. option:: --mtp_draft_model_dir Path to the draft model for MTP multi-prediction functionality diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index baf3fd2fc2..2f8011f7a6 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -14,7 +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.req_manager import ReqManager, ReqManagerForWindowedMTP from lightllm.common.infer_utils import init_req_to_token_indexes from lightllm.common.build_utils import repair_config from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import copy_kv_index_to_req @@ -237,7 +237,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_draft_kv_mode == "window" 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): diff --git a/lightllm/common/basemodel/triton_kernel/hybrid_cpu_cache_copy.py b/lightllm/common/basemodel/triton_kernel/hybrid_cpu_cache_copy.py new file mode 100644 index 0000000000..f9fb32d26a --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/hybrid_cpu_cache_copy.py @@ -0,0 +1,156 @@ +"""Stream-ordered copies for token KV and arbitrary pinned checkpoint tensors.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _copy_checkpoint_pages( + SRC, + DST, + SRC_IDS, + DST_IDS, + READIES, + SRC_STRIDE, + DST_STRIDE, + WIDTH, + PAGE_NUM, + SKIP_READY: tl.constexpr, + BLOCK: tl.constexpr, +): + for page in range(PAGE_NUM): + src_id = tl.load(SRC_IDS + page).to(tl.int64) + dst_id = tl.load(DST_IDS + page).to(tl.int64) + valid = (src_id >= 0) & (dst_id >= 0) + if SKIP_READY: + valid = valid & ~tl.load(READIES + page) + if valid: + for block in range(tl.program_id(0), tl.cdiv(WIDTH, BLOCK), tl.num_programs(0)): + offsets = block * BLOCK + tl.arange(0, BLOCK) + data = tl.load(SRC + src_id * SRC_STRIDE + offsets, offsets < WIDTH, other=0) + tl.store(DST + dst_id * DST_STRIDE + offsets, data, offsets < WIDTH, 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(sources) == len(destinations) + 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] + assert source.stride(1) == destination.stride(1) == 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, + ) + + +@triton.jit +def _copy_checkpoint_tensor(SRC, DST, SIZE, BLOCK: tl.constexpr): + for block in range(tl.program_id(0), tl.cdiv(SIZE, BLOCK), tl.num_programs(0)): + offsets = block * BLOCK + tl.arange(0, BLOCK) + data = tl.load(SRC + offsets, offsets < SIZE, other=0) + tl.store(DST + offsets, data, offsets < SIZE, cache_modifier=".wt") + + +def copy_checkpoint_state(sources, destinations): + """Copy a tail checkpoint on the CUDA stream, including CPU-to-CPU copies. + + A host-side tensor copy would race prior asynchronous saves and transfers. + """ + assert len(sources) == len(destinations) + for source, destination in zip(sources, destinations): + assert source.shape == destination.shape and source.dtype == destination.dtype + assert source.is_contiguous() and destination.is_contiguous() + source = source.reshape(-1).view(torch.uint8) + destination = destination.reshape(-1).view(torch.uint8) + _copy_checkpoint_tensor[(12,)](source, destination, source.numel(), BLOCK=4096) + + +@triton.jit +def _copy_full_att_pages( + KV, + CACHE, + MEM_IDS, + PAGE_IDS, + READIES, + KV_STRIDE_L, + KV_STRIDE_T, + PAGE_STRIDE, + RANK_OFFSET, + WIDTH: tl.constexpr, + LAYERS: tl.constexpr, + PAGE_TOKENS: tl.constexpr, + PAGE_NUM, + WRITE_RANK: tl.constexpr, + OFFLOAD: tl.constexpr, + BLOCK: tl.constexpr, +): + page_width = PAGE_TOKENS * LAYERS * WIDTH + for page in range(PAGE_NUM): + page_id = tl.load(PAGE_IDS + page).to(tl.int64) + valid = page_id >= 0 + if OFFLOAD: + valid = valid & ~tl.load(READIES + page) & WRITE_RANK + if valid: + for block in range(tl.program_id(0), tl.cdiv(page_width, BLOCK), tl.num_programs(0)): + offsets = block * BLOCK + tl.arange(0, BLOCK) + token = offsets // (LAYERS * WIDTH) + layer = offsets // WIDTH % LAYERS + dim = offsets % WIDTH + mem_id = tl.load(MEM_IDS + page * PAGE_TOKENS + token, offsets < page_width, other=-1).to(tl.int64) + mask = (offsets < page_width) & (mem_id >= 0) + kv_ptr = KV + layer.to(tl.int64) * KV_STRIDE_L + mem_id * KV_STRIDE_T + dim + cache_ptr = CACHE + page_id * PAGE_STRIDE + RANK_OFFSET + offsets + if OFFLOAD: + value = tl.load(kv_ptr, mask, other=0) + tl.store(cache_ptr, value, mask, cache_modifier=".wt") + else: + value = tl.load(cache_ptr, mask, other=0) + tl.store(kv_ptr, value, mask) + + +def copy_full_att_pages(kv, cpu_cache, mem_indexes, page_indexes, config, tp_rank, page_tokens, page_readies=None): + """Pack ordinary target KV as [page, TP shard, token, layer, KV bytes].""" + kv = kv.view(kv.shape[0], kv.shape[1], -1).view(torch.uint8) + pages = cpu_cache.view(cpu_cache.shape[0], -1).view(torch.uint8) + world_size = config.tp_world_size + heads = config.full_att_all_num_kv_heads + if heads >= world_size: + assert heads % world_size == 0 + replicas = 1 + else: + assert world_size % heads == 0 + replicas = world_size // heads + shard_bytes = page_tokens * kv.shape[0] * kv.shape[2] + assert shard_bytes * (world_size // replicas) == config.get_cpu_cache_full_att_bytes() + assert len(mem_indexes) == len(page_indexes) * page_tokens + _copy_full_att_pages[(12,)]( + kv, + pages, + mem_indexes, + page_indexes, + page_readies, + kv.stride(0), + kv.stride(1), + pages.stride(0), + (tp_rank // replicas) * shard_bytes, + kv.shape[2], + kv.shape[0], + page_tokens, + len(page_indexes), + WRITE_RANK=tp_rank % replicas == 0, + OFFLOAD=page_readies is not None, + BLOCK=4096, + ) diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index b179857d2a..326e270052 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -8,7 +8,8 @@ from lightllm.utils.log_utils import init_logger from lightllm.server.router.dynamic_prompt.shared_arr import SharedInt from .allocator import KvCacheAllocator -from .windowed_mtp import WindowKVStore +from .windowed_mtp import WindowKVStore, WindowKVPageHelper +from lightllm.common.state_cache_manager.windowed_mtp import load_window_state_config from lightllm.utils.profile_max_tokens import get_available_gpu_memory, get_total_gpu_memory from lightllm.utils.dist_utils import ( get_current_device_id, @@ -41,6 +42,8 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False self.dtype = dtype self.windowed_draft_kv = None self._windowed_draft_kv_params = None + args = get_env_start_args() + self.window_state_config = load_window_state_config(args) if args.mtp_draft_kv_mode == "window" else None # profile the max total token num if the size is None self.profile_size(mem_fraction) @@ -56,7 +59,12 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False self.HOLD_TOKEN_MEMINDEX = self.size # 构建对外的操作类接口 - self.operator: BaseMemManagerOperator = self.operator_class(self) + operator_class = self.operator_class + if self.window_state_config is not None and operator_class is NormalMemOperator: + from .operator.windowed_mtp import WindowedMTPMemOperator + + operator_class = WindowedMTPMemOperator + self.operator: BaseMemManagerOperator = operator_class(self) def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: k = self.kv_buffer[layer_index][:, : self.head_num, :] @@ -117,9 +125,29 @@ def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): self.kv_buffer = torch.empty((layer_num, size + 1, 2 * head_num, head_dim), dtype=dtype, device="cuda") if self._windowed_draft_kv_params is not None: self.init_windowed_draft_kv(**self._windowed_draft_kv_params) + self._init_state_cache_buffers() + + def _init_state_cache_buffers(self): + self.big_page_buffers = None + if self.window_state_config is not None: + from lightllm.common.state_cache_manager import WindowStateCacheManager + + args = get_env_start_args() + page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num + self.big_page_buffers = WindowStateCacheManager( + size=(self.size + page_tokens - 1) // page_tokens + 2, + config=self.window_state_config, + dtype=self.dtype, + keep_num=2, + ) + self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 2 + self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 1 def init_windowed_draft_kv(self, **cache_params): """Create draft window storage from dimensions supplied by the draft model.""" + if self.window_state_config is not None: + for name in ("layers", "kv_heads", "head_dim", "window"): + assert cache_params[name] == getattr(self.window_state_config, name), f"draft window {name} mismatch" self._windowed_draft_kv_params = cache_params self.windowed_draft_kv = WindowKVStore(**cache_params, dtype=self.dtype, device=self.kv_buffer.device) @@ -130,8 +158,13 @@ def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: self._buffer_mem_indexes_tensors = [ torch.empty((page_size,), dtype=torch.int64, device="cpu", pin_memory=True) for _ in range(page_num) ] + self._assert_att_state_page_size() return self.kv_move_buffer + def _assert_att_state_page_size(self): + if self.windowed_draft_kv is not None: + WindowKVPageHelper(self).assert_page_size() + def write_mem_to_page_kv_move_buffer( self, mem_indexes: List[int], @@ -142,6 +175,11 @@ def write_mem_to_page_kv_move_buffer( page_kind: str = "kv", req_idx: int = None, ): + if page_kind == "att_state" and self.windowed_draft_kv is not None: + dp_mems = mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size] + assert len(dp_mems) == dp_world_size + WindowKVPageHelper(self).copy_req_page(page_index, req_idx, dp_mems, mode="write") + return assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" cur_page = self.kv_move_buffer[page_index] pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] @@ -174,6 +212,11 @@ def read_page_kv_move_buffer_to_mem( page_kind: str = "kv", req_idx: int = None, ): + if page_kind == "att_state" and self.windowed_draft_kv is not None: + dp_mems = mem_managers[dp_index * dp_world_size : (dp_index + 1) * dp_world_size] + assert len(dp_mems) == dp_world_size + WindowKVPageHelper(self).copy_req_page(page_index, req_idx, dp_mems, mode="read") + return assert page_kind == "kv", f"{type(self).__name__} does not support page_kind={page_kind}" cur_page = self.kv_move_buffer[page_index] pin_mem_indexes = self._buffer_mem_indexes_tensors[page_index][0 : len(mem_indexes)] @@ -199,6 +242,7 @@ def read_page_kv_move_buffer_to_mem( def _free_buffers(self): self.kv_buffer = None self.windowed_draft_kv = None + self.big_page_buffers = None def alloc(self, need_size) -> torch.Tensor: return self.allocator.alloc(need_size) @@ -255,7 +299,14 @@ def write_to_shm(self, req_manager): with lock: node_world_size = get_node_world_size() shm_name = f"mem_manager_{get_current_rank_in_node()}" - obj_bytes_array = [ForkingPickler.dumps(self).tobytes() for _ in range(node_world_size * 2)] + # Only GPU running state is consumed by PD. Pickling CPU checkpoints + # would move their storage into ordinary shared memory and lose pinning. + big_page_buffers = self.big_page_buffers + self.big_page_buffers = None + try: + obj_bytes_array = [ForkingPickler.dumps(self).tobytes() for _ in range(node_world_size * 2)] + finally: + self.big_page_buffers = big_page_buffers obj_size = len(obj_bytes_array[0]) shm = create_or_link_shm( name=shm_name, expected_size=obj_size * (node_world_size * 2) + 4 + 4, force_mode="create" diff --git a/lightllm/common/kv_cache_mem_manager/operator/hybrid.py b/lightllm/common/kv_cache_mem_manager/operator/hybrid.py new file mode 100644 index 0000000000..655459b092 --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/operator/hybrid.py @@ -0,0 +1,179 @@ +import torch +import triton +from .base import BaseMemManagerOperator +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.dist_utils import get_current_rank_in_dp +from lightllm.common.state_cache_manager import get_hybrid_cache_config +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient + from lightllm.server.router.model_infer.infer_batch import InferReq + + +class HybridAttMemOperator(BaseMemManagerOperator): + """Shared CPU transfer lifecycle for token KV with checkpoint state.""" + + def __init__(self, mem_manager): + super().__init__(mem_manager) + self.cpu_cache_config = get_hybrid_cache_config(linear_config=getattr(mem_manager, "linear_config", None)) + + def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): + from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer + + copy_kv_buffer_to_kv_buffer( + src_mem_index.cuda(non_blocking=True), dst_mem_index.cuda(non_blocking=True), self.mem_manager.kv_buffer + ) + return + + def load_cpu_cache_to_gpu( + self, + mem_indexes: torch.Tensor, + page_indexes: torch.Tensor, + cpu_cache_client: "CpuKvCacheClient", + req: "InferReq", + ): + assert mem_indexes.is_cuda and page_indexes.is_cuda + args = get_env_start_args() + assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) + assert len(mem_indexes) % args.linear_att_hash_page_size == 0 + assert args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num + mem_manager = self.mem_manager + + big_page_num = len(mem_indexes) // args.cpu_cache_token_page_size + max_kv_len = (req.cur_kv_len // args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size + assert max_kv_len % args.cpu_cache_token_page_size == 0 + + big_page_buffer_ids_cpu = [] + for i in range(big_page_num): + page_id = mem_manager.big_page_buffers.alloc_one_state_cache() + assert page_id is not None + req.hybrid_len_to_big_page_id[max_kv_len] = page_id + big_page_buffer_ids_cpu.append(page_id) + max_kv_len -= args.cpu_cache_token_page_size + assert max_kv_len % args.cpu_cache_token_page_size == 0 + + big_page_buffer_ids_cpu.reverse() + + # 碎页情况的处理 + has_tail_page = len(mem_indexes) % args.cpu_cache_token_page_size != 0 + if has_tail_page: + padded_token_num = triton.cdiv( + len(mem_indexes), args.cpu_cache_token_page_size + ) * args.cpu_cache_token_page_size - len(mem_indexes) + mem_indexes = torch.nn.functional.pad(mem_indexes, (0, padded_token_num), mode="constant", value=-1) + + # 将对应的小叶数据拷贝到临时的大页上，再从大页上拷贝到对应的运行态页面上 + big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID) + + big_page_buffer_ids_gpu = torch.tensor(big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu").cuda( + non_blocking=True + ) + + assert len(big_page_buffer_ids_gpu) == len(page_indexes) + + self._load_target_pages(mem_indexes, big_page_buffer_ids_gpu, page_indexes, cpu_cache_client) + self._copy_window_pages(big_page_buffer_ids_gpu, page_indexes, cpu_cache_client) + + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + g_infer_context.req_manager.restore_big_page_state( + big_page_buffer_idx=big_page_buffer_ids_cpu[-1], + req=req, + ) + + return + + def offload_gpu_kv_to_cpu_cache( + self, + mem_indexes: torch.Tensor, + page_indexes: torch.Tensor, + page_readies: torch.Tensor, + cpu_cache_client: "CpuKvCacheClient", + req: "InferReq", + ): + args = get_env_start_args() + if not hasattr(self, "big_page_ids_buffer_store"): + self.big_page_ids_buffer_store = torch.empty((1024 * 1024 * 4,), dtype=torch.int64, device="cuda") + # 多申请3个cpu cache token page size，用于处理碎页情况，碎页情况需要将对应的大页数据拷贝到临时的大页上， + # 再从大页上拷贝到对应的运行态页面上 + self.mem_indexes_buffer = torch.empty( + (args.max_req_total_len + 3 * args.cpu_cache_token_page_size,), dtype=torch.int32, device="cuda" + ) + + assert mem_indexes.is_cuda and page_indexes.is_cuda and page_readies.is_cuda + + assert len(mem_indexes) % args.linear_att_hash_page_size == 0 + assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) + + mem_manager = self.mem_manager + + from lightllm.server.router.model_infer.infer_batch import g_infer_context + + big_page_buffer_ids_cpu = g_infer_context.radix_cache.get_big_page_ids_by_node(req.shared_kv_node) + max_kv_len = (len(mem_indexes) // args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size + start_kv_len = (len(big_page_buffer_ids_cpu) + 1) * args.cpu_cache_token_page_size + for seq_len in range(start_kv_len, max_kv_len + 1, args.cpu_cache_token_page_size): + page_id = req.hybrid_len_to_big_page_id[seq_len] + big_page_buffer_ids_cpu.append(page_id) + + if len(mem_indexes) % args.cpu_cache_token_page_size != 0: + # 存在不满大页的碎页的页面存在需要复制的情况 + dst_len = triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size + assert dst_len <= self.mem_indexes_buffer.shape[0] + dst_mem_indexes = self.mem_indexes_buffer[0:dst_len].fill_(-1) + dst_mem_indexes[0 : len(mem_indexes)].copy_(mem_indexes, non_blocking=True) + mem_indexes = dst_mem_indexes + assert req.tail_small_page_buffer_id is not None + self._copy_tail_state( + g_infer_context.radix_cache.small_page_buffers, + req.tail_small_page_buffer_id, + mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID, + ) + big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID) + + assert len(big_page_buffer_ids_cpu) == len(page_indexes) == len(page_readies) + + big_page_buffer_ids_cpu = torch.tensor( + big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu", pin_memory=True + ) + assert len(big_page_buffer_ids_cpu) <= self.big_page_ids_buffer_store.shape[0] + big_page_buffer_ids_gpu = self.big_page_ids_buffer_store[0 : len(big_page_buffer_ids_cpu)] + big_page_buffer_ids_gpu.copy_(big_page_buffer_ids_cpu, non_blocking=True) + + self._offload_target_pages(mem_indexes, big_page_buffer_ids_gpu, page_indexes, page_readies, cpu_cache_client) + self._copy_window_pages(big_page_buffer_ids_gpu, page_indexes, cpu_cache_client, page_readies) + return + + def _load_target_pages(self, mem_indexes, buffer_ids, page_indexes, cpu_cache_client): + raise NotImplementedError() + + def _offload_target_pages(self, mem_indexes, buffer_ids, page_indexes, page_readies, cpu_cache_client): + raise NotImplementedError() + + def _copy_target_state(self, source, source_id, destination_id): + # A full-attention target has no additional checkpoint payload. + return + + def _copy_tail_state(self, source, source_id, destination_id): + self._copy_target_state(source, source_id, destination_id) + window = getattr(self.mem_manager.big_page_buffers, "draft_window", None) + if window is not None: + from lightllm.common.basemodel.triton_kernel.hybrid_cpu_cache_copy import copy_checkpoint_state + + copy_checkpoint_state( + source.draft_window.get_state_cache(source_id), window.get_state_cache(destination_id) + ) + + def _copy_window_pages(self, buffer_ids, page_indexes, cpu_cache_client, page_readies=None): + window = getattr(self.mem_manager.big_page_buffers, "draft_window", None) + if window is None: + return + from lightllm.common.basemodel.triton_kernel.hybrid_cpu_cache_copy import copy_checkpoint_pages + + buffers = (window.kv, window.ends, window.counts) + pages = self.cpu_cache_config.get_window_views(cpu_cache_client.cpu_kv_cache_tensor, get_current_rank_in_dp()) + if page_readies is None: + copy_checkpoint_pages(pages, buffers, page_indexes, buffer_ids) + else: + copy_checkpoint_pages(buffers, pages, buffer_ids, page_indexes, page_readies) diff --git a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py b/lightllm/common/kv_cache_mem_manager/operator/linear_att.py index 71158ac97a..ebcf5e37a0 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/linear_att.py +++ b/lightllm/common/kv_cache_mem_manager/operator/linear_att.py @@ -1,77 +1,24 @@ import torch -import triton -from typing import List -from typing import TYPE_CHECKING -from .base import BaseMemManagerOperator +from .hybrid import HybridAttMemOperator from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size from lightllm.utils.log_utils import init_logger -from lightllm.common.state_cache_manager import LinearAttCacheConfig - -if TYPE_CHECKING: - from lightllm.server.multi_level_kv_cache.cpu_cache_client import CpuKvCacheClient - from lightllm.server.router.model_infer.infer_batch import InferReq logger = init_logger(__name__) -class LinearAttMemOperator(BaseMemManagerOperator): +class LinearAttMemOperator(HybridAttMemOperator): """ 只用于非量化的linear att 混合 full att的模型，列入 qwen3.5 """ def __init__(self, mem_manager): super().__init__(mem_manager) - self.linear_config = LinearAttCacheConfig.load_from_args() + self.linear_config = mem_manager.linear_config - def load_cpu_cache_to_gpu( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ): - assert mem_indexes.is_cuda and page_indexes.is_cuda + def _load_target_pages(self, mem_indexes, big_page_buffer_ids_gpu, page_indexes, cpu_cache_client): args = get_env_start_args() - assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) - assert len(mem_indexes) % args.linear_att_hash_page_size == 0 - assert args.cpu_cache_token_page_size == args.linear_att_hash_page_size * args.linear_att_page_block_num - from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextMemManager - - mem_manager: Qwen3NextMemManager = self.mem_manager - - big_page_num = len(mem_indexes) // args.cpu_cache_token_page_size - max_kv_len = (req.cur_kv_len // args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size - assert max_kv_len % args.cpu_cache_token_page_size == 0 - - big_page_buffer_ids_cpu = [] - for i in range(big_page_num): - page_id = mem_manager.big_page_buffers.alloc_one_state_cache() - assert page_id is not None - req.hybrid_len_to_big_page_id[max_kv_len] = page_id - big_page_buffer_ids_cpu.append(page_id) - max_kv_len -= args.cpu_cache_token_page_size - assert max_kv_len % args.cpu_cache_token_page_size == 0 - - big_page_buffer_ids_cpu.reverse() - - # 碎页情况的处理 - has_tail_page = len(mem_indexes) % args.cpu_cache_token_page_size != 0 - if has_tail_page: - padded_token_num = triton.cdiv( - len(mem_indexes), args.cpu_cache_token_page_size - ) * args.cpu_cache_token_page_size - len(mem_indexes) - mem_indexes = torch.nn.functional.pad(mem_indexes, (0, padded_token_num), mode="constant", value=-1) - - # 将对应的小叶数据拷贝到临时的大页上，再从大页上拷贝到对应的运行态页面上 - big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID) - - big_page_buffer_ids_gpu = torch.tensor(big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu").cuda( - non_blocking=True - ) - - assert len(big_page_buffer_ids_gpu) == len(page_indexes) - + mem_manager = self.mem_manager from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( copy_cpu_cache_to_kv_buffer, ) @@ -90,85 +37,9 @@ def load_cpu_cache_to_gpu( linear_config=self.linear_config, ) - from lightllm.server.router.model_infer.infer_batch import g_infer_context - - g_infer_context.req_manager.restore_big_page_state( - big_page_buffer_idx=big_page_buffer_ids_cpu[-1], - req=req, - ) - - return - - def offload_gpu_kv_to_cpu_cache( - self, - mem_indexes: torch.Tensor, - page_indexes: torch.Tensor, - page_readies: torch.Tensor, - cpu_cache_client: "CpuKvCacheClient", - req: "InferReq", - ): + def _offload_target_pages(self, mem_indexes, big_page_buffer_ids_gpu, page_indexes, page_readies, cpu_cache_client): args = get_env_start_args() - if not hasattr(self, "big_page_ids_buffer_store"): - self.big_page_ids_buffer_store = torch.empty((1024 * 1024 * 4,), dtype=torch.int64, device="cuda") - # 多申请3个cpu cache token page size，用于处理碎页情况，碎页情况需要将对应的大页数据拷贝到临时的大页上， - # 再从大页上拷贝到对应的运行态页面上 - self.mem_indexes_buffer = torch.empty( - (args.max_req_total_len + 3 * args.cpu_cache_token_page_size,), dtype=torch.int32, device="cuda" - ) - - assert mem_indexes.is_cuda and page_indexes.is_cuda and page_readies.is_cuda - - assert len(mem_indexes) % args.linear_att_hash_page_size == 0 - assert triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) == len(page_indexes) - - from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextMemManager - - mem_manager: Qwen3NextMemManager = self.mem_manager - - from lightllm.server.router.model_infer.infer_batch import g_infer_context - - big_page_buffer_ids_cpu = g_infer_context.radix_cache.get_big_page_ids_by_node(req.shared_kv_node) - max_kv_len = (len(mem_indexes) // args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size - start_kv_len = (len(big_page_buffer_ids_cpu) + 1) * args.cpu_cache_token_page_size - for seq_len in range(start_kv_len, max_kv_len + 1, args.cpu_cache_token_page_size): - page_id = req.hybrid_len_to_big_page_id[seq_len] - big_page_buffer_ids_cpu.append(page_id) - - if len(mem_indexes) % args.cpu_cache_token_page_size != 0: - # 存在不满大页的碎页的页面存在需要复制的情况 - dst_len = triton.cdiv(len(mem_indexes), args.cpu_cache_token_page_size) * args.cpu_cache_token_page_size - assert dst_len <= self.mem_indexes_buffer.shape[0] - dst_mem_indexes = self.mem_indexes_buffer[0:dst_len].fill_(-1) - dst_mem_indexes[0 : len(mem_indexes)].copy_(mem_indexes, non_blocking=True) - mem_indexes = dst_mem_indexes - assert req.tail_small_page_buffer_id is not None - from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( - copy_linear_att_state_to_linear_att_state, - ) - - src_conv_state, src_ssm_state = g_infer_context.radix_cache.small_page_buffers.get_state_cache( - buffer_idx=req.tail_small_page_buffer_id - ) - dst_conv_state, dst_ssm_state = mem_manager.big_page_buffers.get_state_cache( - buffer_idx=mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID, - ) - copy_linear_att_state_to_linear_att_state( - src_conv_state=src_conv_state, - src_ssm_state=src_ssm_state, - dst_conv_state=dst_conv_state, - dst_ssm_state=dst_ssm_state, - ) - big_page_buffer_ids_cpu.append(mem_manager.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID) - - assert len(big_page_buffer_ids_cpu) == len(page_indexes) == len(page_readies) - - big_page_buffer_ids_cpu = torch.tensor( - big_page_buffer_ids_cpu, dtype=torch.int64, device="cpu", pin_memory=True - ) - assert len(big_page_buffer_ids_cpu) <= self.big_page_ids_buffer_store.shape[0] - big_page_buffer_ids_gpu = self.big_page_ids_buffer_store[0 : len(big_page_buffer_ids_cpu)] - big_page_buffer_ids_gpu.copy_(big_page_buffer_ids_cpu, non_blocking=True) - + mem_manager = self.mem_manager from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( copy_kv_buffer_to_cpu_cache, ) @@ -187,7 +58,20 @@ def offload_gpu_kv_to_cpu_cache( big_page_token_num=args.cpu_cache_token_page_size, linear_config=self.linear_config, ) - return + + def _copy_target_state(self, source, source_id, destination_id): + from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( + copy_linear_att_state_to_linear_att_state, + ) + + src_conv_state, src_ssm_state = source.get_state_cache(source_id) + dst_conv_state, dst_ssm_state = self.mem_manager.big_page_buffers.get_state_cache(destination_id) + copy_linear_att_state_to_linear_att_state( + src_conv_state=src_conv_state, + src_ssm_state=src_ssm_state, + dst_conv_state=dst_conv_state, + dst_ssm_state=dst_ssm_state, + ) def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: torch.Tensor): # Qwen3Next 需要调整 layer_index @@ -201,11 +85,3 @@ def copy_kv_to_mem_manager(self, layer_index: int, mem_index: torch.Tensor, kv: destindex_copy_kv(kv, mem_index, mem_manager.kv_buffer[layer_index]) return - - def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): - from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer - - copy_kv_buffer_to_kv_buffer( - src_mem_index.cuda(non_blocking=True), dst_mem_index.cuda(non_blocking=True), self.mem_manager.kv_buffer - ) - return diff --git a/lightllm/common/kv_cache_mem_manager/operator/windowed_mtp.py b/lightllm/common/kv_cache_mem_manager/operator/windowed_mtp.py new file mode 100644 index 0000000000..daee314e79 --- /dev/null +++ b/lightllm/common/kv_cache_mem_manager/operator/windowed_mtp.py @@ -0,0 +1,32 @@ +from .hybrid import HybridAttMemOperator +from .normal import NormalMemOperator +from lightllm.utils.envs_utils import get_env_start_args +from lightllm.utils.dist_utils import get_current_rank_in_dp +from lightllm.common.basemodel.triton_kernel.hybrid_cpu_cache_copy import copy_full_att_pages + + +class WindowedMTPMemOperator(HybridAttMemOperator, NormalMemOperator): + """Ordinary target KV plus request-local draft windows in hybrid CPU pages.""" + + def _load_target_pages(self, mem_indexes, buffer_ids, page_indexes, cpu_cache_client): + copy_full_att_pages( + self.mem_manager.kv_buffer, + cpu_cache_client.cpu_kv_cache_tensor, + mem_indexes, + page_indexes, + self.cpu_cache_config, + get_current_rank_in_dp(), + get_env_start_args().cpu_cache_token_page_size, + ) + + def _offload_target_pages(self, mem_indexes, buffer_ids, page_indexes, page_readies, cpu_cache_client): + copy_full_att_pages( + self.mem_manager.kv_buffer, + cpu_cache_client.cpu_kv_cache_tensor, + mem_indexes, + page_indexes, + self.cpu_cache_config, + get_current_rank_in_dp(), + get_env_start_args().cpu_cache_token_page_size, + page_readies=page_readies, + ) diff --git a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py index e03e9f08da..64c51daa69 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -5,6 +5,7 @@ from lightllm.utils.envs_utils import get_env_start_args from lightllm.common.state_cache_manager import LinearAttCacheConfig, LinearAttCacheManager from .operator import LinearAttMemOperator +from .windowed_mtp import WindowKVPageHelper from typing import Tuple, Any, List logger = init_logger(__name__) @@ -32,13 +33,7 @@ def get_att_input_params(self, layer_index: int) -> Tuple[Any, Any]: layer_index = self.linear_config.get_full_att_kv_layer_index(layer_index) return super().get_att_input_params(layer_index) - def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): - super()._init_buffers(size, dtype, head_num, head_dim, layer_num) - # TODO 初始化线性 att 对应的部分 buffer. - self._init_linear_att_buffers() - return - - def _init_linear_att_buffers(self): + def _init_state_cache_buffers(self): big_page_token_num = ( get_env_start_args().linear_att_page_block_num * get_env_start_args().linear_att_hash_page_size ) @@ -49,40 +44,24 @@ def _init_linear_att_buffers(self): size=triton.cdiv(self.size, big_page_token_num) + 2, linear_config=self.linear_config, keep_num=2, + window_config=self.window_state_config, ) self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 2 self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 1 return - def _free_buffers(self): - super()._free_buffers() - self._free_linear_att_buffers() - return - - def _free_linear_att_buffers(self): - self.big_page_buffers = None - return - def write_to_shm(self, req_manager): self.req_to_conv_state = req_manager.req_to_conv_state self.req_to_ssm_state = req_manager.req_to_ssm_state - # super().write_to_shm() 会用 ForkingPickler 序列化本对象，torch 在 dump 时会把 - # CPU tensor 的 storage 原地迁到共享内存，使本进程大页 state cache 原本 - # pinned(cudaHostAlloc) 的内存退化为普通 shm mmap，之后 Triton kernel 携带该指针 - # 启动会报 "Pointer argument cannot be accessed from Triton (cpu tensor?)"。 - # 跨进程消费方并不使用 cpu 侧大页 state cache，序列化期间临时剔除以保住 pinned。 - big_page_buffers = self.big_page_buffers - self.big_page_buffers = None - try: - return super().write_to_shm(req_manager) - finally: - self.big_page_buffers = big_page_buffers - - def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: - kv_move_buffer = super().alloc_paged_kv_move_buffer(page_num, page_size) - Qwen3NextLinearAttPageHelper(self).assert_page_size() - return kv_move_buffer + return super().write_to_shm(req_manager) + + def _assert_att_state_page_size(self): + helper = Qwen3NextLinearAttPageHelper(self) + if self.windowed_draft_kv is None: + helper.assert_page_size() + else: + WindowKVPageHelper(self, offset=helper.state_nbytes).assert_page_size() def write_mem_to_page_kv_move_buffer( self, @@ -109,6 +88,10 @@ def write_mem_to_page_kv_move_buffer( helper = Qwen3NextLinearAttPageHelper(self) dp_mems = helper.get_dp_mems(mem_managers, dp_index, dp_world_size) helper.write_req_to_page(page_index=page_index, req_idx=req_idx, dp_mems=dp_mems) + if self.windowed_draft_kv is not None: + WindowKVPageHelper(self, offset=helper.state_nbytes).copy_req_page( + page_index, req_idx, dp_mems, mode="write" + ) return def read_page_kv_move_buffer_to_mem( @@ -136,6 +119,10 @@ def read_page_kv_move_buffer_to_mem( helper = Qwen3NextLinearAttPageHelper(self) dp_mems = helper.get_dp_mems(mem_managers, dp_index, dp_world_size) helper.read_page_to_req(page_index=page_index, req_idx=req_idx, dp_mems=dp_mems) + if self.windowed_draft_kv is not None: + WindowKVPageHelper(self, offset=helper.state_nbytes).copy_req_page( + page_index, req_idx, dp_mems, mode="read" + ) return diff --git a/lightllm/common/kv_cache_mem_manager/windowed_mtp.py b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py index 7ef7eb89cc..8230652418 100644 --- a/lightllm/common/kv_cache_mem_manager/windowed_mtp.py +++ b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py @@ -110,6 +110,28 @@ def reset(self): self.ends.zero_() self.counts.zero_() + def reset_req(self, req_idx): + self.ends[req_idx] = 0 + self.counts[req_idx] = 0 + + def save_checkpoint(self, req_idx, buffers, slot): + # Copy contiguous per-layer windows, preserving the physical p % W layout. + # The prefill stream event covers these D2H copies before radix publication. + for layer in range(self.kv.shape[0]): + buffers.kv[slot, layer].copy_(self.kv[layer, req_idx], non_blocking=True) + buffers.ends[slot : slot + 1].copy_(self.ends[req_idx : req_idx + 1], non_blocking=True) + buffers.counts[slot : slot + 1].copy_(self.counts[req_idx : req_idx + 1], non_blocking=True) + + def restore_checkpoint(self, req_idx, buffers, slot): + for layer in range(self.kv.shape[0]): + self.kv[layer, req_idx].copy_(buffers.kv[slot, layer], non_blocking=True) + self.ends[req_idx : req_idx + 1].copy_(buffers.ends[slot : slot + 1], non_blocking=True) + self.counts[req_idx : req_idx + 1].copy_(buffers.counts[slot : slot + 1], non_blocking=True) + # Small-page matching can release the source node immediately after return. + # Finish H2D before its CPU slot can be evicted and overwritten. + if self.kv.is_cuda: + torch.cuda.current_stream(self.kv.device).synchronize() + def prepare(self, reqs, features, starts, first, lengths, max_new): n = min(max_new, self.capacity) packed = features.new_empty((reqs.numel() * n, features.shape[-1])) @@ -166,3 +188,68 @@ def retained_positions(self, reqs): slots = torch.arange(self.capacity, device=reqs.device)[None, :] positions = (end - self.capacity).clamp_min(0) + slots return positions.masked_fill(slots >= end.clamp_max(self.capacity), -1) + + +class WindowKVPageHelper: + """Append a request's draft window to the existing PD att_state page. + + The page uses global KV heads, just like ordinary PD KV pages, so P and D + may use different TP sizes. The ring's physical slot order is preserved. + """ + + def __init__(self, mem_manager, offset=0): + from lightllm.utils.envs_utils import get_env_start_args + + self.mem_manager = mem_manager + store = mem_manager.windowed_draft_kv + args = get_env_start_args() + tp_world_size = args.tp // args.dp + layers, _, window, heads, dim = store.kv.shape + self.shape = (window, layers, heads * tp_world_size, dim) + self.kv_offset = (offset + 7) // 8 * 8 + self.kv_nbytes = window * layers * heads * tp_world_size * dim * store.kv.element_size() + self.end_offset = (self.kv_offset + self.kv_nbytes + 7) // 8 * 8 + self.state_nbytes = self.end_offset + 12 + + def assert_page_size(self): + page = self.mem_manager.kv_move_buffer[0] + page_nbytes = page.numel() * page.element_size() + assert page_nbytes >= self.state_nbytes, ( + f"PD page bytes {page_nbytes} is smaller than attention state bytes {self.state_nbytes}; " + "increase --pd_kv_page_size on both P and D nodes" + ) + + def copy_req_page(self, page_index, req_idx, dp_mems, mode): + from lightllm.common.kv_trans_kernel.nixl_kv_trans import page_io + + assert mode in ("read", "write") + assert req_idx is not None + page = self.mem_manager.kv_move_buffer[page_index].view(torch.uint8).reshape(-1) + kv_page = ( + page[self.kv_offset : self.kv_offset + self.kv_nbytes] + .view(self.mem_manager.windowed_draft_kv.kv.dtype) + .view(self.shape) + ) + end_page = page[self.end_offset : self.end_offset + 8].view(torch.int64) + count_page = page[self.end_offset + 8 : self.end_offset + 12].view(torch.int32) + window = self.shape[0] + indexes = torch.arange(req_idx * window, (req_idx + 1) * window, dtype=torch.int64, device=page.device) + for tp_index, mem in enumerate(dp_mems): + store = mem.windowed_draft_kv + layers, requests, capacity, heads, dim = store.kv.shape + assert (capacity, layers, heads * len(dp_mems), dim) == self.shape + assert store.kv.dtype == kv_page.dtype + page_io( + mem_indexes=indexes, + page_tensor=kv_page, + kv_buffer=store.kv.view(layers, requests * capacity, heads, dim), + tp_index=tp_index, + tp_world_size=len(dp_mems), + mode=mode, + ) + if mode == "read": + store.ends[req_idx : req_idx + 1].copy_(end_page, non_blocking=True) + store.counts[req_idx : req_idx + 1].copy_(count_page, non_blocking=True) + elif tp_index == 0: + end_page.copy_(store.ends[req_idx : req_idx + 1], non_blocking=True) + count_page.copy_(store.counts[req_idx : req_idx + 1], non_blocking=True) diff --git a/lightllm/common/req_manager/__init__.py b/lightllm/common/req_manager/__init__.py index abcd4f491e..fb7eba1211 100644 --- a/lightllm/common/req_manager/__init__.py +++ b/lightllm/common/req_manager/__init__.py @@ -1,6 +1,13 @@ from .base import ReqManager from .linear_att import ReqManagerForMamba from .hybrid_base import HybridAttentionReqManager +from .windowed_mtp import ReqManagerForWindowedMTP from .req_sampling_params import ReqSamplingParamsManager -__all__ = ["ReqManager", "HybridAttentionReqManager", "ReqManagerForMamba", "ReqSamplingParamsManager"] +__all__ = [ + "ReqManager", + "HybridAttentionReqManager", + "ReqManagerForMamba", + "ReqManagerForWindowedMTP", + "ReqSamplingParamsManager", +] diff --git a/lightllm/common/req_manager/linear_att.py b/lightllm/common/req_manager/linear_att.py index ccab8a06f6..acceefa7be 100644 --- a/lightllm/common/req_manager/linear_att.py +++ b/lightllm/common/req_manager/linear_att.py @@ -57,15 +57,20 @@ def init_hybrid_attention_state(self, req: "InferReq"): self.req_to_ssm_state.buffer[:, ssm_start : ssm_start + (self.mtp_step + 1), ...].fill_(0) if self.req_to_mtp_state_index is not None: self.req_to_mtp_state_index[req.req_idx] = 0 + if self.mem_manager.windowed_draft_kv is not None: + self.mem_manager.windowed_draft_kv.reset_req(req.req_idx) return def create_small_page_cache_manager(self, size: int): - self.small_page_buffers = LinearAttCacheManager(size=size, linear_config=self.linear_config) + self.small_page_buffers = LinearAttCacheManager( + size=size, linear_config=self.linear_config, window_config=self.mem_manager.window_state_config + ) return self.small_page_buffers def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], buffer_indexes: List[int]): from lightllm.common.basemodel.triton_kernel.linear_att_copy import copy_linear_att_state_to_kv_buffer + buffer_indexes_cpu = buffer_indexes buffer_indexes = torch.tensor(buffer_indexes, dtype=torch.int32, device="cpu").cuda(non_blocking=True) state_cache_manager = self.big_page_buffers copy_linear_att_state_to_kv_buffer( @@ -77,6 +82,12 @@ def save_big_page_states(self, b_req_idx: torch.Tensor, req_indexes: List[int], cpu_kv_ssm_state=state_cache_manager.ssm_state_cache.buffer, mtp_step=self.mtp_step, ) + if state_cache_manager.draft_window is not None: + for req_idx, buffer_idx in zip(req_indexes, buffer_indexes_cpu): + if buffer_idx != -1: + self.mem_manager.windowed_draft_kv.save_checkpoint( + req_idx, state_cache_manager.draft_window, buffer_idx + ) return def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: LinearAttCacheManager): @@ -87,6 +98,8 @@ def save_state(self, req_idx: int, buffer_idx: int, state_cache_manager: LinearA dst_conv_state, dst_ssm_state = state_cache_manager.get_state_cache(buffer_idx=buffer_idx) dst_conv_state.copy_(gpu_conv_state, non_blocking=True) dst_ssm_state.copy_(gpu_ssm_state, non_blocking=True) + if state_cache_manager.draft_window is not None: + self.mem_manager.windowed_draft_kv.save_checkpoint(req_idx, state_cache_manager.draft_window, buffer_idx) def get_mamba_cache(self, layer_idx_in_all: int): assert ( @@ -116,6 +129,10 @@ def restore_state(self, req: "InferReq", state_cache_manager: LinearAttCacheMana conv_cache_width = conv_state.shape[-1] self.req_to_conv_state.buffer[:, conv_dest, ..., :conv_cache_width] = conv_state self.req_to_ssm_state.buffer[:, ssm_dest, ...] = ssm_state + if state_cache_manager.draft_window is not None: + self.mem_manager.windowed_draft_kv.restore_checkpoint( + req.req_idx, state_cache_manager.draft_window, buffer_idx + ) if self.req_to_mtp_state_index is not None: self.req_to_mtp_state_index[req.req_idx] = 0 return diff --git a/lightllm/common/req_manager/windowed_mtp.py b/lightllm/common/req_manager/windowed_mtp.py new file mode 100644 index 0000000000..a20591d302 --- /dev/null +++ b/lightllm/common/req_manager/windowed_mtp.py @@ -0,0 +1,21 @@ +from lightllm.common.state_cache_manager import WindowStateCacheManager +from .hybrid_base import HybridAttentionReqManager + + +class ReqManagerForWindowedMTP(HybridAttentionReqManager): + """Full-attention target with request-local draft state in hybrid checkpoints.""" + + def create_small_page_cache_manager(self, size: int): + self.small_page_buffers = WindowStateCacheManager( + size=size, config=self.mem_manager.window_state_config, dtype=self.mem_manager.dtype + ) + return self.small_page_buffers + + def init_hybrid_attention_state(self, req): + self.mem_manager.windowed_draft_kv.reset_req(req.req_idx) + + def save_state(self, req_idx, buffer_idx, state_cache_manager): + self.mem_manager.windowed_draft_kv.save_checkpoint(req_idx, state_cache_manager.draft_window, buffer_idx) + + def restore_state(self, req, state_cache_manager, buffer_idx): + self.mem_manager.windowed_draft_kv.restore_checkpoint(req.req_idx, state_cache_manager.draft_window, buffer_idx) diff --git a/lightllm/common/state_cache_manager/__init__.py b/lightllm/common/state_cache_manager/__init__.py index af0479a158..d7055dd96d 100644 --- a/lightllm/common/state_cache_manager/__init__.py +++ b/lightllm/common/state_cache_manager/__init__.py @@ -1,13 +1,20 @@ from .base import StateCacheManager from .layer_cache import LayerCache from .linear_att import LinearAttCacheConfig, LinearAttCacheManager +from .windowed_mtp import WindowStateCacheManager -def get_hybrid_cache_config(): +def get_hybrid_cache_config(linear_config=None): """Return the model-specific layout used by hybrid CPU/disk cache pages.""" from lightllm.utils.config_utils import is_linear_att_mixed_model from lightllm.utils.envs_utils import get_env_start_args + from .windowed_mtp import WindowedMTPCacheConfig - if is_linear_att_mixed_model(get_env_start_args().model_dir): - return LinearAttCacheConfig.load_from_args() + args = get_env_start_args() + if linear_config is None and is_linear_att_mixed_model(args.model_dir): + linear_config = LinearAttCacheConfig.load_from_args() + if args.mtp_draft_kv_mode == "window": + return WindowedMTPCacheConfig.load_from_args(linear_config) + if linear_config is not None: + return linear_config raise ValueError("No hybrid state-cache layout registered for this model") diff --git a/lightllm/common/state_cache_manager/linear_att.py b/lightllm/common/state_cache_manager/linear_att.py index aa7f5e439a..cecbc4da46 100644 --- a/lightllm/common/state_cache_manager/linear_att.py +++ b/lightllm/common/state_cache_manager/linear_att.py @@ -8,6 +8,7 @@ from .base import StateCacheManager from .layer_cache import LayerCache +from .windowed_mtp import WindowStateBuffers, WindowStateCacheConfig logger = init_logger(__name__) @@ -158,9 +159,13 @@ def __init__( size: int, linear_config: LinearAttCacheConfig, keep_num: int = 0, # 用于记录需要保留的缓存数量，用于支持含有 linear_att 的如qwen3.5 模型的cpu cache的碎页处理。 + window_config: WindowStateCacheConfig = None, ): super().__init__(size, keep_num) self.linear_config = linear_config + self.draft_window = ( + WindowStateBuffers(size, window_config, linear_config.full_att_dtype) if window_config is not None else None + ) # init the layer cache self.conv_state_cache = LayerCache( size=self.size, @@ -187,5 +192,7 @@ def get_state_cache(self, buffer_idx: int): def clear_to_init_state(self): self.conv_state_cache.buffer.zero_() self.ssm_state_cache.buffer.zero_() + if self.draft_window is not None: + self.draft_window.clear_to_init_state() super().clear_to_init_state() return diff --git a/lightllm/common/state_cache_manager/windowed_mtp.py b/lightllm/common/state_cache_manager/windowed_mtp.py new file mode 100644 index 0000000000..c0b3a36cce --- /dev/null +++ b/lightllm/common/state_cache_manager/windowed_mtp.py @@ -0,0 +1,171 @@ +"""Window state configuration and checkpoints shared by hybrid attention caches.""" + +import json +from dataclasses import dataclass +from pathlib import Path + +import torch + +from lightllm.utils.log_utils import init_logger +from lightllm.common.build_utils import repair_config +from .base import StateCacheManager +from .layer_cache import LayerCache + + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class WindowStateCacheConfig: + layers: int + kv_heads: int + head_dim: int + window: int + + @classmethod + def from_model_config(cls, config, window, tp_world_size): + return cls( + layers=config["n_layer"], + kv_heads=config["num_key_value_heads"] // tp_world_size, + head_dim=config.get("head_dim", config["n_embed"] // config["num_attention_heads"]), + window=window, + ) + + def get_state_shape(self): + return (self.window, 2 * self.kv_heads, self.head_dim) + + def get_state_bytes(self, element_size): + return self.layers * self.window * 2 * self.kv_heads * self.head_dim * element_size + 12 + + +def load_window_state_config(args): + config = json.loads((Path(args.mtp_draft_model_dir[0]) / "config.json").read_text()) + # Match the existing draft initialization order: base aliases, Llama KV + # defaults, then the overrides applied by Qwen3.5 adapters. + # Do not normalize again after merging: runtime also retains those aliases. + merge_nested = config.get("model_type") in ("qwen3_5", "qwen3_5_text") + repair_config(config, same_names=["num_attention_heads", "n_head"]) + repair_config(config, same_names=["hidden_size", "n_embd", "n_embed"]) + repair_config(config, same_names=["num_hidden_layers", "n_layer"]) + config.setdefault("num_key_value_heads", config["num_attention_heads"]) + if merge_nested: + config.update(config.get("dflash_config", {})) + return WindowStateCacheConfig.from_model_config(config, args.mtp_draft_window_size, args.tp // args.dp) + + +class WindowStateBuffers: + """Pinned CPU payload; the containing StateCacheManager owns all slot IDs.""" + + def __init__(self, size: int, config: WindowStateCacheConfig, dtype: torch.dtype): + self.config = config + self.kv = LayerCache( + size=size, + dtype=dtype, + shape=config.get_state_shape(), + layer_num=config.layers, + device="cpu", + size_first=True, + ).buffer + self.ends = torch.zeros(size, dtype=torch.int64, device="cpu", pin_memory=True) + self.counts = torch.zeros(size, dtype=torch.int32, device="cpu", pin_memory=True) + logger.info( + f"draft window checkpoint slots={size}, pinned_bytes={size * config.get_state_bytes(dtype.itemsize)}" + ) + + def get_state_cache(self, buffer_idx: int): + return self.kv[buffer_idx], self.ends[buffer_idx : buffer_idx + 1], self.counts[buffer_idx : buffer_idx + 1] + + def clear_to_init_state(self): + self.kv.zero_() + self.ends.zero_() + self.counts.zero_() + + +class WindowStateCacheManager(StateCacheManager): + def __init__(self, size: int, config: WindowStateCacheConfig, dtype: torch.dtype, keep_num: int = 0): + super().__init__(size, keep_num) + self.draft_window = WindowStateBuffers(size, config, dtype) + + def get_state_cache(self, buffer_idx: int): + return self.draft_window.get_state_cache(buffer_idx) + + def clear_to_init_state(self): + self.draft_window.clear_to_init_state() + super().clear_to_init_state() + + +@dataclass +class WindowedMTPCacheConfig: + """CPU pages contain target KV/state followed by one draft window per TP rank.""" + + window_config: WindowStateCacheConfig + dtype: torch.dtype + tp_world_size: int + full_att_all_num_kv_heads: int + full_att_head_dim: int + full_att_layer_num: int + linear_config: object = None + + @classmethod + def load_from_args(cls, linear_config=None): + from lightllm.utils.config_utils import get_num_key_value_heads, get_head_dim, get_layer_num + from lightllm.utils.envs_utils import get_env_start_args + from lightllm.utils.torch_dtype_utils import get_torch_dtype + + args = get_env_start_args() + return cls( + window_config=load_window_state_config(args), + dtype=get_torch_dtype(args.data_type), + tp_world_size=args.tp // args.dp, + full_att_all_num_kv_heads=get_num_key_value_heads(args.model_dir), + full_att_head_dim=get_head_dim(args.model_dir), + full_att_layer_num=( + linear_config.get_full_att_kv_layer_num_with_draft_model() + if linear_config is not None + else get_layer_num(args.model_dir) + ), + linear_config=linear_config, + ) + + def get_cpu_cache_full_att_bytes(self): + if self.linear_config is not None: + return self.linear_config.get_cpu_cache_full_att_bytes() + from lightllm.utils.envs_utils import get_env_start_args + + args = get_env_start_args() + page_tokens = args.linear_att_hash_page_size * args.linear_att_page_block_num + assert page_tokens == args.cpu_cache_token_page_size + return ( + page_tokens + * self.full_att_layer_num + * 2 + * self.full_att_all_num_kv_heads + * self.full_att_head_dim + * self.dtype.itemsize + ) + + def get_cpu_cache_window_offset(self): + if self.linear_config is not None: + return self.linear_config.get_cpu_cache_big_page_bytes() + return (self.get_cpu_cache_full_att_bytes() + 15) // 16 * 16 + + def get_cpu_cache_window_rank_bytes(self): + kv_bytes = self.window_config.get_state_bytes(self.dtype.itemsize) - 12 + # Align int64 ends and each rank's payload; counts is int32. + return (((kv_bytes + 7) // 8 * 8 + 12) + 15) // 16 * 16 + + def get_cpu_cache_big_page_bytes(self): + return self.get_cpu_cache_window_offset() + self.tp_world_size * self.get_cpu_cache_window_rank_bytes() + + def get_window_views(self, cpu_cache_tensor, tp_rank): + assert 0 <= tp_rank < self.tp_world_size + pages = cpu_cache_tensor.view(cpu_cache_tensor.shape[0], -1).view(torch.uint8) + assert pages.shape[1] == self.get_cpu_cache_big_page_bytes() + start = self.get_cpu_cache_window_offset() + tp_rank * self.get_cpu_cache_window_rank_bytes() + kv_bytes = self.window_config.get_state_bytes(self.dtype.itemsize) - 12 + end_offset = start + (kv_bytes + 7) // 8 * 8 + return ( + pages[:, start : start + kv_bytes].view(self.dtype), + pages[:, end_offset : end_offset + 8].view(torch.int64), + pages[:, end_offset + 8 : end_offset + 12].view(torch.int32), + ) diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index b63a6beadc..01f4e6add4 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -6,7 +6,6 @@ ) from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.attention.fa3.windowed_mtp import WindowedMTPAttBackend -from lightllm.utils.windowed_mtp import validate_windowed_mtp, window_capacity from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput from lightllm.models.llama.model import LlamaTpPartModel from lightllm.models.draft_registry import DraftModelRegistry @@ -46,9 +45,7 @@ def _verify_params(self): assert self.args.mtp_step <= self.config["block_size"] self.config["block_size"] = self.args.mtp_step - self.uses_windowed_draft_kv = getattr(self.args, "mtp_draft_kv_mode", "full") == "window" - if self.uses_windowed_draft_kv: - validate_windowed_mtp(self.args) + self.uses_windowed_draft_kv = self.args.mtp_draft_kv_mode == "window" def _init_custom(self): self._cos_cached = self.main_model._cos_cached @@ -68,11 +65,15 @@ def _init_mem_manager(self): layers=self.config["n_layer"], kv_heads=self.config["num_key_value_heads"] // self.tp_world_size_, head_dim=self.config.get("head_dim", self.config["n_embed"] // self.config["num_attention_heads"]), - window=window_capacity(self.args), + window=self.args.mtp_draft_window_size, ) def _init_att_backend(self): if self.uses_windowed_draft_kv: + from lightllm.utils.sgl_utils import flash_attn_with_kvcache + if flash_attn_with_kvcache is None: + raise NotImplementedError("Windowed MTP requires FA3") + self.prefill_att_backend = self.decode_att_backend = WindowedMTPAttBackend(model=self) return super()._init_att_backend() diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 2734ac3b5b..859be530e1 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -163,11 +163,6 @@ def _launch_subprocesses(args: StartArgs): f"{sorted(allowed_ep_decode_att_backends)}; flashinfer is not supported." ) - if args.mtp_draft_kv_mode == "window": - from lightllm.utils.windowed_mtp import validate_windowed_mtp - - validate_windowed_mtp(args) - # mtp params check if args.mtp_mode is not None: if args.mtp_draft_model_dir is None: @@ -182,6 +177,12 @@ def _launch_subprocesses(args: StartArgs): assert args.mtp_draft_model_dir is None assert args.mtp_step == 0 + if args.mtp_draft_kv_mode == "window": + assert args.mtp_draft_window_size > 0 + assert args.mtp_mode in ("dflash", "dspark"), "windowed draft KV requires dflash or dspark" + if args.llm_kv_type not in (None, "None"): + raise ValueError("windowed draft KV currently requires unquantized KV") + # automatically set visual_dp based on visual_tp and tp. # In visual proxy mode keep the caller-provided visual_dp / visual_tp. if not args.visual_use_proxy_mode and args.visual_tp < args.tp and args.tp % args.visual_tp == 0: @@ -275,7 +276,7 @@ def _launch_subprocesses(args: StartArgs): # 避免请求释放时将不完整的大页 state 写入 radix cache 并触发断言。 args.linear_att_page_block_num = 10000000 - if args.enable_cpu_cache and is_hybrid_att_model(args.model_dir): + if args.enable_cpu_cache and is_hybrid_att_model(args.model_dir, args): args.cpu_cache_token_page_size = args.linear_att_hash_page_size * args.linear_att_page_block_num logger.info(f"set cpu_cache_token_page_size to {args.cpu_cache_token_page_size} for hybrid att model") diff --git a/lightllm/server/core/objs/req.py b/lightllm/server/core/objs/req.py index b5a5bf685f..835c2b388d 100644 --- a/lightllm/server/core/objs/req.py +++ b/lightllm/server/core/objs/req.py @@ -215,7 +215,7 @@ def init( self.post_init() args = get_env_start_args() - if is_hybrid_att_model(args.model_dir): + if is_hybrid_att_model(args.model_dir, args): self._fill_hybrid_token_hash() if args.enable_cpu_cache: cpu_cache_hash_list, cpu_cache_page_len_list = self._calcu_hybrid_cpu_cache_page_len_list() diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index fe77e5f37b..74f63d7c86 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -218,12 +218,6 @@ def init_model(self, kvargs): [rank for rank in range(self.global_world_size)], backend="nccl" ) - if self.args.run_mode in ["prefill", "decode"] or self.args.enable_dp_prompt_cache_fetch: - # 如果存在需要跨进程使用mem manger的特性，则将mem manager写入到 shm中，方便 - # 读取 - self.model.mem_manager.write_to_shm(req_manager=self.model.req_manager) - dist.barrier(group=self.node_nccl_group) - # 同一 DP 组内只需主 rank 初始化真实的 capture buffer 并执行后续相关操作； # 非主 rank 不需要分配 buffer，避免重复占用内存。 if self.is_master_in_dp: @@ -239,9 +233,6 @@ def init_model(self, kvargs): self.init_custom() - if self.args.enable_dp_prompt_cache_fetch: - self.init_dp_kv_shared() - self.shm_reqs_io_buffer = ShmObjsIOBuffer() # 只会在 pd pd 模式下才会使用，用于上传分块传输任务是否成功。 self.shm_pd_trans_io_buffer = ShmObjsIOBuffer(tail_str="pd") @@ -250,6 +241,14 @@ def init_model(self, kvargs): self.init_mtp_draft_model(model_kvargs) self.init_spec_engine() + # 主模型和 draft 的缓存均就绪后，再发布给跨进程 KV 访问方。 + if self.is_pd_mode or self.args.enable_dp_prompt_cache_fetch: + self.model.mem_manager.write_to_shm(req_manager=self.model.req_manager) + dist.barrier(group=self.node_nccl_group) + + if self.args.enable_dp_prompt_cache_fetch: + self.init_dp_kv_shared() + if self.args.enable_cpu_cache: self.multi_level_cache_module = MultiLevelKvCacheModule(self) diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 4bd78d887f..3bf766ffd9 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -468,9 +468,9 @@ def is_linear_att_mixed_model(model_path: str) -> bool: return False -def is_hybrid_att_model(model_path: str) -> bool: - """Models whose non-full attention state follows hybrid checkpoint pages.""" - return is_linear_att_mixed_model(model_path) +def is_hybrid_att_model(model_path: str, args) -> bool: + """Whether target linear state or draft window snapshots require hybrid pages.""" + return is_linear_att_mixed_model(model_path) or args.mtp_draft_kv_mode == "window" def get_model_type(model_path: str) -> Optional[str]: diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index b3047b488e..11c8f8d1cc 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -62,7 +62,7 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": args = get_env_start_args() assert args.enable_cpu_cache - is_hybrid_model = is_hybrid_att_model(args.model_dir) + is_hybrid_model = is_hybrid_att_model(args.model_dir, args) mem_manager_class = None if is_hybrid_model else select_mem_manager_class() if is_hybrid_model: hybrid_config = get_hybrid_cache_config() diff --git a/lightllm/utils/profile_max_tokens.py b/lightllm/utils/profile_max_tokens.py index a1778e1769..7559e9e348 100644 --- a/lightllm/utils/profile_max_tokens.py +++ b/lightllm/utils/profile_max_tokens.py @@ -77,9 +77,9 @@ def profile_mtp_weight_memory(model): target_weight_bytes = torch.cuda.memory_allocated() - weight_memory_before draft_cache_bytes = 0 if getattr(model.args, "mtp_draft_kv_mode", "full") == "window": - from lightllm.utils.windowed_mtp import window_kv_pool_bytes - - draft_cache_bytes = window_kv_pool_bytes(model.args, torch.tensor([], dtype=model.data_type).element_size()) + from lightllm.common.state_cache_manager.windowed_mtp import load_window_state_config + config = load_window_state_config(model.args) + draft_cache_bytes = (model.args.running_max_req_size + 1) * config.get_state_bytes(model.data_type.itemsize) model.mem_fraction = get_mtp_adjusted_mem_fraction( mem_fraction=model.mem_fraction, target_weight_bytes=target_weight_bytes, diff --git a/lightllm/utils/windowed_mtp.py b/lightllm/utils/windowed_mtp.py deleted file mode 100644 index d7eb568cc6..0000000000 --- a/lightllm/utils/windowed_mtp.py +++ /dev/null @@ -1,53 +0,0 @@ -"""Opt-in bounded draft KV policy, independent of the proposal algorithm.""" -import json -from pathlib import Path - -from lightllm.common.build_utils import repair_config - - -def window_capacity(args): - if args.mtp_draft_window_size <= 0: - raise ValueError("mtp_draft_window_size must be positive") - return args.mtp_draft_window_size - - -def window_kv_pool_bytes(args, element_size): - config = json.loads((Path(args.mtp_draft_model_dir[0]) / "config.json").read_text()) - # Match the existing draft initialization order: base aliases, Llama KV - # defaults, then the overrides applied by Qwen3.5 adapters. - # Do not normalize again after merging: runtime also retains those aliases. - merge_nested = config.get("model_type") in ("qwen3_5", "qwen3_5_text") - repair_config(config, same_names=["num_attention_heads", "n_head"]) - repair_config(config, same_names=["hidden_size", "n_embd", "n_embed"]) - repair_config(config, same_names=["num_hidden_layers", "n_layer"]) - config.setdefault("num_key_value_heads", config["num_attention_heads"]) - if merge_nested: - config.update(config.get("dflash_config", {})) - layers = config["n_layer"] - heads = config["num_key_value_heads"] // args.tp - dim = config.get("head_dim", config["n_embed"] // config["num_attention_heads"]) - return (args.running_max_req_size + 1) * (layers * window_capacity(args) * 2 * heads * dim * element_size + 12) - - -def validate_windowed_mtp(args): - window_capacity(args) - if args.mtp_mode not in ("dflash", "dspark"): - raise ValueError("windowed draft KV requires dflash or dspark") - if args.run_mode != "normal" or args.dp != 1: - raise ValueError("windowed draft KV currently requires normal deployment and dp=1") - if not args.disable_dynamic_prompt_cache: - raise ValueError("windowed draft KV requires --disable_dynamic_prompt_cache") - if args.llm_kv_type not in (None, "None"): - raise ValueError("windowed draft KV currently requires unquantized KV") - for option in ( - "enable_cpu_cache", - "enable_dp_prompt_cache_fetch", - "diverse_mode", - "enable_tpsp_mix_mode", - "enable_prefill_decode_mixed", - "enable_decode_microbatch_overlap", - "enable_prefill_microbatch_overlap", - "mtp_dynamic_verify", - ): - if getattr(args, option, False): - raise ValueError(f"windowed draft KV does not yet support --{option}") diff --git a/unit_tests/common/test_windowed_mtp.py b/unit_tests/common/test_windowed_mtp.py new file mode 100644 index 0000000000..4f099667e3 --- /dev/null +++ b/unit_tests/common/test_windowed_mtp.py @@ -0,0 +1,134 @@ +from types import SimpleNamespace + +import torch + +from lightllm.common.kv_cache_mem_manager.windowed_mtp import WindowKVStore +from lightllm.common.state_cache_manager.windowed_mtp import ( + WindowedMTPCacheConfig, + WindowStateCacheConfig, +) + + +def _checkpoint_buffers(slots, layers, window, kv_heads, head_dim, dtype): + return SimpleNamespace( + kv=torch.zeros((slots, layers, window, 2 * kv_heads, head_dim), dtype=dtype), + ends=torch.zeros(slots, dtype=torch.int64), + counts=torch.zeros(slots, dtype=torch.int32), + ) + + +def test_retained_positions_follow_logical_window_order(): + store = WindowKVStore( + requests=3, + layers=1, + kv_heads=1, + head_dim=2, + dtype=torch.float32, + device="cpu", + window=4, + ) + store.ends.copy_(torch.tensor([0, 3, 7], dtype=torch.int64)) + store.counts.copy_(torch.tensor([0, 3, 4], dtype=torch.int32)) + + positions = store.retained_positions(torch.tensor([0, 1, 2], dtype=torch.int64)) + + assert torch.equal( + positions, + torch.tensor( + [ + [-1, -1, -1, -1], + [0, 1, 2, -1], + [3, 4, 5, 6], + ], + dtype=torch.int64, + ), + ) + + +def test_checkpoint_round_trip_preserves_physical_ring_layout(): + layers = 2 + window = 4 + kv_heads = 2 + head_dim = 3 + source_req = 1 + destination_req = 2 + checkpoint_slot = 1 + + store = WindowKVStore( + requests=3, + layers=layers, + kv_heads=kv_heads, + head_dim=head_dim, + dtype=torch.float32, + device="cpu", + window=window, + ) + expected_kv = torch.arange( + layers * window * 2 * kv_heads * head_dim, + dtype=torch.float32, + ).view(layers, window, 2 * kv_heads, head_dim) + store.kv[:, source_req].copy_(expected_kv) + store.ends[source_req] = 9 + store.counts[source_req] = window + + buffers = _checkpoint_buffers( + slots=2, + layers=layers, + window=window, + kv_heads=kv_heads, + head_dim=head_dim, + dtype=store.kv.dtype, + ) + store.save_checkpoint(source_req, buffers, checkpoint_slot) + + store.kv[:, destination_req].fill_(-1) + store.ends[destination_req] = -1 + store.counts[destination_req] = -1 + store.restore_checkpoint(destination_req, buffers, checkpoint_slot) + + assert torch.equal(buffers.kv[checkpoint_slot], expected_kv) + assert torch.equal(store.kv[:, destination_req], expected_kv) + assert store.ends[destination_req].item() == 9 + assert store.counts[destination_req].item() == window + + +class _TargetStateLayout: + def __init__(self, page_bytes): + self.page_bytes = page_bytes + + def get_cpu_cache_big_page_bytes(self): + return self.page_bytes + + +def test_cpu_page_keeps_tp_rank_window_states_separate(): + target_page_bytes = 32 + window_config = WindowStateCacheConfig(layers=2, kv_heads=1, head_dim=3, window=4) + config = WindowedMTPCacheConfig( + window_config=window_config, + dtype=torch.float16, + tp_world_size=2, + full_att_all_num_kv_heads=0, + full_att_head_dim=0, + full_att_layer_num=0, + linear_config=_TargetStateLayout(target_page_bytes), + ) + pages = torch.zeros((2, config.get_cpu_cache_big_page_bytes()), dtype=torch.uint8) + + rank0_kv, rank0_ends, rank0_counts = config.get_window_views(pages, tp_rank=0) + rank1_kv, rank1_ends, rank1_counts = config.get_window_views(pages, tp_rank=1) + rank0_kv.fill_(1) + rank1_kv.fill_(2) + rank0_ends.copy_(torch.tensor([[11], [12]], dtype=torch.int64)) + rank1_ends.copy_(torch.tensor([[21], [22]], dtype=torch.int64)) + rank0_counts.copy_(torch.tensor([[3], [4]], dtype=torch.int32)) + rank1_counts.copy_(torch.tensor([[1], [2]], dtype=torch.int32)) + + restored_rank0 = config.get_window_views(pages, tp_rank=0) + restored_rank1 = config.get_window_views(pages, tp_rank=1) + assert torch.all(restored_rank0[0] == 1) + assert torch.all(restored_rank1[0] == 2) + assert restored_rank0[1].tolist() == [[11], [12]] + assert restored_rank1[1].tolist() == [[21], [22]] + assert restored_rank0[2].tolist() == [[3], [4]] + assert restored_rank1[2].tolist() == [[1], [2]] + assert torch.count_nonzero(pages[:, :target_page_bytes]).item() == 0 From be685da6b922cc23592ea4d1ecdb99158924cf21 Mon Sep 17 00:00:00 2001 From: Zhao Xintong Date: Thu, 15 Oct 2026 11:37:40 +0000 Subject: [PATCH 5/9] fix --- .../basemodel/attention/fa3/windowed_mtp.py | 7 +- .../triton_kernel/destindex_copy_kv.py | 5 +- .../triton_kernel/hybrid_cpu_cache_copy.py | 136 ++------- .../triton_kernel/kv_cache_offload.py | 24 +- .../basemodel/triton_kernel/windowed_mtp.py | 184 ++++++++++++ .../kv_cache_mem_manager/mem_manager.py | 2 +- .../kv_cache_mem_manager/operator/hybrid.py | 17 +- .../operator/windowed_mtp.py | 41 +-- .../kv_cache_mem_manager/windowed_mtp.py | 242 ++++------------ .../state_cache_manager/windowed_mtp.py | 48 +-- lightllm/models/qwen3_5_dflash/model.py | 2 + lightllm/models/qwen3_5_dspark/model.py | 2 + .../layer_infer/transformer_layer_infer.py | 6 +- lightllm/models/qwen3_dflash/model.py | 62 ++-- lightllm/server/api_start.py | 4 +- lightllm/server/core/objs/req.py | 2 +- .../server/router/model_infer/infer_batch.py | 2 +- .../mtp_speculative/proposers/dflash.py | 9 +- .../mtp_speculative/proposers/dspark.py | 9 +- lightllm/utils/config_utils.py | 6 +- lightllm/utils/kv_cache_utils.py | 2 +- lightllm/utils/profile_max_tokens.py | 14 +- unit_tests/common/test_windowed_mtp.py | 274 +++++++++++++++--- .../mtp_speculative/test_windowed_draft.py | 95 ------ 24 files changed, 647 insertions(+), 548 deletions(-) create mode 100644 lightllm/common/basemodel/triton_kernel/windowed_mtp.py delete mode 100644 unit_tests/server/router/model_infer/mtp_speculative/test_windowed_draft.py diff --git a/lightllm/common/basemodel/attention/fa3/windowed_mtp.py b/lightllm/common/basemodel/attention/fa3/windowed_mtp.py index ad613b30de..510a23eb36 100644 --- a/lightllm/common/basemodel/attention/fa3/windowed_mtp.py +++ b/lightllm/common/basemodel/attention/fa3/windowed_mtp.py @@ -15,10 +15,13 @@ 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].long() + 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.counts.index_select(0, self.b_req_idx) + self.block_size + 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 decode_att(self, q, k, v, att_control: AttControl = AttControl(), alloc_func=torch.empty): return flash_attn_with_kvcache( diff --git a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py b/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py index 060a92bf7c..98ebe3ab2c 100644 --- a/lightllm/common/basemodel/triton_kernel/destindex_copy_kv.py +++ b/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 diff --git a/lightllm/common/basemodel/triton_kernel/hybrid_cpu_cache_copy.py b/lightllm/common/basemodel/triton_kernel/hybrid_cpu_cache_copy.py index f9fb32d26a..3ba6c30b4a 100644 --- a/lightllm/common/basemodel/triton_kernel/hybrid_cpu_cache_copy.py +++ b/lightllm/common/basemodel/triton_kernel/hybrid_cpu_cache_copy.py @@ -1,4 +1,4 @@ -"""Stream-ordered copies for token KV and arbitrary pinned checkpoint tensors.""" +"""Stream-ordered copies for pinned checkpoint pages.""" import torch import triton @@ -7,40 +7,38 @@ @triton.jit def _copy_checkpoint_pages( - SRC, - DST, - SRC_IDS, - DST_IDS, - READIES, - SRC_STRIDE, - DST_STRIDE, - WIDTH, - PAGE_NUM, + 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_IDS + page).to(tl.int64) - dst_id = tl.load(DST_IDS + page).to(tl.int64) + 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(READIES + page) + valid = valid & (tl.load(page_readies + page) == 0) if valid: - for block in range(tl.program_id(0), tl.cdiv(WIDTH, BLOCK), tl.num_programs(0)): + 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 + src_id * SRC_STRIDE + offsets, offsets < WIDTH, other=0) - tl.store(DST + dst_id * DST_STRIDE + offsets, data, offsets < WIDTH, cache_modifier=".wt") + 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(sources) == len(destinations) 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] - assert source.stride(1) == destination.stride(1) == 1 _copy_checkpoint_pages[(12,)]( source, destination, @@ -54,103 +52,3 @@ def copy_checkpoint_pages(sources, destinations, source_ids, destination_ids, pa SKIP_READY=page_readies is not None, BLOCK=4096, ) - - -@triton.jit -def _copy_checkpoint_tensor(SRC, DST, SIZE, BLOCK: tl.constexpr): - for block in range(tl.program_id(0), tl.cdiv(SIZE, BLOCK), tl.num_programs(0)): - offsets = block * BLOCK + tl.arange(0, BLOCK) - data = tl.load(SRC + offsets, offsets < SIZE, other=0) - tl.store(DST + offsets, data, offsets < SIZE, cache_modifier=".wt") - - -def copy_checkpoint_state(sources, destinations): - """Copy a tail checkpoint on the CUDA stream, including CPU-to-CPU copies. - - A host-side tensor copy would race prior asynchronous saves and transfers. - """ - assert len(sources) == len(destinations) - for source, destination in zip(sources, destinations): - assert source.shape == destination.shape and source.dtype == destination.dtype - assert source.is_contiguous() and destination.is_contiguous() - source = source.reshape(-1).view(torch.uint8) - destination = destination.reshape(-1).view(torch.uint8) - _copy_checkpoint_tensor[(12,)](source, destination, source.numel(), BLOCK=4096) - - -@triton.jit -def _copy_full_att_pages( - KV, - CACHE, - MEM_IDS, - PAGE_IDS, - READIES, - KV_STRIDE_L, - KV_STRIDE_T, - PAGE_STRIDE, - RANK_OFFSET, - WIDTH: tl.constexpr, - LAYERS: tl.constexpr, - PAGE_TOKENS: tl.constexpr, - PAGE_NUM, - WRITE_RANK: tl.constexpr, - OFFLOAD: tl.constexpr, - BLOCK: tl.constexpr, -): - page_width = PAGE_TOKENS * LAYERS * WIDTH - for page in range(PAGE_NUM): - page_id = tl.load(PAGE_IDS + page).to(tl.int64) - valid = page_id >= 0 - if OFFLOAD: - valid = valid & ~tl.load(READIES + page) & WRITE_RANK - if valid: - for block in range(tl.program_id(0), tl.cdiv(page_width, BLOCK), tl.num_programs(0)): - offsets = block * BLOCK + tl.arange(0, BLOCK) - token = offsets // (LAYERS * WIDTH) - layer = offsets // WIDTH % LAYERS - dim = offsets % WIDTH - mem_id = tl.load(MEM_IDS + page * PAGE_TOKENS + token, offsets < page_width, other=-1).to(tl.int64) - mask = (offsets < page_width) & (mem_id >= 0) - kv_ptr = KV + layer.to(tl.int64) * KV_STRIDE_L + mem_id * KV_STRIDE_T + dim - cache_ptr = CACHE + page_id * PAGE_STRIDE + RANK_OFFSET + offsets - if OFFLOAD: - value = tl.load(kv_ptr, mask, other=0) - tl.store(cache_ptr, value, mask, cache_modifier=".wt") - else: - value = tl.load(cache_ptr, mask, other=0) - tl.store(kv_ptr, value, mask) - - -def copy_full_att_pages(kv, cpu_cache, mem_indexes, page_indexes, config, tp_rank, page_tokens, page_readies=None): - """Pack ordinary target KV as [page, TP shard, token, layer, KV bytes].""" - kv = kv.view(kv.shape[0], kv.shape[1], -1).view(torch.uint8) - pages = cpu_cache.view(cpu_cache.shape[0], -1).view(torch.uint8) - world_size = config.tp_world_size - heads = config.full_att_all_num_kv_heads - if heads >= world_size: - assert heads % world_size == 0 - replicas = 1 - else: - assert world_size % heads == 0 - replicas = world_size // heads - shard_bytes = page_tokens * kv.shape[0] * kv.shape[2] - assert shard_bytes * (world_size // replicas) == config.get_cpu_cache_full_att_bytes() - assert len(mem_indexes) == len(page_indexes) * page_tokens - _copy_full_att_pages[(12,)]( - kv, - pages, - mem_indexes, - page_indexes, - page_readies, - kv.stride(0), - kv.stride(1), - pages.stride(0), - (tp_rank // replicas) * shard_bytes, - kv.shape[2], - kv.shape[0], - page_tokens, - len(page_indexes), - WRITE_RANK=tp_rank % replicas == 0, - OFFLOAD=page_readies is not None, - BLOCK=4096, - ) diff --git a/lightllm/common/basemodel/triton_kernel/kv_cache_offload.py b/lightllm/common/basemodel/triton_kernel/kv_cache_offload.py index 0fdc43ab9f..ab88b16462 100644 --- a/lightllm/common/basemodel/triton_kernel/kv_cache_offload.py +++ b/lightllm/common/basemodel/triton_kernel/kv_cache_offload.py @@ -57,10 +57,11 @@ def _offload_gpu_kv_to_cpu( ready_state = tl.load(page_readies_ptr + block_index) mask_layer_num = tl.where(cpu_page_index == -1, 0, 1) - mask_layer_num = tl.where(ready_state, 0, mask_layer_num) + mask_layer_num = tl.where(ready_state != 0, 0, mask_layer_num) 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: @@ -100,7 +101,9 @@ def _offload_gpu_kv_to_cpu( + gpu_k_head_index.to(tl.int64) * gpu_scale_stride2 + head_dim_range[None, :] ) - gpu_scale_data = tl.load(gpu_scale_ptr, mask=scale_head_dim_mask[None, :], other=0.0) + gpu_scale_data = tl.load( + gpu_scale_ptr, mask=token_mask[:, None] & scale_head_dim_mask[None, :], other=0.0 + ) cpu_scale_ptr = ( cpu_kv_cache_scale_ptr + cpu_page_index * cpu_scale_stride0 @@ -112,7 +115,7 @@ def _offload_gpu_kv_to_cpu( tl.store( cpu_scale_ptr, gpu_scale_data, - mask=scale_head_dim_mask[None, :], + mask=token_mask[:, None] & scale_head_dim_mask[None, :], cache_modifier=".wt", ) @@ -127,7 +130,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 +142,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: @@ -150,7 +153,9 @@ def _offload_gpu_kv_to_cpu( + gpu_v_head_index.to(tl.int64) * gpu_scale_stride2 + head_dim_range[None, :] ) - gpu_scale_data = tl.load(gpu_scale_ptr, mask=scale_head_dim_mask[None, :], other=0.0) + gpu_scale_data = tl.load( + gpu_scale_ptr, mask=token_mask[:, None] & scale_head_dim_mask[None, :], other=0.0 + ) cpu_scale_ptr = ( cpu_kv_cache_scale_ptr + cpu_page_index * cpu_scale_stride0 @@ -162,7 +167,7 @@ def _offload_gpu_kv_to_cpu( tl.store( cpu_scale_ptr, gpu_scale_data, - mask=scale_head_dim_mask[None, :], + mask=token_mask[:, None] & scale_head_dim_mask[None, :], cache_modifier=".wt", ) @@ -407,6 +412,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 diff --git a/lightllm/common/basemodel/triton_kernel/windowed_mtp.py b/lightllm/common/basemodel/triton_kernel/windowed_mtp.py new file mode 100644 index 0000000000..b562fe913a --- /dev/null +++ b/lightllm/common/basemodel/triton_kernel/windowed_mtp.py @@ -0,0 +1,184 @@ +import torch +import triton +import triton.language as tl + + +@triton.jit +def _gather_window_hidden( + target_hidden, + b_req_idx, + b_start_loc, + b_ready_cache_len, + b_q_seq_len, + out_hidden, + out_position_ids, + out_mem_index, + req_to_seq_len, + stride_hidden_bs, + stride_hidden_d, + stride_out_bs, + stride_out_d, + hidden_dim: tl.constexpr, + max_q_seq_len: tl.constexpr, + window_size: tl.constexpr, + BLOCK: 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_ready_cache_len = tl.load(b_ready_cache_len + cur_batch) + cur_q_seq_len = tl.load(b_q_seq_len + cur_batch) + cur_seq_len = cur_ready_cache_len + cur_q_seq_len + cur_start_loc = tl.load(b_start_loc + cur_batch) + + # A chunk larger than the ring only needs its most recent positions. + cur_window_start = tl.maximum(cur_ready_cache_len, cur_seq_len - window_size) + cur_position = cur_window_start + cur_token + token_mask = cur_position < cur_seq_len + src_index = cur_start_loc + cur_position - cur_ready_cache_len + out_index = cur_batch * max_q_seq_len + cur_token + offs_d = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) + hidden = tl.load( + target_hidden + src_index.to(tl.int64) * stride_hidden_bs + offs_d * stride_hidden_d, + mask=token_mask & (offs_d < hidden_dim), + other=0, + ) + tl.store(out_hidden + out_index * stride_out_bs + offs_d * stride_out_d, hidden, mask=offs_d < hidden_dim) + if tl.program_id(2) == 0: + tl.store(out_position_ids + out_index, tl.where(token_mask, cur_position, -1)) + # This is a logical token index into the standard [token, head, dim] KV view. + mem_index = cur_req_idx * window_size + cur_position % window_size + tl.store(out_mem_index + out_index, tl.where(token_mask, mem_index, -1)) + if cur_token == 0: + tl.store(req_to_seq_len + cur_req_idx, cur_seq_len) + + +def gather_window_hidden( + target_hidden, + b_req_idx, + b_start_loc, + b_ready_cache_len, + b_q_seq_len, + req_to_seq_len, + max_q_seq_len, + window_size, + alloc_func=torch.empty, +): + max_q_seq_len = min(max_q_seq_len, window_size) + token_num = b_req_idx.numel() * max_q_seq_len + out_hidden = alloc_func((token_num, target_hidden.shape[1]), dtype=target_hidden.dtype, device=target_hidden.device) + position_ids = alloc_func((token_num,), dtype=torch.int64, device=target_hidden.device) + mem_index = alloc_func((token_num,), dtype=torch.int32, device=target_hidden.device) + if token_num: + _gather_window_hidden[(max_q_seq_len, b_req_idx.numel(), triton.cdiv(target_hidden.shape[1], 512))]( + target_hidden, + b_req_idx, + b_start_loc, + b_ready_cache_len, + b_q_seq_len, + out_hidden, + position_ids, + mem_index, + req_to_seq_len, + stride_hidden_bs=target_hidden.stride(0), + stride_hidden_d=target_hidden.stride(1), + stride_out_bs=out_hidden.stride(0), + stride_out_d=out_hidden.stride(1), + hidden_dim=target_hidden.shape[1], + max_q_seq_len=max_q_seq_len, + window_size=window_size, + BLOCK=512, + ) + return out_hidden, position_ids, 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 diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 280fd6af89..fb8167afc5 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -312,7 +312,7 @@ def write_to_shm(self, req_manager): shm_name = f"mem_manager_{get_current_rank_in_node()}" # Only GPU running state is consumed by PD. Pickling CPU checkpoints # would move their storage into ordinary shared memory and lose pinning. - big_page_buffers = self.big_page_buffers + big_page_buffers = getattr(self, "big_page_buffers", None) self.big_page_buffers = None try: obj_bytes_array = [ForkingPickler.dumps(self).tobytes() for _ in range(node_world_size * 2)] diff --git a/lightllm/common/kv_cache_mem_manager/operator/hybrid.py b/lightllm/common/kv_cache_mem_manager/operator/hybrid.py index 655459b092..c4774b6dda 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/hybrid.py +++ b/lightllm/common/kv_cache_mem_manager/operator/hybrid.py @@ -157,21 +157,24 @@ def _copy_target_state(self, source, source_id, destination_id): def _copy_tail_state(self, source, source_id, destination_id): self._copy_target_state(source, source_id, destination_id) - window = getattr(self.mem_manager.big_page_buffers, "draft_window", None) + window = self.mem_manager.big_page_buffers.draft_window if window is not None: - from lightllm.common.basemodel.triton_kernel.hybrid_cpu_cache_copy import copy_checkpoint_state - - copy_checkpoint_state( - source.draft_window.get_state_cache(source_id), window.get_state_cache(destination_id) + from lightllm.common.basemodel.triton_kernel.linear_att_cpu_cache_copy import ( + copy_linear_att_state_to_linear_att_state, ) + # The existing copier handles both contiguous checkpoint tensors as bytes. + src_kv, src_seq_lens = source.draft_window.get_state_cache(source_id) + dst_kv, dst_seq_lens = window.get_state_cache(destination_id) + copy_linear_att_state_to_linear_att_state(src_kv, src_seq_lens, dst_kv, dst_seq_lens) + def _copy_window_pages(self, buffer_ids, page_indexes, cpu_cache_client, page_readies=None): - window = getattr(self.mem_manager.big_page_buffers, "draft_window", None) + window = self.mem_manager.big_page_buffers.draft_window if window is None: return from lightllm.common.basemodel.triton_kernel.hybrid_cpu_cache_copy import copy_checkpoint_pages - buffers = (window.kv, window.ends, window.counts) + buffers = (window.kv_buffer, window.seq_lens) pages = self.cpu_cache_config.get_window_views(cpu_cache_client.cpu_kv_cache_tensor, get_current_rank_in_dp()) if page_readies is None: copy_checkpoint_pages(pages, buffers, page_indexes, buffer_ids) diff --git a/lightllm/common/kv_cache_mem_manager/operator/windowed_mtp.py b/lightllm/common/kv_cache_mem_manager/operator/windowed_mtp.py index daee314e79..7f71145b20 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/windowed_mtp.py +++ b/lightllm/common/kv_cache_mem_manager/operator/windowed_mtp.py @@ -1,32 +1,35 @@ from .hybrid import HybridAttMemOperator from .normal import NormalMemOperator -from lightllm.utils.envs_utils import get_env_start_args -from lightllm.utils.dist_utils import get_current_rank_in_dp -from lightllm.common.basemodel.triton_kernel.hybrid_cpu_cache_copy import copy_full_att_pages +from lightllm.utils.dist_utils import get_current_rank_in_dp, get_dp_world_size +from lightllm.common.basemodel.triton_kernel.kv_cache_offload import load_cpu_kv_to_gpu, offload_gpu_kv_to_cpu class WindowedMTPMemOperator(HybridAttMemOperator, NormalMemOperator): """Ordinary target KV plus request-local draft windows in hybrid CPU pages.""" def _load_target_pages(self, mem_indexes, buffer_ids, page_indexes, cpu_cache_client): - copy_full_att_pages( - self.mem_manager.kv_buffer, - cpu_cache_client.cpu_kv_cache_tensor, - mem_indexes, - page_indexes, - self.cpu_cache_config, - get_current_rank_in_dp(), - get_env_start_args().cpu_cache_token_page_size, + load_cpu_kv_to_gpu( + gpu_mem_indexes=mem_indexes, + gpu_kv_cache=self.mem_manager.kv_buffer, + gpu_kv_cache_scale=None, + cpu_kv_cache=self.cpu_cache_config.get_full_att_kv_view(cpu_cache_client.cpu_kv_cache_tensor), + cpu_kv_cache_scale=None, + page_indexes=page_indexes, + tp_index=get_current_rank_in_dp(), + tp_world_size=get_dp_world_size(), + grid_num=16, ) def _offload_target_pages(self, mem_indexes, buffer_ids, page_indexes, page_readies, cpu_cache_client): - copy_full_att_pages( - self.mem_manager.kv_buffer, - cpu_cache_client.cpu_kv_cache_tensor, - mem_indexes, - page_indexes, - self.cpu_cache_config, - get_current_rank_in_dp(), - get_env_start_args().cpu_cache_token_page_size, + offload_gpu_kv_to_cpu( + token_indexes=mem_indexes, + gpu_kv_cache=self.mem_manager.kv_buffer, + gpu_kv_cache_scale=None, + cpu_kv_cache=self.cpu_cache_config.get_full_att_kv_view(cpu_cache_client.cpu_kv_cache_tensor), + cpu_kv_cache_scale=None, + page_indexes=page_indexes, page_readies=page_readies, + tp_index=get_current_rank_in_dp(), + tp_world_size=get_dp_world_size(), + grid_num=16, ) diff --git a/lightllm/common/kv_cache_mem_manager/windowed_mtp.py b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py index 8230652418..24505925e4 100644 --- a/lightllm/common/kv_cache_mem_manager/windowed_mtp.py +++ b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py @@ -1,193 +1,80 @@ -"""Bounded KV storage for parallel-block speculative decoding.""" +"""Request-local draft windows and their existing checkpoint/PD interfaces.""" import torch -import triton -import triton.language as tl - - -@triton.jit -def _gather_new( - FEATURES, - REQS, - STARTS, - FIRST, - LENGTHS, - OUT, - POS, - ENDS, - COUNTS, - N: tl.constexpr, - H: tl.constexpr, - C: tl.constexpr, - BLOCK: tl.constexpr, -): - row = tl.program_id(0) - batch, offset = row // N, row % N - req = tl.load(REQS + batch).to(tl.int64) - first = tl.load(FIRST + batch).to(tl.int64) - length = tl.load(LENGTHS + batch) - end = first + length - recent = tl.maximum(first, end - C) - count = end - recent - position = recent + offset - valid = offset < count - dim = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) - source = tl.load(STARTS + batch).to(tl.int64) + position - first - value = tl.load(FEATURES + source * H + dim, valid & (dim < H), other=0) - tl.store(OUT + row * H + dim, value, dim < H) - if tl.program_id(1) == 0: - tl.store(POS + row, tl.where(valid, position, -1)) - if offset == 0: - tl.store(ENDS + req, end) - tl.store(COUNTS + req, tl.minimum(end, C)) - - -@triton.jit -def _write_ring( - POOL, - NEW, - REQS, - POS, - N: tl.constexpr, - C: tl.constexpr, - WIDTH: tl.constexpr, - BLOCK: tl.constexpr, -): - row = tl.program_id(0) - req = tl.load(REQS + row // N).to(tl.int64) - position = tl.load(POS + row) - slot = position % C - dim = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) - valid = (position >= 0) & (dim < WIDTH) - value = tl.load(NEW + row * WIDTH + dim, valid, other=0) - tl.store(POOL + (req * C + slot) * WIDTH + dim, value, valid) - - -@triton.jit -def _pack_ring( - POOL, - NOISE, - REQS, - ENDS, - OUT, - C: tl.constexpr, - B: tl.constexpr, - WIDTH: tl.constexpr, - BLOCK: tl.constexpr, -): - row = tl.program_id(0) - batch, slot = row // (C + B), row % (C + B) - req = tl.load(REQS + batch).to(tl.int64) - end = tl.load(ENDS + req) - count = tl.minimum(end, C) - position = tl.maximum(0, end - C) + slot - ring_slot = position % C - dim = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) - context = tl.load(POOL + (req * C + ring_slot) * WIDTH + dim, (slot < count) & (dim < WIDTH), other=0) - noise = tl.load( - NOISE + (batch * B + slot - count) * WIDTH + dim, (slot >= count) & (slot < count + B) & (dim < WIDTH), other=0 - ) - tl.store(OUT + row * WIDTH + dim, context + noise, dim < WIDTH) class WindowKVStore: - """Per-layer recent-window K/V, indexed independently of target KV. + """Use the normal [layer, token, head, dim] KV layout for each request's ring. - Only newly accepted positions are projected and written. Each window slot - has at most one writer, including prefill chunks larger than the window. - On request reuse, the new end/count hides stale slots; all visible positions - are overwritten by the first prefill. Draft noise K/V never enters the ring. + Request r owns token slots [r * window, (r + 1) * window). Its committed + sequence length determines the visible positions; noise KV stays temporary. """ def __init__(self, requests, layers, kv_heads, head_dim, dtype, device, window): self.capacity = window - self.kv = torch.zeros((layers, requests, window, 2 * kv_heads, head_dim), dtype=dtype, device=device) - self.ends = torch.zeros(requests, dtype=torch.int64, device=device) - self.counts = torch.zeros(requests, dtype=torch.int32, device=device) + self.kv_buffer = torch.empty((layers, requests * window, 2 * kv_heads, head_dim), dtype=dtype, device=device) + self.req_to_seq_len = torch.zeros(requests, dtype=torch.int64, device=device) def reset(self): - """Invalidate all request windows without changing captured buffer addresses.""" - self.ends.zero_() - self.counts.zero_() + self.req_to_seq_len.zero_() def reset_req(self, req_idx): - self.ends[req_idx] = 0 - self.counts[req_idx] = 0 + self.req_to_seq_len[req_idx] = 0 def save_checkpoint(self, req_idx, buffers, slot): - # Copy contiguous per-layer windows, preserving the physical p % W layout. - # The prefill stream event covers these D2H copies before radix publication. - for layer in range(self.kv.shape[0]): - buffers.kv[slot, layer].copy_(self.kv[layer, req_idx], non_blocking=True) - buffers.ends[slot : slot + 1].copy_(self.ends[req_idx : req_idx + 1], non_blocking=True) - buffers.counts[slot : slot + 1].copy_(self.counts[req_idx : req_idx + 1], non_blocking=True) + start = req_idx * self.capacity + buffers.kv_buffer[slot].copy_(self.kv_buffer[:, start : start + self.capacity], non_blocking=True) + buffers.seq_lens[slot : slot + 1].copy_(self.req_to_seq_len[req_idx : req_idx + 1], non_blocking=True) def restore_checkpoint(self, req_idx, buffers, slot): - for layer in range(self.kv.shape[0]): - self.kv[layer, req_idx].copy_(buffers.kv[slot, layer], non_blocking=True) - self.ends[req_idx : req_idx + 1].copy_(buffers.ends[slot : slot + 1], non_blocking=True) - self.counts[req_idx : req_idx + 1].copy_(buffers.counts[slot : slot + 1], non_blocking=True) - # Small-page matching can release the source node immediately after return. - # Finish H2D before its CPU slot can be evicted and overwritten. - if self.kv.is_cuda: - torch.cuda.current_stream(self.kv.device).synchronize() - - def prepare(self, reqs, features, starts, first, lengths, max_new): - n = min(max_new, self.capacity) - packed = features.new_empty((reqs.numel() * n, features.shape[-1])) - positions = torch.empty(reqs.numel() * n, dtype=torch.int64, device=features.device) - _gather_new[(reqs.numel() * n, triton.cdiv(features.shape[-1], 512))]( - features, - reqs, - starts, - first, - lengths, - packed, - positions, - self.ends, - self.counts, - n, - features.shape[-1], + start = req_idx * self.capacity + self.kv_buffer[:, start : start + self.capacity].copy_(buffers.kv_buffer[slot], non_blocking=True) + self.req_to_seq_len[req_idx : req_idx + 1].copy_(buffers.seq_lens[slot : slot + 1], non_blocking=True) + # The CPU slot can be recycled as soon as prefix matching returns. + if self.kv_buffer.is_cuda: + torch.cuda.current_stream(self.kv_buffer.device).synchronize() + + def prepare( + self, + b_req_idx, + target_hidden, + b_start_loc, + b_ready_cache_len, + b_q_seq_len, + max_q_seq_len, + alloc_func=torch.empty, + ): + from lightllm.common.basemodel.triton_kernel.windowed_mtp import gather_window_hidden + + return gather_window_hidden( + target_hidden, + b_req_idx, + b_start_loc, + b_ready_cache_len, + b_q_seq_len, + self.req_to_seq_len, + max_q_seq_len, self.capacity, - 512, - ) - return packed, positions - - def write(self, layer, reqs, positions, new_kv): - width = new_kv.shape[-2] * new_kv.shape[-1] - _write_ring[(positions.numel(), triton.cdiv(width, 512))]( - self.kv[layer], - new_kv, - reqs, - positions, - positions.numel() // reqs.numel(), - self.capacity, - width, - 512, + alloc_func=alloc_func, ) - def pack(self, layer, reqs, noise_kv, block): - heads, dim = self.kv.shape[-2:] - output = noise_kv.new_empty((reqs.numel(), self.capacity + block, heads, dim)) - _pack_ring[(reqs.numel() * (self.capacity + block), triton.cdiv(heads * dim, 512))]( - self.kv[layer], + def write(self, layer_index, mem_index, kv): + from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv + + destindex_copy_kv(kv, mem_index, self.kv_buffer[layer_index]) + + def pack(self, layer_index, b_req_idx, noise_kv, block_size, alloc_func=torch.empty): + from lightllm.common.basemodel.triton_kernel.windowed_mtp import pack_window_kv + + return pack_window_kv( + self.kv_buffer[layer_index], noise_kv, - reqs, - self.ends, - output, + b_req_idx, + self.req_to_seq_len, self.capacity, - block, - heads * dim, - 512, + block_size, + alloc_func=alloc_func, ) - return output - - def retained_positions(self, reqs): - """Diagnostic only: logical positions, in attention order.""" - end = self.ends.index_select(0, reqs)[:, None] - slots = torch.arange(self.capacity, device=reqs.device)[None, :] - positions = (end - self.capacity).clamp_min(0) + slots - return positions.masked_fill(slots >= end.clamp_max(self.capacity), -1) class WindowKVPageHelper: @@ -204,12 +91,13 @@ def __init__(self, mem_manager, offset=0): store = mem_manager.windowed_draft_kv args = get_env_start_args() tp_world_size = args.tp // args.dp - layers, _, window, heads, dim = store.kv.shape + layers, _, heads, dim = store.kv_buffer.shape + window = store.capacity self.shape = (window, layers, heads * tp_world_size, dim) self.kv_offset = (offset + 7) // 8 * 8 - self.kv_nbytes = window * layers * heads * tp_world_size * dim * store.kv.element_size() + self.kv_nbytes = window * layers * heads * tp_world_size * dim * store.kv_buffer.element_size() self.end_offset = (self.kv_offset + self.kv_nbytes + 7) // 8 * 8 - self.state_nbytes = self.end_offset + 12 + self.state_nbytes = self.end_offset + 8 def assert_page_size(self): page = self.mem_manager.kv_move_buffer[0] @@ -222,34 +110,26 @@ def assert_page_size(self): def copy_req_page(self, page_index, req_idx, dp_mems, mode): from lightllm.common.kv_trans_kernel.nixl_kv_trans import page_io - assert mode in ("read", "write") - assert req_idx is not None page = self.mem_manager.kv_move_buffer[page_index].view(torch.uint8).reshape(-1) kv_page = ( page[self.kv_offset : self.kv_offset + self.kv_nbytes] - .view(self.mem_manager.windowed_draft_kv.kv.dtype) + .view(self.mem_manager.windowed_draft_kv.kv_buffer.dtype) .view(self.shape) ) end_page = page[self.end_offset : self.end_offset + 8].view(torch.int64) - count_page = page[self.end_offset + 8 : self.end_offset + 12].view(torch.int32) window = self.shape[0] indexes = torch.arange(req_idx * window, (req_idx + 1) * window, dtype=torch.int64, device=page.device) for tp_index, mem in enumerate(dp_mems): store = mem.windowed_draft_kv - layers, requests, capacity, heads, dim = store.kv.shape - assert (capacity, layers, heads * len(dp_mems), dim) == self.shape - assert store.kv.dtype == kv_page.dtype page_io( mem_indexes=indexes, page_tensor=kv_page, - kv_buffer=store.kv.view(layers, requests * capacity, heads, dim), + kv_buffer=store.kv_buffer, tp_index=tp_index, tp_world_size=len(dp_mems), mode=mode, ) if mode == "read": - store.ends[req_idx : req_idx + 1].copy_(end_page, non_blocking=True) - store.counts[req_idx : req_idx + 1].copy_(count_page, non_blocking=True) + store.req_to_seq_len[req_idx : req_idx + 1].copy_(end_page, non_blocking=True) elif tp_index == 0: - end_page.copy_(store.ends[req_idx : req_idx + 1], non_blocking=True) - count_page.copy_(store.counts[req_idx : req_idx + 1], non_blocking=True) + end_page.copy_(store.req_to_seq_len[req_idx : req_idx + 1], non_blocking=True) diff --git a/lightllm/common/state_cache_manager/windowed_mtp.py b/lightllm/common/state_cache_manager/windowed_mtp.py index c0b3a36cce..933a66dc74 100644 --- a/lightllm/common/state_cache_manager/windowed_mtp.py +++ b/lightllm/common/state_cache_manager/windowed_mtp.py @@ -1,13 +1,12 @@ """Window state configuration and checkpoints shared by hybrid attention caches.""" -import json from dataclasses import dataclass -from pathlib import Path import torch from lightllm.utils.log_utils import init_logger from lightllm.common.build_utils import repair_config +from lightllm.utils.config_utils import get_config_json from .base import StateCacheManager from .layer_cache import LayerCache @@ -35,11 +34,11 @@ def get_state_shape(self): return (self.window, 2 * self.kv_heads, self.head_dim) def get_state_bytes(self, element_size): - return self.layers * self.window * 2 * self.kv_heads * self.head_dim * element_size + 12 + return self.layers * self.window * 2 * self.kv_heads * self.head_dim * element_size + 8 def load_window_state_config(args): - config = json.loads((Path(args.mtp_draft_model_dir[0]) / "config.json").read_text()) + config = get_config_json(args.mtp_draft_model_dir[0]) # Match the existing draft initialization order: base aliases, Llama KV # defaults, then the overrides applied by Qwen3.5 adapters. # Do not normalize again after merging: runtime also retains those aliases. @@ -57,8 +56,7 @@ class WindowStateBuffers: """Pinned CPU payload; the containing StateCacheManager owns all slot IDs.""" def __init__(self, size: int, config: WindowStateCacheConfig, dtype: torch.dtype): - self.config = config - self.kv = LayerCache( + self.kv_buffer = LayerCache( size=size, dtype=dtype, shape=config.get_state_shape(), @@ -66,19 +64,17 @@ def __init__(self, size: int, config: WindowStateCacheConfig, dtype: torch.dtype device="cpu", size_first=True, ).buffer - self.ends = torch.zeros(size, dtype=torch.int64, device="cpu", pin_memory=True) - self.counts = torch.zeros(size, dtype=torch.int32, device="cpu", pin_memory=True) + self.seq_lens = torch.zeros(size, dtype=torch.int64, device="cpu", pin_memory=True) logger.info( f"draft window checkpoint slots={size}, pinned_bytes={size * config.get_state_bytes(dtype.itemsize)}" ) def get_state_cache(self, buffer_idx: int): - return self.kv[buffer_idx], self.ends[buffer_idx : buffer_idx + 1], self.counts[buffer_idx : buffer_idx + 1] + return self.kv_buffer[buffer_idx], self.seq_lens[buffer_idx : buffer_idx + 1] def clear_to_init_state(self): - self.kv.zero_() - self.ends.zero_() - self.counts.zero_() + self.kv_buffer.zero_() + self.seq_lens.zero_() class WindowStateCacheManager(StateCacheManager): @@ -144,28 +140,42 @@ def get_cpu_cache_full_att_bytes(self): * self.dtype.itemsize ) + def get_full_att_kv_view(self, cpu_cache_tensor): + from lightllm.utils.envs_utils import get_env_start_args + + page_tokens = get_env_start_args().cpu_cache_token_page_size + pages = cpu_cache_tensor.view(cpu_cache_tensor.shape[0], -1).view(torch.uint8) + return ( + pages[:, : self.get_cpu_cache_full_att_bytes()] + .view(self.dtype) + .view( + pages.shape[0], + self.full_att_layer_num, + page_tokens, + 2 * self.full_att_all_num_kv_heads, + self.full_att_head_dim, + ) + ) + def get_cpu_cache_window_offset(self): if self.linear_config is not None: return self.linear_config.get_cpu_cache_big_page_bytes() return (self.get_cpu_cache_full_att_bytes() + 15) // 16 * 16 def get_cpu_cache_window_rank_bytes(self): - kv_bytes = self.window_config.get_state_bytes(self.dtype.itemsize) - 12 - # Align int64 ends and each rank's payload; counts is int32. - return (((kv_bytes + 7) // 8 * 8 + 12) + 15) // 16 * 16 + kv_bytes = self.window_config.get_state_bytes(self.dtype.itemsize) - 8 + # Align the int64 sequence length and each rank's payload. + return (((kv_bytes + 7) // 8 * 8 + 8) + 15) // 16 * 16 def get_cpu_cache_big_page_bytes(self): return self.get_cpu_cache_window_offset() + self.tp_world_size * self.get_cpu_cache_window_rank_bytes() def get_window_views(self, cpu_cache_tensor, tp_rank): - assert 0 <= tp_rank < self.tp_world_size pages = cpu_cache_tensor.view(cpu_cache_tensor.shape[0], -1).view(torch.uint8) - assert pages.shape[1] == self.get_cpu_cache_big_page_bytes() start = self.get_cpu_cache_window_offset() + tp_rank * self.get_cpu_cache_window_rank_bytes() - kv_bytes = self.window_config.get_state_bytes(self.dtype.itemsize) - 12 + kv_bytes = self.window_config.get_state_bytes(self.dtype.itemsize) - 8 end_offset = start + (kv_bytes + 7) // 8 * 8 return ( pages[:, start : start + kv_bytes].view(self.dtype), pages[:, end_offset : end_offset + 8].view(torch.int64), - pages[:, end_offset + 8 : end_offset + 12].view(torch.int32), ) diff --git a/lightllm/models/qwen3_5_dflash/model.py b/lightllm/models/qwen3_5_dflash/model.py index 422ccf61cf..32fd283537 100644 --- a/lightllm/models/qwen3_5_dflash/model.py +++ b/lightllm/models/qwen3_5_dflash/model.py @@ -29,6 +29,8 @@ def _init_custom(self): self.mask_token_id = self.config["mask_token_id"] def _init_mem_manager(self): + if self.uses_windowed_draft_kv: + return super()._init_mem_manager() target_mem_manager = self.main_model.mem_manager draft_kv_shape = (self.config["num_key_value_heads"], self.config["head_dim"]) target_kv_shape = ( diff --git a/lightllm/models/qwen3_5_dspark/model.py b/lightllm/models/qwen3_5_dspark/model.py index 896949d1cf..0a513f66a7 100644 --- a/lightllm/models/qwen3_5_dspark/model.py +++ b/lightllm/models/qwen3_5_dspark/model.py @@ -26,6 +26,8 @@ def _init_custom(self): self.mask_token_id = self.config["mask_token_id"] def _init_mem_manager(self): + if self.uses_windowed_draft_kv: + return super()._init_mem_manager() target_mem_manager = self.main_model.mem_manager draft_kv_shape = (self.config["num_key_value_heads"], self.config["head_dim"]) target_kv_shape = ( diff --git a/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py index 4573a43944..0deb385c0f 100644 --- a/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py @@ -58,7 +58,11 @@ def token_attention_forward(self, input_embdings, infer_state, layer_weight): q, noise_kv = self._get_qkv(input_embdings, infer_state, layer_weight) model = att_state.backend.model scratch = infer_state.mem_manager.windowed_draft_kv.pack( - self.layer_num_ - model.draft_layer_start, att_state.b_req_idx, noise_kv, att_state.block_size + self.layer_num_ - model.draft_layer_start, + att_state.b_req_idx, + noise_kv, + att_state.block_size, + alloc_func=self.alloc_tensor, ) output = att_state.decode_att( q.view(-1, self.tp_q_head_num_, self.head_dim_), diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index a5001bfae9..1c0e9257f4 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -14,6 +14,7 @@ from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer from lightllm.models.qwen3_dflash.layer_weights.pre_and_post_layer_weight import Qwen3DFlashPreAndPostLayerWeight from lightllm.models.qwen3_dflash.layer_weights.transformer_layer_weight import Qwen3DFlashTransformerLayerWeight +from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager class Qwen3DFlashModel(LlamaTpPartModel): @@ -59,7 +60,7 @@ def _init_mem_manager(self): self.mem_manager = self.main_model.mem_manager if self.uses_windowed_draft_kv: self.mem_manager.init_windowed_draft_kv( - requests=self.args.running_max_req_size + 1, + requests=self.req_manager.max_request_num + 1, layers=self.config["n_layer"], kv_heads=self.config["num_key_value_heads"] // self.tp_world_size_, head_dim=self.config.get("head_dim", self.config["n_embed"] // self.config["num_attention_heads"]), @@ -69,6 +70,7 @@ def _init_mem_manager(self): def _init_att_backend(self): if self.uses_windowed_draft_kv: from lightllm.utils.sgl_utils import flash_attn_with_kvcache + if flash_attn_with_kvcache is None: raise NotImplementedError("Windowed MTP requires FA3") @@ -104,42 +106,56 @@ def _init_weights(self, start_layer_index=None): ] @torch.no_grad() - def commit_features(self, model_input, hidden, starts=None, accept_len=None): + def commit_window_kv(self, model_input, hidden, b_req_mtp_start_loc=None, accept_len=None): """Commit accepted target features to the bounded draft KV store.""" - # The fused feature is transient; only its per-layer K/V is retained. - projected = self.pre_post_weight.fc_weight_.mm(hidden, use_custom_tensor_mananger=False) - projected = self.pre_post_weight.hidden_norm_weight_(projected, eps=self.pre_infer.eps_, alloc_func=torch.empty) + if hidden.numel() == 0: + return if model_input.is_prefill: - reqs = model_input.b_req_idx.long() - starts = model_input.b_prefill_start_loc - first = model_input.b_ready_cache_len - lengths = model_input.b_seq_len - first - max_new = model_input.max_q_seq_len + b_req_idx = model_input.b_req_idx + b_start_loc = model_input.b_prefill_start_loc + b_ready_cache_len = model_input.b_ready_cache_len + b_q_seq_len = model_input.b_seq_len - b_ready_cache_len + max_q_seq_len = model_input.max_q_seq_len else: - reqs = model_input.b_req_idx.index_select(0, starts.long()).long() - first = model_input.b_seq_len.index_select(0, starts.long()) - 1 - lengths = accept_len - max_new = self.args.mtp_step + 1 + b_start_loc = b_req_mtp_start_loc + b_req_idx = model_input.b_req_idx.index_select(0, b_start_loc) + b_ready_cache_len = model_input.b_seq_len.index_select(0, b_start_loc) - 1 + b_q_seq_len = accept_len + max_q_seq_len = self.args.mtp_step + 1 store = self.mem_manager.windowed_draft_kv - features, positions = store.prepare(reqs, projected, starts, first, lengths, max_new) + features, position_ids, mem_index = store.prepare( + b_req_idx, + hidden, + b_start_loc, + b_ready_cache_len, + b_q_seq_len, + max_q_seq_len, + alloc_func=self.pre_infer.alloc_tensor, + ) infer_state = self.infer_state_class() - infer_state.position_cos = self._cos_cached.index_select(0, positions.clamp_min(0)) - infer_state.position_sin = self._sin_cached.index_select(0, positions.clamp_min(0)) + infer_state.mtp_draft_input_hiddens = features + features = self.pre_infer.context_forward(None, infer_state, self.pre_post_weight) + position_ids = position_ids.clamp_min(0) + infer_state.position_cos = self._cos_cached.index_select(0, position_ids) + infer_state.position_sin = self._sin_cached.index_select(0, position_ids) for i, (layer, weight) in enumerate(zip(self.layers_infer, self.trans_layers_weight)): kv = layer._get_context_kv(features, infer_state, weight) - store.write(i, reqs, positions, kv) + store.write(i, mem_index, kv) def _select_mem_indexes(self, model_input: ModelInput): if self.uses_windowed_draft_kv and not model_input.is_prefill: # Windowed draft blocks do not write into the target's request pages. - return torch.full_like(model_input.b_req_idx, self.mem_manager.HOLD_TOKEN_MEMINDEXES[0]) + hold_mem_index = self.mem_manager.HOLD_TOKEN_MEMINDEXES[0] + return g_pin_mem_manager.get_const_gpu_tensor( + key=f"window_mtp_mem_indexes_{hold_mem_index}_{model_input.b_req_idx.dtype}", + shape=model_input.b_req_idx.shape, + fill_value=hold_mem_index, + dtype=model_input.b_req_idx.dtype, + ) return super()._select_mem_indexes(model_input) def _decode(self, model_input: ModelInput) -> ModelOutput: - if self.uses_windowed_draft_kv: - assert model_input.batch_size > 0 and model_input.batch_size % self.block_size == 0 - return super()._decode(model_input) - if model_input.mtp_draft_input_hiddens is None: + if self.uses_windowed_draft_kv or model_input.mtp_draft_input_hiddens is None: return super()._decode(model_input) assert model_input.mtp_draft_input_hiddens.shape[0] == model_input.batch_size diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 24f226f867..eabc9b66a6 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -136,7 +136,7 @@ def _launch_subprocesses(args: StartArgs): args.graph_max_len_in_batch = args.max_req_total_len # mode setting check. - is_hybrid_model = is_hybrid_att_model(args.model_dir, args) + is_hybrid_model = is_hybrid_att_model(args.model_dir) or args.mtp_draft_kv_mode == "window" if is_hybrid_model: # 混合注意力模型的状态缓存依赖动态缓存和分块预填充，与是否开启输出约束无关。 assert not args.disable_dynamic_prompt_cache, "hybrid attention models require dynamic prompt cache" @@ -305,7 +305,7 @@ def _launch_subprocesses(args: StartArgs): # 避免请求释放时将不完整的大页 state 写入 radix cache 并触发断言。 args.linear_att_page_block_num = 10000000 - if args.enable_cpu_cache and is_hybrid_att_model(args.model_dir, args): + if args.enable_cpu_cache and is_hybrid_model: args.cpu_cache_token_page_size = args.linear_att_hash_page_size * args.linear_att_page_block_num logger.info(f"set cpu_cache_token_page_size to {args.cpu_cache_token_page_size} for hybrid att model") if args.enable_cpu_cache: diff --git a/lightllm/server/core/objs/req.py b/lightllm/server/core/objs/req.py index f35127c53d..28b5c623c1 100644 --- a/lightllm/server/core/objs/req.py +++ b/lightllm/server/core/objs/req.py @@ -213,7 +213,7 @@ def init( self.post_init() args = get_env_start_args() - if is_hybrid_att_model(args.model_dir, args): + if is_hybrid_att_model(args.model_dir) or getattr(args, "mtp_draft_kv_mode", "full") == "window": self._fill_hybrid_token_hash() if args.enable_cpu_cache: cpu_cache_hash_list, cpu_cache_page_len_list = self._calcu_hybrid_cpu_cache_page_len_list() diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 45d300a12e..03b90bf18b 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -424,7 +424,7 @@ def get_can_alloc_token_num(self): def save_hybrid_state_to_cache(self, b_req_idx: torch.Tensor, reqs: List["InferReq"]): """Snapshot request-level attention state at big/small-page boundaries.""" - if not self.is_hybrid_att_model or self.radix_cache is None: + if not self.is_hybrid_att_model: return # Request-state snapshot at a big-page boundary. diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py index e5a2c9a72c..8450433896 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py @@ -39,8 +39,8 @@ def fill_draft_model_kv_state( # DFlash prefill 直接复用 target prompt 的 token 布局，并注入 target # hidden 初始化唯一 draft model 的 KV。使用浅副本避免在 target 输入上 # 保留 draft 专用状态。 - if getattr(self.backend.draft_models[0], "uses_windowed_draft_kv", False) is True: - self.backend.draft_models[0].commit_features(target_model_input, target_hidden) + if getattr(self.backend.draft_models[0], "uses_windowed_draft_kv", False): + self.backend.draft_models[0].commit_window_kv(target_model_input, target_hidden) return draft_input = copy.copy(target_model_input) draft_input.mtp_draft_input_hiddens = target_hidden @@ -83,9 +83,8 @@ def propose_next( # full 模式通过 draft forward 写入 verify 行对应的 KV；windowed 模式 # 只将接受的 target hidden 投影并写入独立窗口，不保留被拒绝的行。 - windowed = getattr(draft_model, "uses_windowed_draft_kv", False) is True - if windowed: - draft_model.commit_features( + if getattr(draft_model, "uses_windowed_draft_kv", False): + draft_model.commit_window_kv( target_model_input, target_model_output.mtp_collector.spec_hidden, b_req_mtp_start_loc, accept_len ) else: diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py index 6b432791f3..7dc44d7e6f 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py @@ -40,8 +40,8 @@ def fill_draft_model_kv_state( # DSpark prefill 直接使用 target prompt 的 token 布局和 hidden，将 prompt # KV 写入唯一的 parallel-block draft model。使用浅副本，避免把 draft # 专用 hidden 挂到后续流程仍可能读取的 target ModelInput 上。 - if getattr(self.backend.draft_models[0], "uses_windowed_draft_kv", False) is True: - self.backend.draft_models[0].commit_features(target_model_input, target_hidden) + if getattr(self.backend.draft_models[0], "uses_windowed_draft_kv", False): + self.backend.draft_models[0].commit_window_kv(target_model_input, target_hidden) return draft_input = copy.copy(target_model_input) draft_input.mtp_draft_input_hiddens = target_hidden @@ -86,9 +86,8 @@ def propose_next( # full 模式通过 draft forward 写入 verify 行对应的 KV；windowed 模式 # 只将接受的 target hidden 投影并写入独立窗口，不保留被拒绝的行。 - windowed = getattr(draft_model, "uses_windowed_draft_kv", False) is True - if windowed: - draft_model.commit_features( + if getattr(draft_model, "uses_windowed_draft_kv", False): + draft_model.commit_window_kv( target_model_input, target_model_output.mtp_collector.spec_hidden, b_req_mtp_start_loc, accept_len ) else: diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 2b408940d8..8f76b4a53c 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -486,9 +486,9 @@ def is_linear_att_mixed_model(model_path: str) -> bool: return False -def is_hybrid_att_model(model_path: str, args=None) -> bool: - """Whether target linear state or draft window snapshots require hybrid pages.""" - return is_linear_att_mixed_model(model_path) or (args is not None and args.mtp_draft_kv_mode == "window") +def is_hybrid_att_model(model_path: str) -> bool: + """Models whose non-full attention state follows hybrid checkpoint pages.""" + return is_linear_att_mixed_model(model_path) def get_model_type(model_path: str) -> Optional[str]: diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index 11c8f8d1cc..af056c7cea 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -62,7 +62,7 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": args = get_env_start_args() assert args.enable_cpu_cache - is_hybrid_model = is_hybrid_att_model(args.model_dir, args) + is_hybrid_model = is_hybrid_att_model(args.model_dir) or args.mtp_draft_kv_mode == "window" mem_manager_class = None if is_hybrid_model else select_mem_manager_class() if is_hybrid_model: hybrid_config = get_hybrid_cache_config() diff --git a/lightllm/utils/profile_max_tokens.py b/lightllm/utils/profile_max_tokens.py index 7559e9e348..f1cdcd20e8 100644 --- a/lightllm/utils/profile_max_tokens.py +++ b/lightllm/utils/profile_max_tokens.py @@ -38,11 +38,10 @@ def get_mtp_adjusted_mem_fraction( target_weight_bytes: int, target_layer_num: int, mtp_layer_num: int, - draft_cache_bytes: int = 0, ) -> float: mtp_weight_bytes = target_weight_bytes * mtp_layer_num / target_layer_num total_gpu_bytes = torch.cuda.get_device_properties(get_current_device_id()).total_memory - adjusted_mem_fraction = mem_fraction - (mtp_weight_bytes + draft_cache_bytes) / total_gpu_bytes + adjusted_mem_fraction = mem_fraction - mtp_weight_bytes / total_gpu_bytes # 不同 rank 的权重分片大小和 GPU 总显存可能不同。取全局最小值，保证所有 # rank 使用相同且能够安全预留 MTP 权重显存的 KV cache 比例。 @@ -75,17 +74,20 @@ def profile_mtp_weight_memory(model): weight_memory_before = torch.cuda.memory_allocated() yield target_weight_bytes = torch.cuda.memory_allocated() - weight_memory_before - draft_cache_bytes = 0 + mem_fraction = model.mem_fraction if getattr(model.args, "mtp_draft_kv_mode", "full") == "window": from lightllm.common.state_cache_manager.windowed_mtp import load_window_state_config + config = load_window_state_config(model.args) - draft_cache_bytes = (model.args.running_max_req_size + 1) * config.get_state_bytes(model.data_type.itemsize) + draft_cache_bytes = (model.max_req_num + 1) * config.get_state_bytes(model.data_type.itemsize) + total_gpu_bytes = torch.cuda.get_device_properties(get_current_device_id()).total_memory + # Include the window reservation in the existing per-rank budget reduction. + mem_fraction -= draft_cache_bytes / total_gpu_bytes model.mem_fraction = get_mtp_adjusted_mem_fraction( - mem_fraction=model.mem_fraction, + mem_fraction=mem_fraction, target_weight_bytes=target_weight_bytes, target_layer_num=model.config["n_layer"], mtp_layer_num=get_mtp_weight_layer_num(), - draft_cache_bytes=draft_cache_bytes, ) diff --git a/unit_tests/common/test_windowed_mtp.py b/unit_tests/common/test_windowed_mtp.py index 4f099667e3..c50faf886a 100644 --- a/unit_tests/common/test_windowed_mtp.py +++ b/unit_tests/common/test_windowed_mtp.py @@ -1,48 +1,33 @@ from types import SimpleNamespace +import pytest import torch from lightllm.common.kv_cache_mem_manager.windowed_mtp import WindowKVStore +from lightllm.common.kv_cache_mem_manager.windowed_mtp import WindowKVPageHelper from lightllm.common.state_cache_manager.windowed_mtp import ( WindowedMTPCacheConfig, + WindowStateBuffers, WindowStateCacheConfig, ) def _checkpoint_buffers(slots, layers, window, kv_heads, head_dim, dtype): return SimpleNamespace( - kv=torch.zeros((slots, layers, window, 2 * kv_heads, head_dim), dtype=dtype), - ends=torch.zeros(slots, dtype=torch.int64), - counts=torch.zeros(slots, dtype=torch.int32), + kv_buffer=torch.zeros((slots, layers, window, 2 * kv_heads, head_dim), dtype=dtype), + seq_lens=torch.zeros(slots, dtype=torch.int64), ) -def test_retained_positions_follow_logical_window_order(): - store = WindowKVStore( - requests=3, - layers=1, - kv_heads=1, - head_dim=2, - dtype=torch.float32, - device="cpu", - window=4, - ) - store.ends.copy_(torch.tensor([0, 3, 7], dtype=torch.int64)) - store.counts.copy_(torch.tensor([0, 3, 4], dtype=torch.int32)) - - positions = store.retained_positions(torch.tensor([0, 1, 2], dtype=torch.int64)) - - assert torch.equal( - positions, - torch.tensor( - [ - [-1, -1, -1, -1], - [0, 1, 2, -1], - [3, 4, 5, 6], - ], - dtype=torch.int64, - ), - ) +@pytest.fixture +def tensor_cache(): + from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager + + g_cache_manager.cache_env_in() + try: + yield g_cache_manager.empty + finally: + g_cache_manager.cache_env_out() def test_checkpoint_round_trip_preserves_physical_ring_layout(): @@ -67,9 +52,8 @@ def test_checkpoint_round_trip_preserves_physical_ring_layout(): layers * window * 2 * kv_heads * head_dim, dtype=torch.float32, ).view(layers, window, 2 * kv_heads, head_dim) - store.kv[:, source_req].copy_(expected_kv) - store.ends[source_req] = 9 - store.counts[source_req] = window + store.kv_buffer[:, source_req * window : (source_req + 1) * window].copy_(expected_kv) + store.req_to_seq_len[source_req] = 9 buffers = _checkpoint_buffers( slots=2, @@ -77,19 +61,17 @@ def test_checkpoint_round_trip_preserves_physical_ring_layout(): window=window, kv_heads=kv_heads, head_dim=head_dim, - dtype=store.kv.dtype, + dtype=store.kv_buffer.dtype, ) store.save_checkpoint(source_req, buffers, checkpoint_slot) - store.kv[:, destination_req].fill_(-1) - store.ends[destination_req] = -1 - store.counts[destination_req] = -1 + store.kv_buffer[:, destination_req * window : (destination_req + 1) * window].fill_(-1) + store.req_to_seq_len[destination_req] = -1 store.restore_checkpoint(destination_req, buffers, checkpoint_slot) - assert torch.equal(buffers.kv[checkpoint_slot], expected_kv) - assert torch.equal(store.kv[:, destination_req], expected_kv) - assert store.ends[destination_req].item() == 9 - assert store.counts[destination_req].item() == window + assert torch.equal(buffers.kv_buffer[checkpoint_slot], expected_kv) + assert torch.equal(store.kv_buffer[:, destination_req * window : (destination_req + 1) * window], expected_kv) + assert store.req_to_seq_len[destination_req].item() == 9 class _TargetStateLayout: @@ -114,14 +96,12 @@ def test_cpu_page_keeps_tp_rank_window_states_separate(): ) pages = torch.zeros((2, config.get_cpu_cache_big_page_bytes()), dtype=torch.uint8) - rank0_kv, rank0_ends, rank0_counts = config.get_window_views(pages, tp_rank=0) - rank1_kv, rank1_ends, rank1_counts = config.get_window_views(pages, tp_rank=1) + rank0_kv, rank0_ends = config.get_window_views(pages, tp_rank=0) + rank1_kv, rank1_ends = config.get_window_views(pages, tp_rank=1) rank0_kv.fill_(1) rank1_kv.fill_(2) rank0_ends.copy_(torch.tensor([[11], [12]], dtype=torch.int64)) rank1_ends.copy_(torch.tensor([[21], [22]], dtype=torch.int64)) - rank0_counts.copy_(torch.tensor([[3], [4]], dtype=torch.int32)) - rank1_counts.copy_(torch.tensor([[1], [2]], dtype=torch.int32)) restored_rank0 = config.get_window_views(pages, tp_rank=0) restored_rank1 = config.get_window_views(pages, tp_rank=1) @@ -129,6 +109,208 @@ def test_cpu_page_keeps_tp_rank_window_states_separate(): assert torch.all(restored_rank1[0] == 2) assert restored_rank0[1].tolist() == [[11], [12]] assert restored_rank1[1].tolist() == [[21], [22]] - assert restored_rank0[2].tolist() == [[3], [4]] - assert restored_rank1[2].tolist() == [[1], [2]] assert torch.count_nonzero(pages[:, :target_page_bytes]).item() == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("window", [1, 4, 8]) +def test_window_commit_and_pack_handles_ragged_strided_inputs_and_rejections(window, tensor_cache): + store = WindowKVStore(3, 2, 1, 3, torch.float32, "cuda", window) + reqs = torch.tensor([0, 2], dtype=torch.int32, device="cuda") + # Neither the hidden rows nor the noise head dimensions are contiguous. + hidden = torch.arange(20 * 12, device="cuda", dtype=torch.float32).view(20, 12)[::2, ::2] + starts = torch.tensor([0, 7], dtype=torch.int32, device="cuda") + first = torch.tensor([0, 0], dtype=torch.int32, device="cuda") + lengths = torch.tensor([7, 3], dtype=torch.int32, device="cuda") + features, positions, mem_index = store.prepare(reqs, hidden, starts, first, lengths, 7, alloc_func=tensor_cache) + for layer in range(2): + store.write(layer, mem_index, features.view(-1, 2, 3) + layer) + + accepted_hidden = torch.arange(5 * 6, device="cuda", dtype=torch.float32).view(5, 6) + 1000 + starts = torch.tensor([0, 3], dtype=torch.int32, device="cuda") + first = torch.tensor([7, 3], dtype=torch.int32, device="cuda") + lengths = torch.tensor([2, 1], dtype=torch.int32, device="cuda") + features, positions, mem_index = store.prepare( + reqs, accepted_hidden, starts, first, lengths, 3, alloc_func=tensor_cache + ) + for layer in range(2): + store.write(layer, mem_index, features.view(-1, 2, 3) + layer) + + noise = torch.randn(4, 2, 6, device="cuda")[:, :, ::2] + history = [torch.cat((hidden[:7], accepted_hidden[:2])), torch.cat((hidden[7:10], accepted_hidden[3:4]))] + for layer in range(2): + packed = store.pack(layer, reqs, noise, 2, alloc_func=tensor_cache) + expected = torch.zeros_like(packed) + for batch, values in enumerate(history): + retained = values[-window:].view(-1, 2, 3) + layer + expected[batch, : len(retained)] = retained + expected[batch, len(retained) : len(retained) + 2] = noise[batch * 2 : (batch + 1) * 2] + torch.testing.assert_close(packed, expected, rtol=0, atol=0) + assert store.req_to_seq_len.tolist() == [9, 0, 4] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_window_pack_graph_replay_reads_current_request_and_sequence_length(tensor_cache): + store = WindowKVStore(2, 1, 1, 8, torch.float32, "cuda", 4) + store.kv_buffer.copy_(torch.arange(8 * 2 * 8, device="cuda").view(1, 8, 2, 8)) + store.req_to_seq_len.copy_(torch.tensor([7, 1], device="cuda")) + reqs = torch.tensor([0], dtype=torch.int32, device="cuda") + noise = torch.randn(2, 2, 8, device="cuda") + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + store.pack(0, reqs, noise, 2, alloc_func=tensor_cache) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + packed = store.pack(0, reqs, noise, 2, alloc_func=tensor_cache) + reqs.fill_(1) + graph.replay() + expected = torch.zeros_like(packed) + expected[0, 0] = store.kv_buffer[0, 4] + expected[0, 1:3] = noise + torch.testing.assert_close(packed, expected, rtol=0, atol=0) + store.reset_req(1) + graph.replay() + expected.zero_() + expected[0, :2] = noise + torch.testing.assert_close(packed, expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_empty_window_batch_uses_existing_empty_token_shapes(tensor_cache): + store = WindowKVStore(1, 1, 1, 8, torch.float32, "cuda", 4) + reqs = torch.empty(0, dtype=torch.int32, device="cuda") + features, positions, mem_index = store.prepare( + reqs, torch.empty(0, 16, device="cuda"), reqs, reqs, reqs, 3, alloc_func=tensor_cache + ) + assert features.shape == (0, 16) + assert positions.shape == mem_index.shape == (0,) + assert store.pack(0, reqs, torch.empty(0, 2, 8, device="cuda"), 2, alloc_func=tensor_cache).shape == (0, 6, 2, 8) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("tp_world_size", [1, 2, 4]) +def test_window_target_pages_reuse_normal_cpu_transfer_with_partial_pages(monkeypatch, tp_world_size): + from lightllm.common.basemodel.triton_kernel.kv_cache_offload import load_cpu_kv_to_gpu, offload_gpu_kv_to_cpu + + args = SimpleNamespace(cpu_cache_token_page_size=4, linear_att_hash_page_size=2, linear_att_page_block_num=2) + monkeypatch.setattr("lightllm.utils.envs_utils.get_env_start_args", lambda: args) + config = WindowedMTPCacheConfig(WindowStateCacheConfig(1, 1, 8, 4), torch.float32, tp_world_size, 2, 8, 2) + pages = torch.zeros((1, config.get_cpu_cache_big_page_bytes()), dtype=torch.uint8, pin_memory=True) + cpu_kv = config.get_full_att_kv_view(pages) + replicas = max(1, tp_world_size // 2) + kvs = [torch.randn(2, 8, 2 * max(1, 2 // tp_world_size), 8, device="cuda") for _ in range(tp_world_size)] + for rank in range(tp_world_size): + kvs[rank].copy_(kvs[rank // replicas * replicas]) + expected = [kv.clone() for kv in kvs] + mem_index = torch.tensor([1, 5, -1, -1], dtype=torch.int32, device="cuda") + page_index = torch.zeros(1, dtype=torch.int64, device="cuda") + readies = torch.zeros(1, dtype=torch.int32, device="cuda") + for rank, kv in enumerate(kvs): + offload_gpu_kv_to_cpu(mem_index, kv, None, cpu_kv, None, page_index, readies, rank, tp_world_size, 16) + for rank, kv in enumerate(kvs): + kv.zero_() + load_cpu_kv_to_gpu(mem_index, kv, None, cpu_kv, None, page_index, rank, tp_world_size, 16) + torch.testing.assert_close(kv[:, [1, 5]], expected[rank][:, [1, 5]], rtol=0, atol=0) + assert not kv[:, [0, 2, 3, 4, 6, 7]].any() + assert pages[:, config.get_cpu_cache_window_offset() :].count_nonzero() == 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (1, 2), (2, 1)]) +def test_window_pd_page_reuses_page_io_across_tp_sizes(monkeypatch, prefill_tp, decode_tp): + def make_mems(tp): + monkeypatch.setattr("lightllm.utils.envs_utils.get_env_start_args", lambda: SimpleNamespace(tp=tp, dp=1)) + mems = [ + SimpleNamespace( + windowed_draft_kv=WindowKVStore(3, 2, 2 // tp, 8, torch.float32, "cuda", 4), + kv_move_buffer=torch.zeros(1, 2048, dtype=torch.uint8, device="cuda"), + ) + for _ in range(tp) + ] + helper = WindowKVPageHelper(mems[0], offset=16) + helper.assert_page_size() + return mems, helper + + source, write_helper = make_mems(prefill_tp) + destination, read_helper = make_mems(decode_tp) + global_kv = torch.randn(2, 4, 4, 8, device="cuda") + + def shard_heads(rank, tp): + start = rank * (2 // tp) + return list(range(start, start + 2 // tp)) + list(range(2 + start, 2 + start + 2 // tp)) + + for rank, mem in enumerate(source): + store = mem.windowed_draft_kv + store.kv_buffer[:, :4].copy_(global_kv[:, :, shard_heads(rank, prefill_tp)]) + store.req_to_seq_len[0] = 7 + source[0].kv_move_buffer[:, :16].fill_(42) + write_helper.copy_req_page(0, 0, source, "write") + destination[0].kv_move_buffer.copy_(source[0].kv_move_buffer) + for mem in destination: + mem.windowed_draft_kv.kv_buffer.zero_() + read_helper.copy_req_page(0, 2, destination, "read") + for rank, mem in enumerate(destination): + store = mem.windowed_draft_kv + torch.testing.assert_close(store.kv_buffer[:, 8:12], global_kv[:, :, shard_heads(rank, decode_tp)]) + assert store.req_to_seq_len.tolist() == [0, 0, 7] + assert not store.kv_buffer[:, :8].any() + assert destination[0].kv_move_buffer[:, :16].eq(42).all() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("linear_target", [False, True]) +def test_tail_checkpoint_copy_preserves_stream_order_and_neighbor_slots(linear_target): + from lightllm.common.kv_cache_mem_manager.operator.windowed_mtp import WindowedMTPMemOperator + from lightllm.common.kv_cache_mem_manager.operator.linear_att import LinearAttMemOperator + + config = WindowStateCacheConfig(layers=2, kv_heads=1, head_dim=8, window=4) + source_window = WindowStateBuffers(2, config, torch.bfloat16) + destination_window = WindowStateBuffers(2, config, torch.bfloat16) + source_target = tuple( + torch.zeros((2, 4, 8), dtype=dtype, pin_memory=True) for dtype in (torch.bfloat16, torch.float32) + ) + destination_target = tuple(torch.zeros_like(state, pin_memory=True) for state in source_target) + source = SimpleNamespace( + draft_window=source_window, + get_state_cache=lambda slot: tuple(state[slot] for state in source_target), + ) + destination = SimpleNamespace( + draft_window=destination_window, + get_state_cache=lambda slot: tuple(state[slot] for state in destination_target), + ) + operator_class = LinearAttMemOperator if linear_target else WindowedMTPMemOperator + operator = operator_class.__new__(operator_class) + operator.mem_manager = SimpleNamespace(big_page_buffers=destination) + + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + # Keep the save pending so a host-side tail copy would see stale data. + torch.cuda._sleep(1_000_000) + expected_kv = torch.randn(source_window.kv_buffer[1].shape, dtype=torch.bfloat16, device="cuda") + expected_seq_len = torch.tensor([2 ** 35 + 7], dtype=torch.int64, device="cuda") + expected_target = [torch.randn(state.shape[1:], dtype=state.dtype, device="cuda") for state in source_target] + source_window.kv_buffer[1].copy_(expected_kv, non_blocking=True) + source_window.seq_lens[1:2].copy_(expected_seq_len, non_blocking=True) + for state, expected in zip(source_target, expected_target): + state[1].copy_(expected, non_blocking=True) + + operator._copy_tail_state(source, source_id=1, destination_id=0) + restored_kv = torch.empty_like(expected_kv) + restored_seq_len = torch.empty_like(expected_seq_len) + restored_kv.copy_(destination_window.kv_buffer[0], non_blocking=True) + restored_seq_len.copy_(destination_window.seq_lens[:1], non_blocking=True) + stream.synchronize() + + torch.testing.assert_close(restored_kv, expected_kv, rtol=0, atol=0) + torch.testing.assert_close(restored_seq_len, expected_seq_len, rtol=0, atol=0) + torch.testing.assert_close(destination_window.kv_buffer[0], expected_kv.cpu(), rtol=0, atol=0) + assert not destination_window.kv_buffer[1].any() + assert destination_window.seq_lens.tolist() == [2 ** 35 + 7, 0] + for state, expected in zip(destination_target, expected_target): + if linear_target: + torch.testing.assert_close(state[0], expected.cpu(), rtol=0, atol=0) + else: + assert not state[0].any() + assert not state[1].any() diff --git a/unit_tests/server/router/model_infer/mtp_speculative/test_windowed_draft.py b/unit_tests/server/router/model_infer/mtp_speculative/test_windowed_draft.py deleted file mode 100644 index 1632d9fb5a..0000000000 --- a/unit_tests/server/router/model_infer/mtp_speculative/test_windowed_draft.py +++ /dev/null @@ -1,95 +0,0 @@ -from types import SimpleNamespace - -import pytest -import torch - -from lightllm.server.router.model_infer.mtp_speculative.proposers.dflash import DFlashProposer -from lightllm.server.router.model_infer.mtp_speculative.proposers.dspark import DSparkProposer -from lightllm.server.router.model_infer.pin_mem_manager import g_pin_mem_manager - - -@pytest.mark.parametrize("proposer_class", [DFlashProposer, DSparkProposer]) -def test_windowed_prefill_commits_features_without_token_cache_forward(proposer_class): - committed = [] - draft = SimpleNamespace(uses_windowed_draft_kv=True, commit_features=lambda *args: committed.append(args)) - proposer = proposer_class(backend=SimpleNamespace(draft_models=[draft]), enable_dynmaic_mtp=False) - model_input = SimpleNamespace( - is_prefill=True, - b_position_delta=None, - b_req_idx=torch.tensor([3, 5], dtype=torch.int32), - input_ids=torch.arange(5), - ) - hidden = torch.randn(5, 8) - proposer.fill_draft_model_kv_state( - target_model_input=model_input, - target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=hidden)), - target_next_token_ids=torch.tensor([11, 13]), - ) - assert len(committed) == 1 - assert committed[0][0] is model_input - assert committed[0][1] is hidden - - -@pytest.mark.parametrize("proposer_class", [DFlashProposer, DSparkProposer]) -def test_windowed_proposal_commits_accepted_rows_and_uses_page_managed_input(monkeypatch, proposer_class): - committed = [] - forwarded = [] - token_ids = torch.tensor([30, 31, 32, 40, 41, 42]) - block_output = SimpleNamespace(mtp_collector=SimpleNamespace(draft_token_ids=token_ids)) - - def forward(model_input): - forwarded.append(model_input) - return block_output - - draft = SimpleNamespace( - uses_windowed_draft_kv=True, - block_size=3, - mask_token_id=99, - commit_features=lambda *args: committed.append(args), - forward=forward, - ) - backend = SimpleNamespace(draft_models=[draft], _gen_argmax_token_ids=lambda _: token_ids) - proposer = proposer_class(backend=backend, enable_dynmaic_mtp=False) - monkeypatch.setattr( - g_pin_mem_manager, - "get_const_gpu_tensor", - lambda *, shape, fill_value, dtype, **_: torch.full(shape, fill_value, dtype=dtype), - ) - hidden = torch.randn(5, 8) - starts = torch.tensor([0, 3], dtype=torch.int32) - accepted = torch.tensor([2, 1], dtype=torch.int32) - model_input = SimpleNamespace( - is_prefill=False, - batch_size=5, - max_kv_seq_len=9, - input_ids=torch.tensor([10, 11, 12, 20, 21]), - b_req_idx=torch.tensor([7, 7, 7, 9, 9], dtype=torch.int32), - b_mtp_index=torch.arange(5, dtype=torch.int32), - b_seq_len=torch.tensor([4, 5, 6, 8, 9], dtype=torch.int32), - b_position_delta=torch.tensor([1, 1, 1, 2, 2], dtype=torch.int32), - b_shared_seq_len=torch.tensor([3, 3, 3, 6, 6], dtype=torch.int32), - b_shared_radix_node_id=torch.tensor([17, 17, 17, 19, 19]), - mtp_draft_input_hiddens=None, - ) - - proposal = proposer.propose_next( - target_model_input=model_input, - target_model_output=SimpleNamespace(mtp_collector=SimpleNamespace(spec_hidden=hidden)), - target_next_token_ids=model_input.input_ids, - b_req_mtp_start_loc=starts, - draft_step=2, - accept_len=accepted, - ) - - assert len(committed) == len(forwarded) == 1 - assert committed[0][0] is model_input - assert committed[0][1] is hidden - assert committed[0][2] is starts - assert committed[0][3] is accepted - block_input = forwarded[0] - assert block_input is not model_input - assert block_input.mtp_draft_input_hiddens is None - assert model_input.mtp_draft_input_hiddens is None - assert torch.equal(block_input.input_ids, torch.tensor([11, 99, 99, 20, 99, 99])) - assert torch.equal(block_input.b_seq_len, torch.tensor([6, 7, 8, 9, 10, 11], dtype=torch.int32)) - assert torch.equal(proposal.token_ids, torch.tensor([[30, 31], [40, 41]])) From 407a3dce32b660b9c82c2909b38f391bc8bcd61b Mon Sep 17 00:00:00 2001 From: Zhao Xintong Date: Thu, 15 Oct 2026 12:18:36 +0000 Subject: [PATCH 6/9] fix --- docs/CN/source/tutorial/api_server_args.rst | 11 ----------- docs/EN/source/tutorial/api_server_args.rst | 11 ----------- lightllm/common/basemodel/basemodel.py | 2 +- lightllm/common/kv_cache_mem_manager/mem_manager.py | 2 +- lightllm/common/state_cache_manager/__init__.py | 2 +- lightllm/common/state_cache_manager/windowed_mtp.py | 2 +- lightllm/models/qwen3_dflash/model.py | 4 ++-- lightllm/server/api_cli.py | 13 ++++++------- lightllm/server/api_start.py | 6 +++--- lightllm/server/core/objs/req.py | 2 +- lightllm/server/core/objs/start_args_type.py | 4 ++-- lightllm/utils/envs_utils.py | 4 ++-- lightllm/utils/kv_cache_utils.py | 2 +- lightllm/utils/profile_max_tokens.py | 2 +- 14 files changed, 22 insertions(+), 45 deletions(-) diff --git a/docs/CN/source/tutorial/api_server_args.rst b/docs/CN/source/tutorial/api_server_args.rst index c0471bf70d..7411b17a9c 100644 --- a/docs/CN/source/tutorial/api_server_args.rst +++ b/docs/CN/source/tutorial/api_server_args.rst @@ -724,17 +724,6 @@ MTP 多预测参数 * ``eagle_no_att`` * ``None``: 不启用 mtp（默认） -.. option:: --mtp_draft_kv_mode - - 草稿 KV 缓存策略，默认 ``full``。dflash style 草稿模型支持 ``window``， - 仅保留最近窗口内的 KV。target 仍使用完整上下文。 - -.. option:: --mtp_draft_window_size - - ``window`` 策略保留的最近位置数，默认 ``512``，必须大于 0。 - - 每个请求的历史容量上限为 ``mtp_draft_window_size``。 - .. option:: --mtp_draft_model_dir MTP 多预测功能的草稿模型路径 diff --git a/docs/EN/source/tutorial/api_server_args.rst b/docs/EN/source/tutorial/api_server_args.rst index 8c7dc01846..f238a5adb5 100644 --- a/docs/EN/source/tutorial/api_server_args.rst +++ b/docs/EN/source/tutorial/api_server_args.rst @@ -741,17 +741,6 @@ MTP Multi-Prediction Parameters * ``eagle_no_att`` * ``None``: Do not enable mtp (default) -.. option:: --mtp_draft_kv_mode - - Draft KV cache policy; defaults to ``full``. dflash style draft models support - ``window``, which retains only recent-window KV. The target still - uses its full context. - -.. option:: --mtp_draft_window_size - - Number of recent positions retained by ``window``; defaults to ``512`` and - must be positive. - .. option:: --mtp_draft_model_dir Path to the draft model for MTP multi-prediction functionality diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 1b924007a5..7e8e850904 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/lightllm/common/basemodel/basemodel.py @@ -237,7 +237,7 @@ 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) - req_manager_class = ReqManagerForWindowedMTP if self.args.mtp_draft_kv_mode == "window" else ReqManager + 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 diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index fb8167afc5..204ce8b81b 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -43,7 +43,7 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False self.windowed_draft_kv = None self._windowed_draft_kv_params = None args = get_env_start_args() - self.window_state_config = load_window_state_config(args) if args.mtp_draft_kv_mode == "window" else None + self.window_state_config = load_window_state_config(args) if args.mtp_windowed else None # profile the max total token num if the size is None self.profile_size(mem_fraction) diff --git a/lightllm/common/state_cache_manager/__init__.py b/lightllm/common/state_cache_manager/__init__.py index e63a80d961..40feb63211 100644 --- a/lightllm/common/state_cache_manager/__init__.py +++ b/lightllm/common/state_cache_manager/__init__.py @@ -19,7 +19,7 @@ def get_hybrid_cache_config(linear_config=None): linear_config = Glm5NextCacheConfig.from_model_config(model_cfg, args) elif model_type in ("qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"): linear_config = LinearAttCacheConfig.load_from_args() - if args.mtp_draft_kv_mode == "window": + if args.mtp_windowed: return WindowedMTPCacheConfig.load_from_args(linear_config) if linear_config is not None: return linear_config diff --git a/lightllm/common/state_cache_manager/windowed_mtp.py b/lightllm/common/state_cache_manager/windowed_mtp.py index 933a66dc74..fcf502e959 100644 --- a/lightllm/common/state_cache_manager/windowed_mtp.py +++ b/lightllm/common/state_cache_manager/windowed_mtp.py @@ -49,7 +49,7 @@ def load_window_state_config(args): config.setdefault("num_key_value_heads", config["num_attention_heads"]) if merge_nested: config.update(config.get("dflash_config", {})) - return WindowStateCacheConfig.from_model_config(config, args.mtp_draft_window_size, args.tp // args.dp) + return WindowStateCacheConfig.from_model_config(config, args.mtp_window_size, args.tp // args.dp) class WindowStateBuffers: diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index 1c0e9257f4..8e390af1ff 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -44,7 +44,7 @@ def _verify_params(self): assert self.args.mtp_step <= self.config["block_size"] self.config["block_size"] = self.args.mtp_step - self.uses_windowed_draft_kv = getattr(self.args, "mtp_draft_kv_mode", "full") == "window" + self.uses_windowed_draft_kv = getattr(self.args, "mtp_windowed", False) def _init_custom(self): self._cos_cached = self.main_model._cos_cached @@ -64,7 +64,7 @@ def _init_mem_manager(self): layers=self.config["n_layer"], kv_heads=self.config["num_key_value_heads"] // self.tp_world_size_, head_dim=self.config.get("head_dim", self.config["n_embed"] // self.config["num_attention_heads"]), - window=self.args.mtp_draft_window_size, + window=self.args.mtp_window_size, ) def _init_att_backend(self): diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 976cadb760..a178de5e3b 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -795,13 +795,6 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: action="store_true", help="""Whether to enable fused shared experts for supported MoE models. It is auto-enabled when supported.""", ) - parser.add_argument( - "--mtp_draft_kv_mode", - choices=["full", "window"], - default="full", - help="Draft KV retention policy for parallel block MTP", - ) - parser.add_argument("--mtp_draft_window_size", type=int, default=512) parser.add_argument( "--mtp_mode", choices=[ @@ -837,6 +830,12 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: help="""Number of additional draft tokens per request. For DSpark and DFlash this value is derived from the draft checkpoint block_size.""", ) + parser.add_argument( + "--mtp_windowed", + action="store_true", + help="Use a sliding window for draft KV in DFlash-style MTP.", + ) + parser.add_argument("--mtp_window_size", type=int, default=512) parser.add_argument( "--mtp_dynamic_verify", action="store_true", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index eabc9b66a6..309629ecb3 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -136,7 +136,7 @@ def _launch_subprocesses(args: StartArgs): args.graph_max_len_in_batch = args.max_req_total_len # mode setting check. - is_hybrid_model = is_hybrid_att_model(args.model_dir) or args.mtp_draft_kv_mode == "window" + is_hybrid_model = is_hybrid_att_model(args.model_dir) or args.mtp_windowed if is_hybrid_model: # 混合注意力模型的状态缓存依赖动态缓存和分块预填充，与是否开启输出约束无关。 assert not args.disable_dynamic_prompt_cache, "hybrid attention models require dynamic prompt cache" @@ -205,8 +205,8 @@ def _launch_subprocesses(args: StartArgs): assert args.mtp_draft_model_dir is None assert args.mtp_step == 0 - if args.mtp_draft_kv_mode == "window": - assert args.mtp_draft_window_size > 0 + if args.mtp_windowed: + assert args.mtp_window_size > 0 assert args.mtp_mode in ("dflash", "dspark"), "windowed draft KV requires dflash or dspark" if args.llm_kv_type not in (None, "None"): raise ValueError("windowed draft KV currently requires unquantized KV") diff --git a/lightllm/server/core/objs/req.py b/lightllm/server/core/objs/req.py index 28b5c623c1..22e992413e 100644 --- a/lightllm/server/core/objs/req.py +++ b/lightllm/server/core/objs/req.py @@ -213,7 +213,7 @@ def init( self.post_init() args = get_env_start_args() - if is_hybrid_att_model(args.model_dir) or getattr(args, "mtp_draft_kv_mode", "full") == "window": + if is_hybrid_att_model(args.model_dir) or getattr(args, "mtp_windowed", False): self._fill_hybrid_token_hash() if args.enable_cpu_cache: cpu_cache_hash_list, cpu_cache_page_len_list = self._calcu_hybrid_cpu_cache_page_len_list() diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index bef6eaf6de..5d04a97bbc 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -208,8 +208,8 @@ class StartArgs: }, ) mtp_draft_model_dir: Optional[List[str]] = field(default=None) - mtp_draft_kv_mode: str = field(default="full", metadata={"choices": ["full", "window"]}) - mtp_draft_window_size: int = field(default=512) + mtp_windowed: bool = field(default=False) + mtp_window_size: int = field(default=512) mtp_step: int = field(default=0) mtp_dynamic_verify: bool = field(default=False) kv_quant_calibration_config_path: Optional[str] = field(default=None) diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 1c110a2dbb..fadf2febee 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -271,7 +271,7 @@ def enable_cpu_cache_numa_interleave() -> bool: @lru_cache(maxsize=None) def get_added_mtp_kv_layer_num() -> int: args = get_env_start_args() - if getattr(args, "mtp_draft_kv_mode", "full") == "window": + if getattr(args, "mtp_windowed", False): return 0 mtp_mode = args.mtp_mode @@ -308,7 +308,7 @@ def get_mtp_weight_layer_num() -> int: return args.mtp_step if mtp_mode == "eagle_no_att": return 1 - if getattr(args, "mtp_draft_kv_mode", "full") == "window": + if getattr(args, "mtp_windowed", False): return _get_mtp_draft_backbone_layer_num(args.mtp_draft_model_dir[0]) return get_added_mtp_kv_layer_num() diff --git a/lightllm/utils/kv_cache_utils.py b/lightllm/utils/kv_cache_utils.py index af056c7cea..c3533e2c8f 100644 --- a/lightllm/utils/kv_cache_utils.py +++ b/lightllm/utils/kv_cache_utils.py @@ -62,7 +62,7 @@ def calcu_cpu_cache_meta() -> "CpuKVCacheMeta": args = get_env_start_args() assert args.enable_cpu_cache - is_hybrid_model = is_hybrid_att_model(args.model_dir) or args.mtp_draft_kv_mode == "window" + is_hybrid_model = is_hybrid_att_model(args.model_dir) or args.mtp_windowed mem_manager_class = None if is_hybrid_model else select_mem_manager_class() if is_hybrid_model: hybrid_config = get_hybrid_cache_config() diff --git a/lightllm/utils/profile_max_tokens.py b/lightllm/utils/profile_max_tokens.py index f1cdcd20e8..bbbfc700cd 100644 --- a/lightllm/utils/profile_max_tokens.py +++ b/lightllm/utils/profile_max_tokens.py @@ -75,7 +75,7 @@ def profile_mtp_weight_memory(model): yield target_weight_bytes = torch.cuda.memory_allocated() - weight_memory_before mem_fraction = model.mem_fraction - if getattr(model.args, "mtp_draft_kv_mode", "full") == "window": + if getattr(model.args, "mtp_windowed", False): from lightllm.common.state_cache_manager.windowed_mtp import load_window_state_config config = load_window_state_config(model.args) From 6fdc174a37ccf9c87d303af0d719420658ea5989 Mon Sep 17 00:00:00 2001 From: Zhao Xintong Date: Thu, 15 Oct 2026 14:59:52 +0000 Subject: [PATCH 7/9] fix --- .../triton_kernel/kv_cache_offload.py | 14 +- .../kv_cache_mem_manager/mem_manager.py | 26 +- .../kv_cache_mem_manager/operator/hybrid.py | 7 +- .../qwen3next_mem_manager.py | 6 +- .../kv_cache_mem_manager/windowed_mtp.py | 10 +- .../common/state_cache_manager/__init__.py | 16 +- .../state_cache_manager/windowed_mtp.py | 4 +- lightllm/models/qwen3_dflash/model.py | 8 +- .../model_infer/mode_backend/base_backend.py | 18 +- lightllm/utils/envs_utils.py | 2 +- lightllm/utils/profile_max_tokens.py | 7 +- unit_tests/common/test_windowed_mtp.py | 251 +++++------------- 12 files changed, 127 insertions(+), 242 deletions(-) diff --git a/lightllm/common/basemodel/triton_kernel/kv_cache_offload.py b/lightllm/common/basemodel/triton_kernel/kv_cache_offload.py index ab88b16462..6ec91e008d 100644 --- a/lightllm/common/basemodel/triton_kernel/kv_cache_offload.py +++ b/lightllm/common/basemodel/triton_kernel/kv_cache_offload.py @@ -57,7 +57,7 @@ def _offload_gpu_kv_to_cpu( ready_state = tl.load(page_readies_ptr + block_index) mask_layer_num = tl.where(cpu_page_index == -1, 0, 1) - mask_layer_num = tl.where(ready_state != 0, 0, mask_layer_num) + mask_layer_num = tl.where(ready_state, 0, mask_layer_num) token_range = block_index * TOKEN_BLOCK + tl.arange(0, TOKEN_BLOCK) token_indexes = tl.load(token_indexes_ptr + token_range).to(tl.int64) @@ -101,9 +101,7 @@ def _offload_gpu_kv_to_cpu( + gpu_k_head_index.to(tl.int64) * gpu_scale_stride2 + head_dim_range[None, :] ) - gpu_scale_data = tl.load( - gpu_scale_ptr, mask=token_mask[:, None] & scale_head_dim_mask[None, :], other=0.0 - ) + gpu_scale_data = tl.load(gpu_scale_ptr, mask=scale_head_dim_mask[None, :], other=0.0) cpu_scale_ptr = ( cpu_kv_cache_scale_ptr + cpu_page_index * cpu_scale_stride0 @@ -115,7 +113,7 @@ def _offload_gpu_kv_to_cpu( tl.store( cpu_scale_ptr, gpu_scale_data, - mask=token_mask[:, None] & scale_head_dim_mask[None, :], + mask=scale_head_dim_mask[None, :], cache_modifier=".wt", ) @@ -153,9 +151,7 @@ def _offload_gpu_kv_to_cpu( + gpu_v_head_index.to(tl.int64) * gpu_scale_stride2 + head_dim_range[None, :] ) - gpu_scale_data = tl.load( - gpu_scale_ptr, mask=token_mask[:, None] & scale_head_dim_mask[None, :], other=0.0 - ) + gpu_scale_data = tl.load(gpu_scale_ptr, mask=scale_head_dim_mask[None, :], other=0.0) cpu_scale_ptr = ( cpu_kv_cache_scale_ptr + cpu_page_index * cpu_scale_stride0 @@ -167,7 +163,7 @@ def _offload_gpu_kv_to_cpu( tl.store( cpu_scale_ptr, gpu_scale_data, - mask=token_mask[:, None] & scale_head_dim_mask[None, :], + mask=scale_head_dim_mask[None, :], cache_modifier=".wt", ) diff --git a/lightllm/common/kv_cache_mem_manager/mem_manager.py b/lightllm/common/kv_cache_mem_manager/mem_manager.py index 204ce8b81b..e79c7408e8 100755 --- a/lightllm/common/kv_cache_mem_manager/mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/mem_manager.py @@ -41,7 +41,7 @@ def __init__(self, size, dtype, head_num, head_dim, layer_num, always_copy=False self.always_copy = always_copy self.dtype = dtype self.windowed_draft_kv = None - self._windowed_draft_kv_params = None + self._windowed_draft_max_req_num = None args = get_env_start_args() self.window_state_config = load_window_state_config(args) if args.mtp_windowed else None # profile the max total token num if the size is None @@ -131,8 +131,8 @@ def _init_buffers(self, size, dtype, head_num, head_dim, layer_num): self.kv_buffer = torch.empty( (layer_num, size + self.page_size, 2 * head_num, head_dim), dtype=dtype, device="cuda" ) - if self._windowed_draft_kv_params is not None: - self.init_windowed_draft_kv(**self._windowed_draft_kv_params) + if self._windowed_draft_max_req_num is not None: + self.init_windowed_draft_kv(self._windowed_draft_max_req_num) self._init_state_cache_buffers() def _init_state_cache_buffers(self): @@ -151,13 +151,19 @@ def _init_state_cache_buffers(self): self.CPU_CACHE_BIG_PAGE_LOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 2 self.CPU_CACHE_BIG_PAGE_OFFLOAD_TEMP_BUFFER_ID = self.big_page_buffers.size - 1 - def init_windowed_draft_kv(self, **cache_params): - """Create draft window storage from dimensions supplied by the draft model.""" - if self.window_state_config is not None: - for name in ("layers", "kv_heads", "head_dim", "window"): - assert cache_params[name] == getattr(self.window_state_config, name), f"draft window {name} mismatch" - self._windowed_draft_kv_params = cache_params - self.windowed_draft_kv = WindowKVStore(**cache_params, dtype=self.dtype, device=self.kv_buffer.device) + def init_windowed_draft_kv(self, max_request_num): + """Allocate request windows from the shared draft-state configuration.""" + config = self.window_state_config + self._windowed_draft_max_req_num = max_request_num + self.windowed_draft_kv = WindowKVStore( + requests=max_request_num + 1, + layers=config.layers, + kv_heads=config.kv_heads, + head_dim=config.head_dim, + dtype=self.dtype, + device=self.kv_buffer.device, + window=config.window, + ) def alloc_paged_kv_move_buffer(self, page_num, page_size) -> torch.Tensor: self.kv_move_buffer = torch.empty( diff --git a/lightllm/common/kv_cache_mem_manager/operator/hybrid.py b/lightllm/common/kv_cache_mem_manager/operator/hybrid.py index c4774b6dda..9b80012bb8 100644 --- a/lightllm/common/kv_cache_mem_manager/operator/hybrid.py +++ b/lightllm/common/kv_cache_mem_manager/operator/hybrid.py @@ -3,7 +3,7 @@ from .base import BaseMemManagerOperator from lightllm.utils.envs_utils import get_env_start_args from lightllm.utils.dist_utils import get_current_rank_in_dp -from lightllm.common.state_cache_manager import get_hybrid_cache_config +from lightllm.common.state_cache_manager.windowed_mtp import WindowedMTPCacheConfig from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -16,7 +16,10 @@ class HybridAttMemOperator(BaseMemManagerOperator): def __init__(self, mem_manager): super().__init__(mem_manager) - self.cpu_cache_config = get_hybrid_cache_config(linear_config=getattr(mem_manager, "linear_config", None)) + if mem_manager.window_state_config is not None: + self.cpu_cache_config = WindowedMTPCacheConfig.load_from_args( + linear_config=getattr(mem_manager, "linear_config", None) + ) def copy_mem_to_mem(self, src_mem_index: torch.Tensor, dst_mem_index: torch.Tensor): from lightllm.common.basemodel.triton_kernel.kv_move import copy_kv_buffer_to_kv_buffer diff --git a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py index bf704cda67..8cda62e43a 100644 --- a/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py +++ b/lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py @@ -63,7 +63,11 @@ def _init_linear_att_buffers(self): def write_to_shm(self, req_manager): self.req_to_conv_state = req_manager.req_to_conv_state self.req_to_ssm_state = req_manager.req_to_ssm_state - # Keep pinned CPU checkpoints out of cross-process GPU state exports. + # super().write_to_shm() 会用 ForkingPickler 序列化本对象，torch 在 dump 时会把 + # CPU tensor 的 storage 原地迁到共享内存，使本进程大页 state cache 原本 + # pinned(cudaHostAlloc) 的内存退化为普通 shm mmap，之后 Triton kernel 携带该指针 + # 启动会报 "Pointer argument cannot be accessed from Triton (cpu tensor?)"。 + # 跨进程消费方并不使用 cpu 侧大页 state cache，序列化期间临时剔除以保住 pinned。 big_page_buffers = self.big_page_buffers self.big_page_buffers = None try: diff --git a/lightllm/common/kv_cache_mem_manager/windowed_mtp.py b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py index 24505925e4..a2d047a7f5 100644 --- a/lightllm/common/kv_cache_mem_manager/windowed_mtp.py +++ b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py @@ -96,8 +96,8 @@ def __init__(self, mem_manager, offset=0): self.shape = (window, layers, heads * tp_world_size, dim) self.kv_offset = (offset + 7) // 8 * 8 self.kv_nbytes = window * layers * heads * tp_world_size * dim * store.kv_buffer.element_size() - self.end_offset = (self.kv_offset + self.kv_nbytes + 7) // 8 * 8 - self.state_nbytes = self.end_offset + 8 + self.seq_len_offset = (self.kv_offset + self.kv_nbytes + 7) // 8 * 8 + self.state_nbytes = self.seq_len_offset + 8 def assert_page_size(self): page = self.mem_manager.kv_move_buffer[0] @@ -116,7 +116,7 @@ def copy_req_page(self, page_index, req_idx, dp_mems, mode): .view(self.mem_manager.windowed_draft_kv.kv_buffer.dtype) .view(self.shape) ) - end_page = page[self.end_offset : self.end_offset + 8].view(torch.int64) + seq_len_page = page[self.seq_len_offset : self.seq_len_offset + 8].view(torch.int64) window = self.shape[0] indexes = torch.arange(req_idx * window, (req_idx + 1) * window, dtype=torch.int64, device=page.device) for tp_index, mem in enumerate(dp_mems): @@ -130,6 +130,6 @@ def copy_req_page(self, page_index, req_idx, dp_mems, mode): mode=mode, ) if mode == "read": - store.req_to_seq_len[req_idx : req_idx + 1].copy_(end_page, non_blocking=True) + store.req_to_seq_len[req_idx : req_idx + 1].copy_(seq_len_page, non_blocking=True) elif tp_index == 0: - end_page.copy_(store.req_to_seq_len[req_idx : req_idx + 1], non_blocking=True) + seq_len_page.copy_(store.req_to_seq_len[req_idx : req_idx + 1], non_blocking=True) diff --git a/lightllm/common/state_cache_manager/__init__.py b/lightllm/common/state_cache_manager/__init__.py index 40feb63211..5e6f0c6dab 100644 --- a/lightllm/common/state_cache_manager/__init__.py +++ b/lightllm/common/state_cache_manager/__init__.py @@ -5,20 +5,20 @@ from .windowed_mtp import WindowStateCacheManager -def get_hybrid_cache_config(linear_config=None): +def get_hybrid_cache_config(): """Return the model-specific layout used by hybrid CPU/disk cache pages.""" from transformers.configuration_utils import PretrainedConfig from lightllm.utils.envs_utils import get_env_start_args from .windowed_mtp import WindowedMTPCacheConfig args = get_env_start_args() - if linear_config is None: - model_cfg, _ = PretrainedConfig.get_config_dict(args.model_dir) - model_type = model_cfg["model_type"] - if model_type in ("glm5_next", "glm5_next_text"): - linear_config = Glm5NextCacheConfig.from_model_config(model_cfg, args) - elif model_type in ("qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"): - linear_config = LinearAttCacheConfig.load_from_args() + model_cfg, _ = PretrainedConfig.get_config_dict(args.model_dir) + model_type = model_cfg["model_type"] + linear_config = None + if model_type in ("glm5_next", "glm5_next_text"): + linear_config = Glm5NextCacheConfig.from_model_config(model_cfg, args) + elif model_type in ("qwen3_5", "qwen3_5_moe", "qwen3_5_text", "qwen3_5_moe_text"): + linear_config = LinearAttCacheConfig.load_from_args() if args.mtp_windowed: return WindowedMTPCacheConfig.load_from_args(linear_config) if linear_config is not None: diff --git a/lightllm/common/state_cache_manager/windowed_mtp.py b/lightllm/common/state_cache_manager/windowed_mtp.py index fcf502e959..40d9b9ef04 100644 --- a/lightllm/common/state_cache_manager/windowed_mtp.py +++ b/lightllm/common/state_cache_manager/windowed_mtp.py @@ -174,8 +174,8 @@ def get_window_views(self, cpu_cache_tensor, tp_rank): pages = cpu_cache_tensor.view(cpu_cache_tensor.shape[0], -1).view(torch.uint8) start = self.get_cpu_cache_window_offset() + tp_rank * self.get_cpu_cache_window_rank_bytes() kv_bytes = self.window_config.get_state_bytes(self.dtype.itemsize) - 8 - end_offset = start + (kv_bytes + 7) // 8 * 8 + seq_len_offset = start + (kv_bytes + 7) // 8 * 8 return ( pages[:, start : start + kv_bytes].view(self.dtype), - pages[:, end_offset : end_offset + 8].view(torch.int64), + pages[:, seq_len_offset : seq_len_offset + 8].view(torch.int64), ) diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index 8e390af1ff..38fbbf2563 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -59,13 +59,7 @@ def _init_mem_manager(self): # The shared manager owns target KV and, when enabled, draft window KV. self.mem_manager = self.main_model.mem_manager if self.uses_windowed_draft_kv: - self.mem_manager.init_windowed_draft_kv( - requests=self.req_manager.max_request_num + 1, - layers=self.config["n_layer"], - kv_heads=self.config["num_key_value_heads"] // self.tp_world_size_, - head_dim=self.config.get("head_dim", self.config["n_embed"] // self.config["num_attention_heads"]), - window=self.args.mtp_window_size, - ) + self.mem_manager.init_windowed_draft_kv(self.req_manager.max_request_num) def _init_att_backend(self): if self.uses_windowed_draft_kv: diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 4b4ee84c68..a3306e092f 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -215,6 +215,14 @@ def init_model(self, kvargs): [rank for rank in range(self.global_world_size)], backend="nccl" ) + if not self.args.mtp_windowed and ( + self.args.run_mode in ["prefill", "decode"] or self.args.enable_dp_prompt_cache_fetch + ): + # 如果存在需要跨进程使用mem manger的特性，则将mem manager写入到 shm中，方便 + # 读取 + self.model.mem_manager.write_to_shm(req_manager=self.model.req_manager) + dist.barrier(group=self.node_nccl_group) + # 同一 DP 组内只需主 rank 初始化真实的 capture buffer 并执行后续相关操作； # 非主 rank 不需要分配 buffer，避免重复占用内存。 if self.is_master_in_dp: @@ -231,6 +239,9 @@ def init_model(self, kvargs): self.init_custom() + if self.args.enable_dp_prompt_cache_fetch: + self.init_dp_kv_shared() + self.shm_reqs_io_buffer = ShmObjsIOBuffer() # 只会在 pd pd 模式下才会使用，用于上传分块传输任务是否成功。 self.shm_pd_trans_io_buffer = ShmObjsIOBuffer(tail_str="pd") @@ -239,14 +250,11 @@ def init_model(self, kvargs): self.init_mtp_draft_model(model_kvargs) self.init_spec_engine() - # 主模型和 draft 的缓存均就绪后，再发布给跨进程 KV 访问方。 - if self.is_pd_mode or self.args.enable_dp_prompt_cache_fetch: + # window PD 需要在 draft 分配独立缓存后再导出。 + if self.args.mtp_windowed and self.is_pd_mode: self.model.mem_manager.write_to_shm(req_manager=self.model.req_manager) dist.barrier(group=self.node_nccl_group) - if self.args.enable_dp_prompt_cache_fetch: - self.init_dp_kv_shared() - if self.args.enable_cpu_cache: self.multi_level_cache_module = MultiLevelKvCacheModule(self) diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index fadf2febee..373bfec43c 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -308,7 +308,7 @@ def get_mtp_weight_layer_num() -> int: return args.mtp_step if mtp_mode == "eagle_no_att": return 1 - if getattr(args, "mtp_windowed", False): + if args.mtp_windowed: return _get_mtp_draft_backbone_layer_num(args.mtp_draft_model_dir[0]) return get_added_mtp_kv_layer_num() diff --git a/lightllm/utils/profile_max_tokens.py b/lightllm/utils/profile_max_tokens.py index bbbfc700cd..e9adc79743 100644 --- a/lightllm/utils/profile_max_tokens.py +++ b/lightllm/utils/profile_max_tokens.py @@ -74,17 +74,16 @@ def profile_mtp_weight_memory(model): weight_memory_before = torch.cuda.memory_allocated() yield target_weight_bytes = torch.cuda.memory_allocated() - weight_memory_before - mem_fraction = model.mem_fraction - if getattr(model.args, "mtp_windowed", False): + if model.args.mtp_windowed: from lightllm.common.state_cache_manager.windowed_mtp import load_window_state_config config = load_window_state_config(model.args) draft_cache_bytes = (model.max_req_num + 1) * config.get_state_bytes(model.data_type.itemsize) total_gpu_bytes = torch.cuda.get_device_properties(get_current_device_id()).total_memory # Include the window reservation in the existing per-rank budget reduction. - mem_fraction -= draft_cache_bytes / total_gpu_bytes + model.mem_fraction -= draft_cache_bytes / total_gpu_bytes model.mem_fraction = get_mtp_adjusted_mem_fraction( - mem_fraction=mem_fraction, + mem_fraction=model.mem_fraction, target_weight_bytes=target_weight_bytes, target_layer_num=model.config["n_layer"], mtp_layer_num=get_mtp_weight_layer_num(), diff --git a/unit_tests/common/test_windowed_mtp.py b/unit_tests/common/test_windowed_mtp.py index c50faf886a..82a392eeb1 100644 --- a/unit_tests/common/test_windowed_mtp.py +++ b/unit_tests/common/test_windowed_mtp.py @@ -3,20 +3,8 @@ import pytest import torch -from lightllm.common.kv_cache_mem_manager.windowed_mtp import WindowKVStore -from lightllm.common.kv_cache_mem_manager.windowed_mtp import WindowKVPageHelper -from lightllm.common.state_cache_manager.windowed_mtp import ( - WindowedMTPCacheConfig, - WindowStateBuffers, - WindowStateCacheConfig, -) - - -def _checkpoint_buffers(slots, layers, window, kv_heads, head_dim, dtype): - return SimpleNamespace( - kv_buffer=torch.zeros((slots, layers, window, 2 * kv_heads, head_dim), dtype=dtype), - seq_lens=torch.zeros(slots, dtype=torch.int64), - ) +from lightllm.common.kv_cache_mem_manager.windowed_mtp import WindowKVStore, WindowKVPageHelper +from lightllm.common.state_cache_manager.windowed_mtp import WindowedMTPCacheConfig, WindowStateCacheConfig @pytest.fixture @@ -30,114 +18,54 @@ def tensor_cache(): g_cache_manager.cache_env_out() -def test_checkpoint_round_trip_preserves_physical_ring_layout(): - layers = 2 - window = 4 - kv_heads = 2 - head_dim = 3 - source_req = 1 - destination_req = 2 - checkpoint_slot = 1 - - store = WindowKVStore( - requests=3, - layers=layers, - kv_heads=kv_heads, - head_dim=head_dim, - dtype=torch.float32, - device="cpu", - window=window, - ) - expected_kv = torch.arange( - layers * window * 2 * kv_heads * head_dim, - dtype=torch.float32, - ).view(layers, window, 2 * kv_heads, head_dim) - store.kv_buffer[:, source_req * window : (source_req + 1) * window].copy_(expected_kv) - store.req_to_seq_len[source_req] = 9 - - buffers = _checkpoint_buffers( - slots=2, - layers=layers, - window=window, - kv_heads=kv_heads, - head_dim=head_dim, - dtype=store.kv_buffer.dtype, - ) - store.save_checkpoint(source_req, buffers, checkpoint_slot) - - store.kv_buffer[:, destination_req * window : (destination_req + 1) * window].fill_(-1) - store.req_to_seq_len[destination_req] = -1 - store.restore_checkpoint(destination_req, buffers, checkpoint_slot) - - assert torch.equal(buffers.kv_buffer[checkpoint_slot], expected_kv) - assert torch.equal(store.kv_buffer[:, destination_req * window : (destination_req + 1) * window], expected_kv) - assert store.req_to_seq_len[destination_req].item() == 9 - - -class _TargetStateLayout: - def __init__(self, page_bytes): - self.page_bytes = page_bytes - - def get_cpu_cache_big_page_bytes(self): - return self.page_bytes - - -def test_cpu_page_keeps_tp_rank_window_states_separate(): - target_page_bytes = 32 - window_config = WindowStateCacheConfig(layers=2, kv_heads=1, head_dim=3, window=4) - config = WindowedMTPCacheConfig( - window_config=window_config, - dtype=torch.float16, - tp_world_size=2, - full_att_all_num_kv_heads=0, - full_att_head_dim=0, - full_att_layer_num=0, - linear_config=_TargetStateLayout(target_page_bytes), - ) - pages = torch.zeros((2, config.get_cpu_cache_big_page_bytes()), dtype=torch.uint8) - - rank0_kv, rank0_ends = config.get_window_views(pages, tp_rank=0) - rank1_kv, rank1_ends = config.get_window_views(pages, tp_rank=1) - rank0_kv.fill_(1) - rank1_kv.fill_(2) - rank0_ends.copy_(torch.tensor([[11], [12]], dtype=torch.int64)) - rank1_ends.copy_(torch.tensor([[21], [22]], dtype=torch.int64)) - - restored_rank0 = config.get_window_views(pages, tp_rank=0) - restored_rank1 = config.get_window_views(pages, tp_rank=1) - assert torch.all(restored_rank0[0] == 1) - assert torch.all(restored_rank1[0] == 2) - assert restored_rank0[1].tolist() == [[11], [12]] - assert restored_rank1[1].tolist() == [[21], [22]] - assert torch.count_nonzero(pages[:, :target_page_bytes]).item() == 0 +def test_checkpoint_and_cpu_rank_layout(): + store = WindowKVStore(3, 2, 1, 3, torch.float32, "cpu", 4) + expected = torch.arange(48, dtype=torch.float32).view(2, 4, 2, 3) + store.kv_buffer[:, :4].copy_(expected) + store.req_to_seq_len[0] = 9 + buffers = SimpleNamespace(kv_buffer=torch.zeros(2, 2, 4, 2, 3), seq_lens=torch.zeros(2, dtype=torch.int64)) + store.save_checkpoint(0, buffers, 1) + store.restore_checkpoint(2, buffers, 1) + torch.testing.assert_close(store.kv_buffer[:, 8:12], expected, rtol=0, atol=0) + assert store.req_to_seq_len.tolist() == [9, 0, 9] + + config = WindowStateCacheConfig(2, 1, 3, 4) + target = SimpleNamespace(get_cpu_cache_big_page_bytes=lambda: 32) + layout = WindowedMTPCacheConfig(config, torch.float16, 2, 0, 0, 0, target) + pages = torch.zeros(2, layout.get_cpu_cache_big_page_bytes(), dtype=torch.uint8) + for rank in range(2): + kv, seq_lens = layout.get_window_views(pages, rank) + kv.fill_(rank + 1) + seq_lens.fill_(2 ** 35 + rank) + for rank in range(2): + kv, seq_lens = layout.get_window_views(pages, rank) + assert kv.eq(rank + 1).all() + assert seq_lens.eq(2 ** 35 + rank).all() + assert not pages[:, :32].any() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("window", [1, 4, 8]) -def test_window_commit_and_pack_handles_ragged_strided_inputs_and_rejections(window, tensor_cache): +def test_commit_pack_and_graph_replay(window, tensor_cache): store = WindowKVStore(3, 2, 1, 3, torch.float32, "cuda", window) reqs = torch.tensor([0, 2], dtype=torch.int32, device="cuda") - # Neither the hidden rows nor the noise head dimensions are contiguous. hidden = torch.arange(20 * 12, device="cuda", dtype=torch.float32).view(20, 12)[::2, ::2] starts = torch.tensor([0, 7], dtype=torch.int32, device="cuda") - first = torch.tensor([0, 0], dtype=torch.int32, device="cuda") + first = torch.zeros(2, dtype=torch.int32, device="cuda") lengths = torch.tensor([7, 3], dtype=torch.int32, device="cuda") - features, positions, mem_index = store.prepare(reqs, hidden, starts, first, lengths, 7, alloc_func=tensor_cache) + features, _, mem_index = store.prepare(reqs, hidden, starts, first, lengths, 7, alloc_func=tensor_cache) for layer in range(2): store.write(layer, mem_index, features.view(-1, 2, 3) + layer) - accepted_hidden = torch.arange(5 * 6, device="cuda", dtype=torch.float32).view(5, 6) + 1000 + accepted = torch.arange(5 * 6, device="cuda", dtype=torch.float32).view(5, 6) + 1000 starts = torch.tensor([0, 3], dtype=torch.int32, device="cuda") first = torch.tensor([7, 3], dtype=torch.int32, device="cuda") lengths = torch.tensor([2, 1], dtype=torch.int32, device="cuda") - features, positions, mem_index = store.prepare( - reqs, accepted_hidden, starts, first, lengths, 3, alloc_func=tensor_cache - ) + features, _, mem_index = store.prepare(reqs, accepted, starts, first, lengths, 3, alloc_func=tensor_cache) for layer in range(2): store.write(layer, mem_index, features.view(-1, 2, 3) + layer) - + history = [torch.cat((hidden[:7], accepted[:2])), torch.cat((hidden[7:10], accepted[3:4]))] noise = torch.randn(4, 2, 6, device="cuda")[:, :, ::2] - history = [torch.cat((hidden[:7], accepted_hidden[:2])), torch.cat((hidden[7:10], accepted_hidden[3:4]))] for layer in range(2): packed = store.pack(layer, reqs, noise, 2, alloc_func=tensor_cache) expected = torch.zeros_like(packed) @@ -148,50 +76,31 @@ def test_window_commit_and_pack_handles_ragged_strided_inputs_and_rejections(win torch.testing.assert_close(packed, expected, rtol=0, atol=0) assert store.req_to_seq_len.tolist() == [9, 0, 4] - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_window_pack_graph_replay_reads_current_request_and_sequence_length(tensor_cache): - store = WindowKVStore(2, 1, 1, 8, torch.float32, "cuda", 4) - store.kv_buffer.copy_(torch.arange(8 * 2 * 8, device="cuda").view(1, 8, 2, 8)) - store.req_to_seq_len.copy_(torch.tensor([7, 1], device="cuda")) - reqs = torch.tensor([0], dtype=torch.int32, device="cuda") - noise = torch.randn(2, 2, 8, device="cuda") stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - store.pack(0, reqs, noise, 2, alloc_func=tensor_cache) + store.pack(1, reqs, noise, 2, alloc_func=tensor_cache) torch.cuda.current_stream().wait_stream(stream) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - packed = store.pack(0, reqs, noise, 2, alloc_func=tensor_cache) + packed = store.pack(1, reqs, noise, 2, alloc_func=tensor_cache) reqs.fill_(1) graph.replay() - expected = torch.zeros_like(packed) - expected[0, 0] = store.kv_buffer[0, 4] - expected[0, 1:3] = noise + expected.zero_() + expected[:, :2] = noise.reshape(2, 2, 2, 3) torch.testing.assert_close(packed, expected, rtol=0, atol=0) - store.reset_req(1) + store.kv_buffer[:, window : window + 1].fill_(42) + store.req_to_seq_len[1] = 1 graph.replay() expected.zero_() - expected[0, :2] = noise + expected[:, 0] = 42 + expected[:, 1:3] = noise.reshape(2, 2, 2, 3) torch.testing.assert_close(packed, expected, rtol=0, atol=0) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -def test_empty_window_batch_uses_existing_empty_token_shapes(tensor_cache): - store = WindowKVStore(1, 1, 1, 8, torch.float32, "cuda", 4) - reqs = torch.empty(0, dtype=torch.int32, device="cuda") - features, positions, mem_index = store.prepare( - reqs, torch.empty(0, 16, device="cuda"), reqs, reqs, reqs, 3, alloc_func=tensor_cache - ) - assert features.shape == (0, 16) - assert positions.shape == mem_index.shape == (0,) - assert store.pack(0, reqs, torch.empty(0, 2, 8, device="cuda"), 2, alloc_func=tensor_cache).shape == (0, 6, 2, 8) - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("tp_world_size", [1, 2, 4]) -def test_window_target_pages_reuse_normal_cpu_transfer_with_partial_pages(monkeypatch, tp_world_size): +def test_cpu_transfer_partial_pages(monkeypatch, tp_world_size): from lightllm.common.basemodel.triton_kernel.kv_cache_offload import load_cpu_kv_to_gpu, offload_gpu_kv_to_cpu args = SimpleNamespace(cpu_cache_token_page_size=4, linear_att_hash_page_size=2, linear_att_page_block_num=2) @@ -206,7 +115,7 @@ def test_window_target_pages_reuse_normal_cpu_transfer_with_partial_pages(monkey expected = [kv.clone() for kv in kvs] mem_index = torch.tensor([1, 5, -1, -1], dtype=torch.int32, device="cuda") page_index = torch.zeros(1, dtype=torch.int64, device="cuda") - readies = torch.zeros(1, dtype=torch.int32, device="cuda") + readies = torch.zeros(1, dtype=torch.bool, device="cuda") for rank, kv in enumerate(kvs): offload_gpu_kv_to_cpu(mem_index, kv, None, cpu_kv, None, page_index, readies, rank, tp_world_size, 16) for rank, kv in enumerate(kvs): @@ -214,12 +123,12 @@ def test_window_target_pages_reuse_normal_cpu_transfer_with_partial_pages(monkey load_cpu_kv_to_gpu(mem_index, kv, None, cpu_kv, None, page_index, rank, tp_world_size, 16) torch.testing.assert_close(kv[:, [1, 5]], expected[rank][:, [1, 5]], rtol=0, atol=0) assert not kv[:, [0, 2, 3, 4, 6, 7]].any() - assert pages[:, config.get_cpu_cache_window_offset() :].count_nonzero() == 0 + assert not pages[:, config.get_cpu_cache_window_offset() :].any() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") @pytest.mark.parametrize("prefill_tp,decode_tp", [(1, 1), (1, 2), (2, 1)]) -def test_window_pd_page_reuses_page_io_across_tp_sizes(monkeypatch, prefill_tp, decode_tp): +def test_pd_page_round_trip(monkeypatch, prefill_tp, decode_tp): def make_mems(tp): monkeypatch.setattr("lightllm.utils.envs_utils.get_env_start_args", lambda: SimpleNamespace(tp=tp, dp=1)) mems = [ @@ -233,8 +142,8 @@ def make_mems(tp): helper.assert_page_size() return mems, helper - source, write_helper = make_mems(prefill_tp) - destination, read_helper = make_mems(decode_tp) + source, writer = make_mems(prefill_tp) + destination, reader = make_mems(decode_tp) global_kv = torch.randn(2, 4, 4, 8, device="cuda") def shard_heads(rank, tp): @@ -242,15 +151,14 @@ def shard_heads(rank, tp): return list(range(start, start + 2 // tp)) + list(range(2 + start, 2 + start + 2 // tp)) for rank, mem in enumerate(source): - store = mem.windowed_draft_kv - store.kv_buffer[:, :4].copy_(global_kv[:, :, shard_heads(rank, prefill_tp)]) - store.req_to_seq_len[0] = 7 + mem.windowed_draft_kv.kv_buffer[:, :4].copy_(global_kv[:, :, shard_heads(rank, prefill_tp)]) + mem.windowed_draft_kv.req_to_seq_len[0] = 7 source[0].kv_move_buffer[:, :16].fill_(42) - write_helper.copy_req_page(0, 0, source, "write") + writer.copy_req_page(0, 0, source, "write") destination[0].kv_move_buffer.copy_(source[0].kv_move_buffer) for mem in destination: mem.windowed_draft_kv.kv_buffer.zero_() - read_helper.copy_req_page(0, 2, destination, "read") + reader.copy_req_page(0, 2, destination, "read") for rank, mem in enumerate(destination): store = mem.windowed_draft_kv torch.testing.assert_close(store.kv_buffer[:, 8:12], global_kv[:, :, shard_heads(rank, decode_tp)]) @@ -260,57 +168,24 @@ def shard_heads(rank, tp): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("linear_target", [False, True]) -def test_tail_checkpoint_copy_preserves_stream_order_and_neighbor_slots(linear_target): +def test_tail_checkpoint_copy_preserves_stream_order(): from lightllm.common.kv_cache_mem_manager.operator.windowed_mtp import WindowedMTPMemOperator - from lightllm.common.kv_cache_mem_manager.operator.linear_att import LinearAttMemOperator + from lightllm.common.state_cache_manager.windowed_mtp import WindowStateBuffers - config = WindowStateCacheConfig(layers=2, kv_heads=1, head_dim=8, window=4) - source_window = WindowStateBuffers(2, config, torch.bfloat16) - destination_window = WindowStateBuffers(2, config, torch.bfloat16) - source_target = tuple( - torch.zeros((2, 4, 8), dtype=dtype, pin_memory=True) for dtype in (torch.bfloat16, torch.float32) - ) - destination_target = tuple(torch.zeros_like(state, pin_memory=True) for state in source_target) - source = SimpleNamespace( - draft_window=source_window, - get_state_cache=lambda slot: tuple(state[slot] for state in source_target), - ) - destination = SimpleNamespace( - draft_window=destination_window, - get_state_cache=lambda slot: tuple(state[slot] for state in destination_target), - ) - operator_class = LinearAttMemOperator if linear_target else WindowedMTPMemOperator - operator = operator_class.__new__(operator_class) + config = WindowStateCacheConfig(2, 1, 8, 4) + source = SimpleNamespace(draft_window=WindowStateBuffers(2, config, torch.bfloat16)) + destination = SimpleNamespace(draft_window=WindowStateBuffers(2, config, torch.bfloat16)) + operator = WindowedMTPMemOperator.__new__(WindowedMTPMemOperator) operator.mem_manager = SimpleNamespace(big_page_buffers=destination) - stream = torch.cuda.Stream() with torch.cuda.stream(stream): - # Keep the save pending so a host-side tail copy would see stale data. torch.cuda._sleep(1_000_000) - expected_kv = torch.randn(source_window.kv_buffer[1].shape, dtype=torch.bfloat16, device="cuda") - expected_seq_len = torch.tensor([2 ** 35 + 7], dtype=torch.int64, device="cuda") - expected_target = [torch.randn(state.shape[1:], dtype=state.dtype, device="cuda") for state in source_target] - source_window.kv_buffer[1].copy_(expected_kv, non_blocking=True) - source_window.seq_lens[1:2].copy_(expected_seq_len, non_blocking=True) - for state, expected in zip(source_target, expected_target): - state[1].copy_(expected, non_blocking=True) - + expected = torch.randn(source.draft_window.kv_buffer[1].shape, dtype=torch.bfloat16, device="cuda") + seq_len = torch.tensor([2 ** 35 + 7], dtype=torch.int64, device="cuda") + source.draft_window.kv_buffer[1].copy_(expected, non_blocking=True) + source.draft_window.seq_lens[1:2].copy_(seq_len, non_blocking=True) operator._copy_tail_state(source, source_id=1, destination_id=0) - restored_kv = torch.empty_like(expected_kv) - restored_seq_len = torch.empty_like(expected_seq_len) - restored_kv.copy_(destination_window.kv_buffer[0], non_blocking=True) - restored_seq_len.copy_(destination_window.seq_lens[:1], non_blocking=True) stream.synchronize() - - torch.testing.assert_close(restored_kv, expected_kv, rtol=0, atol=0) - torch.testing.assert_close(restored_seq_len, expected_seq_len, rtol=0, atol=0) - torch.testing.assert_close(destination_window.kv_buffer[0], expected_kv.cpu(), rtol=0, atol=0) - assert not destination_window.kv_buffer[1].any() - assert destination_window.seq_lens.tolist() == [2 ** 35 + 7, 0] - for state, expected in zip(destination_target, expected_target): - if linear_target: - torch.testing.assert_close(state[0], expected.cpu(), rtol=0, atol=0) - else: - assert not state[0].any() - assert not state[1].any() + torch.testing.assert_close(destination.draft_window.kv_buffer[0], expected.cpu(), rtol=0, atol=0) + assert destination.draft_window.seq_lens.tolist() == [2 ** 35 + 7, 0] + assert not destination.draft_window.kv_buffer[1].any() From e7b327d5e9df1d8a4416ae51a2e55efdd0758074 Mon Sep 17 00:00:00 2001 From: Zhao Xintong Date: Thu, 29 Oct 2026 15:45:07 +0000 Subject: [PATCH 8/9] fix --- .../basemodel/attention/fa3/windowed_mtp.py | 13 ++ lightllm/common/basemodel/basemodel.py | 10 +- .../basemodel/triton_kernel/windowed_mtp.py | 83 +++------ .../kv_cache_mem_manager/windowed_mtp.py | 25 ++- .../state_cache_manager/windowed_mtp.py | 14 +- .../layer_infer/transformer_layer_infer.py | 30 ++-- lightllm/models/qwen3_dflash/model.py | 23 +-- .../layer_infer/transformer_layer_infer.py | 13 +- lightllm/models/qwen3_dflash2/model.py | 8 +- lightllm/server/api_cli.py | 9 +- lightllm/server/api_start.py | 2 +- .../mtp_speculative/proposers/dflash.py | 2 +- .../mtp_speculative/proposers/dspark.py | 2 +- lightllm/utils/config_utils.py | 10 ++ lightllm/utils/profile_max_tokens.py | 19 +- unit_tests/common/test_windowed_mtp.py | 164 +++++++++++------- 16 files changed, 226 insertions(+), 201 deletions(-) diff --git a/lightllm/common/basemodel/attention/fa3/windowed_mtp.py b/lightllm/common/basemodel/attention/fa3/windowed_mtp.py index 510a23eb36..603ba4c132 100644 --- a/lightllm/common/basemodel/attention/fa3/windowed_mtp.py +++ b/lightllm/common/basemodel/attention/fa3/windowed_mtp.py @@ -23,6 +23,18 @@ def init_state(self): + 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]), @@ -30,5 +42,6 @@ def decode_att(self, q, k, v, att_control: AttControl = AttControl(), alloc_func 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) diff --git a/lightllm/common/basemodel/basemodel.py b/lightllm/common/basemodel/basemodel.py index 7e8e850904..9a167a53c2 100755 --- a/lightllm/common/basemodel/basemodel.py +++ b/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 @@ -16,7 +15,6 @@ 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, ReqManagerForWindowedMTP -from lightllm.common.build_utils import repair_config 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 diff --git a/lightllm/common/basemodel/triton_kernel/windowed_mtp.py b/lightllm/common/basemodel/triton_kernel/windowed_mtp.py index b562fe913a..5030ad7ec5 100644 --- a/lightllm/common/basemodel/triton_kernel/windowed_mtp.py +++ b/lightllm/common/basemodel/triton_kernel/windowed_mtp.py @@ -4,92 +4,59 @@ @triton.jit -def _gather_window_hidden( - target_hidden, +def _prepare_window_write_indices( b_req_idx, b_start_loc, b_ready_cache_len, - b_q_seq_len, - out_hidden, - out_position_ids, + b_commit_len, out_mem_index, req_to_seq_len, - stride_hidden_bs, - stride_hidden_d, - stride_out_bs, - stride_out_d, - hidden_dim: tl.constexpr, - max_q_seq_len: tl.constexpr, window_size: tl.constexpr, BLOCK: tl.constexpr, ): - cur_token = tl.program_id(0) - cur_batch = tl.program_id(1) + 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_q_seq_len = tl.load(b_q_seq_len + cur_batch) - cur_seq_len = cur_ready_cache_len + cur_q_seq_len + 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) - # A chunk larger than the ring only needs its most recent positions. - cur_window_start = tl.maximum(cur_ready_cache_len, cur_seq_len - window_size) - cur_position = cur_window_start + cur_token - token_mask = cur_position < cur_seq_len - src_index = cur_start_loc + cur_position - cur_ready_cache_len - out_index = cur_batch * max_q_seq_len + cur_token - offs_d = tl.program_id(2) * BLOCK + tl.arange(0, BLOCK) - hidden = tl.load( - target_hidden + src_index.to(tl.int64) * stride_hidden_bs + offs_d * stride_hidden_d, - mask=token_mask & (offs_d < hidden_dim), - other=0, - ) - tl.store(out_hidden + out_index * stride_out_bs + offs_d * stride_out_d, hidden, mask=offs_d < hidden_dim) - if tl.program_id(2) == 0: - tl.store(out_position_ids + out_index, tl.where(token_mask, cur_position, -1)) - # This is a logical token index into the standard [token, head, dim] KV view. - mem_index = cur_req_idx * window_size + cur_position % window_size - tl.store(out_mem_index + out_index, tl.where(token_mask, mem_index, -1)) - if cur_token == 0: - tl.store(req_to_seq_len + cur_req_idx, cur_seq_len) + # 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 gather_window_hidden( - target_hidden, +def prepare_window_write_indices( b_req_idx, b_start_loc, b_ready_cache_len, - b_q_seq_len, + b_commit_len, req_to_seq_len, - max_q_seq_len, window_size, + token_num, alloc_func=torch.empty, ): - max_q_seq_len = min(max_q_seq_len, window_size) - token_num = b_req_idx.numel() * max_q_seq_len - out_hidden = alloc_func((token_num, target_hidden.shape[1]), dtype=target_hidden.dtype, device=target_hidden.device) - position_ids = alloc_func((token_num,), dtype=torch.int64, device=target_hidden.device) - mem_index = alloc_func((token_num,), dtype=torch.int32, device=target_hidden.device) - if token_num: - _gather_window_hidden[(max_q_seq_len, b_req_idx.numel(), triton.cdiv(target_hidden.shape[1], 512))]( - target_hidden, + """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_q_seq_len, - out_hidden, - position_ids, + b_commit_len, mem_index, req_to_seq_len, - stride_hidden_bs=target_hidden.stride(0), - stride_hidden_d=target_hidden.stride(1), - stride_out_bs=out_hidden.stride(0), - stride_out_d=out_hidden.stride(1), - hidden_dim=target_hidden.shape[1], - max_q_seq_len=max_q_seq_len, window_size=window_size, - BLOCK=512, + BLOCK=256, ) - return out_hidden, position_ids, mem_index + return mem_index @triton.jit diff --git a/lightllm/common/kv_cache_mem_manager/windowed_mtp.py b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py index a2d047a7f5..e867ce77c8 100644 --- a/lightllm/common/kv_cache_mem_manager/windowed_mtp.py +++ b/lightllm/common/kv_cache_mem_manager/windowed_mtp.py @@ -34,27 +34,34 @@ def restore_checkpoint(self, req_idx, buffers, slot): if self.kv_buffer.is_cuda: torch.cuda.current_stream(self.kv_buffer.device).synchronize() - def prepare( + def prepare_write_indices( self, b_req_idx, - target_hidden, b_start_loc, b_ready_cache_len, - b_q_seq_len, - max_q_seq_len, + b_commit_len, + token_num, alloc_func=torch.empty, ): - from lightllm.common.basemodel.triton_kernel.windowed_mtp import gather_window_hidden + """Map committed window rows to cache slots without changing the source layout. - return gather_window_hidden( - target_hidden, + token_num is the source hidden row count; b_start_loc indexes that same layout. + b_commit_len counts each request's new prompt rows or accepted verify rows. + + This also updates req_to_seq_len before any layer's KV is written. The caller + must enqueue all layer writes on the same CUDA stream before reading or + checkpointing the window state; this is not a side-effect-free index query. + """ + from lightllm.common.basemodel.triton_kernel.windowed_mtp import prepare_window_write_indices + + return prepare_window_write_indices( b_req_idx, b_start_loc, b_ready_cache_len, - b_q_seq_len, + b_commit_len, self.req_to_seq_len, - max_q_seq_len, self.capacity, + token_num=token_num, alloc_func=alloc_func, ) diff --git a/lightllm/common/state_cache_manager/windowed_mtp.py b/lightllm/common/state_cache_manager/windowed_mtp.py index 40d9b9ef04..99f8512691 100644 --- a/lightllm/common/state_cache_manager/windowed_mtp.py +++ b/lightllm/common/state_cache_manager/windowed_mtp.py @@ -5,8 +5,7 @@ import torch from lightllm.utils.log_utils import init_logger -from lightllm.common.build_utils import repair_config -from lightllm.utils.config_utils import get_config_json +from lightllm.utils.config_utils import load_config from .base import StateCacheManager from .layer_cache import LayerCache @@ -38,16 +37,9 @@ def get_state_bytes(self, element_size): def load_window_state_config(args): - config = get_config_json(args.mtp_draft_model_dir[0]) - # Match the existing draft initialization order: base aliases, Llama KV - # defaults, then the overrides applied by Qwen3.5 adapters. - # Do not normalize again after merging: runtime also retains those aliases. - merge_nested = config.get("model_type") in ("qwen3_5", "qwen3_5_text") - repair_config(config, same_names=["num_attention_heads", "n_head"]) - repair_config(config, same_names=["hidden_size", "n_embd", "n_embed"]) - repair_config(config, same_names=["num_hidden_layers", "n_layer"]) + config = load_config(args.mtp_draft_model_dir[0]) config.setdefault("num_key_value_heads", config["num_attention_heads"]) - if merge_nested: + if config.get("model_type") in ("qwen3_5", "qwen3_5_text") or args.mtp_mode == "dflash2": config.update(config.get("dflash_config", {})) return WindowStateCacheConfig.from_model_config(config, args.mtp_window_size, args.tp // args.dp) diff --git a/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py index 0deb385c0f..08fd0aa494 100644 --- a/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash/layer_infer/transformer_layer_infer.py @@ -1,7 +1,5 @@ import torch -from lightllm.common.basemodel.attention.fa3.windowed_mtp import WindowedMTPDecodeAttState - from lightllm.common.basemodel.triton_kernel.norm.qk_norm import qk_rmsnorm_forward from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd @@ -52,24 +50,26 @@ def _get_context_kv(self, input_embdings, infer_state, layer_weight): def token_attention_forward(self, input_embdings, infer_state, layer_weight): att_state = infer_state.decode_att_state - if not isinstance(att_state, WindowedMTPDecodeAttState): + if not att_state.backend.model.uses_windowed_draft_kv: return super().token_attention_forward(input_embdings, infer_state, layer_weight) q, noise_kv = self._get_qkv(input_embdings, infer_state, layer_weight) - model = att_state.backend.model - scratch = infer_state.mem_manager.windowed_draft_kv.pack( - self.layer_num_ - model.draft_layer_start, - att_state.b_req_idx, - noise_kv, - att_state.block_size, + windowed_kv = att_state.prepare_kv(self.layer_num_, noise_kv, alloc_func=self.alloc_tensor) + output = self._token_attention_kernel(q, infer_state, layer_weight, windowed_kv=windowed_kv) + return self._get_o(output, infer_state, layer_weight) + + def _token_attention_kernel(self, q, infer_state, layer_weight, windowed_kv=None): + # Windowed KV is prepared by the backend; ordinary KV stays in the memory manager. + if windowed_kv is None: + return super()._token_attention_kernel(q, infer_state, layer_weight) + k, v = windowed_kv + output = infer_state.decode_att_state.decode_att( + q=q.view(-1, self.tp_q_head_num_, self.head_dim_), + k=k, + v=v, alloc_func=self.alloc_tensor, ) - output = att_state.decode_att( - q.view(-1, self.tp_q_head_num_, self.head_dim_), - scratch[:, :, : self.tp_k_head_num_], - scratch[:, :, self.tp_k_head_num_ :], - ) - return self._get_o(output.reshape(-1, self.tp_q_head_num_ * self.head_dim_), infer_state, layer_weight) + return output.view(q.shape) def _get_qkv(self, input, infer_state: Qwen3DFlashInferStateInfo, layer_weight: Qwen3DFlashTransformerLayerWeight): q = layer_weight.q_proj.mm(input, use_custom_tensor_mananger=False) diff --git a/lightllm/models/qwen3_dflash/model.py b/lightllm/models/qwen3_dflash/model.py index 38fbbf2563..79fcd98cb3 100644 --- a/lightllm/models/qwen3_dflash/model.py +++ b/lightllm/models/qwen3_dflash/model.py @@ -7,6 +7,7 @@ from lightllm.common.basemodel.basemodel import TpPartBaseModel from lightllm.common.basemodel.attention.fa3.windowed_mtp import WindowedMTPAttBackend from lightllm.common.basemodel.batch_objs import ModelInput, ModelOutput +from lightllm.common.basemodel.triton_kernel.gen_prefill_params import gen_prefill_params from lightllm.models.llama.model import LlamaTpPartModel from lightllm.models.qwen3_dflash.infer_struct import Qwen3DFlashInferStateInfo from lightllm.models.qwen3_dflash.layer_infer.post_layer_infer import Qwen3DFlashPostLayerInfer @@ -104,32 +105,32 @@ def commit_window_kv(self, model_input, hidden, b_req_mtp_start_loc=None, accept """Commit accepted target features to the bounded draft KV store.""" if hidden.numel() == 0: return + store = self.mem_manager.windowed_draft_kv if model_input.is_prefill: b_req_idx = model_input.b_req_idx b_start_loc = model_input.b_prefill_start_loc b_ready_cache_len = model_input.b_ready_cache_len - b_q_seq_len = model_input.b_seq_len - b_ready_cache_len - max_q_seq_len = model_input.max_q_seq_len + b_commit_len, _, _, _, position_ids = gen_prefill_params( + hidden.shape[0], b_ready_cache_len, model_input.b_seq_len + ) else: b_start_loc = b_req_mtp_start_loc b_req_idx = model_input.b_req_idx.index_select(0, b_start_loc) b_ready_cache_len = model_input.b_seq_len.index_select(0, b_start_loc) - 1 - b_q_seq_len = accept_len - max_q_seq_len = self.args.mtp_step + 1 - store = self.mem_manager.windowed_draft_kv - features, position_ids, mem_index = store.prepare( + b_commit_len = accept_len + position_ids = model_input.b_seq_len - 1 + # Preserve the target's packed prefill or verify rows, including dynamic verify layouts. + mem_index = store.prepare_write_indices( b_req_idx, - hidden, b_start_loc, b_ready_cache_len, - b_q_seq_len, - max_q_seq_len, + b_commit_len, + token_num=hidden.shape[0], alloc_func=self.pre_infer.alloc_tensor, ) infer_state = self.infer_state_class() - infer_state.mtp_draft_input_hiddens = features + infer_state.mtp_draft_input_hiddens = hidden features = self.pre_infer.context_forward(None, infer_state, self.pre_post_weight) - position_ids = position_ids.clamp_min(0) infer_state.position_cos = self._cos_cached.index_select(0, position_ids) infer_state.position_sin = self._sin_cached.index_select(0, position_ids) for i, (layer, weight) in enumerate(zip(self.layers_infer, self.trans_layers_weight)): diff --git a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py index b03fd77cc2..6286c9269e 100644 --- a/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py +++ b/lightllm/models/qwen3_dflash2/layer_infer/transformer_layer_infer.py @@ -12,10 +12,15 @@ def __init__(self, layer_num, network_config): self.conv_group_size_ = int(network_config["conv_group_size"]) self.sliding_window_ = int(network_config.get("sliding_window", 0) or 0) - def _token_attention_kernel(self, q, infer_state, layer_weight): - k, v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) - k = k.view(k.shape[0], self.tp_k_head_num_, self.head_dim_) - v = v.view(v.shape[0], self.tp_k_head_num_, self.head_dim_) + def _token_attention_kernel(self, q, infer_state, layer_weight, windowed_kv=None): + if windowed_kv is None: + # Ordinary cache: reshape the shared storage layout to the draft model's KV heads. + k, v = infer_state.mem_manager.get_att_input_params(layer_index=self.layer_num_) + k = k.view(k.shape[0], self.tp_k_head_num_, self.head_dim_) + v = v.view(v.shape[0], self.tp_k_head_num_, self.head_dim_) + else: + # Windowed cache: history + current block, already packed as [batch, seq, heads, dim]. + k, v = windowed_kv q = q.view(-1, self.tp_q_head_num_, self.head_dim_) use_sliding_window = self.sliding_window_ > 0 # DFlash2 attention is non-causal, so its local window must retain both diff --git a/lightllm/models/qwen3_dflash2/model.py b/lightllm/models/qwen3_dflash2/model.py index 3542b37188..63c3f80268 100644 --- a/lightllm/models/qwen3_dflash2/model.py +++ b/lightllm/models/qwen3_dflash2/model.py @@ -20,10 +20,7 @@ class Qwen3DFlash2Model(Qwen3DFlashModel): def _init_config(self): super()._init_config() - dflash_config = self.config.get("dflash_config", {}) - if not isinstance(dflash_config, dict): - raise ValueError("dflash_config must be an object in the DFlash2 checkpoint config") - self.config.update(dflash_config) + self.config.update(self.config.get("dflash_config", {})) rope_parameters = self.config.get("rope_parameters", {}) if "rope_theta" in rope_parameters and "rope_theta" not in self.config: @@ -54,6 +51,7 @@ def _verify_params(self): physical_block_size = self.args.mtp_step + 1 assert physical_block_size <= self.config["block_size"] self.config["block_size"] = physical_block_size + self.uses_windowed_draft_kv = getattr(self.args, "mtp_windowed", False) def _init_custom(self): # The released Qwen3.8 drafter uses its own Qwen3 rotary layout. @@ -62,6 +60,8 @@ def _init_custom(self): self.mask_token_id = int(self.config["mask_token_id"]) def _init_mem_manager(self): + if self.uses_windowed_draft_kv: + return super()._init_mem_manager() main_mem_manager = self.main_model.mem_manager draft_head_num = max(self.config["num_key_value_heads"] // self.tp_world_size_, 1) draft_width = draft_head_num * self.config["head_dim"] diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index a178de5e3b..ca13803b3d 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -833,9 +833,14 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument( "--mtp_windowed", action="store_true", - help="Use a sliding window for draft KV in DFlash-style MTP.", + help="""Use sliding-window draft KV storage for DFlash, DSpark or DFlash2.""", + ) + parser.add_argument( + "--mtp_window_size", + type=int, + default=512, + help="""Number of tokens retained in the draft KV cache when mtp_windowed is enabled.""", ) - parser.add_argument("--mtp_window_size", type=int, default=512) parser.add_argument( "--mtp_dynamic_verify", action="store_true", diff --git a/lightllm/server/api_start.py b/lightllm/server/api_start.py index 309629ecb3..747dfd284f 100644 --- a/lightllm/server/api_start.py +++ b/lightllm/server/api_start.py @@ -207,7 +207,7 @@ def _launch_subprocesses(args: StartArgs): if args.mtp_windowed: assert args.mtp_window_size > 0 - assert args.mtp_mode in ("dflash", "dspark"), "windowed draft KV requires dflash or dspark" + assert args.mtp_mode in ("dflash", "dspark", "dflash2"), "windowed draft KV requires dflash, dspark or dflash2" if args.llm_kv_type not in (None, "None"): raise ValueError("windowed draft KV currently requires unquantized KV") diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py index 8450433896..c287d00028 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/dflash.py @@ -82,7 +82,7 @@ def propose_next( accepted_tail_rows = (b_req_mtp_start_loc + accept_len - 1).long() # full 模式通过 draft forward 写入 verify 行对应的 KV；windowed 模式 - # 只将接受的 target hidden 投影并写入独立窗口，不保留被拒绝的行。 + # 沿用 verify 行进行投影，仅将接受且仍在窗口内的行写入独立缓存。 if getattr(draft_model, "uses_windowed_draft_kv", False): draft_model.commit_window_kv( target_model_input, target_model_output.mtp_collector.spec_hidden, b_req_mtp_start_loc, accept_len diff --git a/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py b/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py index 7dc44d7e6f..f5c72c793d 100644 --- a/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py +++ b/lightllm/server/router/model_infer/mtp_speculative/proposers/dspark.py @@ -85,7 +85,7 @@ def propose_next( accepted_tail_rows = (b_req_mtp_start_loc + accept_len - 1).long() # full 模式通过 draft forward 写入 verify 行对应的 KV；windowed 模式 - # 只将接受的 target hidden 投影并写入独立窗口，不保留被拒绝的行。 + # 沿用 verify 行进行投影，仅将接受且仍在窗口内的行写入独立缓存。 if getattr(draft_model, "uses_windowed_draft_kv", False): draft_model.commit_window_kv( target_model_input, target_model_output.mtp_collector.spec_hidden, b_req_mtp_start_loc, accept_len diff --git a/lightllm/utils/config_utils.py b/lightllm/utils/config_utils.py index 8f76b4a53c..d40a58fa1f 100644 --- a/lightllm/utils/config_utils.py +++ b/lightllm/utils/config_utils.py @@ -3,6 +3,7 @@ from typing import Optional, List from functools import lru_cache from .envs_utils import get_env_start_args +from lightllm.common.build_utils import repair_config from lightllm.utils.log_utils import init_logger logger = init_logger(__name__) @@ -14,6 +15,15 @@ def get_config_json(model_path: str): return json_obj +def load_config(weight_dir_): + """Read model configuration and normalize the common architecture aliases.""" + config = get_config_json(weight_dir_) + repair_config(config, same_names=["num_attention_heads", "n_head"]) + repair_config(config, same_names=["hidden_size", "n_embd", "n_embed"]) + repair_config(config, same_names=["num_hidden_layers", "n_layer"]) + return config + + def get_running_max_req_size_per_dp(args) -> int: """Return the request capacity for one local DP rank.""" if args.running_max_req_size < 1: diff --git a/lightllm/utils/profile_max_tokens.py b/lightllm/utils/profile_max_tokens.py index e9adc79743..1694dd4ae1 100644 --- a/lightllm/utils/profile_max_tokens.py +++ b/lightllm/utils/profile_max_tokens.py @@ -1,11 +1,10 @@ import os import gc -import json import torch from contextlib import contextmanager from transformers import AutoModelForCausalLM import argparse -from lightllm.common.build_utils import repair_config +from lightllm.utils.config_utils import load_config from lightllm.utils.dist_utils import get_current_device_id from lightllm.utils.envs_utils import get_mtp_weight_layer_num @@ -90,22 +89,6 @@ def profile_mtp_weight_memory(model): ) -def load_config(weight_dir_): - """ - Load model configuration from the specified directory - Args: - weight_dir_: Path to the weight directory - Returns: - config: Model configuration - """ - with open(os.path.join(weight_dir_, "config.json"), "r") as json_file: - config = json.load(json_file) - repair_config(config, same_names=["num_attention_heads", "n_head"]) - repair_config(config, same_names=["hidden_size", "n_embd", "n_embed"]) - repair_config(config, same_names=["num_hidden_layers", "n_layer"]) - return config - - def load_model(model_dir, tp_size, data_type): """ Load the model and apply parallelization diff --git a/unit_tests/common/test_windowed_mtp.py b/unit_tests/common/test_windowed_mtp.py index 82a392eeb1..983dc1de9f 100644 --- a/unit_tests/common/test_windowed_mtp.py +++ b/unit_tests/common/test_windowed_mtp.py @@ -45,85 +45,133 @@ def test_checkpoint_and_cpu_rank_layout(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("window", [1, 4, 8]) -def test_commit_pack_and_graph_replay(window, tensor_cache): +def test_write_indices_preserve_source_layout(tensor_cache): + store = WindowKVStore(3, 1, 1, 3, torch.float32, "cuda", 4) + reqs = torch.tensor([2, 1, 0], dtype=torch.int32, device="cuda") + q_lengths = torch.tensor([7, 0, 2], dtype=torch.int32, device="cuda") + ready = torch.tensor([5, 9, 7], dtype=torch.int32, device="cuda") + starts = torch.tensor([0, 7, 7], dtype=torch.int32, device="cuda") + mem_index = store.prepare_write_indices(reqs, starts, ready, q_lengths, token_num=9, alloc_func=tensor_cache) + # Retain four rows from the long query, skip the empty query, then keep both short-query rows. + assert mem_index.tolist() == [-1, -1, -1, 8, 9, 10, 11, 3, 0] + assert store.req_to_seq_len.tolist() == [9, 9, 12] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("window", [1, 4, 8, 257]) +def test_commit_and_pack(window, tensor_cache): store = WindowKVStore(3, 2, 1, 3, torch.float32, "cuda", window) reqs = torch.tensor([0, 2], dtype=torch.int32, device="cuda") - hidden = torch.arange(20 * 12, device="cuda", dtype=torch.float32).view(20, 12)[::2, ::2] - starts = torch.tensor([0, 7], dtype=torch.int32, device="cuda") + first_len = max(7, window + 3) + hidden = torch.arange((first_len + 3) * 24, device="cuda", dtype=torch.float32).view(-1, 12)[::2, ::2] + starts = torch.tensor([0, first_len], dtype=torch.int32, device="cuda") first = torch.zeros(2, dtype=torch.int32, device="cuda") - lengths = torch.tensor([7, 3], dtype=torch.int32, device="cuda") - features, _, mem_index = store.prepare(reqs, hidden, starts, first, lengths, 7, alloc_func=tensor_cache) + lengths = torch.tensor([first_len, 3], dtype=torch.int32, device="cuda") + mem_index = store.prepare_write_indices( + reqs, starts, first, lengths, token_num=hidden.shape[0], alloc_func=tensor_cache + ) for layer in range(2): - store.write(layer, mem_index, features.view(-1, 2, 3) + layer) + store.write(layer, mem_index, hidden.view(-1, 2, 3) + layer) accepted = torch.arange(5 * 6, device="cuda", dtype=torch.float32).view(5, 6) + 1000 starts = torch.tensor([0, 3], dtype=torch.int32, device="cuda") - first = torch.tensor([7, 3], dtype=torch.int32, device="cuda") + first = torch.tensor([first_len, 3], dtype=torch.int32, device="cuda") lengths = torch.tensor([2, 1], dtype=torch.int32, device="cuda") - features, _, mem_index = store.prepare(reqs, accepted, starts, first, lengths, 3, alloc_func=tensor_cache) + mem_index = store.prepare_write_indices( + reqs, starts, first, lengths, token_num=accepted.shape[0], alloc_func=tensor_cache + ) + assert mem_index.shape == (5,) + assert mem_index[[2, 4]].eq(-1).all() for layer in range(2): - store.write(layer, mem_index, features.view(-1, 2, 3) + layer) - history = [torch.cat((hidden[:7], accepted[:2])), torch.cat((hidden[7:10], accepted[3:4]))] + store.write(layer, mem_index, accepted.view(-1, 2, 3) + layer) + history = [torch.cat((hidden[:first_len], accepted[:2])), torch.cat((hidden[first_len:], accepted[3:4]))] noise = torch.randn(4, 2, 6, device="cuda")[:, :, ::2] for layer in range(2): packed = store.pack(layer, reqs, noise, 2, alloc_func=tensor_cache) - expected = torch.zeros_like(packed) for batch, values in enumerate(history): retained = values[-window:].view(-1, 2, 3) + layer - expected[batch, : len(retained)] = retained - expected[batch, len(retained) : len(retained) + 2] = noise[batch * 2 : (batch + 1) * 2] - torch.testing.assert_close(packed, expected, rtol=0, atol=0) - assert store.req_to_seq_len.tolist() == [9, 0, 4] + expected = torch.cat((retained, noise[batch * 2 : (batch + 1) * 2])) + torch.testing.assert_close(packed[batch, : len(expected)], expected, rtol=0, atol=0) + assert store.req_to_seq_len.tolist() == [first_len + 2, 0, 4] + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize("sliding_window", [0, 4]) +def test_fa3_graph_refreshes_request_state(tensor_cache, sliding_window): + from lightllm.common.basemodel.attention.fa3.windowed_mtp import WindowedMTPDecodeAttState + from lightllm.common.basemodel.cuda_graph import CudaGraph + from lightllm.common.basemodel.infer_struct import InferStateInfo + from lightllm.models.qwen3_dflash.layer_infer.transformer_layer_infer import Qwen3DFlashTransformerLayerInfer + from lightllm.models.qwen3_dflash2.layer_infer.transformer_layer_infer import Qwen3DFlash2TransformerLayerInfer + from lightllm.utils.sgl_utils import flash_attn_with_kvcache + + if flash_attn_with_kvcache is None: + pytest.skip("FA3 required") + torch.manual_seed(0) + window, block_size, head_dim = 4, 2, 64 + store = WindowKVStore(3, 1, 1, head_dim, torch.bfloat16, "cuda", window) + store.kv_buffer.normal_() + store.req_to_seq_len.copy_(torch.tensor([7, 0, 2], device="cuda")) + backend = SimpleNamespace( + model=SimpleNamespace(block_size=block_size, draft_layer_start=0, uses_windowed_draft_kv=True) + ) + + def make_state(req_ids): + state = InferStateInfo() + state.is_prefill = False + state.input_ids = torch.zeros(4, dtype=torch.int64, device="cuda") + state.b_req_idx = torch.tensor(req_ids, dtype=torch.int32, device="cuda").repeat_interleave(block_size) + state.mem_manager = SimpleNamespace(windowed_draft_kv=store) + state.query = torch.randn(4, 2, head_dim, dtype=torch.bfloat16, device="cuda") + state.noise_kv = torch.randn(4, 2, head_dim, dtype=torch.bfloat16, device="cuda") + state.decode_att_state = WindowedMTPDecodeAttState(backend=backend, infer_state=state) + state.init_att_state() + return state + + layer_class = Qwen3DFlash2TransformerLayerInfer if sliding_window else Qwen3DFlashTransformerLayerInfer + layer = layer_class.__new__(layer_class) + layer.sliding_window_ = sliding_window + layer.layer_num_, layer.tp_q_head_num_, layer.tp_k_head_num_, layer.head_dim_ = 0, 2, 1, head_dim + layer._get_qkv = lambda input, state, weight: (state.query, state.noise_kv) + layer._get_o = lambda output, state, weight: output.flatten(1) + weight = SimpleNamespace() + captured = make_state([0, 2]) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - store.pack(1, reqs, noise, 2, alloc_func=tensor_cache) + for _ in range(3): + layer.token_attention_forward(captured.query.flatten(1), captured, weight) torch.cuda.current_stream().wait_stream(stream) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - packed = store.pack(1, reqs, noise, 2, alloc_func=tensor_cache) - reqs.fill_(1) - graph.replay() - expected.zero_() - expected[:, :2] = noise.reshape(2, 2, 2, 3) - torch.testing.assert_close(packed, expected, rtol=0, atol=0) - store.kv_buffer[:, window : window + 1].fill_(42) - store.req_to_seq_len[1] = 1 - graph.replay() - expected.zero_() - expected[:, 0] = 42 - expected[:, 1:3] = noise.reshape(2, 2, 2, 3) - torch.testing.assert_close(packed, expected, rtol=0, atol=0) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") -@pytest.mark.parametrize("tp_world_size", [1, 2, 4]) -def test_cpu_transfer_partial_pages(monkeypatch, tp_world_size): - from lightllm.common.basemodel.triton_kernel.kv_cache_offload import load_cpu_kv_to_gpu, offload_gpu_kv_to_cpu - - args = SimpleNamespace(cpu_cache_token_page_size=4, linear_att_hash_page_size=2, linear_att_page_block_num=2) - monkeypatch.setattr("lightllm.utils.envs_utils.get_env_start_args", lambda: args) - config = WindowedMTPCacheConfig(WindowStateCacheConfig(1, 1, 8, 4), torch.float32, tp_world_size, 2, 8, 2) - pages = torch.zeros((1, config.get_cpu_cache_big_page_bytes()), dtype=torch.uint8, pin_memory=True) - cpu_kv = config.get_full_att_kv_view(pages) - replicas = max(1, tp_world_size // 2) - kvs = [torch.randn(2, 8, 2 * max(1, 2 // tp_world_size), 8, device="cuda") for _ in range(tp_world_size)] - for rank in range(tp_world_size): - kvs[rank].copy_(kvs[rank // replicas * replicas]) - expected = [kv.clone() for kv in kvs] - mem_index = torch.tensor([1, 5, -1, -1], dtype=torch.int32, device="cuda") - page_index = torch.zeros(1, dtype=torch.int64, device="cuda") - readies = torch.zeros(1, dtype=torch.bool, device="cuda") - for rank, kv in enumerate(kvs): - offload_gpu_kv_to_cpu(mem_index, kv, None, cpu_kv, None, page_index, readies, rank, tp_world_size, 16) - for rank, kv in enumerate(kvs): - kv.zero_() - load_cpu_kv_to_gpu(mem_index, kv, None, cpu_kv, None, page_index, rank, tp_world_size, 16) - torch.testing.assert_close(kv[:, [1, 5]], expected[rank][:, [1, 5]], rtol=0, atol=0) - assert not kv[:, [0, 2, 3, 4, 6, 7]].any() - assert not pages[:, config.get_cpu_cache_window_offset() :].any() + output = layer.token_attention_forward(captured.query.flatten(1), captured, weight) + runner = CudaGraph.__new__(CudaGraph) + runner.graph = {4: (graph, captured, output)} + + for req_ids, lengths in [([2, 1], [7, 6, 0]), ([1, 0], [2, 1, 0])]: + store.kv_buffer.normal_() + store.req_to_seq_len.copy_(torch.tensor(lengths, device="cuda")) + runtime = make_state(req_ids) + att_state = runtime.decode_att_state + expected_lengths = [min(lengths[req], window) + block_size for req in req_ids] + assert att_state.cache_seqlens.tolist() == expected_lengths + expected = [] + for batch, req in enumerate(req_ids): + positions = torch.arange(max(0, lengths[req] - window), lengths[req], device="cuda") + history = store.kv_buffer[0, req * window + positions % window] + noise = runtime.noise_kv[batch * block_size : (batch + 1) * block_size] + kv = torch.cat((history, noise)).float() + q = runtime.query[batch * block_size : (batch + 1) * block_size].float() + scores = torch.einsum("qhd,kd->hqk", q, kv[:, 0]) * head_dim**-0.5 + if sliding_window: + q_pos = torch.arange(block_size, device="cuda") + len(history) + k_pos = torch.arange(len(kv), device="cuda") + mask = (q_pos[:, None] - k_pos[None, :]).abs() >= sliding_window + scores.masked_fill_(mask[None], float("-inf")) + expected.append(torch.einsum("hqk,kd->qhd", scores.softmax(-1), kv[:, 1])) + expected = torch.cat(expected).flatten(1).to(runtime.query.dtype) + actual = runner._replay(runtime) + torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") From 195aa51916f16a9dd56b5bbc942911232637270a Mon Sep 17 00:00:00 2001 From: Zhao Xintong Date: Thu, 29 Oct 2026 16:08:31 +0000 Subject: [PATCH 9/9] fix formatting --- unit_tests/common/test_windowed_mtp.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/unit_tests/common/test_windowed_mtp.py b/unit_tests/common/test_windowed_mtp.py index 5b1e037511..734b92772f 100644 --- a/unit_tests/common/test_windowed_mtp.py +++ b/unit_tests/common/test_windowed_mtp.py @@ -27,7 +27,9 @@ def test_checkpoint_and_cpu_rank_layout(): store.req_to_seq_len[0] = 9 buffers = SimpleNamespace(kv_buffer=torch.zeros(2, 2, 4, 2, 3), seq_lens=torch.zeros(2, dtype=torch.int64)) manager = ReqManagerForWindowedMTP.__new__(ReqManagerForWindowedMTP) - manager.mem_manager = SimpleNamespace(windowed_draft_kv=store, big_page_buffers=SimpleNamespace(draft_window=buffers)) + manager.mem_manager = SimpleNamespace( + windowed_draft_kv=store, big_page_buffers=SimpleNamespace(draft_window=buffers) + ) manager.save_big_page_states(torch.tensor([0]), [SimpleNamespace(req_idx=0)], [1]) store.restore_checkpoint(2, buffers, 1) torch.testing.assert_close(store.kv_buffer[:, 8:12], expected, rtol=0, atol=0) @@ -166,7 +168,7 @@ def make_state(req_ids): noise = runtime.noise_kv[batch * block_size : (batch + 1) * block_size] kv = torch.cat((history, noise)).float() q = runtime.query[batch * block_size : (batch + 1) * block_size].float() - scores = torch.einsum("qhd,kd->hqk", q, kv[:, 0]) * head_dim**-0.5 + scores = torch.einsum("qhd,kd->hqk", q, kv[:, 0]) * head_dim ** -0.5 if sliding_window: q_pos = torch.arange(block_size, device="cuda") + len(history) k_pos = torch.arange(len(kv), device="cuda")