[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>
This commit is contained in:
vstone-w
2026-08-13 11:23:14 +08:00
committed by GitHub
co-authored by ClownBin huangzhenyu clown badmer Liangsheng Yin YAMY Shangming Cai Mick Brayden Zhong Brayden Zhong Jimmy Shong Zijie Xia Thomas Wang siyu Yuang Chen Yuang Chen 黄孝君 Xinyuan Tong Xiaoyu Zhang Cursor Claude Fable 5 TobyMint TobyMint Cheng Wan Tan Trinh Lifan Shen Justin Tong Qiaolin Yu AMD-yanfeiwang QIN2DIM Zhiyao Jiang Brayden Zhong DevashishLal-CB Devashish Lal Alex Nails Mohammad Miadh Angkad Baizhou Zhang Michael Gschwind weireweire weireweire Khoa Pham milesial elvischenv cctry Jae B. Michael forrestl EchO Tanmay patil ybyang Hsiu-Chun, Hung Hung HaiShaw Bingxu Chen YC Yen-Ching Tseng Cherry_ming Even Zhou sglang-npu-bot Zhiqiang Xie Tingwei Huang Kaixi github-actions[bot] Faradawn Yang Ryan Stewart gjsheu Jinyan Yi Ke Bao huangtingwei Hanming Lu Jeremy Zhang Dmitrii Sergeev Hao Zhang zhisbug Douglas Yang gongwei1027 ilyasher-harmonic Hanming Lu sglang-bot sglang-bot Jimmy Shong hnyls2002 Meng, Hengyu Shu Wang Yanbin Jiang zijiec
parent d44c836cfd
commit bca8ed4afc
18 changed files with 2314 additions and 175 deletions
+19 -7
View File
@@ -4629,6 +4629,18 @@ class MHATokenToKOnlyPool(KVCache):
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
return self._get_key_buffer(layer_id)
def set_k_buffer(
self,
layer_id: int,
loc: torch.Tensor,
cache_k: torch.Tensor,
) -> None:
if cache_k.dtype != self.dtype:
cache_k = cache_k.to(self.dtype)
if self.store_dtype != self.dtype:
cache_k = cache_k.view(self.store_dtype)
self.k_buffer[layer_id][loc] = cache_k
def get_value_buffer(self, layer_id: int) -> torch.Tensor:
raise NotImplementedError("MHATokenToKOnlyPool does not allocate V")
@@ -4673,6 +4685,9 @@ class MiniMaxSparseKVPool(KVCache):
index_dtype: Optional[torch.dtype] = None,
start_layer: Optional[int] = None,
end_layer: Optional[int] = None,
main_pool_cls=MHATokenToKVPool,
index_kv_pool_cls=MHATokenToKVPool,
index_k_pool_cls=MHATokenToKOnlyPool,
):
# Do not call super().__init__() — delegate to sub-pools instead.
self.size = size
@@ -4714,7 +4729,7 @@ class MiniMaxSparseKVPool(KVCache):
gid: i for i, gid in enumerate(local_k_only_sparse_layer_ids)
}
self.main_pool = MHATokenToKVPool(
self.main_pool = main_pool_cls(
size=size,
page_size=page_size,
dtype=dtype,
@@ -4728,7 +4743,7 @@ class MiniMaxSparseKVPool(KVCache):
)
self.index_kv_pool: Optional[MHATokenToKVPool] = (
MHATokenToKVPool(
index_kv_pool_cls(
size=size,
page_size=page_size,
dtype=index_dtype,
@@ -4743,7 +4758,7 @@ class MiniMaxSparseKVPool(KVCache):
)
self.index_k_pool: Optional[MHATokenToKOnlyPool] = (
MHATokenToKOnlyPool(
index_k_pool_cls(
size=size,
page_size=page_size,
dtype=index_dtype,
@@ -4891,10 +4906,7 @@ class MiniMaxSparseKVPool(KVCache):
if cache_idx_k.dtype != sub_pool.dtype:
if k_scale is not None:
cache_idx_k = cache_idx_k / k_scale
cache_idx_k = cache_idx_k.to(sub_pool.dtype)
if sub_pool.store_dtype != sub_pool.dtype:
cache_idx_k = cache_idx_k.view(sub_pool.store_dtype)
sub_pool.k_buffer[mapped_id][loc] = cache_idx_k
sub_pool.set_k_buffer(mapped_id, loc, cache_idx_k)
def _can_fuse_kv_index_store(
self,