diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index ac474a024..f72b35c18 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1297,6 +1297,20 @@ class Envs: SGLANG_MINIMAX_M3_FUSED_SWIGLU_MXFP8 = EnvBool(False) SGLANG_MINIMAX_M3_FUSED_MOE_COMBINE = EnvBool(False) + # MiniMax M3 NPU prefill MAIN-attention: route the sparse main attention through + # the native Ascend FA op `torch.ops.npu.npu_fused_infer_attention_score` (FIA) + # with a per-query CUSTOM block_table + SGLANG_MINIMAX_NPU_PREFILL_FIA = EnvBool(True) + + # MiniMax-M3 NPU sparse INDEXER (decode + verify topk block selection): route + # through the native AscendC packed indexer op instead of the Triton indexer. + SGLANG_MINIMAX_NPU_NATIVE_INDEXER = EnvBool(False) + + # MiniMax-M3 NPU sparse MAIN-attention (decode-main + verify-main): route the + # sparse main attention through the native AscendC sparse-attention op with the + # cached block_table override. + SGLANG_MINIMAX_NPU_NATIVE_ATTN = EnvBool(False) + # MiniMax-M3 on ROCm force-disables custom all-reduce in its model override # (arg_groups/overrides.py) when aiter all-reduce fusion is off. Set this to # opt back in and keep custom/quick all-reduce enabled -- e.g. to run the diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index bfda4ebf5..b787b033e 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -5,7 +5,9 @@ import torch from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE from sglang.srt.environ import envs from sglang.srt.mem_cache.memory_pool import ( + MHATokenToKOnlyPool, MHATokenToKVPool, + MiniMaxSparseKVPool, MLATokenToKVPool, get_tensor_size_bytes, unwrap_write_loc, @@ -390,6 +392,135 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): torch.npu.synchronize() +class NPUMHATokenToKOnlyPool(MHATokenToKOnlyPool): + """NPU paged K-only cache used by MiniMax sparse index-only layers.""" + + def __init__( + self, + size: int, + page_size: int, + dtype: torch.dtype, + head_num: int, + head_dim: int, + layer_num: int, + device: str, + enable_memory_saver: bool, + start_layer: Optional[int] = None, + end_layer: Optional[int] = None, + ): + self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False") + super(MHATokenToKOnlyPool, self).__init__( + size=size, + page_size=page_size, + dtype=dtype, + layer_num=layer_num, + device=device, + enable_memory_saver=enable_memory_saver, + start_layer=start_layer, + end_layer=end_layer, + ) + self.head_num = head_num + self.head_dim = head_dim + + with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): + self.k_buffer = torch.zeros( + ( + self.layer_num, + self.size // self.page_size + 1, + self.page_size, + self.head_num, + self.head_dim, + ), + dtype=self.store_dtype, + device=self.device, + ) + if self.use_fia: + self.k_buffer = [ + self.k_buffer[i].view(-1, 1, self.head_num, self.head_dim) + for i in range(self.layer_num) + ] + + self._finalize_allocation_log(size) + + def _get_key_buffer(self, layer_id: int): + k_buffer = self.k_buffer[layer_id - self.start_layer] + if self.store_dtype != self.dtype: + return k_buffer.view(self.dtype) + return k_buffer + + def set_k_buffer( + self, + layer_id: int, + loc_info, + cache_k: torch.Tensor, + ) -> None: + loc, _, _ = unwrap_write_loc(loc_info) + 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) + + k_buffer_layer = self.k_buffer[layer_id - self.start_layer].view( + -1, self.head_num, self.head_dim + ) + loc = loc.to(device=cache_k.device, dtype=torch.int32).contiguous() + torch_npu.npu_scatter_nd_update_( + k_buffer_layer, + loc.view(-1, 1), + cache_k.contiguous().view(-1, self.head_num, self.head_dim), + ) + + def get_contiguous_buf_infos(self): + data_ptrs = [ + self.get_key_buffer(i).data_ptr() + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + data_lens = [ + self.get_key_buffer(i).nbytes + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + if self.use_fia: + item_lens = [ + self.get_key_buffer(i)[0].nbytes * self.page_size + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + else: + item_lens = [ + self.get_key_buffer(i)[0].nbytes + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + return data_ptrs, data_lens, item_lens + + def get_kv_size_bytes(self): + return get_tensor_size_bytes(self.k_buffer), 0 + + +class NPUMiniMaxSparseKVPool(MiniMaxSparseKVPool): + """MiniMax sparse wrapper backed by NPU paged MHA/index pools.""" + + def __init__(self, *args, **kwargs): + super().__init__( + *args, + main_pool_cls=NPUMHATokenToKVPool, + index_kv_pool_cls=NPUMHATokenToKVPool, + index_k_pool_cls=NPUMHATokenToKOnlyPool, + **kwargs, + ) + + def get_index_k_state_buf_infos(self): + pool = self.index_k_pool + n = pool.layer_num + data_ptrs = [pool.get_key_buffer(i).data_ptr() for i in range(n)] + data_lens = [pool.get_key_buffer(i).nbytes for i in range(n)] + if pool.use_fia: + item_lens = [ + pool.get_key_buffer(i)[0].nbytes * pool.page_size for i in range(n) + ] + else: + item_lens = [pool.get_key_buffer(i)[0].nbytes for i in range(n)] + return data_ptrs, data_lens, item_lens + + class NPUMLATokenToKVPool(MLATokenToKVPool): def __init__( diff --git a/python/sglang/srt/hardware_backend/npu/modules/minimax_m3_processor.py b/python/sglang/srt/hardware_backend/npu/modules/minimax_m3_processor.py new file mode 100644 index 000000000..68459980a --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/modules/minimax_m3_processor.py @@ -0,0 +1,301 @@ +"""NPU patch for MiniMax M3 VL image and video preprocessing. + +The MiniMax M3 VL image processor (MiniMaxM3VLImageProcessor) and video +processor (MiniMaxM3VLVideoProcessor) create 10-dimensional tensors during +patch extraction, which exceeds Ascend NPU's 8-dimension limit. + +This patch restructures the computation using transform_patches_to_flatten +to stay within 8 dimensions, following the same pattern as the Qwen VL and +GLM-4.6V NPU patches. +""" + +import math +from typing import List + +import torch +from torchvision.transforms import InterpolationMode +from transformers.image_processing_utils import BatchFeature +from transformers.image_processing_utils_fast import ( + group_images_by_shape, + reorder_images, +) +from transformers.image_utils import PILImageResampling, SizeDict +from transformers.utils import TensorType +from transformers.video_utils import group_videos_by_shape, reorder_videos + +from sglang.srt.hardware_backend.npu.modules.qwen_vl_processor import ( + transform_patches_to_flatten, +) + +MAX_RATIO = 200 + + +def _round_by_factor(number: int, factor: int) -> int: + return round(number / factor) * factor + + +def _ceil_by_factor(number: int, factor: int) -> int: + return math.ceil(number / factor) * factor + + +def _floor_by_factor(number: int, factor: int) -> int: + return math.floor(number / factor) * factor + + +def _smart_resize( + height: int, + width: int, + factor: int = 28, + min_pixels: int = 4 * 28 * 28, + max_pixels: int = 451584, +) -> tuple[int, int]: + if max(height, width) / min(height, width) > MAX_RATIO: + raise ValueError( + f"absolute aspect ratio must be smaller than {MAX_RATIO}, " + f"got {max(height, width) / min(height, width)}" + ) + h_bar = max(factor, _round_by_factor(height, factor)) + w_bar = max(factor, _round_by_factor(width, factor)) + if h_bar * w_bar > max_pixels: + beta = math.sqrt((height * width) / max_pixels) + h_bar = _floor_by_factor(height / beta, factor) + w_bar = _floor_by_factor(width / beta, factor) + elif h_bar * w_bar < min_pixels: + beta = math.sqrt(min_pixels / (height * width)) + h_bar = _ceil_by_factor(height * beta, factor) + w_bar = _ceil_by_factor(width * beta, factor) + return h_bar, w_bar + + +def npu_wrapper_minimax_m3_image_preprocess(func): + + def _preprocess( + self, + images: List[torch.Tensor], + do_resize: bool, + size: SizeDict, + resample: PILImageResampling | InterpolationMode | int | None, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | List[float] | None, + image_std: float | List[float] | None, + patch_size: int, + temporal_patch_size: int, + merge_size: int, + max_pixels: int, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + grouped_images, grouped_images_index = group_images_by_shape( + images, disable_grouping=disable_grouping + ) + resized_images_grouped = {} + factor = patch_size * merge_size + for shape, stacked_images in grouped_images.items(): + height, width = stacked_images.shape[-2:] + if do_resize: + resized_height, resized_width = _smart_resize( + height, + width, + factor=factor, + max_pixels=max_pixels, + ) + stacked_images = self.resize( + stacked_images, + size=SizeDict(height=resized_height, width=resized_width), + resample=resample, + ) + resized_images_grouped[shape] = stacked_images + + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + + grouped_images, grouped_images_index = group_images_by_shape( + resized_images, disable_grouping=disable_grouping + ) + processed_images_grouped = {} + processed_grids = {} + + for shape, stacked_images in grouped_images.items(): + resized_height, resized_width = stacked_images.shape[-2:] + + patches = self.rescale_and_normalize( + stacked_images, + do_rescale, + rescale_factor, + do_normalize, + image_mean, + image_std, + ) + if patches.ndim == 4: + patches = patches.unsqueeze(1) + + if patches.shape[1] % temporal_patch_size != 0: + repeats = patches[:, -1:].repeat( + 1, + temporal_patch_size - (patches.shape[1] % temporal_patch_size), + 1, + 1, + 1, + ) + patches = torch.cat([patches, repeats], dim=1) + + batch_size, grid_t, channel = patches.shape[:3] + grid_t = grid_t // temporal_patch_size + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size + + flatten_patches = transform_patches_to_flatten( + patches, + batch_size, + grid_t, + temporal_patch_size, + channel, + grid_h, + grid_w, + patch_size, + merge_size, + ) + + processed_images_grouped[shape] = flatten_patches + processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size + + processed_images = reorder_images( + processed_images_grouped, grouped_images_index + ) + processed_grids = reorder_images(processed_grids, grouped_images_index) + + pixel_values = torch.cat(processed_images, dim=0) + image_grid_thw = torch.tensor(processed_grids, dtype=torch.long) + + return BatchFeature( + data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, + tensor_type=return_tensors, + ) + + return _preprocess + + +def npu_wrapper_minimax_m3_video_preprocess(func): + + def _preprocess( + self, + videos: List[torch.Tensor], + do_convert_rgb: bool, + do_resize: bool, + size: SizeDict, + resample: PILImageResampling | InterpolationMode | int | None, + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | List[float] | None, + image_std: float | List[float] | None, + patch_size: int, + temporal_patch_size: int, + merge_size: int, + min_pixels: int, + max_pixels: int, + return_tensors: str | TensorType | None = None, + **kwargs, + ) -> BatchFeature: + grouped_videos, grouped_videos_index = group_videos_by_shape(videos) + resized_videos_grouped = {} + factor = patch_size * merge_size + for shape, stacked_videos in grouped_videos.items(): + batch_size, num_frames, channels, height, width = stacked_videos.shape + resized_height, resized_width = height, width + if do_resize: + resized_height, resized_width = _smart_resize( + height, + width, + factor=factor, + min_pixels=min_pixels, + max_pixels=max_pixels, + ) + stacked_videos = stacked_videos.view( + batch_size * num_frames, channels, height, width + ) + stacked_videos = self.resize( + stacked_videos, + size=SizeDict(height=resized_height, width=resized_width), + resample=resample, + ) + stacked_videos = stacked_videos.view( + batch_size, + num_frames, + channels, + resized_height, + resized_width, + ) + resized_videos_grouped[shape] = stacked_videos + resized_videos = reorder_videos(resized_videos_grouped, grouped_videos_index) + + grouped_videos, grouped_videos_index = group_videos_by_shape(resized_videos) + processed_videos_grouped = {} + processed_grids = {} + for shape, stacked_videos in grouped_videos.items(): + resized_height, resized_width = stacked_videos.shape[-2:] + patches = self.rescale_and_normalize( + stacked_videos, + do_rescale, + rescale_factor, + do_normalize, + image_mean, + image_std, + ) + + if pad := -patches.shape[1] % temporal_patch_size: + repeats = patches[:, -1:].expand(-1, pad, -1, -1, -1) + patches = torch.cat([patches, repeats], dim=1) + + batch_size, grid_t, channels = patches.shape[:3] + grid_t = grid_t // temporal_patch_size + grid_h, grid_w = resized_height // patch_size, resized_width // patch_size + + flatten_patches = transform_patches_to_flatten( + patches, + batch_size, + grid_t, + temporal_patch_size, + channels, + grid_h, + grid_w, + patch_size, + merge_size, + ) + + processed_videos_grouped[shape] = flatten_patches + processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size + + processed_videos = reorder_videos( + processed_videos_grouped, grouped_videos_index + ) + processed_grids = reorder_videos(processed_grids, grouped_videos_index) + pixel_values_videos = torch.cat(processed_videos, dim=0) + video_grid_thw = torch.tensor(processed_grids, dtype=torch.long) + + return BatchFeature( + data={ + "pixel_values_videos": pixel_values_videos, + "video_grid_thw": video_grid_thw, + }, + tensor_type=return_tensors, + ) + + return _preprocess + + +def npu_apply_minimax_m3_image_preprocess_patch(image_processor): + cls = type(image_processor) + if getattr(cls, "_sglang_npu_patched", False): + return + cls._preprocess = npu_wrapper_minimax_m3_image_preprocess(cls._preprocess) + cls._sglang_npu_patched = True + + +def npu_apply_minimax_m3_video_preprocess_patch(video_processor): + cls = type(video_processor) + if getattr(cls, "_sglang_npu_video_patched", False): + return + cls._preprocess = npu_wrapper_minimax_m3_video_preprocess(cls._preprocess) + cls._sglang_npu_video_patched = True diff --git a/python/sglang/srt/hardware_backend/npu/moe/activation.py b/python/sglang/srt/hardware_backend/npu/moe/activation.py index bfad1c65d..3ec8819de 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/activation.py +++ b/python/sglang/srt/hardware_backend/npu/moe/activation.py @@ -22,7 +22,7 @@ class BaseActivation(ABC): # ============================================================================= -# Concrete activation implementations (unchanged except removed 8.) +# Concrete activation implementations # ============================================================================= class NPUSwiglu(BaseActivation): def _apply_activation(self, hidden_states: torch.Tensor): @@ -65,11 +65,31 @@ class NPUSwigluQuantWithScales(BaseActivation): class NPUSwigluDeepEPKernel(BaseActivation): - def __init__(self, need_quant: bool = True): - from sgl_kernel_npu.activation.swiglu_quant import swiglu_quant + """DeepEP grouped SwiGLU for the Ascend MoE runner; picks ``swiglu_quant`` vs the MiniMax + SwiGLU-OAI variant (``swiglu_oai_quant``: ``gate*sigmoid(gate*alpha)*(up+1)`` w/ clamping) + based on whether ``alpha``/``limit`` are given. The runner must forward + ``gemm1_alpha``/``gemm1_clamp_limit`` here or experts fall back to wrong SwiGLU.""" - self._kernel = swiglu_quant + def __init__( + self, + need_quant: bool = True, + alpha: Optional[float] = None, + limit: Optional[float] = None, + ): self.need_quant = need_quant + self.alpha = alpha + self.limit = limit + self._use_oai = alpha is not None and limit is not None + if self._use_oai: + from sgl_kernel_npu.activation.swiglu_oai_quant import ( + swiglu_oai_quant, + ) + + self._kernel = swiglu_oai_quant + else: + from sgl_kernel_npu.activation.swiglu_quant import swiglu_quant + + self._kernel = swiglu_quant def _apply_activation( self, @@ -77,9 +97,19 @@ class NPUSwigluDeepEPKernel(BaseActivation): group_list: torch.Tensor, group_list_type: int, ): - hidden_states, per_token_scale = self._kernel( - hidden_states, group_list, group_list_type, need_quant=self.need_quant - ) + if self._use_oai: + hidden_states, per_token_scale = self._kernel( + hidden_states, + self.alpha, + self.limit, + need_quant=self.need_quant, + group_list=group_list, + group_list_type=group_list_type, + ) + else: + hidden_states, per_token_scale = self._kernel( + hidden_states, group_list, group_list_type, need_quant=self.need_quant + ) if self.need_quant: return hidden_states, per_token_scale return hidden_states, None diff --git a/python/sglang/srt/hardware_backend/npu/moe/topk.py b/python/sglang/srt/hardware_backend/npu/moe/topk.py index aa4d85299..75085959b 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/topk.py +++ b/python/sglang/srt/hardware_backend/npu/moe/topk.py @@ -99,7 +99,7 @@ def fused_topk_npu( k_group=topk_config.topk_group if use_grouped_topk else 1, group_count=topk_config.num_expert_group if use_grouped_topk else 1, group_select_mode=(1 if use_grouped_topk else 0), - renorm=0, + renorm=renormalize, # 1 for sigmoid, 0 for softmax norm_type=(0 if topk_config.scoring_func == "softmax" else 1), routed_scaling_factor=( diff --git a/python/sglang/srt/layers/attention/minimax_sparse_backend.py b/python/sglang/srt/layers/attention/minimax_sparse_backend.py index c09e7da12..f68ad4b56 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_backend.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_backend.py @@ -1,7 +1,9 @@ from __future__ import annotations import logging -from typing import TYPE_CHECKING, Optional +import os +from types import SimpleNamespace +from typing import TYPE_CHECKING, Optional, Tuple import torch @@ -11,14 +13,34 @@ from sglang.srt.configs.model_config import ( get_minimax_sparse_layer_ids, get_minimax_sparse_score_type, ) +from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.attention.minimax_sparse_ops.minimax_sparse import ( - minimax_sparse_decode, - minimax_sparse_prefill, -) from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.server_args import m3_fp8_attn_gemm_enabled +from sglang.srt.utils import is_npu + +if is_npu(): + from sglang.kernels.ops.attention.minimax_sparse.common.index import ( + topk_index_reduce, + ) + +# Adaptive block_size_q thresholds (cut K-cache traffic; affects only the serial loop). +_BSQ_THRESHOLD_64 = 4096 # max_seqlen_k >= 4K -> BSQ=64 +_BSQ_THRESHOLD_32 = 1024 # max_seqlen_k >= 1K -> BSQ=32 +_BSQ_THRESHOLD_16 = 512 # max_seqlen_k >= 512 -> BSQ=16 +# BSQ<=64 is UB-safe for the prefill indexer (Q tile up to 8KB at BSQ=64). + + +def _native_indexer_enabled() -> bool: + # Native AscendC packed indexer switch (default off). + return envs.SGLANG_MINIMAX_NPU_NATIVE_INDEXER.get() + + +def _native_attn_enabled() -> bool: + # Native AscendC sparse MAIN-attention switch (default off). + return envs.SGLANG_MINIMAX_NPU_NATIVE_ATTN.get() + if TYPE_CHECKING: from sglang.srt.model_executor.model_runner import ModelRunner @@ -26,6 +48,52 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def _kv_cache_to_bnsd( + k_cache: torch.Tensor, v_cache: torch.Tensor, page_size: int +) -> Tuple[torch.Tensor, torch.Tensor, int, int, int]: + """Reshape NHD slot-major KV caches to BNSD [pages, page_size, heads, dim]. + + Already-paged 4D inputs pass through unchanged. + """ + if k_cache.dim() == 4: + num_pages, _, num_kv_heads, head_dim = k_cache.shape + return k_cache, v_cache, num_pages, num_kv_heads, head_dim + num_pages = k_cache.shape[0] // page_size + num_kv_heads = k_cache.shape[1] + head_dim = k_cache.shape[2] + return ( + k_cache.view(num_pages, page_size, num_kv_heads, head_dim), + v_cache.view(num_pages, page_size, num_kv_heads, head_dim), + num_pages, + num_kv_heads, + head_dim, + ) + + +def _idx_cache_to_bnsd( + idx_k_cache: torch.Tensor, + idx_v_cache: Optional[torch.Tensor], + page_size: int, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], int, int]: + """Reshape NHD slot-major index caches to BNSD; already-paged 4D passes through.""" + if idx_k_cache.dim() == 4: + return idx_k_cache, idx_v_cache, idx_k_cache.shape[2], idx_k_cache.shape[3] + num_pages = idx_k_cache.shape[0] // page_size + idx_kv_heads = idx_k_cache.shape[1] + idx_dim = idx_k_cache.shape[2] + idx_v_bnsd = ( + None + if idx_v_cache is None + else idx_v_cache.view(num_pages, page_size, idx_kv_heads, idx_dim) + ) + return ( + idx_k_cache.view(num_pages, page_size, idx_kv_heads, idx_dim), + idx_v_bnsd, + idx_kv_heads, + idx_dim, + ) + + def _quant_q_fp8(q: torch.Tensor, q_scale: Optional[float]) -> torch.Tensor: # Same convention as the KV pools: the fp8 tensor stores value/scale and # the attention kernels multiply the logits back by the scale (None = unit). @@ -37,9 +105,14 @@ def _quant_q_fp8(q: torch.Tensor, q_scale: Optional[float]) -> torch.Tensor: class MiniMaxSparseAttnBackend(AttentionBackend): def __init__(self, runner: ModelRunner): assert isinstance(runner.token_to_kv_pool, MiniMaxSparseKVPool) + self.is_npu = is_npu() self.kv_pool = runner.token_to_kv_pool + self.token_to_kv_pool = runner.token_to_kv_pool # alias for TboAttnBackend + self.req_to_token_pool = runner.req_to_token_pool # pool obj for TboAttnBackend self.req_to_token = runner.req_to_token_pool.req_to_token self.max_context_len = int(runner.model_config.context_len) + # Per-forward cache for the native decode block table (rebuilt each forward). + self._native_decode_bt: dict = {} self.fp8_attn_gemm = m3_fp8_attn_gemm_enabled(runner.server_args) if self.fp8_attn_gemm: assert self.kv_pool.main_pool.dtype == torch.float8_e4m3fn, ( @@ -62,6 +135,13 @@ class MiniMaxSparseAttnBackend(AttentionBackend): self._max_seqlen_q: int = 1 self._max_seqlen_k: int = 1 + # NPU: per-forward cached metadata for the triton paths (rebuilt each forward). + self._prefill_meta: Optional[SimpleNamespace] = None + self._extend_meta: Optional[SimpleNamespace] = None + self._extend_meta_key: Optional[int] = None + self._decode_seq_lens_i32_cg: dict[int, torch.Tensor] = {} + self._verify_meta_cg: dict[tuple, SimpleNamespace] = {} + self.block_size_q = 1 self.block_size_k = sparse_cfg["sparse_block_size"] if "sparse_init_block" in sparse_cfg: @@ -82,47 +162,58 @@ class MiniMaxSparseAttnBackend(AttentionBackend): # MSA (fmha_sm100) is SM100-only; fall back to the Triton sparse path when # the kernel is unavailable or its constraints don't hold. - from sglang.srt.environ import envs - from sglang.srt.layers.attention.minimax_sparse_ops.msa import ( - msa_available, - ) - - # MSA (fmha_sm100) runs bf16, or uniform fp8_e4m3 under fp8 attn-GEMM mode - # (which also casts q to fp8). An fp8 main KV cache WITHOUT the flag - # would pair a bf16 q with fp8 K/V — unsupported by fmha_sm100's - # uniform-dtype kernels — so it stays on the Triton sparse path (which - # dequants fp8 on load). e5m2 is never allowed into MSA (fmha_sm100's - # variant lookup would silently dispatch the e4m3 kernel). - _main_kv_is_fp8 = self.kv_pool.main_pool.dtype in ( - torch.float8_e4m3fn, - torch.float8_e5m2, - ) - _msa_fp8_ok = ( - self.fp8_attn_gemm and self.kv_pool.main_pool.dtype == torch.float8_e4m3fn - ) - self.use_msa = ( - not envs.SGLANG_DISABLE_MSA.get() - and msa_available() - and self.block_size_k == 128 - and self.kv_pool.page_size == self.block_size_k - and self.topk_blocks in (4, 8, 16, 32) - and (not _main_kv_is_fp8 or _msa_fp8_ok) - ) - if ( - not self.use_msa - and not envs.SGLANG_DISABLE_MSA.get() - and msa_available() - and self.block_size_k == 128 - and self.kv_pool.page_size != self.block_size_k - ): - logger.warning( - "MiniMax-M3 MSA decode disabled: page_size=%d != sparse block size " - "%d. Pass --page-size 128 (with an attention backend that allows it, " - "e.g. fa4 or trtllm_mha) to enable the faster MSA kernel; falling " - "back to the Triton sparse path.", - self.kv_pool.page_size, - self.block_size_k, + if self.is_npu: + self.use_msa = False + # Prime the native sparse op probe before cuda-graph capture. + from sgl_kernel_npu.attention.gqa_share_sparse_attention import ( + _get_native_sparse_op, ) + + self._native_sparse_ok = _get_native_sparse_op() is not None + else: + self._native_sparse_ok = False + from sglang.srt.layers.attention.minimax_sparse_ops.msa import ( + msa_available, + ) + + # MSA (fmha_sm100) runs bf16, or uniform fp8_e4m3 under fp8 attn-GEMM mode + # (which also casts q to fp8). An fp8 main KV cache WITHOUT the flag + # would pair a bf16 q with fp8 K/V — unsupported by fmha_sm100's + # uniform-dtype kernels — so it stays on the Triton sparse path (which + # dequants fp8 on load). e5m2 is never allowed into MSA (fmha_sm100's + # variant lookup would silently dispatch the e4m3 kernel). + _main_kv_is_fp8 = self.kv_pool.main_pool.dtype in ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + _msa_fp8_ok = ( + self.fp8_attn_gemm + and self.kv_pool.main_pool.dtype == torch.float8_e4m3fn + ) + self.use_msa = ( + not envs.SGLANG_DISABLE_MSA.get() + and msa_available() + and self.block_size_k == 128 + and self.kv_pool.page_size == self.block_size_k + and self.topk_blocks in (4, 8, 16, 32) + and (not _main_kv_is_fp8 or _msa_fp8_ok) + ) + if ( + not self.use_msa + and not envs.SGLANG_DISABLE_MSA.get() + and msa_available() + and self.block_size_k == 128 + and self.kv_pool.page_size != self.block_size_k + ): + logger.warning( + "MiniMax-M3 MSA decode disabled: page_size=%d != sparse block size " + "%d. Pass --page-size 128 (with an attention backend that allows it, " + "e.g. fa4 or trtllm_mha) to enable the faster MSA kernel; falling " + "back to the Triton sparse path.", + self.kv_pool.page_size, + self.block_size_k, + ) + self._msa_dec_meta = None if self.use_msa: from sglang.srt.runtime_context import get_parallel @@ -138,7 +229,8 @@ class MiniMaxSparseAttnBackend(AttentionBackend): self.page_size = self.kv_pool.page_size self.use_dense_sparse_decode = ( - envs.SGLANG_OPT_USE_MINIMAX_DENSE_SPARSE_DECODE.get() + (not self.is_npu) + and envs.SGLANG_OPT_USE_MINIMAX_DENSE_SPARSE_DECODE.get() and self.block_size_k % self.page_size == 0 # _dense_sparse_main_decode calls trtllm decode with a bf16 q and # unit bmm scales — no fp8 handling yet (follow-up). @@ -151,6 +243,9 @@ class MiniMaxSparseAttnBackend(AttentionBackend): ) _sa = getattr(runner, "server_args", None) + self.speculative_num_draft_tokens = getattr( + _sa, "speculative_num_draft_tokens", None + ) _decode_cuda_graph = not check_cuda_graph_backend( Phase.DECODE, Backend.DISABLED ) @@ -183,6 +278,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend): f"msa_owns_decode={self._msa_owns_decode}, " f"decode_cuda_graph={_decode_cuda_graph}, " f"fp8_attn_gemm={self.fp8_attn_gemm}, " + f"npu_native_attn={'on' if (self._native_sparse_ok and _native_attn_enabled()) else 'off'}, " f"disable_value_layers={sorted(self.disable_value_layer_ids)})" ) if self.fp8_attn_gemm and self.use_msa: @@ -192,18 +288,58 @@ class MiniMaxSparseAttnBackend(AttentionBackend): "take minutes; compiles serialize across TP ranks)." ) + @staticmethod + def _choose_decode_score_max_chunks(batch_size: int) -> int: + """Score chunk count per graph bucket. + + bs=1 uses 16 chunks; larger buckets keep 32. Verify has its own tuning. + """ + return 16 if int(batch_size) == 1 else 32 + + @staticmethod + def _choose_block_size_q(max_seqlen_k: int) -> int: + """Pick block_size_q from max KV length (MINIMAX_NPU_PREFILL_BSQ overrides).""" + _forced = os.environ.get("MINIMAX_NPU_PREFILL_BSQ") + if _forced: + try: + _v = int(_forced) + if _v > 0: + return _v + except ValueError: + pass + if max_seqlen_k >= _BSQ_THRESHOLD_64: + return 64 + if max_seqlen_k >= _BSQ_THRESHOLD_32: + return 32 + if max_seqlen_k >= _BSQ_THRESHOLD_16: + return 16 + return 1 + + # ------------------------------------------------------------------ + # Delegation helpers + # ------------------------------------------------------------------ + def init_forward_metadata_out_graph( self, forward_batch: ForwardBatch, in_capture: bool = False ): - # cuda-graph replay views are a SimpleNamespace without extend_seq_lens_cpu, - # and TARGET_VERIFY sets it to None despite is_extend() — getattr covers both. + # getattr covers replay views lacking extend_seq_lens_cpu and TARGET_VERIFY. self._msa_dec_meta = None + if self.is_npu: + # Invalidate cached prefill/extend metadata; rebuilt on first sparse layer. + self._prefill_meta = None + self._extend_meta = None + self._extend_meta_key = None extend_lens = getattr(forward_batch, "extend_seq_lens_cpu", None) if extend_lens is not None: self._max_seqlen_q = int(max(extend_lens)) else: self._max_seqlen_q = 1 - if in_capture and forward_batch.forward_mode.is_decode_or_idle(): + if in_capture and ( + forward_batch.forward_mode.is_decode_or_idle() + or (self.is_npu and forward_batch.forward_mode.is_target_verify()) + ): + # Capture uses tiny dummy seq_lens; bound by full context so replay + # (longer sequences) does not miss KV blocks. self._max_seqlen_k = self.max_context_len else: self._max_seqlen_k = int(forward_batch.seq_lens_cpu.max().item()) @@ -213,7 +349,48 @@ class MiniMaxSparseAttnBackend(AttentionBackend): if self._msa_owns_decode and forward_batch.forward_mode.is_decode_or_idle(): self._prepare_msa_decode_meta(forward_batch) + # ---- REPLAY-FRESH native verify block_table ---- + if ( + self.is_npu + and forward_batch.forward_mode.is_target_verify() + and self.speculative_num_draft_tokens + ): + _ndt = self.speculative_num_draft_tokens + _bs = forward_batch.seq_lens.shape[0] + _key = (_bs, int(_ndt)) + _vmeta = self._verify_meta_cg.get(_key) + if _vmeta is not None: + _vmeta.per_query_req.copy_( + forward_batch.req_pool_indices.long().repeat_interleave(int(_ndt)) + ) + _prefix = (forward_batch.seq_lens.to(torch.long) - int(_ndt)).clamp( + min=0 + ) + _offs = torch.arange( + 1, + int(_ndt) + 1, + device=forward_batch.seq_lens.device, + dtype=torch.long, + ) + _vmeta.per_query_seq_lens.copy_( + (_prefix.unsqueeze(1) + _offs.unsqueeze(0)) + .reshape(-1) + .to(torch.int32) + ) + _mb = self.req_to_token.shape[1] // self.page_size + _bt_cols = ( + torch.arange( + _mb, device=_vmeta.per_query_req.device, dtype=torch.long + ) + * self.page_size + ).clamp(max=self.req_to_token.shape[1] - 1) + _vmeta.native_bt = ( + self.req_to_token[_vmeta.per_query_req][:, _bt_cols] + // self.page_size + ).to(torch.int32) + def _prepare_msa_decode_meta(self, forward_batch: ForwardBatch): + """Refresh the persistent per-batch-size MSA decode plan + page table in place.""" from sglang.srt.layers.attention.minimax_sparse_ops.msa import ( build_msa_decode_cg_plan, update_msa_decode_cg_meta, @@ -254,7 +431,48 @@ class MiniMaxSparseAttnBackend(AttentionBackend): self._msa_dec_meta = (kv_indices_buf, plan) def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): - pass + if not self.is_npu: + return + # Layer-invariant decode/verify metadata as captured ops (re-read at replay). + fm = forward_batch.forward_mode + if fm.is_target_verify(): + ndt = self.speculative_num_draft_tokens + if ndt: + prefix = (forward_batch.seq_lens.to(torch.long) - int(ndt)).clamp(min=0) + offsets = torch.arange( + 1, + int(ndt) + 1, + device=forward_batch.seq_lens.device, + dtype=torch.long, + ) + per_query_seq_lens = ( + (prefix.unsqueeze(1) + offsets.unsqueeze(0)) + .reshape(-1) + .to(torch.int32) + ) + per_query_req = forward_batch.req_pool_indices.long().repeat_interleave( + int(ndt) + ) + # Captured block_table for the native verify op (re-runs at replay). + _mb = self.req_to_token.shape[1] // self.page_size + _bt_cols = ( + torch.arange(_mb, device=per_query_req.device, dtype=torch.long) + * self.page_size + ).clamp(max=self.req_to_token.shape[1] - 1) + _native_bt = ( + self.req_to_token[per_query_req][:, _bt_cols] // self.page_size + ).to(torch.int32) + self._verify_meta_cg[(forward_batch.seq_lens.shape[0], int(ndt))] = ( + SimpleNamespace( + per_query_seq_lens=per_query_seq_lens, + per_query_req=per_query_req, + native_bt=_native_bt, + ) + ) + elif fm.is_decode_or_idle(): + self._decode_seq_lens_i32_cg[forward_batch.seq_lens.shape[0]] = ( + forward_batch.seq_lens.to(torch.int32) + ) def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): pass @@ -262,6 +480,749 @@ class MiniMaxSparseAttnBackend(AttentionBackend): def get_cuda_graph_seq_len_fill_value(self): return 1 + def _merge_sparse_blocks( + self, + topk_blocks: torch.Tensor, + query_positions: torch.Tensor, + num_blocks: int, + ) -> torch.Tensor: + """Append forced init/local blocks to top-k block ids and deduplicate.""" + total = self.topk_blocks + self.init_blocks + self.local_blocks + if self.init_blocks <= 0 and self.local_blocks <= 0: + return topk_blocks + + block_size = self.block_size_k + q_len = query_positions.shape[0] + num_idx_heads = topk_blocks.shape[1] + qcol = query_positions[:, None, None] + + if self.init_blocks == 0 and self.local_blocks == 1: + local = (query_positions // block_size).clamp( + min=0, max=max(num_blocks - 1, 0) + ) + local = ( + local.to(topk_blocks.dtype) + .view(q_len, 1, 1) + .expand(-1, num_idx_heads, -1) + ) + valid_topk = (topk_blocks >= 0) & (topk_blocks < num_blocks) + valid_topk = valid_topk & (topk_blocks * block_size <= qcol) + local_duplicate = ((topk_blocks == local) & valid_topk).any( + dim=-1, keepdim=True + ) + valid_local = (local >= 0) & (local < num_blocks) + valid_local = valid_local & (local * block_size <= qcol) & ~local_duplicate + return torch.cat( + [ + torch.where( + valid_topk, topk_blocks, torch.full_like(topk_blocks, -1) + ), + torch.where(valid_local, local, torch.full_like(local, -1)), + ], + dim=-1, + ) + + forced_parts = [] + if self.init_blocks > 0: + forced_parts.append( + torch.arange( + self.init_blocks, + device=topk_blocks.device, + dtype=topk_blocks.dtype, + ) + .view(1, 1, -1) + .expand(q_len, num_idx_heads, -1) + ) + if self.local_blocks > 0: + offsets = torch.arange( + self.local_blocks, + device=topk_blocks.device, + dtype=query_positions.dtype, + ) + block_ids = query_positions // block_size + first = (block_ids - self.local_blocks + 1).clamp(min=0) + forced_parts.append( + (first[:, None] + offsets[None, :]) + .to(topk_blocks.dtype) + .view(q_len, 1, -1) + .expand(-1, num_idx_heads, -1) + ) + + forced = torch.cat(forced_parts, dim=-1) + candidates = torch.cat([forced, topk_blocks], dim=-1) + valid = (candidates >= 0) & (candidates < num_blocks) + valid = valid & (candidates * block_size <= qcol) + invalid_value = torch.full_like(candidates, num_blocks) + sorted_candidates = torch.sort( + torch.where(valid, candidates, invalid_value), dim=-1 + ).values + sorted_valid = sorted_candidates < num_blocks + previous = torch.cat( + [ + torch.full_like(sorted_candidates[..., :1], -1), + sorted_candidates[:, :, :-1], + ], + dim=-1, + ) + keep = sorted_valid & (sorted_candidates != previous) + ranks = torch.cumsum(keep.to(torch.int32), dim=-1) - 1 + output = torch.full( + (q_len, num_idx_heads, total + 1), + -1, + dtype=topk_blocks.dtype, + device=topk_blocks.device, + ) + overflow_rank = torch.full_like(ranks, total) + scatter_index = torch.where(keep & (ranks < total), ranks, overflow_rank).long() + scatter_src = torch.where(keep, sorted_candidates, -1) + output.scatter_(2, scatter_index, scatter_src) + return output[:, :, :total] + + def _prepare_npu_triton_topk_idx( + self, + topk_idx: torch.Tensor, + seq_lens: torch.Tensor, + num_idx_heads: int, + num_kv_heads: int, + max_blocks: int, + ) -> torch.Tensor: + """Prepare NPU triton top-k ids in the GQA kernel layout. MiniMax-M3 (TP=16) emits it directly, skipping transpose+append+dedup.""" + if ( + self.init_blocks == 0 + and self.local_blocks == 1 + and num_idx_heads == num_kv_heads + and topk_idx.shape[0] == num_kv_heads + and topk_idx.dtype == torch.int32 + and topk_idx.is_contiguous() + and seq_lens.is_contiguous() + ): + # Fused prefill topk already appended the causal local block ([..., topk+1]); decode/verify still need the append. + if topk_idx.shape[2] == self.topk_blocks + 1: + return topk_idx + from sgl_kernel_npu.indexer.flash_block_score_decode import ( + append_local_block_to_topk_idx, + ) + + return append_local_block_to_topk_idx( + topk_idx, seq_lens, self.block_size_k, max_blocks + ) + + if num_idx_heads > num_kv_heads: + idx_group_size = num_idx_heads // num_kv_heads + topk_idx = topk_index_reduce( + topk_idx.view(num_kv_heads, idx_group_size, -1, self.topk_blocks), + dim=1, + ) + + topk_2d = topk_idx.permute(1, 0, 2).contiguous() + query_positions = (seq_lens.to(torch.long) - 1).clamp(min=0) + topk_merged = self._merge_sparse_blocks(topk_2d, query_positions, max_blocks) + return topk_merged.permute(1, 0, 2).contiguous() + + def _build_native_block_table( + self, req_indices: torch.Tensor, max_blocks: int, device + ) -> torch.Tensor: + """Logical->physical page table for the native sparse main op.""" + blk_cols = ( + torch.arange(max_blocks, device=device, dtype=torch.long) * self.page_size + ).clamp(max=self.req_to_token.shape[1] - 1) + return (self.req_to_token[req_indices][:, blk_cols] // self.page_size).to( + torch.int32 + ) + + def _forward_npu_triton_decode( + self, + q: torch.Tensor, # [B, num_q_heads, head_dim] + k_cache: torch.Tensor, # [num_slots, num_kv_heads, head_dim] (NHD) + v_cache: torch.Tensor, # [num_slots, num_kv_heads, head_dim] + idx_q: torch.Tensor, # [B, num_idx_heads, idx_dim] + idx_k_cache: torch.Tensor, # [num_slots, idx_kv_heads, idx_dim] + idx_v_cache: Optional[ + torch.Tensor + ], # [num_slots, idx_kv_heads, idx_dim] or None + forward_batch: ForwardBatch, + ): + """NPU decode via the ported triton kernels (BNSD paged). + NHD paged KV reshapes to [pages, block_size, H, D]; block table from + req_to_token. + """ + from sgl_kernel_npu.attention.gqa_share_sparse_attention import ( + flash_decode_bnsd_with_gqa_share_sparse, + ) + from sgl_kernel_npu.indexer.flash_block_score_decode import ( + flash_decode_bnsd_with_topk_idx, + ) + + page_size = self.page_size # == block_size_k + num_q_heads = q.shape[1] + head_dim = q.shape[2] + num_idx_heads = idx_q.shape[1] + idx_dim = idx_q.shape[2] + + k_bnsd, v_bnsd, num_pages, num_kv_heads, head_dim = _kv_cache_to_bnsd( + k_cache, v_cache, page_size + ) + idx_k_bnsd, idx_v_bnsd, idx_kv_heads, idx_dim = _idx_cache_to_bnsd( + idx_k_cache, idx_v_cache, page_size + ) + + # int32 seq_lens is layer-invariant: read the per-bs buffer built once + # per forward (captured op), with an inline eager fallback. + bs = forward_batch.seq_lens.shape[0] + seq_lens = self._decode_seq_lens_i32_cg.get(bs) + if seq_lens is None: + seq_lens = forward_batch.seq_lens.to(torch.int32) + max_seqlen = ( + int(self._max_seqlen_k) + if self._max_seqlen_k + else int(seq_lens.max().item()) + ) + max_blocks = (max_seqlen + page_size - 1) // page_size + disable_index_value = idx_v_cache is None + # Native main op takes a logical->physical block table, hoisted to + # once per forward (cached by id(forward_batch)); triton falls back to req_to_token. + _native_main_kwargs = None + if self._native_sparse_ok and _native_attn_enabled(): + try: + _fb_id = id(forward_batch) + _bt = self._native_decode_bt.get(_fb_id) + if _bt is None or _bt.shape[0] != q.shape[0]: + _bt = self._build_native_block_table( + forward_batch.req_pool_indices.long(), max_blocks, q.device + ) + self._native_decode_bt = {_fb_id: _bt} # single-entry: drop stale + _native_main_kwargs = {"block_table": _bt} + except Exception: + _native_main_kwargs = None + if disable_index_value: + page_source_kwargs = dict( + block_table=None, + req_to_token=self.req_to_token, + req_pool_indices=forward_batch.req_pool_indices, + max_num_blocks=max_blocks, + num_pages=num_pages, + sanitize_page_ids=False, + ) + else: + # Legacy score+index-value contract for non-MiniMax-M3 sparse layouts. + req_idx = forward_batch.req_pool_indices.long() + max_cols = self.req_to_token.shape[1] + blk_cols = ( + torch.arange(max_blocks, device=q.device, dtype=torch.long) * page_size + ).clamp(max=max_cols - 1) + token_slots = self.req_to_token[req_idx][:, blk_cols] + page_source_kwargs = dict( + block_table=(token_slots // page_size).to(torch.int32) + ) + + # 1) indexer: score idx_k + index attention + topk (init/local=0; + # forced blocks are re-appended by _prepare_npu_triton_topk_idx). + idx_o, topk_idx = flash_decode_bnsd_with_topk_idx( + q=idx_q, + sink=None, + k_cache_bnsd=idx_k_bnsd, + v_cache_bnsd=idx_v_bnsd, + **page_source_kwargs, + seq_lens=seq_lens, + max_seqlen=max_seqlen, + block_size=page_size, + topk=self.topk_blocks, + init_blocks=0, + local_blocks=0, + sm_scale=idx_dim**-0.5, + score_type=self.score_type, + disable_index_value=disable_index_value, + runtime_fill_only=True, + score_max_chunks=self._choose_decode_score_max_chunks(bs), + fused_append_local=True, + use_native=_native_indexer_enabled(), + ) + + # 2) Reduce heads and append forced blocks. + topk_idx = self._prepare_npu_triton_topk_idx( + topk_idx, seq_lens, num_idx_heads, num_kv_heads, max_blocks + ) + + # 4) Main sparse attention; native op uses the cached block table override. + _main_kwargs = ( + {**page_source_kwargs, **_native_main_kwargs} + if _native_main_kwargs is not None + else page_source_kwargs + ) + o = flash_decode_bnsd_with_gqa_share_sparse( + q=q, + sink=None, + k_cache_bnsd=k_bnsd, + v_cache_bnsd=v_bnsd, + **_main_kwargs, + seq_lens=seq_lens, + block_size=page_size, + topk_idx=topk_idx, + sm_scale=head_dim**-0.5, + use_native=_native_attn_enabled(), + ) + + return idx_o, o + + def _forward_npu_triton_verify( + self, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + idx_q: torch.Tensor, + idx_k_cache: torch.Tensor, + idx_v_cache: Optional[torch.Tensor], + forward_batch: ForwardBatch, + prefix_lens: torch.Tensor, + ): + """Capture-safe sparse attention for TARGET_VERIFY. + ndt queries per request, each causal (j attends KV[0:prefix+j+1]). Flatten + to per-query rows, reuse the decode kernels (device ops only, no .item()). + """ + from sgl_kernel_npu.attention.gqa_share_sparse_attention import ( + flash_decode_bnsd_with_gqa_share_sparse, + ) + from sgl_kernel_npu.indexer.flash_block_score_decode import ( + flash_decode_bnsd_with_topk_idx, + ) + + page_size = self.page_size # == block_size_k + num_q_heads = q.shape[1] + head_dim = q.shape[2] + num_idx_heads = idx_q.shape[1] + idx_dim = idx_q.shape[2] + num_tokens = q.shape[0] + bs = forward_batch.seq_lens.shape[0] + ndt = num_tokens // max(bs, 1) + + k_bnsd, v_bnsd, num_pages, num_kv_heads, head_dim = _kv_cache_to_bnsd( + k_cache, v_cache, page_size + ) + idx_k_bnsd, idx_v_bnsd, idx_kv_heads, idx_dim = _idx_cache_to_bnsd( + idx_k_cache, idx_v_cache, page_size + ) + + # Per-query causal seq_lens + req are layer-invariant, built once per + # forward as captured ops; inline fallback for the eager path. + vmeta = self._verify_meta_cg.get((bs, ndt)) + if vmeta is None: + prefix = (forward_batch.seq_lens.to(torch.long) - int(ndt)).clamp(min=0) + offsets = torch.arange(1, int(ndt) + 1, device=q.device, dtype=torch.long) + per_query_seq_lens = ( + (prefix.unsqueeze(1) + offsets.unsqueeze(0)).reshape(-1).to(torch.int32) + ) + per_query_req = forward_batch.req_pool_indices.long().repeat_interleave( + int(ndt) + ) + else: + per_query_seq_lens = vmeta.per_query_seq_lens + per_query_req = vmeta.per_query_req + + # ``max_seqlen`` comes from the capture-safe ``_max_seqlen_k`` (host-derived + # in init_forward_metadata_out_graph) so no device->host sync here. + max_seqlen = ( + int(self._max_seqlen_k) + if self._max_seqlen_k + else int(per_query_seq_lens.max().item()) + ) + max_blocks = (max_seqlen + page_size - 1) // page_size + disable_index_value = idx_v_cache is None + # Native verify-main: per-query block_table. CUDA-graph path uses the + # captured vmeta.native_bt (refreshed on replay); eager builds it per call. + _native_main_kwargs = None + if self._native_sparse_ok and _native_attn_enabled(): + try: + _bt = ( + vmeta.native_bt + if ( + vmeta is not None + and getattr(vmeta, "native_bt", None) is not None + ) + else None + ) + if _bt is None: + _bt = self._build_native_block_table( + per_query_req.long(), max_blocks, q.device + ) + _native_main_kwargs = {"block_table": _bt} + except Exception: + _native_main_kwargs = None + if disable_index_value: + # Keep verify's page-id range guard in the direct-map kernel. + page_source_kwargs = dict( + block_table=None, + req_to_token=self.req_to_token, + req_pool_indices=per_query_req, + max_num_blocks=max_blocks, + num_pages=num_pages, + sanitize_page_ids=True, + ) + else: + max_cols = self.req_to_token.shape[1] + blk_cols = ( + torch.arange(max_blocks, device=q.device, dtype=torch.long) * page_size + ).clamp(max=max_cols - 1) + token_slots = self.req_to_token[per_query_req][:, blk_cols] + block_table = (token_slots // page_size).to(torch.int32) + block_table = block_table.clamp(min=0, max=num_pages - 1) + page_source_kwargs = dict(block_table=block_table) + + # 1) indexer: score idx_k + index attention + topk (init/local=0). + # Pack each request's ndt draft queries into the gqa row dim. + pack_verify = ( + disable_index_value and int(ndt) > 1 and num_idx_heads == idx_kv_heads + ) + if pack_verify: + idx_q_score = idx_q.reshape(bs, ndt * num_idx_heads, idx_dim) + if num_idx_heads == 1: + # Row order == flat query order (request-major), so the + # per-query lengths double as the packed per-row lengths. + score_seq_lens = per_query_seq_lens + else: + score_seq_lens = ( + per_query_seq_lens.view(bs, ndt, 1) + .expand(bs, ndt, num_idx_heads) + .reshape(-1) + ) + score_page_source_kwargs = dict( + block_table=None, + req_to_token=self.req_to_token, + req_pool_indices=forward_batch.req_pool_indices, + max_num_blocks=max_blocks, + num_pages=num_pages, + sanitize_page_ids=True, + ) + else: + idx_q_score = idx_q + score_seq_lens = per_query_seq_lens + score_page_source_kwargs = page_source_kwargs + idx_o, topk_idx = flash_decode_bnsd_with_topk_idx( + q=idx_q_score, + sink=None, + k_cache_bnsd=idx_k_bnsd, + v_cache_bnsd=idx_v_bnsd, + **score_page_source_kwargs, + seq_lens=score_seq_lens, + max_seqlen=max_seqlen, + block_size=page_size, + topk=self.topk_blocks, + init_blocks=0, + local_blocks=0, + sm_scale=idx_dim**-0.5, + score_type=self.score_type, + disable_index_value=disable_index_value, + packed_seq_lens=pack_verify, + # 64-chunk graph for long contexts; runtime uses 16 chunks while + # <=256 blocks to cut short-context score work. Runtime direct-fill + # removes register TopK maintenance. + score_blocks_per_chunk=8 if pack_verify else 16, + score_max_chunks=64 if pack_verify else 32, + runtime_fill_only=pack_verify, + runtime_score_short_max_blocks=256 if pack_verify else 0, + runtime_score_short_chunks=16 if pack_verify else 0, + fused_append_local=True, + use_native=_native_indexer_enabled(), + ) + if pack_verify: + # [ndt*H, bs, K] -> [H, bs*ndt, K] (request-major rows). + k_last = topk_idx.shape[-1] + topk_idx = ( + topk_idx.view(ndt, num_idx_heads, bs, k_last) + .permute(1, 2, 0, 3) + .reshape(num_idx_heads, bs * ndt, k_last) + .contiguous() + ) + + # 2) Reduce heads and append forced blocks in the GQA kernel layout. + topk_idx = self._prepare_npu_triton_topk_idx( + topk_idx, + per_query_seq_lens, + num_idx_heads, + num_kv_heads, + max_blocks, + ) + + # 4) Main sparse attention; native op uses the cached block table override. + _vmain_kwargs = ( + {**page_source_kwargs, **_native_main_kwargs} + if _native_main_kwargs is not None + else page_source_kwargs + ) + o = flash_decode_bnsd_with_gqa_share_sparse( + q=q, + sink=None, + k_cache_bnsd=k_bnsd, + v_cache_bnsd=v_bnsd, + **_vmain_kwargs, + seq_lens=per_query_seq_lens, + block_size=page_size, + topk_idx=topk_idx, + sm_scale=head_dim**-0.5, + use_native=_native_attn_enabled(), + ) + return idx_o, o + + def _build_prefill_meta( + self, + forward_batch: ForwardBatch, + cu_seqlens: torch.Tensor, + seq_lens: torch.Tensor, + prefix_lens: torch.Tensor, + device, + page_size: int, + num_pages: int, + total_q: int, + ) -> SimpleNamespace: + """Build layer-invariant prefill metadata once per forward. + Depends only on batch shape + req_to_token (invariant across layers). + per_query_req is the direct-page-lookup map, kept live (no per-query table). + """ + seq_lens_l = seq_lens.to(device=device, dtype=torch.long) + prefix_lens_l = prefix_lens.to(device=device, dtype=torch.long) + cu_q = cu_seqlens.to(device=device, dtype=torch.long) + extend_lens = (seq_lens_l - prefix_lens_l).clamp(min=0) # [bs] + per_query_req = forward_batch.req_pool_indices.long().repeat_interleave( + extend_lens + ) # [total_q] + # Query j of request r sits at position prefix_r + j and causally attends to + # KV[0 : prefix_r + j + 1], so its seq_len = prefix_r + j + 1. + per_query_prefix = prefix_lens_l.repeat_interleave(extend_lens) # [total_q] + per_query_within = torch.arange( + total_q, device=device, dtype=torch.long + ) - cu_q[:-1].repeat_interleave( + extend_lens + ) # 0-indexed within each request + per_query_seq_lens = (per_query_prefix + per_query_within + 1).to(torch.int32) + + max_seqlen = ( + int(self._max_seqlen_k) + if self._max_seqlen_k + else int(per_query_seq_lens.max().item()) + ) + max_blocks = (max_seqlen + page_size - 1) // page_size + block_size_q = self._choose_block_size_q(max_seqlen) + + # Score-path qblock mappings (layer-invariant), built once per forward. + from sgl_kernel_npu.indexer.flash_block_score_prefill import ( + _build_qblock_mappings as _build_score_qblock_mappings, + ) + + qblock_mappings = _build_score_qblock_mappings( + cu_seqlens, + seq_lens, + self.req_to_token, + forward_batch.req_pool_indices, + block_size_q, + page_size, + max_blocks, + device, + ) + + # FIA prep workspace (layer-invariant shape, reused across layers). + topk1 = self.topk_blocks + 1 + fia_block_table_ws = torch.empty( + (total_q, topk1), dtype=torch.int32, device=device + ) + fia_actual_kvlen_ws = torch.empty((total_q,), dtype=torch.int32, device=device) + + return SimpleNamespace( + per_query_req=per_query_req, + # Pre-cast int32 for the FIA prep kernel (avoids a per-layer cast). + per_query_req_i32=per_query_req.to(torch.int32), + per_query_seq_lens=per_query_seq_lens, + max_seqlen=max_seqlen, + max_blocks=max_blocks, + block_size_q=block_size_q, + qblock_mappings=qblock_mappings, + fia_block_table_ws=fia_block_table_ws, + fia_actual_kvlen_ws=fia_actual_kvlen_ws, + ) + + def _forward_npu_triton_prefill( + self, + q: torch.Tensor, # [total_extend_tokens, num_q_heads, head_dim] + k_cache: torch.Tensor, + v_cache: torch.Tensor, + idx_q: torch.Tensor, # [total_extend_tokens, num_idx_heads, idx_dim] + idx_k_cache: torch.Tensor, + idx_v_cache: Optional[torch.Tensor], + forward_batch: ForwardBatch, + cu_seqlens: torch.Tensor, + seq_lens: torch.Tensor, + prefix_lens: torch.Tensor, + # Prefill main-attention launch tuning (decode path is unaffected). + main_num_warps: int = 4, + main_num_stages: int = 2, + # Fuse this many selected blocks per loop step of the main kernel. Same + # block set per query -> same math, only online-softmax regrouping. Use + # num_stages=1 when >1 (larger K/V tiles pressure the UB). + main_blocks_per_step: int = 1, + ): + """NPU block-sparse PREFILL via the ported triton decode kernels. + Generalizes verify to variable per-request extend lengths: each token becomes + a per-query row with a causal seq_len; decode kernels attend selected blocks. + """ + from sgl_kernel_npu.attention.gqa_share_sparse_attention import ( + flash_decode_bnsd_with_gqa_share_sparse, + ) + + page_size = self.page_size # == block_size_k + num_q_heads = q.shape[1] + head_dim = q.shape[2] + num_idx_heads = idx_q.shape[1] + idx_dim = idx_q.shape[2] + total_q = q.shape[0] + + k_bnsd, v_bnsd, num_pages, num_kv_heads, head_dim = _kv_cache_to_bnsd( + k_cache, v_cache, page_size + ) + idx_k_bnsd, idx_v_bnsd, idx_kv_heads, idx_dim = _idx_cache_to_bnsd( + idx_k_cache, idx_v_cache, page_size + ) + + # Layer-invariant metadata: built once per forward (first layer builds). + meta = self._prefill_meta + if meta is None: + meta = self._build_prefill_meta( + forward_batch, + cu_seqlens, + seq_lens, + prefix_lens, + q.device, + page_size, + num_pages, + total_q, + ) + self._prefill_meta = meta + per_query_seq_lens = meta.per_query_seq_lens + max_seqlen = meta.max_seqlen + max_blocks = meta.max_blocks + block_size_q = meta.block_size_q + per_query_req = meta.per_query_req + + disable_index_value = idx_v_cache is None + + # 1) indexer: score idx_k + index attention + topk (init/local=0). + # Batched varlen indexer tiles queries into block_size_q blocks and + # scores every query-block x kv-block in one 2D dot. Fused topk + + # causal-local append yields [..., topk+1]; the prepare helper skips + # the duplicate append. + from sgl_kernel_npu.indexer.flash_block_score_prefill import ( + flash_prefill_bnsd_indexer, + ) + from sgl_kernel_npu.indexer.flash_block_score_prefill import ( + flash_prefill_bnsd_with_topk_idx as _flash_prefill_score_topk, + ) + + topk_per_query_seq_lens = per_query_seq_lens + + if disable_index_value: + idx_o = None + topk_idx = _flash_prefill_score_topk( + idx_q, + idx_k_bnsd, + cu_seqlens, + seq_lens, + self.req_to_token, + forward_batch.req_pool_indices, + block_size_q, + page_size, + self.topk_blocks, + idx_dim**-0.5, + self.score_type, + qblock_mappings=meta.qblock_mappings, + per_query_seq_lens=topk_per_query_seq_lens, + ) + else: + idx_o, topk_idx = flash_prefill_bnsd_indexer( + idx_q, + idx_k_bnsd, + idx_v_bnsd, + cu_seqlens, + seq_lens, + self.req_to_token, + forward_batch.req_pool_indices, + block_size_q, + page_size, + self.topk_blocks, + idx_dim**-0.5, + self.score_type, + qblock_mappings=meta.qblock_mappings, + per_query_seq_lens=topk_per_query_seq_lens, + ) + + # 2) Reduce heads and append forced blocks in the GQA kernel layout. + topk_idx = self._prepare_npu_triton_topk_idx( + topk_idx, + per_query_seq_lens, + num_idx_heads, + num_kv_heads, + max_blocks, + ) + # No range/dtype guard needed: _prepare_npu_triton_topk_idx emits {-1} U + # [0, max_blocks-1] as int32 on both paths, and the main kernel masks + # logical_block < 0 and sanitizes physical ids to [0, num_pages-1]. + + # 4) main sparse attention over the selected blocks. + # BPS>1 fuses blocks per step of the decode-main kernel. + main_bps = int( + os.environ.get( + "SGLANG_MINIMAX_NPU_PREFILL_MAIN_BPS", str(main_blocks_per_step) + ) + ) + main_ns = main_num_stages if main_bps == 1 else min(main_num_stages, 1) + + def _decode_main(): + # Use the request-token map directly in the decode-main kernel. This + # avoids materializing a [total_q, max_blocks] page table for every + # sparse layer and keeps the per-query mapping live for graph replay. + return flash_decode_bnsd_with_gqa_share_sparse( + q=q, + sink=None, + k_cache_bnsd=k_bnsd, + v_cache_bnsd=v_bnsd, + block_table=None, + req_to_token=self.req_to_token, + req_pool_indices=per_query_req, + max_num_blocks=max_blocks, + num_pages=num_pages, + sanitize_page_ids=True, + seq_lens=per_query_seq_lens, + block_size=page_size, + topk_idx=topk_idx, + sm_scale=head_dim**-0.5, + topk_blocks_per_step=main_bps, + num_warps=main_num_warps, + num_stages=main_ns, + ) + + def _fia_main(): + # Native Ascend FA (FIA) with a per-query custom block_table. + from sgl_kernel_npu.attention.fia_blockq_attention import ( + flash_prefill_bnsd_blockq_sparse_fia, + ) + + return flash_prefill_bnsd_blockq_sparse_fia( + q=q, + k_cache_bnsd=k_bnsd, + v_cache_bnsd=v_bnsd, + topk_idx=topk_idx, + seq_lens=per_query_seq_lens, + per_query_req=meta.per_query_req_i32, + req_to_token=self.req_to_token, + block_size=page_size, + sm_scale=head_dim**-0.5, + num_pages=num_pages, + max_num_blocks=max_blocks, + block_table_out=meta.fia_block_table_ws, + actual_kvlen_out=meta.fia_actual_kvlen_ws, + ) + + use_fia = envs.SGLANG_MINIMAX_NPU_PREFILL_FIA.get() and num_kv_heads == 1 + o = _fia_main() if use_fia else _decode_main() + + return idx_o, o + @staticmethod def _is_sparse_kv_cached_by_fusion( forward_batch: ForwardBatch, layer_id: int @@ -295,6 +1256,56 @@ class MiniMaxSparseAttnBackend(AttentionBackend): q, k, v, layer, forward_batch, save_kv_cache, **kwargs ) + def _resolve_extend_meta(self, forward_batch: ForwardBatch, q: torch.Tensor): + """Return (cu_seqlens, seq_lens, prefix_lens); NPU caches per-forward casts.""" + # NPU TARGET_VERIFY has extend_seq_lens=None (seq_lens=prefix+draft); + # reconstruct per-seq extend lengths + prefix_lens for cu_seqlens. + if self.is_npu and forward_batch.extend_seq_lens is None: + _bs = forward_batch.seq_lens.shape[0] + _ndt = self.speculative_num_draft_tokens or (q.shape[0] // max(_bs, 1)) + forward_batch.extend_seq_lens = torch.full( + (_bs,), + int(_ndt), + dtype=torch.int32, + device=forward_batch.seq_lens.device, + ) + forward_batch.extend_seq_lens_cpu = [int(_ndt)] * _bs + if forward_batch.extend_prefix_lens is None: + forward_batch.extend_prefix_lens = ( + forward_batch.seq_lens.to(torch.int32) - int(_ndt) + ).clamp(min=0) + + # NPU cache hit (same forward_batch). + if ( + self.is_npu + and self._extend_meta_key == id(forward_batch) + and self._extend_meta is not None + ): + m = self._extend_meta + return m.cu_seqlens, m.seq_lens, m.prefix_lens + + cu_seqlens = torch.cat( + [ + torch.zeros( + 1, dtype=torch.int32, device=forward_batch.extend_seq_lens.device + ), + forward_batch.extend_seq_lens.to(torch.int32).cumsum(0).to(torch.int32), + ] + ) + seq_lens = forward_batch.seq_lens.to(torch.int32) + if forward_batch.extend_prefix_lens is not None: + prefix_lens = forward_batch.extend_prefix_lens.to(torch.int32) + else: + prefix_lens = torch.zeros_like(seq_lens) + + # NPU cache write. + if self.is_npu: + self._extend_meta = SimpleNamespace( + cu_seqlens=cu_seqlens, seq_lens=seq_lens, prefix_lens=prefix_lens + ) + self._extend_meta_key = id(forward_batch) + return cu_seqlens, seq_lens, prefix_lens + def forward_extend( self, q: torch.Tensor, @@ -332,22 +1343,9 @@ class MiniMaxSparseAttnBackend(AttentionBackend): else: idx_k_cache, idx_v_cache = self.kv_pool.get_index_kv_buffer(layer.layer_id) - cu_seqlens = torch.cat( - [ - torch.zeros( - 1, dtype=torch.int32, device=forward_batch.extend_seq_lens.device - ), - forward_batch.extend_seq_lens.to(torch.int32).cumsum(0).to(torch.int32), - ] - ) - seq_lens = forward_batch.seq_lens.to(torch.int32) - if forward_batch.extend_prefix_lens is not None: - prefix_lens = forward_batch.extend_prefix_lens.to(torch.int32) - else: - prefix_lens = torch.zeros_like(seq_lens) + cu_seqlens, seq_lens, prefix_lens = self._resolve_extend_meta(forward_batch, q) - # DP attention pads q beyond the real token count for collective alignment; - # trim to actual tokens so the sparse kernel sees consistent shapes. + # DP attention pads q beyond real tokens; trim (CPU list avoids a sync). if forward_batch.extend_seq_lens_cpu is not None: actual_num_tokens = int(sum(forward_batch.extend_seq_lens_cpu)) else: @@ -357,45 +1355,77 @@ class MiniMaxSparseAttnBackend(AttentionBackend): q = q[:actual_num_tokens] idx_q = idx_q[:actual_num_tokens] - # fp8 attention GEMMs: quantize q/idx_q AFTER the KV store (which reads - # the bf16 k/v) and the DP trim. - if self.fp8_attn_gemm: - q = _quant_q_fp8(q, layer.q_scale_float) - idx_q = _quant_q_fp8(idx_q, layer.idx_q_scale_float) + if self.is_npu: + if forward_batch.forward_mode.is_target_verify(): + # TARGET_VERIFY runs under cuda-graph capture; use the + # capture-safe verify path (decode kernels, no .item()). + idx_o, o = self._forward_npu_triton_verify( + q, + k_cache, + v_cache, + idx_q, + idx_k_cache, + idx_v_cache, + forward_batch, + prefix_lens, + ) + else: + idx_o, o = self._forward_npu_triton_prefill( + q, + k_cache, + v_cache, + idx_q, + idx_k_cache, + idx_v_cache, + forward_batch, + cu_seqlens, + seq_lens, + prefix_lens, + ) + else: + # fp8 attention GEMMs: quantize q/idx_q AFTER the KV store (which reads + # the bf16 k/v) and the DP trim. + if self.fp8_attn_gemm: + q = _quant_q_fp8(q, layer.q_scale_float) + idx_q = _quant_q_fp8(idx_q, layer.idx_q_scale_float) - idx_o, o = minimax_sparse_prefill( - q, - k_cache, - v_cache, - None, - idx_q, - idx_k_cache, - idx_v_cache, - None, - self.req_to_token, - forward_batch.req_pool_indices, - cu_seqlens, - seq_lens, - prefix_lens, - self._max_seqlen_q, - self._max_seqlen_k, - self.block_size_q, - self.block_size_k, - self.topk_blocks, - self.init_blocks, - self.local_blocks, - score_type=self.score_type, - disable_index_value=disable_value, - use_msa=self.use_msa, - seqlens_cpu=forward_batch.extend_seq_lens_cpu, - q_scale=layer.q_scale_float, - k_scale=layer.k_scale_float, - v_scale=layer.v_scale_float, - idx_q_scale=layer.idx_q_scale_float, - idx_k_scale=layer.idx_k_scale_float, - idx_v_scale=layer.idx_v_scale_float, - ) + # GPU (CUDA/ROCm) sparse path; imported here so NPU never touches it. + from sglang.srt.layers.attention.minimax_sparse_ops.minimax_sparse import ( + minimax_sparse_prefill, + ) + idx_o, o = minimax_sparse_prefill( + q, + k_cache, + v_cache, + None, + idx_q, + idx_k_cache, + idx_v_cache, + None, + self.req_to_token, + forward_batch.req_pool_indices, + cu_seqlens, + seq_lens, + prefix_lens, + self._max_seqlen_q, + self._max_seqlen_k, + self.block_size_q, + self.block_size_k, + self.topk_blocks, + self.init_blocks, + self.local_blocks, + score_type=self.score_type, + disable_index_value=disable_value, + use_msa=self.use_msa, + seqlens_cpu=forward_batch.extend_seq_lens_cpu, + q_scale=layer.q_scale_float, + k_scale=layer.k_scale_float, + v_scale=layer.v_scale_float, + idx_q_scale=layer.idx_q_scale_float, + idx_k_scale=layer.idx_k_scale_float, + idx_v_scale=layer.idx_v_scale_float, + ) if actual_num_tokens < original_num_tokens: pad_len = original_num_tokens - actual_num_tokens o = torch.cat([o, o.new_zeros(pad_len, *o.shape[1:])], dim=0) @@ -510,44 +1540,59 @@ class MiniMaxSparseAttnBackend(AttentionBackend): "did not prepare the plan for this forward (gate mismatch)." ) - # fp8 attention GEMMs: quantize q/idx_q AFTER the KV store (which reads - # the bf16 k/v). - if self.fp8_attn_gemm: - q = _quant_q_fp8(q, layer.q_scale_float) - idx_q = _quant_q_fp8(idx_q, layer.idx_q_scale_float) + if self.is_npu: + idx_o, o = self._forward_npu_triton_decode( + q, + k_cache, + v_cache, + idx_q, + idx_k_cache, + idx_v_cache, + forward_batch, + ) + else: + # fp8 attn-GEMM: quantize q/idx_q after the KV store (reads bf16 k/v). + if self.fp8_attn_gemm: + q = _quant_q_fp8(q, layer.q_scale_float) + idx_q = _quant_q_fp8(idx_q, layer.idx_q_scale_float) - idx_o, o = minimax_sparse_decode( - q, - None, - k_cache, - v_cache, - idx_q, - None, - idx_k_cache, - idx_v_cache, - self.req_to_token, - forward_batch.req_pool_indices, - forward_batch.seq_lens, - self._max_seqlen_k, - 1, - self.block_size_k, - self.topk_blocks, - self.init_blocks, - self.local_blocks, - score_type=self.score_type, - disable_index_value=disable_value, - dense_main_attn_fn=attn_fn, - page_size=self.page_size, - use_msa=self._use_msa_decode, - msa_kv_indices=msa_kv_indices, - msa_plan=msa_plan, - q_scale=layer.q_scale_float, - k_scale=layer.k_scale_float, - v_scale=layer.v_scale_float, - idx_q_scale=layer.idx_q_scale_float, - idx_k_scale=layer.idx_k_scale_float, - idx_v_scale=layer.idx_v_scale_float, - ) + # GPU (CUDA/ROCm) sparse path; imported here so NPU never touches it. + from sglang.srt.layers.attention.minimax_sparse_ops.minimax_sparse import ( + minimax_sparse_decode, + ) + + idx_o, o = minimax_sparse_decode( + q, + None, + k_cache, + v_cache, + idx_q, + None, + idx_k_cache, + idx_v_cache, + self.req_to_token, + forward_batch.req_pool_indices, + forward_batch.seq_lens, + self._max_seqlen_k, + 1, + self.block_size_k, + self.topk_blocks, + self.init_blocks, + self.local_blocks, + score_type=self.score_type, + disable_index_value=disable_value, + dense_main_attn_fn=attn_fn, + page_size=self.page_size, + use_msa=self._use_msa_decode, + msa_kv_indices=msa_kv_indices, + msa_plan=msa_plan, + q_scale=layer.q_scale_float, + k_scale=layer.k_scale_float, + v_scale=layer.v_scale_float, + idx_q_scale=layer.idx_q_scale_float, + idx_k_scale=layer.idx_k_scale_float, + idx_v_scale=layer.idx_v_scale_float, + ) return ( None if idx_o is None else idx_o.reshape(q.shape[0], -1).contiguous(), o.reshape(q.shape[0], -1).contiguous(), @@ -555,6 +1600,8 @@ class MiniMaxSparseAttnBackend(AttentionBackend): class MiniMaxHybridAttnBackend(AttentionBackend): + """Combines a dense backend and a sparse backend, routing by call site.""" + def __init__( self, dense_backend: AttentionBackend, @@ -564,12 +1611,14 @@ class MiniMaxHybridAttnBackend(AttentionBackend): self.dense = dense_backend self.sparse = sparse_backend self.sparse_layer_ids = sparse_layer_ids + # Let the sparse decode reuse the dense paged backend (page table + workspace). self.sparse.dense_backend = dense_backend self.extend_dummy_seqs_capped_by_req_pool = getattr( dense_backend, "extend_dummy_seqs_capped_by_req_pool", False ) or getattr(sparse_backend, "extend_dummy_seqs_capped_by_req_pool", False) def init_forward_metadata(self, forward_batch: ForwardBatch): + # delegate so the dense (FlashInfer) backend keeps its own eager init. self.sparse.init_forward_metadata(forward_batch) self.dense.init_forward_metadata(forward_batch) @@ -590,6 +1639,17 @@ class MiniMaxHybridAttnBackend(AttentionBackend): def get_cuda_graph_seq_len_fill_value(self): return self.sparse.get_cuda_graph_seq_len_fill_value() + def get_verify_buffers_to_fill_after_draft(self): + # EAGLE3 verify buffer interface: the dense (ascend) backend owns the + # tree-mask/position buffers consumed by the verify forward. The base + # AttentionBackend raises NotImplementedError, so delegate to dense. + return self.dense.get_verify_buffers_to_fill_after_draft() + + def update_verify_buffers_to_fill_after_draft(self, spec_info, cuda_graph_bs=None): + return self.dense.update_verify_buffers_to_fill_after_draft( + spec_info, cuda_graph_bs + ) + def forward( self, q, diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py index 73ac933c7..a3ee03580 100644 --- a/python/sglang/srt/layers/layernorm.py +++ b/python/sglang/srt/layers/layernorm.py @@ -174,6 +174,8 @@ if _is_npu: import torch_npu from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm +_NPU_GEMMA_RMS_NORM_TRITON_MAX_HIDDEN_SIZE = 5120 + @lru_cache(maxsize=1) def _get_aiter_per_group_quant(): @@ -1143,9 +1145,15 @@ class GemmaRMSNorm(BaseFusedOp): if residual is not None: if post_residual_addition is not None: residual = residual + post_residual_addition - norm_out, residual = add_gemma_rms_norm( - x, self.weight, residual, self.variance_epsilon - ) + if x.shape[-1] > _NPU_GEMMA_RMS_NORM_TRITON_MAX_HIDDEN_SIZE: + gamma = self.gemma_weight.to(x.dtype) + norm_out, _, residual = torch_npu.npu_add_rms_norm( + residual, x, gamma, self.variance_epsilon + ) + else: + norm_out, residual = add_gemma_rms_norm( + x, self.weight, residual, self.variance_epsilon + ) return norm_out, residual x, _ = torch_npu.npu_gemma_rms_norm(x, self.weight, self.variance_epsilon) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index cf5203f35..f9bcb949c 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -419,6 +419,13 @@ class FusedMoE(torch.nn.Module): if expert_mask is not None: self.register_buffer("expert_mask_gpu", expert_mask, persistent=False) self._use_ascend_fuseep = get_moe_a2a_backend().is_ascend_fuseep() + # Expose swigluoai alpha/clamp on the layer so deepep's W8A8 apply + # (apply_without_routing_weights) picks swiglu_oai_quant instead of plain + # npu_swiglu. fuseep injects these via fuseep_activation (aclnnFusedDeepMoe + # internal); deepep reads them here via getattr(layer, "swiglu_alpha"). + # Default None (no-op for non-swigluoai models). + self.swiglu_alpha = gemm1_alpha + self.swiglu_clamp_limit = gemm1_clamp_limit if ( get_moe_runner_backend().is_flashinfer_trtllm_routed() diff --git a/python/sglang/srt/layers/moe/moe_runner/ascend.py b/python/sglang/srt/layers/moe/moe_runner/ascend.py index c72c00eee..7bc97d17f 100644 --- a/python/sglang/srt/layers/moe/moe_runner/ascend.py +++ b/python/sglang/srt/layers/moe/moe_runner/ascend.py @@ -111,7 +111,11 @@ class AscendRunnerCore(MoeRunnerCore): linear_beta=config.gemm1_clamp_limit, ) else: - self.activation = NPUSwigluDeepEPKernel(need_quant=is_quant_kernel) + self.activation = NPUSwigluDeepEPKernel( + need_quant=is_quant_kernel, + alpha=config.gemm1_alpha, + limit=config.gemm1_clamp_limit, + ) else: # Non‑DeepEP (ascend_tp) path # 1. Choose the base activation according to the quant method diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index 036c7392d..7a02601b1 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -105,6 +105,18 @@ class ModelSlimConfig(QuantizationConfig): for k, v in quant_config.items() } + # Add an mlp.* alias for each block_sparse_moe.* key but KEEP the original, + # so both module namings resolve + for k in list(quant_config.keys()): + if not isinstance(k, str): + continue + if "block_sparse_moe" in k: + quant_config[ + k.replace("block_sparse_moe.experts", "mlp.experts").replace( + "block_sparse_moe.shared_experts", "mlp.shared_experts" + ) + ] = quant_config[k] + self.quant_description = quant_config ignore = cast(List[str], quant_config.get("ignore", [])) self.ignore = ignore if ignore is not None else [] diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 4e3e4df65..6e36069e5 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -993,6 +993,10 @@ class KVCacheConfigurator: full_max_total_num_tokens=sizes.full_max_total_num_tokens, swa_max_total_num_tokens=sizes.swa_max_total_num_tokens, ) + elif is_minimax_sparse(self.model_config.hf_config): + token_to_kv_pool = self._build_ascend_minimax_sparse_kv_pool( + max_total_num_tokens=sizes.max_total_num_tokens, + ) elif self.use_mla_backend: token_to_kv_pool = self._build_ascend_mla_kv_pool( max_total_num_tokens=sizes.max_total_num_tokens, @@ -1238,6 +1242,37 @@ class KVCacheConfigurator: ) return token_to_kv_pool + def _build_ascend_minimax_sparse_kv_pool( + self, *, max_total_num_tokens: int + ) -> KVCache: + _hf_config = self.model_config.hf_config + sparse_cfg = get_minimax_sparse_attention_config(_hf_config) + dense_layer_ids, sparse_layer_ids = get_minimax_sparse_layer_ids(sparse_cfg) + disable_value_sparse_layer_ids = get_minimax_sparse_disable_value_layer_ids( + sparse_cfg + ) + from sglang.srt.hardware_backend.npu.memory_pool_npu import ( + NPUMiniMaxSparseKVPool, + ) + + token_to_kv_pool = NPUMiniMaxSparseKVPool( + size=max_total_num_tokens, + page_size=self.server_args.page_size, + dtype=self.kv_cache_dtype, + index_dtype=self.model_dtype, + head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size), + head_dim=self.model_config.head_dim, + idx_head_dim=sparse_cfg["sparse_index_dim"], + dense_layer_ids=dense_layer_ids, + sparse_layer_ids=sparse_layer_ids, + disable_value_sparse_layer_ids=disable_value_sparse_layer_ids, + device=self.device, + enable_memory_saver=self.server_args.enable_memory_saver, + start_layer=self.layer_info.start_layer, + end_layer=self.layer_info.end_layer, + ) + return token_to_kv_pool + def _build_ascend_mla_kv_pool( self, *, max_total_num_tokens: int, is_dsa_model: bool ) -> KVCache: diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 261b2018e..3696ad38a 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -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, diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index ebe6c516f..3dc54970d 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -242,6 +242,7 @@ def _get_quantization_config( "q_a_proj", "kv_a_proj_with_mqa", ], + "index_qkv_proj": ["index_q_proj", "index_k_proj"], }, } ) diff --git a/python/sglang/srt/models/minimax_m3.py b/python/sglang/srt/models/minimax_m3.py index 6fbca7dcd..7ff7c0b36 100644 --- a/python/sglang/srt/models/minimax_m3.py +++ b/python/sglang/srt/models/minimax_m3.py @@ -42,7 +42,9 @@ from sglang.srt.layers.communicator import ( ScatterMode, enable_moe_dense_fully_dp, ) -from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.layers.dp_attention import ( + is_dp_attention_enabled, +) from sglang.srt.layers.layernorm import GemmaRMSNorm, RMSNorm from sglang.srt.layers.linear import ( MergedColumnParallelLinear, @@ -89,6 +91,7 @@ from sglang.srt.utils import ( get_device_sm, is_cuda, is_hip, + is_npu, log_info_on_rank0, make_layers, ) @@ -96,6 +99,7 @@ from sglang.srt.utils.hf_transformers_utils import get_rope_config _is_cuda = is_cuda() _is_hip = is_hip() +_is_npu = is_npu() _device_sm = get_device_sm() _FP8_KV_DTYPES = ( @@ -121,6 +125,16 @@ if _is_hip: except ImportError: _has_rocm_qk_norm_rope = False +if _is_npu: + from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope_pos_cache_half_npu import ( + split_qkv_rmsnorm_rope_pos_cache_half_npu, + ) + + from sglang.srt.hardware_backend.npu.utils import ( + process_shared_expert, + wait_share_stream, + ) + logger = logging.getLogger(__name__) @@ -213,6 +227,14 @@ def build_minimax_fused_qkv_index(model: nn.Module) -> None: class MiniMaxM3MLP(nn.Module): + @staticmethod + def _swigluoai_fused(x: torch.Tensor, alpha: float, limit: float) -> torch.Tensor: + """swiglu_oai using fused Triton kernel (sgl_kernel_npu), no quant.""" + from sgl_kernel_npu.activation.swiglu_oai_quant import swiglu_oai_quant + + out, _ = swiglu_oai_quant(x, alpha, limit, need_quant=False) + return out + def __init__( self, config: PretrainedConfig, @@ -249,13 +271,18 @@ class MiniMaxM3MLP(nn.Module): if hidden_act == "silu": self.act_fn = SiluAndMul() elif hidden_act == "swigluoai": - from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import ( - swiglu_no_interleaved_with_alpha_and_limit, - ) + if _is_npu: + self.act_fn = lambda x: self._swigluoai_fused( + x, config.swiglu_alpha, config.swiglu_limit + ) + else: + from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import ( + swiglu_no_interleaved_with_alpha_and_limit, + ) - self.act_fn = lambda x: swiglu_no_interleaved_with_alpha_and_limit( - x, config.swiglu_alpha, config.swiglu_limit - ) + self.act_fn = lambda x: swiglu_no_interleaved_with_alpha_and_limit( + x, config.swiglu_alpha, config.swiglu_limit + ) else: raise ValueError( f"Unsupported activation: {hidden_act}. Only silu is supported for now." @@ -264,6 +291,7 @@ class MiniMaxM3MLP(nn.Module): def forward( self, x, + forward_batch: Optional[ForwardBatch] = None, should_allreduce_fusion: bool = False, use_reduce_scatter: bool = False, ): @@ -416,10 +444,22 @@ class MiniMaxM3MoE(nn.Module): def forward_deepep( self, hidden_states: torch.Tensor, forward_batch: ForwardBatch ) -> torch.Tensor: + """DeepEP MoE forward: routed experts via a2a, shared experts replicated.""" shared_output = None + enable_npu_dual_stream = _is_npu and ( + forward_batch.forward_mode.is_extend() + or forward_batch.forward_mode.is_target_verify() + or forward_batch.forward_mode.is_decode() + ) if hidden_states.shape[0] > 0: - shared_output = self._forward_shared_experts(hidden_states) router_logits = self._compute_router_logits(hidden_states) + if enable_npu_dual_stream: + # Overlap shared experts with router/experts on a separate stream. + shared_output = process_shared_expert( + hidden_states, self._forward_shared_experts + ) + else: + shared_output = self._forward_shared_experts(hidden_states) topk_output = self.topk( hidden_states, router_logits, @@ -435,6 +475,9 @@ class MiniMaxM3MoE(nn.Module): # shared experts are replicated (tp_size=1), so both add directly. final_hidden_states = self.experts(hidden_states, topk_output) + if enable_npu_dual_stream: + wait_share_stream() + if shared_output is not None: final_hidden_states = final_hidden_states + shared_output @@ -442,6 +485,9 @@ class MiniMaxM3MoE(nn.Module): def _compute_router_logits(self, hidden_states: torch.Tensor) -> torch.Tensor: if self.bf16_router_gemm: + if _is_npu: + # NPU lacks aten::mm.dtype; bf16 mm then cast keeps topk semantics. + return torch.mm(hidden_states, self.gate.weight.t()).float() return torch.mm( hidden_states, self.gate.weight.t(), out_dtype=torch.float32 ) @@ -501,6 +547,7 @@ class MiniMaxM3Attention(nn.Module): self.max_position_embeddings = getattr(config, "max_position_embeddings", 8192) self.rotary_dim = getattr(config, "rotary_dim", self.head_dim) + self.use_qk_norm = getattr(config, "use_qk_norm", False) self.qk_norm_type = getattr(config, "qk_norm_type", "per_layer") self.use_gemma_norm = getattr(config, "use_gemma_norm", False) @@ -994,6 +1041,44 @@ class MiniMaxM3Attention(nn.Module): return q, k, idx_q, idx_k return self._sparse_qk_index_norm_rope(positions, q, k, idx_q, idx_k) + def forward_prepare_npu( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ): + """NPU qkv projection + fused norm/RoPE/split; returns (None, fb, inner_state).""" + if hidden_states.shape[0] == 0: + assert ( + not self.o_proj.reduce_results + ), "short-circuiting allreduce will lead to hangs" + return hidden_states, forward_batch, None + + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = split_qkv_rmsnorm_rope_pos_cache_half_npu( + input_tensor=qkv, + positions=positions.reshape(-1), + cos_sin_cache=self.rotary_emb.cos_sin_cache, + q_hidden_size=self.q_size, + kv_hidden_size=self.kv_size, + head_dim=self.head_dim, + eps=self.q_norm.variance_epsilon, + q_weight=self.q_norm.gemma_weight, + k_weight=self.k_norm.gemma_weight, + rope_dim=self.rotary_dim, + cast_norm_to_bf16=True, + ) + if self.is_sparse_attention_layer: + idx_qkv, _ = self.index_qkv_proj(hidden_states) + # Index attention disables the V head on all M3 sparse layers, so + # index_qkv_proj emits a 2-way [q|k] tensor. + idx_q, idx_k, idx_v = self._split_index_qkv(idx_qkv) + idx_q, idx_k = self._index_qk_norm_rope(positions, idx_q, idx_k) + inner_state = (q, k, v, idx_q, idx_k, idx_v, forward_batch) + else: + inner_state = (q, k, v, forward_batch) + return None, forward_batch, inner_state + def forward_prepare( self, positions: torch.Tensor, @@ -1110,8 +1195,7 @@ class MiniMaxM3Attention(nn.Module): output, _ = self.o_proj(attn_output) if self.disable_index_value: return output - # idx_replica_size ranks produce identical idx_o; pre-divide idx_o (not the - # o_proj weight) so the TP all-reduce sums right and stays FP8-quant-safe. + # Pre-divide idx_o (not the weight) so the TP all-reduce sums right. if self.idx_replica_size > 1: idx_o = idx_o / self.idx_replica_size idx_output, _ = self.index_o_proj(idx_o) @@ -1128,11 +1212,18 @@ class MiniMaxM3Attention(nn.Module): hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: - s = self.forward_prepare( - positions=positions, - hidden_states=hidden_states, - forward_batch=forward_batch, - ) + if _is_npu: + s = self.forward_prepare_npu( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) + else: + s = self.forward_prepare( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) return self.forward_core(s) @@ -1282,8 +1373,9 @@ class MiniMaxM3DecoderLayer(nn.Module): if self.is_layer_sparse or hidden_states.shape[0] != 0: hidden_states = self.mlp( hidden_states, - should_allreduce_fusion, - use_reduce_scatter, + forward_batch=forward_batch, + should_allreduce_fusion=should_allreduce_fusion, + use_reduce_scatter=use_reduce_scatter, ) if should_allreduce_fusion: @@ -1513,6 +1605,12 @@ class MiniMaxM3SparseForCausalLM(nn.Module): else: self.model.layers_to_capture = [val + 1 for val in layer_ids] + # forward checks the per-layer ``_is_layer_to_capture`` flag, not the id + # list, so set it explicitly (mirrors qwen3_next/qwen2_moe). + for layer_id in self.model.layers_to_capture: + if 0 <= layer_id < len(self.model.layers): + setattr(self.model.layers[layer_id], "_is_layer_to_capture", True) + def get_embed_and_head(self): return self.model.embed_tokens.weight, self.lm_head.weight diff --git a/python/sglang/srt/models/minimax_m3_vl.py b/python/sglang/srt/models/minimax_m3_vl.py index d9008e5b0..2fdff0370 100644 --- a/python/sglang/srt/models/minimax_m3_vl.py +++ b/python/sglang/srt/models/minimax_m3_vl.py @@ -135,6 +135,9 @@ class MiniMaxM3SparseForConditionalGeneration(nn.Module): self.logits_processor = LogitsProcessor(text_config) + # For EAGLE3 support + self.capture_aux_hidden_states = False + @classmethod def shared_experts_fusion_disable_reason(cls, hf_config, quant_config): """Why this checkpoint cannot fuse its shared expert, or None. Asked by @@ -197,6 +200,36 @@ class MiniMaxM3SparseForConditionalGeneration(nn.Module): def get_input_embeddings(self): return self.model.embed_tokens + def get_embed_and_head(self): + # EAGLE3 target interface: share the text embed + lm_head with the draft. + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_eagle3_layers_to_capture(self, layer_ids: Optional[list[int]] = None): + # EAGLE3 target interface: select which decoder layers' hidden states the + # draft consumes. Mirrors MiniMaxM3SparseForCausalLM; operates on the inner + # MiniMaxM3Model (self.model), whose forward returns (hidden, aux) once set. + if not self.pp_group.is_last_rank: + return + + self.capture_aux_hidden_states = True + # MiniMaxM3Model.forward captures at layer ENTRY (= previous layer's + # output), so to capture layer L's output we must mark layer L+1. Apply + # +1 on both paths so EAGLE3 works out-of-the-box even when the draft + # config omits ``eagle_aux_hidden_state_layer_ids`` (the upstream + # Inferact/MiniMax-M3-EAGLE3 checkpoint does not ship it); otherwise the + # default-path layers are off by one and draft accept collapses. + if layer_ids is None: + num_layers = self.config.text_config.num_hidden_layers + layer_ids = [2, num_layers // 2, num_layers - 3] + self.model.layers_to_capture = [val + 1 for val in layer_ids] + + # MiniMaxM3Model.forward checks each layer's ``_is_layer_to_capture`` + # attribute (not ``i in layers_to_capture``); set it explicitly so the + # (hidden, aux) tuple is actually returned during capture-enabled forwards. + for layer_id in self.model.layers_to_capture: + if 0 <= layer_id < len(self.model.layers): + setattr(self.model.layers[layer_id], "_is_layer_to_capture", True) + def forward( self, input_ids: torch.Tensor, @@ -217,12 +250,20 @@ class MiniMaxM3SparseForConditionalGeneration(nn.Module): pp_proxy_tensors=pp_proxy_tensors, ) + # EAGLE3: when layers_to_capture is set, MiniMaxM3Model.forward returns + # (hidden_states, aux_hidden_states) once aux is non-empty; on idle/warmup + # forwards with no captured tokens it returns a bare hidden tensor. + aux_hidden_states = None + if self.capture_aux_hidden_states and isinstance(hidden_states, tuple): + hidden_states, aux_hidden_states = hidden_states + if self.pp_group.is_last_rank and not get_embedding: return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch, + aux_hidden_states, ) return hidden_states diff --git a/python/sglang/srt/multimodal/processors/base_processor.py b/python/sglang/srt/multimodal/processors/base_processor.py index 8da0a632b..f7ffde730 100644 --- a/python/sglang/srt/multimodal/processors/base_processor.py +++ b/python/sglang/srt/multimodal/processors/base_processor.py @@ -512,6 +512,22 @@ class BaseMultimodalProcessor(ABC): return "xpu" if not _is_npu: return f"cuda:{server_args.base_gpu_id}" + if processor.__class__.__name__ == "MiniMaxVLProcessor": + # MiniMax's image/video processors create 10-dim tensors during + # patch extraction, exceeding the Ascend 8-dim limit; patch them + # (same pattern as qwen-vl / GLM-4.6V) and run on NPU. + from sglang.srt.hardware_backend.npu.modules.minimax_m3_processor import ( + npu_apply_minimax_m3_image_preprocess_patch, + npu_apply_minimax_m3_video_preprocess_patch, + ) + + npu_apply_minimax_m3_image_preprocess_patch(processor.image_processor) + if ( + hasattr(processor, "video_processor") + and processor.video_processor is not None + ): + npu_apply_minimax_m3_video_preprocess_patch(processor.video_processor) + return "npu" if processor.__class__.__name__ not in {"Glm4vProcessor", "Glm46VProcessor"}: # For qwen-vl, the processor hits a reshape issue from the Ascend # dims restriction. diff --git a/test/manual/minimax_m3/test_npu_memory_pool.py b/test/manual/minimax_m3/test_npu_memory_pool.py new file mode 100644 index 000000000..509d5d379 --- /dev/null +++ b/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 diff --git a/test/manual/minimax_m3/test_npu_topk.py b/test/manual/minimax_m3/test_npu_topk.py new file mode 100644 index 000000000..06c8ba1ef --- /dev/null +++ b/test/manual/minimax_m3/test_npu_topk.py @@ -0,0 +1,193 @@ +import importlib.util +import sys +import types +from pathlib import Path +from typing import NamedTuple + +import torch + + +def _install_fake_modules(): + for name in ( + "sglang", + "sglang.srt", + "sglang.srt.eplb", + "sglang.srt.layers", + "sglang.srt.layers.moe", + "sglang.srt.state_capturer", + ): + sys.modules.setdefault(name, types.ModuleType(name)) + + root = types.ModuleType("sgl_kernel_npu") + norm = types.ModuleType("sgl_kernel_npu.norm") + l1_norm_mod = types.ModuleType("sgl_kernel_npu.norm.l1_norm") + + def l1_norm(x): + return x / x.sum(dim=-1, keepdim=True) + + l1_norm_mod.l1_norm = l1_norm + sys.modules.setdefault("sgl_kernel_npu", root) + sys.modules.setdefault("sgl_kernel_npu.norm", norm) + sys.modules["sgl_kernel_npu.norm.l1_norm"] = l1_norm_mod + + expert_distribution = types.ModuleType("sglang.srt.eplb.expert_distribution") + + class Recorder: + @staticmethod + def on_select_experts(topk_ids): + pass + + expert_distribution.get_global_expert_distribution_recorder = lambda: Recorder() + sys.modules["sglang.srt.eplb.expert_distribution"] = expert_distribution + + expert_location = types.ModuleType("sglang.srt.eplb.expert_location_dispatch") + expert_location.topk_ids_logical_to_physical = lambda topk_ids, info: topk_ids + sys.modules["sglang.srt.eplb.expert_location_dispatch"] = expert_location + + moe_topk = types.ModuleType("sglang.srt.layers.moe.topk") + + class StandardTopKOutput(NamedTuple): + topk_weights: torch.Tensor + topk_ids: torch.Tensor + router_logits: torch.Tensor + + def select_experts(*args, **kwargs): + raise AssertionError("fallback select_experts should not be used") + + def capture_routed_experts_if_allowed(*args, **kwargs): + return None + + moe_topk.StandardTopKOutput = StandardTopKOutput + moe_topk.select_experts = select_experts + moe_topk.capture_routed_experts_if_allowed = capture_routed_experts_if_allowed + sys.modules["sglang.srt.layers.moe.topk"] = moe_topk + + routed_experts = types.ModuleType("sglang.srt.state_capturer.routed_experts") + routed_experts.get_global_experts_capturer = lambda: None + sys.modules["sglang.srt.state_capturer.routed_experts"] = routed_experts + + +def _load_npu_topk_module(): + _install_fake_modules() + module_path = ( + Path(__file__).resolve().parents[3] + / "python/sglang/srt/hardware_backend/npu/moe/topk.py" + ) + spec = importlib.util.spec_from_file_location("_npu_topk_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 _make_topk_config( + correction_bias, routed_scaling_factor, renormalize=True +) -> types.SimpleNamespace: + """M3-shaped TopKConfig: sigmoid scoring, no grouped routing.""" + return types.SimpleNamespace( + top_k=2, + use_grouped_topk=False, + correction_bias=correction_bias, + topk_group=None, + num_expert_group=None, + renormalize=renormalize, + scoring_func="sigmoid", + num_fused_shared_experts=0, + routed_scaling_factor=routed_scaling_factor, + apply_routed_scaling_factor_on_output=True, + ) + + +def _run_fused_topk_npu(npu_topk, topk_config, router_logits): + return npu_topk.fused_topk_npu( + hidden_states=torch.zeros((1, 4), dtype=torch.bfloat16), + router_logits=router_logits, + topk_config=topk_config, + ) + + +class _SigmoidFakeNpuOps: + """npu_moe_gating_top_k mirroring the real sigmoid contract (norm_type=1).""" + + def __init__(self, router_logits, routed_scaling_factor, expect_bias=None): + self.router_logits = router_logits + self.routed_scaling_factor = routed_scaling_factor + self.expect_bias = expect_bias + + def npu_moe_gating_top_k_softmax(self, *args, **kwargs): + raise AssertionError("sigmoid routing must not use the softmax top-k op") + + def npu_moe_gating_top_k( + self, + router_logits, + *, + k, + bias, + renorm, + norm_type, + routed_scaling_factor, + **kwargs, + ): + # Contract: sigmoid scoring -> norm_type=1; bias (if any) must reach the + # op; renorm and the routed scaling factor are applied inside the op. + assert norm_type == 1 + if self.expect_bias is not None: + assert bias is not None and bias.shape == self.expect_bias.shape + scores = ( + (router_logits + bias).sigmoid() + if bias is not None + else router_logits.sigmoid() + ) + values, ids = torch.topk(scores, k=k, dim=-1) + if renorm: + values = values / values.sum(dim=-1, keepdim=True) + values = values * routed_scaling_factor + return values, ids.to(torch.int32), None + + +def test_npu_sigmoid_topk_without_bias_uses_sigmoid_op(monkeypatch): + """Sigmoid routing without correction bias must NOT fall into the softmax fast path. + + Guards fused_topk_npu's fast-path branch: it previously matched + ``not use_grouped_topk and correction_bias is None`` without excluding + sigmoid scoring, routing sigmoid models through the softmax op. + """ + npu_topk = _load_npu_topk_module() + + router_logits = torch.tensor([[0.0, 1.0, 2.0]], dtype=torch.float32) + routed_scaling_factor = 2.5 + fake = _SigmoidFakeNpuOps(router_logits, routed_scaling_factor, expect_bias=None) + monkeypatch.setattr(torch.ops, "npu", fake, raising=False) + + topk_output = _run_fused_topk_npu( + npu_topk, + _make_topk_config(None, routed_scaling_factor), + router_logits, + ) + + raw = router_logits.sigmoid().topk(2, dim=-1).values + expected = raw / raw.sum(dim=-1, keepdim=True) * routed_scaling_factor + torch.testing.assert_close(topk_output.topk_weights, expected) + + +def test_npu_sigmoid_topk_with_routing_bias_matches_m3_config(monkeypatch): + """M3 real config (use_routing_bias=True): bias must reach the sigmoid op.""" + npu_topk = _load_npu_topk_module() + + router_logits = torch.tensor([[0.0, 1.0, 2.0, 3.0]], dtype=torch.float32) + routed_scaling_factor = 2.5 + correction_bias = torch.tensor([0.1, -0.2, 0.3, 0.05], dtype=torch.float32) + fake = _SigmoidFakeNpuOps( + router_logits, routed_scaling_factor, expect_bias=correction_bias + ) + monkeypatch.setattr(torch.ops, "npu", fake, raising=False) + + topk_output = _run_fused_topk_npu( + npu_topk, + _make_topk_config(correction_bias, routed_scaling_factor), + router_logits, + ) + + raw = (router_logits + correction_bias).sigmoid().topk(2, dim=-1).values + expected = raw / raw.sum(dim=-1, keepdim=True) * routed_scaling_factor + torch.testing.assert_close(topk_output.topk_weights, expected)