fix

Author: sufubao

lightllm/models/qwen3next/mem_manager.py

@@ -104,6 +104,12 @@ def free_all(self):
         self.ssm_state_mem_manager.free_all()
         return
 
+    @override
+    def get_cell_size(self):
+        # Only full attention layers and MTP layers have KV cache
+        kv_cache_layer_num = self.full_attn_layer_num + self.mtp_layer_num
+        return 2 * self.head_num * self.head_dim * kv_cache_layer_num * torch._utils._element_size(self.dtype)
+
     @override
     def get_buffer(self, layer_index) -> Tuple[torch.Tensor, torch.Tensor]:
         assert layer_index < self.layer_num, "layer_index is out of range"