[minimax m3][npu]Adaptation of Minimax M3(w8a8) for NPU platforms [1/2] (#32941)

Signed-off-by: Devashish Lal <devcode@fb.com>
Signed-off-by: Alexandre Milesi <milesial@users.noreply.github.com>
Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com>
Signed-off-by: Ryan Stewart <rystewart@nvidia.com>
Co-authored-by: ClownBin <chaobin1993@126.com>
Co-authored-by: huangzhenyu <q_m_p@qq.com>
Co-authored-by: clown <17490516+ClownBin@users.noreply.github.com>
Co-authored-by: badmer <374330057@qq.com.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: YAMY <74099316+YAMY1234@users.noreply.github.com>
Co-authored-by: Shangming Cai <csmthu@gmail.com>
Co-authored-by: Mick <mickjagger19@icloud.com>
Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
Co-authored-by: Brayden Zhong <brayden.zhong@radixark.ai>
Co-authored-by: Jimmy Shong <jimmysh341@gmail.com>
Co-authored-by: Zijie Xia <zijie.xia@radixark.ai>
Co-authored-by: Thomas Wang <thomawan@amd.com>
Co-authored-by: siyu <liusy58@linux.alibaba.com>
Co-authored-by: Yuang Chen <cya539102@antgroup.com>
Co-authored-by: Yuang Chen <1131578721@qq.com>
Co-authored-by: 黄孝君 <dingfangsu23@gmail.com>
Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>
Co-authored-by: Xiaoyu Zhang <1182563586@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: TobyMint <130973409+TobyMint@users.noreply.github.com>
Co-authored-by: TobyMint <tobymint@users.noreply.github.com>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
Co-authored-by: Tan Trinh <84185999+tanth47@users.noreply.github.com>
Co-authored-by: Lifan Shen <draftbks@gmail.com>
Co-authored-by: Justin Tong <justintong0323@outlook.com>
Co-authored-by: Qiaolin Yu <liin1211@outlook.com>
Co-authored-by: AMD-yanfeiwang <yanfei.wang@amd.com>
Co-authored-by: QIN2DIM <62018067+QIN2DIM@users.noreply.github.com>
Co-authored-by: Zhiyao Jiang <jessicajiang324@gmail.com>
Co-authored-by: Brayden Zhong <brayden@radixark.ai>
Co-authored-by: DevashishLal-CB <devashish@rivosinc.com>
Co-authored-by: Devashish Lal <devcode@fb.com>
Co-authored-by: Alex Nails <alex.nails@radixark.ai>
Co-authored-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Michael Gschwind <mkgschwind+private@gmail.com>
Co-authored-by: weireweire <weiliangl@nvidia.com>
Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
Co-authored-by: Khoa Pham <khoa.pham@radixark.ai>
Co-authored-by: milesial <milesial@users.noreply.github.com>
Co-authored-by: elvischenv <219235043+elvischenv@users.noreply.github.com>
Co-authored-by: cctry <csycfl@gmail.com>
Co-authored-by: Jae B. <jlee5814@gmail.com>
Co-authored-by: Michael <13900043+michaelzhang-ai@users.noreply.github.com>
Co-authored-by: forrestl <16055533+forrestl111@users.noreply.github.com>
Co-authored-by: EchO <117733745+CyberSecurityErial@users.noreply.github.com>
Co-authored-by: Tanmay patil <tanmaypatil3151@gmail.com>
Co-authored-by: ybyang <10629930+whybeyoung@users.noreply.github.com>
Co-authored-by: Hsiu-Chun, Hung <160560375+Emmanuel0612@users.noreply.github.com>
Co-authored-by: Hung <Emmanuel0612@users.noreply.github.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>
Co-authored-by: Bingxu Chen <bingxche@amd.com>
Co-authored-by: YC Yen-Ching Tseng <yctseng@amd.com>
Co-authored-by: Cherry_ming <136634645@qq.com>
Co-authored-by: Even Zhou <even.y.zhou@outlook.com>
Co-authored-by: sglang-npu-bot <sglangnpu@163.com>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
Co-authored-by: Tingwei Huang <huangtingwei9988@gmail.com>
Co-authored-by: Kaixi <kaiximatteoc@nvidia.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com>
Co-authored-by: Ryan Stewart <rystewart@nvidia.com>
Co-authored-by: gjsheu <gjsheu@163.com>
Co-authored-by: Jinyan Yi <yjy20010615@gmail.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
Co-authored-by: huangtingwei <141888744+huangtingwei9988@users.noreply.github.com>
Co-authored-by: Hanming Lu <hanminglu@meta.com>
Co-authored-by: Jeremy Zhang <jeremy.zhang866@gmail.com>
Co-authored-by: Dmitrii Sergeev <dmi.sergeev@gmail.com>
Co-authored-by: Hao Zhang <zhisbug@users.noreply.github.com>
Co-authored-by: zhisbug <1654062+zhisbug@users.noreply.github.com>
Co-authored-by: Douglas Yang <dyang@college.harvard.edu>
Co-authored-by: gongwei1027 <gongwei833x@gmail.com>
Co-authored-by: ilyasher-harmonic <ilya.sherstyuk@harmonic.fun>
Co-authored-by: Hanming Lu <69857889+hanming-lu@users.noreply.github.com>
Co-authored-by: sglang-bot <sglangbot@gmail.com>
Co-authored-by: sglang-bot <232288953+sglang-bot@users.noreply.github.com>
Co-authored-by: Jimmy Shong <69131491+Jiminator@users.noreply.github.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Meng, Hengyu <hengyu.meng@intel.com>
Co-authored-by: Shu Wang <shuw@nvidia.com>
Co-authored-by: Yanbin Jiang <jybsuper@gmail.com>
Co-authored-by: zijiec <zijie.chen@amd.com>

test/manual/minimax_m3/test_npu_memory_pool.py

@@ -0,0 +1,176 @@
+import importlib.util
+import sys
+import types
+from contextlib import nullcontext
+from pathlib import Path
+
+import torch
+
+
+class _FakeMemorySaverAdapter:
+    def region(self, _memory_type):
+        return nullcontext()
+
+
+class _FakeKVCache:
+    def __init__(
+        self,
+        size,
+        page_size,
+        dtype,
+        layer_num,
+        device,
+        enable_memory_saver,
+        start_layer=None,
+        end_layer=None,
+    ):
+        self.size = size
+        self.page_size = page_size
+        self.dtype = dtype
+        self.store_dtype = dtype
+        self.layer_num = layer_num
+        self.device = device
+        self.enable_memory_saver = enable_memory_saver
+        self.start_layer = start_layer or 0
+        self.end_layer = end_layer or layer_num - 1
+        self.memory_saver_adapter = _FakeMemorySaverAdapter()
+        self.mem_usage = 0
+
+    def _finalize_allocation_log(self, _num_tokens):
+        pass
+
+
+class _FakeMHATokenToKVPool(_FakeKVCache):
+    def __init__(
+        self,
+        size,
+        page_size,
+        dtype,
+        head_num,
+        head_dim,
+        layer_num,
+        device,
+        enable_memory_saver,
+        v_head_dim=None,
+        swa_head_num=None,
+        swa_head_dim=None,
+        swa_v_head_dim=None,
+        start_layer=None,
+        end_layer=None,
+        **_kwargs,
+    ):
+        super().__init__(
+            size,
+            page_size,
+            dtype,
+            layer_num,
+            device,
+            enable_memory_saver,
+            start_layer,
+            end_layer,
+        )
+        self.head_num = swa_head_num if swa_head_num is not None else head_num
+        self.head_dim = swa_head_dim if swa_head_dim is not None else head_dim
+        self.v_head_dim = (
+            swa_v_head_dim
+            if swa_v_head_dim is not None
+            else v_head_dim if v_head_dim is not None else head_dim
+        )
+        self._create_buffers()
+
+
+class _FakeMHATokenToKOnlyPool(_FakeKVCache):
+    pass
+
+
+class _FakeMiniMaxSparseKVPool:
+    def __init__(self, *args, **kwargs):
+        pass
+
+
+class _FakeMLATokenToKVPool(_FakeKVCache):
+    pass
+
+
+def _load_npu_memory_pool_module():
+    for name in (
+        "sglang",
+        "sglang.srt",
+        "sglang.srt.constants",
+        "sglang.srt.mem_cache",
+        "sglang.srt.mem_cache.memory_pool",
+        "sglang.srt.utils",
+        "sglang.srt.utils.common",
+    ):
+        sys.modules.setdefault(name, types.ModuleType(name))
+
+    constants = sys.modules["sglang.srt.constants"]
+    constants.GPU_MEMORY_TYPE_KV_CACHE = "kv_cache"
+
+    memory_pool = sys.modules["sglang.srt.mem_cache.memory_pool"]
+    memory_pool.MHATokenToKVPool = _FakeMHATokenToKVPool
+    memory_pool.MHATokenToKOnlyPool = _FakeMHATokenToKOnlyPool
+    memory_pool.MiniMaxSparseKVPool = _FakeMiniMaxSparseKVPool
+    memory_pool.MLATokenToKVPool = _FakeMLATokenToKVPool
+    memory_pool.get_tensor_size_bytes = lambda tensor: tensor.nbytes
+    memory_pool.maybe_detect_oob = lambda *args, **kwargs: None
+    memory_pool.unwrap_write_loc = lambda loc_info: (loc_info, None, None)
+
+    utils = sys.modules["sglang.srt.utils"]
+    utils.get_bool_env_var = lambda _name, default: default == "True"
+
+    common = sys.modules["sglang.srt.utils.common"]
+    common.is_npu = lambda: False
+
+    module_path = (
+        Path(__file__).resolve().parents[3]
+        / "python/sglang/srt/hardware_backend/npu/memory_pool_npu.py"
+    )
+    spec = importlib.util.spec_from_file_location(
+        "_npu_memory_pool_under_test", module_path
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    spec.loader.exec_module(module)
+    return module
+
+
+def test_npu_minimax_k_only_index_cache_uses_scatter_writer():
+    npu_memory_pool = _load_npu_memory_pool_module()
+
+    calls = []
+
+    class FakeTorchNpu:
+        @staticmethod
+        def npu_scatter_nd_update_(cache, indices, updates):
+            assert cache.shape == (10, 1, 4)
+            assert indices.shape == (2, 1)
+            assert updates.shape == (2, 1, 4)
+            calls.append((cache, indices, updates))
+
+        @staticmethod
+        def _npu_reshape_and_cache(*, key, value, key_cache, value_cache, slot_indices):
+            raise AssertionError("K-only MiniMax index cache should use scatter")
+
+    npu_memory_pool.torch_npu = FakeTorchNpu
+
+    pool = npu_memory_pool.NPUMHATokenToKOnlyPool(
+        size=8,
+        page_size=2,
+        dtype=torch.bfloat16,
+        head_num=1,
+        head_dim=4,
+        layer_num=1,
+        device="cpu",
+        enable_memory_saver=False,
+    )
+
+    loc = torch.tensor([1, 3], dtype=torch.int64)
+    cache_k = torch.randn((2, 1, 4), dtype=torch.bfloat16)
+
+    pool.set_k_buffer(0, loc, cache_k)
+
+    assert len(calls) == 1
+    k_size, v_size = pool.get_kv_size_bytes()
+    assert k_size > 0
+    assert v_size == 0