Fix Qwen3Next MTP linear-att page moves

Author: sufubao

lightllm/common/kv_cache_mem_manager/qwen3next_mem_manager.py

@@ -208,9 +208,9 @@ def write_req_to_page(
         dp_mems: List["Qwen3NextMemManager"],
     ):
         conv_page, ssm_page = self.view_page_to_linear_att_state(page_index)
-        req_buffer_idx = req_idx * (get_env_start_args().mtp_step + 1)
+        conv_req_idx, ssm_req_idx = self._get_req_state_indexes(req_idx)
         for tp_index, mem in enumerate(dp_mems):
-            self._write_one_rank(mem, tp_index, req_buffer_idx, conv_page, ssm_page)
+            self._write_one_rank(mem, tp_index, conv_req_idx, ssm_req_idx, conv_page, ssm_page)
         return
 
     def read_page_to_req(
@@ -220,21 +220,27 @@ def read_page_to_req(
         dp_mems: List["Qwen3NextMemManager"],
     ):
         conv_page, ssm_page = self.view_page_to_linear_att_state(page_index)
-        req_buffer_idx = req_idx * (get_env_start_args().mtp_step + 1)
+        conv_req_idx, ssm_req_idx = self._get_req_state_indexes(req_idx)
         for tp_index, mem in enumerate(dp_mems):
-            self._read_one_rank(mem, tp_index, req_buffer_idx, conv_page, ssm_page)
+            self._read_one_rank(mem, tp_index, conv_req_idx, ssm_req_idx, conv_page, ssm_page)
         return
 
+    def _get_req_state_indexes(self, req_idx: int):
+        mtp_size = get_env_start_args().mtp_step + 1
+        # Conv is one widened slot per request; SSM keeps the historical S+1 block layout.
+        return req_idx, req_idx * mtp_size
+
     def _write_one_rank(
         self,
         mem: "Qwen3NextMemManager",
         tp_index: int,
-        req_buffer_idx: int,
+        conv_req_idx: int,
+        ssm_req_idx: int,
         conv_page: torch.Tensor,
         ssm_page: torch.Tensor,
     ):
-        conv_state = mem.req_to_conv_state.buffer[:, req_buffer_idx, ...]
-        ssm_state = mem.req_to_ssm_state.buffer[:, req_buffer_idx, ...]
+        conv_state = mem.req_to_conv_state.buffer[:, conv_req_idx, ..., : self.conv_shape[-1]]
+        ssm_state = mem.req_to_ssm_state.buffer[:, ssm_req_idx, ...]
         self._copy_conv_state_to_page(conv_state, conv_page, mem, tp_index)
         self._copy_ssm_state_to_page(ssm_state, ssm_page, mem, tp_index)
         return
@@ -408,12 +414,13 @@ def _read_one_rank(
         self,
         mem: "Qwen3NextMemManager",
         tp_index: int,
-        req_buffer_idx: int,
+        conv_req_idx: int,
+        ssm_req_idx: int,
         conv_page: torch.Tensor,
         ssm_page: torch.Tensor,
     ):
-        conv_state = mem.req_to_conv_state.buffer[:, req_buffer_idx, ...]
-        ssm_state = mem.req_to_ssm_state.buffer[:, req_buffer_idx, ...]
+        conv_state = mem.req_to_conv_state.buffer[:, conv_req_idx, ..., : self.conv_shape[-1]]
+        ssm_state = mem.req_to_ssm_state.buffer[:, ssm_req_idx, ...]
         self._copy_page_to_conv_state(conv_page, conv_state, mem, tp_index)
         self._copy_page_to_ssm_state(ssm_page, ssm_state, mem, tp_index)
         return

unit_tests/common/basemodel/test_mtp_decode_cuda_graph.py

@@ -103,7 +103,7 @@ class Qwen3_5MOETpPartModel:
         pass
 
     class Qwen3_5MoeMTPModel:
-        pass
+        is_mtp_draft_model = True
 
     graph = CudaGraph.__new__(CudaGraph)
     graph.mtp_step = 2

unit_tests/common/test_qwen3next_linear_att_page_helper.py

@@ -0,0 +1,112 @@
+from types import SimpleNamespace
+
+import torch
+
+
+class _Buf:
+    def __init__(self, tensor):
+        self.buffer = tensor
+
+
+def _make_config():
+    return SimpleNamespace(
+        tp_world_size=1,
+        linear_layer_num=1,
+        conv_kernel_size=4,
+        global_linear_k_heads=1,
+        global_linear_v_heads=1,
+        num_linear_k_heads=1,
+        num_linear_v_heads=1,
+        head_linear_k_dim=2,
+        head_linear_v_dim=3,
+    )
+
+
+def _make_mem(mtp_step=2, req_slots=4):
+    config = _make_config()
+    conv_dim = (
+        2 * config.num_linear_k_heads * config.head_linear_k_dim
+        + config.num_linear_v_heads * config.head_linear_v_dim
+    )
+    narrow_w = config.conv_kernel_size - 1
+    conv = torch.full(
+        (config.linear_layer_num, req_slots, conv_dim, narrow_w + mtp_step),
+        -9.0,
+        dtype=torch.float32,
+    )
+    ssm = torch.full(
+        (
+            config.linear_layer_num,
+            req_slots * (mtp_step + 1),
+            config.num_linear_v_heads,
+            config.head_linear_k_dim,
+            config.head_linear_v_dim,
+        ),
+        -11.0,
+        dtype=torch.float32,
+    )
+    return SimpleNamespace(
+        linear_config=config,
+        req_to_conv_state=_Buf(conv),
+        req_to_ssm_state=_Buf(ssm),
+        kv_move_buffer=torch.zeros((1, 4096), dtype=torch.uint8),
+    )
+
+
+def test_page_helper_writes_req_conv_slot_and_narrow_width(monkeypatch):
+    import lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager as qwen3next_mem_manager
+    from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextLinearAttPageHelper
+
+    mtp_step = 2
+    req_idx = 2
+    monkeypatch.setattr(qwen3next_mem_manager, "get_env_start_args", lambda: SimpleNamespace(mtp_step=mtp_step))
+
+    mem = _make_mem(mtp_step=mtp_step)
+    helper = Qwen3NextLinearAttPageHelper(mem)
+    mem.kv_move_buffer = torch.zeros((1, helper.state_nbytes), dtype=torch.uint8)
+
+    narrow_w = helper.conv_shape[-1]
+    marker_conv = torch.arange(
+        helper.conv_shape[0] * helper.conv_shape[1] * narrow_w,
+        dtype=torch.float32,
+    ).view(helper.conv_shape)
+    marker_ssm = torch.arange(
+        helper.ssm_shape[0] * helper.ssm_shape[1] * helper.ssm_shape[2] * helper.ssm_shape[3],
+        dtype=torch.float32,
+    ).view(helper.ssm_shape)
+
+    mem.req_to_conv_state.buffer[:, req_idx, :, :narrow_w] = marker_conv
+    mem.req_to_conv_state.buffer[:, req_idx, :, narrow_w:] = 999.0
+    mem.req_to_ssm_state.buffer[:, req_idx * (mtp_step + 1), ...] = marker_ssm
+
+    helper.write_req_to_page(page_index=0, req_idx=req_idx, dp_mems=[mem])
+
+    conv_page, ssm_page = helper.view_page_to_linear_att_state(page_index=0)
+    torch.testing.assert_close(conv_page, marker_conv)
+    torch.testing.assert_close(ssm_page, marker_ssm)
+
+
+def test_page_helper_restores_narrow_conv_to_req_slot(monkeypatch):
+    import lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager as qwen3next_mem_manager
+    from lightllm.common.kv_cache_mem_manager.qwen3next_mem_manager import Qwen3NextLinearAttPageHelper
+
+    mtp_step = 2
+    req_idx = 2
+    monkeypatch.setattr(qwen3next_mem_manager, "get_env_start_args", lambda: SimpleNamespace(mtp_step=mtp_step))
+
+    mem = _make_mem(mtp_step=mtp_step)
+    helper = Qwen3NextLinearAttPageHelper(mem)
+    mem.kv_move_buffer = torch.zeros((1, helper.state_nbytes), dtype=torch.uint8)
+    conv_page, ssm_page = helper.view_page_to_linear_att_state(page_index=0)
+
+    marker_conv = torch.arange(conv_page.numel(), dtype=torch.float32).view_as(conv_page)
+    marker_ssm = torch.arange(ssm_page.numel(), dtype=torch.float32).view_as(ssm_page)
+    conv_page.copy_(marker_conv)
+    ssm_page.copy_(marker_ssm)
+
+    helper.read_page_to_req(page_index=0, req_idx=req_idx, dp_mems=[mem])
+
+    narrow_w = helper.conv_shape[-1]
+    torch.testing.assert_close(mem.req_to_conv_state.buffer[:, req_idx, :, :narrow_w], marker_conv)
+    assert torch.all(mem.req_to_conv_state.buffer[:, req_idx, :, narrow_w:] == -9.0)
+    torch.testing.assert_close(mem.req_to_ssm_state.buffer[:, req_idx * (mtp_step + 1), ...], marker_ssm)