MemCache 介绍

MemCache是面向 AI 推理场景的分布式 KV Cache 缓存系统，针对昇腾超节点和服务器进行了性能优化，可作为 vLLM/SGLang 等推理框架的 KV Pool 后端，为长文本和高前缀复用场景提供跨请求 KV Cache 存储与复用能力。 在SGLang上，MemCache作为L3 KV Cache后端，用于存储和复用KV cache，后端名称是npu_memcache。

MemCache 安装

发布镜像自带MemCache，也可自行安装：

pip install memcache_hybrid

MemCache 部署

启动命令：

python -m sglang.srt.mem_cache.storage.npu_memcache.start_meta_service --config_path "metaservice_config.json"

metaservice_config.json

{
    "meta_service_url": "tcp://127.0.0.1:5000",
    "config_store_url": "tcp://127.0.0.1:6000",
    "metrics_url": "http://127.0.0.1:8000"
}

SGLang 开启L3 hicache

• 开启L1 cache需要去掉参数--disable-radix-cache
• 开启L2 cache需：
  • 在P节点增加参数：

--enable-hierarchical-cache
--hicache-io-backend kernel_ascend

• 开启L3 cache需：
  • 启动memcache：参考章节：MemCache部署
  • 在P节点增加参数：

--hicache-storage-backend npu_memcache
--hicache-storage-backend-extra-config '{"meta_service_url":"tcp://127.0.0.1:5000", "config_store_url":"tcp://127.0.0.1:6000", "log_level":"info", "world_size":256, "protocol": "device_sdma", "prefetch_threshold": 1, "dram_size": "50GB"}'

protocol：A3 上使用device_sdma，A5上使用device_urma

代码

kernel：https://github.com/sgl-project/sgl-kernel-npu/pull/713

sglang：https://github.com/sgl-project/sglang/pull/36188

相关资料

• Ascend MemCache Introduction
• Ascend MemCache Config Guide
• Ascend MemCache Python API
• SGLang HiCache Design
• Ascend MemFabric
• Ascend MemCache

测试方法

import random
import os
import time
import unittest
import requests
from sglang.test.kl_test_utils import (
    _flush_cache,
    _generate,
)
from sglang.test.test_utils import (
    CustomTestCase,
)

URL_FOR_TEST = "http://127.0.0.1:8869"

class TestNPUHiCache(CustomTestCase):
    """Exercise load-back on NPU."""

    page_size = 128
    prompt_len = 3200
    max_total_tokens = 12800
    routed_dp_rank = 0

    @classmethod
    def _clear_storage_backend ( cls ):
        requests.post( f" {cls.base_url}/hicache/storage-backend/clear" , timeout= 60 ).raise_for_status()
    
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.base_url = URL_FOR_TEST
        _flush_cache(cls.base_url)
        cls._clear_storage_backend()

    @classmethod
    def _prompt(cls, seed: int) -> list[int]:
        nonce = int (time.time()) & 0xFFFF
        rng = random.Random((seed * 1_000_003 ) ^ nonce)
        return [rng.randint( 1 , 30000 ) for _ in range (cls.prompt_len)]

    def _generate_one(self, input_ids, max_new_tokens, return_logprob=False):
        results = _generate(
            self.base_url,
            [input_ids],
            max_new_tokens=max_new_tokens,
            return_logprob=return_logprob,
            temperature=0,
            routed_dp_rank=self.routed_dp_rank,
        )
        self.assertEqual(len(results), 1)
        return results[0]

    def test_npu_restores_from_l3(self):
        """ L3 restore test """
        _flush_cache(self.base_url)
        
        base_prompt = self._prompt(5)
        first = self._generate_one(base_prompt, max_new_tokens=16)
        first_meta = first["meta_info"]
        first_details = first_meta.get("cached_tokens_details") or {}
        print(f"[L3][1] first_meta.cached_tokens_details={first_details}")
        self.assertEqual(int(first_meta.get("cached_tokens", 0)), 0)

        # The first cache hit reaches write_through's hit-count
        # threshold and starts the asynchronous L1 -> L2 -> L3 copy.
        warm = self._generate_one(base_prompt, max_new_tokens=16)
        warm_details = warm["meta_info"].get("cached_tokens_details") or {}
        print(f"[L3][2] warm_meta.cached_tokens_details={warm_details}")

        device_cached_tokens = int(warm_details.get("device", 0))
        host_cached_tokens = int(warm_details.get("host", 0))
        storage_cached_tokens = int(warm_details.get("storage", 0))
        self.assertGreater(
            device_cached_tokens,
            0,
            f"Expected an L1 device hit: {warm_details}",
        )
        self.assertEqual(
            host_cached_tokens,
            0,
            f"Expected no L2 contribution: {warm_details}",
        )
        self.assertEqual(
            storage_cached_tokens,
            0,
            f"Expected no L3 contribution: {warm_details}",
        )

        time.sleep(5)

        # Remove L1/L2 residency without clearing the file backend.
        _flush_cache(self.base_url)

        restored = self._generate_one(base_prompt, max_new_tokens=16)
        restored_meta = restored["meta_info"]
        restored_details = restored_meta.get("cached_tokens_details") or {}
        print(f"[L3][3] restored_meta.cached_tokens_details={restored_details}")

        device_cached_tokens = int(restored_details.get("device", 0))
        host_cached_tokens = int(restored_details.get("host", 0))
        storage_cached_tokens = int(restored_details.get("storage", 0))
        self.assertEqual(
            device_cached_tokens,
            0,
            f"Expected no L1 device contribution: {restored_details}",
        )
        self.assertEqual(
            host_cached_tokens,
            0,
            f"Expected no L2 contribution: {restored_details}",
        )
        self.assertGreater(
            storage_cached_tokens,
            0,
            f"Expected an L3 storage hit; got meta_info={restored_meta}",
        )
        self.assertEqual(
            restored["text"],
            first["text"],
            "L3 restore changed deterministic generation output.",
        )

        # Hit restored device again.
        restored = self._generate_one(base_prompt, max_new_tokens=16)
        restored_meta = restored["meta_info"]
        details = restored_meta.get("cached_tokens_details") or {}
        print(f"[L3][4] restored_meta.cached_tokens_details={details}")
        self.assertGreater(
            int(details.get("device", 0) or 0),
            0,
            f"Expected an L1 device hit; got meta_info={restored_meta}",
        )


if __name__ == "__main__":
    unittest.main()